5f46939321
The gpu lane is now a pool of nodes -- this Mac plus any remote workers in nodes.json (gitignored, primary-only). A gpu job acquires the first free node that supports its operator; local runs as before, a remote node has its inputs rsynced out, runs the operator's run.py over ssh (keys sourced from the node's .env.remote, kept off the process table), then its outputs are rsynced back and registered locally exactly as for a local job. One job per node (one Metal device each). - server/remote.py (new): node loading, cached health checks, prepare/cmd/collect - runner: _acquire_gpu_node + a remote branch in _run_job. lane_of now treats 'gpu' as a real lane -- it was dropped from LANE_LIMITS (which holds only the semaphore lanes cpu/net), so gpu jobs were silently falling back to the cpu lane and the pool was never reached. A node is reserved before the health await so two concurrent jobs can't grab the same one. - sysinfo: gpu lane capacity = pool size, plus per-node status for the dashboard Verified: all 7 gpu operators route to the pool; remote_cmd/node_supports and nodes.json parsing correct; acquisition is race-free; a scratch-DB server boots clean and logs the pool. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
328 lines
15 KiB
Python
328 lines
15 KiB
Python
"""Job runner: executes operators as subprocesses across concurrency lanes.
|
|
|
|
Operator contract:
|
|
<python> run.py --input <path> [--input <path> ...] --outdir <dir> --params '<json>'
|
|
- inputs are passed in asset order (single-input ops get one --input)
|
|
- stdout/stderr are streamed into the job log (secret values redacted)
|
|
- operator writes outputs into outdir and (optionally) outdir/result.json:
|
|
{"outputs": [{"path": "rel/or/abs", "name": "...", "meta": {...}}, ...],
|
|
"summary": {...}}
|
|
- if result.json is absent, every top-level file in outdir is registered
|
|
Concurrency lanes (from manifest "resources"): gpu=1 (Metal contention),
|
|
cpu=3, net=6. Env vars from settings are injected; "requires_env" gates a job.
|
|
"""
|
|
import asyncio
|
|
import json
|
|
import os
|
|
import signal
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
from . import db, remote, settings as settings_mod, store
|
|
|
|
JOBS_DIR = db.DATA / "jobs"
|
|
# gpu is a NODE POOL (this Mac + remote workers), not a plain semaphore; cpu/net
|
|
# are semaphores. Each Mac has ONE Metal device, so each gpu node runs one job.
|
|
LANE_LIMITS = {"cpu": 3, "net": 6}
|
|
|
|
|
|
class Runner:
|
|
def __init__(self):
|
|
self.queue: asyncio.Queue[str] = asyncio.Queue()
|
|
self.subscribers: dict = {} # websocket -> user dict (for per-user scoping)
|
|
self.operators: dict[str, dict] = {}
|
|
self.lanes: dict[str, asyncio.Semaphore] = {}
|
|
self.gpu_nodes: list[dict] = [] # [{local}, {ssh,root,operators}, ...]
|
|
self.procs: dict = {} # job_id -> subprocess
|
|
self.cancelled: set[str] = set()
|
|
self._tasks: set = set()
|
|
self.get_settings = lambda: {} # set by main.py
|
|
|
|
# -- pubsub -------------------------------------------------------------
|
|
async def broadcast(self, message: dict):
|
|
# "job" events carry a payload injected straight into the client's list, so
|
|
# scope them: a guest only receives their own jobs (owner receives all).
|
|
# "assets_changed" just triggers a per-user re-fetch, so send to everyone.
|
|
job_uid = message.get("job", {}).get("user_id") if message.get("type") == "job" else None
|
|
dead = []
|
|
for ws, u in list(self.subscribers.items()):
|
|
if job_uid is not None and u.get("role") != "owner" and job_uid != u.get("id"):
|
|
continue
|
|
try:
|
|
await ws.send_json(message)
|
|
except Exception:
|
|
dead.append(ws)
|
|
for ws in dead:
|
|
self.subscribers.pop(ws, None)
|
|
|
|
def lane_of(self, operator: str) -> str:
|
|
op = self.operators.get(operator, {})
|
|
lane = op.get("resources", "cpu")
|
|
# gpu is a valid lane (the node pool) even though it isn't in LANE_LIMITS,
|
|
# which holds only the semaphore lanes (cpu/net).
|
|
return lane if lane in LANE_LIMITS or lane == "gpu" else "cpu"
|
|
|
|
# -- job lifecycle --------------------------------------------------------
|
|
def create_job(self, con, operator: str, asset_ids: list[str], params: dict,
|
|
user_id: str | None = None) -> dict:
|
|
job_id = db.new_id()
|
|
outdir = JOBS_DIR / job_id
|
|
outdir.mkdir(parents=True, exist_ok=True)
|
|
primary = asset_ids[0] if asset_ids else None
|
|
con.execute(
|
|
"INSERT INTO jobs (id, operator, status, asset_id, asset_ids, params, outdir, created_at, user_id) "
|
|
"VALUES (?, ?, 'queued', ?, ?, ?, ?, ?, ?)",
|
|
(job_id, operator, primary, json.dumps(asset_ids),
|
|
json.dumps(params), str(outdir), db.now(), user_id),
|
|
)
|
|
con.commit()
|
|
self.queue.put_nowait(job_id)
|
|
return self.get_job(con, job_id)
|
|
|
|
def active_job_count(self, con, user_id: str) -> int:
|
|
row = con.execute(
|
|
"SELECT COUNT(*) c FROM jobs WHERE user_id = ? AND status IN ('queued','running')",
|
|
(user_id,)).fetchone()
|
|
return row["c"]
|
|
|
|
def get_job(self, con, job_id: str) -> dict | None:
|
|
row = con.execute("SELECT * FROM jobs WHERE id = ?", (job_id,)).fetchone()
|
|
return db.row_to_dict(row) if row else None
|
|
|
|
def list_jobs(self, con) -> list[dict]:
|
|
rows = con.execute("SELECT * FROM jobs ORDER BY created_at DESC LIMIT 200").fetchall()
|
|
return [db.row_to_dict(r) for r in rows]
|
|
|
|
def input_paths(self, con, job: dict) -> list[str]:
|
|
ids = job.get("asset_ids") or ([job["asset_id"]] if job.get("asset_id") else [])
|
|
paths = []
|
|
for aid in ids:
|
|
asset = store.get_asset(con, aid)
|
|
if asset:
|
|
paths.append(asset["path"])
|
|
return paths
|
|
|
|
async def cancel(self, con, job_id: str) -> bool:
|
|
job = self.get_job(con, job_id)
|
|
if not job or job["status"] in ("done", "error", "cancelled"):
|
|
return False
|
|
self.cancelled.add(job_id)
|
|
proc = self.procs.get(job_id)
|
|
if proc and proc.returncode is None:
|
|
try:
|
|
os.killpg(os.getpgid(proc.pid), signal.SIGTERM)
|
|
except (ProcessLookupError, PermissionError):
|
|
pass
|
|
else: # queued but not yet started
|
|
await self._update(con, job_id, status="cancelled", finished_at=db.now())
|
|
return True
|
|
|
|
def delete_job(self, con, job_id: str) -> bool:
|
|
job = self.get_job(con, job_id)
|
|
if not job or job_id in self.procs:
|
|
return False
|
|
import shutil
|
|
if job.get("outdir"):
|
|
shutil.rmtree(job["outdir"], ignore_errors=True)
|
|
con.execute("DELETE FROM jobs WHERE id = ?", (job_id,))
|
|
con.commit()
|
|
return True
|
|
|
|
async def _update(self, con, job_id: str, **fields):
|
|
sets = ", ".join(f"{k} = ?" for k in fields)
|
|
con.execute(f"UPDATE jobs SET {sets} WHERE id = ?", (*fields.values(), job_id))
|
|
con.commit()
|
|
job = self.get_job(con, job_id)
|
|
await self.broadcast({"type": "job", "job": job})
|
|
|
|
# -- worker -----------------------------------------------------------------
|
|
async def worker(self):
|
|
con = db.connect()
|
|
for key in LANE_LIMITS:
|
|
self.lanes[key] = asyncio.Semaphore(LANE_LIMITS[key])
|
|
self.gpu_nodes = remote.load_gpu_nodes()
|
|
names = ", ".join(n.get("name", n.get("ssh", "?")) for n in self.gpu_nodes)
|
|
print(f"[runner] gpu pool: {names}", flush=True)
|
|
# re-queue jobs left running/queued by a previous process
|
|
for row in con.execute("SELECT id FROM jobs WHERE status IN ('queued','running')"):
|
|
con.execute("UPDATE jobs SET status='queued' WHERE id = ?", (row["id"],))
|
|
self.queue.put_nowait(row["id"])
|
|
con.commit()
|
|
while True:
|
|
job_id = await self.queue.get()
|
|
task = asyncio.create_task(self._run_lane(con, job_id))
|
|
self._tasks.add(task)
|
|
task.add_done_callback(self._tasks.discard)
|
|
|
|
async def _run_lane(self, con, job_id: str):
|
|
job = self.get_job(con, job_id)
|
|
if not job: # deleted while queued
|
|
self.cancelled.discard(job_id)
|
|
return
|
|
if job_id in self.cancelled:
|
|
self.cancelled.discard(job_id)
|
|
await self._update(con, job_id, status="cancelled", finished_at=db.now())
|
|
return
|
|
lane = self.lane_of(job["operator"])
|
|
if lane == "gpu":
|
|
node = await self._acquire_gpu_node(job["operator"])
|
|
try:
|
|
if job_id in self.cancelled:
|
|
self.cancelled.discard(job_id)
|
|
await self._update(con, job_id, status="cancelled", finished_at=db.now())
|
|
return
|
|
await self._run_job(con, job_id, node=node)
|
|
except Exception as e:
|
|
await self._update(con, job_id, status="error", error=str(e), finished_at=db.now())
|
|
finally:
|
|
node["busy"] = False
|
|
return
|
|
async with self.lanes[lane]:
|
|
if job_id in self.cancelled:
|
|
self.cancelled.discard(job_id)
|
|
await self._update(con, job_id, status="cancelled", finished_at=db.now())
|
|
return
|
|
try:
|
|
await self._run_job(con, job_id)
|
|
except Exception as e:
|
|
await self._update(con, job_id, status="error", error=str(e), finished_at=db.now())
|
|
|
|
async def _acquire_gpu_node(self, op_id: str) -> dict:
|
|
"""Return the first free gpu node that supports this operator (local always
|
|
does; a remote must list it + be reachable). Polls — gpu jobs are coarse."""
|
|
while True:
|
|
for n in self.gpu_nodes:
|
|
if n["busy"] or not remote.node_supports(n, op_id):
|
|
continue
|
|
n["busy"] = True # reserve synchronously — the health check below
|
|
# awaits, and without reserving first two jobs could grab one node
|
|
if not n.get("local") and not await remote.healthy(n):
|
|
n["busy"] = False # unreachable right now; release and keep looking
|
|
continue
|
|
return n
|
|
await asyncio.sleep(0.25)
|
|
|
|
async def _run_job(self, con, job_id: str, node: dict | None = None):
|
|
job = self.get_job(con, job_id)
|
|
if not job:
|
|
return
|
|
op = self.operators.get(job["operator"])
|
|
if not op:
|
|
await self._update(con, job_id, status="error",
|
|
error=f"unknown operator {job['operator']}", finished_at=db.now())
|
|
return
|
|
|
|
cur_settings = self.get_settings()
|
|
env = os.environ.copy()
|
|
# least privilege: only this operator's declared paid keys (+ HF/infra) reach
|
|
# the subprocess, so guest local jobs never get the owner's fal/OpenRouter keys
|
|
env.update(settings_mod.env_for_operator(cur_settings, op.get("requires_env")))
|
|
secrets = settings_mod.secret_values(cur_settings)
|
|
|
|
missing = [e for e in op.get("requires_env", []) if not env.get(e)]
|
|
if missing:
|
|
keys = ", ".join(sorted({k for k, v in settings_mod.ENV_MAP.items() if v in missing}))
|
|
await self._update(con, job_id, status="error",
|
|
error=f"missing setting(s): {keys} — add them in Settings",
|
|
finished_at=db.now())
|
|
return
|
|
|
|
outdir = Path(job["outdir"])
|
|
params_json = json.dumps(job["params"])
|
|
local_inputs = self.input_paths(con, job)
|
|
entry = Path(op["dir"]) / op.get("entry", "run.py")
|
|
# manifest "python" may be repo-relative (portable) or absolute (legacy)
|
|
op_py_rel = op.get("python")
|
|
op_py_abs = (str(db.ROOT / op_py_rel)
|
|
if op_py_rel and not Path(op_py_rel).is_absolute() else op_py_rel)
|
|
|
|
node = node or {"local": True, "name": "local"}
|
|
if node.get("local"):
|
|
python = op_py_abs or sys.executable
|
|
if op_py_abs and not Path(python).exists():
|
|
await self._update(con, job_id, status="error", finished_at=db.now(),
|
|
error=f"operator not installed — {python} missing. "
|
|
f"Run its install script under scripts/.")
|
|
return
|
|
cmd = [python, str(entry), "--outdir", str(outdir), "--params", params_json]
|
|
for p in local_inputs:
|
|
cmd += ["--input", p]
|
|
run_env = env
|
|
else: # remote node: rsync inputs out, run over ssh, rsync outputs back
|
|
if not op_py_rel or Path(op_py_rel).is_absolute():
|
|
await self._update(con, job_id, status="error", finished_at=db.now(),
|
|
error="operator has no repo-relative python — can't run remotely")
|
|
return
|
|
try:
|
|
remote_inputs = await remote.prepare(node, job_id, local_inputs)
|
|
except Exception as e:
|
|
await self._update(con, job_id, status="error", finished_at=db.now(),
|
|
error=f"remote prepare failed: {e}")
|
|
return
|
|
entry_rel = os.path.relpath(str(entry), str(db.ROOT))
|
|
cmd = remote.remote_cmd(node, op["id"], op_py_rel, entry_rel,
|
|
job_id, params_json, remote_inputs)
|
|
run_env = os.environ.copy()
|
|
|
|
await self._update(con, job_id, status="running", started_at=db.now())
|
|
proc = await asyncio.create_subprocess_exec(
|
|
*cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.STDOUT,
|
|
start_new_session=True, env=run_env)
|
|
self.procs[job_id] = proc
|
|
log_lines: list[str] = [f"[node: {node.get('name', 'local')}]\n"]
|
|
assert proc.stdout
|
|
try:
|
|
async for raw in proc.stdout:
|
|
log_lines.append(raw.decode(errors="replace"))
|
|
if len(log_lines) % 5 == 0:
|
|
log = settings_mod.redact("".join(log_lines)[-100_000:], secrets)
|
|
await self._update(con, job_id, log=log)
|
|
code = await proc.wait()
|
|
finally:
|
|
self.procs.pop(job_id, None)
|
|
log = settings_mod.redact("".join(log_lines)[-100_000:], secrets)
|
|
|
|
if job_id in self.cancelled:
|
|
self.cancelled.discard(job_id)
|
|
await self._update(con, job_id, status="cancelled", log=log, finished_at=db.now())
|
|
return
|
|
if code != 0:
|
|
await self._update(con, job_id, status="error", log=log,
|
|
error=f"exit code {code}", finished_at=db.now())
|
|
return
|
|
|
|
if not node.get("local"): # pull the remote outputs into the local job dir
|
|
try:
|
|
await remote.collect(node, job_id, outdir)
|
|
except Exception as e:
|
|
await self._update(con, job_id, status="error", log=log,
|
|
error=f"remote collect failed: {e}", finished_at=db.now())
|
|
return
|
|
registered = self._register_outputs(con, outdir, job_id)
|
|
await self._update(con, job_id, status="done", log=log, finished_at=db.now())
|
|
await self.broadcast({"type": "assets_changed", "job_id": job_id, "assets": registered})
|
|
|
|
def _register_outputs(self, con, outdir: Path, job_id: str) -> list[str]:
|
|
result_path = outdir / "result.json"
|
|
registered = []
|
|
if result_path.exists():
|
|
result = json.loads(result_path.read_text())
|
|
for out in result.get("outputs", []):
|
|
p = Path(out["path"])
|
|
if not p.is_absolute():
|
|
p = outdir / p
|
|
if p.exists():
|
|
a = store.register_file(con, p, name=out.get("name"),
|
|
parent_job=job_id, move=True, meta=out.get("meta"))
|
|
registered.append(a["id"])
|
|
else:
|
|
for p in sorted(outdir.iterdir()):
|
|
if p.name == "result.json" or p.name.startswith("."):
|
|
continue
|
|
a = store.register_file(con, p, parent_job=job_id, move=True)
|
|
registered.append(a["id"])
|
|
return registered
|
|
|
|
|
|
runner = Runner()
|