set a different triton cache for each test to avoid blocking writes to cache (#2843)

* set a different triton cache for each test to avoid blocking writes to cache

* set log level

* disable debug logging for filelock
This commit is contained in:
Wing Lian
2025-06-29 22:05:21 -04:00
committed by GitHub
parent 81893c775c
commit 7563e1bd30
2 changed files with 10 additions and 0 deletions

View File

@@ -32,6 +32,8 @@ df_args = {
"NIGHTLY_BUILD": os.environ.get("NIGHTLY_BUILD", ""), "NIGHTLY_BUILD": os.environ.get("NIGHTLY_BUILD", ""),
"CODECOV_TOKEN": os.environ.get("CODECOV_TOKEN", ""), "CODECOV_TOKEN": os.environ.get("CODECOV_TOKEN", ""),
"HF_HOME": "/workspace/data/huggingface-cache/hub", "HF_HOME": "/workspace/data/huggingface-cache/hub",
"PYTHONUNBUFFERED": os.environ.get("PYTHONUNBUFFERED", "1"),
"DEEPSPEED_LOG_LEVEL": os.environ.get("DEEPSPEED_LOG_LEVEL", "WARNING"),
} }
dockerfile_contents = df_template.render(**df_args) dockerfile_contents = df_template.render(**df_args)

View File

@@ -4,6 +4,7 @@ shared pytest fixtures
import functools import functools
import importlib import importlib
import logging
import os import os
import shutil import shutil
import sys import sys
@@ -25,6 +26,8 @@ from tests.hf_offline_utils import (
hf_offline_context, hf_offline_context,
) )
logging.getLogger("filelock").setLevel(logging.CRITICAL)
def retry_on_request_exceptions(max_retries=3, delay=1): def retry_on_request_exceptions(max_retries=3, delay=1):
# pylint: disable=duplicate-code # pylint: disable=duplicate-code
@@ -420,6 +423,11 @@ def temp_dir() -> Generator[str, None, None]:
shutil.rmtree(_temp_dir) shutil.rmtree(_temp_dir)
@pytest.fixture(scope="function", autouse=True)
def unique_triton_cache_dir(temp_dir):
os.environ["TRITON_CACHE_DIR"] = temp_dir + "/~.triton/cache"
@pytest.fixture(scope="function", autouse=True) @pytest.fixture(scope="function", autouse=True)
def cleanup_monkeypatches(): def cleanup_monkeypatches():
from transformers import Trainer from transformers import Trainer