Initial vLLM agent strategy

app/scheduler/downloader.py

@@ -0,0 +1,53 @@
+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