badgr-cli 1.1.6 → 1.1.7

package/src/gpuDoctor/gpuInfo.js

@@ -4,11 +4,19 @@ function round1(n) {
   return Math.round(n * 10) / 10;
 }
 
+function bytesToGb(n) {
+  return round1(Number(n) / (1024 ** 3));
+}
+
 /**
- * Read-only nvidia-smi probe. Never mutates the machine.
+ * Read-only nvidia-smi probe. Returns null when nvidia-smi isn't installed
+ * (ENOENT) so the caller can fall back to rocm-smi; any other failure
+ * (permission denied, driver broken, etc.) is a real NVIDIA-specific error
+ * and is returned as such, not treated as "try the other vendor".
  */
-export function detectGpus(execImpl = execFileSync) {
+function _probeNvidia(execImpl) {
   const result = {
+    vendor: 'nvidia',
     available: false,
     error: null,
     driverVersion: null,
@@ -21,7 +29,8 @@ export function detectGpus(execImpl = execFileSync) {
   try {
     smiOut = execImpl('nvidia-smi', [], { encoding: 'utf8', timeout: 5000 });
   } catch (err) {
-    result.error = err && err.code === 'ENOENT' ? 'nvidia-smi not found' : (err && err.message) || 'nvidia-smi failed';
+    if (err && err.code === 'ENOENT') return null;
+    result.error = (err && err.message) || 'nvidia-smi failed';
     return result;
   }
 
@@ -33,12 +42,20 @@ export function detectGpus(execImpl = execFileSync) {
   if (cudaMatch) result.cudaVersion = cudaMatch[1];
 
   try {
+    // compute_cap (e.g. "8.9" for Ada Lovelace, "9.0" for Hopper) is the
+    // real, stable NVIDIA architecture identifier -- the same value used in
+    // TORCH_CUDA_ARCH_LIST -- and is what byo_preflight.py's
+    // required_gpu_architecture check compares node_hardware.gpu_architecture
+    // against. Supported by nvidia-smi on any reasonably current driver;
+    // older drivers that reject the field fall into the catch below and
+    // simply leave architecture null per GPU, same as any other unsupported
+    // CSV field here.
     const csv = execImpl('nvidia-smi', [
-      '--query-gpu=index,name,memory.total,memory.free,memory.used,temperature.gpu,power.draw',
+      '--query-gpu=index,name,memory.total,memory.free,memory.used,temperature.gpu,power.draw,compute_cap',
       '--format=csv,noheader,nounits',
     ], { encoding: 'utf8', timeout: 5000 });
     result.gpus = csv.trim().split(/\r?\n/).filter(Boolean).map((line) => {
-      const [index, name, total, free, used, temp, power] = line.split(',').map((s) => s.trim());
+      const [index, name, total, free, used, temp, power, computeCap] = line.split(',').map((s) => s.trim());
       return {
         index: Number(index),
         name,
@@ -47,6 +64,7 @@ export function detectGpus(execImpl = execFileSync) {
         vramUsedGb: round1(Number(used) / 1024),
         temperatureC: temp && temp !== '[N/A]' ? Number(temp) : null,
         powerDrawW: power && power !== '[N/A]' ? Number(power) : null,
+        architecture: computeCap && computeCap !== '[N/A]' ? computeCap : null,
       };
     });
   } catch {
@@ -68,3 +86,124 @@ export function detectGpus(execImpl = execFileSync) {
 
   return result;
 }
+
+/**
+ * Read-only rocm-smi probe (AMD/ROCm — e.g. Radeon Pro / Instinct cards on
+ * gfx9xx/gfx10xx). Same "return null on ENOENT" contract as _probeNvidia so
+ * detectGpus can tell "not installed" apart from "installed but failing".
+ */
+function _probeRocm(execImpl) {
+  const result = {
+    vendor: 'amd',
+    available: false,
+    error: null,
+    driverVersion: null,
+    cudaVersion: null,
+    gpus: [],
+    processes: [],
+  };
+
+  let raw;
+  try {
+    raw = execImpl('rocm-smi', [
+      '--showproductname', '--showmeminfo', 'vram', '--showtemp', '--showpower',
+      '--showdriverversion', '--showtarget', '--json',
+    ], { encoding: 'utf8', timeout: 5000 });
+  } catch (err) {
+    if (err && err.code === 'ENOENT') return null;
+    result.error = (err && err.message) || 'rocm-smi failed';
+    return result;
+  }
+
+  result.available = true;
+
+  let parsed;
+  try {
+    parsed = JSON.parse(raw);
+  } catch {
+    // rocm-smi found but this build doesn't support --json (older ROCm) —
+    // still "available", just no per-GPU detail to report.
+    return result;
+  }
+
+  const cardKeys = Object.keys(parsed)
+    .filter((k) => /^card\d+$/.test(k))
+    .sort((a, b) => Number(a.slice(4)) - Number(b.slice(4)));
+
+  const findField = (card, ...needles) => {
+    const hit = Object.keys(card).find((k) => {
+      const lower = k.toLowerCase();
+      return needles.some((n) => lower.includes(n));
+    });
+    return hit ? card[hit] : undefined;
+  };
+
+  // The exact --showtarget key name has drifted across ROCm releases (e.g.
+  // "GPU Target Graphics Version" vs "Card Series" being reused, etc.), so
+  // key-name matching alone (as used for every other field above) isn't
+  // reliable here. A gfx architecture string (gfx900, gfx1030, gfx1100, ...)
+  // is distinctive and never appears as any other field's value, so scan
+  // every value on the card for that pattern regardless of which key it
+  // landed under -- robust to ROCm's key-naming drift by construction.
+  const _GFX_VALUE_RE = /\bgfx[0-9a-fA-F]{3,4}\b/;
+  const findArchitecture = (card) => {
+    for (const v of Object.values(card)) {
+      const m = String(v).match(_GFX_VALUE_RE);
+      if (m) return m[0].toLowerCase();
+    }
+    return null;
+  };
+
+  result.gpus = cardKeys.map((key, i) => {
+    const card = parsed[key] || {};
+    const totalBytes = findField(card, 'vram total memory');
+    const usedBytes = findField(card, 'vram total used memory');
+    const temp = findField(card, 'temperature');
+    const power = findField(card, 'power');
+    const totalGb = totalBytes != null ? bytesToGb(totalBytes) : null;
+    const usedGb = usedBytes != null ? bytesToGb(usedBytes) : null;
+    return {
+      index: i,
+      name: findField(card, 'card series', 'card model', 'product name') || 'AMD GPU',
+      vramTotalGb: totalGb,
+      vramFreeGb: totalGb != null && usedGb != null ? round1(totalGb - usedGb) : null,
+      vramUsedGb: usedGb,
+      temperatureC: temp != null && !Number.isNaN(parseFloat(temp)) ? Number(parseFloat(temp)) : null,
+      powerDrawW: power != null && !Number.isNaN(parseFloat(power)) ? Number(parseFloat(power)) : null,
+      architecture: findArchitecture(card),
+    };
+  });
+
+  const systemFields = parsed.system || {};
+  const driverKey = Object.keys(systemFields).find((k) => k.toLowerCase().includes('driver version'));
+  if (driverKey) result.driverVersion = systemFields[driverKey];
+  // rocm-smi has no fleet-wide "ROCm version" equivalent to nvidia-smi's CUDA
+  // Version header — a workload's own required HSA_OVERRIDE_GFX_VERSION /
+  // PYTORCH_ROCM_ARCH stays a launch-command concern, not a hardware fact
+  // this probe can report. cudaVersion stays null here on purpose.
+
+  return result;
+}
+
+/**
+ * Read-only GPU probe. Tries nvidia-smi first, falls back to rocm-smi
+ * (AMD/ROCm) when NVIDIA tooling isn't installed. Never mutates the
+ * machine. `vendor` is 'nvidia' | 'amd' | null (neither tool present).
+ */
+export function detectGpus(execImpl = execFileSync) {
+  const nvidia = _probeNvidia(execImpl);
+  if (nvidia) return nvidia;
+
+  const rocm = _probeRocm(execImpl);
+  if (rocm) return rocm;
+
+  return {
+    vendor: null,
+    available: false,
+    error: 'nvidia-smi and rocm-smi not found',
+    driverVersion: null,
+    cudaVersion: null,
+    gpus: [],
+    processes: [],
+  };
+}

package/src/nodeWorker/node_worker.py

@@ -0,0 +1,326 @@
+"""
+badgr-node — the BYO GPU (Phase 1) worker daemon.
+
+Runs on a customer's own machine (registered via `badgr node connect`) and
+implements the loop the locked spec requires: register -> heartbeat ->
+receive job -> execute -> stream status -> complete. Talks only to
+the backend's own /v1/nodes/{node_id}/* endpoints.
+
+Execution is plain `docker run`, with GPU passthrough flags chosen from a
+one-time local vendor probe (`_detect_gpu_vendor`): `--gpus all` (the
+standard NVIDIA Container Toolkit flag) on NVIDIA, or the ROCm container
+device flags (`--device=/dev/kfd --device=/dev/dri --group-add video
+--group-add render --security-opt seccomp=unconfined`) on AMD — no
+Slurm/Kubernetes, per this PR's locked scope.
+A "job" workload runs to completion and reports its real exit code. An
+"endpoint" workload (badgr serve) runs detached and is reported "running"
+only after a real readiness check (GET health_path) and, when a model was
+requested, a real minimal completion request (POST /v1/completions) both
+succeed against the container locally — the same two checks
+the backend's own readiness/inference verification logic makes, same URL
+paths and payload shape, just run from here instead of from the backend. That's not a stylistic choice: Badgr's
+backend cannot reach a customer's own machine directly (outbound-only,
+often behind NAT — see this PR's locked "no P2P networking" scope), so
+the backend trusts this worker's self-reported "running" for a
+customer_node target instead of re-probing an address it could never
+route to (see the backend's own deployment logic for the
+customer-node target case).
+"""
+import argparse
+import json
+import re
+import shutil
+import subprocess
+import sys
+import time
+import traceback
+import urllib.error
+import urllib.request
+
+POLL_INTERVAL_SECONDS = 3
+HEARTBEAT_INTERVAL_SECONDS = 20
+LOG_TAIL_LINES = 200
+READINESS_TIMEOUT_SECONDS = 300
+READINESS_POLL_SECONDS = 3
+# A plan's own max_runtime_seconds always wins; this is only the backstop
+# for a "job" workload that didn't set one, so a hung diagnostic can't wedge
+# the worker's single-job-at-a-time loop forever (see module docstring).
+DEFAULT_JOB_TIMEOUT_SECONDS = 900
+
+# A diagnostic script prints this on its own last line to declare a
+# structured pass/fail verdict distinct from the container's exit code
+# (see report_job_status's `result` field). Anything else is just a normal
+# job with no verdict to report.
+_RESULT_RE = re.compile(r"BADGR_RESULT=(PASS|FAIL)")
+# How often (in poll iterations) to log "still waiting on readiness" during
+# a long model load, so a tail -f doesn't look hung for minutes at a time.
+READINESS_LOG_EVERY_N_POLLS = 5
+
+
+def _detect_gpu_vendor() -> str | None:
+    """One-time, read-only vendor probe (mirrors the CLI's own
+    gpuDoctor/gpuInfo.js `detectGpus`, kept independent here since the
+    worker runs as a separate long-lived process and can't import JS).
+    `badgr node connect` already told the backend which vendor it saw at
+    connect time, but the worker itself needs a local answer too, to pick
+    the right `docker run` device flags per job — see module docstring.
+    NVIDIA is checked first only because that's the incumbent path; a
+    machine only ever has one or the other. Returns "nvidia", "amd", or
+    None (no GPU tooling found — jobs will still be attempted, docker run
+    just gets no device flags and will fail visibly instead of silently
+    picking a vendor that isn't there)."""
+    if shutil.which("nvidia-smi"):
+        return "nvidia"
+    if shutil.which("rocm-smi"):
+        return "amd"
+    return None
+
+
+def _gpu_docker_args(gpu_vendor: str | None) -> list[str]:
+    """GPU passthrough flags for `docker run`, per vendor. See
+    module docstring."""
+    if gpu_vendor == "amd":
+        return [
+            "--device=/dev/kfd", "--device=/dev/dri",
+            "--group-add", "video", "--group-add", "render",
+            "--security-opt", "seccomp=unconfined",
+        ]
+    # Default to the NVIDIA Container Toolkit flag -- matches this worker's
+    # pre-ROCm behavior when vendor detection itself fails to find either
+    # tool (rare: connect already required one to register the node).
+    return ["--gpus", "all"]
+
+
+def _log(level: str, msg: str) -> None:
+    ts = time.strftime("%Y-%m-%dT%H:%M:%S", time.gmtime())
+    stream = sys.stderr if level in ("WARN", "ERROR") else sys.stdout
+    print(f"{ts}Z [badgr-node] [{level}] {msg}", file=stream, flush=True)
+
+
+def _api(base_url: str, api_key: str, method: str, path: str, body: dict | None = None) -> dict:
+    url = f"{base_url}{path}"
+    data = json.dumps(body).encode("utf-8") if body is not None else None
+    req = urllib.request.Request(url, data=data, method=method, headers={
+        "Content-Type": "application/json",
+        "Authorization": f"Bearer {api_key}",
+    })
+    _log("DEBUG", f"-> {method} {path}" + (f" body={body}" if body else ""))
+    try:
+        with urllib.request.urlopen(req, timeout=15) as resp:
+            raw = resp.read()
+            _log("DEBUG", f"<- {method} {path} HTTP {resp.status}")
+            return json.loads(raw) if raw else {}
+    except urllib.error.HTTPError as e:
+        detail = e.read().decode("utf-8", "replace")
+        _log("ERROR", f"<- {method} {path} HTTP {e.code}: {detail}")
+        raise RuntimeError(f"{method} {path} -> HTTP {e.code}: {detail}") from e
+
+
+def _http_ok(method: str, url: str, body: dict | None = None, timeout: float = 5.0) -> tuple[bool, str]:
+    data = json.dumps(body).encode("utf-8") if body is not None else None
+    req = urllib.request.Request(url, data=data, method=method,
+                                  headers={"Content-Type": "application/json"} if data else {})
+    try:
+        with urllib.request.urlopen(req, timeout=timeout) as resp:
+            return resp.status == 200, ""
+    except urllib.error.HTTPError as e:
+        return False, f"HTTP {e.code}"
+    except Exception as exc:  # noqa: BLE001 — a probe failure is evidence ("not ready yet"), not a crash
+        return False, str(exc)
+
+
+def _verify_endpoint_locally(port: int, health_path: str, model: str | None) -> tuple[bool, str]:
+    """Same two checks the backend's own readiness/inference verification
+    logic makes (same paths, same minimal completion payload) — run locally
+    since this is the only place that can actually reach the container.
+    See module docstring."""
+    base = f"http://127.0.0.1:{port}"
+    _log("INFO", f"readiness: polling GET {base}{health_path} (timeout {READINESS_TIMEOUT_SECONDS}s)")
+    deadline = time.monotonic() + READINESS_TIMEOUT_SECONDS
+    ready, detail = False, "readiness_timeout"
+    attempt = 0
+    while time.monotonic() < deadline:
+        attempt += 1
+        ready, detail = _http_ok("GET", base + health_path)
+        if ready:
+            _log("INFO", f"readiness: OK after {attempt} attempt(s)")
+            break
+        if attempt % READINESS_LOG_EVERY_N_POLLS == 0:
+            _log("INFO", f"readiness: still waiting (attempt {attempt}, last result: {detail})")
+        time.sleep(READINESS_POLL_SECONDS)
+    if not ready:
+        _log("ERROR", f"readiness: failed after {attempt} attempt(s): {detail}")
+        return False, f"readiness check failed: {detail}"
+    if not model:
+        _log("INFO", "inference verification: skipped (no model on plan)")
+        return True, ""
+    _log("INFO", f"inference verification: POST {base}/v1/completions model={model}")
+    ok, detail = _http_ok("POST", base + "/v1/completions", {"model": model, "prompt": "hi", "max_tokens": 1}, timeout=30.0)
+    if not ok:
+        _log("ERROR", f"inference verification: failed: {detail}")
+        return False, f"inference verification failed: {detail}"
+    _log("INFO", "inference verification: OK — real completion request succeeded")
+    return True, ""
+
+
+def _parse_diagnostic_result(log_tail: str) -> str | None:
+    """Pulls the workload's own BADGR_RESULT= verdict out of its log, if it
+    printed one. Scans from the end since it's always the last thing a
+    diagnostic script emits (see the module's diagnostic scripts under
+    scripts/diagnostics/)."""
+    for line in reversed(log_tail.splitlines()):
+        m = _RESULT_RE.search(line)
+        if m:
+            return m.group(1)
+    return None
+
+
+def _run_job(base_url: str, api_key: str, node_id: str, job_id: str, plan: dict, gpu_vendor: str | None = None) -> None:
+    image = plan.get("image")
+    command = plan.get("command") or []
+    env = plan.get("env") or {}
+    port = plan.get("port")
+    is_endpoint = plan.get("workload_type") == "endpoint"
+    health_path = plan.get("health_path") or "/v1/models"
+    model = plan.get("model")
+    privileged = bool(plan.get("privileged"))
+    max_runtime_seconds = plan.get("max_runtime_seconds") or DEFAULT_JOB_TIMEOUT_SECONDS
+
+    container_name = f"badgr-job-{job_id}"
+    docker_cmd = ["docker", "run", *_gpu_docker_args(gpu_vendor), "--name", container_name]
+    docker_cmd += ["--rm"] if not is_endpoint else ["-d"]
+    if privileged:
+        # Bounded host diagnostics (NVLink/P2P/FabricManager/dmesg checks)
+        # need real host device + syslog access that a plain --gpus all
+        # container never gets -- see module docstring's "no P2P networking"
+        # note for why this still only ever talks to localhost/the host
+        # itself, never anything the backend can't already see logs from.
+        docker_cmd += ["--privileged", "--pid=host", "-v", "/dev:/dev", "-v", "/var/log:/var/log:ro"]
+    for k, v in env.items():
+        docker_cmd += ["-e", f"{k}={v}"]
+    if is_endpoint and port:
+        docker_cmd += ["-p", f"{port}:{port}"]
+    docker_cmd.append(image)
+    docker_cmd += [str(c) for c in command]
+
+    def _report(**kwargs) -> None:
+        _log("INFO", f"reporting status={kwargs.get('status')} for job {job_id}")
+        try:
+            _api(base_url, api_key, "POST", f"/nodes/{node_id}/jobs/{job_id}/status", kwargs)
+        except Exception as exc:  # noqa: BLE001 — best-effort status reporting must never crash the worker loop
+            _log("ERROR", f"status report failed for job {job_id}: {exc}")
+
+    _log("INFO", f"starting job {job_id} (workload_type={'endpoint' if is_endpoint else 'job'}, image={image})")
+    _log("DEBUG", f"job {job_id} docker command: {' '.join(docker_cmd)}")
+
+    if is_endpoint:
+        _run_endpoint_job(job_id, docker_cmd, port, health_path, model, _report)
+        return
+
+    proc = subprocess.Popen(docker_cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
+    _log("INFO", f"job {job_id} container started, pid={proc.pid}, streaming logs, "
+                 f"max_runtime_seconds={max_runtime_seconds}")
+    deadline = time.monotonic() + max_runtime_seconds
+    log_lines: list[str] = []
+    timed_out = False
+    while proc.poll() is None:
+        if time.monotonic() > deadline:
+            timed_out = True
+            _log("ERROR", f"job {job_id} exceeded max_runtime_seconds={max_runtime_seconds}, killing container")
+            # `docker kill` on the named container, not proc.pid -- pid is
+            # the `docker run` CLI process, killing it alone would leave the
+            # container itself running.
+            subprocess.run(["docker", "kill", container_name], capture_output=True)
+            proc.wait()
+            break
+        line = proc.stdout.readline() if proc.stdout else ""
+        if line:
+            line = line.rstrip("\n")
+            log_lines.append(line)
+            log_lines = log_lines[-LOG_TAIL_LINES:]
+            _log("DEBUG", f"job {job_id} | {line}")
+        else:
+            time.sleep(0.5)
+
+    tail = "\n".join(log_lines)
+    result = _parse_diagnostic_result(tail)
+    if timed_out:
+        _log("ERROR", f"job {job_id} timed out after {max_runtime_seconds}s")
+        _report(status="error", log_tail=tail, error=f"timed out after {max_runtime_seconds}s", result=result or "FAIL")
+        return
+
+    exit_code = proc.returncode
+    if exit_code == 0:
+        _log("INFO", f"job {job_id} exited 0" + (f", result={result}" if result else ""))
+        _report(status="exited", exit_code=exit_code, log_tail=tail, result=result)
+    else:
+        _log("ERROR", f"job {job_id} failed (exit {exit_code}); last {len(log_lines)} log line(s) attached to report")
+        _report(status="error", exit_code=exit_code, log_tail=tail, error=f"container exited {exit_code}", result=result or "FAIL")
+
+
+def _run_endpoint_job(job_id: str, docker_cmd: list, port: int | None, health_path: str, model: str | None, report) -> None:
+    """`badgr serve --target <node>`: launch detached, then block until the
+    workload is actually verified (readiness + real completion request) —
+    never report "running" on "container started" or "port open" alone."""
+    result = subprocess.run(docker_cmd, capture_output=True, text=True)
+    if result.returncode != 0:
+        _log("ERROR", f"job {job_id}: docker run failed (exit {result.returncode}): {result.stderr[-2000:]}")
+        report(status="error", error=f"docker run failed: {result.stderr[-2000:]}")
+        return
+    container_id = result.stdout.strip()
+    _log("INFO", f"job {job_id}: container started ({container_id[:12]}), port={port}, verifying before reporting running")
+    if not port:
+        _log("ERROR", f"job {job_id}: endpoint workload has no port to verify")
+        report(status="error", error="endpoint workload has no port to verify")
+        return
+
+    ok, detail = _verify_endpoint_locally(port, health_path, model)
+    if not ok:
+        _log("ERROR", f"job {job_id}: endpoint verification failed: {detail} — run `docker logs {container_id[:12]}` on this machine to debug")
+        report(status="error", error=detail)
+        return
+
+    _log("INFO", f"job {job_id}: endpoint verified — model={model or '(no model check)'} health_path={health_path}")
+    report(status="running", endpoint_url=f"http://127.0.0.1:{port}", log_tail=f"verified: {detail or 'ok'}")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="badgr-node — BYO GPU worker daemon")
+    parser.add_argument("--node-id", required=True, help="node_… id from `badgr node connect`")
+    parser.add_argument("--api-key", required=True, help="badgr CLI API key (same one `badgr login` stored)")
+    parser.add_argument("--base-url", default="https://aibadgr.com/v1")
+    args = parser.parse_args()
+
+    gpu_vendor = _detect_gpu_vendor()
+    _log("INFO", f"worker starting for node_id={args.node_id} against {args.base_url}")
+    _log("INFO", f"gpu_vendor={gpu_vendor or 'unknown'} (docker run device flags: {' '.join(_gpu_docker_args(gpu_vendor))})")
+    _log("INFO", f"heartbeat_interval={HEARTBEAT_INTERVAL_SECONDS}s poll_interval={POLL_INTERVAL_SECONDS}s")
+    last_heartbeat = 0.0
+    while True:
+        now = time.time()
+        if now - last_heartbeat >= HEARTBEAT_INTERVAL_SECONDS:
+            try:
+                _api(args.base_url, args.api_key, "POST", f"/nodes/{args.node_id}/heartbeat")
+                last_heartbeat = now
+                _log("DEBUG", "heartbeat OK")
+            except Exception as exc:  # noqa: BLE001 — a missed heartbeat must not crash the worker loop
+                _log("WARN", f"heartbeat failed (will retry next cycle): {exc}")
+
+        try:
+            job = _api(args.base_url, args.api_key, "GET", f"/nodes/{args.node_id}/jobs/next")
+        except Exception as exc:  # noqa: BLE001 — a missed poll must not crash the worker loop
+            _log("WARN", f"job poll failed (will retry next cycle): {exc}")
+            job = {}
+
+        if job.get("job_id"):
+            _log("INFO", f"received job {job['job_id']} from poll")
+            try:
+                _run_job(args.base_url, args.api_key, args.node_id, job["job_id"], job.get("plan") or {}, gpu_vendor)
+            except Exception:  # noqa: BLE001 — a bug in job handling must not silently kill the worker loop
+                _log("ERROR", f"unhandled exception while running job {job['job_id']}:\n{traceback.format_exc()}")
+        else:
+            time.sleep(POLL_INTERVAL_SECONDS)
+
+
+if __name__ == "__main__":
+    main()