From 77e9ac6b7ed2bae23a5e1f5627db08556146db3a Mon Sep 17 00:00:00 2001 From: cat-shark Date: Thu, 13 Aug 2026 22:01:42 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E5=A2=9E=E5=8A=A0=E5=A4=B1=E8=B4=A5?= =?UTF-8?q?=E4=BB=BB=E5=8A=A1=E9=87=8D=E8=AF=95=E4=B8=8E=20CUDA=20?= =?UTF-8?q?=E5=8A=A8=E6=80=81=E5=BA=93=E8=B7=AF=E5=BE=84=E6=B3=A8=E5=85=A5?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- README.md | 7 ++ app/config.py | 1 + app/db.py | 13 +++ app/node_manager.py | 43 ++++++++- app/routers/apps.py | 11 +++ tests/test_apps_api.py | 54 +++++++++++ tests/test_db.py | 41 +++++++++ tests/test_node_manager.py | 179 ++++++++++++++++++++++++++++++++++++- 8 files changed, 345 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index 6e673fd..24b806b 100644 --- a/README.md +++ b/README.md @@ -72,6 +72,13 @@ $env:LLM_MODEL="default" 该接口无需 key；如不使用当前默认接口，再按需设置 `LLM_API_BASE` 和 `LLM_API_KEY`。 +LLM 较慢时按需放宽超时（默认单次请求 600 秒、整节点调用 3600 秒）： + +```powershell +$env:LLM_TIMEOUT_SECONDS="600" +$env:WOV_NODE_INVOKE_TIMEOUT_SECONDS="3600" +``` + ## 使用 1. 打开 `http://127.0.0.1:8000/`。 diff --git a/app/config.py b/app/config.py index 67b0d88..362b7ff 100644 --- a/app/config.py +++ b/app/config.py @@ -16,3 +16,4 @@ STORAGE_DIR = Path(os.getenv("WOV_STORAGE_DIR", str(DATA_DIR / "storage"))) NODE_REAP_INTERVAL_SECONDS = float(os.getenv("WOV_REAP_INTERVAL_SECONDS", "3")) NODE_READY_TIMEOUT_SECONDS = float(os.getenv("WOV_READY_TIMEOUT_SECONDS", "12")) +NODE_INVOKE_TIMEOUT_SECONDS = float(os.getenv("WOV_NODE_INVOKE_TIMEOUT_SECONDS", "3600")) diff --git a/app/db.py b/app/db.py index 220d37e..e5b9b4f 100644 --- a/app/db.py +++ b/app/db.py @@ -330,6 +330,19 @@ class Database: with self._connect() as conn: conn.execute(f"UPDATE workflow_runs SET {assignments} WHERE id = ?", values) + def reset_run(self, run_id: str, updated_at: str) -> None: + with self._connect() as conn: + conn.execute( + """ + UPDATE workflow_runs + SET status = 'QUEUED', current_node_id = NULL, progress = 0, + error = NULL, updated_at = ? + WHERE id = ? + """, + (updated_at, run_id), + ) + conn.execute("DELETE FROM artifacts WHERE run_id = ?", (run_id,)) + def next_queued_run(self) -> dict[str, Any] | None: with self._connect() as conn: row = conn.execute( diff --git a/app/node_manager.py b/app/node_manager.py index 661de4b..0647491 100644 --- a/app/node_manager.py +++ b/app/node_manager.py @@ -15,7 +15,13 @@ from pathlib import Path from wov_sdk.models import InvokeRequest, InvokeResponse, NodeManifest -from app.config import NODE_READY_TIMEOUT_SECONDS, NODE_REAP_INTERVAL_SECONDS, SDK_SRC, WORKSPACE_ROOT +from app.config import ( + NODE_INVOKE_TIMEOUT_SECONDS, + NODE_READY_TIMEOUT_SECONDS, + NODE_REAP_INTERVAL_SECONDS, + SDK_SRC, + WORKSPACE_ROOT, +) from app.db import Database @@ -33,6 +39,36 @@ def _idle_seconds(last_used_at: str) -> float: return 0.0 +def _cuda_library_dirs(repo_dir: Path) -> list[Path]: + venv = repo_dir / ".venv" + site_packages: list[Path] = [] + windows_site = venv / "Lib" / "site-packages" + if windows_site.is_dir(): + site_packages.append(windows_site) + site_packages.extend((venv / "lib").glob("python3*/site-packages")) + + dirs: list[Path] = [] + for site in site_packages: + for vendor in ("cublas", "cudnn"): + for subdir in ("lib", "bin"): + lib_dir = site / "nvidia" / vendor / subdir + if lib_dir.is_dir(): + dirs.append(lib_dir) + return dirs + + +def _with_cuda_library_path(env: dict[str, str], repo_dir: Path) -> dict[str, str]: + dirs = [str(path) for path in _cuda_library_dirs(repo_dir)] + if not dirs: + return env + var = "PATH" if os.name == "nt" else "LD_LIBRARY_PATH" + existing = [item for item in env.get(var, "").split(os.pathsep) if item] + additions = [path for path in dirs if path not in existing] + if additions: + env[var] = os.pathsep.join(additions + existing) + return env + + @dataclass class NodeRuntime: instance_id: str @@ -107,7 +143,8 @@ class NodeManager: item for item in [str(SDK_SRC), current_pythonpath] if item ) env.update(manifest.env) - return env + repo_dir = WORKSPACE_ROOT / manifest.repo_dir + return _with_cuda_library_path(env, repo_dir) def _start_locked(self, manifest: NodeManifest) -> NodeRuntime: runtime = NodeRuntime( @@ -270,7 +307,7 @@ class NodeManager: method="POST", ) try: - with urllib.request.urlopen(http_request, timeout=300) as response: + with urllib.request.urlopen(http_request, timeout=NODE_INVOKE_TIMEOUT_SECONDS) as response: payload = json.loads(response.read().decode("utf-8")) return InvokeResponse.from_dict(payload) except urllib.error.HTTPError as exc: diff --git a/app/routers/apps.py b/app/routers/apps.py index 52713a9..55a3083 100644 --- a/app/routers/apps.py +++ b/app/routers/apps.py @@ -100,6 +100,17 @@ def get_run(run_id: str, db: Database = Depends(_get_db)) -> dict: return run +@router.post("/api/runs/{run_id}/retry") +def retry_run(run_id: str, db: Database = Depends(_get_db)) -> dict: + run = db.get_run(run_id) + if run is None: + raise HTTPException(status_code=404, detail="run not found") + if run["status"] != "FAILED": + raise HTTPException(status_code=422, detail="only failed runs can be retried") + db.reset_run(run_id, _now_iso()) + return {"id": run_id, "status": "QUEUED"} + + @router.get("/api/runs/{run_id}/artifacts") def list_run_artifacts(run_id: str, db: Database = Depends(_get_db)) -> list[dict]: if db.get_run(run_id) is None: diff --git a/tests/test_apps_api.py b/tests/test_apps_api.py index 08441d4..fd10c78 100644 --- a/tests/test_apps_api.py +++ b/tests/test_apps_api.py @@ -139,3 +139,57 @@ def test_upload_rejects_workflow_without_version() -> None: files={"file": ("x.txt", b"x", "text/plain")}, ) assert response.status_code == 422 + + +def test_retry_failed_run_requeues_and_reruns() -> None: + with TestClient(app) as client: + workflow_id = _create_published_echo_workflow(client) + uploaded = client.post( + f"/api/apps/{workflow_id}/runs", + files={"file": ("sample.txt", b"hello retry", "text/plain")}, + ) + run_id = uploaded.json()["id"] + db = app.state.db + db.update_run( + run_id, + status="FAILED", + error="boom", + updated_at="2026-01-01T00:00:00+00:00", + ) + db.create_artifact( + { + "run_id": run_id, + "node_id": "step", + "name": "stale", + "uri": "stale.txt", + "mime_type": "text/plain", + "size": 1, + } + ) + + response = client.post(f"/api/runs/{run_id}/retry") + + assert response.status_code == 200 + assert response.json() == {"id": run_id, "status": "QUEUED"} + run = client.get(f"/api/runs/{run_id}").json() + assert run["status"] == "QUEUED" + assert run["error"] is None + assert run["artifacts"] == [] + + app.state.scheduler.execute_run(run_id) + completed = client.get(f"/api/runs/{run_id}").json() + assert completed["status"] == "COMPLETED" + assert any(item["name"] == "result" for item in completed["artifacts"]) + + +def test_retry_rejects_non_failed_and_missing_runs() -> None: + with TestClient(app) as client: + workflow_id = _create_published_echo_workflow(client) + uploaded = client.post( + f"/api/apps/{workflow_id}/runs", + files={"file": ("sample.txt", b"x", "text/plain")}, + ) + run_id = uploaded.json()["id"] + + assert client.post(f"/api/runs/{run_id}/retry").status_code == 422 + assert client.post("/api/runs/missing/retry").status_code == 404 diff --git a/tests/test_db.py b/tests/test_db.py index 27fec59..0fecbe1 100644 --- a/tests/test_db.py +++ b/tests/test_db.py @@ -142,3 +142,44 @@ def test_run_and_artifact_crud(tmp_path) -> None: db.delete_run_artifacts("run_1") assert db.list_artifacts("run_1") == [] + + +def test_reset_run_clears_error_and_artifacts(tmp_path) -> None: + db = Database(tmp_path / "wov.db") + db.upsert_workflow({"id": "demo", "name": "Demo", "published": 1, "latest_version": 1}) + now = "2026-01-01T00:00:00+00:00" + db.create_run( + { + "id": "run_1", + "workflow_id": "demo", + "workflow_version": 1, + "status": "FAILED", + "progress": 0.75, + "current_node_id": "translate", + "error": "timed out", + "input_uri": "in.txt", + "created_at": now, + "updated_at": now, + } + ) + db.create_artifact( + { + "run_id": "run_1", + "node_id": "asr", + "name": "asr.srt_uri", + "uri": "out.srt", + "mime_type": "application/x-subrip", + "size": 3, + } + ) + + db.reset_run("run_1", "2026-01-02T00:00:00+00:00") + + run = db.get_run("run_1") + assert run["status"] == "QUEUED" + assert run["progress"] == 0 + assert run["current_node_id"] is None + assert run["error"] is None + assert run["updated_at"] == "2026-01-02T00:00:00+00:00" + assert run["created_at"] == now + assert db.list_artifacts("run_1") == [] diff --git a/tests/test_node_manager.py b/tests/test_node_manager.py index 9b156b1..9a46fa3 100644 --- a/tests/test_node_manager.py +++ b/tests/test_node_manager.py @@ -1,3 +1,4 @@ +import os import subprocess import sys import time @@ -7,7 +8,14 @@ from pathlib import Path import pytest from app.db import Database -from app.node_manager import NodeManager, NodeRuntime, _idle_seconds, _now_iso +from app.node_manager import ( + NodeManager, + NodeRuntime, + _cuda_library_dirs, + _idle_seconds, + _now_iso, + _with_cuda_library_path, +) from wov_sdk.models import InvokeRequest, NodeManifest WORKSPACE = Path(__file__).resolve().parent.parent.parent @@ -86,6 +94,130 @@ def test_resolve_command_prefers_node_venv(db: Database, tmp_path, monkeypatch) assert manager._resolve_command(manifest)[0] == str(unix_python) +def test_cuda_library_dirs_finds_unix_nvidia_libs(tmp_path) -> None: + cublas = ( + tmp_path / ".venv" / "lib" / "python3.12" / "site-packages" / "nvidia" / "cublas" / "lib" + ) + cudnn = ( + tmp_path / ".venv" / "lib" / "python3.12" / "site-packages" / "nvidia" / "cudnn" / "lib" + ) + cublas.mkdir(parents=True) + cudnn.mkdir(parents=True) + + dirs = _cuda_library_dirs(tmp_path) + + assert cublas in dirs + assert cudnn in dirs + + +def test_cuda_library_dirs_finds_windows_nvidia_bins(tmp_path) -> None: + cublas = tmp_path / ".venv" / "Lib" / "site-packages" / "nvidia" / "cublas" / "bin" + cublas.mkdir(parents=True) + + dirs = _cuda_library_dirs(tmp_path) + + assert cublas in dirs + + +def test_with_cuda_library_path_prepends_on_posix(monkeypatch) -> None: + monkeypatch.setattr(os, "name", "posix") + monkeypatch.setattr( + "app.node_manager._cuda_library_dirs", + lambda repo_dir: ["/opt/nvidia/cublas/lib", "/opt/nvidia/cudnn/lib"], + ) + + env = _with_cuda_library_path({"LD_LIBRARY_PATH": "/usr/lib/foo"}, Path("/repo")) + + entries = env["LD_LIBRARY_PATH"].split(os.pathsep) + assert entries[:2] == ["/opt/nvidia/cublas/lib", "/opt/nvidia/cudnn/lib"] + assert entries[-1] == "/usr/lib/foo" + + +def test_with_cuda_library_path_uses_path_on_windows(monkeypatch) -> None: + monkeypatch.setattr(os, "name", "nt") + monkeypatch.setattr( + "app.node_manager._cuda_library_dirs", + lambda repo_dir: ["/opt/nvidia/cublas/bin"], + ) + + env = _with_cuda_library_path({"PATH": "/usr/local/bin"}, Path("/repo")) + + entries = env["PATH"].split(os.pathsep) + assert entries[0] == "/opt/nvidia/cublas/bin" + assert entries[-1] == "/usr/local/bin" + + +def test_with_cuda_library_path_without_existing(monkeypatch) -> None: + monkeypatch.setattr(os, "name", "posix") + monkeypatch.setattr( + "app.node_manager._cuda_library_dirs", + lambda repo_dir: ["/opt/nvidia/cublas/lib"], + ) + + env = _with_cuda_library_path({}, Path("/repo")) + + assert env["LD_LIBRARY_PATH"] == "/opt/nvidia/cublas/lib" + + +def test_with_cuda_library_path_deduplicates(monkeypatch) -> None: + monkeypatch.setattr(os, "name", "posix") + monkeypatch.setattr( + "app.node_manager._cuda_library_dirs", + lambda repo_dir: ["/opt/nvidia/cublas/lib", "/opt/nvidia/cudnn/lib"], + ) + + env = _with_cuda_library_path( + {"LD_LIBRARY_PATH": "/opt/nvidia/cublas/lib:/usr/lib/foo"}, + Path("/repo"), + ) + + entries = env["LD_LIBRARY_PATH"].split(os.pathsep) + assert entries == ["/opt/nvidia/cudnn/lib", "/opt/nvidia/cublas/lib", "/usr/lib/foo"] + + +def test_with_cuda_library_path_all_present(monkeypatch) -> None: + monkeypatch.setattr(os, "name", "posix") + monkeypatch.setattr( + "app.node_manager._cuda_library_dirs", + lambda repo_dir: ["/opt/nvidia/cublas/lib"], + ) + original = {"LD_LIBRARY_PATH": "/opt/nvidia/cublas/lib:/usr/lib/foo"} + + env = _with_cuda_library_path(dict(original), Path("/repo")) + + assert env["LD_LIBRARY_PATH"] == original["LD_LIBRARY_PATH"] + + +def test_with_cuda_library_path_without_libs(monkeypatch) -> None: + monkeypatch.setattr("app.node_manager._cuda_library_dirs", lambda repo_dir: []) + + env = _with_cuda_library_path({"LD_LIBRARY_PATH": "/usr/lib/foo"}, Path("/repo")) + + assert env == {"LD_LIBRARY_PATH": "/usr/lib/foo"} + + +def test_node_env_injects_cuda_library_path(db: Database, tmp_path, monkeypatch) -> None: + monkeypatch.setattr("app.node_manager.WORKSPACE_ROOT", tmp_path) + repo = tmp_path / "wov-node-whisper" + cublas = repo / ".venv" / "lib" / "python3.12" / "site-packages" / "nvidia" / "cublas" / "lib" + cublas.mkdir(parents=True) + manifest = NodeManifest( + id="x", + name="x", + version="1", + capability="x", + command=["python", "-m", "x"], + repo_dir="wov-node-whisper", + env={"WOV_NODE_PORT": "0"}, + ) + + env = NodeManager(db)._node_env(manifest) + + assert str(cublas) in env["LD_LIBRARY_PATH"].split(os.pathsep) + assert env["WOV_NODE_PORT"] == "0" + assert "wov-sdk" in env["PYTHONPATH"] + + def test_acquire_unregistered(db: Database) -> None: manager = NodeManager(db) with pytest.raises(ValueError): @@ -266,6 +398,51 @@ def test_invoke_network_error(db: Database, tmp_path) -> None: manager.shutdown() +def test_invoke_uses_configured_timeout(db: Database, monkeypatch) -> None: + _register(db, echo_manifest()) + manager = NodeManager(db) + captured = {} + + class FakeRuntime: + instance_id = "ni_x" + status = "ready" + busy_count = 0 + + class FakeResponse: + def read(self) -> bytes: + return b'{"status":"completed","outputs":{"text":"ok"}}' + + def __enter__(self): + return self + + def __exit__(self, *args) -> bool: + return False + + def fake_acquire(node_id): + return FakeRuntime(), "http://127.0.0.1:1" + + def fake_urlopen(request, timeout): + captured["timeout"] = timeout + return FakeResponse() + + monkeypatch.setattr("app.node_manager.NODE_INVOKE_TIMEOUT_SECONDS", 123) + monkeypatch.setattr(manager, "acquire", fake_acquire) + monkeypatch.setattr("app.node_manager.urllib.request.urlopen", fake_urlopen) + + response = manager.invoke( + "echo", + InvokeRequest( + run_id="run_t", + node_instance_id="", + inputs={}, + output_dir=".", + ), + ) + + assert response.status == "completed" + assert captured["timeout"] == 123 + + def test_release_unknown(db: Database) -> None: manager = NodeManager(db) manager.release("missing")