aec41d17edAdds the Studio Hub and download manager: browse Hugging Face models and datasets, download GGUF and safetensors with live progress and cancellation, and manage on-device inventory. The Hub does not require a GPU, so it is available on chat-only hosts. CI: all substantive checks pass, including the three Core jobs after unsloth-zoo#736. The two red checks are non-code flakes, a transient npm-registry DNS resolution failure in the package scan and one quantized vision-model output assertion whose sibling shards passed.
161 lines
4.6 KiB
Python
161 lines
4.6 KiB
Python
# 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
|