2026-07-21 10:55:44 +08:00
|
|
|
|
from __future__ import annotations
|
2026-07-21 09:23:43 +08:00
|
|
|
|
|
2026-07-21 10:55:44 +08:00
|
|
|
|
import hashlib
|
|
|
|
|
|
import hmac
|
2026-07-27 12:26:09 +08:00
|
|
|
|
import json
|
2026-07-21 09:23:43 +08:00
|
|
|
|
import math
|
2026-07-23 19:32:42 +08:00
|
|
|
|
import re
|
2026-07-21 10:55:44 +08:00
|
|
|
|
import secrets
|
2026-07-21 09:23:43 +08:00
|
|
|
|
import time
|
|
|
|
|
|
import uuid
|
|
|
|
|
|
from contextlib import contextmanager
|
2026-07-23 19:32:42 +08:00
|
|
|
|
from datetime import datetime, timedelta, timezone
|
2026-07-21 09:23:43 +08:00
|
|
|
|
from pathlib import Path
|
|
|
|
|
|
from typing import Any, Iterator
|
|
|
|
|
|
|
2026-07-21 10:55:44 +08:00
|
|
|
|
import psycopg
|
2026-08-03 09:34:08 +08:00
|
|
|
|
from psycopg_pool import ConnectionPool
|
2026-07-21 10:55:44 +08:00
|
|
|
|
|
2026-07-21 09:23:43 +08:00
|
|
|
|
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=(",", ":"))
|
|
|
|
|
|
|
|
|
|
|
|
|
2026-07-22 17:32:59 +08:00
|
|
|
|
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
|
|
|
|
|
|
|
|
|
|
|
|
|
2026-07-27 12:26:09 +08:00
|
|
|
|
_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]))
|
|
|
|
|
|
|
|
|
|
|
|
|
2026-08-03 15:49:21 +08:00
|
|
|
|
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()])
|
|
|
|
|
|
|
|
|
|
|
|
|
2026-08-03 17:34:21 +08:00
|
|
|
|
_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)
|
|
|
|
|
|
|
|
|
|
|
|
|
2026-07-27 12:26:09 +08:00
|
|
|
|
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),
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
|
2026-07-23 19:32:42 +08:00
|
|
|
|
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] = {}
|
2026-08-04 16:59:34 +08:00
|
|
|
|
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))
|
2026-07-23 19:32:42 +08:00
|
|
|
|
for key in ["loss", "grad_norm", "learning_rate", "epoch"]:
|
2026-08-04 16:59:34 +08:00
|
|
|
|
match = re.search(rf"['\"]?{key}['\"]?\s*(?:=|:)\s*{number_pattern}", line, re.I)
|
2026-07-23 19:32:42 +08:00
|
|
|
|
if match:
|
|
|
|
|
|
result[key] = float(match.group(1))
|
|
|
|
|
|
return result or None
|
|
|
|
|
|
|
|
|
|
|
|
|
2026-07-21 09:23:43 +08:00
|
|
|
|
def new_id(prefix: str) -> str:
|
|
|
|
|
|
return f"{prefix}_{uuid.uuid4().hex[:12]}"
|
|
|
|
|
|
|
|
|
|
|
|
|
2026-07-23 19:32:42 +08:00
|
|
|
|
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] = {}
|
2026-07-28 13:10:53 +08:00
|
|
|
|
fmt = str(formatting).lower()
|
2026-07-23 19:32:42 +08:00
|
|
|
|
for key, file_name in zip(llama_dataset_keys(dataset_key, file_names), file_names):
|
2026-07-28 13:10:53 +08:00
|
|
|
|
if fmt == "sharegpt":
|
2026-07-23 19:32:42 +08:00
|
|
|
|
result[key] = {
|
|
|
|
|
|
"file_name": file_name,
|
|
|
|
|
|
"formatting": "sharegpt",
|
|
|
|
|
|
"columns": {"messages": "messages"},
|
|
|
|
|
|
}
|
2026-07-28 13:10:53 +08:00
|
|
|
|
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",
|
|
|
|
|
|
},
|
|
|
|
|
|
}
|
2026-07-23 19:32:42 +08:00
|
|
|
|
return result
|
|
|
|
|
|
|
|
|
|
|
|
|
2026-07-21 10:55:44 +08:00
|
|
|
|
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()
|
|
|
|
|
|
|
|
|
|
|
|
|
2026-07-21 09:23:43 +08:00
|
|
|
|
class PlatformStore:
|
2026-07-21 10:55:44 +08:00
|
|
|
|
"""PostgreSQL-backed store for the first runnable platform version.
|
2026-07-21 09:23:43 +08:00
|
|
|
|
|
2026-07-21 10:55:44 +08:00
|
|
|
|
This store mirrors the API-facing subset needed by the first system
|
|
|
|
|
|
iteration while using the same PostgreSQL dependency as later production
|
|
|
|
|
|
development.
|
2026-07-21 09:23:43 +08:00
|
|
|
|
"""
|
|
|
|
|
|
|
2026-07-21 10:55:44 +08:00
|
|
|
|
def __init__(self, database_url: str | None = None) -> None:
|
2026-07-21 09:23:43 +08:00
|
|
|
|
settings = get_settings()
|
2026-07-21 10:55:44 +08:00
|
|
|
|
self.database_url = _psycopg_url(database_url or settings.database_url)
|
2026-08-03 09:34:08 +08:00
|
|
|
|
# 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 = {
|
2026-08-03 15:49:21 +08:00
|
|
|
|
"connect_timeout": 5,
|
2026-08-03 09:34:08 +08:00
|
|
|
|
"keepalives": 1,
|
2026-08-03 16:20:21 +08:00
|
|
|
|
"keepalives_idle": 10,
|
|
|
|
|
|
"keepalives_interval": 5,
|
|
|
|
|
|
"keepalives_count": 3,
|
2026-08-03 09:34:08 +08:00
|
|
|
|
}
|
|
|
|
|
|
self._pool = ConnectionPool(
|
|
|
|
|
|
conninfo=self.database_url,
|
|
|
|
|
|
kwargs=pool_kwargs,
|
|
|
|
|
|
min_size=2,
|
|
|
|
|
|
max_size=10,
|
|
|
|
|
|
# 借出前校验连接可用性,避免执行 SQL 时才发现 [BAD] 再重建。
|
|
|
|
|
|
check=ConnectionPool.check_connection,
|
|
|
|
|
|
# 不主动回收空闲连接(远程库约 10s 断,由 keepalive 维持),
|
|
|
|
|
|
# 减少无谓的重建握手。
|
|
|
|
|
|
max_idle=0,
|
|
|
|
|
|
# 请求最多排队等待 5s,避免雪崩时无限堆积。
|
|
|
|
|
|
max_waiting=16,
|
|
|
|
|
|
open=False,
|
|
|
|
|
|
)
|
|
|
|
|
|
# 注意:不要在此调用 pool.wait(),它会阻塞等待 min_size 个连接就绪,
|
|
|
|
|
|
# 在远程库响应慢/超时时会卡死 uvicorn worker 进程,导致所有请求无响应。
|
|
|
|
|
|
self._pool.open()
|
2026-07-21 09:23:43 +08:00
|
|
|
|
self.ensure_schema()
|
|
|
|
|
|
self.ensure_seed_data()
|
2026-07-28 17:29:16 +08:00
|
|
|
|
# Track which compute nodes have an active inference model loaded
|
|
|
|
|
|
self._inference_nodes: set[str] = set()
|
2026-08-03 17:34:21 +08:00
|
|
|
|
self._last_runtime_refresh = 0.0
|
2026-07-28 17:29:16 +08:00
|
|
|
|
|
|
|
|
|
|
# ── 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
|
2026-07-21 09:23:43 +08:00
|
|
|
|
|
|
|
|
|
|
@contextmanager
|
2026-07-21 10:55:44 +08:00
|
|
|
|
def connect(self) -> Iterator["PgConnection"]:
|
2026-08-03 09:34:08 +08:00
|
|
|
|
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."""
|
2026-07-21 09:23:43 +08:00
|
|
|
|
try:
|
2026-08-03 09:34:08 +08:00
|
|
|
|
self._pool.close()
|
2026-07-21 10:55:44 +08:00
|
|
|
|
except Exception:
|
2026-08-03 09:34:08 +08:00
|
|
|
|
pass
|
2026-07-21 09:23:43 +08:00
|
|
|
|
|
|
|
|
|
|
def ensure_schema(self) -> None:
|
2026-07-21 10:55:44 +08:00
|
|
|
|
schema_path = Path(__file__).with_name("sql") / "001_platform_runtime.sql"
|
2026-07-21 09:23:43 +08:00
|
|
|
|
with self.connect() as conn:
|
2026-07-21 10:55:44 +08:00
|
|
|
|
conn.executescript(schema_path.read_text(encoding="utf-8"))
|
2026-07-22 17:32:59 +08:00
|
|
|
|
user_columns = self._column_names(conn, "users")
|
|
|
|
|
|
if "password" in user_columns and "password_hash" not in user_columns:
|
2026-07-21 10:55:44 +08:00
|
|
|
|
conn.execute("ALTER TABLE users RENAME COLUMN password TO password_hash")
|
2026-07-22 17:32:59 +08:00
|
|
|
|
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"})
|
2026-08-04 16:59:34 +08:00
|
|
|
|
self._ensure_columns(
|
|
|
|
|
|
conn,
|
|
|
|
|
|
"trained_models",
|
|
|
|
|
|
{
|
|
|
|
|
|
"artifact_dir": "TEXT",
|
|
|
|
|
|
"compute_node_id": "TEXT",
|
|
|
|
|
|
"compute_node_name": "TEXT",
|
|
|
|
|
|
},
|
|
|
|
|
|
)
|
2026-07-23 19:32:42 +08:00
|
|
|
|
self._ensure_columns(
|
|
|
|
|
|
conn,
|
|
|
|
|
|
"resource_replicas",
|
|
|
|
|
|
{
|
|
|
|
|
|
"checksum_sha256": "TEXT",
|
|
|
|
|
|
"byte_size": "BIGINT NOT NULL DEFAULT 0",
|
|
|
|
|
|
"last_checked_at": "TEXT",
|
|
|
|
|
|
"last_error": "TEXT",
|
|
|
|
|
|
},
|
|
|
|
|
|
)
|
2026-08-03 09:34:08 +08:00
|
|
|
|
schema_dir = Path(__file__).with_name("sql")
|
|
|
|
|
|
for extra in ("002_governance.sql", "003_tenant_quota.sql"):
|
|
|
|
|
|
extra_path = schema_dir / extra
|
|
|
|
|
|
if extra_path.exists():
|
|
|
|
|
|
conn.executescript(extra_path.read_text(encoding="utf-8"))
|
2026-07-22 17:32:59 +08:00
|
|
|
|
|
|
|
|
|
|
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}")
|
2026-07-21 09:23:43 +08:00
|
|
|
|
|
|
|
|
|
|
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
|
2026-07-21 10:55:44 +08:00
|
|
|
|
(id, username, password_hash, display_name, role, status, permissions, create_time, protected)
|
2026-07-21 09:23:43 +08:00
|
|
|
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
|
|
|
|
|
|
""",
|
2026-07-21 10:55:44 +08:00
|
|
|
|
[(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],
|
2026-07-21 09:23:43 +08:00
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
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:
|
2026-07-21 10:55:44 +08:00
|
|
|
|
if get_settings().compute_mode != "simulator":
|
|
|
|
|
|
return
|
2026-08-03 17:34:21 +08:00
|
|
|
|
now_ts = time.monotonic()
|
|
|
|
|
|
if now_ts - self._last_runtime_refresh < 2:
|
|
|
|
|
|
return
|
|
|
|
|
|
self._last_runtime_refresh = now_ts
|
2026-07-21 10:55:44 +08:00
|
|
|
|
|
2026-07-21 09:23:43 +08:00
|
|
|
|
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"]),
|
|
|
|
|
|
)
|
|
|
|
|
|
|
2026-07-28 13:10:53 +08:00
|
|
|
|
def _ensure_trained_model(self, conn: PgConnection, task: dict[str, Any], job: dict[str, Any] | None = None) -> None:
|
2026-07-21 09:23:43 +08:00
|
|
|
|
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:
|
2026-08-04 16:59:34 +08:00
|
|
|
|
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"]),
|
|
|
|
|
|
)
|
2026-07-21 09:23:43 +08:00
|
|
|
|
return
|
|
|
|
|
|
model = conn.execute("SELECT path FROM models WHERE id=?", (task.get("base_model"),)).fetchone()
|
2026-07-23 19:32:42 +08:00
|
|
|
|
output_dir = task.get("output_dir") or f"/data/yg-ft/outputs/{task['name']}"
|
|
|
|
|
|
trained_model_id = new_id("tm")
|
2026-07-21 09:23:43 +08:00
|
|
|
|
conn.execute(
|
|
|
|
|
|
"""
|
|
|
|
|
|
INSERT INTO trained_models
|
2026-08-04 16:59:34 +08:00
|
|
|
|
(id, name, train_methods, base_model_path, create_time, merged, merging, merged_path, artifact_dir, compute_node_id, compute_node_name)
|
|
|
|
|
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
2026-07-21 09:23:43 +08:00
|
|
|
|
""",
|
|
|
|
|
|
(
|
2026-07-23 19:32:42 +08:00
|
|
|
|
trained_model_id,
|
2026-07-21 09:23:43 +08:00
|
|
|
|
name,
|
|
|
|
|
|
json_dumps([{"name": task.get("train_method", "lora")}]),
|
|
|
|
|
|
model["path"] if model else "",
|
|
|
|
|
|
utcnow(),
|
|
|
|
|
|
0,
|
|
|
|
|
|
0,
|
2026-07-23 19:32:42 +08:00
|
|
|
|
output_dir,
|
2026-07-28 13:10:53 +08:00
|
|
|
|
output_dir,
|
2026-08-04 16:59:34 +08:00
|
|
|
|
task.get("compute_node_id"),
|
|
|
|
|
|
task.get("compute_node_code") or task.get("compute_node_name"),
|
2026-07-23 19:32:42 +08:00
|
|
|
|
),
|
|
|
|
|
|
)
|
2026-07-28 13:10:53 +08:00
|
|
|
|
# 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
|
2026-07-23 19:32:42 +08:00
|
|
|
|
self._upsert_model_artifact(
|
|
|
|
|
|
conn,
|
|
|
|
|
|
trained_model_id,
|
|
|
|
|
|
"trained_model",
|
|
|
|
|
|
"adapter",
|
|
|
|
|
|
output_dir,
|
2026-07-28 13:10:53 +08:00
|
|
|
|
total_size,
|
|
|
|
|
|
combined_checksum,
|
2026-07-23 19:32:42 +08:00
|
|
|
|
{
|
|
|
|
|
|
"task_id": task.get("id"),
|
|
|
|
|
|
"train_method": task.get("train_method", "lora"),
|
|
|
|
|
|
"base_model": task.get("base_model"),
|
2026-07-28 13:10:53 +08:00
|
|
|
|
"file_count": len(artifacts),
|
2026-07-23 19:32:42 +08:00
|
|
|
|
},
|
|
|
|
|
|
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(),
|
2026-07-21 09:23:43 +08:00
|
|
|
|
),
|
|
|
|
|
|
)
|
|
|
|
|
|
|
2026-07-23 19:32:42 +08:00
|
|
|
|
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,
|
2026-08-04 16:59:34 +08:00
|
|
|
|
int(metric.get("step") or line_number),
|
2026-07-23 19:32:42 +08:00
|
|
|
|
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]
|
|
|
|
|
|
|
2026-07-21 09:23:43 +08:00
|
|
|
|
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()
|
2026-07-21 10:55:44 +08:00
|
|
|
|
if not row or row["status"] != "active":
|
|
|
|
|
|
return None
|
|
|
|
|
|
matched, legacy_plaintext = verify_password(password, row["password_hash"])
|
|
|
|
|
|
if not matched:
|
2026-07-21 09:23:43 +08:00
|
|
|
|
return None
|
|
|
|
|
|
last_login = utcnow()
|
2026-07-21 10:55:44 +08:00
|
|
|
|
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"]))
|
2026-07-21 09:23:43 +08:00
|
|
|
|
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
|
2026-07-21 10:55:44 +08:00
|
|
|
|
(id, username, password_hash, display_name, role, status, permissions, create_time, protected)
|
2026-07-21 09:23:43 +08:00
|
|
|
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?, 0)
|
|
|
|
|
|
""",
|
|
|
|
|
|
(
|
|
|
|
|
|
user_id,
|
|
|
|
|
|
payload["username"],
|
2026-07-21 10:55:44 +08:00
|
|
|
|
hash_password(payload.get("password", "platform123")),
|
2026-07-21 09:23:43 +08:00
|
|
|
|
payload.get("display_name") or payload["username"],
|
|
|
|
|
|
payload.get("role", "viewer"),
|
|
|
|
|
|
payload.get("status", "active"),
|
|
|
|
|
|
json_dumps(permissions),
|
|
|
|
|
|
utcnow(),
|
|
|
|
|
|
),
|
|
|
|
|
|
)
|
|
|
|
|
|
return self._user(conn.execute("SELECT * FROM users WHERE id=?", (user_id,)).fetchone())
|
|
|
|
|
|
|
|
|
|
|
|
def update_user(self, user_id: str, payload: dict[str, Any]) -> dict[str, Any]:
|
|
|
|
|
|
with self.connect() as conn:
|
|
|
|
|
|
row = conn.execute("SELECT * FROM users WHERE id=?", (user_id,)).fetchone()
|
|
|
|
|
|
if not row:
|
|
|
|
|
|
raise KeyError(user_id)
|
|
|
|
|
|
values = {
|
|
|
|
|
|
"role": payload.get("role", row["role"]),
|
|
|
|
|
|
"status": payload.get("status", row["status"]),
|
|
|
|
|
|
"permissions": json_dumps(payload.get("permissions", json_loads(row["permissions"], []))),
|
|
|
|
|
|
}
|
|
|
|
|
|
conn.execute(
|
|
|
|
|
|
"UPDATE users SET role=?, status=?, permissions=? WHERE id=?",
|
|
|
|
|
|
(values["role"], values["status"], values["permissions"], user_id),
|
|
|
|
|
|
)
|
|
|
|
|
|
return self._user(conn.execute("SELECT * FROM users WHERE id=?", (user_id,)).fetchone())
|
|
|
|
|
|
|
|
|
|
|
|
def delete_user(self, user_id: str) -> None:
|
|
|
|
|
|
with self.connect() as conn:
|
|
|
|
|
|
row = conn.execute("SELECT protected FROM users WHERE id=?", (user_id,)).fetchone()
|
|
|
|
|
|
if not row:
|
|
|
|
|
|
raise KeyError(user_id)
|
|
|
|
|
|
if row["protected"]:
|
|
|
|
|
|
raise ValueError("protected user cannot be deleted")
|
|
|
|
|
|
conn.execute("DELETE FROM users WHERE id=?", (user_id,))
|
|
|
|
|
|
|
2026-08-03 09:34:08 +08:00
|
|
|
|
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),
|
|
|
|
|
|
)
|
|
|
|
|
|
|
2026-07-21 10:55:44 +08:00
|
|
|
|
def _user(self, row: PgRow) -> dict[str, Any]:
|
2026-07-21 09:23:43 +08:00
|
|
|
|
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")
|
2026-07-28 13:10:53 +08:00
|
|
|
|
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
|
2026-07-21 09:23:43 +08:00
|
|
|
|
with self.connect() as conn:
|
|
|
|
|
|
conn.execute(
|
|
|
|
|
|
"""
|
|
|
|
|
|
INSERT INTO models
|
2026-07-28 13:10:53 +08:00
|
|
|
|
(id, name, type, purpose, model_source, description, path, api_url, api_key, online_model_name, can_train, create_time)
|
|
|
|
|
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
2026-07-21 09:23:43 +08:00
|
|
|
|
""",
|
|
|
|
|
|
(
|
|
|
|
|
|
model_id,
|
|
|
|
|
|
payload["name"],
|
|
|
|
|
|
payload.get("type", "LLM"),
|
|
|
|
|
|
payload.get("purpose", "training"),
|
2026-07-28 13:10:53 +08:00
|
|
|
|
model_source,
|
2026-07-21 09:23:43 +08:00
|
|
|
|
payload.get("description"),
|
2026-07-28 13:10:53 +08:00
|
|
|
|
path,
|
2026-07-21 09:23:43 +08:00
|
|
|
|
payload.get("api_url"),
|
|
|
|
|
|
payload.get("api_key"),
|
|
|
|
|
|
payload.get("online_model_name"),
|
2026-07-28 13:10:53 +08:00
|
|
|
|
can_train,
|
2026-07-21 09:23:43 +08:00
|
|
|
|
utcnow(),
|
|
|
|
|
|
),
|
|
|
|
|
|
)
|
2026-07-22 17:32:59 +08:00
|
|
|
|
return dict(conn.execute("SELECT * FROM models WHERE id=?", (model_id,)).fetchone())
|
2026-07-21 09:23:43 +08:00
|
|
|
|
|
|
|
|
|
|
def update_model(self, model_id: str, payload: dict[str, Any]) -> dict[str, Any]:
|
|
|
|
|
|
current = self.model(model_id)
|
|
|
|
|
|
merged = {**current, **payload}
|
2026-07-28 13:10:53 +08:00
|
|
|
|
# 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
|
2026-07-21 09:23:43 +08:00
|
|
|
|
with self.connect() as conn:
|
|
|
|
|
|
conn.execute(
|
|
|
|
|
|
"""
|
|
|
|
|
|
UPDATE models
|
2026-07-28 13:10:53 +08:00
|
|
|
|
SET name=?, type=?, purpose=?, model_source=?, description=?, path=?, api_url=?, api_key=?, online_model_name=?, can_train=?
|
2026-07-21 09:23:43 +08:00
|
|
|
|
WHERE id=?
|
|
|
|
|
|
""",
|
|
|
|
|
|
(
|
|
|
|
|
|
merged["name"],
|
|
|
|
|
|
merged.get("type", "LLM"),
|
|
|
|
|
|
merged.get("purpose", "training"),
|
2026-07-28 13:10:53 +08:00
|
|
|
|
model_source,
|
2026-07-21 09:23:43 +08:00
|
|
|
|
merged.get("description"),
|
2026-07-28 13:10:53 +08:00
|
|
|
|
path,
|
2026-07-21 09:23:43 +08:00
|
|
|
|
merged.get("api_url"),
|
|
|
|
|
|
merged.get("api_key"),
|
|
|
|
|
|
merged.get("online_model_name"),
|
2026-07-28 13:10:53 +08:00
|
|
|
|
can_train,
|
2026-07-21 09:23:43 +08:00
|
|
|
|
model_id,
|
|
|
|
|
|
),
|
|
|
|
|
|
)
|
2026-07-22 17:32:59 +08:00
|
|
|
|
return dict(conn.execute("SELECT * FROM models WHERE id=?", (model_id,)).fetchone())
|
2026-07-21 09:23:43 +08:00
|
|
|
|
|
|
|
|
|
|
def delete_model(self, model_id: str) -> None:
|
|
|
|
|
|
with self.connect() as conn:
|
|
|
|
|
|
conn.execute("DELETE FROM models WHERE id=?", (model_id,))
|
|
|
|
|
|
|
|
|
|
|
|
def trained_models(self) -> list[dict[str, Any]]:
|
|
|
|
|
|
self.refresh_runtime_state()
|
|
|
|
|
|
with self.connect() as conn:
|
|
|
|
|
|
rows = conn.execute("SELECT * FROM trained_models ORDER BY create_time DESC").fetchall()
|
2026-08-04 16:59:34 +08:00
|
|
|
|
items = []
|
|
|
|
|
|
for row in rows:
|
|
|
|
|
|
item = {
|
2026-07-21 09:23:43 +08:00
|
|
|
|
**dict(row),
|
|
|
|
|
|
"train_methods": json_loads(row["train_methods"], []),
|
|
|
|
|
|
"merged": bool(row["merged"]),
|
|
|
|
|
|
"merging": bool(row["merging"]),
|
|
|
|
|
|
}
|
2026-08-04 16:59:34 +08:00
|
|
|
|
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
|
2026-07-21 09:23:43 +08:00
|
|
|
|
|
2026-07-22 17:32:59 +08:00
|
|
|
|
def delete_trained_model(self, model_id: str) -> None:
|
|
|
|
|
|
with self.connect() as conn:
|
|
|
|
|
|
conn.execute("DELETE FROM trained_models WHERE id=? OR name=?", (model_id, model_id))
|
|
|
|
|
|
|
2026-07-21 09:23:43 +08:00
|
|
|
|
def datasets(self) -> list[dict[str, Any]]:
|
|
|
|
|
|
with self.connect() as conn:
|
2026-07-27 12:44:40 +08:00
|
|
|
|
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)
|
|
|
|
|
|
ORDER BY dataset.create_time DESC"""
|
|
|
|
|
|
).fetchall()
|
2026-07-21 09:23:43 +08:00
|
|
|
|
return [self._dataset(conn, row) for row in rows]
|
|
|
|
|
|
|
|
|
|
|
|
def dataset(self, dataset_id: str) -> dict[str, Any]:
|
|
|
|
|
|
with self.connect() as conn:
|
2026-07-27 12:44:40 +08:00
|
|
|
|
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()
|
2026-07-21 09:23:43 +08:00
|
|
|
|
if not row:
|
|
|
|
|
|
raise KeyError(dataset_id)
|
|
|
|
|
|
return self._dataset(conn, row)
|
|
|
|
|
|
|
2026-07-21 10:55:44 +08:00
|
|
|
|
def _dataset(self, conn: PgConnection, row: PgRow) -> dict[str, Any]:
|
2026-07-21 09:23:43 +08:00
|
|
|
|
files = conn.execute(
|
2026-07-27 12:26:09 +08:00
|
|
|
|
"""SELECT id, name, size, size_bytes, active_version_id,
|
|
|
|
|
|
current_version_id, version_no, versions, create_time,
|
2026-07-24 20:43:47 +08:00
|
|
|
|
record_count, metadata
|
|
|
|
|
|
FROM dataset_files WHERE dataset_id=? ORDER BY create_time, id""",
|
2026-07-21 09:23:43 +08:00
|
|
|
|
(row["id"],),
|
|
|
|
|
|
).fetchall()
|
2026-07-27 12:26:09 +08:00
|
|
|
|
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"))
|
2026-08-03 15:49:21 +08:00
|
|
|
|
file_record_count = int(file_row.get("record_count") or 0)
|
2026-07-27 12:26:09 +08:00
|
|
|
|
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"],
|
2026-08-03 15:49:21 +08:00
|
|
|
|
"record_count": file_record_count,
|
2026-07-27 12:26:09 +08:00
|
|
|
|
"split": metadata.get("file_split"),
|
|
|
|
|
|
}
|
|
|
|
|
|
)
|
2026-07-24 20:43:47 +08:00
|
|
|
|
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}
|
2026-07-27 12:26:09 +08:00
|
|
|
|
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"))
|
2026-08-03 15:49:21 +08:00
|
|
|
|
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)
|
2026-07-27 12:26:09 +08:00
|
|
|
|
current_version_nos = sorted(
|
|
|
|
|
|
{
|
|
|
|
|
|
int(item["current_version_no"])
|
|
|
|
|
|
for item in decoded_files
|
|
|
|
|
|
if item.get("current_version_no")
|
|
|
|
|
|
}
|
|
|
|
|
|
)
|
2026-07-21 09:23:43 +08:00
|
|
|
|
return {
|
|
|
|
|
|
**dict(row),
|
2026-08-03 15:49:21 +08:00
|
|
|
|
"count": total_record_count,
|
|
|
|
|
|
"record_count": total_record_count,
|
2026-07-27 12:26:09 +08:00
|
|
|
|
"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),
|
2026-07-24 20:43:47 +08:00
|
|
|
|
"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),
|
|
|
|
|
|
},
|
2026-07-27 12:26:09 +08:00
|
|
|
|
"files": decoded_files,
|
2026-07-21 09:23:43 +08:00
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
def create_dataset(self, payload: dict[str, Any]) -> dict[str, Any]:
|
|
|
|
|
|
dataset_id = payload.get("id") or new_id("ds")
|
|
|
|
|
|
with self.connect() as conn:
|
|
|
|
|
|
conn.execute(
|
|
|
|
|
|
"""
|
|
|
|
|
|
INSERT INTO datasets
|
|
|
|
|
|
(id, name, type, storage_type, source, task_id, size, count, description, create_time)
|
|
|
|
|
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
|
|
|
|
|
""",
|
|
|
|
|
|
(
|
|
|
|
|
|
dataset_id,
|
|
|
|
|
|
payload["name"],
|
|
|
|
|
|
payload.get("type", "train"),
|
|
|
|
|
|
payload.get("storage_type", "local"),
|
|
|
|
|
|
payload.get("source", "upload"),
|
|
|
|
|
|
payload.get("task_id"),
|
|
|
|
|
|
payload.get("size", "0 KB"),
|
|
|
|
|
|
payload.get("count", 0),
|
|
|
|
|
|
payload.get("description"),
|
|
|
|
|
|
utcnow(),
|
|
|
|
|
|
),
|
|
|
|
|
|
)
|
|
|
|
|
|
return self._dataset(conn, conn.execute("SELECT * FROM datasets WHERE id=?", (dataset_id,)).fetchone())
|
|
|
|
|
|
|
|
|
|
|
|
def update_dataset(self, dataset_id: str, payload: dict[str, Any]) -> dict[str, Any]:
|
|
|
|
|
|
current = self.dataset(dataset_id)
|
|
|
|
|
|
merged = {**current, **payload}
|
|
|
|
|
|
with self.connect() as conn:
|
|
|
|
|
|
conn.execute(
|
|
|
|
|
|
"""
|
|
|
|
|
|
UPDATE datasets
|
|
|
|
|
|
SET name=?, type=?, storage_type=?, source=?, task_id=?, size=?, count=?, description=?
|
|
|
|
|
|
WHERE id=?
|
|
|
|
|
|
""",
|
|
|
|
|
|
(
|
|
|
|
|
|
merged["name"],
|
|
|
|
|
|
merged.get("type", "train"),
|
|
|
|
|
|
merged.get("storage_type", "local"),
|
|
|
|
|
|
merged.get("source", "upload"),
|
|
|
|
|
|
merged.get("task_id"),
|
|
|
|
|
|
merged.get("size", "0 KB"),
|
|
|
|
|
|
merged.get("count", 0),
|
|
|
|
|
|
merged.get("description"),
|
|
|
|
|
|
dataset_id,
|
|
|
|
|
|
),
|
|
|
|
|
|
)
|
|
|
|
|
|
return self._dataset(conn, conn.execute("SELECT * FROM datasets WHERE id=?", (dataset_id,)).fetchone())
|
|
|
|
|
|
|
|
|
|
|
|
def delete_dataset(self, dataset_id: str) -> None:
|
|
|
|
|
|
with self.connect() as conn:
|
|
|
|
|
|
conn.execute("DELETE FROM dataset_files WHERE dataset_id=?", (dataset_id,))
|
|
|
|
|
|
conn.execute("DELETE FROM datasets WHERE id=?", (dataset_id,))
|
|
|
|
|
|
|
2026-07-21 10:55:44 +08:00
|
|
|
|
def add_dataset_file(self, conn: PgConnection, dataset_id: str, name: str, content: str) -> dict[str, Any]:
|
2026-07-21 09:23:43 +08:00
|
|
|
|
now = utcnow()
|
|
|
|
|
|
file_id = new_id("file")
|
|
|
|
|
|
version_id = f"{file_id}_v1"
|
2026-07-27 12:26:09 +08:00
|
|
|
|
size_bytes = len(content.encode("utf-8"))
|
|
|
|
|
|
size = f"{size_bytes} B"
|
2026-08-03 15:49:21 +08:00
|
|
|
|
record_count = count_dataset_records(content)
|
2026-07-27 12:26:09 +08:00
|
|
|
|
version = {
|
|
|
|
|
|
"id": version_id,
|
|
|
|
|
|
"version": 1,
|
|
|
|
|
|
"version_no": 1,
|
|
|
|
|
|
"create_time": now,
|
|
|
|
|
|
"description": "uploaded",
|
|
|
|
|
|
"size_bytes": size_bytes,
|
|
|
|
|
|
"record_count": record_count,
|
|
|
|
|
|
}
|
2026-07-21 09:23:43 +08:00
|
|
|
|
conn.execute(
|
|
|
|
|
|
"""
|
|
|
|
|
|
INSERT INTO dataset_files
|
2026-07-27 12:26:09 +08:00
|
|
|
|
(id, dataset_id, name, size, content, active_version_id, versions, create_time,
|
|
|
|
|
|
current_version_id, size_bytes, record_count, version_no)
|
|
|
|
|
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 1)
|
2026-07-21 09:23:43 +08:00
|
|
|
|
""",
|
|
|
|
|
|
(
|
|
|
|
|
|
file_id,
|
|
|
|
|
|
dataset_id,
|
|
|
|
|
|
name,
|
|
|
|
|
|
size,
|
|
|
|
|
|
content,
|
|
|
|
|
|
version_id,
|
2026-07-27 12:26:09 +08:00
|
|
|
|
json_dumps([version]),
|
2026-07-21 09:23:43 +08:00
|
|
|
|
now,
|
2026-07-27 12:26:09 +08:00
|
|
|
|
version_id,
|
|
|
|
|
|
size_bytes,
|
|
|
|
|
|
record_count,
|
2026-07-21 09:23:43 +08:00
|
|
|
|
),
|
|
|
|
|
|
)
|
|
|
|
|
|
conn.execute(
|
2026-07-27 12:26:09 +08:00
|
|
|
|
"""UPDATE datasets
|
2026-08-03 15:49:21 +08:00
|
|
|
|
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
|
2026-07-27 12:26:09 +08:00
|
|
|
|
WHERE id=?""",
|
2026-08-03 15:49:21 +08:00
|
|
|
|
(dataset_id, dataset_id),
|
2026-07-21 09:23:43 +08:00
|
|
|
|
)
|
2026-07-27 12:26:09 +08:00
|
|
|
|
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)
|
|
|
|
|
|
]
|
2026-07-21 09:23:43 +08:00
|
|
|
|
|
2026-07-21 10:55:44 +08:00
|
|
|
|
def dataset_file(self, file_id: str) -> PgRow:
|
2026-07-21 09:23:43 +08:00
|
|
|
|
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
|
|
|
|
|
|
|
2026-07-23 19:32:42 +08:00
|
|
|
|
def training_dataset_files(self, dataset_id: str) -> list[dict[str, Any]]:
|
|
|
|
|
|
with self.connect() as conn:
|
2026-07-25 17:04:14 +08:00
|
|
|
|
dataset = conn.execute(
|
|
|
|
|
|
"SELECT id, metadata FROM datasets WHERE id=?", (dataset_id,)
|
|
|
|
|
|
).fetchone()
|
2026-07-23 19:32:42 +08:00
|
|
|
|
if not dataset:
|
|
|
|
|
|
raise KeyError(dataset_id)
|
2026-07-25 17:04:14 +08:00
|
|
|
|
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))
|
2026-07-23 19:32:42 +08:00
|
|
|
|
rows = conn.execute(
|
|
|
|
|
|
"""
|
2026-07-24 20:43:47 +08:00
|
|
|
|
SELECT id, dataset_id, name, size, content, active_version_id,
|
|
|
|
|
|
create_time, record_count, metadata
|
2026-07-23 19:32:42 +08:00
|
|
|
|
FROM dataset_files
|
2026-07-25 17:04:14 +08:00
|
|
|
|
WHERE dataset_id = ANY(%s)
|
|
|
|
|
|
ORDER BY CASE WHEN dataset_id=%s THEN 0 ELSE 1 END, create_time, id
|
2026-07-23 19:32:42 +08:00
|
|
|
|
""",
|
2026-07-25 17:04:14 +08:00
|
|
|
|
(runtime_dataset_ids, dataset_id),
|
2026-07-23 19:32:42 +08:00
|
|
|
|
).fetchall()
|
2026-07-24 20:43:47 +08:00
|
|
|
|
return [
|
|
|
|
|
|
{
|
|
|
|
|
|
**dict(row),
|
|
|
|
|
|
"metadata": json_loads(row.get("metadata"), {}),
|
|
|
|
|
|
"split": json_loads(row.get("metadata"), {}).get("file_split"),
|
|
|
|
|
|
}
|
|
|
|
|
|
for row in rows
|
|
|
|
|
|
]
|
2026-07-23 19:32:42 +08:00
|
|
|
|
|
2026-07-21 09:23:43 +08:00
|
|
|
|
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)
|
2026-08-03 15:49:21 +08:00
|
|
|
|
content = payload.get("content", "")
|
|
|
|
|
|
size_bytes = len(content.encode("utf-8"))
|
|
|
|
|
|
record_count = count_dataset_records(content)
|
2026-07-21 09:23:43 +08:00
|
|
|
|
versions = json_loads(row["versions"], [])
|
|
|
|
|
|
version = {
|
|
|
|
|
|
"id": f"{file_id}_v{len(versions) + 1}",
|
|
|
|
|
|
"version": len(versions) + 1,
|
2026-08-03 15:49:21 +08:00
|
|
|
|
"version_no": len(versions) + 1,
|
2026-07-21 09:23:43 +08:00
|
|
|
|
"create_time": utcnow(),
|
|
|
|
|
|
"description": payload.get("description", "online edit"),
|
2026-08-03 15:49:21 +08:00
|
|
|
|
"size_bytes": size_bytes,
|
|
|
|
|
|
"record_count": record_count,
|
2026-07-21 09:23:43 +08:00
|
|
|
|
}
|
|
|
|
|
|
versions.append(version)
|
|
|
|
|
|
conn.execute(
|
2026-08-03 15:49:21 +08:00
|
|
|
|
"""
|
|
|
|
|
|
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,
|
|
|
|
|
|
),
|
2026-07-21 09:23:43 +08:00
|
|
|
|
)
|
2026-08-03 15:49:21 +08:00
|
|
|
|
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}
|
2026-07-21 09:23:43 +08:00
|
|
|
|
|
|
|
|
|
|
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"]}
|
|
|
|
|
|
|
2026-07-22 17:32:59 +08:00
|
|
|
|
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,
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2026-07-21 09:23:43 +08:00
|
|
|
|
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)
|
|
|
|
|
|
|
2026-07-21 10:55:44 +08:00
|
|
|
|
def _task(self, row: PgRow) -> dict[str, Any]:
|
2026-07-21 09:23:43 +08:00
|
|
|
|
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"],
|
2026-07-22 17:32:59 +08:00
|
|
|
|
"compute_job_id": row.get("compute_job_id"),
|
|
|
|
|
|
"completed_at": row.get("completed_at"),
|
2026-07-21 09:23:43 +08:00
|
|
|
|
}
|
|
|
|
|
|
)
|
|
|
|
|
|
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:]}"
|
2026-07-21 10:55:44 +08:00
|
|
|
|
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")
|
2026-07-25 17:04:14 +08:00
|
|
|
|
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")
|
2026-07-21 09:23:43 +08:00
|
|
|
|
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"),
|
2026-07-22 17:32:59 +08:00
|
|
|
|
"engine": payload.get("engine", payload.get("training_engine", "llama_factory")),
|
2026-07-21 09:23:43 +08:00
|
|
|
|
"template": payload.get("template", "qwen"),
|
2026-07-21 10:55:44 +08:00
|
|
|
|
"base_model": base_model,
|
|
|
|
|
|
"train_dataset_id": train_dataset_id,
|
2026-07-21 09:23:43 +08:00
|
|
|
|
"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),
|
2026-07-23 19:32:42 +08:00
|
|
|
|
"quantization_bit": payload.get("quantization_bit", 0),
|
2026-07-21 09:23:43 +08:00
|
|
|
|
"export_quantized": bool(payload.get("export_quantized", False)),
|
|
|
|
|
|
"quant_method": payload.get("quant_method", "bnb"),
|
|
|
|
|
|
"quant_bits": payload.get("quant_bits", 4),
|
|
|
|
|
|
"quant_group_size": payload.get("quant_group_size", 128),
|
|
|
|
|
|
"export_format": payload.get("export_format", "safetensors"),
|
|
|
|
|
|
"progress": 0,
|
|
|
|
|
|
"process_id": None,
|
|
|
|
|
|
"train_duration": "",
|
|
|
|
|
|
"create_time": now,
|
|
|
|
|
|
}
|
|
|
|
|
|
with self.connect() as conn:
|
|
|
|
|
|
conn.execute(
|
|
|
|
|
|
"""
|
|
|
|
|
|
INSERT INTO fine_tune_tasks
|
|
|
|
|
|
(id, name, payload, status, progress, process_id, create_time, gpus)
|
|
|
|
|
|
VALUES (?, ?, ?, 'pending', 0, NULL, ?, ?)
|
|
|
|
|
|
""",
|
|
|
|
|
|
(task_id, name, json_dumps(task), now, json_dumps(task["gpus"])),
|
|
|
|
|
|
)
|
|
|
|
|
|
return task
|
|
|
|
|
|
|
|
|
|
|
|
def update_task(self, task_id: str, payload: dict[str, Any]) -> dict[str, Any]:
|
|
|
|
|
|
current = self.task(task_id)
|
|
|
|
|
|
merged = {**current, **payload, "id": task_id}
|
|
|
|
|
|
with self.connect() as conn:
|
|
|
|
|
|
conn.execute(
|
|
|
|
|
|
"UPDATE fine_tune_tasks SET name=?, payload=?, gpus=? WHERE id=?",
|
|
|
|
|
|
(merged["name"], json_dumps(merged), json_dumps(merged.get("gpus", [])), task_id),
|
|
|
|
|
|
)
|
|
|
|
|
|
return self.task(task_id)
|
|
|
|
|
|
|
|
|
|
|
|
def start_task(self, payload: dict[str, Any]) -> dict[str, Any]:
|
|
|
|
|
|
task_id = str(payload.get("task_id") or payload.get("id"))
|
|
|
|
|
|
current = self.task(task_id)
|
|
|
|
|
|
merged = {**current, **payload, "id": task_id, "status": "syncing", "progress": 8}
|
|
|
|
|
|
selected_gpus = payload.get("gpus") or merged.get("gpus") or [0]
|
|
|
|
|
|
process_id = int(43000 + (time.time() % 10000))
|
|
|
|
|
|
with self.connect() as conn:
|
2026-07-23 19:32:42 +08:00
|
|
|
|
owner = f"start:{task_id}:{uuid.uuid4().hex[:8]}"
|
|
|
|
|
|
if not self._acquire_scheduler_lock(conn, "compute-scheduler", owner):
|
|
|
|
|
|
raise RuntimeError("compute scheduler is busy, please retry")
|
|
|
|
|
|
node = self._schedule_node_locked(conn, payload)
|
|
|
|
|
|
sync_job_id = new_id("sync")
|
|
|
|
|
|
conn.execute(
|
|
|
|
|
|
"""
|
|
|
|
|
|
INSERT INTO resource_sync_jobs
|
|
|
|
|
|
(id, target_node_id, resources, status, progress, create_time)
|
|
|
|
|
|
VALUES (?, ?, ?, 'pending', 0, ?)
|
|
|
|
|
|
""",
|
|
|
|
|
|
(
|
|
|
|
|
|
sync_job_id,
|
|
|
|
|
|
node["id"],
|
|
|
|
|
|
json_dumps(
|
|
|
|
|
|
current.get("resources")
|
|
|
|
|
|
or [
|
|
|
|
|
|
{"resource_type": "model", "resource_id": current.get("base_model")},
|
|
|
|
|
|
{"resource_type": "dataset", "resource_id": current.get("train_dataset_id")},
|
|
|
|
|
|
]
|
|
|
|
|
|
),
|
|
|
|
|
|
utcnow(),
|
|
|
|
|
|
),
|
|
|
|
|
|
)
|
2026-07-21 09:23:43 +08:00
|
|
|
|
conn.execute(
|
|
|
|
|
|
"""
|
|
|
|
|
|
UPDATE fine_tune_tasks
|
|
|
|
|
|
SET payload=?, status='syncing', progress=8, process_id=?, start_time=?,
|
2026-07-22 17:32:59 +08:00
|
|
|
|
compute_node_id=?, gpus=?, sync_job_id=?, compute_job_id=NULL
|
2026-07-21 09:23:43 +08:00
|
|
|
|
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,
|
|
|
|
|
|
),
|
|
|
|
|
|
)
|
2026-07-23 19:32:42 +08:00
|
|
|
|
self._reserve_gpu_allocations(
|
|
|
|
|
|
conn,
|
|
|
|
|
|
{**merged, "id": task_id, "gpus": selected_gpus, "compute_node_id": node["id"]},
|
|
|
|
|
|
node["id"],
|
|
|
|
|
|
selected_gpus,
|
|
|
|
|
|
)
|
2026-07-21 09:23:43 +08:00
|
|
|
|
return self.task(task_id)
|
|
|
|
|
|
|
2026-07-22 17:32:59 +08:00
|
|
|
|
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}
|
2026-07-24 10:27:52 +08:00
|
|
|
|
node = self.select_compute_node(merged)
|
2026-07-22 17:32:59 +08:00
|
|
|
|
selected_gpus = merged.get("gpus") or [0]
|
|
|
|
|
|
return node, self._compute_job_payload_from_task_node(merged, node, selected_gpus)
|
|
|
|
|
|
|
2026-07-24 10:27:52 +08:00
|
|
|
|
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 [0]
|
|
|
|
|
|
return node, self._compute_job_payload_from_task_node(transient_task, node, selected_gpus)
|
|
|
|
|
|
|
2026-07-22 17:32:59 +08:00
|
|
|
|
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]:
|
2026-07-24 10:27:52 +08:00
|
|
|
|
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 "")
|
2026-07-22 17:32:59 +08:00
|
|
|
|
with self.connect() as conn:
|
2026-07-24 10:27:52 +08:00
|
|
|
|
model = conn.execute("SELECT * FROM models WHERE id=?", (base_model_id,)).fetchone()
|
2026-07-28 13:10:53 +08:00
|
|
|
|
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}),请选择其他模型"
|
|
|
|
|
|
)
|
2026-07-24 10:27:52 +08:00
|
|
|
|
dataset = conn.execute("SELECT * FROM datasets WHERE id=?", (dataset_id,)).fetchone()
|
2026-07-23 19:32:42 +08:00
|
|
|
|
files = conn.execute(
|
2026-07-24 20:43:47 +08:00
|
|
|
|
"""SELECT id, name, size, active_version_id, create_time, metadata
|
|
|
|
|
|
FROM dataset_files WHERE dataset_id=? ORDER BY create_time, id""",
|
2026-07-24 10:27:52 +08:00
|
|
|
|
(dataset_id,),
|
2026-07-23 19:32:42 +08:00
|
|
|
|
).fetchall()
|
2026-07-25 17:04:14 +08:00
|
|
|
|
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(),
|
|
|
|
|
|
]
|
2026-07-24 10:27:52 +08:00
|
|
|
|
model_path = (model and model.get("path")) or task.get("model_name_or_path") or base_model_id
|
2026-07-25 17:04:14 +08:00
|
|
|
|
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}")
|
2026-07-23 19:32:42 +08:00
|
|
|
|
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))
|
2026-07-24 20:43:47 +08:00
|
|
|
|
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) :]
|
2026-07-23 19:32:42 +08:00
|
|
|
|
dataset_format = str(task.get("dataset_format") or (dataset and dataset.get("formatting")) or "alpaca").lower()
|
2026-07-28 13:10:53 +08:00
|
|
|
|
# P0-2: Validate dataset content against declared format
|
|
|
|
|
|
train_type = str(task.get("train_type", task.get("train_method", ""))).upper()
|
|
|
|
|
|
expected_format = {
|
|
|
|
|
|
"DPO": "dpo",
|
|
|
|
|
|
"CPT": "cpt",
|
|
|
|
|
|
}.get(train_type)
|
|
|
|
|
|
if expected_format:
|
|
|
|
|
|
dataset_format = expected_format
|
|
|
|
|
|
format_errors: list[str] = []
|
|
|
|
|
|
for file_entry in training_files:
|
|
|
|
|
|
content = file_entry.get("content") 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]))
|
2026-07-22 17:32:59 +08:00
|
|
|
|
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")
|
2026-07-23 19:32:42 +08:00
|
|
|
|
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}"
|
2026-07-22 17:32:59 +08:00
|
|
|
|
return {
|
|
|
|
|
|
**task,
|
|
|
|
|
|
"id": task["id"],
|
|
|
|
|
|
"name": task["name"],
|
|
|
|
|
|
"base_model": model_path,
|
|
|
|
|
|
"model_name_or_path": model_path,
|
2026-07-24 20:43:47 +08:00
|
|
|
|
"dataset": ",".join(training_keys),
|
2026-07-23 19:32:42 +08:00
|
|
|
|
"dataset_key": dataset_key,
|
2026-07-24 20:43:47 +08:00
|
|
|
|
"dataset_keys": training_keys,
|
|
|
|
|
|
"eval_dataset": ",".join(validation_keys) or None,
|
|
|
|
|
|
"eval_dataset_keys": validation_keys,
|
2026-07-23 19:32:42 +08:00
|
|
|
|
"dataset_display_name": (dataset and dataset.get("name")) or dataset_id,
|
|
|
|
|
|
"dataset_dir": dataset_dir,
|
2026-07-24 20:43:47 +08:00
|
|
|
|
"dataset_info": llama_dataset_info(dataset_key, runtime_file_names, dataset_format),
|
2026-07-23 19:32:42 +08:00
|
|
|
|
"dataset_files": [
|
|
|
|
|
|
{
|
2026-07-24 20:43:47 +08:00
|
|
|
|
"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"],
|
2026-07-23 19:32:42 +08:00
|
|
|
|
}
|
2026-07-24 20:43:47 +08:00
|
|
|
|
for item in runtime_files
|
2026-07-23 19:32:42 +08:00
|
|
|
|
],
|
2026-07-22 17:32:59 +08:00
|
|
|
|
"output_dir": output_dir,
|
|
|
|
|
|
"gpus": selected_gpus or task.get("gpus") or [0],
|
|
|
|
|
|
"compute_node_id": node["id"],
|
|
|
|
|
|
"compute_node_code": node["code"],
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
def build_compute_job_payload(self, task_id: str) -> tuple[dict[str, Any], dict[str, Any]]:
|
|
|
|
|
|
task = self.task(task_id)
|
|
|
|
|
|
node = next((item for item in self.compute_nodes() if item["id"] == task.get("compute_node_id")), None)
|
|
|
|
|
|
if not node:
|
|
|
|
|
|
raise RuntimeError("compute node not found")
|
|
|
|
|
|
return node, self._compute_job_payload_from_task_node(task, node, task.get("gpus") or [0])
|
|
|
|
|
|
|
|
|
|
|
|
def apply_compute_job(self, task_id: str, job: dict[str, Any]) -> dict[str, Any]:
|
|
|
|
|
|
status_map = {
|
|
|
|
|
|
"queued": "queued",
|
|
|
|
|
|
"running": "running",
|
|
|
|
|
|
"completed": "completed",
|
|
|
|
|
|
"failed": "failed",
|
|
|
|
|
|
"stopped": "stopped",
|
|
|
|
|
|
}
|
|
|
|
|
|
current = self.task(task_id)
|
|
|
|
|
|
status = status_map.get(str(job.get("status")), str(job.get("status") or current["status"]))
|
|
|
|
|
|
progress = int(job.get("progress", current.get("progress", 0)) or 0)
|
|
|
|
|
|
payload = {
|
|
|
|
|
|
**current,
|
|
|
|
|
|
"status": status,
|
|
|
|
|
|
"progress": progress,
|
|
|
|
|
|
"process_id": job.get("pid") or current.get("process_id"),
|
|
|
|
|
|
"compute_job_id": job.get("id") or current.get("compute_job_id"),
|
|
|
|
|
|
"output_dir": job.get("output_dir") or current.get("output_dir"),
|
|
|
|
|
|
"log_file": job.get("log_file") or current.get("log_file"),
|
|
|
|
|
|
"artifacts": job.get("artifacts") or current.get("artifacts") or [],
|
|
|
|
|
|
}
|
|
|
|
|
|
if status == "failed":
|
|
|
|
|
|
payload["failure_reason"] = job.get("error") or job.get("message") or current.get("failure_reason") or "compute job failed"
|
|
|
|
|
|
elif status in {"queued", "running", "completed"}:
|
|
|
|
|
|
payload.pop("failure_reason", None)
|
|
|
|
|
|
completed_at = utcnow() if status in {"completed", "failed", "stopped"} and not current.get("completed_at") else None
|
|
|
|
|
|
with self.connect() as conn:
|
|
|
|
|
|
conn.execute(
|
|
|
|
|
|
"""
|
|
|
|
|
|
UPDATE fine_tune_tasks
|
|
|
|
|
|
SET payload=?, status=?, progress=?, process_id=?, compute_job_id=?, completed_at=COALESCE(?, completed_at)
|
|
|
|
|
|
WHERE id=?
|
|
|
|
|
|
""",
|
|
|
|
|
|
(
|
|
|
|
|
|
json_dumps(payload),
|
|
|
|
|
|
status,
|
|
|
|
|
|
progress,
|
|
|
|
|
|
payload.get("process_id"),
|
|
|
|
|
|
payload.get("compute_job_id"),
|
|
|
|
|
|
completed_at,
|
|
|
|
|
|
task_id,
|
|
|
|
|
|
),
|
|
|
|
|
|
)
|
2026-07-23 19:32:42 +08:00
|
|
|
|
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 [])
|
2026-07-22 17:32:59 +08:00
|
|
|
|
if status == "completed":
|
2026-07-28 13:10:53 +08:00
|
|
|
|
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),
|
|
|
|
|
|
)
|
2026-07-22 17:32:59 +08:00
|
|
|
|
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]:
|
2026-07-21 09:23:43 +08:00
|
|
|
|
task = self.task(task_id)
|
2026-07-22 17:32:59 +08:00
|
|
|
|
task.update({"status": "failed", "progress": min(task.get("progress", 0), 99), "failure_reason": reason})
|
2026-07-21 09:23:43 +08:00
|
|
|
|
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),
|
|
|
|
|
|
)
|
2026-07-23 19:32:42 +08:00
|
|
|
|
conn.execute(
|
|
|
|
|
|
"UPDATE gpu_allocations SET status='released', released_at=COALESCE(released_at, ?) WHERE task_id=? AND status IN ('allocated','running')",
|
|
|
|
|
|
(utcnow(), task_id),
|
|
|
|
|
|
)
|
2026-07-21 09:23:43 +08:00
|
|
|
|
return self.task(task_id)
|
|
|
|
|
|
|
2026-07-22 17:32:59 +08:00
|
|
|
|
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),
|
|
|
|
|
|
)
|
2026-07-23 19:32:42 +08:00
|
|
|
|
conn.execute(
|
|
|
|
|
|
"UPDATE gpu_allocations SET status='released', released_at=COALESCE(released_at, ?) WHERE task_id=? AND status IN ('allocated','running')",
|
|
|
|
|
|
(utcnow(), task_id),
|
|
|
|
|
|
)
|
2026-07-22 17:32:59 +08:00
|
|
|
|
return self.task(task_id)
|
|
|
|
|
|
|
2026-07-21 09:23:43 +08:00
|
|
|
|
def delete_task(self, task_id: str) -> None:
|
|
|
|
|
|
with self.connect() as conn:
|
|
|
|
|
|
conn.execute("DELETE FROM fine_tune_tasks WHERE id=?", (task_id,))
|
|
|
|
|
|
|
2026-07-22 17:32:59 +08:00
|
|
|
|
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"]})
|
2026-08-03 17:34:21 +08:00
|
|
|
|
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"]
|
2026-07-22 17:32:59 +08:00
|
|
|
|
return payload
|
|
|
|
|
|
|
2026-08-03 17:34:21 +08:00
|
|
|
|
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
|
|
|
|
|
|
|
2026-07-22 17:32:59 +08:00
|
|
|
|
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()
|
2026-08-03 17:34:21 +08:00
|
|
|
|
return [self._enrich_eval_payload(conn, self._json_payload_row(row)) for row in rows]
|
2026-07-22 17:32:59 +08:00
|
|
|
|
|
|
|
|
|
|
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)
|
2026-08-03 17:34:21 +08:00
|
|
|
|
payload = self._enrich_eval_payload(conn, self._json_payload_row(row))
|
2026-07-22 17:32:59 +08:00
|
|
|
|
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:
|
2026-08-03 17:34:21 +08:00
|
|
|
|
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()
|
2026-07-22 17:32:59 +08:00
|
|
|
|
dataset = conn.execute("SELECT name FROM datasets WHERE id=?", (str(payload.get("dataset_id")),)).fetchone()
|
2026-08-03 17:34:21 +08:00
|
|
|
|
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"], {})
|
2026-07-22 17:32:59 +08:00
|
|
|
|
if model:
|
|
|
|
|
|
data.setdefault("model_name", model["name"])
|
2026-08-03 17:34:21 +08:00
|
|
|
|
elif trained_model:
|
|
|
|
|
|
data.setdefault("model_name", trained_model["name"])
|
2026-07-22 17:32:59 +08:00
|
|
|
|
if dataset:
|
|
|
|
|
|
data.setdefault("dataset", dataset["name"])
|
2026-08-03 17:34:21 +08:00
|
|
|
|
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"]
|
2026-07-22 17:32:59 +08:00
|
|
|
|
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)
|
|
|
|
|
|
|
2026-07-28 19:34:41 +08:00
|
|
|
|
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)
|
|
|
|
|
|
|
2026-07-22 17:32:59 +08:00
|
|
|
|
def delete_eval_task(self, task_id: str) -> None:
|
|
|
|
|
|
with self.connect() as conn:
|
|
|
|
|
|
conn.execute("DELETE FROM eval_tasks WHERE id=?", (task_id,))
|
|
|
|
|
|
|
2026-07-28 19:34:41 +08:00
|
|
|
|
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),
|
2026-08-03 15:49:21 +08:00
|
|
|
|
"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(),
|
2026-07-28 19:34:41 +08:00
|
|
|
|
})
|
|
|
|
|
|
return self.update_eval_task(task_id, updates)
|
|
|
|
|
|
|
2026-07-22 17:32:59 +08:00
|
|
|
|
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,))
|
|
|
|
|
|
|
2026-07-23 19:32:42 +08:00
|
|
|
|
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}
|
|
|
|
|
|
|
2026-08-03 15:49:21 +08:00
|
|
|
|
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))
|
|
|
|
|
|
|
2026-07-23 19:32:42 +08:00
|
|
|
|
def _schedule_node_locked(self, conn: PgConnection, payload: dict[str, Any]) -> dict[str, Any]:
|
2026-07-21 09:23:43 +08:00
|
|
|
|
requested = payload.get("requested_node_id") or payload.get("compute_node_id")
|
2026-07-23 19:32:42 +08:00
|
|
|
|
requested_gpus = [int(item) for item in payload.get("gpus") or []]
|
|
|
|
|
|
nodes = self._compute_nodes_locked(conn)
|
2026-07-21 09:23:43 +08:00
|
|
|
|
candidates = [
|
|
|
|
|
|
n
|
|
|
|
|
|
for n in nodes
|
2026-08-03 15:49:21 +08:00
|
|
|
|
if n["enabled"] and n["scheduler_status"] == "online" and n["current_running_jobs"] < self._node_capacity(n)
|
2026-07-21 09:23:43 +08:00
|
|
|
|
]
|
2026-07-23 19:32:42 +08:00
|
|
|
|
if requested_gpus:
|
2026-08-03 15:49:21 +08:00
|
|
|
|
requested_gpu_set = set(requested_gpus)
|
2026-07-23 19:32:42 +08:00
|
|
|
|
candidates = [
|
|
|
|
|
|
node
|
|
|
|
|
|
for node in candidates
|
2026-08-03 15:49:21 +08:00
|
|
|
|
if requested_gpu_set.issubset(self._node_gpu_indexes(conn, node))
|
|
|
|
|
|
and not requested_gpu_set.intersection(self._active_gpu_indexes(conn, node["id"]))
|
2026-07-23 19:32:42 +08:00
|
|
|
|
]
|
2026-07-21 09:23:43 +08:00
|
|
|
|
if requested:
|
|
|
|
|
|
selected = next((n for n in candidates if n["id"] == requested), None)
|
|
|
|
|
|
if selected:
|
|
|
|
|
|
return selected
|
|
|
|
|
|
if not candidates:
|
2026-07-22 17:32:59 +08:00
|
|
|
|
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']}"
|
2026-08-03 15:49:21 +08:00
|
|
|
|
elif node["current_running_jobs"] >= self._node_capacity(node):
|
|
|
|
|
|
reason = f"capacity full {node['current_running_jobs']}/{self._node_capacity(node)}"
|
2026-07-22 17:32:59 +08:00
|
|
|
|
else:
|
|
|
|
|
|
reason = "not selected"
|
|
|
|
|
|
reasons.append(f"{node['code']}({reason})")
|
|
|
|
|
|
raise RuntimeError(f"no available compute node: {', '.join(reasons)}")
|
2026-07-21 09:23:43 +08:00
|
|
|
|
return sorted(candidates, key=lambda n: (-n["scheduler_weight"], n["current_running_jobs"], n["code"]))[0]
|
|
|
|
|
|
|
2026-07-23 19:32:42 +08:00
|
|
|
|
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)
|
|
|
|
|
|
|
2026-07-24 10:27:52 +08:00
|
|
|
|
def select_compute_node(self, payload: dict[str, Any]) -> dict[str, Any]:
|
|
|
|
|
|
with self.connect() as conn:
|
|
|
|
|
|
return self._schedule_node_locked(conn, payload)
|
|
|
|
|
|
|
2026-07-23 19:32:42 +08:00
|
|
|
|
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()),
|
|
|
|
|
|
)
|
|
|
|
|
|
|
2026-07-21 09:23:43 +08:00
|
|
|
|
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(
|
2026-07-22 17:32:59 +08:00
|
|
|
|
task.get("resources")
|
|
|
|
|
|
or [
|
2026-07-21 09:23:43 +08:00
|
|
|
|
{"resource_type": "model", "resource_id": task.get("base_model")},
|
|
|
|
|
|
{"resource_type": "dataset", "resource_id": task.get("train_dataset_id")},
|
|
|
|
|
|
]
|
|
|
|
|
|
),
|
|
|
|
|
|
utcnow(),
|
|
|
|
|
|
),
|
|
|
|
|
|
)
|
|
|
|
|
|
return sync_id
|
|
|
|
|
|
|
2026-07-22 17:32:59 +08:00
|
|
|
|
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())
|
|
|
|
|
|
|
2026-07-23 19:32:42 +08:00
|
|
|
|
def update_resource_replica_check(
|
|
|
|
|
|
self,
|
|
|
|
|
|
replica_id: str,
|
|
|
|
|
|
exists: bool,
|
|
|
|
|
|
byte_size: int = 0,
|
|
|
|
|
|
error: str = "",
|
|
|
|
|
|
) -> dict[str, Any]:
|
|
|
|
|
|
status = "available" if exists else "missing"
|
|
|
|
|
|
sync_status = "synced" if exists else "drifted"
|
|
|
|
|
|
with self.connect() as conn:
|
|
|
|
|
|
conn.execute(
|
|
|
|
|
|
"""
|
|
|
|
|
|
UPDATE resource_replicas
|
|
|
|
|
|
SET status=?, sync_status=?, byte_size=?, last_checked_at=?, last_error=?
|
|
|
|
|
|
WHERE id=?
|
|
|
|
|
|
""",
|
|
|
|
|
|
(status, sync_status, byte_size, utcnow(), error, replica_id),
|
|
|
|
|
|
)
|
|
|
|
|
|
row = conn.execute("SELECT * FROM resource_replicas WHERE id=?", (replica_id,)).fetchone()
|
|
|
|
|
|
if not row:
|
|
|
|
|
|
raise KeyError(replica_id)
|
|
|
|
|
|
return dict(row)
|
|
|
|
|
|
|
|
|
|
|
|
def resource_replicas_by_ids(self, replica_ids: list[str]) -> list[dict[str, Any]]:
|
|
|
|
|
|
if not replica_ids:
|
|
|
|
|
|
return []
|
|
|
|
|
|
with self.connect() as conn:
|
|
|
|
|
|
result: list[dict[str, Any]] = []
|
|
|
|
|
|
for replica_id in replica_ids:
|
|
|
|
|
|
row = conn.execute("SELECT * FROM resource_replicas WHERE id=?", (replica_id,)).fetchone()
|
|
|
|
|
|
if row:
|
|
|
|
|
|
result.append(dict(row))
|
|
|
|
|
|
return result
|
|
|
|
|
|
|
|
|
|
|
|
def update_resource_replica_sync_result(
|
|
|
|
|
|
self,
|
|
|
|
|
|
replica_id: str,
|
|
|
|
|
|
success: bool,
|
|
|
|
|
|
local_path: str | None = None,
|
|
|
|
|
|
byte_size: int = 0,
|
|
|
|
|
|
checksum_sha256: str = "",
|
|
|
|
|
|
error: str = "",
|
|
|
|
|
|
) -> dict[str, Any]:
|
|
|
|
|
|
with self.connect() as conn:
|
|
|
|
|
|
conn.execute(
|
|
|
|
|
|
"""
|
|
|
|
|
|
UPDATE resource_replicas
|
|
|
|
|
|
SET local_path=COALESCE(?, local_path),
|
|
|
|
|
|
status=?,
|
|
|
|
|
|
sync_status=?,
|
|
|
|
|
|
byte_size=?,
|
|
|
|
|
|
checksum_sha256=COALESCE(NULLIF(?, ''), checksum_sha256),
|
|
|
|
|
|
last_checked_at=?,
|
|
|
|
|
|
last_error=?
|
|
|
|
|
|
WHERE id=?
|
|
|
|
|
|
""",
|
|
|
|
|
|
(
|
|
|
|
|
|
local_path,
|
|
|
|
|
|
"available" if success else "missing",
|
|
|
|
|
|
"synced" if success else "failed",
|
|
|
|
|
|
byte_size,
|
|
|
|
|
|
checksum_sha256,
|
|
|
|
|
|
utcnow(),
|
|
|
|
|
|
error,
|
|
|
|
|
|
replica_id,
|
|
|
|
|
|
),
|
|
|
|
|
|
)
|
|
|
|
|
|
row = conn.execute("SELECT * FROM resource_replicas WHERE id=?", (replica_id,)).fetchone()
|
|
|
|
|
|
if not row:
|
|
|
|
|
|
raise KeyError(replica_id)
|
|
|
|
|
|
return dict(row)
|
|
|
|
|
|
|
|
|
|
|
|
def mark_resource_replica_repair_pending(self, replica_ids: list[str]) -> list[dict[str, Any]]:
|
|
|
|
|
|
if not replica_ids:
|
|
|
|
|
|
return []
|
|
|
|
|
|
updated: list[dict[str, Any]] = []
|
|
|
|
|
|
with self.connect() as conn:
|
|
|
|
|
|
for replica_id in replica_ids:
|
|
|
|
|
|
conn.execute(
|
|
|
|
|
|
"""
|
|
|
|
|
|
UPDATE resource_replicas
|
|
|
|
|
|
SET sync_status='repair_pending', last_checked_at=?, last_error=''
|
|
|
|
|
|
WHERE id=?
|
|
|
|
|
|
""",
|
|
|
|
|
|
(utcnow(), replica_id),
|
|
|
|
|
|
)
|
|
|
|
|
|
row = conn.execute("SELECT * FROM resource_replicas WHERE id=?", (replica_id,)).fetchone()
|
|
|
|
|
|
if row:
|
|
|
|
|
|
updated.append(dict(row))
|
|
|
|
|
|
return updated
|
|
|
|
|
|
|
2026-07-21 09:23:43 +08:00
|
|
|
|
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",
|
2026-07-21 10:55:44 +08:00
|
|
|
|
"running": "training with LLaMA-Factory",
|
2026-07-21 09:23:43 +08:00
|
|
|
|
"completed": "training completed",
|
|
|
|
|
|
"failed": "training stopped",
|
2026-07-22 17:32:59 +08:00
|
|
|
|
"stopped": "training stopped",
|
2026-07-21 09:23:43 +08:00
|
|
|
|
}
|
|
|
|
|
|
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),
|
2026-07-21 10:55:44 +08:00
|
|
|
|
"speed": task.get("train_speed") or "--",
|
2026-07-21 09:23:43 +08:00
|
|
|
|
"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}
|
2026-08-04 18:46:53 +08:00
|
|
|
|
# 评测任务同样占用算力节点,纳入运行任务统计
|
|
|
|
|
|
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
|
2026-07-21 09:23:43 +08:00
|
|
|
|
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"], []),
|
2026-07-22 17:32:59 +08:00
|
|
|
|
"capabilities": json_loads(row.get("capabilities"), []),
|
2026-07-21 09:23:43 +08:00
|
|
|
|
"health_detail": json_loads(row["health_detail"], {}),
|
|
|
|
|
|
"current_running_jobs": running_map.get(row["id"], 0),
|
|
|
|
|
|
}
|
|
|
|
|
|
for row in rows
|
|
|
|
|
|
]
|
|
|
|
|
|
|
2026-07-22 17:32:59 +08:00
|
|
|
|
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"},
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2026-07-21 09:23:43 +08:00
|
|
|
|
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)
|
2026-07-22 17:32:59 +08:00
|
|
|
|
merged = self._normalize_compute_node_payload(payload, current)
|
2026-07-21 09:23:43 +08:00
|
|
|
|
with self.connect() as conn:
|
|
|
|
|
|
conn.execute(
|
|
|
|
|
|
"""
|
|
|
|
|
|
UPDATE compute_nodes
|
|
|
|
|
|
SET name=?, api_base_url=?, file_gateway_url=?, enabled=?, scheduler_status=?,
|
2026-07-22 17:32:59 +08:00
|
|
|
|
scheduler_weight=?, tags=?, max_parallel_jobs=?, data_root=?, model_root=?, log_root=?,
|
|
|
|
|
|
api_version=?, capabilities=?, description=?, last_health_check_at=?, health_detail=?
|
2026-07-21 09:23:43 +08:00
|
|
|
|
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"],
|
2026-07-22 17:32:59 +08:00
|
|
|
|
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"]),
|
2026-07-21 09:23:43 +08:00
|
|
|
|
node_id,
|
|
|
|
|
|
),
|
|
|
|
|
|
)
|
|
|
|
|
|
return next(n for n in self.compute_nodes() if n["id"] == node_id)
|
|
|
|
|
|
|
2026-07-21 10:55:44 +08:00
|
|
|
|
def create_compute_node(self, payload: dict[str, Any]) -> dict[str, Any]:
|
2026-07-22 17:32:59 +08:00
|
|
|
|
payload = self._normalize_compute_node_payload(payload)
|
|
|
|
|
|
if not payload["code"]:
|
|
|
|
|
|
raise ValueError("code is required")
|
2026-07-21 10:55:44 +08:00
|
|
|
|
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,
|
2026-07-22 17:32:59 +08:00
|
|
|
|
data_root, model_root, log_root, api_version, capabilities, description,
|
|
|
|
|
|
last_health_check_at, health_detail)
|
|
|
|
|
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 0, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
2026-07-21 10:55:44 +08:00
|
|
|
|
""",
|
|
|
|
|
|
(
|
|
|
|
|
|
node_id,
|
|
|
|
|
|
payload["code"],
|
2026-07-22 17:32:59 +08:00
|
|
|
|
payload["name"] or payload["code"],
|
2026-07-21 10:55:44 +08:00
|
|
|
|
payload["api_base_url"],
|
2026-07-22 17:32:59 +08:00
|
|
|
|
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"],
|
2026-07-21 10:55:44 +08:00
|
|
|
|
now,
|
2026-07-22 17:32:59 +08:00
|
|
|
|
json_dumps(payload["health_detail"]),
|
2026-07-21 10:55:44 +08:00
|
|
|
|
),
|
|
|
|
|
|
)
|
|
|
|
|
|
return next(node for node in self.compute_nodes() if node["id"] == node_id)
|
|
|
|
|
|
|
2026-08-03 15:49:21 +08:00
|
|
|
|
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}
|
|
|
|
|
|
|
2026-07-22 17:32:59 +08:00
|
|
|
|
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))
|
|
|
|
|
|
|
2026-07-21 09:23:43 +08:00
|
|
|
|
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()
|
|
|
|
|
|
]
|
2026-08-04 18:46:53 +08:00
|
|
|
|
# 评测任务同样占用节点 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"])
|
2026-07-21 09:23:43 +08:00
|
|
|
|
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,
|
|
|
|
|
|
)
|
2026-08-04 18:46:53 +08:00
|
|
|
|
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"}
|
|
|
|
|
|
)
|
2026-07-28 17:29:16 +08:00
|
|
|
|
# Also mark GPU as busy if an inference model is loaded on this node
|
2026-08-04 18:46:53 +08:00
|
|
|
|
if row["node_id"] in inference_node_ids and not busy:
|
2026-07-28 17:29:16 +08:00
|
|
|
|
busy = True
|
|
|
|
|
|
reserved = False
|
2026-07-21 09:23:43 +08:00
|
|
|
|
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
|
2026-07-22 17:32:59 +08:00
|
|
|
|
memory_total = float(row["memory_total_gb"] or 0)
|
2026-07-21 09:23:43 +08:00
|
|
|
|
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,
|
2026-07-22 17:32:59 +08:00
|
|
|
|
"memory_total_gb": memory_total,
|
|
|
|
|
|
"memory_percent": round(memory_used / memory_total * 100, 1) if memory_total else 0,
|
2026-07-21 09:23:43 +08:00
|
|
|
|
"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
|
2026-08-04 18:46:53 +08:00
|
|
|
|
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
|
2026-07-21 09:23:43 +08:00
|
|
|
|
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]:
|
2026-08-03 17:24:45 +08:00
|
|
|
|
# 轻量健康检查:采集真实 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
|
2026-07-21 09:23:43 +08:00
|
|
|
|
return {
|
2026-08-03 17:24:45 +08:00
|
|
|
|
"cpu_percent": round(cpu_percent, 1),
|
|
|
|
|
|
"memory_percent": round(memory_percent, 1),
|
|
|
|
|
|
"disk_percent": round(disk_percent, 1),
|
2026-07-21 09:23:43 +08:00
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
def queue(self) -> list[dict[str, Any]]:
|
2026-07-22 17:32:59 +08:00
|
|
|
|
priority_score = {"urgent": 3, "high": 2, "normal": 1, "low": 0}
|
|
|
|
|
|
items = [
|
2026-07-21 09:23:43 +08:00
|
|
|
|
{
|
|
|
|
|
|
"id": task["id"],
|
|
|
|
|
|
"name": task["name"],
|
|
|
|
|
|
"status": task["status"],
|
|
|
|
|
|
"progress": task.get("progress", 0),
|
2026-07-22 17:32:59 +08:00
|
|
|
|
"priority": task.get("priority", "normal"),
|
2026-07-21 09:23:43 +08:00
|
|
|
|
"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"}
|
|
|
|
|
|
]
|
2026-07-22 17:32:59 +08:00
|
|
|
|
return sorted(items, key=lambda item: (-priority_score.get(item["priority"], 1), item["create_time"]), reverse=False)
|
2026-07-21 09:23:43 +08:00
|
|
|
|
|
|
|
|
|
|
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()
|
2026-07-21 10:55:44 +08:00
|
|
|
|
return [dict(row) for row in rows]
|
2026-07-21 09:23:43 +08:00
|
|
|
|
|
|
|
|
|
|
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 = [
|
2026-07-21 10:55:44 +08:00
|
|
|
|
f"[INFO] task={task['name']} engine=llama_factory status={task['status']}",
|
2026-07-21 09:23:43 +08:00
|
|
|
|
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"}
|
|
|
|
|
|
|
2026-08-03 09:34:08 +08:00
|
|
|
|
# ===================== 平台治理:角色 =====================
|
|
|
|
|
|
|
|
|
|
|
|
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, create_time) "
|
|
|
|
|
|
"VALUES (%s, %s, (SELECT username FROM users WHERE id=%s), %s, %s)",
|
|
|
|
|
|
(sid, user_id, user_id, login_at, 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,)
|
|
|
|
|
|
)
|
|
|
|
|
|
|
2026-07-21 09:23:43 +08:00
|
|
|
|
|
|
|
|
|
|
_store: PlatformStore | None = None
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def get_platform_store() -> PlatformStore:
|
|
|
|
|
|
global _store
|
|
|
|
|
|
if _store is None:
|
|
|
|
|
|
_store = PlatformStore()
|
|
|
|
|
|
return _store
|
2026-08-03 09:34:08 +08:00
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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)
|