from __future__ import annotations import json import hashlib import hmac import math import re import secrets import time import uuid from contextlib import contextmanager from datetime import datetime, timedelta, timezone from pathlib import Path from typing import Any, Iterator import psycopg from app.core.config import get_settings ALL_PERMISSIONS = [ "dashboard", "fine-tune", "model-eval", "model-inference", "model-manage", "dataset", "data-process", "data-convert", "compute", "hardware", "logs", "user-settings", ] def utcnow() -> str: return datetime.now(timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z") def parse_time(value: str | None) -> datetime | None: if not value: return None return datetime.fromisoformat(value.replace("Z", "+00:00")) def json_loads(value: str | None, default: Any) -> Any: if not value: return default return json.loads(value) def json_dumps(value: Any) -> str: return json.dumps(value, ensure_ascii=False, separators=(",", ":")) def safe_float(value: Any, default: float = 0) -> float: try: return float(str(value).replace("[N/A]", "").strip() or default) except (TypeError, ValueError): return default def parse_training_metric_line(line: str) -> dict[str, float] | None: if "loss" not in line and "learning_rate" not in line: return None result: dict[str, float] = {} for key in ["loss", "grad_norm", "learning_rate", "epoch"]: match = re.search(rf"['\"]?{key}['\"]?\s*:\s*([-+]?\d+(?:\.\d+)?(?:[eE][-+]?\d+)?)", line) if match: result[key] = float(match.group(1)) return result or None def new_id(prefix: str) -> str: return f"{prefix}_{uuid.uuid4().hex[:12]}" def llama_dataset_key(dataset_id: str) -> str: safe = re.sub(r"[^0-9A-Za-z_]+", "_", dataset_id).strip("_").lower() return f"ygft_{safe or 'dataset'}" def llama_dataset_keys(dataset_key: str, file_names: list[str]) -> list[str]: if len(file_names) <= 1: return [dataset_key] return [f"{dataset_key}_{index + 1}" for index, _ in enumerate(file_names)] def llama_dataset_info(dataset_key: str, file_names: list[str], formatting: str = "alpaca") -> dict[str, Any]: result: dict[str, Any] = {} for key, file_name in zip(llama_dataset_keys(dataset_key, file_names), file_names): if formatting == "sharegpt": result[key] = { "file_name": file_name, "formatting": "sharegpt", "columns": {"messages": "messages"}, } continue result[key] = { "file_name": file_name, "formatting": "alpaca", "columns": { "prompt": "instruction", "query": "input", "response": "output", }, } return result PASSWORD_HASH_ITERATIONS = 390_000 def hash_password(password: str) -> str: salt = secrets.token_hex(16) digest = hashlib.pbkdf2_hmac("sha256", password.encode("utf-8"), salt.encode("utf-8"), PASSWORD_HASH_ITERATIONS) return f"pbkdf2_sha256${PASSWORD_HASH_ITERATIONS}${salt}${digest.hex()}" def verify_password(password: str, stored: str) -> tuple[bool, bool]: if not stored.startswith("pbkdf2_sha256$"): return hmac.compare_digest(password, stored), True try: _, iterations, salt, expected = stored.split("$", 3) digest = hashlib.pbkdf2_hmac("sha256", password.encode("utf-8"), salt.encode("utf-8"), int(iterations)).hex() return hmac.compare_digest(digest, expected), False except ValueError: return False, False def _psycopg_url(database_url: str) -> str: return database_url.replace("postgresql+psycopg://", "postgresql://") def _pg_sql(sql: str) -> str: return sql.replace("?", "%s") class PgRow(dict): def __init__(self, columns: list[str], values: tuple[Any, ...]) -> None: super().__init__(zip(columns, values)) self._values = values def __getitem__(self, key: str | int) -> Any: if isinstance(key, int): return self._values[key] return super().__getitem__(key) class PgCursor: def __init__(self, cursor: psycopg.Cursor[Any]) -> None: self.cursor = cursor def execute(self, sql: str, params: tuple[Any, ...] | list[Any] | None = None) -> "PgCursor": self.cursor.execute(_pg_sql(sql), params) return self def fetchone(self) -> PgRow | None: row = self.cursor.fetchone() if row is None: return None return PgRow(self._columns(), tuple(row)) def fetchall(self) -> list[PgRow]: columns = self._columns() return [PgRow(columns, tuple(row)) for row in self.cursor.fetchall()] def _columns(self) -> list[str]: return [col.name for col in self.cursor.description or []] class PgConnection: def __init__(self, conn: psycopg.Connection[Any]) -> None: self.conn = conn def execute(self, sql: str, params: tuple[Any, ...] | list[Any] | None = None) -> PgCursor: cursor = PgCursor(self.conn.cursor()) return cursor.execute(sql, params) def executemany(self, sql: str, params_seq: list[tuple[Any, ...]] | list[list[Any]]) -> None: with self.conn.cursor() as cursor: cursor.executemany(_pg_sql(sql), params_seq) def executescript(self, sql: str) -> None: with self.conn.cursor() as cursor: for statement in sql.split(";"): statement = statement.strip() if statement: cursor.execute(statement) def commit(self) -> None: self.conn.commit() def rollback(self) -> None: self.conn.rollback() def close(self) -> None: self.conn.close() class PlatformStore: """PostgreSQL-backed store for the first runnable platform version. This store mirrors the API-facing subset needed by the first system iteration while using the same PostgreSQL dependency as later production development. """ def __init__(self, database_url: str | None = None) -> None: settings = get_settings() self.database_url = _psycopg_url(database_url or settings.database_url) self.ensure_schema() self.ensure_seed_data() @contextmanager def connect(self) -> Iterator["PgConnection"]: raw_conn = psycopg.connect(self.database_url) conn = PgConnection(raw_conn) try: yield conn conn.commit() except Exception: conn.rollback() raise finally: conn.close() def ensure_schema(self) -> None: schema_path = Path(__file__).with_name("sql") / "001_platform_runtime.sql" with self.connect() as conn: conn.executescript(schema_path.read_text(encoding="utf-8")) user_columns = self._column_names(conn, "users") if "password" in user_columns and "password_hash" not in user_columns: conn.execute("ALTER TABLE users RENAME COLUMN password TO password_hash") self._ensure_columns( conn, "compute_nodes", { "api_version": "TEXT NOT NULL DEFAULT 'v1'", "capabilities": "TEXT NOT NULL DEFAULT '[]'", "description": "TEXT", }, ) self._ensure_columns(conn, "gpus", {"last_seen_at": "TEXT"}) self._ensure_columns(conn, "fine_tune_tasks", {"compute_job_id": "TEXT"}) self._ensure_columns( conn, "resource_replicas", { "checksum_sha256": "TEXT", "byte_size": "BIGINT NOT NULL DEFAULT 0", "last_checked_at": "TEXT", "last_error": "TEXT", }, ) def _column_names(self, conn: PgConnection, table_name: str) -> set[str]: columns = conn.execute( "SELECT column_name FROM information_schema.columns WHERE table_name=?", (table_name,), ).fetchall() return {row["column_name"] for row in columns} def _ensure_columns(self, conn: PgConnection, table_name: str, columns: dict[str, str]) -> None: existing = self._column_names(conn, table_name) for column, definition in columns.items(): if column not in existing: conn.execute(f"ALTER TABLE {table_name} ADD COLUMN {column} {definition}") def ensure_seed_data(self) -> None: with self.connect() as conn: if conn.execute("SELECT COUNT(*) FROM users").fetchone()[0] > 0: return now = utcnow() users = [ ("u_admin", "admin", "admin123", "Platform Admin", "admin", "active", ALL_PERMISSIONS, 1), ( "u_operator", "operator", "operator123", "Platform Operator", "operator", "active", [p for p in ALL_PERMISSIONS if p != "user-settings"], 0, ), ] conn.executemany( """ INSERT INTO users (id, username, password_hash, display_name, role, status, permissions, create_time, protected) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) """, [(u[0], u[1], hash_password(u[2]), u[3], u[4], u[5], json_dumps(u[6]), now, u[7]) for u in users], ) def _duration(self, start_time: str | None, end_time: str | None = None) -> str: start = parse_time(start_time) if not start: return "" end = parse_time(end_time) or datetime.now(timezone.utc) seconds = max(0, int((end - start).total_seconds())) minutes, sec = divmod(seconds, 60) hours, minutes = divmod(minutes, 60) if hours: return f"{hours}h {minutes}m {sec}s" if minutes: return f"{minutes}m {sec}s" return f"{sec}s" def refresh_runtime_state(self) -> None: if get_settings().compute_mode != "simulator": return with self.connect() as conn: rows = conn.execute( "SELECT * FROM fine_tune_tasks WHERE status IN ('syncing','queued','running')" ).fetchall() now_dt = datetime.now(timezone.utc) for row in rows: start = parse_time(row["start_time"]) if not start: continue age = max(0, int((now_dt - start).total_seconds())) if age < 4: status, progress = "syncing", 8 + age elif age < 8: status, progress = "queued", 18 + age elif age < 70: status = "running" progress = min(96, 25 + int((age - 8) / 62 * 70)) else: status, progress = "completed", 100 payload = json_loads(row["payload"], {}) payload.update( { "status": status, "progress": progress, "train_duration": self._duration(row["start_time"], utcnow() if status == "completed" else None), } ) completed_at = row["completed_at"] or (utcnow() if status == "completed" else None) conn.execute( """ UPDATE fine_tune_tasks SET status=?, progress=?, payload=?, completed_at=? WHERE id=? """, (status, progress, json_dumps(payload), completed_at, row["id"]), ) if status == "completed": self._ensure_trained_model(conn, payload) sync_rows = conn.execute( "SELECT * FROM resource_sync_jobs WHERE status IN ('pending','running')" ).fetchall() for row in sync_rows: created = parse_time(row["create_time"]) age = int((now_dt - created).total_seconds()) if created else 0 status = "completed" if age >= 6 else "running" progress = 100 if status == "completed" else min(95, 15 + age * 12) completed_at = row["completed_at"] or (utcnow() if status == "completed" else None) conn.execute( "UPDATE resource_sync_jobs SET status=?, progress=?, completed_at=? WHERE id=?", (status, progress, completed_at, row["id"]), ) def _ensure_trained_model(self, conn: PgConnection, task: dict[str, Any]) -> None: name = task.get("output_model_name") or f"{task['name']}-lora" exists = conn.execute("SELECT id FROM trained_models WHERE name=?", (name,)).fetchone() if exists: return model = conn.execute("SELECT path FROM models WHERE id=?", (task.get("base_model"),)).fetchone() output_dir = task.get("output_dir") or f"/data/yg-ft/outputs/{task['name']}" trained_model_id = new_id("tm") conn.execute( """ INSERT INTO trained_models (id, name, train_methods, base_model_path, create_time, merged, merging, merged_path) VALUES (?, ?, ?, ?, ?, ?, ?, ?) """, ( trained_model_id, name, json_dumps([{"name": task.get("train_method", "lora")}]), model["path"] if model else "", utcnow(), 0, 0, output_dir, ), ) self._upsert_model_artifact( conn, trained_model_id, "trained_model", "adapter", output_dir, 0, "", { "task_id": task.get("id"), "train_method": task.get("train_method", "lora"), "base_model": task.get("base_model"), }, task.get("compute_job_id"), ) self._insert_model_lineage( conn, "trained_model", trained_model_id, "base_model", str(task.get("base_model") or ""), "fine_tuned_from", task.get("compute_job_id"), {"task_id": task.get("id"), "output_dir": output_dir}, ) def _upsert_compute_job(self, conn: PgConnection, task: dict[str, Any], job: dict[str, Any], status: str) -> None: job_id = str(job.get("id") or task.get("compute_job_id") or task["id"]) now = utcnow() command = job.get("command") or [] command_text = " ".join(str(part) for part in command) if isinstance(command, list) else str(command or "") payload = json_dumps({**job, "task_id": task["id"]}) existing = conn.execute("SELECT id FROM compute_jobs WHERE id=?", (job_id,)).fetchone() if existing: conn.execute( """ UPDATE compute_jobs SET task_id=?, node_id=?, engine=?, status=?, command=?, output_dir=?, log_file=?, payload=?, update_time=?, completed_at=COALESCE(?, completed_at) WHERE id=? """, ( task["id"], task.get("compute_node_id"), str(task.get("engine") or job.get("engine") or "llama_factory"), status, command_text, job.get("output_dir") or task.get("output_dir"), job.get("log_file") or task.get("log_file"), payload, now, now if status in {"completed", "failed", "stopped"} else None, job_id, ), ) return conn.execute( """ INSERT INTO compute_jobs (id, task_id, node_id, engine, status, command, output_dir, log_file, payload, create_time, update_time, completed_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( job_id, task["id"], task.get("compute_node_id"), str(task.get("engine") or job.get("engine") or "llama_factory"), status, command_text, job.get("output_dir") or task.get("output_dir"), job.get("log_file") or task.get("log_file"), payload, now, now, now if status in {"completed", "failed", "stopped"} else None, ), ) def _sync_gpu_allocations(self, conn: PgConnection, task: dict[str, Any], job: dict[str, Any], status: str) -> None: terminal = status in {"completed", "failed", "stopped"} if terminal: conn.execute( "UPDATE gpu_allocations SET status='released', released_at=COALESCE(released_at, ?) WHERE task_id=? AND status IN ('allocated','running')", (utcnow(), task["id"]), ) return job_id = str(job.get("id") or task.get("compute_job_id") or task["id"]) allocation_status = "running" if status == "running" else "allocated" for gpu_index in [int(item) for item in task.get("gpus") or job.get("gpus") or []]: existing = conn.execute( "SELECT id FROM gpu_allocations WHERE task_id=? AND node_id=? AND gpu_index=? AND status IN ('allocated','running')", (task["id"], task.get("compute_node_id"), gpu_index), ).fetchone() if existing: conn.execute("UPDATE gpu_allocations SET status=?, compute_job_id=? WHERE id=?", (allocation_status, job_id, existing["id"])) continue conn.execute( """ INSERT INTO gpu_allocations (id, task_id, compute_job_id, node_id, gpu_index, status, create_time) VALUES (?, ?, ?, ?, ?, ?, ?) """, (new_id("gpu_alloc"), task["id"], job_id, task.get("compute_node_id"), gpu_index, allocation_status, utcnow()), ) def _upsert_checkpoints(self, conn: PgConnection, task_id: str, checkpoints: list[dict[str, Any]]) -> None: for item in checkpoints: path = str(item.get("path") or "") if not path: continue step = int(item.get("step") or 0) name = str(item.get("name") or Path(path).name) size_bytes = int(item.get("size_bytes") or item.get("size") or 0) existing = conn.execute("SELECT id FROM fine_tune_checkpoints WHERE task_id=? AND path=?", (task_id, path)).fetchone() if existing: conn.execute( "UPDATE fine_tune_checkpoints SET step=?, name=?, size_bytes=? WHERE id=?", (step, name, size_bytes, existing["id"]), ) continue conn.execute( """ INSERT INTO fine_tune_checkpoints (id, task_id, step, name, path, size_bytes, create_time) VALUES (?, ?, ?, ?, ?, ?, ?) """, (new_id("ckpt"), task_id, step, name, path, size_bytes, utcnow()), ) def _upsert_model_artifact( self, conn: PgConnection, model_id: str, model_kind: str, artifact_type: str, path: str, size_bytes: int = 0, checksum_sha256: str = "", metadata: dict[str, Any] | None = None, compute_job_id: str | None = None, ) -> str: existing = conn.execute( "SELECT id FROM model_artifacts WHERE model_kind=? AND model_id=? AND artifact_type=? AND path=?", (model_kind, model_id, artifact_type, path), ).fetchone() if existing: conn.execute( """ UPDATE model_artifacts SET size_bytes=?, checksum_sha256=?, metadata=?, compute_job_id=COALESCE(?, compute_job_id) WHERE id=? """, (size_bytes, checksum_sha256, json_dumps(metadata or {}), compute_job_id, existing["id"]), ) return str(existing["id"]) artifact_id = new_id("artifact") conn.execute( """ INSERT INTO model_artifacts (id, model_id, model_kind, artifact_type, path, size_bytes, checksum_sha256, metadata, compute_job_id, create_time) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( artifact_id, model_id, model_kind, artifact_type, path, size_bytes, checksum_sha256, json_dumps(metadata or {}), compute_job_id, utcnow(), ), ) return artifact_id def _insert_model_lineage( self, conn: PgConnection, child_resource_type: str, child_resource_id: str, parent_resource_type: str, parent_resource_id: str, relation_type: str, compute_job_id: str | None = None, payload: dict[str, Any] | None = None, ) -> None: if not child_resource_id or not parent_resource_id: return existing = conn.execute( """ SELECT id FROM model_lineage WHERE child_resource_type=? AND child_resource_id=? AND parent_resource_type=? AND parent_resource_id=? AND relation_type=? """, (child_resource_type, child_resource_id, parent_resource_type, parent_resource_id, relation_type), ).fetchone() if existing: conn.execute( "UPDATE model_lineage SET compute_job_id=COALESCE(?, compute_job_id), payload=? WHERE id=?", (compute_job_id, json_dumps(payload or {}), existing["id"]), ) return conn.execute( """ INSERT INTO model_lineage (id, child_resource_type, child_resource_id, parent_resource_type, parent_resource_id, relation_type, compute_job_id, payload, create_time) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( new_id("lineage"), child_resource_type, child_resource_id, parent_resource_type, parent_resource_id, relation_type, compute_job_id, json_dumps(payload or {}), utcnow(), ), ) def record_training_log_metrics(self, task_id: str, content: str) -> int: rows: list[tuple[Any, ...]] = [] for line_number, line in enumerate(content.splitlines(), start=1): metric = parse_training_metric_line(line) if not metric: continue rows.append( ( new_id("metric"), task_id, line_number, metric.get("epoch"), metric.get("loss"), metric.get("grad_norm"), metric.get("learning_rate"), line[:2000], utcnow(), ) ) with self.connect() as conn: conn.execute("DELETE FROM fine_tune_metrics WHERE task_id=?", (task_id,)) if rows: conn.executemany( """ INSERT INTO fine_tune_metrics (id, task_id, step, epoch, loss, grad_norm, learning_rate, raw, create_time) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) """, rows, ) return len(rows) def task_metrics(self, task_id: str) -> list[dict[str, Any]]: with self.connect() as conn: rows = conn.execute( """ SELECT step, epoch, loss, grad_norm, learning_rate, raw, create_time FROM fine_tune_metrics WHERE task_id=? ORDER BY step """, (task_id,), ).fetchall() return [dict(row) for row in rows] def task_checkpoints(self, task_id: str) -> list[dict[str, Any]]: with self.connect() as conn: rows = conn.execute( """ SELECT id, step, name, path, size_bytes, create_time FROM fine_tune_checkpoints WHERE task_id=? ORDER BY step, create_time """, (task_id,), ).fetchall() return [dict(row) for row in rows] def compute_job(self, job_id: str) -> dict[str, Any]: with self.connect() as conn: row = conn.execute("SELECT * FROM compute_jobs WHERE id=?", (job_id,)).fetchone() if not row: raise KeyError(job_id) payload = json_loads(row["payload"], {}) return {**dict(row), "payload": payload} def active_standalone_compute_jobs(self) -> list[dict[str, Any]]: with self.connect() as conn: rows = conn.execute( """ SELECT * FROM compute_jobs WHERE task_id IS NULL AND status IN ('queued','running') ORDER BY create_time """ ).fetchall() return [{**dict(row), "payload": json_loads(row["payload"], {})} for row in rows] def record_model_merge_job( self, node: dict[str, Any], payload: dict[str, Any], job: dict[str, Any], trained_model_id: str | None = None, ) -> dict[str, Any]: job_id = str(job.get("id") or payload.get("id") or new_id("merge")) now = utcnow() command = job.get("command") or [] command_text = " ".join(str(part) for part in command) if isinstance(command, list) else str(command or "") with self.connect() as conn: conn.execute( """ INSERT INTO compute_jobs (id, task_id, node_id, engine, status, command, output_dir, log_file, payload, create_time, update_time, completed_at) VALUES (?, NULL, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT (id) DO UPDATE SET node_id=EXCLUDED.node_id, engine=EXCLUDED.engine, status=EXCLUDED.status, command=EXCLUDED.command, output_dir=EXCLUDED.output_dir, log_file=EXCLUDED.log_file, payload=EXCLUDED.payload, update_time=EXCLUDED.update_time, completed_at=EXCLUDED.completed_at """, ( job_id, node["id"], str(payload.get("engine") or "merge"), str(job.get("status") or "queued"), command_text, job.get("output_dir") or payload.get("output_dir"), job.get("log_file"), json_dumps({**payload, "job": job}), now, now, now if str(job.get("status")) in {"completed", "failed", "stopped"} else None, ), ) if trained_model_id: conn.execute("UPDATE trained_models SET merging=1 WHERE id=? OR name=?", (trained_model_id, trained_model_id)) conn.execute( """ INSERT INTO model_export_jobs (id, trained_model_id, compute_job_id, node_id, export_type, quantization_bit, status, output_dir, payload, create_time, completed_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT (id) DO UPDATE SET status=EXCLUDED.status, output_dir=EXCLUDED.output_dir, payload=EXCLUDED.payload, completed_at=EXCLUDED.completed_at """, ( job_id, trained_model_id, job_id, node["id"], str(payload.get("engine") or "merge"), int(payload.get("export_quantization_bit", payload.get("quantization_bit", 0)) or 0), str(job.get("status") or "queued"), job.get("output_dir") or payload.get("output_dir"), json_dumps(payload), now, now if str(job.get("status")) in {"completed", "failed", "stopped"} else None, ), ) return self.compute_job(job_id) def sync_model_merge_job(self, job_id: str, job: dict[str, Any]) -> dict[str, Any]: current = self.compute_job(job_id) payload = current.get("payload") or {} job_payload = payload.get("job") if isinstance(payload.get("job"), dict) else {} merged_payload = {**payload, "job": {**job_payload, **job}} status = str(job.get("status") or current.get("status") or "queued") command = job.get("command") or current.get("command") or [] command_text = " ".join(str(part) for part in command) if isinstance(command, list) else str(command or "") output_dir = job.get("output_dir") or payload.get("output_dir") or current.get("output_dir") with self.connect() as conn: conn.execute( """ UPDATE compute_jobs SET status=?, command=?, output_dir=?, log_file=?, payload=?, update_time=?, completed_at=COALESCE(?, completed_at) WHERE id=? """, ( status, command_text, output_dir, job.get("log_file") or current.get("log_file"), json_dumps(merged_payload), utcnow(), utcnow() if status in {"completed", "failed", "stopped"} else None, job_id, ), ) trained_model_id = payload.get("trained_model_id") or payload.get("model_name") if trained_model_id and status == "completed": conn.execute( "UPDATE trained_models SET merged=1, merging=0, merged_path=? WHERE id=? OR name=?", (output_dir or "", trained_model_id, trained_model_id), ) model_row = conn.execute( "SELECT id, name FROM trained_models WHERE id=? OR name=?", (trained_model_id, trained_model_id), ).fetchone() artifact_model_id = model_row["id"] if model_row else str(trained_model_id) artifact_items = job.get("artifacts") or [] artifact_size = sum(int(item.get("size") or item.get("size_bytes") or 0) for item in artifact_items) artifact_checksums = [ str(item.get("checksum_sha256") or "") for item in artifact_items if item.get("checksum_sha256") ] checksum_sha256 = artifact_checksums[0] if len(artifact_checksums) == 1 else "" artifact_id = self._upsert_model_artifact( conn, artifact_model_id, "trained_model", "merged_model", str(output_dir or ""), artifact_size, checksum_sha256, { "export_type": payload.get("engine") or "merge", "quantization_bit": payload.get("export_quantization_bit", payload.get("quantization_bit", 0)) or 0, "artifacts": artifact_items, "checksums": artifact_checksums, }, job_id, ) self._insert_model_lineage( conn, "model_artifact", artifact_id, "trained_model", artifact_model_id, "merged_from_adapter", job_id, {"adapter_path": payload.get("adapter_name_or_path"), "output_dir": output_dir}, ) elif trained_model_id and status in {"failed", "stopped"}: conn.execute("UPDATE trained_models SET merging=0 WHERE id=? OR name=?", (trained_model_id, trained_model_id)) if trained_model_id: conn.execute( """ UPDATE model_export_jobs SET status=?, output_dir=?, payload=?, completed_at=COALESCE(?, completed_at) WHERE compute_job_id=? """, ( status, output_dir, json_dumps(merged_payload), utcnow() if status in {"completed", "failed", "stopped"} else None, job_id, ), ) return self.compute_job(job_id) def model_artifacts(self, model_id: str, model_kind: str = "trained_model") -> list[dict[str, Any]]: with self.connect() as conn: rows = conn.execute( """ SELECT * FROM model_artifacts WHERE model_id=? AND model_kind=? ORDER BY create_time DESC """, (model_id, model_kind), ).fetchall() return [{**dict(row), "metadata": json_loads(row["metadata"], {})} for row in rows] def model_artifact(self, artifact_id: str) -> dict[str, Any]: with self.connect() as conn: row = conn.execute("SELECT * FROM model_artifacts WHERE id=?", (artifact_id,)).fetchone() if not row: raise KeyError(artifact_id) return {**dict(row), "metadata": json_loads(row["metadata"], {})} def model_lineage(self, model_id: str) -> dict[str, Any]: with self.connect() as conn: parents = conn.execute( """ SELECT * FROM model_lineage WHERE child_resource_id=? ORDER BY create_time DESC """, (model_id,), ).fetchall() children = conn.execute( """ SELECT * FROM model_lineage WHERE parent_resource_id=? ORDER BY create_time DESC """, (model_id,), ).fetchall() return { "model_id": model_id, "parents": [{**dict(row), "payload": json_loads(row["payload"], {})} for row in parents], "children": [{**dict(row), "payload": json_loads(row["payload"], {})} for row in children], } def model_export_jobs(self, trained_model_id: str | None = None) -> list[dict[str, Any]]: with self.connect() as conn: if trained_model_id: rows = conn.execute( "SELECT * FROM model_export_jobs WHERE trained_model_id=? ORDER BY create_time DESC", (trained_model_id,), ).fetchall() else: rows = conn.execute("SELECT * FROM model_export_jobs ORDER BY create_time DESC").fetchall() return [{**dict(row), "payload": json_loads(row["payload"], {})} for row in rows] def users(self) -> list[dict[str, Any]]: with self.connect() as conn: rows = conn.execute("SELECT * FROM users ORDER BY create_time").fetchall() return [self._user(row) for row in rows] def login(self, username: str, password: str) -> dict[str, Any] | None: with self.connect() as conn: row = conn.execute("SELECT * FROM users WHERE username=?", (username,)).fetchone() if not row or row["status"] != "active": return None matched, legacy_plaintext = verify_password(password, row["password_hash"]) if not matched: return None last_login = utcnow() if legacy_plaintext: conn.execute( "UPDATE users SET password_hash=?, last_login=? WHERE id=?", (hash_password(password), last_login, row["id"]), ) else: conn.execute("UPDATE users SET last_login=? WHERE id=?", (last_login, row["id"])) data = self._user(row) data["last_login"] = last_login return data def create_user(self, payload: dict[str, Any]) -> dict[str, Any]: user_id = new_id("u") permissions = payload.get("permissions") or (ALL_PERMISSIONS if payload.get("role") == "admin" else ["dashboard"]) with self.connect() as conn: conn.execute( """ INSERT INTO users (id, username, password_hash, display_name, role, status, permissions, create_time, protected) VALUES (?, ?, ?, ?, ?, ?, ?, ?, 0) """, ( user_id, payload["username"], hash_password(payload.get("password", "platform123")), payload.get("display_name") or payload["username"], payload.get("role", "viewer"), payload.get("status", "active"), json_dumps(permissions), utcnow(), ), ) return self._user(conn.execute("SELECT * FROM users WHERE id=?", (user_id,)).fetchone()) def update_user(self, user_id: str, payload: dict[str, Any]) -> dict[str, Any]: with self.connect() as conn: row = conn.execute("SELECT * FROM users WHERE id=?", (user_id,)).fetchone() if not row: raise KeyError(user_id) values = { "role": payload.get("role", row["role"]), "status": payload.get("status", row["status"]), "permissions": json_dumps(payload.get("permissions", json_loads(row["permissions"], []))), } conn.execute( "UPDATE users SET role=?, status=?, permissions=? WHERE id=?", (values["role"], values["status"], values["permissions"], user_id), ) return self._user(conn.execute("SELECT * FROM users WHERE id=?", (user_id,)).fetchone()) def delete_user(self, user_id: str) -> None: with self.connect() as conn: row = conn.execute("SELECT protected FROM users WHERE id=?", (user_id,)).fetchone() if not row: raise KeyError(user_id) if row["protected"]: raise ValueError("protected user cannot be deleted") conn.execute("DELETE FROM users WHERE id=?", (user_id,)) def _user(self, row: PgRow) -> dict[str, Any]: return { "id": row["id"], "username": row["username"], "display_name": row["display_name"], "role": row["role"], "status": row["status"], "permissions": json_loads(row["permissions"], []), "create_time": row["create_time"], "last_login": row["last_login"], "protected": bool(row["protected"]), } def models(self) -> list[dict[str, Any]]: with self.connect() as conn: return [dict(row) for row in conn.execute("SELECT * FROM models ORDER BY create_time DESC").fetchall()] def model(self, model_id: str) -> dict[str, Any]: with self.connect() as conn: row = conn.execute("SELECT * FROM models WHERE id=?", (model_id,)).fetchone() if not row: raise KeyError(model_id) return dict(row) def model_by_name(self, name: str) -> dict[str, Any]: with self.connect() as conn: row = conn.execute("SELECT * FROM models WHERE name=?", (name,)).fetchone() if not row: raise KeyError(name) return dict(row) def create_model(self, payload: dict[str, Any]) -> dict[str, Any]: model_id = payload.get("id") or new_id("m") with self.connect() as conn: conn.execute( """ INSERT INTO models (id, name, type, purpose, model_source, description, path, api_url, api_key, online_model_name, create_time) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( model_id, payload["name"], payload.get("type", "LLM"), payload.get("purpose", "training"), payload.get("model_source", "local"), payload.get("description"), payload.get("path"), payload.get("api_url"), payload.get("api_key"), payload.get("online_model_name"), utcnow(), ), ) return dict(conn.execute("SELECT * FROM models WHERE id=?", (model_id,)).fetchone()) def update_model(self, model_id: str, payload: dict[str, Any]) -> dict[str, Any]: current = self.model(model_id) merged = {**current, **payload} with self.connect() as conn: conn.execute( """ UPDATE models SET name=?, type=?, purpose=?, model_source=?, description=?, path=?, api_url=?, api_key=?, online_model_name=? WHERE id=? """, ( merged["name"], merged.get("type", "LLM"), merged.get("purpose", "training"), merged.get("model_source", "local"), merged.get("description"), merged.get("path"), merged.get("api_url"), merged.get("api_key"), merged.get("online_model_name"), model_id, ), ) return dict(conn.execute("SELECT * FROM models WHERE id=?", (model_id,)).fetchone()) def delete_model(self, model_id: str) -> None: with self.connect() as conn: conn.execute("DELETE FROM models WHERE id=?", (model_id,)) def trained_models(self) -> list[dict[str, Any]]: self.refresh_runtime_state() with self.connect() as conn: rows = conn.execute("SELECT * FROM trained_models ORDER BY create_time DESC").fetchall() return [ { **dict(row), "train_methods": json_loads(row["train_methods"], []), "merged": bool(row["merged"]), "merging": bool(row["merging"]), } for row in rows ] def delete_trained_model(self, model_id: str) -> None: with self.connect() as conn: conn.execute("DELETE FROM trained_models WHERE id=? OR name=?", (model_id, model_id)) def datasets(self) -> list[dict[str, Any]]: with self.connect() as conn: rows = conn.execute("SELECT * FROM datasets ORDER BY create_time DESC").fetchall() return [self._dataset(conn, row) for row in rows] def dataset(self, dataset_id: str) -> dict[str, Any]: with self.connect() as conn: row = conn.execute("SELECT * FROM datasets WHERE id=?", (dataset_id,)).fetchone() if not row: raise KeyError(dataset_id) return self._dataset(conn, row) def _dataset(self, conn: PgConnection, row: PgRow) -> dict[str, Any]: files = conn.execute( "SELECT id, name, size, active_version_id, create_time FROM dataset_files WHERE dataset_id=? ORDER BY create_time", (row["id"],), ).fetchall() return { **dict(row), "files": [ { "id": f["id"], "name": f["name"], "size": f["size"], "active_version_id": f["active_version_id"], "create_time": f["create_time"], } for f in files ], } def create_dataset(self, payload: dict[str, Any]) -> dict[str, Any]: dataset_id = payload.get("id") or new_id("ds") with self.connect() as conn: conn.execute( """ INSERT INTO datasets (id, name, type, storage_type, source, task_id, size, count, description, create_time) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( dataset_id, payload["name"], payload.get("type", "train"), payload.get("storage_type", "local"), payload.get("source", "upload"), payload.get("task_id"), payload.get("size", "0 KB"), payload.get("count", 0), payload.get("description"), utcnow(), ), ) return self._dataset(conn, conn.execute("SELECT * FROM datasets WHERE id=?", (dataset_id,)).fetchone()) def update_dataset(self, dataset_id: str, payload: dict[str, Any]) -> dict[str, Any]: current = self.dataset(dataset_id) merged = {**current, **payload} with self.connect() as conn: conn.execute( """ UPDATE datasets SET name=?, type=?, storage_type=?, source=?, task_id=?, size=?, count=?, description=? WHERE id=? """, ( merged["name"], merged.get("type", "train"), merged.get("storage_type", "local"), merged.get("source", "upload"), merged.get("task_id"), merged.get("size", "0 KB"), merged.get("count", 0), merged.get("description"), dataset_id, ), ) return self._dataset(conn, conn.execute("SELECT * FROM datasets WHERE id=?", (dataset_id,)).fetchone()) def delete_dataset(self, dataset_id: str) -> None: with self.connect() as conn: conn.execute("DELETE FROM dataset_files WHERE dataset_id=?", (dataset_id,)) conn.execute("DELETE FROM datasets WHERE id=?", (dataset_id,)) def add_dataset_file(self, conn: PgConnection, dataset_id: str, name: str, content: str) -> dict[str, Any]: now = utcnow() file_id = new_id("file") version_id = f"{file_id}_v1" size = f"{max(1, len(content.encode('utf-8')) // 1024)} KB" conn.execute( """ INSERT INTO dataset_files (id, dataset_id, name, size, content, active_version_id, versions, create_time) VALUES (?, ?, ?, ?, ?, ?, ?, ?) """, ( file_id, dataset_id, name, size, content, version_id, json_dumps([{"id": version_id, "version": 1, "create_time": now, "description": "uploaded"}]), now, ), ) count = len([line for line in content.splitlines() if line.strip()]) conn.execute( "UPDATE datasets SET count=count+?, size=? WHERE id=?", (count, size, dataset_id), ) return {"id": file_id, "name": name, "size": size} def dataset_file(self, file_id: str) -> PgRow: with self.connect() as conn: row = conn.execute("SELECT * FROM dataset_files WHERE id=?", (file_id,)).fetchone() if not row: raise KeyError(file_id) return row def training_dataset_files(self, dataset_id: str) -> list[dict[str, Any]]: with self.connect() as conn: dataset = conn.execute("SELECT id FROM datasets WHERE id=?", (dataset_id,)).fetchone() if not dataset: raise KeyError(dataset_id) rows = conn.execute( """ SELECT id, dataset_id, name, size, content, active_version_id, create_time FROM dataset_files WHERE dataset_id=? ORDER BY create_time """, (dataset_id,), ).fetchall() return [dict(row) for row in rows] def file_versions(self, file_id: str) -> dict[str, Any]: row = self.dataset_file(file_id) versions = json_loads(row["versions"], []) return { "versions": versions, "active_version_id": row["active_version_id"], "next_version_number": len(versions) + 1, } def create_file_version(self, file_id: str, payload: dict[str, Any]) -> dict[str, Any]: with self.connect() as conn: row = conn.execute("SELECT * FROM dataset_files WHERE id=?", (file_id,)).fetchone() if not row: raise KeyError(file_id) versions = json_loads(row["versions"], []) version = { "id": f"{file_id}_v{len(versions) + 1}", "version": len(versions) + 1, "create_time": utcnow(), "description": payload.get("description", "online edit"), } versions.append(version) conn.execute( "UPDATE dataset_files SET content=?, active_version_id=?, versions=? WHERE id=?", (payload.get("content", ""), version["id"], json_dumps(versions), file_id), ) return {"version": version, "content": payload.get("content", "")} def activate_file_version(self, file_id: str, version_id: str) -> dict[str, Any]: with self.connect() as conn: row = conn.execute("SELECT * FROM dataset_files WHERE id=?", (file_id,)).fetchone() if not row: raise KeyError(file_id) versions = json_loads(row["versions"], []) version = next((item for item in versions if item["id"] == version_id), None) if not version: raise KeyError(version_id) conn.execute("UPDATE dataset_files SET active_version_id=? WHERE id=?", (version_id, file_id)) return {"version": version, "content": row["content"]} def delete_file_version(self, file_id: str, version_id: str) -> dict[str, Any]: with self.connect() as conn: row = conn.execute("SELECT * FROM dataset_files WHERE id=?", (file_id,)).fetchone() if not row: raise KeyError(file_id) versions = json_loads(row["versions"], []) if row["active_version_id"] == version_id: raise ValueError("active dataset version cannot be deleted") if len(versions) <= 1: raise ValueError("last dataset version cannot be deleted") next_versions = [item for item in versions if item["id"] != version_id] if len(next_versions) == len(versions): raise KeyError(version_id) conn.execute("UPDATE dataset_files SET versions=? WHERE id=?", (json_dumps(next_versions), file_id)) return { "versions": next_versions, "active_version_id": row["active_version_id"], "next_version_number": max(item.get("version", 0) for item in next_versions) + 1, } def tasks(self) -> list[dict[str, Any]]: self.refresh_runtime_state() with self.connect() as conn: rows = conn.execute("SELECT * FROM fine_tune_tasks ORDER BY create_time DESC").fetchall() return [self._task(row) for row in rows] def task(self, task_id: str) -> dict[str, Any]: self.refresh_runtime_state() with self.connect() as conn: row = conn.execute("SELECT * FROM fine_tune_tasks WHERE id=?", (task_id,)).fetchone() if not row: raise KeyError(task_id) return self._task(row) def _task(self, row: PgRow) -> dict[str, Any]: payload = json_loads(row["payload"], {}) payload.update( { "id": row["id"], "status": row["status"], "progress": row["progress"], "process_id": row["process_id"], "create_time": row["create_time"], "gpus": json_loads(row["gpus"], payload.get("gpus", [])), "train_duration": self._duration(row["start_time"], row["completed_at"]) if row["start_time"] else "", "compute_node_id": row["compute_node_id"], "sync_job_id": row["sync_job_id"], "compute_job_id": row.get("compute_job_id"), "completed_at": row.get("completed_at"), } ) return payload def create_task(self, payload: dict[str, Any]) -> dict[str, Any]: task_id = str(payload.get("task_id") or payload.get("id") or new_id("ft")) name = payload.get("name") or f"fine-tune-{task_id[-6:]}" base_model = payload.get("base_model") or payload.get("base_model_id") train_dataset_id = payload.get("train_dataset_id") if not base_model: raise ValueError("base_model or base_model_id is required") if not train_dataset_id: raise ValueError("train_dataset_id is required") now = utcnow() task = { "id": task_id, "name": name, "description": payload.get("description", ""), "status": "pending", "train_type": payload.get("train_type", "SFT"), "train_method": payload.get("train_method", "lora"), "engine": payload.get("engine", payload.get("training_engine", "llama_factory")), "template": payload.get("template", "qwen"), "base_model": base_model, "train_dataset_id": train_dataset_id, "auto_merge": bool(payload.get("auto_merge", False)), "output_model_name": payload.get("output_model_name") or f"{name}-lora", "gpus": payload.get("gpus") or [], "batch_size": payload.get("batch_size", 2), "learning_rate": payload.get("learning_rate", 0.0002), "n_epochs": payload.get("n_epochs", 3), "save_steps": payload.get("save_steps", 50), "lr_scheduler_type": payload.get("lr_scheduler_type", "cosine"), "max_length": payload.get("max_length", 2048), "warmup_ratio": payload.get("warmup_ratio", 0.03), "weight_decay": payload.get("weight_decay", 0.01), "lora_alpha": payload.get("lora_alpha", 16), "lora_dropout": payload.get("lora_dropout", 0.05), "lora_rank": payload.get("lora_rank", 8), "quantization_bit": payload.get("quantization_bit", 0), "export_quantized": bool(payload.get("export_quantized", False)), "quant_method": payload.get("quant_method", "bnb"), "quant_bits": payload.get("quant_bits", 4), "quant_group_size": payload.get("quant_group_size", 128), "export_format": payload.get("export_format", "safetensors"), "progress": 0, "process_id": None, "train_duration": "", "create_time": now, } with self.connect() as conn: conn.execute( """ INSERT INTO fine_tune_tasks (id, name, payload, status, progress, process_id, create_time, gpus) VALUES (?, ?, ?, 'pending', 0, NULL, ?, ?) """, (task_id, name, json_dumps(task), now, json_dumps(task["gpus"])), ) return task def update_task(self, task_id: str, payload: dict[str, Any]) -> dict[str, Any]: current = self.task(task_id) merged = {**current, **payload, "id": task_id} with self.connect() as conn: conn.execute( "UPDATE fine_tune_tasks SET name=?, payload=?, gpus=? WHERE id=?", (merged["name"], json_dumps(merged), json_dumps(merged.get("gpus", [])), task_id), ) return self.task(task_id) def start_task(self, payload: dict[str, Any]) -> dict[str, Any]: task_id = str(payload.get("task_id") or payload.get("id")) current = self.task(task_id) merged = {**current, **payload, "id": task_id, "status": "syncing", "progress": 8} selected_gpus = payload.get("gpus") or merged.get("gpus") or [0] process_id = int(43000 + (time.time() % 10000)) with self.connect() as conn: owner = f"start:{task_id}:{uuid.uuid4().hex[:8]}" if not self._acquire_scheduler_lock(conn, "compute-scheduler", owner): raise RuntimeError("compute scheduler is busy, please retry") node = self._schedule_node_locked(conn, payload) sync_job_id = new_id("sync") conn.execute( """ INSERT INTO resource_sync_jobs (id, target_node_id, resources, status, progress, create_time) VALUES (?, ?, ?, 'pending', 0, ?) """, ( sync_job_id, node["id"], json_dumps( current.get("resources") or [ {"resource_type": "model", "resource_id": current.get("base_model")}, {"resource_type": "dataset", "resource_id": current.get("train_dataset_id")}, ] ), utcnow(), ), ) conn.execute( """ UPDATE fine_tune_tasks SET payload=?, status='syncing', progress=8, process_id=?, start_time=?, compute_node_id=?, gpus=?, sync_job_id=?, compute_job_id=NULL WHERE id=? """, ( json_dumps({**merged, "process_id": process_id, "gpus": selected_gpus}), process_id, utcnow(), node["id"], json_dumps(selected_gpus), sync_job_id, task_id, ), ) self._reserve_gpu_allocations( conn, {**merged, "id": task_id, "gpus": selected_gpus, "compute_node_id": node["id"]}, node["id"], selected_gpus, ) return self.task(task_id) def reset_task_for_retry(self, task_id: str, payload: dict[str, Any] | None = None) -> dict[str, Any]: current = self.task(task_id) override = payload or {} merged = { **current, **override, "id": task_id, "status": "pending", "progress": 0, "process_id": None, "compute_job_id": None, } for runtime_key in ["failure_reason", "log_file", "artifacts"]: merged.pop(runtime_key, None) with self.connect() as conn: conn.execute( """ UPDATE fine_tune_tasks SET payload=?, status='pending', progress=0, process_id=NULL, start_time=NULL, completed_at=NULL, compute_node_id=NULL, gpus=?, sync_job_id=NULL, compute_job_id=NULL WHERE id=? """, (json_dumps(merged), json_dumps(merged.get("gpus", [])), task_id), ) return self.task(task_id) def update_task_priority(self, task_id: str, priority: str) -> dict[str, Any]: current = self.task(task_id) priority = priority if priority in {"low", "normal", "high", "urgent"} else "normal" merged = {**current, "priority": priority} with self.connect() as conn: conn.execute("UPDATE fine_tune_tasks SET payload=? WHERE id=?", (json_dumps(merged), task_id)) return self.task(task_id) def prepare_compute_job_payload(self, task_id: str, payload: dict[str, Any] | None = None) -> tuple[dict[str, Any], dict[str, Any]]: task = self.task(task_id) merged = {**task, **(payload or {}), "id": task_id} node = self.schedule_node(merged) selected_gpus = merged.get("gpus") or [0] return node, self._compute_job_payload_from_task_node(merged, node, selected_gpus) def _compute_job_payload_from_task_node( self, task: dict[str, Any], node: dict[str, Any], selected_gpus: list[int] | list[Any] | None = None, ) -> dict[str, Any]: with self.connect() as conn: model = conn.execute("SELECT * FROM models WHERE id=?", (task.get("base_model"),)).fetchone() dataset = conn.execute("SELECT * FROM datasets WHERE id=?", (task.get("train_dataset_id"),)).fetchone() files = conn.execute( "SELECT id, name, size, active_version_id, create_time FROM dataset_files WHERE dataset_id=? ORDER BY create_time", (task.get("train_dataset_id"),), ).fetchall() model_path = (model and model.get("path")) or task.get("base_model") dataset_id = str(task.get("train_dataset_id") or "") if not files: raise RuntimeError(f"dataset has no uploaded file: {dataset_id}") dataset_key = str(task.get("dataset_key") or llama_dataset_key(dataset_id)) dataset_file_names = [Path(str(row["name"] or row["id"])).name for row in files] dataset_keys = llama_dataset_keys(dataset_key, dataset_file_names) dataset_format = str(task.get("dataset_format") or (dataset and dataset.get("formatting")) or "alpaca").lower() health_detail = node.get("health_detail") or {} dataset_root = str(health_detail.get("dataset_root") or f"{node['data_root'].rstrip('/')}/datasets") output_root = str(health_detail.get("output_root") or f"{node['data_root'].rstrip('/')}/outputs") dataset_dir = f"{dataset_root.rstrip('/')}/{dataset_id}" output_name = task.get("output_model_name") or task["name"] output_dir = task.get("output_dir") or f"{output_root.rstrip('/')}/{output_name}" return { **task, "id": task["id"], "name": task["name"], "base_model": model_path, "model_name_or_path": model_path, "dataset": ",".join(dataset_keys), "dataset_key": dataset_key, "dataset_keys": dataset_keys, "dataset_display_name": (dataset and dataset.get("name")) or dataset_id, "dataset_dir": dataset_dir, "dataset_info": llama_dataset_info(dataset_key, dataset_file_names, dataset_format), "dataset_files": [ { "id": row["id"], "name": Path(str(row["name"] or row["id"])).name, "relative_path": f"{dataset_id}/{Path(str(row['name'] or row['id'])).name}", "local_path": f"{dataset_dir.rstrip('/')}/{Path(str(row['name'] or row['id'])).name}", "active_version_id": row["active_version_id"], "size": row["size"], "create_time": row["create_time"], } for row in files ], "output_dir": output_dir, "gpus": selected_gpus or task.get("gpus") or [0], "compute_node_id": node["id"], "compute_node_code": node["code"], } def build_compute_job_payload(self, task_id: str) -> tuple[dict[str, Any], dict[str, Any]]: task = self.task(task_id) node = next((item for item in self.compute_nodes() if item["id"] == task.get("compute_node_id")), None) if not node: raise RuntimeError("compute node not found") return node, self._compute_job_payload_from_task_node(task, node, task.get("gpus") or [0]) def apply_compute_job(self, task_id: str, job: dict[str, Any]) -> dict[str, Any]: status_map = { "queued": "queued", "running": "running", "completed": "completed", "failed": "failed", "stopped": "stopped", } current = self.task(task_id) status = status_map.get(str(job.get("status")), str(job.get("status") or current["status"])) progress = int(job.get("progress", current.get("progress", 0)) or 0) payload = { **current, "status": status, "progress": progress, "process_id": job.get("pid") or current.get("process_id"), "compute_job_id": job.get("id") or current.get("compute_job_id"), "output_dir": job.get("output_dir") or current.get("output_dir"), "log_file": job.get("log_file") or current.get("log_file"), "artifacts": job.get("artifacts") or current.get("artifacts") or [], } if status == "failed": payload["failure_reason"] = job.get("error") or job.get("message") or current.get("failure_reason") or "compute job failed" elif status in {"queued", "running", "completed"}: payload.pop("failure_reason", None) completed_at = utcnow() if status in {"completed", "failed", "stopped"} and not current.get("completed_at") else None with self.connect() as conn: conn.execute( """ UPDATE fine_tune_tasks SET payload=?, status=?, progress=?, process_id=?, compute_job_id=?, completed_at=COALESCE(?, completed_at) WHERE id=? """, ( json_dumps(payload), status, progress, payload.get("process_id"), payload.get("compute_job_id"), completed_at, task_id, ), ) self._upsert_compute_job(conn, payload, job, status) self._sync_gpu_allocations(conn, payload, job, status) self._upsert_checkpoints(conn, task_id, job.get("checkpoints") or []) if status == "completed": self._ensure_trained_model(conn, payload) return self.task(task_id) def running_compute_tasks(self) -> list[dict[str, Any]]: return [ task for task in self.tasks() if task.get("compute_job_id") and task.get("compute_node_id") and task["status"] in {"syncing", "queued", "running"} ] def mark_task_failed(self, task_id: str, reason: str) -> dict[str, Any]: task = self.task(task_id) task.update({"status": "failed", "progress": min(task.get("progress", 0), 99), "failure_reason": reason}) with self.connect() as conn: conn.execute( "UPDATE fine_tune_tasks SET status='failed', payload=?, completed_at=? WHERE id=?", (json_dumps(task), utcnow(), task_id), ) conn.execute( "UPDATE gpu_allocations SET status='released', released_at=COALESCE(released_at, ?) WHERE task_id=? AND status IN ('allocated','running')", (utcnow(), task_id), ) return self.task(task_id) def stop_task(self, task_id: str, status: str = "stopped") -> dict[str, Any]: task = self.task(task_id) task.update({"status": status, "progress": min(task.get("progress", 0), 99)}) with self.connect() as conn: conn.execute( "UPDATE fine_tune_tasks SET status=?, payload=?, completed_at=? WHERE id=?", (status, json_dumps(task), utcnow(), task_id), ) conn.execute( "UPDATE gpu_allocations SET status='released', released_at=COALESCE(released_at, ?) WHERE task_id=? AND status IN ('allocated','running')", (utcnow(), task_id), ) return self.task(task_id) def delete_task(self, task_id: str) -> None: with self.connect() as conn: conn.execute("DELETE FROM fine_tune_tasks WHERE id=?", (task_id,)) def _json_payload_row(self, row: PgRow) -> dict[str, Any]: payload = json_loads(row["payload"], {}) payload.update({"id": row["id"], "status": row.get("status"), "create_time": row["create_time"]}) return payload def eval_tasks(self) -> list[dict[str, Any]]: with self.connect() as conn: rows = conn.execute("SELECT * FROM eval_tasks ORDER BY create_time DESC").fetchall() return [self._json_payload_row(row) for row in rows] def eval_task(self, task_id: str) -> dict[str, Any]: with self.connect() as conn: row = conn.execute("SELECT * FROM eval_tasks WHERE id=?", (task_id,)).fetchone() if not row: raise KeyError(task_id) payload = self._json_payload_row(row) payload.setdefault("sample_count", 0) payload.setdefault("completed_count", 0) payload.setdefault("passed_count", 0) payload.setdefault("overall_score", payload.get("score") or 0) payload.setdefault("overall_score_max", 100) payload.setdefault("overall_evaluation", "") payload.setdefault("improvement_suggestions", []) payload.setdefault("dimension_summary", []) payload.setdefault("samples", []) return payload def create_eval_task(self, payload: dict[str, Any]) -> dict[str, Any]: task_id = str(payload.get("id") or payload.get("task_id") or new_id("eval")) name = str(payload.get("eval_task_name") or payload.get("name") or f"eval-{task_id[-6:]}") status = str(payload.get("status") or "pending") now = payload.get("create_time") or utcnow() data = { **payload, "id": task_id, "eval_task_name": name, "status": status, "create_time": now, "metric": payload.get("metric") or "custom", } with self.connect() as conn: model = conn.execute("SELECT name FROM models WHERE id=?", (str(payload.get("model_id")),)).fetchone() dataset = conn.execute("SELECT name FROM datasets WHERE id=?", (str(payload.get("dataset_id")),)).fetchone() if model: data.setdefault("model_name", model["name"]) if dataset: data.setdefault("dataset", dataset["name"]) conn.execute( "INSERT INTO eval_tasks (id, name, payload, status, create_time) VALUES (?, ?, ?, ?, ?)", (task_id, name, json_dumps(data), status, now), ) return self.eval_task(task_id) def delete_eval_task(self, task_id: str) -> None: with self.connect() as conn: conn.execute("DELETE FROM eval_tasks WHERE id=?", (task_id,)) def dimensions(self) -> list[dict[str, Any]]: with self.connect() as conn: rows = conn.execute("SELECT * FROM eval_dimensions ORDER BY create_time DESC").fetchall() return [ { **json_loads(row["payload"], {}), "id": row["id"], "name": row["name"], "is_active": bool(row["is_active"]), "is_default": bool(row["is_default"]), "create_time": row["create_time"], } for row in rows ] def dimension(self, dimension_id: str) -> dict[str, Any]: with self.connect() as conn: row = conn.execute("SELECT * FROM eval_dimensions WHERE id=?", (dimension_id,)).fetchone() if not row: raise KeyError(dimension_id) return { **json_loads(row["payload"], {}), "id": row["id"], "name": row["name"], "is_active": bool(row["is_active"]), "is_default": bool(row["is_default"]), "create_time": row["create_time"], } def create_dimension(self, payload: dict[str, Any]) -> dict[str, Any]: dimension_id = str(payload.get("id") or new_id("dim")) name = str(payload.get("name") or f"dimension-{dimension_id[-6:]}") now = payload.get("create_time") or utcnow() data = {**payload, "id": dimension_id, "name": name, "create_time": now} with self.connect() as conn: conn.execute( "INSERT INTO eval_dimensions (id, name, payload, is_active, is_default, create_time) VALUES (?, ?, ?, ?, ?, ?)", (dimension_id, name, json_dumps(data), 1 if data.get("is_active", True) else 0, 1 if data.get("is_default") else 0, now), ) return self.dimension(dimension_id) def update_dimension(self, dimension_id: str, payload: dict[str, Any]) -> dict[str, Any]: current = self.dimension(dimension_id) merged = {**current, **payload, "id": dimension_id} with self.connect() as conn: conn.execute( "UPDATE eval_dimensions SET name=?, payload=?, is_active=?, is_default=? WHERE id=?", ( merged["name"], json_dumps(merged), 1 if merged.get("is_active", True) else 0, 1 if merged.get("is_default") else 0, dimension_id, ), ) return self.dimension(dimension_id) def delete_dimension(self, dimension_id: str) -> None: with self.connect() as conn: conn.execute("DELETE FROM eval_dimensions WHERE id=?", (dimension_id,)) def compare_tasks(self) -> list[dict[str, Any]]: with self.connect() as conn: rows = conn.execute("SELECT * FROM compare_tasks ORDER BY create_time DESC").fetchall() return [self._json_payload_row(row) for row in rows] def compare_task(self, task_id: str) -> dict[str, Any]: with self.connect() as conn: row = conn.execute("SELECT * FROM compare_tasks WHERE id=?", (task_id,)).fetchone() if not row: raise KeyError(task_id) return self._json_payload_row(row) def create_compare_task(self, payload: dict[str, Any]) -> dict[str, Any]: task_id = str(payload.get("id") or new_id("cmp")) name = str(payload.get("name") or payload.get("model_name") or f"compare-{task_id[-6:]}") status = str(payload.get("status") or "pending") now = payload.get("create_time") or utcnow() data = {**payload, "id": task_id, "name": name, "model_name": payload.get("model_name") or name, "status": status, "create_time": now} data.setdefault("load_status", json_dumps({"loaded_models": []})) with self.connect() as conn: conn.execute( "INSERT INTO compare_tasks (id, name, payload, status, create_time) VALUES (?, ?, ?, ?, ?)", (task_id, name, json_dumps(data), status, now), ) return self.compare_task(task_id) def update_compare_task(self, task_id: str, payload: dict[str, Any]) -> dict[str, Any]: current = self.compare_task(task_id) merged = {**current, **payload, "id": task_id} status = str(merged.get("status") or current.get("status") or "pending") with self.connect() as conn: conn.execute( "UPDATE compare_tasks SET name=?, payload=?, status=? WHERE id=?", (merged.get("name") or merged.get("model_name") or task_id, json_dumps(merged), status, task_id), ) return self.compare_task(task_id) def delete_compare_task(self, task_id: str) -> None: with self.connect() as conn: conn.execute("DELETE FROM compare_tasks WHERE id=?", (task_id,)) def _acquire_scheduler_lock( self, conn: PgConnection, lock_key: str, owner: str, ttl_seconds: int = 30, ) -> bool: now = utcnow() expires_at = (datetime.now(timezone.utc) + timedelta(seconds=ttl_seconds)).replace(microsecond=0).isoformat().replace("+00:00", "Z") row = conn.execute("SELECT * FROM scheduler_locks WHERE lock_key=? FOR UPDATE", (lock_key,)).fetchone() if not row: conn.execute( "INSERT INTO scheduler_locks (lock_key, owner, expires_at, create_time, update_time) VALUES (?, ?, ?, ?, ?)", (lock_key, owner, expires_at, now, now), ) return True if str(row["owner"]) == owner or str(row["expires_at"]) <= now: conn.execute( "UPDATE scheduler_locks SET owner=?, expires_at=?, update_time=? WHERE lock_key=?", (owner, expires_at, now, lock_key), ) return True return False def _compute_nodes_locked(self, conn: PgConnection) -> list[dict[str, Any]]: running = conn.execute( """ SELECT node_id, COUNT(*) AS cnt FROM compute_jobs WHERE status IN ('queued','running') GROUP BY node_id """ ).fetchall() running_map = {r["node_id"]: r["cnt"] for r in running} task_running = conn.execute( """ SELECT compute_node_id, COUNT(*) AS cnt FROM fine_tune_tasks WHERE status IN ('syncing','queued','running') GROUP BY compute_node_id """ ).fetchall() for row in task_running: running_map[row["compute_node_id"]] = max(running_map.get(row["compute_node_id"], 0), row["cnt"]) rows = conn.execute("SELECT * FROM compute_nodes ORDER BY scheduler_weight DESC, code").fetchall() return [ { **dict(row), "enabled": bool(row["enabled"]), "tags": json_loads(row["tags"], []), "capabilities": json_loads(row.get("capabilities"), []), "health_detail": json_loads(row["health_detail"], {}), "current_running_jobs": running_map.get(row["id"], 0), } for row in rows ] def _active_gpu_indexes(self, conn: PgConnection, node_id: str) -> set[int]: rows = conn.execute( "SELECT gpu_index FROM gpu_allocations WHERE node_id=? AND status IN ('allocated','running')", (node_id,), ).fetchall() return {int(row["gpu_index"]) for row in rows} def _schedule_node_locked(self, conn: PgConnection, payload: dict[str, Any]) -> dict[str, Any]: requested = payload.get("requested_node_id") or payload.get("compute_node_id") requested_gpus = [int(item) for item in payload.get("gpus") or []] nodes = self._compute_nodes_locked(conn) candidates = [ n for n in nodes if n["enabled"] and n["scheduler_status"] == "online" and n["current_running_jobs"] < n["max_parallel_jobs"] ] if requested_gpus: candidates = [ node for node in candidates if not set(requested_gpus).intersection(self._active_gpu_indexes(conn, node["id"])) ] if requested: selected = next((n for n in candidates if n["id"] == requested), None) if selected: return selected if not candidates: if not nodes: raise RuntimeError("no available compute node: no compute node configured") reasons = [] for node in nodes: if not node["enabled"]: reason = "disabled" elif node["scheduler_status"] != "online": reason = f"status={node['scheduler_status']}" elif node["current_running_jobs"] >= node["max_parallel_jobs"]: reason = f"capacity full {node['current_running_jobs']}/{node['max_parallel_jobs']}" else: reason = "not selected" reasons.append(f"{node['code']}({reason})") raise RuntimeError(f"no available compute node: {', '.join(reasons)}") return sorted(candidates, key=lambda n: (-n["scheduler_weight"], n["current_running_jobs"], n["code"]))[0] def schedule_node(self, payload: dict[str, Any]) -> dict[str, Any]: with self.connect() as conn: owner = f"schedule:{uuid.uuid4().hex[:8]}" if not self._acquire_scheduler_lock(conn, "compute-scheduler", owner): raise RuntimeError("compute scheduler is busy, please retry") return self._schedule_node_locked(conn, payload) def _reserve_gpu_allocations( self, conn: PgConnection, task: dict[str, Any], node_id: str, gpus: list[int] | list[Any], ) -> None: for gpu_index in [int(item) for item in gpus or []]: existing = conn.execute( "SELECT id FROM gpu_allocations WHERE task_id=? AND node_id=? AND gpu_index=? AND status IN ('allocated','running')", (task["id"], node_id, gpu_index), ).fetchone() if existing: continue conn.execute( """ INSERT INTO gpu_allocations (id, task_id, compute_job_id, node_id, gpu_index, status, create_time) VALUES (?, ?, ?, ?, ?, 'allocated', ?) """, (new_id("gpu_alloc"), task["id"], task.get("compute_job_id"), node_id, gpu_index, utcnow()), ) def create_sync_job(self, node_id: str, task: dict[str, Any]) -> str: sync_id = new_id("sync") with self.connect() as conn: conn.execute( """ INSERT INTO resource_sync_jobs (id, target_node_id, resources, status, progress, create_time) VALUES (?, ?, ?, 'pending', 0, ?) """, ( sync_id, node_id, json_dumps( task.get("resources") or [ {"resource_type": "model", "resource_id": task.get("base_model")}, {"resource_type": "dataset", "resource_id": task.get("train_dataset_id")}, ] ), utcnow(), ), ) return sync_id def update_sync_job(self, sync_id: str, status: str, progress: int, completed: bool = False) -> dict[str, Any]: with self.connect() as conn: conn.execute( "UPDATE resource_sync_jobs SET status=?, progress=?, completed_at=COALESCE(?, completed_at) WHERE id=?", (status, progress, utcnow() if completed else None, sync_id), ) return self.sync_job(sync_id) def upsert_resource_replica( self, node_id: str, resource_type: str, resource_id: str, local_path: str, status: str = "available", sync_status: str = "synced", ) -> dict[str, Any]: with self.connect() as conn: row = conn.execute( "SELECT * FROM resource_replicas WHERE node_id=? AND resource_type=? AND resource_id=?", (node_id, resource_type, resource_id), ).fetchone() if row: conn.execute( "UPDATE resource_replicas SET local_path=?, status=?, sync_status=? WHERE id=?", (local_path, status, sync_status, row["id"]), ) replica_id = row["id"] else: replica_id = new_id("replica") conn.execute( """ INSERT INTO resource_replicas (id, node_id, resource_type, resource_id, local_path, status, sync_status, create_time) VALUES (?, ?, ?, ?, ?, ?, ?, ?) """, (replica_id, node_id, resource_type, resource_id, local_path, status, sync_status, utcnow()), ) return dict(conn.execute("SELECT * FROM resource_replicas WHERE id=?", (replica_id,)).fetchone()) def update_resource_replica_check( self, replica_id: str, exists: bool, byte_size: int = 0, error: str = "", ) -> dict[str, Any]: status = "available" if exists else "missing" sync_status = "synced" if exists else "drifted" with self.connect() as conn: conn.execute( """ UPDATE resource_replicas SET status=?, sync_status=?, byte_size=?, last_checked_at=?, last_error=? WHERE id=? """, (status, sync_status, byte_size, utcnow(), error, replica_id), ) row = conn.execute("SELECT * FROM resource_replicas WHERE id=?", (replica_id,)).fetchone() if not row: raise KeyError(replica_id) return dict(row) def resource_replicas_by_ids(self, replica_ids: list[str]) -> list[dict[str, Any]]: if not replica_ids: return [] with self.connect() as conn: result: list[dict[str, Any]] = [] for replica_id in replica_ids: row = conn.execute("SELECT * FROM resource_replicas WHERE id=?", (replica_id,)).fetchone() if row: result.append(dict(row)) return result def update_resource_replica_sync_result( self, replica_id: str, success: bool, local_path: str | None = None, byte_size: int = 0, checksum_sha256: str = "", error: str = "", ) -> dict[str, Any]: with self.connect() as conn: conn.execute( """ UPDATE resource_replicas SET local_path=COALESCE(?, local_path), status=?, sync_status=?, byte_size=?, checksum_sha256=COALESCE(NULLIF(?, ''), checksum_sha256), last_checked_at=?, last_error=? WHERE id=? """, ( local_path, "available" if success else "missing", "synced" if success else "failed", byte_size, checksum_sha256, utcnow(), error, replica_id, ), ) row = conn.execute("SELECT * FROM resource_replicas WHERE id=?", (replica_id,)).fetchone() if not row: raise KeyError(replica_id) return dict(row) def mark_resource_replica_repair_pending(self, replica_ids: list[str]) -> list[dict[str, Any]]: if not replica_ids: return [] updated: list[dict[str, Any]] = [] with self.connect() as conn: for replica_id in replica_ids: conn.execute( """ UPDATE resource_replicas SET sync_status='repair_pending', last_checked_at=?, last_error='' WHERE id=? """, (utcnow(), replica_id), ) row = conn.execute("SELECT * FROM resource_replicas WHERE id=?", (replica_id,)).fetchone() if row: updated.append(dict(row)) return updated def progress(self, task_id: str) -> dict[str, Any]: task = self.task(task_id) status = task.get("status", "pending") labels = { "pending": "waiting for start", "syncing": "syncing model and dataset to compute node", "queued": "waiting for GPU slot", "running": "training with LLaMA-Factory", "completed": "training completed", "failed": "training stopped", "stopped": "training stopped", } progress = int(task.get("progress", 0) or 0) eta = "--" if status in {"completed", "failed"} else f"{max(1, math.ceil((100 - progress) / 10))} min" return { "status": status, "progress": progress, "step": labels.get(status, status), "speed": task.get("train_speed") or "--", "eta": eta, } def compute_nodes(self) -> list[dict[str, Any]]: self.refresh_runtime_state() with self.connect() as conn: running = conn.execute( "SELECT compute_node_id, COUNT(*) AS cnt FROM fine_tune_tasks WHERE status IN ('syncing','queued','running') GROUP BY compute_node_id" ).fetchall() running_map = {r["compute_node_id"]: r["cnt"] for r in running} rows = conn.execute("SELECT * FROM compute_nodes ORDER BY scheduler_weight DESC, code").fetchall() return [ { **dict(row), "enabled": bool(row["enabled"]), "tags": json_loads(row["tags"], []), "capabilities": json_loads(row.get("capabilities"), []), "health_detail": json_loads(row["health_detail"], {}), "current_running_jobs": running_map.get(row["id"], 0), } for row in rows ] def _normalize_tags(self, value: Any) -> list[str]: if isinstance(value, str): parts = value.replace("，", ",").split(",") return [item.strip() for item in parts if item.strip()] if isinstance(value, list): return [str(item).strip() for item in value if str(item).strip()] return [] def _normalize_compute_node_payload(self, payload: dict[str, Any], current: dict[str, Any] | None = None) -> dict[str, Any]: merged = {**(current or {}), **payload} api_base_url = str(merged.get("api_base_url") or "").rstrip("/") if not api_base_url: raise ValueError("api_base_url is required") file_gateway_url = str(merged.get("file_gateway_url") or api_base_url).rstrip("/") weight = max(0, min(1000, int(merged.get("scheduler_weight", 100)))) max_jobs = max(1, int(merged.get("max_parallel_jobs", 1))) return { **merged, "code": str(merged.get("code") or "").strip(), "name": str(merged.get("name") or merged.get("code") or "").strip(), "api_base_url": api_base_url, "file_gateway_url": file_gateway_url, "enabled": bool(merged.get("enabled", True)), "scheduler_status": str(merged.get("scheduler_status") or "offline"), "scheduler_weight": weight, "tags": self._normalize_tags(merged.get("tags")), "gpu_count": max(0, int(merged.get("gpu_count", 0) or 0)), "max_parallel_jobs": max_jobs, "data_root": str(merged.get("data_root") or "/data/yg-ft"), "model_root": str(merged.get("model_root") or "/data/yg-ft/models"), "log_root": str(merged.get("log_root") or "/opt/yg-ft/logs/training"), "api_version": str(merged.get("api_version") or "v1"), "capabilities": merged.get("capabilities") or [], "description": merged.get("description") or "", "health_detail": merged.get("health_detail") or {"status": "registered"}, } def update_compute_node(self, node_id: str, payload: dict[str, Any]) -> dict[str, Any]: current = next((n for n in self.compute_nodes() if n["id"] == node_id), None) if not current: raise KeyError(node_id) merged = self._normalize_compute_node_payload(payload, current) with self.connect() as conn: conn.execute( """ UPDATE compute_nodes SET name=?, api_base_url=?, file_gateway_url=?, enabled=?, scheduler_status=?, scheduler_weight=?, tags=?, max_parallel_jobs=?, data_root=?, model_root=?, log_root=?, api_version=?, capabilities=?, description=?, last_health_check_at=?, health_detail=? WHERE id=? """, ( merged["name"], merged["api_base_url"], merged["file_gateway_url"], 1 if merged["enabled"] else 0, merged["scheduler_status"], merged["scheduler_weight"], json_dumps(merged["tags"]), merged["max_parallel_jobs"], merged["data_root"], merged["model_root"], merged["log_root"], merged["api_version"], json_dumps(merged["capabilities"]), merged["description"], payload.get("last_health_check_at") or current.get("last_health_check_at"), json_dumps(merged["health_detail"]), node_id, ), ) return next(n for n in self.compute_nodes() if n["id"] == node_id) def create_compute_node(self, payload: dict[str, Any]) -> dict[str, Any]: payload = self._normalize_compute_node_payload(payload) if not payload["code"]: raise ValueError("code is required") node_id = payload.get("id") or new_id("node") now = utcnow() with self.connect() as conn: conn.execute( """ INSERT INTO compute_nodes (id, code, name, api_base_url, file_gateway_url, enabled, scheduler_status, scheduler_weight, tags, gpu_count, current_running_jobs, max_parallel_jobs, data_root, model_root, log_root, api_version, capabilities, description, last_health_check_at, health_detail) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 0, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( node_id, payload["code"], payload["name"] or payload["code"], payload["api_base_url"], payload["file_gateway_url"], 1 if payload["enabled"] else 0, payload["scheduler_status"], payload["scheduler_weight"], json_dumps(payload["tags"]), payload["gpu_count"], payload["max_parallel_jobs"], payload["data_root"], payload["model_root"], payload["log_root"], payload["api_version"], json_dumps(payload["capabilities"]), payload["description"], now, json_dumps(payload["health_detail"]), ), ) return next(node for node in self.compute_nodes() if node["id"] == node_id) def update_compute_node_health(self, node_id: str, health: dict[str, Any], success: bool, error: str | None = None) -> dict[str, Any]: current = next((n for n in self.compute_nodes() if n["id"] == node_id), None) if not current: raise KeyError(node_id) status = "online" if success and current.get("enabled") else "offline" if current.get("scheduler_status") == "draining" and success: status = "draining" detail = { **(current.get("health_detail") or {}), **health, "status": "ok" if success else "failed", "last_error": error or "", "checked_at": utcnow(), } return self.update_compute_node( node_id, { "scheduler_status": status, "last_health_check_at": detail["checked_at"], "health_detail": detail, "data_root": health.get("data_root") or current.get("data_root"), "api_version": str(health.get("api_version") or current.get("api_version") or "v1"), "capabilities": health.get("capabilities") or current.get("capabilities") or [], }, ) def replace_node_gpus(self, node_id: str, gpus: list[dict[str, Any]]) -> None: now = utcnow() with self.connect() as conn: conn.execute("DELETE FROM gpus WHERE node_id=?", (node_id,)) for index, item in enumerate(gpus): gpu_index = int(item.get("gpu_index", item.get("id", index)) or 0) memory_total = safe_float(item.get("memory_total_gb") or item.get("memory_total")) if not memory_total and item.get("memory_total_mb") is not None: memory_total = round(safe_float(item.get("memory_total_mb")) / 1024, 2) power_limit = safe_float(item.get("power_limit_w") or item.get("power_limit")) temperature = int(safe_float(item.get("temperature") or item.get("base_temperature"), 35)) conn.execute( """ INSERT INTO gpus (id, node_id, gpu_index, uuid, name, memory_total_gb, power_limit_w, base_temperature, last_seen_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( f"{node_id}_gpu_{gpu_index}", node_id, gpu_index, str(item.get("uuid") or f"{node_id}-GPU-{gpu_index}"), str(item.get("name") or "Unknown GPU"), memory_total or 0, power_limit or 0, temperature, now, ), ) conn.execute("UPDATE compute_nodes SET gpu_count=? WHERE id=?", (len(gpus), node_id)) def gpus(self) -> list[dict[str, Any]]: self.refresh_runtime_state() with self.connect() as conn: rows = conn.execute( """ SELECT g.*, n.code AS node_code, n.name AS node_name FROM gpus g JOIN compute_nodes n ON n.id = g.node_id ORDER BY n.code, g.gpu_index """ ).fetchall() running_tasks = [ self._task(row) for row in conn.execute( "SELECT * FROM fine_tune_tasks WHERE status IN ('syncing','queued','running')" ).fetchall() ] items = [] for row in rows: task = next( ( t for t in running_tasks if t.get("compute_node_id") == row["node_id"] and row["gpu_index"] in (t.get("gpus") or []) ), None, ) busy = task is not None and task.get("status") == "running" reserved = task is not None and task.get("status") in {"syncing", "queued"} memory_used = round(row["memory_total_gb"] * (0.72 if busy else 0.18 if reserved else 0.04), 1) gpu_percent = 86 if busy else 22 if reserved else 3 memory_total = float(row["memory_total_gb"] or 0) items.append( { "id": row["gpu_index"], "node_id": row["node_id"], "node_code": row["node_code"], "node_name": row["node_name"], "name": row["name"], "uuid": row["uuid"], "gpu_percent": gpu_percent, "memory_used_gb": memory_used, "memory_total_gb": memory_total, "memory_percent": round(memory_used / memory_total * 100, 1) if memory_total else 0, "temperature": row["base_temperature"] + (21 if busy else 6 if reserved else 0), "power_w": round(row["power_limit_w"] * (0.7 if busy else 0.25 if reserved else 0.08), 1), "power_limit_w": row["power_limit_w"], "status": "busy" if busy else "reserved" if reserved else "idle", "processes": [ { "pid": task["process_id"], "name": "llamafactory-cli", "memory_used_gb": memory_used, "task_name": task["name"], "user": "admin", } ] if task else [], } ) return items def system_info(self) -> dict[str, Any]: gpus = self.gpus() busy = len([g for g in gpus if g["status"] in {"busy", "reserved"}]) cpu_percent = 18 + busy * 9 memory_percent = 37 + busy * 4 return { "timestamp": utcnow(), "cpu": { "percent": min(cpu_percent, 95), "cores": 32, "percents": [min(cpu_percent + (i % 7) - 3, 99) for i in range(32)], "model": "Platform x86_64 CPU", "frequency_mhz": 2600, "load_1m": round(cpu_percent / 10, 2), }, "memory": { "used_gb": round(256 * memory_percent / 100, 1), "total_gb": 256, "percent": min(memory_percent, 95), "available_gb": round(256 * (100 - memory_percent) / 100, 1), "cached_gb": 32, }, "disk": { "used_gb": 840, "total_gb": 2048, "percent": 41, "read_mb_s": 120 if busy else 8, "write_mb_s": 95 if busy else 5, }, "gpu": gpus, "network": { "download_mb_s": 12 if busy else 1.2, "upload_mb_s": 7 if busy else 0.8, "download_mb": 8024, "upload_mb": 1732, }, "system": { "uptime_seconds": int(time.time() % 100000), "process_count": 248 + busy, "os": "Linux platform compute image", }, } def health_metrics(self) -> dict[str, float]: info = self.system_info() return { "cpu_percent": info["cpu"]["percent"], "memory_percent": info["memory"]["percent"], "disk_percent": info["disk"]["percent"], } def queue(self) -> list[dict[str, Any]]: priority_score = {"urgent": 3, "high": 2, "normal": 1, "low": 0} items = [ { "id": task["id"], "name": task["name"], "status": task["status"], "progress": task.get("progress", 0), "priority": task.get("priority", "normal"), "compute_node_id": task.get("compute_node_id"), "gpus": task.get("gpus", []), "create_time": task.get("create_time"), } for task in self.tasks() if task["status"] in {"pending", "syncing", "queued", "running"} ] return sorted(items, key=lambda item: (-priority_score.get(item["priority"], 1), item["create_time"]), reverse=False) def replicas(self, node_id: str) -> list[dict[str, Any]]: with self.connect() as conn: rows = conn.execute( "SELECT * FROM resource_replicas WHERE node_id=? ORDER BY create_time DESC", (node_id,), ).fetchall() return [dict(row) for row in rows] def sync_job(self, sync_id: str) -> dict[str, Any]: self.refresh_runtime_state() with self.connect() as conn: row = conn.execute("SELECT * FROM resource_sync_jobs WHERE id=?", (sync_id,)).fetchone() if not row: raise KeyError(sync_id) return {**dict(row), "resources": json_loads(row["resources"], [])} def training_log_files(self) -> list[dict[str, Any]]: return [ { "file": f"train_{task['id']}_pid{task['process_id'] or 0}.log", "name": task["name"], "pid": task.get("process_id") or 0, "size": f"{max(1, int((task.get('progress', 0) or 0) * 1.5))} KB", "date": task.get("create_time", "")[:10], } for task in self.tasks() ] def training_log_content(self, file_name: str) -> dict[str, Any]: task = next((t for t in self.tasks() if t["id"] in file_name), None) if not task: raise KeyError(file_name) content = self.generate_training_log(task) return {"file": file_name, "content": content, "size": f"{max(1, len(content.encode('utf-8')) // 1024)} KB"} def generate_training_log(self, task: dict[str, Any]) -> str: progress = int(task.get("progress", 0) or 0) points = max(1, min(80, progress)) lines = [ f"[INFO] task={task['name']} engine=llama_factory status={task['status']}", f"[INFO] base_model={task.get('base_model')} dataset={task.get('train_dataset_id')} gpus={task.get('gpus', [])}", "[INFO] command=llamafactory-cli train --stage sft --finetuning_type lora --do_train true", ] for step in range(1, points + 1): if step % 3 != 0 and step != points: continue loss = max(0.12, 2.4 * math.exp(-step / 42)) grad_norm = 0.45 + (step % 8) * 0.03 lr = float(task.get("learning_rate") or 0.0002) * max(0.05, 1 - step / 120) epoch = round(step / max(1, points) * float(task.get("n_epochs") or 3), 4) lines.append( "{" f"'loss': {loss:.4f}, 'grad_norm': {grad_norm:.4f}, " f"'learning_rate': {lr:.8f}, 'epoch': {epoch:.4f}" "}" ) if task.get("status") == "completed": lines.extend( [ "***** train metrics *****", f"epoch = {task.get('n_epochs', 3)}", "train_loss = 0.1248", f"train_runtime = {task.get('train_duration') or '1m 10s'}", "***** train metrics end *****", ] ) return "\n".join(lines) def log_files(self, date: str | None = None) -> list[dict[str, Any]]: today = date or utcnow()[:10] return [ {"file": f"backend-{today}.log", "name": f"backend-{today}.log", "size": "32 KB", "date": today}, {"file": f"error-{today}.log", "name": f"error-{today}.log", "size": "1 KB", "date": today}, ] def log_content(self, file_name: str) -> dict[str, Any]: lines = [ json_dumps( { "timestamp": utcnow(), "level": "INFO", "logger": "platform", "file": "backend/app/api/v1/endpoints/platform.py", "line": 1, "message": "Platform log stream is available.", } ) ] return {"file": file_name, "content": "\n".join(lines), "size": "1 KB"} _store: PlatformStore | None = None def get_platform_store() -> PlatformStore: global _store if _store is None: _store = PlatformStore() return _store