test: isolate cache dirs for tests running in parallel

.claude/skills/litdata/reference/contributing.md

@@ -68,7 +68,7 @@ pytest tests/path/test_x.py::test_name -v --capture=no    # see testing.md
 **`ci-testing.yml`** — job `pytester` (matrix, `fail-fast: false`):
 
 - OS: `ubuntu-22.04`, `macos-14`, `windows-2022`. Python: `3.10–3.14` (windows excludes 3.13/3.14).
-- Two phases: fast tests parallel (`pytest tests --ignore=tests/processing --ignore=tests/raw -n 2 --dist=loadgroup ...`), then processing/raw sequential. `--dist=loadgroup` is **required** (see testing.md).
+- Two phases: fast tests parallel (`pytest tests --ignore=tests/processing --ignore=tests/raw -n 2 ...`), then processing/raw sequential.
 - `--timeout=120` per test; FFmpeg installed for video tests; `UV_TORCH_BACKEND=cpu`.
 - `testing-guardian` job is the required gate — fails the PR if any leg fails.
 

.claude/skills/litdata/reference/testing.md

@@ -13,11 +13,11 @@ pytest tests/streaming/ -v
 
 # How CI runs it (authoritative). Two phases because processing tests must be serial:
 pytest tests --ignore=tests/processing --ignore=tests/raw \
-    -n 2 --dist=loadgroup --cov=litdata --durations=0 --timeout=120 --capture=no --verbose
+    -n 2 --cov=litdata --durations=0 --timeout=120 --capture=no --verbose
 pytest tests/processing tests/raw --cov=litdata --cov-append --durations=0 --timeout=120 --capture=no --verbose
 ```
 
-**`--dist=loadgroup` is required, not optional, when running with `-n>1`.** `tests/conftest.py:165` (`pytest_collection_modifyitems`) pins every test using the `clean_pq_index_cache` fixture onto a single xdist worker via `xdist_group`, because that fixture `shutil.rmtree`s `~/.lightning/chunks` and would otherwise race across workers.
+Tests are safe to run under `-n>1`: the autouse `isolated_cache_dirs` fixture gives every test its own cache dirs, so parallel workers can't delete each other's chunks.
 
 The `Makefile test` target is stale (`--flake8` flag, points at `src`) — prefer the CI commands above.
 
@@ -75,8 +75,8 @@ Use stdlib `unittest.mock` (`Mock`, `MagicMock`, `patch`) + the pytest `monkeypa
 ## Useful fixtures (`tests/conftest.py`)
 
 - **Autouse/session**: `teardown_process_group` (destroys torch.distributed group), `_thread_police` (fails a test leaking a non-daemon "zombie" thread — special-cases `PrepareChunksThread`, `QueueFeederThread`, `pytest_timeout`).
-- **Autouse/per-test**: `disable_signals` (no-ops `signal.signal` so worker signal handlers don't break tests).
-- **Data**: `mosaic_mds_index_data`, `prepare_combined_dataset` (two on-disk 50-item `Cache` datasets), `combined_dataset` (a `CombinedStreamingDataset`), `pq_data`, `write_pq_data` (5 polars parquet files), `clean_pq_index_cache` (wipes the default cache dir — the one that triggers xdist grouping).
+- **Autouse/per-test**: `isolated_cache_dirs` (points `LITDATA_CACHE_DIR`, `DATA_OPTIMIZER_CACHE_FOLDER` and `DATA_OPTIMIZER_DATA_CACHE_FOLDER` at a per-test temp dir and unsets `LITDATA_HF_CACHE_DIR`; drop the env var with `monkeypatch.delenv` to test defaults), `disable_signals` (no-ops `signal.signal` so worker signal handlers don't break tests).
+- **Data**: `mosaic_mds_index_data`, `prepare_combined_dataset` (two on-disk 50-item `Cache` datasets), `combined_dataset` (a `CombinedStreamingDataset`), `pq_data`, `write_pq_data` (5 polars parquet files).
 
 Local idiom: most files define their own `seed_everything(seed)`. CLI tests use a `run_cli(args_list)` helper (patches `sys.argv`, captures stdout). `tests/streaming/utils.py` gives `filter_lock_files` / `get_lock_files` to assert on chunk dirs ignoring `.lock`/`.cnt` artifacts.
 

.github/workflows/ci-testing.yml

@@ -107,7 +107,7 @@ jobs:
           pytest tests \
               --ignore=tests/processing \
               --ignore=tests/raw \
-              -n 2 --dist=loadgroup --cov=litdata --durations=0 --timeout=120 --capture=no --verbose
+              -n 2 --cov=litdata --durations=0 --timeout=120 --capture=no --verbose
 
       - name: Run processing tests sequentially
         run: |

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,25 @@ 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")
+
+    monkeypatch.setenv("LITDATA_CACHE_DIR", str(root / "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 +189,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 +214,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):