88 lines
2.4 KiB
Python
88 lines
2.4 KiB
Python
|
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()
|