From d931745699cfd2c55f03d9b0e6f312138661d0a2 Mon Sep 17 00:00:00 2001 From: Kostiantyn Dementiev <68417998+konstde00@users.noreply.github.com> Date: Wed, 14 Oct 2026 19:09:24 +0200 Subject: [PATCH] Remove only own temp folders in clean_tmpdir Currently clean_tmpdir() removes the whole cache_dir/tmp, so when one job finishes a build, it deletes the temp folders of another job which is still building on the same cache_dir. Now each dataset keeps the list of temp folders it created and removes only them, and tmp/ itself is removed only when it's empty. rmtree uses ignore_errors=True, because it runs in finally blocks and a cleanup error should not hide the build's own error. Added 4 tests (TestTmpDirOwnership), a paragraph in docs/api/datasets.rst and examples/dataset_cache_shared_tmpdir.py. --- docs/api/datasets.rst | 5 +++ examples/dataset_cache_shared_tmpdir.py | 49 +++++++++++++++++++++++ pyhealth/datasets/base_dataset.py | 21 ++++++++-- tests/core/test_caching.py | 53 +++++++++++++++++++++++++ 4 files changed, 124 insertions(+), 4 deletions(-) create mode 100644 examples/dataset_cache_shared_tmpdir.py diff --git a/docs/api/datasets.rst b/docs/api/datasets.rst index f629bef8a..8ff11e34d 100644 --- a/docs/api/datasets.rst +++ b/docs/api/datasets.rst @@ -38,6 +38,11 @@ tables read. Editing the YAML or rewriting a source file (even at the same path) therefore builds a fresh cache instead of silently reusing a stale one. Old cache folders are left in place; delete them to reclaim disk space. +Temporary files (converted source tables, Dask scratch space) go to folders +under ``tmp/`` in the cache directory. Each dataset removes only the folders it +created, so a job that finishes does not delete the temporary files of another +job using the same ``cache_dir``. + The result is a :class:`~pyhealth.datasets.BaseDataset` — a structured patient→event tree. It is different from a PyTorch Dataset: it has no integer length and you cannot index into it with ``dataset[i]``. Think of it as a diff --git a/examples/dataset_cache_shared_tmpdir.py b/examples/dataset_cache_shared_tmpdir.py new file mode 100644 index 000000000..11f653193 --- /dev/null +++ b/examples/dataset_cache_shared_tmpdir.py @@ -0,0 +1,49 @@ +"""Jobs that share a cache_dir do not delete each other's temporary files. + +While it builds, a dataset keeps temporary files (converted source tables, Dask +scratch space) in folders under ``tmp/`` in its cache directory, and removes +them when the build ends. Two jobs started with the same configuration, such as +two runs of a hyperparameter sweep, share that cache directory. Each job removes +only the folders it created, so the job that finishes first no longer deletes +files the other job is still reading. + +Runs in a few seconds on CPU; nothing is built or downloaded. + +Usage: + python examples/dataset_cache_shared_tmpdir.py +""" + +import tempfile + +from pyhealth.datasets import BaseDataset + + +def open_dataset(cache: str) -> BaseDataset: + return BaseDataset( + root="", tables=[], dataset_name="shared_tmpdir_demo", cache_dir=cache + ) + + +def main(): + with tempfile.TemporaryDirectory() as cache: + job_a = open_dataset(cache) + job_b = open_dataset(cache) + print("same cache directory:", job_a.cache_dir == job_b.cache_dir) + + # Job A is mid-build and has converted a source table into tmp/. + table = job_a.create_tmpdir() / "table.parquet" + table.write_text("job A's converted table") + + # Job B finishes its own build, which ends with clean_tmpdir(). + job_b.create_tmpdir() + job_b.clean_tmpdir() + kept = table.exists() + print("job A's file after job B's cleanup:", "kept" if kept else "deleted") + + job_a.clean_tmpdir() + tmp_root = job_a.cache_dir / "tmp" + print("tmp/ removed once both jobs are done:", not tmp_root.exists()) + + +if __name__ == "__main__": + main() diff --git a/pyhealth/datasets/base_dataset.py b/pyhealth/datasets/base_dataset.py index 98ff1c12f..64782492f 100644 --- a/pyhealth/datasets/base_dataset.py +++ b/pyhealth/datasets/base_dataset.py @@ -466,6 +466,7 @@ def __init__( # Cached attributes self.cache_dir = self._init_cache_dir(cache_dir) + self._tmp_dirs: list[Path] = [] # created by create_tmpdir() self._global_event_df = None self._unique_patient_ids = None @@ -540,13 +541,25 @@ def create_tmpdir(self) -> Path: """ tmp_dir = self.cache_dir / "tmp" / str(uuid.uuid4()) tmp_dir.mkdir(parents=True, exist_ok=True) + self._tmp_dirs.append(tmp_dir) return tmp_dir def clean_tmpdir(self) -> None: - """Cleans up the temporary directory within the cache.""" - tmp_dir = self.cache_dir / "tmp" - if tmp_dir.exists(): - shutil.rmtree(tmp_dir) + """Removes the temporary directories this dataset created. + + Other jobs may share the same cache directory and still be using their + own directories under ``tmp/``, so those are left alone. ``tmp/`` itself + is removed once it is empty. + """ + for tmp_dir in self._tmp_dirs: + # Called from ``finally`` blocks: a failed cleanup must not hide the + # error that ended the build. + shutil.rmtree(tmp_dir, ignore_errors=True) + self._tmp_dirs.clear() + try: + (self.cache_dir / "tmp").rmdir() + except OSError: + pass # missing, or still holds another job's directories def _scan_table(self, source_path: str) -> dd.DataFrame: """Routes a table source to the appropriate scanner based on its format. diff --git a/tests/core/test_caching.py b/tests/core/test_caching.py index 65556ab0d..9eacc0058 100644 --- a/tests/core/test_caching.py +++ b/tests/core/test_caching.py @@ -335,5 +335,58 @@ def test_incomplete_parquet_cache_triggers_rebuild(self): self.assertFalse(stale_file.exists()) +class TestTmpDirOwnership(BaseTestCase): + """Jobs sharing a cache_dir must not delete each other's temporary files. + + Before, clean_tmpdir() removed all of tmp/, so a job that finished a build + deleted the converted tables and Dask scratch of a job still building. + """ + + def setUp(self): + temp_dir = tempfile.TemporaryDirectory() + self.addCleanup(temp_dir.cleanup) + # Same configuration, so one cache directory: like two jobs of a sweep. + self.mine = MockDataset(cache_dir=temp_dir.name) + self.other = MockDataset(cache_dir=temp_dir.name) + self.assertEqual(self.mine.cache_dir, self.other.cache_dir) + + def test_clean_tmpdir_keeps_other_datasets_tmpdirs(self): + mine_dir = self.mine.create_tmpdir() + other_file = self.other.create_tmpdir() / "table.parquet" + other_file.write_text("in use") + + self.mine.clean_tmpdir() + + self.assertFalse(mine_dir.exists()) + self.assertTrue(other_file.exists()) + + def test_failed_build_keeps_other_datasets_tmpdirs(self): + # _event_transform calls clean_tmpdir() in its finally block. + other_file = self.other.create_tmpdir() / "table.parquet" + other_file.write_text("in use") + + with ( + patch.object(self.mine, "load_data", side_effect=RuntimeError("boom")), + self.assertRaises(RuntimeError), + ): + self.mine._event_transform(self.mine.cache_dir / "out.parquet") + + self.assertTrue(other_file.exists()) + + def test_tmp_root_is_removed_once_empty(self): + tmp_root = self.mine.cache_dir / "tmp" + self.mine.create_tmpdir() + self.other.create_tmpdir() + + self.mine.clean_tmpdir() + self.assertTrue(tmp_root.exists()) + self.other.clean_tmpdir() + self.assertFalse(tmp_root.exists()) + + def test_clean_tmpdir_without_tmpdirs(self): + self.mine.clean_tmpdir() + self.assertFalse((self.mine.cache_dir / "tmp").exists()) + + if __name__ == "__main__": unittest.main()