54 lines
3.2 KiB
Python
54 lines
3.2 KiB
Python
|
from app.clients.modelhub import ModelHubClient
|
||
|
|
from app.storage.repositories import Repository, future_seconds
|
||
|
|
|
||
|
|
|
||
|
|
class Downloader:
|
||
|
|
def __init__(self, repo: Repository, client: ModelHubClient, max_self_downloads: int, poll_interval_seconds: int):
|
||
|
|
self.repo = repo
|
||
|
|
self.client = client
|
||
|
|
self.max_self_downloads = max_self_downloads
|
||
|
|
self.poll_interval_seconds = poll_interval_seconds
|
||
|
|
|
||
|
|
def poll_active_downloads(self) -> int:
|
||
|
|
changed = 0
|
||
|
|
for download in self.repo.active_downloads_due(limit=20):
|
||
|
|
status = self.client.get_download_status(download["model_id"])
|
||
|
|
active = status in {"WAITING", "RUNNING", "QUERY_FAILED"}
|
||
|
|
self.repo.upsert_download(
|
||
|
|
download["model_id"],
|
||
|
|
download["source"],
|
||
|
|
status,
|
||
|
|
"self",
|
||
|
|
active_self_slot=active,
|
||
|
|
next_poll_at=future_seconds(self.poll_interval_seconds) if active else None,
|
||
|
|
)
|
||
|
|
changed += 1
|
||
|
|
return changed
|
||
|
|
|
||
|
|
def maybe_start_self_downloads(self, limit: int = 5) -> int:
|
||
|
|
started = 0
|
||
|
|
for task in self.repo.due_tasks(("waiting_download",), limit=limit):
|
||
|
|
if self.repo.count_active_self_downloads() >= self.max_self_downloads:
|
||
|
|
self.repo.mark_task_status(task["id"], "waiting_download", next_attempt_at=future_seconds(300))
|
||
|
|
continue
|
||
|
|
sync_result = self.client.sync_from_huggingface(task["model_id"])
|
||
|
|
if sync_result.result != "success":
|
||
|
|
self.repo.mark_task_status(task["id"], "waiting_download", sync_result.code, sync_result.message, future_seconds(600), increment_attempt=True)
|
||
|
|
continue
|
||
|
|
dl = self.client.create_download_task(task["model_id"], task["source"] or "HUGGING_FACE")
|
||
|
|
if dl.result == "created":
|
||
|
|
self.repo.upsert_download(task["model_id"], task["source"] or "HUGGING_FACE", "WAITING", "self", True, dl.code, dl.message, future_seconds(self.poll_interval_seconds))
|
||
|
|
self.repo.mark_task_status(task["id"], "waiting_download", dl.code, dl.message, future_seconds(self.poll_interval_seconds), increment_attempt=True)
|
||
|
|
started += 1
|
||
|
|
elif dl.result == "already_downloaded":
|
||
|
|
self.repo.upsert_download(task["model_id"], task["source"] or "HUGGING_FACE", "SUCCESS", "others", False, dl.code, dl.message)
|
||
|
|
self.repo.mark_task_status(task["id"], "ready_to_submit", dl.code, dl.message)
|
||
|
|
elif dl.result == "running_elsewhere":
|
||
|
|
self.repo.upsert_download(task["model_id"], task["source"] or "HUGGING_FACE", "RUNNING", "others", False, dl.code, dl.message, future_seconds(self.poll_interval_seconds))
|
||
|
|
self.repo.mark_task_status(task["id"], "waiting_download", dl.code, dl.message, future_seconds(self.poll_interval_seconds), increment_attempt=True)
|
||
|
|
elif dl.result == "download_pool_full":
|
||
|
|
self.repo.mark_task_status(task["id"], "waiting_download", dl.code, dl.message, future_seconds(300), increment_attempt=True)
|
||
|
|
else:
|
||
|
|
self.repo.mark_task_status(task["id"], "failed", dl.code, dl.message, increment_attempt=True)
|
||
|
|
return started
|