45 lines
2.2 KiB
Python
45 lines
2.2 KiB
Python
|
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
|