fix(prebuilt): don't sleep after the final nvidia-smi retry attempt

Cosmetic follow-up to a4ad50e (review nit): the retry loop slept 2s even
after the last attempt, adding ~2s only when nvidia-smi is permanently hung.
Sleep only between attempts.
This commit is contained in:
Daniel Han 2026-06-19 02:55:09 -07:00
commit c1a997b438

View file

@ -2873,12 +2873,14 @@ def _nvidia_smi_capture(
CPU-only hosts never incur this wait.
"""
last_exc: Exception | None = None
for _attempt in range(max(1, attempts)):
_attempts = max(1, attempts)
for _attempt in range(_attempts):
try:
return run_capture(command, timeout = timeout)
except subprocess.TimeoutExpired as exc:
last_exc = exc
time.sleep(2)
if _attempt + 1 < _attempts: # don't sleep after the final attempt
time.sleep(2)
raise last_exc if last_exc is not None else RuntimeError("nvidia-smi capture failed")