Remove only own temp folders in clean_tmpdir

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

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()

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.

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()