267 lines
9.3 KiB
Python
267 lines
9.3 KiB
Python
|
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()
|