pre-refactor 041426

This commit is contained in:
rpotter6298
2026-04-14 19:42:16 +02:00
parent eb9eafe715
commit 13290575d5
75 changed files with 8900 additions and 77 deletions
View File
+379
View File
@@ -0,0 +1,379 @@
"""
HyperTower distributed job CLI — submit jobs, view status.
Usage:
# View all connected clients
python -m v3.distributed.cli clients
# View a specific client
python -m v3.distributed.cli clients <client_id>
# View jobs (optionally filter by state)
python -m v3.distributed.cli jobs [--state pending|running|done|failed]
# Submit a job
python -m v3.distributed.cli submit \\
--run-name phase2/leaky \\
-- --bridge-mode image_only --backbone resnet50 --reps 10 ...
# Submit all jobs from a batch file (JSON)
python -m v3.distributed.cli submit-batch jobs.json
# Cancel a pending job
python -m v3.distributed.cli cancel <job_id>
Global flags (can also be set via env vars):
--server HT_SERVER e.g. http://apollo:8765
--token HT_TOKEN
"""
from __future__ import annotations
import argparse
import json
import os
import sys
import time
from datetime import datetime, timezone, timedelta
from typing import Optional
import requests
# ──────────────────────────────────────────────────────────────
# HTTP helpers
# ──────────────────────────────────────────────────────────────
class _API:
def __init__(self, base_url: str, token: str):
self.base_url = base_url.rstrip("/")
self._h = {"x-token": token}
def get(self, path: str, **params) -> object:
r = requests.get(f"{self.base_url}{path}", headers=self._h,
params=params, timeout=10)
r.raise_for_status()
return r.json()
def post(self, path: str, body: dict) -> object:
r = requests.post(f"{self.base_url}{path}", headers=self._h,
json=body, timeout=10)
r.raise_for_status()
return r.json()
def delete(self, path: str) -> object:
r = requests.delete(f"{self.base_url}{path}", headers=self._h, timeout=10)
r.raise_for_status()
return r.json()
# ──────────────────────────────────────────────────────────────
# Formatting helpers
# ──────────────────────────────────────────────────────────────
def _ago(ts: Optional[str]) -> str:
if not ts:
return "-"
try:
dt = datetime.fromisoformat(ts)
delta = datetime.now(timezone.utc) - dt
secs = int(delta.total_seconds())
if secs < 60:
return f"{secs}s ago"
elif secs < 3600:
return f"{secs//60}m ago"
else:
return f"{secs//3600}h{(secs%3600)//60}m ago"
except Exception:
return ts
def _col(text: str, width: int) -> str:
text = str(text) if text is not None else "-"
return text[:width].ljust(width)
def _table(rows: list[list[str]], headers: list[str]):
widths = [max(len(str(r[i])) for r in ([headers] + rows)) for i in range(len(headers))]
sep = " "
def _row(r):
return sep.join(str(r[i]).ljust(widths[i]) for i in range(len(r)))
print(_row(headers))
print("-" * (sum(widths) + len(sep) * (len(widths) - 1)))
for r in rows:
print(_row(r))
# ──────────────────────────────────────────────────────────────
# Subcommands
# ──────────────────────────────────────────────────────────────
def _clients_table(api: _API):
"""Render one clients snapshot. Returns the printed lines as a string."""
clients = api.get("/clients")
if not clients:
return "No clients connected."
rows = []
for c in clients:
s = c["status"]
parts = []
if s.get("rep") is not None:
parts.append(f"rep{s['rep']:02d}")
if s.get("fold") is not None:
parts.append(f"f{s['fold']}")
if s.get("epoch") is not None:
parts.append(f"ep{s['epoch']}/{s.get('total_epochs','?')}")
parts.append(f"auc={s.get('last_auc','?')}")
prog = " ".join(parts) if parts else "-"
rows.append([
c["client_id"],
c["hostname"],
c["gpu_info"][:30],
s["state"],
s.get("run_name") or "-",
prog,
_ago(c["last_seen"]),
])
headers = ["ID", "HOST", "GPU", "STATE", "RUN", "PROGRESS", "SEEN"]
widths = [max(len(str(r[i])) for r in ([headers] + rows)) for i in range(len(headers))]
sep = " "
lines = []
lines.append(sep.join(str(h).ljust(widths[i]) for i, h in enumerate(headers)))
lines.append("-" * (sum(widths) + len(sep) * (len(widths) - 1)))
for r in rows:
lines.append(sep.join(str(r[i]).ljust(widths[i]) for i in range(len(r))))
return "\n".join(lines)
def cmd_clients(api: _API, args):
if hasattr(args, "client_id") and args.client_id:
data = api.get(f"/clients/{args.client_id}")
s = data["status"]
print(f"client_id : {data['client_id']}")
print(f"hostname : {data['hostname']}")
print(f"gpu : {data['gpu_info']}")
print(f"last_seen : {_ago(data['last_seen'])}")
print(f"state : {s['state']}")
if s.get("job_id"):
print(f"job : {s['job_id']} ({s.get('run_name', '')})")
if s.get("rep") is not None:
print(f"progress : rep {s['rep']} fold {s.get('fold', '?')} "
f"ep {s.get('epoch', '?')}/{s.get('total_epochs', '?')} "
f"({s.get('last_epoch_secs', '?')}s/ep) "
f"auc={s.get('last_auc', '?')}")
return
watch = getattr(args, "watch", False)
interval = getattr(args, "interval", 5)
if not watch:
print(_clients_table(api))
return
try:
while True:
now = datetime.now().strftime("%H:%M:%S")
print(f"\033[H\033[2J", end="") # clear screen
print(f"HyperTower clients [{now}] (Ctrl-C to exit)\n")
print(_clients_table(api))
time.sleep(interval)
except KeyboardInterrupt:
print("\nStopped.")
def _rep_label(job: dict) -> str:
"""Extract --rep-index from job args if present."""
try:
a = json.loads(job["args"]) if isinstance(job.get("args"), str) else job.get("args", [])
if "--rep-index" in a:
idx = a[a.index("--rep-index") + 1]
return f"rep{int(idx):02d}"
except Exception:
pass
return "-"
def cmd_jobs(api: _API, args):
params = {}
if hasattr(args, "state") and args.state:
params["state"] = args.state
jobs = api.get("/jobs", **params)
if not jobs:
print("No jobs.")
return
rows = []
for j in jobs:
attempts = j.get("attempts", 0)
rows.append([
j["job_id"][:12],
j["run_name"],
_rep_label(j),
j["state"],
f"{attempts}" if attempts else "-",
j.get("assigned_to") or "-",
_ago(j["created_at"]),
_ago(j.get("started_at")),
_ago(j.get("completed_at")),
])
_table(rows, ["JOB_ID", "RUN_NAME", "REP", "STATE", "TRIES", "CLIENT", "CREATED", "STARTED", "DONE"])
pending = sum(1 for j in jobs if j["state"] == "pending")
running = sum(1 for j in jobs if j["state"] == "running")
done = sum(1 for j in jobs if j["state"] == "done")
failed = sum(1 for j in jobs if j["state"] == "failed")
print(f"\n {len(jobs)} total | {pending} pending {running} running {done} done {failed} failed")
def cmd_submit(api: _API, args):
body = {
"run_name": args.run_name,
"module": args.module,
"args": args.run_args,
"output_dir": args.output_dir,
"priority": args.priority,
}
resp = api.post("/jobs", body)
print(f"Queued job {resp['job_id']} ({args.run_name})")
def cmd_submit_batch(api: _API, args):
with open(args.batch_file) as f:
jobs = json.load(f)
for job in jobs:
resp = api.post("/jobs", job)
print(f"Queued {resp['job_id']} ({job['run_name']})")
def cmd_submit_cv(api: _API, args):
"""Submit one job per rep, each writing to repNN under the same run-name."""
seed_start = args.rep_seed_start
seed_step = args.rep_seed_step
submitted = []
for i in range(args.reps):
seed = seed_start + i * seed_step
rep_args = [a for a in args.run_args
if a not in ("--reps", "--rep-seed-start", "--rep-seed-step")]
rep_args += [
"--reps", "1",
"--rep-seed-start", str(seed),
"--rep-index", str(i),
]
body = {
"run_name": args.run_name,
"module": args.module,
"args": rep_args,
"output_dir": args.output_dir,
"priority": args.priority,
}
resp = api.post("/jobs", body)
submitted.append(resp["job_id"])
print(f"Queued rep{i:02d} seed={seed} job_id={resp['job_id']}")
print(f"\n{len(submitted)} jobs queued for run '{args.run_name}'")
def cmd_cancel(api: _API, args):
resp = api.delete(f"/jobs/{args.job_id}")
print(f"Cancelled {args.job_id}" if resp.get("ok") else resp)
def cmd_clear(api: _API, args):
body: dict = {}
if args.all:
body["all"] = True
elif args.run_name:
body["run_name"] = args.run_name
else:
body["states"] = args.states or ["done", "failed", "cancelled"]
resp = api.post("/jobs/clear", body)
print(f"Cleared {resp['cleared']} jobs.")
# ──────────────────────────────────────────────────────────────
# Parser
# ──────────────────────────────────────────────────────────────
def main():
ap = argparse.ArgumentParser(description=__doc__,
formatter_class=argparse.RawDescriptionHelpFormatter)
ap.add_argument("--server", default=os.environ.get("HT_SERVER", ""),
help="Server URL (or set HT_SERVER)")
ap.add_argument("--token", default=os.environ.get("HT_TOKEN", ""),
help="Shared secret (or set HT_TOKEN)")
sub = ap.add_subparsers(dest="cmd", required=True)
# clients
p_cl = sub.add_parser("clients", help="List clients or inspect one")
p_cl.add_argument("client_id", nargs="?")
p_cl.add_argument("--watch", "-w", action="store_true",
help="Live monitoring mode — refresh every --interval seconds")
p_cl.add_argument("--interval", "-n", type=int, default=5,
help="Refresh interval in seconds for --watch (default: 5)")
# jobs
p_j = sub.add_parser("jobs", help="List jobs")
p_j.add_argument("--state", choices=["pending", "running", "done", "failed", "cancelled"])
# submit
p_s = sub.add_parser("submit", help="Submit a single job")
p_s.add_argument("--run-name", required=True)
p_s.add_argument("--module", default="v3.scripts.main.run_cv")
p_s.add_argument("--output-dir", default="v3/results")
p_s.add_argument("--priority", type=int, default=0)
p_s.add_argument("run_args", nargs=argparse.REMAINDER,
help="Args after '--' are forwarded to the module")
# submit-cv
p_cv = sub.add_parser("submit-cv",
help="Submit one job per rep (distributed 10x5 etc.)")
p_cv.add_argument("--run-name", required=True)
p_cv.add_argument("--reps", type=int, required=True)
p_cv.add_argument("--rep-seed-start", type=int, default=100)
p_cv.add_argument("--rep-seed-step", type=int, default=100)
p_cv.add_argument("--module", default="v3.scripts.main.run_cv")
p_cv.add_argument("--output-dir", default="v3/results")
p_cv.add_argument("--priority", type=int, default=0)
p_cv.add_argument("run_args", nargs=argparse.REMAINDER,
help="Args after '--' forwarded to run_cv (omit --reps/--rep-seed-*)")
# submit-batch
p_b = sub.add_parser("submit-batch", help="Submit jobs from a JSON file")
p_b.add_argument("batch_file")
# cancel
p_c = sub.add_parser("cancel", help="Cancel a pending job")
p_c.add_argument("job_id")
# clear
p_cl = sub.add_parser("clear", help="Delete jobs by run-name, state, or everything")
p_cl.add_argument("--run-name", default=None, help="Delete all jobs with this run-name")
p_cl.add_argument("--states", nargs="+",
default=None,
choices=["done", "failed", "cancelled", "pending", "running"],
help="Delete jobs in these states (default: done+failed+cancelled)")
p_cl.add_argument("--all", action="store_true", help="Delete ALL jobs")
args = ap.parse_args()
if not args.server:
ap.error("--server is required (or set HT_SERVER)")
if not args.token:
ap.error("--token is required (or set HT_TOKEN)")
# strip leading "--" from run_args if present
if hasattr(args, "run_args") and args.run_args and args.run_args[0] == "--":
args.run_args = args.run_args[1:]
api = _API(args.server, args.token)
dispatch = {
"clients": cmd_clients,
"jobs": cmd_jobs,
"submit": cmd_submit,
"submit-cv": cmd_submit_cv,
"submit-batch": cmd_submit_batch,
"cancel": cmd_cancel,
"clear": cmd_clear,
}
dispatch[args.cmd](api, args)
if __name__ == "__main__":
main()
+435
View File
@@ -0,0 +1,435 @@
"""
HyperTower distributed job client daemon.
Registers with the server, polls for jobs, syncs code, runs training,
uploads results, and loops. Parses stdout to stream live status.
Usage:
python -m v3.distributed.client \\
--server http://apollo:8765 \\
--token <secret> \\
--server-ssh rpotter@apollo \\
--server-path /home/rpotter/hypertower \\
[--local-path ~/hypertower] \\
[--poll-interval 15]
Compatibility test (verify GPU env, 1 epoch dry-run):
python -m v3.distributed.client ... --test [extra run_cv args]
"""
from __future__ import annotations
import argparse
import os
import re
import shutil
import signal
import socket
import subprocess
import sys
import tempfile
import threading
import time
from pathlib import Path
from typing import Optional
import requests
from .protocol import (
JobResult,
JobSpec,
PollResponse,
RegisterRequest,
RegisterResponse,
StatusPush,
)
# ──────────────────────────────────────────────────────────────
# Server HTTP wrapper
# ──────────────────────────────────────────────────────────────
class _Server:
def __init__(self, base_url: str, token: str):
self.base_url = base_url.rstrip("/")
self._h = {"x-token": token}
self.client_id: str = ""
def _post(self, path: str, **kw) -> dict:
r = requests.post(f"{self.base_url}{path}", headers=self._h, timeout=15, **kw)
r.raise_for_status()
return r.json()
def register(self, hostname: str, gpu_info: str) -> str:
data = self._post("/register",
json={"hostname": hostname, "gpu_info": gpu_info})
self.client_id = data["client_id"]
self.hostname = hostname
self.gpu_info = gpu_info
return self.client_id
def _reregister(self):
try:
self._post("/register",
json={"hostname": self.hostname, "gpu_info": self.gpu_info},
params={"reuse_id": self.client_id})
except Exception:
pass
def poll(self) -> Optional[JobSpec]:
data = self._post("/poll", params={"client_id": self.client_id})
if data.get("please_reregister"):
self._reregister()
return JobSpec(**data["job"]) if data.get("job") else None
def push_status(self, status: StatusPush):
try:
r = requests.post(
f"{self.base_url}/status/{self.client_id}",
json=status.model_dump(),
headers=self._h,
timeout=5,
)
if r.ok and r.json().get("please_reregister"):
self._reregister()
except Exception:
pass # don't crash job on status push failure
def complete(self, job_id: str, success: bool, error_msg: Optional[str] = None):
self._post("/complete",
json={"job_id": job_id, "success": success, "error_msg": error_msg})
# ──────────────────────────────────────────────────────────────
# GPU info
# ──────────────────────────────────────────────────────────────
def _gpu_info() -> str:
# NVIDIA
try:
out = subprocess.check_output(
["nvidia-smi", "--query-gpu=name,memory.total", "--format=csv,noheader"],
text=True, stderr=subprocess.DEVNULL,
).strip()
if out:
return " | ".join(out.splitlines())
except Exception:
pass
# AMD
try:
out = subprocess.check_output(
["rocm-smi", "--showproductname", "--csv"],
text=True, stderr=subprocess.DEVNULL,
).strip().splitlines()
names = [l for l in out if l and not l.startswith("device")]
if names:
return "AMD: " + " | ".join(names)
except Exception:
pass
# AMD fallback
try:
out = subprocess.check_output(
["rocminfo"],
text=True, stderr=subprocess.DEVNULL,
)
names = [l.split(":", 1)[1].strip() for l in out.splitlines()
if "Marketing Name:" in l]
if names:
return "AMD: " + " | ".join(names)
except Exception:
pass
return "no-gpu"
# ──────────────────────────────────────────────────────────────
# rsync helpers
# ──────────────────────────────────────────────────────────────
def _rsync(src: str, dst: str, delete: bool = False):
cmd = ["rsync", "-az", "--info=progress2"]
if delete:
cmd.append("--delete")
cmd += [src, dst]
subprocess.run(cmd, check=True)
def _sync_code(server_ssh: str, server_path: str, local_path: str):
"""Pull v3/ source from server → local (overwrites local changes)."""
src = f"{server_ssh}:{server_path}/v3/"
dst = f"{local_path}/v3/"
Path(dst).mkdir(parents=True, exist_ok=True)
_rsync(src, dst, delete=True)
def _upload_results(server_ssh: str, server_path: str, local_path: str,
run_name: str, output_dir: str):
src = f"{local_path}/{output_dir}/{run_name}/"
dst = f"{server_ssh}:{server_path}/{output_dir}/{run_name}/"
# Ensure parent directory exists on server before rsyncing
remote_parent = f"{server_path}/{output_dir}/{Path(run_name).parent}"
subprocess.run(["ssh", server_ssh, f"mkdir -p '{remote_parent}'"], check=True)
_rsync(src, dst)
def _clean_local(local_path: str, run_name: str, output_dir: str):
target = Path(local_path) / output_dir / run_name
if target.exists():
shutil.rmtree(target)
print(f"[client] cleaned {target}", flush=True)
# ──────────────────────────────────────────────────────────────
# Stdout parsers (match v3_hypertower.py print format)
# ──────────────────────────────────────────────────────────────
# " ep 3/ 30 (6.5s) auc=0.75 acc=0.82"
_EP_RE = re.compile(
r"ep\s+(\d+)/\s*(\d+)\s+\(([0-9.]+)s\)\s+auc=([0-9.nan]+)\s+acc=([0-9.nan]+)"
)
# "[binary:single] fold 2/5"
_FOLD_RE = re.compile(r"fold\s+(\d+)/\d+", re.IGNORECASE)
# "Rep 3 fold_seed=..." or "Rep 3/10 fold_seed=..."
_REP_RE = re.compile(r"Rep\s+(\d+)", re.IGNORECASE)
def _parse_line(line: str) -> dict:
"""Return any structured fields found in a stdout line."""
out = {}
m = _FOLD_RE.search(line)
if m:
out["fold"] = int(m.group(1))
m = _REP_RE.search(line)
if m:
out["rep"] = int(m.group(1))
m = _EP_RE.search(line)
if m:
out["epoch"] = int(m.group(1))
out["total_epochs"] = int(m.group(2))
out["last_epoch_secs"] = float(m.group(3))
try:
out["last_auc"] = float(m.group(4))
except ValueError:
pass
return out
# ──────────────────────────────────────────────────────────────
# Core job runner
# ──────────────────────────────────────────────────────────────
def _run_job(job: JobSpec, server: _Server,
server_ssh: str, server_path: str, local_path: str,
no_sync: bool = False, num_workers: Optional[int] = None,
extra_args: list[str] | None = None) -> bool:
extra_args = extra_args or []
# 1. Sync code
if no_sync:
print(f"[client] skipping sync (--no-sync)", flush=True)
else:
print(f"[client] syncing v3/ from server...", flush=True)
server.push_status(StatusPush(state="syncing", job_id=job.job_id, run_name=job.run_name))
_sync_code(server_ssh, server_path, local_path)
# 2. Launch training subprocess
job_args = job.args
if num_workers is not None and "--num-workers" not in job_args:
job_args = job_args + ["--num-workers", str(num_workers)]
cmd = [sys.executable, "-m", job.module] + job_args + extra_args
print(f"[client] running: {' '.join(cmd)}", flush=True)
server.push_status(StatusPush(state="running", job_id=job.job_id, run_name=job.run_name))
# Write stdout to a temp file so the child's fd table is unmodified —
# this lets Python 3.14's forkserver communicate over its own pipes without
# interference. We tail the file from a thread for live status.
log_dir = Path(local_path) / "v3" / "distributed" / "logs"
log_dir.mkdir(parents=True, exist_ok=True)
log_file = log_dir / f"job_{job.job_id}.log"
ctx: dict = {}
def _tail(path: Path):
with open(path, "r") as f:
while True:
raw = f.readline()
if raw:
print(raw, end="", flush=True)
info = _parse_line(raw)
ctx.update(info)
if "epoch" in info:
server.push_status(StatusPush(
state="running",
job_id=job.job_id,
run_name=job.run_name,
rep=ctx.get("rep"),
fold=ctx.get("fold"),
epoch=ctx.get("epoch"),
total_epochs=ctx.get("total_epochs"),
last_epoch_secs=ctx.get("last_epoch_secs"),
last_auc=ctx.get("last_auc"),
))
elif proc.poll() is not None:
# Drain any remaining output then exit
for raw in f:
print(raw, end="", flush=True)
break
else:
time.sleep(0.05)
with open(log_file, "w") as logf:
proc = subprocess.Popen(
cmd,
stdout=logf,
stderr=logf,
cwd=local_path,
start_new_session=True, # own process group so we can kill all workers on failure
)
tailer = threading.Thread(target=_tail, args=(log_file,), daemon=True)
tailer.start()
proc.wait()
tailer.join(timeout=5)
log_file.unlink(missing_ok=True)
success = proc.returncode == 0
if not success:
# Kill any surviving worker processes in the group (e.g. DataLoader workers
# that outlived the main process after an OOM kill)
try:
os.killpg(os.getpgid(proc.pid), signal.SIGKILL)
except ProcessLookupError:
pass # already gone
if not success:
print(f"[client] job FAILED (rc={proc.returncode})", flush=True)
return False
# 3. Upload results (skip when client IS the server — results already in place)
if no_sync:
print(f"[client] skipping upload (--no-sync, results already local)", flush=True)
else:
print(f"[client] uploading results...", flush=True)
server.push_status(StatusPush(state="uploading", job_id=job.job_id, run_name=job.run_name))
_upload_results(server_ssh, server_path, local_path, job.run_name, job.output_dir)
# 4. Clean local (only remote clients need cleanup)
_clean_local(local_path, job.run_name, job.output_dir)
print(f"[client] job {job.job_id} complete.", flush=True)
return True
# ──────────────────────────────────────────────────────────────
# Compatibility test
# ──────────────────────────────────────────────────────────────
def _compat_test(server: _Server, server_ssh: str, server_path: str,
local_path: str, extra_args: list[str], no_sync: bool = False):
"""Run 1 epoch on 1 rep to verify the env works end-to-end."""
print("[client] === compatibility test ===", flush=True)
job = JobSpec(
job_id="compat-test",
run_name="_compat_test",
module="v3.scripts.main.run_cv",
args=[
"--bridge-mode", "image_only",
"--backbone", "resnet50",
"--epochs", "1",
"--reps", "1",
"--rep-seed-start", "42",
"--in-memory-cache",
"--num-workers", "0",
"--output-root", "v3/results",
"--run-name", "_compat_test",
] + extra_args,
output_dir="v3/results",
)
ok = _run_job(job, server, server_ssh, server_path, local_path,
no_sync=no_sync, extra_args=[a for a in extra_args if a != "--"])
# always clean up test results
_clean_local(local_path, "_compat_test", "v3/results")
if ok:
print("[client] compatibility test PASSED ✓", flush=True)
else:
print("[client] compatibility test FAILED ✗", flush=True)
return ok
# ──────────────────────────────────────────────────────────────
# Main daemon
# ──────────────────────────────────────────────────────────────
def main():
ap = argparse.ArgumentParser(
description=__doc__,
formatter_class=argparse.RawDescriptionHelpFormatter,
)
ap.add_argument("--server", required=True,
help="Server URL, e.g. http://apollo:8765")
ap.add_argument("--token", default=os.environ.get("HT_TOKEN", ""),
help="Shared secret (or set HT_TOKEN env var)")
ap.add_argument("--server-ssh", required=True,
help="SSH target for rsync, e.g. rpotter@apollo")
ap.add_argument("--server-path", required=True,
help="Absolute path to hypertower root on server")
ap.add_argument("--local-path",
default=str(Path.home() / "hypertower"),
help="Absolute path to hypertower root on this machine")
ap.add_argument("--poll-interval", type=int, default=15,
help="Seconds to wait between polls when idle")
ap.add_argument("--num-workers", type=int, default=None,
help="Override DataLoader num_workers for all jobs on this client.")
ap.add_argument("--extra-args", nargs=argparse.REMAINDER, default=[],
help="Extra args appended to every job on this client (e.g. --cache-workers 0). "
"Use -- to separate from client args: --extra-args -- --cache-workers 0")
ap.add_argument("--no-sync", action="store_true",
help="Skip rsync of v3/ before each job (use when client IS the server)")
ap.add_argument("--test", action="store_true",
help="Run 1-epoch compatibility test and exit")
ap.add_argument("extra_args", nargs="*",
help="Extra args forwarded to run_cv in --test mode")
args = ap.parse_args()
if not args.token:
ap.error("--token is required (or set HT_TOKEN)")
hostname = socket.gethostname()
gpu_info = _gpu_info()
server = _Server(args.server, args.token)
client_id = server.register(hostname, gpu_info)
print(f"[client] registered as {client_id} ({hostname} | {gpu_info})", flush=True)
if args.test:
sys.exit(0 if _compat_test(
server, args.server_ssh, args.server_path,
args.local_path, args.extra_args,
no_sync=args.no_sync,
) else 1)
print(f"[client] polling every {args.poll_interval}s...", flush=True)
while True:
try:
job = server.poll()
if job is None:
server.push_status(StatusPush(state="idle"))
time.sleep(args.poll_interval)
continue
success = _run_job(
job, server,
args.server_ssh, args.server_path, args.local_path,
no_sync=args.no_sync,
num_workers=args.num_workers,
extra_args=[a for a in args.extra_args if a != "--"],
)
server.complete(job.job_id, success,
error_msg=None if success else "non-zero exit code")
server.push_status(StatusPush(state="idle"))
except KeyboardInterrupt:
print("\n[client] shutting down", flush=True)
break
except Exception as exc:
print(f"[client] error: {exc}", flush=True)
time.sleep(args.poll_interval)
if __name__ == "__main__":
main()
Binary file not shown.
+62
View File
@@ -0,0 +1,62 @@
"""Shared data models for server/client communication."""
from __future__ import annotations
from typing import Optional
from pydantic import BaseModel
class RegisterRequest(BaseModel):
hostname: str
gpu_info: str
class RegisterResponse(BaseModel):
client_id: str
class StatusPush(BaseModel):
state: str # idle | syncing | running | uploading | error
job_id: Optional[str] = None
run_name: Optional[str] = None
rep: Optional[int] = None
fold: Optional[int] = None
epoch: Optional[int] = None
total_epochs: Optional[int] = None
last_epoch_secs: Optional[float] = None
last_auc: Optional[float] = None
error: Optional[str] = None
class ClientInfo(BaseModel):
client_id: str
hostname: str
gpu_info: str
status: StatusPush
last_seen: str
class JobSpec(BaseModel):
job_id: str
run_name: str
module: str # e.g. "v3.scripts.main.run_cv"
args: list[str]
output_dir: str = "v3/results" # local subdir for rsync-back
class PollResponse(BaseModel):
job: Optional[JobSpec] = None
please_reregister: bool = False
class JobResult(BaseModel):
job_id: str
success: bool
error_msg: Optional[str] = None
class JobSubmit(BaseModel):
run_name: str
module: str = "v3.scripts.main.run_cv"
args: list[str]
output_dir: str = "v3/results"
priority: int = 0
+402
View File
@@ -0,0 +1,402 @@
"""
HyperTower distributed job server.
Manages a SQLite job queue and a registry of connected clients.
Clients poll for work, push status updates, and report completion.
Usage:
python -m v3.distributed.server --port 8765 --token <secret>
Environment:
HT_TOKEN — fallback if --token is not passed
"""
from __future__ import annotations
import argparse
import json
import os
import sqlite3
import threading
import time
import uuid
from contextlib import contextmanager
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional
from fastapi import Depends, FastAPI, Header, HTTPException
import uvicorn
from .protocol import (
ClientInfo,
JobResult,
JobSpec,
JobSubmit,
PollResponse,
RegisterRequest,
RegisterResponse,
StatusPush,
)
# ──────────────────────────────────────────────────────────────
# Global state
# ──────────────────────────────────────────────────────────────
_TOKEN: str = ""
_DB_PATH: Path = Path("v3/distributed/jobs.db")
_CLIENT_TTL: int = 120 # seconds before a client is considered gone
_MAX_ATTEMPTS: int = 3 # max times a job is retried before being left as failed
_clients: dict[str, ClientInfo] = {}
_clients_lock = threading.Lock()
def _reap_stale_clients():
"""Background thread: remove silent clients and re-queue their running jobs."""
while True:
time.sleep(30)
cutoff = datetime.now(timezone.utc).timestamp() - _CLIENT_TTL
# Step 1: evict timed-out clients from registry
with _clients_lock:
stale = [
cid for cid, c in _clients.items()
if datetime.fromisoformat(c.last_seen).timestamp() < cutoff
]
for cid in stale:
print(f"[server] reaped stale client {cid} ({_clients[cid].hostname})", flush=True)
del _clients[cid]
known_ids = set(_clients.keys())
# Step 2: reset any running job whose assigned client is no longer known
with _db() as conn:
rows = conn.execute(
"SELECT job_id, assigned_to FROM jobs WHERE state='running'"
).fetchall()
for row in rows:
if row["assigned_to"] not in known_ids:
conn.execute(
"UPDATE jobs SET state='pending', assigned_to=NULL, started_at=NULL "
"WHERE job_id=?",
(row["job_id"],)
)
print(f"[server] re-queued job {row['job_id']} "
f"(client {row['assigned_to']} unknown)", flush=True)
# ──────────────────────────────────────────────────────────────
# Database helpers
# ──────────────────────────────────────────────────────────────
@contextmanager
def _db():
conn = sqlite3.connect(str(_DB_PATH))
conn.row_factory = sqlite3.Row
try:
yield conn
conn.commit()
finally:
conn.close()
def _init_db():
_DB_PATH.parent.mkdir(parents=True, exist_ok=True)
with _db() as conn:
conn.execute("""
CREATE TABLE IF NOT EXISTS jobs (
job_id TEXT PRIMARY KEY,
run_name TEXT NOT NULL,
module TEXT NOT NULL,
args TEXT NOT NULL, -- JSON list
output_dir TEXT NOT NULL DEFAULT 'v3/results',
state TEXT NOT NULL DEFAULT 'pending',
priority INTEGER NOT NULL DEFAULT 0,
assigned_to TEXT,
created_at TEXT NOT NULL,
started_at TEXT,
completed_at TEXT,
error_msg TEXT,
attempts INTEGER NOT NULL DEFAULT 0
)
""")
# Add attempts column to existing DBs that predate this field
try:
conn.execute("ALTER TABLE jobs ADD COLUMN attempts INTEGER NOT NULL DEFAULT 0")
except Exception:
pass # column already exists
# Note: running jobs are NOT reset on startup — active clients will re-register
# via poll/status and the reaper will clean up any that don't reconnect within TTL.
def _now() -> str:
return datetime.now(timezone.utc).isoformat()
def _ensure_client(client_id: str, hostname: str = "", gpu_info: str = "") -> bool:
"""Re-register a client that survived a server restart.
Returns True if the client was unknown (placeholder created) so the caller
can ask the client to re-register with full info."""
if client_id not in _clients:
_clients[client_id] = ClientInfo(
client_id=client_id,
hostname=hostname or client_id,
gpu_info=gpu_info or "unknown",
status=StatusPush(state="idle"),
last_seen=_now(),
)
print(f"[server] re-registered {client_id} (survived restart)", flush=True)
return True
return False
# ──────────────────────────────────────────────────────────────
# FastAPI app
# ──────────────────────────────────────────────────────────────
app = FastAPI(title="HyperTower Job Server")
def _check_token(x_token: str = Header(...)):
if x_token != _TOKEN:
raise HTTPException(status_code=403, detail="Invalid token")
# ── Registration ──────────────────────────────────────────────
@app.post("/register", response_model=RegisterResponse,
dependencies=[Depends(_check_token)])
def register(req: RegisterRequest, reuse_id: Optional[str] = None):
with _clients_lock:
client_id = reuse_id if (reuse_id and reuse_id in _clients) else str(uuid.uuid4())[:8]
existing_status = _clients[client_id].status if client_id in _clients else StatusPush(state="idle")
_clients[client_id] = ClientInfo(
client_id=client_id,
hostname=req.hostname,
gpu_info=req.gpu_info,
status=existing_status,
last_seen=_now(),
)
action = "re-registered" if reuse_id else "registered"
print(f"[server] {action} {client_id} ({req.hostname} | {req.gpu_info})", flush=True)
return RegisterResponse(client_id=client_id)
# ── Job polling ───────────────────────────────────────────────
@app.post("/poll", response_model=PollResponse,
dependencies=[Depends(_check_token)])
def poll(client_id: str):
with _clients_lock:
needs_reregister = _ensure_client(client_id)
_clients[client_id].last_seen = _now()
with _db() as conn:
row = conn.execute(
"SELECT * FROM jobs WHERE state='pending' "
"ORDER BY priority DESC, created_at ASC LIMIT 1"
).fetchone()
if row is None:
return PollResponse(job=None, please_reregister=needs_reregister)
job_id = row["job_id"]
# Reset any other running jobs for this client — client can only work on one at a time.
# This cleans up orphans left over from server restarts.
cur = conn.execute(
"UPDATE jobs SET state='pending', assigned_to=NULL, started_at=NULL "
"WHERE assigned_to=? AND state='running' AND job_id!=?",
(client_id, job_id),
)
if cur.rowcount:
print(f"[server] reset {cur.rowcount} orphaned running job(s) for {client_id}", flush=True)
conn.execute(
"UPDATE jobs SET state='running', assigned_to=?, started_at=? WHERE job_id=?",
(client_id, _now(), job_id),
)
job = JobSpec(
job_id=job_id,
run_name=row["run_name"],
module=row["module"],
args=json.loads(row["args"]),
output_dir=row["output_dir"],
)
with _clients_lock:
_clients[client_id].status = StatusPush(
state="syncing", job_id=job_id, run_name=row["run_name"]
)
print(f"[server] dispatched {job_id} ({row['run_name']}) → {client_id}", flush=True)
return PollResponse(job=job, please_reregister=needs_reregister)
# ── Status ────────────────────────────────────────────────────
@app.post("/status/{client_id}", dependencies=[Depends(_check_token)])
def push_status(client_id: str, status: StatusPush):
with _clients_lock:
needs_reregister = _ensure_client(client_id)
_clients[client_id].status = status
_clients[client_id].last_seen = _now()
return {"ok": True, "please_reregister": needs_reregister}
@app.get("/clients", dependencies=[Depends(_check_token)])
def list_clients():
with _clients_lock:
return list(_clients.values())
@app.get("/clients/{client_id}", dependencies=[Depends(_check_token)])
def get_client(client_id: str):
with _clients_lock:
if client_id not in _clients:
raise HTTPException(status_code=404, detail="Unknown client")
return _clients[client_id]
# ── Job completion ────────────────────────────────────────────
@app.post("/complete", dependencies=[Depends(_check_token)])
def complete(result: JobResult):
with _db() as conn:
if result.success:
conn.execute(
"UPDATE jobs SET state='done', completed_at=?, error_msg=NULL WHERE job_id=?",
(_now(), result.job_id),
)
print(f"[server] job {result.job_id} → done", flush=True)
# Auto-clear run if all jobs for this run_name are now done
run_row = conn.execute(
"SELECT run_name FROM jobs WHERE job_id=?", (result.job_id,)
).fetchone()
if run_row:
run_name = run_row["run_name"]
remaining = conn.execute(
"SELECT COUNT(*) FROM jobs WHERE run_name=? AND state != 'done'",
(run_name,)
).fetchone()[0]
if remaining == 0:
total = conn.execute(
"SELECT COUNT(*) FROM jobs WHERE run_name=?", (run_name,)
).fetchone()[0]
conn.execute("DELETE FROM jobs WHERE run_name=?", (run_name,))
print(f"[server] run '{run_name}' complete ({total} jobs) — cleared", flush=True)
else:
row = conn.execute(
"SELECT attempts FROM jobs WHERE job_id=?", (result.job_id,)
).fetchone()
attempts = (row["attempts"] if row else 0) + 1
if attempts < _MAX_ATTEMPTS:
conn.execute(
"UPDATE jobs SET state='pending', assigned_to=NULL, started_at=NULL, "
"attempts=?, error_msg=? WHERE job_id=?",
(attempts, result.error_msg, result.job_id),
)
print(f"[server] job {result.job_id} failed (attempt {attempts}/{_MAX_ATTEMPTS}), "
f"re-queuing", flush=True)
else:
conn.execute(
"UPDATE jobs SET state='failed', completed_at=?, attempts=?, error_msg=? "
"WHERE job_id=?",
(_now(), attempts, result.error_msg, result.job_id),
)
print(f"[server] job {result.job_id} failed permanently after "
f"{attempts} attempts", flush=True)
return {"ok": True}
# ── Job queue management ──────────────────────────────────────
@app.post("/jobs", dependencies=[Depends(_check_token)])
def submit_job(job: JobSubmit):
job_id = str(uuid.uuid4())[:12]
with _db() as conn:
conn.execute(
"INSERT INTO jobs "
"(job_id, run_name, module, args, output_dir, priority, created_at) "
"VALUES (?,?,?,?,?,?,?)",
(job_id, job.run_name, job.module, json.dumps(job.args),
job.output_dir, job.priority, _now()),
)
print(f"[server] queued {job_id} ({job.run_name})", flush=True)
return {"job_id": job_id}
@app.get("/jobs", dependencies=[Depends(_check_token)])
def list_jobs(state: Optional[str] = None):
with _db() as conn:
if state:
rows = conn.execute(
"SELECT * FROM jobs WHERE state=? ORDER BY created_at DESC", (state,)
).fetchall()
else:
rows = conn.execute(
"SELECT * FROM jobs ORDER BY created_at DESC"
).fetchall()
return [dict(r) for r in rows]
@app.post("/jobs/clear", dependencies=[Depends(_check_token)])
def clear_jobs(body: dict):
with _db() as conn:
if body.get("all"):
cur = conn.execute("DELETE FROM jobs")
elif body.get("run_name"):
cur = conn.execute("DELETE FROM jobs WHERE run_name=?", (body["run_name"],))
else:
states = body.get("states", ["done", "failed", "cancelled"])
placeholders = ",".join("?" * len(states))
cur = conn.execute(f"DELETE FROM jobs WHERE state IN ({placeholders})", states)
print(f"[server] cleared {cur.rowcount} jobs", flush=True)
return {"cleared": cur.rowcount}
@app.delete("/jobs/{job_id}", dependencies=[Depends(_check_token)])
def cancel_job(job_id: str):
with _db() as conn:
conn.execute(
"UPDATE jobs SET state='cancelled' WHERE job_id=? AND state='pending'",
(job_id,),
)
return {"ok": True}
# ──────────────────────────────────────────────────────────────
# Entry point
# ──────────────────────────────────────────────────────────────
def main():
ap = argparse.ArgumentParser(description=__doc__,
formatter_class=argparse.RawDescriptionHelpFormatter)
ap.add_argument("--port", type=int, default=8765)
ap.add_argument("--host", default="0.0.0.0")
ap.add_argument("--token", default=os.environ.get("HT_TOKEN", ""),
help="Shared secret (or set HT_TOKEN env var)")
ap.add_argument("--db", default="v3/distributed/jobs.db",
help="Path to SQLite job database")
ap.add_argument("--client-ttl", type=int, default=120,
help="Seconds of silence before a client is reaped (default: 120)")
ap.add_argument("--max-attempts", type=int, default=3,
help="Max times a failed job is retried before being left as failed (default: 3)")
args = ap.parse_args()
if not args.token:
ap.error("--token is required (or set HT_TOKEN)")
global _TOKEN, _DB_PATH, _CLIENT_TTL, _MAX_ATTEMPTS
_TOKEN = args.token
_DB_PATH = Path(args.db)
_CLIENT_TTL = args.client_ttl
_MAX_ATTEMPTS = args.max_attempts
_init_db()
reaper = threading.Thread(target=_reap_stale_clients, daemon=True)
reaper.start()
print(f"[server] listening on {args.host}:{args.port} client_ttl={_CLIENT_TTL}s", flush=True)
uvicorn.run(app, host=args.host, port=args.port, log_level="warning")
if __name__ == "__main__":
main()
+122
View File
@@ -0,0 +1,122 @@
# Distributed Server Cheat Sheet
All commands assume server is running on hades at port 8765.
## Start Server
```bash
python -m v3.distributed.server --token hypertower
```
Run inside tmux so it survives disconnects:
```bash
tmux new -s htserver
python -m v3.distributed.server --token hypertower
# Ctrl-B D to detach
tmux attach -t htserver # reattach later
```
## Start Clients
**Hades (server-local, no sync):**
```bash
python -m v3.distributed.client \
--server http://hades:8765 --token hypertower \
--server-ssh ignored --server-path ignored \
--local-path /home/rpotter/hypertower \
--no-sync
```
**Apollo (remote client):**
```bash
python -m v3.distributed.client \
--server http://hades:8765 --token hypertower \
--server-ssh rpotter@hades \
--server-path /home/rpotter/hypertower \
--local-path /home/odin/hypertower
```
---
## Monitoring
**Live client monitor (refreshes every 5s):**
```bash
python -m v3.distributed.cli --server http://hades:8765 --token hypertower clients --watch
```
**Faster refresh:**
```bash
python -m v3.distributed.cli --server http://hades:8765 --token hypertower clients --watch --interval 2
```
**Inspect a single client:**
```bash
python -m v3.distributed.cli --server http://hades:8765 --token hypertower clients <client_id>
```
**View job queue:**
```bash
python -m v3.distributed.cli --server http://hades:8765 --token hypertower jobs
```
**Filter by state:**
```bash
python -m v3.distributed.cli --server http://hades:8765 --token hypertower jobs --state pending
python -m v3.distributed.cli --server http://hades:8765 --token hypertower jobs --state running
python -m v3.distributed.cli --server http://hades:8765 --token hypertower jobs --state failed
```
---
## Submitting Jobs
**Phase 3 main grid:**
```bash
python -m v3.scripts.main.phase3.dispatch_phase3 \
--server http://hades:8765 --token hypertower
```
**Dry run (check what would be submitted):**
```bash
python -m v3.scripts.main.phase3.dispatch_phase3 \
--server http://hades:8765 --token hypertower --dry-run
```
**Other grids:**
```bash
python -m v3.scripts.main.phase3.dispatch_phase3 \
--server http://hades:8765 --token hypertower \
--grid v3/scripts/main/phase3/epoch_grid.json
python -m v3.scripts.main.phase3.dispatch_phase3 \
--server http://hades:8765 --token hypertower \
--grid v3/scripts/main/phase3/phase35_grid.json
```
---
## Queue Management
**Clear failed jobs:**
```bash
python -m v3.distributed.cli --server http://hades:8765 --token hypertower clear --states failed
```
**Clear running jobs (orphan cleanup):**
```bash
python -m v3.distributed.cli --server http://hades:8765 --token hypertower clear --states running
```
**Clear all jobs:**
```bash
python -m v3.distributed.cli --server http://hades:8765 --token hypertower clear --all
```
**Clear by run name:**
```bash
python -m v3.distributed.cli --server http://hades:8765 --token hypertower clear --run-name phase3/baseline
```
**Cancel a specific job:**
```bash
python -m v3.distributed.cli --server http://hades:8765 --token hypertower cancel <job_id>
```