Initial vLLM agent strategy
This commit is contained in:
1
app/storage/__init__.py
Normal file
1
app/storage/__init__.py
Normal file
@@ -0,0 +1 @@
|
||||
"""SQLite storage helpers."""
|
||||
18
app/storage/db.py
Normal file
18
app/storage/db.py
Normal file
@@ -0,0 +1,18 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import sqlite3
|
||||
|
||||
from app.storage.schema import init_schema
|
||||
|
||||
|
||||
def connect(db_path: str) -> sqlite3.Connection:
|
||||
parent = os.path.dirname(os.path.abspath(db_path))
|
||||
if parent and not os.path.exists(parent):
|
||||
os.makedirs(parent, exist_ok=True)
|
||||
conn = sqlite3.connect(db_path, timeout=30)
|
||||
conn.row_factory = sqlite3.Row
|
||||
conn.execute("PRAGMA journal_mode=WAL")
|
||||
conn.execute("PRAGMA busy_timeout=30000")
|
||||
init_schema(conn)
|
||||
return conn
|
||||
266
app/storage/repositories.py
Normal file
266
app/storage/repositories.py
Normal file
@@ -0,0 +1,266 @@
|
||||
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()
|
||||
87
app/storage/schema.py
Normal file
87
app/storage/schema.py
Normal file
@@ -0,0 +1,87 @@
|
||||
SCHEMA_SQL = """
|
||||
CREATE TABLE IF NOT EXISTS candidates (
|
||||
model_id TEXT PRIMARY KEY,
|
||||
source TEXT NOT NULL DEFAULT 'HUGGING_FACE',
|
||||
origin TEXT NOT NULL,
|
||||
priority INTEGER NOT NULL,
|
||||
is_bounty INTEGER NOT NULL DEFAULT 0,
|
||||
is_promote INTEGER NOT NULL DEFAULT 0,
|
||||
known_downloaded_by_others INTEGER NOT NULL DEFAULT 0,
|
||||
self_download_allowed INTEGER NOT NULL DEFAULT 0,
|
||||
max_model_len INTEGER NOT NULL DEFAULT 1024,
|
||||
target_gpu_aliases TEXT,
|
||||
notes TEXT,
|
||||
created_at TEXT NOT NULL,
|
||||
updated_at TEXT NOT NULL
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS tasks (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
model_id TEXT NOT NULL,
|
||||
gpu_alias TEXT NOT NULL,
|
||||
gpu_type TEXT NOT NULL,
|
||||
engine TEXT NOT NULL DEFAULT 'vllm',
|
||||
priority INTEGER NOT NULL,
|
||||
status TEXT NOT NULL DEFAULT 'pending',
|
||||
attempt_count INTEGER NOT NULL DEFAULT 0,
|
||||
submit_count INTEGER NOT NULL DEFAULT 0,
|
||||
last_error_code TEXT,
|
||||
last_error_message TEXT,
|
||||
next_attempt_at TEXT,
|
||||
submitted_at TEXT,
|
||||
completed_at TEXT,
|
||||
created_at TEXT NOT NULL,
|
||||
updated_at TEXT NOT NULL,
|
||||
UNIQUE(model_id, gpu_type, engine)
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS downloads (
|
||||
model_id TEXT PRIMARY KEY,
|
||||
source TEXT NOT NULL,
|
||||
status TEXT NOT NULL DEFAULT 'unknown',
|
||||
owner TEXT NOT NULL DEFAULT 'unknown',
|
||||
active_self_slot INTEGER NOT NULL DEFAULT 0,
|
||||
download_code TEXT,
|
||||
download_message TEXT,
|
||||
waiting_since TEXT,
|
||||
last_polled_at TEXT,
|
||||
next_poll_at TEXT,
|
||||
created_at TEXT NOT NULL,
|
||||
updated_at TEXT NOT NULL
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS submission_attempts (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
task_id INTEGER NOT NULL,
|
||||
model_id TEXT NOT NULL,
|
||||
gpu_type TEXT NOT NULL,
|
||||
strategy_id TEXT,
|
||||
payload_strategy_field TEXT,
|
||||
response_code TEXT,
|
||||
response_message TEXT,
|
||||
result TEXT NOT NULL,
|
||||
created_at TEXT NOT NULL
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS seed_imports (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
file_path TEXT NOT NULL,
|
||||
file_hash TEXT,
|
||||
imported_count INTEGER NOT NULL,
|
||||
created_at TEXT NOT NULL
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_tasks_sched
|
||||
ON tasks(status, priority, next_attempt_at, created_at);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_tasks_model
|
||||
ON tasks(model_id);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_downloads_poll
|
||||
ON downloads(status, next_poll_at);
|
||||
"""
|
||||
|
||||
|
||||
def init_schema(conn):
|
||||
conn.executescript(SCHEMA_SQL)
|
||||
conn.commit()
|
||||
Reference in New Issue
Block a user