ec7d8c0a3d
- compute_nodes() 的 current_running_jobs 纳入运行中的评测任务（eval_tasks） 与已加载的推理模型（compare_tasks 持久化状态 + 内存标记兜底）， 多节点时 GPU 被评测/推理占用不再显示 0/1 - gpus() 直接按 eval_tasks / compare_tasks 派生 GPU busy/reserved 状态， 修复 gpu_id=0 时 or -1 导致匹配失败；删除评测后 GPU 不再残留 busy - 评测不再复用 mark_inference_loaded 内存标记，GPU 占用由任务数据驱动 - requirements.txt 补充 pytest / ruff（此前仅声明在 pyproject dev 可选依赖） Co-Authored-By: Claude <noreply@anthropic.com>
192 lines
9.2 KiB
Python
192 lines
9.2 KiB
Python
from __future__ import annotations
|
||
|
||
import json
|
||
import time
|
||
from typing import Any
|
||
|
||
from app.db.platform_store import get_platform_store
|
||
from app.modules.compute_gateway.client import ComputeNodeClient
|
||
|
||
# starting 状态允许的最大轮询次数（约 40 * 3s ≈ 2 分钟），超过即判定节点不可达
|
||
MAX_STARTING_ATTEMPTS = 40
|
||
|
||
|
||
def _node_for_task(task: dict[str, Any]) -> dict[str, Any] | None:
|
||
return next((node for node in get_platform_store().compute_nodes() if node["id"] == task.get("compute_node_id")), None)
|
||
|
||
|
||
def _parse_inference_load_status(task: dict[str, Any]) -> tuple[list[dict[str, Any]], dict[str, Any]]:
|
||
load_status = task.get("load_status") or {}
|
||
if isinstance(load_status, str):
|
||
try:
|
||
load_status = json.loads(load_status)
|
||
except (json.JSONDecodeError, TypeError):
|
||
load_status = {}
|
||
return load_status.get("loaded_models") or [], load_status
|
||
|
||
|
||
async def reconcile_inference_loads(store: Any) -> list[dict[str, Any]]:
|
||
"""推进处于 starting 状态的推理加载。
|
||
|
||
模型加载已改为异步派发：/model-compare/{id}/load 立即返回，这里在每次
|
||
轮询时查询对应计算节点的 /inference/status，把任务从 starting 推进到
|
||
ready/error。使用短超时，单节点不可达不会阻塞整轮轮询。
|
||
"""
|
||
reconciled: list[dict[str, Any]] = []
|
||
now = time.time()
|
||
for task in store.compare_tasks():
|
||
items, _ = _parse_inference_load_status(task)
|
||
if not any(item.get("status") == "starting" for item in items):
|
||
continue
|
||
# dirty 只要处理过任一 starting 项就置位：load_attempts / last_polled_at
|
||
# 必须落库，否则节点不可达时计数不会累积，封顶逻辑永远触发不了
|
||
dirty = False
|
||
for item in items:
|
||
if item.get("status") != "starting":
|
||
continue
|
||
# 节流：同一 item 每 3s 只查询一次
|
||
if now - float(item.get("last_polled_at") or 0) < 3:
|
||
continue
|
||
item["last_polled_at"] = now
|
||
item["load_attempts"] = int(item.get("load_attempts") or 0) + 1
|
||
dirty = True
|
||
node = next((n for n in store.compute_nodes() if n["id"] == item.get("node_id")), None)
|
||
if not node:
|
||
item["status"] = "error"
|
||
item["error"] = "compute node deleted"
|
||
store.mark_inference_unloaded(item.get("node_id") or "")
|
||
continue
|
||
if not node.get("enabled") or node.get("scheduler_status") != "online":
|
||
item["status"] = "error"
|
||
item["error"] = "compute node offline"
|
||
store.mark_inference_unloaded(node["id"])
|
||
continue
|
||
try:
|
||
status = await ComputeNodeClient(node["api_base_url"]).inference_status()
|
||
except Exception as exc: # noqa: BLE001 - node unreachable; keep retrying until cap
|
||
if int(item.get("load_attempts") or 0) >= MAX_STARTING_ATTEMPTS:
|
||
item["status"] = "error"
|
||
item["error"] = f"compute node unreachable: {exc}"
|
||
store.mark_inference_unloaded(node["id"])
|
||
continue
|
||
node_status = status.get("status")
|
||
if node_status == "ready":
|
||
item["status"] = "ready"
|
||
item.pop("error", None)
|
||
store.mark_inference_loaded(node["id"])
|
||
elif node_status == "error":
|
||
item["status"] = "error"
|
||
item["error"] = status.get("error") or "model load failed on compute node"
|
||
store.mark_inference_unloaded(node["id"])
|
||
elif node_status == "idle":
|
||
# 节点重启导致已加载模型丢失
|
||
item["status"] = "error"
|
||
item["error"] = "model disappeared from compute node (node may have restarted)"
|
||
store.mark_inference_unloaded(node["id"])
|
||
# node_status == "loading" -> 保持 starting，下轮再查
|
||
if dirty:
|
||
if any(i.get("status") in {"ready", "running"} for i in items):
|
||
new_status = "loaded"
|
||
elif any(i.get("status") == "starting" for i in items):
|
||
new_status = "starting" # 仍在加载中，保持 starting
|
||
else:
|
||
new_status = "failed"
|
||
store.update_compare_task(task["id"], {"status": new_status, "load_status": {"loaded_models": items}})
|
||
reconciled.append({"task_id": task["id"], "status": new_status})
|
||
return reconciled
|
||
|
||
|
||
async def fetch_eval_result_content(client: ComputeNodeClient, node: dict[str, Any], job: dict[str, Any]) -> dict[str, Any] | None:
|
||
output_dir = job.get("output_dir")
|
||
if not output_dir:
|
||
return None
|
||
full_path = f"{str(output_dir).rstrip('/')}/eval_results.json"
|
||
data_root = "/data/yg-ft/"
|
||
if full_path.startswith(data_root):
|
||
full_path = full_path[len(data_root):]
|
||
rel_path = full_path.lstrip("/")
|
||
import httpx
|
||
url = f"{node['api_base_url'].rstrip('/')}/modelTF/compute/files/read"
|
||
async with httpx.AsyncClient(timeout=30, headers=client.headers()) as http:
|
||
response = await http.get(url, params={"path": rel_path})
|
||
response.raise_for_status()
|
||
payload = response.json()
|
||
return payload if isinstance(payload, dict) else None
|
||
|
||
|
||
async def poll_compute_jobs_once() -> dict[str, Any]:
|
||
store = get_platform_store()
|
||
synced: list[dict[str, Any]] = []
|
||
failed: list[dict[str, str]] = []
|
||
for task in store.running_compute_tasks():
|
||
node = _node_for_task(task)
|
||
if not node:
|
||
failed.append({"task_id": task["id"], "error": "compute node not found"})
|
||
continue
|
||
try:
|
||
client = ComputeNodeClient(node["api_base_url"])
|
||
job = await client.get_job(task["compute_job_id"])
|
||
try:
|
||
logs = await client.job_logs(task["compute_job_id"], tail_lines=5000)
|
||
store.record_training_log_metrics(task["id"], str(logs.get("content") or ""))
|
||
except Exception:
|
||
pass
|
||
# P0-4: Force-fetch last log snippet when job reaches terminal state
|
||
if job.get("status") in {"failed", "stopped"}:
|
||
try:
|
||
last_logs = await client.job_logs(task["compute_job_id"], tail_lines=200)
|
||
job["log_snippet"] = str(last_logs.get("content") or "")[:8192]
|
||
except Exception:
|
||
pass
|
||
synced.append(store.apply_compute_job(task["id"], job))
|
||
except Exception as exc: # noqa: BLE001 - keep polling other jobs
|
||
failed.append({"task_id": task["id"], "error": str(exc)})
|
||
standalone_synced: list[dict[str, Any]] = []
|
||
for record in store.active_standalone_compute_jobs():
|
||
node = next((item for item in store.compute_nodes() if item["id"] == record.get("node_id")), None)
|
||
if not node:
|
||
failed.append({"job_id": record["id"], "error": "compute node not found"})
|
||
continue
|
||
try:
|
||
job = await ComputeNodeClient(node["api_base_url"]).get_job(record["id"])
|
||
standalone_synced.append(store.sync_model_merge_job(record["id"], job))
|
||
except Exception as exc: # noqa: BLE001 - keep polling other jobs
|
||
failed.append({"job_id": record["id"], "error": str(exc)})
|
||
|
||
# ── Eval job sync ────────────────────────────────────────────────
|
||
eval_synced = 0
|
||
for eval_task in store.running_eval_tasks():
|
||
node = next(
|
||
(item for item in store.compute_nodes() if item["id"] == eval_task.get("compute_node_id")),
|
||
None,
|
||
)
|
||
if not node:
|
||
failed.append({"eval_task_id": eval_task["id"], "error": "compute node not found"})
|
||
continue
|
||
try:
|
||
client = ComputeNodeClient(node["api_base_url"])
|
||
job = await client.get_job(eval_task["compute_job_id"])
|
||
result_content = None
|
||
# Try to read eval_results.json from the job output directory
|
||
if job.get("status") == "completed" and job.get("output_dir"):
|
||
try:
|
||
result_content = await fetch_eval_result_content(client, node, job)
|
||
except Exception:
|
||
pass
|
||
store.apply_eval_job_result(eval_task["id"], job, result_content)
|
||
# 评测 GPU 占用由 eval_tasks 状态派生，无需维护推理内存标记
|
||
eval_synced += 1
|
||
except Exception as exc: # noqa: BLE001
|
||
failed.append({"eval_task_id": eval_task["id"], "error": str(exc)})
|
||
|
||
# ── Inference load reconciliation ─────────────────────────────────────
|
||
try:
|
||
inference_reconciled = await reconcile_inference_loads(store)
|
||
except Exception as exc: # noqa: BLE001 - keep polling alive
|
||
failed.append({"inference_reconcile": str(exc)})
|
||
inference_reconciled = []
|
||
|
||
return {"synced": len(synced) + len(standalone_synced) + eval_synced, "failed": failed,
|
||
"items": synced, "standalone": standalone_synced, "eval_synced": eval_synced,
|
||
"inference_reconciled": inference_reconciled}
|