fusion-bench 0.2.28__py3-none-any.whl → 0.2.30__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/instantiate_utils.py

@@ -14,8 +14,8 @@ from lightning_utilities.core.rank_zero import rank_zero_only
 from omegaconf import DictConfig, OmegaConf, SCMode
 from omegaconf._utils import is_structured_config
 from rich import print
-from rich.panel import Panel
-from rich.syntax import Syntax
+
+from fusion_bench.utils.rich_utils import print_bordered
 
 PRINT_FUNCTION_CALL = True
 """
@@ -67,12 +67,22 @@ def _resolve_callable_name(f: Callable[..., Any]) -> str:
     return full_name
 
 
-def _format_args_kwargs(args, kwargs):
+def _get_obj_str(obj: Any) -> str:
+    if isinstance(obj, (str, int, float, bool, type(None))):
+        return repr(obj)
+    else:
+        return f"'<{type(obj).__name__} object>'"
+
+
+def _format_args_kwargs(args: Tuple[Any, ...], kwargs: Dict[str, Any]) -> str:
     result_strings = []
     if len(args) > 0:
-        result_strings.append(", ".join(repr(arg) for arg in args))
+        result_strings.append(", ".join(_get_obj_str(arg) for arg in args))
+
     if len(kwargs) > 0:
-        result_strings.append(", ".join(f"{k}={repr(v)}" for k, v in kwargs.items()))
+        result_strings.append(
+            ", ".join(f"{k}={_get_obj_str(v)}" for k, v in kwargs.items())
+        )
 
     if len(result_strings) == 0:
         return ""
@@ -145,14 +155,14 @@ def _call_target(
     if _partial_:
         if PRINT_FUNCTION_CALL and getattr(rank_zero_only, "rank", 0) == 0:
             call_str = f"functools.partial({_resolve_callable_name(_target_)}, {_format_args_kwargs(args, kwargs)})"
-            PRINT_FUNCTION_CALL_FUNC(
-                Panel(
-                    Syntax(call_str, "python", theme="monokai", word_wrap=True),
-                    title="Instantiate by calling partial",
-                    border_style="cyan",
-                )
+            print_bordered(
+                call_str,
+                code_style="python",
+                title=f"Instantiate by calling {'function' if not isinstance(_target_, type) else 'class'}",
+                style="cyan",
+                expand=False,
+                print_fn=PRINT_FUNCTION_CALL_FUNC,
             )
-
         if CATCH_EXCEPTION:
             try:
                 return functools.partial(_target_, *args, **kwargs)
@@ -169,12 +179,13 @@ def _call_target(
     else:
         if PRINT_FUNCTION_CALL and getattr(rank_zero_only, "rank", 0) == 0:
             call_str = f"{_resolve_callable_name(_target_)}({_format_args_kwargs(args, kwargs)})"
-            PRINT_FUNCTION_CALL_FUNC(
-                Panel(
-                    Syntax(call_str, "python", theme="monokai", word_wrap=True),
-                    title="Instantiate by calling function",
-                    border_style="green",
-                )
+            print_bordered(
+                call_str,
+                code_style="python",
+                title=f"Instantiate by calling {'function' if not isinstance(_target_, type) else 'class'}",
+                style="green",
+                expand=False,
+                print_fn=PRINT_FUNCTION_CALL_FUNC,
             )
         if CATCH_EXCEPTION:
             try:

fusion_bench/utils/json.py

@@ -2,6 +2,8 @@ import json
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Union
 
+from fusion_bench.utils.validation import validate_file_exists
+
 if TYPE_CHECKING:
     from pyarrow.fs import FileSystem
 
@@ -49,6 +51,9 @@ def load_from_json(
 
     Returns:
         Union[dict, list]: the loaded object
+
+    Raises:
+        ValidationError: If the file doesn't exist (when using local filesystem)
     """
     if filesystem is not None:
         # Check if it's an fsspec-based filesystem (like s3fs)
@@ -65,6 +70,7 @@ def load_from_json(
                 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)
 

fusion_bench/utils/misc.py

@@ -178,3 +178,19 @@ def validate_and_suggest_corrections(
     if matches:
         msg += f". Did you mean {', '.join(repr(m) for m in matches)}?"
     raise ValueError(msg)
+
+
+class DeprecationWarningMeta(type):
+    """
+    Metaclass that issues a deprecation warning whenever a class using it is instantiated.
+    """
+
+    def __call__(cls, *args, **kwargs):
+        import warnings
+
+        warnings.warn(
+            f"{cls.__name__} is deprecated and will be removed in a future version. ",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        return super(DeprecationWarningMeta, cls).__call__(*args, **kwargs)

fusion_bench/utils/rich_utils.py

@@ -1,6 +1,6 @@
 import logging
 from pathlib import Path
-from typing import Sequence
+from typing import Optional, Sequence
 
 import rich
 import rich.syntax
@@ -19,6 +19,9 @@ from rich.text import Text
 from rich.traceback import install as install_rich_traceback
 
 from fusion_bench.utils import pylogger
+from fusion_bench.utils.packages import _is_package_available
+
+install_rich_traceback()
 
 log = pylogger.RankedLogger(__name__, rank_zero_only=True)
 
@@ -61,7 +64,31 @@ def display_available_styles():
     console.print(Columns(style_samples, equal=True, expand=False))
 
 
-def print_bordered(message, title=None, style="blue", code_style=None):
+def format_code_str(message: str, code_style="python"):
+    if code_style.lower() == "python" and _is_package_available("black"):
+        # Use black formatting for python code if black is available
+        import black
+
+        try:
+            message = black.format_str(message, mode=black.Mode())
+        except black.InvalidInput:
+            pass  # If black fails, use the original message
+
+    return message.strip()
+
+
+def print_bordered(
+    message,
+    title=None,
+    style="blue",
+    code_style=None,
+    *,
+    expand: bool = True,
+    theme: str = "monokai",
+    background_color: Optional[str] = "default",
+    print_fn=print,
+    format_code: bool = True,
+):
     """
     Print a message with a colored border.
 
@@ -73,12 +100,63 @@ def print_bordered(message, title=None, style="blue", code_style=None):
                                 Set to None for plain text. Defaults to "python".
     """
     if code_style:
-        content = Syntax(message, code_style, theme="monokai", word_wrap=True)
+        if format_code:
+            message = format_code_str(message, code_style)
+        content = Syntax(
+            message,
+            code_style,
+            word_wrap=True,
+            theme=theme,
+            background_color=background_color,
+        )
     else:
         content = Text(message)
 
-    panel = Panel(content, title=title, border_style=style)
-    print(panel)
+    panel = Panel(content, title=title, border_style=style, expand=expand)
+    print_fn(panel)
+
+
+def print_code(
+    message,
+    title=None,
+    code_style=None,
+    *,
+    expand: bool = True,
+    theme: str = "monokai",
+    background_color: Optional[str] = "default",
+    print_fn=print,
+):
+    """
+    Print code or plain text with optional syntax highlighting.
+
+    Args:
+        message (str): The message or code to print.
+        title (str, optional): Optional title associated with this output. Currently
+            not used by this function, but kept for API compatibility. Defaults to None.
+        code_style (str, optional): The language/lexer name for syntax highlighting
+            (for example, ``"python"``). If ``None``, the message is rendered as plain
+            text without syntax highlighting. Defaults to ``None``.
+        expand (bool, optional): Placeholder flag for API symmetry with other printing
+            helpers. It is not used in the current implementation. Defaults to True.
+        theme (str, optional): Name of the Rich syntax highlighting theme to use when
+            ``code_style`` is provided. Defaults to ``"monokai"``.
+        background_color (str, optional): Background color style to apply to the code
+            block when using syntax highlighting. Defaults to ``"default"``.
+        print_fn (Callable, optional): Function used to render the resulting Rich
+            object. Defaults to :func:`rich.print`.
+    """
+    if code_style:
+        content = Syntax(
+            message,
+            code_style,
+            word_wrap=True,
+            theme=theme,
+            background_color=background_color,
+        )
+    else:
+        content = Text(message)
+
+    print_fn(content)
 
 
 @rank_zero_only
@@ -95,6 +173,9 @@ def print_config_tree(
     ),
     resolve: bool = False,
     save_to_file: bool = False,
+    *,
+    theme: str = "monokai",
+    background_color: Optional[str] = "default",
 ) -> None:
     """Prints the contents of a DictConfig as a tree structure using the Rich library.
 
@@ -134,7 +215,14 @@ def print_config_tree(
         else:
             branch_content = str(config_group)
 
-        branch.add(rich.syntax.Syntax(branch_content, "yaml"))
+        branch.add(
+            rich.syntax.Syntax(
+                branch_content,
+                "yaml",
+                theme=theme,
+                background_color=background_color,
+            )
+        )
 
     # print config tree
     rich.print(tree)
@@ -145,6 +233,35 @@ def print_config_tree(
             rich.print(tree, file=file)
 
 
+@rank_zero_only
+def print_config_yaml(
+    cfg: DictConfig,
+    resolve: bool = False,
+    output_path: Optional[str] = False,
+    *,
+    theme: str = "monokai",
+    background_color: Optional[str] = "default",
+) -> None:
+    """
+    Prints the contents of a DictConfig as a YAML string using the Rich library.
+
+    Args:
+        cfg: A DictConfig composed by Hydra.
+        resolve: Whether to resolve reference fields of DictConfig. Default is ``False``.
+        output_path: Optional path to export the config YAML to. If provided, the file is written to this path.
+    """
+    config_yaml = OmegaConf.to_yaml(cfg, resolve=resolve)
+    syntax = rich.syntax.Syntax(
+        config_yaml, "yaml", theme=theme, background_color=background_color
+    )
+    rich.print(syntax)
+
+    if output_path:
+        Path(output_path).parent.mkdir(parents=True, exist_ok=True)
+        with open(Path(output_path), "w") as file:
+            rich.print(syntax, file=file)
+
+
 @rank_zero_only
 def enforce_tags(cfg: DictConfig, save_to_file: bool = False) -> None:
     """Prompts user to input tags from command line if no tags are provided in config.