from __future__ import annotations from datetime import datetime, timezone, timedelta import sqlite3 def utcnow() -> str: return datetime.now(timezone.utc).isoformat() def future_seconds(seconds: int) -> str: return (datetime.now(timezone.utc) + timedelta(seconds=seconds)).isoformat() class Repository: def __init__(self, conn: sqlite3.Connection): self.conn = conn def upsert_candidate( self, model_id: str, source: str, origin: str, priority: int, is_bounty: bool = False, is_promote: bool = False, known_downloaded_by_others: bool = False, self_download_allowed: bool = False, max_model_len: int = 1024, target_gpu_aliases: str | None = None, notes: str | None = None, ) -> None: now = utcnow() existing = self.conn.execute( "SELECT target_gpu_aliases FROM candidates WHERE model_id = ?", (model_id,), ).fetchone() if existing and target_gpu_aliases: old_aliases = [item.strip() for item in (existing["target_gpu_aliases"] or "").split(",") if item.strip()] new_aliases = [item.strip() for item in target_gpu_aliases.split(",") if item.strip()] merged = [] for alias in old_aliases + new_aliases: if alias not in merged: merged.append(alias) target_gpu_aliases = ",".join(merged) self.conn.execute( """ INSERT INTO candidates ( model_id, source, origin, priority, is_bounty, is_promote, known_downloaded_by_others, self_download_allowed, max_model_len, target_gpu_aliases, notes, created_at, updated_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT(model_id) DO UPDATE SET source=excluded.source, origin=excluded.origin, priority=min(candidates.priority, excluded.priority), is_bounty=max(candidates.is_bounty, excluded.is_bounty), is_promote=max(candidates.is_promote, excluded.is_promote), known_downloaded_by_others=max(candidates.known_downloaded_by_others, excluded.known_downloaded_by_others), self_download_allowed=max(candidates.self_download_allowed, excluded.self_download_allowed), max_model_len=excluded.max_model_len, target_gpu_aliases=excluded.target_gpu_aliases, notes=excluded.notes, updated_at=excluded.updated_at """, ( model_id, source, origin, priority, int(is_bounty), int(is_promote), int(known_downloaded_by_others), int(self_download_allowed), max_model_len, target_gpu_aliases, notes, now, now, ), ) self.conn.commit() def list_candidates(self) -> list[sqlite3.Row]: cursor = self.conn.execute("SELECT * FROM candidates ORDER BY priority, created_at, model_id") return list(cursor.fetchall()) def insert_task_if_absent(self, model_id: str, gpu_alias: str, gpu_type: str, priority: int) -> bool: now = utcnow() cur = self.conn.execute( """ INSERT OR IGNORE INTO tasks ( model_id, gpu_alias, gpu_type, engine, priority, status, created_at, updated_at ) VALUES (?, ?, ?, 'vllm', ?, 'pending', ?, ?) """, (model_id, gpu_alias, gpu_type, priority, now, now), ) self.conn.commit() return cur.rowcount > 0 def count_budgeted_tasks(self) -> int: cursor = self.conn.execute("SELECT COUNT(*) FROM tasks WHERE status != 'skipped'") return int(cursor.fetchone()[0]) def mark_task_status( self, task_id: int, status: str, error_code: str | None = None, error_message: str | None = None, next_attempt_at: str | None = None, increment_attempt: bool = False, increment_submit: bool = False, ) -> None: now = utcnow() self.conn.execute( """ UPDATE tasks SET status = ?, last_error_code = ?, last_error_message = ?, next_attempt_at = ?, attempt_count = attempt_count + ?, submit_count = submit_count + ?, submitted_at = CASE WHEN ? = 'submitted' THEN ? ELSE submitted_at END, completed_at = CASE WHEN ? IN ('submitted', 'conflict', 'failed', 'skipped') THEN ? ELSE completed_at END, updated_at = ? WHERE id = ? """, ( status, error_code, error_message, next_attempt_at, 1 if increment_attempt else 0, 1 if increment_submit else 0, status, now, status, now, now, task_id, ), ) self.conn.commit() def due_tasks(self, statuses: tuple[str, ...], limit: int = 50) -> list[sqlite3.Row]: placeholders = ",".join("?" for _ in statuses) now = utcnow() cursor = self.conn.execute( f""" SELECT t.*, c.is_bounty, c.is_promote, c.known_downloaded_by_others, c.self_download_allowed, c.max_model_len, c.source FROM tasks t JOIN candidates c ON c.model_id = t.model_id WHERE t.status IN ({placeholders}) AND (t.next_attempt_at IS NULL OR t.next_attempt_at <= ?) ORDER BY t.priority, t.attempt_count, t.created_at, t.model_id, t.gpu_type LIMIT ? """, (*statuses, now, limit), ) return list(cursor.fetchall()) def count_active_self_downloads(self) -> int: cursor = self.conn.execute( "SELECT COUNT(*) FROM downloads WHERE active_self_slot = 1 AND status IN ('WAITING', 'RUNNING', 'unknown')" ) return int(cursor.fetchone()[0]) def upsert_download( self, model_id: str, source: str, status: str, owner: str, active_self_slot: bool = False, download_code: str | None = None, download_message: str | None = None, next_poll_at: str | None = None, ) -> None: now = utcnow() self.conn.execute( """ INSERT INTO downloads ( model_id, source, status, owner, active_self_slot, download_code, download_message, last_polled_at, next_poll_at, created_at, updated_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT(model_id) DO UPDATE SET source=excluded.source, status=excluded.status, owner=excluded.owner, active_self_slot=excluded.active_self_slot, download_code=excluded.download_code, download_message=excluded.download_message, last_polled_at=excluded.last_polled_at, next_poll_at=excluded.next_poll_at, updated_at=excluded.updated_at """, ( model_id, source, status, owner, int(active_self_slot), download_code, download_message, now, next_poll_at, now, now, ), ) self.conn.commit() def get_download(self, model_id: str): cursor = self.conn.execute("SELECT * FROM downloads WHERE model_id = ?", (model_id,)) return cursor.fetchone() def active_downloads_due(self, limit: int = 20) -> list[sqlite3.Row]: now = utcnow() cursor = self.conn.execute( """ SELECT * FROM downloads WHERE active_self_slot = 1 AND status IN ('WAITING', 'RUNNING', 'unknown') AND (next_poll_at IS NULL OR next_poll_at <= ?) ORDER BY next_poll_at, created_at LIMIT ? """, (now, limit), ) return list(cursor.fetchall()) def record_submission_attempt( self, task_id: int, model_id: str, gpu_type: str, strategy_id: str, payload_strategy_field: str, response_code: str | None, response_message: str | None, result: str, ) -> None: self.conn.execute( """ INSERT INTO submission_attempts ( task_id, model_id, gpu_type, strategy_id, payload_strategy_field, response_code, response_message, result, created_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( task_id, model_id, gpu_type, strategy_id, payload_strategy_field, response_code, response_message, result, utcnow(), ), ) self.conn.commit()