From cf0027e35e8b0956141f24e721a5ee506bb955de Mon Sep 17 00:00:00 2001 From: Gia Chi Dang Date: Fri, 10 Jul 2026 00:33:26 +0200 Subject: [PATCH] fix cleanup_raw --- graphbench/datasets/_algoreas.py | 2 +- graphbench/datasets/_base.py | 18 +++++++----------- graphbench/datasets/_bluesky.py | 2 +- graphbench/datasets/_chipdesign.py | 2 +- .../datasets/_combinatorial_optimization.py | 15 +++++++++++++-- graphbench/datasets/_electroniccircuits.py | 2 +- graphbench/datasets/_sat.py | 2 +- 7 files changed, 25 insertions(+), 18 deletions(-) diff --git a/graphbench/datasets/_algoreas.py b/graphbench/datasets/_algoreas.py index 1d22244..2279cf2 100644 --- a/graphbench/datasets/_algoreas.py +++ b/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: diff --git a/graphbench/datasets/_base.py b/graphbench/datasets/_base.py index 1bba3d7..4fa91a0 100644 --- a/graphbench/datasets/_base.py +++ b/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: diff --git a/graphbench/datasets/_bluesky.py b/graphbench/datasets/_bluesky.py index aa5a7bd..ba3a81e 100644 --- a/graphbench/datasets/_bluesky.py +++ b/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, diff --git a/graphbench/datasets/_chipdesign.py b/graphbench/datasets/_chipdesign.py index 35440b9..f0044dc 100644 --- a/graphbench/datasets/_chipdesign.py +++ b/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: diff --git a/graphbench/datasets/_combinatorial_optimization.py b/graphbench/datasets/_combinatorial_optimization.py index 106c201..bb4482d 100644 --- a/graphbench/datasets/_combinatorial_optimization.py +++ b/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() diff --git a/graphbench/datasets/_electroniccircuits.py b/graphbench/datasets/_electroniccircuits.py index 9087bb9..a6b9729 100644 --- a/graphbench/datasets/_electroniccircuits.py +++ b/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", ): diff --git a/graphbench/datasets/_sat.py b/graphbench/datasets/_sat.py index 46142c4..1858dea 100644 --- a/graphbench/datasets/_sat.py +++ b/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, ): """