feat: 增加失败任务重试与 CUDA 动态库路径注入

2026-08-13 22:01:42 +08:00
parent a8d11eafc6
commit 77e9ac6b7e
8 changed files with 345 additions and 4 deletions
@@ -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
@@ -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") == []
@@ -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")