vllm_agent_strategy/app/scheduler/submitter.py

from app.clients.modelhub import ModelHubClient
from app.domain.vllm_configs import gen_vllm_config
from app.storage.repositories import Repository, future_seconds
import logging

LOG = logging.getLogger(__name__)


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)
            LOG.info("submit %s / %s / %s -> %s (%s)", task["model_id"], task["gpu_type"], task["source"], result.result, result.message)
            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