pre-refactor 041426
This commit is contained in:
@@ -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()
|
||||
@@ -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.
@@ -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
|
||||
@@ -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()
|
||||
@@ -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>
|
||||
```
|
||||
Reference in New Issue
Block a user