[controllers] fix: defer keep allocation while busy

AGENTS.md

@@ -84,7 +84,7 @@ This file defines how coding agents should work in this repository.
 - Keep custom `job_id` validation centralized in `session_config.py`: only `None` means omitted/all-sessions, and custom IDs must be non-empty URL-path-safe strings before any session state changes.
 - Stop requests must not miss starting sessions; wait for startup to settle before returning `not found` or taking a stop-all snapshot.
 - Stop-all may release independent sessions concurrently, but must not duplicate release work for `stopping` sessions and must keep deterministic additive result fields.
-- Keep utilization backoff eco-safe: valid `busy_threshold` values are `-1` or `0..100`; public defaults must use the shared `DEFAULT_BUSY_THRESHOLD` (`25`), and when telemetry is unavailable with `busy_threshold >= 0`, controllers should sleep instead of running keepalive compute. Only `busy_threshold=-1` is the explicit unconditional mode.
+- Keep utilization backoff eco-safe: valid `busy_threshold` values are `-1` or `0..100`; public defaults must use the shared `DEFAULT_BUSY_THRESHOLD` (`25`), and when telemetry is unavailable with `busy_threshold >= 0`, controllers should sleep instead of allocating keep tensors or running keepalive compute. Only `busy_threshold=-1` is the explicit unconditional mode.
 - Treat public `gpu_ids` as visible device ordinals after user-supplied CUDA or ROCm visibility filtering. Do not rewrite visibility masks inside KeepGPU command paths; reject explicit CUDA/ROCm ordinals outside the current visible device count before starting keep workers.
 - Keep CUDA telemetry aligned with visible CUDA ordinals: `get_gpu_utilization(index)` receives the visible rank used by `CudaGPUController`, and `gpu_monitor.py` resolves `CUDA_VISIBLE_DEVICES` numeric/UUID tokens to the correct NVML handle. If that mapping is ambiguous or unsupported, return `None` rather than falling back to a possibly wrong physical index.
 - Keep ROCm telemetry aligned with visible ROCm ordinals: resolve `ROCR_VISIBLE_DEVICES` as the base mask and one matching `HIP_VISIBLE_DEVICES`/`CUDA_VISIBLE_DEVICES` overlay before querying ROCm SMI. If the mapping is malformed, conflicting, unsupported, or out of range, return unavailable utilization rather than querying a guessed SMI index.

docs/concepts/architecture.md

@@ -37,9 +37,11 @@ CLI args ──▶ GlobalGPUController ──▶ [CudaGPUController rank=0]
    If a later worker fails to start, already-started workers are released before
    the original start error is re-raised.
 3. Each CUDA worker:
-   - Allocates a tensor sized by way of `vram_to_keep`.
    - Starts a daemon thread that performs intervalled lightweight elementwise batches.
-   - Calls `_monitor_utilization` (by way of NVML) to detect real activity.
+   - Calls `_monitor_utilization` (by way of NVML) to detect real activity
+     before allocating the keep tensor.
+   - Allocates a tensor sized by way of `vram_to_keep` only when backoff allows
+     work.
      The monitor receives the CUDA visible rank and resolves
      `CUDA_VISIBLE_DEVICES` numeric or UUID tokens before querying NVML, so
      telemetry follows the same device the worker keeps.
@@ -50,9 +52,10 @@ CLI args ──▶ GlobalGPUController ──▶ [CudaGPUController rank=0]
    unavailable for that cycle.
 5. If utilization exceeds `busy_threshold`, or if utilization is unavailable
    while `busy_threshold` is non-negative, the worker just sleeps for one more
-   `interval`. Otherwise it runs a new batch of ops. Public defaults use
-   `busy_threshold=25`. Valid thresholds are `-1` or `0..100`;
-   `busy_threshold=-1` is the explicit unconditional mode.
+   `interval` before allocating or running ops. Otherwise it allocates the keep
+   tensor when needed and runs a batch. Public defaults use `busy_threshold=25`.
+   Valid thresholds are `-1` or `0..100`; `busy_threshold=-1` is the explicit
+   unconditional mode.
 6. When you call `release()` (or exit the context), every worker sets a stop
    event, joins the thread, and clears the device cache. Release attempts every
    worker and then raises a summary if any worker failed to stop.

docs/guides/cli.md

@@ -124,4 +124,4 @@ of `keep-gpu status`.
 - **Start cannot reach service**: run `keep-gpu serve --host 127.0.0.1 --port 8765`.
 - **Need to close background service**: run `keep-gpu stop --all` first, then `keep-gpu service-stop`. Use `keep-gpu service-stop --force` only for an unresponsive auto-started daemon; it still refuses to signal a PID that KeepGPU cannot verify as its own.
 - **OOM during keep**: reduce `--vram` or free GPU memory before starting.
-- **No utilization data**: on CUDA, ensure `nvidia-ml-py` works and `nvidia-smi` is available; on ROCm, check the optional `rocm-smi` extra and avoid conflicting `HIP_VISIBLE_DEVICES`/`CUDA_VISIBLE_DEVICES` masks. ROCm telemetry resolves visible ranks through `ROCR_VISIBLE_DEVICES` plus one HIP/CUDA overlay and reports unavailable utilization rather than guessing when the mapping is malformed or ambiguous. Valid `busy_threshold` values are `-1` or `0..100`, and omitted CLI values default to `25`. With non-negative `busy_threshold`, KeepGPU sleeps when utilization is unavailable. On Mac M series, utilization is expected to be `null`, so use `--busy-threshold -1` only when you intentionally want unconditional keepalive compute.
+- **No utilization data**: on CUDA, ensure `nvidia-ml-py` works and `nvidia-smi` is available; on ROCm, check the optional `rocm-smi` extra and avoid conflicting `HIP_VISIBLE_DEVICES`/`CUDA_VISIBLE_DEVICES` masks. ROCm telemetry resolves visible ranks through `ROCR_VISIBLE_DEVICES` plus one HIP/CUDA overlay and reports unavailable utilization rather than guessing when the mapping is malformed or ambiguous. Valid `busy_threshold` values are `-1` or `0..100`, and omitted CLI values default to `25`. With non-negative `busy_threshold`, KeepGPU sleeps before allocation or compute when utilization is unavailable. On Mac M series, utilization is expected to be `null`, so use `--busy-threshold -1` only when you intentionally want unconditional keepalive compute.

docs/guides/mcp.md

@@ -162,7 +162,8 @@ available. Mac M series devices report best-effort MPS memory counters and use
 `null` for unsupported fields such as utilization.
 Valid `busy_threshold` values are `-1` or `0..100`, and omitted API values
 default to `25`. When utilization is unavailable and `busy_threshold` is
-non-negative, controllers sleep instead of running keepalive compute;
+non-negative, controllers sleep instead of allocating keep tensors or running
+keepalive compute;
 `busy_threshold=-1` is the explicit unconditional mode.
 Stop controls show timed-out or failed releases instead of claiming success when
 the backend keeps a session visible for follow-up cleanup. Retained session cards

docs/guides/python.md

@@ -66,8 +66,8 @@ with GlobalGPUController(
 - `busy_threshold` defaults to `25` and accepts `-1` or a percentage in
   `0..100`. Non-negative thresholds throttle the keep-alive loop when
   utilization spikes. When utilization telemetry is unavailable, non-negative
-  thresholds sleep instead of running compute; use `busy_threshold=-1` only for
-  explicit unconditional keepalive work.
+  thresholds sleep before allocating the keep tensor or running compute; use
+  `busy_threshold=-1` only for explicit unconditional keepalive work.
 - `release()` uses threads too, so all GPUs free up quickly.
 
 ## Combine with schedulers or callbacks

docs/plans/defer-allocation-while-busy.md

@@ -0,0 +1,69 @@
+# Defer Allocation While Busy Plan
+
+## Background
+
+Utilization backoff currently gates only the keep-alive compute batch. CUDA,
+ROCm, and MPS controllers still allocate and initialize their keep tensor before
+checking whether telemetry says the device is busy or unavailable. That startup
+allocation can reserve VRAM and run a random-fill kernel even when the default
+eco-safe threshold should make the controller sleep.
+
+## Goal
+
+Check backoff before initial keep tensor allocation, so non-negative
+`busy_threshold` values avoid allocation and GPU work while the device is busy
+or telemetry is unavailable. `busy_threshold=-1` remains the explicit
+unconditional allocation/compute mode.
+
+## Solution
+
+- Add RED no-GPU tests proving CUDA, ROCm, and MPS do not call `torch.rand`
+  before backoff allows allocation.
+- Add no-GPU positive-path tests proving CUDA/ROCm allocate after telemetry
+  becomes idle, and that `busy_threshold=-1` still allocates unconditionally.
+- Add ROCm coverage proving busy deferrals do not consume allocation retries.
+- Move each controller's initial allocation attempt behind its existing
+  `_should_run_batch()` decision.
+- Keep retry behavior simple: if telemetry says busy/unavailable, wait one
+  interval and try the backoff decision again.
+- Update architecture and user docs to state that startup allocation is also
+  deferred by non-negative utilization backoff.
+
+## Tasks
+
+- [x] Add RED allocation-before-backoff tests for CUDA, ROCm, and MPS.
+- [x] Implement allocation deferral in each controller.
+- [x] Add positive-path coverage for idle-after-busy and unconditional
+      allocation modes.
+- [x] Add ROCm retry accounting coverage for busy deferrals.
+- [x] Update `AGENTS.md`, architecture/docs, and this plan.
+- [x] Run targeted tests, full tests, docs build, pre-commit, and local
+      subagent review before PR.
+- [ ] Open PR, resolve review comments, wait for clean checks, squash merge, and
+      clean the worktree.
+
+## Verification
+
+Baseline:
+
+- `PYTHONPATH=$PWD/src pytest tests/cuda_controller/test_throttle.py tests/rocm_controller/test_rocm_backoff.py tests/macm_controller/test_macm_backoff.py tests/global_controller/global_keep_test.py tests/single_gpu_controller/test_release_contract.py -q`:
+  `18 passed, 2 skipped`.
+
+Completed:
+
+- RED focused regression:
+  `PYTHONPATH=$PWD/src pytest tests/cuda_controller/test_throttle.py::test_cuda_busy_utilization_defers_initial_allocation tests/rocm_controller/test_rocm_backoff.py::test_rocm_busy_utilization_defers_initial_allocation tests/macm_controller/test_macm_backoff.py::test_macm_unavailable_utilization_defers_initial_allocation -q`
+  failed with all three tests hitting `allocation should wait for idle telemetry`.
+- GREEN focused regression: same command, `3 passed`.
+- `PYTHONPATH=$PWD/src pytest tests/cuda_controller/test_throttle.py tests/rocm_controller/test_rocm_backoff.py tests/macm_controller/test_macm_backoff.py -q`:
+  `17 passed, 1 skipped`.
+- `PYTHONPATH=$PWD/src pytest tests/cuda_controller/test_throttle.py tests/rocm_controller/test_rocm_backoff.py tests/macm_controller/test_macm_backoff.py tests/global_controller/global_keep_test.py tests/single_gpu_controller/test_release_contract.py -q`:
+  `27 passed, 2 skipped`.
+- `PYTHONPATH=$PWD/src pytest tests -q`: `263 passed, 11 skipped`.
+- `PYTHONPATH=$PWD/src mkdocs build`: passed. Existing Material for MkDocs
+  version warning and docs-nav notices were emitted.
+- `pre-commit run --all-files`: passed.
+- `git diff --check`: passed.
+- Local subagent code review: passed. Follow-up review confirmed the CLI docs
+  parity fix and ROCm retry regression, with no remaining Critical or Important
+  blockers.

docs/reference/api.md

@@ -18,8 +18,9 @@ querying ROCm SMI. Unresolved mappings report unavailable utilization instead
 of falling back to a possibly wrong physical device.
 
 Public controller, CLI, REST, JSON-RPC, and MCP defaults use
-`busy_threshold=25`. Pass `busy_threshold=-1` only when you intentionally want
-unconditional keepalive compute without utilization backoff.
+`busy_threshold=25`, so busy or unavailable telemetry sleeps before allocating
+keep tensors or running compute. Pass `busy_threshold=-1` only when you
+intentionally want unconditional keepalive compute without utilization backoff.
 
 For service session IDs, `job_id=None` is the only omitted/all-sessions
 sentinel. Custom IDs must be non-empty strings containing only letters, digits,

docs/reference/cli.md

@@ -23,7 +23,7 @@ These options apply when you run `keep-gpu` without subcommands.
 | `--interval INTEGER` | seconds | Finite positive sleep duration between utilization checks and keep-alive batches. |
 | `--gpu-ids TEXT` | comma-separated unique non-negative ints | Subset of visible device ordinals to guard (for example, `0,2`). Omit to use all visible GPUs; startup fails if that resolves to none or if an explicit ordinal is out of range. |
 | `--vram TEXT` | human size or bare bytes | Amount of memory each GPU controller allocates (`512MB`, `1GiB`, `1073741824`). |
-| `--busy-threshold INTEGER` / `--util-threshold INTEGER` | percent | `0..100` backs off when utilization is above this value or unavailable; `-1` disables utilization backoff. |
+| `--busy-threshold INTEGER` / `--util-threshold INTEGER` | percent | `0..100` backs off before allocation/compute when utilization is above this value or unavailable; `-1` disables utilization backoff. |
 | `--threshold TEXT` | deprecated | Legacy alias: numeric values map to busy-threshold, size strings map to vram. |
 
 ## Service mode

src/keep_gpu/single_gpu_controller/cuda_gpu_controller.py

@@ -171,6 +171,16 @@ def _keep_loop(self) -> None:
         matrix = None
         while not stop_evt.is_set():
             try:
+                gpu_utilization = self._monitor_utilization(self.rank)
+                if not self._should_run_batch(gpu_utilization, self.busy_threshold):
+                    logger.debug(
+                        "rank %s: GPU utilization unavailable or busy (%s), deferring allocation",
+                        self.rank,
+                        "n/a" if gpu_utilization is None else f"{gpu_utilization}%",
+                    )
+                    if stop_evt.wait(self.interval):
+                        return
+                    continue
                 matrix = torch.rand(
                     num_elements,
                     device=self.device,

src/keep_gpu/single_gpu_controller/macm_gpu_controller.py

@@ -111,6 +111,15 @@ def _keep_loop(self) -> None:
         tensor = None
         while not stop_evt.is_set():
             try:
+                if not self._should_run_batch(None, self.busy_threshold):
+                    logger.debug(
+                        "rank %s: MPS utilization unavailable; deferring allocation because busy_threshold=%s",
+                        self.rank,
+                        self.busy_threshold,
+                    )
+                    if stop_evt.wait(self.interval):
+                        return
+                    continue
                 tensor = torch.rand(
                     num_elements,
                     device=self.device,

src/keep_gpu/single_gpu_controller/rocm_gpu_controller.py

@@ -143,6 +143,16 @@ def _keep_loop(self) -> None:
             return
         while not stop_evt.is_set():
             try:
+                util = self._query_utilization()
+                if not self._should_run_batch(util, self.busy_threshold):
+                    logger.debug(
+                        "rank %s: GPU utilization unavailable or busy (%s), deferring allocation",
+                        self.rank,
+                        "n/a" if util is None else f"{util}%",
+                    )
+                    if stop_evt.wait(self.interval):
+                        return
+                    continue
                 tensor = torch.rand(
                     num_elements,
                     device=self.device,

tests/cuda_controller/test_throttle.py

@@ -6,6 +6,36 @@
 from keep_gpu.single_gpu_controller.cuda_gpu_controller import CudaGPUController
 
 
+class _StopAfterOneWait:
+    def __init__(self):
+        self.stopped = False
+        self.wait_calls = 0
+
+    def is_set(self):
+        return self.stopped
+
+    def wait(self, _timeout):
+        self.wait_calls += 1
+        self.stopped = True
+        return True
+
+
+class _StopAfterWaits:
+    def __init__(self, stop_after):
+        self.stop_after = stop_after
+        self.stopped = False
+        self.wait_calls = 0
+
+    def is_set(self):
+        return self.stopped
+
+    def wait(self, _timeout):
+        self.wait_calls += 1
+        if self.wait_calls >= self.stop_after:
+            self.stopped = True
+        return self.stopped
+
+
 def test_negative_busy_threshold_disables_backoff_without_gpu():
     assert CudaGPUController._should_run_batch(0, -1) is True
     assert CudaGPUController._should_run_batch(100, -1) is True
@@ -34,6 +64,94 @@ def test_cuda_controller_rejects_busy_threshold_below_minus_one_without_gpu():
         CudaGPUController(rank=0, vram_to_keep="4MB", busy_threshold=-2)
 
 
+def test_cuda_busy_utilization_defers_initial_allocation(monkeypatch):
+    import keep_gpu.single_gpu_controller.cuda_gpu_controller as cuda_module
+
+    ctrl = CudaGPUController.__new__(CudaGPUController)
+    ctrl.rank = 0
+    ctrl.device = "cuda:0"
+    ctrl.interval = 0.01
+    ctrl.busy_threshold = 10
+    ctrl.relu_iterations = 1
+    ctrl._num_elements = 4
+    ctrl._stop_evt = _StopAfterOneWait()
+
+    monkeypatch.setattr(cuda_module.torch.cuda, "set_device", lambda _rank: None)
+    monkeypatch.setattr(
+        cuda_module.torch,
+        "rand",
+        lambda *args, **kwargs: (_ for _ in ()).throw(
+            AssertionError("allocation should wait for idle telemetry")
+        ),
+    )
+    monkeypatch.setattr(ctrl, "_monitor_utilization", lambda _rank: 100)
+
+    ctrl._keep_loop()
+
+    assert ctrl._stop_evt.wait_calls == 1
+
+
+def test_cuda_busy_then_idle_utilization_allows_initial_allocation(monkeypatch):
+    import keep_gpu.single_gpu_controller.cuda_gpu_controller as cuda_module
+
+    ctrl = CudaGPUController.__new__(CudaGPUController)
+    ctrl.rank = 0
+    ctrl.device = "cuda:0"
+    ctrl.interval = 0.01
+    ctrl.busy_threshold = 10
+    ctrl.relu_iterations = 1
+    ctrl._num_elements = 4
+    ctrl._stop_evt = _StopAfterWaits(stop_after=2)
+
+    allocations = []
+    batches = []
+    utilization = iter([100, 0, 0])
+
+    monkeypatch.setattr(cuda_module.torch.cuda, "set_device", lambda _rank: None)
+    monkeypatch.setattr(
+        cuda_module.torch,
+        "rand",
+        lambda *args, **kwargs: allocations.append((args, kwargs)) or object(),
+    )
+    monkeypatch.setattr(ctrl, "_monitor_utilization", lambda _rank: next(utilization))
+    monkeypatch.setattr(ctrl, "_run_relu_batch", lambda matrix: batches.append(matrix))
+
+    ctrl._keep_loop()
+
+    assert len(allocations) == 1
+    assert len(batches) == 1
+    assert ctrl._stop_evt.wait_calls == 2
+
+
+def test_cuda_negative_busy_threshold_allocates_even_when_busy(monkeypatch):
+    import keep_gpu.single_gpu_controller.cuda_gpu_controller as cuda_module
+
+    ctrl = CudaGPUController.__new__(CudaGPUController)
+    ctrl.rank = 0
+    ctrl.device = "cuda:0"
+    ctrl.interval = 0.01
+    ctrl.busy_threshold = -1
+    ctrl.relu_iterations = 1
+    ctrl._num_elements = 4
+    ctrl._stop_evt = _StopAfterWaits(stop_after=1)
+
+    allocations = []
+
+    monkeypatch.setattr(cuda_module.torch.cuda, "set_device", lambda _rank: None)
+    monkeypatch.setattr(
+        cuda_module.torch,
+        "rand",
+        lambda *args, **kwargs: allocations.append((args, kwargs)) or object(),
+    )
+    monkeypatch.setattr(ctrl, "_monitor_utilization", lambda _rank: 100)
+    monkeypatch.setattr(ctrl, "_run_relu_batch", lambda _matrix: None)
+
+    ctrl._keep_loop()
+
+    assert len(allocations) == 1
+    assert ctrl._stop_evt.wait_calls == 1
+
+
 @pytest.mark.skipif(
     not torch.cuda.is_available(),
     reason="Only run CUDA tests when CUDA is available",