Fix: torch.compile errors in threaded multi-GPU runs

src/aind_torch_utils/config.py

@@ -4,6 +4,14 @@
 from pydantic import BaseModel, Field, model_validator
 
 
+CompileMode = Literal[
+    "default",
+    "reduce-overhead",
+    "max-autotune",
+    "max-autotune-no-cudagraphs",
+]
+
+
 class InferenceConfig(BaseModel):
     # Geometry
     patch: Tuple[int, int, int] = Field(
@@ -32,13 +40,20 @@ class InferenceConfig(BaseModel):
         default=False, description="Enable cuDNN benchmarking"
     )
     use_compile: bool = Field(default=False, description="Use torch.compile")
-    compile_mode: str = Field(
-        default="reduce-overhead",
-        description="Torch.compile mode",  # or "max-autotune" if you want extra tuning time
+    compile_mode: CompileMode = Field(
+        default="default",
+        description="Torch.compile mode",
+    )
+    compile_dynamic: Optional[bool] = Field(
+        default=None,
+        description=(
+            "torch.compile dynamic shapes: None (auto) compiles static and "
+            "promotes to dynamic on a shape change; True forces dynamic; "
+            "False recompiles per shape. Input shapes are constant (tail "
+            "batches are padded), so auto yields one static-specialized "
+            "graph."
+        ),
     )
-    compile_dynamic: bool = Field(
-        default=True, description="Torch.compile with dynamic shapes"
-    )  # tolerate last-batch size changes
 
     # Concurrency / queues
     max_inflight_batches: int = Field(default=64, description="Max in-flight batches")
@@ -224,5 +239,25 @@ def _validate(self):
         # Devices
         if not self.devices:
             raise ValueError("devices list must not be empty")
+        cuda_device_count = sum(
+            str(device).lower().startswith("cuda") for device in self.devices
+        )
+        if self.use_compile and cuda_device_count > 1:
+            safe_compile_modes = {
+                "reduce-overhead": "default",
+                "max-autotune": "max-autotune-no-cudagraphs",
+            }
+            safe_mode = safe_compile_modes.get(self.compile_mode)
+            if safe_mode:
+                warnings.warn(
+                    (
+                        f"torch.compile mode '{self.compile_mode}' enables "
+                        "CUDA graphs, "
+                        "which can fail in threaded multi-GPU runs; using "
+                        f"'{safe_mode}' instead."
+                    ),
+                    RuntimeWarning,
+                )
+                self.compile_mode = safe_mode
 
         return self

src/aind_torch_utils/run.py

@@ -484,14 +484,25 @@ def _parse_args(argv: List[str]) -> argparse.Namespace:
     ap.add_argument(
         "--compile-mode",
         type=str,
-        default="reduce-overhead",
-        choices=["reduce-overhead", "max-autotune"],
+        default="default",
+        choices=[
+            "default",
+            "reduce-overhead",
+            "max-autotune",
+            "max-autotune-no-cudagraphs",
+        ],
         help="torch.compile mode",
     )
     ap.add_argument(
-        "--no-compile-dynamic",
-        action="store_true",
-        help="Disable dynamic shape support for torch.compile",
+        "--compile-dynamic",
+        type=str,
+        default="auto",
+        choices=["auto", "true", "false"],
+        help=(
+            "torch.compile dynamic shapes: 'auto' compiles static and "
+            "promotes to dynamic if a shape changes; 'true' forces dynamic; "
+            "'false' recompiles per shape"
+        ),
     )
     ap.add_argument(
         "--max-inflight-batches",
@@ -627,7 +638,9 @@ def main(argv: Optional[List[str]] = None) -> None:
         cudnn_benchmark=args.cudnn_benchmark,
         use_compile=args.compile,
         compile_mode=args.compile_mode,
-        compile_dynamic=not args.no_compile_dynamic,
+        compile_dynamic={"auto": None, "true": True, "false": False}[
+            args.compile_dynamic
+        ],
         max_inflight_batches=args.max_inflight_batches,
         seam_mode=args.seam_mode,
         trim_voxels=args.trim_voxels,

src/aind_torch_utils/workers.py

@@ -34,7 +34,11 @@ class Batch:
         List of (z, y, x) start coordinates for each patch in the batch,
         relative to the expanded block.
     host_in : torch.Tensor
-        The input tensor of patches, pinned to host memory.
+        The input tensor of patches, pinned to host memory. When compiling
+        (cfg.use_compile), tail batches are zero-padded up to batch_size so
+        the model sees a constant input shape, and rows beyond
+        len(starts_in_block) are padding. In eager mode it has exactly
+        len(starts_in_block) rows.
     valid_sizes : List[Tuple[int, int, int]]
         List of (dz, dy, dx) valid dimensions for each patch, handling
         boundary conditions.
@@ -260,10 +264,18 @@ def run(self, stop_event: threading.Event) -> None:
             # batch over those starts
             for i in range(0, total_patches, self.cfg.batch_size):
                 batch_starts = starts[i : i + self.cfg.batch_size]
-                B = len(batch_starts)
+                n_real = len(batch_starts)
                 pin_memory = any("cuda" in d for d in self.cfg.devices)
+                # When compiling, pad the tail batch up to batch_size so the
+                # model always sees a constant input shape; this prevents
+                # torch.compile from recompiling at runtime (not thread-safe
+                # across GPU workers). Writers ignore padded rows since they
+                # only index rows in starts_in_block. In eager mode there is
+                # no shape constraint, so allocate exactly n_real rows and
+                # avoid wasting compute and copy bandwidth on padding.
+                n_rows = self.cfg.batch_size if self.cfg.use_compile else n_real
                 host_in = torch.zeros(
-                    (B, 1, pz, py, px),
+                    (n_rows, 1, pz, py, px),
                     dtype=torch.float16 if self.cfg.amp else torch.float32,
                     pin_memory=pin_memory,
                 )
@@ -351,18 +363,68 @@ def __init__(
         self.copy_stream = torch.cuda.Stream(device=self.device)
 
         if getattr(torch, "compile", None) and self.cfg.use_compile:
+            self._compile_model()
+
+    def _autocast_context(self):
+        return (
+            torch.autocast(device_type="cuda", dtype=torch.float16)
+            if self.cfg.amp
+            else nullcontext()
+        )
+
+    def _compile_model(self) -> None:
+        # Keep a handle to the original module so we can fall back to eager
+        # execution if compilation fails. torch.compile returns a new wrapper
+        # and does not mutate the original, so this reference stays valid.
+        eager_model = self.model
+        try:
             try:
-                # dynamic=True avoids recompiles when the final batch is smaller
+                # PrepWorker pads tail batches to batch_size, so input shapes
+                # are constant and no runtime recompiles are expected
+                # regardless of the `dynamic` setting.
                 self.model = torch.compile(
                     self.model,
                     mode=self.cfg.compile_mode,
                     dynamic=self.cfg.compile_dynamic,
                 )
-                logger.info("Successfully compiled model.")
+                logger.info("Compiled model on %s.", self.device)
             except TypeError:
                 # older PyTorch without `dynamic` kwarg
                 self.model = torch.compile(self.model, mode=self.cfg.compile_mode)
-                logger.info("Successfully compiled model (older pytorch).")
+                logger.info("Compiled model on %s (older pytorch).", self.device)
+
+            # Compilation is lazy: the graph is traced on the first forward,
+            # so tracing/guard errors surface here in warmup, not above.
+            self._warmup_compiled_model()
+        except Exception as exc:
+            # Some models do host-side numpy/Python work in forward that
+            # dynamo cannot trace. Fall back to eager so the run proceeds
+            # instead of aborting. Warmup runs on the main thread, so this
+            # also keeps the failure off the worker threads.
+            logger.warning(
+                "torch.compile failed on %s (%s); falling back to eager "
+                "execution.",
+                self.device,
+                type(exc).__name__,
+                exc_info=True,
+            )
+            self.model = eager_model
+
+    def _warmup_compiled_model(self) -> None:
+        torch.cuda.set_device(self.device)
+        dtype = torch.float16 if self.cfg.amp else torch.float32
+        shape = (self.cfg.batch_size, 1, *self.cfg.patch)
+        warmup_in = torch.zeros(shape, dtype=dtype, device=self.device)
+
+        logger.info(
+            "Warming compiled model on %s with shape %s.", self.device, shape
+        )
+        with torch.inference_mode():
+            with self._autocast_context():
+                warmup_out = self.model(warmup_in)
+        torch.cuda.synchronize(self.device)
+        del warmup_in, warmup_out
+        logger.info("Finished compiled model warmup on %s.", self.device)
 
     def run(self, stop_event: threading.Event) -> None:
         """
@@ -376,11 +438,7 @@ def run(self, stop_event: threading.Event) -> None:
         stop_event : threading.Event
             An event that signals the worker to stop.
         """
-        autocast_ctx = (
-            torch.autocast(device_type="cuda", dtype=torch.float16)
-            if self.cfg.amp
-            else nullcontext()
-        )
+        autocast_ctx = self._autocast_context()
 
         # Ensure the current device matches self.device for streams/events
         torch.cuda.set_device(self.device)

tests/test_config.py

@@ -13,6 +13,8 @@ def test_precision_performance_defaults_to_false():
 
     assert cfg.use_tf32 is False
     assert cfg.cudnn_benchmark is False
+    assert cfg.compile_mode == "default"
+    assert cfg.compile_dynamic is None
 
 
 @pytest.mark.parametrize("field", ["use_tf32", "cudnn_benchmark"])
@@ -26,3 +28,41 @@ def test_precision_performance_config_overrides(field):
     )
 
     assert getattr(cfg, field) is True
+
+
+@pytest.mark.parametrize(
+    ("requested_mode", "expected_mode"),
+    [
+        ("reduce-overhead", "default"),
+        ("max-autotune", "max-autotune-no-cudagraphs"),
+    ],
+)
+def test_compile_cudagraph_modes_are_downgraded_for_multi_cuda(
+    requested_mode, expected_mode
+):
+    with pytest.warns(RuntimeWarning, match="threaded multi-GPU"):
+        cfg = InferenceConfig(
+            patch=(16, 16, 16),
+            overlap=4,
+            trim_voxels=2,
+            block=(32, 32, 32),
+            devices=["cuda:0", "cuda:1"],
+            use_compile=True,
+            compile_mode=requested_mode,
+        )
+
+    assert cfg.compile_mode == expected_mode
+
+
+def test_compile_cudagraph_mode_allowed_for_single_cuda():
+    cfg = InferenceConfig(
+        patch=(16, 16, 16),
+        overlap=4,
+        trim_voxels=2,
+        block=(32, 32, 32),
+        devices=["cuda:0"],
+        use_compile=True,
+        compile_mode="reduce-overhead",
+    )
+
+    assert cfg.compile_mode == "reduce-overhead"