test: isolate cache dirs for tests running in parallel

pyproject.toml

@@ -183,7 +183,6 @@ ini_options.markers = [
   "asyncio: mark a test as a coroutine",
   "cloud: Run the cloud tests for example",
   "timeout: per-test timeout from pytest-timeout",
-  "xdist_group: Pin tests that share ~/.lightning/chunks onto one xdist worker",
 ]
 ini_options.xfail_strict = true
 ini_options.filterwarnings = [ "error::FutureWarning" ]

src/litdata/constants.py

@@ -30,7 +30,6 @@
 _DEFAULT_FAST_DEV_RUN_ITEMS = 10
 _DEFAULT_CACHE_DIR = os.path.join(Path.home(), ".lightning", "chunks")
 _DEFAULT_LIGHTNING_CACHE_DIR = os.path.join("/cache", "chunks")
-_LITDATA_CACHE_DIR = os.getenv("LITDATA_CACHE_DIR", None)
 _SUPPORTED_PROVIDERS = (
     "s3",
     "gs",

src/litdata/utilities/dataset_utilities.py

@@ -10,7 +10,7 @@
 
 import numpy as np
 
-from litdata.constants import _DEFAULT_CACHE_DIR, _DEFAULT_LIGHTNING_CACHE_DIR, _INDEX_FILENAME, _LITDATA_CACHE_DIR
+from litdata.constants import _DEFAULT_CACHE_DIR, _DEFAULT_LIGHTNING_CACHE_DIR, _INDEX_FILENAME
 from litdata.streaming.downloader import get_downloader
 from litdata.streaming.item_loader import BaseItemLoader, TokensLoader
 from litdata.streaming.resolver import Dir, _resolve_dir
@@ -306,8 +306,11 @@ def get_default_cache_dir() -> str:
     Returns:
         str: The resolved default cache root directory.
     """
-    if _LITDATA_CACHE_DIR is not None:
-        return _LITDATA_CACHE_DIR
+    # Read lazily so changes to the environment variable after import are honored.
+    # Useful in tests that modify the variable via monkeypatch.
+    litdata_cache_dir = os.getenv("LITDATA_CACHE_DIR")
+    if litdata_cache_dir is not None:
+        return litdata_cache_dir
     is_lightning_cloud = "LIGHTNING_CLUSTER_ID" in os.environ and "LIGHTNING_CLOUD_PROJECT_ID" in os.environ
     return _DEFAULT_LIGHTNING_CACHE_DIR if is_lightning_cloud else _DEFAULT_CACHE_DIR
 

tests/conftest.py

@@ -1,5 +1,4 @@
 import os
-import shutil
 import signal
 import sys
 import threading
@@ -14,7 +13,6 @@
 from litdata.constants import _POLARS_AVAILABLE
 from litdata.streaming.cache import Cache
 from litdata.streaming.reader import PrepareChunksThread
-from litdata.utilities.dataset_utilities import get_default_cache_dir
 
 
 @pytest.fixture(autouse=True, scope="session")
@@ -25,6 +23,24 @@ def teardown_process_group():
         torch.distributed.destroy_process_group()
 
 
+@pytest.fixture(autouse=True)
+def isolated_cache_dirs(tmp_path_factory, monkeypatch):
+    """Isolate cache directories for each test to prevent cross-test interference.
+
+    ``optimize()`` clears optimizer cache directories, while streaming and Hugging Face
+    index caches live under ``get_default_cache_dir()``. Without isolation, concurrent
+    tests can delete or overwrite each other's cache files.
+
+    Use environment variables rather than module attributes so spawned worker
+    processes inherit the isolated cache configuration.
+    """
+    root = tmp_path_factory.mktemp("litdata-cache")  # not ``tmp_path``: tests list/compare its contents
+    monkeypatch.setenv("LITDATA_CACHE_DIR", str(root / "litdata-cache" / "chunks"))
+    monkeypatch.delenv("LITDATA_HF_CACHE_DIR", raising=False)
+    monkeypatch.setenv("DATA_OPTIMIZER_CACHE_FOLDER", str(root / "optimizer-chunks"))
+    monkeypatch.setenv("DATA_OPTIMIZER_DATA_CACHE_FOLDER", str(root / "optimizer-data"))
+
+
 @pytest.fixture(autouse=True)
 def disable_signals(monkeypatch):
     monkeypatch.setattr(signal, "signal", lambda *args, **kwargs: None)
@@ -172,21 +188,6 @@ def _thread_police():
             raise AssertionError(f"Test left zombie thread: {thread}")
 
 
-@pytest.hookimpl(tryfirst=True)
-def pytest_collection_modifyitems(items):
-    """Force tests sharing the global default cache dir onto the same xdist worker.
-
-    Why: ``clean_pq_index_cache`` shutil.rmtree's ``~/.lightning/chunks`` on setup/teardown.
-    Under ``pytest -n>1`` two such tests on different workers race, deleting each other's
-    chunks mid-iteration. Requires the runner to pass ``--dist=loadgroup``. Must run with
-    ``tryfirst=True`` so the marker is in place before xdist's own hook reads it and
-    appends the ``@group`` nodeid suffix that LoadGroupScheduling uses for routing.
-    """
-    for item in items:
-        if "clean_pq_index_cache" in getattr(item, "fixturenames", ()):
-            item.add_marker(pytest.mark.xdist_group("hf_default_cache"))
-
-
 # ==== fixtures for parquet ====
 @pytest.fixture
 def pq_data():
@@ -212,21 +213,6 @@ def write_pq_data(pq_data, tmp_path):
             df.write_parquet(file_path)
 
 
-@pytest.fixture
-def clean_pq_index_cache():
-    """Ensures the PQ index cache is cleared before and after the test."""
-    cache_path = os.path.join(get_default_cache_dir())
-
-    # Cleanup before the test
-    if os.path.exists(cache_path):
-        shutil.rmtree(cache_path, ignore_errors=True)
-
-    yield
-
-    if os.path.exists(cache_path):
-        shutil.rmtree(cache_path, ignore_errors=True)
-
-
 @pytest.fixture
 def huggingface_hub_mock(monkeypatch):
     huggingface_hub = ModuleType("huggingface_hub")

tests/processing/test_data_processor.py

@@ -32,6 +32,8 @@
     _cache_local_path,
     _chunks_dir,
     _download_data_target,
+    _get_cache_data_dir,
+    _get_cache_dir,
     _get_item_filesizes,
     _is_local_write_through,
     _is_path,
@@ -1045,6 +1047,29 @@ def test_is_local_write_through_and_chunks_dir(tmp_path):
     assert os.path.isdir(local.path)
 
 
+@pytest.mark.parametrize(
+    ("get_dir", "env_var", "default_name"),
+    [
+        (_get_cache_dir, "DATA_OPTIMIZER_CACHE_FOLDER", "chunks"),
+        (_get_cache_data_dir, "DATA_OPTIMIZER_DATA_CACHE_FOLDER", "data"),
+    ],
+)
+def test_get_cache_dirs_defaults_and_env(monkeypatch, tmp_path, get_dir, env_var, default_name):
+    # conftest's isolated_cache_dirs sets the env var for every test; drop it to exercise the defaults
+    monkeypatch.delenv(env_var)
+
+    monkeypatch.setattr(data_processor_module, "_IS_IN_STUDIO", False)
+    assert get_dir() == f"{tempfile.gettempdir()}/{default_name}"
+    assert get_dir("/my/name") == os.path.join(f"{tempfile.gettempdir()}/{default_name}", "my/name")
+
+    monkeypatch.setattr(data_processor_module, "_IS_IN_STUDIO", True)
+    assert get_dir() == f"/cache/{default_name}"
+
+    monkeypatch.setenv(env_var, str(tmp_path))
+    assert get_dir() == str(tmp_path)
+    assert get_dir("name") == os.path.join(str(tmp_path), "name")
+
+
 def test_adaptive_download_concurrency_env_and_jobs(monkeypatch):
     monkeypatch.delenv("LITDATA_OPTIMIZE_DOWNLOAD_CONCURRENCY", raising=False)
     monkeypatch.setenv("DATA_OPTIMIZER_NUM_WORKERS", "4")

tests/streaming/test_parquet.py

@@ -8,7 +8,7 @@
 
 import pytest
 
-from litdata.constants import _DEFAULT_CACHE_DIR, _DEFAULT_LIGHTNING_CACHE_DIR, _INDEX_FILENAME
+from litdata.constants import _INDEX_FILENAME
 from litdata.streaming.dataset import StreamingDataset
 from litdata.streaming.item_loader import ParquetLoader, PyTreeLoader
 from litdata.streaming.writer import index_parquet_dataset
@@ -32,7 +32,6 @@
 
 #! TODO: Fix test failing on windows
 @pytest.mark.skipif(condition=sys.platform == "win32", reason="Fails on windows and test gets cancelled")
-@pytest.mark.usefixtures("clean_pq_index_cache")
 @pytest.mark.parametrize(
     ("pq_dir_url"),
     [
@@ -92,7 +91,6 @@ def test_parquet_index_write(
 
 
 @pytest.mark.skipif(condition=sys.platform == "win32", reason="Fails on windows and test gets cancelled")
-@pytest.mark.usefixtures("clean_pq_index_cache")
 @patch("litdata.utilities.parquet._HF_HUB_AVAILABLE", True)
 def test_index_hf_dataset(monkeypatch, tmp_path, huggingface_hub_fs_mock):
     with pytest.raises(ValueError, match="Invalid Hugging Face dataset URL"):
@@ -122,7 +120,6 @@ def fake_list(url, storage_options=None):
     )
 
 
-@pytest.mark.usefixtures("clean_pq_index_cache")
 def test_optimize_hf(tmp_path, write_pq_data, monkeypatch):
     pq_dir = tmp_path / "pq-dataset"
     files = [str(pq_dir / name) for name in sorted(os.listdir(pq_dir)) if name.endswith(".parquet")]
@@ -432,7 +429,6 @@ def test_cloud_parquet_dir_forwards_storage_options(tmp_path, monkeypatch, fsspe
     fsspec_mock.filesystem.assert_called_once_with("s3", **storage_options)
 
 
-@pytest.mark.usefixtures("clean_pq_index_cache")
 @patch("litdata.utilities.parquet._HF_HUB_AVAILABLE", True)
 @patch("litdata.streaming.downloader._HF_HUB_AVAILABLE", True)
 @pytest.mark.parametrize(("pre_load_chunk"), [False, True])
@@ -465,7 +461,6 @@ def test_stream_hf_parquet_dataset(monkeypatch, huggingface_hub_fs_mock, pq_data
         assert _ds["height"] == pq_data["height"][idx]
 
 
-@pytest.mark.usefixtures("clean_pq_index_cache")
 @patch("litdata.utilities.parquet._HF_HUB_AVAILABLE", True)
 @patch("litdata.streaming.downloader._HF_HUB_AVAILABLE", True)
 @pytest.mark.parametrize(
@@ -486,7 +481,6 @@ def test_input_dir_wildcard(monkeypatch, huggingface_hub_fs_mock, hf_url, length
         assert len(ds) == length  # 5 datasets for 5 loops
 
 
-@pytest.mark.usefixtures("clean_pq_index_cache")
 @patch("litdata.utilities.parquet._HF_HUB_AVAILABLE", True)
 @patch("litdata.streaming.downloader._HF_HUB_AVAILABLE", True)
 @pytest.mark.parametrize("default", [False, True])
@@ -496,7 +490,7 @@ def test_cache_dir_option(monkeypatch, huggingface_hub_fs_mock, default):
         ds = StreamingDataset(hf_url, cache_dir=None if default else tmpdir)
         assert ds.cache_dir.path == (None if default else os.path.realpath(tmpdir))
         assert ds.input_dir.path.startswith(
-            (_DEFAULT_CACHE_DIR, _DEFAULT_LIGHTNING_CACHE_DIR) if default else os.path.realpath(tmpdir)
+            os.path.realpath(os.environ["LITDATA_CACHE_DIR"]) if default else os.path.realpath(tmpdir)
         )
         # check index file is sole file in chunk cache dir
         assert len(os.listdir(ds.input_dir.path)) == 1

tests/utilities/test_dataset_utilities.py

@@ -1,11 +1,9 @@
-import importlib
 import json
 import os
 from unittest import mock
 
 import pytest
 
-import litdata.constants
 import litdata.utilities
 import litdata.utilities.dataset_utilities
 from litdata.constants import _DEFAULT_CACHE_DIR, _DEFAULT_LIGHTNING_CACHE_DIR, _INDEX_FILENAME
@@ -44,7 +42,7 @@ def test_try_create_cache_dir():
 
     # the cache dir creating at /cache requires root privileges, so we need to mock `os.makedirs()`
     with (
-        mock.patch.dict("os.environ", {"LIGHTNING_CLUSTER_ID": "abc", "LIGHTNING_CLOUD_PROJECT_ID": "123"}),
+        mock.patch.dict("os.environ", {"LIGHTNING_CLUSTER_ID": "abc", "LIGHTNING_CLOUD_PROJECT_ID": "123"}, clear=True),
         mock.patch("litdata.streaming.dataset.os.makedirs") as makedirs_mock,
     ):
         cache_dir_1 = _try_create_cache_dir("")
@@ -64,7 +62,7 @@ def test_try_create_cache_dir_with_custom_cache_dir(tmpdir):
         ) in _try_create_cache_dir("any", cache_dir)
 
     with (
-        mock.patch.dict("os.environ", {"LIGHTNING_CLUSTER_ID": "abc", "LIGHTNING_CLOUD_PROJECT_ID": "123"}),
+        mock.patch.dict("os.environ", {"LIGHTNING_CLUSTER_ID": "abc", "LIGHTNING_CLOUD_PROJECT_ID": "123"}, clear=True),
         mock.patch("litdata.streaming.dataset.os.makedirs") as makedirs_mock,
     ):
         cache_dir_1 = _try_create_cache_dir("", cache_dir)
@@ -108,13 +106,11 @@ def test_get_default_cache_dir():
     with mock.patch.dict(os.environ, {}, clear=True):
         assert get_default_cache_dir() == _DEFAULT_CACHE_DIR
 
-    with mock.patch.dict(os.environ, {"LIGHTNING_CLUSTER_ID": "abc", "LIGHTNING_CLOUD_PROJECT_ID": "123"}):
+    with mock.patch.dict(os.environ, {"LIGHTNING_CLUSTER_ID": "abc", "LIGHTNING_CLOUD_PROJECT_ID": "123"}, clear=True):
         assert get_default_cache_dir() == _DEFAULT_LIGHTNING_CACHE_DIR
 
     with mock.patch.dict(os.environ, {"LITDATA_CACHE_DIR": "/custom/cache/dir"}):
-        importlib.reload(litdata.constants)
-        importlib.reload(litdata.utilities.dataset_utilities)
-        assert litdata.utilities.dataset_utilities.get_default_cache_dir() == "/custom/cache/dir"
+        assert get_default_cache_dir() == "/custom/cache/dir"
 
 
 def test_read_updated_at_falls_back_on_truncated_local_index(tmp_path, monkeypatch):