Fix the bug with cleanup_raw

graphbench/datasets/_algoreas.py

@@ -173,7 +173,7 @@ def __init__(
         transform: Optional[Callable[[Data], Data]] = None,
         pre_transform: Optional[Callable[[Data], Data]] = None,
         pre_filter: Optional[Callable[[Data], bool]] = None,
-        cleanup_raw: bool = False,
+        cleanup_raw: bool = True,
     ):
         """
         Args:

graphbench/datasets/_base.py

@@ -1,3 +1,4 @@
+import shutil
 from pathlib import Path
 from abc import ABC, abstractmethod
 
@@ -16,17 +17,10 @@ def _cleanup_path(raw_path, logger=None):
         if logger is not None:
             logger.info(f"Cleaning up: {raw_dir}")
 
-        # remove only the dataset-specific temp folder
-        for path in sorted(raw_dir.rglob("*"), reverse=True):
-            try:
-                path.unlink()
-            except (IsADirectoryError, PermissionError):
-                pass
-
-        try:
-            raw_dir.rmdir()
-        except OSError:
-            pass
+        if raw_dir.is_dir():
+            shutil.rmtree(raw_dir, ignore_errors=False)
+        else:
+            raw_dir.unlink()
 
     @abstractmethod
     def _prepare(self) -> None:
@@ -79,6 +73,8 @@ def _load_cached_or_prepare(
                 if logger is not None:
                     logger.info(f"Loading cached processed data: {processed_path}")
                 self.load(resolved_load_path)
+                if cleanup_raw:
+                    self._cleanup()
                 return
             except Exception as e:
                 if logger is not None:

graphbench/datasets/_bluesky.py

@@ -284,7 +284,7 @@ def __init__(
         transform: Optional[Callable[[Data], Data]] = None,
         pre_transform: Optional[Callable[[Data], Data]] = None,
         pre_filter: Optional[Callable[[Data], bool]] = None,
-        cleanup_raw: bool = False,
+        cleanup_raw: bool = True,
         # TODO: This should be removed in the future -- the user will download these files
         load_preprocessed = True,
         feature_file_name: Union[str, Path] = _FEATURE_PT_PATH,

graphbench/datasets/_chipdesign.py

@@ -104,7 +104,7 @@ def __init__(
         transform: Optional[Callable[[Data], Data]] = None,
         pre_transform: Optional[Callable[[Data], Data]] = None,
         pre_filter: Optional[Callable[[Data], bool]] = None,
-        cleanup_raw: bool = False,  # TODO Disabling this for now since it leads to errors on my machine
+        cleanup_raw: bool = True,  # TODO Disabling this for now since it leads to errors on my machine
     ):
         """
         Args:

graphbench/datasets/_combinatorial_optimization.py

@@ -186,7 +186,7 @@ def __init__(
         pre_filter: Optional[Callable[[Data], bool]] = None,
         generate: Optional[bool] = False,
         num_samples: Optional[int] = None,
-        cleanup_raw: bool = False,
+        cleanup_raw: bool = True,
     ):
         """
         Args:
@@ -303,6 +303,8 @@ def _load_graphs(self):
             return self._generate()
 
         filepaths = self._find_matching_files(task=self.dataset_name, directory=self._raw_dir)
+        if not filepaths:
+            raise FileNotFoundError(f"No matching raw dataset files found in {self._raw_dir}")
         self.load(filepaths[0])
 
         return [self.get(i) for i in range(len(self))]
@@ -311,7 +313,16 @@ def _find_matching_files(self, directory, task, split: Optional[str] = None, siz
         """
         Returns a list of filenames matching the convention in the directory.
         """
-        return [str(self.processed_path)]
+        directory = Path(directory)
+        if not directory.exists():
+            return []
+
+        matches = sorted(
+            str(path)
+            for path in directory.rglob("data.pt")
+            if path.is_file()
+        )
+        return matches
 
     def process(self):
         self._prepare()

graphbench/datasets/_electroniccircuits.py

@@ -142,7 +142,7 @@ def __init__(
         transform: Optional[Callable[[Data], Data]] = None,
         pre_transform: Optional[Callable[[Data], Data]] = None,
         pre_filter: Optional[Callable[[Data], bool]] = None,
-        cleanup_raw: bool = False,
+        cleanup_raw: bool = True,
         target_vout : Optional[float] = None,
         vout_norm_method : Optional[Literal["min-max", "reward", "IQR", "z-score"]] = "min-max",
     ):

graphbench/datasets/_sat.py

@@ -157,7 +157,7 @@ def __init__(
         pre_transform: Optional[Callable[[Data], Data]] = None,
         pre_filter: Optional[Callable[[Data], bool]] = None,
         use_satzilla_features: bool = False,
-        cleanup_raw: bool = False,
+        cleanup_raw: bool = True,
         solver: Optional[str] = None,
     ):
         """