mirror of
https://github.com/rookiestar28/ComfyUI-OpenClaw.git
synced 2026-08-14 00:48:07 +00:00
837 lines
28 KiB
Python
837 lines
28 KiB
Python
"""
|
|
F54/F65 model search/download/import service.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
import logging
|
|
import os
|
|
import tempfile
|
|
import threading
|
|
import time
|
|
from concurrent.futures import Future, ThreadPoolExecutor
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
from typing import Any, Dict, List, Optional
|
|
from urllib.parse import urlparse
|
|
|
|
from .job_events import JobEventType, get_job_event_store
|
|
from .model_manager_catalog import (
|
|
collect_catalog_entries as _collect_catalog_entries_impl,
|
|
)
|
|
from .model_manager_catalog import (
|
|
collect_install_entries as _collect_install_entries_impl,
|
|
)
|
|
from .model_manager_catalog import list_installations as _list_installations_impl
|
|
from .model_manager_catalog import load_installations as _load_installations_impl
|
|
from .model_manager_catalog import save_installations as _save_installations_impl
|
|
from .model_manager_catalog import search_models as _search_models_impl
|
|
from .model_manager_tasks import cancel_download_task as _cancel_download_task_impl
|
|
from .model_manager_tasks import (
|
|
checkpoint_matches_task as _checkpoint_matches_task_impl,
|
|
)
|
|
from .model_manager_tasks import checkpoint_path as _checkpoint_path_impl
|
|
from .model_manager_tasks import emit as _emit_impl
|
|
from .model_manager_tasks import get_download_task as _get_download_task_impl
|
|
from .model_manager_tasks import list_download_tasks as _list_download_tasks_impl
|
|
from .model_manager_tasks import load_checkpoint as _load_checkpoint_impl
|
|
from .model_manager_tasks import load_tasks_from_disk as _load_tasks_from_disk_impl
|
|
from .model_manager_tasks import persist_tasks_locked as _persist_tasks_locked_impl
|
|
from .model_manager_tasks import progress as _progress_impl
|
|
from .model_manager_tasks import (
|
|
recover_incomplete_tasks as _recover_incomplete_tasks_impl,
|
|
)
|
|
from .model_manager_tasks import save_checkpoint as _save_checkpoint_impl
|
|
from .model_manager_tasks import set_resume_status as _set_resume_status_impl
|
|
from .model_manager_tasks import tenant_ok as _tenant_ok_impl
|
|
from .model_manager_tasks import validators_match as _validators_match_impl
|
|
from .model_manager_transfer import assert_budget as _assert_budget_impl
|
|
from .model_manager_transfer import create_download_task as _create_download_task_impl
|
|
from .model_manager_transfer import download as _download_impl
|
|
from .model_manager_transfer import (
|
|
import_downloaded_model as _import_downloaded_model_impl,
|
|
)
|
|
from .model_manager_transfer import normalize_tenant as _normalize_tenant_impl
|
|
from .model_manager_transfer import run_task as _run_task_impl
|
|
from .model_manager_transfer import (
|
|
stream_response_to_part as _stream_response_to_part_impl,
|
|
)
|
|
from .model_manager_transfer import validate_provenance as _validate_provenance_impl
|
|
from .model_manager_transfer import validate_url_policy as _validate_url_policy_impl
|
|
from .safe_io import (
|
|
STANDARD_OUTBOUND_POLICY,
|
|
SSRFError,
|
|
_build_pinned_opener,
|
|
resolve_under_root,
|
|
validate_outbound_url,
|
|
)
|
|
from .state_dir import get_state_dir
|
|
from .tenant_context import (
|
|
DEFAULT_TENANT_ID,
|
|
is_multi_tenant_enabled,
|
|
normalize_tenant_id,
|
|
)
|
|
|
|
logger = logging.getLogger("ComfyUI-OpenClaw.services.model_manager")
|
|
|
|
STATE_SUBDIR = "model_manager"
|
|
CATALOG_SUBDIR = "catalog"
|
|
STAGING_SUBDIR = "staging"
|
|
INSTALLATIONS_FILE = "installations.json"
|
|
TASKS_FILE = "download_tasks.json"
|
|
CHECKPOINT_VERSION = 1
|
|
CHECKPOINT_SUFFIX = ".checkpoint.json"
|
|
DEFAULT_MODEL_TYPE = "checkpoint"
|
|
MODEL_TYPE_TO_SUBDIR = {
|
|
"checkpoint": "checkpoints",
|
|
"lora": "loras",
|
|
"vae": "vae",
|
|
"controlnet": "controlnet",
|
|
"embedding": "embeddings",
|
|
}
|
|
ALLOWED_MODEL_EXTENSIONS = {".safetensors", ".ckpt", ".pt", ".pth", ".bin", ".onnx"}
|
|
|
|
|
|
class ModelManagerError(ValueError):
|
|
def __init__(self, code: str, detail: str, status: int = 400):
|
|
super().__init__(detail)
|
|
self.code = code
|
|
self.detail = detail
|
|
self.status = int(status)
|
|
|
|
|
|
class DownloadCancelled(RuntimeError):
|
|
pass
|
|
|
|
|
|
@dataclass
|
|
class DownloadTask:
|
|
task_id: str
|
|
model_id: str
|
|
name: str
|
|
model_type: str
|
|
source: str
|
|
source_label: str
|
|
download_url: str
|
|
destination_subdir: str
|
|
filename: str
|
|
expected_sha256: str
|
|
provenance: Dict[str, Any]
|
|
tenant_id: str
|
|
state: str = "queued"
|
|
created_at: float = field(default_factory=time.time)
|
|
updated_at: float = field(default_factory=time.time)
|
|
started_at: float = 0.0
|
|
finished_at: float = 0.0
|
|
bytes_downloaded: int = 0
|
|
total_bytes: int = 0
|
|
progress: float = 0.0
|
|
cancel_requested: bool = False
|
|
error: str = ""
|
|
staged_path: str = ""
|
|
computed_sha256: str = ""
|
|
imported: bool = False
|
|
installation_path: str = ""
|
|
installation_record_id: str = ""
|
|
resume_status: str = "not_started"
|
|
recovery_attempts: int = 0
|
|
last_checkpoint_at: float = 0.0
|
|
change_seq: int = 0
|
|
|
|
def is_terminal(self) -> bool:
|
|
return self.state in {"completed", "failed", "cancelled"}
|
|
|
|
def to_dict(self) -> Dict[str, Any]:
|
|
return {
|
|
"task_id": self.task_id,
|
|
"model_id": self.model_id,
|
|
"name": self.name,
|
|
"model_type": self.model_type,
|
|
"source": self.source,
|
|
"source_label": self.source_label,
|
|
"download_url": self.download_url,
|
|
"destination_subdir": self.destination_subdir,
|
|
"filename": self.filename,
|
|
"expected_sha256": self.expected_sha256,
|
|
"provenance": dict(self.provenance),
|
|
"tenant_id": self.tenant_id,
|
|
"state": self.state,
|
|
"created_at": self.created_at,
|
|
"updated_at": self.updated_at,
|
|
"started_at": self.started_at,
|
|
"finished_at": self.finished_at,
|
|
"bytes_downloaded": self.bytes_downloaded,
|
|
"total_bytes": self.total_bytes,
|
|
"progress": self.progress,
|
|
"cancel_requested": self.cancel_requested,
|
|
"error": self.error,
|
|
"staged_path": self.staged_path,
|
|
"computed_sha256": self.computed_sha256,
|
|
"imported": self.imported,
|
|
"installation_path": self.installation_path,
|
|
"installation_record_id": self.installation_record_id,
|
|
"resume_status": self.resume_status,
|
|
"recovery_attempts": self.recovery_attempts,
|
|
"last_checkpoint_at": self.last_checkpoint_at,
|
|
"change_seq": self.change_seq,
|
|
}
|
|
|
|
@classmethod
|
|
def from_dict(cls, payload: Dict[str, Any]) -> "DownloadTask":
|
|
if not isinstance(payload, dict):
|
|
raise ValueError("task payload must be an object")
|
|
return cls(
|
|
task_id=str(payload.get("task_id") or ""),
|
|
model_id=str(payload.get("model_id") or ""),
|
|
name=str(payload.get("name") or ""),
|
|
model_type=_norm_model_type(str(payload.get("model_type") or "")),
|
|
source=_norm_source(str(payload.get("source") or "")),
|
|
source_label=str(payload.get("source_label") or ""),
|
|
download_url=str(payload.get("download_url") or ""),
|
|
destination_subdir=str(payload.get("destination_subdir") or ""),
|
|
filename=str(payload.get("filename") or ""),
|
|
expected_sha256=str(payload.get("expected_sha256") or ""),
|
|
provenance=dict(payload.get("provenance") or {}),
|
|
tenant_id=str(payload.get("tenant_id") or DEFAULT_TENANT_ID),
|
|
state=str(payload.get("state") or "queued"),
|
|
created_at=float(payload.get("created_at") or time.time()),
|
|
updated_at=float(payload.get("updated_at") or time.time()),
|
|
started_at=float(payload.get("started_at") or 0.0),
|
|
finished_at=float(payload.get("finished_at") or 0.0),
|
|
bytes_downloaded=max(0, int(payload.get("bytes_downloaded") or 0)),
|
|
total_bytes=max(0, int(payload.get("total_bytes") or 0)),
|
|
progress=max(0.0, min(1.0, float(payload.get("progress") or 0.0))),
|
|
cancel_requested=bool(payload.get("cancel_requested")),
|
|
error=str(payload.get("error") or ""),
|
|
staged_path=str(payload.get("staged_path") or ""),
|
|
computed_sha256=str(payload.get("computed_sha256") or ""),
|
|
imported=bool(payload.get("imported")),
|
|
installation_path=str(payload.get("installation_path") or ""),
|
|
installation_record_id=str(payload.get("installation_record_id") or ""),
|
|
resume_status=str(payload.get("resume_status") or "not_started"),
|
|
recovery_attempts=max(0, int(payload.get("recovery_attempts") or 0)),
|
|
last_checkpoint_at=float(payload.get("last_checkpoint_at") or 0.0),
|
|
change_seq=max(0, int(payload.get("change_seq") or 0)),
|
|
)
|
|
|
|
|
|
def _truthy(value: str) -> bool:
|
|
return str(value or "").strip().lower() in {"1", "true", "yes", "on"}
|
|
|
|
|
|
def _norm_model_type(model_type: str) -> str:
|
|
text = str(model_type or "").strip().lower()
|
|
if not text:
|
|
return DEFAULT_MODEL_TYPE
|
|
return text if text in MODEL_TYPE_TO_SUBDIR else "other"
|
|
|
|
|
|
def _norm_source(source: str) -> str:
|
|
out = "".join(
|
|
ch
|
|
for ch in str(source or "unknown").strip().lower()
|
|
if ch.isalnum() or ch in {"_", "-", "."}
|
|
)
|
|
return out[:48] or "unknown"
|
|
|
|
|
|
def _parse_hosts(raw: str) -> set[str]:
|
|
out: set[str] = set()
|
|
for item in str(raw or "").replace(";", ",").split(","):
|
|
host = item.strip().lower().rstrip(".")
|
|
if host:
|
|
out.add(host)
|
|
return out
|
|
|
|
|
|
def _is_sha256(value: str) -> bool:
|
|
text = str(value or "").strip().lower()
|
|
return len(text) == 64 and all(ch in "0123456789abcdef" for ch in text)
|
|
|
|
|
|
def _sanitize_subdir(text: str) -> str:
|
|
parts = [
|
|
p
|
|
for p in str(text or "").replace("\\", "/").split("/")
|
|
if p not in {"", ".", ".."}
|
|
]
|
|
if not parts:
|
|
raise ModelManagerError("invalid_destination", "destination_subdir is required")
|
|
cleaned = []
|
|
for part in parts:
|
|
token = "".join(ch for ch in part if ch.isalnum() or ch in {"_", "-", "."})
|
|
if not token:
|
|
raise ModelManagerError(
|
|
"invalid_destination", f"invalid destination segment: {part!r}"
|
|
)
|
|
cleaned.append(token)
|
|
return "/".join(cleaned)
|
|
|
|
|
|
def _sanitize_filename(text: str) -> str:
|
|
clean = (
|
|
str(text or "").strip().replace("\\", "_").replace("/", "_").replace(" ", "_")
|
|
)
|
|
clean = "".join(
|
|
ch for ch in clean if ch.isalnum() or ch in {"_", "-", ".", "(", ")", "[", "]"}
|
|
)
|
|
if not clean or clean in {".", ".."}:
|
|
raise ModelManagerError("invalid_filename", "filename is invalid")
|
|
ext = Path(clean).suffix.lower()
|
|
if ext not in ALLOWED_MODEL_EXTENSIONS:
|
|
raise ModelManagerError(
|
|
"invalid_filename",
|
|
"filename extension must be one of "
|
|
+ ", ".join(sorted(ALLOWED_MODEL_EXTENSIONS)),
|
|
)
|
|
return clean[:180]
|
|
|
|
|
|
def _filename_from_url(url: str) -> str:
|
|
seg = (urlparse(url).path or "").rstrip("/").split("/")[-1]
|
|
if not seg:
|
|
raise ModelManagerError("invalid_filename", "URL has no terminal filename")
|
|
return _sanitize_filename(seg)
|
|
|
|
|
|
def _atomic_json_write(path: Path, payload: Any) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
fd, temp = tempfile.mkstemp(
|
|
prefix=f"{path.name}.tmp.", dir=str(path.parent), text=True
|
|
)
|
|
try:
|
|
with os.fdopen(fd, "w", encoding="utf-8") as fh:
|
|
json.dump(payload, fh, ensure_ascii=False, indent=2)
|
|
os.replace(temp, path)
|
|
except Exception:
|
|
try:
|
|
os.remove(temp)
|
|
except OSError:
|
|
pass
|
|
raise
|
|
|
|
|
|
def _file_sha256(path: Path) -> str:
|
|
digest = hashlib.sha256()
|
|
with open(path, "rb") as fh:
|
|
while True:
|
|
chunk = fh.read(1024 * 1024)
|
|
if not chunk:
|
|
break
|
|
digest.update(chunk)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def _safe_unlink(path: Path) -> None:
|
|
try:
|
|
path.unlink(missing_ok=True) # type: ignore[attr-defined]
|
|
except Exception:
|
|
try:
|
|
if path.exists():
|
|
path.unlink()
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
def _parse_content_range(value: str) -> tuple[int, int, int]:
|
|
text = str(value or "").strip()
|
|
if not text.lower().startswith("bytes "):
|
|
return -1, -1, -1
|
|
body = text[6:]
|
|
if "/" not in body or "-" not in body:
|
|
return -1, -1, -1
|
|
range_part, total_part = body.split("/", 1)
|
|
start_part, end_part = range_part.split("-", 1)
|
|
try:
|
|
start = int(start_part.strip())
|
|
end = int(end_part.strip())
|
|
total = int(total_part.strip()) if total_part.strip() != "*" else -1
|
|
except Exception:
|
|
return -1, -1, -1
|
|
if start < 0 or end < start:
|
|
return -1, -1, -1
|
|
if total != -1 and total <= end:
|
|
return -1, -1, -1
|
|
return start, end, total
|
|
|
|
|
|
class ModelManager:
|
|
def __init__(
|
|
self, *, state_root: Optional[Path] = None, install_root: Optional[Path] = None
|
|
):
|
|
self.state_root = Path(state_root or (Path(get_state_dir()) / STATE_SUBDIR))
|
|
self.catalog_dir = self.state_root / CATALOG_SUBDIR
|
|
self.staging_dir = self.state_root / STAGING_SUBDIR
|
|
self.installations_path = self.state_root / INSTALLATIONS_FILE
|
|
self.tasks_path = self.state_root / TASKS_FILE
|
|
self.state_root.mkdir(parents=True, exist_ok=True)
|
|
self.catalog_dir.mkdir(parents=True, exist_ok=True)
|
|
self.staging_dir.mkdir(parents=True, exist_ok=True)
|
|
root_env = (
|
|
os.environ.get("OPENCLAW_MODEL_INSTALL_ROOT")
|
|
or os.environ.get("MOLTBOT_MODEL_INSTALL_ROOT")
|
|
or ""
|
|
).strip()
|
|
# CRITICAL: keep explicit branch order. A compact inline ternary here can
|
|
# accidentally ignore injected test/runtime install_root overrides.
|
|
if install_root is not None:
|
|
resolved_install_root = Path(install_root)
|
|
elif root_env:
|
|
resolved_install_root = Path(os.path.abspath(root_env))
|
|
else:
|
|
resolved_install_root = Path(get_state_dir()) / "models"
|
|
self.install_root = resolved_install_root
|
|
self.install_root.mkdir(parents=True, exist_ok=True)
|
|
self.allow_hosts = _parse_hosts(
|
|
os.environ.get("OPENCLAW_MODEL_DOWNLOAD_ALLOW_HOSTS")
|
|
or os.environ.get("MOLTBOT_MODEL_DOWNLOAD_ALLOW_HOSTS")
|
|
or ""
|
|
)
|
|
self.allow_any_public = _truthy(
|
|
os.environ.get("OPENCLAW_MODEL_DOWNLOAD_ALLOW_ANY_PUBLIC")
|
|
or os.environ.get("MOLTBOT_MODEL_DOWNLOAD_ALLOW_ANY_PUBLIC")
|
|
or "0"
|
|
)
|
|
self.allow_loopback_hosts = _parse_hosts(
|
|
os.environ.get("OPENCLAW_MODEL_DOWNLOAD_ALLOW_LOOPBACK_HOSTS")
|
|
or os.environ.get("MOLTBOT_MODEL_DOWNLOAD_ALLOW_LOOPBACK_HOSTS")
|
|
or ""
|
|
)
|
|
self.max_workers = self._read_int(
|
|
(
|
|
"OPENCLAW_MODEL_DOWNLOAD_MAX_CONCURRENCY",
|
|
"MOLTBOT_MODEL_DOWNLOAD_MAX_CONCURRENCY",
|
|
),
|
|
2,
|
|
1,
|
|
4,
|
|
)
|
|
self.max_active = self._read_int(
|
|
("OPENCLAW_MODEL_DOWNLOAD_MAX_ACTIVE", "MOLTBOT_MODEL_DOWNLOAD_MAX_ACTIVE"),
|
|
16,
|
|
1,
|
|
128,
|
|
)
|
|
self.timeout_sec = self._read_int(
|
|
(
|
|
"OPENCLAW_MODEL_DOWNLOAD_TIMEOUT_SEC",
|
|
"MOLTBOT_MODEL_DOWNLOAD_TIMEOUT_SEC",
|
|
),
|
|
120,
|
|
5,
|
|
3600,
|
|
)
|
|
self.recovery_replay_limit = self._read_int(
|
|
(
|
|
"OPENCLAW_MODEL_DOWNLOAD_RECOVERY_REPLAY_LIMIT",
|
|
"MOLTBOT_MODEL_DOWNLOAD_RECOVERY_REPLAY_LIMIT",
|
|
),
|
|
32,
|
|
0,
|
|
256,
|
|
)
|
|
self._lock = threading.Lock()
|
|
self._tasks: Dict[str, DownloadTask] = {}
|
|
self._futures: Dict[str, Future] = {}
|
|
self._cancel_events: Dict[str, threading.Event] = {}
|
|
self._task_change_seq = 0
|
|
self._last_tasks_persist_at = 0.0
|
|
self._download_task_cls = DownloadTask
|
|
self._download_cancelled_cls = DownloadCancelled
|
|
self._ssrf_error_cls = SSRFError
|
|
self._default_tenant_id = DEFAULT_TENANT_ID
|
|
self._model_type_to_subdir = MODEL_TYPE_TO_SUBDIR
|
|
self._threading_event_factory = threading.Event
|
|
self._executor = ThreadPoolExecutor(
|
|
max_workers=self.max_workers, thread_name_prefix="openclaw-model-download"
|
|
)
|
|
self._load_tasks_from_disk()
|
|
self._recover_incomplete_tasks()
|
|
|
|
@staticmethod
|
|
def _read_int(
|
|
keys: tuple[str, ...], default: int, minimum: int, maximum: int
|
|
) -> int:
|
|
for key in keys:
|
|
raw = os.environ.get(key)
|
|
if raw is None or str(raw).strip() == "":
|
|
continue
|
|
try:
|
|
val = int(str(raw).strip())
|
|
except Exception:
|
|
return default
|
|
if val < minimum or val > maximum:
|
|
return default
|
|
return val
|
|
return default
|
|
|
|
@staticmethod
|
|
def _error(code: str, detail: str, status: int = 400) -> ModelManagerError:
|
|
return ModelManagerError(code, detail, status)
|
|
|
|
@staticmethod
|
|
def _norm_model_type(value: str) -> str:
|
|
return _norm_model_type(value)
|
|
|
|
@staticmethod
|
|
def _norm_source(value: str) -> str:
|
|
return _norm_source(value)
|
|
|
|
@staticmethod
|
|
def _is_sha256(value: str) -> bool:
|
|
return _is_sha256(value)
|
|
|
|
@staticmethod
|
|
def _sanitize_subdir(value: str) -> str:
|
|
return _sanitize_subdir(value)
|
|
|
|
@staticmethod
|
|
def _sanitize_filename(value: str) -> str:
|
|
return _sanitize_filename(value)
|
|
|
|
@staticmethod
|
|
def _filename_from_url(value: str) -> str:
|
|
return _filename_from_url(value)
|
|
|
|
@staticmethod
|
|
def _file_sha256(path: Path) -> str:
|
|
return _file_sha256(path)
|
|
|
|
@staticmethod
|
|
def _safe_unlink(path: Path) -> None:
|
|
_safe_unlink(path)
|
|
|
|
@staticmethod
|
|
def _parse_content_range(value: str) -> tuple[int, int, int]:
|
|
return _parse_content_range(value)
|
|
|
|
@staticmethod
|
|
def _is_multi_tenant_enabled() -> bool:
|
|
return is_multi_tenant_enabled()
|
|
|
|
@staticmethod
|
|
def _normalize_tenant_id(value: str) -> str:
|
|
return normalize_tenant_id(value)
|
|
|
|
@staticmethod
|
|
def _validate_outbound_download_url(
|
|
url: str,
|
|
*,
|
|
allow_hosts: Optional[set[str]],
|
|
allow_any_public_host: bool,
|
|
allow_loopback_hosts: Optional[set[str]],
|
|
) -> tuple[str, str, int, list[str]]:
|
|
return validate_outbound_url(
|
|
url,
|
|
allow_hosts=allow_hosts,
|
|
allow_any_public_host=allow_any_public_host,
|
|
allow_loopback_hosts=allow_loopback_hosts,
|
|
policy=STANDARD_OUTBOUND_POLICY,
|
|
)
|
|
|
|
@staticmethod
|
|
def _build_pinned_download_opener(pinned_ips: list[str]) -> Any:
|
|
return _build_pinned_opener(pinned_ips)
|
|
|
|
@staticmethod
|
|
def _resolve_install_target(root: str, rel_target: str) -> str:
|
|
return resolve_under_root(root, rel_target)
|
|
|
|
def _persist_tasks_locked(self, *, force: bool = False) -> None:
|
|
_persist_tasks_locked_impl(
|
|
manager=self, force=force, atomic_json_write=_atomic_json_write
|
|
)
|
|
|
|
def _load_tasks_from_disk(self) -> None:
|
|
_load_tasks_from_disk_impl(
|
|
manager=self, task_from_dict=DownloadTask.from_dict, logger=logger
|
|
)
|
|
|
|
def _recover_incomplete_tasks(self) -> None:
|
|
_recover_incomplete_tasks_impl(manager=self)
|
|
|
|
@staticmethod
|
|
def _checkpoint_path(part_path: Path) -> Path:
|
|
return _checkpoint_path_impl(
|
|
part_path=part_path, checkpoint_suffix=CHECKPOINT_SUFFIX
|
|
)
|
|
|
|
def _load_checkpoint(self, checkpoint_path: Path) -> Dict[str, Any]:
|
|
return _load_checkpoint_impl(checkpoint_path=checkpoint_path)
|
|
|
|
@staticmethod
|
|
def _checkpoint_matches_task(
|
|
task: DownloadTask, checkpoint: Dict[str, Any], partial_bytes: int
|
|
) -> bool:
|
|
return _checkpoint_matches_task_impl(
|
|
task=task,
|
|
checkpoint=checkpoint,
|
|
partial_bytes=partial_bytes,
|
|
checkpoint_version=CHECKPOINT_VERSION,
|
|
)
|
|
|
|
@staticmethod
|
|
def _validators_match(
|
|
checkpoint: Dict[str, Any], response_etag: str, response_last_modified: str
|
|
) -> bool:
|
|
return _validators_match_impl(
|
|
checkpoint=checkpoint,
|
|
response_etag=response_etag,
|
|
response_last_modified=response_last_modified,
|
|
)
|
|
|
|
def _save_checkpoint(
|
|
self,
|
|
checkpoint_path: Path,
|
|
task: DownloadTask,
|
|
*,
|
|
bytes_downloaded: int,
|
|
total_bytes: int,
|
|
etag: str,
|
|
last_modified: str,
|
|
) -> None:
|
|
_save_checkpoint_impl(
|
|
manager=self,
|
|
checkpoint_path=checkpoint_path,
|
|
task=task,
|
|
bytes_downloaded=bytes_downloaded,
|
|
total_bytes=total_bytes,
|
|
etag=etag,
|
|
last_modified=last_modified,
|
|
checkpoint_version=CHECKPOINT_VERSION,
|
|
atomic_json_write=_atomic_json_write,
|
|
)
|
|
|
|
def _set_resume_status(self, task_id: str, status: str) -> None:
|
|
_set_resume_status_impl(manager=self, task_id=task_id, status=status)
|
|
|
|
def _tenant_ok(self, record_tenant: str, request_tenant: Optional[str]) -> bool:
|
|
return _tenant_ok_impl(
|
|
record_tenant=record_tenant,
|
|
request_tenant=request_tenant,
|
|
default_tenant_id=DEFAULT_TENANT_ID,
|
|
is_multi_tenant_enabled=is_multi_tenant_enabled,
|
|
normalize_tenant_id=normalize_tenant_id,
|
|
)
|
|
|
|
def _emit(self, task: DownloadTask) -> None:
|
|
_emit_impl(
|
|
task=task,
|
|
event_type_cls=JobEventType,
|
|
event_store_getter=get_job_event_store,
|
|
)
|
|
|
|
def _bump_task_change_seq_locked(self, task: DownloadTask) -> int:
|
|
self._task_change_seq += 1
|
|
task.change_seq = self._task_change_seq
|
|
return task.change_seq
|
|
|
|
def _load_installations(self) -> List[Dict[str, Any]]:
|
|
return _load_installations_impl(installations_path=self.installations_path)
|
|
|
|
def _save_installations(self, rows: List[Dict[str, Any]]) -> None:
|
|
_save_installations_impl(
|
|
installations_path=self.installations_path,
|
|
atomic_json_write=_atomic_json_write,
|
|
rows=rows,
|
|
)
|
|
|
|
def _collect_install_entries(
|
|
self, tenant_id: Optional[str]
|
|
) -> List[Dict[str, Any]]:
|
|
return _collect_install_entries_impl(
|
|
manager=self,
|
|
tenant_id=tenant_id,
|
|
default_tenant_id=DEFAULT_TENANT_ID,
|
|
norm_model_type=_norm_model_type,
|
|
norm_source=_norm_source,
|
|
)
|
|
|
|
def _collect_catalog_entries(
|
|
self, tenant_id: Optional[str]
|
|
) -> List[Dict[str, Any]]:
|
|
return _collect_catalog_entries_impl(
|
|
manager=self,
|
|
tenant_id=tenant_id,
|
|
default_tenant_id=DEFAULT_TENANT_ID,
|
|
norm_model_type=_norm_model_type,
|
|
norm_source=_norm_source,
|
|
)
|
|
|
|
def search_models(
|
|
self,
|
|
*,
|
|
query: str = "",
|
|
source: str = "",
|
|
model_type: str = "",
|
|
installed: Optional[bool] = None,
|
|
limit: int = 50,
|
|
offset: int = 0,
|
|
tenant_id: Optional[str] = None,
|
|
) -> Dict[str, Any]:
|
|
return _search_models_impl(
|
|
manager=self,
|
|
query=query,
|
|
source=source,
|
|
model_type=model_type,
|
|
installed=installed,
|
|
limit=limit,
|
|
offset=offset,
|
|
tenant_id=tenant_id,
|
|
norm_source=_norm_source,
|
|
norm_model_type=_norm_model_type,
|
|
default_tenant_id=DEFAULT_TENANT_ID,
|
|
)
|
|
|
|
def _validate_url_policy(self, url: str) -> None:
|
|
_validate_url_policy_impl(manager=self, url=url)
|
|
|
|
def _validate_provenance(self, provenance: Dict[str, Any]) -> Dict[str, Any]:
|
|
return _validate_provenance_impl(manager=self, provenance=provenance)
|
|
|
|
def _normalize_tenant(self, tenant_id: Optional[str]) -> str:
|
|
return _normalize_tenant_impl(manager=self, tenant_id=tenant_id)
|
|
|
|
def _assert_budget(self) -> None:
|
|
_assert_budget_impl(manager=self)
|
|
|
|
def create_download_task(
|
|
self,
|
|
*,
|
|
model_id: str,
|
|
name: str,
|
|
model_type: str,
|
|
source: str,
|
|
source_label: str,
|
|
download_url: str,
|
|
expected_sha256: str,
|
|
provenance: Dict[str, Any],
|
|
destination_subdir: Optional[str] = None,
|
|
filename: Optional[str] = None,
|
|
tenant_id: Optional[str] = None,
|
|
) -> Dict[str, Any]:
|
|
return _create_download_task_impl(
|
|
manager=self,
|
|
model_id=model_id,
|
|
name=name,
|
|
model_type=model_type,
|
|
source=source,
|
|
source_label=source_label,
|
|
download_url=download_url,
|
|
expected_sha256=expected_sha256,
|
|
provenance=provenance,
|
|
destination_subdir=destination_subdir,
|
|
filename=filename,
|
|
tenant_id=tenant_id,
|
|
)
|
|
|
|
def _run_task(self, task_id: str) -> None:
|
|
_run_task_impl(manager=self, task_id=task_id)
|
|
|
|
def _download(
|
|
self, task: DownloadTask, cancel_event: threading.Event
|
|
) -> tuple[str, str]:
|
|
return _download_impl(manager=self, task=task, cancel_event=cancel_event)
|
|
|
|
def _stream_response_to_part(
|
|
self,
|
|
*,
|
|
opener: Any,
|
|
task: DownloadTask,
|
|
cancel_event: threading.Event,
|
|
part: Path,
|
|
checkpoint: Path,
|
|
digest: "hashlib._Hash",
|
|
resume_from: int,
|
|
checkpoint_data: Dict[str, Any],
|
|
) -> tuple[int, int, str, str, str]:
|
|
return _stream_response_to_part_impl(
|
|
manager=self,
|
|
opener=opener,
|
|
task=task,
|
|
cancel_event=cancel_event,
|
|
part=part,
|
|
checkpoint=checkpoint,
|
|
digest=digest,
|
|
resume_from=resume_from,
|
|
checkpoint_data=checkpoint_data,
|
|
)
|
|
|
|
def _progress(self, task_id: str, downloaded: int, total: int) -> None:
|
|
_progress_impl(
|
|
manager=self, task_id=task_id, downloaded=downloaded, total=total
|
|
)
|
|
|
|
def list_download_tasks(
|
|
self,
|
|
*,
|
|
tenant_id: Optional[str] = None,
|
|
state: str = "",
|
|
limit: int = 100,
|
|
offset: int = 0,
|
|
since_seq: Optional[int] = None,
|
|
) -> Dict[str, Any]:
|
|
return _list_download_tasks_impl(
|
|
manager=self,
|
|
tenant_id=tenant_id,
|
|
state=state,
|
|
limit=limit,
|
|
offset=offset,
|
|
since_seq=since_seq,
|
|
)
|
|
|
|
def get_download_task(
|
|
self, task_id: str, *, tenant_id: Optional[str] = None
|
|
) -> Dict[str, Any]:
|
|
return _get_download_task_impl(
|
|
manager=self, task_id=task_id, tenant_id=tenant_id
|
|
)
|
|
|
|
def cancel_download_task(
|
|
self, task_id: str, *, tenant_id: Optional[str] = None
|
|
) -> Dict[str, Any]:
|
|
return _cancel_download_task_impl(
|
|
manager=self, task_id=task_id, tenant_id=tenant_id
|
|
)
|
|
|
|
def import_downloaded_model(
|
|
self,
|
|
*,
|
|
task_id: str,
|
|
tenant_id: Optional[str] = None,
|
|
destination_subdir: Optional[str] = None,
|
|
filename: Optional[str] = None,
|
|
tags: Optional[List[str]] = None,
|
|
) -> Dict[str, Any]:
|
|
return _import_downloaded_model_impl(
|
|
manager=self,
|
|
task_id=task_id,
|
|
tenant_id=tenant_id,
|
|
destination_subdir=destination_subdir,
|
|
filename=filename,
|
|
tags=tags,
|
|
)
|
|
|
|
def list_installations(
|
|
self,
|
|
*,
|
|
tenant_id: Optional[str] = None,
|
|
model_type: str = "",
|
|
limit: int = 100,
|
|
offset: int = 0,
|
|
) -> Dict[str, Any]:
|
|
return _list_installations_impl(
|
|
manager=self,
|
|
tenant_id=tenant_id,
|
|
model_type=model_type,
|
|
limit=limit,
|
|
offset=offset,
|
|
norm_model_type=_norm_model_type,
|
|
default_tenant_id=DEFAULT_TENANT_ID,
|
|
)
|
|
|
|
|
|
model_manager = ModelManager()
|