Implement 'fresh_triton_cache' fixture through 'compilation.always_compile'

Signed-off-by: Anatoly Myachev <[email protected]>

Author: anmyachev

python/test/regression/conftest.py

@@ -1,6 +1,4 @@
-import os
 import pytest
-import tempfile
 
 
 def pytest_addoption(parser):
@@ -14,9 +12,7 @@ def device(request):
 
 @pytest.fixture
 def fresh_triton_cache():
-    with tempfile.TemporaryDirectory() as tmpdir:
-        try:
-            os.environ["TRITON_CACHE_DIR"] = tmpdir
-            yield tmpdir
-        finally:
-            os.environ.pop("TRITON_CACHE_DIR", None)
+    from triton import knobs
+    with knobs.compilation.scope():
+        knobs.compilation.always_compile = True
+        yield

python/test/unit/conftest.py

@@ -2,8 +2,8 @@
 import sys
 import pathlib
 import pytest
-import tempfile
 from typing import Optional, Set
+import contextlib
 
 
 def pytest_configure(config):
@@ -69,11 +69,23 @@ def device(request):
 
 @pytest.fixture
 def fresh_triton_cache():
-    with tempfile.TemporaryDirectory() as tmpdir:
-        from triton import knobs
-        with knobs.cache.scope():
-            knobs.cache.dir = tmpdir
-            yield tmpdir
+    from triton import knobs
+    with knobs.compilation.scope():
+        knobs.compilation.always_compile = True
+        yield
+
+
+@pytest.fixture
+def fresh_triton_cache_scope():
+    from triton import knobs
+
+    @contextlib.contextmanager
+    def fresh_cache():
+        with knobs.compilation.scope():
+            knobs.compilation.always_compile = True
+            yield
+
+    yield fresh_cache
 
 
 def _fresh_knobs_impl(monkeypatch, skipped_attr: Optional[Set[str]] = None):

python/test/unit/runtime/test_compilation_listener.py

@@ -5,7 +5,7 @@
 from triton.knobs import CompileTimes
 from triton.compiler.compiler import ASTSource, IRSource
 
-from typing import Any, Union
+from typing import Any, Union, Callable
 
 import torch
 
@@ -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_except_libraries: Any, fresh_triton_cache_scope: Callable) -> 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],
@@ -29,7 +29,8 @@ def compile_listener(src: Union[ASTSource, IRSource], metadata: dict[str, str],
     fresh_knobs_except_libraries.compilation.listener = compile_listener
 
     x = torch.randn(4, device=device)
-    cumsum_kernel[(1, )](x)
+    with fresh_triton_cache_scope():
+        cumsum_kernel[(1, )](x)
 
     assert captured is not None
 

third_party/intel/python/test/conftest.py

@@ -1,6 +1,4 @@
-import os
 import pytest
-import tempfile
 
 
 def pytest_addoption(parser):
@@ -14,9 +12,7 @@ def device(request):
 
 @pytest.fixture
 def fresh_triton_cache():
-    with tempfile.TemporaryDirectory() as tmpdir:
-        try:
-            os.environ["TRITON_CACHE_DIR"] = tmpdir
-            yield tmpdir
-        finally:
-            os.environ.pop("TRITON_CACHE_DIR", None)
+    from triton import knobs
+    with knobs.compilation.scope():
+        knobs.compilation.always_compile = True
+        yield