"""Job runner: executes operators as subprocesses across concurrency lanes. Operator contract: run.py --input [--input ...] --outdir --params '' - 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, settings as settings_mod, store JOBS_DIR = db.DATA / "jobs" LANE_LIMITS = {"gpu": 1, "cpu": 3, "net": 6} class Runner: def __init__(self): self.queue: asyncio.Queue[str] = asyncio.Queue() self.subscribers: set = set() self.operators: dict[str, dict] = {} self.lanes: dict[str, asyncio.Semaphore] = {} 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): dead = [] for ws in self.subscribers: try: await ws.send_json(message) except Exception: dead.append(ws) for ws in dead: self.subscribers.discard(ws) def lane_of(self, operator: str) -> str: op = self.operators.get(operator, {}) lane = op.get("resources", "cpu") return lane if lane in LANE_LIMITS 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]) # 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"]) 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 _run_job(self, con, job_id: str): 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() env.update(settings_mod.env_from_settings(cur_settings)) 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"]) entry = Path(op["dir"]) / op.get("entry", "run.py") python = op.get("python") or sys.executable if op.get("python") 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", json.dumps(job["params"])] for p in self.input_paths(con, job): cmd += ["--input", p] 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=env) self.procs[job_id] = proc log_lines: list[str] = [] 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 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()