log memory at the failing block instead of probing

app.py

@@ -146,6 +146,13 @@ def log_worker_tracebacks():
                 return inner(self, *a, **k)
             except Exception:
                 import traceback
+                free, total = torch.cuda.mem_get_info()
+                print(f"[fail] {name}: allocated="
+                      f"{torch.cuda.memory_allocated() / 2**30:.1f}G peak="
+                      f"{torch.cuda.max_memory_allocated() / 2**30:.1f}G reserved="
+                      f"{torch.cuda.memory_reserved() / 2**30:.1f}G "
+                      f"device_free={free / 2**30:.1f}G/{total / 2**30:.1f}G",
+                      flush=True)
                 traceback.print_exc()
                 raise
         setattr(cls, name, outer)
@@ -207,35 +214,6 @@ def write_mp4(frames):
     return path
 
 
-def _probe(tag):
-    """How much VRAM is actually reachable inside the worker, and up to what size?
-
-    The allocation that aborts the run is only ~49 MB, so either something rejects
-    allocations regardless of size or the device has far less headroom than the
-    48 GB a ZeroGPU `large` slice is supposed to give.
-    """
-    import threading
-    free, total = torch.cuda.mem_get_info()
-    got = []
-
-    def work():
-        for mb in (64, 256, 1024, 4096, 8192):
-            try:
-                x = torch.empty(mb * 1024 * 256, dtype=torch.float32, device="cuda")
-                del x
-                got.append(f"{mb}MB ok")
-            except Exception as e:
-                got.append(f"{mb}MB {type(e).__name__}: {str(e)[:90]}")
-                break
-
-    t = threading.Thread(target=work)
-    t.start()
-    t.join()
-    print(f"[probe:{tag}] free={free / 2**30:.1f}G total={total / 2**30:.1f}G "
-          f"allocated={torch.cuda.memory_allocated() / 2**30:.1f}G "
-          f"reserved={torch.cuda.memory_reserved() / 2**30:.1f}G | {got}", flush=True)
-
-
 def _duration(world, steer_to, steer_at, seconds, seed, *args, **kwargs):
     # Worst case is the whole clip generated at the slowest observed rate, plus the
     # world decode at the start. Declared tight on purpose: ZeroGPU compares the
@@ -261,8 +239,6 @@ def run(world: str, steer_to: str, steer_at: float, seconds: float, seed: int):
     steer_frame = int(float(steer_at) * FPS)
     steer_idx = PROMPTS.get(steer_to) if steer_to and steer_to != NO_STEER else None
 
-    _probe("before-start")
-
     t0 = time.perf_counter()
     engine.start(world=WORLDS[world], seed=int(seed))
     # Opening a world decodes all 81 of its pixel frames in one call, which is by far
@@ -272,7 +248,7 @@ def run(world: str, steer_to: str, steer_at: float, seconds: float, seed: int):
     # PyTorch asks NVML which processes hold memory, NVML is not available there, and
     # the run dies on an internal assert instead. Hand the cache back first.
     torch.cuda.empty_cache()
-    _probe("after-world")
+    torch.cuda.reset_peak_memory_stats()
     frames, pending_steer, last_push = [], steer_idx is not None, 0.0
 
     # The worker runs in a thread and reports failure by setting status.state rather