Initial vLLM agent strategy

app/scheduler/__init__.py

@@ -0,0 +1 @@
+"""Scheduler components."""

app/scheduler/budget.py

@@ -0,0 +1,13 @@
+from app.storage.repositories import Repository
+
+
+class Budget:
+    def __init__(self, repo: Repository, max_tasks: int):
+        self.repo = repo
+        self.max_tasks = max_tasks
+
+    def remaining(self) -> int:
+        return max(0, self.max_tasks - self.repo.count_budgeted_tasks())
+
+    def can_create_task(self) -> bool:
+        return self.remaining() > 0

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

app/scheduler/loop.py

@@ -0,0 +1,118 @@
+from __future__ import annotations
+
+import logging
+from datetime import datetime, timezone
+
+from app.clients.modelhub import ModelHubClient
+from app.crawlers.download_success import DownloadSuccessCrawler
+from app.crawlers.not_adapted import NotAdaptedCrawler
+from app.scheduler.budget import Budget
+from app.scheduler.downloader import Downloader
+from app.scheduler.planner import Planner
+from app.scheduler.submitter import Submitter
+from app.settings import Settings
+from app.storage.repositories import Repository, future_seconds
+
+LOG = logging.getLogger(__name__)
+
+
+class StrategyLoop:
+    def __init__(self, settings: Settings, repo: Repository, client: ModelHubClient):
+        self.settings = settings
+        self.repo = repo
+        self.client = client
+        self.budget = Budget(repo, settings.max_user_tasks)
+        self.planner = Planner(repo, self.budget, settings.default_gpu_alias_list)
+        self.submitter = Submitter(repo, client)
+        self.downloader = Downloader(repo, client, settings.max_self_downloads, settings.download_poll_interval_seconds)
+        self.not_adapted_crawler = NotAdaptedCrawler(settings, repo, client)
+        self.download_success_crawler = DownloadSuccessCrawler(settings, repo, client)
+        self._last_crawler_refresh = None
+
+    def refresh_crawlers_if_due(self) -> dict[str, int]:
+        now = datetime.now(timezone.utc)
+        if self._last_crawler_refresh is not None:
+            elapsed = (now - self._last_crawler_refresh).total_seconds()
+            if elapsed < self.settings.crawler_refresh_seconds:
+                return {"not_adapted": 0, "download_success": 0}
+        self._last_crawler_refresh = now
+        stats = {"not_adapted": 0, "download_success": 0}
+        try:
+            stats["not_adapted"] = self.not_adapted_crawler.refresh()
+        except Exception:
+            LOG.exception("not-adapted crawler failed")
+        try:
+            stats["download_success"] = self.download_success_crawler.refresh()
+        except Exception:
+            LOG.exception("download-success crawler failed")
+        return stats
+
+    def mark_ready_tasks(self, limit: int = 50) -> int:
+        changed = 0
+        for task in self.repo.due_tasks(("pending",), limit=limit):
+            if task["is_bounty"]:
+                if self.settings.bounty_allow_direct_submit:
+                    self.repo.mark_task_status(task["id"], "ready_to_submit")
+                    changed += 1
+                else:
+                    status = self.client.get_download_status(task["model_id"])
+                    if status == "SUCCESS":
+                        self.repo.upsert_download(task["model_id"], task["source"], "SUCCESS", "others")
+                        self.repo.mark_task_status(task["id"], "ready_to_submit")
+                        changed += 1
+                    else:
+                        self.repo.mark_task_status(task["id"], "pending", "DOWNLOAD_STATUS", status, future_seconds(300))
+            elif task["is_promote"]:
+                absent = self.client.is_model_absent_from_modelhub(task["model_id"])
+                if not absent:
+                    self.repo.mark_task_status(task["id"], "ready_to_submit")
+                    changed += 1
+                else:
+                    self.repo.mark_task_status(task["id"], "pending", "NOT_IN_MODELHUB", "model not adapted yet", future_seconds(1800))
+            elif task["known_downloaded_by_others"]:
+                status = self.client.get_download_status(task["model_id"])
+                if status == "SUCCESS":
+                    self.repo.upsert_download(task["model_id"], task["source"], "SUCCESS", "others")
+                    self.repo.mark_task_status(task["id"], "ready_to_submit")
+                    changed += 1
+                elif task["self_download_allowed"]:
+                    self.repo.mark_task_status(task["id"], "waiting_download", "DOWNLOAD_STATUS", status, future_seconds(300))
+                else:
+                    self.repo.mark_task_status(task["id"], "pending", "DOWNLOAD_STATUS", status, future_seconds(300))
+            elif task["self_download_allowed"]:
+                self.repo.mark_task_status(task["id"], "waiting_download")
+            else:
+                self.repo.mark_task_status(task["id"], "skipped", "NO_READY_PATH", "candidate is not bounty/promote/downloaded/self-download")
+        return changed
+
+    def run_once(self) -> dict[str, int]:
+        crawler_stats = self.refresh_crawlers_if_due()
+        created = self.planner.materialize_tasks()
+        polled = self.downloader.poll_active_downloads()
+        released = self.submitter.release_queue_full_backoff()
+        ready = self.mark_ready_tasks()
+        submitted = self.submitter.submit_due()
+        started_downloads = 0
+        if submitted == 0 and ready == 0:
+            started_downloads = self.downloader.maybe_start_self_downloads()
+        return {
+            "crawled_not_adapted": crawler_stats["not_adapted"],
+            "crawled_download_success": crawler_stats["download_success"],
+            "created_tasks": created,
+            "polled_downloads": polled,
+            "released_backoff": released,
+            "ready_tasks": ready,
+            "submitted_tasks": submitted,
+            "started_downloads": started_downloads,
+        }
+
+    def run_forever(self, stop_event):
+        LOG.info("strategy loop started")
+        while not stop_event.is_set():
+            try:
+                stats = self.run_once()
+                LOG.info("loop stats: %s", stats)
+            except Exception:
+                LOG.exception("strategy loop iteration failed")
+            stop_event.wait(self.settings.poll_interval_seconds)
+        LOG.info("strategy loop stopped")

app/scheduler/planner.py

@@ -0,0 +1,24 @@
+from app.domain.gpu import expand_gpu_aliases, parse_aliases
+from app.domain.model_type import engine_for_model
+from app.storage.repositories import Repository
+from app.scheduler.budget import Budget
+
+
+class Planner:
+    def __init__(self, repo: Repository, budget: Budget, default_aliases: list[str]):
+        self.repo = repo
+        self.budget = budget
+        self.default_aliases = default_aliases
+
+    def materialize_tasks(self) -> int:
+        created = 0
+        for candidate in self.repo.list_candidates():
+            if engine_for_model(candidate["model_id"]) != "vllm":
+                continue
+            aliases = parse_aliases(candidate["target_gpu_aliases"], self.default_aliases)
+            for gpu_alias, gpu_type in expand_gpu_aliases(aliases):
+                if not self.budget.can_create_task():
+                    return created
+                if self.repo.insert_task_if_absent(candidate["model_id"], gpu_alias, gpu_type, candidate["priority"]):
+                    created += 1
+        return created

app/scheduler/submitter.py

@@ -0,0 +1,44 @@
+from app.clients.modelhub import ModelHubClient
+from app.domain.vllm_configs import gen_vllm_config
+from app.storage.repositories import Repository, future_seconds
+
+
+class Submitter:
+    def __init__(self, repo: Repository, client: ModelHubClient):
+        self.repo = repo
+        self.client = client
+
+    def submit_due(self, limit: int = 20) -> int:
+        submitted = 0
+        for task in self.repo.due_tasks(("ready_to_submit",), limit=limit):
+            config = gen_vllm_config(task["gpu_type"], task["model_id"], task["max_model_len"])
+            result = self.client.create_contest_task(task["model_id"], task["gpu_type"], config)
+            self.repo.record_submission_attempt(
+                task["id"],
+                task["model_id"],
+                task["gpu_type"],
+                self.client.settings.strategy_id,
+                self.client.settings.contest_task_strategy_field if self.client.settings.inject_strategy_id else "",
+                result.code,
+                result.message,
+                result.result,
+            )
+            if result.result == "success":
+                self.repo.mark_task_status(task["id"], "submitted", result.code, result.message, increment_attempt=True, increment_submit=True)
+                submitted += 1
+            elif result.result == "conflict":
+                self.repo.mark_task_status(task["id"], "conflict", result.code, result.message, increment_attempt=True)
+            elif result.result == "queue_full":
+                backoff = min(3600, 300 * (task["attempt_count"] + 1))
+                self.repo.mark_task_status(task["id"], "queue_full_backoff", result.code, result.message, future_seconds(backoff), increment_attempt=True)
+            else:
+                backoff = min(1800, 120 * (task["attempt_count"] + 1))
+                self.repo.mark_task_status(task["id"], "pending", result.code, result.message, future_seconds(backoff), increment_attempt=True)
+        return submitted
+
+    def release_queue_full_backoff(self) -> int:
+        released = 0
+        for task in self.repo.due_tasks(("queue_full_backoff",), limit=50):
+            self.repo.mark_task_status(task["id"], "ready_to_submit")
+            released += 1
+        return released