diff --git a/studio/backend/hub/__init__.py b/studio/backend/hub/__init__.py new file mode 100644 index 0000000000..706dae9224 --- /dev/null +++ b/studio/backend/hub/__init__.py @@ -0,0 +1,8 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Hub + Download Manager feature module. + +Self-contained routes, schemas, utilities, workers, and storage for the model +inventory layer and the HuggingFace download manager. Wired into the FastAPI +app via two routers plus startup/shutdown hooks in main.py.""" diff --git a/studio/backend/hub/dependencies.py b/studio/backend/hub/dependencies.py new file mode 100644 index 0000000000..ff78dce8f6 --- /dev/null +++ b/studio/backend/hub/dependencies.py @@ -0,0 +1,24 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Shared FastAPI dependencies for Hub routes.""" + +from __future__ import annotations + +from typing import Optional + +from fastapi import Header + +HUB_HF_TOKEN_HEADER = "X-Unsloth-HF-Token" +HUB_HF_TOKEN_MAX_LENGTH = 512 + + +def get_hf_token( + hf_token: Optional[str] = Header( + None, + alias = HUB_HF_TOKEN_HEADER, + max_length = HUB_HF_TOKEN_MAX_LENGTH, + ), +) -> Optional[str]: + token = (hf_token or "").strip() + return token or None diff --git a/studio/backend/hub/routes/__init__.py b/studio/backend/hub/routes/__init__.py new file mode 100644 index 0000000000..e9579635b0 --- /dev/null +++ b/studio/backend/hub/routes/__init__.py @@ -0,0 +1,12 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Hub routers exposed at /api/hub/* and /api/hub/datasets/*.""" + +from hub.routes.inventory import router as inventory_router +from hub.routes.datasets import router as datasets_router + +__all__ = [ + "inventory_router", + "datasets_router", +] diff --git a/studio/backend/hub/routes/datasets.py b/studio/backend/hub/routes/datasets.py new file mode 100644 index 0000000000..edf4f36ac0 --- /dev/null +++ b/studio/backend/hub/routes/datasets.py @@ -0,0 +1,138 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Endpoints mounted at /api/hub/datasets/*.""" + +from __future__ import annotations + +from typing import Optional + +from fastapi import APIRouter, Body, Depends, Query, UploadFile + +from auth.authentication import get_current_subject +from hub.dependencies import get_hf_token +from hub.schemas.datasets import ( + AiAssistMappingRequest, + AiAssistMappingResponse, + CachedDatasetsResponse, + CheckFormatRequest, + CheckFormatResponse, + DeleteCachedDatasetResponse, + LocalDatasetsResponse, + UploadDatasetResponse, +) +from hub.schemas.downloads import ( + ActiveDownloadsResponse, + CancelDatasetDownloadRequest, + CancelDatasetDownloadResponse, + DatasetDownloadJobStatus, + DatasetDownloadStartResponse, + DownloadProgressResponse, + DownloadDatasetRequest, + TransportStatusResponse, +) +from hub.services.datasets import cache_inventory, downloads, formatting, local + +router = APIRouter() + + +@router.post("/upload", response_model = UploadDatasetResponse) +async def upload_dataset( + file: UploadFile, current_subject: str = Depends(get_current_subject) +) -> UploadDatasetResponse: + return await local.upload_dataset_response(file) + + +@router.get("/local", response_model = LocalDatasetsResponse) +def list_local_datasets( + current_subject: str = Depends(get_current_subject), +) -> LocalDatasetsResponse: + return local.list_local_datasets_response() + + +@router.get( + "/cached", + response_model = CachedDatasetsResponse, + response_model_exclude_unset = True, +) +async def list_cached_datasets(current_subject: str = Depends(get_current_subject)): + return await cache_inventory.list_cached_datasets_response() + + +@router.delete("/cached", response_model = DeleteCachedDatasetResponse) +async def delete_cached_dataset( + repo_id: str = Body(..., embed = True), current_subject: str = Depends(get_current_subject) +): + return await cache_inventory.delete_cached_dataset_response(repo_id) + + +@router.get("/download-progress", response_model = DownloadProgressResponse) +async def get_dataset_download_progress( + repo_id: str = Query(..., description = "HuggingFace dataset repo ID, e.g. 'unsloth/LaTeX_OCR'"), + expected_bytes: int = Query(0, description = "Expected total download size in bytes"), + hf_token: Optional[str] = Depends(get_hf_token), + current_subject: str = Depends(get_current_subject), +): + return await downloads.get_dataset_download_progress_response( + repo_id, + expected_bytes = expected_bytes, + hf_token = hf_token, + ) + + +@router.post("/download", response_model = DatasetDownloadStartResponse, status_code = 202) +async def download_dataset( + body: DownloadDatasetRequest, + hf_token: Optional[str] = Depends(get_hf_token), + current_subject: str = Depends(get_current_subject), +): + return await downloads.download_dataset_response(body, hf_token) + + +@router.post("/download/cancel", response_model = CancelDatasetDownloadResponse, status_code = 202) +async def cancel_dataset_download( + body: CancelDatasetDownloadRequest, current_subject: str = Depends(get_current_subject) +): + return await downloads.cancel_dataset_download_response(body) + + +@router.get("/download-status", response_model = DatasetDownloadJobStatus) +async def get_dataset_download_status( + repo_id: str = Query(..., description = "HuggingFace dataset repo ID"), + current_subject: str = Depends(get_current_subject), +): + return await downloads.get_dataset_download_status_response(repo_id) + + +@router.get("/active-downloads", response_model = ActiveDownloadsResponse) +async def get_active_dataset_downloads( + repo_id: str = Query("", description = "HuggingFace dataset repo ID"), + current_subject: str = Depends(get_current_subject), +): + return await downloads.get_active_dataset_downloads_response(repo_id) + + +@router.get("/transport-status", response_model = TransportStatusResponse) +async def get_dataset_transport_status( + repo_id: str = Query(..., description = "HuggingFace dataset repo ID"), + current_subject: str = Depends(get_current_subject), +): + return await downloads.get_dataset_transport_status_response(repo_id) + + +@router.post("/check-format", response_model = CheckFormatResponse) +def check_format( + request: CheckFormatRequest, + hf_token: Optional[str] = Depends(get_hf_token), + current_subject: str = Depends(get_current_subject), +): + return formatting.check_format_response(request, hf_token) + + +@router.post("/ai-assist-mapping", response_model = AiAssistMappingResponse) +def ai_assist_mapping( + request: AiAssistMappingRequest, + hf_token: Optional[str] = Depends(get_hf_token), + current_subject: str = Depends(get_current_subject), +): + return formatting.ai_assist_mapping_response(request, hf_token) diff --git a/studio/backend/hub/routes/inventory.py b/studio/backend/hub/routes/inventory.py new file mode 100644 index 0000000000..fcfdd0ad14 --- /dev/null +++ b/studio/backend/hub/routes/inventory.py @@ -0,0 +1,222 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Endpoints mounted at /api/hub/* for the model inventory.""" + +from __future__ import annotations + +from typing import Optional + +from fastapi import APIRouter, Body, Depends, Query + +from auth.authentication import get_current_subject +from hub.dependencies import get_hf_token +from hub.schemas.downloads import ( + ActiveDownloadsResponse, + CancelDownloadResponse, + CancelDownloadRequest, + DownloadProgressResponse, + DownloadJobStatus, + DownloadModelRequest, + DownloadStartResponse, + TransportStatusResponse, +) +from hub.schemas.inventory import ( + AddScanFolderRequest, + BrowseFoldersResponse, + CachedGgufResponse, + CachedModelsResponse, + DeleteCachedModelResponse, + GgufVariantsResponse, + LocalModelListResponse, + RecommendedFoldersResponse, + RemoveScanFolderResponse, + ScanFolderInfo, + ScanFoldersResponse, +) +from hub.services.models import ( + cache_inventory, + deletion, + downloads, + folder_browser, + gguf_variants, + local_inventory, +) + +router = APIRouter() + + +@router.get("/local", response_model = LocalModelListResponse) +async def list_local_models( + models_dir: str = Query( + default = "./models", description = "Directory to scan for local model folders" + ), + current_subject: str = Depends(get_current_subject), +): + return await local_inventory.list_local_models_response(models_dir) + + +# Plain `def` (not async): synchronous SQLite + filesystem work runs in +# FastAPI's thread-pool instead of blocking the event loop. +@router.get("/scan-folders", response_model = ScanFoldersResponse) +def get_scan_folders(current_subject: str = Depends(get_current_subject)): + return local_inventory.get_scan_folders_response() + + +@router.post("/scan-folders", response_model = ScanFolderInfo, status_code = 201) +def add_scan_folder_endpoint( + body: AddScanFolderRequest, current_subject: str = Depends(get_current_subject) +): + return local_inventory.add_scan_folder_response(body.path) + + +@router.delete("/scan-folders/{folder_id}", response_model = RemoveScanFolderResponse) +def remove_scan_folder_endpoint( + folder_id: int, current_subject: str = Depends(get_current_subject) +): + return local_inventory.remove_scan_folder_response(folder_id) + + +@router.get("/recommended-folders", response_model = RecommendedFoldersResponse) +def get_recommended_folders(current_subject: str = Depends(get_current_subject)): + return folder_browser.get_recommended_folders_response() + + +@router.get("/browse-folders", response_model = BrowseFoldersResponse) +def browse_folders( + path: Optional[str] = Query(None), + show_hidden: bool = Query(False), + current_subject: str = Depends(get_current_subject), +): + return folder_browser.browse_folders_response(path, show_hidden) + + +@router.get("/gguf-variants", response_model = GgufVariantsResponse) +async def get_gguf_variants( + repo_id: str = Query( + ..., description = "HuggingFace repo ID (e.g. 'unsloth/gemma-3-4b-it-GGUF')" + ), + prefer_local_cache: bool = Query(False), + offline: bool = Query(False), + local_path: Optional[str] = Query(None), + hf_token: Optional[str] = Depends(get_hf_token), + current_subject: str = Depends(get_current_subject), +): + return await gguf_variants.get_gguf_variants_response( + repo_id, + prefer_local_cache = prefer_local_cache, + offline = offline, + local_path = local_path, + hf_token = hf_token, + ) + + +@router.post("/download", response_model = DownloadStartResponse, status_code = 202) +async def download_model( + body: DownloadModelRequest, + hf_token: Optional[str] = Depends(get_hf_token), + current_subject: str = Depends(get_current_subject), +): + return await downloads.download_model_response(body, hf_token) + + +@router.post("/download/cancel", response_model = CancelDownloadResponse, status_code = 202) +async def cancel_download_model( + body: CancelDownloadRequest, current_subject: str = Depends(get_current_subject) +): + return await downloads.cancel_download_model_response(body) + + +@router.get("/download-status", response_model = DownloadJobStatus) +async def get_download_status( + repo_id: str = Query(..., description = "HuggingFace repo ID"), + gguf_variant: str = Query("", description = "Quantization variant (empty for safetensors)"), + current_subject: str = Depends(get_current_subject), +): + return await downloads.get_download_status_response(repo_id, gguf_variant) + + +@router.get("/active-downloads", response_model = ActiveDownloadsResponse) +async def get_active_downloads( + repo_id: str = Query("", description = "HuggingFace repo ID"), + current_subject: str = Depends(get_current_subject), +): + return await downloads.get_active_downloads_response(repo_id) + + +@router.get("/transport-status", response_model = TransportStatusResponse) +async def get_model_transport_status( + repo_id: str = Query(..., description = "HuggingFace repo ID"), + gguf_variant: str = Query("", description = "Quantization variant (empty for safetensors)"), + hf_token: Optional[str] = Depends(get_hf_token), + current_subject: str = Depends(get_current_subject), +): + return await downloads.get_model_transport_status_response( + repo_id, + gguf_variant, + hf_token, + ) + + +@router.get( + "/gguf-download-progress", + response_model = DownloadProgressResponse, + response_model_exclude_none = True, +) +async def get_gguf_download_progress( + repo_id: str = Query(..., description = "HuggingFace repo ID"), + variant: str = Query("", description = "Quantization variant (e.g. UD-TQ1_0)"), + expected_bytes: int = Query(0, description = "Expected total download size in bytes"), + hf_token: Optional[str] = Depends(get_hf_token), + current_subject: str = Depends(get_current_subject), +): + return await downloads.get_gguf_download_progress_response( + repo_id, + variant = variant, + expected_bytes = expected_bytes, + hf_token = hf_token, + ) + + +@router.get("/download-progress", response_model = DownloadProgressResponse) +async def get_download_progress( + repo_id: str = Query(..., description = "HuggingFace repo ID"), + expected_bytes: int = Query(0, description = "Expected total download size in bytes"), + hf_token: Optional[str] = Depends(get_hf_token), + current_subject: str = Depends(get_current_subject), +): + return await downloads.get_download_progress_response( + repo_id, + expected_bytes = expected_bytes, + hf_token = hf_token, + ) + + +@router.get("/cached-gguf", response_model = CachedGgufResponse) +async def list_cached_gguf( + hf_token: Optional[str] = Depends(get_hf_token), + current_subject: str = Depends(get_current_subject), +): + return await cache_inventory.list_cached_gguf_response(hf_token) + + +@router.get("/cached-models", response_model = CachedModelsResponse) +async def list_cached_models( + hf_token: Optional[str] = Depends(get_hf_token), + current_subject: str = Depends(get_current_subject), +): + return await cache_inventory.list_cached_models_response(hf_token) + + +@router.delete( + "/delete-cached", + response_model = DeleteCachedModelResponse, + response_model_exclude_none = True, +) +async def delete_cached_model( + repo_id: str = Body(...), + variant: Optional[str] = Body(None), + hf_token: Optional[str] = Depends(get_hf_token), + current_subject: str = Depends(get_current_subject), +): + return await deletion.delete_cached_model_response(repo_id, variant, hf_token) diff --git a/studio/backend/hub/schemas/__init__.py b/studio/backend/hub/schemas/__init__.py new file mode 100644 index 0000000000..32014236c6 --- /dev/null +++ b/studio/backend/hub/schemas/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 diff --git a/studio/backend/hub/schemas/datasets.py b/studio/backend/hub/schemas/datasets.py new file mode 100644 index 0000000000..02365b1992 --- /dev/null +++ b/studio/backend/hub/schemas/datasets.py @@ -0,0 +1,108 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from __future__ import annotations + +from typing import Any, Dict, List, Literal, Optional + +from pydantic import BaseModel, Field, model_validator + + +class CheckFormatRequest(BaseModel): + dataset_name: str + is_vlm: bool = False + subset: Optional[str] = None + train_split: Optional[str] = "train" + prefer_local_cache: bool = False + local_path: Optional[str] = None + + @model_validator(mode = "before") + @classmethod + def _compat_split(cls, values: Any) -> Any: + if isinstance(values, dict) and "split" in values: + merged = {**values} + merged.setdefault("train_split", merged.pop("split")) + return merged + return values + + +class CheckFormatResponse(BaseModel): + requires_manual_mapping: bool + detected_format: str + columns: List[str] + is_image: bool = False + is_audio: bool = False + multimodal_columns: Optional[List[str]] = None + suggested_mapping: Optional[Dict[str, str]] = None + detected_image_column: Optional[str] = None + detected_audio_column: Optional[str] = None + detected_text_column: Optional[str] = None + detected_speaker_column: Optional[str] = None + preview_samples: Optional[List[Dict]] = None + total_rows: Optional[int] = None + warning: Optional[str] = None + + +class AiAssistMappingRequest(BaseModel): + columns: List[str] + samples: List[Dict[str, Any]] + dataset_name: Optional[str] = None + model_name: Optional[str] = None + model_type: Optional[str] = None + + +class AiAssistMappingResponse(BaseModel): + success: bool + suggested_mapping: Optional[Dict[str, str]] = None + warning: Optional[str] = None + system_prompt: Optional[str] = None + user_template: Optional[str] = None + assistant_template: Optional[str] = None + label_mapping: Optional[Dict[str, Dict[str, str]]] = None + dataset_type: Optional[str] = None + is_conversational: Optional[bool] = None + user_notification: Optional[str] = None + + +class UploadDatasetResponse(BaseModel): + filename: str = Field(..., description = "Original filename") + stored_path: str = Field(..., description = "Absolute path stored on backend") + + +class LocalDatasetItem(BaseModel): + class Metadata(BaseModel): + actual_num_records: Optional[int] = None + target_num_records: Optional[int] = None + total_num_batches: Optional[int] = None + num_completed_batches: Optional[int] = None + columns: Optional[List[str]] = None + + id: str + label: str + path: str + source: Literal["recipe", "upload"] + rows: Optional[int] = None + updated_at: Optional[float] = None + metadata: Optional[Metadata] = None + + +class LocalDatasetsResponse(BaseModel): + datasets: List[LocalDatasetItem] = Field(default_factory = list) + + +class CachedDatasetItem(BaseModel): + repo_id: str + size_bytes: int = 0 + cache_path: Optional[str] = None + processed_cache: bool = False + partial: bool = False + partial_transport: Optional[str] = None + + +class CachedDatasetsResponse(BaseModel): + cached: List[CachedDatasetItem] = Field(default_factory = list) + + +class DeleteCachedDatasetResponse(BaseModel): + status: str + repo_id: str diff --git a/studio/backend/hub/schemas/downloads.py b/studio/backend/hub/schemas/downloads.py new file mode 100644 index 0000000000..dccd0c6733 --- /dev/null +++ b/studio/backend/hub/schemas/downloads.py @@ -0,0 +1,161 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Pydantic schemas for the Hub download manager (/api/hub/downloads/*).""" + +from pydantic import BaseModel, Field +from typing import List, Literal, Optional + + +DownloadJobState = Literal["idle", "running", "cancelling", "cancelled", "complete", "error"] + + +class DownloadModelRequest(BaseModel): + """Body for POST /api/hub/download. + + The HuggingFace token travels in the internal Hub token header. + """ + + repo_id: str = Field( + ..., + description = "HuggingFace repo ID, e.g. 'unsloth/Qwen3-4B-GGUF'", + ) + gguf_variant: Optional[str] = Field( + None, + description = "Quantization label (e.g. 'Q4_K_M'). Required for GGUF repos.", + ) + use_xet: bool = Field( + False, + description = "Enable Xet parallel chunked transport. Default False uses HTTP Range-resume.", + ) + + +class CancelDownloadRequest(BaseModel): + repo_id: str = Field(..., description = "HuggingFace repo ID") + gguf_variant: Optional[str] = Field( + None, + description = "GGUF variant label; omit for safetensors snapshots", + ) + generation: Optional[int] = Field( + None, + description = "Download generation tag from a prior start; passing it scopes the cancel to that exact run.", + ) + + +class DownloadJobStatus(BaseModel): + """Live state of a background download job.""" + + state: DownloadJobState = Field( + ..., + description = "Current download job state.", + ) + error: Optional[str] = Field(None, description = "Error message if state == 'error'") + generation: int = Field( + 0, + description = "Current run generation; an adopting client stores it so a later cancel is scoped to this exact run.", + ) + + +class DownloadStartResponse(BaseModel): + job_key: str + state: str + accepted: bool + generation: int + + +class CancelDownloadResponse(BaseModel): + job_key: str + state: str + + +class ActiveDownload(BaseModel): + """One in-flight download for a repo. ``variant`` is null for safetensors.""" + + repo_id: Optional[str] = None + variant: Optional[str] = None + transport: Optional[str] = None + state: str + generation: int = Field( + 0, + description = "Current run generation; an adopting client stores it so a later cancel is scoped to this exact run.", + ) + + +class ActiveDownloadsResponse(BaseModel): + downloads: List[ActiveDownload] + + +class TransportCapability(BaseModel): + available: bool + reason: Optional[str] = None + + +class TransportCapabilities(BaseModel): + http: TransportCapability + xet: TransportCapability + + +class TransportStatusResponse(BaseModel): + has_partial: bool + last_transport: Optional[str] = None + resumable: bool + + +class DownloadProgressResponse(BaseModel): + downloaded_bytes: int + # Finalized-blob bytes only (no ``.incomplete``). Registry-loss completion + # fallbacks key off this so a partial isn't mistaken for a finished download. + completed_bytes: int = 0 + complete_on_disk: bool = Field( + False, + description = ( + "True only when the backend verified a usable completed snapshot/variant on disk." + ), + ) + expected_bytes: int + progress: float + cache_path: Optional[str] = None + + +class DownloadDatasetRequest(BaseModel): + """Body for POST /api/hub/datasets/download. + + The HuggingFace token travels in the internal Hub token header. + """ + + repo_id: str = Field(..., description = "HuggingFace dataset repo ID") + use_xet: bool = Field( + False, + description = "Enable Xet parallel chunked transport. Default False uses HTTP Range-resume.", + ) + + +class CancelDatasetDownloadRequest(BaseModel): + repo_id: str = Field(..., description = "HuggingFace dataset repo ID") + generation: Optional[int] = Field(None, description = "Download generation") + + +class DatasetDownloadJobStatus(BaseModel): + """Live state of a background dataset download job.""" + + state: DownloadJobState = Field( + ..., + description = "Current dataset download job state.", + ) + error: Optional[str] = Field(None, description = "Error message if state == 'error'") + generation: int = Field( + 0, + description = "Current run generation; an adopting client stores it so a later cancel is scoped to this exact run.", + ) + + +class DatasetDownloadStartResponse(BaseModel): + repo_id: str + state: str + accepted: bool + generation: int + + +class CancelDatasetDownloadResponse(BaseModel): + repo_id: str + state: str diff --git a/studio/backend/hub/schemas/inventory.py b/studio/backend/hub/schemas/inventory.py new file mode 100644 index 0000000000..c333c7ca89 --- /dev/null +++ b/studio/backend/hub/schemas/inventory.py @@ -0,0 +1,286 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Pydantic schemas for the Hub inventory layer (/api/hub/*). + +Kept independent from upstream models/models.py so the Hub module can ship +without modifying any upstream schema.""" + +from pydantic import BaseModel, Field +from typing import List, Literal, Optional + + +ModelFormat = Literal["gguf", "safetensors", "adapter", "checkpoint", "unknown"] +ModelRuntime = Literal["llama_cpp", "transformers", "adapter", "unknown"] + + +class GgufVariantDetail(BaseModel): + """A single GGUF quantization variant in a HuggingFace repo.""" + + filename: str = Field(..., description = "GGUF filename (e.g., 'gemma-3-4b-it-Q4_K_M.gguf')") + quant: str = Field(..., description = "Quantization label or internal GGUF variant key") + display_label: Optional[str] = Field( + None, description = "Optional user-facing label when quant is an internal key" + ) + size_bytes: int = Field(0, description = "File size in bytes") + download_size_bytes: int = Field(0, description = "Total bytes needed to download this variant") + downloaded: bool = Field( + False, description = "Whether this variant is already in the local HF cache" + ) + partial: bool = Field( + False, + description = "Whether this variant has an in-progress (.incomplete) blob in cache", + ) + partial_transport: Optional[str] = Field( + None, + description = ( + 'Transport recorded for the partial state ("http" or ' + '"xet"), or null if not partial / unknown. Frontend uses ' + "this to pick Resume (http) vs Redownload (xet) labels." + ), + ) + + +class GgufVariantsResponse(BaseModel): + """Response for listing GGUF quantization variants in a HuggingFace repo.""" + + repo_id: str = Field(..., description = "HuggingFace repo ID") + variants: List[GgufVariantDetail] = Field( + default_factory = list, description = "Available GGUF variants" + ) + has_vision: bool = Field( + False, description = "Whether the model has vision support (mmproj files)" + ) + default_variant: Optional[str] = Field( + None, description = "Recommended default quantization variant" + ) + + +class LocalModelCapabilities(BaseModel): + can_train: bool = False + can_chat: bool = False + can_delete: bool = False + can_download: bool = False + requires_variant: bool = False + supports_lora: bool = False + supports_vision: bool = False + + +class LocalModelInfo(BaseModel): + """Discovered local model candidate.""" + + id: str = Field(..., description = "Identifier to use for loading/training") + inventory_id: Optional[str] = Field( + None, description = "Stable semantic inventory row identifier" + ) + load_id: Optional[str] = Field( + None, description = "Identifier/path to pass to load or train APIs" + ) + display_name: str = Field(..., description = "Display label") + path: str = Field(..., description = "Local path where model data was discovered") + size_bytes: int = Field(0, description = "Observed model artifact size in bytes") + model_format: ModelFormat = Field("unknown", description = "Model file format") + runtime: ModelRuntime = Field("unknown", description = "Expected runtime backend") + format_variant: Optional[str] = Field( + None, description = "Format variant label, for example a GGUF quant" + ) + capabilities: LocalModelCapabilities = Field( + default_factory = LocalModelCapabilities, + description = "Declared capabilities for this inventory row", + ) + source: Literal["models_dir", "hf_cache", "lmstudio", "ollama", "custom"] = Field( + ..., + description = "Discovery source", + ) + model_id: Optional[str] = Field( + None, + description = "HF repo id for cached models, e.g. org/model", + ) + base_model: Optional[str] = Field( + None, + description = "Base model from adapter_config.json when this is an adapter", + ) + base_model_source: Optional[Literal["huggingface", "local", "unknown"]] = Field( + None, + description = "Whether the adapter base model is a HF repo id or local path", + ) + adapter_type: Optional[str] = Field( + None, + description = "Adapter type from adapter_config.json, e.g. LORA", + ) + training_method: Optional[str] = Field( + None, + description = "Training method hint from adapter_config.json", + ) + updated_at: Optional[float] = Field( + None, + description = "Unix timestamp of latest observed update", + ) + partial: bool = Field( + False, + description = "True when this hf_cache entry has incomplete blobs", + ) + partial_transport: Optional[str] = Field( + None, + description = ( + 'Transport recorded for the partial state ("http" or ' + '"xet"), or null if not partial / unknown.' + ), + ) + + +class LocalModelListResponse(BaseModel): + """Response schema for listing local/cached models.""" + + models_dir: str = Field(..., description = "Directory scanned for custom local models") + hf_cache_dir: Optional[str] = Field( + None, + description = "HF cache root that was scanned", + ) + lmstudio_dirs: List[str] = Field( + default_factory = list, + description = "LM Studio model directories that were scanned", + ) + ollama_dirs: List[str] = Field( + default_factory = list, + description = "Ollama model directories that were scanned", + ) + models: List[LocalModelInfo] = Field( + default_factory = list, + description = "Discovered local/cached models", + ) + + +class CachedRepoBase(BaseModel): + """Shared shape for a cached HF repo row surfaced under On Device.""" + + repo_id: str + size_bytes: int = 0 + cache_path: Optional[str] = None + partial: bool = False + partial_transport: Optional[str] = None + inventory_id: Optional[str] = None + load_id: Optional[str] = None + model_format: ModelFormat = "unknown" + runtime: ModelRuntime = "unknown" + format_variant: Optional[str] = None + capabilities: LocalModelCapabilities = Field(default_factory = LocalModelCapabilities) + + +class CachedGgufRepo(CachedRepoBase): + model_format: ModelFormat = "gguf" + + +class CachedGgufResponse(BaseModel): + cached: List[CachedGgufRepo] = Field(default_factory = list) + + +class CachedModelRepo(CachedRepoBase): + quant_method: Optional[str] = None + pipeline_tag: Optional[str] = None + library_name: Optional[str] = None + tags: Optional[List[str]] = None + + +class CachedModelsResponse(BaseModel): + cached: List[CachedModelRepo] = Field(default_factory = list) + + +class AddScanFolderRequest(BaseModel): + """Request body for adding a custom scan folder.""" + + path: str = Field( + ..., + description = "Absolute or relative folder path, or a model weight file path", + ) + + +class ScanFolderInfo(BaseModel): + """A registered custom model scan folder.""" + + id: int = Field(..., description = "Database row ID") + path: str = Field(..., description = "Normalized absolute path") + created_at: str = Field(..., description = "ISO 8601 creation timestamp") + + +class ScanFoldersResponse(BaseModel): + folders: List[ScanFolderInfo] = Field(default_factory = list) + + +class RemoveScanFolderResponse(BaseModel): + ok: bool + + +class RecommendedFoldersResponse(BaseModel): + folders: List[str] = Field(default_factory = list) + + +class DeleteCachedModelResponse(BaseModel): + status: str + repo_id: str + variant: Optional[str] = None + + +class BrowseEntry(BaseModel): + """A directory entry surfaced by the folder browser.""" + + name: str = Field(..., description = "Entry name (basename, not full path)") + has_models: bool = Field( + False, + description = ( + "Hint that the directory likely contains models " + "(*.gguf, *.safetensors, config.json, or HF-style " + "`models--*` subfolders). Used by the UI to highlight " + "promising candidates; the scanner itself is authoritative." + ), + ) + hidden: bool = Field( + False, + description = "Name starts with a dot (e.g. `.cache`)", + ) + + +class BrowseFoldersResponse(BaseModel): + """Response schema for the folder browser endpoint.""" + + current: str = Field(..., description = "Absolute path of the directory just listed") + parent: Optional[str] = Field( + None, + description = ( + "Parent directory of `current`, or null if `current` is the " + "filesystem root. The frontend uses this to render an `Up` row." + ), + ) + entries: List[BrowseEntry] = Field( + default_factory = list, + description = ( + "Subdirectories of `current`. Sorted with model-bearing " + "directories first, then alphabetically case-insensitive; " + "hidden entries come last within each group." + ), + ) + suggestions: List[str] = Field( + default_factory = list, + description = ( + "Handy starting points (home, HF cache, already-registered " + "scan folders). Rendered as quick-pick chips above the list." + ), + ) + truncated: bool = Field( + False, + description = ( + "True when the listing was capped because the directory had " + "more subfolders than the server is willing to enumerate in " + "one request. The UI should show a hint telling the user to " + "narrow their path." + ), + ) + model_files_here: int = Field( + 0, + description = ( + "Count of GGUF/safetensors files immediately inside " + "``current``. Used by the UI to surface a hint on leaf " + "model directories (which otherwise look `empty` because " + "they contain only files, no subdirectories)." + ), + ) diff --git a/studio/backend/hub/services/__init__.py b/studio/backend/hub/services/__init__.py new file mode 100644 index 0000000000..e86fcb6f46 --- /dev/null +++ b/studio/backend/hub/services/__init__.py @@ -0,0 +1,29 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Shared helpers for the Hub service layer.""" + +from __future__ import annotations + +from typing import Iterable + +from fastapi import HTTPException + +from hub.utils.hf_cache_state import resolve_destructive_case_matches + + +def resolve_destructive_repo_ids(repo_id: str, candidates: Iterable[str], *, noun: str) -> set[str]: + """Cache-dir repo ids a destructive op on *repo_id* may target. + + Refuses with 409 on ambiguous case-only matches so a delete never removes + the wrong casing. *noun* is the plural shown to the user.""" + resolved = resolve_destructive_case_matches(repo_id, candidates) + if resolved is None: + raise HTTPException( + status_code = 409, + detail = ( + f"Multiple cached {noun} differ only by case. " + "Delete the exact repo casing from On Device." + ), + ) + return resolved diff --git a/studio/backend/hub/services/datasets/__init__.py b/studio/backend/hub/services/datasets/__init__.py new file mode 100644 index 0000000000..c917f7bc02 --- /dev/null +++ b/studio/backend/hub/services/datasets/__init__.py @@ -0,0 +1,4 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Dataset services for Hub routes.""" diff --git a/studio/backend/hub/services/datasets/cache_inventory.py b/studio/backend/hub/services/datasets/cache_inventory.py new file mode 100644 index 0000000000..a180c9df58 --- /dev/null +++ b/studio/backend/hub/services/datasets/cache_inventory.py @@ -0,0 +1,474 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Cached dataset inventory and deletion services.""" + +from __future__ import annotations + +import asyncio +import os +from pathlib import Path +from typing import Optional + +from fastapi import HTTPException +from loggers import get_logger + +from hub.services import resolve_destructive_repo_ids +from hub.services.datasets import downloads +from hub.utils import download_manifest +from hub.utils import inventory_scan as hf_cache_scan +from hub.utils.hf_cache_state import ( + purge_partial_repo, + purge_repo_cache_dirs, + resolve_destructive_case_matches, +) +from hub.utils.paths import ( + hf_default_cache_dir, + is_valid_repo_id as _is_valid_repo_id, + legacy_hf_cache_dir, + resolve_cached_repo_id_case, +) + +logger = get_logger(__name__) + + +def _collect_hf_cache_scans() -> tuple[list, set[str]]: + scans = hf_cache_scan.all_hf_cache_scans() + seen_roots = { + str(cache_dir) + for cache_dir in (getattr(scan, "cache_dir", None) for scan in scans) + if cache_dir is not None + } + return scans, seen_roots + + +def _hf_hub_cache_roots() -> list[Path]: + roots: list[Path] = [] + seen: set[str] = set() + + def _add(path: Optional[Path]) -> None: + if path is None or not path.is_dir(): + return + try: + resolved = str(path.resolve()) + except OSError: + return + if resolved in seen: + return + seen.add(resolved) + roots.append(path) + + try: + from huggingface_hub.constants import HF_HUB_CACHE + _add(Path(HF_HUB_CACHE)) + except Exception: + pass + + hf_hub_cache = os.environ.get("HF_HUB_CACHE") + if hf_hub_cache: + _add(Path(hf_hub_cache).expanduser()) + + hf_home = os.environ.get("HF_HOME") + if hf_home: + _add(Path(hf_home).expanduser() / "hub") + + _add(legacy_hf_cache_dir()) + _add(hf_default_cache_dir()) + return roots + + +def _repo_id_from_hub_dataset_dir(name: str) -> str | None: + if not name.startswith("datasets--"): + return None + encoded = name.removeprefix("datasets--") + owner, sep, repo = encoded.partition("--") + if not sep or not owner or not repo: + return None + repo_id = f"{owner}/{repo}" + return repo_id if _is_valid_repo_id(repo_id) else None + + +def _directory_size(path: Path) -> int: + total = 0 + try: + for entry in path.rglob("*"): + try: + if entry.is_file() and not entry.is_symlink(): + total += entry.stat().st_size + except OSError: + continue + except OSError: + return 0 + return total + + +def _prefer_dataset_cache_row(candidate: dict, existing: Optional[dict]) -> bool: + if existing is None: + return True + candidate_partial = bool(candidate.get("partial")) + existing_partial = bool(existing.get("partial")) + if candidate_partial != existing_partial: + return not candidate_partial + return int(candidate.get("size_bytes") or 0) > int(existing.get("size_bytes") or 0) + + +def _hub_dataset_snapshot_count(path: Path) -> int: + snapshots = path / "snapshots" + try: + return sum(1 for entry in snapshots.iterdir() if entry.is_dir()) + except OSError: + return 0 + + +def _scan_hub_dataset_cache_dirs() -> list[dict]: + """Fallback scanner: ``scan_cache_dir()`` skips repos when one cache entry is partially corrupt, so this keeps On Device matching disk.""" + seen_lower: dict[str, dict] = {} + for root in _hf_hub_cache_roots(): + try: + entries = [entry for entry in root.iterdir() if entry.is_dir()] + except OSError: + continue + for entry in entries: + repo_id = _repo_id_from_hub_dataset_dir(entry.name) + if repo_id is None: + continue + size_bytes = _directory_size(entry / "blobs") + if size_bytes <= 0: + size_bytes = _directory_size(entry) + if size_bytes <= 0: + continue + key = repo_id.lower() + existing = seen_lower.get(key) + snapshot_partial = _hub_dataset_snapshot_count( + entry + ) == 0 or hf_cache_scan.is_snapshot_partial("dataset", repo_id, entry) + row = { + "repo_id": repo_id, + "size_bytes": size_bytes, + "cache_path": str(entry.resolve()), + # snapshot_count == 0 catches blobs-but-no-snapshot; + # is_snapshot_partial adds active-row state checks. + "partial": snapshot_partial, + "partial_transport": ( + hf_cache_scan.partial_transport_for( + "dataset", + repo_id, + repo_cache_dir = entry, + ) + if snapshot_partial + else None + ), + } + if _prefer_dataset_cache_row(row, existing): + seen_lower[key] = row + return sorted(seen_lower.values(), key = lambda c: c["repo_id"]) + + +def _hf_datasets_cache_roots() -> list[Path]: + roots: list[Path] = [] + seen: set[str] = set() + + def _add(path: Optional[Path]) -> None: + if path is None or not path.is_dir(): + return + try: + resolved = str(path.resolve()) + except OSError: + return + if resolved in seen: + return + seen.add(resolved) + roots.append(path) + + env_cache = os.environ.get("HF_DATASETS_CACHE") + if env_cache: + _add(Path(env_cache).expanduser()) + + try: + from datasets import config as datasets_config + _add(Path(datasets_config.HF_DATASETS_CACHE)) + except Exception: + pass + + hf_home = os.environ.get("HF_HOME") + if hf_home: + _add(Path(hf_home).expanduser() / "datasets") + + xdg_cache = Path(os.environ.get("XDG_CACHE_HOME", Path.home() / ".cache")).expanduser() + _add(xdg_cache / "huggingface" / "datasets") + return roots + + +def _repo_id_from_datasets_cache_dir(name: str) -> str | None: + if "___" not in name: + return None + owner, repo = name.split("___", 1) + repo_id = f"{owner}/{repo}" + return repo_id if _is_valid_repo_id(repo_id) else None + + +def _processed_dataset_cache_size(path: Path) -> int: + total = 0 + try: + for entry in path.rglob("*"): + try: + if entry.is_file(): + total += entry.stat().st_size + except OSError: + continue + except OSError: + return 0 + return total + + +def _looks_like_processed_dataset_cache(path: Path) -> bool: + try: + for entry in path.rglob("*"): + if not entry.is_file(): + continue + if entry.name in {"dataset_info.json", "state.json"}: + return True + if entry.suffix == ".arrow": + return True + except OSError: + return False + return False + + +def _scan_processed_dataset_caches() -> list[dict]: + """`load_dataset()` stores processed Arrow caches separately from the Hub snapshot cache, so they're usable on-device but invisible to `scan_cache_dir()`.""" + seen_lower: dict[str, dict] = {} + for root in _hf_datasets_cache_roots(): + try: + entries = [entry for entry in root.iterdir() if entry.is_dir()] + except OSError: + continue + for entry in entries: + repo_id = _repo_id_from_datasets_cache_dir(entry.name) + if repo_id is None: + continue + if not _looks_like_processed_dataset_cache(entry): + continue + size_bytes = _processed_dataset_cache_size(entry) + if size_bytes <= 0: + continue + key = repo_id.lower() + existing = seen_lower.get(key) + if existing is None or size_bytes > existing["size_bytes"]: + seen_lower[key] = { + "repo_id": repo_id, + "size_bytes": size_bytes, + "cache_path": str(entry.resolve()), + "processed_cache": True, + "partial": False, + } + return sorted(seen_lower.values(), key = lambda c: c["repo_id"]) + + +def _scan_hf_dataset_caches() -> list[dict]: + scans, seen_roots = _collect_hf_cache_scans() + + seen_lower: dict[str, dict] = {} + inspected = 0 + for hf_cache in scans: + for repo_info in hf_cache.repos: + inspected += 1 + try: + # str(...) guards against the library switching repo_type to an Enum. + if str(repo_info.repo_type) != "dataset": + continue + total_size = int(getattr(repo_info, "size_on_disk", 0) or 0) + if total_size == 0: + unique_blobs: dict[str, int] = {} + for rev in repo_info.revisions: + rev_id = getattr(rev, "commit_hash", None) or str(id(rev)) + for f in rev.files: + blob_path = getattr(f, "blob_path", None) + key = str(blob_path) if blob_path else f"{rev_id}:{f.file_name}" + unique_blobs[key] = int(f.size_on_disk or 0) + total_size = sum(unique_blobs.values()) + key = repo_info.repo_id.lower() + existing = seen_lower.get(key) + cache_dir = Path(repo_info.repo_path) + snapshot_partial = hf_cache_scan.is_snapshot_partial( + "dataset", + repo_info.repo_id, + cache_dir, + ) + row = { + "repo_id": repo_info.repo_id, + "size_bytes": total_size, + "cache_path": str(repo_info.repo_path), + "partial": snapshot_partial, + "partial_transport": ( + hf_cache_scan.partial_transport_for( + "dataset", + repo_info.repo_id, + repo_cache_dir = cache_dir, + ) + if snapshot_partial + else None + ), + } + if _prefer_dataset_cache_row(row, existing): + seen_lower[key] = row + except Exception as exc: + label = getattr(repo_info, "repo_id", "") + logger.warning("Skipping cached dataset repo %s: %s", label, exc) + for row in _scan_hub_dataset_cache_dirs(): + key = row["repo_id"].lower() + existing = seen_lower.get(key) + if _prefer_dataset_cache_row(row, existing): + seen_lower[key] = row + elif existing is not None and bool(existing.get("partial")) == bool(row.get("partial")): + existing["size_bytes"] = max(existing["size_bytes"], row["size_bytes"]) + existing["cache_path"] = existing.get("cache_path") or row.get("cache_path") + if ( + existing.get("partial") + and not existing.get("partial_transport") + and row.get("partial_transport") + ): + existing["partial_transport"] = row["partial_transport"] + for row in _scan_processed_dataset_caches(): + key = row["repo_id"].lower() + existing = seen_lower.get(key) + if existing is None or (bool(existing.get("partial")) and not bool(row.get("partial"))): + seen_lower[key] = row + else: + existing["size_bytes"] = max(existing["size_bytes"], row["size_bytes"]) + # Keep the processed-cache marker when a repo is both snapshot and + # processed Arrow cache; merging by size alone dropped it. + if row.get("processed_cache"): + existing["processed_cache"] = True + logger.info( + "Cached dataset scan: roots=%d inspected=%d returned=%d", + len(seen_roots) or len(scans), + inspected, + len(seen_lower), + ) + return sorted(seen_lower.values(), key = lambda c: c["repo_id"]) + + +async def list_cached_datasets_response() -> dict: + """List dataset repos already downloaded into the HF cache.""" + try: + return {"cached": await asyncio.to_thread(_scan_hf_dataset_caches)} + except Exception as exc: + logger.error("Error listing cached datasets: %s", exc, exc_info = True) + raise HTTPException( + status_code = 500, + detail = "Failed to read the local dataset cache.", + ) from exc + + +async def delete_cached_dataset_response(repo_id: str) -> dict: + """Remove a cached dataset repo from the HF cache.""" + if not _is_valid_repo_id(repo_id): + raise HTTPException(status_code = 400, detail = "Invalid repo_id format") + + repo_key = await asyncio.to_thread(resolve_cached_repo_id_case, repo_id, repo_type = "dataset") + if not downloads.registry.begin_delete(repo_key): + raise HTTPException( + status_code = 400, + detail = "Cancel the active download before deleting.", + ) + try: + return await asyncio.to_thread(_delete_cached_dataset_blocking, repo_key) + finally: + downloads.registry.end_delete(repo_key) + hf_cache_scan.invalidate_hf_cache_scans() + + +def _delete_cached_dataset_blocking(repo_id: str) -> dict: + scans, _seen_roots = _collect_hf_cache_scans() + + candidate_entries = [] + for hf_cache in scans: + for repo_info in hf_cache.repos: + if str(repo_info.repo_type) != "dataset": + continue + if repo_info.repo_id.lower() == repo_id.lower(): + candidate_entries.append((hf_cache, repo_info)) + matched_repo_ids = resolve_destructive_repo_ids( + repo_id, + [str(repo_info.repo_id) for _hf_cache, repo_info in candidate_entries], + noun = "datasets", + ) + + deleted = False + failures: list[str] = [] + for hf_cache, repo_info in candidate_entries: + if str(repo_info.repo_id) not in matched_repo_ids: + continue + try: + strategy = hf_cache.delete_revisions(*(rev.commit_hash for rev in repo_info.revisions)) + strategy.execute() + deleted = True + except Exception as exc: + failures.append(str(exc)) + logger.error( + "Failed deleting cached dataset %s from %s: %s", + repo_id, + getattr(hf_cache, "cache_dir", ""), + exc, + exc_info = True, + ) + + processed_deleted, processed_failures = _delete_processed_dataset_cache(repo_id) + failures.extend(processed_failures) + if failures: + raise HTTPException( + status_code = 500, + detail = ( + f"Failed to delete dataset from {len(failures)} cache " + "location(s). Some files may remain." + ), + ) + + # ``scan_cache_dir()`` skips blob-only/corrupt repos the revision delete + # can't touch, yet the fallback scanner shows them; purge the whole dir. + cache_purged = purge_repo_cache_dirs("dataset", repo_id) + partial_purged = purge_partial_repo("dataset", repo_id) + state_purged = download_manifest.purge_all_state_for_repo("dataset", repo_id) > 0 + if not (deleted or processed_deleted or cache_purged or partial_purged or state_purged): + raise HTTPException(status_code = 404, detail = "Dataset not found in cache") + return {"status": "deleted", "repo_id": repo_id} + + +def _delete_processed_dataset_cache(repo_id: str) -> tuple[bool, list[str]]: + import shutil + + target = repo_id.replace("/", "___") + folded_target = target.lower() + deleted = False + failures: list[str] = [] + for root in _hf_datasets_cache_roots(): + try: + entries = [ + entry + for entry in root.iterdir() + if entry.is_dir() and entry.name.lower() == folded_target + ] + except OSError: + continue + matched_names = resolve_destructive_case_matches( + target, + (entry.name for entry in entries), + ) + if not matched_names: + continue + for entry in entries: + if entry.name not in matched_names: + continue + try: + shutil.rmtree(entry) + deleted = True + except Exception as exc: + failures.append(str(exc)) + logger.error( + "Failed deleting processed dataset cache %s: %s", + repo_id, + exc, + exc_info = True, + ) + return deleted, failures diff --git a/studio/backend/hub/services/datasets/downloads.py b/studio/backend/hub/services/datasets/downloads.py new file mode 100644 index 0000000000..ac90be8c6f --- /dev/null +++ b/studio/backend/hub/services/datasets/downloads.py @@ -0,0 +1,268 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Start, cancel, and report progress for dataset downloads.""" + +from __future__ import annotations + +import asyncio +import threading +import time +from collections import OrderedDict +from typing import Optional + +from fastapi import HTTPException +from loggers import get_logger + +from hub.schemas.downloads import ( + ActiveDownloadsResponse, + CancelDatasetDownloadRequest, + DatasetDownloadJobStatus, + DownloadDatasetRequest, +) +from hub.services import snapshot_progress +from hub.services import download_lifecycle +from hub.utils import download_manifest +from hub.utils import download_registry +from hub.utils import inventory_scan as hf_cache_scan +from hub.utils.hf_cache_state import has_active_incomplete_blobs +from hub.utils.paths import ( + is_valid_repo_id as _is_valid_repo_id, + resolve_cached_repo_id_case, +) +from hub.utils.snapshot_filters import ( + blob_hashes_for_siblings, + total_size_for_siblings, +) + +logger = get_logger(__name__) + +_dataset_size_cache: "OrderedDict[str, tuple[int, frozenset[str], bool, str, float]]" = ( + OrderedDict() +) +_dataset_size_neg_cache: "OrderedDict[tuple[str, str], float]" = OrderedDict() +_DATASET_SIZE_CACHE_MAX = 256 +_DATASET_SIZE_POS_TTL = 60.0 +_DATASET_SIZE_NEG_TTL = 60.0 +_DATASET_SIZE_TIMEOUT_SECONDS = 5.0 +_dataset_size_cache_lock = threading.Lock() + +_registry = download_registry.get_datasets_registry() + + +def _download_job_key(repo_id: str) -> str: + return download_registry.normalize_repo_key(repo_id) + + +def get_dataset_snapshot_metadata_cached( + repo_id: str, hf_token: Optional[str] = None +) -> tuple[int, frozenset[str]]: + """Raw snapshot size + expected blob hashes for a dataset repo. + + The dataset worker downloads every sibling, so the denominator is the full + sibling-size sum and the hashes cover every file. Consumed by the shared + ``snapshot_progress`` accounting.""" + token_fp = hf_cache_scan.token_fingerprint(hf_token) + cache_key = (repo_id, token_fp) + with _dataset_size_cache_lock: + cached = _dataset_size_cache.get(repo_id) + if cached is not None: + size, hashes, restricted, cached_fp, ts = cached + if (time.monotonic() - ts) >= _DATASET_SIZE_POS_TTL: + del _dataset_size_cache[repo_id] + # A gated/private repo's metadata is only served back to the token + # that fetched it; another token may have no access at all. + elif not restricted or cached_fp == token_fp: + _dataset_size_cache.move_to_end(repo_id) + return size, hashes + neg_ts = _dataset_size_neg_cache.get(cache_key) + if neg_ts is not None and (time.monotonic() - neg_ts) < _DATASET_SIZE_NEG_TTL: + return 0, frozenset() + try: + from huggingface_hub import HfApi + + info = HfApi(token = hf_token).dataset_info( + repo_id, + files_metadata = True, + timeout = _DATASET_SIZE_TIMEOUT_SECONDS, + ) + total = total_size_for_siblings(info.siblings) + hashes = blob_hashes_for_siblings(info.siblings) + restricted = bool(getattr(info, "private", False) or getattr(info, "gated", False)) + except Exception: + with _dataset_size_cache_lock: + _dataset_size_neg_cache[cache_key] = time.monotonic() + _dataset_size_neg_cache.move_to_end(cache_key) + while len(_dataset_size_neg_cache) > _DATASET_SIZE_CACHE_MAX: + _dataset_size_neg_cache.popitem(last = False) + return 0, frozenset() + with _dataset_size_cache_lock: + _dataset_size_cache[repo_id] = ( + total, + hashes, + restricted, + token_fp, + time.monotonic(), + ) + _dataset_size_cache.move_to_end(repo_id) + _dataset_size_neg_cache.pop(cache_key, None) + while len(_dataset_size_cache) > _DATASET_SIZE_CACHE_MAX: + _dataset_size_cache.popitem(last = False) + return total, hashes + + +async def get_dataset_download_progress_response( + repo_id: str, + expected_bytes: int = 0, + hf_token: Optional[str] = None, +) -> dict: + """Return download progress for a HuggingFace dataset repo. + + Scans the ``datasets--owner--name`` cache dir and shares the blob accounting + with the model path via ``snapshot_progress``. Returns ``cache_path`` for the + UI.""" + return await snapshot_progress.snapshot_progress_response( + repo_type = "dataset", + repo_id = repo_id, + job_key = _download_job_key(repo_id), + expected_bytes = expected_bytes, + hf_token = hf_token, + registry = _registry, + metadata_resolver = get_dataset_snapshot_metadata_cached, + ) + + +def _dataset_status(key: str, *, repo_id: Optional[str] = None) -> DatasetDownloadJobStatus: + state, error, generation = download_lifecycle.idle_status( + _registry, + key, + repo_type = "dataset", + repo_id = repo_id, + variant = None, + ) + return DatasetDownloadJobStatus(state = state, error = error, generation = generation) + + +async def download_dataset_response( + body: DownloadDatasetRequest, hf_token: Optional[str] = None +) -> dict: + """Start a background download for a HuggingFace dataset.""" + repo_id = body.repo_id.strip() + if not _is_valid_repo_id(repo_id): + raise HTTPException( + status_code = 400, + detail = f"Invalid repo_id: {repo_id!r}", + ) + # Canonicalize so two different-cased paste-ins share one job + cache dir. + repo_id = await asyncio.to_thread(resolve_cached_repo_id_case, repo_id, repo_type = "dataset") + key = _download_job_key(repo_id) + + transport = download_lifecycle.resolve_transport(body.use_xet) + + claimed, claim_state = _registry.claim( + key, + transport, + repo_type = "dataset", + repo_id = repo_id, + ) + generation = _registry.current_generation(key) + if not claimed: + # Pollable when rejected by this repo's own in-flight job; an + # in-progress delete leaves no job, so flag it via ``adoptable``. + return { + "repo_id": repo_id, + "state": claim_state, + "accepted": _registry.adoptable(key), + "generation": generation, + } + download_manifest.clear_cancel_marker("dataset", repo_id, None) + + state = download_lifecycle.launch_worker( + _registry, + key, + spawn = lambda: download_lifecycle.spawn_worker( + ["--repo-id", repo_id, "--dataset"], + hf_token, + use_xet = body.use_xet, + ), + hf_token = hf_token, + label = repo_id, + log_prefix = "Dataset download", + logger = logger, + repo_type = "dataset", + repo_id = repo_id, + transport = transport, + watch_name = f"hf-dataset-download-watch-{repo_id}", + ) + + return { + "repo_id": repo_id, + "state": state, + "accepted": True, + "generation": generation, + } + + +async def cancel_dataset_download_response(body: CancelDatasetDownloadRequest) -> dict: + """Cancel an in-flight dataset download (SIGKILL; HF cache resumes on next download).""" + repo_id = body.repo_id.strip() + if not _is_valid_repo_id(repo_id): + raise HTTPException( + status_code = 400, + detail = f"Invalid repo_id: {repo_id!r}", + ) + repo_id = await asyncio.to_thread(resolve_cached_repo_id_case, repo_id, repo_type = "dataset") + key = _download_job_key(repo_id) + + state = download_lifecycle.cancel_worker( + _registry, + key, + generation = body.generation, + label = f"dataset {repo_id}", + logger = logger, + ) + return {"repo_id": repo_id, "state": state} + + +async def get_dataset_download_status_response(repo_id: str) -> DatasetDownloadJobStatus: + """Return the latest state of a background dataset download job.""" + repo_id = repo_id.strip() + if not _is_valid_repo_id(repo_id): + return DatasetDownloadJobStatus(state = "idle") + repo_id = await asyncio.to_thread(resolve_cached_repo_id_case, repo_id, repo_type = "dataset") + return _dataset_status(_download_job_key(repo_id), repo_id = repo_id) + + +async def get_active_dataset_downloads_response(repo_id: str = "") -> ActiveDownloadsResponse: + repo_id = repo_id.strip() + if repo_id and not _is_valid_repo_id(repo_id): + return ActiveDownloadsResponse(downloads = []) + canonical_repo_id = ( + await asyncio.to_thread(resolve_cached_repo_id_case, repo_id, repo_type = "dataset") + if repo_id + else None + ) + return ActiveDownloadsResponse( + downloads = download_lifecycle.active_download_refs( + _registry, + canonical_repo_id, + with_variant = False, + ) + ) + + +async def get_dataset_transport_status_response(repo_id: str) -> dict: + """Last transport used, whether partial blobs exist, and whether they + support byte-level resume. XET partials show via ``has_partial`` but are not + byte-level resumable (see ``models.get_model_transport_status``).""" + repo_id = repo_id.strip() + if not _is_valid_repo_id(repo_id): + return {"has_partial": False, "last_transport": None, "resumable": False} + return { + "has_partial": has_active_incomplete_blobs("dataset", repo_id), + "last_transport": download_registry.read_active_transport_marker("dataset", repo_id), + "resumable": download_registry.is_resumable_partial("dataset", repo_id), + } + + +registry = _registry diff --git a/studio/backend/hub/services/datasets/formatting.py b/studio/backend/hub/services/datasets/formatting.py new file mode 100644 index 0000000000..8b0ff39f63 --- /dev/null +++ b/studio/backend/hub/services/datasets/formatting.py @@ -0,0 +1,527 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Dataset preview, format-check, and mapping-assist services.""" + +from __future__ import annotations + +import base64 +import errno +import io +import re +from pathlib import Path +from typing import Optional + +from fastapi import HTTPException +from loggers import get_logger + +from hub.schemas.datasets import ( + AiAssistMappingRequest, + AiAssistMappingResponse, + CheckFormatRequest, + CheckFormatResponse, +) +from hub.services.datasets.local import ( + DATA_EXTS, + _TABULAR_EXTS, + _load_local_preview_slice, + _stream_file_preview_slice, +) +from hub.utils.dataset_cache import ( + cached_dataset_candidates as _shared_cached_dataset_candidates, + latest_cached_dataset_snapshot as _shared_latest_cached_dataset_snapshot, + split_label_matches as _split_label_matches, +) +from hub.utils import download_registry +from hub.utils.dataset_format import check_dataset_format, format_dataset_preview +from hub.utils.hf_errors import hf_error_status +from hub.utils.paths import ( + is_valid_repo_id as _is_valid_repo_id, + resolve_dataset_path, +) + +logger = get_logger(__name__) + +_BINARY_IMAGE_PREVIEW_MAX_BYTES = 10 * 1024 * 1024 +_IMAGE_PREVIEW_MAX_PIXELS = 16_000_000 +_IMAGE_PREVIEW_THUMBNAIL_SIZE = (512, 512) + + +def _image_pixel_count(image) -> int: + width = max(int(getattr(image, "width", 0) or 0), 0) + height = max(int(getattr(image, "height", 0) or 0), 0) + return width * height + + +def _pil_image_has_transparency(image) -> bool: + if "A" in image.getbands(): + extrema = image.getchannel("A").getextrema() + return bool(extrema and extrema[0] < 255) + if image.mode == "P": + transparency = image.info.get("transparency") + if transparency is None: + return False + if isinstance(transparency, bytes): + return any(alpha < 255 for alpha in transparency) + return True + return False + + +def _serialize_pil_image(image): + pixel_count = _image_pixel_count(image) + if pixel_count > _IMAGE_PREVIEW_MAX_PIXELS: + return ( + f"" + ) + + preview = image.copy() + preview.thumbnail(_IMAGE_PREVIEW_THUMBNAIL_SIZE) + buffer = io.BytesIO() + if _pil_image_has_transparency(preview): + preview.save(buffer, format = "PNG") + mime = "image/png" + else: + preview.convert("RGB").save(buffer, format = "JPEG", quality = 85) + mime = "image/jpeg" + return { + "type": "image", + "mime": mime, + "width": preview.width, + "height": preview.height, + "data": base64.b64encode(buffer.getvalue()).decode("ascii"), + } + + +def _serialize_binary_value(data): + if len(data) > _BINARY_IMAGE_PREVIEW_MAX_BYTES: + return ( + f"" + ) + + try: + from PIL import Image as PILImageModule + with PILImageModule.open(io.BytesIO(data)) as image: + return _serialize_pil_image(image) + except Exception: + return f"" + + +def _serialize_preview_value(value): + if value is None or isinstance(value, (str, int, float, bool)): + return value + + if isinstance(value, (bytes, bytearray, memoryview)): + return _serialize_binary_value(value) + + try: + from PIL.Image import Image as PILImage + if isinstance(value, PILImage): + return _serialize_pil_image(value) + except Exception: + pass + + if isinstance(value, dict): + # Undecoded HF Image/Audio cells are {"bytes": b"...", "path": ...}. + raw = value.get("bytes") + if isinstance(raw, (bytes, bytearray, memoryview)) and not ( + value.keys() - {"bytes", "path"} + ): + return _serialize_binary_value(raw) + return {str(key): _serialize_preview_value(item) for key, item in value.items()} + + if isinstance(value, (list, tuple)): + return [_serialize_preview_value(item) for item in value] + + return str(value) + + +def _serialize_preview_rows(rows): + return [ + {str(key): _serialize_preview_value(value) for key, value in dict(row).items()} + for row in rows + ] + + +def _latest_cached_dataset_snapshot( + repo_id: str, local_path: Optional[str] = None +) -> Optional[Path]: + return _shared_latest_cached_dataset_snapshot(repo_id, local_path) + + +def _cached_dataset_candidates( + snapshot: Path, *, subset: Optional[str], train_split: str +) -> list[Path]: + return _shared_cached_dataset_candidates( + snapshot, + subset = subset, + train_split = train_split, + extensions = DATA_EXTS, + preferred_extensions = _TABULAR_EXTS, + ) + + +def _repo_file_label_tokens(path: str) -> set[str]: + return {token for token in re.split(r"[^a-z0-9]+", path.lower()) if token} + + +def _repo_file_matches_label(path: str, label: str) -> bool: + return label.strip().lower() in _repo_file_label_tokens(path) + + +def _repo_file_matches_split(path: str, split: str) -> bool: + return _split_label_matches(path, split) + + +def _select_tier1_repo_file( + files: list[str], *, subset: Optional[str], train_split: str +) -> Optional[str]: + data_files = sorted(f for f in files if any(f.lower().endswith(ext) for ext in DATA_EXTS)) + if not data_files: + return None + tabular_files = [f for f in data_files if any(f.lower().endswith(ext) for ext in _TABULAR_EXTS)] + candidates = tabular_files or data_files + if subset: + candidates = [f for f in candidates if _repo_file_matches_label(f, subset)] + if not candidates: + return None + candidates = [f for f in candidates if _repo_file_matches_split(f, train_split)] + return candidates[0] if candidates else None + + +def _load_cached_hf_preview_slice(request: CheckFormatRequest, preview_size: int): + if not _is_valid_repo_id(request.dataset_name): + return None + snapshot = _latest_cached_dataset_snapshot( + request.dataset_name, + request.local_path, + ) + if snapshot is None: + return None + train_split = request.train_split or "train" + for candidate in _cached_dataset_candidates( + snapshot, + subset = request.subset, + train_split = train_split, + ): + try: + preview = _stream_file_preview_slice(candidate, preview_size) + except Exception as exc: + logger.debug("Cached dataset preview failed for %s: %s", candidate, exc) + continue + if preview is not None: + return preview + return None + + +def _load_processed_hf_preview_slice( + request: CheckFormatRequest, + preview_size: int, + hf_token: Optional[str] = None, +): + if not _is_valid_repo_id(request.dataset_name): + return None + try: + from datasets import DownloadConfig, load_dataset + except Exception: + return None + + load_kwargs = { + "path": request.dataset_name, + "split": request.train_split or "train", + "download_config": DownloadConfig(local_files_only = True), + } + if request.subset: + load_kwargs["name"] = request.subset + if hf_token: + load_kwargs["token"] = hf_token + + dataset = load_dataset(**load_kwargs) + total_rows = len(dataset) + preview_slice = dataset.select(range(min(preview_size, total_rows))) + return preview_slice, total_rows + + +def _load_any_cached_hf_preview_slice( + request: CheckFormatRequest, + preview_size: int, + hf_token: Optional[str] = None, +): + cached_preview = _load_cached_hf_preview_slice(request, preview_size) + if cached_preview is not None: + return cached_preview + try: + return _load_processed_hf_preview_slice(request, preview_size, hf_token) + except Exception as exc: + logger.debug( + "Processed dataset cache preview failed for %s: %s", + request.dataset_name, + exc, + ) + return None + + +def check_format_response( + request: CheckFormatRequest, hf_token: Optional[str] = None +) -> CheckFormatResponse: + """ + Check if a dataset requires manual column mapping. + + HF datasets: tier 1 loads a single requested split/subset file (avoids + resolving thousands of files); tier 2 falls back to full streaming. Local + files load directly. Plain `def` so FastAPI runs the blocking IO in a + thread-pool. + """ + try: + from itertools import islice + + PREVIEW_SIZE = 10 + + logger.info(f"Checking format for dataset: {request.dataset_name}") + + try: + dataset_path = resolve_dataset_path(request.dataset_name) + except ValueError as e: + # Malformed path (null bytes, '..', outside roots) is a client error: + # surface 400 rather than the generic 500 below. + raise HTTPException(status_code = 400, detail = str(e)) from e + total_rows = None + + if dataset_path.exists(): + train_split = request.train_split or "train" + preview_slice, total_rows = _load_local_preview_slice( + dataset_path = dataset_path, + train_split = train_split, + preview_size = PREVIEW_SIZE, + ) + else: + from datasets import Dataset, load_dataset + + # Tier 1: list_repo_files → load only the first data file + cached_preview = ( + _load_any_cached_hf_preview_slice(request, PREVIEW_SIZE, hf_token) + if request.prefer_local_cache + else None + ) + if cached_preview is not None: + preview_slice, total_rows = cached_preview + elif request.prefer_local_cache: + raise HTTPException( + status_code = 404, + detail = "Dataset is not available in the local cache.", + ) + else: + preview_slice = None + + try: + from huggingface_hub import HfApi + + api = HfApi() + repo_files = api.list_repo_files( + request.dataset_name, + repo_type = "dataset", + token = hf_token or None, + ) + train_split = request.train_split or "train" + first_file = _select_tier1_repo_file( + repo_files, + subset = request.subset, + train_split = train_split, + ) + if first_file: + logger.info(f"Tier 1: loading single file {first_file}") + load_kwargs = { + "path": request.dataset_name, + "data_files": {train_split: [first_file]}, + "split": train_split, + "streaming": True, + } + if hf_token: + load_kwargs["token"] = hf_token + + streamed_ds = load_dataset(**load_kwargs) + rows = list(islice(streamed_ds, PREVIEW_SIZE)) + if rows: + preview_slice = Dataset.from_list(rows) + except Exception as e: + logger.warning( + "Tier 1 (single-file) failed: %s", + download_registry.scrub_secrets(str(e), hf_token = hf_token), + ) + + if preview_slice is None: + # Tier 2: full streaming (resolves all files — slow for large repos) + logger.info("Tier 2: falling back to full streaming load_dataset") + try: + load_kwargs = { + "path": request.dataset_name, + "split": request.train_split or "train", + "streaming": True, + } + if request.subset: + load_kwargs["name"] = request.subset + if hf_token: + load_kwargs["token"] = hf_token + + streamed_ds = load_dataset(**load_kwargs) + + rows = list(islice(streamed_ds, PREVIEW_SIZE)) + if not rows: + raise HTTPException( + status_code = 400, + detail = "Dataset appears to be empty or could not be streamed", + ) + + preview_slice = Dataset.from_list(rows) + total_rows = None + except Exception: + cached_preview = _load_any_cached_hf_preview_slice( + request, + PREVIEW_SIZE, + hf_token, + ) + if cached_preview is None: + raise + preview_slice, total_rows = cached_preview + + result = check_dataset_format(preview_slice, is_vlm = request.is_vlm) + + logger.info( + f"Format check result: requires_mapping={result['requires_manual_mapping']}, format={result['detected_format']}, is_image={result.get('is_image', False)}" + ) + + preview_samples = None + if not result["requires_manual_mapping"]: + if result.get("suggested_mapping"): + # Heuristic-detected: show raw data so columns match the response; + # column stripping happens at training time, not preview. + preview_samples = _serialize_preview_rows(preview_slice) + else: + try: + processed = format_dataset_preview(preview_slice) + preview_samples = _serialize_preview_rows(processed) + except Exception as e: + logger.warning(f"Processed preview generation failed (non-fatal): {e}") + preview_samples = _serialize_preview_rows(preview_slice) + else: + preview_samples = _serialize_preview_rows(preview_slice) + + # Collect warnings: from check_dataset_format + URL-based image detection + warning = result.get("warning") + image_col = result.get("detected_image_column") + if image_col and image_col in (result.get("columns") or []): + try: + sample_val = preview_slice[0][image_col] + if isinstance(sample_val, str) and sample_val.startswith(("http://", "https://")): + url_warning = ( + "This dataset contains image URLs instead of embedded images. " + "Images will be downloaded during training, which may be slow for large datasets." + ) + logger.info(f"URL-based image column detected: {image_col}") + warning = f"{warning} {url_warning}" if warning else url_warning + except Exception: + pass + + return CheckFormatResponse( + requires_manual_mapping = result["requires_manual_mapping"], + detected_format = result["detected_format"], + columns = result["columns"], + is_image = result.get("is_image", False), + is_audio = result.get("is_audio", False), + multimodal_columns = result.get("multimodal_columns"), + suggested_mapping = result.get("suggested_mapping"), + detected_image_column = result.get("detected_image_column"), + detected_audio_column = result.get("detected_audio_column"), + detected_text_column = result.get("detected_text_column"), + detected_speaker_column = result.get("detected_speaker_column"), + preview_samples = preview_samples, + total_rows = total_rows, + warning = warning, + ) + + except HTTPException: + raise + except Exception as e: + scrubbed = download_registry.scrub_secrets(str(e), hf_token = hf_token) + # Missing/gated/bad-token and malformed names are client errors, not 500s. + status = hf_error_status(e) + if ( + status is None + and isinstance(e, OSError) + and getattr(e, "errno", None) == errno.ENAMETOOLONG + ): + status, scrubbed = 400, "Invalid dataset name" + elif status is None and isinstance(e, FileNotFoundError): + # datasets raises DatasetNotFoundError (FileNotFoundError) for missing/gated. + status = 404 + elif status is None and isinstance(e, ValueError): + status = 400 + if status is not None: + raise HTTPException(status_code = status, detail = scrubbed) + logger.error("Error checking dataset format: %s", scrubbed) + raise HTTPException( + status_code = 500, + detail = "Failed to check dataset format: " + scrubbed, + ) + + +def ai_assist_mapping_response( + request: AiAssistMappingRequest, hf_token: Optional[str] = None +) -> AiAssistMappingResponse: + """ + Run the LLM-assisted dataset conversion advisor (user-triggered). + + Multi-pass analysis with a 7B helper model: classify dataset type, generate + a conversion strategy, then validate it. Falls back to simple column + classification if the advisor fails. + """ + try: + from hub.utils.llm_assist import llm_conversion_advisor + + truncated = [ + {col: str(s.get(col, ""))[:200] for col in request.columns} for s in request.samples[:5] + ] + + result = llm_conversion_advisor( + column_names = request.columns, + samples = truncated, + dataset_name = request.dataset_name, + hf_token = hf_token, + model_name = request.model_name, + model_type = request.model_type, + ) + + if result and result.get("success"): + return AiAssistMappingResponse( + success = True, + suggested_mapping = result.get("suggested_mapping"), + system_prompt = result.get("system_prompt"), + user_template = result.get("user_template"), + assistant_template = result.get("assistant_template"), + label_mapping = result.get("label_mapping"), + dataset_type = result.get("dataset_type"), + is_conversational = result.get("is_conversational"), + user_notification = result.get("user_notification"), + warning = result.get("warning"), + ) + + return AiAssistMappingResponse( + success = False, + warning = "AI could not determine column roles. Please assign them manually.", + ) + + except Exception as e: + scrubbed = download_registry.scrub_secrets(str(e), hf_token = hf_token) + status = hf_error_status(e) + if status is None and isinstance(e, FileNotFoundError): + status = 404 + elif status is None and isinstance(e, ValueError): + status = 400 + if status is not None: + raise HTTPException(status_code = status, detail = scrubbed) + logger.error("AI assist mapping failed: %s", scrubbed) + raise HTTPException( + status_code = 500, + detail = "AI assist failed: " + scrubbed, + ) diff --git a/studio/backend/hub/services/datasets/local.py b/studio/backend/hub/services/datasets/local.py new file mode 100644 index 0000000000..2d2c7c3a0d --- /dev/null +++ b/studio/backend/hub/services/datasets/local.py @@ -0,0 +1,330 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Local dataset upload and listing services.""" + +from __future__ import annotations + +import asyncio +import json +import uuid +from pathlib import Path + +from fastapi import HTTPException, UploadFile + +from hub.schemas.datasets import ( + LocalDatasetItem, + LocalDatasetsResponse, + UploadDatasetResponse, +) +from hub.utils.paths import dataset_uploads_root, ensure_dir, recipe_datasets_root + +# Tabular formats are preferred over archives for Tier 1 preview: archives +# (e.g. images.zip) load as ImageFolder with synthetic columns that don't +# match the real schema. +_TABULAR_EXTS = (".parquet", ".json", ".jsonl", ".csv", ".tsv", ".arrow") +_ARCHIVE_EXTS = (".tar", ".tar.gz", ".tgz", ".gz", ".zst", ".zip", ".txt") +DATA_EXTS = _TABULAR_EXTS + _ARCHIVE_EXTS +LOCAL_FILE_EXTS = (".json", ".jsonl", ".csv", ".parquet") +LOCAL_UPLOAD_EXTS = {".csv", ".json", ".jsonl", ".parquet"} +LOCAL_UPLOAD_CHUNK_BYTES = 1024 * 1024 +LOCAL_UPLOAD_MAX_BYTES = 500 * 1024 * 1024 +LOCAL_DATASETS_ROOT = recipe_datasets_root() +DATASET_UPLOAD_DIR = dataset_uploads_root() + + +def _safe_read_metadata(path: Path) -> dict | None: + try: + payload = json.loads(path.read_text(encoding = "utf-8")) + except (OSError, ValueError, TypeError): + return None + if not isinstance(payload, dict): + return None + return payload + + +def _safe_read_rows_from_metadata(payload: dict | None) -> int | None: + if not payload: + return None + for key in ("actual_num_records", "target_num_records"): + value = payload.get(key) + if isinstance(value, int): + return value + return None + + +def _safe_read_metadata_summary(payload: dict | None) -> dict | None: + if not payload: + return None + + actual_num_records = ( + payload.get("actual_num_records") + if isinstance(payload.get("actual_num_records"), int) + else None + ) + target_num_records = ( + payload.get("target_num_records") + if isinstance(payload.get("target_num_records"), int) + else actual_num_records + ) + + columns: list[str] | None = None + schema = payload.get("schema") + if isinstance(schema, dict): + columns = [str(key) for key in schema.keys()] + if not columns: + stats = payload.get("column_statistics") + if isinstance(stats, list): + derived = [ + str(item.get("column_name")) + for item in stats + if isinstance(item, dict) and item.get("column_name") + ] + columns = derived or None + + parquet_files_count = None + file_paths = payload.get("file_paths") + if isinstance(file_paths, dict): + parquet_files = file_paths.get("parquet-files") + if isinstance(parquet_files, list): + parquet_files_count = len(parquet_files) + + total_num_batches = ( + payload.get("total_num_batches") + if isinstance(payload.get("total_num_batches"), int) + else parquet_files_count + ) + num_completed_batches = ( + payload.get("num_completed_batches") + if isinstance(payload.get("num_completed_batches"), int) + else total_num_batches + ) + + return { + "actual_num_records": actual_num_records, + "target_num_records": target_num_records, + "total_num_batches": total_num_batches, + "num_completed_batches": num_completed_batches, + "columns": columns, + } + + +def _safe_mtime(path: Path) -> float | None: + try: + return path.stat().st_mtime + except OSError: + return None + + +def _display_uploaded_dataset_name(path: Path) -> str: + stem = path.stem + prefix, sep, rest = stem.partition("_") + if sep and len(prefix) == 32 and all(c in "0123456789abcdef" for c in prefix): + return f"{rest}{path.suffix}" + return path.name + + +def _build_recipe_dataset_items() -> list[LocalDatasetItem]: + if not LOCAL_DATASETS_ROOT.exists(): + return [] + + items: list[LocalDatasetItem] = [] + for entry in LOCAL_DATASETS_ROOT.iterdir(): + if not entry.is_dir() or not entry.name.startswith("recipe_"): + continue + parquet_dir = entry / "parquet-files" + if not parquet_dir.exists() or not any(parquet_dir.glob("*.parquet")): + continue + + rows = None + metadata_summary = None + metadata_path = entry / "metadata.json" + if metadata_path.exists(): + metadata_payload = _safe_read_metadata(metadata_path) + rows = _safe_read_rows_from_metadata(metadata_payload) + metadata_summary = _safe_read_metadata_summary(metadata_payload) + + items.append( + LocalDatasetItem( + id = entry.name, + label = entry.name, + path = str(parquet_dir.resolve()), + source = "recipe", + rows = rows, + updated_at = _safe_mtime(entry), + metadata = metadata_summary, + ) + ) + + return items + + +def _build_uploaded_dataset_items() -> list[LocalDatasetItem]: + if not DATASET_UPLOAD_DIR.exists(): + return [] + + items: list[LocalDatasetItem] = [] + for path in DATASET_UPLOAD_DIR.iterdir(): + if not path.is_file() or path.suffix.lower() not in LOCAL_UPLOAD_EXTS: + continue + try: + if path.stat().st_size == 0: + continue + except OSError: + continue + label = _display_uploaded_dataset_name(path) + items.append( + LocalDatasetItem( + id = path.name, + label = label, + path = str(path.resolve()), + source = "upload", + updated_at = _safe_mtime(path), + ) + ) + return items + + +def _build_local_dataset_items() -> list[LocalDatasetItem]: + items = _build_recipe_dataset_items() + _build_uploaded_dataset_items() + items.sort(key = lambda item: item.updated_at or 0, reverse = True) + return items + + +def _stream_file_preview_slice(path: Path, preview_size: int): + """Stream the first ``preview_size`` rows so a large file is never fully parsed into Arrow; returns ``(Dataset, None)`` or ``None`` if empty/unsupported.""" + from itertools import islice + + from datasets import Dataset, load_dataset + + name = path.name.lower() + if name.endswith((".json", ".jsonl")): + loader = "json" + elif name.endswith((".csv", ".tsv")): + loader = "csv" + elif name.endswith(".parquet"): + loader = "parquet" + elif name.endswith(".arrow"): + loader = "arrow" + elif name.endswith(".txt"): + loader = "text" + else: + return None + + streamed = load_dataset( + loader, + data_files = str(path), + split = "train", + streaming = True, + ) + rows = list(islice(streamed, preview_size)) + if not rows: + return None + return Dataset.from_list(rows), None + + +def _load_local_preview_slice(*, dataset_path: Path, train_split: str, preview_size: int): + from datasets import load_dataset + + if dataset_path.is_dir(): + parquet_dir = ( + dataset_path / "parquet-files" + if (dataset_path / "parquet-files").exists() + else dataset_path + ) + parquet_files = sorted(parquet_dir.glob("*.parquet")) + if parquet_files: + dataset = load_dataset( + "parquet", + data_files = [str(path) for path in parquet_files], + split = train_split, + ) + total_rows = len(dataset) + preview_slice = dataset.select(range(min(preview_size, total_rows))) + return preview_slice, total_rows + + candidate_files: list[Path] = [] + for ext in LOCAL_FILE_EXTS: + candidate_files.extend(sorted(dataset_path.glob(f"*{ext}"))) + if not candidate_files: + raise HTTPException( + status_code = 400, + detail = "Unsupported local dataset directory (expected parquet/json/jsonl/csv files)", + ) + dataset_path = candidate_files[0] + + suffix = dataset_path.suffix.lower() + # Parquet/Arrow give a cheap exact total_rows via len()+select; JSON/CSV + # carry no such metadata, so stream them and report total_rows=None. + if suffix == ".parquet": + dataset = load_dataset("parquet", data_files = str(dataset_path), split = train_split) + total_rows = len(dataset) + preview_slice = dataset.select(range(min(preview_size, total_rows))) + return preview_slice, total_rows + + if suffix in (".json", ".jsonl", ".csv"): + preview = _stream_file_preview_slice(dataset_path, preview_size) + if preview is None: + raise HTTPException( + status_code = 400, + detail = "Dataset appears to be empty or could not be read", + ) + return preview + + raise HTTPException(status_code = 400, detail = f"Unsupported file format: {dataset_path.suffix}") + + +def _sanitize_filename(filename: str) -> str: + name = Path(filename).name.strip().replace("\x00", "") + if not name: + return "dataset_upload" + return name + + +def _upload_too_large(size_bytes: int) -> HTTPException: + return HTTPException( + status_code = 413, + detail = (f"Upload is too large " f"({size_bytes:,} bytes; max {LOCAL_UPLOAD_MAX_BYTES:,})."), + ) + + +async def upload_dataset_response(file: UploadFile) -> UploadDatasetResponse: + filename = _sanitize_filename(file.filename or "dataset_upload") + ext = Path(filename).suffix.lower() + if ext not in LOCAL_UPLOAD_EXTS: + allowed = ", ".join(sorted(LOCAL_UPLOAD_EXTS)) + raise HTTPException( + status_code = 400, + detail = f"Unsupported file type: {ext}. Allowed: {allowed}", + ) + + declared_size = getattr(file, "size", None) + if isinstance(declared_size, int) and declared_size > LOCAL_UPLOAD_MAX_BYTES: + raise _upload_too_large(declared_size) + + ensure_dir(DATASET_UPLOAD_DIR) + stem = Path(filename).stem + stored_name = f"{uuid.uuid4().hex}_{stem}{ext}" + stored_path = DATASET_UPLOAD_DIR / stored_name + + written = 0 + try: + with open(stored_path, "wb") as f: + while chunk := await file.read(LOCAL_UPLOAD_CHUNK_BYTES): + written += len(chunk) + if written > LOCAL_UPLOAD_MAX_BYTES: + raise _upload_too_large(written) + await asyncio.to_thread(f.write, chunk) + except Exception: + stored_path.unlink(missing_ok = True) + raise + + if written == 0: + stored_path.unlink(missing_ok = True) + raise HTTPException(status_code = 400, detail = "Empty upload payload") + + return UploadDatasetResponse(filename = filename, stored_path = str(stored_path)) + + +def list_local_datasets_response() -> LocalDatasetsResponse: + return LocalDatasetsResponse(datasets = _build_local_dataset_items()) diff --git a/studio/backend/hub/services/download_lifecycle.py b/studio/backend/hub/services/download_lifecycle.py new file mode 100644 index 0000000000..8256b00252 --- /dev/null +++ b/studio/backend/hub/services/download_lifecycle.py @@ -0,0 +1,449 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from __future__ import annotations + +import os +import signal +import subprocess +import sys +import threading +from pathlib import Path +from typing import Callable, Optional + +from fastapi import HTTPException + +from hub.schemas.downloads import ActiveDownload, DownloadJobState +from hub.utils import download_manifest +from hub.utils import download_registry +from hub.utils import inventory_scan as hf_cache_scan +from hub.utils.hf_cache_state import EXIT_CANCELLED +from hub.utils.state_dir import RepoType + + +def backend_dir() -> Path: + return Path(__file__).resolve().parent.parent.parent + + +def resolve_transport(use_xet: bool) -> str: + transport = download_registry.TRANSPORT_XET if use_xet else download_registry.TRANSPORT_HTTP + unavailable_reason = download_registry.download_transport_unavailable_reason(transport) + if unavailable_reason is not None: + raise HTTPException(status_code = 400, detail = unavailable_reason) + return transport + + +def spawn_worker( + args: list[str], + hf_token: Optional[str], + *, + use_xet: bool, + protected_blob_hashes: Optional[frozenset[str]] = None, +) -> subprocess.Popen: + """Spawn the download worker. + + XET and ``hf_transfer`` write chunks out of order, so their partials can't + resume under a sequential writer; the HTTP path stays sequential so + SIGKILL -> resume is byte-identical. ``protected_blob_hashes`` are blobs a + concurrent same-repo peer is writing, excluded from the cache-prep purge so a + shared ``.incomplete`` (e.g. bundled mmproj) is never deleted. + """ + cwd = backend_dir() + mode = download_registry.TRANSPORT_XET if use_xet else download_registry.TRANSPORT_HTTP + env = os.environ.copy() + if protected_blob_hashes: + env["UNSLOTH_PROTECTED_BLOB_HASHES"] = ",".join(sorted(protected_blob_hashes)) + else: + env.pop("UNSLOTH_PROTECTED_BLOB_HASHES", None) + env["HF_HUB_DISABLE_PROGRESS_BARS"] = "1" + env["HF_HUB_DISABLE_TELEMETRY"] = "1" + env["HF_HUB_DISABLE_XET"] = "0" if use_xet else "1" + env["HF_HUB_DISABLE_IMPLICIT_TOKEN"] = "0" if hf_token else "1" + # hf_transfer's parallel Range chunks can leave sparse partials even in + # "http" mode; disable so the worker's writer is always sequential. + env["HF_HUB_ENABLE_HF_TRANSFER"] = "0" + for token_key in ( + "HF_TOKEN", + "HF_HUB_TOKEN", + "HUGGING_FACE_HUB_TOKEN", + "HUGGINGFACE_HUB_TOKEN", + "HUGGINGFACEHUB_API_TOKEN", + ): + env.pop(token_key, None) + if hf_token: + env["HF_TOKEN"] = hf_token + existing_path = env.get("PYTHONPATH", "") + env["PYTHONPATH"] = f"{cwd}{os.pathsep}{existing_path}" if existing_path else str(cwd) + return subprocess.Popen( + [ + sys.executable, + "-m", + "hub.workers.hf_download", + *args, + "--parent-pid", + str(os.getpid()), + "--transport", + mode, + ], + env = env, + cwd = str(cwd), + stdout = subprocess.DEVNULL, + stderr = subprocess.PIPE, + start_new_session = sys.platform != "win32", + ) + + +def drain_stderr_excerpt(stream, edge_bytes: int = 500) -> bytes: + """Drain a worker's stderr to EOF, retaining the first and last bytes. + + Incremental reads keep the pipe from filling while bounding memory; long + messages keep both ends since stderr prefixes often name the failing repo.""" + if stream is None: + return b"" + edge_bytes = max(1, edge_bytes) + max_bytes = edge_bytes * 2 + full = bytearray() + head = bytearray() + tail = bytearray() + truncated = False + for chunk in iter(lambda: stream.read(4096), b""): + if not truncated: + full.extend(chunk) + if len(full) <= max_bytes: + continue + truncated = True + head.extend(full[:edge_bytes]) + tail.extend(full[-edge_bytes:]) + full.clear() + continue + tail.extend(chunk) + if len(tail) > edge_bytes: + del tail[:-edge_bytes] + if not truncated: + return bytes(full) + return bytes(head + b"\n...[stderr truncated]...\n" + tail) + + +def _cancellation_return_codes() -> frozenset[int]: + """Returncodes for intentional cancellation only (SIGKILL/SIGTERM/SIGINT); crash signals stay errors, and ``getattr`` tolerates Windows where these signals are absent.""" + codes: set[int] = set() + for name in ("SIGKILL", "SIGTERM", "SIGINT"): + sig = getattr(signal, name, None) + if sig is not None: + codes.add(-int(sig)) + return frozenset(codes) + + +_CANCELLATION_RETURN_CODES = _cancellation_return_codes() + + +def _sigpipe_return_codes() -> frozenset[int]: + sig = getattr(signal, "SIGPIPE", None) + if sig is None: + return frozenset() + value = int(sig) + return frozenset({-value, 128 + value}) + + +_SIGPIPE_RETURN_CODES = _sigpipe_return_codes() + + +def classify_exit(rc: int, *, cancel_requested: bool = False) -> str: + """Map a worker process exit code to a job state. + + - rc == 0: clean completion. + - rc == EXIT_CANCELLED (130): the worker trapped a stop signal and exited + cleanly with a resumable partial. In-app cancel uses untrappable SIGKILL + and the OOM killer never produces 130, so 130 is always a resumable cancel. + - rc killed by SIGKILL/SIGTERM/SIGINT: a cancel only when *we* asked for it. + The OOM killer also sends SIGKILL, so an unrequested kill surfaces as error. + - rc killed by SIGPIPE (or 128+SIGPIPE): parent pipe is gone; treated as + cancelled. + - any other non-zero rc (incl. crash signals): worker errored out. + + Windows has no POSIX signal exit encoding, so a user cancel can't be told from + an error by code alone; there ``cancel_requested`` decides. + """ + if rc == 0: + return "complete" + if rc == EXIT_CANCELLED: + return "cancelled" + if rc in _SIGPIPE_RETURN_CODES: + return "cancelled" + if rc in _CANCELLATION_RETURN_CODES: + return "cancelled" if cancel_requested else "error" + if cancel_requested and sys.platform == "win32": + return "cancelled" + return "error" + + +def finalize_worker_exit( + registry: download_registry.DownloadRegistry, + key: str, + proc: subprocess.Popen, + *, + hf_token: Optional[str], + label: str, + log_prefix: str, + logger, + repo_type: Optional[RepoType] = None, + repo_id: Optional[str] = None, + transport: Optional[str] = None, +) -> None: + """Block until *proc* exits, then record the job's terminal state in + *registry*. Drains and scrubs stderr first, then classifies the exit code. + A no-op when the process was already dropped (e.g. superseded). + + No stall watchdog: huggingface_hub already times out chunk reads and raises + a resumable error on a dead connection, so the worker's exit code is the + single source of truth.""" + stderr_data = drain_stderr_excerpt(proc.stderr) + rc = proc.wait() + cancel_requested = registry.cancel_requested(key) + if not registry.drop_process(key, proc): + return + stderr_text = download_registry.scrub_secrets( + (stderr_data or b"").decode("utf-8", "replace").strip(), + hf_token = hf_token, + ) + state = classify_exit(rc, cancel_requested = cancel_requested) + if state == "complete": + registry.set_job(key, "complete") + if stderr_text: + if download_manifest.MANIFEST_DEGRADED_MARKER in stderr_text: + logger.warning( + f"{log_prefix} complete with degraded diagnostics for " + f"{label}: {stderr_text}" + ) + else: + logger.info(f"{log_prefix} worker diagnostics for {label}: {stderr_text}") + logger.info(f"{log_prefix} complete: {label}") + # Defensive cleanup: the canonical clear is at download-start; this + # catches the rare case where that failed but the download succeeded. + if repo_type and repo_id: + try: + download_manifest.clear_cancel_marker( + repo_type, + repo_id, + download_registry.variant_from_key(key), + ) + except Exception as exc: + logger.debug(f"clear_cancel_marker failed for {repo_id} (rc=0): {exc}") + elif state == "cancelled": + # Read metadata before the terminal set_job so a concurrent eviction + # can't drop it; the job key is the fallback variant label. + metadata = registry.get_job_metadata(key) + registry.set_job(key, "cancelled") + logger.info(f"{log_prefix} cancelled: {label} (rc={rc})") + download_registry.persist_cancel_marker( + repo_type, + repo_id, + metadata.variant + if metadata is not None and metadata.variant + else download_registry.variant_from_key(key), + transport, + logger = logger, + ) + else: + registry.set_job( + key, + "error", + stderr_text or f"worker exited with code {rc}", + ) + logger.error( + f"{log_prefix} failed for {label} (rc={rc}): {stderr_text}", + ) + + +def kill_and_reap_process( + proc: subprocess.Popen, + *, + label: str, + logger, + timeout: float = 10.0, +) -> None: + try: + proc.kill() + except ProcessLookupError: + pass + except Exception as exc: + logger.warning(f"Cancel SIGKILL for {label} failed: {exc}") + try: + proc.wait(timeout = timeout) + except subprocess.TimeoutExpired: + logger.warning(f"Cancelled worker for {label} did not exit after SIGKILL") + except Exception: + pass + + +def register_worker( + registry: download_registry.DownloadRegistry, + key: str, + proc: subprocess.Popen, + *, + hf_token: Optional[str], + label: str, + log_prefix: str, + logger, + repo_type: RepoType, + repo_id: str, + transport: str, + watch_name: str, +) -> bool: + if not registry.register_process(key, proc): + kill_and_reap_process(proc, label = label, logger = logger) + return False + + worker_token = hf_token + + def _watch() -> None: + finalize_worker_exit( + registry, + key, + proc, + hf_token = worker_token, + label = label, + log_prefix = log_prefix, + logger = logger, + repo_type = repo_type, + repo_id = repo_id, + transport = transport, + ) + if registry.get_job(key).state in ("error", "cancelled"): + download_registry.purge_empty_marker_dir( + repo_type, + repo_id, + download_registry.variant_from_key(key), + ) + hf_cache_scan.invalidate_hf_cache_scans() + + threading.Thread(target = _watch, name = watch_name, daemon = True).start() + return True + + +def launch_worker( + registry: download_registry.DownloadRegistry, + key: str, + *, + spawn: Callable[[], subprocess.Popen], + hf_token: Optional[str], + label: str, + log_prefix: str, + logger, + repo_type: RepoType, + repo_id: str, + transport: str, + watch_name: str, +) -> str: + try: + proc = spawn() + except Exception as e: + scrubbed = download_registry.scrub_secrets(str(e), hf_token = hf_token) + logger.error( + f"Failed to spawn {log_prefix.lower()} worker for {label}: {scrubbed}", + exc_info = True, + ) + registry.set_job(key, "error", scrubbed) + raise HTTPException( + status_code = 500, + detail = f"Failed to start {log_prefix.lower()}: {scrubbed}", + ) from e + register_worker( + registry, + key, + proc, + hf_token = hf_token, + label = label, + log_prefix = log_prefix, + logger = logger, + repo_type = repo_type, + repo_id = repo_id, + transport = transport, + watch_name = watch_name, + ) + return registry.get_job(key).state + + +def cancel_worker( + registry: download_registry.DownloadRegistry, + key: str, + *, + generation: Optional[int], + label: str, + logger, +) -> str: + proc = registry.get_process(key) + # No worker process yet: arm a pending cancel so register_process kills it on + # arrival during the claim-to-register window. + if proc is None: + if registry.mark_pending_cancel(key, generation): + return "cancelling" + return registry.get_job(key).state + # Worker already exited; let its watcher classify the real return code. + # Arming a pending cancel here could mislabel a genuine failure as a cancel. + if proc.poll() is not None: + return registry.get_job(key).state + + if not registry.request_cancel(key, proc, generation): + return registry.get_job(key).state + # No eager marker: finalize_worker_exit writes it on a "cancelled" exit. + # Persisting before the kill races a clean completion and strands a stale marker. + try: + proc.kill() + except ProcessLookupError: + pass + except Exception as e: + logger.warning(f"Cancel SIGKILL for {label} failed: {e}") + + return "cancelling" + + +def idle_status( + registry: download_registry.DownloadRegistry, + key: str, + *, + repo_type: RepoType, + repo_id: Optional[str], + variant: Optional[str], +) -> tuple[DownloadJobState, Optional[str], int]: + state = registry.get_job(key) + generation = registry.current_generation(key) + if ( + state.state == "idle" + and repo_id + and download_manifest.has_cancel_marker( + repo_type, + repo_id, + variant, + ) + ): + return ("cancelled", None, generation) + return (state.state, state.error, generation) + + +def active_download_refs( + registry: download_registry.DownloadRegistry, repo_id: Optional[str], *, with_variant: bool +) -> list[ActiveDownload]: + downloads: list[ActiveDownload] = [] + for ref in registry.active_job_refs(repo_id): + metadata = ref.metadata + if with_variant: + ref_repo_id = metadata.repo_id if metadata is not None else ref.key.split("::", 1)[0] + if metadata is not None: + variant = metadata.variant + else: + _repo, sep, raw_variant = ref.key.partition("::") + variant = raw_variant if sep and raw_variant else None + else: + ref_repo_id = metadata.repo_id if metadata is not None else ref.key + variant = None + downloads.append( + ActiveDownload( + repo_id = ref_repo_id, + variant = variant, + transport = metadata.transport if metadata is not None else None, + state = ref.state, + generation = ref.generation, + ) + ) + return downloads diff --git a/studio/backend/hub/services/models/__init__.py b/studio/backend/hub/services/models/__init__.py new file mode 100644 index 0000000000..707a7c633a --- /dev/null +++ b/studio/backend/hub/services/models/__init__.py @@ -0,0 +1,4 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Model service layer.""" diff --git a/studio/backend/hub/services/models/cache_inventory.py b/studio/backend/hub/services/models/cache_inventory.py new file mode 100644 index 0000000000..a961a6ae9d --- /dev/null +++ b/studio/backend/hub/services/models/cache_inventory.py @@ -0,0 +1,461 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Cached model inventory.""" + +from __future__ import annotations + +import json +import asyncio +import threading +import time +from collections import OrderedDict +from pathlib import Path +from typing import NamedTuple, Optional + +from fastapi import HTTPException +from loggers import get_logger + +from hub.schemas.inventory import ModelFormat +from hub.utils import inventory_scan as hf_cache_scan +from hub.utils import download_registry +from hub.utils.snapshot_filters import ( + snapshot_download_blob_hashes, + snapshot_download_size, +) +from hub.services.models.common import ( + _capabilities_for_format, + _classify_non_gguf_model_format, + _gguf_variant_state_summary, + _is_adapter_weight_name, + _is_checkpoint_weight_name, + _is_gguf_filename, + _is_main_gguf_filename, + _is_transformers_safetensors_weight_name, + _local_inventory_id, + _prefer_complete_larger, + _runtime_for_format, +) + +logger = get_logger(__name__) + +_repo_size_cache: "OrderedDict[tuple[str, str], tuple[int, frozenset[str], float]]" = OrderedDict() +_repo_size_neg_cache: "OrderedDict[tuple[str, str], float]" = OrderedDict() +_REPO_SIZE_CACHE_MAX = 256 +_REPO_SIZE_POS_TTL = 60.0 +_REPO_SIZE_NEG_TTL = 60.0 +_MODEL_METADATA_TIMEOUT_SECONDS = 5.0 +_repo_size_cache_lock = threading.Lock() + + +def get_repo_snapshot_metadata_cached( + repo_id: str, hf_token: Optional[str] = None +) -> tuple[int, frozenset[str]]: + token_fp = hf_cache_scan.token_fingerprint(hf_token) + cache_key = (repo_id, token_fp) + with _repo_size_cache_lock: + cached = _repo_size_cache.get(cache_key) + if cached is not None: + total, blob_hashes, ts = cached + if (time.monotonic() - ts) < _REPO_SIZE_POS_TTL: + _repo_size_cache.move_to_end(cache_key) + return total, blob_hashes + del _repo_size_cache[cache_key] + neg_ts = _repo_size_neg_cache.get(cache_key) + if neg_ts is not None and (time.monotonic() - neg_ts) < _REPO_SIZE_NEG_TTL: + return 0, frozenset() + try: + from huggingface_hub import HfApi + + info = HfApi(token = hf_token).model_info( + repo_id, + files_metadata = True, + timeout = _MODEL_METADATA_TIMEOUT_SECONDS, + ) + total = snapshot_download_size(info.siblings) + blob_hashes = snapshot_download_blob_hashes(info.siblings) + except Exception as e: + logger.warning( + "Failed to get repo size for %s: %s", + repo_id, + download_registry.scrub_secrets(str(e), hf_token = hf_token), + ) + with _repo_size_cache_lock: + _repo_size_neg_cache[cache_key] = time.monotonic() + _repo_size_neg_cache.move_to_end(cache_key) + while len(_repo_size_neg_cache) > _REPO_SIZE_CACHE_MAX: + _repo_size_neg_cache.popitem(last = False) + return 0, frozenset() + with _repo_size_cache_lock: + _repo_size_cache[cache_key] = (total, blob_hashes, time.monotonic()) + _repo_size_cache.move_to_end(cache_key) + _repo_size_neg_cache.pop(cache_key, None) + while len(_repo_size_cache) > _REPO_SIZE_CACHE_MAX: + _repo_size_cache.popitem(last = False) + return total, blob_hashes + + +def all_hf_cache_scans(): + return hf_cache_scan.all_hf_cache_scans() + + +def _repo_gguf_size_bytes(repo_info) -> int: + """Sum primary GGUF blob sizes across revisions, deduped by blob path (HF hardlinks shared blobs); mmproj is excluded so a vision-adapter-only repo isn't classed as GGUF.""" + unique_blobs: dict[str, int] = {} + for revision in repo_info.revisions: + rev_id = getattr(revision, "commit_hash", None) or str(id(revision)) + for f in revision.files: + if _is_main_gguf_filename(f.file_name): + blob_path = getattr(f, "blob_path", None) + size = f.size_on_disk or 0 + if blob_path: + unique_blobs[str(blob_path)] = size + else: + unique_blobs[f"{rev_id}:{f.file_name}"] = size + return sum(unique_blobs.values()) + + +def _repo_has_gguf_files(repo_info) -> bool: + return _repo_gguf_size_bytes(repo_info) > 0 + + +def _prefer_cache_row(candidate: dict, existing: Optional[dict]) -> bool: + if existing is None: + return True + return _prefer_complete_larger( + bool(candidate.get("partial")), + int(candidate.get("size_bytes") or 0), + bool(existing.get("partial")), + int(existing.get("size_bytes") or 0), + ) + + +def _cache_inventory_fields( + repo_id: str, + model_format: ModelFormat, + *, + partial: bool = False, + requires_variant: bool = False, +) -> dict: + return { + "inventory_id": _local_inventory_id("cache", model_format, repo_id), + "load_id": repo_id, + "model_format": model_format, + "runtime": _runtime_for_format(model_format), + "format_variant": None, + "capabilities": _capabilities_for_format( + model_format, + "hf_cache", + partial = partial, + requires_variant = requires_variant, + ).model_dump(), + } + + +def invalidate_hf_cache_scans() -> None: + hf_cache_scan.invalidate_hf_cache_scans() + + +def _scan_cached_gguf() -> list[dict]: + """Synchronous HF-cache disk walk for GGUF repos; runs in a worker thread.""" + cache_scans = all_hf_cache_scans() + + seen_lower: dict[str, dict] = {} + for hf_cache in cache_scans: + for repo_info in hf_cache.repos: + try: + if str(repo_info.repo_type) != "model": + continue + repo_id = repo_info.repo_id + total_size = _repo_gguf_size_bytes(repo_info) + has_variant_state, variant_state_size = _gguf_variant_state_summary(repo_id) + if total_size == 0 and not has_variant_state: + continue + partial = hf_cache_scan.is_gguf_repo_partial( + repo_id, + Path(repo_info.repo_path), + ) + if total_size == 0 and not partial: + continue + key = repo_id.lower() + existing = seen_lower.get(key) + row = { + "repo_id": repo_id, + "size_bytes": max(total_size, variant_state_size), + "cache_path": str(repo_info.repo_path), + "partial": partial, + # GGUF row-level transport is ambiguous (variants may differ); + # per-variant detail lives on GgufVariantDetail. + "partial_transport": None, + } + row.update( + _cache_inventory_fields( + repo_id, + "gguf", + partial = bool(row["partial"]), + requires_variant = True, + ) + ) + if _prefer_cache_row(row, existing): + seen_lower[key] = row + except Exception as e: + repo_label = getattr(repo_info, "repo_id", "") + logger.warning(f"Skipping cached GGUF repo {repo_label}: {e}") + continue + return sorted(seen_lower.values(), key = lambda c: c["repo_id"]) + + +async def list_cached_gguf_response(hf_token: Optional[str] = None): + """List GGUF repos downloaded to HF cache, legacy Unsloth cache, and HF default cache.""" + try: + cached = await asyncio.to_thread(_scan_cached_gguf) + return {"cached": cached} + except Exception as e: + logger.error( + "Error listing cached GGUF repos: %s", + download_registry.scrub_secrets(str(e), hf_token = hf_token), + ) + raise HTTPException( + status_code = 500, + detail = "Failed to read the local model cache.", + ) from e + + +class _CachedNonGgufPayload(NamedTuple): + size_bytes: int + has_runnable_weights: bool + model_format: ModelFormat + + +def _repo_non_gguf_model_payload(repo_info) -> _CachedNonGgufPayload: + all_weight_blobs: dict[str, int] = {} + adapter_blobs: dict[str, int] = {} + safetensors_blobs: dict[str, int] = {} + checkpoint_blobs: dict[str, int] = {} + has_config = False + has_adapter_config = False + has_adapter_weights = False + has_safetensors = False + has_transformers_safetensors = False + has_checkpoint = False + + def _record_blob(target: dict[str, int], file_obj, rev_id: str, file_name: str) -> None: + blob_path = getattr(file_obj, "blob_path", None) + size = int(file_obj.size_on_disk or 0) + key = str(blob_path) if blob_path else f"{rev_id}:{file_name}" + target[key] = size + all_weight_blobs[key] = size + + for revision in repo_info.revisions: + rev_id = getattr(revision, "commit_hash", None) or str(id(revision)) + for f in revision.files: + file_name = str(f.file_name) + lower = file_name.lower() + name = lower.replace("\\", "/").rsplit("/", 1)[-1] + if _is_gguf_filename(lower): + continue + if name == "config.json": + has_config = True + continue + if name == "adapter_config.json": + has_adapter_config = True + continue + is_adapter = _is_adapter_weight_name(name) + is_safetensors = name.endswith(".safetensors") and not is_adapter + is_checkpoint = _is_checkpoint_weight_name(name) + if is_adapter: + has_adapter_weights = True + _record_blob(adapter_blobs, f, rev_id, file_name) + if is_safetensors: + has_safetensors = True + if _is_transformers_safetensors_weight_name(name): + has_transformers_safetensors = True + _record_blob(safetensors_blobs, f, rev_id, file_name) + if is_checkpoint: + has_checkpoint = True + _record_blob(checkpoint_blobs, f, rev_id, file_name) + + model_format = ( + _classify_non_gguf_model_format( + has_config = has_config, + has_adapter_config = has_adapter_config, + has_adapter_weights = has_adapter_weights, + has_safetensors = has_safetensors, + has_transformers_safetensors = has_transformers_safetensors, + has_checkpoint_weights = has_checkpoint, + trusted_hf_cache_repo = True, + ) + or "unknown" + ) + if model_format == "adapter": + size_bytes = sum(adapter_blobs.values()) + elif model_format == "safetensors": + size_bytes = sum(safetensors_blobs.values()) + elif model_format == "checkpoint": + size_bytes = sum(checkpoint_blobs.values()) + else: + size_bytes = sum(all_weight_blobs.values()) + + return _CachedNonGgufPayload( + size_bytes = size_bytes, + has_runnable_weights = model_format != "unknown", + model_format = model_format, + ) + + +def _cached_model_snapshot_path(repo_path: Path) -> Optional[Path]: + resolved = hf_cache_scan.resolve_hf_cache_realpath(repo_path) + if not resolved: + return None + path = Path(resolved) + return path if path.is_dir() else None + + +def _read_json_object(path: Path) -> dict: + try: + with open(path, "r", encoding = "utf-8") as f: + data = json.load(f) + return data if isinstance(data, dict) else {} + except Exception: + return {} + + +def _read_model_card_frontmatter(path: Path) -> dict: + try: + text = path.read_text(encoding = "utf-8") + except Exception: + return {} + lines = text.splitlines() + if not lines or lines[0].strip() != "---": + return {} + body: list[str] = [] + for line in lines[1:]: + if line.strip() == "---": + break + body.append(line) + if not body: + return {} + try: + import yaml + data = yaml.safe_load("\n".join(body)) or {} + return data if isinstance(data, dict) else {} + except Exception: + return {} + + +def _cached_model_local_metadata(repo_path: Path) -> dict: + snapshot = _cached_model_snapshot_path(repo_path) + if snapshot is None: + return {} + + result: dict = {} + config = _read_json_object(snapshot / "config.json") + quant_method = ( + config.get("quantization_config", {}).get("quant_method") + if isinstance(config.get("quantization_config"), dict) + else None + ) + if isinstance(quant_method, str) and quant_method.strip(): + result["quant_method"] = quant_method.strip() + + card = _read_model_card_frontmatter(snapshot / "README.md") + pipeline_tag = card.get("pipeline_tag") + if isinstance(pipeline_tag, str) and pipeline_tag.strip(): + result["pipeline_tag"] = pipeline_tag.strip() + library_name = card.get("library_name") + if isinstance(library_name, str) and library_name.strip(): + result["library_name"] = library_name.strip() + tags = card.get("tags") + if isinstance(tags, list): + clean_tags = [tag.strip() for tag in tags if isinstance(tag, str) and tag.strip()] + if clean_tags: + result["tags"] = clean_tags + return result + + +def _scan_cached_models() -> list[dict]: + """Synchronous HF-cache disk walk for non-GGUF model repos; runs in a worker thread.""" + cache_scans = all_hf_cache_scans() + + seen_lower: dict[str, dict] = {} + inspected = 0 + skipped_gguf = 0 + skipped_no_weights = 0 + for hf_cache in cache_scans: + for repo_info in hf_cache.repos: + inspected += 1 + try: + if str(repo_info.repo_type) != "model": + continue + repo_id = repo_info.repo_id + has_main_gguf = _repo_has_gguf_files(repo_info) + payload = _repo_non_gguf_model_payload(repo_info) + if payload.size_bytes == 0: + if has_main_gguf: + skipped_gguf += 1 + continue + if not payload.has_runnable_weights: + skipped_no_weights += 1 + continue + key = repo_id.lower() + existing = seen_lower.get(key) + repo_path = Path(repo_info.repo_path) + snapshot_partial = hf_cache_scan.is_snapshot_partial( + "model", + repo_id, + repo_path, + ) + row = { + "repo_id": repo_id, + "size_bytes": payload.size_bytes, + "cache_path": str(repo_info.repo_path), + "partial": snapshot_partial, + "partial_transport": ( + hf_cache_scan.partial_transport_for( + "model", + repo_id, + repo_cache_dir = repo_path, + ) + if snapshot_partial + else None + ), + **_cached_model_local_metadata(repo_path), + } + row.update( + _cache_inventory_fields( + repo_id, + payload.model_format, + partial = bool(row["partial"]), + ) + ) + if _prefer_cache_row(row, existing): + seen_lower[key] = row + except Exception as e: + repo_label = getattr(repo_info, "repo_id", "") + logger.warning(f"Skipping cached model repo {repo_label}: {e}") + continue + cached = sorted(seen_lower.values(), key = lambda c: c["repo_id"]) + logger.info( + "Cached model scan: inspected=%d skipped_gguf=%d skipped_no_weights=%d returned=%d", + inspected, + skipped_gguf, + skipped_no_weights, + len(cached), + ) + return cached + + +async def list_cached_models_response(hf_token: Optional[str] = None): + """List non-GGUF model repos downloaded to HF cache, legacy Unsloth cache, and HF default cache.""" + try: + cached = await asyncio.to_thread(_scan_cached_models) + return {"cached": cached} + except Exception as e: + logger.error( + "Error listing cached models: %s", + download_registry.scrub_secrets(str(e), hf_token = hf_token), + ) + raise HTTPException( + status_code = 500, + detail = "Failed to read the local model cache.", + ) from e diff --git a/studio/backend/hub/services/models/common.py b/studio/backend/hub/services/models/common.py new file mode 100644 index 0000000000..688d324a60 --- /dev/null +++ b/studio/backend/hub/services/models/common.py @@ -0,0 +1,610 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Shared model inventory helpers for the Hub service layer.""" + +from __future__ import annotations + +import json +import os +from pathlib import Path +from typing import List, Literal, Optional +from urllib.parse import quote + +from hub.schemas.inventory import ( + LocalModelCapabilities, + LocalModelInfo, + ModelFormat, + ModelRuntime, +) +from hub.utils.gguf import ( + extract_quant_label, + is_gguf_filename as _is_gguf_filename, + is_mmproj_filename as _is_mmproj_filename, +) +from hub.utils.paths import is_valid_repo_id as _is_valid_repo_id + +ModelType = Literal["text", "vision", "audio", "embeddings"] +LocalModelSource = Literal["models_dir", "hf_cache", "lmstudio", "ollama", "custom"] + + +def _safe_is_dir(path) -> bool: + # Py >= 3.12 propagates PermissionError (EACCES) from is_dir(); folder scans + # probe root-owned system dirs, so treat un-stat-able paths as not-a-dir. + try: + return Path(path).is_dir() + except OSError: + return False + + +_LOCAL_CHECKPOINT_EXTENSIONS = ( + ".bin", + ".pt", + ".pth", + ".ckpt", + ".h5", + ".msgpack", + ".npz", +) + +_LOCAL_BASE_MODEL_PREFIXES = { + "checkpoint", + "checkpoints", + "export", + "exports", + "model", + "models", + "output", + "outputs", + "run", + "runs", + "train", +} +_HF_CACHE_MODEL_FILE_PROBE_LIMIT = 2000 + + +def _is_model_directory(d: Path) -> bool: + """True when *d* has a config plus real weights; excludes mmproj GGUFs and non-weight ``.bin`` files (``tokenizer.bin``) to avoid false positives.""" + + def _is_weight_file(f: Path) -> bool: + suffix = f.suffix.lower() + if suffix == ".safetensors": + return True + if suffix == ".gguf": + return "mmproj" not in f.name.lower() + if suffix == ".bin": + name = f.name.lower() + return ( + name.startswith("pytorch_model") + or name.startswith("model") + or name.startswith("adapter_model") + or name.startswith("consolidated") + ) + return False + + try: + has_config = (d / "config.json").exists() or (d / "adapter_config.json").exists() + if not has_config: + return False + return any(_is_weight_file(f) for f in d.iterdir() if f.is_file()) + except OSError: + return False + + +def _local_inventory_id( + source: str, + model_format: ModelFormat, + semantic_id: str, + variant: Optional[str] = None, +) -> str: + parts = [ + source, + model_format, + quote(semantic_id, safe = ""), + ] + if variant: + parts.append(quote(variant, safe = "")) + return ":".join(parts) + + +def _runtime_for_format(model_format: ModelFormat) -> ModelRuntime: + if model_format == "gguf": + return "llama_cpp" + if model_format == "adapter": + return "adapter" + if model_format in {"safetensors", "checkpoint"}: + return "transformers" + return "unknown" + + +def _capabilities_for_format( + model_format: ModelFormat, + source: str, + *, + partial: bool = False, + requires_variant: bool = False, +) -> LocalModelCapabilities: + is_complete = not partial + can_chat = model_format in {"gguf", "safetensors", "adapter", "checkpoint"} + can_train = model_format in {"safetensors", "checkpoint"} and is_complete + return LocalModelCapabilities( + can_train = can_train, + can_chat = can_chat and is_complete, + can_delete = source == "hf_cache", + can_download = False, + requires_variant = requires_variant, + supports_lora = model_format in {"safetensors", "checkpoint"} and is_complete, + supports_vision = False, + ) + + +def _prefer_complete_larger( + candidate_partial: bool, + candidate_size_bytes: int, + existing_partial: bool, + existing_size_bytes: int, +) -> bool: + if candidate_partial != existing_partial: + return not candidate_partial + return candidate_size_bytes > existing_size_bytes + + +def _gguf_variant_state_summary(repo_id: str) -> tuple[bool, int]: + """Whether GGUF variant-scoped state exists and its expected size; a cancelled/in-progress variant may have only manifests/markers/`.incomplete` blobs, which inventory needs to avoid a generic fallback row.""" + from hub.utils import download_manifest + + variant_keys: set[str] = set() + size_by_variant: dict[str, int] = {} + for variant, _path in download_manifest.iter_variant_manifests( + "model", + repo_id, + ): + key = variant.lower() + variant_keys.add(key) + manifest = download_manifest.read_manifest("model", repo_id, variant) + if manifest is None: + continue + size_by_variant[key] = max( + size_by_variant.get(key, 0), + sum(max(0, int(file.size or 0)) for file in manifest.expected_files), + ) + for variant, _path in download_manifest.iter_variant_markers( + "model", + repo_id, + ): + variant_keys.add(variant.lower()) + return bool(variant_keys), sum(size_by_variant.values()) + + +def _apply_format_aware_partial( + rows: List[LocalModelInfo], + *, + snapshot_partial: bool, + gguf_partial: bool, + snapshot_partial_transport: Optional[str] = None, +) -> List[LocalModelInfo]: + """Rewrite each row's partial flag with format-aware predicates so a hybrid (gguf + safetensors) repo's broken format doesn't taint the clean one; capabilities are recomputed from the new flag.""" + rewritten: List[LocalModelInfo] = [] + for row in rows: + target = gguf_partial if row.model_format == "gguf" else snapshot_partial + if not target: + rewritten.append(row) + continue + # GGUF row-level transport is ambiguous (variants may differ); per-variant + # detail lives on GgufVariantDetail.partial_transport via the variants endpoint. + partial_transport = None if row.model_format == "gguf" else snapshot_partial_transport + rewritten.append( + row.model_copy( + update = { + "partial": True, + "partial_transport": partial_transport, + "capabilities": _capabilities_for_format( + row.model_format, + row.source, + partial = True, + requires_variant = row.capabilities.requires_variant, + ), + } + ) + ) + return rewritten + + +def _weight_basename(name: str) -> str: + return name.replace("\\", "/").rsplit("/", 1)[-1].lower() + + +def _is_adapter_weight_name(name: str) -> bool: + lower = _weight_basename(name) + return lower.startswith("adapter_model") and lower.endswith((".safetensors", ".bin")) + + +def _is_transformers_safetensors_weight_name(name: str) -> bool: + lower = _weight_basename(name) + return lower.endswith(".safetensors") and lower.startswith( + ("model", "pytorch_model", "consolidated") + ) + + +def _is_transformers_bin_weight_name(name: str) -> bool: + lower = _weight_basename(name) + if not lower.endswith(".bin"): + return False + return lower.startswith(("pytorch_model", "model", "consolidated", "adapter_model")) + + +def _is_checkpoint_weight_name(name: str) -> bool: + lower = _weight_basename(name) + if lower.endswith(".bin"): + return _is_transformers_bin_weight_name(lower) + return lower.endswith(_LOCAL_CHECKPOINT_EXTENSIONS) + + +def _is_adapter_weight_file(path: Path) -> bool: + return _is_adapter_weight_name(path.name) + + +def _is_transformers_safetensors_weight_file(path: Path) -> bool: + return _is_transformers_safetensors_weight_name(path.name) + + +def _is_transformers_bin_weight_file(path: Path) -> bool: + return _is_transformers_bin_weight_name(path.name) + + +def _is_checkpoint_weight_file(path: Path) -> bool: + return _is_checkpoint_weight_name(path.name) + + +def _classify_non_gguf_model_format( + *, + has_config: bool, + has_adapter_config: bool, + has_adapter_weights: bool, + has_safetensors: bool, + has_transformers_safetensors: bool, + has_checkpoint_weights: bool, + trusted_hf_cache_repo: bool = False, +) -> Optional[ModelFormat]: + if has_safetensors and (has_config or (trusted_hf_cache_repo and has_transformers_safetensors)): + return "safetensors" + if has_adapter_config and has_adapter_weights: + return "adapter" + if has_config and has_checkpoint_weights: + return "checkpoint" + return None + + +def _is_main_gguf_filename(name: str) -> bool: + return _is_gguf_filename(name) and not _is_mmproj_filename(name) + + +def _iter_gguf_paths(root: Path): + stack = [root] + while stack: + current = stack.pop() + try: + entries = list(current.iterdir()) + except OSError: + continue + for path in entries: + try: + if path.is_dir() and not path.is_symlink(): + stack.append(path) + elif path.is_file() and _is_gguf_filename(path.name): + yield path + except OSError: + continue + + +def _iter_immediate_files(path: Path, *, include_symlinks: bool = False) -> list[Path]: + if path.is_file(): + return [path] + if not path.is_dir(): + return [] + try: + return [ + entry + for entry in path.iterdir() + if entry.is_file() or (include_symlinks and entry.is_symlink()) + ] + except OSError: + return [] + + +def _iter_hf_cache_model_files(path: Path) -> list[Path]: + files = _iter_immediate_files(path, include_symlinks = True) + if not path.is_dir(): + return files + if any( + _is_main_gguf_filename(entry.name) + or _is_transformers_safetensors_weight_file(entry) + or _is_checkpoint_weight_file(entry) + for entry in files + ): + return files + try: + bounded: list[Path] = [] + for index, entry in enumerate(path.rglob("*"), start = 1): + if index > _HF_CACHE_MODEL_FILE_PROBE_LIMIT: + break + if entry.is_file() or entry.is_symlink(): + bounded.append(entry) + return bounded + except OSError: + return [] + + +def _file_size_bytes(path: Path) -> int: + try: + if path.is_file() or path.is_symlink(): + return path.stat().st_size + except OSError: + return 0 + return 0 + + +def _sum_file_sizes(paths) -> int: + return sum(_file_size_bytes(path) for path in paths) + + +def _main_gguf_files(path: Path, *, include_symlinks: bool = False) -> list[Path]: + return [ + entry + for entry in _iter_immediate_files(path, include_symlinks = include_symlinks) + if _is_main_gguf_filename(entry.name) + ] + + +def _format_label(model_format: ModelFormat) -> str: + if model_format == "gguf": + return "GGUF" + if model_format == "safetensors": + return "Safetensors" + if model_format == "adapter": + return "Adapter" + if model_format == "checkpoint": + return "Checkpoint" + return "Unknown" + + +def _read_adapter_config(path: Path) -> dict: + if not path.is_dir(): + return {} + try: + with (path / "adapter_config.json").open("r", encoding = "utf-8") as f: + data = json.load(f) + except Exception: + return {} + return data if isinstance(data, dict) else {} + + +def _clean_optional_string(value: object) -> Optional[str]: + return value.strip() if isinstance(value, str) and value.strip() else None + + +def _base_model_looks_local(value: str) -> bool: + raw = value.strip() + normalized = raw.replace("\\", "/") + if raw.startswith(("/", "./", "../", "~", "\\\\")) or ( + len(raw) >= 3 and raw[1] == ":" and raw[0].isalpha() + ): + return True + first = normalized.split("/", 1)[0].lower() + return "/" in normalized and first in _LOCAL_BASE_MODEL_PREFIXES + + +def _base_model_source(value: Optional[str], adapter_dir: Path) -> Optional[str]: + if not value: + return None + candidates = [value, value.replace("\\", "/")] + for candidate in candidates: + try: + expanded = Path(os.path.expanduser(candidate)) + if expanded.exists() or (adapter_dir / candidate).exists(): + return "local" + except (OSError, ValueError): + return "unknown" + if _base_model_looks_local(value): + return "local" + if _is_valid_repo_id(value): + return "huggingface" + return "unknown" + + +def _local_model_info( + *, + scan_path: Path, + load_path: Path, + source: LocalModelSource, + model_format: ModelFormat, + display_name: Optional[str] = None, + model_id: Optional[str] = None, + updated_at: Optional[float] = None, + partial: bool = False, + requires_variant: bool = False, + format_variant: Optional[str] = None, + size_bytes: int = 0, + base_model: Optional[str] = None, + base_model_source: Optional[str] = None, + adapter_type: Optional[str] = None, + training_method: Optional[str] = None, +) -> LocalModelInfo: + load_id = model_id if source == "hf_cache" and model_id else str(load_path) + semantic_id = model_id or str(load_path) + return LocalModelInfo( + id = load_id, + inventory_id = _local_inventory_id( + source, + model_format, + semantic_id, + format_variant, + ), + load_id = load_id, + model_id = model_id, + display_name = display_name or (scan_path.stem if scan_path.is_file() else scan_path.name), + path = str(load_path), + size_bytes = max(0, int(size_bytes or 0)), + source = source, + base_model = base_model, + base_model_source = base_model_source, + adapter_type = adapter_type, + training_method = training_method, + updated_at = updated_at, + partial = partial, + model_format = model_format, + runtime = _runtime_for_format(model_format), + format_variant = format_variant, + capabilities = _capabilities_for_format( + model_format, + source, + partial = partial, + requires_variant = requires_variant, + ), + ) + + +def _classify_local_path( + scan_path: Path, + source: LocalModelSource, + *, + load_path: Optional[Path] = None, + display_name: Optional[str] = None, + model_id: Optional[str] = None, + updated_at: Optional[float] = None, + partial: bool = False, +) -> list[LocalModelInfo]: + load_path = load_path or scan_path + files = ( + _iter_hf_cache_model_files(scan_path) + if source == "hf_cache" + else _iter_immediate_files(scan_path) + ) + if not files: + return [] + + rows: list[LocalModelInfo] = [] + include_broken_snapshot_symlinks = source == "hf_cache" + gguf_files = _main_gguf_files( + scan_path, + include_symlinks = include_broken_snapshot_symlinks, + ) + if gguf_files: + gguf_size_bytes = _sum_file_sizes(gguf_files) + variant = ( + extract_quant_label(gguf_files[0].name) + if scan_path.is_file() and len(gguf_files) == 1 + else None + ) + rows.append( + _local_model_info( + scan_path = scan_path, + load_path = load_path, + source = source, + model_format = "gguf", + display_name = display_name, + model_id = model_id, + updated_at = updated_at, + partial = partial, + requires_variant = scan_path.is_dir(), + format_variant = variant, + size_bytes = gguf_size_bytes, + ) + ) + + has_config = (scan_path / "config.json").is_file() if scan_path.is_dir() else False + has_adapter_config = ( + (scan_path / "adapter_config.json").is_file() if scan_path.is_dir() else False + ) + adapter_config = _read_adapter_config(scan_path) if has_adapter_config else {} + adapter_base_model = _clean_optional_string(adapter_config.get("base_model_name_or_path")) + adapter_type = _clean_optional_string(adapter_config.get("peft_type")) + training_method = _clean_optional_string(adapter_config.get("unsloth_training_method")) + has_adapter_weights = any(_is_adapter_weight_file(f) for f in files) + has_safetensors = any( + f.suffix.lower() == ".safetensors" and not _is_adapter_weight_file(f) for f in files + ) + has_transformers_safetensors = any( + _is_transformers_safetensors_weight_file(f) and not _is_adapter_weight_file(f) + for f in files + ) + has_checkpoint_weights = any(_is_checkpoint_weight_file(f) for f in files) + trusted_hf_cache_repo = source == "hf_cache" and bool(model_id) + + model_format = _classify_non_gguf_model_format( + has_config = has_config, + has_adapter_config = has_adapter_config, + has_adapter_weights = has_adapter_weights, + has_safetensors = has_safetensors, + has_transformers_safetensors = has_transformers_safetensors, + has_checkpoint_weights = has_checkpoint_weights, + trusted_hf_cache_repo = trusted_hf_cache_repo, + ) + + if model_format is not None: + if model_format == "adapter": + size_bytes = _sum_file_sizes(f for f in files if _is_adapter_weight_file(f)) + elif model_format == "safetensors": + size_bytes = _sum_file_sizes( + f + for f in files + if f.suffix.lower() == ".safetensors" and not _is_adapter_weight_file(f) + ) + else: + size_bytes = _sum_file_sizes(f for f in files if _is_checkpoint_weight_file(f)) + rows.append( + _local_model_info( + scan_path = scan_path, + load_path = load_path, + source = source, + model_format = model_format, + display_name = display_name, + model_id = model_id, + updated_at = updated_at, + partial = partial, + size_bytes = size_bytes, + base_model = adapter_base_model if model_format == "adapter" else None, + base_model_source = ( + _base_model_source(adapter_base_model, scan_path) + if model_format == "adapter" + else None + ), + adapter_type = adapter_type if model_format == "adapter" else None, + training_method = training_method if model_format == "adapter" else None, + ) + ) + elif not rows: + fallback_format: ModelFormat = ( + "safetensors" if trusted_hf_cache_repo and has_config else "unknown" + ) + size_bytes = _sum_file_sizes(files) + rows.append( + _local_model_info( + scan_path = scan_path, + load_path = load_path, + source = source, + model_format = fallback_format, + display_name = display_name, + model_id = model_id, + updated_at = updated_at, + partial = partial or trusted_hf_cache_repo, + size_bytes = size_bytes, + ) + ) + + if len(rows) > 1: + rows = [ + row.model_copy( + update = { + "display_name": f"{row.display_name} ({_format_label(row.model_format)})", + "inventory_id": _local_inventory_id( + row.source, + row.model_format, + row.model_id or row.path, + row.format_variant, + ), + } + ) + for row in rows + ] + return rows diff --git a/studio/backend/hub/services/models/deletion.py b/studio/backend/hub/services/models/deletion.py new file mode 100644 index 0000000000..d4aed0d59f --- /dev/null +++ b/studio/backend/hub/services/models/deletion.py @@ -0,0 +1,455 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Cached model deletion.""" + +from __future__ import annotations + +import asyncio +from pathlib import Path +from typing import Optional + +from fastapi import HTTPException +from loggers import get_logger + +from hub.utils import download_manifest +from hub.utils import download_registry +from hub.utils import inventory_scan as hf_cache_scan +from hub.utils.gguf import extract_quant_label +from hub.utils.hf_cache_state import ( + INCOMPLETE_SUFFIX, + purge_partial_repo, + purge_repo_cache_dirs, +) +from hub.utils.paths import ( + is_valid_gguf_variant as _is_valid_gguf_variant, + is_valid_repo_id as _is_valid_repo_id, + resolve_cached_repo_id_case, +) +from hub.services import resolve_destructive_repo_ids +from hub.services.models import cache_inventory, downloads, gguf_variants +from hub.services.models.common import ( + _is_gguf_filename, + _is_main_gguf_filename, + _is_mmproj_filename, +) + +logger = get_logger(__name__) + + +def _snapshot_blob_reference_counts(repo_dir: Optional[Path]) -> dict[Path, int]: + """Map each blob's realpath to its live snapshot symlink count, so per-variant deletion never unlinks a blob another revision still references (call after the target variant's own symlinks are removed).""" + counts: dict[Path, int] = {} + if repo_dir is None: + return counts + snapshots = repo_dir / "snapshots" + if not snapshots.is_dir(): + return counts + try: + entries = list(snapshots.rglob("*")) + except OSError: + return counts + for link in entries: + try: + if not link.is_symlink(): + continue + target = link.resolve() + except OSError: + continue + counts[target] = counts.get(target, 0) + 1 + return counts + + +def _blob_hash_from_path(blob: Path) -> Optional[str]: + name = blob.name + if not name or name.endswith(INCOMPLETE_SUFFIX): + return None + return name + + +def _path_exists_or_symlink(path: Path) -> bool: + try: + return path.is_symlink() or path.exists() + except OSError: + return False + + +def _repo_file_matches(target_repo, predicate) -> list[tuple[Path, Optional[Path], str]]: + matches: list[tuple[Path, Optional[Path], str]] = [] + for rev in getattr(target_repo, "revisions", ()): + for f in getattr(rev, "files", ()): + name = str(getattr(f, "file_name", "")) + if not predicate(name): + continue + file_path = getattr(f, "file_path", None) + if not file_path: + continue + blob_path = getattr(f, "blob_path", None) + matches.append( + ( + Path(file_path), + Path(blob_path) if blob_path else None, + name, + ) + ) + return matches + + +def _has_remaining_main_gguf(target_repo) -> bool: + return any( + _path_exists_or_symlink(snap) + for snap, _blob, _name in _repo_file_matches( + target_repo, + _is_main_gguf_filename, + ) + ) + + +def _delete_gguf_variant_from_repos( + repo_id: str, + variant: str, + target_repos: list, + hf_token: Optional[str], + *, + sibling_active: bool = False, +) -> dict: + failures: list[str] = [] + removed_snapshots = 0 + deleted_bytes = 0 + deleted_blobs = 0 + completed_hashes: set[str] = set() + + for target_repo in target_repos: + repo_dir = Path(target_repo.repo_path) if getattr(target_repo, "repo_path", None) else None + matched = _repo_file_matches( + target_repo, + lambda name: _is_main_gguf_filename(name) + and extract_quant_label(name).lower() == variant.lower(), + ) + + for snap, _blob, name in matched: + try: + if _path_exists_or_symlink(snap): + snap.unlink() + removed_snapshots += 1 + except OSError as e: + failures.append(f"{name}: {e}") + + companion_matches: list[tuple[Path, Optional[Path], str]] = [] + if matched and not sibling_active and not _has_remaining_main_gguf(target_repo): + companion_matches = _repo_file_matches( + target_repo, + lambda name: _is_gguf_filename(name) and _is_mmproj_filename(name), + ) + for snap, _blob, name in companion_matches: + try: + if _path_exists_or_symlink(snap): + snap.unlink() + removed_snapshots += 1 + except OSError as e: + failures.append(f"{name}: {e}") + + ref_counts = _snapshot_blob_reference_counts(repo_dir) + seen_blobs: set[Path] = set() + for _snap, blob, name in [*matched, *companion_matches]: + if blob is None: + continue + blob_hash = _blob_hash_from_path(blob) + if blob_hash: + completed_hashes.add(blob_hash) + try: + blob_key = blob.resolve() + except OSError: + blob_key = blob + if blob_key in seen_blobs: + continue + seen_blobs.add(blob_key) + if ref_counts.get(blob_key, 0) > 0: + continue + try: + if blob.exists(): + deleted_bytes += blob.stat().st_size + blob.unlink() + deleted_blobs += 1 + except OSError as e: + failures.append(f"{name}: {e}") + + if failures: + raise HTTPException( + status_code = 409, + detail = ( + f"Couldn't fully delete {variant} for {repo_id}: " + f"{len(failures)} file(s) are in use. " + "Unload the model and try again." + ), + ) + + incomplete_result = gguf_variants.delete_variant_incomplete_blobs_result( + repo_id, + variant, + hf_token, + extra_hashes = frozenset(completed_hashes), + companions = not sibling_active, + ) + if incomplete_result.unresolved: + raise HTTPException( + status_code = 409, + detail = ( + f"Couldn't fully delete {variant} for {repo_id}: partial " + "download bytes exist but this variant's blob hashes are unavailable. " + "Reconnect or provide access to the repo, then try again." + ), + ) + + state_purged = download_manifest.purge_state("model", repo_id, variant) + if ( + removed_snapshots == 0 + and deleted_blobs == 0 + and incomplete_result.deleted == 0 + and not state_purged + ): + raise HTTPException( + status_code = 404, + detail = f"Variant {variant} not found in cache for {repo_id}", + ) + + freed_mb = deleted_bytes / (1024 * 1024) + logger.info( + f"Deleted {removed_snapshots} file(s) for {repo_id} variant {variant}: " + f"{freed_mb:.1f} MB freed" + ) + return {"status": "deleted", "repo_id": repo_id, "variant": variant} + + +def _loaded_id_matches_repo(loaded_id: str, repo_id: str) -> bool: + """True when *loaded_id* is *repo_id* or a file within it; ``/``-boundary aware so ``org/model`` doesn't match sibling ``org/model-v2``.""" + rid = repo_id.lower() + lid = loaded_id.lower() + return lid == rid or lid.startswith(f"{rid}/") + + +def _loaded_repo_variant_blocks_delete( + loaded_id: str, repo_id: str, delete_variant: Optional[str], loaded_variant: Optional[str] +) -> bool: + if not _loaded_id_matches_repo(loaded_id, repo_id): + return False + if not delete_variant: + return True + if not loaded_variant: + return True + return loaded_variant.lower() == delete_variant.lower() + + +_LOAD_STATE_UNVERIFIABLE_DETAIL = ( + "Couldn't verify whether this model is still loaded for inference. " + "Unload it if it is active, then try deleting again." +) + + +def _llama_cpp_blocks_delete(repo_id: str, variant: Optional[str]) -> bool: + """Whether the llama.cpp backend holds *repo_id* (/variant). Acquiring fails open (import error means nothing loaded); reading load state is unguarded so a raise propagates and the caller fails closed rather than delete a live model.""" + try: + from routes.inference import get_llama_cpp_backend + backend = get_llama_cpp_backend() + except Exception as e: + logger.debug(f"llama.cpp backend unavailable during delete guard for {repo_id}: {e}") + return False + loaded_id = backend.model_identifier + loaded_variant = getattr(backend, "hf_variant", None) + if backend.is_active and not backend.is_loaded and loaded_id: + return _loaded_repo_variant_blocks_delete( + loaded_id, + repo_id, + variant, + loaded_variant, + ) + if backend.is_loaded and loaded_id: + return _loaded_repo_variant_blocks_delete( + loaded_id, + repo_id, + variant, + loaded_variant, + ) + return False + + +def _inference_backend_blocks_delete(repo_id: str) -> bool: + """Whether the subprocess inference backend holds *repo_id*; same fail-open-on-acquire / surface-on-query contract as :func:`_llama_cpp_blocks_delete`.""" + try: + from core.inference import get_inference_backend + backend = get_inference_backend() + except Exception as e: + logger.debug(f"Inference backend unavailable during delete guard for {repo_id}: {e}") + return False + active_name = backend.active_model_name + return bool(active_name) and _loaded_id_matches_repo(active_name, repo_id) + + +async def delete_cached_model_response( + repo_id: str, + variant: Optional[str] = None, + hf_token: Optional[str] = None, +): + """Delete a cached model repo (or a specific GGUF variant) from the HF cache. + + When *variant* is provided, only the GGUF files matching that quant label + are removed (e.g. ``UD-Q4_K_XL``). Otherwise the entire repo is deleted. + Refuses if the model is currently loaded for inference. + """ + if not _is_valid_repo_id(repo_id): + raise HTTPException(status_code = 400, detail = "Invalid repo_id format") + variant = (variant or "").strip() or None + if variant is not None and not _is_valid_gguf_variant(variant): + raise HTTPException( + status_code = 400, + detail = f"Invalid gguf_variant: {variant!r}", + ) + + # Guard fails closed: if a live backend's load state can't be read, abort + # with 503 rather than risk unlinking weights under a running process. + try: + blocks_delete = _llama_cpp_blocks_delete(repo_id, variant) or ( + _inference_backend_blocks_delete(repo_id) + ) + except Exception as e: + logger.warning(f"Load-state verification failed for {repo_id}; refusing delete: {e}") + raise HTTPException( + status_code = 503, + detail = _LOAD_STATE_UNVERIFIABLE_DETAIL, + ) + if blocks_delete: + raise HTTPException( + status_code = 400, + detail = "Unload the model before deleting", + ) + + repo_key = await asyncio.to_thread(resolve_cached_repo_id_case, repo_id, repo_type = "model") + if not downloads.registry.begin_delete(repo_key, variant): + detail = ( + f"Cancel the {variant} download before deleting it." + if variant is not None + else "Cancel the active downloads before deleting." + ) + raise HTTPException(status_code = 400, detail = detail) + try: + return await asyncio.to_thread(_delete_cached_model_blocking, repo_id, variant, hf_token) + finally: + downloads.registry.end_delete(repo_key, variant) + cache_inventory.invalidate_hf_cache_scans() + + +def _delete_cached_model_blocking( + repo_id: str, variant: Optional[str], hf_token: Optional[str] +) -> dict: + try: + # If a sibling quant is downloading concurrently, restrict this delete to + # the variant's own files and leave the shared mmproj companion for it. + sibling_active = bool( + variant and downloads.registry.has_active_peer_variant(repo_id, variant) + ) + + cache_scans = cache_inventory.all_hf_cache_scans() + + candidate_entries = [] + for hf_cache in cache_scans: + for repo_info in hf_cache.repos: + if str(repo_info.repo_type) != "model": + continue + if repo_info.repo_id.lower() == repo_id.lower(): + candidate_entries.append((hf_cache, repo_info)) + + matched_repo_ids = resolve_destructive_repo_ids( + repo_id, + [str(repo_info.repo_id) for _hf_cache, repo_info in candidate_entries], + noun = "models", + ) + target_entries = [ + (hf_cache, repo_info) + for hf_cache, repo_info in candidate_entries + if str(repo_info.repo_id) in matched_repo_ids + ] + + if not target_entries: + if variant is None: + cache_purged = purge_repo_cache_dirs("model", repo_id) or purge_partial_repo( + "model", repo_id + ) + state_purged = download_manifest.purge_all_state_for_repo("model", repo_id) > 0 + if cache_purged or state_purged: + return {"status": "deleted", "repo_id": repo_id} + if variant: + incomplete_result = gguf_variants.delete_variant_incomplete_blobs_result( + repo_id, + variant, + hf_token, + companions = not sibling_active, + ) + if incomplete_result.unresolved: + raise HTTPException( + status_code = 409, + detail = ( + f"Couldn't fully delete {variant} for {repo_id}: partial " + "download bytes exist but this variant's blob hashes are unavailable. " + "Reconnect or provide access to the repo, then try again." + ), + ) + state_purged = download_manifest.purge_state( + "model", + repo_id, + variant, + ) + if incomplete_result.deleted > 0 or state_purged: + return { + "status": "deleted", + "repo_id": repo_id, + "variant": variant, + } + raise HTTPException(status_code = 404, detail = "Model not found in cache") + + if variant: + return _delete_gguf_variant_from_repos( + repo_id, + variant, + [repo for _cache, repo in target_entries], + hf_token, + sibling_active = sibling_active, + ) + + deleted_revisions = False + for hf_cache, repo_info in target_entries: + revision_hashes = [ + rev.commit_hash for rev in repo_info.revisions if getattr(rev, "commit_hash", None) + ] + if not revision_hashes: + continue + delete_strategy = hf_cache.delete_revisions(*revision_hashes) + logger.info( + f"Deleting cached model {repo_id} from " + f"{getattr(hf_cache, 'cache_dir', '')}: " + f"{delete_strategy.expected_freed_size_str} will be freed" + ) + delete_strategy.execute() + deleted_revisions = True + + cache_purged = purge_repo_cache_dirs("model", repo_id) + partial_purged = purge_partial_repo("model", repo_id) + state_purged = download_manifest.purge_all_state_for_repo("model", repo_id) > 0 + + if not (deleted_revisions or cache_purged or partial_purged or state_purged): + raise HTTPException(status_code = 404, detail = "No revisions found for model") + + return {"status": "deleted", "repo_id": repo_id} + + except HTTPException: + raise + except Exception as e: + logger.error( + "Error deleting cached model %s: %s", + repo_id, + download_registry.scrub_secrets(str(e), hf_token = hf_token), + ) + raise HTTPException( + status_code = 500, + detail = "Failed to delete cached model: " + + download_registry.scrub_secrets(str(e), hf_token = hf_token), + ) diff --git a/studio/backend/hub/services/models/downloads.py b/studio/backend/hub/services/models/downloads.py new file mode 100644 index 0000000000..db95b82c95 --- /dev/null +++ b/studio/backend/hub/services/models/downloads.py @@ -0,0 +1,411 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Download orchestration.""" + +from __future__ import annotations + +import asyncio +from typing import Optional, TYPE_CHECKING + +from fastapi import HTTPException +from loggers import get_logger + +from hub.schemas.downloads import ( + ActiveDownloadsResponse, + CancelDownloadRequest, + DownloadJobStatus, + DownloadModelRequest, +) +from hub.utils import download_registry +from hub.utils import download_manifest +from hub.utils import inventory_scan as hf_cache_scan +from hub.utils.hf_cache_state import has_active_incomplete_blobs +from hub.utils.paths import ( + is_valid_gguf_variant as _is_valid_gguf_variant, + is_valid_repo_id as _is_valid_repo_id, + resolve_cached_repo_id_case, +) +from hub.services import snapshot_progress +from hub.services import download_lifecycle +from hub.services.models import cache_inventory, gguf_variants + +logger = get_logger(__name__) + +if TYPE_CHECKING: + import subprocess + +_registry = download_registry.get_models_registry() + + +def _download_job_key(repo_id: str, variant: Optional[str]) -> str: + return download_registry.normalize_job_key( + f"{download_registry.normalize_repo_key(repo_id)}::{variant or ''}" + ) + + +def _job_status( + key: str, + *, + repo_id: Optional[str] = None, + variant: Optional[str] = None, +) -> DownloadJobStatus: + state, error, generation = download_lifecycle.idle_status( + _registry, + key, + repo_type = "model", + repo_id = repo_id, + variant = variant, + ) + return DownloadJobStatus(state = state, error = error, generation = generation) + + +def _spawn_download_worker( + repo_id: str, + variant: Optional[str], + hf_token: Optional[str], + use_xet: bool = False, + protected_blob_hashes: Optional[frozenset[str]] = None, +) -> subprocess.Popen: + args = ["--repo-id", repo_id] + if variant: + args.extend(["--variant", variant]) + return download_lifecycle.spawn_worker( + args, + hf_token, + use_xet = use_xet, + protected_blob_hashes = protected_blob_hashes, + ) + + +async def download_model_response(body: DownloadModelRequest, hf_token: Optional[str] = None): + """Start a background download for a HuggingFace model.""" + repo_id = body.repo_id.strip() + if not _is_valid_repo_id(repo_id): + raise HTTPException( + status_code = 400, + detail = f"Invalid repo_id: {repo_id!r}", + ) + # Canonicalize so two different-cased paste-ins share one job + cache dir. + repo_id = await asyncio.to_thread(resolve_cached_repo_id_case, repo_id, repo_type = "model") + + variant = (body.gguf_variant or "").strip() or None + if variant is not None and not _is_valid_gguf_variant(variant): + raise HTTPException( + status_code = 400, + detail = f"Invalid gguf_variant: {variant!r}", + ) + key = _download_job_key(repo_id, variant) + transport = download_lifecycle.resolve_transport(body.use_xet) + variant_blob_hashes = frozenset() + variant_progress_blob_hashes = frozenset() + completed_baseline_bytes = 0 + if variant is not None: + try: + variant_blob_hashes = await asyncio.to_thread( + gguf_variants.gguf_variant_blob_hashes, + repo_id, + variant, + hf_token, + include_companions = False, + ) + variant_progress_blob_hashes = await asyncio.to_thread( + gguf_variants.gguf_variant_blob_hashes, + repo_id, + variant, + hf_token, + include_companions = True, + ) + except Exception as e: + logger.warning( + "GGUF hash pre-resolution failed for %s [%s]; continuing without " + "a completed-bytes baseline or peer-protection hashes (the worker " + "re-resolves its own blobs before purging): %s", + repo_id, + variant, + download_registry.scrub_secrets(str(e), hf_token = hf_token), + ) + has_variant_resume_state = ( + download_manifest.has_cancel_marker("model", repo_id, variant) + or download_manifest.read_manifest("model", repo_id, variant) is not None + ) + if variant_progress_blob_hashes and not has_variant_resume_state: + completed_baseline_bytes = await asyncio.to_thread( + download_registry.completed_blob_bytes, + "model", + repo_id, + variant_progress_blob_hashes, + ) + + claimed, claim_state = _registry.claim( + key, + transport, + repo_type = "model", + repo_id = repo_id, + variant = variant, + blob_hashes = variant_blob_hashes, + progress_blob_hashes = variant_progress_blob_hashes, + completed_baseline_bytes = completed_baseline_bytes, + ) + generation = _registry.current_generation(key) + if not claimed: + # claim_state is the blocking job's state. The client can attach only + # when the blocker is this key's own in-flight job (adoptable); a + # cross-variant conflict or in-progress delete is not accepted. + return { + "job_key": key, + "state": claim_state, + "accepted": _registry.adoptable(key), + "generation": generation, + } + download_manifest.clear_cancel_marker("model", repo_id, variant) + # Blobs a concurrent same-repo variant is already writing (e.g. a shared + # mmproj). The worker must not purge these during cache preparation. + protected_blob_hashes = _registry.peer_blob_hashes(key) if variant else frozenset() + + label = f"{repo_id}{f' [{variant}]' if variant else ''}" + state = download_lifecycle.launch_worker( + _registry, + key, + spawn = lambda: _spawn_download_worker( + repo_id, + variant, + hf_token, + use_xet = body.use_xet, + protected_blob_hashes = protected_blob_hashes, + ), + hf_token = hf_token, + label = label, + log_prefix = "Download", + logger = logger, + repo_type = "model", + repo_id = repo_id, + transport = transport, + watch_name = f"hf-download-watch-{repo_id}", + ) + + return { + "job_key": key, + "state": state, + "accepted": True, + "generation": generation, + } + + +async def cancel_download_model_response(body: CancelDownloadRequest): + """Cancel an in-flight model download (SIGKILL; HF cache resumes on next download).""" + repo_id = body.repo_id.strip() + if not _is_valid_repo_id(repo_id): + raise HTTPException( + status_code = 400, + detail = f"Invalid repo_id: {repo_id!r}", + ) + repo_id = await asyncio.to_thread(resolve_cached_repo_id_case, repo_id, repo_type = "model") + variant = (body.gguf_variant or "").strip() or None + if variant is not None and not _is_valid_gguf_variant(variant): + raise HTTPException( + status_code = 400, + detail = f"Invalid gguf_variant: {variant!r}", + ) + key = _download_job_key(repo_id, variant) + + state = download_lifecycle.cancel_worker( + _registry, + key, + generation = body.generation, + label = repo_id, + logger = logger, + ) + return {"job_key": key, "state": state} + + +async def get_download_status_response(repo_id: str, gguf_variant: str = "") -> DownloadJobStatus: + """Return the latest state of a background download job.""" + repo_id = repo_id.strip() + if not _is_valid_repo_id(repo_id): + return DownloadJobStatus(state = "idle") + repo_id = await asyncio.to_thread(resolve_cached_repo_id_case, repo_id, repo_type = "model") + variant = (gguf_variant or "").strip() or None + key = _download_job_key(repo_id, variant) + return _job_status(key, repo_id = repo_id, variant = variant) + + +async def get_active_downloads_response(repo_id: str = "") -> ActiveDownloadsResponse: + """Return every in-flight download for a repo in a single call.""" + repo_id = repo_id.strip() + if repo_id and not _is_valid_repo_id(repo_id): + return ActiveDownloadsResponse(downloads = []) + canonical_repo_id = ( + await asyncio.to_thread(resolve_cached_repo_id_case, repo_id, repo_type = "model") + if repo_id + else None + ) + return ActiveDownloadsResponse( + downloads = download_lifecycle.active_download_refs( + _registry, + canonical_repo_id, + with_variant = True, + ) + ) + + +def _variant_transport_status(repo_id: str, variant: str, hf_token: Optional[str]) -> dict: + incomplete_hashes = download_registry.incomplete_blob_hashes( + "model", + repo_id, + active_only = True, + ) + variant_hashes = gguf_variants.gguf_variant_blob_hashes( + repo_id, + variant, + hf_token, + allow_remote = False, + ) + has_partial = hf_cache_scan.is_variant_partial( + repo_id, + variant, + incomplete_blob_hashes = incomplete_hashes, + variant_blob_hashes = variant_hashes, + ) + last_transport = hf_cache_scan.partial_transport_for("model", repo_id, variant) + if ( + last_transport is None + and has_partial + and incomplete_hashes + and variant_hashes + and incomplete_hashes.intersection(variant_hashes) + ): + last_transport = download_registry.read_active_transport_marker( + "model", + repo_id, + variant, + ) + has_matching_incomplete = bool( + incomplete_hashes and variant_hashes and incomplete_hashes.intersection(variant_hashes) + ) + return { + "has_partial": has_partial, + "last_transport": last_transport, + "resumable": ( + has_matching_incomplete and last_transport == download_registry.TRANSPORT_HTTP + ), + } + + +async def get_model_transport_status_response( + repo_id: str, + gguf_variant: str = "", + hf_token: Optional[str] = None, +) -> dict: + """Return last transport used for this repo + whether any partial blobs + exist + whether that partial supports byte-level resume. + + ``resumable`` is True only when an HTTP partial exists. XET partials + are reported via ``has_partial`` but always have ``resumable=False`` + because ``hf_xet`` rewrites the destination from scratch on every + call (network resume happens transparently via its chunk cache). + """ + repo_id = repo_id.strip() + if not _is_valid_repo_id(repo_id): + return {"has_partial": False, "last_transport": None, "resumable": False} + variant = (gguf_variant or "").strip() + if variant: + if not _is_valid_gguf_variant(variant): + return {"has_partial": False, "last_transport": None, "resumable": False} + return _variant_transport_status(repo_id, variant, hf_token) + return { + "has_partial": has_active_incomplete_blobs("model", repo_id), + "last_transport": download_registry.read_active_transport_marker("model", repo_id), + "resumable": download_registry.is_resumable_partial("model", repo_id), + } + + +async def get_gguf_download_progress_response( + repo_id: str, + variant: str = "", + expected_bytes: int = 0, + hf_token: Optional[str] = None, +) -> dict: + """Return download progress for a specific GGUF variant.""" + expected_total = max(expected_bytes, 0) + progress_variant = variant.strip() or None + if progress_variant is not None and not _is_valid_gguf_variant(progress_variant): + return { + "downloaded_bytes": 0, + "completed_bytes": 0, + "complete_on_disk": False, + "expected_bytes": expected_total, + "progress": 0, + "cache_path": None, + } + + def _metadata_resolver( + resolved_repo_id: str, token: Optional[str] + ) -> tuple[int, frozenset[str]]: + if progress_variant is None: + return expected_total, frozenset() + requirement = gguf_variants.gguf_variant_requirements( + resolved_repo_id, + progress_variant, + token, + ) + if requirement is not None: + return requirement.download_size_bytes, requirement.required_hashes + manifest = download_manifest.read_manifest( + "model", + resolved_repo_id, + progress_variant, + ) + if manifest is not None: + return ( + sum(max(0, int(file.size or 0)) for file in manifest.expected_files), + frozenset(file.sha256 for file in manifest.expected_files if file.sha256), + ) + return ( + expected_total, + gguf_variants.gguf_variant_blob_hashes( + resolved_repo_id, + progress_variant, + token, + allow_remote = False, + ), + ) + + return await snapshot_progress.snapshot_progress_response( + repo_type = "model", + repo_id = repo_id, + job_key = _download_job_key(repo_id, progress_variant), + expected_bytes = expected_total, + hf_token = hf_token, + registry = _registry, + metadata_resolver = _metadata_resolver, + variant = progress_variant, + ) + + +async def get_download_progress_response( + repo_id: str, + expected_bytes: int = 0, + hf_token: Optional[str] = None, +) -> dict: + """Return download progress for any HuggingFace model repo. + + Checks the local HF cache for completed blobs and in-progress + (.incomplete) downloads. Uses the caller-supplied expected total + when available; otherwise queries HF metadata and caches it. + Also returns ``cache_path``: the realpath of the snapshot directory + (or the cache repo root if no snapshot exists yet) so the UI can + show users where the weights actually live on disk. + """ + return await snapshot_progress.snapshot_progress_response( + repo_type = "model", + repo_id = repo_id, + job_key = _download_job_key(repo_id, None), + expected_bytes = expected_bytes, + hf_token = hf_token, + registry = _registry, + metadata_resolver = cache_inventory.get_repo_snapshot_metadata_cached, + ) + + +registry = _registry diff --git a/studio/backend/hub/services/models/folder_browser.py b/studio/backend/hub/services/models/folder_browser.py new file mode 100644 index 0000000000..9b0b46509b --- /dev/null +++ b/studio/backend/hub/services/models/folder_browser.py @@ -0,0 +1,518 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Model folder recommendation and browsing services.""" + +from __future__ import annotations + +import os +from pathlib import Path +from typing import Optional + +from fastapi import HTTPException +from loggers import get_logger + +from hub.schemas.inventory import BrowseEntry, BrowseFoldersResponse +from hub.storage.scan_folders import ( + contains_sensitive_path_component, + list_scan_folders, +) +from hub.utils.paths import ( + exports_root, + hf_default_cache_dir, + legacy_hf_cache_dir, + lmstudio_model_dirs, + normalize_path, + outputs_root, + studio_root, + well_known_model_dirs, +) +from hub.services.models.common import _safe_is_dir +from hub.services.models.local_inventory import _resolve_hf_cache_dir + +logger = get_logger(__name__) + + +def get_recommended_folders_response() -> dict: + """Return well-known model directories that exist on this machine.""" + folders: list[str] = [] + seen: set[str] = set() + + def _add(p: Optional[Path]) -> None: + if p is None: + return + try: + resolved = str(p.resolve()) + except OSError: + return + if resolved in seen: + return + if _safe_is_dir(resolved) and os.access(resolved, os.R_OK | os.X_OK): + seen.add(resolved) + folders.append(resolved) + + try: + for p in lmstudio_model_dirs(): + _add(p) + except Exception as e: + logger.warning("Failed to scan for LM Studio model directories: %s", e) + + ollama_env = os.environ.get("OLLAMA_MODELS") + if ollama_env: + _add(Path(normalize_path(ollama_env)).expanduser()) + for candidate in ( + Path.home() / ".ollama" / "models", + Path("/usr/share/ollama/.ollama/models"), + Path("/var/lib/ollama/.ollama/models"), + ): + _add(candidate) + + return {"folders": folders} + + +# Ceiling on children to stat when guessing if a directory holds models. +_BROWSE_MODEL_HINT_PROBE = 64 +# Hard cap on returned subdirectory entries so pointing at ``/usr/lib`` or +# ``/proc`` can't stat-storm the process or flood the client. +_BROWSE_ENTRY_CAP = 2000 + + +def _count_model_files(directory: Path, cap: int = 200) -> int: + """Count GGUF/safetensors files immediately inside *directory*, bounded by visited entries (not matches) so the hint never costs more than ``cap`` stats.""" + n = 0 + visited = 0 + try: + for f in directory.iterdir(): + visited += 1 + if visited > cap: + break + try: + if f.is_file(): + low = f.name.lower() + if low.endswith((".gguf", ".safetensors")): + n += 1 + except OSError: + continue + except PermissionError as e: + logger.debug("browse-folders: permission denied counting %s: %s", directory, e) + return 0 + except OSError as e: + logger.debug("browse-folders: OS error counting %s: %s", directory, e) + return 0 + return n + + +def _has_direct_model_signal(directory: Path) -> bool: + """True if an immediate child signals a model (GGUF/safetensors/config file or ``models--*`` HF-cache subdir); bounded by the hint probe.""" + try: + it = directory.iterdir() + except OSError: + return False + try: + for i, child in enumerate(it): + if i >= _BROWSE_MODEL_HINT_PROBE: + break + try: + name = child.name + if child.is_file(): + low = name.lower() + if low.endswith((".gguf", ".safetensors")): + return True + if low in ("config.json", "adapter_config.json"): + return True + elif child.is_dir() and name.startswith("models--"): + return True + except OSError: + continue + except OSError: + return False + return False + + +def _looks_like_model_dir(directory: Path) -> bool: + """Bounded heuristic flagging dirs worth exploring (false negatives are fine; the scanner is authoritative). Three signals, cheapest first: a ``models--*`` name, a direct child signal, or a grandchild signal (LM Studio / Ollama ``publisher/model/weights.gguf`` layout).""" + if directory.name.startswith("models--"): + return True + if _has_direct_model_signal(directory): + return True + try: + it = directory.iterdir() + except OSError: + return False + try: + for i, child in enumerate(it): + if i >= _BROWSE_MODEL_HINT_PROBE: + break + try: + if not child.is_dir(): + continue + except OSError: + continue + if child.name.startswith("models--"): + return True + if _has_direct_model_signal(child): + return True + except OSError: + return False + return False + + +def _build_browse_allowlist() -> list[Path]: + """Root directories the browser may walk (also seeds the suggestion chips): HOME, resolved HF cache dirs, Studio outputs/exports/root, registered scan folders, and well-known local-LLM dirs. Each is added only if it resolves to a real directory so the sandbox has no dead boundary.""" + from hub.storage.scan_folders import list_scan_folders + + candidates: list[Path] = [] + + def _add(p: Optional[Path | str]) -> None: + if p is None: + return + try: + p = Path(normalize_path(str(p))).expanduser() + resolved = p.resolve() + except (OSError, RuntimeError, ValueError): + return + if _safe_is_dir(resolved): + candidates.append(resolved) + + _add(Path.home()) + _add(_resolve_hf_cache_dir()) + try: + _add(hf_default_cache_dir()) + except Exception: # noqa: BLE001 -- best-effort + pass + try: + _add(legacy_hf_cache_dir()) + except Exception: # noqa: BLE001 -- best-effort + pass + try: + _add(studio_root()) + _add(outputs_root()) + _add(exports_root()) + except Exception as exc: # noqa: BLE001 -- best-effort + logger.debug("browse-folders: studio roots unavailable: %s", exc) + try: + for folder in list_scan_folders(): + p = folder.get("path") + if p: + _add(p) + except Exception as exc: # noqa: BLE001 -- best-effort + logger.debug("browse-folders: could not load scan folders: %s", exc) + try: + for p in well_known_model_dirs(): + _add(p) + except Exception as exc: # noqa: BLE001 -- best-effort + logger.debug("browse-folders: well-known dirs unavailable: %s", exc) + + seen: set[str] = set() + deduped: list[Path] = [] + for p in candidates: + key = os.path.normcase(os.path.realpath(str(p))) + if key in seen: + continue + seen.add(key) + deduped.append(p) + return deduped + + +def _is_path_inside_allowlist(target: Path, allowed_roots: list[Path]) -> bool: + """True if *target* equals or descends from any allowed root; uses ``os.path.realpath`` so symlinks cannot escape the sandbox.""" + try: + target_real = os.path.normcase(os.path.realpath(str(target))) + except OSError: + return False + for root in allowed_roots: + try: + root_real = os.path.normcase(os.path.realpath(str(root))) + except OSError: + continue + try: + if os.path.commonpath([target_real, root_real]) == root_real: + return True + except ValueError: + continue + if target_real == root_real: + return True + return False + + +def _normalize_browse_request_path(path: Optional[str], *, relative_root: Path) -> str: + """Normalize the browse request path lexically, without touching the FS.""" + if path is None or not path.strip(): + return os.path.normpath(str(Path.home())) + + expanded = os.path.expanduser(normalize_path(path.strip())) + if not os.path.isabs(expanded): + expanded = os.path.join(str(relative_root), expanded) + return os.path.normpath(expanded) + + +def _browse_relative_parts(requested_path: str, root: Path) -> Optional[list[str]]: + if "\x00" in requested_path: + raise HTTPException( + status_code = 400, + detail = "Path cannot contain null bytes", + ) + root_text = os.path.normcase(os.path.normpath(str(root))) + requested_text = os.path.normcase(os.path.normpath(requested_path)) + try: + rel_text = os.path.relpath(requested_text, root_text) + except ValueError: + return None + + if rel_text == ".": + return [] + if rel_text == ".." or rel_text.startswith(f"..{os.sep}"): + return None + + parts = [part for part in rel_text.split(os.sep) if part not in ("", ".")] + altsep = os.altsep + for part in parts: + if part == ".." or "\x00" in part or os.sep in part or (altsep and altsep in part): + return None + return parts + + +def _match_browse_child(current: Path, name: str) -> Optional[Path]: + """Immediate child named ``name`` under ``current``, or None. ``name`` is pre-validated as a safe single component, so the join is O(1); case resolution follows OS filesystem semantics.""" + child = current / name + try: + child.stat() + except (FileNotFoundError, NotADirectoryError): + return None + except PermissionError: + raise HTTPException( + status_code = 403, + detail = f"Permission denied reading {current}", + ) from None + except OSError as exc: + raise HTTPException( + status_code = 500, + detail = f"Could not read {current}: {exc}", + ) from exc + return child + + +def _resolve_browse_target(path: Optional[str], allowed_roots: list[Path]) -> Path: + """Resolve a requested browse path by walking from trusted allowlist roots.""" + requested_path = _normalize_browse_request_path(path, relative_root = Path.home()) + resolved_roots: list[Path] = [] + seen_roots: set[str] = set() + for root in sorted(allowed_roots, key = lambda p: len(str(p)), reverse = True): + try: + resolved = root.resolve() + except OSError: + continue + key = os.path.normcase(os.path.realpath(str(resolved))) + if key in seen_roots: + continue + seen_roots.add(key) + resolved_roots.append(resolved) + + for root in resolved_roots: + parts = _browse_relative_parts(requested_path, root) + if parts is None: + continue + + current = root + for part in parts: + child = _match_browse_child(current, part) + if child is None: + raise HTTPException( + status_code = 404, + detail = f"Path does not exist: {requested_path}", + ) + try: + resolved_child = child.resolve() + except OSError as exc: + raise HTTPException( + status_code = 400, + detail = f"Invalid path: {exc}", + ) from exc + if not _is_path_inside_allowlist(resolved_child, resolved_roots): + raise HTTPException( + status_code = 403, + detail = ( + "Path is not in the browseable allowlist. Register it via " + "POST /api/hub/scan-folders first, or pick a directory " + "under your home folder." + ), + ) + # HOME is in the allowlist, so without this denylist (same one + # registration enforces) a user could browse into ~/.ssh, ~/.aws, etc. + if contains_sensitive_path_component(str(resolved_child)): + raise HTTPException( + status_code = 403, + detail = "Credential or configuration directories are not browseable.", + ) + current = resolved_child + + if not current.is_dir(): + raise HTTPException( + status_code = 400, + detail = f"Not a directory: {current}", + ) + return current + + raise HTTPException( + status_code = 403, + detail = ( + "Path is not in the browseable allowlist. Register it via " + "POST /api/hub/scan-folders first, or pick a directory " + "under your home folder." + ), + ) + + +def browse_folders_response( + path: Optional[str] = None, show_hidden: bool = False +) -> BrowseFoldersResponse: + """List immediate subdirectories of *path* for the Custom Folders picker. + + Requests are bounded to the :func:`_build_browse_allowlist` roots; paths + outside it return 403 (symlinks resolved via realpath first, so traversal + can't escape). Sorting: model-bearing dirs first, then plain, then hidden. + """ + from hub.storage.scan_folders import list_scan_folders + + # Build the allowlist once -- the sandbox check and suggestion chips share + # it so chips are always navigable. + allowed_roots = _build_browse_allowlist() + + try: + target = _resolve_browse_target(path, allowed_roots) + except HTTPException: + requested_path = _normalize_browse_request_path( + path, + relative_root = Path.home(), + ) + if path is not None and path.strip(): + logger.warning( + "browse-folders: rejected path %r (normalized=%s)", + path, + requested_path, + ) + raise + + # Enumerate immediate subdirectories with a bounded cap so a stray + # query against ``/usr/lib`` or ``/proc`` can't stat-storm the process. + entries: list[BrowseEntry] = [] + truncated = False + visited = 0 + try: + it = target.iterdir() + except PermissionError: + raise HTTPException( + status_code = 403, + detail = f"Permission denied reading {target}", + ) + except OSError as exc: + raise HTTPException( + status_code = 500, + detail = f"Could not read {target}: {exc}", + ) + + try: + for child in it: + # Bound by visited entries, not appended ones, so a directory full + # of files still caps work at ``_BROWSE_ENTRY_CAP`` stats. + visited += 1 + if visited > _BROWSE_ENTRY_CAP: + truncated = True + break + try: + if not child.is_dir(): + continue + except OSError: + continue + name = child.name + is_hidden = name.startswith(".") + if is_hidden and not show_hidden: + continue + # Don't surface credential/config dirs even with show_hidden: + # descending into them is refused and registration rejects them. + if contains_sensitive_path_component(name): + continue + entries.append( + BrowseEntry( + name = name, + has_models = _looks_like_model_dir(child), + hidden = is_hidden, + ) + ) + except PermissionError as exc: + logger.debug( + "browse-folders: permission denied during enumeration of %s: %s", + target, + exc, + ) + except OSError as exc: + # Rare: iterdir succeeded but reading a specific entry failed. + logger.warning("browse-folders: partial enumeration of %s: %s", target, exc) + + # Model-bearing dirs first, then plain, then hidden; case-insensitive + # alphabetical within each bucket. + def _sort_key(e: BrowseEntry) -> tuple[int, str]: + bucket = 0 if e.has_models else (2 if e.hidden else 1) + return (bucket, e.name.lower()) + + entries.sort(key = _sort_key) + + # Parent is None at the FS root and when it would step outside the sandbox, + # so the up-row never 403s on click. + parent: Optional[str] + if target.parent == target or not _is_path_inside_allowlist(target.parent, allowed_roots): + parent = None + else: + parent = str(target.parent) + + # Handy starting points for the quick-pick chips. + suggestions: list[str] = [] + seen_sug: set[str] = set() + + def _add_sug(p: Optional[Path | str]) -> None: + if p is None: + return + try: + p = Path(normalize_path(str(p))).expanduser() + resolved = str(p.resolve()) + except (OSError, RuntimeError, ValueError): + return + if resolved in seen_sug: + return + if _safe_is_dir(resolved): + seen_sug.add(resolved) + suggestions.append(resolved) + + # Home first as the safe fallback. + _add_sug(Path.home()) + # The HF cache root in use (honors HF_HOME / HF_HUB_CACHE), then the default. + try: + _add_sug(_resolve_hf_cache_dir()) + except Exception: + pass + try: + _add_sug(hf_default_cache_dir()) + except Exception: + pass + # Already-registered scan folders (what the user has curated). + try: + for folder in list_scan_folders(): + _add_sug(folder.get("path", "")) + except Exception as exc: + logger.debug("browse-folders: could not load scan folders: %s", exc) + # Well-known third-party dirs (LM Studio, Ollama, ~/models). Each helper + # only returns existing paths so we never show dead chips. + try: + for p in well_known_model_dirs(): + _add_sug(p) + except Exception as exc: + logger.debug("browse-folders: could not load well-known dirs: %s", exc) + + return BrowseFoldersResponse( + current = str(target), + parent = parent, + entries = entries, + suggestions = suggestions, + truncated = truncated, + model_files_here = _count_model_files(target), + ) diff --git a/studio/backend/hub/services/models/gguf_variants.py b/studio/backend/hub/services/models/gguf_variants.py new file mode 100644 index 0000000000..efe6bdf6c4 --- /dev/null +++ b/studio/backend/hub/services/models/gguf_variants.py @@ -0,0 +1,643 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""GGUF variant resolution.""" + +from __future__ import annotations + +import asyncio +import threading +import time +from collections import OrderedDict +from typing import NamedTuple, Optional + +from fastapi import HTTPException +from loggers import get_logger + +from hub.schemas.inventory import GgufVariantDetail, GgufVariantsResponse +from hub.utils import download_manifest +from hub.utils import download_registry +from hub.utils import inventory_scan as hf_cache_scan +from hub.utils.hf_errors import hf_error_status +from hub.utils.hf_cache_state import ( + INCOMPLETE_SUFFIX, + iter_destructive_repo_cache_dirs, +) +from hub.utils.gguf import ( + extract_quant_label, + iter_hf_cache_snapshots, + list_gguf_variants, + list_gguf_variants_from_hf_cache, + list_local_gguf_variants, + list_partial_gguf_variants_from_state, + pick_best_gguf, +) +from hub.utils.paths import ( + is_local_path, + is_valid_repo_id as _is_valid_repo_id, +) +from hub.services.models.common import ( + _is_mmproj_filename, + _iter_gguf_paths, +) +from hub.utils.gguf_plan import ( + GgufVariantPlan as _GgufVariantRequirement, + build_gguf_variant_plans, + is_main_gguf_variant_path, +) + +logger = get_logger(__name__) + +_VARIANT_HASH_CACHE: "OrderedDict[tuple[str, str, str, bool], tuple[frozenset[str], float]]" = ( + OrderedDict() +) +_VARIANT_REQUIREMENT_CACHE: "OrderedDict[tuple[str, str, str], tuple[_GgufVariantRequirement, float]]" = OrderedDict() +_VARIANT_REQUIREMENT_NEG_CACHE: "OrderedDict[tuple[str, str], float]" = OrderedDict() +_VARIANT_HASH_MAX = 512 +# Blob hashes are derived from the same mutable remote revision metadata as +# variant requirements, so they must not outlive that freshness window. +_VARIANT_HASH_POS_TTL = 60.0 +# Refresh resolved variant requirements so a moved repo revision is picked up +# within the session instead of being pinned for the backend's lifetime. +_VARIANT_REQUIREMENT_POS_TTL = 60.0 +# Suppress retries on a metadata-fetch failure so a slow/flaky link doesn't +# re-hammer the API on every page refresh. +_VARIANT_REQUIREMENT_NEG_TTL = 60.0 +# Fail fast on a slow link so the variant render isn't blocked for seconds. +_GGUF_METADATA_TIMEOUT_SECONDS = 5.0 +_VARIANT_HASH_LOCK = threading.Lock() + + +class VariantIncompleteDeleteResult(NamedTuple): + deleted: int + unresolved: bool + + +def _variant_hash_cache_key( + repo_id: str, variant: str, hf_token: Optional[str] +) -> tuple[str, str, str]: + return ( + repo_id.lower(), + variant.lower(), + hf_cache_scan.token_fingerprint(hf_token), + ) + + +def _variant_blob_hash_cache_key( + repo_id: str, variant: str, hf_token: Optional[str], include_companions: bool +) -> tuple[str, str, str, bool]: + base = _variant_hash_cache_key(repo_id, variant, hf_token) + return (*base, include_companions) + + +def _variant_repo_cache_key(repo_id: str, hf_token: Optional[str]) -> tuple[str, str]: + return (repo_id.lower(), hf_cache_scan.token_fingerprint(hf_token)) + + +def _variant_requirement_neg_cache_active(key: tuple[str, str]) -> bool: + with _VARIANT_HASH_LOCK: + cached_at = _VARIANT_REQUIREMENT_NEG_CACHE.get(key) + if cached_at is None: + return False + if (time.monotonic() - cached_at) < _VARIANT_REQUIREMENT_NEG_TTL: + _VARIANT_REQUIREMENT_NEG_CACHE.move_to_end(key) + return True + _VARIANT_REQUIREMENT_NEG_CACHE.pop(key, None) + return False + + +def _variant_requirement_neg_cache_set(key: tuple[str, str]) -> None: + with _VARIANT_HASH_LOCK: + _VARIANT_REQUIREMENT_NEG_CACHE[key] = time.monotonic() + _VARIANT_REQUIREMENT_NEG_CACHE.move_to_end(key) + while len(_VARIANT_REQUIREMENT_NEG_CACHE) > _VARIANT_HASH_MAX: + _VARIANT_REQUIREMENT_NEG_CACHE.popitem(last = False) + + +def _variant_requirement_neg_cache_clear(key: tuple[str, str]) -> None: + with _VARIANT_HASH_LOCK: + _VARIANT_REQUIREMENT_NEG_CACHE.pop(key, None) + + +def _variant_hash_cache_get(key: tuple[str, str, str, bool]) -> Optional[frozenset[str]]: + with _VARIANT_HASH_LOCK: + cached = _VARIANT_HASH_CACHE.get(key) + if cached is None: + return None + hashes, ts = cached + if (time.monotonic() - ts) >= _VARIANT_HASH_POS_TTL: + _VARIANT_HASH_CACHE.pop(key, None) + return None + _VARIANT_HASH_CACHE.move_to_end(key) + return hashes + + +def _variant_hash_cache_set(key: tuple[str, str, str, bool], hashes: frozenset[str]) -> None: + with _VARIANT_HASH_LOCK: + _VARIANT_HASH_CACHE[key] = (hashes, time.monotonic()) + _VARIANT_HASH_CACHE.move_to_end(key) + while len(_VARIANT_HASH_CACHE) > _VARIANT_HASH_MAX: + _VARIANT_HASH_CACHE.popitem(last = False) + + +def _variant_requirement_cache_get(key: tuple[str, str, str]) -> Optional[_GgufVariantRequirement]: + with _VARIANT_HASH_LOCK: + cached = _VARIANT_REQUIREMENT_CACHE.get(key) + if cached is None: + return None + requirement, ts = cached + if (time.monotonic() - ts) >= _VARIANT_REQUIREMENT_POS_TTL: + _VARIANT_REQUIREMENT_CACHE.pop(key, None) + return None + _VARIANT_REQUIREMENT_CACHE.move_to_end(key) + return requirement + + +def _variant_requirement_cache_set_many( + repo_id: str, hf_token: Optional[str], requirements: dict[str, _GgufVariantRequirement] +) -> None: + with _VARIANT_HASH_LOCK: + now = time.monotonic() + for quant, requirement in requirements.items(): + key = _variant_hash_cache_key(repo_id, quant, hf_token) + _VARIANT_REQUIREMENT_CACHE[key] = (requirement, now) + _VARIANT_REQUIREMENT_CACHE.move_to_end(key) + while len(_VARIANT_REQUIREMENT_CACHE) > _VARIANT_HASH_MAX: + _VARIANT_REQUIREMENT_CACHE.popitem(last = False) + + +def _build_gguf_variant_requirements(siblings: list) -> dict[str, _GgufVariantRequirement]: + return build_gguf_variant_plans(siblings) + + +def gguf_variant_requirements( + repo_id: str, + variant: str, + hf_token: Optional[str] = None, +) -> Optional[_GgufVariantRequirement]: + key = _variant_hash_cache_key(repo_id, variant, hf_token) + cached = _variant_requirement_cache_get(key) + if cached is not None: + return cached + requirements = _fetch_gguf_variant_requirements(repo_id, hf_token) + return requirements.get(variant.lower()) + + +def _fetch_gguf_variant_requirements( + repo_id: str, + hf_token: Optional[str] = None, + *, + siblings: Optional[list] = None, +) -> dict[str, _GgufVariantRequirement]: + repo_key = _variant_repo_cache_key(repo_id, hf_token) + if siblings is None: + if _variant_requirement_neg_cache_active(repo_key): + return {} + try: + from huggingface_hub import HfApi + info = HfApi(token = hf_token).model_info( + repo_id, + files_metadata = True, + timeout = _GGUF_METADATA_TIMEOUT_SECONDS, + ) + except Exception as e: + logger.warning( + "model_info failed resolving GGUF files for %s: %s", + repo_id, + download_registry.scrub_secrets(str(e), hf_token = hf_token), + ) + _variant_requirement_neg_cache_set(repo_key) + return {} + siblings = list(info.siblings) + requirements = _build_gguf_variant_requirements(siblings) + if requirements: + _variant_requirement_cache_set_many(repo_id, hf_token, requirements) + _variant_requirement_neg_cache_clear(repo_key) + return requirements + + +def _gguf_all_variant_requirements( + repo_id: str, + hf_token: Optional[str] = None, + *, + siblings: Optional[list] = None, +) -> dict[str, _GgufVariantRequirement]: + return _fetch_gguf_variant_requirements(repo_id, hf_token, siblings = siblings) + + +def _manifest_variant_blob_hashes( + repo_id: str, + variant: str, + *, + include_companions: bool = True, +) -> frozenset[str]: + manifest = download_manifest.read_manifest("model", repo_id, variant) + if manifest is None: + return frozenset() + variant_key = variant.lower() + hashes: set[str] = set() + for expected in manifest.expected_files: + if not expected.sha256: + continue + if include_companions: + hashes.add(expected.sha256) + continue + if is_main_gguf_variant_path(expected.path, variant_key): + hashes.add(expected.sha256) + return frozenset(hashes) + + +def gguf_variant_blob_hashes( + repo_id: str, + variant: str, + hf_token: Optional[str] = None, + *, + include_companions: bool = True, + allow_remote: bool = True, +) -> frozenset[str]: + key = _variant_blob_hash_cache_key( + repo_id, + variant, + hf_token, + include_companions, + ) + cached = _variant_hash_cache_get(key) + if cached is not None: + return cached + hashes = _manifest_variant_blob_hashes( + repo_id, + variant, + include_companions = include_companions, + ) + if hashes: + _variant_hash_cache_set(key, hashes) + return hashes + requirement_key = _variant_hash_cache_key(repo_id, variant, hf_token) + requirement = _variant_requirement_cache_get(requirement_key) + if requirement is None and allow_remote: + requirement = gguf_variant_requirements(repo_id, variant, hf_token) + if requirement is not None: + hashes = requirement.required_hashes if include_companions else requirement.main_hashes + if hashes: + _variant_hash_cache_set(key, hashes) + return hashes + return frozenset() + + +def _partial_transport_for_variant(repo_id: str, variant: str) -> Optional[str]: + return hf_cache_scan.partial_transport_for("model", repo_id, variant) + + +def delete_variant_incomplete_blobs_result( + repo_id: str, + variant: str, + hf_token: Optional[str], + *, + extra_hashes: frozenset[str] = frozenset(), + companions: bool = True, +) -> VariantIncompleteDeleteResult: + # With a sibling still downloading, ``companions=False`` keeps a shared mmproj + # from being unlinked out from under it; the repo's last delete reclaims it. + target_hashes = ( + gguf_variant_blob_hashes(repo_id, variant, hf_token, include_companions = companions) + | extra_hashes + ) + if not target_hashes: + has_variant_partial_state = hf_cache_scan.is_variant_partial( + repo_id, + variant, + incomplete_blob_hashes = set(), + variant_blob_hashes = frozenset(), + ) + has_repo_partials = bool(download_registry.incomplete_blob_hashes("model", repo_id)) + return VariantIncompleteDeleteResult( + deleted = 0, + unresolved = has_variant_partial_state and has_repo_partials, + ) + deleted = 0 + # Destructive iterator: only the exact-case match (or abort if ambiguous), + # so a case-variant sibling repo's partials are never unlinked. + for entry in iter_destructive_repo_cache_dirs("model", repo_id): + blobs_dir = entry / "blobs" + if not blobs_dir.is_dir(): + continue + for h in target_hashes: + incomplete = blobs_dir / f"{h}{INCOMPLETE_SUFFIX}" + if incomplete.exists(): + try: + incomplete.unlink() + deleted += 1 + except OSError as e: + logger.warning(f"Failed to unlink {incomplete}: {e}") + return VariantIncompleteDeleteResult(deleted = deleted, unresolved = False) + + +async def get_gguf_variants_response( + repo_id: str, + prefer_local_cache: bool = False, + offline: bool = False, + local_path: Optional[str] = None, + hf_token: Optional[str] = None, +): + """ + List available GGUF quantization variants for a HuggingFace repo + or a local directory (e.g. LM Studio model folder). + + Returns all available quantization variants (Q4_K_M, Q8_0, BF16, etc.) + with file sizes, whether the model supports vision, and the recommended + default variant. + """ + + def _compute() -> GgufVariantsResponse: + def _local_response( + response_repo_id: str, variants, has_vision: bool + ) -> GgufVariantsResponse: + filenames = [v.filename for v in variants] + best = pick_best_gguf(filenames) + default_variant = extract_quant_label(best) if best else None + return GgufVariantsResponse( + repo_id = response_repo_id, + variants = [ + GgufVariantDetail( + filename = v.filename, + quant = v.quant, + display_label = v.display_label, + size_bytes = v.size_bytes, + download_size_bytes = v.size_bytes, + downloaded = True, + ) + for v in variants + ], + has_vision = has_vision, + default_variant = default_variant, + ) + + def _partial_local_response( + response_repo_id: str, variants, has_vision: bool + ) -> GgufVariantsResponse: + filenames = [v.filename for v in variants] + best = pick_best_gguf(filenames) + default_variant = extract_quant_label(best) if best else None + return GgufVariantsResponse( + repo_id = response_repo_id, + variants = [ + GgufVariantDetail( + filename = v.filename, + quant = v.quant, + display_label = v.display_label, + size_bytes = v.size_bytes, + download_size_bytes = v.download_size_bytes or v.size_bytes, + downloaded = False, + partial = True, + partial_transport = _partial_transport_for_variant( + response_repo_id, + v.quant, + ), + ) + for v in variants + ], + has_vision = has_vision, + default_variant = default_variant, + ) + + # Local directory path (e.g. LM Studio models) — scan filesystem + if is_local_path(repo_id): + variants, has_vision = list_local_gguf_variants(repo_id) + + return _local_response(repo_id, variants, has_vision) + + # Reject invalid remote repo_ids up front (like download/delete) so a + # malformed id returns 400 instead of a 500 from the HF client. + if not _is_valid_repo_id(repo_id): + raise HTTPException(status_code = 400, detail = f"Invalid repo_id: {repo_id!r}") + + local_only = prefer_local_cache or offline + if local_only: + cached = list_gguf_variants_from_hf_cache(repo_id) + if cached is not None: + variants, has_vision = cached + return _local_response(repo_id, variants, has_vision) + if local_path and is_local_path(local_path): + variants, has_vision = list_local_gguf_variants(local_path) + if variants or has_vision: + return _local_response(repo_id, variants, has_vision) + partial = list_partial_gguf_variants_from_state(repo_id) + if partial is not None: + variants, has_vision = partial + return _partial_local_response(repo_id, variants, has_vision) + if local_path and offline: + return GgufVariantsResponse( + repo_id = repo_id, + variants = [], + has_vision = False, + default_variant = None, + ) + if offline: + raise HTTPException( + status_code = 404, + detail = "No cached GGUF variants available while offline.", + ) + + try: + variants, has_vision, siblings = list_gguf_variants(repo_id, hf_token = hf_token) + except Exception: + cached = list_gguf_variants_from_hf_cache(repo_id) + if cached is not None: + variants, has_vision = cached + return _local_response(repo_id, variants, has_vision) + partial = list_partial_gguf_variants_from_state(repo_id) + if partial is not None: + variants, has_vision = partial + return _partial_local_response(repo_id, variants, has_vision) + raise + + filenames = [v.filename for v in variants] + best = pick_best_gguf(filenames) + default_variant = extract_quant_label(best) if best else None + + # Per-snapshot accounting: a variant counts as present only when one + # snapshot holds all its files (split GGUFs need every shard together), + # sizes are max across snapshots so shared blobs aren't double-counted, + # and keys are lowercased since cache dir casing can differ from repo_id. + cached_filenames_by_snapshot: list[dict[str, int]] = [] + cached_quant_bytes_by_snapshot: list[dict[str, int]] = [] + if _is_valid_repo_id(repo_id): + for snap in iter_hf_cache_snapshots(repo_id): + try: + gguf_paths = list(_iter_gguf_paths(snap)) + except (OSError, RuntimeError, ValueError) as e: + logger.debug("Skipping GGUF cache snapshot %s: %s", snap, e) + continue + by_filename: dict[str, int] = {} + by_quant: dict[str, int] = {} + for f in gguf_paths: + try: + rel = f.relative_to(snap).as_posix() + size = f.stat().st_size + except (OSError, RuntimeError, ValueError) as e: + logger.debug("Skipping GGUF cache file %s: %s", f, e) + continue + key = rel.lower() + by_filename[key] = max(by_filename.get(key, 0), size) + if _is_mmproj_filename(f.name): + continue + q = extract_quant_label(rel).lower() + by_quant[q] = by_quant.get(q, 0) + size + if by_filename: + cached_filenames_by_snapshot.append(by_filename) + if by_quant: + cached_quant_bytes_by_snapshot.append(by_quant) + + requirements_by_quant = { + v.quant.lower(): _variant_requirement_cache_get( + _variant_hash_cache_key(repo_id, v.quant, hf_token) + ) + for v in variants + } + if any(req is None for req in requirements_by_quant.values()): + fetched_requirements = _gguf_all_variant_requirements( + repo_id, hf_token, siblings = siblings + ) + for v in variants: + key = v.quant.lower() + if requirements_by_quant.get(key) is None: + requirements_by_quant[key] = fetched_requirements.get(key) + + def _filenames_cached(filenames: frozenset[str], expected_size: int) -> bool: + if not filenames: + return False + wanted = [name.lower() for name in filenames] + # All files must live in a single snapshot, not spread across several. + for by_filename in cached_filenames_by_snapshot: + cached = 0 + for name in wanted: + size = by_filename.get(name) + if size is None: + break + cached += size + else: + return expected_size <= 0 or cached >= expected_size * 0.99 + return False + + def _any_mmproj_cached(filenames: frozenset[str]) -> bool: + return any( + by_filename.get(name.lower()) is not None + for by_filename in cached_filenames_by_snapshot + for name in filenames + ) + + def _is_fully_downloaded(variant) -> bool: + requirement = requirements_by_quant.get(variant.quant.lower()) + if requirement is None: + if variant.size_bytes == 0: + return False + quant = variant.quant.lower() + # Allow small rounding tolerance (symlinks vs real sizes). + return any( + by_quant.get(quant, 0) >= variant.size_bytes * 0.99 + for by_quant in cached_quant_bytes_by_snapshot + ) + if not _filenames_cached( + requirement.main_filenames, + requirement.main_size_bytes, + ): + return False + # Vision repos ship an mmproj adapter per variant. Any mmproj + # precision on disk suffices (the loader picks whichever is present); + # requiring the API-preferred one would falsely demote variants. + if requirement.mmproj_filenames and not _any_mmproj_cached( + requirement.mmproj_filenames, + ): + return False + return True + + partial_quants: set[str] = set() + partial_quant_transports: dict[str, Optional[str]] = {} + try: + incomplete_hashes = download_registry.incomplete_blob_hashes("model", repo_id) + except Exception as e: + logger.warning(f"Failed to compute partial GGUF variants for {repo_id}: {e}") + incomplete_hashes = set() + scan_snapshot_dir = hf_cache_scan.resolve_snapshot_dir_for_scan("model", repo_id) + # Manifest + marker + main incomplete-blob check: catches variants whose + # download was cancelled or whose expected shards are missing/undersized. + for variant in variants: + try: + requirement = requirements_by_quant.get(variant.quant.lower()) + variant_hashes = requirement.main_hashes if requirement is not None else None + if variant_hashes is None and incomplete_hashes: + variant_hashes = gguf_variant_blob_hashes( + repo_id, + variant.quant, + hf_token, + include_companions = False, + ) + if hf_cache_scan.is_variant_partial( + repo_id, + variant.quant, + scan_snapshot_dir, + incomplete_blob_hashes = incomplete_hashes, + variant_blob_hashes = variant_hashes, + ): + partial_quants.add(variant.quant) + partial_quant_transports[variant.quant] = _partial_transport_for_variant( + repo_id, + variant.quant, + ) + except Exception as e: + logger.warning( + f"Manifest-based partial check failed for " f"{repo_id}/{variant.quant}: {e}" + ) + if incomplete_hashes: + for variant in variants: + requirement = requirements_by_quant.get(variant.quant.lower()) + if requirement is None: + continue + if requirement.mmproj_hashes & incomplete_hashes and _filenames_cached( + requirement.main_filenames, + requirement.main_size_bytes, + ): + partial_quants.add(variant.quant) + partial_quant_transports.setdefault( + variant.quant, + _partial_transport_for_variant(repo_id, variant.quant), + ) + + def _variant_detail(v) -> GgufVariantDetail: + is_partial = v.quant in partial_quants + requirement = requirements_by_quant.get(v.quant.lower()) + return GgufVariantDetail( + filename = v.filename, + quant = v.quant, + display_label = v.display_label, + size_bytes = v.size_bytes, + download_size_bytes = ( + requirement.download_size_bytes if requirement is not None else v.size_bytes + ), + downloaded = _is_fully_downloaded(v) and not is_partial, + partial = is_partial, + partial_transport = (partial_quant_transports.get(v.quant) if is_partial else None), + ) + + return GgufVariantsResponse( + repo_id = repo_id, + variants = [_variant_detail(v) for v in variants], + has_vision = has_vision, + default_variant = default_variant, + ) + + try: + return await asyncio.to_thread(_compute) + except HTTPException: + raise + except Exception as e: + scrubbed = download_registry.scrub_secrets(str(e), hf_token = hf_token) + # Client-side HF error (missing repo, gated, bad token): pass the status through. + status = hf_error_status(e) + if status is not None: + raise HTTPException(status_code = status, detail = scrubbed) + logger.error("Error listing GGUF variants for %s: %s", repo_id, scrubbed) + raise HTTPException( + status_code = 500, + detail = "Failed to list GGUF variants: " + scrubbed, + ) diff --git a/studio/backend/hub/services/models/local_inventory.py b/studio/backend/hub/services/models/local_inventory.py new file mode 100644 index 0000000000..6e6a28b919 --- /dev/null +++ b/studio/backend/hub/services/models/local_inventory.py @@ -0,0 +1,679 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Local model, HF cache, LM Studio, and Ollama inventory services. + +Ollama logic lives in :mod:`hub.services.models.ollama`; this module +orchestrates all on-device sources and exposes the route handlers. +""" + +from __future__ import annotations + +import asyncio +import os +from pathlib import Path +from typing import List, Optional + +from fastapi import HTTPException +from loggers import get_logger + +from hub.schemas.inventory import LocalModelInfo, LocalModelListResponse, ModelFormat +from hub.storage.scan_folders import ( + add_scan_folder, + list_scan_folders, + remove_scan_folder, +) +from hub.utils import inventory_scan as hf_cache_scan +from hub.utils.paths import ( + hf_default_cache_dir, + legacy_hf_cache_dir, + lmstudio_model_dirs, + normalize_path, + ollama_model_dirs, + outputs_root, + path_is_same_or_child, + studio_root, +) +from hub.services.models import common as model_common +from hub.services.models.ollama import scan_ollama_dir + +logger = get_logger(__name__) +_MAX_MODELS_PER_CUSTOM_FOLDER = 200 +_MAX_CUSTOM_FOLDER_ENTRIES = 2000 +_MODEL_SIGNAL_PROBE_LIMIT = 200 + +# Local aliases keep the extracted code close to the original implementation. +_is_model_directory = model_common._is_model_directory +_local_inventory_id = model_common._local_inventory_id +_local_model_info = model_common._local_model_info +_capabilities_for_format = model_common._capabilities_for_format +_apply_format_aware_partial = model_common._apply_format_aware_partial +_classify_local_path = model_common._classify_local_path +_is_main_gguf_filename = model_common._is_main_gguf_filename +_is_transformers_bin_weight_file = model_common._is_transformers_bin_weight_file +_prefer_complete_larger = model_common._prefer_complete_larger +_gguf_variant_state_summary = model_common._gguf_variant_state_summary + + +def _is_immediate_model_weight_file(path: Path) -> bool: + suffix = path.suffix.lower() + if suffix == ".safetensors": + return True + if suffix == ".gguf": + return _is_main_gguf_filename(path.name) + if suffix == ".bin": + return _is_transformers_bin_weight_file(path) + return False + + +def _has_immediate_model_weight( + path: Path, *, probe_limit: int = _MODEL_SIGNAL_PROBE_LIMIT +) -> bool: + try: + for index, entry in enumerate(path.iterdir(), start = 1): + if index > probe_limit: + break + try: + if entry.is_file() and _is_immediate_model_weight_file(entry): + return True + except OSError: + continue + except OSError: + return False + return False + + +def _has_immediate_model_signal( + path: Path, *, probe_limit: int = _MODEL_SIGNAL_PROBE_LIMIT +) -> bool: + try: + if (path / "config.json").exists() or (path / "adapter_config.json").exists(): + return True + except OSError: + return False + return _has_immediate_model_weight(path, probe_limit = probe_limit) + + +def _is_model_directory_for_scan(path: Path, *, entry_limit: int | None) -> bool: + if entry_limit is None: + return _is_model_directory(path) + try: + has_config = (path / "config.json").exists() or (path / "adapter_config.json").exists() + except OSError: + return False + return has_config and _has_immediate_model_weight(path) + + +def _resolve_hf_cache_dir() -> Path: + try: + from huggingface_hub.constants import HF_HUB_CACHE + return Path(HF_HUB_CACHE) + except Exception: + return Path.home() / ".cache" / "huggingface" / "hub" + + +def _scan_models_dir( + models_dir: Path, + *, + limit: int | None = None, + entry_limit: int | None = None, +) -> List[LocalModelInfo]: + if not models_dir.exists() or not models_dir.is_dir(): + return [] + + _is_self_model = _is_model_directory_for_scan( + models_dir, + entry_limit = entry_limit, + ) + + if _is_self_model: + try: + updated_at = models_dir.stat().st_mtime + except OSError: + updated_at = None + return _classify_local_path( + models_dir, + "models_dir", + updated_at = updated_at, + ) + + found: List[LocalModelInfo] = [] + visited = 0 + try: + children = models_dir.iterdir() + except OSError: + return found + for child in children: + if limit is not None and len(found) >= limit: + break + visited += 1 + if entry_limit is not None and visited > entry_limit: + break + try: + is_dir = child.is_dir() + is_gguf_file = not is_dir and child.suffix.lower() == ".gguf" and child.is_file() + if not is_dir and not is_gguf_file: + continue + has_model_files = is_gguf_file or _has_immediate_model_signal(child) + except OSError: + # Skip individual children that are unreadable (permissions, broken + # symlinks, etc.) rather than failing the entire scan. + continue + if not has_model_files: + continue + try: + updated_at = child.stat().st_mtime + except OSError: + updated_at = None + rows = _classify_local_path( + child, + "models_dir", + updated_at = updated_at, + ) + if limit is not None: + rows = rows[: max(0, limit - len(found))] + found.extend(rows) + + return found + + +def _hf_repo_dir_has_content(repo_dir: Path) -> bool: + blobs_dir = repo_dir / "blobs" + if not blobs_dir.is_dir(): + return False + try: + for entry in blobs_dir.iterdir(): + if entry.is_file() or entry.is_symlink(): + return True + except OSError: + return False + return False + + +def _scan_hf_cache(cache_dir: Path, *, entry_limit: int | None = None) -> List[LocalModelInfo]: + if not cache_dir.exists() or not cache_dir.is_dir(): + return [] + + discovered: List[tuple[Path, str, Optional[float]]] = [] + visited = 0 + try: + entries = cache_dir.iterdir() + except OSError: + return [] + for repo_dir in entries: + visited += 1 + if entry_limit is not None and visited > entry_limit: + break + if not repo_dir.name.startswith("models--"): + continue + if not repo_dir.is_dir(): + continue + if not _hf_repo_dir_has_content(repo_dir): + continue + repo_name = repo_dir.name[len("models--") :] + if not repo_name: + continue + model_id = repo_name.replace("--", "/") + try: + updated_at = repo_dir.stat().st_mtime + except OSError: + updated_at = None + discovered.append((repo_dir, model_id, updated_at)) + + found: list[LocalModelInfo] = [] + for repo_dir, model_id, updated_at in discovered: + snapshot_partial = hf_cache_scan.is_snapshot_partial( + "model", + model_id, + repo_dir, + ) + gguf_partial = hf_cache_scan.is_gguf_repo_partial(model_id, repo_dir) + has_gguf_variant_state, gguf_variant_state_size = _gguf_variant_state_summary(model_id) + snapshot_partial_transport = ( + hf_cache_scan.partial_transport_for( + "model", + model_id, + repo_cache_dir = repo_dir, + ) + if snapshot_partial + else None + ) + resolved = hf_cache_scan.resolve_hf_cache_realpath(repo_dir) + scan_path = Path(resolved) if resolved else repo_dir + # partial=False here; _apply_format_aware_partial below rewrites per-row + # so a hybrid repo's gguf row doesn't taint its safetensors row. + rows = _classify_local_path( + scan_path, + "hf_cache", + load_path = repo_dir, + display_name = model_id.split("/")[-1], + model_id = model_id, + updated_at = updated_at, + partial = False, + ) + if not rows: + if has_gguf_variant_state and gguf_partial: + rows = [ + _local_model_info( + scan_path = repo_dir, + load_path = repo_dir, + source = "hf_cache", + model_format = "gguf", + display_name = model_id.split("/")[-1], + model_id = model_id, + updated_at = updated_at, + partial = True, + requires_variant = True, + size_bytes = gguf_variant_state_size, + ) + ] + else: + # Fallback row's model_format is "unknown"; either signal + # applies because we can't dispatch to a specific predicate. + rows = [ + _local_model_info( + scan_path = repo_dir, + load_path = repo_dir, + source = "hf_cache", + model_format = "unknown", + display_name = model_id.split("/")[-1], + model_id = model_id, + updated_at = updated_at, + partial = snapshot_partial or gguf_partial, + ) + ] + elif ( + has_gguf_variant_state + and gguf_partial + and not any(row.model_format == "gguf" for row in rows) + ): + rows.append( + _local_model_info( + scan_path = repo_dir, + load_path = repo_dir, + source = "hf_cache", + model_format = "gguf", + display_name = model_id.split("/")[-1], + model_id = model_id, + updated_at = updated_at, + partial = True, + requires_variant = True, + size_bytes = gguf_variant_state_size, + ) + ) + rows = _apply_format_aware_partial( + rows, + snapshot_partial = snapshot_partial, + gguf_partial = gguf_partial, + snapshot_partial_transport = snapshot_partial_transport, + ) + found.extend(rows) + return found + + +def _scan_lmstudio_dir(lm_dir: Path, *, entry_limit: int | None = None) -> List[LocalModelInfo]: + """Scan an LM Studio models dir (``publisher/model-name`` folders of GGUFs, or top-level standalone GGUFs).""" + if not lm_dir.exists() or not lm_dir.is_dir(): + return [] + + # If the dir is itself a model dir (config + weights), it's not an LM Studio + # publisher structure -- return it as a single entry rather than descend. + if _is_model_directory(lm_dir): + try: + updated_at = lm_dir.stat().st_mtime + except OSError: + updated_at = None + return _classify_local_path( + lm_dir, + "lmstudio", + updated_at = updated_at, + ) + + found: List[LocalModelInfo] = [] + visited = 0 + exhausted = False + + def _consume_visit() -> bool: + nonlocal visited + visited += 1 + return entry_limit is not None and visited > entry_limit + + try: + children = lm_dir.iterdir() + except OSError: + return found + for child in children: + if _consume_visit(): + break + try: + if not child.is_dir(): + if child.suffix == ".gguf" and child.is_file(): + try: + updated_at = child.stat().st_mtime + except OSError: + updated_at = None + found.extend( + _classify_local_path( + child, + "lmstudio", + updated_at = updated_at, + ) + ) + continue + + # Child is itself a model dir: surface it directly, not as a publisher. + if _is_model_directory(child): + try: + updated_at = child.stat().st_mtime + except OSError: + updated_at = None + found.extend( + _classify_local_path( + child, + "lmstudio", + updated_at = updated_at, + ) + ) + continue + + # child is a publisher directory -- scan its sub-directories + for model_dir in child.iterdir(): + if _consume_visit(): + exhausted = True + break + try: + if model_dir.is_dir(): + has_model = _has_immediate_model_signal(model_dir) + if not has_model: + continue + model_id = f"{child.name}/{model_dir.name}" + try: + updated_at = model_dir.stat().st_mtime + except OSError: + updated_at = None + found.extend( + _classify_local_path( + model_dir, + "lmstudio", + display_name = model_dir.name, + model_id = model_id, + updated_at = updated_at, + ) + ) + elif model_dir.suffix == ".gguf" and model_dir.is_file(): + try: + updated_at = model_dir.stat().st_mtime + except OSError: + updated_at = None + found.extend( + _classify_local_path( + model_dir, + "lmstudio", + model_id = f"{child.name}/{model_dir.stem}", + updated_at = updated_at, + ) + ) + except OSError: + continue + if exhausted: + break + except OSError: + continue + return found + + +def _resolve_allowed_models_dir(models_dir: str, allowed_roots: list[Path]) -> Path: + """Resolve a requested model scan directory without widening subpaths.""" + if not models_dir or not models_dir.strip(): + raise ValueError("Directory not allowed") + + requested = Path(os.path.realpath(os.path.expanduser(normalize_path(models_dir.strip())))) + if any(path_is_same_or_child(requested, root) for root in allowed_roots): + return requested + + raise ValueError("Directory not allowed") + + +def _coerce_scan_folder_path(raw_path: str) -> str: + """Normalize a scan registration target; the registry stores directories, so a pasted weight-file path is reduced to its parent folder.""" + if not raw_path or not raw_path.strip(): + raise ValueError("Path cannot be empty") + raw = raw_path.strip() + if "\x00" in raw: + raise ValueError("Path cannot contain null bytes") + + def normalize(value: str) -> Path: + return Path(os.path.realpath(os.path.expanduser(normalize_path(value)))) + + try: + normalized = normalize(raw) + except (OSError, ValueError) as e: + raise ValueError(f"Path is not readable: {e}") from e + try: + exists = normalized.exists() + is_dir = normalized.is_dir() + is_file = normalized.is_file() + except (OSError, ValueError) as e: + raise ValueError(f"Path is not readable: {e}") from e + + if not exists and "\\" in raw: + try: + slash_normalized = normalize(raw.replace("\\", "/")) + slash_exists = slash_normalized.exists() + except (OSError, ValueError) as e: + raise ValueError(f"Path is not readable: {e}") from e + if slash_exists: + normalized = slash_normalized + try: + is_dir = normalized.is_dir() + is_file = normalized.is_file() + except (OSError, ValueError) as e: + raise ValueError(f"Path is not readable: {e}") from e + exists = True + + if not exists: + return str(normalized) + if is_dir: + return str(normalized) + if is_file: + suffix = normalized.suffix.lower() + if suffix not in {".gguf", ".safetensors", ".bin"}: + raise ValueError("Path must be a folder or model weight file") + return str(normalized.parent) + return str(normalized) + + +async def _scan_source(label: str, scanner, path: Path) -> List[LocalModelInfo]: + try: + return await asyncio.to_thread(scanner, path) + except Exception as e: + logger.warning("Skipping %s scan for %s: %s", label, path, e) + return [] + + +async def _collect_models_from_default_sources( + models_root: Path, + hf_cache_dir: Path, + legacy_hf: Path, + hf_default: Path, + lm_dirs: list[Path], + ollama_dirs: list[Path], +) -> List[LocalModelInfo]: + local_models = await _scan_source("models directory", _scan_models_dir, models_root) + local_models += await _scan_source("HF cache", _scan_hf_cache, hf_cache_dir) + + if legacy_hf.is_dir() and legacy_hf.resolve() != hf_cache_dir.resolve(): + local_models += await _scan_source("legacy HF cache", _scan_hf_cache, legacy_hf) + + if ( + hf_default.is_dir() + and hf_default.resolve() != hf_cache_dir.resolve() + and hf_default.resolve() != legacy_hf.resolve() + ): + local_models += await _scan_source("default HF cache", _scan_hf_cache, hf_default) + + for lm_dir in lm_dirs: + local_models += await _scan_source("LM Studio", _scan_lmstudio_dir, lm_dir) + + for ollama_dir in ollama_dirs: + local_models += await _scan_source("Ollama", scan_ollama_dir, ollama_dir) + + return local_models + + +def _scan_custom_folder(folder_path: Path) -> List[LocalModelInfo]: + supported_formats: set[ModelFormat] = {"gguf", "safetensors", "adapter"} + generic = [ + m + for m in ( + _scan_models_dir( + folder_path, + limit = _MAX_MODELS_PER_CUSTOM_FOLDER, + entry_limit = _MAX_CUSTOM_FOLDER_ENTRIES, + ) + + _scan_hf_cache(folder_path, entry_limit = _MAX_CUSTOM_FOLDER_ENTRIES) + + _scan_lmstudio_dir(folder_path, entry_limit = _MAX_CUSTOM_FOLDER_ENTRIES) + ) + if m.model_format in supported_formats + if not any(p in (".studio_links", "ollama_links") for p in Path(m.path).parts) + ] + return generic[:_MAX_MODELS_PER_CUSTOM_FOLDER] + + +def _promote_to_custom_source(model: LocalModelInfo) -> LocalModelInfo: + if model.source == "hf_cache": + return model + return model.model_copy( + update = { + "source": "custom", + "model_id": None, + "inventory_id": _local_inventory_id( + "custom", + model.model_format, + model.path, + model.format_variant, + ), + "capabilities": _capabilities_for_format( + model.model_format, + "custom", + partial = model.partial, + requires_variant = model.capabilities.requires_variant, + ), + } + ) + + +async def _collect_models_from_custom_folders() -> List[LocalModelInfo]: + try: + custom_folders = await asyncio.to_thread(list_scan_folders) + except Exception as e: + logger.warning("Could not load custom scan folders: %s", e) + return [] + + local_models: List[LocalModelInfo] = [] + for folder in custom_folders: + folder_path = Path(normalize_path(folder["path"])).expanduser() + try: + custom_models = await asyncio.to_thread(_scan_custom_folder, folder_path) + except Exception as e: + logger.warning("Skipping unreadable scan folder %s: %s", folder_path, e) + continue + local_models.extend(_promote_to_custom_source(m) for m in custom_models) + return local_models + + +def _dedupe_local_models(local_models: List[LocalModelInfo]) -> list[LocalModelInfo]: + deduped: dict[str, LocalModelInfo] = {} + for model in local_models: + if model.source == "hf_cache" and model.model_id: + key = "\x00".join( + ( + "hf_cache", + model.model_id.strip().lower(), + model.model_format, + model.format_variant or "", + ) + ) + else: + row_key = model.inventory_id or model.id + key = f"{row_key}\x00custom" if model.source == "custom" else row_key + existing = deduped.get(key) + if existing is None or _prefer_complete_larger( + model.partial, + model.size_bytes, + existing.partial, + existing.size_bytes, + ): + deduped[key] = model + return sorted( + deduped.values(), + key = lambda item: (item.updated_at or 0), + reverse = True, + ) + + +async def list_local_models_response(models_dir: str = "./models") -> LocalModelListResponse: + """List local model candidates from every supported on-device source.""" + hf_cache_dir = _resolve_hf_cache_dir() + legacy_hf = legacy_hf_cache_dir() + hf_default = hf_default_cache_dir() + lm_dirs = lmstudio_model_dirs() + ollama_dirs = ollama_model_dirs() + + allowed_roots: list[Path] = [Path("./models").resolve(), hf_cache_dir] + if legacy_hf.is_dir(): + allowed_roots.append(legacy_hf) + if hf_default.is_dir(): + allowed_roots.append(hf_default) + allowed_roots.extend([studio_root(), outputs_root()]) + + try: + models_root = _resolve_allowed_models_dir(models_dir, allowed_roots) + except ValueError: + raise HTTPException(status_code = 403, detail = "Directory not allowed") + + try: + local_models = await _collect_models_from_default_sources( + models_root, + hf_cache_dir, + legacy_hf, + hf_default, + lm_dirs, + ollama_dirs, + ) + local_models += await _collect_models_from_custom_folders() + models = _dedupe_local_models(local_models) + + return LocalModelListResponse( + models_dir = str(models_root), + hf_cache_dir = str(hf_cache_dir), + lmstudio_dirs = [str(d) for d in lm_dirs], + ollama_dirs = [str(d) for d in ollama_dirs], + models = models, + ) + except Exception as e: + logger.error(f"Error listing local models: {e}", exc_info = True) + raise HTTPException( + status_code = 500, + detail = f"Failed to list local models: {str(e)}", + ) + + +def get_scan_folders_response() -> dict: + return {"folders": list_scan_folders()} + + +def add_scan_folder_response(path: str) -> dict: + try: + folder = add_scan_folder(_coerce_scan_folder_path(path)) + except ValueError as e: + logger.warning("Scan folder rejected: %s (path=%s)", e, path) + raise HTTPException(status_code = 400, detail = str(e)) + logger.info("Scan folder added: %s", folder.get("path")) + return folder + + +def remove_scan_folder_response(folder_id: int) -> dict: + remove_scan_folder(folder_id) + logger.info("Scan folder removed: id=%s", folder_id) + return {"ok": True} diff --git a/studio/backend/hub/services/models/ollama.py b/studio/backend/hub/services/models/ollama.py new file mode 100644 index 0000000000..96a4114620 --- /dev/null +++ b/studio/backend/hub/services/models/ollama.py @@ -0,0 +1,394 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Ollama model inventory: manifest parsing and writable-symlink materialization. + +Ollama stores models content-addressed under ``/manifests/`` and +``/blobs/``. Inventory scans read the manifests directly (no writes), +returning rows whose ``id`` is an opaque ``ollama-manifest:`` reference. The +load path then calls :func:`materialize_ollama_model_ref`, which creates a +``.gguf``-named symlink (or hardlink) so that downstream loaders see a path +with the GGUF suffix without copying multi-GB blobs inside an API request. +""" + +from __future__ import annotations + +import hashlib +import json +import os +import uuid +from pathlib import Path +from typing import List, Optional +from urllib.parse import quote, unquote + +from loggers import get_logger + +from hub.schemas.inventory import LocalModelInfo +from hub.services.models.common import ( + _capabilities_for_format, + _local_inventory_id, +) +from hub.utils.paths import ( + cache_root, + ollama_model_dirs, + path_is_same_or_child, + tmp_root, +) + +logger = get_logger(__name__) + +_OLLAMA_MANIFEST_REF_PREFIX = "ollama-manifest:" +_OLLAMA_BLOB_NAME_CHARS = frozenset( + "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789._+-" +) + + +def _ollama_manifest_ref(tag_file: Path) -> str: + return f"{_OLLAMA_MANIFEST_REF_PREFIX}{quote(str(tag_file), safe = '')}" + + +def _safe_is_file(path: Path) -> bool: + try: + return path.is_file() + except OSError: + return False + + +def _ollama_blob_path(blobs_dir: Path, digest: object) -> Optional[Path]: + if not isinstance(digest, str): + return None + algorithm, separator, value = digest.partition(":") + if separator != ":" or not algorithm or not value: + return None + name = f"{algorithm}-{value}" + if ( + not name + or name in (".", "..") + or any(char not in _OLLAMA_BLOB_NAME_CHARS for char in name) + or not name.isprintable() + ): + return None + return blobs_dir / name + + +def _contained_link_path(link_dir: Path, link_name: str) -> Optional[Path]: + """Resolve *link_name* to a direct child of *link_dir*, or ``None``. ``link_name`` derives from manifest fields, so requiring a direct child keeps a crafted value with separators, ``..``, or a drive prefix from escaping the links dir.""" + if not link_name or link_name in (".", ".."): + return None + link_path = link_dir / link_name + try: + if link_path.parent.resolve() != link_dir.resolve(): + return None + except OSError: + return None + return link_path + + +def _ollama_links_dir(ollama_dir: Path) -> Optional[Path]: + """Writable directory for Ollama ``.gguf`` symlinks. Prefers ``/.studio_links/`` next to the blobs; falls back to Studio's cache (read-only system installs), then the temp dir (sandboxed installs).""" + + def _ensure_writable_dir(path: Path) -> Optional[Path]: + try: + path.mkdir(parents = True, exist_ok = True) + probe = path / f".write-test-{uuid.uuid4().hex[:8]}" + probe.mkdir() + probe.rmdir() + return path + except OSError as e: + logger.debug("Ollama link dir %s is not writable: %s", path, e) + return None + + primary = ollama_dir / ".studio_links" + if _ensure_writable_dir(primary) is not None: + return primary + + # Namespace by a hash of the ollama_dir so two different Ollama roots + # don't collide. This is a cache path, not a security boundary. + try: + digest = hashlib.sha256(str(ollama_dir.resolve()).encode()).hexdigest()[:12] + except (OSError, RuntimeError): + digest = "default" + + fallback = cache_root() / "ollama_links" / digest + if _ensure_writable_dir(fallback) is not None: + return fallback + + tmp_fallback = tmp_root() / "ollama_links" / digest + if _ensure_writable_dir(tmp_fallback) is not None: + return tmp_fallback + + logger.warning( + "Could not create a writable Ollama link directory for %s", + ollama_dir, + ) + return None + + +def _make_ollama_blob_link(link_dir: Path, link_name: str, target: Path) -> Optional[str]: + """Create a .gguf-named link to an Ollama blob: tries symlink then hardlink, skips the model if neither works (a full multi-GB copy would block the API). Idempotent.""" + try: + link_dir.mkdir(parents = True, exist_ok = True) + except OSError as e: + logger.warning( + "Could not create Ollama link directory %s: %s", + link_dir, + e, + ) + return None + link_path = _contained_link_path(link_dir, link_name) + if link_path is None: + logger.warning("Refusing unsafe Ollama link name %r under %s", link_name, link_dir) + return None + try: + resolved = target.resolve() + except OSError as e: + logger.debug("Could not resolve Ollama blob %s: %s", target, e) + return None + + # Skip if the link already points at the same blob. Use samefile, not size: + # `ollama pull` can swap a tag to a same-sized blob, leaving a stale link. + try: + if link_path.exists() and os.path.samefile(str(link_path), str(resolved)): + return str(link_path) + except OSError as e: + logger.debug("Error checking existing link %s: %s", link_path, e) + + tmp_path = link_dir / f".{link_name}.tmp-{uuid.uuid4().hex[:8]}" + try: + if tmp_path.is_symlink() or tmp_path.exists(): + tmp_path.unlink() + try: + tmp_path.symlink_to(resolved) + except OSError: + try: + os.link(str(resolved), str(tmp_path)) + except OSError: + logger.warning( + "Could not create link for Ollama blob %s " + "(symlinks and hardlinks both failed). " + "Skipping model to avoid blocking the API.", + target, + ) + return None + os.replace(str(tmp_path), str(link_path)) + return str(link_path) + except OSError as e: + logger.debug("Could not create Ollama link %s: %s", link_path, e) + try: + if tmp_path.is_symlink() or tmp_path.exists(): + tmp_path.unlink() + except OSError as cleanup_err: + logger.debug("Could not clean up tmp path %s: %s", tmp_path, cleanup_err) + return None + + +def _ollama_model_info_from_manifest( + ollama_dir: Path, + tag_file: Path, + *, + materialize_links: bool = False, + links_root: Optional[Path] = None, +) -> Optional[LocalModelInfo]: + manifests_root = ollama_dir / "manifests" + blobs_dir = ollama_dir / "blobs" + + try: + rel = tag_file.relative_to(manifests_root) + except ValueError: + return None + parts = rel.parts + if len(parts) < 3: + return None + + host = parts[0] + repo_parts = list(parts[1:-1]) + tag = parts[-1] + + if host == "registry.ollama.ai" and repo_parts and repo_parts[0] == "library": + repo_name = "/".join(repo_parts[1:]) + elif host == "registry.ollama.ai": + repo_name = "/".join(repo_parts) + else: + repo_name = "/".join([host] + repo_parts) + + if not repo_name: + return None + + try: + manifest = json.loads(tag_file.read_text()) + except (json.JSONDecodeError, OSError) as e: + logger.debug("Skipping unreadable/invalid Ollama manifest %s: %s", tag_file, e) + return None + + config = manifest.get("config", {}) + config_digest = config.get("digest", "") if isinstance(config, dict) else "" + model_type = "" + file_type = "" + if config_digest and blobs_dir.is_dir(): + config_blob = _ollama_blob_path(blobs_dir, config_digest) + if config_blob is not None and _safe_is_file(config_blob): + try: + cfg = json.loads(config_blob.read_text()) + model_type = cfg.get("model_type", "") + file_type = cfg.get("file_type", "") + except (json.JSONDecodeError, OSError) as e: + logger.debug("Could not parse Ollama config blob %s: %s", config_blob, e) + + layers = manifest.get("layers") or [] + if not isinstance(layers, list): + return None + + model_blob: Optional[Path] = None + gguf_link_path: Optional[str] = None + stem_hash = hashlib.sha256(rel.as_posix().encode()).hexdigest()[:10] + model_link_dir = links_root / stem_hash if links_root is not None else None + safe_name = repo_name.replace("/", "-") + quant = f"-{file_type}" if file_type else "" + + for layer in layers: + if not isinstance(layer, dict): + continue + media = layer.get("mediaType", "") + digest = layer.get("digest", "") + if not digest: + continue + + if media == "application/vnd.ollama.image.model": + candidate = _ollama_blob_path(blobs_dir, digest) + if candidate is None or not _safe_is_file(candidate): + continue + model_blob = candidate + if materialize_links and model_link_dir is not None: + link_name = f"{safe_name}-{tag}{quant}.gguf" + gguf_link_path = _make_ollama_blob_link(model_link_dir, link_name, candidate) + + elif materialize_links and media == "application/vnd.ollama.image.projector": + candidate = _ollama_blob_path(blobs_dir, digest) + if candidate is not None and _safe_is_file(candidate) and model_link_dir is not None: + mmproj_name = f"{safe_name}-{tag}-mmproj.gguf" + _make_ollama_blob_link(model_link_dir, mmproj_name, candidate) + + if model_blob is None: + return None + if materialize_links and not gguf_link_path: + return None + + suffix = "" + if model_type: + suffix += f" ({model_type}" + if file_type: + suffix += f" {file_type}" + suffix += ")" + + try: + updated_at = tag_file.stat().st_mtime + except OSError: + updated_at = None + + display = f"{repo_name}:{tag}" + model_id = f"ollama/{repo_name}:{tag}" + path = gguf_link_path if materialize_links and gguf_link_path else str(model_blob) + load_id = path if materialize_links else _ollama_manifest_ref(tag_file) + return LocalModelInfo( + id = load_id, + inventory_id = _local_inventory_id("ollama", "gguf", model_id), + load_id = load_id, + model_id = model_id, + display_name = display + suffix, + path = path, + source = "ollama", + updated_at = updated_at, + model_format = "gguf", + runtime = "llama_cpp", + capabilities = _capabilities_for_format("gguf", "ollama"), + ) + + +def scan_ollama_dir( + ollama_dir: Path, + *, + limit: Optional[int] = None, + materialize_links: bool = False, +) -> List[LocalModelInfo]: + """Scan an Ollama models directory for downloaded models. + + Ollama uses a content-addressable layout + (``manifests////`` + ``blobs/sha256-...``), + iterated via ``rglob`` to find every depth. Each manifest's ``model`` layer + holds the GGUF weights (vision models add a projector layer). + + Scans are read-only by default and return an opaque manifest reference; + the load route later calls :func:`materialize_ollama_model_ref` to create a + ``.gguf`` symlink/hardlink, keeping GET /local free of filesystem writes. + """ + manifests_root = ollama_dir / "manifests" + if not manifests_root.is_dir(): + return [] + + found: List[LocalModelInfo] = [] + links_root = _ollama_links_dir(ollama_dir) if materialize_links else None + if materialize_links and links_root is None: + logger.warning( + "Skipping Ollama scan for %s: no writable location for .gguf links", + ollama_dir, + ) + return [] + + try: + for tag_file in manifests_root.rglob("*"): + if not _safe_is_file(tag_file): + continue + + info = _ollama_model_info_from_manifest( + ollama_dir, + tag_file, + materialize_links = materialize_links, + links_root = links_root, + ) + if info is None: + continue + found.append(info) + if limit is not None and len(found) >= limit: + return found + except OSError as e: + logger.warning("Error scanning Ollama directory %s: %s", ollama_dir, e) + return found + + +def _ollama_dir_for_manifest(tag_file: Path) -> Optional[Path]: + """Discovered Ollama root whose ``manifests/`` contains *tag_file*, or ``None``. Validating against known roots keeps a crafted reference from driving materialization to an arbitrary path.""" + for ollama_dir in ollama_model_dirs(): + if path_is_same_or_child(tag_file, ollama_dir / "manifests"): + return ollama_dir + return None + + +def materialize_ollama_model_ref(ref: str) -> str: + """Resolve an ``ollama-manifest:`` reference to a loadable ``.gguf`` path, + creating the writable symlink/hardlink on demand. + + Raises ``ValueError`` if the reference is malformed, points outside a + discovered Ollama models directory, or cannot be materialized. + """ + if not ref.startswith(_OLLAMA_MANIFEST_REF_PREFIX): + raise ValueError("Not an Ollama manifest reference") + + tag_file = Path(unquote(ref[len(_OLLAMA_MANIFEST_REF_PREFIX) :])) + + ollama_dir = _ollama_dir_for_manifest(tag_file) + if ollama_dir is None: + raise ValueError("Reference is outside any known Ollama models directory") + + links_root = _ollama_links_dir(ollama_dir) + if links_root is None: + raise ValueError("No writable location for Ollama .gguf links") + + info = _ollama_model_info_from_manifest( + ollama_dir, + tag_file, + materialize_links = True, + links_root = links_root, + ) + if info is None or not info.path: + raise ValueError("Could not materialize Ollama model from manifest") + return info.path diff --git a/studio/backend/hub/services/snapshot_progress.py b/studio/backend/hub/services/snapshot_progress.py new file mode 100644 index 0000000000..9c8bc9a891 --- /dev/null +++ b/studio/backend/hub/services/snapshot_progress.py @@ -0,0 +1,260 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Shared snapshot download-progress computation for models and datasets. + +Both scan the cache's ``blobs/`` dir, split finalized vs ``.incomplete`` bytes, +filter to the target revision's expected hashes, and divide by its total size; +only the ``metadata_resolver`` differs. One copy keeps the two from drifting (a +prior hash-filter fix once landed only on the model copy, leaving datasets +summing stale blobs against the wrong total).""" + +from __future__ import annotations + +import asyncio +from pathlib import Path +from typing import Callable, Optional + +from loggers import get_logger + +from hub.utils import download_manifest +from hub.utils import download_registry +from hub.utils import inventory_scan as hf_cache_scan +from hub.utils.state_dir import RepoType +from hub.utils.hf_cache_state import ( + INCOMPLETE_SUFFIX, + blob_bytes_present, + latest_snapshot_dir, + preferred_repo_cache_dirs, +) +from hub.utils.paths import is_valid_repo_id as _is_valid_repo_id + +logger = get_logger(__name__) + +# (repo_id, hf_token) -> (expected_total_bytes, expected_blob_hashes) +SnapshotMetadataResolver = Callable[[str, Optional[str]], "tuple[int, frozenset[str]]"] + + +def _empty_progress(expected_bytes: int) -> dict: + return { + "downloaded_bytes": 0, + "completed_bytes": 0, + "complete_on_disk": False, + "expected_bytes": max(expected_bytes, 0), + "progress": 0, + "cache_path": None, + } + + +def _snapshot_complete_on_disk( + *, + repo_type: RepoType, + repo_id: str, + variant: Optional[str], + entry: Path, + expected_total: int, + completed_bytes: int, + in_progress_bytes: int, +) -> bool: + if expected_total <= 0 or completed_bytes < expected_total or in_progress_bytes > 0: + return False + snapshot_dir = latest_snapshot_dir(entry) + if snapshot_dir is None: + return False + if variant is None and hf_cache_scan.repo_cache_dir_has_incomplete_blobs(entry): + return False + if download_manifest.has_cancel_marker(repo_type, repo_id, variant): + return False + manifest = download_manifest.read_manifest(repo_type, repo_id, variant) + if manifest is None: + return False + return download_manifest.verify_against_disk(manifest, snapshot_dir).ok + + +def compute_snapshot_progress( + *, + repo_type: RepoType, + repo_id: str, + job_key: str, + expected_bytes: int, + hf_token: Optional[str], + registry, + metadata_resolver: SnapshotMetadataResolver, + variant: Optional[str] = None, +) -> dict: + """Synchronous progress reading. Safe to run under ``asyncio.to_thread``.""" + empty = _empty_progress(expected_bytes) + if not _is_valid_repo_id(repo_id): + return empty + + job_state = registry.get_job(job_key).state + force_active = job_state in {"running", "cancelling"} + get_job_metadata = getattr(registry, "get_job_metadata", None) + metadata = get_job_metadata(job_key) if callable(get_job_metadata) else None + completed_baseline_bytes = max( + 0, + int(getattr(metadata, "completed_baseline_bytes", 0) or 0), + ) + + expected_total = max(expected_bytes, 0) + # Always resolve the revision's blob hashes so stale blobs from a superseded + # revision can't inflate the count; hashes degrade to empty (count-all) only + # when metadata is unavailable (e.g. offline). Take the larger total so a low + # caller hint can't cap the bar below the revision's real size. + meta_total, expected_hashes = metadata_resolver(repo_id, hf_token) + expected_total = max(expected_total, meta_total) + + # Without resolved hashes, a variant must not count unscoped blobs: sibling + # quants share one blobs/ dir, so a sibling's bytes (or .incomplete) would be + # misattributed and make the bar jump backward. A no-variant snapshot owns + # the whole dir, so it counts unscoped. + count_finalized_unscoped = variant is None + + readings: list[tuple[int, int, Optional[str], bool]] = [] + for entry in preferred_repo_cache_dirs( + repo_type, + repo_id, + force_active = force_active, + ): + completed_bytes = 0 + in_progress_bytes = 0 + cache_path = hf_cache_scan.resolve_hf_cache_realpath(entry) + blobs_dir = entry / "blobs" + if blobs_dir.is_dir(): + try: + blob_entries = list(blobs_dir.iterdir()) + except OSError: + blob_entries = [] + for f in blob_entries: + # Skip a blob that vanished mid-poll rather than zeroing the reading. + try: + if not f.is_file(): + continue + if f.name.endswith(INCOMPLETE_SUFFIX): + blob_hash = f.name[: -len(INCOMPLETE_SUFFIX)] + if expected_hashes: + if blob_hash not in expected_hashes: + continue + elif not count_finalized_unscoped: + continue + in_progress_bytes += blob_bytes_present(f) + else: + if expected_hashes: + if f.name not in expected_hashes: + continue + elif not count_finalized_unscoped: + continue + completed_bytes += f.stat().st_size + except OSError: + continue + readings.append( + ( + completed_bytes, + in_progress_bytes, + cache_path, + _snapshot_complete_on_disk( + repo_type = repo_type, + repo_id = repo_id, + variant = variant, + entry = entry, + expected_total = expected_total, + completed_bytes = completed_bytes, + in_progress_bytes = in_progress_bytes, + ), + ) + ) + + selected = max( + readings, + key = lambda item: (item[0] + item[1], item[0]), + default = None, + ) + if selected is None: + return empty + + completed_bytes, in_progress_bytes, cache_path, complete_on_disk = selected + downloaded_bytes = completed_bytes + in_progress_bytes + # Subtract the companion baseline only while still counted in completed_bytes + # and the variant is not yet verified complete, else genuine progress reads as + # 0-byte. + effective_baseline_bytes = ( + completed_baseline_bytes + if not complete_on_disk and completed_baseline_bytes <= completed_bytes + else 0 + ) + display_completed_bytes = max(0, completed_bytes - effective_baseline_bytes) + display_downloaded_bytes = max(0, downloaded_bytes - effective_baseline_bytes) + + if expected_total <= 0: + # Cannot determine total; report bytes only, no percentage. + return { + "downloaded_bytes": display_downloaded_bytes, + "completed_bytes": display_completed_bytes, + "complete_on_disk": False, + "expected_bytes": 0, + "progress": 0, + "cache_path": cache_path, + } + + display_expected_total = max(0, expected_total - effective_baseline_bytes) + if downloaded_bytes == 0: + return { + **empty, + "expected_bytes": display_expected_total, + "cache_path": cache_path, + } + + # Cap at 0.99 until the manifest-backed disk check verifies completion: on + # resume, completed bytes can sit above the threshold while files still download. + progress = ( + 1.0 + if complete_on_disk + else ( + min(display_downloaded_bytes / display_expected_total, 0.99) + if display_expected_total > 0 + else 0 + ) + ) + return { + "downloaded_bytes": display_downloaded_bytes, + "completed_bytes": display_completed_bytes, + "complete_on_disk": complete_on_disk, + "expected_bytes": display_expected_total, + "progress": round(progress, 3), + "cache_path": cache_path, + } + + +async def snapshot_progress_response( + *, + repo_type: RepoType, + repo_id: str, + job_key: str, + expected_bytes: int, + hf_token: Optional[str], + registry, + metadata_resolver: SnapshotMetadataResolver, + variant: Optional[str] = None, +) -> dict: + """Async wrapper: offloads the blocking cache walk and never raises.""" + try: + return await asyncio.to_thread( + compute_snapshot_progress, + repo_type = repo_type, + repo_id = repo_id, + job_key = job_key, + expected_bytes = expected_bytes, + hf_token = hf_token, + registry = registry, + metadata_resolver = metadata_resolver, + variant = variant, + ) + except Exception as e: + logger.warning( + "Error checking %s download progress for %s: %s: %s", + repo_type, + repo_id, + type(e).__name__, + download_registry.scrub_secrets(str(e), hf_token = hf_token), + ) + return _empty_progress(expected_bytes) diff --git a/studio/backend/hub/storage/__init__.py b/studio/backend/hub/storage/__init__.py new file mode 100644 index 0000000000..32014236c6 --- /dev/null +++ b/studio/backend/hub/storage/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 diff --git a/studio/backend/hub/storage/scan_folders.py b/studio/backend/hub/storage/scan_folders.py new file mode 100644 index 0000000000..85f515da00 --- /dev/null +++ b/studio/backend/hub/storage/scan_folders.py @@ -0,0 +1,194 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Persistence for user-registered custom model scan folders. + +Self-bootstrapping table inside the existing studio SQLite so the Hub module +doesn't have to modify upstream studio_db.py's schema init.""" + +from __future__ import annotations + +import os +import platform +import sqlite3 +import threading +from datetime import datetime, timezone + +from storage.studio_db import get_connection +from hub.utils.paths import normalize_path + + +_schema_lock = threading.Lock() +_schema_ready = False +_SENSITIVE_PATH_COMPONENTS = { + ".aws", + ".azure", + ".config", + ".docker", + ".gcloud", + ".gnupg", + ".huggingface", + ".kaggle", + ".kube", + ".modelscope", + ".ngc", + ".local", + ".mozilla", + ".pki", + ".thunderbird", + ".ssh", + ".1password", + ".bitwarden", + ".password-store", + "1password", + "bitwarden", + "keychains", + "keyrings", + "mozilla", + "thunderbird", +} + + +def _denied_path_prefixes() -> list[str]: + system = platform.system() + if system == "Linux": + return ["/proc", "/sys", "/dev", "/etc", "/boot", "/run"] + if system == "Darwin": + # realpath() resolves /etc -> /private/etc, /tmp -> /private/tmp on macOS, + # so include the /private variants to avoid bypasses. + return [ + "/System", + "/Library", + "/dev", + "/etc", + "/private/etc", + "/tmp", + "/private/tmp", + "/var", + "/private/var", + ] + if system == "Windows": + win = os.environ.get("SystemRoot", r"C:\Windows") + pf = os.environ.get("ProgramFiles", r"C:\Program Files") + pf86 = os.environ.get("ProgramFiles(x86)", r"C:\Program Files (x86)") + return [os.path.normcase(p) for p in [win, pf, pf86]] + return [] + + +def _contains_sensitive_path_component(path: str) -> bool: + parts = os.path.normpath(path).split(os.sep) + return any(part.lower() in _SENSITIVE_PATH_COMPONENTS for part in parts) + + +def contains_sensitive_path_component(path: str) -> bool: + """Public predicate for the credential/config denylist (.ssh, .aws, ...). + + Shared with the folder browser so browse and register enforce one policy.""" + return _contains_sensitive_path_component(path) + + +def _ensure_schema(conn: sqlite3.Connection) -> None: + global _schema_ready + if _schema_ready: + return + with _schema_lock: + if _schema_ready: + return + collation = "COLLATE NOCASE" if platform.system() == "Windows" else "" + conn.execute( + f""" + CREATE TABLE IF NOT EXISTS scan_folders ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + path TEXT NOT NULL UNIQUE {collation}, + created_at TEXT NOT NULL + ) + """ + ) + conn.commit() + _schema_ready = True + + +def list_scan_folders() -> list[dict]: + conn = get_connection() + try: + _ensure_schema(conn) + rows = conn.execute( + "SELECT id, path, created_at FROM scan_folders ORDER BY created_at" + ).fetchall() + return [dict(row) for row in rows] + finally: + conn.close() + + +def add_scan_folder(path: str) -> dict: + """Add a readable directory for the local OS user; not a multi-user sandbox.""" + if not path or not path.strip(): + raise ValueError("Path cannot be empty") + normalized = os.path.realpath(os.path.expanduser(normalize_path(path.strip()))) + + if not os.path.exists(normalized): + raise ValueError("Path does not exist") + if not os.path.isdir(normalized): + raise ValueError("Path must be a directory, not a file") + if not os.access(normalized, os.R_OK | os.X_OK): + raise ValueError("Path is not readable") + if os.path.dirname(normalized) == normalized: + # Registering a filesystem root would expose denied system dirs via browse. + raise ValueError("The filesystem root cannot be registered") + if _contains_sensitive_path_component(normalized): + raise ValueError("Credential or configuration directories are not allowed") + + is_win = platform.system() == "Windows" + check = os.path.normcase(normalized) if is_win else normalized + for prefix in _denied_path_prefixes(): + if check == prefix or check.startswith(prefix + os.sep): + raise ValueError(f"Path under {prefix} is not allowed") + + conn = get_connection() + try: + _ensure_schema(conn) + now = datetime.now(timezone.utc).isoformat() + if is_win: + existing = conn.execute( + "SELECT id, path, created_at FROM scan_folders WHERE path = ? COLLATE NOCASE", + (normalized,), + ).fetchone() + else: + existing = conn.execute( + "SELECT id, path, created_at FROM scan_folders WHERE path = ?", + (normalized,), + ).fetchone() + if existing is not None: + return dict(existing) + try: + conn.execute( + "INSERT INTO scan_folders (path, created_at) VALUES (?, ?)", + (normalized, now), + ) + conn.commit() + except sqlite3.IntegrityError: + pass + fallback_sql = ( + "SELECT id, path, created_at FROM scan_folders WHERE path = ? COLLATE NOCASE" + if is_win + else "SELECT id, path, created_at FROM scan_folders WHERE path = ?" + ) + row = conn.execute(fallback_sql, (normalized,)).fetchone() + if row is None: + raise ValueError("Folder was concurrently removed") + return dict(row) + finally: + conn.close() + + +def remove_scan_folder(id: int) -> None: + # sqlite INTEGER is signed 64-bit; ids outside that range cannot exist. + if not -(2**63) <= id < 2**63: + return + conn = get_connection() + try: + _ensure_schema(conn) + conn.execute("DELETE FROM scan_folders WHERE id = ?", (id,)) + conn.commit() + finally: + conn.close() diff --git a/studio/backend/hub/tests/conftest.py b/studio/backend/hub/tests/conftest.py new file mode 100644 index 0000000000..38b6bed938 --- /dev/null +++ b/studio/backend/hub/tests/conftest.py @@ -0,0 +1,106 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import sys +import types + + +class _BaseModel: + def __init__(self, **kwargs): + for name, value in self.__class__.__dict__.items(): + if name.startswith("_") or callable(value): + continue + if name not in kwargs: + setattr(self, name, value) + for key, value in kwargs.items(): + setattr(self, key, value) + + def model_dump(self): + return dict(self.__dict__) + + def model_copy(self, update = None): + data = self.model_dump() + if update: + data.update(update) + return self.__class__(**data) + + +def _field(default = ..., **kwargs): + if "default_factory" in kwargs: + return kwargs["default_factory"]() + return None if default is ... else default + + +def _model_validator(*args, **kwargs): + def decorator(fn): + return fn + + return decorator + + +class _HTTPException(Exception): + def __init__( + self, + status_code: int, + detail = None, + ): + super().__init__(detail) + self.status_code = status_code + self.detail = detail + + +class _APIRouter: + def get(self, *args, **kwargs): + return lambda fn: fn + + def post(self, *args, **kwargs): + return lambda fn: fn + + def delete(self, *args, **kwargs): + return lambda fn: fn + + +def _fastapi_marker( + default = None, + *args, + **kwargs, +): + return default + + +class _DummyLogger: + def __getattr__(self, _name): + return lambda *args, **kwargs: None + + +sys.modules.setdefault( + "pydantic", + types.SimpleNamespace( + BaseModel = _BaseModel, + Field = _field, + model_validator = _model_validator, + ), +) +sys.modules.setdefault( + "fastapi", + types.SimpleNamespace( + APIRouter = _APIRouter, + Body = _fastapi_marker, + Depends = _fastapi_marker, + Header = _fastapi_marker, + HTTPException = _HTTPException, + Query = _fastapi_marker, + UploadFile = object, + ), +) +sys.modules.setdefault( + "loggers", + types.SimpleNamespace(get_logger = lambda *args, **kwargs: _DummyLogger()), +) +sys.modules.setdefault( + "structlog", + types.SimpleNamespace( + BoundLogger = _DummyLogger, + get_logger = lambda *args, **kwargs: _DummyLogger(), + ), +) diff --git a/studio/backend/hub/tests/test_dataset_services.py b/studio/backend/hub/tests/test_dataset_services.py new file mode 100644 index 0000000000..4890714cd0 --- /dev/null +++ b/studio/backend/hub/tests/test_dataset_services.py @@ -0,0 +1,356 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import asyncio +from pathlib import Path +from types import SimpleNamespace + +import pytest +from fastapi import HTTPException + +from hub.schemas.datasets import CheckFormatRequest, LocalDatasetItem +from hub.services.datasets import cache_inventory, downloads, formatting, local +from hub.utils import download_manifest, download_registry, state_dir + + +class _Upload: + def __init__(self, filename: str, payload: bytes): + self.filename = filename + self._payload = payload + self._offset = 0 + + async def read(self, size: int) -> bytes: + if self._offset >= len(self._payload): + return b"" + chunk = self._payload[self._offset : self._offset + size] + self._offset += len(chunk) + return chunk + + +def test_dataset_cache_scan_merges_raw_and_processed_rows(monkeypatch): + raw_repo = SimpleNamespace( + repo_id = "Org/Data", + repo_type = "dataset", + repo_path = "/cache/datasets--Org--Data", + size_on_disk = 100, + revisions = [SimpleNamespace(files = [], commit_hash = "abc")], + ) + monkeypatch.setattr( + cache_inventory, + "_collect_hf_cache_scans", + lambda: ([SimpleNamespace(repos = [raw_repo])], {"/cache"}), + ) + monkeypatch.setattr( + cache_inventory.hf_cache_scan, + "is_snapshot_partial", + lambda _repo_type, _repo_id, _cache_dir: False, + ) + monkeypatch.setattr( + cache_inventory, + "_scan_hub_dataset_cache_dirs", + lambda: [], + ) + monkeypatch.setattr( + cache_inventory, + "_scan_processed_dataset_caches", + lambda: [ + { + "repo_id": "org/data", + "size_bytes": 250, + "cache_path": "/processed/org___data", + "processed_cache": True, + "partial": False, + } + ], + ) + + rows = cache_inventory._scan_hf_dataset_caches() + + assert len(rows) == 1 + assert rows[0]["repo_id"] == "Org/Data" + assert rows[0]["size_bytes"] == 250 + assert rows[0]["partial"] is False + + +def test_delete_cached_dataset_attempts_all_roots_before_raising(monkeypatch): + calls = [] + purged_state = [] + + class _DeleteStrategy: + def __init__(self, label: str, fail: bool): + self.label = label + self.fail = fail + + def execute(self): + calls.append(self.label) + if self.fail: + raise RuntimeError(f"{self.label} failed") + + class _Cache: + def __init__(self, label: str, fail: bool): + self.cache_dir = label + self.repos = [ + SimpleNamespace( + repo_type = "dataset", + repo_id = "Org/Data", + revisions = [SimpleNamespace(commit_hash = f"{label}-rev")], + ) + ] + self.fail = fail + + def delete_revisions(self, *_revisions): + return _DeleteStrategy(self.cache_dir, self.fail) + + monkeypatch.setattr( + cache_inventory, + "_collect_hf_cache_scans", + lambda: ([_Cache("first", True), _Cache("second", False)], set()), + ) + monkeypatch.setattr( + cache_inventory, + "_delete_processed_dataset_cache", + lambda _repo_id: (True, []), + ) + monkeypatch.setattr( + cache_inventory.download_manifest, + "purge_all_state_for_repo", + lambda *_args: purged_state.append(True) or 1, + ) + + with pytest.raises(HTTPException) as exc_info: + cache_inventory._delete_cached_dataset_blocking("Org/Data") + + assert exc_info.value.status_code == 500 + assert calls == ["first", "second"] + assert purged_state == [] + + +def test_delete_cached_dataset_purges_blob_only_repo_dir(monkeypatch): + """A blob-only ``datasets--owner--repo`` dir (no usable snapshot/refs) is + fully removable: purge_partial_repo alone clears only ``.incomplete`` files + and would leave the complete blobs and the row.""" + purged_dirs: list[str] = [] + + monkeypatch.setattr( + cache_inventory, + "_collect_hf_cache_scans", + lambda: ([], set()), + ) + monkeypatch.setattr( + cache_inventory, + "_delete_processed_dataset_cache", + lambda _repo_id: (False, []), + ) + monkeypatch.setattr( + cache_inventory, + "purge_repo_cache_dirs", + lambda _repo_type, repo_id: purged_dirs.append(repo_id) or True, + ) + monkeypatch.setattr( + cache_inventory, + "purge_partial_repo", + lambda *_args: False, + ) + monkeypatch.setattr( + cache_inventory.download_manifest, + "purge_all_state_for_repo", + lambda *_args: 0, + ) + + result = cache_inventory._delete_cached_dataset_blocking("Org/Data") + + assert result == {"status": "deleted", "repo_id": "Org/Data"} + assert purged_dirs == ["Org/Data"] + + +def test_delete_cached_dataset_absent_everywhere_raises_404(monkeypatch): + monkeypatch.setattr( + cache_inventory, + "_collect_hf_cache_scans", + lambda: ([], set()), + ) + monkeypatch.setattr( + cache_inventory, + "_delete_processed_dataset_cache", + lambda _repo_id: (False, []), + ) + monkeypatch.setattr( + cache_inventory, + "purge_repo_cache_dirs", + lambda *_args: False, + ) + monkeypatch.setattr( + cache_inventory, + "purge_partial_repo", + lambda *_args: False, + ) + monkeypatch.setattr( + cache_inventory.download_manifest, + "purge_all_state_for_repo", + lambda *_args: 0, + ) + + with pytest.raises(HTTPException) as exc_info: + cache_inventory._delete_cached_dataset_blocking("Org/Missing") + + assert exc_info.value.status_code == 404 + + +def test_check_format_rejects_invalid_path_as_400(): + with pytest.raises(HTTPException) as exc_info: + formatting.check_format_response(CheckFormatRequest(dataset_name = "../../etc/passwd")) + + assert exc_info.value.status_code == 400 + + +def test_dataset_download_status_preserves_idle_shape(): + status = downloads._dataset_status("Org/Data") + + assert status.state == "idle" + assert status.error is None + + +def test_dataset_download_registry_key_is_case_insensitive(): + registry = download_registry.DownloadRegistry() + + claimed, state = registry.claim( + "Org/Data", + download_registry.TRANSPORT_HTTP, + repo_type = "dataset", + repo_id = "Org/Data", + ) + duplicate_claimed, duplicate_state = registry.claim( + "org/data", + download_registry.TRANSPORT_HTTP, + repo_type = "dataset", + repo_id = "org/data", + ) + + assert claimed is True + assert state == "running" + assert duplicate_claimed is False + assert duplicate_state == "running" + assert registry.active_jobs("ORG/DATA") == {"org/data": "running"} + + +def test_dataset_idle_status_uses_cancel_marker_after_restart(monkeypatch, tmp_path): + monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path) + monkeypatch.setattr(downloads, "_registry", download_registry.DownloadRegistry()) + assert download_manifest.write_cancel_marker("dataset", "Owner/Data", None, "http") + + status = asyncio.run(downloads.get_dataset_download_status_response("owner/data")) + + assert status.state == "cancelled" + assert status.error is None + + +def test_dataset_claim_register_cancel_uses_registry_marker_owner(monkeypatch): + killed = [] + + class _Registry: + def claim(self, *_args, **_kwargs): + return True, "running" + + def current_generation(self, _key): + return 1 + + def register_process(self, _key, _proc): + return False + + def persist_cancel_for_key(self, *_args, **_kwargs): + raise AssertionError("register_process owns pending-cancel markers") + + def get_job(self, _key): + return SimpleNamespace(state = "cancelled", error = None) + + monkeypatch.setattr(downloads, "_registry", _Registry()) + monkeypatch.setattr( + downloads, + "resolve_cached_repo_id_case", + lambda repo_id, **_kwargs: repo_id, + ) + monkeypatch.setattr( + downloads.download_registry, + "download_transport_unavailable_reason", + lambda _transport: None, + ) + monkeypatch.setattr( + downloads.download_lifecycle, + "spawn_worker", + lambda *_args, **_kwargs: object(), + ) + monkeypatch.setattr( + downloads.download_lifecycle, + "kill_and_reap_process", + lambda proc, **_kwargs: killed.append(proc), + ) + + result = asyncio.run( + downloads.download_dataset_response(SimpleNamespace(repo_id = "Org/Data", use_xet = False)) + ) + + assert result["state"] == "cancelled" + assert killed + + +def test_dataset_cancel_pending_spawn_arms_pending_cancel(monkeypatch): + events = [] + + class _Registry: + def get_process(self, _key): + return None + + def mark_pending_cancel(self, key, generation): + events.append(("pending", key, generation)) + return True + + def get_job(self, _key): + return SimpleNamespace(state = "running") + + monkeypatch.setattr(downloads, "_registry", _Registry()) + monkeypatch.setattr( + downloads, + "resolve_cached_repo_id_case", + lambda repo_id, **_kwargs: repo_id, + ) + + result = asyncio.run( + downloads.cancel_dataset_download_response( + SimpleNamespace(repo_id = "Org/Data", generation = 5) + ) + ) + + assert result == {"repo_id": "Org/Data", "state": "cancelling"} + assert events == [("pending", "org/data", 5)] + + +def test_upload_dataset_response_writes_non_empty_file(monkeypatch, tmp_path): + payload = b'{"text":"hello"}\n' + monkeypatch.setattr(local, "DATASET_UPLOAD_DIR", tmp_path) + + response = asyncio.run(local.upload_dataset_response(_Upload("../train.jsonl", payload))) + + stored_path = Path(response.stored_path) + assert response.filename == "train.jsonl" + assert stored_path.parent == tmp_path + assert stored_path.name.endswith("_train.jsonl") + assert stored_path.read_bytes() == payload + + +def test_local_dataset_items_expose_recipe_and_upload_source(monkeypatch, tmp_path): + recipe_root = tmp_path / "recipes" + upload_root = tmp_path / "uploads" + parquet_dir = recipe_root / "recipe_alpha" / "parquet-files" + parquet_dir.mkdir(parents = True) + (parquet_dir / "part.parquet").write_bytes(b"parquet") + upload_root.mkdir() + (upload_root / "manual.jsonl").write_text('{"text":"hello"}\n', encoding = "utf-8") + monkeypatch.setattr(local, "LOCAL_DATASETS_ROOT", recipe_root) + monkeypatch.setattr(local, "DATASET_UPLOAD_DIR", upload_root) + + response = local.list_local_datasets_response() + + assert "source" in LocalDatasetItem.__annotations__ + by_id = {item.id: item for item in response.datasets} + assert by_id["recipe_alpha"].source == "recipe" + assert by_id["manual.jsonl"].source == "upload" diff --git a/studio/backend/hub/tests/test_model_services.py b/studio/backend/hub/tests/test_model_services.py new file mode 100644 index 0000000000..44ab4b80d0 --- /dev/null +++ b/studio/backend/hub/tests/test_model_services.py @@ -0,0 +1,2962 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import asyncio +import sys +from pathlib import Path +from types import SimpleNamespace + +import pytest +from fastapi import HTTPException + +from hub.dependencies import get_hf_token +from hub.storage import scan_folders +from hub.services import download_lifecycle +from hub.services import snapshot_progress +from hub.services.datasets import downloads as dataset_downloads +from hub.services.models import ( + cache_inventory, + common as model_common, + deletion, + downloads, + folder_browser, + gguf_variants, + local_inventory, + ollama, +) +from hub.utils import ( + download_manifest, + download_registry, + gguf, + hf_cache_state, + inventory_scan, + paths, + state_dir, +) +from hub.workers import hf_download + + +def _repo(repo_id: str, files: list[SimpleNamespace], repo_path: Path): + return SimpleNamespace( + repo_id = repo_id, + repo_type = "model", + repo_path = repo_path, + revisions = [SimpleNamespace(files = files)], + ) + + +def _file( + name: str, + size: int, + blob_path: str | None = None, +): + return SimpleNamespace(file_name = name, size_on_disk = size, blob_path = blob_path) + + +def _sibling(name: str, size: int, sha: str): + return SimpleNamespace(rfilename = name, size = size, lfs = {"sha256": sha}) + + +class TestExtractQuantToken: + def test_trailing_precision_is_kept(self): + assert gguf.extract_quant_token("model-it-F16.gguf") == "F16" + assert gguf.extract_quant_token("model-BF16.gguf") == "BF16" + + def test_real_quant_wins_over_infix_precision(self): + assert gguf.extract_quant_token("Foo-BF16-Q4_K_M.gguf") == "Q4_K_M" + assert gguf.extract_quant_token("Foo-F16-Q8_0.gguf") == "Q8_0" + assert gguf.extract_quant_token("Foo-F32-IQ4_XS.gguf") == "IQ4_XS" + + def test_ud_prefix_preserved(self): + assert gguf.extract_quant_token("Foo-BF16-UD-Q4_K_XL.gguf") == "UD-Q4_K_XL" + + def test_precision_infix_variants_do_not_collapse(self): + labels = { + gguf.extract_quant_label("Foo-BF16-Q4_K_M.gguf"), + gguf.extract_quant_label("Foo-BF16-Q8_0.gguf"), + } + assert labels == {"Q4_K_M", "Q8_0"} + + +@pytest.mark.parametrize("repo_id", ["bert-base-uncased", "owner/repo"]) +def test_repo_id_validation_accepts_hf_repo_id_contract(repo_id): + assert paths.is_valid_repo_id(repo_id) + + +@pytest.mark.parametrize( + "repo_id", + [ + "datasets/foo/bar", + ".repo", + "repo.git", + "foo..bar", + "foo--bar", + "../repo", + "owner/../repo", + ], +) +def test_repo_id_validation_rejects_unsafe_or_invalid_ids(repo_id): + assert not paths.is_valid_repo_id(repo_id) + + +class _RecordingLogger: + def __init__(self): + self.warnings = [] + + def warning(self, *args, **kwargs): + self.warnings.append((args, kwargs)) + + +def test_resolve_browse_target_preserves_allowlist_and_symlink_safety(tmp_path): + home = tmp_path / "home" + scan = tmp_path / "scan" + target = scan / "nested" + home.mkdir() + target.mkdir(parents = True) + (home / "scan-link").symlink_to(scan, target_is_directory = True) + + resolved = folder_browser._resolve_browse_target( + str(home / "scan-link" / "nested"), + [home, scan], + ) + + assert resolved == target.resolve() + + +def test_resolve_browse_target_rejects_outside_allowlist(tmp_path): + allowed = tmp_path / "allowed" + outside = tmp_path / "outside" + allowed.mkdir() + outside.mkdir() + + with pytest.raises(HTTPException) as exc_info: + folder_browser._resolve_browse_target(str(outside), [allowed]) + + assert exc_info.value.status_code == 403 + + +def test_resolve_browse_target_rejects_sensitive_dir(tmp_path): + home = tmp_path / "home" + ssh = home / ".ssh" + ssh.mkdir(parents = True) + + with pytest.raises(HTTPException) as exc_info: + folder_browser._resolve_browse_target(str(ssh), [home]) + + assert exc_info.value.status_code == 403 + + +def test_browse_folders_hides_sensitive_dirs(monkeypatch, tmp_path): + home = tmp_path / "home" + (home / ".ssh").mkdir(parents = True) + (home / "models").mkdir() + monkeypatch.setattr(folder_browser, "_build_browse_allowlist", lambda: [home]) + + response = folder_browser.browse_folders_response(str(home), show_hidden = True) + + names = {entry.name for entry in response.entries} + assert "models" in names + assert ".ssh" not in names + + +def test_contained_link_path_confines_to_link_dir(tmp_path): + link_dir = tmp_path / "ollama" / ".studio_links" / "abc123" + + legit = ollama._contained_link_path(link_dir, "llama3-latest-Q4_K_M.gguf") + assert legit == link_dir / "llama3-latest-Q4_K_M.gguf" + + for unsafe in ( + "", + ".", + "..", + "a/b.gguf", + "../evil.gguf", + "/etc/passwd", + "model-tag-../../../pwned.gguf", + ): + assert ollama._contained_link_path(link_dir, unsafe) is None + + +def test_make_ollama_blob_link_refuses_escaping_name(tmp_path): + root = tmp_path / "ollama" + link_dir = root / ".studio_links" / "abc123" + blob = root / "blobs" / "sha256-deadbeef" + blob.parent.mkdir(parents = True) + blob.write_bytes(b"weights") + + escaped = ollama._make_ollama_blob_link(link_dir, "model-tag-../../../pwned.gguf", blob) + assert escaped is None + assert not list(tmp_path.rglob("pwned.gguf")) + + safe = ollama._make_ollama_blob_link(link_dir, "model-tag.gguf", blob) + assert safe == str(link_dir / "model-tag.gguf") + assert (link_dir / "model-tag.gguf").exists() + + +def test_cached_gguf_scan_dedupes_and_excludes_mmproj_only(monkeypatch, tmp_path): + smaller = _repo("Org/Dupe", [_file("Q4_K_M.gguf", 100)], tmp_path / "small") + larger = _repo( + "org/dupe", + [_file("Q4_K_M.gguf", 300), _file("Q8_0.gguf", 200)], + tmp_path / "large", + ) + mmproj_only = _repo("Org/VisionAdapter", [_file("mmproj-F16.gguf", 900)], tmp_path / "mmproj") + monkeypatch.setattr( + cache_inventory, + "all_hf_cache_scans", + lambda: [SimpleNamespace(repos = [smaller, larger, mmproj_only])], + ) + monkeypatch.setattr( + cache_inventory.hf_cache_scan, + "is_gguf_repo_partial", + lambda _repo_id, _path: False, + ) + + result = {"cached": cache_inventory._scan_cached_gguf()} + + assert [row["repo_id"] for row in result["cached"]] == ["org/dupe"] + assert result["cached"][0]["size_bytes"] == 500 + assert result["cached"][0]["model_format"] == "gguf" + assert result["cached"][0]["capabilities"]["requires_variant"] is True + + +def test_cached_gguf_scan_preserves_partial_flag(monkeypatch, tmp_path): + partial = _repo("Org/Partial", [_file("Q4_K_M.gguf", 100)], tmp_path / "partial") + monkeypatch.setattr( + cache_inventory, + "all_hf_cache_scans", + lambda: [SimpleNamespace(repos = [partial])], + ) + monkeypatch.setattr( + cache_inventory.hf_cache_scan, + "is_gguf_repo_partial", + lambda _repo_id, _path: True, + ) + + result = {"cached": cache_inventory._scan_cached_gguf()} + row = result["cached"][0] + + assert row["partial"] is True + assert row["partial_transport"] is None + assert row["capabilities"]["can_chat"] is False + + +def test_cached_gguf_scan_includes_variant_state_without_completed_gguf(monkeypatch, tmp_path): + monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state") + repo_path = tmp_path / "hub" / "models--Org--PartialGguf" + repo_path.mkdir(parents = True) + partial = _repo( + "Org/PartialGguf", + [_file("config.json", 12)], + repo_path, + ) + assert download_manifest.write_manifest( + "model", + "Org/PartialGguf", + "Q4_K_M", + [download_manifest.ExpectedFile(path = "model-Q4_K_M.gguf", size = 4096)], + "http", + ) + assert download_manifest.write_cancel_marker("model", "Org/PartialGguf", "Q4_K_M", "http") + monkeypatch.setattr( + cache_inventory, + "all_hf_cache_scans", + lambda: [SimpleNamespace(repos = [partial])], + ) + monkeypatch.setattr( + cache_inventory.hf_cache_scan, + "is_gguf_repo_partial", + lambda _repo_id, _path: True, + ) + + result = {"cached": cache_inventory._scan_cached_gguf()} + row = result["cached"][0] + + assert row["repo_id"] == "Org/PartialGguf" + assert row["model_format"] == "gguf" + assert row["size_bytes"] == 4096 + assert row["partial"] is True + assert row["capabilities"]["requires_variant"] is True + + +def test_gguf_variant_requirements_include_split_files_and_preferred_mmproj(): + requirements = gguf_variants._build_gguf_variant_requirements( + [ + _sibling("model-Q4_K_M-00001-of-00002.gguf", 10, "main-a"), + _sibling("model-Q4_K_M-00002-of-00002.gguf", 20, "main-b"), + _sibling("mmproj-BF16.gguf", 7, "mm-bf16"), + _sibling("mmproj-F16.gguf", 5, "mm-f16"), + ] + ) + + req = requirements["q4_k_m"] + + assert req.main_size_bytes == 30 + assert req.download_size_bytes == 35 + assert req.main_hashes == frozenset({"main-a", "main-b"}) + assert req.required_hashes == frozenset({"main-a", "main-b", "mm-f16"}) + assert req.companion_hashes == frozenset({"mm-f16"}) + assert req.mmproj_hashes == frozenset({"mm-bf16", "mm-f16"}) + assert req.target_filenames == ( + "model-Q4_K_M-00001-of-00002.gguf", + "model-Q4_K_M-00002-of-00002.gguf", + "mmproj-F16.gguf", + ) + + +def test_worker_gguf_variant_plan_matches_service_requirement(monkeypatch): + siblings = [ + _sibling("model-Q4_K_M-00001-of-00002.gguf", 10, "main-a"), + _sibling("model-Q4_K_M-00002-of-00002.gguf", 20, "main-b"), + _sibling("mmproj-BF16.gguf", 7, "mm-bf16"), + _sibling("mmproj-F16.gguf", 5, "mm-f16"), + ] + monkeypatch.setattr( + hf_download, + "_model_info_with_retry", + lambda *_args, **_kwargs: SimpleNamespace(siblings = siblings), + ) + + service_req = gguf_variants._build_gguf_variant_requirements(siblings)["q4_k_m"] + worker_plan = hf_download._gguf_variant_target_plan("Org/Vision", "Q4_K_M", None) + + assert worker_plan == service_req + + +def test_gguf_variant_blob_hashes_accept_dict_lfs_fallback(monkeypatch): + with gguf_variants._VARIANT_HASH_LOCK: + gguf_variants._VARIANT_HASH_CACHE.clear() + gguf_variants._VARIANT_REQUIREMENT_CACHE.clear() + monkeypatch.setitem( + sys.modules, + "huggingface_hub", + SimpleNamespace( + HfApi = lambda *_args, **_kwargs: SimpleNamespace( + model_info = lambda *_a, **_k: SimpleNamespace( + siblings = [ + _sibling("model-Q4_K_M.gguf", 10, "main-dict"), + _sibling("model-Q8_0.gguf", 20, "other"), + _sibling("mmproj-F16.gguf", 5, "mmproj"), + ] + ) + ) + ), + ) + + result = gguf_variants.gguf_variant_blob_hashes("Org/DictLfs", "Q4_K_M", None) + main_only = gguf_variants.gguf_variant_blob_hashes( + "Org/DictLfs", + "Q4_K_M", + None, + include_companions = False, + ) + + assert result == frozenset({"main-dict", "mmproj"}) + assert main_only == frozenset({"main-dict"}) + + +def test_gguf_variant_blob_hashes_skip_missing_rfilename(monkeypatch): + with gguf_variants._VARIANT_HASH_LOCK: + gguf_variants._VARIANT_HASH_CACHE.clear() + gguf_variants._VARIANT_REQUIREMENT_CACHE.clear() + siblings = [ + SimpleNamespace(rfilename = None, size = 1, lfs = {"sha256": "bad"}), + _sibling("model-Q4_K_M.gguf", 10, "main"), + ] + monkeypatch.setattr( + gguf_variants, + "_fetch_gguf_variant_requirements", + lambda _repo_id, _hf_token = None: gguf_variants._build_gguf_variant_requirements(siblings), + ) + + result = gguf_variants.gguf_variant_blob_hashes("Org/Malformed", "Q4_K_M", None) + + assert result == frozenset({"main"}) + + +def test_worker_gguf_variant_targets_skip_missing_rfilename(monkeypatch): + monkeypatch.setattr( + hf_download, + "_model_info_with_retry", + lambda *_args, **_kwargs: SimpleNamespace( + siblings = [ + SimpleNamespace(rfilename = None, size = 1), + _sibling("model-Q4_K_M.gguf", 10, "main"), + _sibling("mmproj-F16.gguf", 5, "mm"), + ] + ), + ) + + result = hf_download._gguf_variant_target_plan("Org/Malformed", "Q4_K_M", None) + + assert list(result.target_filenames) == ["model-Q4_K_M.gguf", "mmproj-F16.gguf"] + + +def test_download_gguf_variant_purges_only_main_quant_hashes(monkeypatch, tmp_path): + prepare_calls = [] + snapshot_calls = [] + written = [] + verified = [] + + monkeypatch.setattr( + hf_download, + "_model_info_with_retry", + lambda *_args, **_kwargs: SimpleNamespace( + siblings = [ + _sibling("model-Q4_K_M.gguf", 10, "q4-main"), + _sibling("model-Q8_0.gguf", 20, "q8-main"), + _sibling("mmproj-F16.gguf", 5, "shared-mmproj"), + ] + ), + ) + monkeypatch.setattr( + hf_download, "_verify_completed_download", lambda *args, **kwargs: verified.append(args) + ) + monkeypatch.setattr( + download_registry, + "prepare_cache_for_transport", + lambda *args, **kwargs: prepare_calls.append((args, kwargs)) or 0, + ) + monkeypatch.setattr(download_manifest, "clear_cancel_marker", lambda *_args: None) + monkeypatch.setattr( + download_manifest, "write_manifest", lambda *args: written.append(args) or True + ) + monkeypatch.setitem( + sys.modules, + "huggingface_hub", + SimpleNamespace( + snapshot_download = lambda **kwargs: snapshot_calls.append(kwargs) or str(tmp_path) + ), + ) + + hf_download._download_gguf_variant("Org/Vision", "Q4_K_M", None, "http") + + assert prepare_calls == [ + ( + ("model", "Org/Vision", "http", "Q4_K_M"), + { + "only_blob_hashes": frozenset({"q4-main"}), + "companion_blob_hashes": frozenset({"shared-mmproj"}), + "protected_blob_hashes": frozenset(), + }, + ) + ] + assert [file.path for file in written[0][3]] == ["model-Q4_K_M.gguf", "mmproj-F16.gguf"] + assert snapshot_calls[0]["allow_patterns"] == ["model-Q4_K_M.gguf", "mmproj-F16.gguf"] + assert verified == [("model", "Org/Vision", "Q4_K_M", str(tmp_path))] + + +def test_download_gguf_variant_manifest_resume_purges_only_main_quant_hashes(monkeypatch, tmp_path): + prepare_calls = [] + snapshot_calls = [] + + def _metadata_unavailable(*_args, **_kwargs): + raise RuntimeError("metadata down") + + manifest = download_manifest.Manifest( + repo_type = "model", + repo_id = "Org/Vision", + variant = "Q4_K_M", + started_at = "", + expected_files = ( + download_manifest.ExpectedFile( + path = "model-Q4_K_M.gguf", + size = 10, + sha256 = "q4-main", + ), + download_manifest.ExpectedFile( + path = "mmproj-F16.gguf", + size = 5, + sha256 = "shared-mmproj", + ), + ), + transport = "http", + ) + monkeypatch.setattr( + hf_download, + "_gguf_variant_target_plan", + _metadata_unavailable, + ) + monkeypatch.setattr(download_manifest, "read_manifest", lambda *_args: manifest) + monkeypatch.setattr(download_manifest, "clear_cancel_marker", lambda *_args: None) + monkeypatch.setattr( + download_registry, + "prepare_cache_for_transport", + lambda *args, **kwargs: prepare_calls.append((args, kwargs)) or 0, + ) + monkeypatch.setattr(hf_download, "_verify_completed_download", lambda *_args, **_kwargs: None) + monkeypatch.setitem( + sys.modules, + "huggingface_hub", + SimpleNamespace( + snapshot_download = lambda **kwargs: snapshot_calls.append(kwargs) or str(tmp_path) + ), + ) + + hf_download._download_gguf_variant("Org/Vision", "Q4_K_M", None, "http") + + assert prepare_calls == [ + ( + ("model", "Org/Vision", "http", "Q4_K_M"), + { + "only_blob_hashes": frozenset({"q4-main"}), + "companion_blob_hashes": frozenset({"shared-mmproj"}), + "protected_blob_hashes": frozenset(), + }, + ) + ] + assert snapshot_calls[0]["allow_patterns"] == ["model-Q4_K_M.gguf", "mmproj-F16.gguf"] + + +def test_download_snapshot_recovers_manifest_after_metadata_fallback(monkeypatch, tmp_path): + metadata_calls = [] + written = [] + cleared = [] + verified = [] + + def _metadata(*_args, **_kwargs): + metadata_calls.append(True) + if len(metadata_calls) == 1: + raise RuntimeError("metadata down") + return SimpleNamespace(siblings = [SimpleNamespace(rfilename = "config.json", size = 12)]) + + monkeypatch.setattr(hf_download, "_model_info_with_retry", _metadata) + monkeypatch.setattr( + hf_download, "_verify_completed_download", lambda *args, **kwargs: verified.append(args) + ) + monkeypatch.setattr( + download_registry, "prepare_cache_for_transport", lambda *_args, **_kwargs: 0 + ) + monkeypatch.setattr( + download_manifest, "clear_cancel_marker", lambda *args: cleared.append(args) + ) + monkeypatch.setattr( + download_manifest, "write_manifest", lambda *args: written.append(args) or True + ) + monkeypatch.setitem( + sys.modules, + "huggingface_hub", + SimpleNamespace(snapshot_download = lambda **_kwargs: str(tmp_path)), + ) + + hf_download._download_snapshot("Org/Model", None, "http") + + assert len(metadata_calls) == 2 + assert cleared == [("model", "Org/Model", None)] + assert written[0][0:3] == ("model", "Org/Model", None) + assert written[0][3][0].path == "config.json" + assert verified == [("model", "Org/Model", None, str(tmp_path))] + + +def test_download_dataset_continues_without_metadata_manifest(monkeypatch, tmp_path): + metadata_calls = [] + snapshot_calls = [] + written = [] + cleared = [] + verified = [] + + def _metadata(*_args, **_kwargs): + metadata_calls.append(True) + raise RuntimeError("metadata down") + + monkeypatch.setattr(hf_download, "_dataset_info_with_retry", _metadata) + monkeypatch.setattr( + hf_download, "_verify_completed_download", lambda *args, **kwargs: verified.append(args) + ) + monkeypatch.setattr( + download_registry, "prepare_cache_for_transport", lambda *_args, **_kwargs: 0 + ) + monkeypatch.setattr( + download_manifest, "clear_cancel_marker", lambda *args: cleared.append(args) + ) + monkeypatch.setattr( + download_manifest, "write_manifest", lambda *args: written.append(args) or True + ) + monkeypatch.setattr( + hf_cache_state, "has_active_incomplete_blobs", lambda *_args, **_kwargs: False + ) + monkeypatch.setitem( + sys.modules, + "huggingface_hub", + SimpleNamespace( + snapshot_download = lambda **kwargs: snapshot_calls.append(kwargs) or str(tmp_path) + ), + ) + + hf_download._download_dataset("Org/Data", None, "http") + + assert len(metadata_calls) == 2 + assert cleared == [("dataset", "Org/Data", None)] + assert written == [] + assert snapshot_calls == [ + { + "repo_id": "Org/Data", + "token": False, + "repo_type": "dataset", + "max_workers": 1, + } + ] + assert verified == [("dataset", "Org/Data", None, str(tmp_path))] + + +def test_download_snapshot_fails_when_metadata_unavailable_and_partial_remains( + monkeypatch, tmp_path +): + """No prior manifest + metadata unavailable + leftover .incomplete blobs means + a cached partial was returned without downloading: the worker must exit 1, not + derive a self-certifying manifest from the finalized subset.""" + written = [] + verified = [] + + def _metadata(*_args, **_kwargs): + raise RuntimeError("metadata down") + + monkeypatch.setattr(hf_download, "_model_info_with_retry", _metadata) + monkeypatch.setattr( + hf_download, "_verify_completed_download", lambda *args, **kwargs: verified.append(args) + ) + monkeypatch.setattr( + download_registry, "prepare_cache_for_transport", lambda *_args, **_kwargs: 0 + ) + monkeypatch.setattr(download_manifest, "clear_cancel_marker", lambda *_args: None) + monkeypatch.setattr(download_manifest, "read_manifest", lambda *_args, **_kwargs: None) + monkeypatch.setattr( + download_manifest, "write_manifest", lambda *args: written.append(args) or True + ) + monkeypatch.setattr( + hf_cache_state, "has_active_incomplete_blobs", lambda *_args, **_kwargs: True + ) + monkeypatch.setitem( + sys.modules, + "huggingface_hub", + SimpleNamespace(snapshot_download = lambda **_kwargs: str(tmp_path)), + ) + + with pytest.raises(SystemExit) as excinfo: + hf_download._download_snapshot("Org/Model", None, "http") + + assert excinfo.value.code == 1 + assert written == [] + assert verified == [] + + +def test_purge_repo_cache_dirs_skips_top_level_symlink(monkeypatch, tmp_path): + root = tmp_path / "hub" + target = tmp_path / "target" + root.mkdir() + target.mkdir() + link = root / "models--Org--Repo" + link.symlink_to(target, target_is_directory = True) + monkeypatch.setattr(hf_cache_state, "hf_cache_roots", lambda: [root]) + + removed = hf_cache_state.purge_repo_cache_dirs("model", "Org/Repo") + + assert removed is False + assert link.is_symlink() + assert target.is_dir() + + +def test_gguf_download_progress_fallback_logs_warning(monkeypatch): + token = "hf_12345678901234567890" + logger = _RecordingLogger() + + async def _run_inline(fn, *args, **kwargs): + return fn(*args, **kwargs) + + def _raise_permission_error(*_args, **_kwargs): + raise PermissionError(f"denied {token}") + + monkeypatch.setattr(snapshot_progress, "logger", logger) + monkeypatch.setattr(downloads.asyncio, "to_thread", _run_inline) + monkeypatch.setattr( + downloads.gguf_variants, + "gguf_variant_requirements", + lambda *_args, **_kwargs: None, + ) + monkeypatch.setattr( + downloads.gguf_variants, + "gguf_variant_blob_hashes", + lambda *_args, **_kwargs: frozenset(), + ) + monkeypatch.setattr( + snapshot_progress, + "preferred_repo_cache_dirs", + _raise_permission_error, + ) + monkeypatch.setattr( + downloads, + "_registry", + SimpleNamespace(get_job = lambda _key: SimpleNamespace(state = "running")), + ) + + result = asyncio.run( + downloads.get_gguf_download_progress_response( + "Org/Model", + variant = "Q4_K_M", + expected_bytes = -1, + hf_token = token, + ) + ) + + assert result == { + "downloaded_bytes": 0, + "completed_bytes": 0, + "complete_on_disk": False, + "expected_bytes": 0, + "progress": 0, + "cache_path": None, + } + assert logger.warnings + args, kwargs = logger.warnings[0] + assert args[:4] == ( + "Error checking %s download progress for %s: %s: %s", + "model", + "Org/Model", + "PermissionError", + ) + assert token not in args[4] + assert "***" in args[4] + assert kwargs == {} + + +def test_gguf_progress_counts_completed_mmproj_with_expected_bytes(monkeypatch, tmp_path): + """A finished mmproj companion keeps counting toward progress once the caller + supplies expected bytes; resolving the variant requirement credits it.""" + entry = tmp_path / "models--Org--Model-GGUF" + snap = entry / "snapshots" / "rev0" + blobs = entry / "blobs" + snap.mkdir(parents = True) + blobs.mkdir(parents = True) + (snap / "model-Q4_K_M.gguf").write_bytes(b"x" * 100) + (snap / "mmproj-F16.gguf").write_bytes(b"y" * 30) + (blobs / "mainhash").write_bytes(b"x" * 100) + (blobs / "mmprojhash").write_bytes(b"y" * 30) + monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state") + assert download_manifest.write_manifest( + "model", + "Org/Model-GGUF", + "Q4_K_M", + [ + download_manifest.ExpectedFile( + path = "model-Q4_K_M.gguf", + size = 100, + sha256 = "mainhash", + ), + download_manifest.ExpectedFile( + path = "mmproj-F16.gguf", + size = 30, + sha256 = "mmprojhash", + ), + ], + "http", + ) + + requirement = gguf_variants._GgufVariantRequirement( + main_filenames = frozenset({"model-Q4_K_M.gguf"}), + target_filenames = ("model-Q4_K_M.gguf", "mmproj-F16.gguf"), + main_hashes = frozenset({"mainhash"}), + required_hashes = frozenset({"mainhash", "mmprojhash"}), + companion_hashes = frozenset({"mmprojhash"}), + mmproj_filenames = frozenset({"mmproj-F16.gguf"}), + mmproj_hashes = frozenset({"mmprojhash"}), + expected_files = ( + download_manifest.ExpectedFile( + path = "model-Q4_K_M.gguf", + size = 100, + sha256 = "mainhash", + ), + download_manifest.ExpectedFile( + path = "mmproj-F16.gguf", + size = 30, + sha256 = "mmprojhash", + ), + ), + main_size_bytes = 100, + download_size_bytes = 130, + ) + + async def _run_inline(fn, *args, **kwargs): + return fn(*args, **kwargs) + + monkeypatch.setattr(downloads.asyncio, "to_thread", _run_inline) + monkeypatch.setattr( + downloads.gguf_variants, + "gguf_variant_requirements", + lambda *_args, **_kwargs: requirement, + ) + monkeypatch.setattr( + snapshot_progress, + "preferred_repo_cache_dirs", + lambda *_args, **_kwargs: [entry], + ) + monkeypatch.setattr( + downloads, + "_registry", + SimpleNamespace(get_job = lambda _key: SimpleNamespace(state = "idle")), + ) + + result = asyncio.run( + downloads.get_gguf_download_progress_response( + "Org/Model-GGUF", + variant = "Q4_K_M", + expected_bytes = 130, + ) + ) + + assert result["completed_bytes"] == 130 + assert result["downloaded_bytes"] == 130 + assert result["complete_on_disk"] is True + assert result["progress"] == 1.0 + + +def test_gguf_progress_subtracts_new_job_completed_baseline(monkeypatch, tmp_path): + entry = tmp_path / "models--Org--Model-GGUF" + snap = entry / "snapshots" / "rev0" + blobs = entry / "blobs" + snap.mkdir(parents = True) + blobs.mkdir(parents = True) + (snap / "mmproj-F16.gguf").write_bytes(b"y" * 30) + (blobs / "mmprojhash").write_bytes(b"y" * 30) + monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state") + assert download_manifest.write_manifest( + "model", + "Org/Model-GGUF", + "Q4_K_M", + [ + download_manifest.ExpectedFile( + path = "model-Q4_K_M.gguf", + size = 100, + sha256 = "mainhash", + ), + download_manifest.ExpectedFile( + path = "mmproj-F16.gguf", + size = 30, + sha256 = "mmprojhash", + ), + ], + "http", + ) + + requirement = gguf_variants._GgufVariantRequirement( + main_filenames = frozenset({"model-Q4_K_M.gguf"}), + target_filenames = ("model-Q4_K_M.gguf", "mmproj-F16.gguf"), + main_hashes = frozenset({"mainhash"}), + required_hashes = frozenset({"mainhash", "mmprojhash"}), + companion_hashes = frozenset({"mmprojhash"}), + mmproj_filenames = frozenset({"mmproj-F16.gguf"}), + mmproj_hashes = frozenset({"mmprojhash"}), + expected_files = ( + download_manifest.ExpectedFile( + path = "model-Q4_K_M.gguf", + size = 100, + sha256 = "mainhash", + ), + download_manifest.ExpectedFile( + path = "mmproj-F16.gguf", + size = 30, + sha256 = "mmprojhash", + ), + ), + main_size_bytes = 100, + download_size_bytes = 130, + ) + + async def _run_inline(fn, *args, **kwargs): + return fn(*args, **kwargs) + + monkeypatch.setattr(downloads.asyncio, "to_thread", _run_inline) + monkeypatch.setattr( + downloads.gguf_variants, + "gguf_variant_requirements", + lambda *_args, **_kwargs: requirement, + ) + monkeypatch.setattr( + snapshot_progress, + "preferred_repo_cache_dirs", + lambda *_args, **_kwargs: [entry], + ) + monkeypatch.setattr( + downloads, + "_registry", + SimpleNamespace( + get_job = lambda _key: SimpleNamespace(state = "running"), + get_job_metadata = lambda _key: SimpleNamespace( + completed_baseline_bytes = 30, + ), + ), + ) + + result = asyncio.run( + downloads.get_gguf_download_progress_response( + "Org/Model-GGUF", + variant = "Q4_K_M", + expected_bytes = 130, + ) + ) + + assert result["completed_bytes"] == 0 + assert result["downloaded_bytes"] == 0 + assert result["expected_bytes"] == 100 + assert result["complete_on_disk"] is False + assert result["progress"] == 0 + + +def test_gguf_progress_shows_main_when_companion_left_the_count(monkeypatch, tmp_path): + # The mmproj companion that seeded the baseline is gone, so completed_bytes + # is main-only and below the baseline; it must not be subtracted to 0. + entry = tmp_path / "models--Org--Model-GGUF" + blobs = entry / "blobs" + blobs.mkdir(parents = True) + (blobs / "mainhash").write_bytes(b"x" * 20) + monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state") + + requirement = gguf_variants._GgufVariantRequirement( + main_filenames = frozenset({"model-Q4_K_M.gguf"}), + target_filenames = ("model-Q4_K_M.gguf", "mmproj-F16.gguf"), + main_hashes = frozenset({"mainhash"}), + required_hashes = frozenset({"mainhash", "mmprojhash"}), + companion_hashes = frozenset({"mmprojhash"}), + mmproj_filenames = frozenset({"mmproj-F16.gguf"}), + mmproj_hashes = frozenset({"mmprojhash"}), + expected_files = ( + download_manifest.ExpectedFile( + path = "model-Q4_K_M.gguf", + size = 100, + sha256 = "mainhash", + ), + download_manifest.ExpectedFile( + path = "mmproj-F16.gguf", + size = 30, + sha256 = "mmprojhash", + ), + ), + main_size_bytes = 100, + download_size_bytes = 130, + ) + + async def _run_inline(fn, *args, **kwargs): + return fn(*args, **kwargs) + + monkeypatch.setattr(downloads.asyncio, "to_thread", _run_inline) + monkeypatch.setattr( + downloads.gguf_variants, + "gguf_variant_requirements", + lambda *_args, **_kwargs: requirement, + ) + monkeypatch.setattr( + snapshot_progress, + "preferred_repo_cache_dirs", + lambda *_args, **_kwargs: [entry], + ) + monkeypatch.setattr( + downloads, + "_registry", + SimpleNamespace( + get_job = lambda _key: SimpleNamespace(state = "running"), + get_job_metadata = lambda _key: SimpleNamespace( + completed_baseline_bytes = 30, + ), + ), + ) + + result = asyncio.run( + downloads.get_gguf_download_progress_response( + "Org/Model-GGUF", + variant = "Q4_K_M", + expected_bytes = 130, + ) + ) + + assert result["completed_bytes"] == 20 + assert result["downloaded_bytes"] == 20 + assert result["expected_bytes"] == 130 + assert result["complete_on_disk"] is False + + +def test_gguf_progress_complete_on_disk_ignores_full_baseline(monkeypatch, tmp_path): + # A variant already complete on disk carries a baseline equal to its full + # size; subtracting it would report 0/0 for a finished variant (frontend + # evicts it as gone). Once complete_on_disk is verified, the full figures + # must survive. + entry = tmp_path / "models--Org--Model-GGUF" + snap = entry / "snapshots" / "rev0" + blobs = entry / "blobs" + snap.mkdir(parents = True) + blobs.mkdir(parents = True) + (snap / "model-Q4_K_M.gguf").write_bytes(b"x" * 100) + (snap / "mmproj-F16.gguf").write_bytes(b"y" * 30) + (blobs / "mainhash").write_bytes(b"x" * 100) + (blobs / "mmprojhash").write_bytes(b"y" * 30) + monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state") + assert download_manifest.write_manifest( + "model", + "Org/Model-GGUF", + "Q4_K_M", + [ + download_manifest.ExpectedFile( + path = "model-Q4_K_M.gguf", + size = 100, + sha256 = "mainhash", + ), + download_manifest.ExpectedFile( + path = "mmproj-F16.gguf", + size = 30, + sha256 = "mmprojhash", + ), + ], + "http", + ) + + requirement = gguf_variants._GgufVariantRequirement( + main_filenames = frozenset({"model-Q4_K_M.gguf"}), + target_filenames = ("model-Q4_K_M.gguf", "mmproj-F16.gguf"), + main_hashes = frozenset({"mainhash"}), + required_hashes = frozenset({"mainhash", "mmprojhash"}), + companion_hashes = frozenset({"mmprojhash"}), + mmproj_filenames = frozenset({"mmproj-F16.gguf"}), + mmproj_hashes = frozenset({"mmprojhash"}), + expected_files = ( + download_manifest.ExpectedFile( + path = "model-Q4_K_M.gguf", + size = 100, + sha256 = "mainhash", + ), + download_manifest.ExpectedFile( + path = "mmproj-F16.gguf", + size = 30, + sha256 = "mmprojhash", + ), + ), + main_size_bytes = 100, + download_size_bytes = 130, + ) + + async def _run_inline(fn, *args, **kwargs): + return fn(*args, **kwargs) + + monkeypatch.setattr(downloads.asyncio, "to_thread", _run_inline) + monkeypatch.setattr( + downloads.gguf_variants, + "gguf_variant_requirements", + lambda *_args, **_kwargs: requirement, + ) + monkeypatch.setattr( + snapshot_progress, + "preferred_repo_cache_dirs", + lambda *_args, **_kwargs: [entry], + ) + monkeypatch.setattr( + downloads, + "_registry", + SimpleNamespace( + get_job = lambda _key: SimpleNamespace(state = "running"), + get_job_metadata = lambda _key: SimpleNamespace( + completed_baseline_bytes = 130, + ), + ), + ) + + result = asyncio.run( + downloads.get_gguf_download_progress_response( + "Org/Model-GGUF", + variant = "Q4_K_M", + expected_bytes = 130, + ) + ) + + assert result["complete_on_disk"] is True + assert result["completed_bytes"] == 130 + assert result["downloaded_bytes"] == 130 + assert result["expected_bytes"] == 130 + assert result["progress"] == 1.0 + + +def test_gguf_progress_scoped_hashes_exclude_sibling_quant(monkeypatch, tmp_path): + # The "instant ~900 MB" bug: a sibling quant is fully cached when a different + # variant starts. With this variant's hashes resolved, progress counts ONLY + # its in-progress blob, never the sibling's finalized bytes in the shared + # blobs/ dir. + entry = tmp_path / "models--Org--Model-GGUF" + blobs = entry / "blobs" + blobs.mkdir(parents = True) + (blobs / "siblinghash").write_bytes(b"z" * 900) # other quant, complete + (blobs / "mainhash.incomplete").write_bytes(b"x" * 5) # this variant, started + + requirement = gguf_variants._GgufVariantRequirement( + main_filenames = frozenset({"model-Q4_K_M.gguf"}), + target_filenames = ("model-Q4_K_M.gguf",), + main_hashes = frozenset({"mainhash"}), + required_hashes = frozenset({"mainhash"}), + companion_hashes = frozenset(), + mmproj_filenames = frozenset(), + mmproj_hashes = frozenset(), + expected_files = ( + download_manifest.ExpectedFile( + path = "model-Q4_K_M.gguf", + size = 100, + sha256 = "mainhash", + ), + ), + main_size_bytes = 100, + download_size_bytes = 100, + ) + + async def _run_inline(fn, *args, **kwargs): + return fn(*args, **kwargs) + + monkeypatch.setattr(downloads.asyncio, "to_thread", _run_inline) + monkeypatch.setattr( + downloads.gguf_variants, + "gguf_variant_requirements", + lambda *_args, **_kwargs: requirement, + ) + monkeypatch.setattr( + snapshot_progress, + "preferred_repo_cache_dirs", + lambda *_args, **_kwargs: [entry], + ) + monkeypatch.setattr( + downloads, + "_registry", + SimpleNamespace( + get_job = lambda _key: SimpleNamespace(state = "running"), + get_job_metadata = lambda _key: SimpleNamespace( + completed_baseline_bytes = 0, + ), + ), + ) + + result = asyncio.run( + downloads.get_gguf_download_progress_response( + "Org/Model-GGUF", + variant = "Q4_K_M", + expected_bytes = 100, + ) + ) + + assert result["completed_bytes"] == 0 + assert result["downloaded_bytes"] == 5 + + +def test_gguf_progress_unknown_hashes_does_not_count_foreign_blobs(monkeypatch, tmp_path): + # With a variant's hashes unresolved (metadata flaked, no manifest), the + # shared blobs/ dir's FINALIZED blobs must NOT be counted wholesale: a cached + # sibling quant (``siblinghash``) alongside is the "instant ~900 MB" bug. + # With no .incomplete present, downloaded must be 0. + entry = tmp_path / "models--Org--Model-GGUF" + snap = entry / "snapshots" / "rev0" + blobs = entry / "blobs" + snap.mkdir(parents = True) + blobs.mkdir(parents = True) + (blobs / "mainhash").write_bytes(b"x" * 100) + (blobs / "mmprojhash").write_bytes(b"y" * 30) + (blobs / "siblinghash").write_bytes(b"z" * 900) + + async def _run_inline(fn, *args, **kwargs): + return fn(*args, **kwargs) + + monkeypatch.setattr(downloads.asyncio, "to_thread", _run_inline) + monkeypatch.setattr( + downloads.gguf_variants, + "gguf_variant_requirements", + lambda *_args, **_kwargs: None, + ) + monkeypatch.setattr( + downloads.gguf_variants, + "gguf_variant_blob_hashes", + lambda *_args, **_kwargs: frozenset(), + ) + monkeypatch.setattr( + snapshot_progress, + "preferred_repo_cache_dirs", + lambda *_args, **_kwargs: [entry], + ) + monkeypatch.setattr( + downloads, + "_registry", + SimpleNamespace(get_job = lambda _key: SimpleNamespace(state = "running")), + ) + + result = asyncio.run( + downloads.get_gguf_download_progress_response( + "Org/Model-GGUF", + variant = "Q4_K_M", + expected_bytes = 130, + ) + ) + + assert result["completed_bytes"] == 0 + assert result["downloaded_bytes"] == 0 + assert result["complete_on_disk"] is False + + +def test_gguf_progress_unknown_hashes_drops_unscoped_incomplete_blob(monkeypatch, tmp_path): + # With hashes unresolved, an .incomplete in the shared blobs/ dir can't be + # attributed to this variant (it may be a concurrent sibling's active write), + # so it is dropped, mirroring the finalized-blob guard. In production the + # worker writes the manifest before any .incomplete exists, so hashes resolve + # via the manifest backstop and this window never suppresses real progress. + entry = tmp_path / "models--Org--Model-GGUF" + blobs = entry / "blobs" + blobs.mkdir(parents = True) + (blobs / "activehash.incomplete").write_bytes(b"x" * 50) # unattributable + (blobs / "siblinghash").write_bytes(b"z" * 900) # finalized sibling + + async def _run_inline(fn, *args, **kwargs): + return fn(*args, **kwargs) + + monkeypatch.setattr(downloads.asyncio, "to_thread", _run_inline) + monkeypatch.setattr( + downloads.gguf_variants, + "gguf_variant_requirements", + lambda *_args, **_kwargs: None, + ) + monkeypatch.setattr( + downloads.gguf_variants, + "gguf_variant_blob_hashes", + lambda *_args, **_kwargs: frozenset(), + ) + monkeypatch.setattr( + snapshot_progress, + "preferred_repo_cache_dirs", + lambda *_args, **_kwargs: [entry], + ) + monkeypatch.setattr( + downloads, + "_registry", + SimpleNamespace(get_job = lambda _key: SimpleNamespace(state = "running")), + ) + + result = asyncio.run( + downloads.get_gguf_download_progress_response( + "Org/Model-GGUF", + variant = "Q4_K_M", + expected_bytes = 1000, + ) + ) + + assert result["downloaded_bytes"] == 0 # unscoped .incomplete not leaked + assert result["completed_bytes"] == 0 # finalized sibling still ignored + + +def test_gguf_progress_unknown_hashes_no_backward_dip_when_variant_finalizes(monkeypatch, tmp_path): + # Regression for the two-variant dip: with hashes unresolved, the first quant + # finalizes while the sibling still writes its .incomplete. The sibling's + # bytes used to leak into this numerator, dipping the bar ~99% -> ~78% for + # one poll. The unscoped .incomplete must be dropped so the reading stays 0. + entry = tmp_path / "models--unsloth--SmolLM2-360M-Instruct-GGUF" + blobs = entry / "blobs" + snap = entry / "snapshots" / "rev0" + blobs.mkdir(parents = True) + snap.mkdir(parents = True) + own_total = 218_673_760 # Q2_K finished blob size (denominator) + sibling_total = 234_686_560 # Q3_K_M total + + def _sparse_file(path: Path, size: int) -> None: + with path.open("wb") as handle: + handle.truncate(size) + + own_finalized = blobs / "q2hash" + _sparse_file(own_finalized, own_total) + # ~72.7% of the sibling => sibling_partial / own_total == 0.78 pre-fix. + sibling_incomplete = blobs / "q3hash.incomplete" + _sparse_file(sibling_incomplete, int(sibling_total * 0.727)) + + async def _run_inline(fn, *args, **kwargs): + return fn(*args, **kwargs) + + monkeypatch.setattr(downloads.asyncio, "to_thread", _run_inline) + monkeypatch.setattr( + downloads.gguf_variants, + "gguf_variant_requirements", + lambda *_args, **_kwargs: None, + ) + monkeypatch.setattr( + downloads.gguf_variants, + "gguf_variant_blob_hashes", + lambda *_args, **_kwargs: frozenset(), + ) + monkeypatch.setattr( + snapshot_progress, + "preferred_repo_cache_dirs", + lambda *_args, **_kwargs: [entry], + ) + monkeypatch.setattr( + downloads, + "_registry", + SimpleNamespace(get_job = lambda _key: SimpleNamespace(state = "running")), + ) + + result = asyncio.run( + downloads.get_gguf_download_progress_response( + "unsloth/SmolLM2-360M-Instruct-GGUF", + variant = "Q2_K", + expected_bytes = own_total, + ) + ) + + assert result["downloaded_bytes"] == 0 # sibling .incomplete did not leak + assert result["progress"] == 0 # no ~0.78 backward dip + + +def test_hf_cache_model_file_probe_is_bounded(monkeypatch, tmp_path): + snapshot = tmp_path / "snapshot" + snapshot.mkdir() + first = tmp_path / "README.md" + second = tmp_path / "notes.txt" + model = tmp_path / "model.safetensors" + first.write_text("readme", encoding = "utf-8") + second.write_text("notes", encoding = "utf-8") + model.write_bytes(b"weights") + entries = [first, second, model] + + monkeypatch.setattr(model_common.Path, "rglob", lambda _self, _pattern: iter(entries)) + monkeypatch.setattr(model_common, "_HF_CACHE_MODEL_FILE_PROBE_LIMIT", 2) + + bounded = model_common._iter_hf_cache_model_files(snapshot) + + assert bounded == [first, second] + + monkeypatch.setattr(model_common, "_HF_CACHE_MODEL_FILE_PROBE_LIMIT", 3) + + unbounded = model_common._iter_hf_cache_model_files(snapshot) + + assert unbounded == [first, second, model] + + +def test_download_state_lookup_is_repo_case_insensitive(monkeypatch, tmp_path): + monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path) + + assert download_manifest.write_manifest( + "model", + "Owner/Repo", + None, + [download_manifest.ExpectedFile(path = "config.json", size = 12)], + ) + assert download_manifest.write_cancel_marker("model", "Owner/Repo", "Q4_K_M", "http") + + manifest = download_manifest.read_manifest("model", "owner/repo", None) + + assert manifest is not None + assert manifest.repo_id == "Owner/Repo" + assert manifest.expected_files[0].path == "config.json" + assert download_manifest.has_cancel_marker("model", "owner/repo", "Q4_K_M") + assert ( + download_manifest.read_cancel_marker_transport( + "model", + "owner/repo", + "Q4_K_M", + ) + == "http" + ) + assert [ + variant + for variant, _path in download_manifest.iter_variant_markers( + "model", + "owner/repo", + ) + ] == ["Q4_K_M"] + assert download_manifest.purge_all_state_for_repo("model", "owner/repo") == 2 + assert download_manifest.read_manifest("model", "owner/repo", None) is None + + +def test_hf_cache_scan_fallback_row_uses_local_model_info_alias(monkeypatch, tmp_path): + cache_dir = tmp_path / "hub" + repo_dir = cache_dir / "models--Org--Broken" + blobs_dir = repo_dir / "blobs" + blobs_dir.mkdir(parents = True) + (blobs_dir / "blob").write_bytes(b"content") + monkeypatch.setattr(local_inventory, "_classify_local_path", lambda *_args, **_kwargs: []) + monkeypatch.setattr( + local_inventory.hf_cache_scan, + "is_snapshot_partial", + lambda *_args, **_kwargs: False, + ) + monkeypatch.setattr( + local_inventory.hf_cache_scan, + "is_gguf_repo_partial", + lambda *_args, **_kwargs: False, + ) + monkeypatch.setattr( + local_inventory.hf_cache_scan, + "resolve_hf_cache_realpath", + lambda *_args, **_kwargs: None, + ) + + rows = local_inventory._scan_hf_cache(cache_dir) + + assert len(rows) == 1 + assert rows[0].model_id == "Org/Broken" + assert rows[0].source == "hf_cache" + assert rows[0].model_format == "unknown" + + +def test_hf_cache_scan_uses_gguf_partial_row_for_variant_state(monkeypatch, tmp_path): + monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state") + cache_dir = tmp_path / "hub" + repo_dir = cache_dir / "models--Org--PartialGguf" + blobs_dir = repo_dir / "blobs" + blobs_dir.mkdir(parents = True) + (blobs_dir / "partial").write_bytes(b"content") + assert download_manifest.write_manifest( + "model", + "Org/PartialGguf", + "Q4_K_M", + [download_manifest.ExpectedFile(path = "model-Q4_K_M.gguf", size = 8192)], + "http", + ) + assert download_manifest.write_cancel_marker("model", "Org/PartialGguf", "Q4_K_M", "http") + monkeypatch.setattr(local_inventory, "_classify_local_path", lambda *_args, **_kwargs: []) + monkeypatch.setattr( + local_inventory.hf_cache_scan, + "is_snapshot_partial", + lambda *_args, **_kwargs: False, + ) + monkeypatch.setattr( + local_inventory.hf_cache_scan, + "is_gguf_repo_partial", + lambda *_args, **_kwargs: True, + ) + monkeypatch.setattr( + local_inventory.hf_cache_scan, + "resolve_hf_cache_realpath", + lambda *_args, **_kwargs: None, + ) + + rows = local_inventory._scan_hf_cache(cache_dir) + + assert len(rows) == 1 + assert rows[0].model_id == "Org/PartialGguf" + assert rows[0].source == "hf_cache" + assert rows[0].model_format == "gguf" + assert rows[0].partial is True + assert rows[0].size_bytes == 8192 + assert rows[0].capabilities.requires_variant is True + + +def test_model_download_job_helpers_preserve_idle_shape(): + key = downloads._download_job_key("Org/Model", None) + status = downloads._job_status(key) + + assert key == "org/model::" + assert status.state == "idle" + assert status.error is None + + +def test_gguf_repo_partial_treats_completed_disk_variant_as_clean(monkeypatch, tmp_path): + monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state") + snapshot = tmp_path / "cache" / "models--Org--Repo" / "snapshots" / "abc" + snapshot.mkdir(parents = True) + (snapshot / "model-Q8_0.gguf").write_bytes(b"complete") + assert download_manifest.write_cancel_marker("model", "Org/Repo", "Q4_K_M", "xet") + monkeypatch.setattr( + inventory_scan, + "resolve_snapshot_dir_for_scan", + lambda *_args: snapshot, + ) + + assert inventory_scan.is_gguf_repo_partial("Org/Repo", snapshot.parents[1]) is False + + +def test_gguf_repo_partial_flags_vision_variant_missing_mmproj(monkeypatch, tmp_path): + monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state") + snapshot = tmp_path / "cache" / "models--Org--Vision" / "snapshots" / "abc" + snapshot.mkdir(parents = True) + (snapshot / "model-Q4_K_M.gguf").write_bytes(b"complete-weight") + assert download_manifest.write_manifest( + "model", + "Org/Vision", + "Q4_K_M", + [ + download_manifest.ExpectedFile(path = "model-Q4_K_M.gguf", size = 15), + download_manifest.ExpectedFile(path = "mmproj-F16.gguf", size = 8), + ], + "http", + ) + monkeypatch.setattr( + inventory_scan, + "resolve_snapshot_dir_for_scan", + lambda *_args: snapshot, + ) + + assert inventory_scan.is_gguf_repo_partial("Org/Vision") is True + + +def test_cancel_worker_leaves_exited_process_to_watcher(): + calls: list = [] + + class _Registry: + def get_process(self, _key): + return SimpleNamespace(poll = lambda: 1) + + def get_job(self, _key): + return SimpleNamespace(state = "running") + + def mark_pending_cancel(self, key, generation): + calls.append(("pending", key, generation)) + return True + + def request_cancel(self, key, proc, generation): + calls.append(("request", key, generation)) + return True + + def cancel_requested(self, _key): + return False + + state = download_lifecycle.cancel_worker( + _Registry(), + "org/model::", + generation = 3, + label = "Org/Model", + logger = SimpleNamespace(warning = lambda *_a, **_k: None), + ) + + assert state == "running" + assert calls == [] + + +def test_completed_gguf_split_variant_requires_all_shards(tmp_path): + snapshot = tmp_path / "snapshot" + snapshot.mkdir() + first = snapshot / "model-Q8_0-00001-of-00002.gguf" + second = snapshot / "model-Q8_0-00002-of-00002.gguf" + first.write_bytes(b"first") + + assert "Q8_0" not in inventory_scan._completed_gguf_variants(snapshot) + + second.write_bytes(b"second") + assert "Q8_0" in inventory_scan._completed_gguf_variants(snapshot) + + +def test_variant_partial_accepts_variant_filtered_legacy_hashes(monkeypatch, tmp_path): + monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path) + + assert inventory_scan.is_variant_partial( + "Org/Repo", + "Q4_K_M", + incomplete_blob_hashes = {"main-q4", "main-q8"}, + variant_blob_hashes = frozenset({"main-q4"}), + ) + assert not inventory_scan.is_variant_partial( + "Org/Repo", + "Q5_K_M", + incomplete_blob_hashes = {"main-q4"}, + variant_blob_hashes = frozenset({"main-q5"}), + ) + + +def test_gguf_variants_partial_marker_overrides_size_only_downloaded(monkeypatch, tmp_path): + async def _run_inline(fn, *args, **kwargs): + return fn(*args, **kwargs) + + monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state") + monkeypatch.setattr(gguf_variants.asyncio, "to_thread", _run_inline) + assert download_manifest.write_cancel_marker("model", "Org/PartialRepo", "Q4_K_M", "http") + snapshot = tmp_path / "cache" / "models--Org--PartialRepo" / "snapshots" / "rev0" + snapshot.mkdir(parents = True) + (snapshot / "model-Q4_K_M.gguf").write_bytes(b"x" * 100) + + monkeypatch.setattr( + gguf_variants, + "list_gguf_variants", + lambda *_args, **_kwargs: ( + [ + SimpleNamespace( + filename = "model-Q4_K_M.gguf", + quant = "Q4_K_M", + display_label = None, + size_bytes = 100, + ) + ], + False, + None, + ), + ) + monkeypatch.setattr( + gguf_variants, + "iter_hf_cache_snapshots", + lambda _repo_id: [snapshot], + ) + monkeypatch.setattr( + gguf_variants, + "_gguf_all_variant_requirements", + lambda *_args, **_kwargs: {}, + ) + monkeypatch.setattr( + gguf_variants.download_registry, + "incomplete_blob_hashes", + lambda *_args, **_kwargs: set(), + ) + + result = asyncio.run(gguf_variants.get_gguf_variants_response("Org/PartialRepo")) + + assert result.variants[0].downloaded is False + assert result.variants[0].partial is True + + +def test_download_registry_repo_keys_are_case_insensitive(): + registry = download_registry.DownloadRegistry() + + claimed, state = registry.claim( + "Org/Repo::Q8_0", + download_registry.TRANSPORT_HTTP, + repo_type = "model", + repo_id = "Org/Repo", + variant = "Q8_0", + ) + # The same variant under a different-cased repo id resolves to the same + # job, so the second claim attaches to the running one instead of starting + # a duplicate. + duplicate_claimed, duplicate_state = registry.claim( + "org/repo::Q8_0", + download_registry.TRANSPORT_HTTP, + repo_type = "model", + repo_id = "org/repo", + variant = "Q8_0", + ) + + assert claimed is True + assert state == "running" + assert duplicate_claimed is False + assert duplicate_state == "running" + assert registry.active_jobs("ORG/REPO") == {"org/repo::Q8_0": "running"} + + +def test_download_registry_allows_disjoint_gguf_variant_downloads(): + registry = download_registry.DownloadRegistry() + + claimed, state = registry.claim( + "Org/Repo::Q8_0", + download_registry.TRANSPORT_HTTP, + repo_type = "model", + repo_id = "Org/Repo", + variant = "Q8_0", + blob_hashes = frozenset({"q8-main"}), + progress_blob_hashes = frozenset({"q8-main", "shared-mmproj"}), + ) + second_claimed, second_state = registry.claim( + "Org/Repo::Q4_K_M", + download_registry.TRANSPORT_HTTP, + repo_type = "model", + repo_id = "Org/Repo", + variant = "Q4_K_M", + blob_hashes = frozenset({"q4-main"}), + progress_blob_hashes = frozenset({"q4-main", "shared-mmproj"}), + ) + + assert claimed is True + assert state == "running" + assert second_claimed is True + assert second_state == "running" + assert registry.active_jobs("org/repo") == { + "org/repo::Q8_0": "running", + "org/repo::Q4_K_M": "running", + } + + +def test_download_registry_allows_overlapping_same_transport_variant_downloads(): + # Two variants sharing one mmproj blob still download together on one + # transport: huggingface_hub's per-blob lock serializes the shared write and + # prepare_cache_for_transport never purges a blob a peer is writing. + registry = download_registry.DownloadRegistry() + + claimed, state = registry.claim( + "Org/Repo::Q8_0", + download_registry.TRANSPORT_HTTP, + repo_type = "model", + repo_id = "Org/Repo", + variant = "Q8_0", + blob_hashes = frozenset({"q8-main"}), + progress_blob_hashes = frozenset({"q8-main", "shared-mmproj"}), + ) + second_claimed, second_state = registry.claim( + "Org/Repo::Q4_K_M", + download_registry.TRANSPORT_HTTP, + repo_type = "model", + repo_id = "Org/Repo", + variant = "Q4_K_M", + blob_hashes = frozenset({"q4-main"}), + progress_blob_hashes = frozenset({"q4-main", "shared-mmproj"}), + ) + + assert claimed is True + assert state == "running" + assert second_claimed is True + assert second_state == "running" + + +def test_download_registry_variant_delete_does_not_block_sibling_download(): + # Deleting one quant's partial must be allowed while a different quant of the + # same repo is downloading, and must protect every blob the live sibling is + # writing (including a shared mmproj companion). + registry = download_registry.DownloadRegistry() + registry.claim( + "Org/Repo::Q8_0", + download_registry.TRANSPORT_HTTP, + repo_type = "model", + repo_id = "Org/Repo", + variant = "Q8_0", + blob_hashes = frozenset({"q8-main"}), + progress_blob_hashes = frozenset({"q8-main", "shared-mmproj"}), + ) + + # A sibling variant delete is allowed; deleting the in-flight variant is not. + assert registry.begin_delete("Org/Repo", "Q4_K_M") is True + assert registry.begin_delete("Org/Repo", "Q8_0") is False + # A whole-repo delete still waits for every active download. + assert registry.begin_delete("Org/Repo") is False + + # The live sibling is detected so the delete keeps the shared companion. + assert registry.has_active_peer_variant("Org/Repo", "Q4_K_M") is True + assert registry.has_active_peer_variant("Org/Repo", "Q8_0") is False + + # While Q4_K_M is being deleted, re-downloading it is blocked but an + # untouched third variant may still start. + blocked, blocked_state = registry.claim( + "Org/Repo::Q4_K_M", + download_registry.TRANSPORT_HTTP, + repo_type = "model", + repo_id = "Org/Repo", + variant = "Q4_K_M", + ) + assert blocked is False + assert blocked_state == "deleting" + started, started_state = registry.claim( + "Org/Repo::Q5_K_M", + download_registry.TRANSPORT_HTTP, + repo_type = "model", + repo_id = "Org/Repo", + variant = "Q5_K_M", + ) + assert started is True + assert started_state == "running" + + registry.end_delete("Org/Repo", "Q4_K_M") + assert registry.begin_delete("Org/Repo", "Q4_K_M") is True + + +def test_partial_gguf_reconstruction_dedupes_variant_casing(monkeypatch): + # The manifest keeps original casing while the marker is lowercased; offline + # reconstruction must collapse them to ONE entry (manifest's casing), not two. + monkeypatch.setattr( + download_manifest, + "iter_variant_manifests", + lambda _repo_type, _repo_id: iter([("Q4_K_M", Path("manifest.json"))]), + ) + monkeypatch.setattr( + download_manifest, + "iter_variant_markers", + lambda _repo_type, _repo_id: iter([("q4_k_m", Path("marker.json"))]), + ) + monkeypatch.setattr(download_manifest, "read_manifest", lambda *_a, **_k: None) + + result = gguf.list_partial_gguf_variants_from_state("Org/Repo") + + assert result is not None + variants, _has_vision = result + assert [variant.quant for variant in variants] == ["Q4_K_M"] + + +def test_download_registry_serializes_cross_transport_variant_downloads(): + # An HTTP append-resume and an XET rewrite of the same shared blob would + # corrupt each other, so different-transport variants are serialized. + registry = download_registry.DownloadRegistry() + + claimed, state = registry.claim( + "Org/Repo::Q8_0", + download_registry.TRANSPORT_HTTP, + repo_type = "model", + repo_id = "Org/Repo", + variant = "Q8_0", + blob_hashes = frozenset({"q8-main"}), + progress_blob_hashes = frozenset({"q8-main", "shared-mmproj"}), + ) + second_claimed, second_state = registry.claim( + "Org/Repo::Q4_K_M", + download_registry.TRANSPORT_XET, + repo_type = "model", + repo_id = "Org/Repo", + variant = "Q4_K_M", + blob_hashes = frozenset({"q4-main"}), + progress_blob_hashes = frozenset({"q4-main", "shared-mmproj"}), + ) + + assert claimed is True + assert state == "running" + assert second_claimed is False + assert second_state == "running" + + +def test_download_registry_allows_unknown_hash_gguf_variant_downloads(): + # Resolved blob hashes are NOT required to run two same-transport variants + # concurrently: on-disk safety comes from each worker purging only its own + # main-quant blobs plus huggingface_hub's per-etag lock. Requiring them here + # used to reject the second variant whenever a metadata fetch flaked. + registry = download_registry.DownloadRegistry() + + claimed, state = registry.claim( + "Org/Repo::Q8_0", + download_registry.TRANSPORT_HTTP, + repo_type = "model", + repo_id = "Org/Repo", + variant = "Q8_0", + ) + second_claimed, second_state = registry.claim( + "Org/Repo::Q4_K_M", + download_registry.TRANSPORT_HTTP, + repo_type = "model", + repo_id = "Org/Repo", + variant = "Q4_K_M", + blob_hashes = frozenset({"q4-main"}), + progress_blob_hashes = frozenset({"q4-main", "shared-mmproj"}), + ) + + assert claimed is True + assert state == "running" + assert second_claimed is True + assert second_state == "running" + assert registry.active_jobs("org/repo") == { + "org/repo::Q8_0": "running", + "org/repo::Q4_K_M": "running", + } + + +def test_finalize_worker_exit_never_kills_a_healthy_worker(monkeypatch, tmp_path): + # finalize_worker_exit relies solely on the worker's exit code and never kills + # a live process: huggingface_hub already bounds reads with timeouts, so a + # liveness kill could only false-cancel a healthy download. + import inspect + import io + import logging + + monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state") + + class _Proc: + pid = 4242 + + def __init__(self): + self.killed = False + self.stderr = io.BytesIO(b"") + + def poll(self): + return 0 + + def wait(self, timeout = None): + return 0 + + def kill(self): + self.killed = True + + registry = download_registry.DownloadRegistry() + proc = _Proc() + key = "Org/Repo::Q4_K_M" + registry.claim( + key, + download_registry.TRANSPORT_HTTP, + repo_type = "model", + repo_id = "Org/Repo", + variant = "Q4_K_M", + ) + registry.register_process(key, proc) + + download_lifecycle.finalize_worker_exit( + registry, + key, + proc, + hf_token = None, + label = "Org/Repo [Q4_K_M]", + log_prefix = "Download", + logger = logging.getLogger("test"), + repo_type = "model", + repo_id = "Org/Repo", + transport = "http", + ) + + assert proc.killed is False + assert registry.get_job(key).state == "complete" + # The stall-watchdog knob is gone entirely; no caller may re-enable it. + assert ( + "enable_stall_watchdog" + not in inspect.signature(download_lifecycle.finalize_worker_exit).parameters + ) + + +def test_prepare_cache_for_transport_purges_only_requested_hashes(monkeypatch, tmp_path): + root = tmp_path / "hub" + blobs = root / "models--Org--Repo" / "blobs" + blobs.mkdir(parents = True) + (blobs / "variant-main.incomplete").write_bytes(b"x") + (blobs / "shared-mmproj.incomplete").write_bytes(b"y") + monkeypatch.setattr(download_registry, "hf_cache_root", lambda create = False: root) + + purged = download_registry.prepare_cache_for_transport( + "model", + "Org/Repo", + download_registry.TRANSPORT_XET, + "Q4_K_M", + frozenset({"variant-main"}), + ) + + assert purged == 1 + assert not (blobs / "variant-main.incomplete").exists() + assert (blobs / "shared-mmproj.incomplete").exists() + + +def _vision_cache_root(monkeypatch, tmp_path): + root = tmp_path / "hub" + blobs = root / "models--Org--Vision" / "blobs" + blobs.mkdir(parents = True) + monkeypatch.setattr(download_registry, "hf_cache_root", lambda create = False: root) + return blobs + + +def test_prepare_cache_for_transport_purges_cross_transport_companion(monkeypatch, tmp_path): + blobs = _vision_cache_root(monkeypatch, tmp_path) + companion = frozenset({"shared-mmproj"}) + + # An interrupted XET download stamps the companion marker "xet" and leaves a + # sparse partial. A later HTTP download of a different variant must purge it, + # else the HTTP resumer appends to the sparse bytes and corrupts the blob. + download_registry.prepare_cache_for_transport( + "model", + "Org/Vision", + download_registry.TRANSPORT_XET, + "Q4_K_M", + only_blob_hashes = frozenset({"q4-main"}), + companion_blob_hashes = companion, + ) + (blobs / "shared-mmproj.incomplete").write_bytes(b"sparse") + + purged = download_registry.prepare_cache_for_transport( + "model", + "Org/Vision", + download_registry.TRANSPORT_HTTP, + "Q8_0", + only_blob_hashes = frozenset({"q8-main"}), + companion_blob_hashes = companion, + ) + + assert purged == 1 + assert not (blobs / "shared-mmproj.incomplete").exists() + + +def test_prepare_cache_for_transport_preserves_same_transport_companion(monkeypatch, tmp_path): + blobs = _vision_cache_root(monkeypatch, tmp_path) + companion = frozenset({"shared-mmproj"}) + + download_registry.prepare_cache_for_transport( + "model", + "Org/Vision", + download_registry.TRANSPORT_HTTP, + "Q4_K_M", + only_blob_hashes = frozenset({"q4-main"}), + companion_blob_hashes = companion, + ) + (blobs / "shared-mmproj.incomplete").write_bytes(b"resumable") + + purged = download_registry.prepare_cache_for_transport( + "model", + "Org/Vision", + download_registry.TRANSPORT_HTTP, + "Q4_K_M", + only_blob_hashes = frozenset({"q4-main"}), + companion_blob_hashes = companion, + ) + + assert purged == 0 + assert (blobs / "shared-mmproj.incomplete").exists() + + +def test_prepare_cache_for_transport_protects_peer_companion(monkeypatch, tmp_path): + blobs = _vision_cache_root(monkeypatch, tmp_path) + companion = frozenset({"shared-mmproj"}) + + download_registry.prepare_cache_for_transport( + "model", + "Org/Vision", + download_registry.TRANSPORT_XET, + "Q4_K_M", + only_blob_hashes = frozenset({"q4-main"}), + companion_blob_hashes = companion, + ) + (blobs / "shared-mmproj.incomplete").write_bytes(b"sparse") + + purged = download_registry.prepare_cache_for_transport( + "model", + "Org/Vision", + download_registry.TRANSPORT_HTTP, + "Q8_0", + only_blob_hashes = frozenset({"q8-main"}), + companion_blob_hashes = companion, + protected_blob_hashes = companion, + ) + + assert purged == 0 + assert (blobs / "shared-mmproj.incomplete").exists() + + +def test_model_download_records_completed_baseline_for_new_gguf_variant(monkeypatch, tmp_path): + async def _run_inline(fn, *args, **kwargs): + return fn(*args, **kwargs) + + monkeypatch.setattr(downloads.asyncio, "to_thread", _run_inline) + monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state") + monkeypatch.setattr( + downloads, + "resolve_cached_repo_id_case", + lambda repo_id, repo_type = "model": repo_id, + ) + monkeypatch.setattr( + downloads.gguf_variants, + "gguf_variant_blob_hashes", + lambda _repo, _variant, _token = None, include_companions = True, **_kwargs: ( + frozenset({"mainhash", "mmprojhash"}) if include_companions else frozenset({"mainhash"}) + ), + ) + monkeypatch.setattr( + downloads.download_registry, + "completed_blob_bytes", + lambda *_args, **_kwargs: 30, + ) + + class _Registry: + claim_kwargs = None + + def claim(self, _key, _transport, **kwargs): + self.claim_kwargs = kwargs + return True, "running" + + def current_generation(self, _key): + return 1 + + def get_job(self, _key): + return SimpleNamespace(state = "running") + + def register_process(self, _key, _proc): + return False + + def peer_blob_hashes(self, _key): + return frozenset() + + class _Proc: + pid = 123 + stderr = None + + def poll(self): + return None + + def kill(self): + return None + + def wait(self, timeout = None): + return 0 + + registry = _Registry() + monkeypatch.setattr(downloads, "_registry", registry) + monkeypatch.setattr(downloads, "_spawn_download_worker", lambda *_args, **_kwargs: _Proc()) + + asyncio.run( + downloads.download_model_response( + SimpleNamespace(repo_id = "Org/Model", gguf_variant = "Q4_K_M", use_xet = False) + ) + ) + + assert registry.claim_kwargs["blob_hashes"] == frozenset({"mainhash"}) + assert registry.claim_kwargs["progress_blob_hashes"] == frozenset({"mainhash", "mmprojhash"}) + assert registry.claim_kwargs["completed_baseline_bytes"] == 30 + + +def test_gguf_model_download_skips_completed_baseline_for_variant_resume_state( + monkeypatch, tmp_path +): + async def _run_inline(fn, *args, **kwargs): + return fn(*args, **kwargs) + + monkeypatch.setattr(downloads.asyncio, "to_thread", _run_inline) + monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state") + assert download_manifest.write_manifest( + "model", + "Org/Model", + "Q4_K_M", + [ + download_manifest.ExpectedFile( + path = "model-Q4_K_M.gguf", + size = 100, + sha256 = "mainhash", + ) + ], + "http", + ) + monkeypatch.setattr( + downloads, + "resolve_cached_repo_id_case", + lambda repo_id, repo_type = "model": repo_id, + ) + monkeypatch.setattr( + downloads.gguf_variants, + "gguf_variant_blob_hashes", + lambda _repo, _variant, _token = None, include_companions = True, **_kwargs: ( + frozenset({"mainhash", "mmprojhash"}) if include_companions else frozenset({"mainhash"}) + ), + ) + monkeypatch.setattr( + downloads.download_registry, + "completed_blob_bytes", + lambda *_args, **_kwargs: 30, + ) + + class _Registry: + claim_kwargs = None + + def claim(self, _key, _transport, **kwargs): + self.claim_kwargs = kwargs + return True, "running" + + def current_generation(self, _key): + return 1 + + def get_job(self, _key): + return SimpleNamespace(state = "running") + + def register_process(self, _key, _proc): + return False + + def peer_blob_hashes(self, _key): + return frozenset() + + class _Proc: + pid = 123 + stderr = None + + def poll(self): + return None + + def kill(self): + return None + + def wait(self, timeout = None): + return 0 + + registry = _Registry() + monkeypatch.setattr(downloads, "_registry", registry) + monkeypatch.setattr(downloads, "_spawn_download_worker", lambda *_args, **_kwargs: _Proc()) + + asyncio.run( + downloads.download_model_response( + SimpleNamespace(repo_id = "Org/Model", gguf_variant = "Q4_K_M", use_xet = False) + ) + ) + + assert registry.claim_kwargs["completed_baseline_bytes"] == 0 + + +def test_model_idle_status_uses_cancel_marker_after_restart(monkeypatch, tmp_path): + monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path) + monkeypatch.setattr(downloads, "_registry", download_registry.DownloadRegistry()) + assert download_manifest.write_cancel_marker("model", "Owner/Repo", "Q4_K_M", "http") + + status = asyncio.run(downloads.get_download_status_response("owner/repo", "Q4_K_M")) + + assert status.state == "cancelled" + assert status.error is None + + +def test_shutdown_kills_all_workers_before_shared_deadline_reap(monkeypatch): + events = [] + now = [100.0] + + class _Proc: + def __init__(self, name): + self.name = name + + def poll(self): + return None + + def kill(self): + events.append(("kill", self.name)) + + def wait(self, timeout): + events.append(("wait", self.name, timeout)) + now[0] += 7.0 + + registry = download_registry.DownloadRegistry() + proc_a = _Proc("a") + proc_b = _Proc("b") + registry.claim( + "Org/A", + download_registry.TRANSPORT_HTTP, + repo_type = "dataset", + repo_id = "Org/A", + ) + registry.claim( + "Org/B", + download_registry.TRANSPORT_HTTP, + repo_type = "dataset", + repo_id = "Org/B", + ) + assert registry.register_process("org/a", proc_a) + assert registry.register_process("org/b", proc_b) + monkeypatch.setattr( + download_registry, + "persist_cancel_marker", + lambda *args, **kwargs: events.append(("marker", args[1])), + ) + monkeypatch.setattr(download_registry.time, "monotonic", lambda: now[0]) + + registry.terminate_all("dataset download") + + assert events == [ + ("kill", "a"), + ("kill", "b"), + ("wait", "a", 10.0), + ("marker", "Org/A"), + ("wait", "b", 3.0), + ("marker", "Org/B"), + ] + + +def test_shutdown_skips_marker_for_worker_that_exits_cleanly(monkeypatch): + markers = [] + + class _Proc: + def __init__(self, final_rc): + self._final_rc = final_rc + self._exited = False + + def poll(self): + return self._final_rc if self._exited else None + + def kill(self): + pass + + def wait(self, timeout): + self._exited = True + + registry = download_registry.DownloadRegistry() + clean = _Proc(0) + interrupted = _Proc(-9) + registry.claim( + "Org/Clean", + download_registry.TRANSPORT_HTTP, + repo_type = "dataset", + repo_id = "Org/Clean", + ) + registry.claim( + "Org/Cut", + download_registry.TRANSPORT_HTTP, + repo_type = "dataset", + repo_id = "Org/Cut", + ) + assert registry.register_process("org/clean", clean) + assert registry.register_process("org/cut", interrupted) + monkeypatch.setattr( + download_registry, + "persist_cancel_marker", + lambda *args, **kwargs: markers.append(args[1]), + ) + + registry.terminate_all("dataset download") + + assert markers == ["Org/Cut"] + + +def test_model_claim_register_cancel_uses_registry_marker_owner(monkeypatch): + killed = [] + + class _Registry: + def claim(self, *_args, **_kwargs): + return True, "running" + + def current_generation(self, _key): + return 1 + + def register_process(self, _key, _proc): + return False + + def persist_cancel_for_key(self, *_args, **_kwargs): + raise AssertionError("register_process owns pending-cancel markers") + + def get_job(self, _key): + return SimpleNamespace(state = "cancelled", error = None) + + monkeypatch.setattr(downloads, "_registry", _Registry()) + monkeypatch.setattr( + downloads, + "resolve_cached_repo_id_case", + lambda repo_id, **_kwargs: repo_id, + ) + monkeypatch.setattr( + downloads.download_registry, + "download_transport_unavailable_reason", + lambda _transport: None, + ) + monkeypatch.setattr( + downloads, + "_spawn_download_worker", + lambda *_args, **_kwargs: object(), + ) + monkeypatch.setattr( + downloads.download_lifecycle, + "kill_and_reap_process", + lambda proc, **_kwargs: killed.append(proc), + ) + + result = asyncio.run( + downloads.download_model_response( + SimpleNamespace(repo_id = "Org/Model", gguf_variant = None, use_xet = False) + ) + ) + + assert result["state"] == "cancelled" + assert killed + + +def test_model_cancel_registered_worker_requests_and_kills(monkeypatch): + events = [] + + class _Proc: + def poll(self): + return None + + def kill(self): + events.append(("kill",)) + + class _Registry: + def get_process(self, _key): + return _Proc() + + def request_cancel(self, key, _proc, generation): + events.append(("request", key, generation)) + return True + + def persist_cancel_for_key(self, *_args, **_kwargs): + raise AssertionError( + "cancel_worker must leave marker persistence to the exit watcher; " + "an eager persist races a clean completion and strands a stale marker" + ) + + def get_job(self, _key): + return SimpleNamespace(state = "running") + + monkeypatch.setattr(downloads, "_registry", _Registry()) + monkeypatch.setattr( + downloads, + "resolve_cached_repo_id_case", + lambda repo_id, **_kwargs: repo_id, + ) + + result = asyncio.run( + downloads.cancel_download_model_response( + SimpleNamespace(repo_id = "Org/Model", gguf_variant = "Q4_K_M", generation = 7) + ) + ) + + assert result == { + "job_key": downloads._download_job_key("Org/Model", "Q4_K_M"), + "state": "cancelling", + } + assert events == [ + ("request", downloads._download_job_key("Org/Model", "Q4_K_M"), 7), + ("kill",), + ] + + +def test_model_download_watcher_invalidates_hf_cache_scan(monkeypatch): + invalidated = [] + + class _Registry: + def claim(self, *_args, **_kwargs): + return True, "running" + + def current_generation(self, _key): + return 1 + + def register_process(self, _key, _proc): + return True + + def get_job(self, _key): + return SimpleNamespace(state = "complete", error = None) + + class _ImmediateThread: + def __init__(self, *, target, **_kwargs): + self._target = target + + def start(self): + self._target() + + monkeypatch.setattr(downloads, "_registry", _Registry()) + monkeypatch.setattr( + downloads, + "resolve_cached_repo_id_case", + lambda repo_id, **_kwargs: repo_id, + ) + monkeypatch.setattr( + downloads.download_registry, + "download_transport_unavailable_reason", + lambda _transport: None, + ) + monkeypatch.setattr( + downloads.download_lifecycle, + "finalize_worker_exit", + lambda *_args, **_kwargs: None, + ) + monkeypatch.setattr( + downloads, + "_spawn_download_worker", + lambda *_args, **_kwargs: object(), + ) + monkeypatch.setattr(downloads.download_lifecycle.threading, "Thread", _ImmediateThread) + monkeypatch.setattr( + downloads.hf_cache_scan, + "invalidate_hf_cache_scans", + lambda: invalidated.append(True), + ) + + async def _inline_to_thread(func, *args, **kwargs): + return func(*args, **kwargs) + + monkeypatch.setattr(downloads.asyncio, "to_thread", _inline_to_thread) + + result = asyncio.run( + downloads.download_model_response( + SimpleNamespace(repo_id = "Org/Model", gguf_variant = None, use_xet = False) + ) + ) + + assert result["accepted"] is True + assert invalidated == [True] + + +def test_two_concurrent_same_repo_variants_both_complete(monkeypatch, tmp_path): + # End-to-end proof that two GGUF variants of ONE repo download concurrently + # without cancelling each other, with real registry/finalize/subprocess/watch + # threads exercising the claim gate, register, finalize funnel, and + # classify_exit under true concurrency. + import subprocess + import time + + monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state") + monkeypatch.setattr( + downloads, + "_registry", + download_registry.DownloadRegistry(), + ) + monkeypatch.setattr( + downloads, + "resolve_cached_repo_id_case", + lambda repo_id, **_k: repo_id, + ) + monkeypatch.setattr( + downloads.download_registry, + "download_transport_unavailable_reason", + lambda _transport: None, + ) + # Per-variant blob hashes (distinct main shard, shared mmproj companion). + monkeypatch.setattr( + downloads.gguf_variants, + "gguf_variant_blob_hashes", + lambda _repo, variant, _token = None, include_companions = True, **_k: ( + frozenset({f"{variant.lower()}-main", "shared-mmproj"}) + if include_companions + else frozenset({f"{variant.lower()}-main"}) + ), + ) + monkeypatch.setattr( + downloads.download_registry, + "completed_blob_bytes", + lambda *_a, **_k: 0, + ) + monkeypatch.setattr( + downloads.hf_cache_scan, + "invalidate_hf_cache_scans", + lambda: None, + ) + # Real subprocess that exits 0 immediately, with a stderr pipe to drain. + spawned: list[subprocess.Popen] = [] + + def _fake_spawn(*_args, **_kwargs): + proc = subprocess.Popen( + [sys.executable, "-c", "import sys; sys.exit(0)"], + stderr = subprocess.PIPE, + ) + spawned.append(proc) + return proc + + monkeypatch.setattr(downloads, "_spawn_download_worker", _fake_spawn) + + async def _run_both(): + return await asyncio.gather( + downloads.download_model_response( + SimpleNamespace( + repo_id = "Org/Model", + gguf_variant = "Q4_K_M", + use_xet = False, + ) + ), + downloads.download_model_response( + SimpleNamespace( + repo_id = "Org/Model", + gguf_variant = "Q8_0", + use_xet = False, + ) + ), + ) + + results = asyncio.run(_run_both()) + assert all(r["accepted"] is True for r in results), results + + registry = downloads._registry + key_q4 = downloads._download_job_key("Org/Model", "Q4_K_M") + key_q8 = downloads._download_job_key("Org/Model", "Q8_0") + deadline = time.monotonic() + 10.0 + while time.monotonic() < deadline: + s4 = registry.get_job(key_q4).state + s8 = registry.get_job(key_q8).state + if s4 in download_registry.TERMINAL_STATES and s8 in download_registry.TERMINAL_STATES: + break + time.sleep(0.02) + + for p in spawned: + try: + p.wait(timeout = 5) + except Exception: + pass + + assert registry.get_job(key_q4).state == "complete" + assert registry.get_job(key_q8).state == "complete" + + +def test_download_registry_factories_reuse_service_singletons(): + registry_module = downloads.download_registry + before_count = len(registry_module._REGISTRIES) + + assert registry_module.get_models_registry() is downloads.registry + assert registry_module.get_models_registry() is downloads.registry + assert registry_module.get_datasets_registry() is dataset_downloads.registry + assert registry_module.get_datasets_registry() is dataset_downloads.registry + assert len(registry_module._REGISTRIES) == before_count + + +def test_hub_hf_token_header_uses_namespaced_header_only(): + assert get_hf_token("new-token") == "new-token" + assert get_hf_token(None) is None + + +def test_scan_folder_rejects_credential_directories(tmp_path): + sensitive_dir = tmp_path / ".ssh" / "models" + sensitive_dir.mkdir(parents = True) + + with pytest.raises(ValueError, match = "Credential or configuration"): + scan_folders.add_scan_folder(str(sensitive_dir)) + + +def _build_variant_cache_repo(repo_dir, blob_specs, snapshot_links): + """Build a HF cache repo dir with blobs + snapshot symlinks for the + per-variant deletion path. blob_specs: {blob_name: bytes_payload}; + snapshot_links: list of (revision, filename, blob_name).""" + blobs_dir = repo_dir / "blobs" + blobs_dir.mkdir(parents = True) + for blob_name, payload in blob_specs.items(): + (blobs_dir / blob_name).write_bytes(payload) + + files = [] + for revision, filename, blob_name in snapshot_links: + snap_dir = repo_dir / "snapshots" / revision + snap_dir.mkdir(parents = True, exist_ok = True) + blob = blobs_dir / blob_name + link = snap_dir / filename + link.symlink_to(blob) + files.append( + SimpleNamespace( + file_name = filename, + file_path = str(link), + blob_path = str(blob), + size_on_disk = blob.stat().st_size, + ) + ) + repo = SimpleNamespace( + repo_id = "Org/Repo-GGUF", + repo_type = "model", + repo_path = repo_dir, + revisions = [SimpleNamespace(commit_hash = "rev1", files = files)], + ) + return repo + + +def _patch_variant_delete_side_effects(monkeypatch): + monkeypatch.setattr( + deletion.download_manifest, + "purge_state", + lambda *_args, **_kwargs: False, + ) + + +def test_snapshot_progress_filters_stale_blobs(monkeypatch, tmp_path): + """Exclude superseded-revision blobs; count an in-progress blob only when its + hash belongs to the target.""" + entry = tmp_path / "datasets--Org--Data" + blobs = entry / "blobs" + blobs.mkdir(parents = True) + (blobs / "keep1").write_bytes(b"a" * 100) + (blobs / "stale").write_bytes(b"b" * 500) + (blobs / "keep2.incomplete").write_bytes(b"c" * 40) + + monkeypatch.setattr( + snapshot_progress, + "preferred_repo_cache_dirs", + lambda _repo_type, _repo_id, force_active = False: [entry], + ) + + result = snapshot_progress.compute_snapshot_progress( + repo_type = "dataset", + repo_id = "Org/Data", + job_key = "org/data", + expected_bytes = 0, + hf_token = None, + registry = SimpleNamespace( + get_job = lambda _key: SimpleNamespace(state = "running"), + ), + metadata_resolver = lambda _repo_id, _hf_token: ( + 140, + frozenset({"keep1", "keep2"}), + ), + ) + + assert result["completed_bytes"] == 100 + assert result["downloaded_bytes"] == 140 + assert result["complete_on_disk"] is False + assert result["expected_bytes"] == 140 + + +def test_snapshot_progress_confirms_complete_only_with_verified_snapshot(monkeypatch, tmp_path): + entry = tmp_path / "models--Org--Model" + blobs = entry / "blobs" + snap = entry / "snapshots" / "rev0" + blobs.mkdir(parents = True) + snap.mkdir(parents = True) + (blobs / "keep1").write_bytes(b"a" * 100) + (snap / "model.safetensors").write_bytes(b"a" * 100) + + monkeypatch.setattr( + snapshot_progress, + "preferred_repo_cache_dirs", + lambda _repo_type, _repo_id, force_active = False: [entry], + ) + monkeypatch.setattr( + snapshot_progress.download_manifest, + "has_cancel_marker", + lambda *_args, **_kwargs: False, + ) + monkeypatch.setattr( + snapshot_progress.download_manifest, + "read_manifest", + lambda *_args, **_kwargs: SimpleNamespace(), + ) + monkeypatch.setattr( + snapshot_progress.download_manifest, + "verify_against_disk", + lambda *_args, **_kwargs: SimpleNamespace(ok = True), + ) + + result = snapshot_progress.compute_snapshot_progress( + repo_type = "model", + repo_id = "Org/Model", + job_key = "org/model::", + expected_bytes = 100, + hf_token = None, + registry = SimpleNamespace( + get_job = lambda _key: SimpleNamespace(state = "idle"), + ), + metadata_resolver = lambda _repo_id, _hf_token: ( + 100, + frozenset({"keep1"}), + ), + ) + + assert result["completed_bytes"] == 100 + assert result["complete_on_disk"] is True + + +def test_expected_files_from_snapshot_dir_records_relative_paths_and_sizes(tmp_path): + snap = tmp_path / "snapshots" / "rev0" + (snap / "nested").mkdir(parents = True) + (snap / "model.safetensors").write_bytes(b"a" * 12) + (snap / "nested" / "config.json").write_bytes(b"b" * 3) + + files = download_manifest.expected_files_from_snapshot_dir(snap) + + by_path = {f.path: f for f in files} + assert by_path["model.safetensors"].size == 12 + assert by_path["nested/config.json"].size == 3 + assert all(f.sha256 is None for f in files) + + +def test_snapshot_progress_complete_with_manifest_synthesized_from_disk(monkeypatch, tmp_path): + """A finished snapshot whose only manifest was synthesized from on-disk files + still verifies as complete, so a refresh finalizes it instead of capping at + 99% and evicting it as gone.""" + entry = tmp_path / "models--Org--Model" + blobs = entry / "blobs" + snap = entry / "snapshots" / "rev0" + blobs.mkdir(parents = True) + snap.mkdir(parents = True) + (blobs / "keep1").write_bytes(b"a" * 100) + (snap / "model.safetensors").write_bytes(b"a" * 100) + + synthesized = download_manifest.expected_files_from_snapshot_dir(snap) + manifest = download_manifest.Manifest( + repo_type = "model", + repo_id = "Org/Model", + variant = None, + started_at = "", + expected_files = tuple(synthesized), + ) + + monkeypatch.setattr( + snapshot_progress, + "preferred_repo_cache_dirs", + lambda _repo_type, _repo_id, force_active = False: [entry], + ) + monkeypatch.setattr( + snapshot_progress.download_manifest, + "has_cancel_marker", + lambda *_args, **_kwargs: False, + ) + monkeypatch.setattr( + snapshot_progress.download_manifest, + "read_manifest", + lambda *_args, **_kwargs: manifest, + ) + + result = snapshot_progress.compute_snapshot_progress( + repo_type = "model", + repo_id = "Org/Model", + job_key = "org/model::", + expected_bytes = 100, + hf_token = None, + registry = SimpleNamespace( + get_job = lambda _key: SimpleNamespace(state = "idle"), + ), + metadata_resolver = lambda _repo_id, _hf_token: ( + 100, + frozenset({"keep1"}), + ), + ) + + assert result["complete_on_disk"] is True + assert result["progress"] == 1.0 + + +def test_delete_variant_keeps_blob_shared_with_other_snapshot(monkeypatch, tmp_path): + """A blob still referenced by a non-target snapshot symlink survives so that + symlink doesn't dangle (which the scanner reports as partial).""" + repo_dir = tmp_path / "models--Org--Repo-GGUF" + repo = _build_variant_cache_repo( + repo_dir, + blob_specs = {"sharedblob": b"x" * 200, "q8blob": b"y" * 300}, + snapshot_links = [ + ("rev1", "model-Q4_K_M.gguf", "sharedblob"), + ("rev1", "model-Q8_0.gguf", "q8blob"), + # An unrelated file that happens to share Q4's blob content. + ("rev1", "extra-copy.gguf", "sharedblob"), + ], + ) + monkeypatch.setattr( + deletion.cache_inventory, + "all_hf_cache_scans", + lambda: [SimpleNamespace(repos = [repo])], + ) + _patch_variant_delete_side_effects(monkeypatch) + + result = deletion._delete_cached_model_blocking("Org/Repo-GGUF", "Q4_K_M", None) + + assert result["status"] == "deleted" + # Q4 snapshot link gone, but its blob survives (extra-copy still links it). + assert not (repo_dir / "snapshots" / "rev1" / "model-Q4_K_M.gguf").exists() + assert (repo_dir / "blobs" / "sharedblob").exists() + extra = repo_dir / "snapshots" / "rev1" / "extra-copy.gguf" + assert extra.is_symlink() and extra.exists() # not dangling + + +def test_delete_variant_unlinks_unshared_blob(monkeypatch, tmp_path): + repo_dir = tmp_path / "models--Org--Repo-GGUF" + repo = _build_variant_cache_repo( + repo_dir, + blob_specs = {"q4blob": b"x" * 200, "q8blob": b"y" * 300}, + snapshot_links = [ + ("rev1", "model-Q4_K_M.gguf", "q4blob"), + ("rev1", "model-Q8_0.gguf", "q8blob"), + ], + ) + monkeypatch.setattr( + deletion.cache_inventory, + "all_hf_cache_scans", + lambda: [SimpleNamespace(repos = [repo])], + ) + _patch_variant_delete_side_effects(monkeypatch) + + result = deletion._delete_cached_model_blocking("Org/Repo-GGUF", "Q4_K_M", None) + + assert result["status"] == "deleted" + assert not (repo_dir / "blobs" / "q4blob").exists() + # Untouched sibling variant remains fully intact. + assert (repo_dir / "blobs" / "q8blob").exists() + q8 = repo_dir / "snapshots" / "rev1" / "model-Q8_0.gguf" + assert q8.is_symlink() and q8.exists() + + +def test_delete_variant_surfaces_locked_file_as_conflict(monkeypatch, tmp_path): + """A blob unlink that fails (e.g. a Windows file lock on a loaded model) + must raise a clear 409, not report a misleading success.""" + repo_dir = tmp_path / "models--Org--Repo-GGUF" + repo = _build_variant_cache_repo( + repo_dir, + blob_specs = {"lockedblob": b"x" * 200}, + snapshot_links = [("rev1", "model-Q4_K_M.gguf", "lockedblob")], + ) + monkeypatch.setattr( + deletion.cache_inventory, + "all_hf_cache_scans", + lambda: [SimpleNamespace(repos = [repo])], + ) + _patch_variant_delete_side_effects(monkeypatch) + + real_unlink = Path.unlink + + def fake_unlink(self, *args, **kwargs): + if self.name == "lockedblob": + raise PermissionError("file in use") + return real_unlink(self, *args, **kwargs) + + monkeypatch.setattr(Path, "unlink", fake_unlink) + + with pytest.raises(HTTPException) as exc_info: + deletion._delete_cached_model_blocking("Org/Repo-GGUF", "Q4_K_M", None) + + assert exc_info.value.status_code == 409 + + +def test_download_snapshot_writes_manifest_for_xet(monkeypatch, tmp_path): + written = [] + verified = [] + + monkeypatch.setattr( + hf_download, + "_model_info_with_retry", + lambda *_args, **_kwargs: SimpleNamespace( + siblings = [SimpleNamespace(rfilename = "config.json", size = 12)] + ), + ) + monkeypatch.setattr( + hf_download, "_verify_completed_download", lambda *args, **kwargs: verified.append(args) + ) + monkeypatch.setattr( + download_registry, "prepare_cache_for_transport", lambda *_args, **_kwargs: 0 + ) + monkeypatch.setattr(download_manifest, "clear_cancel_marker", lambda *_args: None) + monkeypatch.setattr( + download_manifest, "write_manifest", lambda *args: written.append(args) or True + ) + monkeypatch.setitem( + sys.modules, + "huggingface_hub", + SimpleNamespace(snapshot_download = lambda **_kwargs: str(tmp_path)), + ) + + hf_download._download_snapshot("Org/Model", None, "xet") + + assert written, "XET snapshot download must still record a manifest" + assert written[0][0:3] == ("model", "Org/Model", None) + assert written[0][3][0].path == "config.json" + assert verified == [("model", "Org/Model", None, str(tmp_path))] + + +def test_download_gguf_variant_writes_manifest_for_xet(monkeypatch, tmp_path): + written = [] + verified = [] + + monkeypatch.setattr( + hf_download, + "_model_info_with_retry", + lambda *_args, **_kwargs: SimpleNamespace( + siblings = [_sibling("model-Q4_K_M.gguf", 10, "main")] + ), + ) + monkeypatch.setattr( + hf_download, "_verify_completed_download", lambda *args, **kwargs: verified.append(args) + ) + monkeypatch.setattr( + download_registry, "prepare_cache_for_transport", lambda *_args, **_kwargs: 0 + ) + monkeypatch.setattr(download_manifest, "clear_cancel_marker", lambda *_args: None) + monkeypatch.setattr( + download_manifest, "write_manifest", lambda *args: written.append(args) or True + ) + monkeypatch.setitem( + sys.modules, + "huggingface_hub", + SimpleNamespace(snapshot_download = lambda **_kwargs: str(tmp_path)), + ) + + hf_download._download_gguf_variant("Org/Model", "Q4_K_M", None, "xet") + + assert written, "XET GGUF variant download must still record a manifest" + assert written[0][0:3] == ("model", "Org/Model", "Q4_K_M") + assert written[0][3][0].path == "model-Q4_K_M.gguf" + assert verified == [("model", "Org/Model", "Q4_K_M", str(tmp_path))] + + +def test_download_dataset_writes_manifest_for_xet(monkeypatch, tmp_path): + written = [] + verified = [] + + monkeypatch.setattr( + hf_download, + "_dataset_info_with_retry", + lambda *_args, **_kwargs: SimpleNamespace( + siblings = [SimpleNamespace(rfilename = "data.parquet", size = 30)] + ), + ) + monkeypatch.setattr( + hf_download, "_verify_completed_download", lambda *args, **kwargs: verified.append(args) + ) + monkeypatch.setattr( + download_registry, "prepare_cache_for_transport", lambda *_args, **_kwargs: 0 + ) + monkeypatch.setattr(download_manifest, "clear_cancel_marker", lambda *_args: None) + monkeypatch.setattr( + download_manifest, "write_manifest", lambda *args: written.append(args) or True + ) + monkeypatch.setitem( + sys.modules, + "huggingface_hub", + SimpleNamespace(snapshot_download = lambda **_kwargs: str(tmp_path)), + ) + + hf_download._download_dataset("Org/Data", None, "xet") + + assert written, "XET dataset download must still record a manifest" + assert written[0][0:3] == ("dataset", "Org/Data", None) + assert written[0][3][0].path == "data.parquet" + assert verified == [("dataset", "Org/Data", None, str(tmp_path))] + + +def test_dataset_status_includes_generation(monkeypatch): + class _Registry: + def get_job(self, _key): + return SimpleNamespace(state = "running", error = None) + + def current_generation(self, _key): + return 4 + + monkeypatch.setattr(dataset_downloads, "_registry", _Registry()) + monkeypatch.setattr( + dataset_downloads, + "resolve_cached_repo_id_case", + lambda repo_id, **_kwargs: repo_id, + ) + + result = asyncio.run(dataset_downloads.get_dataset_download_status_response("Org/Data")) + + assert result.state == "running" + assert result.generation == 4 diff --git a/studio/backend/hub/utils/__init__.py b/studio/backend/hub/utils/__init__.py new file mode 100644 index 0000000000..32014236c6 --- /dev/null +++ b/studio/backend/hub/utils/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 diff --git a/studio/backend/hub/utils/dataset_cache.py b/studio/backend/hub/utils/dataset_cache.py new file mode 100644 index 0000000000..1a7a90d8a5 --- /dev/null +++ b/studio/backend/hub/utils/dataset_cache.py @@ -0,0 +1,145 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import re +from pathlib import Path +from typing import Optional + +from hub.utils.hf_cache_state import iter_repo_cache_dirs + + +TRAINING_DATA_EXTS = (".parquet", ".json", ".jsonl", ".csv") + + +def _rel_lower(snapshot: Path, path: Path) -> str: + return path.relative_to(snapshot).as_posix().lower() + + +_SPLIT_ALIASES = { + "validation": frozenset({"validation", "valid", "val"}), + "valid": frozenset({"validation", "valid", "val"}), + "val": frozenset({"validation", "valid", "val"}), + "eval": frozenset({"eval", "validation", "valid", "val"}), +} + + +def _label_tokens(text: str) -> set[str]: + return {token for token in re.split(r"[^a-z0-9]+", text.lower()) if token} + + +def split_label_matches(text: str, split: str) -> bool: + """Match a split name against a file path's tokens, expanding split aliases + (validation/valid/val, eval) so cached and remote selection agree.""" + normalized = split.strip().lower() + if not normalized: + return False + labels = _SPLIT_ALIASES.get(normalized, frozenset({normalized})) + return bool(labels.intersection(_label_tokens(text))) + + +def _matches_label(snapshot: Path, path: Path, label: str) -> bool: + label = label.strip().lower() + if not label: + return False + rel = _rel_lower(snapshot, path) + tokens = [token for token in re.split(r"[^a-z0-9]+", rel) if token] + if label in tokens: + return True + if label in {"train", "test", "validation", "valid", "val", "eval"}: + return False + return label in rel + + +def dataset_snapshot_from_cache_path(local_path: Optional[str], repo_id: str) -> Optional[Path]: + if not local_path or not repo_id: + return None + try: + root = Path(local_path).expanduser() + if not root.exists(): + return None + expected_repo_dir = f"datasets--{repo_id.replace('/', '--')}".lower() + if expected_repo_dir not in {part.lower() for part in root.parts}: + return None + if root.is_dir() and root.parent.name == "snapshots": + return root.resolve() + snapshots = root / "snapshots" if root.is_dir() else None + if snapshots is None or not snapshots.is_dir(): + return None + candidates = [p for p in snapshots.iterdir() if p.is_dir()] + if not candidates: + return None + candidates.sort( + key = lambda path: path.stat().st_mtime if path.exists() else 0, + reverse = True, + ) + return candidates[0].resolve() + except Exception: + return None + + +def latest_cached_dataset_snapshot( + repo_id: str, local_path: Optional[str] = None +) -> Optional[Path]: + local_snapshot = dataset_snapshot_from_cache_path(local_path, repo_id) + if local_snapshot is not None: + return local_snapshot + + newest: Optional[Path] = None + newest_mtime = -1.0 + for entry in iter_repo_cache_dirs("dataset", repo_id): + snapshots = entry / "snapshots" + if not snapshots.is_dir(): + continue + try: + candidates = [s for s in snapshots.iterdir() if s.is_dir()] + except OSError: + continue + for snap in candidates: + try: + mtime = snap.stat().st_mtime + except OSError: + continue + if mtime > newest_mtime: + newest = snap + newest_mtime = mtime + return newest + + +def cached_dataset_candidates( + snapshot: Path, + *, + subset: Optional[str], + train_split: str, + extensions: tuple[str, ...], + preferred_extensions: tuple[str, ...] = TRAINING_DATA_EXTS, +) -> list[Path]: + try: + files = [ + p for p in snapshot.rglob("*") if p.is_file() and p.name.lower().endswith(extensions) + ] + except OSError: + return [] + if not files: + return [] + + subset_lower = subset.lower() if subset else "" + split_lower = train_split.lower() + + def score(path: Path) -> tuple[int, int, str]: + rel = _rel_lower(snapshot, path) + subset_match = bool(subset_lower and _matches_label(snapshot, path, subset_lower)) + split_match = bool(split_lower and split_label_matches(rel, split_lower)) + location_rank = 3 + if split_match and (not subset_lower or subset_match): + location_rank = 0 + elif split_match: + location_rank = 1 + elif subset_match: + location_rank = 2 + return ( + 0 if path.name.lower().endswith(preferred_extensions) else 1, + location_rank, + rel, + ) + + return sorted(files, key = score) diff --git a/studio/backend/hub/utils/dataset_format.py b/studio/backend/hub/utils/dataset_format.py new file mode 100644 index 0000000000..df02035365 --- /dev/null +++ b/studio/backend/hub/utils/dataset_format.py @@ -0,0 +1,749 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from __future__ import annotations + +import re +from typing import Any, Optional + + +def _first_row(dataset) -> Optional[dict]: + try: + row = next(iter(dataset)) + except StopIteration: + return None + return row if isinstance(row, dict) else None + + +def _column_names(dataset, sample: Optional[dict] = None) -> list[str]: + names = getattr(dataset, "column_names", None) + if names is not None: + return list(names) + return list((sample or {}).keys()) + + +def _keyword_in_column(keyword: str, col_name: str) -> bool: + return re.search(r"\b" + re.escape(keyword) + r"\b", col_name, re.IGNORECASE) is not None + + +def _unknown_dataset_format( + chat_column: Optional[str] = None, sample_keys: Optional[list[str]] = None +) -> dict: + return { + "format": "unknown", + "chat_column": chat_column, + "needs_standardization": None, + "sample_keys": sample_keys or [], + } + + +def detect_dataset_format(dataset) -> dict: + sample = _first_row(dataset) + if sample is None: + return _unknown_dataset_format() + column_names = set(sample.keys()) + if {"instruction", "output"}.issubset(column_names): + return { + "format": "alpaca", + "chat_column": None, + "needs_standardization": False, + "sample_keys": [], + } + + chat_column = None + if "messages" in column_names: + chat_column = "messages" + elif "conversations" in column_names: + chat_column = "conversations" + elif "texts" in column_names: + chat_column = "texts" + + if not chat_column: + return _unknown_dataset_format() + + chat_data = sample.get(chat_column) + if not isinstance(chat_data, (list, tuple)) or not chat_data: + return _unknown_dataset_format(chat_column) + first_msg = chat_data[0] + if not isinstance(first_msg, dict): + return _unknown_dataset_format(chat_column) + msg_keys = set(first_msg.keys()) + sample_keys = [str(key) for key in msg_keys] + if "from" in msg_keys or "value" in msg_keys: + return { + "format": "sharegpt", + "chat_column": chat_column, + "needs_standardization": True, + "sample_keys": sample_keys, + } + if "role" in msg_keys and "content" in msg_keys: + return { + "format": "chatml", + "chat_column": chat_column, + "needs_standardization": False, + "sample_keys": sample_keys, + } + return _unknown_dataset_format(chat_column, sample_keys) + + +def detect_custom_format_heuristic(dataset): + sample = _first_row(dataset) + if sample is None: + return None + all_columns = list(sample.keys()) + mapping = {} + assistant_words = [ + "output", + "answer", + "response", + "assistant", + "completion", + "expected", + "recommendation", + "reply", + "result", + "target", + "solution", + "explanation", + "solve", + ] + user_words_high_priority = [ + "input", + "question", + "query", + "prompt", + "instruction", + "request", + "snippet", + "user", + "text", + "problem", + "exercise", + ] + user_words_low_priority = ["task"] + user_words = user_words_high_priority + user_words_low_priority + system_words = [ + "system", + "context", + "description", + "persona", + "role", + "template", + "task", + ] + metadata_exact_match = { + "id", + "idx", + "index", + "key", + "timestamp", + "date", + "metadata", + "source", + "kind", + "type", + "category", + "score", + "label", + "tag", + "inference_mode", + } + metadata_prefix_patterns = [ + "problem_type", + "problem_source", + "generation_model", + "pass_rate", + ] + priority_patterns = { + "generated": 100, + "gen_": 90, + "model_": 80, + "predicted": 70, + "completion": 60, + } + + def has_keyword(col_name, keywords): + col_lower = col_name.lower() + col_normalized = col_lower.replace("_", "").replace("-", "").replace(" ", "") + return any(keyword in col_lower or keyword in col_normalized for keyword in keywords) + + def is_metadata(col_name): + col_lower = col_name.lower() + if col_lower in metadata_exact_match or col_lower in metadata_prefix_patterns: + return True + for pattern in metadata_prefix_patterns: + if col_lower.startswith(pattern.split("_")[0] + "_") and col_lower != pattern: + if "_" in col_lower: + prefix = col_lower.split("_")[0] + if prefix in ["generation", "pass", "inference"]: + return True + return len(col_lower) <= 2 and col_lower not in ["qa", "q", "a"] + + def get_priority_score(col_name): + col_lower = col_name.lower() + return sum(score for pattern, score in priority_patterns.items() if pattern in col_lower) + + def get_content_length(col_name): + try: + return len(str(sample[col_name])) if sample.get(col_name) else 0 + except Exception: + return 0 + + def score_column(col_name, keywords, role_type, num_candidates): + if not has_keyword(col_name, keywords): + return 0 + score = 10 + if role_type == "user": + col_lower = col_name.lower() + if "task" in col_lower and not any(kw in col_lower for kw in user_words_high_priority): + score -= 15 + score += get_priority_score(col_name) + if role_type in ["assistant", "user"]: + avg_length = get_content_length(col_name) + if num_candidates > 1: + if avg_length > 1000: + score += 50 + elif avg_length > 200: + score += 30 + elif avg_length > 50: + score += 10 + elif avg_length < 50: + score -= 20 + else: + if avg_length > 1000: + score += 50 + elif avg_length > 200: + score += 30 + elif avg_length > 50: + score += 10 + return score + + content_columns = [col for col in all_columns if not is_metadata(col)] + assistant_potential = [col for col in content_columns if has_keyword(col, assistant_words)] + user_potential = [col for col in content_columns if has_keyword(col, user_words)] + assistant_candidates = [ + (col, score) + for col in assistant_potential + if (score := score_column(col, assistant_words, "assistant", len(assistant_potential))) > 0 + ] + if assistant_candidates: + assistant_candidates.sort(key = lambda item: item[1], reverse = True) + assistant_col = assistant_candidates[0][0] + mapping[assistant_col] = "assistant" + else: + assistant_col = None + + user_candidates = [] + for col in user_potential: + if col == assistant_col: + continue + score = score_column(col, user_words, "user", len(user_potential)) + if score > 0: + user_candidates.append((col, score)) + if user_candidates: + user_candidates.sort(key = lambda item: item[1], reverse = True) + user_col = user_candidates[0][0] + mapping[user_col] = "user" + else: + user_col = None + + remaining_columns = [col for col in content_columns if col not in mapping] + system_col = None + for col in remaining_columns: + if has_keyword(col, system_words): + mapping[col] = "system" + system_col = col + break + if system_col: + remaining_columns = [col for col in remaining_columns if col != system_col] + if remaining_columns: + remaining_col = remaining_columns[0] + if not has_keyword(remaining_col, user_words + assistant_words): + mapping[remaining_col] = "system" + elif user_col is None: + mapping[remaining_col] = "user" + else: + mapping[remaining_col] = "system" + + has_user = any(role == "user" for role in mapping.values()) + has_assistant = any(role == "assistant" for role in mapping.values()) + if not has_user: + for col in remaining_columns: + if col not in mapping: + mapping[col] = "user" + has_user = True + break + return mapping if has_user and has_assistant else None + + +_AUDIO_EXTENSIONS = ( + ".wav", + ".mp3", + ".flac", + ".ogg", + ".opus", + ".m4a", + ".aac", + ".wma", + ".webm", +) + + +def _is_audio_value(value) -> bool: + if value is None: + return False + if isinstance(value, dict): + if "array" in value and "sampling_rate" in value: + return True + if "bytes" in value or "path" in value: + path = value.get("path") or "" + return isinstance(path, str) and any( + path.lower().endswith(ext) for ext in _AUDIO_EXTENSIONS + ) + return False + + +def _has_image_header(data: bytes) -> bool: + if len(data) < 4: + return False + return ( + data[:2] == b"\xff\xd8" + or data[:4] == b"\x89PNG" + or data[:3] == b"GIF" + or (data[:4] == b"RIFF" and len(data) >= 12 and data[8:12] == b"WEBP") + or data[:2] == b"BM" + ) + + +def _is_image_value(value) -> bool: + if value is None: + return False + try: + from PIL.Image import Image as PILImage + if isinstance(value, PILImage): + return True + except ImportError: + pass + if isinstance(value, dict): + if "array" in value and "sampling_rate" in value: + return False + if "bytes" in value and "path" in value: + path = value.get("path") or "" + if isinstance(path, str) and any( + path.lower().endswith(ext) for ext in _AUDIO_EXTENSIONS + ): + return False + return True + if isinstance(value, (bytes, bytearray)): + return _has_image_header(value) + image_exts = (".png", ".jpg", ".jpeg", ".webp", ".gif", ".bmp", ".tiff", ".svg") + if isinstance(value, str) and len(value) < 1000: + lower = value.strip().lower() + if lower.startswith(("http://", "https://")): + return any(lower.split("?")[0].endswith(ext) for ext in image_exts) + return any(lower.endswith(ext) for ext in image_exts) + return False + + +def detect_multimodal_dataset(dataset): + sample = _first_row(dataset) + if sample is None: + return { + "is_image": False, + "multimodal_columns": [], + "modality_types": [], + "is_audio": False, + "audio_columns": [], + "detected_audio_column": None, + "detected_text_column": None, + "detected_speaker_column": None, + } + column_names = list(sample.keys()) + image_keywords = [ + "image", + "img", + "pixel", + "jpg", + "jpeg", + "png", + "webp", + "bmp", + "gif", + "tiff", + "svg", + "photo", + "pic", + "picture", + "visual", + "file_name", + "filename", + ] + audio_keywords = ["audio", "speech", "wav", "waveform", "sound"] + multimodal_columns = [] + audio_columns = [] + modality_types = set() + for col_name in column_names: + if any(_keyword_in_column(keyword, col_name) for keyword in image_keywords): + multimodal_columns.append(col_name) + modality_types.add("image") + for col_name in column_names: + if col_name not in multimodal_columns and _is_image_value(sample[col_name]): + multimodal_columns.append(col_name) + modality_types.add("image") + for col_name in column_names: + if any(_keyword_in_column(keyword, col_name) for keyword in audio_keywords): + audio_columns.append(col_name) + modality_types.add("audio") + for col_name in column_names: + if col_name not in audio_columns and _is_audio_value(sample[col_name]): + audio_columns.append(col_name) + modality_types.add("audio") + if audio_columns: + multimodal_columns = [col for col in multimodal_columns if col not in set(audio_columns)] + + detected_text_col = None + if audio_columns: + for col_name in column_names: + if col_name.lower() in [ + "text", + "sentence", + "transcript", + "transcription", + "label", + ]: + detected_text_col = col_name + break + detected_speaker_col = None + if audio_columns: + for col_name in column_names: + if col_name.lower() in ["source", "speaker", "speaker_id"]: + detected_speaker_col = col_name + break + return { + "is_image": len(multimodal_columns) > 0, + "multimodal_columns": multimodal_columns, + "modality_types": list(modality_types), + "is_audio": len(audio_columns) > 0, + "audio_columns": audio_columns, + "detected_audio_column": audio_columns[0] if audio_columns else None, + "detected_text_column": detected_text_col, + "detected_speaker_column": detected_speaker_col, + } + + +def detect_vlm_dataset_structure(dataset): + sample = _first_row(dataset) + if sample is None: + return { + "format": "unknown", + "needs_conversion": None, + "image_column": None, + "text_column": None, + "messages_column": None, + } + column_names = set(sample.keys()) + if "messages" in column_names: + messages = sample["messages"] + if messages and len(messages) > 0: + first_msg = messages[0] + if "content" in first_msg: + content = first_msg["content"] + if ( + isinstance(content, list) + and content + and isinstance(content[0], dict) + and "type" in content[0] + ): + has_index = any("index" in item for item in content if isinstance(item, dict)) + if has_index and "images" in column_names: + return { + "format": "vlm_messages_llava", + "needs_conversion": True, + "messages_column": "messages", + "image_column": "images", + "text_column": None, + } + has_image = any("image" in item for item in content if isinstance(item, dict)) + if has_image: + return { + "format": "vlm_messages", + "needs_conversion": False, + "messages_column": "messages", + "image_column": None, + "text_column": None, + } + + for chat_col in ("conversations", "messages"): + if chat_col not in column_names: + continue + chat_data = sample[chat_col] + if not isinstance(chat_data, list) or not chat_data: + continue + has_image_placeholder = any( + "" in str(message.get("value", "") or message.get("content", "")) + for message in chat_data + if isinstance(message, dict) + ) + if not has_image_placeholder: + continue + image_col = next( + ( + col + for col in column_names + if col != chat_col + and (_keyword_in_column("image", col) or _keyword_in_column("img", col)) + ), + None, + ) + if image_col: + return { + "format": "sharegpt_with_images", + "needs_conversion": True, + "image_column": image_col, + "text_column": None, + "messages_column": chat_col, + } + + metadata_suffixes = ( + "_id", + "_url", + "_name", + "_filename", + "_uri", + "_link", + "_key", + "_index", + ) + metadata_prefixes = ( + "id_", + "url_", + "name_", + "filename_", + "uri_", + "link_", + "key_", + "index_", + ) + image_keywords = [ + "image", + "img", + "photo", + "picture", + "pic", + "visual", + "scan", + "file_name", + "filename", + ] + text_keywords = [ + "text", + "caption", + "captions", + "description", + "answer", + "output", + "response", + "label", + ] + + def is_metadata_column(col_name): + lower = col_name.lower() + return any(lower.endswith(suffix) for suffix in metadata_suffixes) or any( + lower.startswith(prefix) for prefix in metadata_prefixes + ) + + image_candidates = [] + for col in column_names: + value = sample[col] + if any(_keyword_in_column(keyword, col) for keyword in image_keywords) or _is_image_value( + value + ): + if hasattr(value, "size") and hasattr(value, "mode"): + score = 100 + elif isinstance(value, dict) and ("bytes" in value or "path" in value): + score = 75 + elif isinstance(value, str): + score = ( + 55 + if is_metadata_column(col) + else 70 + if value.startswith(("http://", "https://")) + else 50 + ) + else: + score = 0 + if score > 0: + image_candidates.append((col, score)) + image_candidates.sort(key = lambda item: item[1], reverse = True) + + text_candidates = [] + for col in column_names: + if is_metadata_column(col) or not any( + _keyword_in_column(keyword, col) for keyword in text_keywords + ): + continue + value = sample[col] + if isinstance(value, str) and value: + text_candidates.append((col, min(len(value), 1000))) + elif isinstance(value, list) and value and isinstance(value[0], str): + text_candidates.append((col, min(len(value[0]), 1000) // 2)) + text_candidates.sort(key = lambda item: item[1], reverse = True) + + found_image = image_candidates[0][0] if image_candidates else None + found_text = text_candidates[0][0] if text_candidates else None + if found_image and found_text: + return { + "format": "simple_image_text", + "needs_conversion": True, + "image_column": found_image, + "text_column": found_text, + "messages_column": None, + } + return { + "format": "unknown", + "needs_conversion": None, + "image_column": found_image, + "text_column": found_text, + "messages_column": None, + } + + +def check_dataset_format(dataset, is_vlm: bool = False) -> dict: + sample = _first_row(dataset) + columns = _column_names(dataset, sample) + multimodal_info = detect_multimodal_dataset(dataset) + is_audio = multimodal_info.get("is_audio", False) + audio_fields = { + "is_audio": is_audio, + "detected_audio_column": multimodal_info.get("detected_audio_column"), + "detected_speaker_column": multimodal_info.get("detected_speaker_column"), + } + + if is_vlm: + vlm_structure = detect_vlm_dataset_structure(dataset) + requires_mapping = vlm_structure["format"] == "unknown" + warning = None + if requires_mapping: + missing = [] + if not vlm_structure.get("image_column"): + missing.append("image") + if not vlm_structure.get("text_column"): + missing.append("text") + if missing: + warning = ( + f"Could not auto-detect {' or '.join(missing)} column. " + "Please assign image and text columns manually." + ) + return { + "requires_manual_mapping": requires_mapping, + "detected_format": vlm_structure["format"], + "columns": columns, + "suggested_mapping": None, + "detected_image_column": vlm_structure.get("image_column"), + "detected_text_column": vlm_structure.get("text_column"), + "is_image": multimodal_info["is_image"], + "multimodal_columns": multimodal_info.get("multimodal_columns"), + "warning": warning, + **audio_fields, + } + + if is_audio: + detected_audio = multimodal_info.get("detected_audio_column") + detected_text = multimodal_info.get("detected_text_column") + return { + "requires_manual_mapping": not detected_audio or not detected_text, + "detected_format": "audio", + "columns": columns, + "suggested_mapping": None, + "detected_image_column": None, + "detected_text_column": detected_text, + "is_image": False, + "multimodal_columns": multimodal_info.get("audio_columns"), + **audio_fields, + } + + detected = detect_dataset_format(dataset) + if detected["format"] == "unknown": + heuristic_mapping = detect_custom_format_heuristic(dataset) + if heuristic_mapping: + return { + "requires_manual_mapping": False, + "detected_format": "custom_heuristic", + "columns": columns, + "suggested_mapping": heuristic_mapping, + "detected_image_column": None, + "detected_text_column": None, + "is_image": multimodal_info["is_image"], + "multimodal_columns": multimodal_info.get("multimodal_columns"), + **audio_fields, + } + return { + "requires_manual_mapping": True, + "detected_format": "unknown", + "columns": columns, + "suggested_mapping": None, + "detected_image_column": None, + "detected_text_column": None, + "is_image": multimodal_info["is_image"], + "multimodal_columns": multimodal_info.get("multimodal_columns"), + "warning": ( + f"Could not auto-detect column roles for columns: {columns}. " + "Please assign roles manually, or use AI Assist." + ), + **audio_fields, + } + + return { + "requires_manual_mapping": False, + "detected_format": detected["format"], + "columns": columns, + "suggested_mapping": None, + "detected_image_column": None, + "detected_text_column": None, + "is_image": multimodal_info["is_image"], + "multimodal_columns": multimodal_info.get("multimodal_columns"), + **audio_fields, + } + + +_ROLE_MAP = { + "human": "user", + "user": "user", + "input": "user", + "gpt": "assistant", + "assistant": "assistant", + "output": "assistant", + "system": "system", +} + + +def _standardize_sharegpt_row(row: dict[str, Any], chat_column: str) -> dict[str, Any]: + chat_data = row.get(chat_column) + if not isinstance(chat_data, list): + return row + messages = [] + for message in chat_data: + if not isinstance(message, dict): + continue + role = message.get("role") or message.get("from") + content = message.get("content") if "content" in message else message.get("value") + messages.append( + { + "role": _ROLE_MAP.get(str(role), str(role or "user")), + "content": "" if content is None else content, + } + ) + return {chat_column: messages} + + +def format_dataset_preview(dataset): + detected = detect_dataset_format(dataset) + if detected.get("format") != "sharegpt": + return dataset + chat_column = detected.get("chat_column") + if not isinstance(chat_column, str): + return dataset + + if hasattr(dataset, "map"): + return dataset.map(lambda row: _standardize_sharegpt_row(row, chat_column)) + return dataset diff --git a/studio/backend/hub/utils/download_manifest.py b/studio/backend/hub/utils/download_manifest.py new file mode 100644 index 0000000000..5366689296 --- /dev/null +++ b/studio/backend/hub/utils/download_manifest.py @@ -0,0 +1,487 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Hub download manifest + cancel-marker primitives. + +Manifests record what a download was supposed to fetch (path + declared +size per expected file). Consumed by: + - the worker post-download, to verify on-disk sizes match what HF + declared, so a resume that no-ops doesn't get classified as success; + - the inventory scanner, to mark a row partial when expected files + are absent or undersized, so a half-finished GGUF/dataset doesn't + masquerade as a complete on-device row. + +Cancel markers record that a user-initiated cancel landed for a +(repo_type, repo_id, variant) triple. *Existence* is the signal the +scanner reads; the body carries debuggability metadata. Markers are +cleared at the start of a new download attempt (supersedes prior cancel) +and on successful completion (defensive, in case the start clear failed). + +I/O contracts: + - Writes are atomic via ``tmp + os.replace``: a SIGKILL mid-write + cannot leave a half-written file readable to the next reader. + - Manifest reads fail *open*: missing/corrupt/schema-mismatched + manifests return ``None`` and the scanner falls through to the + legacy on-disk-only check (matches HF-cache imports and pre-fix + downloads that never wrote a manifest). + - Cancel-marker reads fail *closed*: file existence is the signal + regardless of body parseability, so a corrupt marker still + suppresses the "on device" classification. +""" + +from __future__ import annotations + +import json +import os +import uuid +from dataclasses import dataclass +from datetime import datetime, timezone +from pathlib import Path +from typing import Iterator, Optional, Sequence + +from loggers import get_logger + +from hub.utils.state_dir import ( + RepoType, + cancelled_dir, + manifest_path, + manifests_dir, + marker_path, + variant_filename_prefix, +) + +logger = get_logger(__name__) + + +_MANIFEST_VERSION = 1 +_MARKER_VERSION = 2 +_LEGACY_MARKER_VERSION = 1 + +# Verbatim phrase the worker emits on a degraded completion and the download +# lifecycle escalates to a warning log. Shared so the emit and match stay coupled. +MANIFEST_DEGRADED_MARKER = "completed without a manifest so partial detection is degraded" + + +@dataclass(frozen = True) +class ExpectedFile: + path: str + size: int + sha256: Optional[str] = None + + +@dataclass(frozen = True) +class Manifest: + repo_type: RepoType + repo_id: str + variant: Optional[str] + started_at: str + expected_files: tuple[ExpectedFile, ...] + transport: Optional[str] = None + + +@dataclass(frozen = True) +class VerifyResult: + ok: bool + missing: tuple[str, ...] + size_mismatched: tuple[str, ...] + + +def _atomic_write_json(path: Path, payload: dict) -> bool: + # Per-write uuid suffix so a concurrent caller or a stale tmp from a + # previous crash cannot collide with the in-flight write. + tmp = path.with_name(f".{path.name}.tmp-{uuid.uuid4().hex[:8]}") + try: + with tmp.open("w", encoding = "utf-8") as handle: + handle.write(json.dumps(payload, indent = 2)) + handle.flush() + os.fsync(handle.fileno()) + os.replace(tmp, path) + except OSError as exc: + logger.debug("Atomic write failed for %s: %s", path, exc) + try: + tmp.unlink(missing_ok = True) + except OSError: + pass + return False + if os.name != "nt": + try: + flags = os.O_RDONLY + if hasattr(os, "O_DIRECTORY"): + flags |= os.O_DIRECTORY + parent_fd = os.open(path.parent, flags) + try: + os.fsync(parent_fd) + finally: + os.close(parent_fd) + except OSError as exc: + logger.debug("Parent dir fsync failed for %s: %s", path, exc) + return True + + +def write_manifest( + repo_type: RepoType, + repo_id: str, + variant: Optional[str], + expected_files: Sequence[ExpectedFile], + transport: Optional[str] = None, +) -> bool: + """Write/overwrite the manifest for this triple. Best-effort. + + ``False`` on write failure must not be treated as fatal: the + worst-case fallback is the pre-fix scanner behavior (one missed + partial detection), which is no regression. + """ + path = manifest_path(repo_type, repo_id, variant) + if path is None: + return False + payload = { + "version": _MANIFEST_VERSION, + "repo_type": repo_type, + "repo_id": repo_id, + "variant": variant, + "started_at": datetime.now(timezone.utc).isoformat(), + "expected_files": [ + { + "path": f.path, + "size": int(f.size), + **({"sha256": f.sha256} if f.sha256 else {}), + } + for f in expected_files + ], + "transport": transport, + } + return _atomic_write_json(path, payload) + + +def read_manifest( + repo_type: RepoType, + repo_id: str, + variant: Optional[str] = None, +) -> Optional[Manifest]: + """Return the manifest if present and parseable; ``None`` otherwise. + + Treats missing-file, parse-error, and any schema mismatch all as + ``None`` (fail-open). Scanner callers fall through to on-disk-only + behavior on ``None`` so this never regresses legacy/imported repos + that have no manifest. + + Forward-compat: accepts only ``version == 1``; an unknown version is + treated as no manifest. A future v2 schema MUST either keep v1's + ``expected_files`` shape on the same filename (bump + ``_MANIFEST_VERSION`` and widen this check) or live under a different + filename, so an incompatible payload can never mis-classify rows. + """ + path = manifest_path(repo_type, repo_id, variant) + if path is None or not path.is_file(): + return None + try: + data = json.loads(path.read_text(encoding = "utf-8")) + except (OSError, ValueError) as exc: + logger.debug("Could not read manifest %s: %s", path, exc) + return None + if not isinstance(data, dict): + return None + if data.get("version") != _MANIFEST_VERSION: + logger.debug( + "Manifest %s has unknown version %r; ignoring.", + path, + data.get("version"), + ) + return None + raw_files = data.get("expected_files") + if not isinstance(raw_files, list): + return None + expected: list[ExpectedFile] = [] + for item in raw_files: + if not isinstance(item, dict): + return None + file_path = item.get("path") + size = item.get("size") + if not isinstance(file_path, str) or not isinstance(size, int): + return None + sha256 = item.get("sha256") + expected.append( + ExpectedFile( + path = file_path, + size = size, + sha256 = sha256 if isinstance(sha256, str) and sha256 else None, + ) + ) + raw_variant = data.get("variant") + transport = data.get("transport") + return Manifest( + repo_type = repo_type, + repo_id = str(data.get("repo_id", repo_id)), + variant = raw_variant if raw_variant else None, + started_at = str(data.get("started_at", "")), + expected_files = tuple(expected), + transport = transport if transport in ("http", "xet") else None, + ) + + +def verify_against_disk(manifest: Manifest, snapshot_dir: Path) -> VerifyResult: + """Check every expected file is present in *snapshot_dir* at its declared size. + + Presence + size only, not content integrity: it converts a + no-op-on-cached ``snapshot_download`` into a clear error when shards are + missing or truncated, and marks a scanner row partial when expected bytes + aren't on disk. Byte-level integrity is already covered upstream by + ``huggingface_hub`` (size check on HTTP, content-addressed chunk hashes on + XET), so re-hashing finalized multi-GB weights here would only duplicate + that at a large cost. ``Path.stat()`` follows symlinks, so HF's symlink and + Windows copy cache layouts both verify correctly. + """ + missing: list[str] = [] + mismatched: list[str] = [] + for expected in manifest.expected_files: + target = snapshot_dir / expected.path + try: + actual_size = target.stat().st_size + except OSError: + missing.append(expected.path) + continue + # expected.size == 0 means HF metadata had no declared size: verify + # existence only rather than flagging every such file as mismatched. + if expected.size > 0 and actual_size != expected.size: + mismatched.append(expected.path) + return VerifyResult( + ok = not missing and not mismatched, + missing = tuple(missing), + size_mismatched = tuple(mismatched), + ) + + +def expected_files_from_snapshot_dir(snapshot_dir: Path) -> list[ExpectedFile]: + """Derive expected-file entries from a completed snapshot directory. + + Last-resort manifest source for when HF metadata was unreachable for the + whole download. ``snapshot_download`` has already exited cleanly, so every + regular file is a finished, correctly-sized blob; recording them keeps the + scanner's completion check in agreement with the worker's exit-0 success + instead of leaving a finished repo perpetually partial. ``stat()`` follows + HF's symlink layout and Windows copies, so the recorded sizes match what + ``verify_against_disk`` later reads. + """ + out: list[ExpectedFile] = [] + try: + entries = sorted(snapshot_dir.rglob("*")) + except OSError: + return out + for path in entries: + try: + if not path.is_file(): + continue + relative = path.relative_to(snapshot_dir).as_posix() + out.append( + ExpectedFile( + path = relative, + size = path.stat().st_size, + sha256 = None, + ) + ) + except OSError: + continue + return out + + +def write_cancel_marker( + repo_type: RepoType, + repo_id: str, + variant: Optional[str] = None, + transport: Optional[str] = None, +) -> bool: + """Record that this triple was cancelled. Idempotent across repeated cancels. + + ``transport`` ("http"/"xet") is surfaced via partial_transport on + inventory rows so the UI labels HTTP retries as continuable and XET + retries as full redownloads. None is accepted for forward-compat. + """ + path = marker_path(repo_type, repo_id, variant) + if path is None: + return False + payload = { + "version": _MARKER_VERSION, + "repo_type": repo_type, + "repo_id": repo_id, + "variant": variant, + "transport": transport, + "cancelled_at": datetime.now(timezone.utc).isoformat(), + } + return _atomic_write_json(path, payload) + + +def read_cancel_marker_transport( + repo_type: RepoType, + repo_id: str, + variant: Optional[str] = None, +) -> Optional[str]: + """Return the transport recorded in the cancel marker, or ``None`` if no + marker exists or it is unreadable. + + Cases: + + * No marker on disk → ``None``. + * Legacy v1 marker → ``"http"``: v1 markers were only written by the + HTTP path, so the transport is unambiguous despite the absent field. + * v2 marker with a valid ``"http"`` / ``"xet"`` transport → that value. + * Corrupt, non-dict, or v2-with-missing-transport marker → ``None``. + Defaulting these to ``"http"`` misled the UI into showing a + byte-resume "Continue" label for what may have been an XET cancel; + ``None`` keeps the neutral "Retry" label. + * Unknown future versions → ``None`` (unknown layout, unknown transport). + """ + path = marker_path(repo_type, repo_id, variant) + if path is None or not path.is_file(): + return None + try: + data = json.loads(path.read_text(encoding = "utf-8")) + except (OSError, ValueError) as exc: + logger.debug("Could not read cancel marker %s: %s", path, exc) + return None + if not isinstance(data, dict): + return None + version = data.get("version") + if version == _LEGACY_MARKER_VERSION: + return "http" + if version != _MARKER_VERSION: + return None + transport = data.get("transport") + if isinstance(transport, str) and transport in ("http", "xet"): + return transport + return None + + +def clear_cancel_marker( + repo_type: RepoType, + repo_id: str, + variant: Optional[str] = None, +) -> None: + """Remove the cancel marker for this triple if present. + + Idempotent: a missing marker is not an error. Called at + download-start (a fresh attempt supersedes prior cancel state) and + again at successful completion (cleans up if the start clear failed). + """ + path = marker_path(repo_type, repo_id, variant) + if path is None: + return + try: + path.unlink(missing_ok = True) + except OSError as exc: + logger.debug("Could not clear cancel marker %s: %s", path, exc) + + +def has_cancel_marker( + repo_type: RepoType, + repo_id: str, + variant: Optional[str] = None, +) -> bool: + """File-existence check only. Body is never read. + + Fail-closed: a corrupt marker still returns ``True`` because the + file's existence is the signal (the user once cancelled this + triple, even if the body is unreadable). + """ + path = marker_path(repo_type, repo_id, variant) + if path is None: + return False + try: + return path.is_file() + except OSError: + return False + + +def delete_manifest( + repo_type: RepoType, + repo_id: str, + variant: Optional[str] = None, +) -> bool: + path = manifest_path(repo_type, repo_id, variant) + if path is None: + return False + try: + if not path.is_file(): + return False + path.unlink() + return True + except OSError as exc: + logger.debug("Could not delete manifest %s: %s", path, exc) + return False + + +def purge_state( + repo_type: RepoType, + repo_id: str, + variant: Optional[str] = None, +) -> bool: + """Remove manifest + cancel marker for this triple. Returns ``True`` + when anything was present on disk before the call. Idempotent.""" + marker_existed = has_cancel_marker(repo_type, repo_id, variant) + manifest_removed = delete_manifest(repo_type, repo_id, variant) + clear_cancel_marker(repo_type, repo_id, variant) + return marker_existed or manifest_removed + + +def purge_all_state_for_repo(repo_type: RepoType, repo_id: str) -> int: + """Remove the snapshot-level manifest + marker AND every variant-keyed + manifest + marker for this repo. Used by the route delete handlers so + scanner state never outlives the cache it described. Returns the count + of (repo, variant) triples that had any state on disk.""" + removed = 0 + if purge_state(repo_type, repo_id, None): + removed += 1 + variants: set[str] = set() + for variant, _ in iter_variant_manifests(repo_type, repo_id): + variants.add(variant) + for variant, _ in iter_variant_markers(repo_type, repo_id): + variants.add(variant) + for variant in variants: + if purge_state(repo_type, repo_id, variant): + removed += 1 + return removed + + +def _variant_from_state_file(path: Path, fallback: str) -> str: + try: + data = json.loads(path.read_text(encoding = "utf-8")) + except (OSError, ValueError): + return fallback + if not isinstance(data, dict): + return fallback + variant = data.get("variant") + return variant if isinstance(variant, str) and variant else fallback + + +def _iter_variant_state_files( + parent: Optional[Path], repo_type: RepoType, repo_id: str +) -> Iterator[tuple[str, Path]]: + if parent is None: + return + prefix = variant_filename_prefix(repo_type, repo_id) + try: + entries = list(parent.iterdir()) + except OSError: + return + for entry in entries: + if not entry.is_file() or not entry.name.endswith(".json"): + continue + stem = entry.name[: -len(".json")] + if not stem.lower().startswith(prefix): + continue + variant = stem[len(prefix) :] + if variant: + yield _variant_from_state_file(entry, variant), entry + + +def iter_variant_manifests(repo_type: RepoType, repo_id: str) -> Iterator[tuple[str, Path]]: + """Yield (variant, manifest_path) for every variant-keyed manifest + written for this repo. Used by is_gguf_repo_partial to enumerate all + variants present on disk so the all-variants-broken gate can run.""" + yield from _iter_variant_state_files(manifests_dir(), repo_type, repo_id) + + +def iter_variant_markers(repo_type: RepoType, repo_id: str) -> Iterator[tuple[str, Path]]: + """Yield (variant, marker_path) for every variant-keyed cancel marker. + Companion to iter_variant_manifests: catches variants cancelled + before download-start ever wrote a manifest (very early failures).""" + yield from _iter_variant_state_files(cancelled_dir(), repo_type, repo_id) diff --git a/studio/backend/hub/utils/download_registry.py b/studio/backend/hub/utils/download_registry.py new file mode 100644 index 0000000000..777d63e1b5 --- /dev/null +++ b/studio/backend/hub/utils/download_registry.py @@ -0,0 +1,1263 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""HF cache inspection, download registry state, and orphan-worker reaping. + +Worker spawning and exit handling live in +:mod:`hub.services.download_lifecycle`; this module owns the registry state +machine plus the cache/marker inspection those workers depend on. + +Resume model +------------ +Only the HTTP transport supports true partial-file resume: +huggingface_hub's HTTP resumer opens ``.incomplete`` in append mode +and sends ``Range: bytes={resume_size}-`` to continue from disk. + +The XET transport CANNOT resume from a ``.incomplete`` partial: +``hf_xet.download_files`` rewrites the destination from scratch. +Network-level dedup still happens, but through the separate chunk cache at +``~/.cache/huggingface/xet/chunk-cache``, which these helpers never touch. + +Cross-transport corruption: a partial written by XET (or ``hf_transfer``'s +parallel-Range writer) can be sparse — high reported size, zero-filled +gaps below. Feeding it to the HTTP resumer would produce a correct-sized +blob whose internal bytes are silently wrong. To prevent that, we keep +transport markers at the download's scope (repo for snapshots/datasets, +variant for GGUF) and refuse to inherit an HTTP partial unless the marker +proves the previous writer was the same single-stream sequential writer. + +Marker writes go through tmp+rename in :func:`prepare_cache_for_transport` +before the worker hands off to ``snapshot_download``, so the next process +always sees a consistent provenance signal. +""" + +from __future__ import annotations + +import hashlib +import importlib.util +import json +import os +import re +import shlex +import signal +import subprocess +import sys +import threading +import time +import weakref +from dataclasses import dataclass, field +from pathlib import Path +from typing import Iterator, Literal, Optional + +from loggers import get_logger + +from hub.utils import state_dir +from hub.utils.state_dir import RepoType + +logger = get_logger(__name__) + +from hub.utils.hf_cache_state import ( + INCOMPLETE_SUFFIX, + TRANSPORT_HTTP, + TRANSPORT_XET, + TRANSPORT_MARKER_NAME, + VALID_TRANSPORTS, + has_active_incomplete_blobs, + iter_repo_cache_dirs, + iter_active_repo_cache_dirs, + repo_cache_dir_name, + target_dir_name, + hf_cache_root, +) + + +@dataclass(frozen = True) +class DownloadTransportCapability: + available: bool + reason: Optional[str] = None + + +@dataclass(frozen = True) +class DownloadTransportCapabilities: + http: DownloadTransportCapability + xet: DownloadTransportCapability + + +def get_download_transport_capabilities() -> DownloadTransportCapabilities: + xet_available = importlib.util.find_spec("hf_xet") is not None + return DownloadTransportCapabilities( + http = DownloadTransportCapability(available = True), + xet = DownloadTransportCapability( + available = xet_available, + reason = None + if xet_available + else "Xet transport is unavailable because hf_xet is not installed.", + ), + ) + + +def download_transport_unavailable_reason(transport: str) -> Optional[str]: + if transport == TRANSPORT_HTTP: + return None + if transport == TRANSPORT_XET: + caps = get_download_transport_capabilities().xet + return None if caps.available else caps.reason + return f"Unsupported download transport: {transport}" + + +def _worker_breadcrumb_path(key: str) -> Optional[Path]: + parent = state_dir.workers_dir() + if parent is None: + return None + safe = hashlib.sha256(key.encode("utf-8")).hexdigest()[:32] + return parent / f"{safe}.json" + + +def write_worker_breadcrumb(key: str, pid: int, metadata: Optional["DownloadMetadata"]) -> None: + """Record a live worker's PID so a restarted backend can reap it. Best + effort: a write failure only forfeits boot-time reaping for this worker, + still covered by the worker's own parent-death watchdog.""" + path = _worker_breadcrumb_path(key) + if path is None: + return + payload = { + "pid": int(pid), + "repo_type": metadata.repo_type if metadata is not None else None, + "repo_id": metadata.repo_id if metadata is not None else None, + "variant": metadata.variant if metadata is not None else None, + "transport": metadata.transport if metadata is not None else None, + } + tmp = path.with_name(f".{path.name}.tmp-{pid}") + try: + tmp.write_text(json.dumps(payload), encoding = "utf-8") + os.replace(tmp, path) + except OSError as exc: + logger.debug("Could not write worker breadcrumb %s: %s", path, exc) + try: + tmp.unlink(missing_ok = True) + except OSError: + pass + + +def remove_worker_breadcrumb(key: str) -> None: + path = _worker_breadcrumb_path(key) + if path is None: + return + _safe_unlink(path) + + +def _safe_unlink(path: Path) -> None: + try: + path.unlink(missing_ok = True) + except OSError as exc: + logger.debug("Could not remove %s: %s", path, exc) + + +def _process_alive(pid: int) -> bool: + if sys.platform == "win32": + import ctypes + from ctypes import wintypes + + SYNCHRONIZE = 0x00100000 + ERROR_INVALID_PARAMETER = 87 + kernel32 = ctypes.WinDLL("kernel32", use_last_error = True) + kernel32.OpenProcess.argtypes = [wintypes.DWORD, wintypes.BOOL, wintypes.DWORD] + kernel32.OpenProcess.restype = wintypes.HANDLE + kernel32.CloseHandle.argtypes = [wintypes.HANDLE] + kernel32.CloseHandle.restype = wintypes.BOOL + ctypes.set_last_error(0) + handle = kernel32.OpenProcess(SYNCHRONIZE, False, pid) + if not handle: + return ctypes.get_last_error() != ERROR_INVALID_PARAMETER + kernel32.CloseHandle(handle) + return True + try: + os.kill(pid, 0) + return True + except ProcessLookupError: + return False + except OSError: + return True + + +def _read_process_cmdline(pid: int) -> Optional[str]: + proc_cmdline = Path(f"/proc/{pid}/cmdline") + try: + if proc_cmdline.exists(): + raw = proc_cmdline.read_bytes() + return raw.replace(b"\x00", b" ").decode("utf-8", "replace") + except OSError: + pass + try: + import psutil + return " ".join(psutil.Process(pid).cmdline()) + except Exception: + return None + + +def _cmdline_repo_id(cmdline: str) -> Optional[str]: + try: + args = shlex.split(cmdline) + except ValueError: + args = cmdline.split() + for i, arg in enumerate(args): + if arg == "--repo-id" and i + 1 < len(args): + return args[i + 1] + if arg.startswith("--repo-id="): + return arg.split("=", 1)[1] + return None + + +def _is_our_worker(pid: int, repo_id: Optional[str]) -> bool: + cmdline = _read_process_cmdline(pid) + if cmdline is None: + return False + if "hub.workers.hf_download" not in cmdline: + return False + # Exact --repo-id match: a substring match would let a stale breadcrumb for + # Org/Model reap a live worker for Org/Model-v2. + if isinstance(repo_id, str) and repo_id: + return _cmdline_repo_id(cmdline) == repo_id + return True + + +def _kill_orphan(pid: int) -> None: + try: + os.kill(pid, signal.SIGTERM if sys.platform == "win32" else signal.SIGKILL) + except OSError: + pass + + +def _settle_orphaned_download( + repo_type: Optional[str], + repo_id: Optional[str], + variant: Optional[str], + transport: Optional[str], +) -> None: + """Persist a cancel marker for a reaped orphan still mid-download so the next + launch settles it to a resumable "cancelled" state instead of a phantom-running + row. + + Gated on surviving partial state and on the recorded manifest not already + verifying against an active snapshot, so a download that finished before its + breadcrumb was cleaned up is never mislabeled cancelled. For a GGUF variant + manifest with blob hashes, the partial-state check is scoped to those hashes so + a sibling variant cannot contaminate this orphan's state. The recorded + transport is preserved so the resume affordance stays accurate.""" + if repo_type not in ("model", "dataset") or not repo_id: + return + from hub.utils import download_manifest + + manifest = download_manifest.read_manifest(repo_type, repo_id, variant) + if repo_type == "model" and variant and manifest is None: + return + if manifest is None: + if not has_active_incomplete_blobs(repo_type, repo_id): + return + else: + if _manifest_verifies_against_active_cache(repo_type, repo_id, manifest): + return + if not _manifest_has_active_incomplete_blobs(repo_type, repo_id, manifest): + return + persist_cancel_marker(repo_type, repo_id, variant, transport, logger = logger) + + +def reap_orphan_workers() -> None: + """Kill download workers left running by a previous backend instance. + + Verifies each breadcrumb's PID is alive AND its command line is one of our + workers before terminating, so a recycled PID can't take down an unrelated + process. Partial blobs are never touched, so a reaped download stays + resumable; an interrupted one with bytes on disk is settled to a cancelled + marker (see :func:`_settle_orphaned_download`) so its resume affordance + survives a hard crash like a graceful shutdown's does. Runs once at startup + and never raises.""" + parent = state_dir.workers_dir() + if parent is None: + return + try: + entries = list(parent.iterdir()) + except OSError: + return + for entry in entries: + if not entry.is_file() or not entry.name.endswith(".json"): + continue + try: + data = json.loads(entry.read_text(encoding = "utf-8")) + except (OSError, ValueError): + _safe_unlink(entry) + continue + pid = data.get("pid") if isinstance(data, dict) else None + repo_id = data.get("repo_id") if isinstance(data, dict) else None + if not isinstance(pid, int) or pid <= 0: + _safe_unlink(entry) + continue + try: + if _process_alive(pid) and _is_our_worker(pid, repo_id): + _kill_orphan(pid) + logger.warning( + "Reaped orphaned download worker pid=%s repo=%s from a " + "previous backend instance.", + pid, + repo_id, + ) + _settle_orphaned_download( + data.get("repo_type"), + repo_id, + data.get("variant"), + data.get("transport"), + ) + except Exception as exc: + logger.debug("Reaper failed for breadcrumb %s: %s", entry, exc) + _safe_unlink(entry) + + +def _purge_incomplete_blobs( + entry: Path, + only_hashes: Optional[frozenset[str]] = None, + protected_hashes: Optional[frozenset[str]] = None, +) -> int: + """Delete matching ``*.incomplete`` blobs beneath *entry*; return the count + removed. Per-file failures are swallowed. + + ``only_hashes`` whitelists which partials may be purged; ``None`` means + every partial (full-repo snapshot/dataset). ``protected_hashes`` is honoured + unconditionally, even when ``only_hashes`` is ``None``, so a blob a + concurrent same-repo peer is writing is never purged from under it.""" + blobs_dir = entry / "blobs" + if not blobs_dir.is_dir(): + return 0 + removed = 0 + try: + candidates = list(blobs_dir.iterdir()) + except OSError: + return 0 + for blob in candidates: + try: + if not blob.is_file(): + continue + if not blob.name.endswith(INCOMPLETE_SUFFIX): + continue + blob_hash = blob.name[: -len(INCOMPLETE_SUFFIX)] + if protected_hashes and blob_hash in protected_hashes: + continue + if only_hashes is not None and blob_hash not in only_hashes: + continue + blob.unlink() + removed += 1 + except OSError: + # Swallow; downstream snapshot_download surfaces a precise error if + # it actually can't proceed. + continue + return removed + + +def _iter_active_snapshot_dirs(repo_type: str, repo_id: str) -> Iterator[Path]: + for entry in iter_active_repo_cache_dirs(repo_type, repo_id): + snapshots_dir = entry / "snapshots" + if not snapshots_dir.is_dir(): + continue + try: + snapshots = list(snapshots_dir.iterdir()) + except OSError: + continue + for snapshot in snapshots: + if snapshot.is_dir(): + yield snapshot + + +def _manifest_verifies_against_active_cache(repo_type: str, repo_id: str, manifest) -> bool: + from hub.utils import download_manifest + for snapshot_dir in _iter_active_snapshot_dirs(repo_type, repo_id): + if download_manifest.verify_against_disk(manifest, snapshot_dir).ok: + return True + return False + + +def _manifest_has_active_incomplete_blobs(repo_type: str, repo_id: str, manifest) -> bool: + if not getattr(manifest, "variant", None): + return has_active_incomplete_blobs(repo_type, repo_id) + expected_hashes = frozenset( + expected.sha256 for expected in manifest.expected_files if expected.sha256 + ) + if not expected_hashes: + return has_active_incomplete_blobs(repo_type, repo_id) + return bool( + incomplete_blob_hashes(repo_type, repo_id, active_only = True).intersection(expected_hashes) + ) + + +def _marker_path(entry: Path, variant: Optional[str] = None) -> Path: + if not variant: + return entry / TRANSPORT_MARKER_NAME + digest = hashlib.sha256(variant.strip().lower().encode("utf-8")).hexdigest()[:24] + return entry / f"{TRANSPORT_MARKER_NAME}.gguf-{digest}" + + +def _is_transport_marker_file(path: Path) -> bool: + # Matches ".transport", its tmps, and variant-scoped ".transport.gguf-*". + # Real HF cache entries (blobs/refs/snapshots/.no_exist) never start with + # ".transport.". + return path.name == TRANSPORT_MARKER_NAME or path.name.startswith(f"{TRANSPORT_MARKER_NAME}.") + + +def _companion_marker_path(entry: Path) -> Path: + return entry / f"{TRANSPORT_MARKER_NAME}.companion" + + +def _read_marker_value(marker: Path) -> Optional[str]: + try: + if not marker.exists(): + return None + value = marker.read_text().strip() + except OSError: + return None + return value if value in VALID_TRANSPORTS else None + + +def _write_marker_value(marker: Path, mode: str) -> None: + try: + # tmp + rename so a SIGKILL mid-write can't leave a half-written marker. + # The tmp name is per-process so concurrent writers don't clobber tmps. + tmp = marker.with_name(f"{marker.name}.tmp-{os.getpid()}") + tmp.write_text(mode) + os.replace(tmp, marker) + except OSError: + # Best-effort: a missing marker next run purges the partial defensively, + # the safe failure mode. + pass + + +def _read_marker(entry: Path, variant: Optional[str] = None) -> Optional[str]: + return _read_marker_value(_marker_path(entry, variant)) + + +def _write_marker( + entry: Path, + mode: str, + variant: Optional[str] = None, +) -> None: + _write_marker_value(_marker_path(entry, variant), mode) + + +def _read_companion_marker(entry: Path) -> Optional[str]: + return _read_marker_value(_companion_marker_path(entry)) + + +def _write_companion_marker(entry: Path, mode: str) -> None: + _write_marker_value(_companion_marker_path(entry), mode) + + +def prepare_cache_for_transport( + repo_type: str, + repo_id: str, + mode: str, + variant: Optional[str] = None, + only_blob_hashes: Optional[frozenset[str]] = None, + companion_blob_hashes: Optional[frozenset[str]] = None, + protected_blob_hashes: Optional[frozenset[str]] = None, +) -> int: + """Guarantee any pre-existing ``.incomplete`` blobs are SAFE to resume under + *mode*. Returns the number of partial blobs purged for untrusted provenance. + + Two marker scopes govern GGUF downloads. ``only_blob_hashes`` are the + variant's own (main quant) blobs, judged by the ``variant``-scoped marker; + ``None`` widens the scope to every partial for full-repo snapshots/datasets. + ``companion_blob_hashes`` are blobs shared across sibling variants (a vision + mmproj), judged by a separate repo-scoped companion marker — so a companion + partial is trusted against the transport that wrote it, not against + whichever sibling variant resumes next. + + The contract: + - HTTP mode: a partial is trusted ONLY when its governing marker equals + ``"http"``. Any other case (missing/unreadable/mismatched marker) purges, + since the HTTP resumer would otherwise append to a sparse + XET/parallel-Range partial and silently produce a corrupt blob. + - XET mode: incomplete blobs are purged (``hf_xet.download_files`` rewrites + from scratch, so this only fixes UI accounting — bytes already in CAS are + reused via the chunk-cache). Scoped to ``only_blob_hashes``: companion + blobs fall outside that set and survive (shared, and XET overwrites them). + + ``protected_blob_hashes`` are blobs a concurrent same-repo peer is writing; + they are excluded from every purge so a shared companion is never deleted + mid-write. + + Scope: only the active ``HF_HUB_CACHE`` root is inspected. That suffices for + resume safety because ``snapshot_download`` runs without a ``cache_dir`` + override and so can only read or resume a ``.incomplete`` under this same + active root. Markers are written for the new mode before returning. + """ + if mode not in VALID_TRANSPORTS: + raise ValueError(f"Invalid transport mode: {mode!r}") + root = hf_cache_root(create = True) + if root is None: + return 0 + target = target_dir_name(repo_type, repo_id) + try: + entries = [e for e in root.iterdir() if e.name.lower() == target] + except OSError: + return 0 + if not entries: + # First download: pre-create the repo dir so the marker lands before the + # worker writes any bytes. Otherwise a SIGKILL mid-download leaves a + # partial with no marker that the resume then purges. + canonical = repo_cache_dir_name(repo_type, repo_id) + new_entry = root / canonical + try: + new_entry.mkdir(exist_ok = True) + except OSError: + return 0 + entries = [new_entry] + protected = protected_blob_hashes or frozenset() + has_companion = bool(companion_blob_hashes) + total_purged = 0 + for entry in entries: + if mode == TRANSPORT_XET: + total_purged += _purge_incomplete_blobs(entry, only_blob_hashes, protected) + else: + if _read_marker(entry, variant) != mode: + total_purged += _purge_incomplete_blobs(entry, only_blob_hashes, protected) + if companion_blob_hashes and _read_companion_marker(entry) != mode: + total_purged += _purge_incomplete_blobs(entry, companion_blob_hashes, protected) + _write_marker(entry, mode, variant) + if has_companion: + _write_companion_marker(entry, mode) + return total_purged + + +_HF_TOKEN_RE = re.compile(r"hf_[A-Za-z0-9]{20,}") +_BEARER_RE = re.compile(r"(?i)bearer\s+[A-Za-z0-9._\-]+") + + +def scrub_secrets(text: str, *, hf_token: Optional[str] = None) -> str: + if not text: + return text + cleaned = text + if hf_token: + cleaned = cleaned.replace(hf_token, "***") + cleaned = _BEARER_RE.sub("Bearer ***", cleaned) + cleaned = _HF_TOKEN_RE.sub("hf_***", cleaned) + return cleaned + + +def purge_empty_marker_dir( + repo_type: str, + repo_id: str, + variant: Optional[str] = None, +) -> bool: + """Remove the failed download's own transport marker from a marker-only dir. + + ``prepare_cache_for_transport`` pre-creates the dir + marker before any + download; a failure during validation/auth/network setup leaves the dir as + marker-only litter. Only the failed download's OWN marker is removed (the + repo-scope ``.transport`` or the variant-scoped ``.transport.gguf-*`` plus + its ``.tmp-*`` siblings); a sibling variant's marker and the shared + ``.transport.companion`` are left intact, so cancelling one quant never + strips a peer's provenance. A dir holding ``blobs/``/``snapshots/``/``refs/`` + won't match and is left untouched, so a resumable partial isn't blown away. + """ + cleaned = False + for entry in iter_repo_cache_dirs(repo_type, repo_id): + try: + contents = list(entry.iterdir()) + except OSError: + continue + if not contents or not all(_is_transport_marker_file(item) for item in contents): + continue + own_name = _marker_path(entry, variant).name + own_markers = [ + item + for item in contents + if item.name == own_name or item.name.startswith(f"{own_name}.tmp") + ] + if not own_markers: + continue + try: + for marker in own_markers: + marker.unlink() + except OSError: + continue + cleaned = True + try: + entry.rmdir() + except OSError: + continue + return cleaned + + +def read_active_transport_marker( + repo_type: str, + repo_id: str, + variant: Optional[str] = None, +) -> Optional[str]: + for entry in iter_active_repo_cache_dirs(repo_type, repo_id): + value = _read_marker(entry, variant) + if value is not None: + return value + return None + + +def is_resumable_partial( + repo_type: str, + repo_id: str, + variant: Optional[str] = None, +) -> bool: + """True only when a partial exists AND was produced by a byte-resumable + writer (the HTTP transport). XET partials exist on disk but are discarded on + the next download attempt.""" + if not has_active_incomplete_blobs(repo_type, repo_id): + return False + return read_active_transport_marker(repo_type, repo_id, variant) == TRANSPORT_HTTP + + +def incomplete_blob_hashes( + repo_type: str, + repo_id: str, + *, + active_only: bool = False, +) -> set[str]: + out: set[str] = set() + entries = ( + iter_active_repo_cache_dirs(repo_type, repo_id) + if active_only + else iter_repo_cache_dirs(repo_type, repo_id) + ) + for entry in entries: + blobs_dir = entry / "blobs" + if not blobs_dir.is_dir(): + continue + try: + for blob in blobs_dir.iterdir(): + if blob.is_file() and blob.name.endswith(INCOMPLETE_SUFFIX): + out.add(blob.name[: -len(INCOMPLETE_SUFFIX)]) + except OSError: + continue + return out + + +def completed_blob_bytes(repo_type: str, repo_id: str, blob_hashes: frozenset[str]) -> int: + """Sum finalized blob bytes for *blob_hashes* in the active HF cache root. + + A worker only writes to the active ``HF_HUB_CACHE`` root, so a baseline must + ignore legacy/default roots that ``snapshot_download`` won't reuse this run. + """ + if not blob_hashes: + return 0 + total = 0 + for entry in iter_active_repo_cache_dirs(repo_type, repo_id): + blobs_dir = entry / "blobs" + if not blobs_dir.is_dir(): + continue + for blob_hash in blob_hashes: + blob = blobs_dir / blob_hash + try: + if blob.is_file(): + total += max(0, int(blob.stat().st_size)) + except OSError: + continue + return total + + +def existing_blob_bytes(repo_type: str, repo_id: str, blob_hashes: frozenset[str]) -> int: + """Bytes already on disk (finalized + ``.incomplete``) for *blob_hashes* in + the active HF cache root. A blob is in exactly one state, so summing both + candidate names never double-counts. Used to size what a (possibly resumed) + download still needs to write before the run starts.""" + if not blob_hashes: + return 0 + total = 0 + for entry in iter_active_repo_cache_dirs(repo_type, repo_id): + blobs_dir = entry / "blobs" + if not blobs_dir.is_dir(): + continue + for blob_hash in blob_hashes: + for name in (blob_hash, f"{blob_hash}{INCOMPLETE_SUFFIX}"): + blob = blobs_dir / name + try: + if blob.is_file(): + total += max(0, int(blob.stat().st_size)) + except OSError: + continue + return total + + +JobState = Literal["idle", "running", "cancelling", "cancelled", "complete", "error"] + +TERMINAL_STATES = frozenset({"complete", "cancelled", "error"}) +_ACTIVE_STATES = frozenset({"running", "cancelling"}) + + +@dataclass(frozen = True) +class DownloadState: + state: JobState + error: Optional[str] = None + + +@dataclass(frozen = True) +class DownloadMetadata: + repo_type: RepoType + repo_id: str + variant: Optional[str] + transport: Optional[str] + # GGUF variant main/writable hashes, identifying the variant-specific shards + # for concurrency decisions. + blob_hashes: frozenset[str] = field(default_factory = frozenset) + # Full required hash set for progress/completion (includes the shared mmproj + # companion for vision GGUF repos). + progress_blob_hashes: frozenset[str] = field(default_factory = frozenset) + # Bytes already complete before this job started; not counted as this run's + # progress. + completed_baseline_bytes: int = 0 + + +@dataclass(frozen = True) +class ActiveDownloadRef: + key: str + state: str + metadata: Optional[DownloadMetadata] + generation: int + + +def normalize_repo_key(repo_id: str) -> str: + return repo_id.strip().lower() + + +def normalize_job_key(key: str) -> str: + repo, sep, variant = key.partition("::") + repo_key = normalize_repo_key(repo) + return f"{repo_key}{sep}{variant.strip().lower()}" if sep else repo_key + + +def _repo_of_key(key: str) -> str: + return normalize_repo_key(key.split("::", 1)[0]) + + +def variant_from_key(key: str) -> Optional[str]: + """Parse the variant suffix from a 'repo_id::variant' key. Empty + variant returns None — matches the manifest/marker calling + convention for full-snapshot models and datasets.""" + if "::" not in key: + return None + _, _, variant = key.partition("::") + return variant or None + + +def persist_cancel_marker( + repo_type: Optional[RepoType], + repo_id: Optional[str], + variant: Optional[str], + transport: Optional[str], + *, + logger = logger, +) -> None: + if not repo_type or not repo_id: + return + try: + from hub.utils.download_manifest import write_cancel_marker + if not write_cancel_marker( + repo_type, + repo_id, + variant, + transport = transport, + ): + logger.debug("write_cancel_marker returned False for %s", repo_id) + except Exception as exc: + logger.debug("write_cancel_marker failed for %s: %s", repo_id, exc) + + +_REGISTRIES: "weakref.WeakSet[DownloadRegistry]" = weakref.WeakSet() +_NAMED_REGISTRIES: dict[str, "DownloadRegistry"] = {} +_NAMED_REGISTRIES_LOCK = threading.Lock() + + +def terminate_active_downloads() -> None: + """Best-effort shutdown hook called from the FastAPI lifespan. + + Walks every live DownloadRegistry instance and SIGKILLs any in-flight + workers so the parent exit path doesn't leak zombies. The WeakSet drops + ad-hoc registries (e.g. test fixtures) automatically once their last + strong reference is gone; the long-lived named registries stay reachable + via ``_NAMED_REGISTRIES``. Quiet on its own failures: shutdown must not + raise. + """ + for registry in list(_REGISTRIES): + try: + registry.terminate_all("download") + except Exception as exc: + logger.warning("terminate_active_downloads: %s", exc) + + +class DownloadRegistry: + """Thread-safe state machine for background HF download jobs. + + One instance backs model downloads (keys ``repo_id::variant``) and another + backs dataset downloads (keys ``repo_id``). Repo-scoped tracking serializes + full snapshots, datasets, cross-transport work, and deletes; same-transport + GGUF variants may run concurrently. + """ + + def __init__(self, max_terminal: int = 64) -> None: + self._jobs: dict[str, DownloadState] = {} + self._processes: dict[str, subprocess.Popen] = {} + self._repo_active: dict[str, set[str]] = {} + self._metadata: dict[str, DownloadMetadata] = {} + self._pending_cancel: dict[str, Optional[int]] = {} + self._generations: dict[str, int] = {} + # Monotonic across keys so an evicted then re-claimed key never reuses a + # prior generation (which would let a stale cancel match a new run). + self._generation_seq = 0 + self._deleting: dict[str, set[Optional[str]]] = {} + self._lock = threading.Lock() + _REGISTRIES.add(self) + self._max_terminal = max_terminal + + def _put_terminal_job_locked( + self, + key: str, + state: JobState, + error: Optional[str] = None, + ) -> None: + self._jobs.pop(key, None) + self._jobs[key] = DownloadState(state, error) + if len(self._jobs) > self._max_terminal: + for stale_key, stale in list(self._jobs.items()): + if stale.state in TERMINAL_STATES and stale_key != key: + self._jobs.pop(stale_key, None) + self._metadata.pop(stale_key, None) + self._generations.pop(stale_key, None) + if len(self._jobs) <= self._max_terminal: + break + + def set_job( + self, + key: str, + state: JobState, + error: Optional[str] = None, + ) -> None: + key = normalize_job_key(key) + with self._lock: + if state in TERMINAL_STATES: + self._put_terminal_job_locked(key, state, error) + self._pending_cancel.pop(key, None) + repo = _repo_of_key(key) + active = self._repo_active.get(repo) + if active is not None: + active.discard(key) + if not active: + self._repo_active.pop(repo, None) + else: + self._jobs[key] = DownloadState(state, error) + + def get_job(self, key: str) -> DownloadState: + key = normalize_job_key(key) + with self._lock: + return self._jobs.get(key, DownloadState("idle")) + + def current_generation(self, key: str) -> int: + key = normalize_job_key(key) + with self._lock: + return self._generations.get(key, 0) + + def get_job_metadata(self, key: str) -> Optional[DownloadMetadata]: + key = normalize_job_key(key) + with self._lock: + return self._metadata.get(key) + + def _generation_matches_locked(self, key: str, generation: Optional[int]) -> bool: + key = normalize_job_key(key) + return generation is None or self._generations.get(key, 0) == generation + + def register_process(self, key: str, proc: subprocess.Popen) -> bool: + """Register *proc* for *key*. Returns ``False`` when a cancel was + requested during the claim→register window (the caller must kill + *proc* immediately); ``True`` otherwise.""" + key = normalize_job_key(key) + metadata_to_persist: Optional[DownloadMetadata] = None + registered = False + breadcrumb_metadata: Optional[DownloadMetadata] = None + with self._lock: + has_pending_cancel = key in self._pending_cancel + pending_generation = self._pending_cancel.pop(key, None) + if has_pending_cancel and self._generation_matches_locked( + key, + pending_generation, + ): + self._put_terminal_job_locked(key, "cancelled") + metadata_to_persist = self._metadata.pop(key, None) + repo = _repo_of_key(key) + active = self._repo_active.get(repo) + if active is not None: + active.discard(key) + if not active: + self._repo_active.pop(repo, None) + else: + self._processes[key] = proc + breadcrumb_metadata = self._metadata.get(key) + registered = True + if registered: + try: + write_worker_breadcrumb(key, proc.pid, breadcrumb_metadata) + except Exception as exc: + logger.debug("Could not record worker breadcrumb: %s", exc) + return True + if metadata_to_persist is not None: + persist_cancel_marker( + metadata_to_persist.repo_type, + metadata_to_persist.repo_id, + metadata_to_persist.variant, + metadata_to_persist.transport, + ) + return False + + def mark_pending_cancel( + self, + key: str, + generation: Optional[int] = None, + ) -> bool: + """Record a cancel for an active job whose worker process hasn't + registered yet. Returns ``True`` when the pending cancel was armed, + so :func:`register_process` will kill the process on arrival.""" + key = normalize_job_key(key) + with self._lock: + if self._jobs.get(key, DownloadState("idle")).state not in _ACTIVE_STATES: + return False + if not self._generation_matches_locked(key, generation): + return False + self._pending_cancel[key] = generation + self._jobs[key] = DownloadState("cancelling") + return True + + def cancel_requested(self, key: str) -> bool: + """True when *we* initiated a stop for *key* (a pending cancel armed + before the worker registered, or the job already moved to + ``cancelling``). Lets exit classification tell an intentional kill + apart from an OOM/external SIGKILL.""" + key = normalize_job_key(key) + with self._lock: + if key in self._pending_cancel: + return True + return self._jobs.get(key, DownloadState("idle")).state == "cancelling" + + def get_process(self, key: str) -> Optional[subprocess.Popen]: + key = normalize_job_key(key) + with self._lock: + return self._processes.get(key) + + def drop_process(self, key: str, proc: subprocess.Popen) -> bool: + key = normalize_job_key(key) + with self._lock: + if self._processes.get(key) is not proc: + return False + self._processes.pop(key, None) + remove_worker_breadcrumb(key) + return True + + def claim( + self, + key: str, + transport: str, + *, + repo_type: Optional[RepoType] = None, + repo_id: Optional[str] = None, + variant: Optional[str] = None, + blob_hashes: Optional[frozenset[str]] = None, + progress_blob_hashes: Optional[frozenset[str]] = None, + completed_baseline_bytes: int = 0, + ) -> tuple[bool, str]: + key = normalize_job_key(key) + repo = _repo_of_key(key) + requested_hashes = blob_hashes or frozenset() + requested_progress_hashes = progress_blob_hashes or frozenset() + with self._lock: + deleting_scopes = self._deleting.get(repo) + if deleting_scopes is not None and ( + None in deleting_scopes or variant_from_key(key) in deleting_scopes + ): + return False, "deleting" + active = self._repo_active.get(repo, set()) + stale_keys: list[str] = [] + conflict_state: Optional[str] = None + for other_key in active: + if other_key == key: + continue + other_status = self._jobs.get(other_key) + if other_status is None or other_status.state not in _ACTIVE_STATES: + stale_keys.append(other_key) + continue + other_metadata = self._metadata.get(other_key) + # Same-transport variants of one model run concurrently: each + # worker purges only its own re-resolved main blobs and the + # shared companion is guarded by its marker. Cross-transport + # stays serialized so an HTTP resume and an XET rewrite never + # write one shared blob at once. + concurrent_gguf_variants = ( + repo_type == "model" + and bool(variant) + and other_metadata is not None + and other_metadata.repo_type == "model" + and bool(other_metadata.variant) + and other_metadata.transport == transport + ) + if concurrent_gguf_variants: + continue + conflict_state = other_status.state + break + for stale_key in stale_keys: + active.discard(stale_key) + if conflict_state is not None: + return False, conflict_state + current = self._jobs.get(key, DownloadState("idle")).state + if current in _ACTIVE_STATES: + return False, current + self._generation_seq += 1 + self._generations[key] = self._generation_seq + self._jobs[key] = DownloadState("running") + self._repo_active.setdefault(repo, active).add(key) + if repo_type and repo_id: + self._metadata[key] = DownloadMetadata( + repo_type = repo_type, + repo_id = repo_id, + variant = variant, + transport = transport, + blob_hashes = requested_hashes, + progress_blob_hashes = requested_progress_hashes, + completed_baseline_bytes = max( + 0, + int(completed_baseline_bytes or 0), + ), + ) + else: + self._metadata.pop(key, None) + return True, "running" + + def adoptable(self, key: str) -> bool: + """True when *key* itself has a live job a client can attach to. + + Lets a rejected claim distinguish a collision with this key's own + in-flight job (pollable) from one blocked by a different repo job + or an in-progress delete, where no job exists for this key.""" + key = normalize_job_key(key) + with self._lock: + return self._jobs.get(key, DownloadState("idle")).state in _ACTIVE_STATES + + def _active_job_variant_locked(self, key: str) -> Optional[str]: + metadata = self._metadata.get(key) + if metadata is not None: + return (metadata.variant or "").strip().lower() or None + return variant_from_key(key) + + def _delete_blocked_by_active_locked(self, repo_id: str, variant: Optional[str]) -> bool: + """Whether an active download conflicts with deleting *repo_id*/*variant*. + + A whole-repo delete (``variant is None``) conflicts with any active + download. A variant delete conflicts only with that same variant or a + whole-repo download writing the shared snapshot; other quantizations + download concurrently and never block it.""" + for key in self._repo_active.get(repo_id, set()): + job = self._jobs.get(key) + if job is None or job.state not in _ACTIVE_STATES: + continue + if variant is None: + return True + other_variant = self._active_job_variant_locked(key) + if other_variant is None or other_variant == variant: + return True + return False + + def peer_blob_hashes(self, key: str) -> frozenset[str]: + """Union of the writable blob hashes of every OTHER active download for + this key's repo. A worker excludes these from its purge so it never + deletes an ``.incomplete`` a concurrent same-repo variant is writing + (e.g. a shared mmproj bundled with two GGUF quants).""" + key = normalize_job_key(key) + repo = _repo_of_key(key) + out: set[str] = set() + with self._lock: + for other_key in self._repo_active.get(repo, set()): + if other_key == key: + continue + job = self._jobs.get(other_key) + if job is None or job.state not in _ACTIVE_STATES: + continue + metadata = self._metadata.get(other_key) + if metadata is not None: + out |= set(metadata.progress_blob_hashes or metadata.blob_hashes) + return frozenset(out) + + def active_jobs(self, repo_id: str) -> dict[str, str]: + """Map of every active job key for *repo_id* to its state.""" + repo_id = normalize_repo_key(repo_id) + with self._lock: + result: dict[str, str] = {} + for key in self._repo_active.get(repo_id, set()): + job = self._jobs.get(key) + if job is not None and job.state in _ACTIVE_STATES: + metadata = self._metadata.get(key) + display_key = ( + f"{_repo_of_key(key)}::{metadata.variant}" + if metadata is not None and metadata.variant + else key + ) + result[display_key] = job.state + return result + + def active_job_refs(self, repo_id: Optional[str] = None) -> list[ActiveDownloadRef]: + repo_key = normalize_repo_key(repo_id) if repo_id else None + with self._lock: + if repo_key: + candidate_keys = list(self._repo_active.get(repo_key, set())) + else: + candidate_keys = [key for active in self._repo_active.values() for key in active] + refs: list[ActiveDownloadRef] = [] + for key in candidate_keys: + job = self._jobs.get(key) + if job is None or job.state not in _ACTIVE_STATES: + continue + refs.append( + ActiveDownloadRef( + key = key, + state = job.state, + metadata = self._metadata.get(key), + generation = self._generations.get(key, 0), + ) + ) + return refs + + def begin_delete( + self, + repo_id: str, + variant: Optional[str] = None, + ) -> bool: + """Reserve *repo_id* (or one GGUF *variant* of it) for deletion. Returns + ``False`` when a conflicting download is active (a whole-repo delete vs + any download, a variant delete vs that same variant or a whole-repo + download), so sibling quantizations keep downloading. On success the + scope is marked so :func:`claim` rejects overlapping downloads until + :func:`end_delete` runs, closing the check-then-delete race against a + concurrently spawned worker.""" + repo_id = normalize_repo_key(repo_id) + variant_key = (variant or "").strip().lower() or None + with self._lock: + if self._delete_blocked_by_active_locked(repo_id, variant_key): + return False + self._deleting.setdefault(repo_id, set()).add(variant_key) + return True + + def end_delete( + self, + repo_id: str, + variant: Optional[str] = None, + ) -> None: + repo_id = normalize_repo_key(repo_id) + variant_key = (variant or "").strip().lower() or None + with self._lock: + scopes = self._deleting.get(repo_id) + if scopes is None: + return + scopes.discard(variant_key) + if not scopes: + self._deleting.pop(repo_id, None) + + def has_active_peer_variant(self, repo_id: str, variant: Optional[str]) -> bool: + """Whether a DIFFERENT quantization of *repo_id* is downloading while + *variant* is being deleted. When one is, the delete reclaims only this + variant's files and leaves the shared companion (mmproj) for the live + sibling. Point-in-time (a sibling may claim just after it returns), but + safe: the finalized companion is held by deletion's reference-count + walk and a sibling starting mid-delete re-fetches it, so protection + never depends on the sibling having resolved its blob hashes.""" + repo_id = normalize_repo_key(repo_id) + target = (variant or "").strip().lower() or None + with self._lock: + for key in self._repo_active.get(repo_id, set()): + job = self._jobs.get(key) + if job is None or job.state not in _ACTIVE_STATES: + continue + if self._active_job_variant_locked(key) != target: + return True + return False + + def request_cancel( + self, + key: str, + proc: subprocess.Popen, + generation: Optional[int] = None, + ) -> bool: + """Authorize a SIGKILL for the registered *proc*. Idempotent across an + active job's lifetime: a repeated cancel while already ``cancelling`` + still returns ``True`` so a kill that raced and lost can be re-sent.""" + key = normalize_job_key(key) + with self._lock: + if self._processes.get(key) is not proc: + return False + if not self._generation_matches_locked(key, generation): + return False + if self._jobs.get(key, DownloadState("idle")).state not in _ACTIVE_STATES: + return False + self._jobs[key] = DownloadState("cancelling") + return True + + def terminate_all(self, kind: str = "download") -> None: + with self._lock: + live = [ + (key, proc, self._metadata.get(key)) + for key, proc in self._processes.items() + if proc.poll() is None + ] + # Flag as an intentional stop so the watcher's exit classification + # reports them cancelled rather than an OOM/crash once SIGKILL lands. + for key, _proc, _metadata in live: + if self._jobs.get(key, DownloadState("idle")).state == "running": + self._jobs[key] = DownloadState("cancelling") + reaped: list[tuple[str, subprocess.Popen, Optional[DownloadMetadata]]] = [] + for key, proc, metadata in live: + try: + proc.kill() + except ProcessLookupError: + pass + except Exception as e: + logger.warning(f"shutdown: failed to kill {kind} worker for {key}: {e}") + if metadata is not None: + persist_cancel_marker( + metadata.repo_type, + metadata.repo_id, + metadata.variant, + metadata.transport, + ) + continue + reaped.append((key, proc, metadata)) + deadline = time.monotonic() + 10.0 + for key, proc, metadata in reaped: + try: + proc.wait(timeout = max(0.0, deadline - time.monotonic())) + except subprocess.TimeoutExpired: + logger.warning(f"shutdown: {kind} worker for {key} did not exit after kill") + except Exception: + pass + # Mark only genuinely interrupted workers (rc != 0, or None on wait + # timeout); persisting before the exit is known would strand a stale + # marker on a worker that completed cleanly during shutdown. + if metadata is not None and proc.poll() != 0: + persist_cancel_marker( + metadata.repo_type, + metadata.repo_id, + metadata.variant, + metadata.transport, + ) + + +def _named_registry(name: str) -> DownloadRegistry: + with _NAMED_REGISTRIES_LOCK: + registry = _NAMED_REGISTRIES.get(name) + if registry is None: + registry = DownloadRegistry() + _NAMED_REGISTRIES[name] = registry + return registry + + +def get_models_registry() -> DownloadRegistry: + return _named_registry("models") + + +def get_datasets_registry() -> DownloadRegistry: + return _named_registry("datasets") diff --git a/studio/backend/hub/utils/gguf.py b/studio/backend/hub/utils/gguf.py new file mode 100644 index 0000000000..4bd1961c33 --- /dev/null +++ b/studio/backend/hub/utils/gguf.py @@ -0,0 +1,406 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""GGUF filename helpers. Quantization variants are derived from filenames, not parsed from binary GGUF headers.""" + +from __future__ import annotations + +import os +import re +from dataclasses import dataclass +from pathlib import Path +from typing import Optional + +from loggers import get_logger + +logger = get_logger(__name__) +_GGUF_MODEL_INFO_TIMEOUT_SECONDS = 5.0 + + +@dataclass +class GgufVariantInfo: + filename: str + quant: str + size_bytes: int + display_label: Optional[str] = None + download_size_bytes: int = 0 + + +GGUF_QUANT_PREFERENCE = [ + "UD-Q4_K_XL", + "UD-Q4_K_L", + "UD-Q5_K_XL", + "UD-Q3_K_XL", + "UD-Q6_K_XL", + "UD-Q6_K_S", + "UD-Q8_K_XL", + "UD-Q2_K_XL", + "UD-IQ4_NL", + "UD-IQ4_XS", + "UD-IQ3_S", + "UD-IQ3_XXS", + "UD-IQ2_M", + "UD-IQ2_XXS", + "UD-IQ1_M", + "UD-IQ1_S", + "Q4_K_M", + "Q4_K_S", + "Q5_K_M", + "Q5_K_S", + "Q6_K", + "Q8_0", + "Q3_K_M", + "Q3_K_L", + "Q3_K_S", + "Q2_K", + "Q2_K_L", + "IQ4_NL", + "IQ4_XS", + "IQ3_M", + "IQ3_XXS", + "IQ2_M", + "IQ1_M", + "F16", + "BF16", + "F32", +] + +_GGUF_SPLIT_SUFFIX_RE = re.compile(r"-\d{3,}-of-\d{3,}", re.IGNORECASE) +_GGUF_QUANT_RE = re.compile( + r"(UD-)?" + r"(MXFP[0-9]+(?:_[A-Z0-9]+)*" + r"|IQ[0-9]+_[A-Z]+(?:_[A-Z0-9]+)?" + r"|TQ[0-9]+_[0-9]+" + r"|Q[0-9]+_K_[A-Z]+" + r"|Q[0-9]+_[0-9]+" + r"|Q[0-9]+_K" + r"|BF16|F16|F32)", + re.IGNORECASE, +) + + +def is_mmproj_filename(filename: str) -> bool: + return "mmproj" in filename.lower() + + +def is_gguf_filename(filename: str) -> bool: + return filename.lower().endswith(".gguf") + + +# Cap recursive walks so a huge or system path cannot run unbounded. +_MAX_LOCAL_SCAN_ENTRIES = 100_000 + + +def iter_gguf_files(directory: Path, recursive: bool = False): + if not directory.is_dir(): + return + if recursive: + seen = 0 + # os.walk skips unreadable subdirs instead of raising (e.g. /proc). + for dirpath, dirnames, filenames in os.walk(directory, onerror = lambda _e: None): + for name in filenames: + if is_gguf_filename(name): + yield Path(dirpath) / name + seen += len(dirnames) + len(filenames) + if seen > _MAX_LOCAL_SCAN_ENTRIES: + return + return + try: + entries = list(directory.iterdir()) + except OSError: + return + for file in entries: + try: + if file.is_file() and is_gguf_filename(file.name): + yield file + except OSError: + continue + + +def pick_best_gguf(filenames: list[str]) -> Optional[str]: + gguf_files = [ + name for name in filenames if is_gguf_filename(name) and not is_mmproj_filename(name) + ] + if not gguf_files: + return None + by_quant: dict[str, str] = {} + for name in gguf_files: + by_quant.setdefault(extract_quant_label(name).upper(), name) + for quant in GGUF_QUANT_PREFERENCE: + filename = by_quant.get(quant.upper()) + if filename is not None: + return filename + return gguf_files[0] + + +def _gguf_stem(filename: str) -> str: + basename = filename.rsplit("/", 1)[-1] + return _GGUF_SPLIT_SUFFIX_RE.sub("", basename.rsplit(".", 1)[0]).strip() + + +_FLOAT_PRECISION_QUANTS = frozenset({"BF16", "F16", "F32"}) + + +def _select_quant_match(text: str) -> Optional[re.Match]: + fallback: Optional[re.Match] = None + for match in _GGUF_QUANT_RE.finditer(text): + if match.group(2).upper() in _FLOAT_PRECISION_QUANTS: + if fallback is None: + fallback = match + continue + return match + return fallback + + +def extract_quant_token(filename: str) -> Optional[str]: + stem = _gguf_stem(filename) + match = _select_quant_match(stem) + if not match and "/" in filename: + parents = filename.rsplit("/", 1)[0] + for segment in reversed(parents.split("/")): + parent_match = _select_quant_match(segment) + if parent_match: + match = parent_match + break + if match: + prefix = match.group(1) or "" + return f"{prefix}{match.group(2)}" + return None + + +def _unknown_gguf_variant_key(filename: str) -> str: + stem = _gguf_stem(filename) + if "/" not in filename: + return stem or "gguf" + parents = filename.rsplit("/", 1)[0].strip("/") + return f"{parents}/{stem}" if parents and stem else stem or "gguf" + + +def extract_quant_label(filename: str) -> str: + return extract_quant_token(filename) or _unknown_gguf_variant_key(filename) + + +def _apply_gguf_display_labels(variants: list[GgufVariantInfo]) -> None: + unknown_variants = [ + variant for variant in variants if extract_quant_token(variant.filename) is None + ] + if not unknown_variants: + return + ambiguous = len(unknown_variants) > 1 + for variant in unknown_variants: + variant.display_label = f"GGUF · {variant.filename}" if ambiguous else "GGUF" + + +def _env_offline() -> bool: + return os.environ.get("HF_HUB_OFFLINE", "").lower() in ( + "1", + "true", + "yes", + ) or os.environ.get("TRANSFORMERS_OFFLINE", "").lower() in ("1", "true", "yes") + + +def iter_hf_cache_snapshots(repo_id: str): + from hub.utils.hf_cache_state import iter_repo_cache_dirs + + snapshots: list[Path] = [] + for repo_dir in iter_repo_cache_dirs("model", repo_id): + snapshots_dir = repo_dir / "snapshots" + if not snapshots_dir.is_dir(): + continue + try: + snapshots.extend(snap for snap in snapshots_dir.iterdir() if snap.is_dir()) + except OSError: + continue + + def _mtime(path: Path) -> float: + try: + return path.stat().st_mtime + except OSError: + return 0.0 + + snapshots.sort(key = _mtime, reverse = True) + yield from snapshots + + +def list_gguf_variants_from_hf_cache(repo_id: str) -> Optional[tuple[list[GgufVariantInfo], bool]]: + for snapshot in iter_hf_cache_snapshots(repo_id): + variants, has_vision = list_local_gguf_variants(str(snapshot)) + if variants or has_vision: + return variants, has_vision + return None + + +def list_partial_gguf_variants_from_state( + repo_id: str, +) -> Optional[tuple[list[GgufVariantInfo], bool]]: + """Reconstruct GGUF variants from download manifests/markers alone. + + Used when no completed snapshot exists (download cancelled or interrupted) + and the HF API is unreachable (offline/gated/private). Each variant's + ``quant`` is the stored variant key so a resume passes the matching + ``--variant`` back to the worker. + """ + from hub.utils import download_manifest + + # Variant identity on disk is case-insensitive (_entry_key lowercases it), so + # dedupe on the lowercased key. Manifests are read first to keep their + # original-casing label over a lowercased cancel marker for the same variant. + seen: set[str] = set() + ordered: list[str] = [] + for source in ( + download_manifest.iter_variant_manifests("model", repo_id), + download_manifest.iter_variant_markers("model", repo_id), + ): + for variant, _path in source: + key = variant.lower() + if key not in seen: + seen.add(key) + ordered.append(variant) + if not ordered: + return None + + variants: list[GgufVariantInfo] = [] + has_vision = False + for variant in ordered: + manifest = download_manifest.read_manifest("model", repo_id, variant) + main_filename: Optional[str] = None + size_bytes = 0 + companion_bytes = 0 + if manifest is not None: + for expected in manifest.expected_files: + if not is_gguf_filename(expected.path): + continue + if is_mmproj_filename(expected.path): + has_vision = True + companion_bytes += max(0, int(expected.size or 0)) + continue + if main_filename is None: + main_filename = expected.path + size_bytes += max(0, int(expected.size or 0)) + if main_filename is None: + main_filename = f"{variant}.gguf" + variants.append( + GgufVariantInfo( + filename = main_filename, + quant = variant, + size_bytes = size_bytes, + download_size_bytes = size_bytes + companion_bytes, + ) + ) + + variants.sort(key = lambda variant: -variant.size_bytes) + _apply_gguf_display_labels(variants) + return variants, has_vision + + +def list_gguf_variants( + repo_id: str, hf_token: Optional[str] = None +) -> tuple[list[GgufVariantInfo], bool, Optional[list]]: + from huggingface_hub import HfApi + + if _env_offline(): + cached = list_gguf_variants_from_hf_cache(repo_id) + if cached is not None: + return (*cached, None) + + try: + info = HfApi(token = hf_token).model_info( + repo_id, + files_metadata = True, + timeout = _GGUF_MODEL_INFO_TIMEOUT_SECONDS, + ) + except Exception as exc: + if type(exc).__name__ in ( + "RepositoryNotFoundError", + "GatedRepoError", + "RevisionNotFoundError", + "EntryNotFoundError", + ): + raise + cached = list_gguf_variants_from_hf_cache(repo_id) + if cached is not None: + logger.warning( + "HF API unreachable for %s (%s); using local cache snapshot.", + repo_id, + exc.__class__.__name__, + ) + return (*cached, None) + raise + + variants: list[GgufVariantInfo] = [] + has_vision = False + quant_totals: dict[str, int] = {} + quant_first_file: dict[str, str] = {} + + for sibling in info.siblings: + filename = getattr(sibling, "rfilename", None) + if not isinstance(filename, str) or not is_gguf_filename(filename): + continue + if is_mmproj_filename(filename): + has_vision = True + continue + quant = extract_quant_label(filename) + quant_totals[quant] = quant_totals.get(quant, 0) + int(getattr(sibling, "size", 0) or 0) + quant_first_file.setdefault(quant, filename) + + for quant, total_size in quant_totals.items(): + variants.append( + GgufVariantInfo( + filename = quant_first_file[quant], + quant = quant, + size_bytes = total_size, + ) + ) + + variants.sort(key = lambda variant: -variant.size_bytes) + _apply_gguf_display_labels(variants) + return variants, has_vision, list(info.siblings) + + +def _resolve_gguf_dir(path: Path) -> Optional[Path]: + if path.is_dir(): + return path + if path.is_file() and path.suffix.lower() == ".gguf": + parent = path.parent + if ( + (parent / "config.json").exists() + or (parent / "adapter_config.json").exists() + or (parent / "export_metadata.json").exists() + ): + return parent + return None + + +def list_local_gguf_variants(directory: str) -> tuple[list[GgufVariantInfo], bool]: + root = _resolve_gguf_dir(Path(directory)) + if root is None: + return [], False + + quant_totals: dict[str, int] = {} + quant_first_file: dict[str, str] = {} + has_vision = False + + for file in sorted(iter_gguf_files(root, recursive = True)): + if is_mmproj_filename(file.name): + has_vision = True + continue + try: + size = file.stat().st_size + except OSError: + size = 0 + rel = file.relative_to(root).as_posix() + quant = extract_quant_label(rel) + quant_totals[quant] = quant_totals.get(quant, 0) + size + quant_first_file.setdefault(quant, rel) + + variants = [ + GgufVariantInfo( + filename = quant_first_file[quant], + quant = quant, + size_bytes = size, + ) + for quant, size in quant_totals.items() + ] + variants.sort(key = lambda variant: -variant.size_bytes) + _apply_gguf_display_labels(variants) + return variants, has_vision diff --git a/studio/backend/hub/utils/gguf_plan.py b/studio/backend/hub/utils/gguf_plan.py new file mode 100644 index 0000000000..03ec847f4c --- /dev/null +++ b/studio/backend/hub/utils/gguf_plan.py @@ -0,0 +1,158 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional, Sequence + +from hub.utils.download_manifest import ExpectedFile +from hub.utils.gguf import extract_quant_label, is_gguf_filename, is_mmproj_filename + + +@dataclass(frozen = True) +class GgufVariantPlan: + main_filenames: frozenset[str] + target_filenames: tuple[str, ...] + main_hashes: frozenset[str] + required_hashes: frozenset[str] + companion_hashes: frozenset[str] + mmproj_filenames: frozenset[str] + mmproj_hashes: frozenset[str] + expected_files: tuple[ExpectedFile, ...] + main_size_bytes: int + download_size_bytes: int + + +def sibling_sha256(sibling) -> Optional[str]: + lfs = getattr(sibling, "lfs", None) + if isinstance(lfs, dict): + value = lfs.get("sha256") + else: + value = getattr(lfs, "sha256", None) + return value if isinstance(value, str) and value else None + + +def sibling_size(sibling) -> int: + size = getattr(sibling, "size", 0) or 0 + try: + return int(size) + except (TypeError, ValueError): + return 0 + + +def expected_file_from_sibling(sibling) -> Optional[ExpectedFile]: + name = getattr(sibling, "rfilename", None) + if not isinstance(name, str): + return None + return ExpectedFile( + path = name, + size = sibling_size(sibling), + sha256 = sibling_sha256(sibling), + ) + + +def is_companion_gguf_path(path: str) -> bool: + return is_gguf_filename(path) and is_mmproj_filename(path) + + +def is_main_gguf_variant_path(path: str, variant: str) -> bool: + return ( + is_gguf_filename(path) + and not is_mmproj_filename(path) + and extract_quant_label(path).lower() == variant.lower() + ) + + +def mmproj_siblings(siblings: Sequence) -> list: + return [ + s + for s in siblings + if isinstance(getattr(s, "rfilename", None), str) + and is_companion_gguf_path(getattr(s, "rfilename")) + ] + + +def preferred_mmproj_sibling(siblings: Sequence) -> Optional[object]: + candidates = mmproj_siblings(siblings) + if not candidates: + return None + return next( + (s for s in candidates if extract_quant_label(getattr(s, "rfilename")).upper() == "F16"), + candidates[0], + ) + + +def build_gguf_variant_plans(siblings: Sequence) -> dict[str, GgufVariantPlan]: + main: dict[str, list] = {} + all_mmproj = mmproj_siblings(siblings) + all_mmproj_filenames = frozenset( + getattr(s, "rfilename") + for s in all_mmproj + if isinstance(getattr(s, "rfilename", None), str) + ) + all_mmproj_hashes = frozenset(h for h in (sibling_sha256(s) for s in all_mmproj) if h) + companion = preferred_mmproj_sibling(siblings) + companion_expected = expected_file_from_sibling(companion) if companion is not None else None + + for sibling in siblings: + name = getattr(sibling, "rfilename", None) + if not isinstance(name, str) or not is_gguf_filename(name): + continue + if is_mmproj_filename(name): + continue + quant = extract_quant_label(name).lower() + main.setdefault(quant, []).append(sibling) + + plans: dict[str, GgufVariantPlan] = {} + for quant, target_main_siblings in main.items(): + main_expected = tuple( + file + for sibling in target_main_siblings + if (file := expected_file_from_sibling(sibling)) is not None + ) + expected_files = ( + (*main_expected, companion_expected) + if companion_expected is not None + else main_expected + ) + plans[quant] = plan_from_expected_files( + quant, + expected_files, + all_mmproj_filenames = all_mmproj_filenames, + all_mmproj_hashes = all_mmproj_hashes, + ) + return plans + + +def plan_from_expected_files( + variant: str, + expected_files: Sequence[ExpectedFile], + *, + all_mmproj_filenames: frozenset[str] | None = None, + all_mmproj_hashes: frozenset[str] | None = None, +) -> GgufVariantPlan: + expected = tuple(expected_files) + main_files = tuple(file for file in expected if is_main_gguf_variant_path(file.path, variant)) + companion_files = tuple(file for file in expected if is_companion_gguf_path(file.path)) + main_hashes = frozenset(file.sha256 for file in main_files if file.sha256) + companion_hashes = frozenset(file.sha256 for file in companion_files if file.sha256) + required_hashes = frozenset(file.sha256 for file in expected if file.sha256) + main_size = sum(max(0, int(file.size or 0)) for file in main_files) + download_size = sum(max(0, int(file.size or 0)) for file in expected) + return GgufVariantPlan( + main_filenames = frozenset(file.path for file in main_files), + target_filenames = tuple(file.path for file in expected), + main_hashes = main_hashes, + required_hashes = required_hashes, + companion_hashes = companion_hashes, + mmproj_filenames = ( + all_mmproj_filenames + if all_mmproj_filenames is not None + else frozenset(file.path for file in companion_files) + ), + mmproj_hashes = (all_mmproj_hashes if all_mmproj_hashes is not None else companion_hashes), + expected_files = expected, + main_size_bytes = main_size, + download_size_bytes = download_size, + ) diff --git a/studio/backend/hub/utils/hf_cache_state.py b/studio/backend/hub/utils/hf_cache_state.py new file mode 100644 index 0000000000..a1ac372abb --- /dev/null +++ b/studio/backend/hub/utils/hf_cache_state.py @@ -0,0 +1,293 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from __future__ import annotations + +import errno +import shutil +import sys +from pathlib import Path +from typing import Iterable, Iterator, Optional + + +EXIT_CANCELLED = 130 + +TRANSPORT_HTTP = "http" +TRANSPORT_XET = "xet" +VALID_TRANSPORTS = frozenset({TRANSPORT_HTTP, TRANSPORT_XET}) +TRANSPORT_MARKER_NAME = ".transport" +INCOMPLETE_SUFFIX = ".incomplete" + + +def hf_cache_root(*, create: bool = False) -> Optional[Path]: + try: + from huggingface_hub import constants as hf_constants + except ImportError: + return None + root = Path(hf_constants.HF_HUB_CACHE) + if create: + try: + root.mkdir(parents = True, exist_ok = True) + except OSError: + return None + return root + return root if root.is_dir() else None + + +def hf_cache_roots() -> list[Path]: + from hub.utils.paths import hf_default_cache_dir, legacy_hf_cache_dir + + roots: list[Path] = [] + seen: set[str] = set() + + def _add(path: Optional[Path]) -> None: + if path is None or not path.is_dir(): + return + try: + key = str(path.resolve()) + except OSError: + return + if key in seen: + return + seen.add(key) + roots.append(path) + + _add(hf_cache_root()) + _add(legacy_hf_cache_dir()) + _add(hf_default_cache_dir()) + return roots + + +def target_dir_name(repo_type: str, repo_id: str) -> str: + return repo_cache_dir_name(repo_type, repo_id).lower() + + +def repo_cache_dir_name(repo_type: str, repo_id: str) -> str: + return f"{repo_type}s--{repo_id.replace('/', '--')}" + + +def resolve_destructive_case_matches(target: str, candidates: Iterable[str]) -> Optional[set[str]]: + values = list(candidates) + exact = {candidate for candidate in values if candidate == target} + if exact: + return exact + folded = {candidate for candidate in values if candidate.lower() == target.lower()} + if len(folded) <= 1: + return folded + return None + + +def _blob_dir_is_partial(blobs_dir: Path) -> bool: + try: + for blob in blobs_dir.iterdir(): + if blob.is_file() and blob.name.endswith(INCOMPLETE_SUFFIX): + return True + except OSError: + return False + return False + + +def blob_bytes_present(path: Path) -> int: + """Sparse-aware on-disk size: XET/``hf_transfer`` ``.incomplete`` partials + report a full ``st_size`` while only some blocks are allocated, so prefer + ``st_blocks``, falling back to ``st_size`` where it is unreported (Windows, + some network filesystems).""" + st = path.stat() + blocks = getattr(st, "st_blocks", 0) + if blocks > 0: + return min(blocks * 512, st.st_size) + if sys.platform == "win32": + allocated = _windows_allocated_size(path) + if allocated is not None: + return min(allocated, st.st_size) + return st.st_size + + +def _windows_allocated_size(path: Path) -> Optional[int]: + """Best-effort allocated-byte count for sparse files on Windows.""" + if sys.platform != "win32": + return None + try: + import ctypes + from ctypes import wintypes + + kernel32 = ctypes.WinDLL("kernel32", use_last_error = True) + get_compressed_file_size = kernel32.GetCompressedFileSizeW + get_compressed_file_size.argtypes = [ + wintypes.LPCWSTR, + ctypes.POINTER(wintypes.DWORD), + ] + get_compressed_file_size.restype = wintypes.DWORD + + high = wintypes.DWORD(0) + ctypes.set_last_error(0) + low = get_compressed_file_size(str(path), ctypes.byref(high)) + if low == 0xFFFFFFFF and ctypes.get_last_error() != 0: + return None + return (int(high.value) << 32) + int(low) + except Exception: + return None + + +def latest_snapshot_dir(repo_dir: Path) -> Optional[Path]: + """Newest immediate child of ``repo_dir/snapshots`` by mtime, or None. + + mtime is the signal huggingface_hub's from_pretrained resolves to, so this + points at whatever snapshot most recently landed on disk. + """ + snapshots_dir = repo_dir / "snapshots" + try: + if not snapshots_dir.is_dir(): + return None + snapshots = [entry for entry in snapshots_dir.iterdir() if entry.is_dir()] + if not snapshots: + return None + return max(snapshots, key = lambda entry: entry.stat().st_mtime) + except OSError: + return None + + +def _repo_dir_has_broken_snapshot_symlinks(repo_dir: Path) -> bool: + latest = latest_snapshot_dir(repo_dir) + if latest is None: + return False + try: + for entry in latest.rglob("*"): + if entry.is_symlink() and not entry.exists(): + return True + except OSError: + return False + return False + + +def iter_repo_cache_dirs(repo_type: str, repo_id: str) -> Iterator[Path]: + target = target_dir_name(repo_type, repo_id) + for root in hf_cache_roots(): + try: + for entry in root.iterdir(): + if entry.name.lower() == target: + yield entry + except OSError: + continue + + +def iter_destructive_repo_cache_dirs(repo_type: str, repo_id: str) -> Iterator[Path]: + target = repo_cache_dir_name(repo_type, repo_id) + folded_target = target.lower() + for root in hf_cache_roots(): + try: + entries = [entry for entry in root.iterdir() if entry.name.lower() == folded_target] + except OSError: + continue + matched_names = resolve_destructive_case_matches( + target, + (entry.name for entry in entries), + ) + if not matched_names: + continue + for entry in entries: + if entry.name in matched_names: + yield entry + + +def iter_active_repo_cache_dirs(repo_type: str, repo_id: str) -> Iterator[Path]: + root = hf_cache_root() + if root is None: + return + target = target_dir_name(repo_type, repo_id) + try: + for entry in root.iterdir(): + if entry.name.lower() == target: + yield entry + except OSError: + return + + +def preferred_repo_cache_dirs( + repo_type: str, + repo_id: str, + *, + force_active: bool = False, +) -> list[Path]: + active_entries = list(iter_active_repo_cache_dirs(repo_type, repo_id)) + if active_entries: + return active_entries + if force_active: + root = hf_cache_root() + if root is not None: + canonical = repo_cache_dir_name(repo_type, repo_id) + return [root / canonical] + return list(iter_repo_cache_dirs(repo_type, repo_id)) + + +def has_incomplete_blobs(repo_type: str, repo_id: str) -> bool: + for entry in iter_repo_cache_dirs(repo_type, repo_id): + if repo_cache_dir_has_incomplete_blobs(entry): + return True + return False + + +def has_active_incomplete_blobs(repo_type: str, repo_id: str) -> bool: + for entry in iter_active_repo_cache_dirs(repo_type, repo_id): + if repo_cache_dir_has_incomplete_blobs(entry): + return True + return False + + +def repo_cache_dir_has_incomplete_blobs(repo_dir: Path) -> bool: + blobs_dir = repo_dir / "blobs" + return (blobs_dir.is_dir() and _blob_dir_is_partial(blobs_dir)) or ( + _repo_dir_has_broken_snapshot_symlinks(repo_dir) + ) + + +def _prune_empty_dirs(root: Path) -> bool: + removed = False + try: + dirs = sorted( + (path for path in root.rglob("*") if path.is_dir()), + key = lambda path: len(path.parts), + reverse = True, + ) + except OSError: + dirs = [] + for directory in [*dirs, root]: + try: + directory.rmdir() + removed = True + except FileNotFoundError: + continue + except OSError as exc: + if exc.errno not in (errno.ENOTEMPTY, errno.EEXIST): + raise + return removed + + +def purge_partial_repo(repo_type: str, repo_id: str) -> bool: + removed = False + for entry in iter_destructive_repo_cache_dirs(repo_type, repo_id): + blobs_dir = entry / "blobs" + if blobs_dir.is_dir(): + for blob in blobs_dir.iterdir(): + if blob.is_file() and blob.name.endswith(INCOMPLETE_SUFFIX): + try: + blob.unlink() + removed = True + except FileNotFoundError: + continue + if _prune_empty_dirs(entry): + removed = True + return removed + + +def purge_repo_cache_dirs(repo_type: str, repo_id: str) -> bool: + removed = False + for entry in iter_destructive_repo_cache_dirs(repo_type, repo_id): + try: + if entry.is_symlink() or not entry.is_dir(): + continue + shutil.rmtree(entry) + removed = True + except FileNotFoundError: + continue + return removed diff --git a/studio/backend/hub/utils/hf_errors.py b/studio/backend/hub/utils/hf_errors.py new file mode 100644 index 0000000000..b2569758c6 --- /dev/null +++ b/studio/backend/hub/utils/hf_errors.py @@ -0,0 +1,28 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Map Hugging Face Hub client-side errors to HTTP status codes.""" + +from __future__ import annotations + +from typing import Optional + + +def hf_error_status(exc: Exception) -> Optional[int]: + # Client-side HF errors should surface as 4xx, not a generic 500. + name = type(exc).__name__ + if name in ( + "RepositoryNotFoundError", + "RevisionNotFoundError", + "EntryNotFoundError", + ): + return 404 + if name == "GatedRepoError": + return 403 + if name == "HFValidationError": + return 400 + # HfHubHTTPError subclasses carry the upstream response status. + code = getattr(getattr(exc, "response", None), "status_code", None) + if isinstance(code, int) and 400 <= code < 500: + return code + return None diff --git a/studio/backend/hub/utils/inventory_scan.py b/studio/backend/hub/utils/inventory_scan.py new file mode 100644 index 0000000000..a7281ad8b3 --- /dev/null +++ b/studio/backend/hub/utils/inventory_scan.py @@ -0,0 +1,533 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""HF cache inventory scanner. + +Read-only walks of the HuggingFace hub cache plus legacy/default +cache locations. Builds the foundation that Hub inventory endpoints +and the DownloadRegistry both consume. + +The worker spawn / transport-marker preparation / DownloadRegistry +layers built on top of these primitives live in download_registry.py. +""" + +from __future__ import annotations + +import hashlib +import re +import threading +import time +from dataclasses import dataclass +from pathlib import Path +from typing import Callable, Optional + +from loggers import get_logger + +logger = get_logger(__name__) + +from hub.utils.gguf import extract_quant_label, is_gguf_filename, is_mmproj_filename +from hub.utils.state_dir import RepoType + +from hub.utils.hf_cache_state import ( + INCOMPLETE_SUFFIX, + has_incomplete_blobs, + hf_cache_root, + iter_repo_cache_dirs, + latest_snapshot_dir, + repo_cache_dir_has_incomplete_blobs, +) + +# Inventory is invalidated explicitly on every app-driven cache mutation, so +# this TTL only bounds staleness from out-of-band edits while skipping re-walks +# on rapid UI navigation. +_HF_CACHE_SCANS_TTL_SECONDS = 15.0 +_GGUF_SPLIT_RE = re.compile(r"-(\d{3,})-of-(\d{3,})(?=\.gguf$)", re.IGNORECASE) +_hf_cache_scans_lock = threading.Lock() + + +@dataclass +class _HfCacheScanFlight: + event: threading.Event + epoch: int + result: Optional[list] = None + error: Optional[BaseException] = None + + +_hf_cache_scans_flight: Optional[_HfCacheScanFlight] = None +_hf_cache_scans_result: Optional[list] = None +_hf_cache_scans_cached_at: float = 0.0 +# Bumped on every invalidation. A scan tags itself with the epoch it began +# under; an invalidation mid-scan changes the epoch so the in-flight result is +# neither cached nor served to callers that arrived after the mutation. +_hf_cache_scans_epoch: int = 0 + + +def invalidate_hf_cache_scans() -> None: + global _hf_cache_scans_result, _hf_cache_scans_cached_at, _hf_cache_scans_epoch + with _hf_cache_scans_lock: + _hf_cache_scans_result = None + _hf_cache_scans_cached_at = 0.0 + _hf_cache_scans_epoch += 1 + + +def all_hf_cache_scans() -> list: + global _hf_cache_scans_flight, _hf_cache_scans_result, _hf_cache_scans_cached_at + + now = time.monotonic() + with _hf_cache_scans_lock: + if ( + _hf_cache_scans_result is not None + and (now - _hf_cache_scans_cached_at) < _HF_CACHE_SCANS_TTL_SECONDS + ): + return list(_hf_cache_scans_result) + start_epoch = _hf_cache_scans_epoch + flight = _hf_cache_scans_flight + # Only coalesce onto an in-flight scan from the current epoch; one that + # began before an intervening invalidation is superseded so + # post-mutation callers never receive pre-mutation data. + if flight is None or flight.epoch != start_epoch: + flight = _HfCacheScanFlight(event = threading.Event(), epoch = start_epoch) + _hf_cache_scans_flight = flight + owner = True + else: + owner = False + + if not owner: + flight.event.wait() + if flight.error is not None: + raise flight.error + return list(flight.result or []) + + try: + scans = _compute_all_hf_cache_scans() + with _hf_cache_scans_lock: + flight.result = scans + if _hf_cache_scans_epoch == flight.epoch: + _hf_cache_scans_result = scans + _hf_cache_scans_cached_at = time.monotonic() + return scans + except Exception as exc: + with _hf_cache_scans_lock: + if _hf_cache_scans_epoch == flight.epoch: + _hf_cache_scans_result = None + _hf_cache_scans_cached_at = 0.0 + flight.error = exc + raise + finally: + with _hf_cache_scans_lock: + if _hf_cache_scans_flight is flight: + _hf_cache_scans_flight = None + flight.event.set() + + +def _compute_all_hf_cache_scans() -> list: + from huggingface_hub import scan_cache_dir + from hub.utils.paths import legacy_hf_cache_dir, hf_default_cache_dir + + scans: list = [] + seen: set[str] = set() + try: + from huggingface_hub.constants import HF_HUB_CACHE + + active = Path(HF_HUB_CACHE).resolve() + seen.add(str(active)) + if active.is_dir(): + scans.append(scan_cache_dir()) + except Exception as exc: + logger.warning("Could not scan active HF cache: %s", exc) + + for extra_fn in (legacy_hf_cache_dir, hf_default_cache_dir): + extra = extra_fn() + if extra.is_dir() and str(extra.resolve()) not in seen: + seen.add(str(extra.resolve())) + try: + scans.append(scan_cache_dir(cache_dir = str(extra))) + except Exception as exc: + logger.warning("Could not scan HF cache %s: %s", extra, exc) + return scans + + +def token_fingerprint(hf_token: Optional[str]) -> str: + """16-char SHA256 prefix used as a cache-key qualifier for gated repos. + + Lets per-token size/snapshot caches refuse to serve a previously + fetched value back to a different token (a private/gated repo's + metadata is only valid for the credential that fetched it). + """ + if not hf_token: + return "" + return hashlib.sha256(hf_token.encode()).hexdigest()[:16] + + +def resolve_hf_cache_realpath(repo_dir: Path) -> Optional[str]: + """Pick the most useful on-disk path for a HF cache repo dir. + + Prefers the most-recent snapshot dir (what ``from_pretrained`` + actually points at). Falls back to the cache repo root. Returns the + resolved realpath so symlinks under ``snapshots/`` are followed back + to ``blobs/``. + """ + try: + latest = latest_snapshot_dir(repo_dir) + if latest is not None: + return str(latest.resolve()) + return str(repo_dir.resolve()) + except Exception: + return None + + +def resolve_snapshot_dir_for_scan( + repo_type: str, + repo_id: str, + repo_cache_dir: Optional[Path] = None, +) -> Optional[Path]: + """Latest snapshot dir for a cache row, or the first populated HF cache root. + + Scanner-side counterpart to snapshot_download()'s return value (which the + scanner cannot access). With a *repo_cache_dir*, returns its newest + snapshot. Otherwise scans roots in priority order (active, legacy, default) + and returns the newest snapshot in the first root that holds one; active is + where snapshot_download writes, so it is authoritative. Within a root, + picks by mtime (what from_pretrained resolves to) rather than refs/main, + since the user may have downloaded a non-main commit. + """ + if repo_cache_dir is not None: + latest = latest_snapshot_dir(repo_cache_dir) + if latest is None: + return None + try: + return latest.resolve() + except OSError: + return None + for repo_dir in iter_repo_cache_dirs(repo_type, repo_id): + latest = latest_snapshot_dir(repo_dir) + if latest is None: + continue + try: + return latest.resolve() + except OSError: + continue + return None + + +def _compose_partial(*signals: Callable[[], bool]) -> bool: + return any(signal() for signal in signals) + + +def _state_applies_to_repo_cache_dir(repo_cache_dir: Optional[Path]) -> bool: + if repo_cache_dir is None: + return True + root = hf_cache_root() + if root is None: + return False + try: + return repo_cache_dir.resolve().parent == root.resolve() + except OSError: + return False + + +def _legacy_partial( + repo_type: str, + repo_id: str, + repo_cache_dir: Optional[Path] = None, +) -> bool: + if repo_cache_dir is not None: + return repo_cache_dir_has_incomplete_blobs(repo_cache_dir) + return has_incomplete_blobs(repo_type, repo_id) + + +def _repo_cache_dir_incomplete_hashes(repo_cache_dir: Path) -> set[str]: + blobs_dir = repo_cache_dir / "blobs" + if not blobs_dir.is_dir(): + return set() + hashes: set[str] = set() + try: + entries = list(blobs_dir.iterdir()) + except OSError: + return hashes + for blob in entries: + try: + if blob.is_file() and blob.name.endswith(INCOMPLETE_SUFFIX): + hashes.add(blob.name[: -len(INCOMPLETE_SUFFIX)]) + except OSError: + continue + return hashes + + +def _repo_cache_dir_has_non_gguf_broken_snapshot_symlinks(repo_cache_dir: Path) -> bool: + latest = latest_snapshot_dir(repo_cache_dir) + if latest is None: + return False + try: + entries = list(latest.rglob("*")) + except OSError: + return False + for entry in entries: + try: + if not entry.is_symlink() or entry.exists(): + continue + rel = entry.relative_to(latest).as_posix() + if is_gguf_filename(rel): + continue + return True + except OSError: + continue + return False + + +def _gguf_variant_manifest_blob_hashes(repo_id: str) -> frozenset[str]: + from hub.utils import download_manifest + + hashes: set[str] = set() + for variant, _path in download_manifest.iter_variant_manifests("model", repo_id): + manifest = download_manifest.read_manifest("model", repo_id, variant) + if manifest is None: + continue + for expected in manifest.expected_files: + if expected.sha256 and is_gguf_filename(expected.path): + hashes.add(expected.sha256) + return frozenset(hashes) + + +def _repo_cache_dir_has_snapshot_legacy_partial( + repo_cache_dir: Path, *, ignored_blob_hashes: frozenset[str] +) -> bool: + incomplete_hashes = _repo_cache_dir_incomplete_hashes(repo_cache_dir) + if any(blob_hash not in ignored_blob_hashes for blob_hash in incomplete_hashes): + return True + return _repo_cache_dir_has_non_gguf_broken_snapshot_symlinks(repo_cache_dir) + + +def _snapshot_legacy_partial( + repo_type: str, + repo_id: str, + repo_cache_dir: Optional[Path] = None, +) -> bool: + if repo_type != "model": + return _legacy_partial(repo_type, repo_id, repo_cache_dir) + ignored_hashes = _gguf_variant_manifest_blob_hashes(repo_id) + if repo_cache_dir is not None: + return _repo_cache_dir_has_snapshot_legacy_partial( + repo_cache_dir, + ignored_blob_hashes = ignored_hashes, + ) + return any( + _repo_cache_dir_has_snapshot_legacy_partial( + entry, + ignored_blob_hashes = ignored_hashes, + ) + for entry in iter_repo_cache_dirs(repo_type, repo_id) + ) + + +def _completed_gguf_variants(snapshot_dir: Optional[Path]) -> set[str]: + if snapshot_dir is None: + return set() + complete: set[str] = set() + split_groups: dict[str, dict[int, set[int]]] = {} + try: + paths = list(snapshot_dir.rglob("*")) + except OSError: + return set() + for path in paths: + try: + if not path.is_file() or path.stat().st_size <= 0: + continue + except OSError: + continue + rel = path.relative_to(snapshot_dir).as_posix() + if not is_gguf_filename(rel) or is_mmproj_filename(rel): + continue + quant = extract_quant_label(rel) + split = _GGUF_SPLIT_RE.search(path.name) + if split is None: + complete.add(quant) + continue + index = int(split.group(1)) + total = int(split.group(2)) + if index <= 0 or total <= 0 or index > total: + continue + split_groups.setdefault(quant, {}).setdefault(total, set()).add(index) + for quant, groups in split_groups.items(): + for total, indices in groups.items(): + if indices == set(range(1, total + 1)): + complete.add(quant) + break + return complete + + +def _manifest_partial( + repo_type: RepoType, + repo_id: str, + variant: Optional[str] = None, + snapshot_dir: Optional[Path] = None, + repo_cache_dir: Optional[Path] = None, +) -> bool: + from hub.utils import download_manifest + + if not _state_applies_to_repo_cache_dir(repo_cache_dir): + return False + manifest = download_manifest.read_manifest(repo_type, repo_id, variant) + if manifest is None: + return False + resolved = ( + snapshot_dir + if snapshot_dir is not None + else resolve_snapshot_dir_for_scan(repo_type, repo_id, repo_cache_dir) + ) + if resolved is None: + return True + return not download_manifest.verify_against_disk(manifest, resolved).ok + + +def is_snapshot_partial( + repo_type: RepoType, + repo_id: str, + repo_cache_dir: Optional[Path] = None, +) -> bool: + """Repo-row partial flag for snapshot-style downloads (full-snapshot + models — safetensors/adapter/checkpoint — and all datasets). + + Composes three signals, cheapest first: + 1. Cancel marker (single stat). + 2. Snapshot-attributed legacy .incomplete blob / broken-symlink check. + 3. Manifest walk (stat per expected file under the latest snapshot). + + A manifest without a resolvable snapshot is partial: the worker got + far enough to record expectations but did not leave a usable snapshot.""" + from hub.utils import download_manifest + + state_applies = _state_applies_to_repo_cache_dir(repo_cache_dir) + return _compose_partial( + lambda: state_applies and download_manifest.has_cancel_marker(repo_type, repo_id, None), + lambda: _snapshot_legacy_partial(repo_type, repo_id, repo_cache_dir), + lambda: _manifest_partial( + repo_type, + repo_id, + None, + None, + repo_cache_dir, + ), + ) + + +def is_variant_partial( + repo_id: str, + variant: str, + snapshot_dir: Optional[Path] = None, + *, + incomplete_blob_hashes: Optional[set[str]] = None, + variant_blob_hashes: Optional[frozenset[str]] = None, + repo_cache_dir: Optional[Path] = None, +) -> bool: + """Per-variant partial detection. Owns its manifest, owns its marker. + Used by the GGUF variants endpoint to flag a specific quant as broken + without contaminating other quants in the same repo. + + snapshot_dir is an optional hint to avoid re-walking the cache when a + caller is checking many variants of the same repo (see + is_gguf_repo_partial for that usage).""" + from hub.utils import download_manifest + + state_applies = _state_applies_to_repo_cache_dir(repo_cache_dir) + return _compose_partial( + lambda: state_applies and download_manifest.has_cancel_marker("model", repo_id, variant), + lambda: bool( + incomplete_blob_hashes + and variant_blob_hashes + and incomplete_blob_hashes.intersection(variant_blob_hashes) + ), + lambda: _manifest_partial( + "model", + repo_id, + variant, + snapshot_dir, + repo_cache_dir, + ), + ) + + +def is_gguf_repo_partial(repo_id: str, repo_cache_dir: Optional[Path] = None) -> bool: + """Repo-row partial flag for a GGUF repo. The inventory shows ONE row per + GGUF repo (requires_variant=True); per-variant detail lives in + GET /api/models/gguf-variants and uses is_variant_partial. + + *** DO NOT simplify this to "any variant partial -> repo partial" *** + + Tripwire scenario: user downloads Q8_0 fully, then starts Q4_K_M and + cancels. Both variants share ONE inventory row. If row.partial flips True, + _capabilities_for_format flips can_chat=False, so the user can no longer + chat with the perfectly-good Q8_0 because of an unrelated cancelled Q4_K_M. + + Correct semantics: partial=True only when at least one variant is broken + AND no other variant is clean. "Simplifying" to the obvious "any broken" + form re-introduces this Q8+Q4 mixed-state regression. + + Composes signals: + 1. Cheap legacy fast-path (.incomplete blobs / broken symlinks). + 2. Per-variant manifest + marker enumeration, gated on "all broken". + """ + from hub.utils import download_manifest + + has_legacy_partial = _legacy_partial("model", repo_id, repo_cache_dir) + state_applies = _state_applies_to_repo_cache_dir(repo_cache_dir) + snapshot_dir = resolve_snapshot_dir_for_scan( + "model", + repo_id, + repo_cache_dir, + ) + variants: set[str] = set(_completed_gguf_variants(snapshot_dir)) + if state_applies: + for variant, _path in download_manifest.iter_variant_manifests( + "model", + repo_id, + ): + variants.add(variant) + for variant, _path in download_manifest.iter_variant_markers( + "model", + repo_id, + ): + variants.add(variant) + if not variants: + return has_legacy_partial + has_clean = False + has_broken = has_legacy_partial + for variant in variants: + if is_variant_partial( + repo_id, + variant, + snapshot_dir, + repo_cache_dir = repo_cache_dir, + ): + has_broken = True + else: + has_clean = True + return has_broken and not has_clean + + +def partial_transport_for( + repo_type: RepoType, + repo_id: str, + variant: Optional[str] = None, + repo_cache_dir: Optional[Path] = None, +) -> Optional[str]: + """Transport to surface on a partial row's resume affordance. + + Prefers the cancel marker's transport, then the manifest's. The fallback + matters for rows partial without a marker (an errored/interrupted download + leaves the manifest but no marker) so the UI can still show HTTP-resume vs + XET-redownload instead of the neutral retry label. ``None`` when neither is + available.""" + from hub.utils import download_manifest + + if not _state_applies_to_repo_cache_dir(repo_cache_dir): + return None + marker_transport = download_manifest.read_cancel_marker_transport( + repo_type, + repo_id, + variant, + ) + if marker_transport is not None: + return marker_transport + manifest = download_manifest.read_manifest(repo_type, repo_id, variant) + return manifest.transport if manifest is not None else None diff --git a/studio/backend/hub/utils/llm_assist.py b/studio/backend/hub/utils/llm_assist.py new file mode 100644 index 0000000000..00edb204c5 --- /dev/null +++ b/studio/backend/hub/utils/llm_assist.py @@ -0,0 +1,440 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from __future__ import annotations + +import json +import os +import re +import textwrap +import time +from typing import Any, Optional + +from loggers import get_logger + +from hub.utils import download_registry + +logger = get_logger(__name__) + +DEFAULT_HELPER_MODEL_REPO = "unsloth/gemma-4-E2B-it-GGUF" +DEFAULT_HELPER_MODEL_VARIANT = "UD-Q4_K_XL" +README_MAX_CHARS = 1500 + + +def _helper_disabled() -> bool: + return os.environ.get("UNSLOTH_HELPER_MODEL_DISABLE", "").strip().lower() in { + "1", + "true", + } + + +def _strip_think_tags(text: str) -> str: + if "" not in text: + return text + stripped = re.sub(r".*?\s*", "", text, flags = re.DOTALL).strip() + if stripped: + return stripped + matches = re.findall(r"(.*?)", text, flags = re.DOTALL) + return matches[-1].strip() if matches else text + + +def _parse_json_response(text: str) -> Optional[dict[str, Any]]: + cleaned = (text or "").strip() + if not cleaned: + return None + if cleaned.startswith("```"): + lines = cleaned.splitlines() + end = -1 if lines and lines[-1].strip().startswith("```") else len(lines) + cleaned = "\n".join(lines[1:end]).strip() + try: + parsed = json.loads(cleaned) + return parsed if isinstance(parsed, dict) else None + except json.JSONDecodeError: + pass + match = re.search(r"\{.*\}", cleaned, re.DOTALL) + if not match: + return None + try: + parsed = json.loads(match.group()) + except json.JSONDecodeError: + return None + return parsed if isinstance(parsed, dict) else None + + +def _generate_with_backend(backend, messages: list[dict[str, str]], max_tokens: int) -> str: + cumulative = "" + for chunk in backend.generate_chat_completion( + messages = messages, + temperature = 0.1, + top_p = 0.9, + top_k = 20, + max_tokens = max_tokens, + repetition_penalty = 1.0, + enable_thinking = False, + ): + if isinstance(chunk, dict): + continue + cumulative = chunk + return _strip_think_tags(cumulative.strip()) + + +def _fetch_hf_dataset_card( + dataset_name: str, hf_token: Optional[str] +) -> tuple[Optional[str], Optional[dict[str, Any]]]: + try: + from huggingface_hub import DatasetCard + + card = DatasetCard.load(dataset_name, token = hf_token) + readme = card.text or "" + if len(readme) > README_MAX_CHARS: + cut = readme[:README_MAX_CHARS].rfind(".") + if cut > README_MAX_CHARS // 2: + readme = readme[: cut + 1] + "\n[...truncated]" + else: + readme = readme[:README_MAX_CHARS] + "\n[...truncated]" + metadata: dict[str, Any] = {} + if card.data: + for key in ( + "task_categories", + "task_ids", + "language", + "size_categories", + "tags", + "license", + "pretty_name", + ): + value = getattr(card.data, key, None) + if value is not None: + metadata[key] = value + return readme, metadata + except Exception as exc: + logger.warning( + "Could not fetch dataset card for %s: %s", + dataset_name, + download_registry.scrub_secrets(str(exc), hf_token = hf_token), + ) + return None, None + + +def _is_gemma_3n(model_name: Optional[str]) -> bool: + normalized = (model_name or "").lower().replace("_", "-") + return "gemma-3n" in normalized or "gemma3n" in normalized + + +def _sample_text(columns: list[str], samples: list[dict[str, Any]]) -> str: + rows: list[str] = [] + for index, row in enumerate(samples[:5], 1): + parts = [f" {col}: {str(row.get(col, ''))[:200]}" for col in columns] + rows.append(f"Row {index}:\n" + "\n".join(parts)) + return "\n".join(rows) + + +def _target_hints(model_name: Optional[str], model_type: Optional[str]) -> str: + if model_type == "audio" and not _is_gemma_3n(model_name): + return ( + "\n\nHINT: The user is training an AUDIO model. The dataset must contain " + "a column with audio files or paths and one such column should be selected " + "as part of the input." + ) + if model_type == "embeddings": + return ( + "\n\nHINT: The user is training an EMBEDDING model. Prefer dataset formats " + "such as text pairs for STS, premise/hypothesis/label for NLI, or query " + "and document columns for retrieval." + ) + return "" + + +def _run_multi_pass_advisor( + *, + columns: list[str], + samples: list[dict[str, Any]], + dataset_name: Optional[str], + dataset_card: Optional[str], + dataset_metadata: Optional[dict[str, Any]], + model_name: Optional[str], + model_type: Optional[str], +) -> Optional[dict[str, Any]]: + if _helper_disabled(): + return None + + repo = os.environ.get("UNSLOTH_HELPER_MODEL_REPO", DEFAULT_HELPER_MODEL_REPO) + variant = os.environ.get("UNSLOTH_HELPER_MODEL_VARIANT", DEFAULT_HELPER_MODEL_VARIANT) + backend = None + try: + from core.inference.llama_cpp import LlamaCppBackend + + backend = LlamaCppBackend() + started = time.monotonic() + if not backend.load_model( + hf_repo = repo, + hf_variant = variant, + model_identifier = f"hub-advisor:{repo}:{variant}", + is_vision = False, + n_ctx = 2048, + n_gpu_layers = -1, + ): + return None + logger.info("Hub advisor model loaded in %.1fs", time.monotonic() - started) + + samples_text = _sample_text(columns, samples) + metadata_text = ( + json.dumps(dataset_metadata, indent = 2, default = str)[:500] if dataset_metadata else "N/A" + ) + card_excerpt = (dataset_card or "")[:1200] or "N/A" + hints = _target_hints(model_name, model_type) + + pass1_raw = _generate_with_backend( + backend, + [ + { + "role": "system", + "content": ( + "You are a dataset analyst. Classify the dataset and respond " + "with only a valid JSON object." + f"{hints}" + ), + }, + { + "role": "user", + "content": textwrap.dedent(f"""\ + Dataset: {dataset_name or "unknown"} + + DATASET CARD: + {card_excerpt} + + METADATA: + {metadata_text} + + COLUMNS: {columns} + + SAMPLE DATA: + {samples_text} + + Return this JSON shape: + {{ + "dataset_type": "", + "is_conversational": , + "needs_conversion": , + "description": "", + "task_description": "" + }}"""), + }, + ], + 256, + ) + pass1 = _parse_json_response(pass1_raw) + if not pass1: + return None + + if pass1.get("is_conversational") and not pass1.get("needs_conversion"): + return { + "success": True, + "dataset_type": pass1.get("dataset_type"), + "is_conversational": True, + "user_notification": ( + "This dataset is already in conversational format. No conversion is needed." + ), + } + + pass2_raw = _generate_with_backend( + backend, + [ + { + "role": "system", + "content": ( + "Assign each dataset column to user, assistant, or skip for " + "LLM fine-tuning. The target/output/answer/label column must be " + "assistant. Return only valid JSON." + f"{hints}" + ), + }, + { + "role": "user", + "content": textwrap.dedent(f"""\ + CLASSIFICATION: + {json.dumps(pass1, indent = 2)} + + COLUMNS: {columns} + + SAMPLE DATA: + {samples_text} + + Return this JSON shape: + {{ + "column_roles": {{"": ""}}, + "label_mapping": null, + "notes": "" + }}"""), + }, + ], + 512, + ) + pass2 = _parse_json_response(pass2_raw) + if not pass2: + return None + column_roles = pass2.get("column_roles") + if not isinstance(column_roles, dict): + return None + roles_present = set(column_roles.values()) + if "user" not in roles_present or "assistant" not in roles_present: + return None + + label_mapping = pass2.get("label_mapping") or None + system_prompt = "" + if not pass1.get("is_conversational"): + user_cols = [col for col, role in column_roles.items() if role == "user"] + assistant_cols = [col for col, role in column_roles.items() if role == "assistant"] + prompt_raw = _generate_with_backend( + backend, + [ + { + "role": "user", + "content": textwrap.dedent(f"""\ + Write a concise system prompt for fine-tuning. + + Dataset type: {pass1.get("dataset_type", "other")} + Task: {pass1.get("task_description") or pass1.get("description") or ""} + User input columns: {user_cols} + Assistant output columns: {assistant_cols} + + Write only the system prompt text."""), + }, + ], + 256, + ) + cleaned = prompt_raw.strip().strip('"').strip("'").strip() + if 20 <= len(cleaned) <= 800 and cleaned.lower() not in {"null", "none"}: + system_prompt = cleaned + + suggested_mapping = { + col: role + for col, role in column_roles.items() + if col in columns and role in {"user", "assistant", "system"} + } + if ( + "user" not in suggested_mapping.values() + or "assistant" not in suggested_mapping.values() + ): + return None + + dtype = str(pass1.get("dataset_type") or "other") + notification_parts = [f"This is a {dtype} dataset."] + description = pass1.get("task_description") or pass1.get("description") + if description: + notification_parts.append(str(description)) + notification_parts.append("Columns were mapped to conversation roles.") + + return { + "success": True, + "suggested_mapping": suggested_mapping, + "system_prompt": system_prompt, + "label_mapping": label_mapping if isinstance(label_mapping, dict) else None, + "dataset_type": dtype, + "is_conversational": bool(pass1.get("is_conversational")), + "user_notification": " ".join(notification_parts), + } + except Exception as exc: + logger.warning("Hub advisor failed: %s", exc) + return None + finally: + if backend is not None: + try: + backend.unload_model() + except Exception: + pass + + +def _heuristic_mapping(columns: list[str]) -> Optional[dict[str, str]]: + if not columns: + return None + lowered = {col: col.lower().replace("-", "_") for col in columns} + metadata_terms = ("id", "uuid", "url", "source", "date", "time", "score", "index") + assistant_terms = ( + "assistant", + "answer", + "response", + "output", + "completion", + "target", + "label", + "summary", + "translation", + ) + user_terms = ( + "user", + "human", + "prompt", + "instruction", + "input", + "question", + "query", + "context", + "document", + "article", + "problem", + "text", + ) + mapping: dict[str, str] = {} + for col, name in lowered.items(): + if any(term == name or name.endswith(f"_{term}") for term in metadata_terms): + continue + if any(term in name for term in assistant_terms): + mapping[col] = "assistant" + elif any(term in name for term in user_terms): + mapping[col] = "user" + + if "assistant" not in mapping.values(): + candidates = [col for col in columns if col not in mapping] + if candidates: + mapping[candidates[-1]] = "assistant" + elif columns: + mapping[columns[-1]] = "assistant" + if "user" not in mapping.values(): + for col in columns: + if mapping.get(col) != "assistant": + mapping[col] = "user" + break + if "user" not in mapping.values() or "assistant" not in mapping.values(): + return None + return mapping + + +def llm_conversion_advisor( + column_names: list[str], + samples: list[dict[str, Any]], + dataset_name: Optional[str] = None, + hf_token: Optional[str] = None, + model_name: Optional[str] = None, + model_type: Optional[str] = None, +) -> Optional[dict[str, Any]]: + dataset_card = None + dataset_metadata = None + if dataset_name and "/" in dataset_name: + dataset_card, dataset_metadata = _fetch_hf_dataset_card(dataset_name, hf_token) + + result = _run_multi_pass_advisor( + columns = column_names, + samples = samples, + dataset_name = dataset_name, + dataset_card = dataset_card, + dataset_metadata = dataset_metadata, + model_name = model_name, + model_type = model_type, + ) + if result and result.get("success"): + return result + + mapping = _heuristic_mapping(column_names) + if mapping: + return { + "success": True, + "suggested_mapping": mapping, + "dataset_type": None, + "is_conversational": None, + "warning": ( + "The helper model was unavailable, so Hub used column-name heuristics. " + "Review the suggested mapping before training." + ), + } + return None diff --git a/studio/backend/hub/utils/paths.py b/studio/backend/hub/utils/paths.py new file mode 100644 index 0000000000..afcb0b41dc --- /dev/null +++ b/studio/backend/hub/utils/paths.py @@ -0,0 +1,522 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Path validators and storage roots for the Hub layer.""" + +from __future__ import annotations + +import json +import os +import re +import sys +import tempfile +import threading +from collections import OrderedDict +from pathlib import Path +from typing import Optional + +from loggers import get_logger + +logger = get_logger(__name__) + + +def _infer_studio_home_from_venv() -> Optional[Path]: + try: + prefix = Path(sys.prefix).resolve() + except (OSError, ValueError): + return None + if prefix.name != "unsloth_studio": + return None + candidate = prefix.parent + shim_name = "unsloth.exe" if os.name == "nt" else "unsloth" + try: + if (candidate / "share" / "studio.conf").is_file() or ( + candidate / "bin" / shim_name + ).is_file(): + return candidate + except OSError: + return None + return None + + +def studio_root() -> Path: + override = (os.environ.get("UNSLOTH_STUDIO_HOME") or "").strip() + if not override: + override = (os.environ.get("STUDIO_HOME") or "").strip() + if override: + try: + return Path(override).expanduser().resolve() + except (OSError, ValueError): + return Path(override).expanduser() + inferred = _infer_studio_home_from_venv() + if inferred is not None: + return inferred + return Path.home() / ".unsloth" / "studio" + + +def cache_root() -> Path: + return studio_root() / "cache" + + +def assets_root() -> Path: + return studio_root() / "assets" + + +def datasets_root() -> Path: + return assets_root() / "datasets" + + +def dataset_uploads_root() -> Path: + return datasets_root() / "uploads" + + +def recipe_datasets_root() -> Path: + return datasets_root() / "recipes" + + +def outputs_root() -> Path: + return studio_root() / "outputs" + + +def exports_root() -> Path: + return studio_root() / "exports" + + +def tmp_root() -> Path: + return Path(tempfile.gettempdir()) / "unsloth-studio" + + +def ensure_dir(path: Path) -> Path: + path.mkdir(parents = True, exist_ok = True) + return path + + +def legacy_hf_cache_dir() -> Path: + return cache_root() / "huggingface" / "hub" + + +def hf_default_cache_dir() -> Path: + return Path.home() / ".cache" / "huggingface" / "hub" + + +def _is_wsl() -> bool: + if sys.platform == "win32": + return False + try: + return "microsoft" in Path("/proc/version").read_text().lower() + except Exception: + return False + + +_IS_WSL = _is_wsl() + + +def _wsl_automount_root() -> str: + """DrvFs root under which WSL maps Windows drives, with a trailing slash. + + Defaults to ``/mnt/`` but is user-configurable via ``/etc/wsl.conf`` + (``[automount] root``), so hard-coding ``/mnt/`` mistranslates Windows paths + on a host with a custom root (e.g. ``root = /`` → ``C:`` at ``/c/``).""" + default = "/mnt/" + if not _IS_WSL: + return default + try: + import configparser + + parser = configparser.ConfigParser(inline_comment_prefixes = ("#", ";")) + parser.read("/etc/wsl.conf") + root = parser.get("automount", "root", fallback = "").strip().strip("\"'") + except Exception: + return default + if not root: + return default + return root if root.endswith("/") else f"{root}/" + + +_WSL_AUTOMOUNT_ROOT = _wsl_automount_root() + + +def normalize_path(path: str) -> str: + if not path: + return path + if len(path) >= 3 and path[1] == ":" and path[2] in ("\\", "/"): + if _IS_WSL: + drive = path[0].lower() + rest = path[3:].replace("\\", "/") + return f"{_WSL_AUTOMOUNT_ROOT}{drive}/{rest}" + return path.replace("\\", "/") + return path.replace("\\", "/") + + +def _host_path(path: str | Path) -> Path: + return Path(normalize_path(str(path))).expanduser() + + +def is_local_path(path: str) -> bool: + if not path: + return False + normalized = normalize_path(path) + has_local_syntax = ( + path.startswith(("/", ".", "~")) + or ":" in path + or "\\" in path + or os.path.isabs(path) + or os.path.isabs(normalized) + ) + if path.count("/") == 1 and not has_local_syntax: + return False + try: + if has_local_syntax and Path(normalized).expanduser().exists(): + return True + except Exception: + pass + return has_local_syntax + + +_VALID_REPO_ID_SEGMENT = re.compile(r"^[A-Za-z0-9](?:[A-Za-z0-9._-]*[A-Za-z0-9])?$") +_MAX_REPO_ID_LENGTH = 96 + + +def is_valid_repo_id(repo_id: str) -> bool: + """Validate Hugging Face ``repo_name`` or ``namespace/repo_name`` IDs.""" + if not repo_id or repo_id != repo_id.strip(): + return False + if len(repo_id) > _MAX_REPO_ID_LENGTH or repo_id.endswith(".git"): + return False + if "--" in repo_id or ".." in repo_id: + return False + segments = repo_id.split("/") + if len(segments) not in (1, 2): + return False + return all( + segment not in ("", ".", "..") and _VALID_REPO_ID_SEGMENT.fullmatch(segment) is not None + for segment in segments + ) + + +_GGUF_VARIANT_CONTROL_CHARS = re.compile(r"[\x00-\x1f\x7f]") +_MAX_GGUF_VARIANT_LENGTH = 512 + + +def is_valid_gguf_variant(variant: str) -> bool: + """Validate Hub GGUF variant keys. + + Known quant labels are short tokens (``Q4_K_M``), but unknown GGUF layouts + use a snapshot-relative key derived from the filename and may contain + slashes or spaces. + """ + if not variant or variant != variant.strip(): + return False + if len(variant) > _MAX_GGUF_VARIANT_LENGTH: + return False + if _GGUF_VARIANT_CONTROL_CHARS.search(variant) or not variant.isprintable(): + return False + normalized = variant.replace("\\", "/") + return all(segment not in ("", ".", "..") for segment in normalized.split("/")) + + +def ollama_model_dirs() -> list[Path]: + """Return Ollama model directories that exist on disk.""" + dirs: list[Path] = [] + seen: set[str] = set() + + def _add(p: Path | str) -> None: + try: + expanded = _host_path(p) + resolved = expanded.resolve() + is_dir = expanded.is_dir() + except (OSError, RuntimeError, ValueError): + return + key = str(resolved) + if key in seen or not is_dir: + return + seen.add(key) + dirs.append(expanded) + + ollama_env = os.environ.get("OLLAMA_MODELS") + if ollama_env: + _add(ollama_env) + _add(Path.home() / ".ollama" / "models") + _add(Path("/usr/share/ollama/.ollama/models")) + _add(Path("/var/lib/ollama/.ollama/models")) + return dirs + + +# Per-process memo for resolve_cached_repo_id_case. Bounded LRU so a long-lived +# process touching many repo ids can't grow it without limit; evicted cold +# entries simply recompute on next use. +_CACHE_CASE_RESOLUTION_MEMO_MAX = 512 +_CACHE_CASE_RESOLUTION_MEMO: "OrderedDict[tuple[str, str], str]" = OrderedDict() +_CACHE_CASE_RESOLUTION_LOCK = threading.Lock() + + +def _memo_get(memo_key: tuple[str, str]) -> Optional[str]: + with _CACHE_CASE_RESOLUTION_LOCK: + value = _CACHE_CASE_RESOLUTION_MEMO.get(memo_key) + if value is not None: + _CACHE_CASE_RESOLUTION_MEMO.move_to_end(memo_key) + return value + + +def _memo_set(memo_key: tuple[str, str], value: str) -> None: + with _CACHE_CASE_RESOLUTION_LOCK: + _CACHE_CASE_RESOLUTION_MEMO[memo_key] = value + _CACHE_CASE_RESOLUTION_MEMO.move_to_end(memo_key) + while len(_CACHE_CASE_RESOLUTION_MEMO) > _CACHE_CASE_RESOLUTION_MEMO_MAX: + _CACHE_CASE_RESOLUTION_MEMO.popitem(last = False) + + +def _memo_drop(memo_key: tuple[str, str]) -> None: + with _CACHE_CASE_RESOLUTION_LOCK: + _CACHE_CASE_RESOLUTION_MEMO.pop(memo_key, None) + + +def _hf_hub_cache_dir() -> Path: + try: + from huggingface_hub.constants import HF_HUB_CACHE + return Path(HF_HUB_CACHE) + except Exception as exc: + logger.debug("Could not read huggingface_hub HF_HUB_CACHE, using default: %s", exc) + return Path.home() / ".cache" / "huggingface" / "hub" + + +def _hf_hub_cache_dirs() -> list[Path]: + roots: list[Path] = [] + seen: set[str] = set() + + def _add(path: Path) -> None: + try: + resolved = path.resolve() + except OSError: + return + key = str(resolved) + if key in seen or not resolved.is_dir(): + return + seen.add(key) + roots.append(resolved) + + _add(_hf_hub_cache_dir()) + try: + _add(legacy_hf_cache_dir()) + _add(hf_default_cache_dir()) + except Exception as exc: + logger.debug("Could not enumerate secondary HF cache roots: %s", exc) + return roots + + +def lmstudio_model_dirs() -> list[Path]: + dirs: list[Path] = [] + seen: set[str] = set() + + def _add(path: Path | str) -> None: + try: + expanded = _host_path(path) + resolved = expanded.resolve() + except (OSError, RuntimeError, ValueError): + return + key = str(resolved) + if key in seen or not expanded.is_dir(): + return + seen.add(key) + dirs.append(expanded) + + settings_path = Path.home() / ".lmstudio" / "settings.json" + if settings_path.is_file(): + try: + settings = json.loads(settings_path.read_text(encoding = "utf-8")) + downloads = settings.get("downloadsFolder", "") + if downloads: + _add(downloads) + except Exception: + pass + _add(Path.home() / ".lmstudio" / "models") + _add(Path.home() / ".cache" / "lm-studio" / "models") + return dirs + + +def well_known_model_dirs() -> list[Path]: + candidates: list[Path] = [] + candidates.extend(lmstudio_model_dirs()) + candidates.extend(ollama_model_dirs()) + candidates.append(Path.home() / ".cache" / "huggingface" / "hub") + candidates.append(Path.home() / "models") + candidates.append(Path.home() / "Models") + + out: list[Path] = [] + seen: set[str] = set() + for path in candidates: + try: + resolved = path.resolve() + except OSError: + continue + key = str(resolved) + if key in seen or not resolved.is_dir(): + continue + seen.add(key) + out.append(resolved) + return out + + +def _assert_contained(resolved: Path, root: Path) -> None: + try: + resolved_real = Path(os.path.realpath(resolved)) + root_real = Path(os.path.realpath(root)) + except OSError as exc: + raise ValueError(f"path resolution failed: {exc}") from exc + try: + resolved_real.relative_to(root_real) + except ValueError as exc: + raise ValueError(f"path escapes root: {resolved!s}") from exc + + +def path_is_same_or_child(path: Path, root: Path) -> bool: + """True when *path* is *root* or lives beneath it. + + Compares real (symlink-resolved, case-normalized) paths so the check holds + through symlinks and on case-insensitive filesystems, where a plain + ``Path.is_relative_to`` would miss a casing-only match. Returns False on any + resolution error rather than raising. + """ + try: + path_real = os.path.normcase(os.path.realpath(str(path))) + root_real = os.path.normcase(os.path.realpath(str(root))) + return os.path.commonpath([path_real, root_real]) == root_real + except (OSError, ValueError): + return False + + +def resolve_dataset_path(path_value: str) -> Path: + raw = str(path_value or "").strip() + if "\x00" in raw: + raise ValueError("dataset path may not contain null bytes") + # Normalize first so Windows/UNC and backslash paths resolve like the rest + # of the Hub path layer (e.g. C:\data -> /mnt/c/data on WSL) and a + # backslashed '..' is caught by the traversal guard below. + normalized = normalize_path(raw) + path = Path(normalized).expanduser() + if ".." in path.parts: + raise ValueError(f"dataset path may not contain '..' segments: {raw!r}") + if path.is_absolute(): + for root in (datasets_root(), dataset_uploads_root(), recipe_datasets_root()): + try: + _assert_contained(path, root) + return path + except ValueError: + continue + raise ValueError(f"dataset path must be relative or under a dataset root: {raw!r}") + + parts = [part for part in Path(normalized).parts if part not in ("", ".")] + if parts[:2] == ["assets", "datasets"]: + parts = parts[2:] + if parts and parts[0] == "uploads": + cleaned = Path(*parts[1:]) if len(parts) > 1 else Path() + return dataset_uploads_root() / cleaned + if parts and parts[0] == "recipes": + cleaned = Path(*parts[1:]) if len(parts) > 1 else Path() + return recipe_datasets_root() / cleaned + + cleaned = Path(*parts) if parts else Path() + candidates = [ + dataset_uploads_root() / cleaned, + recipe_datasets_root() / cleaned, + datasets_root() / cleaned, + dataset_uploads_root() / cleaned.name, + recipe_datasets_root() / cleaned.name, + ] + for candidate in candidates: + if candidate.exists(): + return candidate + return candidates[0] + + +def resolve_cached_repo_id_case( + model_name: str, + use_memo: bool = True, + repo_type: str = "model", +) -> str: + """Resolve repo_id to the exact casing already present in local HF cache. + + Prefers the requested casing, but if a case-variant already exists in + local HF cache, reuses that exact cached spelling so we don't trigger + a duplicate download. + """ + if not model_name or "/" not in model_name: + return model_name + + cache_dirs = _hf_hub_cache_dirs() + if not cache_dirs: + return model_name + + prefix = f"{repo_type}s--" + expected_dir = f"{prefix}{model_name.replace('/', '--')}" + memo_key = (repo_type, model_name) + + for cache_dir in cache_dirs: + exact_path = cache_dir / expected_dir + if exact_path.is_dir(): + if use_memo: + _memo_set(memo_key, model_name) + return model_name + + if use_memo: + cached = _memo_get(memo_key) + if cached is not None: + if any( + (cache_dir / f"{prefix}{cached.replace('/', '--')}").is_dir() + for cache_dir in cache_dirs + ): + return cached + _memo_drop(memo_key) + + expected_lower = expected_dir.lower() + try: + candidates: set[str] = set() + for cache_dir in cache_dirs: + for entry in cache_dir.iterdir(): + if not entry.is_dir(): + continue + if entry.name.lower() != expected_lower: + continue + # The lowercased full-name match already proves the prefix + # matches; a case-sensitive startswith would reject a mixed-case + # imported dir such as Models--Org--Repo. + repo_part = entry.name[len(prefix) :] + if not repo_part: + continue + candidates.add(repo_part.replace("--", "/")) + + if candidates: + resolved = sorted(candidates)[0] + if use_memo: + _memo_set(memo_key, resolved) + return resolved + except Exception as exc: + logger.debug(f"resolve_cached_repo_id_case failed for {model_name!r}: {exc}") + + return model_name + + +__all__ = [ + "assets_root", + "cache_root", + "dataset_uploads_root", + "datasets_root", + "ensure_dir", + "exports_root", + "hf_default_cache_dir", + "is_local_path", + "is_valid_gguf_variant", + "is_valid_repo_id", + "legacy_hf_cache_dir", + "lmstudio_model_dirs", + "normalize_path", + "ollama_model_dirs", + "outputs_root", + "path_is_same_or_child", + "recipe_datasets_root", + "resolve_cached_repo_id_case", + "resolve_dataset_path", + "studio_root", + "tmp_root", + "well_known_model_dirs", +] diff --git a/studio/backend/hub/utils/snapshot_filters.py b/studio/backend/hub/utils/snapshot_filters.py new file mode 100644 index 0000000000..20674db4f0 --- /dev/null +++ b/studio/backend/hub/utils/snapshot_filters.py @@ -0,0 +1,110 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from __future__ import annotations + +from fnmatch import fnmatchcase +from typing import Iterable + + +SNAPSHOT_IGNORE_PATTERNS: tuple[str, ...] = ( + "*.gguf", + "*.onnx", + "onnx/*", + "openvino/*", + "mlx/*", + "*.bin.index.json.bak", +) +CONSOLIDATED_PATTERN = "consolidated*" +SNAPSHOT_WEIGHT_EXTENSIONS = ( + ".safetensors", + ".bin", + ".pt", + ".pth", + ".ckpt", + ".h5", + ".msgpack", + ".npz", +) +SNAPSHOT_NON_BIN_WEIGHT_EXTENSIONS = tuple( + ext for ext in SNAPSHOT_WEIGHT_EXTENSIONS if ext != ".bin" +) +SNAPSHOT_BIN_WEIGHT_PREFIXES = ("model", "pytorch_model", "adapter_model") + + +def _filename(sibling) -> str: + value = getattr(sibling, "rfilename", "") + return value if isinstance(value, str) else "" + + +def _size(sibling) -> int: + value = getattr(sibling, "size", None) + return int(value) if isinstance(value, int) and value > 0 else 0 + + +def repo_ships_transformers_weights(filenames: Iterable[str]) -> bool: + for name in filenames: + base = name.rsplit("/", 1)[-1].lower() + if base.startswith("consolidated"): + continue + if base.endswith(SNAPSHOT_NON_BIN_WEIGHT_EXTENSIONS): + return True + if base.endswith(".bin") and base.startswith(SNAPSHOT_BIN_WEIGHT_PREFIXES): + return True + return False + + +def resolve_snapshot_ignore_patterns_for_files(filenames: Iterable[str]) -> list[str]: + names = list(filenames) + ignore = list(SNAPSHOT_IGNORE_PATTERNS) + if repo_ships_transformers_weights(names): + ignore.append(CONSOLIDATED_PATTERN) + return ignore + + +def sibling_matches_ignore(filename: str, ignore_patterns: Iterable[str]) -> bool: + return any(fnmatchcase(filename, pattern) for pattern in ignore_patterns) + + +def snapshot_download_siblings(siblings: Iterable) -> list: + items = list(siblings) + ignore_patterns = resolve_snapshot_ignore_patterns_for_files( + _filename(sibling) for sibling in items + ) + return [ + sibling + for sibling in items + if not sibling_matches_ignore(_filename(sibling), ignore_patterns) + ] + + +def snapshot_download_size(siblings: Iterable) -> int: + return sum(_size(sibling) for sibling in snapshot_download_siblings(siblings)) + + +def total_size_for_siblings(siblings: Iterable) -> int: + """Sum of declared sizes across siblings verbatim (no ignore filter). + + Use for repo types that download every file (datasets); models go + through ``snapshot_download_size`` so the ignore patterns apply.""" + return sum(_size(sibling) for sibling in siblings) + + +def blob_hashes_for_siblings(siblings: Iterable) -> frozenset[str]: + # Blob filename == file etag (LFS sha256, else git blob id). Collecting both + # lets progress count exactly this revision's files without summing stale + # blobs from other revisions. + hashes: set[str] = set() + for sibling in siblings: + sha = getattr(getattr(sibling, "lfs", None), "sha256", None) + if isinstance(sha, str) and sha: + hashes.add(sha) + continue + blob_id = getattr(sibling, "blob_id", None) + if isinstance(blob_id, str) and blob_id: + hashes.add(blob_id) + return frozenset(hashes) + + +def snapshot_download_blob_hashes(siblings: Iterable) -> frozenset[str]: + return blob_hashes_for_siblings(snapshot_download_siblings(siblings)) diff --git a/studio/backend/hub/utils/state_dir.py b/studio/backend/hub/utils/state_dir.py new file mode 100644 index 0000000000..a304477a3d --- /dev/null +++ b/studio/backend/hub/utils/state_dir.py @@ -0,0 +1,154 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Filesystem layout for Hub download state. + +State directory sits beside HF's cache (under Studio's own cache root) +so it survives ``huggingface-cli delete-cache`` and any other HF-side +cache lifecycle. Two subdirectories: + + /hub-state/ + manifests/ .json per-download expected-files manifest + cancelled/ .json per-download cancel marker + +The ```` mirrors HF's cache dir naming so a state file can be +eyeballed next to the on-disk repo it describes: + + models---- full snapshot + models------variant-- GGUF variant + datasets---- dataset snapshot + +All path accessors return ``Optional[Path]`` and yield ``None`` when +the directory can't be created (read-only FS, permission error). +Callers must treat ``None`` as "no state available" and fall through +to existing on-disk-only behavior; this module never raises on a +configuration failure. +""" + +from __future__ import annotations + +import hashlib +import re +from pathlib import Path +from typing import Literal, Optional, get_args + +from loggers import get_logger + +from hub.utils.paths import cache_root + +logger = get_logger(__name__) + + +RepoType = Literal["model", "dataset"] + +_VALID_REPO_TYPES: tuple[RepoType, ...] = get_args(RepoType) + + +_HUB_STATE_DIRNAME = "hub-state" +_MANIFESTS_SUBDIR = "manifests" +_CANCELLED_SUBDIR = "cancelled" +_WORKERS_SUBDIR = "workers" +_SAFE_VARIANT_FRAGMENT = re.compile(r"^[a-z0-9._-]{1,64}$") + + +def state_root() -> Optional[Path]: + """Return the Hub state root, creating it if needed. ``None`` on failure.""" + root = cache_root() / _HUB_STATE_DIRNAME + try: + root.mkdir(parents = True, exist_ok = True) + except OSError as exc: + logger.debug("Could not create hub state root %s: %s", root, exc) + return None + return root + + +def _subdir(name: str) -> Optional[Path]: + root = state_root() + if root is None: + return None + path = root / name + try: + path.mkdir(parents = True, exist_ok = True) + except OSError as exc: + logger.debug("Could not create hub state subdir %s: %s", path, exc) + return None + return path + + +def repo_cache_basename(repo_type: RepoType, repo_id: str) -> str: + # Reject a bad repo_type at runtime: a wrong value would silently produce a + # wrong filename and a misclassified scanner row (the Literal only guards + # statically; dynamic/JSON-sourced values slip past it). + if repo_type not in _VALID_REPO_TYPES: + raise ValueError(f"repo_type must be one of {_VALID_REPO_TYPES}, got {repo_type!r}") + return f"{repo_type}s--{repo_id.replace('/', '--')}".lower() + + +def variant_filename_prefix(repo_type: RepoType, repo_id: str) -> str: + """Lowercased prefix every variant-keyed state file for this repo shares. + + The single source the download_manifest enumerators match against, so the + scheme in :func:`_entry_key` cannot drift from them silently.""" + return f"{repo_cache_basename(repo_type, repo_id)}--variant--" + + +def _entry_key(repo_type: RepoType, repo_id: str, variant: Optional[str]) -> str: + base = repo_cache_basename(repo_type, repo_id) + if not variant: + return base + normalized_variant = variant.strip().lower() + if _SAFE_VARIANT_FRAGMENT.fullmatch(normalized_variant): + variant_fragment = normalized_variant + else: + digest = hashlib.sha256(normalized_variant.encode("utf-8")).hexdigest()[:32] + variant_fragment = f"sha256-{digest}" + return f"{variant_filename_prefix(repo_type, repo_id)}{variant_fragment}" + + +def manifest_path( + repo_type: RepoType, + repo_id: str, + variant: Optional[str] = None, +) -> Optional[Path]: + """Path to the manifest file for this triple. May or may not exist.""" + parent = _subdir(_MANIFESTS_SUBDIR) + if parent is None: + return None + return parent / f"{_entry_key(repo_type, repo_id, variant)}.json" + + +def marker_path( + repo_type: RepoType, + repo_id: str, + variant: Optional[str] = None, +) -> Optional[Path]: + """Path to the cancel-marker file for this triple. May or may not exist.""" + parent = _subdir(_CANCELLED_SUBDIR) + if parent is None: + return None + return parent / f"{_entry_key(repo_type, repo_id, variant)}.json" + + +def manifests_dir() -> Optional[Path]: + """Manifests subdirectory, created on demand. ``None`` on failure. + + Exposed for iter_variant_manifests, which enumerates the directory to find + every variant-keyed manifest for a repo (the path helpers above answer + "where would key X go" but not "what keys exist").""" + return _subdir(_MANIFESTS_SUBDIR) + + +def cancelled_dir() -> Optional[Path]: + """Cancel-marker subdirectory, created on demand. ``None`` on failure. + + See manifests_dir for why this iteration entry point is needed.""" + return _subdir(_CANCELLED_SUBDIR) + + +def workers_dir() -> Optional[Path]: + """Worker PID-breadcrumb subdirectory, created on demand. ``None`` on failure. + + Each live download worker drops one breadcrumb here so a backend that + restarts after a hard crash can reap workers it can no longer reach through + its in-memory registry.""" + return _subdir(_WORKERS_SUBDIR) diff --git a/studio/backend/hub/workers/__init__.py b/studio/backend/hub/workers/__init__.py new file mode 100644 index 0000000000..32014236c6 --- /dev/null +++ b/studio/backend/hub/workers/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 diff --git a/studio/backend/hub/workers/hf_download.py b/studio/backend/hub/workers/hf_download.py new file mode 100644 index 0000000000..42a8ca52b3 --- /dev/null +++ b/studio/backend/hub/workers/hf_download.py @@ -0,0 +1,753 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""HuggingFace Hub download worker, spawned as a subprocess so SIGKILL stops all chunk threads. + +Resume safety +------------- +Downloads here MUST be single-stream sequential writers so the parent's +SIGKILL → restart loop can rely on ``os.path.getsize(.incomplete)`` to +compute the correct resume offset. + +Enforced by: +- Setting ``HF_HUB_DISABLE_XET=1`` and ``HF_HUB_ENABLE_HF_TRANSFER=0`` on + the spawning side (see :mod:`hub.utils.download_registry`) for transport=http. +- Passing ``max_workers=1`` to ``snapshot_download`` so files download + serially, making the at-most-one-active-`.incomplete` invariant hold + globally and simplifying reasoning about partial state during a SIGKILL. +- Letting ``prepare_cache_for_transport`` purge any pre-existing + ``.incomplete`` blobs not provably from the same sequential writer. + +If the final byte count doesn't match what HF declared, huggingface_hub +raises ``EnvironmentError`` ("Consistency check failed: …"); we surface +that on stderr so the watcher can show the exact message to the user. +""" + +from __future__ import annotations + +import argparse +import os +import signal +import sys +import threading +import time +from pathlib import Path + +_HERE = Path(__file__).resolve().parent +_BACKEND = _HERE.parent.parent +if str(_BACKEND) not in sys.path: + sys.path.insert(0, str(_BACKEND)) + +from hub.utils.snapshot_filters import ( + SNAPSHOT_IGNORE_PATTERNS, +) +from hub.utils.gguf_plan import ( + GgufVariantPlan, + build_gguf_variant_plans, + plan_from_expected_files, + sibling_sha256, +) +from hub.utils.state_dir import RepoType + +HfTokenArg = str | bool | None + + +# Bound the metadata fetch so a stalled connection fails the worker (exit 1) +# instead of hanging at 0%. The file download itself is governed separately by +# huggingface_hub's own timeout. +_METADATA_REQUEST_TIMEOUT = 10.0 +_METADATA_RETRY_TIMEOUT = 30.0 +_METADATA_RETRY_DELAY = 1.0 + + +def _on_signal(signum, frame): + # 130 is what `classify_exit` maps to the "cancelled" job state. + sys.exit(130) + + +def _install_signal_handlers() -> None: + signal.signal(signal.SIGTERM, _on_signal) + signal.signal(signal.SIGINT, _on_signal) + sigpipe = getattr(signal, "SIGPIPE", None) + if sigpipe is not None: + signal.signal(sigpipe, _on_signal) + + +def _parent_poll_seconds() -> float: + raw = os.environ.get("UNSLOTH_HF_WORKER_PARENT_POLL_SECONDS") + if raw: + try: + value = float(raw) + if value > 0: + return value + except ValueError: + pass + return 2.0 + + +def _protected_blob_hashes() -> frozenset[str]: + """Blob hashes a concurrent same-repo peer is writing (passed by the backend + as a plain env list). Excluded from this worker's purge so a shared + ``.incomplete`` (e.g. a bundled mmproj) is never deleted under the peer.""" + raw = os.environ.get("UNSLOTH_PROTECTED_BLOB_HASHES", "") + return frozenset(h for h in raw.split(",") if h) + + +def _parent_is_alive(parent_pid: int) -> bool: + """Whether the recorded parent (the backend) is still running. + + Liveness ONLY: ``os.kill(pid, 0)`` on POSIX, an ``OpenProcess`` handle on + Windows, against the *recorded* PID (never os.getppid(), so POSIX + reparenting to init after the backend dies still resolves as dead). Probe + ambiguity is treated as alive so a transient error never kills a healthy + download. + + We deliberately do NOT compare psutil ``create_time()`` for PID-reuse + detection: it isn't stable across reads on some platforms, so an exact match + can spuriously kill a live download. PID-reuse after parent death is covered + by the boot-time orphan reaper. + """ + if sys.platform == "win32": + import ctypes + from ctypes import wintypes + + SYNCHRONIZE = 0x00100000 + WAIT_OBJECT_0 = 0x0 + ERROR_INVALID_PARAMETER = 87 + kernel32 = ctypes.WinDLL("kernel32", use_last_error = True) + kernel32.OpenProcess.argtypes = [wintypes.DWORD, wintypes.BOOL, wintypes.DWORD] + kernel32.OpenProcess.restype = wintypes.HANDLE + kernel32.WaitForSingleObject.argtypes = [wintypes.HANDLE, wintypes.DWORD] + kernel32.WaitForSingleObject.restype = wintypes.DWORD + kernel32.CloseHandle.argtypes = [wintypes.HANDLE] + kernel32.CloseHandle.restype = wintypes.BOOL + ctypes.set_last_error(0) + handle = kernel32.OpenProcess(SYNCHRONIZE, False, parent_pid) + if not handle: + return ctypes.get_last_error() != ERROR_INVALID_PARAMETER + try: + return kernel32.WaitForSingleObject(handle, 0) != WAIT_OBJECT_0 + finally: + kernel32.CloseHandle(handle) + try: + os.kill(parent_pid, 0) + except ProcessLookupError: + return False + except OSError: + return True + return True + + +def _terminate_orphaned_self() -> None: + # Hard exit from the watchdog thread: a self-SIGTERM would be deferred while + # the main thread is GIL-blocked in a C socket read. The partial .incomplete + # resumes byte-exact and marker/manifest writes are atomic, so cancelled code + # 130 is safe. The diagnostic is best-effort: a dead parent's closed stderr + # pipe can raise BrokenPipeError, which must never preempt the exit. + try: + print( + "Parent process exited; stopping orphaned download worker.", + file = sys.stderr, + ) + sys.stderr.flush() + except Exception: + pass + os._exit(130) + + +def _install_parent_death_watchdog(parent_pid: int | None) -> None: + if not parent_pid or parent_pid <= 0: + return + interval = _parent_poll_seconds() + + def _watch() -> None: + while True: + try: + alive = _parent_is_alive(parent_pid) + except Exception: + alive = True + if not alive: + _terminate_orphaned_self() + return + time.sleep(interval) + + threading.Thread( + target = _watch, + name = "parent-death-watchdog", + daemon = True, + ).start() + + +def _hf_token_arg(hf_token: str | None) -> HfTokenArg: + return hf_token if hf_token else False + + +def _retry_metadata_fetch(repo_id: str, fetch, *, label: str): + for attempt, timeout in enumerate((_METADATA_REQUEST_TIMEOUT, _METADATA_RETRY_TIMEOUT)): + try: + return fetch(timeout) + except Exception as e: + if attempt == 1: + raise + print( + f"{label} request failed for {repo_id} " f"({type(e).__name__}: {e}); retrying.", + file = sys.stderr, + ) + time.sleep(_METADATA_RETRY_DELAY) + raise RuntimeError(f"{label} unavailable for {repo_id}") + + +def _model_info_with_retry(repo_id: str, hf_token: str | None): + from huggingface_hub import model_info as hf_model_info + return _retry_metadata_fetch( + repo_id, + lambda timeout: hf_model_info( + repo_id, + token = _hf_token_arg(hf_token), + timeout = timeout, + files_metadata = True, + ), + label = "Metadata", + ) + + +def _dataset_info_with_retry(repo_id: str, hf_token: str | None): + from huggingface_hub import HfApi + api = HfApi(token = _hf_token_arg(hf_token)) + return _retry_metadata_fetch( + repo_id, + lambda timeout: api.dataset_info( + repo_id, + timeout = timeout, + files_metadata = True, + ), + label = "Dataset metadata", + ) + + +# Tied to drain_stderr_excerpt's 500-byte head/tail window in the parent (see +# hub/utils/download_registry.py): listing every expected file would blow past +# it and lose the diagnostic. Cap the preview so the summary line survives. +_VERIFY_PATH_LIST_CAP = 10 + + +def _format_path_list(paths: tuple[str, ...], cap: int = _VERIFY_PATH_LIST_CAP) -> str: + if len(paths) <= cap: + return ", ".join(paths) + head = ", ".join(paths[:cap]) + return f"{head}, ... and {len(paths) - cap} more" + + +def _verify_completed_download( + repo_type: RepoType, + repo_id: str, + variant: str | None, + snapshot_path: str, + *, + metadata_unavailable: bool = False, +) -> None: + """Verify every manifest file is on disk at its declared size; exit nonzero + with a diagnostic if not. + + No-op when no manifest exists: the manifest write is best-effort, so absence + means "verification unavailable, trust snapshot_download's exit code". + """ + from hub.utils import download_manifest + + manifest = download_manifest.read_manifest(repo_type, repo_id, variant) + if manifest is None: + return + result = download_manifest.verify_against_disk( + manifest, + Path(snapshot_path), + ) + if result.ok: + return + label = f"{repo_id}{f' [{variant}]' if variant else ''}" + if metadata_unavailable: + print( + f"Could not reach Hugging Face for {label} and the copy on disk is " + f"incomplete ({len(result.missing)} file(s) missing, " + f"{len(result.size_mismatched)} the wrong size). Access to a private " + "or restricted repo may have been lost (HF token removed or " + "changed), the connection dropped, or Hugging Face is temporarily " + "unavailable. Set a valid HF token or reconnect, then resume the " + "download.", + file = sys.stderr, + ) + else: + print( + f"Verification failed for {label}: snapshot_download completed but " + f"{len(result.missing)} expected file(s) are missing and " + f"{len(result.size_mismatched)} have incorrect size on disk.", + file = sys.stderr, + ) + if result.missing: + print( + f"Missing: {_format_path_list(result.missing)}", + file = sys.stderr, + ) + if result.size_mismatched: + print( + f"Size mismatched: {_format_path_list(result.size_mismatched)}", + file = sys.stderr, + ) + sys.exit(1) + + +def _preflight_disk_space(repo_type: str, repo_id: str, expected_files: list) -> None: + """Fail fast when the active HF cache filesystem can't hold what's left to + download. Fail-open: any inability to size the work or read free space skips + the check, so a real download is never blocked by an estimation gap.""" + import shutil + + from hub.utils.download_registry import existing_blob_bytes + from hub.utils.hf_cache_state import hf_cache_root + + try: + size_by_hash: dict[str, int] = {} + unhashed_bytes = 0 + for expected in expected_files: + size = int(getattr(expected, "size", 0) or 0) + if size <= 0: + continue + blob_hash = getattr(expected, "sha256", None) + if blob_hash: + # Dedup by content hash: a blob listed under two filenames is + # written once, so it must be counted once. + size_by_hash[blob_hash] = size + else: + unhashed_bytes += size + total_expected = sum(size_by_hash.values()) + unhashed_bytes + if total_expected <= 0: + return + already_have = existing_blob_bytes( + repo_type, + repo_id, + frozenset(size_by_hash), + ) + remaining = max(0, total_expected - already_have) + if remaining <= 0: + return + root = hf_cache_root(create = True) + if root is None: + return + free = shutil.disk_usage(root).free + except Exception: + return + + if free < remaining: + print( + f"Not enough disk space to download {repo_id}: need about " + f"{remaining / (1024 ** 3):.1f} GB free in {root}, but only " + f"{free / (1024 ** 3):.1f} GB is available. Free up space and " + "try again.", + file = sys.stderr, + ) + sys.exit(1) + + +def _snapshot_download_plan(info) -> tuple[list[str], list]: + from hub.utils.download_manifest import ExpectedFile + from hub.utils.snapshot_filters import ( + resolve_snapshot_ignore_patterns_for_files, + snapshot_download_siblings, + ) + + filenames = [s.rfilename for s in info.siblings if isinstance(s.rfilename, str)] + filtered = snapshot_download_siblings(info.siblings) + expected_files = [ + ExpectedFile( + path = s.rfilename, + size = int(getattr(s, "size", 0) or 0), + sha256 = sibling_sha256(s), + ) + for s in filtered + if isinstance(s.rfilename, str) + ] + return resolve_snapshot_ignore_patterns_for_files(filenames), expected_files + + +def _dataset_expected_files(info) -> list: + from hub.utils.download_manifest import ExpectedFile + return [ + ExpectedFile( + path = s.rfilename, + size = int(getattr(s, "size", 0) or 0), + sha256 = sibling_sha256(s), + ) + for s in info.siblings + if isinstance(s.rfilename, str) + ] + + +def _recover_manifest_after_download( + repo_type: RepoType, + repo_id: str, + snapshot_path: str, + mode: str, + *, + fetch_info, + expected_files_from_info, + label: str = "", +) -> None: + """Best-effort manifest write for a download whose metadata was unavailable + at start: re-fetch and record the expected files, else fall back to the + on-disk file list. Shared by the model and dataset workers. + + A pre-existing manifest is authoritative and is preserved untouched. This is + load-bearing: when access is lost on resume (token revoked/changed on a + gated/private repo), snapshot_download returns the cached partial snapshot + WITHOUT downloading, so rebuilding the manifest from on-disk files would record + the partial set as expected and let _verify_completed_download certify a + half-finished download as complete. + + The same hazard exists with NO prior manifest. When metadata is still + unavailable here, leftover ``.incomplete`` blobs prove a cached partial was + returned without downloading, so we fail (exit 1) instead of deriving a + self-certifying manifest, leaving the partial intact for a later resume. That + signal misses a file that never started (no ``.incomplete``), so a kill + between files is accepted optimistically from the on-disk subset; a later + metadata-bearing attempt writes the true manifest and catches any shortfall.""" + from hub.utils import download_manifest + from hub.utils.hf_cache_state import has_active_incomplete_blobs + + if download_manifest.read_manifest(repo_type, repo_id, None) is not None: + return + + try: + if download_manifest.write_manifest( + repo_type, + repo_id, + None, + expected_files_from_info(fetch_info()), + mode, + ): + return + reason = "manifest write failed" + except Exception as e: + reason = f"{type(e).__name__}: {e}" + + if has_active_incomplete_blobs(repo_type, repo_id): + print( + f"{label}could not reach Hugging Face for {repo_id} and the copy on " + "disk is still incomplete. Access to a private or restricted repo may " + "have been lost (HF token removed or changed), the connection dropped, " + "or Hugging Face is temporarily unavailable. Set a valid HF token or " + "reconnect, then resume the download.", + file = sys.stderr, + ) + sys.exit(1) + + fallback_files = download_manifest.expected_files_from_snapshot_dir(Path(snapshot_path)) + if fallback_files and download_manifest.write_manifest( + repo_type, + repo_id, + None, + fallback_files, + mode, + ): + print( + f"{label}could not record the metadata manifest for {repo_id}, " + "recorded one from the downloaded files so completion " + f"is tracked ({reason})", + file = sys.stderr, + ) + else: + print( + f"{label}could not record the metadata manifest for {repo_id}, " + f"{download_manifest.MANIFEST_DEGRADED_MARKER} ({reason})", + file = sys.stderr, + ) + + +def _download_snapshot(repo_id: str, hf_token: str | None, mode: str) -> None: + from huggingface_hub import snapshot_download + from hub.utils.download_registry import prepare_cache_for_transport + from hub.utils import download_manifest + + # One metadata fetch powers both the ignore-pattern decision (drop + # consolidated.* when transformers weights exist) and the manifest's + # expected_files. A failure is non-fatal: fall back to the legacy + # ignore-pattern set (keeping consolidated) and skip the manifest, so + # download proceeds without the verification + partial detection it enables. + try: + info = _model_info_with_retry(repo_id, hf_token) + except Exception as e: + print( + f"metadata unavailable, downloading full snapshot for {repo_id} " + f"({type(e).__name__}: {e})", + file = sys.stderr, + ) + info = None + + download_manifest.clear_cancel_marker("model", repo_id, None) + if info is not None: + ignore_patterns, expected_files = _snapshot_download_plan(info) + # Written for every transport. The manifest verifies the finalized + # files under snapshots/, which both transports produce identically + # (XET also renames a full, correctly-sized blob into place). XET's + # block-level dedup lives only in the chunk-cache the manifest never + # inspects, so per-file size verification is valid regardless of transport. + download_manifest.write_manifest("model", repo_id, None, expected_files, mode) + else: + ignore_patterns = list(SNAPSHOT_IGNORE_PATTERNS) + expected_files = [] + + purged = prepare_cache_for_transport("model", repo_id, mode) + if purged: + print( + f"Purged {purged} untrusted partial blob(s) for {repo_id} " + f"before starting {mode} download.", + file = sys.stderr, + ) + _preflight_disk_space("model", repo_id, expected_files) + snapshot_path = snapshot_download( + repo_id = repo_id, + token = _hf_token_arg(hf_token), + ignore_patterns = ignore_patterns, + max_workers = 1, + ) + if info is None: + _recover_manifest_after_download( + "model", + repo_id, + snapshot_path, + mode, + fetch_info = lambda: _model_info_with_retry(repo_id, hf_token), + expected_files_from_info = lambda recovered: _snapshot_download_plan(recovered)[1], + ) + _verify_completed_download( + "model", + repo_id, + None, + snapshot_path, + metadata_unavailable = info is None, + ) + + +def _gguf_variant_target_plan( + repo_id: str, variant: str, hf_token: str | None +) -> GgufVariantPlan | None: + try: + info = _model_info_with_retry(repo_id, hf_token) + except Exception as e: + print( + f"metadata unavailable, cannot resolve GGUF variant '{variant}' " + f"for {repo_id} ({type(e).__name__}: {e})", + file = sys.stderr, + ) + raise RuntimeError( + f"Metadata unavailable while resolving GGUF variant '{variant}' " f"for {repo_id}" + ) from e + return build_gguf_variant_plans(list(info.siblings)).get(variant.lower()) + + +def _download_gguf_variant(repo_id: str, variant: str, hf_token: str | None, mode: str) -> None: + from huggingface_hub import snapshot_download + from hub.utils.download_registry import prepare_cache_for_transport + from hub.utils.hf_cache_state import has_active_incomplete_blobs + from hub.utils import download_manifest + + metadata_unavailable = False + try: + plan = _gguf_variant_target_plan(repo_id, variant, hf_token) + except RuntimeError: + plan = None + metadata_unavailable = True + + if not metadata_unavailable: + if plan is None: + print( + f"No GGUF shards matching variant '{variant}' in {repo_id}", + file = sys.stderr, + ) + sys.exit(1) + targets = list(plan.target_filenames) + expected_files = list(plan.expected_files) + main_blob_hashes = plan.main_hashes + companion_blob_hashes = plan.companion_hashes + download_manifest.write_manifest( + "model", + repo_id, + variant, + expected_files, + mode, + ) + else: + # Metadata unreachable (offline / gated / private). Resume the exact + # shards the original attempt recorded so snapshot_download can range + # over the surviving .incomplete blobs without a model_info call. + manifest = download_manifest.read_manifest("model", repo_id, variant) + if manifest is None or not manifest.expected_files: + print( + f"Metadata unavailable and no manifest to resume GGUF " + f"variant '{variant}' for {repo_id}", + file = sys.stderr, + ) + sys.exit(1) + plan = plan_from_expected_files(variant, manifest.expected_files) + targets = list(plan.target_filenames) + expected_files = list(plan.expected_files) + download_manifest.write_manifest( + "model", + repo_id, + variant, + expected_files, + mode, + ) + main_blob_hashes = plan.main_hashes + companion_blob_hashes = plan.companion_hashes + print( + f"Metadata unavailable; resuming GGUF variant '{variant}' for " + f"{repo_id} from the existing manifest.", + file = sys.stderr, + ) + + download_manifest.clear_cancel_marker("model", repo_id, variant) + purge_blob_hashes = main_blob_hashes + if not main_blob_hashes: + if has_active_incomplete_blobs("model", repo_id): + print( + f"GGUF variant '{variant}' for {repo_id} has partial cache state " + "but no resolvable blob hashes; delete the partial download or " + "retry when metadata is available.", + file = sys.stderr, + ) + sys.exit(1) + purge_blob_hashes = frozenset() + print( + f"GGUF variant '{variant}' for {repo_id} has no resolvable blob " + "hashes; starting without partial cache reuse.", + file = sys.stderr, + ) + # Main quant blobs are owned by this variant (variant-scoped marker). The + # shared vision companion (mmproj) is judged by a separate companion marker + # and never purged while a concurrent peer is writing it. + purged = prepare_cache_for_transport( + "model", + repo_id, + mode, + variant, + only_blob_hashes = purge_blob_hashes, + companion_blob_hashes = companion_blob_hashes, + protected_blob_hashes = _protected_blob_hashes(), + ) + if purged: + print( + f"Purged {purged} untrusted partial blob(s) for {repo_id} " + f"before starting {mode} download.", + file = sys.stderr, + ) + _preflight_disk_space("model", repo_id, expected_files) + snapshot_path = snapshot_download( + repo_id = repo_id, + token = _hf_token_arg(hf_token), + allow_patterns = targets, + max_workers = 1, + ) + _verify_completed_download( + "model", + repo_id, + variant, + snapshot_path, + metadata_unavailable = metadata_unavailable, + ) + + +def _download_dataset(repo_id: str, hf_token: str | None, mode: str) -> None: + from huggingface_hub import snapshot_download + from hub.utils.download_registry import prepare_cache_for_transport + from hub.utils import download_manifest + + try: + info = _dataset_info_with_retry(repo_id, hf_token) + except Exception as e: + print( + f"dataset metadata unavailable, downloading full dataset for {repo_id} " + f"({type(e).__name__}: {e})", + file = sys.stderr, + ) + info = None + # Cancel-marker clear and manifest write run on every transport. See + # _download_snapshot for why per-file size verification is valid under XET. + download_manifest.clear_cancel_marker("dataset", repo_id, None) + if info is not None: + expected_files = _dataset_expected_files(info) + download_manifest.write_manifest( + "dataset", + repo_id, + None, + expected_files, + mode, + ) + else: + expected_files = [] + purged = prepare_cache_for_transport("dataset", repo_id, mode) + if purged: + print( + f"Purged {purged} untrusted partial blob(s) for {repo_id} " + f"before starting {mode} download.", + file = sys.stderr, + ) + _preflight_disk_space("dataset", repo_id, expected_files) + snapshot_path = snapshot_download( + repo_id = repo_id, + token = _hf_token_arg(hf_token), + repo_type = "dataset", + max_workers = 1, + ) + if info is None: + _recover_manifest_after_download( + "dataset", + repo_id, + snapshot_path, + mode, + fetch_info = lambda: _dataset_info_with_retry(repo_id, hf_token), + expected_files_from_info = _dataset_expected_files, + label = "dataset ", + ) + _verify_completed_download( + "dataset", + repo_id, + None, + snapshot_path, + metadata_unavailable = info is None, + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description = "HuggingFace Hub download worker") + parser.add_argument("--repo-id", required = True) + parser.add_argument("--variant", default = None) + parser.add_argument("--dataset", action = "store_true") + parser.add_argument("--transport", choices = ("http", "xet"), default = "http") + parser.add_argument("--parent-pid", type = int, default = None) + args = parser.parse_args() + + _install_signal_handlers() + _install_parent_death_watchdog(args.parent_pid) + + hf_token = os.environ.get("HF_TOKEN") or None + + try: + if args.dataset: + _download_dataset(args.repo_id, hf_token, args.transport) + elif args.variant: + _download_gguf_variant(args.repo_id, args.variant, hf_token, args.transport) + else: + _download_snapshot(args.repo_id, hf_token, args.transport) + sys.exit(0) + except SystemExit: + raise + except Exception as e: + # Surface a precise message so the UI doesn't show a generic "worker + # exited with code 1". huggingface_hub's consistency check recommends + # force_download=True to recover, which our "Restart" UI maps to a fresh + # start by purging the partial via prepare_cache_for_transport. + print(f"{type(e).__name__}: {e}", file = sys.stderr) + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/studio/backend/main.py b/studio/backend/main.py index 104d557eea..a9f7004df7 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -8,6 +8,8 @@ Main FastAPI application for Unsloth UI Backend import os import sys from pathlib import Path as _Path +import asyncio +from dataclasses import asdict # Suppress C-level dependency warnings globally os.environ["PYTHONWARNINGS"] = "ignore" @@ -217,6 +219,16 @@ from routes import ( training_history_router, training_router, ) +from hub.routes import ( + inventory_router as hub_inventory_router, + datasets_router as hub_datasets_router, +) +from hub.schemas.downloads import TransportCapabilities +from hub.utils.download_registry import ( + get_download_transport_capabilities, + reap_orphan_workers as reap_hub_orphan_workers, + terminate_active_downloads as terminate_hub_downloads, +) from routes.settings import router as settings_router from auth import storage from auth.authentication import get_current_subject @@ -293,6 +305,9 @@ async def lifespan(app: FastAPI): # Detect hardware first — sets the DEVICE global used everywhere. detect_hardware() + # Reap download workers orphaned by a previous crash before new downloads start. + reap_hub_orphan_workers() + # llama.cpp probes: capability (MTP support) + freshness (release age). # Both cached; freshness has a 24h disk TTL. try: @@ -367,6 +382,7 @@ async def lifespan(app: FastAPI): else: app.state.bootstrap_password = storage.get_bootstrap_password() yield + await asyncio.to_thread(terminate_hub_downloads) _hw_module.DEVICE = None clear_unsloth_compiled_cache() @@ -389,8 +405,8 @@ logger = LogConfig.setup_logging( app.add_middleware(LoggingMiddleware) -# Citation favicons load from www.google.com/s2/favicons; *.gstatic.com is -# kept for legacy web-search faviconV2 paths. All else is same-origin. +# img/media-src allow any https origin so HF model-card assets render (mirrors +# tauri.conf.json); scripts/frames/connect-src stay same-origin + HF. from starlette.middleware.base import BaseHTTPMiddleware # noqa: E402 from starlette.requests import Request as _StarletteRequest # noqa: E402 @@ -435,9 +451,8 @@ def _build_csp(script_nonce: "str | None" = None) -> str: return ( "default-src 'self'; " - "img-src 'self' data: blob: https://t0.gstatic.com " - "https://t1.gstatic.com https://t2.gstatic.com " - "https://t3.gstatic.com https://www.google.com; " + "img-src 'self' data: blob: https:; " + "media-src 'self' data: blob: https:; " f"connect-src {connect_src}; " "style-src 'self' 'unsafe-inline'; " f"{script_src}; " @@ -491,6 +506,7 @@ _BODY_PROTECTED_PREFIXES = ( "/api/inference", "/api/data-recipe", "/api/datasets", + "/api/hub", "/api/chat", "/api/settings", "/api/train", @@ -714,6 +730,8 @@ app.include_router(datasets_router, prefix = "/api/datasets", tags = ["datasets" app.include_router(data_recipe_router, prefix = "/api/data-recipe", tags = ["data-recipe"]) app.include_router(export_router, prefix = "/api/export", tags = ["export"]) app.include_router(training_history_router, prefix = "/api/train", tags = ["training-history"]) +app.include_router(hub_inventory_router, prefix = "/api/hub", tags = ["hub"]) +app.include_router(hub_datasets_router, prefix = "/api/hub/datasets", tags = ["hub"]) # ============ Health and System Endpoints ============ @@ -781,6 +799,14 @@ def studio_update_status(_current_subject: str = Depends(get_current_subject)): return get_studio_update_status(UNSLOTH_VERSION) +@app.get( + "/api/studio/download-transport-capabilities", + response_model = TransportCapabilities, +) +def studio_download_transport_capabilities(_current_subject: str = Depends(get_current_subject)): + return asdict(get_download_transport_capabilities()) + + @app.post("/api/shutdown") async def shutdown_server(request: Request, current_subject: str = Depends(get_current_subject)): """Gracefully shut down the Unsloth Studio server. @@ -788,7 +814,6 @@ async def shutdown_server(request: Request, current_subject: str = Depends(get_c Called by the frontend quit dialog so users can stop the server from the UI without the CLI or killing the process manually. """ - import asyncio async def _delayed_shutdown(): await asyncio.sleep(0.2) # Let the HTTP response return first diff --git a/studio/backend/tests/test_middleware.py b/studio/backend/tests/test_middleware.py index 16a7bbc46d..1005431926 100644 --- a/studio/backend/tests/test_middleware.py +++ b/studio/backend/tests/test_middleware.py @@ -276,25 +276,22 @@ class TestSecurityHeadersMiddleware: nonced = main_module._build_csp("XYZ") assert "script-src 'self' 'nonce-XYZ';" in nonced - def test_img_src_allows_google_favicons(self, main_module): - # sources.tsx fetches https://www.google.com/s2/favicons?... ; without - # this allowlist entry, citation favicons fall back to gray initials. + def test_img_and_media_allow_https_sources(self, main_module): + # Model-card READMEs and citation favicons pull images/media from many + # https origins (HF LFS/XET CDNs, shields/badge hosts, GitHub-hosted + # assets, audio/video samples). img-src/media-src allow any https source + # so they render; this mirrors the desktop CSP in tauri.conf.json. csp = main_module._build_csp() - img_directive = next( - chunk.strip() for chunk in csp.split(";") if chunk.strip().startswith("img-src ") - ) - # Tokenise and compare with `==` so CodeQL's URL-substring rule - # doesn't read directive-string `in` membership as URL sanitisation. - img_sources = img_directive.split() - assert any(src == "https://www.google.com" for src in img_sources) - # Pre-existing favicon CDNs stay allowed. - for host in ( - "https://t0.gstatic.com", - "https://t1.gstatic.com", - "https://t2.gstatic.com", - "https://t3.gstatic.com", - ): - assert any(src == host for src in img_sources) + directives = { + chunk.strip().split()[0]: chunk.strip().split() + for chunk in csp.split(";") + if chunk.strip() + } + for name in ("img-src", "media-src"): + assert name in directives, f"missing {name} directive" + # Tokenise and compare with `==` so CodeQL's URL-substring rule does + # not read directive-string `in` membership as URL sanitisation. + assert any(src == "https:" for src in directives[name]) # /api/health auth gate diff --git a/studio/frontend/package-lock.json b/studio/frontend/package-lock.json index f36f2d8e79..7d1283e94d 100644 --- a/studio/frontend/package-lock.json +++ b/studio/frontend/package-lock.json @@ -31,6 +31,7 @@ "@tailwindcss/vite": "^4.2.2", "@tanstack/react-router": "1.169.2", "@tanstack/react-table": "^8.21.3", + "@tanstack/react-virtual": "3.13.25", "@tauri-apps/api": "^2.10.1", "@tauri-apps/plugin-clipboard-manager": "^2.3.2", "@tauri-apps/plugin-notification": "^2.3.3", @@ -6112,6 +6113,23 @@ "react-dom": ">=16.8" } }, + "node_modules/@tanstack/react-virtual": { + "version": "3.13.25", + "resolved": "https://registry.npmjs.org/@tanstack/react-virtual/-/react-virtual-3.13.25.tgz", + "integrity": "sha512-bmNoqMu6gcAW9JGrKVB0Q1tN1i5RONZF8r1fW0bbE4Oyf3DwEGnzzQJ2OW+Ozg1P4s8PyugkHg2ULZoFQN+cqw==", + "license": "MIT", + "dependencies": { + "@tanstack/virtual-core": "3.15.0" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/tannerlinsley" + }, + "peerDependencies": { + "react": "^16.8.0 || ^17.0.0 || ^18.0.0 || ^19.0.0", + "react-dom": "^16.8.0 || ^17.0.0 || ^18.0.0 || ^19.0.0" + } + }, "node_modules/@tanstack/router-core": { "version": "1.169.2", "resolved": "https://registry.npmjs.org/@tanstack/router-core/-/router-core-1.169.2.tgz", @@ -6154,6 +6172,16 @@ "url": "https://github.com/sponsors/tannerlinsley" } }, + "node_modules/@tanstack/virtual-core": { + "version": "3.15.0", + "resolved": "https://registry.npmjs.org/@tanstack/virtual-core/-/virtual-core-3.15.0.tgz", + "integrity": "sha512-0AwPGx0I8QxPYjAxShT/+z+ZOe9u8mW5rsXvivCTjRfRmz9a43+3mRyi4wwlyoUqOC56q/jatKa0Bh9M99BEHQ==", + "license": "MIT", + "funding": { + "type": "github", + "url": "https://github.com/sponsors/tannerlinsley" + } + }, "node_modules/@tauri-apps/api": { "version": "2.11.0", "resolved": "https://registry.npmjs.org/@tauri-apps/api/-/api-2.11.0.tgz", diff --git a/studio/frontend/package.json b/studio/frontend/package.json index 5ba0db143f..8537b0c076 100644 --- a/studio/frontend/package.json +++ b/studio/frontend/package.json @@ -40,6 +40,7 @@ "@tailwindcss/vite": "^4.2.2", "@tanstack/react-router": "1.169.2", "@tanstack/react-table": "^8.21.3", + "@tanstack/react-virtual": "3.13.25", "@tauri-apps/api": "^2.10.1", "@tauri-apps/plugin-clipboard-manager": "^2.3.2", "@tauri-apps/plugin-notification": "^2.3.3", diff --git a/studio/frontend/public/hub/profile/logo/anthropic.svg b/studio/frontend/public/hub/profile/logo/anthropic.svg new file mode 100644 index 0000000000..7545cc8f3e --- /dev/null +++ b/studio/frontend/public/hub/profile/logo/anthropic.svg @@ -0,0 +1,6 @@ + + + + + + \ No newline at end of file diff --git a/studio/frontend/public/hub/profile/logo/cohere.png b/studio/frontend/public/hub/profile/logo/cohere.png new file mode 100644 index 0000000000..99eabbb54f Binary files /dev/null and b/studio/frontend/public/hub/profile/logo/cohere.png differ diff --git a/studio/frontend/public/hub/profile/logo/deepseek.svg b/studio/frontend/public/hub/profile/logo/deepseek.svg new file mode 100644 index 0000000000..d1ba06b942 --- /dev/null +++ b/studio/frontend/public/hub/profile/logo/deepseek.svg @@ -0,0 +1,14 @@ + + + + + + + + + \ No newline at end of file diff --git a/studio/frontend/public/hub/profile/logo/google.png b/studio/frontend/public/hub/profile/logo/google.png new file mode 100644 index 0000000000..01bb81206e Binary files /dev/null and b/studio/frontend/public/hub/profile/logo/google.png differ diff --git a/studio/frontend/public/hub/profile/logo/hf.svg b/studio/frontend/public/hub/profile/logo/hf.svg new file mode 100644 index 0000000000..ab959d165f --- /dev/null +++ b/studio/frontend/public/hub/profile/logo/hf.svg @@ -0,0 +1,8 @@ + + + + + + + + diff --git a/studio/frontend/public/hub/profile/logo/ibm.png b/studio/frontend/public/hub/profile/logo/ibm.png new file mode 100644 index 0000000000..31c965f0b3 Binary files /dev/null and b/studio/frontend/public/hub/profile/logo/ibm.png differ diff --git a/studio/frontend/public/hub/profile/logo/meta.svg b/studio/frontend/public/hub/profile/logo/meta.svg new file mode 100644 index 0000000000..9fa656bd6b --- /dev/null +++ b/studio/frontend/public/hub/profile/logo/meta.svg @@ -0,0 +1,19 @@ + + +Logo of Meta Platforms -- Graphic created by Detmar Owen + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/studio/frontend/public/hub/profile/logo/microsoft.svg b/studio/frontend/public/hub/profile/logo/microsoft.svg new file mode 100644 index 0000000000..5334aa7ca6 --- /dev/null +++ b/studio/frontend/public/hub/profile/logo/microsoft.svg @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/studio/frontend/public/hub/profile/logo/minimax-color.png b/studio/frontend/public/hub/profile/logo/minimax-color.png new file mode 100644 index 0000000000..e9472c676d Binary files /dev/null and b/studio/frontend/public/hub/profile/logo/minimax-color.png differ diff --git a/studio/frontend/public/hub/profile/logo/mistral.svg b/studio/frontend/public/hub/profile/logo/mistral.svg new file mode 100644 index 0000000000..40c2591b31 --- /dev/null +++ b/studio/frontend/public/hub/profile/logo/mistral.svg @@ -0,0 +1,19 @@ + + + + + + + + + + + + + + + + + + + diff --git a/studio/frontend/public/hub/profile/logo/moonshot.jpg b/studio/frontend/public/hub/profile/logo/moonshot.jpg new file mode 100644 index 0000000000..956a5b58b1 Binary files /dev/null and b/studio/frontend/public/hub/profile/logo/moonshot.jpg differ diff --git a/studio/frontend/public/hub/profile/logo/nvidia.svg b/studio/frontend/public/hub/profile/logo/nvidia.svg new file mode 100644 index 0000000000..ae65b09a2b --- /dev/null +++ b/studio/frontend/public/hub/profile/logo/nvidia.svg @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/studio/frontend/public/hub/profile/logo/openai.svg b/studio/frontend/public/hub/profile/logo/openai.svg new file mode 100644 index 0000000000..74d9b1b44b --- /dev/null +++ b/studio/frontend/public/hub/profile/logo/openai.svg @@ -0,0 +1,5 @@ + + + + + \ No newline at end of file diff --git a/studio/frontend/public/hub/profile/logo/qwen.png b/studio/frontend/public/hub/profile/logo/qwen.png new file mode 100644 index 0000000000..67d2258f40 Binary files /dev/null and b/studio/frontend/public/hub/profile/logo/qwen.png differ diff --git a/studio/frontend/public/hub/profile/logo/xai.svg b/studio/frontend/public/hub/profile/logo/xai.svg new file mode 100644 index 0000000000..0c83eb3d9b --- /dev/null +++ b/studio/frontend/public/hub/profile/logo/xai.svg @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/studio/frontend/public/hub/profile/logo/zai.svg b/studio/frontend/public/hub/profile/logo/zai.svg new file mode 100644 index 0000000000..28ca7280a1 --- /dev/null +++ b/studio/frontend/public/hub/profile/logo/zai.svg @@ -0,0 +1,215 @@ + + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/studio/frontend/src/app/provider.tsx b/studio/frontend/src/app/provider.tsx index 4bd348effc..25dbfdc780 100644 --- a/studio/frontend/src/app/provider.tsx +++ b/studio/frontend/src/app/provider.tsx @@ -9,7 +9,9 @@ import { shouldUseCustomWindowTitlebar, } from "@/components/tauri/window-titlebar"; import { Toaster } from "@/components/ui/sonner"; +import { TooltipProvider } from "@/components/ui/tooltip"; import { WebUpdateBanner } from "@/components/web/update-banner"; +import { DownloadManagerPanel } from "@/features/hub/download-manager"; import { getTauriAuthFailure, tauriAutoAuth } from "@/features/auth"; import { NativeIntentDrain } from "@/features/native-intents/native-intent-drain"; import { useTauriBackend, type BackendStatus } from "@/hooks/use-tauri-backend"; @@ -255,6 +257,7 @@ function TauriWrapper({ children }: { children: ReactNode }) { return ( <> {children} + ); @@ -272,6 +275,7 @@ function TauriWrapper({ children }: { children: ReactNode }) { {children} + ) : ( - - {children} - - + + + {children} + + + ); } diff --git a/studio/frontend/src/app/router.tsx b/studio/frontend/src/app/router.tsx index a0ca1e8cdb..dbb74ee1a1 100644 --- a/studio/frontend/src/app/router.tsx +++ b/studio/frontend/src/app/router.tsx @@ -12,6 +12,7 @@ import { Route as exportRoute } from "./routes/export"; import { Route as gridTestRoute } from "./routes/grid-test"; import { Route as indexRoute } from "./routes/index"; import { Route as loginRoute } from "./routes/login"; +import { Route as hubRoute } from "./routes/hub"; import { Route as onboardingRoute } from "./routes/onboarding"; import { Route as projectsRoute } from "./routes/projects"; import { Route as changePasswordRoute } from "./routes/change-password"; @@ -24,6 +25,7 @@ const routeTree = rootRoute.addChildren([ loginRoute, changePasswordRoute, gridTestRoute, + hubRoute, settingsRoute, studioRoute, chatRoute, diff --git a/studio/frontend/src/app/routes/__root.tsx b/studio/frontend/src/app/routes/__root.tsx index 200200d190..c8e47902b9 100644 --- a/studio/frontend/src/app/routes/__root.tsx +++ b/studio/frontend/src/app/routes/__root.tsx @@ -42,6 +42,7 @@ const CHAT_ONLY_ALLOWED = new Set([ "/", "/chat", "/projects", + "/hub", "/login", "/signup", "/change-password", diff --git a/studio/frontend/src/app/routes/hub.tsx b/studio/frontend/src/app/routes/hub.tsx new file mode 100644 index 0000000000..dcd6617ec8 --- /dev/null +++ b/studio/frontend/src/app/routes/hub.tsx @@ -0,0 +1,29 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { createRoute } from "@tanstack/react-router"; +import { lazy } from "react"; +import { requireAuth } from "../auth-guards"; +import { Route as rootRoute } from "./__root"; + +const ModelsPage = lazy(() => + import("@/features/hub/hub-page").then((m) => ({ + default: m.ModelsPage, + })), +); + +export interface ModelsSearch { + tab?: "discover" | "downloaded"; +} + +export const Route = createRoute({ + getParentRoute: () => rootRoute, + path: "/hub", + beforeLoad: () => requireAuth(), + component: ModelsPage, + validateSearch: (search: Record): ModelsSearch => { + const raw = search.tab; + if (raw === "discover" || raw === "downloaded") return { tab: raw }; + return {}; + }, +}); diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx index 8d56d90f98..2abc4d190b 100644 --- a/studio/frontend/src/components/app-sidebar.tsx +++ b/studio/frontend/src/components/app-sidebar.tsx @@ -47,6 +47,7 @@ import { cn } from "@/lib/utils"; import { ChefHatIcon, CursorInfo02Icon, + DashboardCircleIcon, Delete02Icon, DownloadSquare01Icon, Edit03Icon, @@ -791,6 +792,15 @@ export function AppSidebar() { closeMobileIfOpen(); }} /> + { + navigate({ to: "/hub" }); + closeMobileIfOpen(); + }} + /> {/* Train has a labelled section when expanded; plain icon here only when collapsed. */} - - - - + + + ); } diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts index 320cf4c962..ebaa756501 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -23,6 +23,7 @@ import { } from "../utils/chat-settings-storage"; const HF_TOKEN_KEY = "unsloth_hf_token"; +const HF_TOKEN_CHANGED_EVENT = "unsloth:hf-token-changed"; export const CHAT_REASONING_ENABLED_KEY = "unsloth_chat_reasoning_enabled"; export const CHAT_TOOLS_ENABLED_KEY = "unsloth_chat_tools_enabled"; export const CHAT_CODE_TOOLS_ENABLED_KEY = "unsloth_chat_code_tools_enabled"; @@ -231,6 +232,15 @@ function saveString(key: string, value: string): void { } } +function notifyHfTokenChanged(value: string): void { + if (!canUseStorage()) return; + try { + window.dispatchEvent(new CustomEvent(HF_TOKEN_CHANGED_EVENT, { detail: value })); + } catch { + // ignore + } +} + type ChatRuntimeStore = { settingsHydrated: boolean; params: InferenceParams; @@ -766,11 +776,11 @@ export const useChatRuntimeStore = create((set, get) => ({ setScalarSettingVersion("autoTitle", autoTitle, state.autoTitle); return { autoTitle }; }), - setHfToken: (hfToken) => - set(() => { - saveString(HF_TOKEN_KEY, hfToken); - return { hfToken }; - }), + setHfToken: (hfToken) => { + saveString(HF_TOKEN_KEY, hfToken); + set({ hfToken }); + notifyHfTokenChanged(hfToken); + }, setModelsError: (modelsError) => set({ modelsError }), setCheckpoint: (modelId, ggufVariant) => set((state) => { diff --git a/studio/frontend/src/features/hub/catalog/catalog-states.tsx b/studio/frontend/src/features/hub/catalog/catalog-states.tsx new file mode 100644 index 0000000000..e4c8554773 --- /dev/null +++ b/studio/frontend/src/features/hub/catalog/catalog-states.tsx @@ -0,0 +1,300 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { + CloudOffIcon, + CubeIcon, + FilterIcon, + RefreshIcon, + WifiDisconnected02Icon, +} from "@hugeicons/core-free-icons"; +import type { IconSvgElement } from "@hugeicons/react"; +import { HugeiconsIcon } from "@hugeicons/react"; +import { useLayoutEffect, useRef, useState } from "react"; + +export function NetworkErrorState({ + online, + message, + onRetry, + onSwitchDevice, + resourceLabel = "models", +}: { + online: boolean; + message: string; + onRetry: () => void; + onSwitchDevice?: () => void; + resourceLabel?: "models" | "datasets"; +}) { + const title = online ? "Couldn't reach Hugging Face" : "You're offline"; + const body = online + ? "The discovery feed couldn't load. Check your connection or try again." + : `Reconnect to the internet to browse ${resourceLabel} from Hugging Face.`; + const icon = online ? CloudOffIcon : WifiDisconnected02Icon; + + return ( +
+
+ +
+
+

+ {title} +

+

+ {body} +

+

{message}

+
+
+ {onSwitchDevice ? ( + + ) : null} + +
+
+ ); +} + +export function DiscoverFetchMoreState({ + scannedCount, + hasActiveFilters, + isLoadingMore, + onFetchMore, + onClearFilters, +}: { + scannedCount: number; + hasActiveFilters: boolean; + isLoadingMore: boolean; + onFetchMore: () => void; + onClearFilters: () => void; +}) { + return ( +
+
+ +
+
+

+ No matches yet +

+

+ Scanned {scannedCount.toLocaleString()} results. Load another page to + keep searching Hugging Face. +

+
+
+ {hasActiveFilters && ( + + )} + +
+
+ ); +} + +export function DiscoverFetchMoreFooter({ + scannedCount, + manualFetchAvailable, + hasActiveFilters, + isLoadingMore, + onFetchMore, +}: { + scannedCount: number; + manualFetchAvailable: boolean; + hasActiveFilters: boolean; + isLoadingMore: boolean; + onFetchMore: () => void; +}) { + return ( +
+

+ {hasActiveFilters + ? "Some results may be hidden by your filters." + : manualFetchAvailable + ? `Scanned ${scannedCount.toLocaleString()} results. Load more to continue.` + : "More results are available."} +

+ +
+ ); +} + +export function InventoryErrorState({ + isDataset, + onRetry, +}: { + isDataset: boolean; + onRetry: () => void; +}) { + return ( +
+
+ +
+
+

+ Couldn't load your library +

+

+ Something went wrong reading your downloaded{" "} + {isDataset ? "datasets" : "models"}. Check that the backend is running + and try again. +

+
+ +
+ ); +} + +export function EmptyState({ + title, + body, + icon = CubeIcon, +}: { + title: string; + body: string; + icon?: IconSvgElement; +}) { + return ( +
+
+ +
+
+

+ {title} +

+

+ {body} +

+
+
+ ); +} + +function SkeletonRow() { + return ( +
+
+
+
+
+
+
+ ); +} + +const SKELETON_ROW_ESTIMATE_PX = 56; +const MIN_SKELETON_ROWS = 4; +const MAX_SKELETON_ROWS = 24; +const DEFAULT_SKELETON_ROWS = 6; + +function clampSkeletonCount(height: number): number { + if (!Number.isFinite(height) || height <= 0) return DEFAULT_SKELETON_ROWS; + return Math.max( + MIN_SKELETON_ROWS, + Math.min(MAX_SKELETON_ROWS, Math.ceil(height / SKELETON_ROW_ESTIMATE_PX)), + ); +} + +export function SkeletonList({ count }: { count?: number }) { + const ref = useRef(null); + const [autoCount, setAutoCount] = useState(count ?? DEFAULT_SKELETON_ROWS); + const rowCount = count ?? autoCount; + + useLayoutEffect(() => { + if (count != null) return; + const container = ref.current?.parentElement; + if (!container || typeof window === "undefined") return; + + let frame: number | null = null; + const update = () => { + frame = null; + setAutoCount(clampSkeletonCount(container.clientHeight)); + }; + const schedule = () => { + if (frame !== null) return; + frame = window.requestAnimationFrame(update); + }; + schedule(); + + if (typeof ResizeObserver === "undefined") { + window.addEventListener("resize", schedule); + return () => { + if (frame !== null) window.cancelAnimationFrame(frame); + window.removeEventListener("resize", schedule); + }; + } + + const observer = new ResizeObserver(schedule); + observer.observe(container); + return () => { + if (frame !== null) window.cancelAnimationFrame(frame); + observer.disconnect(); + }; + }, [count]); + + return ( + + ); +} diff --git a/studio/frontend/src/features/hub/catalog/dataset-download-section.tsx b/studio/frontend/src/features/hub/catalog/dataset-download-section.tsx new file mode 100644 index 0000000000..253d818116 --- /dev/null +++ b/studio/frontend/src/features/hub/catalog/dataset-download-section.tsx @@ -0,0 +1,193 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { + Tooltip, + TooltipContent, + TooltipTrigger, +} from "@/components/ui/tooltip"; +import { useRepoDownload } from "../download-manager"; +import { deleteCachedDataset } from "../inventory"; +import { cn } from "@/lib/utils"; +import { TrainIcon } from "../components/train-icon"; +import { HUB_POST_DOWNLOAD_ACTIONS_VISIBLE } from "../lib/hub-feature-flags"; +import { DotTag } from "./dot-tag"; +import { PathInfoButton } from "./path-info-button"; +import { HugeiconsIcon } from "@hugeicons/react"; +import { useState } from "react"; +import { useHfTokenStore } from "../stores/hf-token-store"; +import { formatBytes } from "../lib/format"; +import { useDatasetSize } from "../hooks/use-dataset-size"; +import { + CardDivider, + CardDeleteButton, + DeleteConfirmDialog, + DownloadActionButton, + DownloadCard, +} from "./download-card"; +import { useCardDelete } from "./use-card-delete"; +import { useDownloadCardState } from "./use-download-card-state"; + +export function DatasetDownloadSection({ + repoId, + isDownloaded, + isPartial = false, + partialTransport = null, + cachePath, + knownBytes, + onTrain, + onChange, +}: { + repoId: string; + isDownloaded: boolean; + isPartial?: boolean; + partialTransport?: string | null; + cachePath?: string | null; + knownBytes?: number | null; + onTrain?: () => void; + onChange?: () => void; +}) { + const hfToken = useHfTokenStore((s) => s.token); + const [deleteOpen, setDeleteOpen] = useState(false); + const { deleting, runDelete } = useCardDelete({ + action: () => deleteCachedDataset(repoId), + resourceName: "dataset", + successMessage: () => `Deleted ${repoId}`, + onSuccess: () => { + setDeleteOpen(false); + onChange?.(); + }, + }); + + const job = useRepoDownload({ + kind: "dataset", + repoId, + autoAdopt: true, + }); + + const progress = job.progress; + const cancelling = job.cancelling; + const upstreamSize = useDatasetSize(repoId, { + enabled: + progress === null && !isDownloaded && !(knownBytes && knownBytes > 0), + token: hfToken || undefined, + }); + const upstreamBytes = + upstreamSize?.numBytesParquet ?? upstreamSize?.numBytesOriginal ?? null; + const progressBytes = + progress && progress.expectedBytes > 0 ? progress.expectedBytes : null; + const totalBytes = + progressBytes && progressBytes > 0 + ? progressBytes + : knownBytes && knownBytes > 0 + ? knownBytes + : upstreamBytes; + + const downloading = progress !== null; + const canDelete = + (isDownloaded || isPartial) && !downloading && !cancelling && !deleting; + const downloadAction = useDownloadCardState({ + job, + variant: null, + // The datasets-server size above is a parquet/original estimate, not the raw + // repo bytes snapshot_download fetches; 0 lets the backend resolve the true total. + expectedBytes: 0, + downloading, + disabled: cancelling || deleting, + isPartial, + partialTransport, + }); + + return ( + { + if (!o && !deleting) setDeleteOpen(false); + }} + title="Delete cached dataset?" + deleting={deleting} + onConfirm={() => void runDelete()} + description={ + <> + This will remove{" "} + {repoId} and + its downloaded files + {totalBytes && totalBytes > 0 + ? ` (${formatBytes(totalBytes)})` + : ""}{" "} + from disk. You can re-download it later. + + } + /> + } + > +
+ + {isDownloaded && } + {!isDownloaded && isPartial && !downloading && ( + + + + + + + + Partial download. Click to continue. + + + )} + {totalBytes && totalBytes > 0 && ( + {formatBytes(totalBytes)} + )} + +
+ {canDelete && ( + setDeleteOpen(true)} + /> + )} + {isDownloaded && cachePath && ( + + )} +
+
+ {/* Train CTA hidden until Hub->train picker ships; divider pairs with it. */} + {(!isDownloaded || downloading || HUB_POST_DOWNLOAD_ACTIONS_VISIBLE) && ( + + )} + {isDownloaded && !downloading ? ( + + ) : ( + + )} + + ); +} diff --git a/studio/frontend/src/features/hub/catalog/dot-tag.tsx b/studio/frontend/src/features/hub/catalog/dot-tag.tsx new file mode 100644 index 0000000000..77b0a73d7b --- /dev/null +++ b/studio/frontend/src/features/hub/catalog/dot-tag.tsx @@ -0,0 +1,46 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { cn } from "@/lib/utils"; + +type DotTagTone = + | "success" + | "warning" + | "danger" + | "gguf" + | "checkpoint" + | "adapter"; + +const TONE_CLASS: Record = { + success: "bg-status-success", + warning: "bg-status-warning", + danger: "bg-status-danger", + gguf: "bg-format-gguf", + checkpoint: "bg-format-checkpoint", + adapter: "bg-format-adapter", +}; + +export function DotTag({ + tone, + label, + className, +}: { + tone: DotTagTone; + label: string; + className?: string; +}) { + return ( + + + ); +} diff --git a/studio/frontend/src/features/hub/catalog/download-cancel-indicator.tsx b/studio/frontend/src/features/hub/catalog/download-cancel-indicator.tsx new file mode 100644 index 0000000000..d1b8671bfe --- /dev/null +++ b/studio/frontend/src/features/hub/catalog/download-cancel-indicator.tsx @@ -0,0 +1,25 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { Spinner } from "@/components/ui/spinner"; +import { Cancel01Icon } from "@hugeicons/core-free-icons"; +import { HugeiconsIcon } from "@hugeicons/react"; + +/** + * Inspector action-button affordance during a download: spinner that cross-fades + * to a cancel glyph on `.hub-action-btn` hover, in the same 16x16 slot so the + * percentage label never shifts. The swap is pure CSS; the component only carries + * the marker classes. + */ +export function DownloadCancelIndicator() { + return ( + + + + + ); +} diff --git a/studio/frontend/src/features/hub/catalog/download-card.tsx b/studio/frontend/src/features/hub/catalog/download-card.tsx new file mode 100644 index 0000000000..e5253526ab --- /dev/null +++ b/studio/frontend/src/features/hub/catalog/download-card.tsx @@ -0,0 +1,217 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import type { ReactNode } from "react"; +import { + AlertDialog, + AlertDialogAction, + AlertDialogCancel, + AlertDialogContent, + AlertDialogDescription, + AlertDialogFooter, + AlertDialogHeader, + AlertDialogTitle, +} from "@/components/ui/alert-dialog"; +import { Spinner } from "@/components/ui/spinner"; +import { + Tooltip, + TooltipContent, + TooltipTrigger, +} from "@/components/ui/tooltip"; +import { cn } from "@/lib/utils"; +import { Delete02Icon, Download01Icon } from "@hugeicons/core-free-icons"; +import { HugeiconsIcon } from "@hugeicons/react"; +import { + DownloadProgressBar, + type DownloadJob, + type DownloadJobProgress, +} from "../download-manager"; +import { DownloadCancelIndicator } from "./download-cancel-indicator"; +import { TransportConflictDialog } from "./transport-conflict-dialog"; +import { + downloadActionAriaLabel, + downloadActionLabel, +} from "./use-download-card-state"; + +/** + * Shared shell for every download surface (safetensors, GGUF, dataset): card frame, + * progress bar, transport-conflict dialog, plus card-specific `dialogs` and children. + */ +export function DownloadCard({ + job, + progress, + children, + dialogs, +}: { + job: DownloadJob; + progress: DownloadJobProgress | null; + children: ReactNode; + dialogs?: ReactNode; +}) { + return ( + <> +
+
{children}
+ {progress && ( + + )} +
+ + {dialogs} + + ); +} + +/** Vertical hairline that fades out on row hover, separating info from actions. */ +export function CardDivider() { + return ( +