572 lines
17 KiB
Python
572 lines
17 KiB
Python
import os
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
from datetime import datetime, timedelta, timezone
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
from app.db import Database
|
|
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
|
|
_BASE_ECHO_MANIFEST = NodeManifest.load(WORKSPACE / "wov-node-echo" / "node.manifest.json")
|
|
FIXTURES = Path(__file__).resolve().parent / "fixtures"
|
|
|
|
|
|
@pytest.fixture
|
|
def db(tmp_path) -> Database:
|
|
return Database(tmp_path / "wov.db")
|
|
|
|
|
|
def _register(db: Database, manifest: NodeManifest) -> None:
|
|
db.upsert_node(manifest)
|
|
|
|
|
|
def echo_manifest() -> NodeManifest:
|
|
return NodeManifest.from_dict(_BASE_ECHO_MANIFEST.to_dict())
|
|
|
|
|
|
def test_time_helpers() -> None:
|
|
now = _now_iso()
|
|
assert datetime.fromisoformat(now)
|
|
assert _idle_seconds(now) < 0.1
|
|
old = (datetime.now(timezone.utc) - timedelta(seconds=10)).isoformat()
|
|
assert _idle_seconds(old) >= 9
|
|
assert _idle_seconds("") == 0
|
|
assert _idle_seconds("invalid") == 0
|
|
|
|
|
|
def test_resolve_command_and_env(db: Database) -> None:
|
|
manager = NodeManager(db)
|
|
python_command = NodeManifest(
|
|
id="x",
|
|
name="x",
|
|
version="1",
|
|
capability="x",
|
|
command=["python", "-m", "x"],
|
|
)
|
|
assert manager._resolve_command(python_command)[0] == sys.executable
|
|
|
|
fixed_command = NodeManifest(
|
|
id="x",
|
|
name="x",
|
|
version="1",
|
|
capability="x",
|
|
command=["node", "index.js"],
|
|
)
|
|
assert manager._resolve_command(fixed_command) == ["node", "index.js"]
|
|
|
|
env = manager._node_env(python_command)
|
|
assert "wov-sdk" in env["PYTHONPATH"]
|
|
|
|
|
|
def test_resolve_command_prefers_node_venv(db: Database, tmp_path, monkeypatch) -> None:
|
|
monkeypatch.setattr("app.node_manager.WORKSPACE_ROOT", tmp_path)
|
|
repo = tmp_path / "wov-node-demo"
|
|
python_exe = repo / ".venv" / "Scripts" / "python.exe"
|
|
python_exe.parent.mkdir(parents=True)
|
|
python_exe.write_bytes(b"")
|
|
manifest = NodeManifest(
|
|
id="x",
|
|
name="x",
|
|
version="1",
|
|
capability="x",
|
|
command=["python", "-m", "x"],
|
|
repo_dir="wov-node-demo",
|
|
)
|
|
manager = NodeManager(db)
|
|
assert manager._resolve_command(manifest)[0] == str(python_exe)
|
|
|
|
python_exe.unlink()
|
|
unix_python = repo / ".venv" / "bin" / "python"
|
|
unix_python.parent.mkdir(parents=True)
|
|
unix_python.write_bytes(b"")
|
|
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):
|
|
manager.acquire("missing")
|
|
|
|
|
|
def test_start_missing_repo(db: Database) -> None:
|
|
manifest = echo_manifest()
|
|
manifest.repo_dir = "missing-repo"
|
|
_register(db, manifest)
|
|
manager = NodeManager(db)
|
|
with pytest.raises(RuntimeError, match="repo_dir not found"):
|
|
manager.acquire("echo")
|
|
assert db.list_instances()[0]["status"] == "error"
|
|
|
|
|
|
def test_start_command_not_found(db: Database) -> None:
|
|
manifest = echo_manifest()
|
|
manifest.command = ["definitely-not-a-real-wov-command"]
|
|
_register(db, manifest)
|
|
manager = NodeManager(db)
|
|
with pytest.raises(RuntimeError):
|
|
manager.acquire("echo")
|
|
assert db.list_instances()[0]["status"] == "error"
|
|
|
|
|
|
def test_start_no_ready_timeout(db: Database, monkeypatch) -> None:
|
|
monkeypatch.setattr("app.node_manager.NODE_READY_TIMEOUT_SECONDS", 0.2)
|
|
manifest = echo_manifest()
|
|
manifest.command = ["python", "-c", "import time; time.sleep(0.5)"]
|
|
_register(db, manifest)
|
|
manager = NodeManager(db)
|
|
with pytest.raises(RuntimeError, match="did not report readiness"):
|
|
manager.acquire("echo")
|
|
assert db.list_instances()[0]["status"] == "error"
|
|
|
|
|
|
def test_start_bad_ready_line(db: Database) -> None:
|
|
manifest = echo_manifest()
|
|
manifest.command = ["python", "-u", str(FIXTURES / "bad_ready_node.py")]
|
|
_register(db, manifest)
|
|
manager = NodeManager(db)
|
|
with pytest.raises(RuntimeError, match="did not report readiness"):
|
|
manager.acquire("echo")
|
|
assert db.list_instances()[0]["status"] == "error"
|
|
|
|
|
|
def test_start_health_check_failure(db: Database) -> None:
|
|
manifest = echo_manifest()
|
|
manifest.command = ["python", "-u", str(FIXTURES / "bad_node.py")]
|
|
_register(db, manifest)
|
|
manager = NodeManager(db)
|
|
with pytest.raises(RuntimeError, match="health check failed"):
|
|
manager.acquire("echo")
|
|
assert db.list_instances()[0]["status"] == "error"
|
|
|
|
|
|
def test_start_health_check_bad_status(db: Database) -> None:
|
|
manifest = echo_manifest()
|
|
manifest.command = ["python", "-u", str(FIXTURES / "bad_status_node.py")]
|
|
_register(db, manifest)
|
|
manager = NodeManager(db)
|
|
with pytest.raises(RuntimeError, match="health check returned 201"):
|
|
manager.acquire("echo")
|
|
assert db.list_instances()[0]["status"] == "error"
|
|
|
|
|
|
def test_stderr_is_captured(db: Database) -> None:
|
|
manifest = echo_manifest()
|
|
manifest.command = ["python", "-u", str(FIXTURES / "stderr_node.py")]
|
|
_register(db, manifest)
|
|
manager = NodeManager(db)
|
|
try:
|
|
runtime, _ = manager.acquire("echo")
|
|
time.sleep(0.2)
|
|
runtime = next(iter(manager._runtimes.values()))
|
|
assert "stderr warning" in runtime.stderr_tail
|
|
manager.release(runtime.instance_id)
|
|
finally:
|
|
manager.shutdown()
|
|
|
|
|
|
def test_acquire_reuses_ready_instance(db: Database) -> None:
|
|
_register(db, echo_manifest())
|
|
manager = NodeManager(db)
|
|
try:
|
|
first, first_address = manager.acquire("echo")
|
|
manager.release(first.instance_id)
|
|
second, second_address = manager.acquire("echo")
|
|
assert first.instance_id == second.instance_id
|
|
assert first_address == second_address
|
|
manager.release(second.instance_id)
|
|
finally:
|
|
manager.shutdown()
|
|
|
|
|
|
def test_invoke_success(db: Database, tmp_path) -> None:
|
|
_register(db, echo_manifest())
|
|
manager = NodeManager(db)
|
|
try:
|
|
response = manager.invoke(
|
|
"echo",
|
|
InvokeRequest(
|
|
run_id="run_ok",
|
|
node_instance_id="",
|
|
inputs={"text": "real path"},
|
|
output_dir=str(tmp_path),
|
|
),
|
|
)
|
|
assert response.status == "completed"
|
|
assert response.outputs["text"] == "real path"
|
|
finally:
|
|
manager.shutdown()
|
|
|
|
|
|
def test_invoke_http_error(db: Database, tmp_path) -> None:
|
|
manifest = echo_manifest()
|
|
manifest.command = ["python", "-u", str(FIXTURES / "failing_node.py")]
|
|
_register(db, manifest)
|
|
manager = NodeManager(db)
|
|
try:
|
|
response = manager.invoke(
|
|
"echo",
|
|
InvokeRequest(
|
|
run_id="run_fail",
|
|
node_instance_id="",
|
|
inputs={},
|
|
output_dir=str(tmp_path),
|
|
),
|
|
)
|
|
assert response.status == "failed"
|
|
assert response.error == "boom"
|
|
finally:
|
|
manager.shutdown()
|
|
|
|
|
|
def test_invoke_http_error_invalid_json(db: Database, tmp_path) -> None:
|
|
manifest = echo_manifest()
|
|
manifest.command = ["python", "-u", str(FIXTURES / "failing_node.py")]
|
|
manifest.env["WOV_FAIL_INVALID_JSON"] = "1"
|
|
_register(db, manifest)
|
|
manager = NodeManager(db)
|
|
try:
|
|
response = manager.invoke(
|
|
"echo",
|
|
InvokeRequest(
|
|
run_id="run_bad_json",
|
|
node_instance_id="",
|
|
inputs={},
|
|
output_dir=str(tmp_path),
|
|
),
|
|
)
|
|
assert response.status == "failed"
|
|
assert "node returned 500" in response.error
|
|
finally:
|
|
manager.shutdown()
|
|
|
|
|
|
def test_invoke_network_error(db: Database, tmp_path) -> None:
|
|
_register(db, echo_manifest())
|
|
manager = NodeManager(db)
|
|
try:
|
|
runtime, _ = manager.acquire("echo")
|
|
manager.release(runtime.instance_id)
|
|
runtime.process.terminate()
|
|
runtime.process.wait(timeout=5)
|
|
response = manager.invoke(
|
|
"echo",
|
|
InvokeRequest(
|
|
run_id="run_net",
|
|
node_instance_id="",
|
|
inputs={},
|
|
output_dir=str(tmp_path),
|
|
),
|
|
)
|
|
assert response.status == "failed"
|
|
finally:
|
|
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")
|
|
|
|
|
|
def test_stop_unknown_instance(db: Database) -> None:
|
|
manager = NodeManager(db)
|
|
manager.stop_instance("missing")
|
|
|
|
|
|
def test_stop_all_for_node(db: Database) -> None:
|
|
_register(db, echo_manifest())
|
|
manager = NodeManager(db)
|
|
runtime, _ = manager.acquire("echo")
|
|
manager.stop_all_for_node("echo")
|
|
assert runtime.status == "stopped"
|
|
assert runtime.instance_id not in manager._runtimes
|
|
|
|
|
|
def test_stop_locked_without_process(db: Database) -> None:
|
|
_register(db, echo_manifest())
|
|
manager = NodeManager(db)
|
|
runtime = NodeRuntime(
|
|
instance_id="ni_static",
|
|
node_id="echo",
|
|
manifest=echo_manifest(),
|
|
status="ready",
|
|
)
|
|
manager._runtimes[runtime.instance_id] = runtime
|
|
manager.stop_instance(runtime.instance_id)
|
|
assert runtime.status == "stopped"
|
|
assert runtime.instance_id not in manager._runtimes
|
|
|
|
|
|
def test_stop_process_already_exited() -> None:
|
|
class FakeProcess:
|
|
def poll(self):
|
|
return 0
|
|
|
|
manager = object.__new__(NodeManager)
|
|
manager._stop_process(FakeProcess())
|
|
|
|
|
|
def test_stop_process_timeout() -> None:
|
|
class FakeProcess:
|
|
def poll(self):
|
|
return None
|
|
|
|
def terminate(self):
|
|
self.terminated = True
|
|
|
|
def wait(self, timeout):
|
|
if not self.killed:
|
|
raise subprocess.TimeoutExpired("fake", timeout)
|
|
|
|
def kill(self):
|
|
self.killed = True
|
|
|
|
fake = FakeProcess()
|
|
fake.killed = False
|
|
manager = object.__new__(NodeManager)
|
|
manager._stop_process(fake)
|
|
assert fake.terminated
|
|
assert fake.killed
|
|
|
|
|
|
def test_reaper_recycles_idle(db: Database, tmp_path) -> None:
|
|
manifest = echo_manifest()
|
|
manifest.idle_ttl_seconds = 0
|
|
_register(db, manifest)
|
|
manager = NodeManager(db)
|
|
manager.start_reaper()
|
|
try:
|
|
response = manager.invoke(
|
|
"echo",
|
|
InvokeRequest(
|
|
run_id="run_reap",
|
|
node_instance_id="",
|
|
inputs={"text": "reap"},
|
|
output_dir=str(tmp_path),
|
|
),
|
|
)
|
|
assert response.status == "completed"
|
|
time.sleep(0.4)
|
|
assert db.list_instances()[0]["status"] == "stopped"
|
|
finally:
|
|
manager.shutdown()
|
|
|
|
|
|
def test_reaper_keeps_warm(db: Database, tmp_path) -> None:
|
|
manifest = echo_manifest()
|
|
manifest.idle_ttl_seconds = 0
|
|
manifest.keep_warm = True
|
|
_register(db, manifest)
|
|
manager = NodeManager(db)
|
|
manager.start_reaper()
|
|
try:
|
|
response = manager.invoke(
|
|
"echo",
|
|
InvokeRequest(
|
|
run_id="run_warm",
|
|
node_instance_id="",
|
|
inputs={"text": "warm"},
|
|
output_dir=str(tmp_path),
|
|
),
|
|
)
|
|
assert response.status == "completed"
|
|
time.sleep(0.4)
|
|
assert db.list_instances()[0]["status"] == "ready"
|
|
finally:
|
|
manager.shutdown()
|
|
|
|
|
|
def test_reaper_keeps_busy(db: Database) -> None:
|
|
manifest = echo_manifest()
|
|
manifest.idle_ttl_seconds = 0
|
|
_register(db, manifest)
|
|
manager = NodeManager(db)
|
|
manager.start_reaper()
|
|
try:
|
|
runtime, _ = manager.acquire("echo")
|
|
time.sleep(0.4)
|
|
assert runtime.status == "ready"
|
|
manager.release(runtime.instance_id)
|
|
finally:
|
|
manager.shutdown()
|