144 lines
5.8 KiB
Python
144 lines
5.8 KiB
Python
"""Job runner: executes operators as subprocesses, one heavy job at a time.
|
|
|
|
Operator contract:
|
|
<python> run.py --input <asset path> --outdir <job outdir> --params '<json>'
|
|
- stdout/stderr are streamed into the job log
|
|
- operator writes outputs into outdir and (optionally) outdir/result.json:
|
|
{"outputs": [{"path": "relative/or/absolute", "name": "...", "meta": {...}}, ...],
|
|
"summary": {...}}
|
|
- if result.json is absent, every top-level file in outdir is registered as an asset
|
|
"""
|
|
import asyncio
|
|
import json
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
from . import db, store
|
|
|
|
JOBS_DIR = db.DATA / "jobs"
|
|
|
|
|
|
class Runner:
|
|
def __init__(self):
|
|
self.queue: asyncio.Queue[str] = asyncio.Queue()
|
|
self.subscribers: set = set() # websockets
|
|
self.operators: dict[str, dict] = {}
|
|
|
|
# -- 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)
|
|
|
|
# -- job lifecycle --------------------------------------------------------
|
|
def create_job(self, con, operator: str, asset_id: str | None, params: dict) -> dict:
|
|
job_id = db.new_id()
|
|
outdir = JOBS_DIR / job_id
|
|
outdir.mkdir(parents=True, exist_ok=True)
|
|
con.execute(
|
|
"INSERT INTO jobs (id, operator, status, asset_id, params, outdir, created_at) "
|
|
"VALUES (?, ?, 'queued', ?, ?, ?, ?)",
|
|
(job_id, operator, asset_id, json.dumps(params), str(outdir), db.now()),
|
|
)
|
|
con.commit()
|
|
self.queue.put_nowait(job_id)
|
|
return self.get_job(con, job_id)
|
|
|
|
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]
|
|
|
|
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()
|
|
# re-queue jobs that were 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()
|
|
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
|
|
asset = store.get_asset(con, job["asset_id"]) if job["asset_id"] else None
|
|
outdir = Path(job["outdir"])
|
|
entry = Path(op["dir"]) / op.get("entry", "run.py")
|
|
python = op.get("python") or sys.executable
|
|
cmd = [python, str(entry), "--outdir", str(outdir),
|
|
"--params", json.dumps(job["params"])]
|
|
if asset:
|
|
cmd += ["--input", asset["path"]]
|
|
|
|
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)
|
|
log_lines: list[str] = []
|
|
assert proc.stdout
|
|
async for raw in proc.stdout:
|
|
line = raw.decode(errors="replace")
|
|
log_lines.append(line)
|
|
if len(log_lines) % 5 == 0: # don't hammer the DB on chatty tools
|
|
await self._update(con, job_id, log="".join(log_lines)[-100_000:])
|
|
code = await proc.wait()
|
|
log = "".join(log_lines)[-100_000:]
|
|
|
|
if code != 0:
|
|
await self._update(con, job_id, status="error", log=log,
|
|
error=f"exit code {code}", finished_at=db.now())
|
|
return
|
|
|
|
# register outputs
|
|
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"])
|
|
|
|
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})
|
|
|
|
|
|
runner = Runner()
|