fusion-bench 0.2.27__py3-none-any.whl → 0.2.29__py3-none-any.whl

fusion_bench/utils/fabric.py

@@ -1,24 +1,205 @@
 import time
-from typing import Optional
+from contextlib import contextmanager
+from functools import wraps
+from typing import Any, Callable, Dict, List, Optional, TypeVar, Union
 
 import lightning as L
+import torch
 
 from fusion_bench.utils.pylogger import get_rankzero_logger
 
 log = get_rankzero_logger(__name__)
 
+T = TypeVar("T")
 
-def seed_everything_by_time(fabric: Optional[L.Fabric] = None):
+
+def seed_everything_by_time(fabric: Optional[L.Fabric] = None) -> int:
     """
-    Set seed for all processes by time.
+    Set seed for all processes based on current timestamp.
+
+    This function generates a time-based seed on the global zero process and broadcasts
+    it to all other processes in a distributed setting to ensure reproducibility across
+    all workers. When no fabric instance is provided, it generates a seed locally without
+    synchronization.
+
+    Args:
+        fabric: Optional Lightning Fabric instance for distributed synchronization.
+                If None, seed is generated locally without broadcasting.
+
+    Returns:
+        The seed value used for random number generation.
+
+    Example:
+        ```python
+        import lightning as L
+        from fusion_bench.utils.fabric import seed_everything_by_time
+
+        # With fabric (distributed)
+        fabric = L.Fabric(accelerator="auto", devices=2)
+        fabric.launch()
+        seed = seed_everything_by_time(fabric)
+        print(f"All processes using seed: {seed}")
+
+        # Without fabric (single process)
+        seed = seed_everything_by_time()
+        print(f"Using seed: {seed}")
+        ```
+
+    Note:
+        - In distributed settings, only the global zero process generates the seed
+        - All other processes receive the broadcasted seed for consistency
+        - The seed is based on `time.time()`, so it will differ across runs
     """
-    # set seed for all processes
+    # Generate seed on global zero process, None on others
     if fabric is None or fabric.is_global_zero:
         seed = int(time.time())
+        log.info(f"Generated time-based seed: {seed}")
     else:
         seed = None
+
+    # Broadcast seed to all processes in distributed setting
     if fabric is not None:
         log.debug(f"Broadcasting seed `{seed}` to all processes")
         fabric.barrier()
         seed = fabric.broadcast(seed, src=0)
+
+    # Apply seed to all random number generators
     L.seed_everything(seed)
+    return seed
+
+
+def is_distributed(fabric: Optional[L.Fabric] = None) -> bool:
+    """
+    Check if running in distributed mode (multi-process).
+
+    Args:
+        fabric: Optional Lightning Fabric instance. If None, returns False.
+
+    Returns:
+        True if running with multiple processes, False otherwise.
+
+    Example:
+        ```python
+        fabric = L.Fabric(accelerator="auto", devices=2)
+        fabric.launch()
+        if is_distributed(fabric):
+            print("Running in distributed mode")
+        ```
+    """
+    return fabric is not None and fabric.world_size > 1
+
+
+def get_world_info(fabric: Optional[L.Fabric] = None) -> Dict[str, Any]:
+    """
+    Get comprehensive information about the distributed setup.
+
+    Args:
+        fabric: Optional Lightning Fabric instance.
+
+    Returns:
+        Dictionary containing:
+        - world_size: Total number of processes
+        - global_rank: Global rank of current process
+        - local_rank: Local rank on current node
+        - is_global_zero: Whether this is the main process
+        - is_distributed: Whether running in distributed mode
+
+    Example:
+        ```python
+        fabric = L.Fabric(accelerator="auto", devices=2)
+        fabric.launch()
+        info = get_world_info(fabric)
+        print(f"Process {info['global_rank']}/{info['world_size']}")
+        ```
+    """
+    if fabric is None:
+        return {
+            "world_size": 1,
+            "global_rank": 0,
+            "local_rank": 0,
+            "is_global_zero": True,
+            "is_distributed": False,
+        }
+
+    return {
+        "world_size": fabric.world_size,
+        "global_rank": fabric.global_rank,
+        "local_rank": fabric.local_rank,
+        "is_global_zero": fabric.is_global_zero,
+        "is_distributed": fabric.world_size > 1,
+    }
+
+
+def wait_for_everyone(
+    fabric: Optional[L.Fabric] = None, message: Optional[str] = None
+) -> None:
+    """
+    Synchronize all processes with optional logging.
+
+    This is a wrapper around fabric.barrier() with optional message logging.
+
+    Args:
+        fabric: Optional Lightning Fabric instance. If None, does nothing.
+        message: Optional message to log before synchronization.
+
+    Example:
+        ```python
+        fabric = L.Fabric(accelerator="auto", devices=2)
+        fabric.launch()
+
+        # Do some work...
+        wait_for_everyone(fabric, "Waiting after model loading")
+        # All processes synchronized
+        ```
+    """
+    if fabric is not None:
+        if message and fabric.is_global_zero:
+            log.info(message)
+        fabric.barrier()
+
+
+@contextmanager
+def rank_zero_only_context(fabric: Optional[L.Fabric] = None):
+    """
+    Context manager to execute code block only on global rank 0.
+
+    Args:
+        fabric: Optional Lightning Fabric instance.
+
+    Example:
+        ```python
+        fabric = L.Fabric(accelerator="auto", devices=2)
+        fabric.launch()
+
+        with rank_zero_only_context(fabric):
+            print("This prints only on rank 0")
+            save_checkpoint(model, "checkpoint.pt")
+        ```
+    """
+    should_execute = fabric is None or fabric.is_global_zero
+    try:
+        yield should_execute
+    finally:
+        pass
+
+
+def print_on_rank_zero(*args, fabric: Optional[L.Fabric] = None, **kwargs) -> None:
+    """
+    Print message only on global rank 0.
+
+    Args:
+        *args: Arguments to pass to print().
+        fabric: Optional Lightning Fabric instance.
+        **kwargs: Keyword arguments to pass to print().
+
+    Example:
+        ```python
+        fabric = L.Fabric(accelerator="auto", devices=2)
+        fabric.launch()
+
+        print_on_rank_zero("Starting training", fabric=fabric)
+        # Prints only on rank 0
+        ```
+    """
+    if fabric is None or fabric.is_global_zero:
+        print(*args, **kwargs)

fusion_bench/utils/json.py

@@ -1,31 +1,78 @@
 import json
 from pathlib import Path
-from typing import Any, Union
+from typing import TYPE_CHECKING, Any, Union
 
+from fusion_bench.utils.validation import validate_file_exists
 
-def save_to_json(obj, path: Union[str, Path]):
+if TYPE_CHECKING:
+    from pyarrow.fs import FileSystem
+
+
+def save_to_json(obj, path: Union[str, Path], filesystem: "FileSystem" = None):
     """
     save an object to a json file
 
     Args:
         obj (Any): the object to save
         path (Union[str, Path]): the path to save the object
+        filesystem (FileSystem, optional): PyArrow FileSystem to use for writing.
+            If None, uses local filesystem via standard Python open().
+            Can also be an s3fs.S3FileSystem or fsspec filesystem.
     """
-    with open(path, "w") as f:
-        json.dump(obj, f)
+    if filesystem is not None:
+        json_str = json.dumps(obj)
+        # Check if it's an fsspec-based filesystem (like s3fs)
+        if hasattr(filesystem, "open"):
+            # Direct fsspec/s3fs usage - more reliable for some endpoints
+            path_str = str(path)
+            with filesystem.open(path_str, "w") as f:
+                f.write(json_str)
+        else:
+            # Use PyArrow filesystem
+            path_str = str(path)
+            with filesystem.open_output_stream(path_str) as f:
+                f.write(json_str.encode("utf-8"))
+    else:
+        # Use standard Python file operations
+        with open(path, "w") as f:
+            json.dump(obj, f)
 
 
-def load_from_json(path: Union[str, Path]) -> Union[dict, list]:
+def load_from_json(
+    path: Union[str, Path], filesystem: "FileSystem" = None
+) -> Union[dict, list]:
     """load an object from a json file
 
     Args:
         path (Union[str, Path]): the path to load the object
+        filesystem (FileSystem, optional): PyArrow FileSystem to use for reading.
+            If None, uses local filesystem via standard Python open().
+            Can also be an s3fs.S3FileSystem or fsspec filesystem.
 
     Returns:
-        dict: the loaded object
+        Union[dict, list]: the loaded object
+
+    Raises:
+        ValidationError: If the file doesn't exist (when using local filesystem)
     """
-    with open(path, "r") as f:
-        return json.load(f)
+    if filesystem is not None:
+        # Check if it's an fsspec-based filesystem (like s3fs)
+        if hasattr(filesystem, "open"):
+            # Direct fsspec/s3fs usage
+            path_str = str(path)
+            with filesystem.open(path_str, "r") as f:
+                return json.load(f)
+        else:
+            # Use PyArrow filesystem
+            path_str = str(path)
+            with filesystem.open_input_stream(path_str) as f:
+                json_data = f.read().decode("utf-8")
+                return json.loads(json_data)
+    else:
+        # Use standard Python file operations
+        validate_file_exists(path)
+        with open(path, "r") as f:
+            return json.load(f)
 
 
 def _is_list_of_dict(obj) -> bool:

fusion_bench/utils/validation.py

@@ -0,0 +1,197 @@
+"""
+Validation utilities for FusionBench.
+
+This module provides robust input validation functions to ensure data integrity
+and provide clear error messages throughout the FusionBench framework.
+"""
+
+import logging
+from pathlib import Path
+from typing import Any, Optional, Union
+
+__all__ = [
+    "ValidationError",
+    "validate_path_exists",
+    "validate_file_exists",
+    "validate_directory_exists",
+    "validate_model_name",
+]
+
+log = logging.getLogger(__name__)
+
+
+class ValidationError(ValueError):
+    """Custom exception for validation errors with detailed context."""
+
+    def __init__(self, message: str, field: Optional[str] = None, value: Any = None):
+        self.field = field
+        self.value = value
+        detailed_message = message
+        if field:
+            detailed_message = f"Validation error for '{field}': {message}"
+        if value is not None:
+            detailed_message += f" (got: {value!r})"
+        super().__init__(detailed_message)
+
+
+def validate_path_exists(
+    path: Union[str, Path],
+    name: str = "path",
+    create_if_missing: bool = False,
+    must_be_file: bool = False,
+    must_be_dir: bool = False,
+) -> Path:
+    """
+    Validate that a path exists and optionally check its type.
+
+    Args:
+        path: Path to validate.
+        name: Name of the path for error messages.
+        create_if_missing: If True and path doesn't exist, create it as a directory.
+        must_be_file: If True, ensure path points to a file.
+        must_be_dir: If True, ensure path points to a directory.
+
+    Returns:
+        Path object of the validated path.
+
+    Raises:
+        ValidationError: If path validation fails.
+
+    Examples:
+        >>> validate_path_exists("./config", name="config_dir", must_be_dir=True)
+        PosixPath('config')
+    """
+    if path is None:
+        raise ValidationError(f"{name} cannot be None", field=name, value=path)
+
+    assert not (
+        create_if_missing and must_be_file
+    ), "create_if_missing and must_be_file cannot both be True. By definition, a created path is a directory."
+
+    path_obj = Path(path).expanduser().resolve()
+
+    if not path_obj.exists():
+        if create_if_missing:
+            log.info(f"Creating missing directory: {path_obj}")
+            path_obj.mkdir(parents=True, exist_ok=True)
+        else:
+            raise ValidationError(
+                f"{name} does not exist: {path_obj}", field=name, value=str(path)
+            )
+
+    if must_be_file and not path_obj.is_file():
+        raise ValidationError(
+            f"{name} must be a file, but got directory: {path_obj}",
+            field=name,
+            value=str(path),
+        )
+
+    if must_be_dir and not path_obj.is_dir():
+        raise ValidationError(
+            f"{name} must be a directory, but got file: {path_obj}",
+            field=name,
+            value=str(path),
+        )
+
+    return path_obj
+
+
+def validate_file_exists(path: Union[str, Path], name: str = "file") -> Path:
+    """
+    Validate that a file exists.
+
+    Args:
+        path: File path to validate.
+        name: Name of the file for error messages.
+
+    Returns:
+        Path object of the validated file.
+
+    Raises:
+        ValidationError: If file doesn't exist or is not a file.
+    """
+    return validate_path_exists(path, name=name, must_be_file=True)
+
+
+def validate_directory_exists(
+    path: Union[str, Path], name: str = "directory", create_if_missing: bool = False
+) -> Path:
+    """
+    Validate that a directory exists.
+
+    Args:
+        path: Directory path to validate.
+        name: Name of the directory for error messages.
+        create_if_missing: If True, create directory if it doesn't exist.
+
+    Returns:
+        Path object of the validated directory.
+
+    Raises:
+        ValidationError: If directory doesn't exist (and not creating) or is not a directory.
+    """
+    return validate_path_exists(
+        path, name=name, must_be_dir=True, create_if_missing=create_if_missing
+    )
+
+
+def validate_model_name(
+    model_name: str, allow_special: bool = True, field: str = "model_name"
+) -> str:
+    """
+    Validate a model name string.
+
+    Args:
+        model_name: Model name to validate.
+        allow_special: If True, allow special names like "_pretrained_". If False,
+            names starting and ending with underscores will be rejected.
+        field: Field name for error messages.
+
+    Returns:
+        The validated model name.
+
+    Raises:
+        ValidationError: If model name is invalid.
+
+    Examples:
+        >>> validate_model_name("openai/clip-vit-base-patch32")
+        'openai/clip-vit-base-patch32'
+        >>> validate_model_name("_pretrained_", allow_special=True)
+        '_pretrained_'
+        >>> validate_model_name("_pretrained_", allow_special=False)
+        Traceback (most recent call last):
+        ...
+        ValidationError: Validation error for 'model_name': Special model names (starting and ending with '_') are not allowed (got: '_pretrained_')
+    """
+    if not model_name or not isinstance(model_name, str):
+        raise ValidationError(
+            "Model name must be a non-empty string", field=field, value=model_name
+        )
+
+    model_name = model_name.strip()
+    if not model_name:
+        raise ValidationError(
+            "Model name cannot be empty or whitespace only",
+            field=field,
+            value=model_name,
+        )
+
+    # Check for special names (e.g., _pretrained_, _base_model_)
+    if not allow_special and model_name.startswith("_") and model_name.endswith("_"):
+        raise ValidationError(
+            "Special model names (starting and ending with '_') are not allowed",
+            field=field,
+            value=model_name,
+        )
+
+    # Check for invalid characters that might cause issues
+    invalid_chars = ["\n", "\r", "\t", "\0"]
+    for char in invalid_chars:
+        if char in model_name:
+            raise ValidationError(
+                f"Model name contains invalid character: {char!r}",
+                field=field,
+                value=model_name,
+            )
+
+    return model_name

{fusion_bench-0.2.27.dist-info → fusion_bench-0.2.29.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: fusion-bench
-Version: 0.2.27
+Version: 0.2.29
 Summary: A Comprehensive Benchmark of Deep Model Fusion
 Author-email: Anke Tang <tang.anke@foxmail.com>
 Project-URL: Repository, https://github.com/tanganke/fusion_bench
@@ -44,12 +44,14 @@ Dynamic: license-file
 
 [![arXiv](https://img.shields.io/badge/arXiv-2406.03280-b31b1b.svg)](http://arxiv.org/abs/2406.03280)
 [![GitHub License](https://img.shields.io/github/license/tanganke/fusion_bench)](https://github.com/tanganke/fusion_bench/blob/main/LICENSE)
-[![PyPI - Version](https://img.shields.io/pypi/v/fusion-bench)](https://pypi.org/project/fusion-bench/)
-[![Downloads](https://static.pepy.tech/badge/fusion-bench/month)](https://pepy.tech/project/fusion-bench)
 [![Static Badge](https://img.shields.io/badge/doc-mkdocs-blue)](https://tanganke.github.io/fusion_bench/)
 [![Static Badge](https://img.shields.io/badge/code%20style-black-black)](https://github.com/psf/black)
 [![Static Badge](https://img.shields.io/badge/code%20style-yamlfmt-black)](https://github.com/google/yamlfmt)
 
+[![CodeFactor](https://www.codefactor.io/repository/github/tanganke/fusion_bench/badge/main)](https://www.codefactor.io/repository/github/tanganke/fusion_bench/overview/main)
+[![PyPI - Version](https://img.shields.io/pypi/v/fusion-bench)](https://pypi.org/project/fusion-bench/)
+[![Downloads](https://static.pepy.tech/badge/fusion-bench/month)](https://pepy.tech/project/fusion-bench)
+
 </div>
 
 > [!TIP]
@@ -205,6 +207,48 @@ The CLI's design allows for easy extension to new fusion methods, model types, a
 
 Read the [CLI documentation](https://tanganke.github.io/fusion_bench/cli/fusion_bench/) for more information.
 
+## The FusionBench Workflow
+
+FusionBench follows a three-component architecture to perform model fusion experiments:
+
+```mermaid
+graph LR
+    CLI[fusion_bench CLI] --> Hydra[Hydra Config]
+    Hydra --> Program[Program]
+    
+    Program --> MP[ModelPool<br/>Manages Models<br/>& Datasets]
+    Program --> Method[Method<br/>Fusion Algorithm]
+    Program --> TP[TaskPool<br/>Evaluation Tasks]
+    
+    MP --> Method
+    Method --> Merged[Merged Model]
+    Merged --> TP
+    TP --> Report[Evaluation Report]
+    
+    style CLI fill:#e1f5e1
+    style Hydra fill:#f0e1ff
+    style Method fill:#ffe1f0
+    style Merged fill:#fff4e1
+    style Report fill:#e1f0ff
+```
+
+**Key Components:**
+
+1. **CLI**: Entry point using Hydra for configuration management
+2. **Program**: Orchestrates the fusion workflow (e.g., `FabricModelFusionProgram`)
+3. **ModelPool**: Manages task-specific models and their datasets
+4. **Method**: Implements the fusion algorithm (e.g., Simple Average, Task Arithmetic, AdaMerging)
+5. **TaskPool**: Evaluates the merged model on benchmark tasks
+
+**Workflow Steps:**
+
+1. User runs `fusion_bench` with config overrides
+2. Hydra loads YAML configs for method, modelpool, and taskpool
+3. Program instantiates all three components
+4. Method executes fusion algorithm on ModelPool
+5. TaskPool evaluates the merged model
+6. Results are saved and reported
+
 ## Implement your own model fusion algorithm
 
 First, create a new Python file for the algorithm in the `fusion_bench/method` directory.
@@ -272,11 +316,24 @@ Click on [<kbd>Use this template</kbd>](https://github.com/fusion-bench/fusion-b
 
 </div>
 
-### FusionBench Command Generator WebUI (for v0.1.x)
+### FusionBench Command Generator WebUI
+
+> [!NOTE]
+> Requires `gradio` package. Install with `pip install gradio`.
+
+For users who prefer a graphical interface, FusionBench provides an interactive web UI for generating commands:
+
+```bash
+fusion_bench_webui
+```
+
+This launches a browser-based interface where you can:
 
-FusionBench Command Generator is a user-friendly web interface for generating FusionBench commands based on configuration files.
-It provides an interactive way to select and customize FusionBench configurations, making it easier to run experiments with different settings.
-[Read more here](https://tanganke.github.io/fusion_bench/cli/fusion_bench_webui/).
+- Select root configurations and components through dropdowns
+- Adjust hyperparameters interactively
+- View real-time YAML configuration updates
+
+The WebUI is particularly useful for exploring available configurations, experimenting with different parameter combinations, and learning the FusionBench configuration structure. [Learn more about the WebUI](https://tanganke.github.io/fusion_bench/cli/fusion_bench_webui/).
 
 ![FusionBench Command Generator Web Interface](docs/cli/images/fusion_bench_webui.png)
 
@@ -296,3 +353,5 @@ If you find this benchmark useful, please consider citing our work:
 ## Star History
 
 [![Star History Chart](https://api.star-history.com/svg?repos=tanganke/fusion_bench&type=Date)](https://www.star-history.com/#tanganke/fusion_bench&Date)
+
+![Alt](https://repobeats.axiom.co/api/embed/83f1f046562e4a4787bdd6ed1190856f9f30bd9f.svg "Repobeats analytics image")