From d233bb58ce68e78d1e7e3cc83b7349cdc0793acb Mon Sep 17 00:00:00 2001 From: Carlo van Driesten Date: Wed, 11 Nov 2026 13:01:37 +0100 Subject: [PATCH] fix(ci): skip the fused encoder where triton is absent The shipped puffer_drive.yaml sets policy.fused_slot_encoder=true, and DriveBackbone raises ImportError when that is set without triton. Two CI jobs build a policy straight from the shipped config, so Unit Tests and Notebook Execution have been failing on every pull request: ImportError: fused_slot_encoder=true requires triton The smoke and eval suites already pin the flag off for exactly this reason; these two call sites were simply missed. test_action_selection only exercises the discrete-to-continuous action table, which the fused kernel has no part in, so it pins the flag off outright. The inference notebook keeps the kernel whenever it has a CUDA device and drops it otherwise, which is also what a reader running the notebook on a laptop needs. Neither change touches the training path or the kernel itself: with triton and a GPU present, both call sites behave exactly as before. --- notebooks/05_inference.py | 1 + tests/eval/test_eval.py | 1 + tests/unit_tests/test_action_selection.py | 3 ++- 3 files changed, 4 insertions(+), 1 deletion(-) diff --git a/notebooks/05_inference.py b/notebooks/05_inference.py index 25999cddc8..470e699004 100644 --- a/notebooks/05_inference.py +++ b/notebooks/05_inference.py @@ -48,6 +48,7 @@ N = env.num_agents device = torch.device("cuda" if torch.cuda.is_available() else "cpu") +config["policy"]["fused_slot_encoder"] = config["policy"]["fused_slot_encoder"] and device.type == "cuda" policy = DrivePolicy(env, **config["policy"]).to(device) if CHECKPOINT_PATH: diff --git a/tests/eval/test_eval.py b/tests/eval/test_eval.py index ae5c261ead..235e92f07c 100644 --- a/tests/eval/test_eval.py +++ b/tests/eval/test_eval.py @@ -877,6 +877,7 @@ def test_policy_builds_from_the_released_driver_env(sdc_replay_map_dir): with patch.object(sys, "argv", ["pufferl.py"]): args = pufferl.load_config("puffer_drive") args["train"]["device"] = "cpu" + args["policy"]["fused_slot_encoder"] = False # the fused encoder needs triton, which is CUDA-only env_kwargs = { "num_agents": SDC_SCENARIO_COUNT, "num_maps": SDC_SCENARIO_COUNT, diff --git a/tests/unit_tests/test_action_selection.py b/tests/unit_tests/test_action_selection.py index 1fa51b0de0..92dbd2e493 100644 --- a/tests/unit_tests/test_action_selection.py +++ b/tests/unit_tests/test_action_selection.py @@ -201,7 +201,8 @@ def test_drive_mean_action_is_exact_in_physical_units(monkeypatch): action_type="continuous", ) try: - policy = ocean_torch.Drive(env, **args["policy"]) + # the fused encoder is a CUDA/triton kernel and plays no part in the action table + policy = ocean_torch.Drive(env, **{**args["policy"], "fused_slot_encoder": False}) num_lat = len(binding.JERK_LAT) lat_zero_idx = binding.JERK_LAT.index(0.0)