ryan/hypertower
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

pre-refactor 041426

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files
v3/distributed/client.py
+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()