[Tuner] Raise a GPU fault from the power-sampling pass

aiter/utility/package_power.py

@@ -210,8 +210,8 @@ def loop():
             fn()
             calls += 1
             in_flight.mark()
-    finally:
         _sync_cuda()
+    finally:
         window_s = time.perf_counter() - window_start
         stop.set()
         thread.join()
@@ -349,12 +349,34 @@ def _positive(value) -> float | None:
 
 
 def _sync_cuda() -> None:
+    """Wait for the GPU. A failed synchronize is a GPU fault and is raised."""
     try:
         import torch
+    except ImportError:
+        return
+    if torch.cuda.is_available() and torch.cuda.is_initialized():
+        torch.cuda.synchronize()
+
 
+def _gpu_fault(exc: BaseException) -> BaseException | None:
+    """The error to raise when ``exc`` left the GPU context unusable, else None.
+
+    An ``AcceleratorError`` is a fault. So is any error after which a
+    synchronize fails. A sensor that cannot be read is neither.
+    """
+    if type(exc).__name__ == "AcceleratorError":
+        return exc
+    try:
+        import torch
+    except ImportError:
+        return None
+    if not (torch.cuda.is_available() and torch.cuda.is_initialized()):
+        return None
+    try:
         torch.cuda.synchronize()
-    except Exception:
-        logger.debug("CUDA synchronize skipped", exc_info=True)
+    except Exception as error:  # noqa: BLE001 - any failure here is the fault
+        return error
+    return None
 
 
 def shared_power_override_rejected() -> str | None:
@@ -405,7 +427,10 @@ def measure_kernel(
     kernel the timer measured and nothing else. ``kernel_us`` is that call's
     timed GPU time; with it the reading is corrected for host launch gaps
     (``active_watts``). NaN means the sensor could not be read, the call
-    raised, or the window was too idle to tell.
+    raised without breaking the GPU, or the window was too idle to tell.
+
+    A GPU fault during the sampling pass is raised, so the tuner charges it
+    to this candidate instead of the next one in the same worker.
     """
     if device_index is None:
         try:
@@ -420,6 +445,11 @@ def measure_kernel(
                 call, device_index=device_index, kernel_us=_positive(kernel_us)
             )
         )
-    except Exception as exc:  # noqa: BLE001
+    except Exception as exc:
+        fault = _gpu_fault(exc)
+        if fault is exc:
+            raise
+        if fault is not None:
+            raise fault from exc
         logger.warning("package power measurement failed: %s", exc)
         return float("nan")

op_tests/tuning_tests/test_fast_tuning.py

@@ -1437,6 +1437,89 @@ def launch():
         self.assertGreater(corrected, 430.0)
         self.assertLess(corrected, 620.0)
 
+    def test_gpu_fault_in_power_pass_is_raised(self):
+        from aiter.utility import package_power
+
+        class AcceleratorError(RuntimeError):
+            pass
+
+        def faulting_launch():
+            raise AcceleratorError("illegal memory access")
+
+        with (
+            mock.patch.object(package_power, "power_input_path", return_value="x"),
+            self.assertRaises(AcceleratorError),
+        ):
+            package_power.measure_kernel(faulting_launch, device_index=0)
+
+    def test_error_that_breaks_synchronize_is_raised_as_the_fault(self):
+        from aiter.utility import package_power
+
+        lost = RuntimeError("HIP error: unspecified launch failure")
+
+        def launch():
+            raise ValueError("kernel rejected its arguments")
+
+        with (
+            mock.patch.object(package_power, "power_input_path", return_value="x"),
+            mock.patch.object(package_power, "_gpu_fault", return_value=lost),
+            self.assertRaises(RuntimeError) as raised,
+        ):
+            package_power.measure_kernel(launch, device_index=0)
+        self.assertIs(raised.exception, lost)
+        self.assertIsInstance(raised.exception.__cause__, ValueError)
+
+    def test_non_fault_error_in_power_pass_reads_as_nan(self):
+        import math
+
+        from aiter.utility import package_power
+
+        def launch():
+            raise ValueError("kernel rejected its arguments")
+
+        with (
+            mock.patch.object(package_power, "power_input_path", return_value="x"),
+            mock.patch.object(package_power, "_gpu_fault", return_value=None),
+        ):
+            watts = package_power.measure_kernel(launch, device_index=0)
+        self.assertTrue(math.isnan(watts))
+
+    def test_gpu_fault_check_uses_synchronize(self):
+        from aiter.utility import package_power
+
+        lost = RuntimeError("HIP error: unspecified launch failure")
+        with (
+            mock.patch("torch.cuda.is_available", return_value=True),
+            mock.patch("torch.cuda.is_initialized", return_value=True),
+            mock.patch("torch.cuda.synchronize", side_effect=lost),
+        ):
+            self.assertIs(package_power._gpu_fault(ValueError("x")), lost)
+        with (
+            mock.patch("torch.cuda.is_available", return_value=True),
+            mock.patch("torch.cuda.is_initialized", return_value=True),
+            mock.patch("torch.cuda.synchronize"),
+        ):
+            self.assertIsNone(package_power._gpu_fault(ValueError("x")))
+
+    def test_failed_synchronize_stops_the_sampling_thread(self):
+        import threading
+
+        from aiter.utility import package_power
+
+        lost = RuntimeError("HIP error: unspecified launch failure")
+        before = threading.active_count()
+        with (
+            mock.patch.object(package_power, "power_input_path", return_value="x"),
+            mock.patch.object(package_power, "_read_watts", return_value=400.0),
+            mock.patch.object(
+                package_power, "_sync_cuda", side_effect=[None, None, lost]
+            ),
+            mock.patch.object(package_power._InFlight, "mark"),
+            self.assertRaises(RuntimeError),
+        ):
+            package_power.measure_package_watts(lambda: None, seconds=0.05)
+        self.assertEqual(threading.active_count(), before)
+
     def test_a16w16_hands_its_policy_to_hipblaslt(self):
         from csrc.gemm_a16w16 import gemm_a16w16_tune