[TEST] Update fresh_knobs fixture default behavior (#9184)

Refactor the default behavior of fresh_knobs in test fixtures.

`fresh_knobs` (default): Now preserves library paths (build, nvidia, amd
knobs)
Most tests need CUDA toolkit paths to compile successfully
Respects environment variables like `TRITON_PTXAS_BLACKWELL_PATH`

`fresh_knobs_including_libraries` (new): Resets ALL knobs including
library paths

Author: CliveUnger

python/test/conftest.py

@@ -27,23 +27,26 @@ def fresh_triton_cache():
 
 @pytest.fixture
 def fresh_knobs():
+    """
+    Resets all knobs except ``build``, ``nvidia``, and ``amd`` (preserves
+    library paths needed to compile kernels).
+    """
     from triton._internal_testing import _fresh_knobs_impl
-    fresh_function, reset_function = _fresh_knobs_impl()
+    fresh_function, reset_function = _fresh_knobs_impl(skipped_attr={"build", "nvidia", "amd"})
     try:
         yield fresh_function()
     finally:
         reset_function()
 
 
 @pytest.fixture
-def fresh_knobs_except_libraries():
+def fresh_knobs_including_libraries():
     """
-    A variant of `fresh_knobs` that keeps library path
-    information from the environment as these may be
-    needed to successfully compile kernels.
+    Resets ALL knobs including ``build``, ``nvidia``, and ``amd``.
+    Use for tests that verify initial values of these knobs.
     """
     from triton._internal_testing import _fresh_knobs_impl
-    fresh_function, reset_function = _fresh_knobs_impl(skipped_attr={"build", "nvidia", "amd"})
+    fresh_function, reset_function = _fresh_knobs_impl()
     try:
         yield fresh_function()
     finally:

python/test/unit/language/test_core.py

@@ -5745,7 +5745,7 @@ def test_dot_max_num_imprecise_acc(M, N, K, BLOCK_M, BLOCK_N, BLOCK_K, in_type_s
 
 @pytest.mark.parametrize("enable_fp_fusion", [False, True])
 @pytest.mark.parametrize("default_override", [False, True])
-def test_enable_fp_fusion(enable_fp_fusion, default_override, device, fresh_knobs_except_libraries):
+def test_enable_fp_fusion(enable_fp_fusion, default_override, device, fresh_knobs):
     # Sequential multiply add can be fused by backend
     @triton.jit
     def mul_add(data):
@@ -5754,7 +5754,7 @@ def mul_add(data):
 
     data = torch.randn((128, ), device=device, dtype=torch.float32)
     if default_override:
-        fresh_knobs_except_libraries.language.default_fp_fusion = enable_fp_fusion
+        fresh_knobs.language.default_fp_fusion = enable_fp_fusion
         h = mul_add.warmup(data, grid=(1, ))
     else:
         h = mul_add.warmup(data, grid=(1, ), enable_fp_fusion=enable_fp_fusion)
@@ -5772,7 +5772,7 @@ def mul_add(data):
 
 @pytest.mark.skipif(not is_cuda(), reason="Requires CUDA")
 @pytest.mark.parametrize("enable_reflect_ftz", [False, True])
-def test_enable_reflect_ftz(enable_reflect_ftz, device, fresh_knobs_except_libraries):
+def test_enable_reflect_ftz(enable_reflect_ftz, device, fresh_knobs):
 
     @triton.jit
     def exp2(data):
@@ -5793,7 +5793,7 @@ def exp2(data):
 
 @pytest.mark.parametrize("arch", ["sm70", "sm80", "sm90", "gfx942", "gfx950", "gfx1200"])
 @pytest.mark.parametrize("env_var_override", [False, True])
-def test_override_arch(arch, env_var_override, device, fresh_knobs_except_libraries):
+def test_override_arch(arch, env_var_override, device, fresh_knobs):
     if arch.startswith("sm") and not is_cuda():
         pytest.skip(f"{arch} arch only for CUDA")
     elif arch.startswith("gfx") and not is_hip():
@@ -5810,7 +5810,7 @@ def simple(data, out):
 
     if is_cuda():
         if env_var_override:
-            fresh_knobs_except_libraries.runtime.override_arch = str(arch)
+            fresh_knobs.runtime.override_arch = str(arch)
             h = simple.warmup(data, out, grid=(1, ))
         else:
             h = simple.warmup(data, out, arch=arch, grid=(1, ))
@@ -5820,7 +5820,7 @@ def simple(data, out):
         # For HIP, the generated kernel is a binary containing the final ISA. So we cannot run
         # them like CUDA side if the chip doesn't match. Here we just check generated ISA.
         if env_var_override:
-            fresh_knobs_except_libraries.runtime.override_arch = str(arch)
+            fresh_knobs.runtime.override_arch = str(arch)
             h = simple.warmup(data, out, grid=(1, ))
         else:
             h = simple.warmup(data, out, arch=arch, grid=(1, ))

python/test/unit/language/test_tuple.py

@@ -269,15 +269,15 @@ def test_passing_nested_tuple_with_constexpr(device):
     _nested_tuple_kernel[(1, )](((1, ), (tl.constexpr(2), )))
 
 
-def test_passing_nested_tuple_with_constexpr_and_jit_hook(device, fresh_knobs_except_libraries):
+def test_passing_nested_tuple_with_constexpr_and_jit_hook(device, fresh_knobs):
     # get the serialized specialization data
     specialization_data = None
 
     def cache_hook(*args, **kwargs):
         nonlocal specialization_data
         specialization_data = kwargs["compile"]["specialization_data"]
 
-    fresh_knobs_except_libraries.runtime.jit_cache_hook = cache_hook
+    fresh_knobs.runtime.jit_cache_hook = cache_hook
 
     device = getattr(torch, device).current_device()
 

python/test/unit/runtime/test_build.py

@@ -65,7 +65,7 @@ def test_compile_module(fresh_triton_cache):
     assert mod2.__file__ == mod.__file__
 
 
-def test_compile_module_bad_cache(fresh_knobs_except_libraries):
+def test_compile_module_bad_cache(fresh_knobs):
     with tempfile.TemporaryDirectory() as tmpd:
         tmp = Path(tmpd)
         called_get_file = False
@@ -79,7 +79,7 @@ def get_file(self, filename: str) -> str | None:
                 return str(tmp / filename)
 
         # First corrupt the cache
-        fresh_knobs_except_libraries.cache.manager_class = InvalidFileCacheManager
+        fresh_knobs.cache.manager_class = InvalidFileCacheManager
 
         mod = compile_module_from_src(TEST_MODULE_C, "test_module")
         assert called_get_file

python/test/unit/runtime/test_compilation_listener.py

@@ -17,7 +17,7 @@ def cumsum_kernel(ptr):
     tl.store(block, tl.cumsum(x, 0))
 
 
-def test_compile_stats(device: str, fresh_knobs_except_libraries: Any, fresh_triton_cache: str) -> None:
+def test_compile_stats(device: str, fresh_knobs: Any, fresh_triton_cache: str) -> None:
     captured: Union[tuple[Union[ASTSource, IRSource], dict[str, Any], dict[str, Any], CompileTimes, bool], None] = None
 
     def compile_listener(src: Union[ASTSource, IRSource], metadata: dict[str, str], metadata_group: dict[str, Any],
@@ -26,7 +26,7 @@ def compile_listener(src: Union[ASTSource, IRSource], metadata: dict[str, str],
         assert captured is None
         captured = (src, metadata, metadata_group, times, cache_hit)
 
-    fresh_knobs_except_libraries.compilation.listener = compile_listener
+    fresh_knobs.compilation.listener = compile_listener
 
     x = torch.randn(4, device=device)
     cumsum_kernel[(1, )](x)

python/test/unit/test_knobs.py

@@ -113,7 +113,8 @@ def test_env_updated(fresh_knobs, monkeypatch):
 
 @pytest.mark.parametrize("truthy, falsey", [("1", "0"), ("true", "false"), ("True", "False"), ("TRUE", "FALSE"),
                                             ("y", "n"), ("YES", "NO"), ("ON", "OFF")])
-def test_read_env(truthy, falsey, fresh_knobs, monkeypatch):
+def test_read_env(truthy, falsey, fresh_knobs_including_libraries, monkeypatch):
+    fresh_knobs = fresh_knobs_including_libraries
     # bool defaulting to False
     assert not fresh_knobs.runtime.debug
     # bool defaulting to True
@@ -170,7 +171,8 @@ def test_triton_home(fresh_knobs, monkeypatch):
     assert fresh_knobs.cache.override_dir == "/tmp/user/triton_home/.triton/override"
 
 
-def test_set_knob_directly(fresh_knobs, monkeypatch):
+def test_set_knob_directly(fresh_knobs_including_libraries, monkeypatch):
+    fresh_knobs = fresh_knobs_including_libraries
     assert fresh_knobs.cache.dir.endswith(".triton/cache")
 
     fresh_knobs.cache.dir = "/tmp/triton_cache"
@@ -279,7 +281,8 @@ def test_nvidia_tool(fresh_knobs, tmp_path, monkeypatch):
     assert fresh_knobs.nvidia.ptxas_options is None
 
 
-def test_opt_bool(fresh_knobs, monkeypatch):
+def test_opt_bool(fresh_knobs_including_libraries, monkeypatch):
+    fresh_knobs = fresh_knobs_including_libraries
     assert fresh_knobs.amd.use_block_pingpong is None
     monkeypatch.setenv("TRITON_HIP_USE_BLOCK_PINGPONG", "0")
     assert not fresh_knobs.amd.use_block_pingpong

python/triton_kernels/tests/conftest.py

@@ -14,6 +14,26 @@ def device(request):
 
 @pytest.fixture
 def fresh_knobs():
+    """
+    Default fresh knobs fixture that preserves library path
+    information from the environment as these are typically
+    needed to successfully compile kernels.
+    """
+    from triton._internal_testing import _fresh_knobs_impl
+    fresh_function, reset_function = _fresh_knobs_impl(skipped_attr={"build", "nvidia", "amd"})
+    try:
+        yield fresh_function()
+    finally:
+        reset_function()
+
+
+@pytest.fixture
+def fresh_knobs_including_libraries():
+    """
+    A variant of `fresh_knobs` that resets ALL knobs including
+    library paths. Use this only for tests that need complete
+    environment isolation.
+    """
     from triton._internal_testing import _fresh_knobs_impl
     fresh_function, reset_function = _fresh_knobs_impl()
     try: