[TEST] Reset jit_cache_hook and jit_cache_hook in the fresh_triton_cache fixture (#7885)

These are not environment variables and thereby cannot be automatically
reset by `knobs.reset`

Author: Jokeren

python/test/conftest.py

@@ -19,7 +19,8 @@ def device(request):
 def fresh_triton_cache():
     with tempfile.TemporaryDirectory() as tmpdir:
         from triton import knobs
-        with knobs.cache.scope():
+
+        with knobs.cache.scope(), knobs.runtime.scope():
             knobs.cache.dir = tmpdir
             yield tmpdir
 

python/test/unit/runtime/test_cache.py

@@ -583,7 +583,6 @@ def inc_counter(*args, **kwargs):
 
     triton.knobs.runtime.jit_cache_hook = inc_counter
     final_kernel = kernel_add.warmup(torch.float32, torch.float32, torch.float32, 32, tl.float32, grid=(1, ))
-    triton.knobs.runtime.jit_cache_hook = None
     assert counter == 0
     assert len(kernel_add.device_caches[device][0]) == 1
     assert final_kernel.hash == hash
@@ -692,8 +691,6 @@ def func4(x, y):
     def kernel(Y, fn: tl.constexpr, fn_args):
         tl.store(Y, fn(*fn_args))
 
-    triton.knobs.runtime.jit_cache_hook = None
-    triton.knobs.runtime.jit_post_compile_hook = None
     y = torch.zeros((5, ), dtype=torch.int32, device=device)
     kernel[(1, )](y[0], func1, tuple())
     kernel[(1, )](y[1], func2, tuple())

python/triton_kernels/tests/conftest.py

@@ -1,4 +1,5 @@
 import pytest
+import tempfile
 
 
 def pytest_addoption(parser):
@@ -18,3 +19,13 @@ def fresh_knobs():
         yield fresh_function()
     finally:
         reset_function()
+
+
+@pytest.fixture
+def fresh_triton_cache():
+    with tempfile.TemporaryDirectory() as tmpdir:
+        from triton import knobs
+
+        with knobs.cache.scope(), knobs.runtime.scope():
+            knobs.cache.dir = tmpdir
+            yield tmpdir

python/triton_kernels/tests/test_specialize.py

@@ -38,7 +38,7 @@ def cacheable_kernel():
     return get_specialized_kernel()
 
 
-def test_cacheable(device, fresh_knobs):
+def test_cacheable(device, fresh_triton_cache):
     specialized_kernel = get_specialized_kernel()
 
     specialization_data = None