init commit, tldr the cache_dir or files should probably indicate tab… (#835)

* init commit, tldr the cache_dir or files should probably indicate table names

* better approach is to hash the filename rather than the directory name

* Fix cache dir

* Simplify new caching behavior code and add unit tests

* Remove unnecessary f-string

* commit to remove cache_dir from task, will update after seeing if test cases pass

* Add proc UUID to samples cache dir docs

* Remove the cache_format param from set_task

* Remove task schema from the proc UUID (since it is now included in the task UUID)

* Add proc UUID to the _init_cache_dir docs

* Remove cache_format param from unit tests

* Remove set_task(cache_dir) from more examples

* Fix cache unit test (call_count increments trigger new caches and we can no longer override the default task cache)

* update cache size directory check to be more in line with our tasks/ subdir setup with how cache_dir works

* Remove unused ensure_empty_dir function

* Correct cache size calculations (the "workers_n" examples do not need to be corrected because they were already using the default task cache path)

---------

Co-authored-by: Yongda Fan <yongdaf2@illinois.edu>
Co-authored-by: EricSchrock <eric.schrock22@gmail.com>

Author: 3 people

examples/benchmark_perf/benchmark_workers_1.py

@@ -117,6 +117,7 @@ def main():
     print("\n[1/2] Loading MIMIC-IV base dataset...")
     dataset_start = time.time()
 
+    base_cache_dir = f"{cache_root}/base_dataset"
     base_dataset = MIMIC4Dataset(
         ehr_root="/srv/local/data/physionet.org/files/mimiciv/2.2/",
         ehr_tables=[
@@ -127,7 +128,7 @@ def main():
             "labevents",
         ],
         dev=dev,
-        cache_dir=f"{cache_root}/base_dataset",
+        cache_dir=base_cache_dir,
     )
 
     dataset_time = time.time() - dataset_start
@@ -140,20 +141,17 @@ def main():
     sample_dataset = base_dataset.set_task(
         MortalityPredictionStageNetMIMIC4(),
         num_workers=1,
-        cache_dir=f"{cache_root}/task_samples",
     )
 
     task_time = time.time() - task_start
     print(f"✓ Task processing completed in {task_time:.2f} seconds")
 
     # Measure cache sizes
     print("\n[3/3] Measuring cache sizes...")
-    base_cache_dir = f"{cache_root}/base_dataset"
-    task_cache_dir = f"{cache_root}/task_samples"
 
-    base_cache_size = get_directory_size(base_cache_dir)
-    task_cache_size = get_directory_size(task_cache_dir)
-    total_cache_size = base_cache_size + task_cache_size
+    total_cache_size = get_directory_size(base_cache_dir)
+    task_cache_size = get_directory_size(f"{base_cache_dir}/tasks")
+    base_cache_size = total_cache_size - task_cache_size
 
     print(f"✓ Base dataset cache: {format_size(base_cache_size)}")
     print(f"✓ Task samples cache: {format_size(task_cache_size)}")

examples/benchmark_perf/benchmark_workers_12.py

@@ -117,6 +117,7 @@ def main():
     print("\n[1/2] Loading MIMIC-IV base dataset...")
     dataset_start = time.time()
 
+    base_cache_dir = f"{cache_root}/base_dataset"
     base_dataset = MIMIC4Dataset(
         ehr_root="/srv/local/data/physionet.org/files/mimiciv/2.2/",
         ehr_tables=[
@@ -127,7 +128,7 @@ def main():
             "labevents",
         ],
         dev=dev,
-        cache_dir=f"{cache_root}/base_dataset",
+        cache_dir=base_cache_dir,
     )
 
     dataset_time = time.time() - dataset_start
@@ -140,20 +141,17 @@ def main():
     sample_dataset = base_dataset.set_task(
         MortalityPredictionStageNetMIMIC4(),
         num_workers=12,
-        cache_dir=f"{cache_root}/task_samples",
     )
 
     task_time = time.time() - task_start
     print(f"✓ Task processing completed in {task_time:.2f} seconds")
 
     # Measure cache sizes
     print("\n[3/3] Measuring cache sizes...")
-    base_cache_dir = f"{cache_root}/base_dataset"
-    task_cache_dir = f"{cache_root}/task_samples"
 
-    base_cache_size = get_directory_size(base_cache_dir)
-    task_cache_size = get_directory_size(task_cache_dir)
-    total_cache_size = base_cache_size + task_cache_size
+    total_cache_size = get_directory_size(base_cache_dir)
+    task_cache_size = get_directory_size(f"{base_cache_dir}/tasks")
+    base_cache_size = total_cache_size - task_cache_size
 
     print(f"✓ Base dataset cache: {format_size(base_cache_size)}")
     print(f"✓ Task samples cache: {format_size(task_cache_size)}")

examples/benchmark_perf/benchmark_workers_4.py

@@ -117,6 +117,7 @@ def main():
     print("\n[1/2] Loading MIMIC-IV base dataset...")
     dataset_start = time.time()
 
+    base_cache_dir = f"{cache_root}/base_dataset"
     base_dataset = MIMIC4Dataset(
         ehr_root="/srv/local/data/physionet.org/files/mimiciv/2.2/",
         ehr_tables=[
@@ -127,7 +128,7 @@ def main():
             "labevents",
         ],
         dev=dev,
-        # cache_dir=f"{cache_root}/base_dataset",
+        cache_dir=base_cache_dir,
     )
 
     dataset_time = time.time() - dataset_start
@@ -140,20 +141,17 @@ def main():
     sample_dataset = base_dataset.set_task(
         MortalityPredictionStageNetMIMIC4(),
         num_workers=4,
-        cache_dir=f"{cache_root}/task_samples_old",
     )
 
     task_time = time.time() - task_start
     print(f"✓ Task processing completed in {task_time:.2f} seconds")
 
     # Measure cache sizes
     print("\n[3/3] Measuring cache sizes...")
-    base_cache_dir = f"{cache_root}/base_dataset"
-    task_cache_dir = f"{cache_root}/task_samples"
 
-    base_cache_size = get_directory_size(base_cache_dir)
-    task_cache_size = get_directory_size(task_cache_dir)
-    total_cache_size = base_cache_size + task_cache_size
+    total_cache_size = get_directory_size(base_cache_dir)
+    task_cache_size = get_directory_size(f"{base_cache_dir}/tasks")
+    base_cache_size = total_cache_size - task_cache_size
 
     print(f"✓ Base dataset cache: {format_size(base_cache_size)}")
     print(f"✓ Task samples cache: {format_size(task_cache_size)}")

examples/benchmark_perf/benchmark_workers_n.py

@@ -159,13 +159,6 @@ def parse_workers(value: str) -> list[int]:
     return workers
 
 
-def ensure_empty_dir(path: str | Path) -> None:
-    p = Path(path)
-    if p.exists():
-        shutil.rmtree(p)
-    p.mkdir(parents=True, exist_ok=True)
-
-
 def remove_dir(path: str | Path, retries: int = 3, delay: float = 1.0) -> None:
     """Remove a directory with retry logic for busy file handles."""
     p = Path(path)
@@ -282,11 +275,8 @@ def main() -> None:
     print("\n[1/1] Sweeping num_workers (each run reloads dataset + task)...")
     for w in args.workers:
         for r in range(args.repeats):
-            task_cache_dir = cache_root / f"task_samples_w{w}"
-
             # Ensure no cache artifacts before this run.
             remove_dir(base_cache_dir)
-            ensure_empty_dir(task_cache_dir)
 
             tracker.reset()
             run_start = time.time()
@@ -311,13 +301,13 @@ def main() -> None:
             sample_dataset = base_dataset.set_task(
                 MortalityPredictionStageNetMIMIC4(),
                 num_workers=w,
-                cache_dir=str(task_cache_dir),
             )
 
             task_process_s = time.time() - task_start
             total_s = time.time() - run_start
             peak_rss_bytes = tracker.peak_bytes()
-            task_cache_bytes = get_directory_size(task_cache_dir)
+            tasks_dir = base_cache_dir / "tasks"
+            task_cache_bytes = get_directory_size(tasks_dir)
 
             # Capture sample count BEFORE cleaning up the cache (litdata needs it).
             num_samples = len(sample_dataset)
@@ -327,7 +317,6 @@ def main() -> None:
             del base_dataset
 
             # Clean up to avoid disk growth across an overnight sweep.
-            remove_dir(task_cache_dir)
             remove_dir(base_cache_dir)
 
             results.append(

examples/benchmark_perf/benchmark_workers_n_drug_recommendation.py

@@ -159,13 +159,6 @@ def parse_workers(value: str) -> list[int]:
     return workers
 
 
-def ensure_empty_dir(path: str | Path) -> None:
-    p = Path(path)
-    if p.exists():
-        shutil.rmtree(p)
-    p.mkdir(parents=True, exist_ok=True)
-
-
 def remove_dir(path: str | Path, retries: int = 3, delay: float = 1.0) -> None:
     """Remove a directory with retry logic for busy file handles."""
     p = Path(path)
@@ -284,11 +277,8 @@ def main() -> None:
     print("\n[1/1] Sweeping num_workers (each run reloads dataset + task)...")
     for w in args.workers:
         for r in range(args.repeats):
-            task_cache_dir = cache_root / f"task_samples_drug_rec_w{w}"
-
             # Ensure no cache artifacts before this run.
             remove_dir(base_cache_dir)
-            ensure_empty_dir(task_cache_dir)
 
             tracker.reset()
             run_start = time.time()
@@ -313,13 +303,13 @@ def main() -> None:
             sample_dataset = base_dataset.set_task(
                 DrugRecommendationMIMIC4(),
                 num_workers=w,
-                cache_dir=str(task_cache_dir),
             )
 
             task_process_s = time.time() - task_start
             total_s = time.time() - run_start
             peak_rss_bytes = tracker.peak_bytes()
-            task_cache_bytes = get_directory_size(task_cache_dir)
+            tasks_dir = base_cache_dir / "tasks"
+            task_cache_bytes = get_directory_size(tasks_dir)
 
             # Capture sample count BEFORE cleaning up the cache (litdata needs it).
             num_samples = len(sample_dataset)
@@ -329,7 +319,6 @@ def main() -> None:
             del base_dataset
 
             # Clean up to avoid disk growth across an overnight sweep.
-            remove_dir(task_cache_dir)
             remove_dir(base_cache_dir)
 
             results.append(

examples/benchmark_perf/benchmark_workers_n_length_of_stay.py

@@ -159,13 +159,6 @@ def parse_workers(value: str) -> list[int]:
     return workers
 
 
-def ensure_empty_dir(path: str | Path) -> None:
-    p = Path(path)
-    if p.exists():
-        shutil.rmtree(p)
-    p.mkdir(parents=True, exist_ok=True)
-
-
 def remove_dir(path: str | Path, retries: int = 3, delay: float = 1.0) -> None:
     """Remove a directory with retry logic for busy file handles."""
     p = Path(path)
@@ -284,11 +277,8 @@ def main() -> None:
     print("\n[1/1] Sweeping num_workers (each run reloads dataset + task)...")
     for w in args.workers:
         for r in range(args.repeats):
-            task_cache_dir = cache_root / f"task_samples_los_w{w}"
-
             # Ensure no cache artifacts before this run.
             remove_dir(base_cache_dir)
-            ensure_empty_dir(task_cache_dir)
 
             tracker.reset()
             run_start = time.time()
@@ -313,13 +303,13 @@ def main() -> None:
             sample_dataset = base_dataset.set_task(
                 LengthOfStayPredictionMIMIC4(),
                 num_workers=w,
-                cache_dir=str(task_cache_dir),
             )
 
             task_process_s = time.time() - task_start
             total_s = time.time() - run_start
             peak_rss_bytes = tracker.peak_bytes()
-            task_cache_bytes = get_directory_size(task_cache_dir)
+            tasks_dir = base_cache_dir / "tasks"
+            task_cache_bytes = get_directory_size(tasks_dir)
 
             # Capture sample count BEFORE cleaning up the cache (litdata needs it).
             num_samples = len(sample_dataset)
@@ -329,7 +319,6 @@ def main() -> None:
             del base_dataset
 
             # Clean up to avoid disk growth across an overnight sweep.
-            remove_dir(task_cache_dir)
             remove_dir(base_cache_dir)
 
             results.append(