optimizer: keep the probe's margin when the whole job fits

docs/config.md

@@ -223,7 +223,7 @@ technical:
 |---|---|---|
 | `device` | `"cuda"` | `"cpu"`, `"cuda"`, or `"cuda:N"` (e.g. `"cuda:0"`). Falls back to `cuda:0` if the requested index doesn't exist |
 | `max_memory_padding` | `0.95` | Fraction of GPU memory the autobatcher is allowed to fill during calibration. Lower = more headroom, smaller batches |
-| `memory_scaling_factor` | `1.75` | Factor by which the autobatcher grows the probe size during calibration. Larger = faster calibration but coarser final batch size; smaller = slower but tighter. Must be > 1 |
+| `memory_scaling_factor` | `1.75` | Factor by which the autobatcher grows the probe size during calibration. Larger = faster calibration but coarser final batch size; smaller = slower but tighter. A job runs as one batch without probing only if its forward-pass peak times this factor fits in `max_memory_padding` of free memory. Must be > 1 |
 | `max_atoms_to_try` | `100000` | Upper bound on the autobatcher's calibration probe size (atoms) |
 | `steps_between_swaps` | `1` | Optimization steps between batch swaps in the in-flight autobatcher. `1` is fastest on this codebase's screenings (uma-s/uma-m/orb); higher values are monotonically slower |
 | `logging_level` | `"INFO"` | Logging verbosity: `"DEBUG"`, `"INFO"`, `"WARNING"`, `"ERROR"` |

src/gpuma/optimizer.py

@@ -377,9 +377,11 @@ def _whole_job_memory_scaler(state: Any, model: Any, batcher: Any) -> float | No
 
     free, _ = torch.cuda.mem_get_info()
     budget_gb = (free + torch.cuda.memory_reserved()) / 1024**3
-    if peak_gb > batcher.max_memory_padding * budget_gb:
+    # Keep the probe's margin: its batch sits one growth step below a size that fit.
+    if peak_gb * batcher.memory_scaling_factor > batcher.max_memory_padding * budget_gb:
         logger.info(
-            "Job peaks at %.2f of %.2f GB, over max_memory_padding; probing GPU memory",
+            "Job peaks at %.2f of %.2f GB, too close to the limit for one batch; "
+            "probing GPU memory",
             peak_gb,
             budget_gb,
         )

tests/optimization/test_batch.py

@@ -390,6 +390,24 @@ def oom(state, model, real):
         assert all(r is not None for r in results)
         assert probes
 
+    def test_peak_within_one_growth_step_falls_back_to_probe(
+        self, monkeypatch, methane, ethanol
+    ):
+        """Used to fail: a job at 0.8 of the budget ran as one batch and OOMed mid-run."""
+        import torch
+
+        def near_limit(state, model, real):
+            real(state, model)
+            free, _ = torch.cuda.mem_get_info()
+            return 0.8 * (free + torch.cuda.memory_reserved()) / 1024**3
+
+        jobs = [methane, ethanol, self._largest()]
+        self._patch_whole_job_forward(monkeypatch, jobs, near_limit)
+        probes = self._count_probes(monkeypatch)
+        results = optimize_structure_batch(jobs, self._config())
+        assert all(r is not None for r in results)
+        assert probes
+
     def test_insufficient_headroom_falls_back_to_probe(self, monkeypatch, methane, ethanol):
         def huge_peak(state, model, real):
             real(state, model)