badgr-cli 1.1.6 → 1.1.7

package/src/nodeWorker/test_node_worker.py

@@ -0,0 +1,188 @@
+"""
+Unit tests for node_worker.py's three additions to the customer_node BYO GPU
+worker (see backend/customer_node_adapter.py / backend/node_routes.py for
+the control-plane side of the same feature):
+  - privileged/host diagnostic execution (docker flags only, no live docker)
+  - hard timeout enforcement in the job-streaming loop
+  - structured BADGR_RESULT= parsing
+
+No real docker/network calls are made -- subprocess.Popen/run are stubbed.
+"""
+import os
+import sys
+import time
+
+sys.path.insert(0, os.path.dirname(__file__))
+
+import node_worker  # noqa: E402
+
+
+def test_parse_diagnostic_result_pass():
+    log = "=== RESULT ===\nBADGR_RESULT=PASS\n"
+    assert node_worker._parse_diagnostic_result(log) == "PASS"
+
+
+def test_parse_diagnostic_result_fail():
+    log = "some noise\nBADGR_RESULT=FAIL\nmore noise after\n"
+    assert node_worker._parse_diagnostic_result(log) == "FAIL"
+
+
+def test_parse_diagnostic_result_absent_for_ordinary_job():
+    assert node_worker._parse_diagnostic_result("epoch 1/1 done\n") is None
+
+
+def test_privileged_plan_adds_host_diagnostic_docker_flags(monkeypatch):
+    captured_cmds = []
+
+    class _FakeProc:
+        pid = 4242
+        returncode = 0
+        stdout = None
+
+        def poll(self):
+            return 0
+
+    def fake_popen(cmd, **kwargs):
+        captured_cmds.append(cmd)
+        return _FakeProc()
+
+    def fake_run(cmd, **kwargs):
+        captured_cmds.append(cmd)
+
+        class _R:
+            returncode = 0
+        return _R()
+
+    reports = []
+    monkeypatch.setattr(node_worker.subprocess, "Popen", fake_popen)
+    monkeypatch.setattr(node_worker.subprocess, "run", fake_run)
+    monkeypatch.setattr(node_worker, "_api", lambda *a, **kw: reports.append(kw) or {})
+
+    node_worker._run_job(
+        "https://x/v1", "key", "node_1", "job_1",
+        {"image": "badgr/diagnostics:latest", "command": ["./nvlink_p2p_check.sh"], "privileged": True},
+    )
+
+    docker_cmd = captured_cmds[0]
+    assert "--privileged" in docker_cmd
+    assert "--pid=host" in docker_cmd
+    assert "/dev:/dev" in docker_cmd
+
+
+def test_ordinary_job_has_no_privileged_flags(monkeypatch):
+    captured_cmds = []
+
+    class _FakeProc:
+        pid = 4242
+        returncode = 0
+        stdout = None
+
+        def poll(self):
+            return 0
+
+    monkeypatch.setattr(node_worker.subprocess, "Popen", lambda cmd, **kw: captured_cmds.append(cmd) or _FakeProc())
+    monkeypatch.setattr(node_worker, "_api", lambda *a, **kw: {})
+
+    node_worker._run_job(
+        "https://x/v1", "key", "node_1", "job_2",
+        {"image": "python:3.11-slim", "command": ["python", "train.py"]},
+    )
+
+    assert "--privileged" not in captured_cmds[0]
+
+
+def test_nvidia_job_uses_gpus_all_flag_by_default(monkeypatch):
+    captured_cmds = []
+
+    class _FakeProc:
+        pid = 4242
+        returncode = 0
+        stdout = None
+
+        def poll(self):
+            return 0
+
+    monkeypatch.setattr(node_worker.subprocess, "Popen", lambda cmd, **kw: captured_cmds.append(cmd) or _FakeProc())
+    monkeypatch.setattr(node_worker, "_api", lambda *a, **kw: {})
+
+    node_worker._run_job(
+        "https://x/v1", "key", "node_1", "job_4",
+        {"image": "python:3.11-slim", "command": ["python", "train.py"]},
+        "nvidia",
+    )
+
+    docker_cmd = captured_cmds[0]
+    assert "--gpus" in docker_cmd and "all" in docker_cmd
+    assert "--device=/dev/kfd" not in docker_cmd
+
+
+def test_amd_job_uses_rocm_device_flags(monkeypatch):
+    captured_cmds = []
+
+    class _FakeProc:
+        pid = 4242
+        returncode = 0
+        stdout = None
+
+        def poll(self):
+            return 0
+
+    monkeypatch.setattr(node_worker.subprocess, "Popen", lambda cmd, **kw: captured_cmds.append(cmd) or _FakeProc())
+    monkeypatch.setattr(node_worker, "_api", lambda *a, **kw: {})
+
+    node_worker._run_job(
+        "https://x/v1", "key", "node_1", "job_5",
+        {"image": "rocm/vllm:latest", "command": ["python", "train.py"]},
+        "amd",
+    )
+
+    docker_cmd = captured_cmds[0]
+    assert "--device=/dev/kfd" in docker_cmd
+    assert "--device=/dev/dri" in docker_cmd
+    assert "--gpus" not in docker_cmd
+
+
+def test_gpu_docker_args_defaults_to_nvidia_flag_when_vendor_unknown():
+    assert node_worker._gpu_docker_args(None) == ["--gpus", "all"]
+
+
+def test_job_exceeding_max_runtime_is_killed_and_reported(monkeypatch):
+    """A hung diagnostic must not wedge the worker forever -- the loop kills
+    the named container once max_runtime_seconds elapses and reports a
+    FAIL result rather than hanging indefinitely."""
+    kill_calls = []
+
+    class _FakeProc:
+        pid = 4242
+        returncode = None
+        stdout = None
+
+        def poll(self):
+            return None  # never exits on its own -- simulates a hang
+
+        def wait(self):
+            self.returncode = -9
+
+    fake_proc = _FakeProc()
+    monkeypatch.setattr(node_worker.subprocess, "Popen", lambda cmd, **kw: fake_proc)
+    monkeypatch.setattr(node_worker.subprocess, "run", lambda cmd, **kw: kill_calls.append(cmd))
+
+    # Force the deadline to already be in the past on the very first loop
+    # check, instead of sleeping in the test.
+    clock = iter([1000.0, 1000.0, 2000.0])
+    monkeypatch.setattr(node_worker.time, "monotonic", lambda: next(clock))
+
+    reports = []
+    monkeypatch.setattr(node_worker, "_api", lambda *a, **kw: reports.append(a[-1]) or {})
+
+    node_worker._run_job(
+        "https://x/v1", "key", "node_1", "job_3",
+        {"image": "badgr/diagnostics:latest", "privileged": True, "max_runtime_seconds": 5},
+    )
+
+    assert any("badgr-job-job_3" in cmd for cmd in kill_calls)
+    assert reports, "worker must report a status even on timeout"
+    report = reports[0]
+    assert report["status"] == "error"
+    assert "timed out" in report["error"]
+    assert report["result"] == "FAIL"

package/src/progress.js

@@ -101,7 +101,7 @@ export async function pollJobUntilTerminal(callApi, config, jobId, { chalk, maxM
 
     // One-time announcement of the automatic retry-on-a-different-route —
     // printed once per retry (not every poll tick) so it reads as an event,
-    // not a repeated status line. See job_progress.py's retrying_different_route
+    // not a repeated status line. See the backend's own retrying_different_route
     // stage — this is what "one safe retry" looks like from the CLI.
     if (detail.stage === 'retrying_different_route' && !announcedRetry) {
       announcedRetry = true;
@@ -148,7 +148,7 @@ function _formatElapsed(seconds) {
  * Prints the `Class:`/`Next:` lines for a failed GpuDeployment (the raw
  * /deployments/{id} path used by `badgr run`/`badgr serve`/`badgr comfyui run`)
  * using the exact same failure_class/next_action fields — computed by the
- * same backend/job_progress.py classify_failure()/next_step_for() functions —
+ * same backend classify_failure()/next_step_for() functions —
  * that renderJobClosingBlock below already prints for the productized
  * /v1/jobs path (comfy.batch, train.lora, custom.run, model.serve). One
  * failure-class vocabulary shown the same way regardless of which API path

package/src/spec.js

@@ -1,6 +1,6 @@
 /**
  * Deployment spec: the single document that describes what `gpu up` should provision.
- * Mirrors the concepts in overflow_dispatch.py (gpu_type, region, workload type).
+ * Mirrors the backend's own deployment-dispatch concepts (gpu_type, region, workload type).
  */
 
 export const WORKLOAD_TYPES = ['endpoint', 'job'];
@@ -26,7 +26,7 @@ export function vmClassForWorkload(workloadName) {
   return workloadName === 'playwright' ? 'browser' : 'small';
 }
 
-// Match canonical GPU IDs used in overflow_providers.py GPU_ALIASES
+// Match canonical GPU IDs used in the backend's own GPU_ALIASES
 export const GPU_TYPE_MAP = {
   'rtx-4090': 'RTX_4090',  'rtx4090': 'RTX_4090',  '4090': 'RTX_4090',
   'rtx-3090': 'RTX_3090',  'rtx3090': 'RTX_3090',  '3090': 'RTX_3090',