diff --git a/server/remote.py b/server/remote.py new file mode 100644 index 0000000..8ef5cb6 --- /dev/null +++ b/server/remote.py @@ -0,0 +1,113 @@ +"""Remote GPU worker dispatch (the unified worker pool). + +The gpu lane is a pool of NODES: the local Mac plus any remote workers listed in +nodes.json. A remote node runs the SAME operator code (a git-synced checkout) with +repo-relative paths. The primary never lets a remote touch its DB — it rsyncs the +inputs out, runs the operator's run.py over ssh, and rsyncs the outputs back, then +registers them locally exactly as for a local job. + +nodes.json (repo root, gitignored — present only on the primary), e.g.: + [{"name": "m1", "ssh": "johnking@100.91.239.7", "root": "/Users/johnking/MODELBEAST", + "operators": ["flux_local","mflux_image_edit","seedvr2_upscale","bg_remove_local","sf3d"]}] +Omit "operators" to allow every gpu op. The implicit local node is always present. +""" +import asyncio +import json +import shlex +import time +from pathlib import Path + +from . import db + +SSH = ["ssh", "-o", "BatchMode=yes", "-o", "ConnectTimeout=8"] +_health: dict[str, tuple[float, bool]] = {} # ssh -> (checked_at, ok) +_HEALTH_TTL = 30.0 + + +def load_gpu_nodes() -> list[dict]: + """Local node first, then remote workers from nodes.json.""" + nodes = [{"name": "local", "local": True, "busy": False}] + f = db.ROOT / "nodes.json" + if f.exists(): + try: + for n in json.loads(f.read_text()): + n["busy"] = False + n["local"] = False + nodes.append(n) + except (ValueError, KeyError): + pass + return nodes + + +def node_supports(node: dict, op_id: str) -> bool: + if node.get("local"): + return True + allow = node.get("operators") + return (not allow) or (op_id in allow) + + +async def _run(cmd: list[str], timeout: float = 300) -> tuple[int, str]: + proc = await asyncio.create_subprocess_exec( + *cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.STDOUT) + try: + out, _ = await asyncio.wait_for(proc.communicate(), timeout) + except asyncio.TimeoutError: + proc.kill() + return 124, "timeout" + return proc.returncode, (out or b"").decode(errors="replace") + + +async def healthy(node: dict) -> bool: + """Reachable + checkout present. Cached for _HEALTH_TTL to avoid per-job ssh.""" + if node.get("local"): + return True + key = node["ssh"] + hit = _health.get(key) + if hit and time.time() - hit[0] < _HEALTH_TTL: + return hit[1] + code, _ = await _run(SSH + [node["ssh"], f"test -d {shlex.quote(node['root'])}"], timeout=10) + ok = code == 0 + _health[key] = (time.time(), ok) + return ok + + +async def prepare(node: dict, job_id: str, local_inputs: list[str]) -> list[str]: + """Make remote dirs + rsync inputs. Returns the remote input paths.""" + root = node["root"] + rjob = f"{root}/data/remote/{job_id}" + await _run(SSH + [node["ssh"], f"mkdir -p {shlex.quote(rjob)}/inputs {shlex.quote(rjob)}/out"]) + remote_inputs = [] + for lp in local_inputs: + name = Path(lp).name + dest = f"{node['ssh']}:{rjob}/inputs/{name}" + code, out = await _run(["rsync", "-a", "-e", "ssh -o BatchMode=yes", + lp.rstrip("/") + ("/" if Path(lp).is_dir() else ""), dest], timeout=600) + if code != 0: + raise RuntimeError(f"rsync input failed: {out[:200]}") + remote_inputs.append(f"{rjob}/inputs/{name}") + return remote_inputs + + +def remote_cmd(node: dict, op_id: str, op_python_rel: str, entry_rel: str, + job_id: str, params_json: str, remote_inputs: list[str]) -> list[str]: + """ssh command that runs the operator on the remote node with repo-relative paths. + HF_TOKEN etc. are sourced from the node's .env.remote (kept off the process table).""" + root = node["root"] + rout = f"{root}/data/remote/{job_id}/out" + inner = (f"cd {shlex.quote(root)} && set -a; [ -f .env.remote ] && . ./.env.remote; set +a; " + f"exec {shlex.quote(op_python_rel)} {shlex.quote(entry_rel)} " + f"--outdir {shlex.quote(rout)} --params {shlex.quote(params_json)}") + for ri in remote_inputs: + inner += f" --input {shlex.quote(ri)}" + return SSH + ["-T", node["ssh"], f"bash -lc {shlex.quote(inner)}"] + + +async def collect(node: dict, job_id: str, local_outdir: Path) -> None: + """rsync the remote outputs back into the local job outdir, then clean up remote.""" + root = node["root"] + rjob = f"{root}/data/remote/{job_id}" + code, out = await _run(["rsync", "-a", "-e", "ssh -o BatchMode=yes", + f"{node['ssh']}:{rjob}/out/", f"{str(local_outdir).rstrip('/')}/"], timeout=600) + if code != 0: + raise RuntimeError(f"rsync output failed: {out[:200]}") + await _run(SSH + [node["ssh"], f"rm -rf {shlex.quote(rjob)}"], timeout=30) diff --git a/server/runner.py b/server/runner.py index 82ef728..43567bb 100644 --- a/server/runner.py +++ b/server/runner.py @@ -18,10 +18,12 @@ import signal import sys from pathlib import Path -from . import db, settings as settings_mod, store +from . import db, remote, settings as settings_mod, store JOBS_DIR = db.DATA / "jobs" -LANE_LIMITS = {"gpu": 1, "cpu": 3, "net": 6} +# 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: @@ -30,6 +32,7 @@ class Runner: 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() @@ -55,7 +58,9 @@ class Runner: 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" + # 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, @@ -135,6 +140,9 @@ class Runner: 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"],)) @@ -156,6 +164,19 @@ class Runner: 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) @@ -166,7 +187,22 @@ class Runner: 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): + 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 @@ -192,28 +228,48 @@ class Runner: 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") - # a manifest "python" may be relative to the repo root (portable across - # machines) or absolute (legacy); resolve relative ones against db.ROOT - op_py = op.get("python") - if op_py and not Path(op_py).is_absolute(): - op_py = str(db.ROOT / op_py) - python = op_py or sys.executable - if op_py 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] + # 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=env) + start_new_session=True, env=run_env) self.procs[job_id] = proc - log_lines: list[str] = [] + log_lines: list[str] = [f"[node: {node.get('name', 'local')}]\n"] assert proc.stdout try: async for raw in proc.stdout: @@ -235,6 +291,13 @@ class Runner: 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}) diff --git a/server/sysinfo.py b/server/sysinfo.py index 989cebe..fe59e5d 100644 --- a/server/sysinfo.py +++ b/server/sysinfo.py @@ -53,8 +53,11 @@ def _du_gb(path: Path): def _lanes(con, runner): - limits = getattr(runner, "LANE_LIMITS", None) or {"gpu": 1, "cpu": 3, "net": 6} - lanes = {k: {"limit": v, "running": 0, "queued": 0} for k, v in limits.items()} + # gpu lane capacity = number of gpu nodes in the pool (this Mac + remote workers) + gpu_limit = max(1, len(getattr(runner, "gpu_nodes", []) or [])) + lanes = {"gpu": {"limit": gpu_limit, "running": 0, "queued": 0}, + "cpu": {"limit": 3, "running": 0, "queued": 0}, + "net": {"limit": 6, "running": 0, "queued": 0}} rows = con.execute("SELECT operator, status FROM jobs WHERE status IN ('queued','running')").fetchall() for r in rows: lane = runner.lane_of(r["operator"]) @@ -64,6 +67,13 @@ def _lanes(con, runner): return lanes +def gpu_pool(runner) -> list[dict]: + """Per-node status for the dashboard: name, remote?, currently busy.""" + return [{"name": n.get("name", n.get("ssh", "?")), + "remote": not n.get("local", False), + "busy": n.get("busy", False)} for n in getattr(runner, "gpu_nodes", [])] + + def _jobs_24h(con): cutoff = time.time() - 86400 rows = con.execute( @@ -99,6 +109,7 @@ def snapshot(con, runner) -> dict: **base, "gpu": _macmon_gpu(), "lanes": _lanes(con, runner), + "nodes": gpu_pool(runner), "jobs_24h": _jobs_24h(con), "data_dir_gb": _cached("data_du", 300.0, lambda: _du_gb(db.DATA)), "hf_cache_gb": _cached("hf_du", 300.0, lambda: _du_gb(hf) if hf.exists() else None),