19cf9ed21e
scripts/agent_token.py mints the 'agents' guest user + token non-interactively, writes data/agent.env (0600, gitignored, never printed). Fixes _register_outputs dropping user_id — guest job outputs were registered ownerless, so guests could never see or download their own results. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
333 lines
15 KiB
Python
333 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 = []
|
|
# outputs belong to whoever ran the job — guests see only their own
|
|
# assets, so registering without user_id hides a guest's results from them
|
|
row = con.execute("SELECT user_id FROM jobs WHERE id = ?", (job_id,)).fetchone()
|
|
owner_id = row["user_id"] if row else None
|
|
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"),
|
|
user_id=owner_id)
|
|
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, user_id=owner_id)
|
|
registered.append(a["id"])
|
|
return registered
|
|
|
|
|
|
runner = Runner()
|