fusion-bench 0.2.28__py3-none-any.whl → 0.2.30__py3-none-any.whl

fusion_bench/scripts/webui.py

@@ -1,5 +1,5 @@
 """
-TODO: Per-session state management (use AppState)
+Web UI for FusionBench Command Generator with per-session state management.
 """
 
 import argparse
@@ -13,39 +13,94 @@ import hydra
 import yaml
 from colorama import Fore, Style  # For cross-platform color support
 from hydra import compose, initialize_config_dir
+from hydra.core.hydra_config import HydraConfig
 from omegaconf import DictConfig, ListConfig, OmegaConf
 
 from fusion_bench.scripts.cli import _get_default_config_path
 
 
+def escape_overrides(value: str) -> str:
+    """
+    Escapes special characters in Hydra command-line override values.
+
+    Adds quotes around values containing spaces and escapes equals signs
+    to prevent them from being interpreted as key-value separators.
+
+    Args:
+        value (str): The override value to escape.
+
+    Returns:
+        str: The escaped value ready for use in command-line overrides.
+    """
+    if " " in value and not (value.startswith('"') or value.startswith("'")):
+        return f"'{value}'"
+    if "=" in value:
+        return value.replace("=", "\\=")
+    return value
+
+
 class ConfigGroupNode:
+    """
+    Represents a node in the configuration directory tree.
+
+    This class recursively builds a tree structure representing the Hydra
+    configuration directory hierarchy, including subdirectories (child groups)
+    and YAML configuration files.
+
+    Attributes:
+        name (str): Name of the configuration group (directory name).
+        path (Path): Full path to the directory.
+        parent (Optional[ConfigGroupNode]): Parent node in the tree.
+        children (List[ConfigGroupNode]): Child directory nodes.
+        configs (List[str]): List of YAML config file names (without extension).
+    """
+
     name: str
     path: Path
-    parent: Optional["ConfigGroupNode"] = None
+    parent: Optional["ConfigGroupNode"]
     children: List["ConfigGroupNode"]
     configs: List[str]
 
-    def __init__(self, path: str | Path):
+    def __init__(self, path: str | Path, parent: Optional["ConfigGroupNode"] = None):
+        """
+        Initialize a ConfigGroupNode.
+
+        Args:
+            path: Path to the configuration directory.
+            parent: Parent node in the tree (None for root).
+        """
         self.path = Path(path)
         assert self.path.is_dir()
         self.name = self.path.stem
+        self.parent = parent
         self.children = []
         self.configs = []
         for child in self.path.iterdir():
             if child.is_dir():
-                child_node = ConfigGroupNode(child)
-                child_node.parent = self
+                child_node = ConfigGroupNode(child, parent=self)
                 self.children.append(child_node)
             elif child.is_file() and child.suffix == ".yaml":
                 self.configs.append(child.stem)
 
     def __repr__(self):
         """
-        Return string of the tree structure
+        Return a colored string representation of the tree structure.
+
+        Returns:
+            str: Tree structure with colored group names.
         """
         return f"{Fore.BLUE}{self.name}{Style.RESET_ALL}\n" + self._repr_indented()
 
     def _repr_indented(self, prefix=""):
+        """
+        Generate indented tree representation recursively.
+
+        Args:
+            prefix: String prefix for indentation.
+
+        Returns:
+            str: Indented tree structure.
+        """
         result = ""
 
         items = self.configs + self.children
@@ -63,9 +118,30 @@ class ConfigGroupNode:
         return result
 
     def has_child_group(self, name: str) -> bool:
+        """
+        Check if this node has a child group with the given name.
+
+        Args:
+            name: Name of the child group to check.
+
+        Returns:
+            bool: True if child group exists, False otherwise.
+        """
         return any(child.name == name for child in self.children)
 
     def __getitem__(self, key: str) -> Union["ConfigGroupNode", str]:
+        """
+        Get a child group or config by name.
+
+        Args:
+            key: Name of the child group or config file.
+
+        Returns:
+            Union[ConfigGroupNode, str]: Child node or config file name.
+
+        Raises:
+            KeyError: If no child group or config with that name exists.
+        """
         for child in self.children:
             if child.name == key:
                 return child
@@ -76,12 +152,28 @@ class ConfigGroupNode:
 
     @functools.cached_property
     def prefix(self) -> str:
+        """
+        Get the dot-separated prefix path from root to this node.
+
+        Returns:
+            str: Prefix path (e.g., "method.modelpool.").
+        """
         if self.parent is None:
             return ""
         return self.parent.prefix + self.name + "."
 
 
 def priority_iterable(iter, priority_keys):
+    """
+    Iterate over items with priority keys first, then remaining items.
+
+    Args:
+        iter: Iterable to process.
+        priority_keys: Keys to yield first.
+
+    Yields:
+        Items from iter, with priority_keys first.
+    """
     items = list(iter)
     for key in priority_keys:
         if key in items:
@@ -93,7 +185,16 @@ def priority_iterable(iter, priority_keys):
 
 class AppState:
     """
-    Pre-session state of the app
+    Per-session state of the app.
+
+    Manages the current configuration state including the selected config name,
+    overrides, and the composed Hydra configuration.
+
+    Attributes:
+        config_name (str): Name of the root configuration file.
+        hydra_options (List[str]): Hydra-specific command line options.
+        overrides (List[str]): List of configuration overrides.
+        config (DictConfig): The composed OmegaConf configuration.
     """
 
     config_name: str
@@ -108,6 +209,15 @@ class AppState:
         hydra_options: List[str] = [],
         overrides: List[str] = [],
     ) -> None:
+        """
+        Initialize the application state.
+
+        Args:
+            config_path: Path to the config directory.
+            config_name: Name of the root config file.
+            hydra_options: Hydra command line options.
+            overrides: Initial configuration overrides.
+        """
         super().__init__()
         self.config_path = config_path
         self.config_name = config_name
@@ -117,6 +227,12 @@ class AppState:
 
     @property
     def config_str(self):
+        """
+        Get the YAML string representation of the current configuration.
+
+        Returns:
+            str: YAML formatted configuration.
+        """
         return OmegaConf.to_yaml(self.config)
 
     def update_config(
@@ -124,6 +240,16 @@ class AppState:
         config_name: str = None,
         overrides: List[str] = None,
     ) -> "AppState":
+        """
+        Update the configuration with new name and/or overrides.
+
+        Args:
+            config_name: New root config name (optional).
+            overrides: New list of overrides (optional).
+
+        Returns:
+            AppState: Self for method chaining.
+        """
         if config_name is not None:
             self.config_name = config_name
         if overrides is not None:
@@ -133,11 +259,21 @@ class AppState:
             self.config = ""
         else:
             self.config = compose(
-                config_name=self.config_name, overrides=self.overrides
+                config_name=self.config_name,
+                overrides=self.overrides,
+                return_hydra_config=True,
             )
+            HydraConfig().set_config(self.config)
+            del self.config.hydra
         return self
 
     def generate_command(self):
+        """
+        Generate the fusion_bench CLI command from current state.
+
+        Returns:
+            str: Complete command ready to execute in shell.
+        """
         # Generate the command according to `config_name` and `overrides` (a list of strings)
         command = "fusion_bench \\\n"
         if self.config_path is not None:
@@ -152,23 +288,66 @@ class AppState:
 
     @property
     def config_str_and_command(self):
+        """
+        Get both config string and command as a tuple.
+
+        Returns:
+            Tuple[str, str]: (YAML config, shell command).
+        """
         return self.config_str, self.generate_command()
 
     def get_override(self, key: str):
+        """
+        Get the override value for a specific key.
+
+        Args:
+            key: Configuration key to look up.
+
+        Returns:
+            Optional[str]: Override value or None if not found.
+        """
         for ov in self.overrides:
             if ov.startswith(f"{key}="):
                 return "".join(ov.split("=")[1:])
         return None
 
     def update_override(self, key: str, value):
+        """
+        Update or add an override for a specific key.
+
+        Args:
+            key: Configuration key to override.
+            value: New value for the key.
+
+        Returns:
+            AppState: Updated state after recomposing config.
+        """
         self.overrides = [ov for ov in self.overrides if not ov.startswith(f"{key}=")]
         if value:
-            self.overrides.append(f"{key}={value}")
+            self.overrides.append(f"{key}={escape_overrides(value)}")
         return self.update_config()
 
 
 class App:
+    """
+    Main application class for the FusionBench WebUI.
+
+    Manages the Gradio interface, configuration tree, and application state.
+
+    Attributes:
+        args: Command line arguments.
+        group_tree (ConfigGroupNode): Root of the config directory tree.
+        init_config_name (str): Initial configuration name.
+        app_state (AppState): Current application state.
+    """
+
     def __init__(self, args):
+        """
+        Initialize the application.
+
+        Args:
+            args: Parsed command line arguments.
+        """
         super().__init__()
         self.args = args
         group_tree = ConfigGroupNode(self.config_path)
@@ -177,8 +356,8 @@ class App:
 
         self.group_tree = group_tree
 
-        if "example_config" in group_tree.configs:
-            self.init_config_name = "example_config"
+        if "fabric_model_fusion" in group_tree.configs:
+            self.init_config_name = "fabric_model_fusion"
         else:
             self.init_config_name = group_tree.configs[0]
 
@@ -197,31 +376,64 @@ class App:
 
     @functools.cached_property
     def config_path(self):
+        """
+        Get the configuration directory path.
+
+        Returns:
+            Path: Path to the config directory.
+        """
         if self.args.config_path:
             return Path(self.args.config_path)
         else:
             return _get_default_config_path()
 
     def __getattr__(self, name):
+        """
+        Delegate attribute access to app_state if not found in App.
+
+        Args:
+            name: Attribute name.
+
+        Returns:
+            Attribute value from app_state.
+
+        Raises:
+            AttributeError: If attribute not found in app_state either.
+        """
         if hasattr(self.app_state, name):
             return getattr(self.app_state, name)
         raise AttributeError(f"App object has no attribute {name}")
 
     def generate_ui(self):
+        """
+        Generate the Gradio user interface.
+
+        Creates interactive UI components for configuration selection,
+        parameter editing, and command generation.
+
+        Returns:
+            gr.Blocks: Gradio application instance.
+        """
         with gr.Blocks() as app:
             gr.Markdown("# FusionBench Command Generator")
 
             # 1. Choose a root config file
-            root_configs = gr.Dropdown(
-                choices=self.group_tree.configs,
-                value=self.config_name,
-                label="Root Config",
-            )
+            with gr.Row(equal_height=True):
+                root_configs = gr.Dropdown(
+                    choices=self.group_tree.configs,
+                    value=self.config_name,
+                    label="Root Config",
+                    scale=4,
+                )
+                reset_button = gr.Button("Reset", scale=1)
 
             with gr.Row():
                 with gr.Column(scale=2):
                     command_output = gr.Code(
-                        language="shell", label="Generated Command"
+                        value=self.app_state.generate_command(),
+                        language="shell",
+                        label="Generated Command",
+                        interactive=False,
                     )
 
                     @gr.render(inputs=[root_configs, command_output])
@@ -352,10 +564,26 @@ class App:
                 outputs=[config_output, command_output],
             )
 
+            def reset_app(config_name):
+                # Reset overrides and update config
+                self.app_state.overrides = []
+                return self.app_state.update_config(config_name).config_str_and_command
+
+            reset_button.click(
+                reset_app,
+                inputs=[root_configs],
+                outputs=[config_output, command_output],
+            )
         return app
 
 
 def parse_args():
+    """
+    Parse command line arguments for the WebUI.
+
+    Returns:
+        argparse.Namespace: Parsed arguments.
+    """
     parser = argparse.ArgumentParser(description="FusionBench Command Generator")
     parser.add_argument(
         "--config-path",
@@ -390,6 +618,11 @@ def parse_args():
 
 
 def main() -> None:
+    """
+    Main entry point for the FusionBench WebUI application.
+
+    Parses arguments, initializes the app, and launches the Gradio interface.
+    """
     args = parse_args()
 
     app = App(args).generate_ui()

fusion_bench/utils/__init__.py

@@ -93,6 +93,13 @@ _import_structure = {
         "StateDictType",
         "TorchModelType",
     ],
+    "validation": [
+        "validate_path_exists",
+        "validate_file_exists",
+        "validate_directory_exists",
+        "validate_model_name",
+        "ValidationError",
+    ],
 }
 
 if TYPE_CHECKING:
@@ -159,6 +166,13 @@ if TYPE_CHECKING:
     )
     from .timer import timeit_context
     from .type import BoolStateDictType, StateDictType, TorchModelType
+    from .validation import (
+        ValidationError,
+        validate_directory_exists,
+        validate_file_exists,
+        validate_model_name,
+        validate_path_exists,
+    )
 
 else:
     sys.modules[__name__] = LazyImporter(

fusion_bench/utils/data.py

@@ -7,6 +7,8 @@ import torch
 import torch.utils.data
 from torch.utils.data import DataLoader, Dataset
 
+from fusion_bench.utils.validation import ValidationError, validate_file_exists
+
 
 class InfiniteDataLoader:
     """
@@ -18,23 +20,105 @@ class InfiniteDataLoader:
 
     Attributes:
         data_loader (DataLoader): The DataLoader to wrap.
-        data_iter (iterator): An iterator over the DataLoader.
+        _data_iter (iterator): An iterator over the DataLoader.
+        _iteration_count (int): Number of complete iterations through the dataset.
+
+    Example:
+        >>> train_loader = DataLoader(dataset, batch_size=32)
+        >>> infinite_loader = InfiniteDataLoader(train_loader)
+        >>> for i, batch in enumerate(infinite_loader):
+        ...     if i >= 1000:  # Train for 1000 steps
+        ...         break
+        ...     train_step(batch)
     """
 
-    def __init__(self, data_loader: DataLoader):
+    def __init__(self, data_loader: DataLoader, max_retries: int = 1):
+        """
+        Initialize the InfiniteDataLoader.
+
+        Args:
+            data_loader: The DataLoader to wrap.
+            max_retries: Maximum number of retry attempts when resetting the data loader (default: 1).
+
+        Raises:
+            ValidationError: If data_loader is None or not a DataLoader instance.
+        """
+        if data_loader is None:
+            raise ValidationError(
+                "data_loader cannot be None", field="data_loader", value=data_loader
+            )
+
         self.data_loader = data_loader
-        self.data_iter = iter(data_loader)
+        self.max_retries = max_retries
+        self._data_iter = iter(data_loader)
+        self._iteration_count = 0
 
     def __iter__(self):
+        """Reset the iterator to the beginning."""
+        self._data_iter = iter(self.data_loader)
+        self._iteration_count = 0
         return self
 
     def __next__(self):
-        try:
-            data = next(self.data_iter)
-        except StopIteration:
-            self.data_iter = iter(self.data_loader)  # Reset the data loader
-            data = next(self.data_iter)
-        return data
+        """
+        Get the next batch, resetting to the beginning when the dataset is exhausted.
+
+        Returns:
+            The next batch from the data loader.
+
+        Raises:
+            RuntimeError: If the data loader consistently fails to produce data.
+        """
+        last_exception = None
+        for attempt in range(self.max_retries):
+            try:
+                data = next(self._data_iter)
+                return data
+            except StopIteration:
+                # Dataset exhausted or dataloader is empty, reset to beginning
+                self._iteration_count += 1
+                try:
+                    self._data_iter = iter(self.data_loader)
+                    data = next(self._data_iter)
+                    return data
+                except Exception as e:
+                    last_exception = e
+                    continue
+            except Exception as e:
+                # Handle other potential errors from the data loader
+                raise RuntimeError(
+                    f"Error retrieving data from data loader: [{type(e).__name__}]{e}"
+                ) from e
+
+        # If we get here, all attempts failed
+        raise RuntimeError(
+            f"Failed to retrieve data from data loader after {self.max_retries} attempts. "
+            f"Last error: [{type(last_exception).__name__}]{last_exception}. "
+            + (
+                f"The data loader appears to be empty."
+                if isinstance(last_exception, StopIteration)
+                else ""
+            )
+        ) from last_exception
+
+    def reset(self):
+        """Manually reset the iterator to the beginning of the dataset."""
+        self._data_iter = iter(self.data_loader)
+        self._iteration_count = 0
+
+    @property
+    def iteration_count(self) -> int:
+        """Get the number of complete iterations through the dataset."""
+        return self._iteration_count
+
+    def __len__(self) -> int:
+        """
+        Return the length of the underlying data loader.
+
+        Returns:
+            The number of batches in one complete iteration.
+        """
+        return len(self.data_loader)
 
 
 def load_tensor_from_file(
@@ -50,7 +134,14 @@ def load_tensor_from_file(
 
     Returns:
         torch.Tensor: The tensor loaded from the file.
+
+    Raises:
+        ValidationError: If the file doesn't exist
+        ValueError: If the file format is unsupported
     """
+    # Validate file exists
+    validate_file_exists(file_path)
+
     if file_path.endswith(".np"):
         tensor = torch.from_numpy(np.load(file_path)).detach_()
     if file_path.endswith((".pt", ".pth")):

fusion_bench/utils/devices.py

@@ -32,11 +32,13 @@ def clear_cuda_cache():
     Clears the CUDA memory cache to free up GPU memory.
     Works only if CUDA is available.
     """
+
     gc.collect()
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
+        gc.collect()
     else:
-        log.warning("CUDA is not available. No cache to clear.")
+        log.debug("CUDA is not available. No cache to clear.")
 
 
 def to_device(