feat: 增加失败任务重试与 CUDA 动态库路径注入
This commit is contained in:
+178
-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")
|
||||
|
||||
Reference in New Issue
Block a user