update
This commit is contained in:
wangjiming
@@ -229,15 +229,21 @@ class PlatformStore:
|
||||
return f"{sec}s"
|
||||
|
||||
def refresh_runtime_state(self) -> None:
|
||||
if get_settings().compute_mode != "simulator":
|
||||
return
|
||||
|
||||
with self.connect() as conn:
|
||||
rows = conn.execute(
|
||||
"SELECT * FROM fine_tune_tasks WHERE status IN ('syncing','queued','running')"
|
||||
).fetchall()
|
||||
now_dt = datetime.now(timezone.utc)
|
||||
for row in rows:
|
||||
payload = json_loads(row["payload"], {})
|
||||
compute_job_id = payload.get("compute_job_id")
|
||||
compute_node_api = payload.get("compute_node_api")
|
||||
if compute_job_id and compute_node_api:
|
||||
# 已派发到算力：状态/进度/日志回传来自算力进程（架构 §1.1）
|
||||
self._sync_task_from_compute(conn, row, payload, compute_job_id, compute_node_api)
|
||||
continue
|
||||
if get_settings().compute_mode != "simulator":
|
||||
continue
|
||||
start = parse_time(row["start_time"])
|
||||
if not start:
|
||||
continue
|
||||
@@ -252,7 +258,6 @@ class PlatformStore:
|
||||
else:
|
||||
status, progress = "completed", 100
|
||||
|
||||
payload = json_loads(row["payload"], {})
|
||||
payload.update(
|
||||
{
|
||||
"status": status,
|
||||
@@ -272,19 +277,63 @@ class PlatformStore:
|
||||
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"]),
|
||||
)
|
||||
if get_settings().compute_mode == "simulator":
|
||||
sync_rows = conn.execute(
|
||||
"SELECT * FROM resource_sync_jobs WHERE status IN ('pending','running')"
|
||||
).fetchall()
|
||||
for row in sync_rows:
|
||||
created = parse_time(row["create_time"])
|
||||
age = int((now_dt - created).total_seconds()) if created else 0
|
||||
status = "completed" if age >= 6 else "running"
|
||||
progress = 100 if status == "completed" else min(95, 15 + age * 12)
|
||||
completed_at = row["completed_at"] or (utcnow() if status == "completed" else None)
|
||||
conn.execute(
|
||||
"UPDATE resource_sync_jobs SET status=?, progress=?, completed_at=? WHERE id=?",
|
||||
(status, progress, completed_at, row["id"]),
|
||||
)
|
||||
|
||||
def _sync_task_from_compute(
|
||||
self,
|
||||
conn,
|
||||
row,
|
||||
payload: dict[str, Any],
|
||||
compute_job_id: str,
|
||||
compute_node_api: str,
|
||||
) -> None:
|
||||
"""从算力节点拉回已派发任务的状态/进度/日志，写回本地任务记录。"""
|
||||
try:
|
||||
from app.modules.compute_gateway.client import ComputeNodeClient
|
||||
|
||||
job = ComputeNodeClient(compute_node_api).get_job(compute_job_id)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
payload["compute_sync_error"] = str(exc)
|
||||
conn.execute(
|
||||
"UPDATE fine_tune_tasks SET payload=? WHERE id=?",
|
||||
(json_dumps(payload), row["id"]),
|
||||
)
|
||||
return
|
||||
status = job.get("status")
|
||||
progress = int(job.get("progress", 0) or 0)
|
||||
logs = job.get("logs") or ""
|
||||
payload.update(
|
||||
{
|
||||
"status": status,
|
||||
"progress": progress,
|
||||
"compute_logs": logs,
|
||||
"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)
|
||||
|
||||
def _ensure_trained_model(self, conn: PgConnection, task: dict[str, Any]) -> None:
|
||||
name = task.get("output_model_name") or f"{task['name']}-lora"
|
||||
@@ -306,7 +355,7 @@ class PlatformStore:
|
||||
utcnow(),
|
||||
0,
|
||||
0,
|
||||
output_dir or f"/data/yg-ft/outputs/{task['name']}/adapter",
|
||||
task.get("output_dir") or f"/data/yg-ft/outputs/{task.get('name')}/adapter",
|
||||
),
|
||||
)
|
||||
|
||||
@@ -402,6 +451,79 @@ class PlatformStore:
|
||||
)
|
||||
return {"id": aid}
|
||||
|
||||
# ---- 登录会话（采集在线时长） ----
|
||||
def create_session(self, user: dict[str, Any]) -> str:
|
||||
sid = new_id("sess")
|
||||
with self.connect() as conn:
|
||||
conn.execute(
|
||||
"""INSERT INTO sessions
|
||||
(id, user_id, username, display_name, role, login_at, logout_at, duration_seconds, create_time)
|
||||
VALUES (?,?,?,?,?,?,?,?,?)""",
|
||||
(
|
||||
sid,
|
||||
user.get("id"),
|
||||
user.get("username"),
|
||||
user.get("display_name"),
|
||||
user.get("role"),
|
||||
utcnow(),
|
||||
None,
|
||||
None,
|
||||
utcnow(),
|
||||
),
|
||||
)
|
||||
return sid
|
||||
|
||||
def close_session(self, session_id: str | None) -> None:
|
||||
if not session_id:
|
||||
return
|
||||
with self.connect() as conn:
|
||||
row = conn.execute(
|
||||
"SELECT login_at FROM sessions WHERE id=? AND logout_at IS NULL", (session_id,)
|
||||
).fetchone()
|
||||
if not row:
|
||||
return
|
||||
now = datetime.now(timezone.utc)
|
||||
start = parse_time(row["login_at"]) or now
|
||||
seconds = max(0, int((now - start).total_seconds()))
|
||||
conn.execute(
|
||||
"UPDATE sessions SET logout_at=?, duration_seconds=? WHERE id=?",
|
||||
(utcnow(), seconds, session_id),
|
||||
)
|
||||
|
||||
def login_duration_rank(self, limit: int = 8) -> list[dict[str, Any]]:
|
||||
"""本月登录时长排行：按用户聚合会话时长（小时）。"""
|
||||
month_start = utcnow()[:7] + "01T00:00:00Z"
|
||||
with self.connect() as conn:
|
||||
rows = conn.execute(
|
||||
"SELECT user_id, username, display_name, role, login_at, logout_at, duration_seconds "
|
||||
"FROM sessions WHERE login_at >= ?",
|
||||
(month_start,),
|
||||
).fetchall()
|
||||
agg: dict[str, dict[str, Any]] = {}
|
||||
now = datetime.now(timezone.utc)
|
||||
for r in rows:
|
||||
uid = r["user_id"]
|
||||
bucket = agg.setdefault(
|
||||
uid,
|
||||
{"user": r["display_name"] or r["username"], "role": r["role"] or "", "total": 0.0, "has": False},
|
||||
)
|
||||
dur = r["duration_seconds"]
|
||||
if dur is None and r["logout_at"] is None:
|
||||
start = parse_time(r["login_at"])
|
||||
if start:
|
||||
dur = max(0, int((now - start).total_seconds()))
|
||||
if dur is None:
|
||||
dur = 0
|
||||
bucket["total"] += dur
|
||||
bucket["has"] = True
|
||||
result = [
|
||||
{"user": b["user"], "role": b["role"], "duration": round(b["total"] / 3600, 1)}
|
||||
for b in agg.values()
|
||||
if b["has"]
|
||||
]
|
||||
result.sort(key=lambda x: x["duration"], reverse=True)
|
||||
return result[:limit]
|
||||
|
||||
# ---- 审批模板 ----
|
||||
def create_approval_template(self, payload: dict[str, Any]) -> dict[str, Any]:
|
||||
tid = payload.get("id") or new_id("tpl")
|
||||
@@ -768,7 +890,7 @@ class PlatformStore:
|
||||
utcnow(),
|
||||
),
|
||||
)
|
||||
return self.model(model_id)
|
||||
return self.model(model_id)
|
||||
|
||||
def update_model(self, model_id: str, payload: dict[str, Any]) -> dict[str, Any]:
|
||||
current = self.model(model_id)
|
||||
@@ -793,7 +915,7 @@ class PlatformStore:
|
||||
model_id,
|
||||
),
|
||||
)
|
||||
return self.model(model_id)
|
||||
return self.model(model_id)
|
||||
|
||||
def delete_model(self, model_id: str) -> None:
|
||||
with self.connect() as conn:
|
||||
@@ -1096,20 +1218,80 @@ class PlatformStore:
|
||||
task_id,
|
||||
),
|
||||
)
|
||||
if get_settings().compute_mode != "simulator":
|
||||
api_base = node.get("api_base_url")
|
||||
if api_base:
|
||||
# GPU 计算派发到算力节点进程执行；后端只做调度编排（架构 §1.1）
|
||||
self._dispatch_to_compute(task_id, node, merged, selected_gpus)
|
||||
elif get_settings().compute_mode != "simulator":
|
||||
# 降级路径：未配置算力节点时后端本机执行（违反 §1.1，待移除）
|
||||
from app.modules.fine_tune.service import launch_training
|
||||
|
||||
launch_training(task_id)
|
||||
return self.task(task_id)
|
||||
|
||||
def stop_task(self, task_id: str) -> dict[str, Any]:
|
||||
task = self.task(task_id)
|
||||
task.update({"status": "failed", "progress": min(task.get("progress", 0), 99)})
|
||||
def _dispatch_to_compute(
|
||||
self,
|
||||
task_id: str,
|
||||
node: dict[str, Any],
|
||||
task: dict[str, Any],
|
||||
gpus: list,
|
||||
) -> None:
|
||||
"""把训练作业派发到算力节点，GPU 计算在算力进程内执行；后端记录算力 job id。"""
|
||||
from app.modules.compute_gateway.client import ComputeNodeClient
|
||||
|
||||
cfg = {
|
||||
"id": f"ft_{task_id}",
|
||||
"name": task.get("name") or task_id,
|
||||
"type": "fine_tune",
|
||||
"gpus": gpus,
|
||||
"stage": str(task.get("train_type") or "SFT").lower(),
|
||||
"base_model": task.get("base_model") or "placeholder-base-model",
|
||||
"dataset": task.get("train_dataset_id") or task.get("dataset") or "placeholder-dataset",
|
||||
"template": task.get("template") or "qwen",
|
||||
"train_method": task.get("train_method") or "lora",
|
||||
"output_dir": f"/data/yg-ft/outputs/{task.get('name') or task_id}/adapter",
|
||||
"batch_size": int(task.get("batch_size", 2) or 2),
|
||||
"learning_rate": float(task.get("learning_rate", 0.0002) or 0.0002),
|
||||
"n_epochs": int(task.get("n_epochs", 3) or 3),
|
||||
}
|
||||
client = ComputeNodeClient(node["api_base_url"])
|
||||
try:
|
||||
job = client.create_job(cfg)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
self.update_task_runtime(task_id, status="failed", extra={"dispatch_error": str(exc)})
|
||||
raise RuntimeError(f"dispatch training job to compute node failed: {exc}") from exc
|
||||
compute_job_id = job.get("id")
|
||||
current = self.task(task_id)
|
||||
updated = {**current, "compute_job_id": compute_job_id, "compute_node_api": node["api_base_url"]}
|
||||
with self.connect() as conn:
|
||||
conn.execute(
|
||||
"UPDATE fine_tune_tasks SET status='failed', payload=?, completed_at=? WHERE id=?",
|
||||
"UPDATE fine_tune_tasks SET payload=? WHERE id=?",
|
||||
(json_dumps(updated), task_id),
|
||||
)
|
||||
|
||||
def stop_task(self, task_id: str) -> dict[str, Any]:
|
||||
task = self.task(task_id)
|
||||
# 如果任务已派发到算力节点，先通知算力停止
|
||||
compute_job_id = task.get("compute_job_id")
|
||||
node_id = task.get("compute_node_id")
|
||||
if compute_job_id and node_id:
|
||||
try:
|
||||
from app.modules.compute_gateway.client import ComputeNodeClient
|
||||
node = next((n for n in self.compute_nodes() if n["id"] == node_id), None)
|
||||
if node and node.get("api_base_url"):
|
||||
ComputeNodeClient(node["api_base_url"]).stop_job(compute_job_id)
|
||||
except Exception: # noqa: BLE001 - best effort stop
|
||||
pass
|
||||
task.update({"status": "stopped", "progress": min(task.get("progress", 0), 99)})
|
||||
with self.connect() as conn:
|
||||
conn.execute(
|
||||
"UPDATE fine_tune_tasks SET status='stopped', payload=?, completed_at=? WHERE id=?",
|
||||
(json_dumps(task), utcnow(), task_id),
|
||||
)
|
||||
conn.execute(
|
||||
"UPDATE gpu_allocations SET status='released', released_at=COALESCE(released_at, ?) WHERE task_id=? AND status IN ('allocated','running')",
|
||||
(utcnow(), task_id),
|
||||
)
|
||||
return self.task(task_id)
|
||||
|
||||
def update_task_runtime(
|
||||
@@ -1934,6 +2116,490 @@ class PlatformStore:
|
||||
content = self.generate_training_log(task)
|
||||
return {"job_id": job_id, "content": content, "lines": len(content.splitlines())}
|
||||
|
||||
# ===================== Model Evaluation =====================
|
||||
|
||||
def _json_payload_row(self, row: PgRow) -> dict[str, Any]:
|
||||
payload = json_loads(row["payload"], {})
|
||||
payload.update({"id": row["id"], "status": row.get("status"), "create_time": row["create_time"]})
|
||||
return payload
|
||||
|
||||
def eval_tasks(self) -> list[dict[str, Any]]:
|
||||
with self.connect() as conn:
|
||||
rows = conn.execute("SELECT * FROM eval_tasks ORDER BY create_time DESC").fetchall()
|
||||
return [self._json_payload_row(row) for row in rows]
|
||||
|
||||
def eval_task(self, task_id: str) -> dict[str, Any]:
|
||||
with self.connect() as conn:
|
||||
row = conn.execute("SELECT * FROM eval_tasks WHERE id=?", (task_id,)).fetchone()
|
||||
if not row:
|
||||
raise KeyError(task_id)
|
||||
payload = self._json_payload_row(row)
|
||||
payload.setdefault("sample_count", 0)
|
||||
payload.setdefault("completed_count", 0)
|
||||
payload.setdefault("passed_count", 0)
|
||||
payload.setdefault("overall_score", payload.get("score") or 0)
|
||||
payload.setdefault("overall_score_max", 100)
|
||||
payload.setdefault("overall_evaluation", "")
|
||||
payload.setdefault("improvement_suggestions", [])
|
||||
payload.setdefault("dimension_summary", [])
|
||||
payload.setdefault("samples", [])
|
||||
return payload
|
||||
|
||||
def create_eval_task(self, payload: dict[str, Any]) -> dict[str, Any]:
|
||||
task_id = str(payload.get("id") or payload.get("task_id") or new_id("eval"))
|
||||
name = str(payload.get("eval_task_name") or payload.get("name") or f"eval-{task_id[-6:]}")
|
||||
status = str(payload.get("status") or "pending")
|
||||
now = payload.get("create_time") or utcnow()
|
||||
data = {
|
||||
**payload,
|
||||
"id": task_id,
|
||||
"eval_task_name": name,
|
||||
"status": status,
|
||||
"create_time": now,
|
||||
"metric": payload.get("metric") or "custom",
|
||||
}
|
||||
with self.connect() as conn:
|
||||
model = conn.execute("SELECT name FROM models WHERE id=?", (str(payload.get("model_id")),)).fetchone()
|
||||
dataset = conn.execute("SELECT name FROM datasets WHERE id=?", (str(payload.get("dataset_id")),)).fetchone()
|
||||
if model:
|
||||
data.setdefault("model_name", model["name"])
|
||||
if dataset:
|
||||
data.setdefault("dataset", dataset["name"])
|
||||
conn.execute(
|
||||
"INSERT INTO eval_tasks (id, name, payload, status, create_time) VALUES (?, ?, ?, ?, ?)",
|
||||
(task_id, name, json_dumps(data), status, now),
|
||||
)
|
||||
return self.eval_task(task_id)
|
||||
|
||||
def delete_eval_task(self, task_id: str) -> None:
|
||||
with self.connect() as conn:
|
||||
conn.execute("DELETE FROM eval_tasks WHERE id=?", (task_id,))
|
||||
|
||||
def dimensions(self) -> list[dict[str, Any]]:
|
||||
with self.connect() as conn:
|
||||
rows = conn.execute("SELECT * FROM eval_dimensions ORDER BY create_time DESC").fetchall()
|
||||
return [
|
||||
{
|
||||
**json_loads(row["payload"], {}),
|
||||
"id": row["id"],
|
||||
"name": row["name"],
|
||||
"is_active": bool(row["is_active"]),
|
||||
"is_default": bool(row["is_default"]),
|
||||
"create_time": row["create_time"],
|
||||
}
|
||||
for row in rows
|
||||
]
|
||||
|
||||
def dimension(self, dimension_id: str) -> dict[str, Any]:
|
||||
with self.connect() as conn:
|
||||
row = conn.execute("SELECT * FROM eval_dimensions WHERE id=?", (dimension_id,)).fetchone()
|
||||
if not row:
|
||||
raise KeyError(dimension_id)
|
||||
return {
|
||||
**json_loads(row["payload"], {}),
|
||||
"id": row["id"],
|
||||
"name": row["name"],
|
||||
"is_active": bool(row["is_active"]),
|
||||
"is_default": bool(row["is_default"]),
|
||||
"create_time": row["create_time"],
|
||||
}
|
||||
|
||||
def create_dimension(self, payload: dict[str, Any]) -> dict[str, Any]:
|
||||
dimension_id = str(payload.get("id") or new_id("dim"))
|
||||
name = str(payload.get("name") or f"dimension-{dimension_id[-6:]}")
|
||||
now = payload.get("create_time") or utcnow()
|
||||
data = {**payload, "id": dimension_id, "name": name, "create_time": now}
|
||||
with self.connect() as conn:
|
||||
conn.execute(
|
||||
"INSERT INTO eval_dimensions (id, name, payload, is_active, is_default, create_time) VALUES (?, ?, ?, ?, ?, ?)",
|
||||
(dimension_id, name, json_dumps(data), 1 if data.get("is_active", True) else 0, 1 if data.get("is_default") else 0, now),
|
||||
)
|
||||
return self.dimension(dimension_id)
|
||||
|
||||
def update_dimension(self, dimension_id: str, payload: dict[str, Any]) -> dict[str, Any]:
|
||||
current = self.dimension(dimension_id)
|
||||
merged = {**current, **payload, "id": dimension_id}
|
||||
with self.connect() as conn:
|
||||
conn.execute(
|
||||
"UPDATE eval_dimensions SET name=?, payload=?, is_active=?, is_default=? WHERE id=?",
|
||||
(
|
||||
merged["name"],
|
||||
json_dumps(merged),
|
||||
1 if merged.get("is_active", True) else 0,
|
||||
1 if merged.get("is_default") else 0,
|
||||
dimension_id,
|
||||
),
|
||||
)
|
||||
return self.dimension(dimension_id)
|
||||
|
||||
def delete_dimension(self, dimension_id: str) -> None:
|
||||
with self.connect() as conn:
|
||||
conn.execute("DELETE FROM eval_dimensions WHERE id=?", (dimension_id,))
|
||||
|
||||
# ===================== Model Compare / Inference =====================
|
||||
|
||||
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,))
|
||||
|
||||
# ===================== Compute Job Sync (派发回传) =====================
|
||||
|
||||
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")
|
||||
existing = conn.execute("SELECT owner, expires_at FROM scheduler_locks WHERE lock_key=?", (lock_key,)).fetchone()
|
||||
if existing:
|
||||
if existing["expires_at"] > now and existing["owner"] != owner:
|
||||
return False
|
||||
conn.execute(
|
||||
"UPDATE scheduler_locks SET owner=?, expires_at=?, update_time=? WHERE lock_key=?",
|
||||
(owner, expires_at, now, lock_key),
|
||||
)
|
||||
else:
|
||||
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
|
||||
|
||||
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 apply_compute_job(self, task_id: str, job: dict[str, Any]) -> dict[str, Any]:
|
||||
status_map = {
|
||||
"queued": "queued",
|
||||
"running": "running",
|
||||
"completed": "completed",
|
||||
"failed": "failed",
|
||||
"stopped": "stopped",
|
||||
}
|
||||
current = self.task(task_id)
|
||||
status = status_map.get(str(job.get("status")), str(job.get("status") or current["status"]))
|
||||
progress = int(job.get("progress", current.get("progress", 0)) or 0)
|
||||
payload = {
|
||||
**current,
|
||||
"status": status,
|
||||
"progress": progress,
|
||||
"process_id": job.get("pid") or current.get("process_id"),
|
||||
"compute_job_id": job.get("id") or current.get("compute_job_id"),
|
||||
"output_dir": job.get("output_dir") or current.get("output_dir"),
|
||||
"log_file": job.get("log_file") or current.get("log_file"),
|
||||
"artifacts": job.get("artifacts") or current.get("artifacts") or [],
|
||||
}
|
||||
if status == "failed":
|
||||
payload["failure_reason"] = job.get("error") or job.get("message") or current.get("failure_reason") or "compute job failed"
|
||||
elif status in {"queued", "running", "completed"}:
|
||||
payload.pop("failure_reason", None)
|
||||
completed_at = utcnow() if status in {"completed", "failed", "stopped"} and not current.get("completed_at") else None
|
||||
with self.connect() as conn:
|
||||
conn.execute(
|
||||
"""
|
||||
UPDATE fine_tune_tasks
|
||||
SET payload=?, status=?, progress=?, process_id=?, compute_job_id=?, completed_at=COALESCE(?, completed_at)
|
||||
WHERE id=?
|
||||
""",
|
||||
(
|
||||
json_dumps(payload),
|
||||
status,
|
||||
progress,
|
||||
payload.get("process_id"),
|
||||
payload.get("compute_job_id"),
|
||||
completed_at,
|
||||
task_id,
|
||||
),
|
||||
)
|
||||
self._upsert_compute_job(conn, payload, job, status)
|
||||
self._sync_gpu_allocations(conn, payload, job, status)
|
||||
self._upsert_checkpoints(conn, task_id, job.get("checkpoints") or [])
|
||||
if status == "completed":
|
||||
self._ensure_trained_model(conn, payload)
|
||||
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),
|
||||
)
|
||||
if log_snippet:
|
||||
conn.execute(
|
||||
"UPDATE fine_tune_tasks SET payload = ? WHERE id = ?",
|
||||
(json_dumps({**payload, "last_log_snippet": log_snippet[:8192]}), task_id),
|
||||
)
|
||||
return self.task(task_id)
|
||||
|
||||
def running_compute_tasks(self) -> list[dict[str, Any]]:
|
||||
return [
|
||||
task
|
||||
for task in self.tasks()
|
||||
if task.get("compute_job_id") and task.get("compute_node_id") and task["status"] in {"syncing", "queued", "running"}
|
||||
]
|
||||
|
||||
def mark_task_failed(self, task_id: str, reason: str) -> dict[str, Any]:
|
||||
task = self.task(task_id)
|
||||
task.update({"status": "failed", "progress": min(task.get("progress", 0), 99), "failure_reason": reason})
|
||||
with self.connect() as conn:
|
||||
conn.execute(
|
||||
"UPDATE fine_tune_tasks SET status='failed', payload=?, completed_at=? WHERE id=?",
|
||||
(json_dumps(task), utcnow(), task_id),
|
||||
)
|
||||
conn.execute(
|
||||
"UPDATE gpu_allocations SET status='released', released_at=COALESCE(released_at, ?) WHERE task_id=? AND status IN ('allocated','running')",
|
||||
(utcnow(), task_id),
|
||||
)
|
||||
return self.task(task_id)
|
||||
|
||||
def 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 = self._parse_training_metric(line)
|
||||
if not metric:
|
||||
continue
|
||||
rows.append(
|
||||
(
|
||||
new_id("metric"),
|
||||
task_id,
|
||||
line_number,
|
||||
metric.get("epoch"),
|
||||
metric.get("loss"),
|
||||
metric.get("grad_norm"),
|
||||
metric.get("learning_rate"),
|
||||
line[:2000],
|
||||
utcnow(),
|
||||
)
|
||||
)
|
||||
with self.connect() as conn:
|
||||
conn.execute("DELETE FROM fine_tune_metrics WHERE task_id=?", (task_id,))
|
||||
if rows:
|
||||
conn.executemany(
|
||||
"""
|
||||
INSERT INTO fine_tune_metrics
|
||||
(id, task_id, step, epoch, loss, grad_norm, learning_rate, raw, create_time)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
""",
|
||||
rows,
|
||||
)
|
||||
return len(rows)
|
||||
|
||||
def _parse_training_metric(self, line: str) -> dict[str, Any] | None:
|
||||
if "loss" not in line or "learning_rate" not in line:
|
||||
return None
|
||||
import re
|
||||
result: dict[str, Any] = {}
|
||||
for key in ["loss", "grad_norm", "learning_rate", "epoch"]:
|
||||
match = re.search(rf"['\"]?{key}['\"]?\s*:\s*([-+]?\d+(?:\.\d+)?(?:[eE][-+]?\d+)?)", line)
|
||||
if match:
|
||||
result[key] = float(match.group(1))
|
||||
return result or None
|
||||
|
||||
def 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 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 [json_loads(row["payload"], {}) if "payload" in row.keys() else dict(row) for row in rows]
|
||||
|
||||
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,
|
||||
),
|
||||
)
|
||||
return self.compute_job(job_id)
|
||||
|
||||
|
||||
_store: PlatformStore | None = None
|
||||
|
||||
|
||||
Reference in New Issue
Block a user