wehub-resource-sync
547 lines
18 KiB
Python
547 lines
18 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""HuggingFace model uploader for oMLX admin panel.
|
|
|
|
Uploads oQ-quantized models to HuggingFace Hub with queue-based sequential
|
|
processing, following the same pattern as hf_downloader.py.
|
|
"""
|
|
|
|
import asyncio
|
|
import enum
|
|
import json
|
|
import logging
|
|
import shutil
|
|
import tempfile
|
|
import time
|
|
import uuid
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def _format_size(size_bytes: int) -> str:
|
|
"""Format size in bytes to human-readable string."""
|
|
if size_bytes < 1024**2:
|
|
return f"{size_bytes / 1024:.1f} KB"
|
|
elif size_bytes < 1024**3:
|
|
return f"{size_bytes / 1024**2:.1f} MB"
|
|
else:
|
|
return f"{size_bytes / 1024**3:.1f} GB"
|
|
|
|
|
|
def _has_meaningful_readme(path: Path) -> bool:
|
|
"""Check if a README.md exists and has content beyond YAML frontmatter.
|
|
|
|
Returns False if the file doesn't exist, is empty, or contains only
|
|
YAML frontmatter (e.g. mlx-lm's default stub).
|
|
"""
|
|
readme = path / "README.md"
|
|
if not readme.exists():
|
|
return False
|
|
try:
|
|
text = readme.read_text(encoding="utf-8").strip()
|
|
except Exception:
|
|
return False
|
|
if not text:
|
|
return False
|
|
# Strip YAML frontmatter and check if anything remains
|
|
if text.startswith("---"):
|
|
parts = text.split("---", 2)
|
|
# parts[0] is empty (before first ---), parts[1] is frontmatter
|
|
if len(parts) >= 3:
|
|
body = parts[2].strip()
|
|
return len(body) > 0
|
|
# Only opening --- or unclosed frontmatter
|
|
return False
|
|
return True
|
|
|
|
|
|
def _is_oq_model(name: str) -> bool:
|
|
"""Check if a model name indicates an oQ-quantized model.
|
|
|
|
Any folder name containing 'oQ' (case-sensitive) is treated as an oQ model,
|
|
e.g. 'Qwen3.5-122B-oQ4', 'Llama-3B-oQ4e', 'Qwen3.6-27B-oQ3.5e'.
|
|
"""
|
|
return "oQ" in name
|
|
|
|
|
|
def _generate_model_card(
|
|
model_name: str, config: dict, redownload_notice: bool = False,
|
|
) -> str:
|
|
"""Generate a minimal HuggingFace model card for an oQ model."""
|
|
from omlx._version import __version__
|
|
|
|
model_type = config.get("model_type", "unknown")
|
|
quant = config.get("quantization", {})
|
|
bits = quant.get("bits", "?")
|
|
group_size = quant.get("group_size", "?")
|
|
|
|
notice = ""
|
|
if redownload_notice:
|
|
from datetime import date
|
|
|
|
today = date.today().strftime("%Y-%m-%d")
|
|
notice = f"""> [!IMPORTANT]
|
|
> This quantization was uploaded on **{today}** and replaces a previous version.
|
|
> If you downloaded this model before this date, please re-download for the updated weights.
|
|
|
|
"""
|
|
|
|
return f"""---
|
|
library_name: mlx
|
|
tags:
|
|
- mlx
|
|
- oq
|
|
- quantized
|
|
---
|
|
|
|
{notice}# {model_name}
|
|
|
|
This model was quantized using [oQ](https://github.com/jundot/omlx) (oMLX v{__version__}) mixed-precision quantization.
|
|
|
|
## Quantization details
|
|
|
|
- **Model type**: {model_type}
|
|
- **Bits**: {bits}
|
|
- **Group size**: {group_size}
|
|
- **Format**: MLX safetensors
|
|
"""
|
|
|
|
|
|
class UploadStatus(str, enum.Enum):
|
|
"""Status of an upload task."""
|
|
|
|
PENDING = "pending"
|
|
UPLOADING = "uploading"
|
|
COMPLETED = "completed"
|
|
FAILED = "failed"
|
|
CANCELLED = "cancelled"
|
|
|
|
|
|
_ACTIVE_STATUSES = {UploadStatus.PENDING, UploadStatus.UPLOADING}
|
|
|
|
|
|
@dataclass
|
|
class UploadTask:
|
|
"""Represents a single model upload task."""
|
|
|
|
task_id: str
|
|
model_name: str
|
|
model_path: str
|
|
repo_id: str
|
|
status: UploadStatus = UploadStatus.PENDING
|
|
progress: float = 0.0
|
|
error: str = ""
|
|
created_at: float = field(default_factory=time.time)
|
|
started_at: float = 0.0
|
|
completed_at: float = 0.0
|
|
total_size: int = 0
|
|
repo_url: str = ""
|
|
|
|
def to_dict(self) -> dict:
|
|
"""Serialize task to a JSON-compatible dict."""
|
|
return {
|
|
"task_id": self.task_id,
|
|
"model_name": self.model_name,
|
|
"model_path": self.model_path,
|
|
"repo_id": self.repo_id,
|
|
"status": self.status.value,
|
|
"progress": round(self.progress, 1),
|
|
"error": self.error,
|
|
"created_at": self.created_at,
|
|
"started_at": self.started_at,
|
|
"completed_at": self.completed_at,
|
|
"total_size": self.total_size,
|
|
"total_size_formatted": _format_size(self.total_size) if self.total_size else "",
|
|
"repo_url": self.repo_url,
|
|
}
|
|
|
|
|
|
class HFUploader:
|
|
"""Manages HuggingFace model uploads with queue-based sequential processing.
|
|
|
|
Uses huggingface_hub's upload_folder() with a semaphore to ensure only
|
|
one upload runs at a time. Multiple uploads can be queued.
|
|
|
|
Args:
|
|
model_dirs: List of model directory paths to scan for oQ models.
|
|
"""
|
|
|
|
def __init__(self, model_dirs: list[str]):
|
|
self._model_dirs = [Path(d) for d in model_dirs]
|
|
self._tasks: dict[str, UploadTask] = {}
|
|
self._active_tasks: dict[str, asyncio.Task] = {}
|
|
self._cancelled: set[str] = set()
|
|
self._upload_sem = asyncio.Semaphore(1)
|
|
|
|
def update_model_dirs(self, model_dirs: list[str]) -> None:
|
|
"""Update model directory paths."""
|
|
self._model_dirs = [Path(d) for d in model_dirs]
|
|
|
|
@staticmethod
|
|
async def validate_token(token: str) -> dict:
|
|
"""Validate a HuggingFace token and return user info.
|
|
|
|
Args:
|
|
token: HuggingFace write-access token.
|
|
|
|
Returns:
|
|
Dict with 'username' and 'orgs' list.
|
|
|
|
Raises:
|
|
ValueError: If the token is invalid or lacks write access.
|
|
"""
|
|
from huggingface_hub import HfApi
|
|
|
|
try:
|
|
api = HfApi()
|
|
info = await asyncio.to_thread(api.whoami, token=token)
|
|
except Exception as e:
|
|
raise ValueError(f"Invalid token: {e}")
|
|
|
|
username = info.get("name", "")
|
|
orgs = [
|
|
{"name": org.get("name", "")}
|
|
for org in info.get("orgs", [])
|
|
if org.get("name")
|
|
]
|
|
|
|
# Check for write access
|
|
auth = info.get("auth", {})
|
|
access_token = auth.get("accessToken", {})
|
|
role = access_token.get("role", "")
|
|
if role == "read":
|
|
raise ValueError(
|
|
"Token has read-only access. A write token is required for uploads."
|
|
)
|
|
|
|
return {"username": username, "orgs": orgs}
|
|
|
|
async def list_oq_models(self) -> list[dict]:
|
|
"""Scan model directories and return oQ-quantized models.
|
|
|
|
Returns:
|
|
List of dicts with model name, path, size info.
|
|
"""
|
|
|
|
def _scan() -> list[dict]:
|
|
models = []
|
|
seen: set[str] = set()
|
|
|
|
for model_dir in self._model_dirs:
|
|
if not model_dir.exists():
|
|
continue
|
|
for subdir in sorted(model_dir.iterdir()):
|
|
if not subdir.is_dir():
|
|
continue
|
|
candidates = []
|
|
if (subdir / "config.json").exists():
|
|
candidates.append(subdir)
|
|
else:
|
|
for child in sorted(subdir.iterdir()):
|
|
if child.is_dir() and (child / "config.json").exists():
|
|
candidates.append(child)
|
|
|
|
for path in candidates:
|
|
if path.name in seen:
|
|
continue
|
|
seen.add(path.name)
|
|
if not _is_oq_model(path.name):
|
|
continue
|
|
try:
|
|
size = sum(
|
|
f.stat().st_size
|
|
for f in path.glob("*.safetensors")
|
|
)
|
|
if size == 0:
|
|
continue
|
|
models.append({
|
|
"name": path.name,
|
|
"path": str(path),
|
|
"size": size,
|
|
"size_formatted": _format_size(size),
|
|
})
|
|
except Exception:
|
|
continue
|
|
return models
|
|
|
|
return await asyncio.to_thread(_scan)
|
|
|
|
async def list_all_models(self) -> list[dict]:
|
|
"""Scan model directories and return all models (for README source selection).
|
|
|
|
Returns:
|
|
List of dicts with model name and path.
|
|
"""
|
|
|
|
def _scan() -> list[dict]:
|
|
models = []
|
|
seen: set[str] = set()
|
|
|
|
for model_dir in self._model_dirs:
|
|
if not model_dir.exists():
|
|
continue
|
|
for subdir in sorted(model_dir.iterdir()):
|
|
if not subdir.is_dir():
|
|
continue
|
|
candidates = []
|
|
if (subdir / "config.json").exists():
|
|
candidates.append(subdir)
|
|
else:
|
|
for child in sorted(subdir.iterdir()):
|
|
if child.is_dir() and (child / "config.json").exists():
|
|
candidates.append(child)
|
|
|
|
for path in candidates:
|
|
if path.name in seen:
|
|
continue
|
|
seen.add(path.name)
|
|
has_readme = _has_meaningful_readme(path)
|
|
models.append({
|
|
"name": path.name,
|
|
"path": str(path),
|
|
"has_readme": has_readme,
|
|
})
|
|
return models
|
|
|
|
return await asyncio.to_thread(_scan)
|
|
|
|
async def start_upload(
|
|
self,
|
|
model_path: str,
|
|
repo_id: str,
|
|
token: str,
|
|
readme_source_path: str = "",
|
|
auto_readme: bool = True,
|
|
redownload_notice: bool = False,
|
|
private: bool = False,
|
|
) -> UploadTask:
|
|
"""Queue a model upload to HuggingFace Hub.
|
|
|
|
Args:
|
|
model_path: Local path to the oQ model directory.
|
|
repo_id: Target HuggingFace repository ID (e.g., 'user/model-oQ4').
|
|
token: HuggingFace write token.
|
|
readme_source_path: Optional path to model whose README.md to copy.
|
|
auto_readme: If True and no readme_source_path, generate a basic README.
|
|
private: If True, create a private repository.
|
|
|
|
Returns:
|
|
The created UploadTask.
|
|
|
|
Raises:
|
|
ValueError: If model path is invalid or upload is already queued.
|
|
"""
|
|
source = Path(model_path)
|
|
if not source.exists() or not source.is_dir():
|
|
raise ValueError(f"Model directory not found: {model_path}")
|
|
|
|
if not (source / "config.json").exists():
|
|
raise ValueError(f"Not a valid model directory (no config.json): {model_path}")
|
|
|
|
repo_id = repo_id.strip()
|
|
if "/" not in repo_id or len(repo_id.split("/")) != 2:
|
|
raise ValueError(
|
|
f"Invalid repository ID: '{repo_id}'. "
|
|
"Expected format: 'owner/model' (e.g., 'user/Llama-3B-oQ4')"
|
|
)
|
|
|
|
# Check for duplicate active uploads
|
|
for task in self._tasks.values():
|
|
if task.repo_id == repo_id and task.status in _ACTIVE_STATUSES:
|
|
raise ValueError(
|
|
f"Upload to '{repo_id}' is already in progress"
|
|
)
|
|
|
|
model_name = source.name
|
|
total_size = sum(
|
|
f.stat().st_size for f in source.rglob("*") if f.is_file()
|
|
)
|
|
|
|
task_id = str(uuid.uuid4())
|
|
task = UploadTask(
|
|
task_id=task_id,
|
|
model_name=model_name,
|
|
model_path=model_path,
|
|
repo_id=repo_id,
|
|
total_size=total_size,
|
|
)
|
|
self._tasks[task_id] = task
|
|
|
|
self._active_tasks[task_id] = asyncio.create_task(
|
|
self._run_upload(task_id, token, readme_source_path, auto_readme, redownload_notice, private)
|
|
)
|
|
|
|
logger.info(f"Upload queued: {model_name} -> {repo_id} (task_id={task_id})")
|
|
return task
|
|
|
|
async def cancel_upload(self, task_id: str) -> bool:
|
|
"""Cancel an active or pending upload.
|
|
|
|
Args:
|
|
task_id: The task ID to cancel.
|
|
|
|
Returns:
|
|
True if the task was found and cancelled.
|
|
"""
|
|
task = self._tasks.get(task_id)
|
|
if task is None:
|
|
return False
|
|
|
|
if task.status not in _ACTIVE_STATUSES:
|
|
return False
|
|
|
|
self._cancelled.add(task_id)
|
|
task.status = UploadStatus.CANCELLED
|
|
|
|
active_task = self._active_tasks.pop(task_id, None)
|
|
if active_task and not active_task.done():
|
|
active_task.cancel()
|
|
|
|
logger.info(f"Upload cancelled: {task.model_name} (task_id={task_id})")
|
|
return True
|
|
|
|
def remove_task(self, task_id: str) -> bool:
|
|
"""Remove a completed, failed, or cancelled task from the list.
|
|
|
|
Args:
|
|
task_id: The task ID to remove.
|
|
|
|
Returns:
|
|
True if the task was found and removed.
|
|
"""
|
|
task = self._tasks.get(task_id)
|
|
if task is None:
|
|
return False
|
|
|
|
if task.status in _ACTIVE_STATUSES:
|
|
return False
|
|
|
|
del self._tasks[task_id]
|
|
self._cancelled.discard(task_id)
|
|
return True
|
|
|
|
def get_tasks(self) -> list[dict]:
|
|
"""Return all tasks as serializable dicts, ordered by creation time."""
|
|
return [
|
|
task.to_dict()
|
|
for task in sorted(self._tasks.values(), key=lambda t: t.created_at)
|
|
]
|
|
|
|
async def shutdown(self) -> None:
|
|
"""Cancel all active uploads and clean up."""
|
|
for task_id, active_task in list(self._active_tasks.items()):
|
|
if not active_task.done():
|
|
active_task.cancel()
|
|
task = self._tasks.get(task_id)
|
|
if task and task.status == UploadStatus.UPLOADING:
|
|
task.status = UploadStatus.CANCELLED
|
|
self._active_tasks.clear()
|
|
logger.info("HF Uploader shut down")
|
|
|
|
async def _run_upload(
|
|
self,
|
|
task_id: str,
|
|
token: str,
|
|
readme_source_path: str,
|
|
auto_readme: bool,
|
|
redownload_notice: bool,
|
|
private: bool,
|
|
) -> None:
|
|
"""Execute an upload task with semaphore-guarded sequential processing."""
|
|
from huggingface_hub import HfApi
|
|
|
|
task = self._tasks[task_id]
|
|
tmp_readme: Optional[Path] = None
|
|
|
|
try:
|
|
async with self._upload_sem:
|
|
if task_id in self._cancelled:
|
|
return
|
|
|
|
task.status = UploadStatus.UPLOADING
|
|
task.started_at = time.time()
|
|
|
|
model_path = Path(task.model_path)
|
|
api = HfApi()
|
|
|
|
# Create repo (exist_ok handles already-existing repos)
|
|
await asyncio.to_thread(
|
|
api.create_repo,
|
|
repo_id=task.repo_id,
|
|
token=token,
|
|
exist_ok=True,
|
|
private=private,
|
|
)
|
|
|
|
if task_id in self._cancelled:
|
|
return
|
|
|
|
# Handle README
|
|
readme_in_model = model_path / "README.md"
|
|
if readme_source_path:
|
|
source_readme = Path(readme_source_path) / "README.md"
|
|
if source_readme.exists():
|
|
shutil.copy2(source_readme, readme_in_model)
|
|
tmp_readme = readme_in_model
|
|
elif auto_readme and not _has_meaningful_readme(model_path):
|
|
try:
|
|
with open(model_path / "config.json") as f:
|
|
config = json.load(f)
|
|
except Exception:
|
|
config = {}
|
|
readme_content = _generate_model_card(
|
|
task.model_name, config,
|
|
redownload_notice=redownload_notice,
|
|
)
|
|
readme_in_model.write_text(readme_content, encoding="utf-8")
|
|
tmp_readme = readme_in_model
|
|
|
|
if task_id in self._cancelled:
|
|
return
|
|
|
|
# Upload the entire model folder
|
|
# upload_folder is blocking; run in thread
|
|
task.progress = 10.0 # Signal that upload has started
|
|
|
|
await asyncio.to_thread(
|
|
api.upload_folder,
|
|
folder_path=str(model_path),
|
|
repo_id=task.repo_id,
|
|
token=token,
|
|
commit_message=f"Upload {task.model_name} via oMLX",
|
|
)
|
|
|
|
if task_id in self._cancelled:
|
|
return
|
|
|
|
# Success
|
|
task.status = UploadStatus.COMPLETED
|
|
task.progress = 100.0
|
|
task.completed_at = time.time()
|
|
task.repo_url = f"https://huggingface.co/{task.repo_id}"
|
|
|
|
elapsed = task.completed_at - task.started_at
|
|
logger.info(
|
|
f"Upload completed: {task.model_name} -> {task.repo_id} "
|
|
f"({elapsed:.0f}s, {_format_size(task.total_size)})"
|
|
)
|
|
|
|
except asyncio.CancelledError:
|
|
if task.status not in (UploadStatus.CANCELLED, UploadStatus.FAILED):
|
|
task.status = UploadStatus.CANCELLED
|
|
except Exception as e:
|
|
if task_id not in self._cancelled:
|
|
task.status = UploadStatus.FAILED
|
|
task.error = str(e)
|
|
logger.error(f"Upload failed for {task.model_name}: {e}")
|
|
finally:
|
|
# Clean up copied/generated README if we created it
|
|
if tmp_readme and tmp_readme.exists():
|
|
try:
|
|
tmp_readme.unlink()
|
|
except Exception:
|
|
pass
|
|
self._active_tasks.pop(task_id, None)
|