336 lines
12 KiB
Python
336 lines
12 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
import threading
|
|
import time
|
|
import urllib.error
|
|
import urllib.request
|
|
import uuid
|
|
from dataclasses import dataclass, field
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
|
|
from wov_sdk.models import InvokeRequest, InvokeResponse, NodeManifest
|
|
|
|
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
|
|
|
|
|
|
def _now_iso() -> str:
|
|
return datetime.now(timezone.utc).isoformat()
|
|
|
|
|
|
def _idle_seconds(last_used_at: str) -> float:
|
|
if not last_used_at:
|
|
return 0.0
|
|
try:
|
|
last = datetime.fromisoformat(last_used_at)
|
|
return max(0.0, (datetime.now(timezone.utc) - last).total_seconds())
|
|
except ValueError:
|
|
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
|
|
node_id: str
|
|
manifest: NodeManifest
|
|
status: str = "starting"
|
|
process: subprocess.Popen | None = None
|
|
pid: int | None = None
|
|
address: str | None = None
|
|
started_at: str = field(default_factory=_now_iso)
|
|
last_used_at: str = field(default_factory=_now_iso)
|
|
busy_count: int = 0
|
|
error: str | None = None
|
|
stderr_tail: list[str] = field(default_factory=list)
|
|
|
|
|
|
class NodeManager:
|
|
def __init__(self, db: Database) -> None:
|
|
self.db = db
|
|
self._runtimes: dict[str, NodeRuntime] = {}
|
|
self._lock = threading.RLock()
|
|
self._reaper_thread: threading.Thread | None = None
|
|
self._stopping = False
|
|
|
|
def start_reaper(self) -> None:
|
|
self._stopping = False
|
|
self._reaper_thread = threading.Thread(
|
|
target=self._reaper_loop,
|
|
name="wov-node-reaper",
|
|
daemon=True,
|
|
)
|
|
self._reaper_thread.start()
|
|
|
|
def shutdown(self) -> None:
|
|
self._stopping = True
|
|
with self._lock:
|
|
for runtime in list(self._runtimes.values()):
|
|
self._stop_locked(runtime)
|
|
|
|
def _reaper_loop(self) -> None:
|
|
while not self._stopping:
|
|
time.sleep(NODE_REAP_INTERVAL_SECONDS)
|
|
with self._lock:
|
|
for runtime in list(self._runtimes.values()):
|
|
if runtime.status != "ready" or runtime.busy_count > 0:
|
|
continue
|
|
if runtime.manifest.keep_warm:
|
|
continue
|
|
ttl = runtime.manifest.idle_ttl_seconds
|
|
if ttl >= 0 and _idle_seconds(runtime.last_used_at) >= ttl:
|
|
runtime.status = "stopping"
|
|
self._stop_locked(runtime)
|
|
|
|
def _resolve_command(self, manifest: NodeManifest) -> list[str]:
|
|
command = list(manifest.command)
|
|
if command and command[0].lower() in {"python", "python3"}:
|
|
repo_dir = WORKSPACE_ROOT / manifest.repo_dir
|
|
windows_python = repo_dir / ".venv" / "Scripts" / "python.exe"
|
|
unix_python = repo_dir / ".venv" / "bin" / "python"
|
|
if windows_python.is_file():
|
|
command[0] = str(windows_python)
|
|
elif unix_python.is_file():
|
|
command[0] = str(unix_python)
|
|
else:
|
|
command[0] = sys.executable
|
|
return command
|
|
|
|
def _node_env(self, manifest: NodeManifest) -> dict[str, str]:
|
|
env = os.environ.copy()
|
|
current_pythonpath = env.get("PYTHONPATH", "")
|
|
env["PYTHONPATH"] = os.pathsep.join(
|
|
item for item in [str(SDK_SRC), current_pythonpath] if item
|
|
)
|
|
env.update(manifest.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(
|
|
instance_id=f"ni_{uuid.uuid4().hex[:12]}",
|
|
node_id=manifest.id,
|
|
manifest=manifest,
|
|
)
|
|
repo_dir = WORKSPACE_ROOT / manifest.repo_dir
|
|
if not repo_dir.is_dir():
|
|
runtime.status = "error"
|
|
runtime.error = f"repo_dir not found: {repo_dir}"
|
|
self.db.upsert_instance(self._instance_row(runtime))
|
|
return runtime
|
|
|
|
ready_event = threading.Event()
|
|
ready_port: dict[str, int] = {}
|
|
|
|
try:
|
|
process = subprocess.Popen(
|
|
self._resolve_command(manifest),
|
|
cwd=str(repo_dir),
|
|
env=self._node_env(manifest),
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
text=True,
|
|
encoding="utf-8",
|
|
bufsize=1,
|
|
)
|
|
except Exception as exc: # noqa: BLE001
|
|
runtime.status = "error"
|
|
runtime.error = str(exc)
|
|
self.db.upsert_instance(self._instance_row(runtime))
|
|
return runtime
|
|
|
|
runtime.process = process
|
|
runtime.pid = process.pid
|
|
|
|
def read_stdout() -> None:
|
|
assert process.stdout is not None
|
|
for line in iter(process.stdout.readline, ""):
|
|
line = line.strip()
|
|
if line.startswith("WOV_NODE_READY"):
|
|
try:
|
|
ready_port["port"] = int(line.split("port=", 1)[1])
|
|
except (IndexError, ValueError):
|
|
ready_port["port"] = 0
|
|
ready_event.set()
|
|
|
|
def read_stderr() -> None:
|
|
assert process.stderr is not None
|
|
for line in iter(process.stderr.readline, ""):
|
|
line = line.strip()
|
|
if line:
|
|
runtime.stderr_tail.append(line)
|
|
runtime.stderr_tail = runtime.stderr_tail[-200:]
|
|
|
|
threading.Thread(target=read_stdout, daemon=True).start()
|
|
threading.Thread(target=read_stderr, daemon=True).start()
|
|
|
|
ready = ready_event.wait(timeout=NODE_READY_TIMEOUT_SECONDS)
|
|
if not ready or "port" not in ready_port or not ready_port["port"]:
|
|
self._stop_process(process)
|
|
runtime.status = "error"
|
|
runtime.error = "node did not report readiness"
|
|
self.db.upsert_instance(self._instance_row(runtime))
|
|
return runtime
|
|
|
|
address = f"http://127.0.0.1:{ready_port['port']}"
|
|
try:
|
|
with urllib.request.urlopen(f"{address}/health", timeout=2) as response:
|
|
if response.status != 200:
|
|
raise RuntimeError(f"health check returned {response.status}")
|
|
except Exception as exc: # noqa: BLE001
|
|
self._stop_process(process)
|
|
runtime.status = "error"
|
|
runtime.error = f"health check failed: {exc}"
|
|
self.db.upsert_instance(self._instance_row(runtime))
|
|
return runtime
|
|
|
|
runtime.address = address
|
|
runtime.status = "ready"
|
|
runtime.started_at = _now_iso()
|
|
runtime.last_used_at = _now_iso()
|
|
self._runtimes[runtime.instance_id] = runtime
|
|
self.db.upsert_instance(self._instance_row(runtime))
|
|
return runtime
|
|
|
|
def _stop_process(self, process: subprocess.Popen) -> None:
|
|
if process.poll() is not None:
|
|
return
|
|
process.terminate()
|
|
try:
|
|
process.wait(timeout=5)
|
|
except subprocess.TimeoutExpired:
|
|
process.kill()
|
|
process.wait(timeout=5)
|
|
|
|
def _stop_locked(self, runtime: NodeRuntime) -> None:
|
|
if runtime.process is not None:
|
|
self._stop_process(runtime.process)
|
|
runtime.status = "stopped"
|
|
runtime.pid = None
|
|
runtime.address = None
|
|
self._runtimes.pop(runtime.instance_id, None)
|
|
self.db.upsert_instance(self._instance_row(runtime))
|
|
|
|
def _instance_row(self, runtime: NodeRuntime) -> dict:
|
|
return {
|
|
"id": runtime.instance_id,
|
|
"node_id": runtime.node_id,
|
|
"status": runtime.status,
|
|
"pid": runtime.pid,
|
|
"address": runtime.address,
|
|
"started_at": runtime.started_at,
|
|
"last_used_at": runtime.last_used_at,
|
|
"busy_since": None,
|
|
"error": runtime.error,
|
|
}
|
|
|
|
def acquire(self, node_id: str) -> tuple[NodeRuntime, str]:
|
|
manifest = self.db.get_node(node_id)
|
|
if manifest is None:
|
|
raise ValueError(f"node not registered: {node_id}")
|
|
|
|
with self._lock:
|
|
for runtime in self._runtimes.values():
|
|
if (
|
|
runtime.node_id == node_id
|
|
and runtime.status == "ready"
|
|
and runtime.busy_count < manifest.max_concurrency
|
|
):
|
|
runtime.busy_count += 1
|
|
runtime.last_used_at = _now_iso()
|
|
return runtime, runtime.address or ""
|
|
|
|
runtime = self._start_locked(manifest)
|
|
if runtime.status != "ready":
|
|
raise RuntimeError(runtime.error or "node failed to start")
|
|
runtime.busy_count += 1
|
|
runtime.last_used_at = _now_iso()
|
|
return runtime, runtime.address or ""
|
|
|
|
def release(self, instance_id: str) -> None:
|
|
with self._lock:
|
|
runtime = self._runtimes.get(instance_id)
|
|
if runtime is None:
|
|
return
|
|
runtime.busy_count = max(0, runtime.busy_count - 1)
|
|
runtime.last_used_at = _now_iso()
|
|
self.db.upsert_instance(self._instance_row(runtime))
|
|
|
|
def invoke(self, node_id: str, request: InvokeRequest) -> InvokeResponse:
|
|
runtime, address = self.acquire(node_id)
|
|
request.node_instance_id = runtime.instance_id
|
|
body = json.dumps(request.to_dict(), ensure_ascii=False).encode("utf-8")
|
|
http_request = urllib.request.Request(
|
|
f"{address}/invoke",
|
|
data=body,
|
|
headers={"Content-Type": "application/json; charset=utf-8"},
|
|
method="POST",
|
|
)
|
|
try:
|
|
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:
|
|
try:
|
|
payload = json.loads(exc.read().decode("utf-8"))
|
|
return InvokeResponse.from_dict(payload)
|
|
except Exception: # noqa: BLE001
|
|
return InvokeResponse(status="failed", error=f"node returned {exc.code}")
|
|
except Exception as exc: # noqa: BLE001
|
|
return InvokeResponse(status="failed", error=str(exc))
|
|
finally:
|
|
self.release(request.node_instance_id)
|
|
|
|
def stop_instance(self, instance_id: str) -> None:
|
|
with self._lock:
|
|
runtime = self._runtimes.get(instance_id)
|
|
if runtime is None:
|
|
return
|
|
self._stop_locked(runtime)
|
|
|
|
def stop_all_for_node(self, node_id: str) -> None:
|
|
with self._lock:
|
|
for runtime in list(self._runtimes.values()):
|
|
if runtime.node_id == node_id:
|
|
self._stop_locked(runtime)
|