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

This commit is contained in:
cat-shark
2026-08-13 22:01:42 +08:00
parent a8d11eafc6
commit 77e9ac6b7e
8 changed files with 345 additions and 4 deletions
+40 -3
View File
@@ -15,7 +15,13 @@ from pathlib import Path
from wov_sdk.models import InvokeRequest, InvokeResponse, NodeManifest
from app.config import NODE_READY_TIMEOUT_SECONDS, NODE_REAP_INTERVAL_SECONDS, SDK_SRC, WORKSPACE_ROOT
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
@@ -33,6 +39,36 @@ def _idle_seconds(last_used_at: str) -> float:
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
@@ -107,7 +143,8 @@ class NodeManager:
item for item in [str(SDK_SRC), current_pythonpath] if item
)
env.update(manifest.env)
return 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(
@@ -270,7 +307,7 @@ class NodeManager:
method="POST",
)
try:
with urllib.request.urlopen(http_request, timeout=300) as response:
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: