413 lines
17 KiB
Python
413 lines
17 KiB
Python
"""节点生命周期管理器。
|
||
|
||
NodeManager 是节点实例生命周期的唯一所有者：负责按 Manifest 启动节点进程、
|
||
等待就绪、复用空闲实例、空闲回收和调用转发。API 与调度器不得绕过本模块
|
||
直接启动或杀死节点进程。
|
||
"""
|
||
|
||
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:
|
||
"""返回当前 UTC 时间的 ISO 格式字符串，用于统一时间戳存储。"""
|
||
return datetime.now(timezone.utc).isoformat()
|
||
|
||
|
||
def _idle_seconds(last_used_at: str) -> float:
|
||
"""计算实例距上次使用的空闲秒数；非法时间按 0 处理。"""
|
||
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:
|
||
# 时间格式损坏时保守返回 0，避免误回收实例。
|
||
return 0.0
|
||
|
||
|
||
def _cuda_library_dirs(repo_dir: Path) -> list[Path]:
|
||
"""在节点虚拟环境中查找 nvidia cublas/cudnn 动态库目录。
|
||
|
||
faster-whisper 通过 pip 安装的 CUDA 库位于
|
||
.venv/.../site-packages/nvidia/<包>/lib（Linux）或 bin（Windows），
|
||
需要把这些目录加入动态库搜索路径才能被加载。
|
||
"""
|
||
venv = repo_dir / ".venv"
|
||
site_packages: list[Path] = []
|
||
# Windows 的 site-packages 位于 Lib 下。
|
||
windows_site = venv / "Lib" / "site-packages"
|
||
if windows_site.is_dir():
|
||
site_packages.append(windows_site)
|
||
# Linux 使用 lib/python3.x/site-packages 结构。
|
||
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]:
|
||
"""把 CUDA 动态库目录注入进程环境变量并去重。"""
|
||
dirs = [str(path) for path in _cuda_library_dirs(repo_dir)]
|
||
if not dirs:
|
||
return env
|
||
# Windows 动态库搜索走 PATH，Linux 走 LD_LIBRARY_PATH。
|
||
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:
|
||
# 新目录放在最前，优先使用虚拟环境内的 CUDA 库。
|
||
env[var] = os.pathsep.join(additions + existing)
|
||
return env
|
||
|
||
|
||
@dataclass
|
||
class NodeRuntime:
|
||
"""运行中的节点实例内存态：进程句柄、状态、地址与使用计数。"""
|
||
|
||
# 实例唯一 ID，格式 ni_xxxx。
|
||
instance_id: str
|
||
# 所属节点注册 ID。
|
||
node_id: str
|
||
# 启动该实例使用的 manifest 快照。
|
||
manifest: NodeManifest
|
||
# 生命周期状态：starting / ready / stopping / stopped / error。
|
||
status: str = "starting"
|
||
# 节点子进程句柄，用于终止与回收。
|
||
process: subprocess.Popen | None = None
|
||
pid: int | None = None
|
||
# 节点 HTTP 服务地址，例如 http://127.0.0.1:xxxx。
|
||
address: str | None = None
|
||
started_at: str = field(default_factory=_now_iso)
|
||
last_used_at: str = field(default_factory=_now_iso)
|
||
# 当前并发占用数，超过 manifest.max_concurrency 不再复用。
|
||
busy_count: int = 0
|
||
error: str | None = None
|
||
# 最近 200 行 stderr，便于排查启动失败。
|
||
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:
|
||
"""后台循环：每隔固定间隔检查并回收满足 TTL 的空闲实例。"""
|
||
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]:
|
||
"""解析启动命令；python 命令替换为节点仓库虚拟环境解释器。"""
|
||
command = list(manifest.command)
|
||
if command and command[0].lower() in {"python", "python3"}:
|
||
repo_dir = WORKSPACE_ROOT / manifest.repo_dir
|
||
# 同时兼容 Windows 与 Linux 的虚拟环境路径。
|
||
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:
|
||
# 虚拟环境不存在时使用 API 自身的 Python，便于演示。
|
||
command[0] = sys.executable
|
||
return command
|
||
|
||
def _node_env(self, manifest: NodeManifest) -> dict[str, str]:
|
||
"""构造节点进程环境：注入 SDK 路径、manifest 环境变量与 CUDA 路径。"""
|
||
env = os.environ.copy()
|
||
current_pythonpath = env.get("PYTHONPATH", "")
|
||
# 把 wov-sdk 源码目录放到 PYTHONPATH 最前。
|
||
env["PYTHONPATH"] = os.pathsep.join(
|
||
item for item in [str(SDK_SRC), current_pythonpath] if item
|
||
)
|
||
env.update(manifest.env)
|
||
# 为 faster-whisper 等节点补充 CUDA 动态库搜索路径。
|
||
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
|
||
|
||
# 子进程 stdout 读取线程解析到端口后设置事件。
|
||
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
|
||
# 命令不存在等启动错误统一记录为 error。
|
||
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:
|
||
"""逐行读取子进程 stdout，解析 WOV_NODE_READY 端口。"""
|
||
assert process.stdout is not None
|
||
for line in iter(process.stdout.readline, ""):
|
||
line = line.strip()
|
||
if line.startswith("WOV_NODE_READY"):
|
||
try:
|
||
# 从 "port=<数字>" 中提取端口。
|
||
ready_port["port"] = int(line.split("port=", 1)[1])
|
||
except (IndexError, ValueError):
|
||
# 格式异常时按 0 处理，随后会判定为未就绪。
|
||
ready_port["port"] = 0
|
||
ready_event.set()
|
||
|
||
def read_stderr() -> None:
|
||
"""缓存子进程 stderr 尾部，供失败诊断使用。"""
|
||
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
|
||
|
||
# 端口有效后做一次真实 HTTP 健康检查。
|
||
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:
|
||
"""先优雅 terminate，超时后强制 kill。"""
|
||
# 进程已退出时无需处理。
|
||
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:
|
||
"""释放一次实例占用；由 invoke 的 finally 保证执行。"""
|
||
with self._lock:
|
||
runtime = self._runtimes.get(instance_id)
|
||
if runtime is None:
|
||
return
|
||
# 计数下限为 0，防止重复释放导致负值。
|
||
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)
|
||
# 回填实际实例 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:
|
||
# 超时来自配置，避免 LLM 等慢节点被过早中断。
|
||
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:
|
||
# 节点返回 500 时 body 仍是协议 JSON，尝试解析失败原因。
|
||
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
|
||
# 网络错误、超时等统一转换为 failed。
|
||
return InvokeResponse(status="failed", error=str(exc))
|
||
finally:
|
||
self.release(request.node_instance_id)
|
||
|
||
def stop_instance(self, instance_id: str) -> None:
|
||
"""按实例 ID 停止指定节点，供管理后台使用。"""
|
||
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)
|