rpotter6298
403 lines
16 KiB
Python
403 lines
16 KiB
Python
"""
|
|
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()
|