diff --git a/backend/app/api/v1/endpoints/platform.py b/backend/app/api/v1/endpoints/platform.py index ba2b499..fa04cca 100644 --- a/backend/app/api/v1/endpoints/platform.py +++ b/backend/app/api/v1/endpoints/platform.py @@ -33,6 +33,15 @@ def _select_first_online_node(store: Any) -> dict[str, Any] | None: return None +def _candidate_online_nodes(store: Any, preferred_node_id: str | None = None) -> list[dict[str, Any]]: + nodes = [node for node in store.compute_nodes() if node.get("enabled") and node.get("scheduler_status") == "online"] + if not preferred_node_id: + return nodes + preferred = [node for node in nodes if node.get("id") == preferred_node_id] + others = [node for node in nodes if node.get("id") != preferred_node_id] + return preferred + others + + def _build_messages_payload(payload: dict[str, Any]) -> dict[str, Any]: """Convert frontend inference payload to compute API messages format. @@ -61,10 +70,47 @@ def _build_messages_payload(payload: dict[str, Any]) -> dict[str, Any]: } +def _node_for_inference_payload(store: Any, payload: dict[str, Any]) -> dict[str, Any] | None: + node_id = payload.get("node_id") or payload.get("compute_node_id") + task_id = payload.get("task_id") or payload.get("compare_task_id") + if task_id and not node_id: + try: + task = store.compare_task(str(task_id)) + load_status = task.get("load_status") or {} + if isinstance(load_status, str): + load_status = json.loads(load_status) + loaded_models = load_status.get("loaded_models") or [] + ready_model = next((item for item in loaded_models if item.get("status") in {"ready", "running"} and item.get("node_id")), None) + if ready_model: + node_id = ready_model.get("node_id") + except Exception: + node_id = None + if node_id: + return next((node for node in store.compute_nodes() if node.get("id") == node_id), None) + return _select_first_online_node(store) + + async def _stream_chat_proxy(payload: dict[str, Any]) -> StreamingResponse: """Common SSE streaming proxy: convert payload → forward to compute node → stream back.""" store = get_platform_store() - node = _select_first_online_node(store) + # 任务仍在加载中时，直接返回明确的加载中提示，避免转发到尚未就绪的节点 + task_id = payload.get("task_id") or payload.get("compare_task_id") + if task_id: + try: + task = store.compare_task(str(task_id)) + load_status = task.get("load_status") or {} + if isinstance(load_status, str): + load_status = json.loads(load_status) + items = load_status.get("loaded_models") or [] + if items and not any(item.get("status") in {"ready", "running"} for item in items): + if any(item.get("status") == "starting" for item in items): + return StreamingResponse( + iter(['data: {"error": "模型加载中，请稍候再试"}\n\n']), + media_type="text/event-stream", + ) + except Exception: # noqa: BLE001 - fall through to normal routing on lookup errors + pass + node = _node_for_inference_payload(store, payload) if not node: return StreamingResponse( iter(['data: {"error": "no online compute node available for inference"}\n\n']), @@ -729,12 +775,17 @@ async def merge_model(payload: dict[str, Any] = Body(...)) -> dict[str, Any]: None, ) base_model_path = payload.get("base_model_path") or (trained_model and trained_model.get("base_model_path")) - adapter_path = payload.get("adapter_path") or payload.get("adapter_name_or_path") or (trained_model and trained_model.get("merged_path")) + adapter_path = ( + payload.get("adapter_path") + or payload.get("adapter_name_or_path") + or (trained_model and (trained_model.get("artifact_dir") or trained_model.get("adapter_path") or trained_model.get("merged_path"))) + ) if not base_model_path: raise fail(400, "base_model_path is required") if not adapter_path: raise fail(400, "adapter_path is required") - node = store.schedule_node({**payload, "gpus": payload.get("gpus") or []}) + requested_node_id = payload.get("requested_node_id") or payload.get("compute_node_id") or (trained_model and trained_model.get("compute_node_id")) + node = store.schedule_node({**payload, "requested_node_id": requested_node_id, "gpus": payload.get("gpus") or []}) health = node.get("health_detail") or {} output_root = str(health.get("output_root") or f"{node['data_root'].rstrip('/')}/outputs") output_name = str(payload.get("output_model_name") or payload.get("merged_model_name") or f"{trained_model_id or 'model'}-merged") @@ -753,11 +804,13 @@ async def merge_model(payload: dict[str, Any] = Body(...)) -> dict[str, Any]: "gpus": payload.get("gpus") or [], "trained_model_id": trained_model["id"] if trained_model else trained_model_id, "model_name": trained_model["name"] if trained_model else payload.get("model_name"), + "compute_node_id": node["id"], + "compute_node_code": node.get("code"), } if get_settings().compute_mode == "simulator": job = {"id": job_payload["id"], "status": "queued", "progress": 10, "command": [], "output_dir": output_dir} else: - client = ComputeNodeClient(node["api_base_url"]) + client = ComputeNodeClient(node["api_base_url"], timeout=900) preview = await client.validate_job(job_payload) if not preview.get("valid", False): raise fail(409, "; ".join(preview.get("errors") or ["merge preflight failed"])) @@ -1537,28 +1590,49 @@ async def model_compare_detail(task_id: str) -> dict[str, Any]: raise fail(404, "compare task not found") -async def _unload_from_compute_node() -> dict[str, Any]: - """Best-effort unload the inference model from the first online compute node.""" - store = get_platform_store() - # Clear all inference tracking — only one model can be loaded at a time - for node in store.compute_nodes(): - store.mark_inference_unloaded(node["id"]) - node = _select_first_online_node(store) - if not node: - return {"unloaded": False, "error": "no online compute node"} - try: - client = ComputeNodeClient(node["api_base_url"]) - result = await client._request("POST", "/inference/unload", json_data={}) - return result - except Exception as exc: - return {"unloaded": False, "error": str(exc)} +async def _unload_from_compute_node(store: Any, task: dict[str, Any] | None = None) -> dict[str, Any]: + """Best-effort unload the inference model from the node(s) that hold it. + + 任务感知：优先卸载 ``task.load_status.loaded_models`` 中记录的节点； + 无任务时回退到平台记录的已加载推理的节点。每个节点使用短超时， + 保证卸载永远不会长时间阻塞调用方（例如删除操作）。 + """ + node_ids: set[str] = set() + if task: + load_status = task.get("load_status") or {} + if isinstance(load_status, str): + try: + load_status = json.loads(load_status) + except json.JSONDecodeError: + load_status = {} + node_ids = {item.get("node_id") for item in load_status.get("loaded_models") or [] if item.get("node_id")} + if not node_ids: + node_ids = {node["id"] for node in store.compute_nodes() if store.is_inference_loaded(node["id"])} + nodes = [node for node in store.compute_nodes() if node["id"] in node_ids] + results: list[dict[str, Any]] = [] + for node in nodes: + try: + result = await ComputeNodeClient(node["api_base_url"]).inference_unload() + results.append({"node_id": node["id"], "node_code": node.get("code"), "success": True, "result": result}) + except Exception as exc: # noqa: BLE001 - best-effort unload must not raise + results.append({"node_id": node["id"], "node_code": node.get("code"), "success": False, "error": str(exc)}) + finally: + store.mark_inference_unloaded(node["id"]) + return {"unloaded": bool(results), "nodes": results} @router.delete("/model-compare/{task_id}") async def model_compare_delete(task_id: str) -> dict[str, Any]: - # 删除前先释放算力节点上的模型 - await _unload_from_compute_node() + # 先删记录（快），再 best-effort 释放算力节点上的模型——删除绝不被卸载阻塞 + try: + task = get_platform_store().compare_task(task_id) + except KeyError: + raise fail(404, "compare task not found") get_platform_store().delete_compare_task(task_id) + try: + await _unload_from_compute_node(get_platform_store(), task=task) + except Exception: # noqa: BLE001 - deletion must succeed even if unload fails + pass return ok({"deleted": task_id}) @@ -1585,9 +1659,44 @@ async def model_compare_update_load_status(task_id: str, payload: dict[str, Any] raise fail(404, "compare task not found") +def _invalidate_superseded_models(store: Any, task_id: str, loaded_models: list[dict[str, Any]]) -> None: + """同一计算节点同一时刻只能加载一个推理模型。 + + 当新任务把模型派发到了某节点后，把其它任务中在该节点上 ready/running + 的模型标记为已被替换，保持平台 DB 与计算节点实际状态一致。 + """ + taken_node_ids = {m.get("node_id") for m in loaded_models if m.get("node_id") and m.get("status") == "starting"} + if not taken_node_ids: + return + for other in store.compare_tasks(): + if str(other.get("id")) == str(task_id): + continue + load_status = other.get("load_status") or {} + if isinstance(load_status, str): + try: + load_status = json.loads(load_status) + except json.JSONDecodeError: + load_status = {} + items = load_status.get("loaded_models") or [] + changed = False + for item in items: + if item.get("node_id") in taken_node_ids and item.get("status") in {"ready", "running"}: + item["status"] = "error" + item["error"] = "模型已被其他推理任务替换" + changed = True + if changed: + new_status = "loaded" if any(i.get("status") in {"ready", "running"} for i in items) else "failed" + store.update_compare_task(other["id"], {"status": new_status, "load_status": {"loaded_models": items}}) + + @router.post("/model-compare/{task_id}/load") async def model_compare_load(task_id: str) -> dict[str, Any]: - """真正加载模型到算力节点（不再使用假 PID/端口）。""" + """异步派发模型加载到算力节点，立即返回。 + + 加载进度由轮询对账器（compute_poller → reconcile_inference_loads）推进： + 任务项先以 status=starting 记录，对账器查询节点 /inference/status 后 + 推进到 ready/error。这里只负责把加载请求派发出去，绝不同步等待加载完成。 + """ try: store = get_platform_store() task = store.compare_task(task_id) @@ -1597,15 +1706,14 @@ async def model_compare_load(task_id: str) -> dict[str, Any]: models = json.loads(models) except json.JSONDecodeError: models = [] - # 选取在线算力节点 - node = _select_first_online_node(store) - if not node: + online_nodes = _candidate_online_nodes(store) + if not online_nodes: return ok({"status": "failed", "error": "no online compute node"}) - client = ComputeNodeClient(node["api_base_url"]) loaded_models = [] for item in models: if not isinstance(item, dict): continue + preferred_node_id = item.get("node_id") or item.get("compute_node_id") model_path = item.get("model_path", "") if not model_path: # 尝试从模型库获取路径 @@ -1614,28 +1722,46 @@ async def model_compare_load(task_id: str) -> dict[str, Any]: db_model = store.model(model_id) model_path = db_model.get("path", "") except KeyError: - pass + trained_model = next((m for m in store.trained_models() if str(m.get("id")) == str(model_id)), None) + if trained_model: + model_path = trained_model.get("merged_path") or trained_model.get("artifact_dir") or "" + preferred_node_id = preferred_node_id or trained_model.get("compute_node_id") if not model_path: loaded_models.append({**item, "status": "error", "error": "model_path not found"}) continue - # 真正调用算力节点加载模型 load_payload = { "model_name_or_path": model_path, "template": item.get("template", "qwen"), } if item.get("adapter_path"): load_payload["adapter_name_or_path"] = item["adapter_path"] - try: - result = await client._request("POST", "/inference/load", json_data=load_payload) - if result.get("loaded"): + if get_settings().compute_mode == "simulator": + loaded_models.append({**item, "status": "ready", "node_id": "", "node_name": ""}) + continue + # 只派发：HTTP 响应成功即视为已接受（节点会异步加载），loaded 字段忽略 + item_dispatched = False + errors = [] + for node in _candidate_online_nodes(store, preferred_node_id): + try: + client = ComputeNodeClient(node["api_base_url"]) + await client.inference_load(load_payload) store.mark_inference_loaded(node["id"]) - loaded_models.append({**item, "status": "ready"}) - else: - loaded_models.append({**item, "status": "error", "error": result.get("error", "load failed")}) - except Exception as exc: - loaded_models.append({**item, "status": "error", "error": str(exc)}) - status = "loaded" if any(m.get("status") == "ready" for m in loaded_models) else "failed" + loaded_models.append({**item, "status": "starting", "node_id": node["id"], "node_name": node.get("name")}) + item_dispatched = True + break + except Exception as exc: # noqa: BLE001 - try next candidate node + errors.append(f"{node.get('name') or node.get('code')}: {exc}") + if not item_dispatched: + loaded_models.append({**item, "status": "error", "error": "; ".join(errors) or "load dispatch failed"}) + if any(m.get("status") == "starting" for m in loaded_models): + status = "starting" + elif any(m.get("status") == "error" for m in loaded_models): + status = "failed" + else: + status = "loaded" updated = store.update_compare_task(task_id, {"status": status, "load_status": {"loaded_models": loaded_models}}) + # 同一节点同一时刻只能有一个推理模型；新任务占用了节点后，把其它任务上该节点的模型标记为已被替换 + _invalidate_superseded_models(store, task_id, loaded_models) return ok(updated) except KeyError: raise fail(404, "compare task not found") @@ -1645,8 +1771,9 @@ async def model_compare_load(task_id: str) -> dict[str, Any]: async def model_compare_unload(task_id: str) -> dict[str, Any]: try: store = get_platform_store() - # 真正释放算力节点上的模型资源 - unload_result = await _unload_from_compute_node() + task = store.compare_task(task_id) + # 任务感知卸载：只释放该任务实际加载到的节点，短超时快速返回 + unload_result = await _unload_from_compute_node(store, task=task) updated = store.update_compare_task(task_id, {"status": "pending", "load_status": {"loaded_models": []}}) return ok({"task": updated, "unload": unload_result}) except KeyError: @@ -1662,7 +1789,7 @@ async def model_compare_start_model(task_id: str, payload: dict[str, Any] = Body async def model_compare_chat_with_port(payload: dict[str, Any] = Body(...)) -> dict[str, Any]: """Proxy non-streaming chat to the compute node running the inference model.""" store = get_platform_store() - node = _select_first_online_node(store) + node = _node_for_inference_payload(store, payload) if not node: return ok({"response": "no online compute node available for inference", "request": payload}) try: @@ -1717,8 +1844,9 @@ async def model_chat_local_preload(payload: dict[str, Any] = Body(...)) -> dict[ return ok({"loaded": False, "error": "no online compute node"}) try: client = ComputeNodeClient(node["api_base_url"]) - result = await client._request("POST", "/inference/load", json_data=payload) - if result.get("loaded"): + # 计算节点现在异步加载：HTTP 接受（loading/ready）即视为派发成功 + result = await client.inference_load(payload) + if result.get("loaded") or result.get("status") in {"loading", "ready"}: store.mark_inference_loaded(node["id"]) return ok(result) except Exception as exc: @@ -1729,18 +1857,19 @@ async def model_chat_local_preload(payload: dict[str, Any] = Body(...)) -> dict[ async def model_chat_local_unload() -> dict[str, Any]: """Unload the inference model from the compute node.""" store = get_platform_store() - # Clear all inference tracking + # 释放所有已加载推理的节点（短超时，best-effort） + results: list[dict[str, Any]] = [] for n in store.compute_nodes(): - store.mark_inference_unloaded(n["id"]) - node = _select_first_online_node(store) - if not node: - return ok({"unloaded": False, "error": "no online compute node"}) - try: - client = ComputeNodeClient(node["api_base_url"]) - result = await client._request("POST", "/inference/unload", json_data={}) - return ok(result) - except Exception as exc: - return ok({"unloaded": False, "error": str(exc)}) + if not store.is_inference_loaded(n["id"]): + continue + try: + result = await ComputeNodeClient(n["api_base_url"]).inference_unload() + results.append({"node_id": n["id"], "success": True, "result": result}) + except Exception as exc: # noqa: BLE001 - best-effort unload + results.append({"node_id": n["id"], "success": False, "error": str(exc)}) + finally: + store.mark_inference_unloaded(n["id"]) + return ok({"unloaded": True, "nodes": results}) @router.get("/model-chat/local/status") @@ -1752,7 +1881,7 @@ async def model_chat_local_status() -> dict[str, Any]: return ok({"loaded": False, "error": "no online compute node"}) try: client = ComputeNodeClient(node["api_base_url"]) - result = await client._request("GET", "/inference/status") + result = await client.inference_status() return ok(result) except Exception as exc: return ok({"loaded": False, "error": str(exc)}) @@ -1770,8 +1899,9 @@ async def model_chat_trained_preload(payload: dict[str, Any] = Body(...)) -> dic return ok({"loaded": False, "error": "no online compute node"}) try: client = ComputeNodeClient(node["api_base_url"]) - result = await client._request("POST", "/inference/load", json_data=payload) - if result.get("loaded"): + # 计算节点现在异步加载：HTTP 接受（loading/ready）即视为派发成功 + result = await client.inference_load(payload) + if result.get("loaded") or result.get("status") in {"loading", "ready"}: store.mark_inference_loaded(node["id"]) return ok(result) except Exception as exc: diff --git a/backend/app/db/platform_store.py b/backend/app/db/platform_store.py index 7a074ea..02835d4 100644 --- a/backend/app/db/platform_store.py +++ b/backend/app/db/platform_store.py @@ -208,8 +208,12 @@ def parse_training_metric_line(line: str) -> dict[str, float] | None: if "loss" not in line and "learning_rate" not in line: return None result: dict[str, float] = {} + number_pattern = r"([-+]?(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][-+]?\d+)?)" + step_match = re.search(rf"(?:^|[\s,{{])['\"]?step['\"]?\s*(?:=|:)\s*{number_pattern}", line, re.I) + if step_match: + result["step"] = float(step_match.group(1)) for key in ["loss", "grad_norm", "learning_rate", "epoch"]: - match = re.search(rf"['\"]?{key}['\"]?\s*:\s*([-+]?\d+(?:\.\d+)?(?:[eE][-+]?\d+)?)", line) + match = re.search(rf"['\"]?{key}['\"]?\s*(?:=|:)\s*{number_pattern}", line, re.I) if match: result[key] = float(match.group(1)) return result or None @@ -452,6 +456,15 @@ class PlatformStore: ) self._ensure_columns(conn, "gpus", {"last_seen_at": "TEXT"}) self._ensure_columns(conn, "fine_tune_tasks", {"compute_job_id": "TEXT"}) + self._ensure_columns( + conn, + "trained_models", + { + "artifact_dir": "TEXT", + "compute_node_id": "TEXT", + "compute_node_name": "TEXT", + }, + ) self._ensure_columns( conn, "resource_replicas", @@ -589,6 +602,15 @@ class PlatformStore: name = task.get("output_model_name") or f"{task['name']}-lora" exists = conn.execute("SELECT id FROM trained_models WHERE name=?", (name,)).fetchone() if exists: + conn.execute( + """ + UPDATE trained_models + SET compute_node_id=COALESCE(compute_node_id, ?), + compute_node_name=COALESCE(compute_node_name, ?) + WHERE id=? + """, + (task.get("compute_node_id"), task.get("compute_node_code") or task.get("compute_node_name"), exists["id"]), + ) return model = conn.execute("SELECT path FROM models WHERE id=?", (task.get("base_model"),)).fetchone() output_dir = task.get("output_dir") or f"/data/yg-ft/outputs/{task['name']}" @@ -596,8 +618,8 @@ class PlatformStore: conn.execute( """ INSERT INTO trained_models - (id, name, train_methods, base_model_path, create_time, merged, merging, merged_path, artifact_dir) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) + (id, name, train_methods, base_model_path, create_time, merged, merging, merged_path, artifact_dir, compute_node_id, compute_node_name) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( trained_model_id, @@ -609,6 +631,8 @@ class PlatformStore: 0, output_dir, output_dir, + task.get("compute_node_id"), + task.get("compute_node_code") or task.get("compute_node_name"), ), ) # Use real artifact data from compute node when available @@ -875,7 +899,7 @@ class PlatformStore: ( new_id("metric"), task_id, - line_number, + int(metric.get("step") or line_number), metric.get("epoch"), metric.get("loss"), metric.get("grad_norm"), @@ -1342,15 +1366,25 @@ class PlatformStore: self.refresh_runtime_state() with self.connect() as conn: rows = conn.execute("SELECT * FROM trained_models ORDER BY create_time DESC").fetchall() - return [ - { + items = [] + for row in rows: + item = { **dict(row), "train_methods": json_loads(row["train_methods"], []), "merged": bool(row["merged"]), "merging": bool(row["merging"]), } - for row in rows - ] + if not item.get("compute_node_id"): + task_rows = conn.execute("SELECT payload FROM fine_tune_tasks ORDER BY create_time DESC").fetchall() + for task in task_rows: + task_payload = json_loads(task["payload"], {}) + output_name = task_payload.get("output_model_name") or f"{task_payload.get('name')}-lora" + if output_name == item["name"]: + item["compute_node_id"] = task_payload.get("compute_node_id") + item["compute_node_name"] = task_payload.get("compute_node_code") or task_payload.get("compute_node_name") + break + items.append(item) + return items def delete_trained_model(self, model_id: str) -> None: with self.connect() as conn: diff --git a/backend/app/db/sql/001_platform_runtime.sql b/backend/app/db/sql/001_platform_runtime.sql index e58518f..84cf674 100644 --- a/backend/app/db/sql/001_platform_runtime.sql +++ b/backend/app/db/sql/001_platform_runtime.sql @@ -33,7 +33,10 @@ CREATE TABLE IF NOT EXISTS trained_models ( create_time TEXT NOT NULL, merged INTEGER NOT NULL DEFAULT 0, merging INTEGER NOT NULL DEFAULT 0, - merged_path TEXT + merged_path TEXT, + artifact_dir TEXT, + compute_node_id TEXT, + compute_node_name TEXT ); CREATE TABLE IF NOT EXISTS model_lineage ( diff --git a/backend/app/modules/compute_gateway/client.py b/backend/app/modules/compute_gateway/client.py index aeeaae6..2e7343b 100644 --- a/backend/app/modules/compute_gateway/client.py +++ b/backend/app/modules/compute_gateway/client.py @@ -33,6 +33,16 @@ def _unwrap_dict(payload: Any) -> dict[str, Any]: return payload if isinstance(payload, dict) else {} +# Inference calls are intentionally short-timeout: +# - load dispatch only confirms the compute node accepted the request +# (the actual model load now runs asynchronously on the node). +# - status/unload must never block the platform for long when a node is +# unreachable but still marked online. +INFERENCE_LOAD_TIMEOUT = httpx.Timeout(30, connect=10) +INFERENCE_STATUS_TIMEOUT = httpx.Timeout(30, connect=5) +INFERENCE_UNLOAD_TIMEOUT = httpx.Timeout(30, connect=5) + + class ComputeNodeClient: """Application-side client for one compute node. @@ -182,10 +192,16 @@ class ComputeNodeClient: response.raise_for_status() return _unwrap_dict(response.json()) - async def _request(self, method: str, path: str, json_data: dict[str, Any] | None = None) -> dict[str, Any]: + async def _request( + self, + method: str, + path: str, + json_data: dict[str, Any] | None = None, + timeout: float | None = None, + ) -> dict[str, Any]: """Generic request method for compute API endpoints.""" url = _join_url(self.api_base_url, f"{self.route_prefix}{path}") - async with httpx.AsyncClient(timeout=300, headers=self.headers()) as client: + async with httpx.AsyncClient(timeout=timeout or 300, headers=self.headers()) as client: if method.upper() == "GET": response = await client.get(url) else: @@ -193,6 +209,19 @@ class ComputeNodeClient: response.raise_for_status() return _unwrap_dict(response.json()) + # ── Inference helpers (short timeouts — see module constants) ────────── + + async def inference_load(self, payload: dict[str, Any]) -> dict[str, Any]: + """Dispatch a model load. Returns as soon as the node accepts the + request; the node now loads asynchronously (status goes 'loading').""" + return await self._request("POST", "/inference/load", json_data=payload, timeout=INFERENCE_LOAD_TIMEOUT) + + async def inference_status(self) -> dict[str, Any]: + return await self._request("GET", "/inference/status", timeout=INFERENCE_STATUS_TIMEOUT) + + async def inference_unload(self) -> dict[str, Any]: + return await self._request("POST", "/inference/unload", json_data={}, timeout=INFERENCE_UNLOAD_TIMEOUT) + async def upload_file( self, filename: str, diff --git a/backend/app/modules/compute_gateway/sync.py b/backend/app/modules/compute_gateway/sync.py index 27f51e1..4019ba2 100644 --- a/backend/app/modules/compute_gateway/sync.py +++ b/backend/app/modules/compute_gateway/sync.py @@ -1,15 +1,101 @@ 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: @@ -95,5 +181,13 @@ async def poll_compute_jobs_once() -> dict[str, Any]: 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} + "items": synced, "standalone": standalone_synced, "eval_synced": eval_synced, + "inference_reconciled": inference_reconciled} diff --git a/backend/tests/test_compare_inference_async.py b/backend/tests/test_compare_inference_async.py new file mode 100644 index 0000000..075e9ac --- /dev/null +++ b/backend/tests/test_compare_inference_async.py @@ -0,0 +1,254 @@ +""" +模型推理异步加载改造的单元测试。 + +覆盖： +- model_compare_load：异步派发，立即返回 starting + 节点信息（不等待加载完成） +- model_compare_delete：先删记录，卸载失败也不阻塞删除 +- reconcile_inference_loads：starting -> ready/error/idle/不可达的状态迁移与封顶 +- _unload_from_compute_node：任务感知，只命中记录中的节点 +""" + +from __future__ import annotations + +import asyncio +from types import SimpleNamespace +from typing import Any + +from app.api.v1.endpoints.platform import model_compare_delete, model_compare_load +import app.api.v1.endpoints.platform as platform +from app.modules.compute_gateway.client import ComputeNodeClient +from app.modules.compute_gateway.sync import MAX_STARTING_ATTEMPTS, reconcile_inference_loads + + +class FakeInferenceStore: + """内存 store，仅实现推理加载/对账用到的接口。""" + + def __init__(self, tasks: list[dict[str, Any]] | None = None, nodes: list[dict[str, Any]] | None = None) -> None: + self._tasks: dict[str, dict[str, Any]] = {t["id"]: dict(t) for t in (tasks or [])} + self._nodes = nodes or [] + self._inference_nodes: set[str] = set() + + def compare_task(self, task_id: str) -> dict[str, Any]: + if task_id not in self._tasks: + raise KeyError(task_id) + return dict(self._tasks[task_id]) + + def compare_tasks(self) -> list[dict[str, Any]]: + return [dict(t) for t in self._tasks.values()] + + def update_compare_task(self, task_id: str, payload: dict[str, Any]) -> dict[str, Any]: + current = self._tasks[task_id] + merged = {**current, **payload, "id": task_id} + self._tasks[task_id] = merged + return dict(merged) + + def delete_compare_task(self, task_id: str) -> None: + self._tasks.pop(task_id, None) + + def compute_nodes(self) -> list[dict[str, Any]]: + return [dict(n) for n in self._nodes] + + def model(self, model_id: str) -> dict[str, Any]: + raise KeyError(model_id) + + def trained_models(self) -> list[dict[str, Any]]: + return [] + + 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 + + +def _node(node_id: str, code: str = "") -> dict[str, Any]: + return { + "id": node_id, + "code": code or node_id, + "name": code or node_id, + "api_base_url": f"http://{code or node_id}:19100", + "enabled": True, + "scheduler_status": "online", + } + + +def _task(task_id: str, *, node_id: str | None = None, load_status: dict[str, Any] | None = None) -> dict[str, Any]: + return { + "id": task_id, + "name": f"task-{task_id}", + "status": "pending", + "models": [ + {"model_id": "m_1", "model_name": "qwen", "model_path": "/models/qwen", "node_id": node_id} + ], + "load_status": load_status or {"loaded_models": []}, + } + + +async def _fake_inference_load(self, payload: dict[str, Any]) -> dict[str, Any]: + return {"loaded": False, "status": "loading", "request_id": "req-1"} + + +async def _fake_inference_unload(self) -> dict[str, Any]: + return {"unloaded": True, "status": "idle"} + + +def _patch_store(monkeypatch, store: FakeInferenceStore) -> None: + monkeypatch.setattr(platform, "get_platform_store", lambda: store) + monkeypatch.setattr(platform, "get_settings", lambda: SimpleNamespace(compute_mode="real")) + + +def test_model_compare_load_dispatches_and_returns_starting(monkeypatch) -> None: + store = FakeInferenceStore(tasks=[_task("t1", node_id="n1")], nodes=[_node("n1")]) + _patch_store(monkeypatch, store) + monkeypatch.setattr(ComputeNodeClient, "inference_load", _fake_inference_load) + + result = asyncio.run(model_compare_load("t1")) + assert result["code"] == 0 + updated = result["data"] + assert updated["status"] == "starting" + items = updated["load_status"]["loaded_models"] + assert items[0]["status"] == "starting" + assert items[0]["node_id"] == "n1" + assert "n1" in store._inference_nodes + + +def test_model_compare_load_marks_error_when_all_nodes_fail(monkeypatch) -> None: + store = FakeInferenceStore(tasks=[_task("t1", node_id="n1")], nodes=[_node("n1")]) + _patch_store(monkeypatch, store) + + async def _raise(self, payload: dict[str, Any]) -> dict[str, Any]: + raise RuntimeError("conn refused") + + monkeypatch.setattr(ComputeNodeClient, "inference_load", _raise) + + result = asyncio.run(model_compare_load("t1")) + updated = result["data"] + assert updated["status"] == "failed" + assert updated["load_status"]["loaded_models"][0]["status"] == "error" + assert "conn refused" in updated["load_status"]["loaded_models"][0]["error"] + + +def test_model_compare_delete_removes_record_even_if_unload_raises(monkeypatch) -> None: + task = _task( + "t1", + node_id="n1", + load_status={"loaded_models": [{"model_id": "m_1", "status": "ready", "node_id": "n1"}]}, + ) + store = FakeInferenceStore(tasks=[task], nodes=[_node("n1")]) + _patch_store(monkeypatch, store) + + async def _raise(self) -> dict[str, Any]: + raise RuntimeError("unload boom") + + monkeypatch.setattr(ComputeNodeClient, "inference_unload", _raise) + + result = asyncio.run(model_compare_delete("t1")) + assert result["data"] == {"deleted": "t1"} + assert "t1" not in store._tasks + # finally 中仍清掉了节点标记 + assert "n1" not in store._inference_nodes + + +def test_unload_from_compute_node_only_hits_recorded_node(monkeypatch) -> None: + task = _task( + "t1", + load_status={"loaded_models": [{"model_id": "m_1", "status": "ready", "node_id": "n1"}]}, + ) + store = FakeInferenceStore(tasks=[task], nodes=[_node("n1"), _node("n2")]) + _patch_store(monkeypatch, store) + monkeypatch.setattr(ComputeNodeClient, "inference_unload", _fake_inference_unload) + + from app.api.v1.endpoints.platform import _unload_from_compute_node + + result = asyncio.run(_unload_from_compute_node(store, task=task)) + assert result["unloaded"] is True + # 只命中任务记录中的节点 n1，n2 未被卸载 + assert [r["node_id"] for r in result["nodes"]] == ["n1"] + assert "n1" not in store._inference_nodes + + +async def _status_ready(self) -> dict[str, Any]: + return {"loaded": True, "status": "ready", "model_name": "qwen"} + + +def test_reconcile_transitions_starting_to_ready(monkeypatch) -> None: + task = _task( + "t1", + node_id="n1", + load_status={"loaded_models": [{"model_id": "m_1", "status": "starting", "node_id": "n1"}]}, + ) + store = FakeInferenceStore(tasks=[task], nodes=[_node("n1")]) + monkeypatch.setattr(ComputeNodeClient, "inference_status", _status_ready) + + reconciled = asyncio.run(reconcile_inference_loads(store)) + assert reconciled == [{"task_id": "t1", "status": "loaded"}] + updated = store._tasks["t1"] + assert updated["status"] == "loaded" + assert updated["load_status"]["loaded_models"][0]["status"] == "ready" + assert "n1" in store._inference_nodes + + +def test_reconcile_transitions_to_error_and_failed(monkeypatch) -> None: + async def _status_error(self) -> dict[str, Any]: + return {"loaded": False, "status": "error", "error": "CUDA out of memory"} + + task = _task( + "t1", + node_id="n1", + load_status={"loaded_models": [{"model_id": "m_1", "status": "starting", "node_id": "n1"}]}, + ) + store = FakeInferenceStore(tasks=[task], nodes=[_node("n1")]) + monkeypatch.setattr(ComputeNodeClient, "inference_status", _status_error) + + reconciled = asyncio.run(reconcile_inference_loads(store)) + assert reconciled == [{"task_id": "t1", "status": "failed"}] + item = store._tasks["t1"]["load_status"]["loaded_models"][0] + assert item["status"] == "error" + assert "CUDA out of memory" in item["error"] + assert "n1" not in store._inference_nodes + + +def test_reconcile_idle_marks_model_disappeared(monkeypatch) -> None: + async def _status_idle(self) -> dict[str, Any]: + return {"loaded": False, "status": "idle"} + + task = _task( + "t1", + node_id="n1", + load_status={"loaded_models": [{"model_id": "m_1", "status": "starting", "node_id": "n1"}]}, + ) + store = FakeInferenceStore(tasks=[task], nodes=[_node("n1")]) + monkeypatch.setattr(ComputeNodeClient, "inference_status", _status_idle) + + asyncio.run(reconcile_inference_loads(store)) + item = store._tasks["t1"]["load_status"]["loaded_models"][0] + assert item["status"] == "error" + assert "disappeared" in item["error"] + assert store._tasks["t1"]["status"] == "failed" + + +def test_reconcile_unreachable_node_flips_to_error_after_cap(monkeypatch) -> None: + async def _raise(self) -> dict[str, Any]: + raise RuntimeError("conn refused") + + task = _task( + "t1", + node_id="n1", + load_status={"loaded_models": [{"model_id": "m_1", "status": "starting", "node_id": "n1"}]}, + ) + store = FakeInferenceStore(tasks=[task], nodes=[_node("n1")]) + monkeypatch.setattr(ComputeNodeClient, "inference_status", _raise) + + # 每次轮询前重置节流时间戳，逐次推进 load_attempts 到封顶 + for _ in range(MAX_STARTING_ATTEMPTS): + item = store._tasks["t1"]["load_status"]["loaded_models"][0] + item["last_polled_at"] = 0 + asyncio.run(reconcile_inference_loads(store)) + + item = store._tasks["t1"]["load_status"]["loaded_models"][0] + assert item["status"] == "error" + assert "unreachable" in item["error"] + assert store._tasks["t1"]["status"] == "failed" diff --git a/compute/api/main.py b/compute/api/main.py index d88d2fd..aa51029 100644 --- a/compute/api/main.py +++ b/compute/api/main.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio import json import os import math @@ -717,7 +718,9 @@ def create_app() -> FastAPI: @app.post(f"{route_prefix}/inference/unload") async def inference_unload() -> dict[str, Any]: """Unload the currently loaded model and free GPU memory.""" - return get_inference_session().unload() + # Teardown (gc.collect + cuda.empty_cache) can take a while; run it off + # the event loop so /health and /inference/status stay responsive. + return await asyncio.to_thread(get_inference_session().unload) @app.get(f"{route_prefix}/inference/status") async def inference_status() -> dict[str, Any]: @@ -737,7 +740,10 @@ def create_app() -> FastAPI: messages = payload.get("messages") or [] if not messages: raise HTTPException(status_code=400, detail="messages is required") - result = get_inference_session().chat( + # Generation is long-running; run it in a thread so the event loop keeps + # serving /inference/status and /health during inference. + result = await asyncio.to_thread( + get_inference_session().chat, messages=messages, temperature=float(payload.get("temperature", 0.95)), top_p=float(payload.get("top_p", 0.7)), diff --git a/compute/engines/llama_factory/inference.py b/compute/engines/llama_factory/inference.py index 0afb036..dee77cd 100644 --- a/compute/engines/llama_factory/inference.py +++ b/compute/engines/llama_factory/inference.py @@ -2,114 +2,208 @@ from __future__ import annotations import threading import time -from typing import Any +import uuid +from typing import Any, Iterator class InferenceSession: - """Manages a loaded model for inference with LLaMA-Factory ChatModel.""" + """Manages a loaded model for inference with LLaMA-Factory ChatModel. + + Model loading is asynchronous: ``load()`` spawns a background daemon thread + and returns immediately with ``status == "loading"``. ``info()`` (served by + ``/inference/status``) is always responsive, so the platform backend can + poll loading progress without being blocked by a minutes-long model load — + which previously froze the whole compute node event loop. + + State machine: idle -> loading -> ready | error, ready -> idle (unload), + loading -> idle (cancelled). Long operations (ChatModel build, teardown, + generation) never run while holding ``_state_lock``; they either run in the + worker thread or under ``_chat_lock`` only. + """ def __init__(self) -> None: + self._state_lock = threading.Lock() # brief state transitions only + self._chat_lock = threading.Lock() # serialize chat/teardown + self._status: str = "idle" + self._error: str = "" + self._request_id: str = "" + self._load_args: dict[str, Any] = {} + self._teardown_old = False # load-while-ready: unload old before loading new + self._cancel_requested = False # unload-while-loading: tear down after load finishes + self._load_thread: threading.Thread | None = None self._model: Any = None self._tokenizer: Any = None self._generating_args: dict[str, Any] = {} self._model_name: str = "" self._adapter_path: str = "" - self._lock = threading.Lock() self._loaded_at: float = 0.0 - self._status: str = "idle" @property def status(self) -> str: - return self._status - - @property - def model_name(self) -> str: - return self._model_name - - @property - def adapter_path(self) -> str: - return self._adapter_path - - @property - def loaded_at(self) -> float: - return self._loaded_at + with self._state_lock: + return self._status def info(self) -> dict[str, Any]: - return { - "loaded": self._status == "ready", - "status": self._status, - "model_name": self._model_name, - "adapter_path": self._adapter_path, - "loaded_at": self._loaded_at, - } + with self._state_lock: + return { + "loaded": self._status == "ready", + "status": self._status, + "model_name": self._model_name, + "adapter_path": self._adapter_path, + "loaded_at": self._loaded_at, + "request_id": self._request_id, + "error": self._error, + } - def load(self, model_name_or_path, adapter_name_or_path="", template="qwen", infer_backend="huggingface", infer_dtype="auto", **kwargs): - with self._lock: + def load( + self, + model_name_or_path, + adapter_name_or_path="", + template="qwen", + infer_backend="huggingface", + infer_dtype="auto", + **kwargs, + ) -> dict[str, Any]: + with self._state_lock: if self._status == "loading": - return {"loaded": False, "error": "model is already loading"} - if self._status == "ready": - self.unload() + # A model is already loading — dedupe, reuse the same request id. + return {"loaded": False, "status": "loading", "request_id": self._request_id} + self._teardown_old = self._status == "ready" self._status = "loading" + self._error = "" + self._request_id = uuid.uuid4().hex[:12] + self._cancel_requested = False + self._load_args = { + "model_name_or_path": model_name_or_path, + "template": template, + "infer_backend": infer_backend, + "infer_dtype": infer_dtype, + } + if adapter_name_or_path: + self._load_args["adapter_name_or_path"] = adapter_name_or_path + self._load_args.update(kwargs) self._model_name = model_name_or_path self._adapter_path = adapter_name_or_path - try: - from llamafactory.chat import ChatModel - from llamafactory.hparams import get_infer_args - args = {"model_name_or_path": model_name_or_path, "template": template, "infer_backend": infer_backend, "infer_dtype": infer_dtype} - if adapter_name_or_path: - args["adapter_name_or_path"] = adapter_name_or_path - args.update(kwargs) - infer_result = get_infer_args(args) - # ChatModel internally re-parses the args dict via get_infer_args, - # so pass the original args (not the parsed dataclass objects). - self._model = ChatModel(args) - self._tokenizer = getattr(self._model, 'tokenizer', None) or self._model.engine.tokenizer - # Extract generating_args (last element) for later use in chat() - generating_args = infer_result[-1] - if hasattr(generating_args, '__dataclass_fields__'): - self._generating_args = {k: v for k, v in vars(generating_args).items() - if not k.startswith('_')} - else: - self._generating_args = dict(generating_args) - self._loaded_at = time.time() - self._status = "ready" - return {"loaded": True, "status": "ready"} - except Exception as exc: - self._status = "error" - self._model = None - return {"loaded": False, "status": "error", "error": str(exc)} + self._load_thread = threading.Thread(target=self._load_worker, daemon=True) + self._load_thread.start() + return {"loaded": False, "status": "loading", "request_id": self._request_id} - def unload(self): - with self._lock: - if self._model is not None: - try: - del self._model - except Exception: - pass + def _load_worker(self) -> None: + """Build the ChatModel off the state lock so info() never blocks.""" + model = None + tokenizer = None + generating_args: dict[str, Any] = {} + error = "" + try: + if self._teardown_old: + self._release_model() + from llamafactory.chat import ChatModel + from llamafactory.hparams import get_infer_args + + args = dict(self._load_args) + infer_result = get_infer_args(args) + model = ChatModel(args) + tokenizer = getattr(model, "tokenizer", None) or model.engine.tokenizer + generating_args = infer_result[-1] + if hasattr(generating_args, "__dataclass_fields__"): + generating_args = { + k: v for k, v in vars(generating_args).items() if not k.startswith("_") + } + else: + generating_args = dict(generating_args) + except Exception as exc: # noqa: BLE001 - surface load failure via status + error = str(exc) + with self._state_lock: + if error: self._model = None self._tokenizer = None + self._status = "error" + self._error = error + return + if self._cancel_requested: + # Unload was requested while loading — drop the fresh model. + model = None + tokenizer = None + self._model = None + self._tokenizer = None + self._status = "idle" + return + self._model = model + self._tokenizer = tokenizer + self._generating_args = generating_args + self._loaded_at = time.time() + self._status = "ready" + + def _release_model(self) -> None: + with self._chat_lock: + with self._state_lock: + self._status = "unloading" + model = self._model + self._model = None + self._tokenizer = None + if model is not None: + try: + del model + except Exception: # noqa: BLE001 - best-effort teardown + pass # 强制释放 PyTorch CUDA 缓存，真正归还 GPU 显存 try: import gc + gc.collect() import torch + if torch.cuda.is_available(): torch.cuda.empty_cache() torch.cuda.synchronize() - except Exception: + except Exception: # noqa: BLE001 - teardown must not raise pass - self._status = "idle" - self._model_name = "" - self._adapter_path = "" - self._loaded_at = 0.0 - return {"unloaded": True} + with self._state_lock: + self._status = "idle" + self._model_name = "" + self._adapter_path = "" + self._loaded_at = 0.0 + self._error = "" - def chat(self, messages, temperature=0.95, top_p=0.7, max_new_tokens=1024, do_sample=True, **kwargs): - with self._lock: - if self._status != "ready" or self._model is None: - return {"error": "model not loaded", "response": ""} + def unload(self) -> dict[str, Any]: + with self._state_lock: + if self._status == "loading": + # Ask the worker to tear down right after the load finishes. + self._cancel_requested = True + return {"unloaded": False, "status": "cancelling", "request_id": self._request_id} + was_ready = self._status == "ready" + if was_ready: + self._release_model() + else: + with self._state_lock: + self._model = None + self._tokenizer = None + self._status = "idle" + self._model_name = "" + self._adapter_path = "" + self._loaded_at = 0.0 + self._error = "" + return {"unloaded": True, "status": "idle"} + + def chat(self, messages, temperature=0.95, top_p=0.7, max_new_tokens=1024, do_sample=True, **kwargs) -> dict[str, Any]: + with self._chat_lock: + with self._state_lock: + if self._status == "loading": + return { + "error": f"model is still loading (request_id={self._request_id}); please retry", + "response": "", + } + if self._status == "error": + return {"error": f"model load failed: {self._error}", "response": ""} + if self._status != "ready" or self._model is None: + return {"error": "model not loaded", "response": ""} try: - generate_kwargs = {"temperature": temperature, "top_p": top_p, "max_new_tokens": max_new_tokens, "do_sample": do_sample} + generate_kwargs = { + "temperature": temperature, + "top_p": top_p, + "max_new_tokens": max_new_tokens, + "do_sample": do_sample, + } generate_kwargs.update(kwargs) system = next((m["content"] for m in messages if m["role"] == "system"), None) user_messages = [m for m in messages if m["role"] != "system"] @@ -118,27 +212,35 @@ class InferenceSession: responses.append(response) full_response = "".join(str(r) for r in responses) return {"response": full_response} - except Exception as exc: + except Exception as exc: # noqa: BLE001 - return generation error to caller return {"error": str(exc), "response": ""} - def chat_stream(self, messages, **kwargs): - with self._lock: - if self._status != "ready" or self._model is None: - yield 'data: {"error": "model not loaded"}\n\n' - return + def chat_stream(self, messages, **kwargs) -> Iterator[str]: + with self._chat_lock: + with self._state_lock: + if self._status == "loading": + yield 'data: {"error": "model is still loading; please retry"}\n\n' + return + if self._status == "error": + yield 'data: {"error": "model load failed: ' + str(self._error) + '"}\n\n' + return + if self._status != "ready" or self._model is None: + yield 'data: {"error": "model not loaded"}\n\n' + return try: generate_kwargs = {**kwargs} system = next((m["content"] for m in messages if m["role"] == "system"), None) user_messages = [m for m in messages if m["role"] != "system"] for new_text in self._model.stream_chat(user_messages, system=system, **generate_kwargs): yield new_text - except Exception as exc: + except Exception as exc: # noqa: BLE001 - stream error as SSE event yield 'data: {"error": "' + str(exc) + '"}\n\n' _inference_session = None -def get_inference_session(): + +def get_inference_session() -> InferenceSession: global _inference_session if _inference_session is None: _inference_session = InferenceSession() diff --git a/compute/tests/test_inference_session.py b/compute/tests/test_inference_session.py new file mode 100644 index 0000000..ab55514 --- /dev/null +++ b/compute/tests/test_inference_session.py @@ -0,0 +1,127 @@ +from __future__ import annotations + +import sys +import time +import types +from typing import Any + +import pytest + +from compute.engines.llama_factory.inference import InferenceSession + +# 模拟模型加载耗时，用于验证 load() 立即返回、info() 不阻塞 +LOAD_DELAY = 0.2 + + +class FakeChatModel: + def __init__(self, args: dict[str, Any]) -> None: + time.sleep(LOAD_DELAY) + self.tokenizer = object() + self.engine = types.SimpleNamespace(tokenizer=object()) + self._output = "hello from model" + + def stream_chat(self, *args, **kwargs): + for _ in range(1): + yield self._output + + +class FailingChatModel: + def __init__(self, args: dict[str, Any]) -> None: + time.sleep(LOAD_DELAY) + raise RuntimeError("boom: fake load failure") + + +def _get_infer_args(args: dict[str, Any]) -> list[Any]: + # 最后一个元素为 generating_args，worker 会转成 dict + return [None, None, {"temperature": 0.7}] + + +def _install_llamafactory(monkeypatch, chat_model: type) -> None: + llmf = types.ModuleType("llamafactory") + chat_mod = types.ModuleType("llamafactory.chat") + hparams_mod = types.ModuleType("llamafactory.hparams") + chat_mod.ChatModel = chat_model + hparams_mod.get_infer_args = _get_infer_args + llmf.chat = chat_mod + llmf.hparams = hparams_mod + monkeypatch.setitem(sys.modules, "llamafactory", llmf) + monkeypatch.setitem(sys.modules, "llamafactory.chat", chat_mod) + monkeypatch.setitem(sys.modules, "llamafactory.hparams", hparams_mod) + + +@pytest.fixture +def stub_llamafactory(monkeypatch) -> None: + _install_llamafactory(monkeypatch, FakeChatModel) + + +@pytest.fixture +def stub_failing_llamafactory(monkeypatch) -> None: + _install_llamafactory(monkeypatch, FailingChatModel) + + +def _wait_for_status(session: InferenceSession, status: str, timeout: float = 3.0) -> bool: + deadline = time.time() + timeout + while time.time() < deadline: + if session.info()["status"] == status: + return True + time.sleep(0.02) + return False + + +def test_load_returns_immediately_then_ready(stub_llamafactory) -> None: + session = InferenceSession() + started = time.time() + result = session.load("/models/qwen") + assert result["status"] == "loading" + assert result["loaded"] is False + assert result["request_id"] + # 在慢加载完成前就返回，且 info() 加载期间可响应 + assert time.time() - started < LOAD_DELAY + assert session.info()["status"] == "loading" + assert _wait_for_status(session, "ready") + info = session.info() + assert info["loaded"] is True + assert info["status"] == "ready" + assert info["model_name"] == "/models/qwen" + + +def test_second_load_while_loading_deduped(stub_llamafactory) -> None: + session = InferenceSession() + r1 = session.load("/models/a") + r2 = session.load("/models/b") + assert r2["status"] == "loading" + assert r2["request_id"] == r1["request_id"] + assert _wait_for_status(session, "ready") + assert session.info()["status"] == "ready" + + +def test_load_error_surfaces_in_status(stub_failing_llamafactory) -> None: + session = InferenceSession() + session.load("/models/bad") + assert _wait_for_status(session, "error") + assert "boom" in session.info()["error"] + + +def test_unload_while_loading_cancels(stub_llamafactory) -> None: + session = InferenceSession() + session.load("/models/qwen") + result = session.unload() + assert result["status"] == "cancelling" + assert _wait_for_status(session, "idle") + + +def test_chat_while_loading_returns_loading_error(stub_llamafactory) -> None: + session = InferenceSession() + session.load("/models/qwen") + out = session.chat([{"role": "user", "content": "hi"}]) + assert "still loading" in (out.get("error") or "") + assert _wait_for_status(session, "ready") + out = session.chat([{"role": "user", "content": "hi"}]) + assert out.get("response") == "hello from model" + + +def test_chat_stream_while_loading_yields_error(stub_llamafactory) -> None: + session = InferenceSession() + session.load("/models/qwen") + chunks = list(session.chat_stream([{"role": "user", "content": "hi"}])) + assert any("still loading" in c for c in chunks) diff --git a/docker/nginx.conf.template b/docker/nginx.conf.template index 93c196a..da30293 100644 --- a/docker/nginx.conf.template +++ b/docker/nginx.conf.template @@ -14,8 +14,8 @@ server { proxy_set_header X-Real-IP $remote_addr; proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for; proxy_set_header X-Forwarded-Proto $scheme; - proxy_read_timeout 300s; - proxy_send_timeout 300s; + proxy_read_timeout 900s; + proxy_send_timeout 900s; } location = /modelTF { @@ -25,8 +25,8 @@ server { proxy_set_header X-Real-IP $remote_addr; proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for; proxy_set_header X-Forwarded-Proto $scheme; - proxy_read_timeout 300s; - proxy_send_timeout 300s; + proxy_read_timeout 900s; + proxy_send_timeout 900s; } location ~* \.(?:js|css|png|jpg|jpeg|gif|ico|svg|woff|woff2|ttf)$ { diff --git a/frontend/src/api/modules/compare.ts b/frontend/src/api/modules/compare.ts index 35f1ee5..c58f195 100644 --- a/frontend/src/api/modules/compare.ts +++ b/frontend/src/api/modules/compare.ts @@ -1,6 +1,8 @@ import { get, post, del } from '../request' import type { CompareTask, CompareModelRef } from '@/types' +const INFERENCE_START_TIMEOUT_MS = 15 * 60 * 1000 + /** 推理/对比任务列表 */ export const getCompareList = () => get('/model-compare') @@ -12,7 +14,7 @@ export const createCompare = (data: Partial) => post<{ id: string | number }>('/model-compare', data) /** 删除任务 */ -export const deleteCompare = (id: string | number) => del(`/model-compare/${id}`) +export const deleteCompare = (id: string | number) => del(`/model-compare/${id}`, undefined, { timeout: 60_000 }) /** 更新任务加载状态 */ export const updateLoadStatus = (id: string | number, load_status: any) => @@ -34,7 +36,8 @@ export const stopModelByPid = (pid: number) => post('/model-compare/stop-by-pid', { pid }) /** 加载任务 */ -export const loadCompare = (id: string | number) => post(`/model-compare/${id}/load`) +export const loadCompare = (id: string | number) => + post(`/model-compare/${id}/load`, undefined, { timeout: INFERENCE_START_TIMEOUT_MS }) /** 卸载任务 */ export const unloadCompare = (id: string | number) => post(`/model-compare/${id}/unload`) @@ -81,6 +84,10 @@ export const streamChatReal = (data: any): Promise => { temperature: data.temperature ?? 0.7, top_p: data.top_p ?? 0.95, max_tokens: data.max_tokens ?? 2048, + // 透传 task_id/node_id，让后端按 load_status 路由到真正加载了模型的算力节点， + // 避免在多节点时回退到“第一个在线节点”导致连接失败 + task_id: data.task_id, + node_id: data.node_id, }), }) } @@ -94,8 +101,8 @@ export const batchChat = (data: any) => post('/model-chat/batch', data) /** 本地 transformers 模型对话 */ export const localChat = (data: any) => post('/model-chat/local/chat', data) -/** 预加载本地模型（模型加载耗时长，超时 5 分钟） */ -export const preloadLocalModel = (data: any) => post('/model-chat/local/preload', data, { timeout: 300000 }) +/** 预加载本地模型（模型加载耗时长，超时 15 分钟） */ +export const preloadLocalModel = (data: any) => post('/model-chat/local/preload', data, { timeout: INFERENCE_START_TIMEOUT_MS }) -/** 预加载已训练模型（超时 5 分钟） */ -export const preloadTrainedModel = (data: any) => post('/model-chat/trained/preload', data, { timeout: 300000 }) +/** 预加载已训练模型（超时 15 分钟） */ +export const preloadTrainedModel = (data: any) => post('/model-chat/trained/preload', data, { timeout: INFERENCE_START_TIMEOUT_MS }) diff --git a/frontend/src/api/modules/fineTune.ts b/frontend/src/api/modules/fineTune.ts index 4570af9..374ed49 100644 --- a/frontend/src/api/modules/fineTune.ts +++ b/frontend/src/api/modules/fineTune.ts @@ -1,6 +1,16 @@ import { get, post, put, del } from '../request' import type { FineTuneStartPayload, FineTuneTask, TrainingProgress, LogContent } from '@/types' +export interface FineTuneMetricPoint { + step: number + epoch?: number | null + loss?: number | null + grad_norm?: number | null + learning_rate?: number | null + raw?: string + create_time?: string +} + export interface TrainingDiagnostic { level: string title: string @@ -80,6 +90,10 @@ export const getFineTuneLogs = ( params: { tail_lines?: number; offset?: number; limit?: number } = {}, ) => get(`/fine-tune/${id}/logs`, params) +/** 获取训练指标曲线数据 */ +export const getFineTuneMetrics = (id: string | number) => + get(`/fine-tune/${id}/metrics`) + /** 启动 TensorBoard */ export const startTensorboard = () => post('/fine-tune/tensorboard/start') diff --git a/frontend/src/api/modules/model.ts b/frontend/src/api/modules/model.ts index 0b61d3e..2b063fe 100644 --- a/frontend/src/api/modules/model.ts +++ b/frontend/src/api/modules/model.ts @@ -88,10 +88,14 @@ export const updateModelPurpose = (id: string | number, purpose: string) => /** 合并 LoRA 权重 */ export const mergeModel = (data: { + trained_model_id?: string | number model_name: string train_method: string base_model_path: string -}) => post('/model-manage/merge', data) + adapter_path?: string + compute_node_id?: string + output_model_name?: string +}) => post('/model-manage/merge', data, { timeout: 15 * 60 * 1000 }) /** 导出已训练模型权重 */ export const exportModelUrl = (modelName: string) => diff --git a/frontend/src/composables/useStreamChat.ts b/frontend/src/composables/useStreamChat.ts index d322529..2e10783 100644 --- a/frontend/src/composables/useStreamChat.ts +++ b/frontend/src/composables/useStreamChat.ts @@ -44,6 +44,24 @@ export function useStreamChat() { }) const loading = ref(false) + /** 从 SSE 帧中提取错误信息（后端/计算节点错误以 data: {"error": "..."} 形式下发） */ + function extractSseError(buffer: string): string | null { + const trimmed = buffer.trim() + if (!trimmed.startsWith('data: ')) return null + const lines = trimmed.split(/\r?\n/) + for (let i = lines.length - 1; i >= 0; i--) { + const line = lines[i].trim() + if (!line.startsWith('data: ')) continue + try { + const obj = JSON.parse(line.slice(6)) + if (obj && typeof obj.error === 'string' && obj.error) return obj.error + } catch { + /* 非 JSON 的 data 行忽略 */ + } + } + return trimmed + } + /** 从内容中解析 think 标签 */ function parseContent(content: string) { const thinkRegex = /([\s\S]*?)(<\/think>)?/g @@ -121,6 +139,16 @@ export function useStreamChat() { } // 最终更新 + // 若整段响应是 SSE 错误帧，提取 error 字段以干净文案展示 + const sseError = extractSseError(buffer) + if (sseError) { + message.value.isThinking = false + message.value.isStreaming = false + message.value.done = true + message.value.error = sseError + message.value.displayContent = sseError + return + } const parsed = parseContent(buffer) message.value.thinkContent = parsed.think message.value.displayContent = parsed.display diff --git a/frontend/src/types/index.ts b/frontend/src/types/index.ts index a4d10df..6bac76a 100644 --- a/frontend/src/types/index.ts +++ b/frontend/src/types/index.ts @@ -34,6 +34,10 @@ export interface TrainedModel { name: string train_methods?: TrainMethod[] base_model_path?: string + artifact_dir?: string + adapter_path?: string + compute_node_id?: string + compute_node_name?: string create_time?: string merged?: boolean merging?: boolean @@ -212,6 +216,9 @@ export interface LoadedModel { status?: string pid?: number port?: number + node_id?: string + node_name?: string + error?: string } export interface CompareTask { @@ -230,6 +237,8 @@ export interface CompareModelRef { model_name: string model_path: string gpu_id: number + node_id?: string + node_name?: string source?: string port?: number } diff --git a/frontend/src/views/inference/InferenceChatView.vue b/frontend/src/views/inference/InferenceChatView.vue index 1700abd..d1176d0 100644 --- a/frontend/src/views/inference/InferenceChatView.vue +++ b/frontend/src/views/inference/InferenceChatView.vue @@ -1,10 +1,10 @@