Support to cleanup disk cache when enabled (#816)

* Support to cleanup disk cache when enabled

* lint errors

* remove broad exception

Author: kashishmittal55

petastorm/local_disk_cache.py

@@ -21,7 +21,7 @@
 
 
 class LocalDiskCache(CacheBase):
-    def __init__(self, path, size_limit_bytes, expected_row_size_bytes, shards=6, cleanup=False, **settings):
+    def __init__(self, path, size_limit_bytes, expected_row_size_bytes, shards=6, cleanup=True, **settings):
         """LocalDiskCache is an adapter to a diskcache implementation.
 
         LocalDiskCache can be used by a petastorm Reader class to temporarily keep parts of the dataset on a local

petastorm/reader.py

@@ -448,7 +448,7 @@ def __init__(self, pyarrow_filesystem, dataset_path, schema_fields=None,
             raise NotImplementedError('Using timestamp_overlap=False is not implemented with'
                                       ' shuffle_options.shuffle_row_drop_partitions > 1')
 
-        cache = cache or NullCache()
+        self.cache = cache or NullCache()
 
         self._workers_pool = reader_pool or ThreadPool(10, shuffle_rows=shuffle_rows, seed=seed)
 
@@ -491,7 +491,7 @@ def __init__(self, pyarrow_filesystem, dataset_path, schema_fields=None,
 
         # 5. Start workers pool
         self._workers_pool.start(worker_class, (pyarrow_filesystem, dataset_path, storage_schema,
-                                                self.ngram, row_groups, cache, transform_spec,
+                                                self.ngram, row_groups, self.cache, transform_spec,
                                                 self.schema, filters, shuffle_rows, seed,
                                                 convert_early_to_numpy),
                                  ventilator=self.ventilator)
@@ -690,6 +690,14 @@ def join(self):
         """Joins all worker threads/processes. Will block until all worker workers have been fully terminated."""
         self._workers_pool.join()
 
+    def cleanup_cache(self):
+        if isinstance(self.cache, LocalDiskCache):
+            try:
+                self.cache.cleanup()
+            except (OSError, IOError, AttributeError) as e:
+                print(f"Error cleaning cache: {e}")
+        print("Cache cleanup complete.")
+
     @property
     def diagnostics(self):
         return self._workers_pool.diagnostics
@@ -719,3 +727,4 @@ def __enter__(self):
     def __exit__(self, exc_type, exc_val, exc_tb):
         self.stop()
         self.join()
+        self.cleanup_cache()

petastorm/tests/test_reader.py

@@ -94,3 +94,83 @@ def test_deprecated_shard_seed(synthetic_dataset, reader_factory):
     match_str = 'shard_seed was deprecated and will be removed in future versions.'
     with pytest.warns(UserWarning, match=match_str):
         reader_factory(synthetic_dataset.url, shard_seed=123)
+
+
+def test_cleanup_cache_with_local_disk_cache(synthetic_dataset, tmpdir):
+    """Test that cleanup_cache properly removes local disk cache directory"""
+    import os
+    cache_location = tmpdir.strpath
+
+    with make_reader(synthetic_dataset.url,
+                     cache_type='local-disk',
+                     cache_location=cache_location,
+                     cache_size_limit=1000000,
+                     cache_row_size_estimate=100) as reader:
+        # Read some data to populate cache
+        next(reader)
+        # Cache directory should exist
+        assert os.path.exists(cache_location)
+
+    # After context manager exit, cache should be cleaned up
+    assert not os.path.exists(cache_location)
+
+
+def test_cleanup_cache_with_null_cache(synthetic_dataset, capsys):
+    """Test that cleanup_cache works properly with null cache"""
+    with make_reader(synthetic_dataset.url, cache_type='null') as reader:
+        next(reader)
+        # Manually call cleanup_cache to test it
+        reader.cleanup_cache()
+
+    # Check that cleanup message was printed
+    captured = capsys.readouterr()
+    assert "Cache cleanup complete." in captured.out
+
+
+def test_cleanup_cache_manual_call(synthetic_dataset, tmpdir):
+    """Test manually calling cleanup_cache method"""
+    import os
+    cache_location = tmpdir.strpath
+
+    reader = make_reader(synthetic_dataset.url,
+                         cache_type='local-disk',
+                         cache_location=cache_location,
+                         cache_size_limit=1000000,
+                         cache_row_size_estimate=100)
+
+    try:
+        next(reader)
+        assert os.path.exists(cache_location)
+
+        # Manually call cleanup_cache
+        reader.cleanup_cache()
+        assert not os.path.exists(cache_location)
+    finally:
+        reader.stop()
+        reader.join()
+
+
+def test_cleanup_cache_exception_handling(synthetic_dataset, tmpdir, capsys, monkeypatch):
+    """Test that cleanup_cache handles exceptions gracefully"""
+    cache_location = tmpdir.strpath
+
+    with make_reader(synthetic_dataset.url,
+                     cache_type='local-disk',
+                     cache_location=cache_location,
+                     cache_size_limit=1000000,
+                     cache_row_size_estimate=100) as reader:
+        next(reader)
+
+        # Mock the cleanup method to raise an exception
+        def mock_cleanup():
+            raise OSError("Simulated cleanup error")
+
+        monkeypatch.setattr(reader.cache, 'cleanup', mock_cleanup)
+
+        # Call cleanup_cache - should handle exception gracefully
+        reader.cleanup_cache()
+
+        # Check that error message was printed
+        captured = capsys.readouterr()
+        assert "Error cleaning cache: Simulated cleanup error" in captured.out
+        assert "Cache cleanup complete." in captured.out