feat: 增加失败任务重试与 CUDA 动态库路径注入
This commit is contained in:
@@ -16,3 +16,4 @@ STORAGE_DIR = Path(os.getenv("WOV_STORAGE_DIR", str(DATA_DIR / "storage")))
|
||||
|
||||
NODE_REAP_INTERVAL_SECONDS = float(os.getenv("WOV_REAP_INTERVAL_SECONDS", "3"))
|
||||
NODE_READY_TIMEOUT_SECONDS = float(os.getenv("WOV_READY_TIMEOUT_SECONDS", "12"))
|
||||
NODE_INVOKE_TIMEOUT_SECONDS = float(os.getenv("WOV_NODE_INVOKE_TIMEOUT_SECONDS", "3600"))
|
||||
|
||||
@@ -330,6 +330,19 @@ class Database:
|
||||
with self._connect() as conn:
|
||||
conn.execute(f"UPDATE workflow_runs SET {assignments} WHERE id = ?", values)
|
||||
|
||||
def reset_run(self, run_id: str, updated_at: str) -> None:
|
||||
with self._connect() as conn:
|
||||
conn.execute(
|
||||
"""
|
||||
UPDATE workflow_runs
|
||||
SET status = 'QUEUED', current_node_id = NULL, progress = 0,
|
||||
error = NULL, updated_at = ?
|
||||
WHERE id = ?
|
||||
""",
|
||||
(updated_at, run_id),
|
||||
)
|
||||
conn.execute("DELETE FROM artifacts WHERE run_id = ?", (run_id,))
|
||||
|
||||
def next_queued_run(self) -> dict[str, Any] | None:
|
||||
with self._connect() as conn:
|
||||
row = conn.execute(
|
||||
|
||||
+40
-3
@@ -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:
|
||||
|
||||
@@ -100,6 +100,17 @@ def get_run(run_id: str, db: Database = Depends(_get_db)) -> dict:
|
||||
return run
|
||||
|
||||
|
||||
@router.post("/api/runs/{run_id}/retry")
|
||||
def retry_run(run_id: str, db: Database = Depends(_get_db)) -> dict:
|
||||
run = db.get_run(run_id)
|
||||
if run is None:
|
||||
raise HTTPException(status_code=404, detail="run not found")
|
||||
if run["status"] != "FAILED":
|
||||
raise HTTPException(status_code=422, detail="only failed runs can be retried")
|
||||
db.reset_run(run_id, _now_iso())
|
||||
return {"id": run_id, "status": "QUEUED"}
|
||||
|
||||
|
||||
@router.get("/api/runs/{run_id}/artifacts")
|
||||
def list_run_artifacts(run_id: str, db: Database = Depends(_get_db)) -> list[dict]:
|
||||
if db.get_run(run_id) is None:
|
||||
|
||||
Reference in New Issue
Block a user