from __future__ import annotations import hashlib import hmac import json 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 psycopg_pool import ConnectionPool 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 _SIZE_UNIT_BYTES = { "B": 1, "KB": 1024, "MB": 1024**2, "GB": 1024**3, "TB": 1024**4, } def parse_size_bytes(value: Any) -> int: """把历史字符串大小统一换算为字节，供接口返回稳定的数值字段。""" if isinstance(value, bool): return 0 if isinstance(value, (int, float)): return max(0, int(value)) match = re.fullmatch( r"\s*([0-9]+(?:\.[0-9]+)?)\s*(B|KB|MB|GB|TB)?\s*", str(value or ""), flags=re.IGNORECASE, ) if not match: return 0 amount = float(match.group(1)) unit = (match.group(2) or "B").upper() return max(0, round(amount * _SIZE_UNIT_BYTES[unit])) def count_dataset_records(content: str) -> int: text = (content or "").strip() if not text: return 0 try: value = json.loads(text) if isinstance(value, list): return len(value) return 1 except (TypeError, ValueError, json.JSONDecodeError): pass return len([line for line in text.splitlines() if line.strip()]) _EVAL_METHOD_LABELS = { "standard": "标准匹配", "metric_standard": "综合评测", "semantic": "语义相似度", "sentiment": "情感分析", "accuracy": "准确性评估", "safety": "安全性评估", "relevance": "相关性评估", "fluency": "流畅性评估", "factuality": "事实性评估", "custom": "自定义评估", } def _eval_method_label(value: Any) -> str: if isinstance(value, list): return "、".join(_eval_method_label(item) for item in value if item) text = str(value or "").strip() return _EVAL_METHOD_LABELS.get(text, text) def _basic_metric_labels(config: dict[str, Any] | None) -> list[str]: cfg = config or {} labels: list[str] = [] bleu = cfg.get("bleu") or {} if bleu.get("enabled"): labels.append(f"BLEU-{int(bleu.get('ngram') or 4)}") rouge = cfg.get("rouge") or {} if rouge.get("enabled"): methods = rouge.get("methods") or [] method_labels = { "rouge1": "ROUGE-1", "rouge2": "ROUGE-2", "rougeL": "ROUGE-L", "rouge_1": "ROUGE-1", "rouge_2": "ROUGE-2", "rouge_l": "ROUGE-L", } labels.extend(method_labels.get(str(item), str(item)) for item in methods) cosine = cfg.get("cosine") or {} if cosine.get("enabled"): labels.append("Cosine") return labels def _build_eval_metric_label(payload: dict[str, Any], dimension: dict[str, Any] | None = None) -> str: parts: list[str] = [] dim = dimension or {} dim_type = str(dim.get("type") or payload.get("dimension_type") or "").strip() method_label = _eval_method_label(dim.get("eval_method") or payload.get("eval_method")) if dim_type in {"classification", "metric"} and method_label: parts.append(f"LLM：{method_label}") elif dim_type == "text_similarity" and method_label: parts.append(method_label) parts.extend(_basic_metric_labels(payload.get("basic_metrics") or {})) if not parts: metric = str(payload.get("metric") or payload.get("eval_type") or "").strip() return "自定义评测" if metric == "custom" else (metric or "-") return " + ".join(parts) def version_number(value: Any, default: int = 0) -> int: try: number = int(value) except (TypeError, ValueError): return default return number if number > 0 else default def dataset_file_version_summary(file_row: PgRow) -> dict[str, Any]: versions = json_loads(file_row.get("versions"), []) versions = versions if isinstance(versions, list) else [] active_version_id = str( file_row.get("active_version_id") or file_row.get("current_version_id") or "" ) active_version = next( ( item for item in versions if isinstance(item, dict) and str(item.get("id") or "") == active_version_id ), None, ) current_version_no = version_number( (active_version or {}).get("version_no") or (active_version or {}).get("version") or file_row.get("version_no"), default=1 if active_version_id or versions else 0, ) return { "active_version_id": active_version_id or None, "current_version_id": active_version_id or None, "current_version_no": current_version_no or None, "version_count": len(versions), } 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] = {} number_pattern = r"([-+]?(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][-+]?\d+)?)" step_match = re.search(rf"(?:^|[\s,{{])['\"]?step['\"]?\s*(?:=|:)\s*{number_pattern}", line, re.I) if step_match: result["step"] = float(step_match.group(1)) for key in ["loss", "grad_norm", "learning_rate", "epoch"]: match = re.search(rf"['\"]?{key}['\"]?\s*(?:=|:)\s*{number_pattern}", line, re.I) 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] = {} fmt = str(formatting).lower() for key, file_name in zip(llama_dataset_keys(dataset_key, file_names), file_names): if fmt == "sharegpt": # 平台校验按 OpenAI 风格消息（role/content），故 tags 用 role/content # 与 LLaMA-Factory 默认的 from/value 不同，需显式声明避免解析失败。 result[key] = { "file_name": file_name, "formatting": "sharegpt", "columns": {"messages": "messages"}, "tags": { "role_tag": "role", "content_tag": "content", "user_tag": "user", "assistant_tag": "assistant", }, } elif fmt == "dpo": result[key] = { "file_name": file_name, "formatting": "dpo", "columns": { "prompt": "system", "chosen": "chosen", "rejected": "rejected", }, } elif fmt in {"cpt", "pt", "pretrain"}: result[key] = { "file_name": file_name, "formatting": "cpt", "columns": {"prompt": "text"}, } else: result[key] = { "file_name": file_name, "formatting": "alpaca", "columns": { "prompt": "instruction", "query": "input", "response": "output", }, } return result def _sniff_dataset_format(sample_text: str, max_samples: int = 20) -> str: """嗅探数据集内容格式（兼容 jsonl），返回 sharegpt / dpo / cpt / alpaca。 按内容而非文件名判断，纯 jsonl 数据集（如 ShareGPT messages、缺省 input 的 Alpaca）都能被正确识别，避免训练任务误按 alpaca 解析而失败。 """ text = (sample_text or "").strip() if not text: return "" records: list[dict[str, Any]] = [] try: value = json.loads(text) except (TypeError, ValueError, json.JSONDecodeError): value = None if isinstance(value, list): records = [item for item in value[:max_samples] if isinstance(item, dict)] elif isinstance(value, dict): records = [value] else: for line in text.splitlines()[:max_samples]: line = line.strip() if not line: continue try: obj = json.loads(line) except json.JSONDecodeError: continue if isinstance(obj, dict): records.append(obj) records = records[:max_samples] if not records: return "" if all("messages" in record for record in records): return "sharegpt" if all(record.get("chosen") and record.get("rejected") for record in records): return "dpo" if all(record.get("text") and not (record.get("instruction") or record.get("output")) for record in records): return "cpt" return "alpaca" 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) # Reuse connections via a pool to avoid the TCP+auth handshake on every # request (notably expensive against the remote PostgreSQL instance). # TCP keepalive 让操作系统持续保活连接，抵抗远程库空闲静默断连。 pool_kwargs = { "connect_timeout": 30, "keepalives": 1, "keepalives_idle": 10, "keepalives_interval": 5, "keepalives_count": 3, } self._pool = ConnectionPool( conninfo=self.database_url, kwargs=pool_kwargs, min_size=2, max_size=20, # 借出前校验连接可用性，避免执行 SQL 时才发现 [BAD] 再重建。 check=ConnectionPool.check_connection, # 不主动回收空闲连接（远程库约 10s 断，由 keepalive 维持）， # 减少无谓的重建握手。 max_idle=0, # 请求最多排队等待，调大以适应远程库慢查询。 max_waiting=50, open=False, ) # 注意：不要在此调用 pool.wait()，它会阻塞等待 min_size 个连接就绪， # 在远程库响应慢/超时时会卡死 uvicorn worker 进程，导致所有请求无响应。 self._pool.open() self.ensure_schema() self.ensure_seed_data() # Track which compute nodes have an active inference model loaded self._inference_nodes: set[str] = set() self._last_runtime_refresh = 0.0 # ── inference node tracking ──────────────────────────────────── def mark_inference_loaded(self, node_id: str) -> None: self._inference_nodes.add(node_id) def mark_inference_unloaded(self, node_id: str) -> None: self._inference_nodes.discard(node_id) def is_inference_loaded(self, node_id: str) -> bool: return node_id in self._inference_nodes @contextmanager def connect(self) -> Iterator["PgConnection"]: with self._pool.connection() as raw_conn: conn = PgConnection(raw_conn) try: yield conn conn.commit() except Exception: conn.rollback() raise finally: conn.close() def close_pool(self) -> None: """Release pooled connections. Safe to call multiple times.""" try: self._pool.close() except Exception: pass 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, "sessions", { "username": "TEXT", "login_at": "TEXT", "logout_at": "TEXT", "duration_seconds": "INTEGER", "issued_at": "TEXT", "expires_at": "TEXT", "ip": "TEXT", "create_time": "TEXT", }, ) self._ensure_columns( conn, "trained_models", { "artifact_dir": "TEXT", "compute_node_id": "TEXT", "compute_node_name": "TEXT", "created_by": "TEXT", "tenant_id": "TEXT", "project_id": "TEXT", "deleted_at": "TEXT", "deleted_by": "TEXT", }, ) for table in ("models", "datasets", "eval_tasks"): self._ensure_columns(conn, table, {"deleted_at": "TEXT", "deleted_by": "TEXT", "tenant_id": "TEXT", "project_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", }, ) schema_dir = Path(__file__).with_name("sql") for extra in ( "002_governance.sql", "003_model_path_governance.sql", "003_tenant_quota.sql", "004_permissions.sql", ): extra_path = schema_dir / extra if extra_path.exists(): conn.executescript(extra_path.read_text(encoding="utf-8")) 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 now_ts = time.monotonic() if now_ts - self._last_runtime_refresh < 2: return self._last_runtime_refresh = now_ts 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], job: dict[str, Any] | None = None) -> 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: conn.execute( """ UPDATE trained_models SET compute_node_id=COALESCE(compute_node_id, ?), compute_node_name=COALESCE(compute_node_name, ?) WHERE id=? """, (task.get("compute_node_id"), task.get("compute_node_code") or task.get("compute_node_name"), exists["id"]), ) 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, artifact_dir, compute_node_id, compute_node_name, created_by) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( trained_model_id, name, json_dumps([{"name": task.get("train_method", "lora")}]), model["path"] if model else "", utcnow(), 0, 0, output_dir, output_dir, task.get("compute_node_id"), task.get("compute_node_code") or task.get("compute_node_name"), task.get("created_by"), ), ) # Use real artifact data from compute node when available artifacts = (job or {}).get("artifacts") or [] if artifacts: total_size = sum(int(a.get("size_bytes") or a.get("size", 0)) for a in artifacts) checksums = [a.get("checksum_sha256", "") for a in artifacts if a.get("checksum_sha256")] combined_checksum = checksums[0] if len(checksums) == 1 else "" # Register individual artifact files for artifact in artifacts[:50]: # limit to 50 file entries artifact_path = artifact.get("path") or artifact.get("name", "") abs_path = artifact_path if artifact_path.startswith("/") else f"{output_dir.rstrip('/')}/{artifact_path.lstrip('/')}" self._upsert_model_artifact( conn, trained_model_id, "trained_model", "adapter_file", abs_path, int(artifact.get("size_bytes") or artifact.get("size", 0)), artifact.get("checksum_sha256", ""), { "task_id": task.get("id"), "train_method": task.get("train_method", "lora"), "base_model": task.get("base_model"), "artifact_name": artifact.get("name", ""), }, task.get("compute_job_id"), ) else: total_size = 0 combined_checksum = "" # Register the top-level adapter directory entry self._upsert_model_artifact( conn, trained_model_id, "trained_model", "adapter", output_dir, total_size, combined_checksum, { "task_id": task.get("id"), "train_method": task.get("train_method", "lora"), "base_model": task.get("base_model"), "file_count": len(artifacts), }, 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, int(metric.get("step") or 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) is_admin = row["role"] == "admin" or bool(row["protected"]) if "permissions" in payload: perms = payload["permissions"] if is_admin: # 管理员权限不可更改，必须是全部 perms = ALL_PERMISSIONS else: # 非 admin 用户不能拥有 user-settings 权限 perms = [p for p in (perms or []) if p != "user-settings"] payload = {**payload, "permissions": perms} 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") # 级联删除该用户关联的数据 tables_to_clean = [ # ACL 授权 ("acls", "principal_type='user' AND principal_id=?", [user_id]), # 审批实例（申请人） ("approval_instances", "applicant_id=?", [user_id]), # 审计日志 ("audit_logs", "actor_id=?", [user_id]), # 项目成员 ("project_members", "user_id=?", [user_id]), # GPU 分配 ("gpu_assignments", "user_id=?", [user_id]), # 数据集 ("datasets", "created_by=?", [user_id]), # 基座模型 ("models", "created_by=?", [user_id]), # 微调产物 ("trained_models", "created_by=?", [user_id]), # 评测任务 ("eval_tasks", "created_by=?", [user_id]), # 对比/推理任务（payload 中 creator） # 训练任务：仅标记为已删除或保留（有 compute_job_id 关联），不清物理数据 ] for table_name, where_clause, params in tables_to_clean: try: conn.execute(f"DELETE FROM {table_name} WHERE {where_clause}", params) except Exception: pass # 表可能不存在或字段不存在，跳过 conn.execute("DELETE FROM users WHERE id=?", (user_id,)) def reset_password(self, user_id: str, new_password: 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 reset password") conn.execute( "UPDATE users SET password_hash=? WHERE id=?", (hash_password(new_password), 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") model_source = payload.get("model_source", "local") path = payload.get("path", "") # Automatically determine can_train: local models with a path can be trained can_train = 1 if (model_source != "api" and path and str(path).strip()) else 0 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, can_train, create_time, created_by) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( model_id, payload["name"], payload.get("type", "LLM"), payload.get("purpose", "training"), model_source, payload.get("description"), path, payload.get("api_url"), payload.get("api_key"), payload.get("online_model_name"), can_train, utcnow(), payload.get("created_by"), ), ) 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} # Recompute can_train when relevant fields change model_source = merged.get("model_source", "local") path = merged.get("path", "") can_train = 1 if (model_source != "api" and path and str(path).strip()) else 0 with self.connect() as conn: conn.execute( """ UPDATE models SET name=?, type=?, purpose=?, model_source=?, description=?, path=?, api_url=?, api_key=?, online_model_name=?, can_train=? WHERE id=? """, ( merged["name"], merged.get("type", "LLM"), merged.get("purpose", "training"), model_source, merged.get("description"), path, merged.get("api_url"), merged.get("api_key"), merged.get("online_model_name"), can_train, 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("UPDATE models SET deleted_at=?, deleted_by=? WHERE id=?", (utcnow(), "system", 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 WHERE deleted_at IS NULL ORDER BY create_time DESC").fetchall() items = [] for row in rows: item = { **dict(row), "train_methods": json_loads(row["train_methods"], []), "merged": bool(row["merged"]), "merging": bool(row["merging"]), } if not item.get("compute_node_id"): task_rows = conn.execute("SELECT payload FROM fine_tune_tasks ORDER BY create_time DESC").fetchall() for task in task_rows: task_payload = json_loads(task["payload"], {}) output_name = task_payload.get("output_model_name") or f"{task_payload.get('name')}-lora" if output_name == item["name"]: item["compute_node_id"] = task_payload.get("compute_node_id") item["compute_node_name"] = task_payload.get("compute_node_code") or task_payload.get("compute_node_name") break items.append(item) return items def delete_trained_model(self, model_id: str) -> None: with self.connect() as conn: conn.execute("UPDATE trained_models SET deleted_at=?, deleted_by=? WHERE id=? OR name=?", (utcnow(), "system", model_id, model_id)) def datasets(self) -> list[dict[str, Any]]: with self.connect() as conn: rows = conn.execute( """SELECT dataset.*, task.name AS task_name FROM datasets dataset LEFT JOIN data_process_tasks task ON task.id=COALESCE(dataset.source_task_id, dataset.task_id) WHERE dataset.deleted_at IS NULL ORDER BY dataset.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 dataset.*, task.name AS task_name FROM datasets dataset LEFT JOIN data_process_tasks task ON task.id=COALESCE(dataset.source_task_id, dataset.task_id) WHERE dataset.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, size_bytes, active_version_id, current_version_id, version_no, versions, create_time, record_count, metadata FROM dataset_files WHERE dataset_id=? ORDER BY create_time, id""", (row["id"],), ).fetchall() decoded_files: list[dict[str, Any]] = [] for file_row in files: metadata = json_loads(file_row.get("metadata"), {}) file_size_bytes = int(file_row.get("size_bytes") or 0) if file_size_bytes <= 0: file_size_bytes = parse_size_bytes(file_row.get("size")) file_record_count = int(file_row.get("record_count") or 0) decoded_files.append( { "id": file_row["id"], "name": file_row["name"], "size": file_row["size"], "size_bytes": file_size_bytes, **dataset_file_version_summary(file_row), "create_time": file_row["create_time"], "record_count": file_record_count, "split": metadata.get("file_split"), } ) dataset_metadata = json_loads(row.get("metadata"), {}) split_counts = dict(dataset_metadata.get("split_counts") or {}) if row.get("source") == "task" and not split_counts: split_rows = conn.execute( """SELECT split, COUNT(*) AS count FROM dataset_records WHERE dataset_id=? GROUP BY split""", (row["id"],), ).fetchall() split_counts = {str(item["split"]): int(item["count"]) for item in split_rows} total_size_bytes = sum(item["size_bytes"] for item in decoded_files) if not decoded_files: total_size_bytes = int(row.get("size_bytes") or 0) if total_size_bytes <= 0: total_size_bytes = parse_size_bytes(row.get("size")) total_record_count = sum(int(item.get("record_count") or 0) for item in decoded_files) if not decoded_files: total_record_count = int(row.get("record_count") or row.get("count") or 0) current_version_nos = sorted( { int(item["current_version_no"]) for item in decoded_files if item.get("current_version_no") } ) return { **dict(row), "count": total_record_count, "record_count": total_record_count, "size_bytes": total_size_bytes, "current_version_no": ( current_version_nos[0] if len(current_version_nos) == 1 else None ), "current_version_nos": current_version_nos, "version_count": sum(int(item["version_count"]) for item in decoded_files), "metadata": dataset_metadata, "split_counts": { "train": int(split_counts.get("train", 0) or 0), "validation": int(split_counts.get("validation", 0) or 0), "test": int(split_counts.get("test", 0) or 0), }, "files": decoded_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, created_by) 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(), payload.get("created_by"), ), ) 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("UPDATE datasets SET deleted_at=?, deleted_by=? WHERE id=?", (utcnow(), "system", 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_bytes = len(content.encode("utf-8")) size = f"{size_bytes} B" record_count = count_dataset_records(content) version = { "id": version_id, "version": 1, "version_no": 1, "create_time": now, "description": "uploaded", "size_bytes": size_bytes, "record_count": record_count, } conn.execute( """ INSERT INTO dataset_files (id, dataset_id, name, size, content, active_version_id, versions, create_time, current_version_id, size_bytes, record_count, version_no) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 1) """, ( file_id, dataset_id, name, size, content, version_id, json_dumps([version]), now, version_id, size_bytes, record_count, ), ) conn.execute( """UPDATE datasets SET count=stats.record_count, record_count=stats.record_count, size_bytes=stats.size_bytes, size=(stats.size_bytes::text || ' B') FROM ( SELECT COALESCE(SUM(record_count), 0) AS record_count, COALESCE(SUM(size_bytes), 0) AS size_bytes FROM dataset_files WHERE dataset_id=? ) stats WHERE id=?""", (dataset_id, dataset_id), ) return { "id": file_id, "name": name, "size": size, "size_bytes": size_bytes, "current_version_no": 1, "version_count": 1, } def dataset_file_record_sources(self, file_id: str) -> list[dict[str, Any]]: """返回发布样本关联的真实原文，供数据集详情核对生成内容。""" with self.connect() as conn: file_row = conn.execute( "SELECT id FROM dataset_files WHERE id=?", (file_id,) ).fetchone() if not file_row: raise KeyError(file_id) rows = conn.execute( """ SELECT records.line_no, records.instruction, records.input, records.output, COALESCE( NULLIF(preview.edited_content, ''), preview.original_content, '' ) AS source_text, CASE WHEN preview.edited_content IS NOT NULL AND preview.edited_content <> '' THEN TRUE ELSE FALSE END AS preprocessed FROM dataset_records AS records LEFT JOIN data_process_preview_items AS preview ON preview.id=records.preview_item_id WHERE records.dataset_file_id=? ORDER BY records.line_no NULLS LAST, records.created_at, records.id """, (file_id,), ).fetchall() return [ { "line_no": int(item.get("line_no") or index + 1), "instruction": str(item.get("instruction") or ""), "input": str(item.get("input") or ""), "output": str(item.get("output") or ""), "source_text": str(item.get("source_text") or ""), "preprocessed": bool(item.get("preprocessed")), } for index, item in enumerate(rows) ] 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, metadata FROM datasets WHERE id=?", (dataset_id,) ).fetchone() if not dataset: raise KeyError(dataset_id) dataset_metadata = json_loads(dataset.get("metadata"), {}) related_ids = dataset_metadata.get("split_dataset_ids") or {} runtime_dataset_ids = [dataset_id] validation_dataset_id = related_ids.get("validation") if dataset_metadata.get("dataset_split") == "train" and validation_dataset_id: runtime_dataset_ids.append(str(validation_dataset_id)) rows = conn.execute( """ SELECT id, dataset_id, name, size, content, active_version_id, create_time, record_count, metadata FROM dataset_files WHERE dataset_id = ANY(%s) ORDER BY CASE WHEN dataset_id=%s THEN 0 ELSE 1 END, create_time, id """, (runtime_dataset_ids, dataset_id), ).fetchall() return [ { **dict(row), "metadata": json_loads(row.get("metadata"), {}), "split": json_loads(row.get("metadata"), {}).get("file_split"), } 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) content = payload.get("content", "") size_bytes = len(content.encode("utf-8")) record_count = count_dataset_records(content) versions = json_loads(row["versions"], []) version = { "id": f"{file_id}_v{len(versions) + 1}", "version": len(versions) + 1, "version_no": len(versions) + 1, "create_time": utcnow(), "description": payload.get("description", "online edit"), "size_bytes": size_bytes, "record_count": record_count, } versions.append(version) conn.execute( """ UPDATE dataset_files SET content=?, active_version_id=?, current_version_id=?, versions=?, size_bytes=?, size=?, record_count=?, version_no=? WHERE id=? """, ( content, version["id"], version["id"], json_dumps(versions), size_bytes, f"{size_bytes} B", record_count, version["version_no"], file_id, ), ) conn.execute( """UPDATE datasets SET count=stats.record_count, record_count=stats.record_count, size_bytes=stats.size_bytes, size=(stats.size_bytes::text || ' B') FROM ( SELECT dataset_id, COALESCE(SUM(record_count), 0) AS record_count, COALESCE(SUM(size_bytes), 0) AS size_bytes FROM dataset_files WHERE dataset_id=(SELECT dataset_id FROM dataset_files WHERE id=?) GROUP BY dataset_id ) stats WHERE datasets.id=stats.dataset_id""", (file_id,), ) return {"version": version, "content": 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") with self.connect() as conn: train_dataset = conn.execute( "SELECT id, type, metadata FROM datasets WHERE id=?", (train_dataset_id,), ).fetchone() if not train_dataset: raise ValueError("training dataset not found") train_metadata = json_loads(train_dataset.get("metadata"), {}) if train_dataset.get("type") != "train" or train_metadata.get( "dataset_split" ) in {"validation", "test"}: raise ValueError("train_dataset_id must reference a training dataset") 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 [] 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) selected_gpus = list(node.get("selected_gpus") or selected_gpus) 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.select_compute_node(merged) selected_gpus = merged.get("gpus") or [] return node, self._compute_job_payload_from_task_node(merged, node, selected_gpus) def prepare_compute_job_payload_from_payload(self, payload: dict[str, Any]) -> tuple[dict[str, Any], dict[str, Any]]: task_id = str(payload.get("task_id") or payload.get("id") or new_id("ft_preview")) name = str(payload.get("name") or task_id) transient_task = { **payload, "id": task_id, "name": name, "status": str(payload.get("status") or "pending"), "progress": int(payload.get("progress") or 0), } node = self.select_compute_node(transient_task) selected_gpus = transient_task.get("gpus") or [] return node, self._compute_job_payload_from_task_node(transient_task, 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]: base_model_id = task.get("base_model") or task.get("model_id") dataset_id = str(task.get("train_dataset_id") or task.get("dataset_id") or "") with self.connect() as conn: model = conn.execute("SELECT * FROM models WHERE id=?", (base_model_id,)).fetchone() if not model: raise RuntimeError(f"base model not found: {base_model_id}") # P0-1: Reject non-trainable models (API models or models without local path) if not model.get("can_train"): model_source = model.get("model_source") or "unknown" model_path = model.get("path") or "" if model_source == "api": raise RuntimeError( f"模型 '{model['name']}' 为 API 模型，不能作为 LLaMA-Factory 本地训练基座，请选择本地路径模型" ) if not model_path or not str(model_path).strip(): raise RuntimeError( f"模型 '{model['name']}' 未配置算力节点可访问路径，请先在模型管理中设置模型本地路径" ) raise RuntimeError( f"模型 '{model['name']}' 不支持本地训练（source={model_source}），请选择其他模型" ) dataset = conn.execute("SELECT * FROM datasets WHERE id=?", (dataset_id,)).fetchone() files = conn.execute( """SELECT id, name, size, active_version_id, create_time, metadata FROM dataset_files WHERE dataset_id=? ORDER BY create_time, id""", (dataset_id,), ).fetchall() dataset_metadata = json_loads(dataset.get("metadata"), {}) if dataset else {} related_ids = dataset_metadata.get("split_dataset_ids") or {} validation_dataset_id = related_ids.get("validation") if dataset_metadata.get("dataset_split") == "train" and validation_dataset_id: files = [ *files, *conn.execute( """SELECT id, name, size, active_version_id, create_time, metadata FROM dataset_files WHERE dataset_id=? ORDER BY create_time, id""", (str(validation_dataset_id),), ).fetchall(), ] model_path = (model and model.get("path")) or task.get("model_name_or_path") or base_model_id dataset_metadata = json_loads(dataset.get("metadata"), {}) if dataset else {} if not dataset or dataset.get("type") != "train" or dataset_metadata.get( "dataset_split" ) in {"validation", "test"}: raise RuntimeError(f"training dataset is invalid: {dataset_id}") 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)) file_entries = [ { **dict(row), "name": Path(str(row["name"] or row["id"])).name, "split": json_loads(row.get("metadata"), {}).get("file_split"), } for row in files ] split_aware = any(item["split"] for item in file_entries) training_files = [ item for item in file_entries if not split_aware or item["split"] == "train" ] validation_files = [ item for item in file_entries if split_aware and item["split"] == "validation" ] if not training_files: raise RuntimeError(f"dataset has no training split: {dataset_id}") runtime_files = [*training_files, *validation_files] runtime_file_names = [str(item["name"]) for item in runtime_files] runtime_keys = llama_dataset_keys(dataset_key, runtime_file_names) training_keys = runtime_keys[: len(training_files)] validation_keys = runtime_keys[len(training_files) :] # P0-2: 推导数据集格式并校验内容（兼容 jsonl：按内容嗅探 ShareGPT/DPO/CPT/Alpaca） train_type = str(task.get("train_type", task.get("train_method", ""))).upper() expected_format = {"DPO": "dpo", "CPT": "cpt"}.get(train_type) raw_format = str( task.get("dataset_format") or dataset_metadata.get("format") or (dataset and dataset.get("formatting")) or "alpaca" ).lower() sniffed_format = "" content_samples: dict[str, str] = {} if training_files: with self.connect() as conn: for file_entry in training_files: sample_row = conn.execute( "SELECT substr(content, 1, 400000) AS sample FROM dataset_files WHERE id=?", (str(file_entry["id"]),), ).fetchone() sample = (sample_row or {}).get("sample") or "" content_samples[str(file_entry["id"])] = sample if not sniffed_format: sniffed_format = _sniff_dataset_format(sample) known_formats = {"sharegpt", "dpo", "cpt", "pt", "pretrain"} if expected_format: dataset_format = expected_format elif raw_format in known_formats: dataset_format = raw_format else: dataset_format = sniffed_format or raw_format or "alpaca" format_errors: list[str] = [] for file_entry in training_files: content = content_samples.get(str(file_entry["id"])) or "" if content: from app.modules.data_process.dataset_format import validate_dataset_format file_errors = validate_dataset_format(dataset_format, content=content) if file_errors: format_errors.extend(file_errors) if format_errors: raise RuntimeError("数据集格式校验失败：\n" + "\n".join(f" - {e}" for e in format_errors[:10])) 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(training_keys), "dataset_key": dataset_key, "dataset_keys": training_keys, "eval_dataset": ",".join(validation_keys) or None, "eval_dataset_keys": validation_keys, "dataset_display_name": (dataset and dataset.get("name")) or dataset_id, "dataset_dir": dataset_dir, "dataset_info": llama_dataset_info(dataset_key, runtime_file_names, dataset_format), "dataset_files": [ { "id": item["id"], "name": item["name"], "relative_path": f"{dataset_id}/{item['name']}", "local_path": f"{dataset_dir.rstrip('/')}/{item['name']}", "active_version_id": item["active_version_id"], "size": item["size"], "create_time": item["create_time"], "split": item["split"], } for item in runtime_files ], "output_dir": output_dir, "gpus": selected_gpus if selected_gpus is not None else (task.get("gpus") or []), "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 []) 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, job) # P0-4: Persist failure info for diagnosis if status in {"failed", "stopped"}: failure_reason = job.get("error") or job.get("message") or "compute job failed" log_snippet = job.get("log_snippet") or "" conn.execute( "UPDATE fine_tune_tasks SET failure_reason = ? WHERE id = ?", (failure_reason[:2000], task_id), ) # Store last log snippet if available (max 8KB) if log_snippet: conn.execute( "UPDATE fine_tune_tasks SET payload = ? WHERE id = ?", (json_dumps({**payload, "last_log_snippet": log_snippet[:8192]}), task_id), ) 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"]}) if not payload.get("metric_label"): payload["metric_label"] = _build_eval_metric_label(payload) if not payload.get("metric") or payload.get("metric") == "custom": payload["metric"] = payload["metric_label"] return payload def _enrich_eval_payload(self, conn: PgConnection, payload: dict[str, Any]) -> dict[str, Any]: data = dict(payload) model_id = str(data.get("model_id") or "") if model_id and not data.get("model_name"): model = conn.execute("SELECT name FROM models WHERE id=?", (model_id,)).fetchone() if not model: model = conn.execute("SELECT name FROM trained_models WHERE id=? OR name=?", (model_id, model_id)).fetchone() if model: data["model_name"] = model["name"] dataset_id = str(data.get("dataset_id") or "") if dataset_id and not data.get("dataset"): dataset = conn.execute("SELECT name FROM datasets WHERE id=?", (dataset_id,)).fetchone() if dataset: data["dataset"] = dataset["name"] dimension = None dimension_id = str(data.get("dimension_id") or "") if dimension_id: dimension_row = conn.execute("SELECT payload FROM eval_dimensions WHERE id=?", (dimension_id,)).fetchone() if dimension_row: dimension = json_loads(dimension_row["payload"], {}) data.setdefault("dimension_type", dimension.get("type")) data.setdefault("eval_method", dimension.get("eval_method")) data.setdefault("evaluator_model", dimension.get("eval_model")) data["metric_label"] = _build_eval_metric_label(data, dimension) if not data.get("metric") or data.get("metric") in {"custom", "自定义评测"}: data["metric"] = data["metric_label"] return data 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._enrich_eval_payload(conn, 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._enrich_eval_payload(conn, 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_id = str(payload.get("model_id") or "") model = conn.execute("SELECT name FROM models WHERE id=?", (model_id,)).fetchone() trained_model = conn.execute("SELECT name FROM trained_models WHERE id=? OR name=?", (model_id, model_id)).fetchone() dataset = conn.execute("SELECT name FROM datasets WHERE id=?", (str(payload.get("dataset_id")),)).fetchone() dimension = None dimension_id = str(payload.get("dimension_id") or "") if dimension_id: dimension_row = conn.execute("SELECT payload FROM eval_dimensions WHERE id=?", (dimension_id,)).fetchone() if dimension_row: dimension = json_loads(dimension_row["payload"], {}) if model: data.setdefault("model_name", model["name"]) elif trained_model: data.setdefault("model_name", trained_model["name"]) if dataset: data.setdefault("dataset", dataset["name"]) if dimension: data.setdefault("dimension_type", dimension.get("type")) data.setdefault("eval_method", dimension.get("eval_method")) data.setdefault("evaluator_model", dimension.get("eval_model")) data["metric_label"] = _build_eval_metric_label(data, dimension) if data.get("metric") == "custom": data["metric"] = data["metric_label"] 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 update_eval_task(self, task_id: str, updates: dict[str, Any]) -> dict[str, Any]: """Update fields in an eval task's payload without replacing the whole record.""" task = self.eval_task(task_id) merged = {**task, **updates} with self.connect() as conn: conn.execute( "UPDATE eval_tasks SET payload=?, status=? WHERE id=?", (json_dumps(merged), merged.get("status", task.get("status", "pending")), task_id), ) return self.eval_task(task_id) def delete_eval_task(self, task_id: str) -> None: with self.connect() as conn: conn.execute("UPDATE eval_tasks SET deleted_at=?, deleted_by=? WHERE id=?", (utcnow(), "system", task_id)) def running_eval_tasks(self) -> list[dict[str, Any]]: """Return eval tasks that have been submitted to a compute node and are still running.""" return [ task for task in self.eval_tasks() if task.get("compute_job_id") and task.get("status") in {"queued", "running"} ] def apply_eval_job_result(self, task_id: str, job: dict[str, Any], result_content: dict[str, Any] | None = None) -> dict[str, Any]: """Sync a compute job status/result back to an eval task.""" task = self.eval_task(task_id) job_status = str(job.get("status", "")) status_map = {"queued": "running", "running": "running", "completed": "completed", "failed": "failed", "stopped": "stopped"} new_status = status_map.get(job_status, job_status or task.get("status", "pending")) updates: dict[str, Any] = { "status": new_status, "progress": int(job.get("progress", 0)), "output_dir": job.get("output_dir", task.get("output_dir", "")), } # On completion, populate results from eval_results.json content if new_status == "completed" and result_content: updates.update({ "overall_score": result_content.get("overall_score", 0), "overall_score_max": result_content.get("overall_score_max", 100), "overall_evaluation": result_content.get("overall_evaluation", ""), "improvement_suggestions": result_content.get("improvement_suggestions", []), "dimension_summary": result_content.get("dimension_summary", []), "samples": result_content.get("samples", []), "sample_count": result_content.get("sample_count", 0), "completed_count": result_content.get("completed_count", 0), "passed_count": result_content.get("passed_count", 0), "basic_metrics": result_content.get("basic_metrics", {}), "score": result_content.get("overall_score", 0), "completed_time": utcnow(), }) elif new_status in {"failed", "stopped"}: updates.update({ "error": job.get("error") or task.get("error") or "", "completed_time": utcnow(), }) return self.update_eval_task(task_id, updates) 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 _node_gpu_indexes(self, conn: PgConnection, node: dict[str, Any]) -> set[int]: rows = conn.execute("SELECT gpu_index FROM gpus WHERE node_id=?", (node["id"],)).fetchall() if rows: return {int(row["gpu_index"]) for row in rows} return set(range(max(0, int(node.get("gpu_count") or 0)))) def _node_capacity(self, node: dict[str, Any]) -> int: return max(1, int(node.get("max_parallel_jobs") or 1), int(node.get("gpu_count") or 0)) @staticmethod def _payload_gpu_indexes(payload: dict[str, Any]) -> list[int]: raw = payload.get("gpu_indices") if raw is None: raw = payload.get("gpus") if raw is None: return [] try: values = [int(item) for item in raw] except (TypeError, ValueError) as exc: raise RuntimeError("invalid GPU index") from exc if any(item < 0 for item in values): raise RuntimeError("GPU index must be non-negative") return sorted(set(values)) def _select_node_gpus( self, conn: PgConnection, node: dict[str, Any], requested: list[int], payload: dict[str, Any], ) -> list[int]: available = self._node_gpu_indexes(conn, node) active = self._active_gpu_indexes(conn, node["id"]) allowed = payload.get("allowed_gpu_indices") if allowed is not None: available &= {int(item) for item in allowed} if requested: selected = set(requested) if not selected.issubset(available): raise RuntimeError(f"requested GPU is not available on compute node {node['code']}") if selected.intersection(active): raise RuntimeError(f"requested GPU is busy on compute node {node['code']}") return sorted(selected) if payload.get("allow_cpu") or payload.get("device") == "cpu": return [] count = max(1, int(payload.get("gpu_count") or 1)) free = sorted(available - active) if len(free) < count: raise RuntimeError(f"compute node {node['code']} has only {len(free)} available GPU(s)") return free[:count] 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 = self._payload_gpu_indexes(payload) 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"] < self._node_capacity(n) ] if requested: selected = next((n for n in candidates if n["id"] == requested), None) if not selected: raise RuntimeError("selected compute node is unavailable") selected["selected_gpus"] = self._select_node_gpus(conn, selected, requested_gpus, payload) return selected filtered = [] for node in candidates: try: node["selected_gpus"] = self._select_node_gpus(conn, node, requested_gpus, payload) filtered.append(node) except RuntimeError: continue candidates = filtered 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"] >= self._node_capacity(node): reason = f"capacity full {node['current_running_jobs']}/{self._node_capacity(node)}" elif requested: reason = "selected node unavailable" 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 select_compute_node(self, payload: dict[str, Any]) -> dict[str, Any]: with self.connect() as conn: 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 create_storage_object(self, payload: dict[str, Any]) -> dict[str, Any]: object_id = str(payload.get("id") or new_id("object")) with self.connect() as conn: conn.execute( """ INSERT INTO storage_objects (id, resource_type, resource_id, version_id, bucket, object_key, file_name, content_type, checksum_sha256, byte_size, status, created_by, create_time) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT (resource_type, resource_id, version_id, object_key) DO UPDATE SET file_name=EXCLUDED.file_name, content_type=EXCLUDED.content_type, checksum_sha256=EXCLUDED.checksum_sha256, byte_size=EXCLUDED.byte_size, status=EXCLUDED.status, created_by=EXCLUDED.created_by """, ( object_id, payload["resource_type"], payload["resource_id"], payload["version_id"], payload["bucket"], payload["object_key"], payload.get("file_name"), payload.get("content_type"), payload.get("checksum_sha256"), int(payload.get("byte_size") or 0), payload.get("status", "pending"), payload.get("created_by"), payload.get("create_time") or utcnow(), ), ) row = conn.execute( "SELECT * FROM storage_objects WHERE resource_type=? AND resource_id=? AND version_id=? AND object_key=?", (payload["resource_type"], payload["resource_id"], payload["version_id"], payload["object_key"]), ).fetchone() return dict(row) def storage_objects_for_resource(self, resource_type: str, resource_id: str) -> list[dict[str, Any]]: with self.connect() as conn: rows = conn.execute( "SELECT * FROM storage_objects WHERE resource_type=? AND resource_id=? AND status='available' ORDER BY version_id, object_key", (resource_type, resource_id), ).fetchall() return [dict(row) for row in rows] def update_storage_object(self, object_id: str, payload: dict[str, Any]) -> dict[str, Any]: allowed = {"status", "checksum_sha256", "byte_size", "content_type"} fields = {key: value for key, value in payload.items() if key in allowed} if fields: assignments = ", ".join(f"{key}=?" for key in fields) with self.connect() as conn: conn.execute(f"UPDATE storage_objects SET {assignments} WHERE id=?", (*fields.values(), object_id)) with self.connect() as conn: row = conn.execute("SELECT * FROM storage_objects WHERE id=?", (object_id,)).fetchone() if not row: raise KeyError(object_id) return dict(row) def create_storage_cache_job(self, payload: dict[str, Any]) -> dict[str, Any]: job_id = str(payload.get("id") or new_id("cache")) with self.connect() as conn: conn.execute( """ INSERT INTO storage_cache_jobs (id, storage_object_id, node_id, direction, status, progress, local_path, error, create_time, completed_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, (job_id, payload["storage_object_id"], payload["node_id"], payload.get("direction", "download"), payload.get("status", "running"), int(payload.get("progress", 0)), payload.get("local_path"), payload.get("error"), payload.get("create_time") or utcnow(), payload.get("completed_at")), ) row = conn.execute("SELECT * FROM storage_cache_jobs WHERE id=?", (job_id,)).fetchone() return dict(row) def update_storage_cache_job(self, job_id: str, payload: dict[str, Any]) -> dict[str, Any]: allowed = {"status", "progress", "local_path", "error", "completed_at"} fields = {key: value for key, value in payload.items() if key in allowed} if fields: assignments = ", ".join(f"{key}=?" for key in fields) with self.connect() as conn: conn.execute(f"UPDATE storage_cache_jobs SET {assignments} WHERE id=?", (*fields.values(), job_id)) with self.connect() as conn: row = conn.execute("SELECT * FROM storage_cache_jobs WHERE id=?", (job_id,)).fetchone() if not row: raise KeyError(job_id) return dict(row) def storage_cache_jobs_for_node(self, node_id: str, limit: int = 100) -> list[dict[str, Any]]: with self.connect() as conn: rows = conn.execute( "SELECT * FROM storage_cache_jobs WHERE node_id=? ORDER BY create_time DESC LIMIT ?", (node_id, max(1, min(limit, 500))), ).fetchall() return [dict(row) for row in rows] 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} # 评测任务同样占用算力节点，纳入运行任务统计 for row in conn.execute( "SELECT payload FROM eval_tasks WHERE status IN ('syncing','queued','running')" ).fetchall(): node_id = json_loads(row["payload"], {}).get("compute_node_id") if node_id: running_map[node_id] = running_map.get(node_id, 0) + 1 # 推理模型占用算力节点同样计入：优先从 compare_tasks 持久化状态派生 # （重启后仍准确），并用内存标记兜底（直接 preload 的模型无 compare 记录） inference_node_ids = set(self._inference_nodes) for ctr in conn.execute("SELECT payload FROM compare_tasks").fetchall(): ls = json_loads(ctr["payload"], {}).get("load_status") or {} if isinstance(ls, str): try: ls = json.loads(ls) except (json.JSONDecodeError, TypeError): ls = {} for m in ls.get("loaded_models") or []: if m.get("status") in {"ready", "running"} and m.get("node_id"): inference_node_ids.add(m["node_id"]) for nid in inference_node_ids: running_map[nid] = running_map.get(nid, 0) + 1 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 delete_compute_node(self, node_id: str) -> dict[str, Any]: with self.connect() as conn: node = conn.execute("SELECT * FROM compute_nodes WHERE id=?", (node_id,)).fetchone() if not node: raise KeyError(node_id) active = conn.execute( """ SELECT COUNT(*) AS cnt FROM fine_tune_tasks WHERE compute_node_id=? AND status IN ('syncing','queued','running') """, (node_id,), ).fetchone() if active and int(active["cnt"] or 0) > 0: raise ValueError("compute node has active training tasks") conn.execute("DELETE FROM compute_nodes WHERE id=?", (node_id,)) return {"deleted": 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() ] # 评测任务同样占用节点 GPU eval_running = [ json_loads(row["payload"], {}) for row in conn.execute( "SELECT payload FROM eval_tasks WHERE status IN ('syncing','queued','running')" ).fetchall() ] # 推理模型占用的节点：优先从 compare_tasks 持久化状态派生（重启后仍准确）， # 内存标记兜底（直接 preload 的模型无 compare 记录） inference_node_ids = set(self._inference_nodes) for ctr in conn.execute("SELECT payload FROM compare_tasks").fetchall(): ls = json_loads(ctr["payload"], {}).get("load_status") or {} if isinstance(ls, str): try: ls = json.loads(ls) except (json.JSONDecodeError, TypeError): ls = {} for m in ls.get("loaded_models") or []: if m.get("status") in {"ready", "running"} and m.get("node_id"): inference_node_ids.add(m["node_id"]) 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, ) eval_task = next( ( t for t in eval_running if t.get("compute_node_id") == row["node_id"] and row["gpu_index"] == (int(t["gpu_id"]) if t.get("gpu_id") is not None else -1) ), None, ) busy = (task is not None and task.get("status") == "running") or eval_task is not None reserved = (task is not None and task.get("status") in {"syncing", "queued"}) or ( eval_task is not None and eval_task.get("status") in {"syncing", "queued"} ) # Also mark GPU as busy if an inference model is loaded on this node if row["node_id"] in inference_node_ids and not busy: busy = True reserved = False 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 [ { "pid": int(eval_task.get("process_id") or 0), "name": "eval_runner", "memory_used_gb": memory_used, "task_name": eval_task.get("eval_task_name") or eval_task.get("name") or "评测任务", "user": "admin", } ] if eval_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]: # 轻量健康检查：采集真实 CPU/内存/磁盘使用率 # 用于顶部栏快速展示与 Docker 健康检查。 try: import psutil # cpu_percent(interval=None) 首次调用返回 0，需要短暂采样 cpu_percent = float(psutil.cpu_percent(interval=0.1)) memory_percent = float(psutil.virtual_memory().percent) # Windows 兼容：尝试当前盘符 try: disk_percent = float(psutil.disk_usage('/').percent) except Exception: disk_percent = float(psutil.disk_usage('C:\\').percent) except Exception: cpu_percent = memory_percent = disk_percent = 0.0 return { "cpu_percent": round(cpu_percent, 1), "memory_percent": round(memory_percent, 1), "disk_percent": round(disk_percent, 1), } 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"} # ===================== 平台治理：角色 ===================== def roles(self) -> list[dict[str, Any]]: with self.connect() as conn: rows = conn.execute("SELECT * FROM roles ORDER BY name").fetchall() return [dict(r) for r in rows] # ===================== 平台治理：审计日志 ===================== def audit_logs( self, *, tenant_id: str | None = None, project_id: str | None = None, actor_id: str | None = None, action: str | None = None, target_type: str | None = None, start_time: str | None = None, end_time: str | None = None, limit: int = 50, offset: int = 0, ) -> dict[str, Any]: clauses: list[str] = [] params: list[Any] = [] if tenant_id: clauses.append("tenant_id=?") params.append(tenant_id) if project_id: clauses.append("project_id=?") params.append(project_id) if actor_id: clauses.append("actor_id=?") params.append(actor_id) if action: clauses.append("action=?") params.append(action) if target_type: clauses.append("target_type=?") params.append(target_type) if start_time: clauses.append("time>=?") params.append(start_time) if end_time: clauses.append("time<=?") params.append(end_time) where = (" WHERE " + " AND ".join(clauses)) if clauses else "" with self.connect() as conn: total = conn.execute(f"SELECT COUNT(*) AS c FROM audit_logs{where}", tuple(params)).fetchone()["c"] params_paged = list(params) + [limit, offset] rows = conn.execute( f"SELECT * FROM audit_logs{where} ORDER BY time DESC LIMIT ? OFFSET ?", tuple(params_paged), ).fetchall() return {"total": total, "items": [dict(r) for r in rows]} def record_audit( self, *, action: str, actor_id: str | None = None, target_type: str | None = None, target_id: str | None = None, tenant_id: str | None = None, project_id: str | None = None, detail: str | None = None, ip: str | None = None, ) -> None: with self.connect() as conn: conn.execute( """ INSERT INTO audit_logs (id, tenant_id, project_id, actor_id, action, target_type, target_id, detail, client_ip, time) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( new_id("log"), tenant_id, project_id, actor_id, action, target_type, target_id, detail, ip, utcnow(), ), ) # ===================== 平台治理：会话 ===================== def create_session(self, user_id: str, *, ip: str | None = None) -> dict[str, Any]: sid = new_id("sess") login_at = utcnow() with self.connect() as conn: conn.execute( "INSERT INTO sessions (id, user_id, username, login_at, issued_at, expires_at, create_time) " "VALUES (%s, %s, (SELECT username FROM users WHERE id=%s), %s, %s, %s, %s)", (sid, user_id, user_id, login_at, login_at, datetime.fromtimestamp(time.time() + 1800, timezone.utc).isoformat(), login_at), ) return {"session_id": sid, "user_id": user_id, "login_at": login_at} def finish_session(self, session_id: str) -> None: """登出时记录 logout_at 与时长（秒）。""" logout_at = utcnow() with self.connect() as conn: conn.execute( "UPDATE sessions SET logout_at=%s, " "duration_seconds=EXTRACT(EPOCH FROM (%s::timestamptz - login_at::timestamptz))::int " "WHERE id=%s AND logout_at IS NULL", (logout_at, logout_at, session_id), ) def active_sessions(self, user_id: str) -> list[dict[str, Any]]: with self.connect() as conn: rows = conn.execute( "SELECT * FROM sessions WHERE user_id=%s AND logout_at IS NULL " "ORDER BY login_at DESC", (user_id,), ).fetchall() return [dict(r) for r in rows] def destroy_session(self, session_id: str) -> None: with self.connect() as conn: conn.execute("DELETE FROM sessions WHERE id=%s", (session_id,)) def extend_session(self, session_id: str, *, expires_in_seconds: int = 3600 * 8) -> dict[str, Any] | None: # 兼容旧调用，仅更新 login_at 之后延长的含义在此简化为 no-op 返回现有记录。 with self.connect() as conn: row = conn.execute("SELECT * FROM sessions WHERE id=%s", (session_id,)).fetchone() if not row: return None return { "session_id": session_id, "user_id": row["user_id"], "login_at": row["login_at"], } def set_session_user(self, session_id: str, user_id: str) -> None: with self.connect() as conn: conn.execute("UPDATE sessions SET user_id=%s WHERE id=%s", (user_id, session_id)) def login_duration_rank(self, limit: int = 8, days: int = 30) -> list[dict[str, Any]]: """登录时长排行：按用户聚合近 N 天的会话时长（小时）。 sessions 表列：login_at(TEXT), logout_at(TEXT), duration_seconds(INT)。 优先用 duration_seconds；为空时回退计算 now-login_at（未登出）或 logout_at-login_at。 """ with self.connect() as conn: rows = conn.execute( "SELECT s.user_id, s.login_at, s.logout_at, s.duration_seconds, " "u.username, u.display_name, u.role " "FROM sessions s LEFT JOIN users u ON s.user_id = u.id " "WHERE s.login_at::timestamptz >= NOW() - make_interval(days => %s)", (days,), ).fetchall() now = datetime.now(timezone.utc) agg: dict[str, dict[str, Any]] = {} for r in rows: uid = r["user_id"] or "" bucket = agg.setdefault( uid, { "user": r["display_name"] or r["username"] or uid, "role": r["role"] or "", "total": 0.0, }, ) dur = r["duration_seconds"] if dur is not None: bucket["total"] += float(dur) continue start = parse_time(r["login_at"]) end = parse_time(r["logout_at"]) if r["logout_at"] else None if start and end: bucket["total"] += max(0, (end - start).total_seconds()) elif start: bucket["total"] += max(0, (now - start).total_seconds()) result = [ {"user": b["user"], "role": b["role"], "duration": round(b["total"] / 3600, 1)} for b in agg.values() ] result.sort(key=lambda x: x["duration"], reverse=True) return result[:limit] # ===================== 平台治理：审批 ===================== def create_approval_template(self, payload: dict[str, Any]) -> dict[str, Any]: with self.connect() as conn: tid = new_id("tpl") conn.execute( "INSERT INTO approval_templates (id, name, steps, create_time) VALUES (?, ?, ?, ?)", (tid, payload["name"], json_dumps(payload.get("steps", [])), utcnow()), ) return self.approval_template(tid) def approval_templates(self) -> list[dict[str, Any]]: with self.connect() as conn: rows = conn.execute("SELECT * FROM approval_templates ORDER BY create_time DESC").fetchall() return [dict(r) for r in rows] def approval_template(self, template_id: str) -> dict[str, Any]: with self.connect() as conn: row = conn.execute("SELECT * FROM approval_templates WHERE id=?", (template_id,)).fetchone() if not row: raise KeyError(template_id) return dict(row) def update_approval_template(self, template_id: str, payload: dict[str, Any]) -> dict[str, Any]: fields = {k: v for k, v in payload.items() if k in ("name", "steps")} if "steps" in fields: fields["steps"] = json_dumps(fields["steps"]) if not fields: return self.approval_template(template_id) set_clause = ", ".join(f"{k}=?" for k in fields) params = list(fields.values()) + [template_id] with self.connect() as conn: conn.execute(f"UPDATE approval_templates SET {set_clause} WHERE id=?", tuple(params)) return self.approval_template(template_id) def delete_approval_template(self, template_id: str) -> dict[str, Any]: with self.connect() as conn: row = conn.execute("SELECT * FROM approval_templates WHERE id=?", (template_id,)).fetchone() if not row: raise KeyError(template_id) conn.execute("DELETE FROM approval_templates WHERE id=?", (template_id,)) return dict(row) def create_approval_instance(self, payload: dict[str, Any]) -> dict[str, Any]: template_id = payload.get("template_id") steps = [] if template_id: tpl = self.approval_template(template_id) steps = json_loads(tpl["steps"]) if tpl.get("steps") else [] with self.connect() as conn: iid = new_id("appr") conn.execute( """ INSERT INTO approval_instances (id, template_id, resource_type, resource_id, applicant_id, status, current_step, create_time) VALUES (?, ?, ?, ?, ?, ?, ?, ?) """, ( iid, template_id, payload["resource_type"], payload["resource_id"], payload["applicant_id"], "pending", 0, utcnow(), ), ) for idx, step in enumerate(steps): conn.execute( "INSERT INTO approval_steps (id, instance_id, step_index, approver_id, status, time) VALUES (?, ?, ?, ?, ?, ?)", (new_id("step"), iid, idx, step.get("approver_id"), "pending", None), ) return self.approval_instance(iid) def approval_instances(self, *, status: str | None = None) -> list[dict[str, Any]]: with self.connect() as conn: if status: rows = conn.execute( "SELECT * FROM approval_instances WHERE status=? ORDER BY create_time DESC", (status,) ).fetchall() else: rows = conn.execute("SELECT * FROM approval_instances ORDER BY create_time DESC").fetchall() return [dict(r) for r in rows] def approval_instance(self, instance_id: str) -> dict[str, Any]: with self.connect() as conn: row = conn.execute("SELECT * FROM approval_instances WHERE id=?", (instance_id,)).fetchone() if not row: raise KeyError(instance_id) steps = conn.execute( "SELECT * FROM approval_steps WHERE instance_id=? ORDER BY step_index", (instance_id,) ).fetchall() result = dict(row) result["steps"] = [dict(s) for s in steps] return result def decide_approval_step(self, instance_id: str, step_index: int, *, approver_id: str, approved: bool, comment: str | None = None) -> dict[str, Any]: with self.connect() as conn: inst = conn.execute("SELECT * FROM approval_instances WHERE id=?", (instance_id,)).fetchone() if not inst: raise KeyError(instance_id) if inst["status"] != "pending": raise ValueError("instance not pending") step = conn.execute( "SELECT * FROM approval_steps WHERE instance_id=? AND step_index=?", (instance_id, step_index), ).fetchone() if not step: raise KeyError("step not found") if step["status"] != "pending": raise ValueError("step already decided") new_status = "approved" if approved else "rejected" conn.execute( "UPDATE approval_steps SET status=?, comment=?, time=? WHERE id=?", (new_status, comment, utcnow(), step["id"]), ) if approved: conn.execute( "UPDATE approval_instances SET current_step=? WHERE id=?", (step_index + 1, instance_id), ) step_rows = conn.execute( "SELECT * FROM approval_steps WHERE instance_id=? ORDER BY step_index", (instance_id,) ).fetchall() if all(s["status"] == "approved" for s in step_rows): conn.execute("UPDATE approval_instances SET status='approved' WHERE id=?", (instance_id,)) else: conn.execute("UPDATE approval_instances SET status='rejected' WHERE id=?", (instance_id,)) return self.approval_instance(instance_id) # ===================== 平台治理：租户 ===================== def tenants(self) -> list[dict[str, Any]]: with self.connect() as conn: rows = conn.execute("SELECT * FROM tenants ORDER BY create_time DESC").fetchall() return [dict(r) for r in rows] def tenant(self, tenant_id: str) -> dict[str, Any]: with self.connect() as conn: row = conn.execute("SELECT * FROM tenants WHERE id=?", (tenant_id,)).fetchone() if not row: raise KeyError(tenant_id) return dict(row) def create_tenant(self, payload: dict[str, Any]) -> dict[str, Any]: with self.connect() as conn: tid = new_id("tnt") conn.execute( """ INSERT INTO tenants (id, name, code, status, owner_user_id, quota, retention_policy_id, create_time) VALUES (?, ?, ?, ?, ?, ?, ?, ?) """, ( tid, payload["name"], payload.get("code"), "active", payload.get("owner_user_id"), json_dumps(payload.get("quota", {})), payload.get("retention_policy_id"), utcnow(), ), ) return self.tenant(tid) def update_tenant(self, tenant_id: str, payload: dict[str, Any]) -> dict[str, Any]: fields = {k: v for k, v in payload.items() if k in ("name", "code", "status", "owner_user_id", "quota", "retention_policy_id")} if "quota" in fields: fields["quota"] = json_dumps(fields["quota"]) if not fields: return self.tenant(tenant_id) set_clause = ", ".join(f"{k}=?" for k in fields) params = list(fields.values()) + [tenant_id] with self.connect() as conn: conn.execute(f"UPDATE tenants SET {set_clause} WHERE id=?", tuple(params)) return self.tenant(tenant_id) def set_tenant_quota(self, tenant_id: str, quota: dict[str, Any]) -> dict[str, Any]: with self.connect() as conn: conn.execute("UPDATE tenants SET quota=? WHERE id=?", (json_dumps(quota), tenant_id)) return self.tenant(tenant_id) def set_tenant_retention(self, tenant_id: str, retention_policy_id: str | None) -> dict[str, Any]: with self.connect() as conn: conn.execute("UPDATE tenants SET retention_policy_id=? WHERE id=?", (retention_policy_id, tenant_id)) return self.tenant(tenant_id) def delete_tenant(self, tenant_id: str) -> dict[str, Any]: with self.connect() as conn: row = conn.execute("SELECT * FROM tenants WHERE id=?", (tenant_id,)).fetchone() if not row: raise KeyError(tenant_id) conn.execute("DELETE FROM tenants WHERE id=?", (tenant_id,)) return dict(row) def get_acl(self, resource_type: str, resource_id: str) -> list[dict[str, Any]]: with self.connect() as conn: rows = conn.execute( "SELECT * FROM acls WHERE resource_type=? AND resource_id=?", (resource_type, resource_id), ).fetchall() return [dict(r) for r in rows] def set_acl(self, resource_type: str, resource_id: str, entries: list[dict[str, Any]]) -> list[dict[str, Any]]: with self.connect() as conn: conn.execute( "DELETE FROM acls WHERE resource_type=? AND resource_id=?", (resource_type, resource_id), ) for e in entries: conn.execute( """ INSERT INTO acls (id, resource_type, resource_id, principal_type, principal_id, permission, create_time) VALUES (?, ?, ?, ?, ?, ?, ?) """, ( new_id("acl"), resource_type, resource_id, e.get("principal_type"), e.get("principal_id"), e.get("permission"), utcnow(), ), ) rows = conn.execute( "SELECT * FROM acls WHERE resource_type=? AND resource_id=?", (resource_type, resource_id), ).fetchall() return [dict(r) for r in rows] # ===================== 平台治理：项目空间 ===================== def projects(self, *, tenant_id: str = "default", status: str | None = None, keyword: str | None = None) -> list[dict[str, Any]]: clauses = ["tenant_id=?"] params: list[Any] = [tenant_id] if status: clauses.append("status=?") params.append(status) if keyword: clauses.append("(name LIKE ? OR code LIKE ?)") params.extend([f"%{keyword}%", f"%{keyword}%"]) with self.connect() as conn: rows = conn.execute( f"SELECT * FROM projects WHERE {' AND '.join(clauses)} ORDER BY create_time DESC", tuple(params), ).fetchall() return [dict(r) for r in rows] def project(self, project_id: str) -> dict[str, Any]: with self.connect() as conn: row = conn.execute("SELECT * FROM projects WHERE id=?", (project_id,)).fetchone() if not row: raise KeyError(project_id) return dict(row) def create_project(self, payload: dict[str, Any]) -> dict[str, Any]: with self.connect() as conn: pid = new_id("prj") conn.execute( """ INSERT INTO projects (id, tenant_id, name, code, description, quota, status, create_time, create_by, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( pid, payload.get("tenant_id", "default"), payload["name"], payload["code"], payload.get("description"), json_dumps(payload.get("quota", {})), "active", utcnow(), payload.get("create_by"), utcnow(), ), ) return self.project(pid) def update_project(self, project_id: str, payload: dict[str, Any]) -> dict[str, Any]: fields = {k: v for k, v in payload.items() if k in ("name", "code", "description", "quota", "status")} if "quota" in fields: fields["quota"] = json_dumps(fields["quota"]) if not fields: return self.project(project_id) set_clause = ", ".join(f"{k}=?" for k in fields) params = list(fields.values()) + [project_id] with self.connect() as conn: conn.execute(f"UPDATE projects SET {set_clause} WHERE id=?", tuple(params)) return self.project(project_id) def archive_project(self, project_id: str) -> dict[str, Any]: with self.connect() as conn: conn.execute("UPDATE projects SET status='archived' WHERE id=?", (project_id,)) return self.project(project_id) def activate_project(self, project_id: str) -> dict[str, Any]: with self.connect() as conn: conn.execute("UPDATE projects SET status='active' WHERE id=?", (project_id,)) return self.project(project_id) def delete_project(self, project_id: str) -> None: with self.connect() as conn: conn.execute("DELETE FROM projects WHERE id=?", (project_id,)) def project_members(self, project_id: str) -> list[dict[str, Any]]: with self.connect() as conn: row = conn.execute("SELECT * FROM projects WHERE id=?", (project_id,)).fetchone() if not row: raise KeyError(project_id) rows = conn.execute( """ SELECT pm.*, u.username, u.display_name FROM project_members pm JOIN users u ON u.id = pm.user_id WHERE pm.project_id=? """, (project_id,), ).fetchall() return [dict(r) for r in rows] def add_project_member(self, project_id: str, payload: dict[str, Any]) -> dict[str, Any]: user_id = payload["user_id"] role = payload.get("role", "member") with self.connect() as conn: conn.execute( "INSERT INTO project_members (project_id, user_id, role, create_time) VALUES (?, ?, ?, ?)", (project_id, user_id, role, utcnow()), ) row = conn.execute( """ SELECT pm.*, u.username, u.display_name FROM project_members pm JOIN users u ON u.id = pm.user_id WHERE pm.project_id=? AND pm.user_id=? """, (project_id, user_id), ).fetchone() return { "project_id": row["project_id"], "user_id": row["user_id"], "username": row["username"], "display_name": row["display_name"], "role": row["role"], "create_time": row["create_time"], } def update_project_member_role(self, project_id: str, user_id: str, role: str) -> dict[str, Any]: with self.connect() as conn: conn.execute( "UPDATE project_members SET role=? WHERE project_id=? AND user_id=?", (role, project_id, user_id), ) row = conn.execute( """ SELECT pm.*, u.username, u.display_name FROM project_members pm JOIN users u ON u.id = pm.user_id WHERE pm.project_id=? AND pm.user_id=? """, (project_id, user_id), ).fetchone() if not row: raise KeyError(user_id) return { "project_id": row["project_id"], "user_id": row["user_id"], "username": row["username"], "display_name": row["display_name"], "role": row["role"], "create_time": row["create_time"], } def remove_project_member(self, project_id: str, user_id: str) -> None: with self.connect() as conn: conn.execute( "DELETE FROM project_members WHERE project_id=? AND user_id=?", (project_id, user_id), ) # ===================== 平台治理：资源 ACL ===================== def resource_acl(self, resource_type: str, resource_id: str) -> list[dict[str, Any]]: """返回资源 ACL，按主体分组，permissions 为数组。""" rows = self.get_acl(resource_type, resource_id) grouped: dict[str, dict[str, Any]] = {} for r in rows: key = f"{r.get('principal_type')}:{r.get('principal_id')}" bucket = grouped.setdefault( key, { "subject_type": r.get("principal_type"), "subject_id": r.get("principal_id"), "permissions": [], }, ) perm = r.get("permission") if perm and perm not in bucket["permissions"]: bucket["permissions"].append(perm) return list(grouped.values()) def set_resource_acl( self, resource_type: str, resource_id: str, entries: list[dict[str, Any]] ) -> list[dict[str, Any]]: """按前端格式设置资源 ACL：entries 为 [{subject_type, subject_id, permissions: []}]。""" flat: list[dict[str, Any]] = [] for e in entries: for perm in e.get("permissions") or []: flat.append( { "principal_type": e.get("subject_type"), "principal_id": e.get("subject_id"), "permission": perm, } ) self.set_acl(resource_type, resource_id, flat) return self.resource_acl(resource_type, resource_id) # ===================== 平台治理：留存策略 ===================== def retention_policies(self) -> list[dict[str, Any]]: with self.connect() as conn: rows = conn.execute( "SELECT * FROM retention_policies ORDER BY create_time DESC" ).fetchall() return [dict(r) for r in rows] def retention_policy(self, policy_id: str) -> dict[str, Any]: with self.connect() as conn: row = conn.execute( "SELECT * FROM retention_policies WHERE id=?", (policy_id,) ).fetchone() if not row: raise KeyError(policy_id) return dict(row) def create_retention_policy(self, payload: dict[str, Any]) -> dict[str, Any]: pid = payload.get("id") or new_id("rpol") with self.connect() as conn: conn.execute( """ INSERT INTO retention_policies (id, name, scope, rule, status, create_time, create_by, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?) """, ( pid, payload["name"], payload.get("scope"), payload.get("rule"), payload.get("status", "active"), utcnow(), payload.get("create_by"), utcnow(), ), ) return self.retention_policy(pid) def update_retention_policy( self, policy_id: str, payload: dict[str, Any] ) -> dict[str, Any]: fields = { k: v for k, v in payload.items() if k in ("name", "scope", "rule", "status") } if not fields: return self.retention_policy(policy_id) fields["updated_at"] = utcnow() set_clause = ", ".join(f"{k}=?" for k in fields) params = list(fields.values()) + [policy_id] with self.connect() as conn: conn.execute( f"UPDATE retention_policies SET {set_clause} WHERE id=?", tuple(params), ) return self.retention_policy(policy_id) def delete_retention_policy(self, policy_id: str) -> None: with self.connect() as conn: conn.execute( "DELETE FROM retention_policies WHERE id=?", (policy_id,) ) # ===================== 平台治理：GPU 算力分配 ===================== def gpu_assignments(self) -> list[dict[str, Any]]: """查询全部分配关系。""" with self.connect() as conn: rows = conn.execute( """ SELECT ga.*, u.username, u.display_name, n.code AS node_code, n.name AS node_name, g.name AS gpu_name FROM gpu_assignments ga LEFT JOIN users u ON u.id = ga.user_id LEFT JOIN compute_nodes n ON n.id = ga.node_id LEFT JOIN gpus g ON g.node_id = ga.node_id AND g.gpu_index = ga.gpu_index ORDER BY ga.assigned_at DESC """ ).fetchall() return [dict(r) for r in rows] def gpu_assignments_for_user(self, user_id: str) -> list[dict[str, Any]]: """查询某用户被分配的 GPU 列表。""" with self.connect() as conn: rows = conn.execute( """ SELECT ga.node_id, ga.gpu_index, n.code AS node_code, n.name AS node_name, g.name AS gpu_name, g.uuid, g.memory_total_gb FROM gpu_assignments ga JOIN compute_nodes n ON n.id = ga.node_id LEFT JOIN gpus g ON g.node_id = ga.node_id AND g.gpu_index = ga.gpu_index WHERE ga.user_id = ? ORDER BY n.code, ga.gpu_index """, (user_id,), ).fetchall() return [dict(r) for r in rows] def assign_gpus(self, assignments: list[dict[str, Any]], assigned_by: str | None = None) -> list[dict[str, Any]]: """批量分配 GPU（幂等：已存在的分配跳过）。""" now = utcnow() with self.connect() as conn: for a in assignments: node_id = a["node_id"] gpu_index = a["gpu_index"] user_id = a["user_id"] existing = conn.execute( "SELECT id FROM gpu_assignments WHERE node_id=? AND gpu_index=? AND user_id=?", (node_id, gpu_index, user_id), ).fetchone() if existing: continue aid = new_id("ga") conn.execute( """ INSERT INTO gpu_assignments (id, node_id, gpu_index, user_id, assigned_by, assigned_at) VALUES (?, ?, ?, ?, ?, ?) """, (aid, node_id, gpu_index, user_id, assigned_by, now), ) return self.gpu_assignments() def unassign_gpu(self, assignment_id: str) -> None: with self.connect() as conn: conn.execute("DELETE FROM gpu_assignments WHERE id=?", (assignment_id,)) def check_gpu_access(self, user_id: str, node_id: str, gpu_indices: list[int]) -> bool: """检查用户是否被分配了指定节点的指定 GPU 卡。""" if not gpu_indices: return True with self.connect() as conn: rows = conn.execute( """ SELECT gpu_index FROM gpu_assignments WHERE user_id=? AND node_id=? """, (user_id, node_id), ).fetchall() assigned = {r["gpu_index"] for r in rows} return all(idx in assigned for idx in gpu_indices) def assigned_gpu_indexes(self, user_id: str, node_id: str) -> list[int]: with self.connect() as conn: rows = conn.execute( "SELECT gpu_index FROM gpu_assignments WHERE user_id=? AND node_id=? ORDER BY gpu_index", (user_id, node_id), ).fetchall() return [int(row["gpu_index"]) for row in rows] # ===================== 平台治理：资源可见性过滤 ===================== def _filter_accessible_ids( self, resource_type: str, all_ids: list[str], user: dict[str, Any] ) -> list[str]: """从全部资源 ID 中过滤出当前用户可访问的 ID 列表。 - admin 直接返回全部。 - 资源所有者可见（需调用方在 all_ids 中提供 owned ids）。 - ACL 授权的用户/角色可见。 """ if user.get("role") == "admin" or user.get("protected"): return all_ids if not all_ids: return [] user_id = user.get("id") user_role = user.get("role") with self.connect() as conn: rows = conn.execute( """ SELECT DISTINCT resource_id FROM acls WHERE resource_type=? AND ( (principal_type='user' AND principal_id=?) OR (principal_type='role' AND principal_id=?) ) """, (resource_type, user_id, user_role), ).fetchall() accessible = {r["resource_id"] for r in rows} return [rid for rid in all_ids if rid in accessible] def change_password(self, user_id: str, old_password: str, new_password: str) -> bool: """用户自行修改密码：验证旧密码后设置新密码。""" with self.connect() as conn: row = conn.execute("SELECT * FROM users WHERE id=?", (user_id,)).fetchone() if not row: raise KeyError(user_id) matched, _ = verify_password(old_password, row["password_hash"]) if not matched: return False conn.execute( "UPDATE users SET password_hash=? WHERE id=?", (hash_password(new_password), user_id), ) return True _store: PlatformStore | None = None def get_platform_store() -> PlatformStore: global _store if _store is None: _store = PlatformStore() return _store import atexit as _atexit def _close_store_pool() -> None: global _store if _store is not None: _store.close_pool() _store = None _atexit.register(_close_store_pool)