Delete lyra_2

lyra_2/__init__.py

@@ -1,14 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.

lyra_2/_ext/__init__.py

@@ -1,4 +0,0 @@
-"""
-Copied from cosmos repository, with necessary modifications.
-Original source: https://github.com/nvidia-cosmos/cosmos-predict2.5/
-"""

lyra_2/_ext/imaginaire/__init__.py

@@ -1,14 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.

lyra_2/_ext/imaginaire/checkpointer/__init__.py

@@ -1,14 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.

lyra_2/_ext/imaginaire/checkpointer/base.py

@@ -1,177 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import os
-from abc import ABC, abstractmethod
-from typing import Optional
-
-import torch
-
-from lyra_2._ext.imaginaire.config import CheckpointConfig, JobConfig
-from lyra_2._ext.imaginaire.model import ImaginaireModel
-from lyra_2._ext.imaginaire.utils import callback
-from lyra_2._ext.imaginaire.utils.easy_io import easy_io
-
-
-class AbstractCheckpointer(ABC):
-    """The checkpointer class. Supports checkpoint saving/loading to both local disk or object store."""
-
-    def __init__(
-        self,
-        config_checkpoint: CheckpointConfig,
-        config_job: JobConfig,
-        callbacks: Optional[callback.CallBackGroup] = None,
-    ):
-        """Constructor of the checkpointer.
-
-        Args:
-            config_checkpoint (CheckpointConfig): The config object for the checkpointer.
-        """
-        self.config_checkpoint = config_checkpoint
-        # Set the callback functions.
-        self.callbacks = callbacks
-        self.save_to_object_store = config_checkpoint.save_to_object_store.enabled
-        self.load_from_object_store = config_checkpoint.load_from_object_store.enabled
-
-        # Set checkpoint directories for local and object store paths
-        self._local_dirname = os.path.join(config_job.path_local, "checkpoints")
-        self._object_store_dirname = os.path.join(config_job.path, "checkpoints")
-
-        self.strict_resume = config_checkpoint.strict_resume
-        self.load_path = config_checkpoint.load_path or None
-        self.load_training_state = config_checkpoint.load_training_state
-        self.only_load_scheduler_state = config_checkpoint.only_load_scheduler_state
-        self.save_thread = None
-        self.verbose = config_checkpoint.verbose
-        self.keys_not_to_resume = config_checkpoint.keys_not_to_resume
-        self.broadcast_via_filesystem = config_checkpoint.broadcast_via_filesystem
-        # Create the object store client interface.
-        if config_checkpoint.load_from_object_store.enabled:
-            self.load_s3_backend_key = "_ckpt_s3_loader"
-            easy_io.set_s3_backend(
-                key="_ckpt_s3_loader",
-                backend_args={
-                    "backend": "s3",
-                    "path_mapping": {
-                        "s3://ckpt/": f"s3://{config_checkpoint.load_from_object_store.bucket}/",
-                    },
-                    "s3_credential_path": config_checkpoint.load_from_object_store.credentials,
-                },
-            )
-        else:
-            self.load_s3_backend_key = None
-
-        if config_checkpoint.save_to_object_store.enabled:
-            self.save_s3_backend_key = "_ckpt_s3_saver"
-            easy_io.set_s3_backend(
-                key="_ckpt_s3_saver",
-                backend_args={
-                    "backend": "s3",
-                    "path_mapping": {
-                        "s3://ckpt/": f"s3://{config_checkpoint.save_to_object_store.bucket}/",
-                    },
-                    "s3_credential_path": config_checkpoint.save_to_object_store.credentials,
-                },
-            )
-        else:
-            self.save_s3_backend_key = None
-
-    @abstractmethod
-    def save(
-        self,
-        model: ImaginaireModel,
-        optimizer: torch.optim.Optimizer,
-        scheduler: torch.optim.lr_scheduler.LRScheduler,
-        grad_scaler: torch.amp.GradScaler,
-        iteration: int,
-    ) -> None:
-        pass
-
-    @abstractmethod
-    def load(
-        self,
-        model: ImaginaireModel,
-        optimizer: Optional[torch.optim.Optimizer] = None,
-        scheduler: Optional[torch.optim.lr_scheduler.LRScheduler] = None,
-        grad_scaler: Optional[torch.amp.GradScaler] = None,
-    ) -> int:
-        pass
-
-    @property
-    def save_bucket(self):
-        """Get the bucket name for saving checkpoints."""
-        return self.config_checkpoint.save_to_object_store.bucket if self.save_to_object_store else None
-
-    @property
-    def load_bucket(self):
-        """Get the bucket name for loading checkpoints."""
-        return self.config_checkpoint.load_from_object_store.bucket if self.load_from_object_store else None
-
-    @property
-    def save_dirname(self):
-        return (
-            f"s3://{self.save_bucket}/{self._object_store_dirname}"
-            if self.save_to_object_store
-            else self._local_dirname
-        )
-
-    @property
-    def load_dirname(self):
-        return (
-            f"s3://{self.load_bucket}/{self._object_store_dirname}"
-            if self.load_from_object_store
-            else self._local_dirname
-        )
-
-    def finalize(self) -> None:
-        """Finalize the checkpointer."""
-        if self.save_thread:
-            self.save_thread.join()
-
-    def _read_latest_checkpoint_file(self) -> str | None:
-        """Get the file name of the latest saved checkpoint. If it doesn't exist, return None.
-
-        Returns:
-            checkpoint_file (str | None): file name of the latest saved checkpoint.
-        """
-        checkpoint_file = None
-        checkpoint_path = os.path.join(self.load_dirname, "latest_checkpoint.txt")
-        if easy_io.exists(f"{checkpoint_path}", backend_key=self.load_s3_backend_key):
-            checkpoint_file = easy_io.load(f"{checkpoint_path}", backend_key=self.load_s3_backend_key).strip()
-
-        return checkpoint_file
-
-    def _write_latest_checkpoint_file(self, checkpoint_file: str) -> None:
-        """Track the file name of the latest saved checkpoint.
-
-        Args:
-            checkpoint_file (str): file name of the latest saved checkpoint.
-        """
-        content = f"{checkpoint_file}\n"
-        checkpoint_path = os.path.join(self.save_dirname, "latest_checkpoint.txt")
-        easy_io.dump(
-            content,
-            checkpoint_path,
-            backend_key=self.save_s3_backend_key,
-        )
-
-    def _check_checkpoint_exists(self, checkpoint_path: str) -> None:
-        """If the file checkpoint_path does not exist, raise an error.
-
-        Args:
-            checkpoint_path (str): full path to the checkpoint.
-        """
-        if not easy_io.exists(f"{checkpoint_path}", backend_key=self.load_s3_backend_key):
-            raise FileNotFoundError(f"File not found (object store): {checkpoint_path}")

lyra_2/_ext/imaginaire/checkpointer/dcp.py

@@ -1,1003 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-
-"""
-Distributed checkpoint (DCP) directory structure and storage backends.
-
-The checkpointer saves model state in a sharded format across multiple processes:
-
-self.save_dirname/
-├── iter_000000005/                    # Checkpoint at iteration 5
-│   ├── model/                         # Model state shards
-│   │   ├── __0_0.distcp              # Shard 0 from rank 0
-│   │   └── __1_0.distcp              # Shard 1 from rank 1
-│   ├── optim/                        # Optimizer state shards
-│   │   ├── __0_0.distcp              # Shard 0 from rank 0
-│   │   └── __1_0.distcp              # Shard 1 from rank 1
-│   ├── scheduler/                    # Learning rate scheduler state
-│   │   ├── __0_0.distcp              # Shard 0 from rank 0
-│   │   └── __1_0.distcp              # Shard 1 from rank 1
-│   └── trainer/                      # Additional training state
-│       ├── __0_0.distcp              # Shard 0 from rank 0
-│       └── __1_0.distcp              # Shard 1 from rank 1
-└── latest_checkpoint.txt             # Points to most recent checkpoint folder, e.g. iter_000000005
-
-Storage path format:
-  self.save_dirname = "{config_job.path_local}/checkpoints"
-
-The sharded format enables efficient distributed saving/loading by:
-1. Parallelizing I/O across processes
-2. Reducing memory usage per process
-3. Supporting both local and cloud storage backends
-"""
-
-import enum
-import functools
-import multiprocessing
-import os
-import queue
-import re
-import time
-import warnings
-from collections import namedtuple
-from multiprocessing import get_context
-from typing import Any, Dict, List, Optional, Set, Tuple, Union, cast
-
-import torch
-import torch.distributed
-import torch.distributed.checkpoint as dcp
-from torch import nn
-from torch.distributed.checkpoint import FileSystemReader, FileSystemWriter
-from torch.distributed.checkpoint._storage_utils import _storage_setup
-from torch.distributed.checkpoint.default_planner import DefaultSavePlanner
-from torch.distributed.checkpoint.logger import _dcp_method_logger
-from torch.distributed.checkpoint.state_dict import (
-    StateDictOptions,
-    get_model_state_dict,
-    get_optimizer_state_dict,
-    set_model_state_dict,
-    set_optimizer_state_dict,
-)
-from torch.distributed.checkpoint.stateful import Stateful
-from torch.distributed.checkpoint.storage import StorageReader
-from torch.distributed.checkpoint.utils import _api_bc_check, _DistWrapper, _profile
-from torch.distributed.tensor import distribute_tensor
-
-from lyra_2._ext.imaginaire.checkpointer.base import AbstractCheckpointer
-from lyra_2._ext.imaginaire.checkpointer.s3_filesystem import S3StorageReader, S3StorageWriter
-from lyra_2._ext.imaginaire.config import CheckpointConfig, JobConfig
-from lyra_2._ext.imaginaire.model import ImaginaireModel
-from lyra_2._ext.imaginaire.utils import callback, distributed, log, misc
-from lyra_2._ext.imaginaire.utils.easy_io import easy_io
-from lyra_2._src.models.wan_t2v_model import WANDiffusionModel as DiffusionModel
-
-try:
-    """
-    We override the default load function to skip loadding _extra_state keys created by transformer-engine.
-    """
-    from torch.distributed.checkpoint.default_planner import DefaultLoadPlanner as _DefaultLoadPlanner
-    from torch.distributed.checkpoint.default_planner import (
-        DTensor,
-        LoadPlan,
-        _create_read_items,
-        _version,
-        flatten_state_dict,
-    )
-    from torch.distributed.checkpoint.metadata import Metadata, TensorStorageMetadata
-
-    def create_default_local_load_plan(
-        state_dict: dict[str, Any], metadata: Metadata, strict: bool = True, dcp_allow_mismatched_size: bool = False
-    ) -> LoadPlan:
-        requests = []
-        """
-        Create the ``LoadPlan`` used by DefaultLoadPlanner.
-
-        It produces one read item per value in ``state_dict`` using the metadata in ``metadata``.
-
-        The default behavior is to match key exactly between state_dict and metadata.
-        It handles resharding by issuing multiple read requests against storage in order to match
-        load requirements.
-        """
-
-        for fqn, obj in state_dict.items():
-            if fqn.endswith("._extra_state"):  # dirty TE attention package!
-                continue
-            # ignore state_dict keys which do not exist in `state_dict` if strict=False
-            if fqn not in metadata.state_dict_metadata:
-                if strict:
-                    raise RuntimeError(f"Missing key in checkpoint state_dict: {fqn}.")
-                else:
-                    log.warning(f"Local load plan: Missing key in checkpoint state_dict: {fqn}.")
-                    continue
-
-            md = metadata.state_dict_metadata[fqn]
-
-            if not dcp_allow_mismatched_size:
-                if (
-                    isinstance(md, TensorStorageMetadata)
-                    and getattr(obj, "size", None) is not None
-                    and md.size != obj.size()
-                ):
-                    if not strict:
-                        log.critical(f"Size mismatch between saved {md.size} and current: {obj.size()} for {fqn}")
-                        continue
-                    else:
-                        raise ValueError(
-                            f"Size mismatch between saved {md.size} and current: {obj.size()} for {fqn}",
-                        )
-            # Since DTensor supports submesh, adding extra check to ensure _create_read_items()
-            # gets called only when the current rank is part of the mesh for the corresponding DTensor.
-            if isinstance(obj, DTensor):
-                if obj.device_mesh.get_coordinate() is not None:
-                    requests += _create_read_items(fqn, md, obj)
-            else:
-                requests += _create_read_items(fqn, md, obj)
-
-        return LoadPlan(requests)
-
-    class DefaultLoadPlanner(_DefaultLoadPlanner):
-        def set_partial_channel_weight(self, dcp_allow_mismatched_size: bool):
-            self.dcp_allow_mismatched_size = dcp_allow_mismatched_size
-
-        def create_local_plan(self) -> LoadPlan:
-            assert self.metadata is not None
-            if self.flatten_state_dict:
-                # To support checkpoints that are saved before v2.4, we have to
-                # differentiate if the missing keys are due to old checkpoints.
-                # The contracts are:
-                # 1. There are 3 cases when we found a missing key.
-                #    1.1 Actual missing key, but allow_partial_load is False
-                #    1.2 Actual missing key, but allow_partial load is True
-                #    1.3 Old checkpoint, but allow_partial_load is False
-                #    1.4 Old checkpoint, but allow_partial_load is True
-                # 2. If we found a missing key, we first convert the keys back to
-                #    the key format of v2.3
-                # 3. If the previous missing keys are in the v2.3 keys, we assume
-                #    this is a old checkpoint.
-                # 4. Pass the state_dict to `create_default_local_load_plan()`,
-                #    which has the logic to check missing for allow_partial_load.
-                # So for 1.2 and 1.4 cases, we delegate allow_partial_load check to
-                # `create_default_local_load_plan()`. The logic here is to determine
-                # whether the checkpoint belong to 2.3 (or before) or 2.4 (or after).
-                current_keys = set(self.state_dict.keys())
-                load_keys = set(self.metadata.state_dict_metadata.keys())
-                missing_keys = load_keys - current_keys
-                if missing_keys:
-                    _version._derived_version = "2_3"
-                    old_state_dict, old_mappings = flatten_state_dict(self.original_state_dict)
-                    old_keys = set(old_state_dict.keys())
-                    if old_keys & missing_keys:
-                        self.state_dict, self.mappings = old_state_dict, old_mappings
-                    # _derived_version is only used by flatten_state_dict now.
-                    # Set it back to None so that later we can save to a new version.
-                    _version._derived_version = None
-
-            return create_default_local_load_plan(
-                self.state_dict,
-                self.metadata,
-                not self.allow_partial_load,
-                getattr(self, "dcp_allow_mismatched_size", False),
-            )
-
-    log.info("for the back comptiable pytorch! New DefaultLoadPlanner class is created.")
-
-    @_dcp_method_logger(log_exceptions=True)
-    @_api_bc_check
-    def load(
-        state_dict: dict[str, Any],
-        *,
-        checkpoint_id: Union[str, os.PathLike, None] = None,
-        storage_reader: Optional[StorageReader] = None,
-        planner: Optional[DefaultLoadPlanner] = None,
-        process_group: Optional[torch.distributed.ProcessGroup] = None,
-        no_dist: bool = False,
-    ) -> None:
-        """
-        We override the default load function to perform a load plan check for mismatched/missing keys.
-        ==========================Original Doc string=====================================
-        Load a checkpoint into a distributed state dict in SPMD style.
-
-        Each rank must have the same keys in their ``state_dict`` provided to this
-        API. Mismatched keys may result in hangs or errors. If unsure, you can use
-        the ``utils._assert_same_keys`` API to check (but may incur communication
-        costs).
-
-        Each rank will try to read the least amount of data necessary
-        to fullfill the requested `state_dict`. When loading :class:`ShardedTensor`
-        or :class:`DTensor` instances, each rank only reads data for their local shards.
-
-        For each ``Stateful`` object (having both a ``state_dict`` and a ``load_state_dict``),
-        load will first call ``state_dict`` before attempting deserialization, followed by
-        ``load_state_dict`` once the deserialization is complete.
-        For each non-``Stateful`` object, load will deserailize the object, and then replace
-        it in the ``state_dict`` with the deserialized object.
-
-        .. warning::
-            All tensors in ``state_dict`` must be allocated on their
-            destination device *prior to* calling this function.
-
-            All non-tensor data is loaded using `torch.load()` and modified in place
-            on state_dict.
-
-        .. warning::
-            Users must call `load_state_dict` on the root module to ensure load
-            pos-processing and non-tensor data properly propagates.
-
-        .. note:
-            If no process group is initialized, this function will assume the intent
-            is to load a checkpoint into the local process. This can be useful in the
-            case of local inference, and when using regular Tensors (as opposed to DTensor
-            or ShardedTensor)
-
-        .. note:
-            Rank 0 is assumed to be the coordinator rank.
-
-        Args:
-            state_dict (Dict[str, Any]): The state_dict to load the checkpoint into.
-            checkpoint_id (Union[str, os.PathLike, None]):
-                The ID of this checkpoint instance. The meaning of the checkpoint_id
-                depends on the storage. It can be a path to a folder or to a file.
-                It can also be a key if the storage is a key-value store.
-                (Default: ``None``)
-            storage_reader (Optional[StorageReader]):
-                Instance of StorageWriter used to perform reads. If this is not
-                specified, DCP will automatically infer the reader based on the
-                checkpoint_id. If checkpoint_id is also None, an exception will
-                be raised. (Default: ``None``)
-            planner (Optional[LoadPlanner]):
-                Instance of LoadPlanner. If this is not specificed, the default
-                planner will be used. (Default: ``None``)
-            process_group (Optional[ProcessGroup]):
-                ProcessGroup to be used for cross-rank synchronization.
-                (Default: ``None``)
-            no_dist (bool): If ``True``, this function will assume the intent is to load
-                a checkpoint without using cross-rank synchronization. (Default: ``False``)
-        Returns:
-            None.
-
-        Examples
-            >>> # xdoctest: +SKIP
-            >>> my_model = MyModule()
-            >>> optimizer = Adagrad(my_model.parameters())
-            >>> model_state_dict = my_model.state_dict()
-            >>> fs_storage_reader = torch.distributed.checkpoint.FileSystemReader(
-            ...     "/checkpoint/1"
-            ... )
-
-            >>> torch.distributed.checkpoint.load_state_dict(
-            >>>     state_dict=model_state_dict,
-            >>>     storage_reader=fs_storage_reader,
-            >>> )
-
-            >>> # module.load_state_dict() function might have customized steps
-            >>> # to flush the state_dict, must call it to
-            >>> # ensure correct behavior.
-            >>> my_model.load_state_dict(model_state_dict)
-
-        .. note::
-            load_state_dict uses collectives to coordinate reads across ranks.
-            For NCCL-based process groups, internal tensor representations of
-            objects must be moved to the GPU device before communication takes place.
-            In this case, the device used is given by ``torch.cuda.current_device()``
-            and it is the user's responsibility to ensure that this is set so that each
-            rank has an individual GPU, via ``torch.cuda.set_device()``.
-        """
-
-        no_dist = no_dist or (not torch.distributed.is_available()) or (not torch.distributed.is_initialized())
-        if no_dist:
-            warnings.warn(
-                "torch.distributed is disabled, unavailable or uninitialized, assuming the intent is to load in a single process."
-            )
-
-        with _profile():
-            storage_reader = cast(StorageReader, _storage_setup(storage_reader, checkpoint_id, reader=True))
-
-            # All ranks must have the same keys in their `state_dict` provided to
-            # this API.  See documentation for more details.
-            # Here we simply sort the keys to ensure that all ranks load values in
-            # the same order.
-            keys = sorted(state_dict.keys())
-
-            statetful_sd = {}
-            for key in keys:
-                if key not in state_dict:
-                    continue
-                elem = state_dict[key]
-                statetful_sd[key] = elem.state_dict() if isinstance(elem, Stateful) else elem
-
-            _load_state_dict(
-                state_dict=statetful_sd,
-                storage_reader=storage_reader,
-                process_group=process_group,
-                no_dist=no_dist,
-                planner=planner,
-            )
-            for key in keys:
-                if key not in state_dict:
-                    continue
-                elem = state_dict[key]
-                if isinstance(elem, Stateful):
-                    # If the state_dict is a Stateful object,
-                    # DCP does an in-place load in the original state dict.
-                    elem.load_state_dict(statetful_sd[key])
-                else:
-                    # Otherwise, replace the state_dict with the loaded state_dict.
-                    state_dict[key] = statetful_sd[key]
-
-    def _load_state_dict(
-        state_dict: dict[str, Any],
-        storage_reader: StorageReader,
-        process_group: Optional[torch.distributed.ProcessGroup] = None,
-        coordinator_rank: int = 0,
-        no_dist: bool = False,
-        planner: Optional[DefaultLoadPlanner] = None,
-    ) -> None:
-        torch._C._log_api_usage_once("torch.distributed.checkpoint.load_state_dict")
-
-        distW = _DistWrapper(process_group, not no_dist, coordinator_rank)
-        if planner is None:
-            planner = DefaultLoadPlanner()
-
-        ckpt_kwargs = {}
-        if (ckpt_id := getattr(storage_reader, "checkpoint_id", None)) is not None:
-            ckpt_kwargs["checkpoint_id"] = ckpt_id
-            ckpt_kwargs["process_group"] = distW.group
-
-        @_dcp_method_logger(**ckpt_kwargs)
-        def local_step():
-            assert planner is not None
-            metadata = storage_reader.read_metadata()
-            planner.set_up_planner(state_dict, metadata, distW.is_coordinator)
-            storage_reader.set_up_storage_reader(metadata, distW.is_coordinator)
-
-            local_plan = planner.create_local_plan()
-            local_plan = storage_reader.prepare_local_plan(local_plan)
-            return local_plan
-
-        @_dcp_method_logger(**ckpt_kwargs)
-        def global_step(all_local_plans):
-            assert planner is not None
-            all_local_plans = planner.create_global_plan(all_local_plans)
-            all_local_plans = storage_reader.prepare_global_plan(all_local_plans)
-            return all_local_plans
-
-        central_plan: LoadPlan = distW.reduce_scatter("plan", local_step, global_step)
-        if distW.is_coordinator:
-            # Compare central_plan items with storage_reader.storage_data keys
-            dest_fqns = set()
-            storage_fqns = set()
-
-            # Extract FQNs from central_plan items
-            for item in central_plan.items:
-                if hasattr(item, "dest_index") and hasattr(item.dest_index, "fqn"):
-                    dest_fqns.add(item.dest_index.fqn)
-                if hasattr(item, "storage_index") and hasattr(item.storage_index, "fqn"):
-                    storage_fqns.add(item.storage_index.fqn)
-
-            # Get storage data keys
-            storage_data_keys = set()
-            if hasattr(storage_reader, "storage_data") and storage_reader.storage_data is not None:
-                storage_data_keys = set(item[0].fqn for item in storage_reader.storage_data.items())
-            state_dict_keys = set(state_dict.keys())
-            # Compare sets and log differences
-            # Remove any item that has "_extra_state" as substring in the sets
-            state_dict_keys = {fqn for fqn in state_dict_keys if "_extra_state" not in fqn}
-            dest_fqns = {fqn for fqn in dest_fqns if "_extra_state" not in fqn}
-            storage_fqns = {fqn for fqn in storage_fqns if "_extra_state" not in fqn}
-            storage_data_keys = {fqn for fqn in storage_data_keys if "_extra_state" not in fqn}
-
-            log.info("=== Load Plan FQN Analysis ===")
-            log.info(f"State Dict FQNs count: {len(state_dict_keys)}")
-            log.info(f"Destination FQNs count (without _extra_state): {len(dest_fqns)}")
-            log.info(f"Loaded FQNs count (without _extra_state): {len(storage_fqns)}")
-            log.info(f"In Storage keys count (without _extra_state): {len(storage_data_keys)}")
-
-            # Find missing keys in each direction
-            state_dict_missing_from_dest = state_dict_keys - dest_fqns
-            storage_data_missing_from_storage_fqns = storage_data_keys - storage_fqns
-
-            if state_dict_missing_from_dest:
-                log.info(
-                    f"State Dict FQNs missing from load plan ({len(state_dict_missing_from_dest)} items): {sorted(state_dict_missing_from_dest)}"
-                )
-            else:
-                log.info("✓ All State Dict FQNs found in storage_data")
-
-            if storage_data_missing_from_storage_fqns:
-                # If there are more than 100 "net_ema" keys in storage_data_missing_from_storage_fqns, summarize them
-                net_ema_keys = {k for k in storage_data_missing_from_storage_fqns if "net_ema" in k}
-                if len(net_ema_keys) > 100:
-                    storage_data_missing_from_storage_fqns = storage_data_missing_from_storage_fqns - net_ema_keys
-                    storage_data_missing_from_storage_fqns = set(
-                        storage_data_missing_from_storage_fqns
-                    )  # ensure set type
-                    storage_data_missing_from_storage_fqns.add("net_ema")
-                    log.info(
-                        f"Summarized {len(net_ema_keys)} 'net_ema' keys as 'net_ema' in missing storage data keys."
-                    )
-                log.info(
-                    f"Storage data keys not loaded by load plan ({len(storage_data_missing_from_storage_fqns)} items): {sorted(storage_data_missing_from_storage_fqns)}"
-                )
-            else:
-                log.info("✓ All storage data keys found in Loaded FQNs")
-
-            log.info("=== End Load Plan FQN Analysis ===")
-
-        @_dcp_method_logger(**ckpt_kwargs)
-        def read_data():
-            assert planner is not None
-            final_local_plan = planner.finish_plan(central_plan)
-            all_reads = storage_reader.read_data(final_local_plan, planner)
-
-            all_reads.wait()
-            return None
-
-        _ = distW.all_gather("read", read_data)
-
-except ImportError as e:
-    from torch.distributed.checkpoint.default_planner import DefaultLoadPlanner
-
-    log.critical(f"{e}, using default planner")
-
-StateDictItemPath = namedtuple("StateDictItemPath", ["state_dict", "save_path"])
-
-
-class ModelWrapper(Stateful):
-    """Wrapper for model state dict handling"""
-
-    def __init__(self, model: Union[nn.Module, List[nn.Module]], load_ema_to_reg: bool = False):
-        self.model = [model] if isinstance(model, nn.Module) else model
-        self.load_ema_to_reg = load_ema_to_reg
-        if self.load_ema_to_reg:
-            assert isinstance(model, DiffusionModel)
-
-    def state_dict(self) -> Dict[str, Any]:
-        _state_dict = {k: v for sd in map(get_model_state_dict, self.model) for k, v in sd.items()}
-        if self.load_ema_to_reg:
-            if not self.model[0].config.ema.enabled:
-                all_keys = list(_state_dict.keys())
-                assert all(k.startswith("net.") for k in all_keys), "All keys must start with net."
-                for k in all_keys:
-                    _state_dict[k.replace("net.", "net_ema.")] = _state_dict.pop(k)
-            else:
-                log.warning("EMA is enabled, will only load EMA weights from checkpoint file.")
-                all_keys = list(_state_dict.keys())
-                for k in all_keys:
-                    if k.startswith("net_ema."):
-                        break
-                else:
-                    raise ValueError("No EMA keys found in state_dict")
-                # do not load .net keys, since we do not need them anyway.
-                _state_dict = {k: _state_dict[k] for k in all_keys if not k.startswith("net.")}
-
-        if hasattr(self.model[0].config, "lora_config") and self.model[0].config.lora_config.enabled:
-            """
-            When using LoRA, `inject_adapter_in_model` modifies the target modules in place.
-            For example, `blocks[0].attn.q_proj.weight` will be modified to `blocks[0].attn.q_proj.base_layer.weight`.
-            This means that the model will have the key `blocks[0].attn.q_proj.base_layer.weight`,
-            but the checkpoint will have the key `blocks[0].attn.q_proj.weight`.
-            We need to map the model key to the checkpoint key.
-            """
-            self.checkpoint_to_model_key = {}
-            mapping_keys = {"base_layer.": "", "base_model.model.": ""}
-            keys_to_update = []
-            for k in _state_dict.keys():
-                new_key = k
-                for from_key, to_key in mapping_keys.items():
-                    new_key = new_key.replace(from_key, to_key)
-                if new_key != k:
-                    keys_to_update.append((k, new_key))
-                    self.checkpoint_to_model_key[new_key] = k
-            for k, new_key in keys_to_update:
-                _state_dict[new_key] = _state_dict.pop(k)
-
-        return _state_dict
-
-    def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
-        if hasattr(self.model[0].config, "lora_config") and self.model[0].config.lora_config.enabled:
-            if hasattr(self, "checkpoint_to_model_key"):
-                for checkpoint_key, model_key in self.checkpoint_to_model_key.items():
-                    state_dict[model_key] = state_dict.pop(checkpoint_key)
-            else:
-                raise ValueError("checkpoint_to_model_key is not set by `state_dict`")
-
-        if self.load_ema_to_reg:
-            if not self.model[0].config.ema.enabled:
-                all_keys = list(state_dict.keys())
-                assert all(k.startswith("net_ema.") for k in all_keys), "All keys must start with net_ema."
-                for k in all_keys:
-                    state_dict[k.replace("net_ema.", "net.")] = state_dict.pop(k)
-            else:
-                log.warning("EMA is enabled, will load EMA weights to regular model weights")
-                all_keys = list(state_dict.keys())
-                assert all(not k.startswith("net.") for k in all_keys), "No .net keys should be in state_dict"
-                for k in all_keys:
-                    if k.startswith("net_ema."):
-                        state_dict[k.replace("net_ema.", "net.")] = torch.clone(state_dict[k])
-        func = functools.partial(
-            set_model_state_dict,
-            model_state_dict=state_dict,
-            options=StateDictOptions(strict=False),
-        )
-        list(map(func, self.model))
-
-
-class OptimizerWrapper(Stateful):
-    def __init__(
-        self,
-        model: Union[nn.Module, List[nn.Module]],
-        optim: Union[torch.optim.Optimizer, List[torch.optim.Optimizer]],
-    ) -> None:
-        self.model = [model] if isinstance(model, nn.Module) else model
-        self.optim = [optim] if isinstance(optim, torch.optim.Optimizer) else optim
-
-    def state_dict(self) -> Dict[str, Any]:
-        func = functools.partial(
-            get_optimizer_state_dict,
-            options=StateDictOptions(flatten_optimizer_state_dict=True),
-        )
-        return {k: v for sd in map(func, self.model, self.optim) for k, v in sd.items()}
-
-    def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
-        func = functools.partial(
-            set_optimizer_state_dict,
-            optim_state_dict=state_dict,
-            options=StateDictOptions(flatten_optimizer_state_dict=True),
-        )
-        list(map(func, self.model, self.optim))
-
-
-class AsyncMode(str, enum.Enum):
-    DISABLED = "disabled"
-    ASYNC_WITH_PINNED_MEM = "async_with_pinned_mem"
-
-
-class Terminate:
-    pass
-
-
-class SaveDone:
-    def __init__(self, iteration: int, elapsed_time: float, succeeded: bool):
-        self.iteration = iteration
-        self.elapsed_time = elapsed_time
-        self.succeeded = succeeded
-
-    def __str__(self):
-        return f"SaveDone(iteration={self.iteration}, elapsed_time={self.elapsed_time}, succeeded={self.succeeded})"
-
-
-def save_checkpoint_in_background(
-    receiver_queue: multiprocessing.Queue,
-    sender_queue: multiprocessing.Queue,
-    checkpoint_config: CheckpointConfig,
-    job_config: JobConfig,
-) -> None:
-    """
-    Handles model checkpoint saving in a separate background process using PyTorch's distributed functionality.
-    This function runs in a dedicated process to avoid blocking the main training loop.
-
-    Args:
-        receiver_queue: Queue to receive state dictionaries and commands from the main process
-        sender_queue: Queue to send completion signals back to the main process
-        checkpoint_config: Configuration settings for checkpoint saving behavior
-        job_config: Configuration settings for the training job
-
-    Flow:
-        1. Initializes distributed processing environment
-        2. Continuously waits for state dictionaries to save
-        3. Saves checkpoints asynchronously
-        4. Signals completion back to main process
-        5. Terminates when receiving a Terminate signal
-
-    Raises:
-        AssertionError: If received object is neither Terminate signal nor valid state dict tuple
-
-    Note:
-        - Uses a different port than the main process to avoid conflicts
-        - Disables TorchElastic agent store for checkpoint operations
-        - Automatically cleans up distributed process group on exit
-    """
-    # Configure distributed environment
-    os.environ["MASTER_PORT"] = str(int(os.environ["MASTER_PORT"]) + 2)
-    os.environ["TORCHELASTIC_USE_AGENT_STORE"] = "False"
-
-    # Set up GPU device and distributed processing
-    torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
-    distributed.init()
-
-    # Initialize checkpointing mechanism
-    checkpoint_handler = DistributedCheckpointer(checkpoint_config, job_config, None, disable_async=True)
-
-    try:
-        while True:
-            log.debug("Checkpoint background process is ready for next task, waiting for new state_dict")
-            received_data = receiver_queue.get()
-            log.debug("Received new state_dict")
-
-            if isinstance(received_data, Terminate):
-                log.info("Received termination signal in checkpoint background process, closing sender queue")
-                sender_queue.put(Terminate())
-                sender_queue.close()
-                return
-
-            assert isinstance(received_data, tuple), "Received data must be a tuple of (state_dict, checkpoint_path)"
-            state_dict, checkpoint_path = received_data
-
-            # Save checkpoint and measure time taken
-            start_time = time.monotonic()
-            iteration = state_dict["trainer"][0]["iteration"]
-            elapsed_time = 0
-            succeeded = False
-            try:
-                checkpoint_handler.save_state_dict_worker(state_dict, checkpoint_path)
-                elapsed_time = time.monotonic() - start_time
-                log.info(
-                    f"Checkpoint saved successfully in background process. Time taken: {elapsed_time:.2f} seconds, iteration: {iteration}"
-                )
-                succeeded = True
-            except Exception as e:
-                log.error(f"Error saving checkpoint to {checkpoint_path}: {e}")
-                # continue because if the thread exits, the main thread keeps on adding to the queue
-            finally:
-                if elapsed_time == 0:
-                    elapsed_time = time.monotonic() - start_time
-                sender_queue.put(SaveDone(iteration, elapsed_time, succeeded))
-
-    finally:
-        log.info("Cleaning up: destroying distributed process group")
-        torch.distributed.destroy_process_group()
-
-
-class DistributedCheckpointer(AbstractCheckpointer):
-    KEYS_TO_SAVE = ["model", "optim", "scheduler", "trainer"]
-
-    def __init__(
-        self,
-        config_checkpoint: CheckpointConfig,
-        config_job: JobConfig,
-        callbacks: Optional[callback.CallBackGroup] = None,
-        disable_async: bool = False,
-    ):
-        super().__init__(config_checkpoint, config_job, callbacks)
-        self.config_checkpoint = config_checkpoint
-        if config_checkpoint.dcp_async_mode_enabled:
-            self.async_mode = AsyncMode.ASYNC_WITH_PINNED_MEM
-        else:
-            self.async_mode = AsyncMode.DISABLED
-
-        if disable_async:
-            self.async_mode = AsyncMode.DISABLED
-
-        if self.async_mode == AsyncMode.ASYNC_WITH_PINNED_MEM:
-            ctx = get_context("spawn")
-            self.mp_queue_send = ctx.Queue()
-            self.mp_queue_recv = ctx.Queue()
-            self.mp = ctx.Process(
-                target=save_checkpoint_in_background,
-                args=(
-                    self.mp_queue_send,
-                    self.mp_queue_recv,
-                    config_checkpoint,
-                    config_job,
-                ),
-                daemon=True,
-            )
-            self.mp.start()
-            self.cpu_offload_state_dict = None
-            self.staging = False
-            self.staging_ckpt_file = None
-            self.staging_stream = torch.cuda.Stream()
-
-    def keys_to_resume_during_load(self) -> Tuple[Set, Union[str, None]]:
-        latest_checkpoint_file = self._read_latest_checkpoint_file()
-
-        resume_keys = []
-
-        if latest_checkpoint_file is not None:
-            # 1. Resume training from latest_checkpoint.txt under the same name.
-            checkpoint_path = os.path.join(self.load_dirname, latest_checkpoint_file)
-            resume_keys.extend(self.KEYS_TO_SAVE)
-        else:
-            if self.load_path and not self.load_path.endswith(".pth"):
-                # 2. Load the module weights specified by config_checkpoint.path.
-                checkpoint_path = self.load_path
-                if self.load_s3_backend_key:
-                    checkpoint_path = f"s3://{self.config_checkpoint.load_from_object_store.bucket}/{checkpoint_path}"
-                    if not re.search(r"/checkpoints/iter_\d{9}/?$", checkpoint_path):
-                        old_ckpt_path = checkpoint_path
-                        # If path doesn't end with specific checkpoint, read latest checkpoint file
-                        latest_ckpt_path = os.path.join(checkpoint_path, "checkpoints/latest_checkpoint.txt")
-                        if easy_io.exists(latest_ckpt_path, backend_key=self.load_s3_backend_key):
-                            checkpoint_file = easy_io.load(
-                                latest_ckpt_path, backend_key=self.load_s3_backend_key
-                            ).strip()
-                            checkpoint_path = f"{checkpoint_path}/checkpoints/{checkpoint_file}"
-                        else:
-                            log.warning(
-                                f"Latest checkpoint file {latest_ckpt_path} not found, load from {old_ckpt_path}"
-                            )
-                            checkpoint_path = old_ckpt_path
-                if self.load_training_state:
-                    resume_keys.extend(self.KEYS_TO_SAVE)
-                else:
-                    resume_keys.append("model")
-                    if self.only_load_scheduler_state:
-                        resume_keys.append("scheduler")
-            elif self.load_path and self.load_path.endswith(".pth"):
-                checkpoint_path = self.load_path
-            else:
-                checkpoint_path = None
-        if len(self.keys_not_to_resume) > 0:
-            for key in self.keys_not_to_resume:
-                assert key in self.KEYS_TO_SAVE, f"Invalid key to resume: {key} not in {self.KEYS_TO_SAVE}"
-            resume_keys = [key for key in resume_keys if key not in self.keys_not_to_resume]
-        return set(resume_keys), checkpoint_path
-
-    @misc.timer("checkpoint loading")
-    def load(
-        self,
-        model: ImaginaireModel,
-        optimizer: torch.optim.Optimizer | None = None,
-        scheduler: torch.optim.lr_scheduler.LRScheduler | None = None,
-        grad_scaler: torch.amp.GradScaler | None = None,
-    ) -> int:
-        if self.callbacks is not None:
-            self.callbacks.on_load_checkpoint_start(model)
-
-        resume_keys, checkpoint_path = self.keys_to_resume_during_load()
-        resume_keys = sorted(resume_keys)
-        log.info(f"Resuming ckpt {checkpoint_path} with keys: {resume_keys}")
-
-        iteration = 0
-
-        if checkpoint_path is not None and not checkpoint_path.endswith(".pth"):
-            self._check_checkpoint_exists(checkpoint_path)
-            for key in resume_keys:
-                load_planner = DefaultLoadPlanner(allow_partial_load=True)
-                if hasattr(load_planner, "set_partial_channel_weight"):
-                    log.info(f"set_partial_channel_weight: {self.config_checkpoint.dcp_allow_mismatched_size}")
-                    load_planner.set_partial_channel_weight(self.config_checkpoint.dcp_allow_mismatched_size)
-                cur_key_ckpt_full_path = os.path.join(checkpoint_path, key)
-                log.info(f"Start loading checkpoint from {checkpoint_path}")
-                storage_reader = self.get_storage_reader(cur_key_ckpt_full_path)
-                torch.distributed.barrier()
-                log.info(f"starting {cur_key_ckpt_full_path}", rank0_only=False)
-                if key == "model":
-                    log.info("- Loading the model...")
-                    _model_wrapper = ModelWrapper(model)
-                    _state_dict = _model_wrapper.state_dict()
-                    load(_state_dict, storage_reader=storage_reader, planner=load_planner)
-                    _model_wrapper.load_state_dict(_state_dict)
-                elif key == "optim":
-                    log.info("- Loading the optimizer...")
-                    _optim_wrapper = OptimizerWrapper(model, optimizer)
-                    _state_dict = _optim_wrapper.state_dict()
-                    dcp.load(
-                        _state_dict,
-                        storage_reader=storage_reader,
-                        planner=load_planner,
-                    )
-                    _optim_wrapper.load_state_dict(_state_dict)
-                elif key == "scheduler":
-                    log.info("- Loading the scheduler...")
-                    _state_dict = scheduler.state_dict()
-                    dcp.load(
-                        _state_dict,
-                        storage_reader=storage_reader,
-                        planner=load_planner,
-                    )
-                    scheduler.load_state_dict(_state_dict)
-                elif key == "trainer":
-                    log.info("- Loading the trainer...")
-                    _state_dict = {
-                        "grad_scaler": grad_scaler.state_dict(),
-                        "iteration": iteration,
-                    }
-                    dcp.load(
-                        _state_dict,
-                        storage_reader=storage_reader,
-                        planner=load_planner,
-                    )
-                    grad_scaler.load_state_dict(_state_dict["grad_scaler"])
-                    iteration = _state_dict["iteration"]
-                else:
-                    raise ValueError(f"Invalid key: {key}. not support to resume.")
-            if self.callbacks is not None:
-                self.callbacks.on_load_checkpoint(model, state_dict=_state_dict)
-            log.info(f"Loaded checkpoint from {checkpoint_path} in iteration {iteration}")
-        elif checkpoint_path is not None and checkpoint_path.endswith(".pth"):
-            state = easy_io.load(checkpoint_path)
-            model_state = model.net.state_dict()
-
-            for k, v in list(state.items()):
-                tgt = model_state.get(k, None)
-                if tgt is None:
-                    continue
-                # If target param/buffer is a DTensor and checkpoint value is not, distribute it
-                if isinstance(tgt, DTensor) and not isinstance(v, DTensor):
-                    # Match device, dtype, and placements from the target DTensor
-                    v = v.to(tgt.device, dtype=tgt.dtype, copy=False)
-                    v = distribute_tensor(v, tgt.device_mesh, tgt.placements)
-                    state[k] = v
-                # If target is a plain Tensor but checkpoint is a DTensor, bring it local
-                if not isinstance(tgt, DTensor) and isinstance(v, DTensor):
-                    state[k] = v.to_local().to(tgt.device, dtype=tgt.dtype, copy=False)
-
-            model.load_state_dict(state, strict=False, pretrain_copy=True)
-            # Clear unused reserved memory from fp32
-            torch.cuda.empty_cache()
-            log.critical(f"Loaded checkpoint from {checkpoint_path}. **** This only happen at iteration 0 **** ")
-        else:
-            log.info("Training from scratch.")
-        torch.cuda.empty_cache()
-
-        if self.callbacks is not None:
-            self.callbacks.on_load_checkpoint_end(model, iteration=iteration, checkpoint_path=checkpoint_path)
-        return iteration
-
-    def _async_with_pinned_memory(self, checkpoint_file: str, state_dict: Dict[str, Tuple[Any, str]]) -> None:
-        try:
-            from torch.distributed._state_dict_utils import _copy_state_dict, _create_cpu_state_dict
-        except ImportError as e:
-            raise ImportError(
-                "Please install the latest PyTorch nightly to use async checkpointing with pinned memory."
-            ) from e
-        if self.cpu_offload_state_dict is None:
-            log.debug(f"Preparing the CPU memory, {time.monotonic()=}.:.2f")
-            self.cpu_offload_state_dict = _create_cpu_state_dict(state_dict, pin_memory=True, share_memory=True)
-
-        log.debug(f"Staging the state_dict, {time.monotonic()=}.:.2f")
-        with torch.cuda.stream(self.staging_stream):
-            self.cpu_offload_state_dict = _copy_state_dict(
-                state_dict,
-                self.cpu_offload_state_dict,
-                non_blocking=True,
-            )
-            self.staging = True
-            self.staging_ckpt_file = checkpoint_file
-
-        self.maybe_wait_for_staging()
-
-    def maybe_wait_for_staging(self) -> None:
-        if self.async_mode == AsyncMode.ASYNC_WITH_PINNED_MEM and self.staging:
-            if not self.staging_stream.query():
-                self.staging_stream.synchronize()
-
-            def sync_func():
-                self.mp_queue_send.put_nowait((self.cpu_offload_state_dict, self.staging_ckpt_file))
-
-            sync_func()
-            self.staging = False
-
-    def get_previous_checkpoint_results(self, wait_for: int = 0) -> None:
-        """Get the results of previously submitted checkpoints and pass them to callbacks if checkpoint succeeded"""
-        if self.async_mode == AsyncMode.ASYNC_WITH_PINNED_MEM:
-            try:
-                start_time = time.monotonic()
-                while not self.mp_queue_recv.empty() or wait_for > 0:
-                    try:
-                        ret = self.mp_queue_recv.get(timeout=1)
-                        if isinstance(ret, Terminate):
-                            log.info("Received termination event from checkpoint background process")
-                            break
-                        save_done: SaveDone = ret
-                        log.logger.info(f"Received checkpoint save result: {save_done}")
-                        if self.callbacks is not None and save_done.succeeded:
-                            self.callbacks.on_save_checkpoint_success(
-                                iteration=save_done.iteration, elapsed_time=save_done.elapsed_time
-                            )
-                    except queue.Empty:
-                        elapsed_time = time.monotonic() - start_time
-                        if elapsed_time > wait_for:
-                            break
-            except (EOFError, BrokenPipeError):
-                log.info("Queue was closed by checkpoint background process")
-
-    def get_storage_writer(self, checkpoint_path: str) -> Union[S3StorageWriter, FileSystemWriter]:
-        if self.save_to_object_store:
-            return S3StorageWriter(
-                credential_path=self.config_checkpoint.save_to_object_store.credentials,
-                path=checkpoint_path,
-            )
-        return FileSystemWriter(path=checkpoint_path)
-
-    def get_storage_reader(self, checkpoint_path: str) -> Union[S3StorageReader, FileSystemReader]:
-        if self.load_from_object_store:
-            return S3StorageReader(
-                credential_path=self.config_checkpoint.load_from_object_store.credentials,
-                path=checkpoint_path,
-            )
-        return FileSystemReader(checkpoint_path)
-
-    def save_state_dict_worker(self, to_save_dict: Dict[str, Tuple[Any, str]], checkpoint_file: str) -> None:
-        for k, (v, full_checkpoint_path) in to_save_dict.items():
-            storage_writer = self.get_storage_writer(full_checkpoint_path)
-            dcp.save(
-                v,
-                storage_writer=storage_writer,
-                planner=DefaultSavePlanner(dedup_save_to_lowest_rank=True),
-            )
-
-        if distributed.is_rank0():
-            print(f"Saving last checkpoint file {checkpoint_file}")
-            self._write_latest_checkpoint_file(checkpoint_file)
-
-        log.critical(f"Saved checkpoint to {os.path.join(self.save_dirname, checkpoint_file)}", rank0_only=True)
-
-    def save(
-        self,
-        model: ImaginaireModel,
-        optimizer: torch.optim.Optimizer,
-        scheduler: torch.optim.lr_scheduler.LRScheduler,
-        grad_scaler: torch.amp.GradScaler,
-        iteration: int,
-    ) -> None:
-        """Save network weights, optimizer parameters, scheduler parameters to a checkpoint.
-
-        Args:
-            model (ImaginaireModel): The PyTorch model.
-            optimizer (torch.optim.Optimizer): The model optimizer.
-            scheduler (torch.optim.lr_scheduler.LRScheduler): The optimization scheduler.
-            grad_scaler (torch.amp.GradScaler): The gradient scaler (for mixed precision training).
-            iteration (int): Current iteration number.
-        """
-        if self.async_mode == AsyncMode.ASYNC_WITH_PINNED_MEM:
-            self.get_previous_checkpoint_results(wait_for=0)
-
-        if self.callbacks is not None:
-            self.callbacks.on_save_checkpoint_start(model, iteration)
-
-        checkpoint_file = f"iter_{iteration:09}"
-        to_save_dict = {
-            "model": ModelWrapper(model).state_dict(),
-            "optim": OptimizerWrapper(model, optimizer).state_dict(),
-            "scheduler": scheduler.state_dict(),
-            "trainer": {
-                "grad_scaler": grad_scaler.state_dict(),
-                "iteration": iteration,
-            },
-        }
-        for k in to_save_dict.keys():
-            output_dirname = os.path.join(self.save_dirname, f"iter_{iteration:09}/{k}")
-            to_save_dict[k] = (to_save_dict[k], output_dirname)
-
-        if self.callbacks is not None:
-            self.callbacks.on_save_checkpoint(model, state_dict=to_save_dict)
-
-        if self.async_mode == AsyncMode.ASYNC_WITH_PINNED_MEM:
-            self._async_with_pinned_memory(checkpoint_file, to_save_dict)
-        else:
-            start_time = time.monotonic()
-            try:
-                self.save_state_dict_worker(to_save_dict, checkpoint_file)
-            finally:
-                if self.callbacks is not None:
-                    self.callbacks.on_save_checkpoint_success(
-                        iteration=iteration, elapsed_time=time.monotonic() - start_time
-                    )
-
-        # This measures exposed (synchronous) checkpoint time, on_save_checkpoint_success()
-        # is instead called to measure the entire duration for asynchronous checkpoint for the async case too.
-        if self.callbacks is not None:
-            self.callbacks.on_save_checkpoint_end(model=None, iteration=iteration)
-
-    def finalize(self) -> None:
-        super().finalize()
-        if self.async_mode == AsyncMode.ASYNC_WITH_PINNED_MEM:
-            if self.mp and self.mp.is_alive():
-                self.mp_queue_send.put(Terminate())
-                self.get_previous_checkpoint_results(wait_for=60)
-                self.mp.join()

lyra_2/_ext/imaginaire/checkpointer/dummy.py

@@ -1,47 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from typing import Optional
-
-import torch
-import torch.distributed
-
-from lyra_2._ext.imaginaire.checkpointer.base import AbstractCheckpointer
-from lyra_2._ext.imaginaire.model import ImaginaireModel
-
-
-class Checkpointer(AbstractCheckpointer):
-    """
-    A dummy checkpointer that does not save or load anything. This is useful for debugging jobs or share workload with collobrators.
-    """
-
-    def save(
-        self,
-        model: ImaginaireModel,
-        optimizer: torch.optim.Optimizer,
-        scheduler: torch.optim.lr_scheduler.LRScheduler,
-        grad_scaler: torch.amp.GradScaler,
-        iteration: int,
-    ) -> None:
-        pass
-
-    def load(
-        self,
-        model: ImaginaireModel,
-        optimizer: Optional[torch.optim.Optimizer] = None,
-        scheduler: Optional[torch.optim.lr_scheduler.LRScheduler] = None,
-        grad_scaler: Optional[torch.amp.GradScaler] = None,
-    ) -> int:
-        return 0

lyra_2/_ext/imaginaire/checkpointer/s3_filesystem.py

@@ -1,323 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import io
-import json
-import os
-import time
-from contextlib import contextmanager
-from typing import Generator, Union
-from urllib.parse import urlparse
-
-import boto3
-from botocore.config import Config
-from botocore.exceptions import ClientError
-from torch.distributed.checkpoint import FileSystemReader, FileSystemWriter
-from torch.distributed.checkpoint.filesystem import FileSystemBase
-
-from lyra_2._ext.imaginaire.utils import log
-
-
-class S3Stream(io.BytesIO):
-    """
-    Workaround for PyTorch manually closing the stream before we can upload it to S3. We override the close() as noop
-    and instead call our own _true_close() method to close the stream after we are done using it.
-    The commit at fault is https://github.com/pytorch/pytorch/commit/9c909bf3bb122db2cce95e2eb7459bbe50dfa15a
-    """
-
-    def close(self):
-        self.flush()
-        # No close
-
-    def _true_close(self):
-        super().close()
-
-
-class S3FileSystem(FileSystemBase):
-    """Implementation of FileSystemBase for AWS S3 storage."""
-
-    def __init__(
-        self,
-        credential_path: str,
-        max_attempts: int = 20,
-        initial_backoff: float = 1.0,
-        max_backoff: float = 30.0,
-        backoff_factor: float = 2.0,
-    ) -> None:
-        """
-        Initialize S3FileSystem with retry configuration.
-
-        Args:
-            credential_path: Path to AWS credentials JSON file
-            max_attempts: Maximum number of retry attempts
-            initial_backoff: Initial backoff time in seconds
-            max_backoff: Maximum backoff time in seconds
-            backoff_factor: Multiplicative factor for backoff time
-        """
-        with open(credential_path, "r") as f:
-            conf = json.load(f)
-
-        # Configure boto3 with retry settings
-        config = Config(
-            retries=dict(max_attempts=max_attempts, mode="adaptive"),  # Adaptive mode automatically handles throttling
-            connect_timeout=60,
-            read_timeout=60,
-            request_checksum_calculation="when_required",  # Data integrity check for uploads and downloads
-            response_checksum_validation="when_required",  # Data integrity check for uploads and downloads
-        )
-
-        self.s3_client = boto3.client("s3", config=config, **conf)
-        self.max_attempts = max_attempts
-        self.initial_backoff = initial_backoff
-        self.max_backoff = max_backoff
-        self.backoff_factor = backoff_factor
-
-    def _retry_with_backoff(self, operation_func, *args, **kwargs):
-        """
-        Execute an operation with exponential backoff retry logic.
-
-        Args:
-            operation_func: Function to execute
-            *args: Positional arguments for the function
-            **kwargs: Keyword arguments for the function
-
-        Returns:
-            Result of the operation function
-
-        Raises:
-            Exception: If all retry attempts fail
-        """
-        last_exception = None
-        backoff = self.initial_backoff
-
-        for attempt in range(self.max_attempts):
-            try:
-                return operation_func(*args, **kwargs)
-            except ClientError as e:
-                error_code = e.response.get("Error", {}).get("Code", "")
-                log.info(f"S3 Filesystem: Received ClientError: {error_code}", rank0_only=False)
-
-                # Handle specific error cases
-                if error_code in ["SlowDown", "ThrottlingException", "RequestLimitExceeded", "InternalError"]:
-                    last_exception = e
-                    if attempt < self.max_attempts - 1:  # Don't sleep on last attempt
-                        current_backoff = min(backoff, self.max_backoff)
-                        log.info(f"S3 Filesystem: Retrying in {current_backoff} seconds", rank0_only=False)
-                        time.sleep(current_backoff)
-                        backoff *= self.backoff_factor
-                        continue
-                # For other client errors, raise immediately
-                raise
-            except Exception as e:
-                log.info(f"S3 Filesystem: Received Exception: {str(e)}", rank0_only=False)
-                last_exception = e
-                if attempt < self.max_attempts - 1:
-                    current_backoff = min(backoff, self.max_backoff)
-                    log.info(f"S3 Filesystem: Retrying in {current_backoff} seconds", rank0_only=False)
-                    time.sleep(current_backoff)
-                    backoff *= self.backoff_factor
-                    continue
-
-        raise last_exception
-
-    @contextmanager
-    def create_stream(self, path: Union[str, os.PathLike], mode: str) -> Generator[io.IOBase, None, None]:
-        """Create a stream for reading from or writing to S3 with retry logic."""
-        path_str = str(path)
-        bucket, key = self._parse_s3_uri(path_str)
-        log.info(f"S3 Filesystem: Creating stream for {key} in bucket {bucket}", rank0_only=False)
-
-        if mode == "rb":
-            stream = io.BytesIO()
-            try:
-
-                def download_operation():
-                    self.s3_client.download_fileobj(bucket, key, stream)
-                    stream.seek(0)
-
-                log.info(f"S3 Filesystem: Downloading {key} from bucket {bucket}", rank0_only=False)
-                self._retry_with_backoff(download_operation)
-                log.info("S3 Filesystem: Download complete", rank0_only=False)
-                yield stream
-            finally:
-                stream.close()
-        elif mode == "wb":
-            stream = S3Stream()
-            try:
-                yield stream
-
-                def upload_operation():
-                    stream.seek(0)
-                    self.s3_client.upload_fileobj(stream, bucket, key)
-
-                log.info(f"S3 Filesystem: Uploading {key} to bucket {bucket}", rank0_only=False)
-                self._retry_with_backoff(upload_operation)
-                log.info("S3 Filesystem: Upload complete", rank0_only=False)
-            finally:
-                stream._true_close()
-        else:
-            raise ValueError(f"Unsupported mode: {mode}")
-
-    def concat_path(self, path: Union[str, os.PathLike], suffix: str) -> Union[str, os.PathLike]:
-        """Concatenate S3 path with suffix."""
-        path_str = str(path)
-        if path_str.endswith("/"):
-            return f"{path_str}{suffix}"
-        return f"{path_str}/{suffix}"
-
-    def init_path(self, path: Union[str, os.PathLike]) -> Union[str, os.PathLike]:
-        """Initialize and validate S3 path."""
-        path_str = str(path)
-        if not path_str.startswith("s3://"):
-            raise ValueError(f"Invalid S3 URI: {path_str}. Must start with 's3://'")
-        return path_str
-
-    def rename(self, path: Union[str, os.PathLike], new_path: Union[str, os.PathLike]) -> None:
-        """Rename (move) an object in S3 with retry logic."""
-        src_bucket, src_key = self._parse_s3_uri(str(path))
-        dst_bucket, dst_key = self._parse_s3_uri(str(new_path))
-
-        def copy_operation():
-            copy_source = {"Bucket": src_bucket, "Key": src_key}
-            self.s3_client.copy(copy_source, dst_bucket, dst_key)
-
-        self._retry_with_backoff(copy_operation)
-
-        def delete_operation():
-            self.s3_client.delete_object(Bucket=src_bucket, Key=src_key)
-
-        self._retry_with_backoff(delete_operation)
-
-    def mkdir(self, path: Union[str, os.PathLike]) -> None:
-        """
-        Create a "directory" in S3.
-
-        Note: S3 doesn't have real directories, but we can create an empty object
-        with a trailing slash to simulate a directory.
-        """
-        # Creating same buckets from different ranks can cause rate limit issues in GCP.
-        # In object store, we don't need to create a directory.
-        pass
-
-    @classmethod
-    def validate_checkpoint_id(cls, checkpoint_id: Union[str, os.PathLike]) -> bool:
-        """Validate if the checkpoint_id is a valid S3 URI."""
-        checkpoint_id_str = str(checkpoint_id)
-        try:
-            if not checkpoint_id_str.startswith("s3://"):
-                return False
-            parsed = urlparse(checkpoint_id_str)
-            return bool(parsed.netloc and parsed.path)  # Must have bucket and key
-        except Exception:
-            return False
-
-    def exists(self, path: Union[str, os.PathLike]) -> bool:
-        """Check if an object exists in S3 with retry logic."""
-        bucket, key = self._parse_s3_uri(str(path))
-        try:
-
-            def head_operation():
-                self.s3_client.head_object(Bucket=bucket, Key=key)
-
-            self._retry_with_backoff(head_operation)
-            return True
-        except ClientError as e:
-            if e.response.get("Error", {}).get("Code", "") == "404":
-                return False
-            raise
-
-    def rm_file(self, path: Union[str, os.PathLike]) -> None:
-        """Remove a file from S3 with retry logic."""
-        bucket, key = self._parse_s3_uri(str(path))
-
-        def delete_operation():
-            self.s3_client.delete_object(Bucket=bucket, Key=key)
-
-        self._retry_with_backoff(delete_operation)
-
-    def _parse_s3_uri(self, uri: str) -> tuple[str, str]:
-        """
-        Parse an S3 URI into bucket and key.
-
-        Args:
-            uri: S3 URI in the format s3://bucket-name/key
-
-        Returns:
-            Tuple of (bucket_name, key)
-
-        Raises:
-            ValueError: If the URI is invalid
-        """
-        uri = uri if isinstance(uri, str) else str(uri)
-        if not uri.startswith("s3://"):
-            raise ValueError(f"Invalid S3 URI: {uri}. Must start with 's3://'")
-
-        parsed = urlparse(uri)
-        bucket = parsed.netloc
-
-        # Remove leading slash from key
-        key = parsed.path.lstrip("/")
-
-        if not bucket:
-            raise ValueError(f"Invalid S3 URI: {uri}. No bucket specified")
-
-        return bucket, key
-
-
-class S3StorageWriter(FileSystemWriter):
-    def __init__(
-        self,
-        credential_path: str,
-        path: str,
-        **kwargs,
-    ) -> None:
-        """
-        Initialize an S3 writer for distributed checkpointing.
-
-        Args:
-            region (str): The AWS region for S3.
-            path (str): The S3 URI to write checkpoints to.
-            kwargs (dict): Keyword arguments to pass to the parent :class:`FileSystemWriter`.
-        """
-        super().__init__(
-            path=path,
-            sync_files=False,
-            **kwargs,
-        )
-        self.fs = S3FileSystem(credential_path)  # type: ignore
-        self.path = self.fs.init_path(path)
-
-    @classmethod
-    def validate_checkpoint_id(cls, checkpoint_id: Union[str, os.PathLike]) -> bool:
-        return S3FileSystem.validate_checkpoint_id(checkpoint_id)
-
-
-class S3StorageReader(FileSystemReader):
-    def __init__(self, credential_path: str, path: Union[str, os.PathLike]) -> None:
-        """
-        Initialize an S3 reader for distributed checkpointing.
-
-        Args:
-            region (str): The AWS region for S3.
-            path (Union[str, os.PathLike]): The S3 path to read checkpoints from.
-        """
-        super().__init__(path)
-        self.fs = S3FileSystem(credential_path)  # type: ignore
-        self.path = self.fs.init_path(path)
-        self.sync_files = False
-
-    @classmethod
-    def validate_checkpoint_id(cls, checkpoint_id: Union[str, os.PathLike]) -> bool:
-        return S3FileSystem.validate_checkpoint_id(checkpoint_id)

lyra_2/_ext/imaginaire/config.py

@@ -1,445 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Training config system for Imaginare4"""
-
-from __future__ import annotations
-
-import os
-from typing import Any, Dict, Optional, Type, TypeVar, Union
-
-import attrs
-import torch
-import torch.utils.data
-import torch.utils.data.distributed
-
-try:
-    from megatron.core import ModelParallelConfig
-
-    USE_MEGATRON = True
-except ImportError:
-    USE_MEGATRON = False
-    print("Megatron-core is not installed.")
-
-from lyra_2._ext.imaginaire.lazy_config import LazyCall as L
-from lyra_2._ext.imaginaire.lazy_config import LazyDict
-from lyra_2._ext.imaginaire.utils import callback, distributed
-from lyra_2._ext.imaginaire.utils.misc import Color
-
-T = TypeVar("T")
-
-
-def _is_attrs_instance(obj: object) -> bool:
-    """
-    Helper function to check if an object is an instance of an attrs-defined class.
-
-    Args:
-        obj: The object to check.
-
-    Returns:
-        bool: True if the object is an instance of an attrs-defined class, False otherwise.
-    """
-    return hasattr(obj, "__attrs_attrs__")
-
-
-def make_freezable(cls: T) -> T:
-    """
-    A decorator that adds the capability to freeze instances of an attrs-defined class.
-
-    NOTE: This requires the wrapped attrs to be defined with attrs.define(slots=False) because we need
-    to hack on a "_is_frozen" attribute.
-
-    This decorator enhances an attrs-defined class with the ability to be "frozen" at runtime.
-    Once an instance is frozen, its attributes cannot be changed. It also recursively freezes
-    any attrs-defined objects that are attributes of the class.
-
-    Usage:
-        @make_freezable
-        @attrs.define(slots=False)
-        class MyClass:
-            attribute1: int
-            attribute2: str
-
-        obj = MyClass(1, 'a')
-        obj.freeze()  # Freeze the instance
-        obj.attribute1 = 2  # Raises AttributeError
-
-    Args:
-        cls: The class to be decorated.
-
-    Returns:
-        The decorated class with added freezing capability.
-    """
-
-    if not hasattr(cls, "__dict__"):
-        raise TypeError(
-            "make_freezable cannot be used with classes that do not define __dict__. Make sure that the wrapped "
-            "class was defined with `@attrs.define(slots=False)`"
-        )
-
-    original_setattr = cls.__setattr__
-
-    def setattr_override(self, key, value) -> None:  # noqa: ANN001
-        """
-        Override __setattr__ to allow modifications during initialization
-        and prevent modifications once the instance is frozen.
-        """
-        if hasattr(self, "_is_frozen") and self._is_frozen and key != "_is_frozen":
-            raise AttributeError("Cannot modify frozen instance")
-        original_setattr(self, key, value)  # type: ignore
-
-    cls.__setattr__ = setattr_override  # type: ignore
-
-    def freeze(self: object) -> None:
-        """
-        Freeze the instance and all its attrs-defined attributes.
-        """
-        for _, value in attrs.asdict(self, recurse=False).items():
-            if _is_attrs_instance(value) and hasattr(value, "freeze"):
-                value.freeze()
-        self._is_frozen = True  # type: ignore
-
-    cls.freeze = freeze  # type: ignore
-
-    return cls
-
-
-def _pretty_print_attrs_instance(obj: object, indent: int = 0, use_color: bool = False) -> str:
-    """
-    Recursively pretty prints attrs objects with color.
-    """
-
-    assert attrs.has(obj.__class__)
-
-    lines: list[str] = []
-    for attribute in attrs.fields(obj.__class__):
-        value = getattr(obj, attribute.name)
-        if attrs.has(value.__class__):
-            if use_color:
-                lines.append("   " * indent + Color.cyan("* ") + Color.green(attribute.name) + ":")
-            else:
-                lines.append("   " * indent + "* " + attribute.name + ":")
-            lines.append(_pretty_print_attrs_instance(value, indent + 1, use_color))
-        else:
-            if use_color:
-                lines.append(
-                    "   " * indent + Color.cyan("* ") + Color.green(attribute.name) + ": " + Color.yellow(value)
-                )
-            else:
-                lines.append("   " * indent + "* " + attribute.name + ": " + str(value))
-    return "\n".join(lines)
-
-
-def pretty_print_overrides(overrides: Optional[list[str]] = None, use_color: bool = False) -> str:
-    """
-    Pretty prints overrides.
-    """
-
-    lines: list[str] = []
-    lines.append(Color.cyan("* ") + Color.green("overrides") + ": ")
-    for override in overrides:
-        if override == "--":
-            continue
-        if override.startswith("~"):
-            attribute_name = override[1:]
-            attribute_value = None
-        else:
-            attribute_name, attribute_value = override.split("=")
-        if use_color:
-            lines.append("   " + Color.cyan("* ") + Color.green(attribute_name) + ": " + Color.yellow(attribute_value))
-        else:
-            lines.append("   " + "* " + attribute_name + ": " + str(attribute_value))
-
-    return "\n".join(lines)
-
-
-@make_freezable
-@attrs.define(slots=False)  # slots=False is required for make_freezable. See the make_freezable notes for more info.
-class ObjectStoreConfig:
-    # Whether the file I/O is from object store instead of local disk.
-    enabled: bool = False
-    # Path to the object store credentials file.
-    credentials: str = ""
-    # Object store bucket to read from / write to the objects.
-    bucket: str = ""
-
-
-@make_freezable
-@attrs.define(slots=False)
-class JobConfig:
-    # Project name.
-    project: str = ""
-    # Experiment name.
-    group: str = ""
-    # Run/job name.
-    name: str = ""
-    @property
-    def path(self) -> str:
-        return f"{self.project}/{self.group}/{self.name}"
-
-    @property
-    def path_local(self) -> str:
-        local_root = os.environ.get("IMAGINAIRE_OUTPUT_ROOT", "/tmp/imaginaire4-output")
-        return f"{local_root}/{self.path}"
-
-
-@make_freezable
-@attrs.define(slots=False)
-class EMAConfig:
-    # Enable tracking a set of exponential moving average (EMA) weights.
-    enabled: bool = False
-    # EMA decay rate.
-    beta: float = 0.9999
-    # Enable removing "_orig_mod-" from buffer names that is added by torch.compile
-    torch_compile_buffer_renaming: bool = False
-
-
-@make_freezable
-@attrs.define(slots=False)
-class PowerEMAConfig:
-    # Enable tracking a set of exponential moving average (EMA) weights.
-    enabled: bool = False
-    # EDM2 paper EMA decay rate.
-    s: float = 0.1
-    # Enable removing "_orig_mod-" from buffer names that is added by torch.compile
-    torch_compile_buffer_renaming: bool = False
-
-
-@make_freezable
-@attrs.define(slots=False)
-class DDPConfig:
-    # Traverse the computation graph to find parameters that don't receive gradients.
-    find_unused_parameters: bool = False
-    # Set to True if the computation graph does not change during the whole training loop.
-    static_graph: bool = True
-    # Set to True if we want to synchronize buffers. Set to False if the sync is going to be handled elsewhere.
-    broadcast_buffers: bool = True
-
-
-@make_freezable
-@attrs.define(slots=False)
-class CuDNNConfig:
-    # Set to True for better reproducibility of the results (only using deterministic cudnn functions).
-    deterministic: bool = False
-    # If set to True, cudnn will benchmark several algorithms and pick the fastest one.
-    benchmark: bool = True
-
-
-@make_freezable
-@attrs.define(slots=False)
-class JITConfig:
-    # Enable exporting a JIT compiled model.
-    enabled: bool = False
-    # Input tensor shape, for example input.
-    input_shape: Union[list[int], None] = None
-    # Device to compile onto.
-    device: str = "cuda"
-    # # Data type to compile onto.
-    dtype: str = "bfloat16"
-    # Strict mode for PyTorch JIT.
-    strict: bool = True
-
-
-@make_freezable
-@attrs.define(slots=False)
-class CheckpointConfig:
-    # possible checkpoint class
-    type: Optional[Dict] = None
-    # for dcp, whether to use async mode
-    dcp_async_mode_enabled: bool = False
-    # Configs for saving the checkpoints to object store.
-    save_to_object_store: ObjectStoreConfig = attrs.field(factory=ObjectStoreConfig)
-    # Save the checkpoint every N iterations.
-    save_iter: int = 999999999
-    # Configs for loading the checkpoints from object store.
-    load_from_object_store: ObjectStoreConfig = attrs.field(factory=ObjectStoreConfig)
-    # Path of model weights to resume the checkpoint from.
-    load_path: str = ""
-    # Whether to load the training states (optimizer/scheduler/grad-scaler) from the checkpoint path.
-    load_training_state: bool = False
-    # Whether to load the scheduler state only from the checkpoint path. If load_training_state is True, this will be ignored.
-    only_load_scheduler_state: bool = False
-    # Load state_dict to the models in strict mode.
-    strict_resume: bool = True
-    # Configs for JIT compiling EMA model.
-    jit: JITConfig = attrs.field(factory=JITConfig)
-    # Print detailed information during checkpoint saving/loading.
-    verbose: bool = True
-    # keys not to resume from the checkpoint, choices: ["model", "optim", "scheduler", "trainer"]
-    keys_not_to_resume: list[str] = []
-    # Whether to use the local filesystem for broadcasting checkpoint data (used for Tensor Parallel Checkpointer).
-    broadcast_via_filesystem: bool = False
-    load_ema_to_reg: bool = False  # used in inference, load EMA weights to regular model
-    # In dcp planner, skip the weight shape check, load weights into the model even weight shape is different
-    dcp_allow_mismatched_size: bool = False
-
-
-@make_freezable
-@attrs.define(slots=False)
-class NVTXConfig:
-    """Config for NVTX ranges used in the main training loop.
-
-    See tutorials/nanogpt for more details on how to integrate profiling into your model."""
-
-    # Enable the NVTX ranges.
-    enabled: bool = False
-    # Synchronize everything in each NVTX range.
-    cuda_synchronize: bool = False
-
-
-@make_freezable
-@attrs.define(slots=False)
-class StragglerDetectionConfig:
-    """Config for Straggler detection tool."""
-
-    # Enable the Straggler Detection.
-    enabled: bool = False
-    # How frequently should the Straggler reports be generated.
-    report_freq: int = 100
-    # How frequently iterations should be profiled
-    profile_freq: int = 1
-    # What is the maximum relative difference between GPUs after they are considered stragglers
-    max_diff: float = 2.0
-    # Should the error be raised when straggler is detected
-    raise_error: bool = True
-    # Analyze kernels in the forward pass.
-    analyze_forward: bool = True
-    # Analyze kernels in the backward pass.
-    analyze_backward: bool = True
-    # Analyze kernels in the optimizer.
-    analyze_optimizer: bool = True
-    # Analyze dataloading time.
-    analyze_dataloading: bool = True
-
-
-@make_freezable
-@attrs.define(slots=False)
-class Profiling:
-    enable_profiling: bool = False
-    enable_memory_snapshot: bool = False
-    save_s3: bool = False
-    profile_freq: int = 1
-    # Target ranks for profiling, each entry must be >=0 and < world_size.
-    target_ranks: list[int] = list(range(8))
-    # Set `record_shape` and `profile_memory` to False to reduce profile size.
-    record_shape: bool = False
-    profile_memory: bool = False
-    with_stack: bool = True
-    with_modules: bool = True
-
-
-@make_freezable
-@attrs.define(slots=False)
-class TrainerConfig:
-    from lyra_2._ext.imaginaire.trainer import ImaginaireTrainer
-
-    type: Type[ImaginaireTrainer] = ImaginaireTrainer
-    # Set the callback class.
-    # Defaults to the callbacks below.
-    callbacks: LazyDict = LazyDict(
-        dict(
-            ema=L(callback.EMAModelCallback)(),
-            progress_bar=L(callback.ProgressBarCallback)(),
-        )
-    )
-    # distributed parallelism strategy
-    distributed_parallelism: str = "ddp"
-    # Distributed data parallel configs.
-    ddp: DDPConfig = attrs.field(factory=DDPConfig)
-    # cuDNN configs.
-    cudnn: CuDNNConfig = attrs.field(factory=CuDNNConfig)
-    # Set the random seed.
-    seed: int = 0
-    # Set the random seed based on current timestamp
-    timestamp_seed: bool = False
-    # Gradient scaler arguments (for torch.amp.GradScaler).
-    grad_scaler_args: dict = attrs.field(factory=lambda: dict(enabled=False))
-    # Maximum number of iterations to train the model.
-    max_iter: int = 999999999
-    # Maximum number of iterations to validate the model. If None, validate on the entire dataset.
-    max_val_iter: int | None = None
-    # How often we log the training stats.
-    logging_iter: int = 100
-    # Whether we want to run the validation routines.
-    run_validation: bool = True
-    # How often we evaluate on the validation set.
-    validation_iter: int = 999999999
-    # Kill the process after N seconds since the last iteration (usually means dead job).
-    timeout_period: int = 999999999
-    # Tensor memory organization format.
-    memory_format: torch.memory_format = torch.preserve_format
-    # Gradient accumulation (update step every N iteration).
-    grad_accum_iter: int = 1
-    # Straggler Detection config
-    straggler_detection: StragglerDetectionConfig = attrs.field(factory=StragglerDetectionConfig)
-    # Profiling config
-    profiling: Profiling = attrs.field(factory=Profiling)
-
-
-@make_freezable
-@attrs.define(slots=False)
-class Config:
-    """Config for an imaginaire4 job.
-
-    See /README.md/Configuration System for more info.
-    """
-
-    # Model configs.
-    model: LazyDict
-    # Optimizer configs.
-    optimizer: LazyDict
-    # Scheduler configs.
-    scheduler: LazyDict
-    # Training data configs.
-    dataloader_train: LazyDict
-    # Validation data configs.
-    dataloader_val: LazyDict
-
-    # Training job configs.
-    job: JobConfig = attrs.field(factory=JobConfig)
-
-    # Trainer configs.
-    trainer: TrainerConfig = attrs.field(factory=TrainerConfig)
-
-    if USE_MEGATRON:
-        # Megatron-Core configs
-        model_parallel: ModelParallelConfig = attrs.field(factory=ModelParallelConfig)
-    else:
-        model_parallel: None = None
-
-    # Checkpointer configs.
-    checkpoint: CheckpointConfig = attrs.field(factory=CheckpointConfig)
-
-    # enable upload reproducible setup to s3
-    upload_reproducible_setup: bool = False
-
-    def pretty_print(self, use_color: bool = False) -> str:
-        return _pretty_print_attrs_instance(self, 0, use_color)
-
-    def to_dict(self) -> dict[str, Any]:
-        return attrs.asdict(self)
-
-    def validate(self) -> None:
-        """Validate that the config has all required fields."""
-
-        # broadcast job.name across all ranks to make sure it is consistent
-        # otherwise, unaligned job names leads unaligned path to save checkpoints
-        job_name_tensor = torch.ByteTensor(bytearray(self.job.name, "utf-8")).cuda()
-        distributed.broadcast(job_name_tensor, 0)
-        self.job.name = job_name_tensor.cpu().numpy().tobytes().decode("utf-8")
-
-        assert self.job.project != ""
-        assert self.job.group != ""
-        assert self.job.name != ""

lyra_2/_ext/imaginaire/flags.py

@@ -1,52 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Feature flags."""
-
-import os
-from dataclasses import dataclass
-
-
-def _parse_bool(value: str) -> bool:
-    """Parse string to a boolean."""
-    return value.lower() in ["true", "1", "yes", "y"]
-
-
-INTERNAL = _parse_bool(os.environ.get("COSMOS_INTERNAL", "0"))
-"""Whether to enable internal features."""
-
-SMOKE = _parse_bool(os.environ.get("COSMOS_SMOKE", "0"))
-"""Whether to enable smoke test.
-
-Disables expensive operations such as checkpoint loading.
-"""
-
-VERBOSE = _parse_bool(os.environ.get("COSMOS_VERBOSE", "0"))
-"""Whether to enable verbose output."""
-
-EXPERIMENTAL_CHECKPOINTS = _parse_bool(os.environ.get("COSMOS_EXPERIMENTAL_CHECKPOINTS", "0"))
-"""Whether to enable experimental checkpoints."""
-
-
-@dataclass
-class Flags:
-    internal: bool = INTERNAL
-    smoke: bool = SMOKE
-    verbose: bool = VERBOSE
-    experimental_checkpoints: bool = EXPERIMENTAL_CHECKPOINTS
-
-
-FLAGS = Flags()
-"""Convenience object for accessing flags."""

lyra_2/_ext/imaginaire/functional/batch_ops.py

@@ -1,61 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-# Functions for performing operations with broadcasting to the right axis
-#
-# Example
-# input1: tensor of size (N1, N2)
-# input2: tensor of size (N1, N2, N3, N4)
-# batch_mul(input1, input2) = input1[:, :, None, None] * input2
-#
-# If the common dimensions don't match, we raise an assertion error.
-
-from torch import Tensor
-
-
-def common_broadcast(x: Tensor, y: Tensor) -> tuple[Tensor, Tensor]:
-    ndims1 = x.ndim
-    ndims2 = y.ndim
-
-    common_ndims = min(ndims1, ndims2)
-    for axis in range(common_ndims):
-        assert x.shape[axis] == y.shape[axis], "Dimensions not equal at axis {}".format(axis)
-
-    if ndims1 < ndims2:
-        x = x.reshape(x.shape + (1,) * (ndims2 - ndims1))
-    elif ndims2 < ndims1:
-        y = y.reshape(y.shape + (1,) * (ndims1 - ndims2))
-
-    return x, y
-
-
-def batch_add(x: Tensor, y: Tensor) -> Tensor:
-    x, y = common_broadcast(x, y)
-    return x + y
-
-
-def batch_mul(x: Tensor, y: Tensor) -> Tensor:
-    x, y = common_broadcast(x, y)
-    return x * y
-
-
-def batch_sub(x: Tensor, y: Tensor) -> Tensor:
-    x, y = common_broadcast(x, y)
-    return x - y
-
-
-def batch_div(x: Tensor, y: Tensor) -> Tensor:
-    x, y = common_broadcast(x, y)
-    return x / y

lyra_2/_ext/imaginaire/functional/lr_scheduler.py

@@ -1,178 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from typing import Optional
-
-import numpy as np
-
-from lyra_2._ext.imaginaire.utils import distributed, log
-
-
-class TeroPolyScheduler:
-    def __init__(
-        self,
-        total_Mimg: int,
-        batch_size: int,
-        ref_Mimg: Optional[int] = None,
-        ref_batches: float = 70e3 / 1024,
-        max_lr_ratio: Optional[float] = 1.0,
-        min_lr_ratio: Optional[float] = None,
-        rampup_Mimg: float = 0,
-        rampdown_Mimg: int = 0,
-        verbosity_interval: int = 0,
-        formula: str = "poly",
-        poly_exp: float = 0.5,
-    ):
-        self.total_Mimg = total_Mimg
-        self.batch_size = batch_size * distributed.get_world_size()
-        self.ref_Mimg = ref_Mimg or ref_batches * batch_size / 1e6
-        self.ref_batches = ref_batches
-        self.max_lr_ratio = max_lr_ratio
-        self.min_lr_ratio = min_lr_ratio
-        self.rampup_Mimg = rampup_Mimg
-        self.rampdown_Mimg = rampdown_Mimg
-        self.verbosity_interval = verbosity_interval
-        self.formula = formula
-        self.poly_exp = poly_exp
-
-        self._model = None
-
-    @property
-    def model(self):
-        return self._model
-
-    @model.setter
-    def model(self, model):
-        self._model = model
-
-    def schedule(self, n, **kwargs):
-        cur_Mimg = getattr(self.model, "sample_counter", 0) / 1e6
-
-        if self.formula == "constant":
-            lr = 1.0
-        elif self.formula == "poly":
-            lr = max(cur_Mimg / self.ref_Mimg, 1e-8) ** -self.poly_exp
-        else:
-            raise ValueError(f'Invalid learning rate formula "{self.formula}"')
-
-        if self.max_lr_ratio is not None:
-            lr = min(lr, self.max_lr_ratio)
-        if self.min_lr_ratio is not None:
-            lr = max(lr, self.min_lr_ratio)
-
-        if self.rampup_Mimg > 0 and cur_Mimg < self.rampup_Mimg:
-            lr *= cur_Mimg / self.rampup_Mimg
-        if self.rampdown_Mimg > 0 and cur_Mimg > self.total_Mimg - self.rampdown_Mimg:
-            lr *= (self.total_Mimg - cur_Mimg) / self.rampdown_Mimg
-
-        return lr
-
-    def __call__(self, n, **kwargs):
-        return self.schedule(n, **kwargs)
-
-
-class LambdaWarmUpCosineScheduler:
-    """
-    A learning rate scheduler that combines warm-up with a cosine decay schedule for multiple cycles.
-    It supports different configurations for each cycle, including the number of warm-up steps, minimum
-    and maximum scaling factors for the learning rate.
-
-    The scheduler is intended to be used with a base learning rate of 1.0, where the actual learning
-    rate at any step is the base learning rate multiplied by the scaling factor computed by the scheduler.
-
-    Parameters:
-        warm_up_steps (list[int]): List of integers where each element represents the number of warm-up
-                                   steps for the corresponding cycle.
-        f_min (list[float]): List of the minimum scaling factors for each cycle after warm-up.
-        f_max (list[float]): List of the maximum scaling factors at the start and end of each cosine cycle.
-        f_start (list[float]): List of starting scaling factors for each warm-up phase.
-        cycle_lengths (list[int]): List of the total lengths of each cycle, including warm-up steps.
-        verbosity_interval (int, optional): Interval of training steps at which to print current step and
-                                            scaling factor information. Set to 0 by default to disable verbosity.
-
-    Examples:
-        >>> scheduler = LambdaWarmUpCosineScheduler2(
-                warm_up_steps=[10, 10],
-                f_min=[0.1, 0.1],
-                f_max=[1.0, 1.0],
-                f_start=[0.01, 0.01],
-                cycle_lengths=[50, 50],
-                verbosity_interval=10)
-        >>> for step in range(100):
-        >>>     lr_multiplier = scheduler(step)
-        >>>     print(f"Step {step}: LR Multiplier = {lr_multiplier}")
-    """
-
-    def __init__(self, warm_up_steps, f_min, f_max, f_start, cycle_lengths, verbosity_interval=0):
-        assert len(warm_up_steps) == len(f_min) == len(f_max) == len(f_start) == len(cycle_lengths)
-        self.lr_warm_up_steps = warm_up_steps
-        self.f_start = f_start
-        self.f_min = f_min
-        self.f_max = f_max
-        self.cycle_lengths = cycle_lengths
-        self.cum_cycles = np.cumsum([0] + list(self.cycle_lengths))
-        self.last_f = 0.0
-        self.verbosity_interval = verbosity_interval
-
-    def find_in_interval(self, n):
-        interval = 0
-        for cl in self.cum_cycles[1:]:
-            if n <= cl:
-                return interval
-            interval += 1
-
-    def schedule(self, n, **kwargs):
-        cycle = self.find_in_interval(n)
-        n = n - self.cum_cycles[cycle]
-        if self.verbosity_interval > 0:
-            if n % self.verbosity_interval == 0:
-                log.info(f"current step: {n}, recent lr-multiplier: {self.last_f}, current cycle {cycle}")
-        if n < self.lr_warm_up_steps[cycle]:
-            f = (self.f_max[cycle] - self.f_start[cycle]) / self.lr_warm_up_steps[cycle] * n + self.f_start[cycle]
-            self.last_f = f
-            return f
-        else:
-            t = (n - self.lr_warm_up_steps[cycle]) / (self.cycle_lengths[cycle] - self.lr_warm_up_steps[cycle])
-            t = min(t, 1.0)
-            f = self.f_min[cycle] + 0.5 * (self.f_max[cycle] - self.f_min[cycle]) * (1 + np.cos(t * np.pi))
-            self.last_f = f
-            return f
-
-    def __call__(self, n, **kwargs):
-        return self.schedule(n, **kwargs)
-
-
-class LambdaLinearScheduler(LambdaWarmUpCosineScheduler):
-    """
-    Linear instead of cosine decay for the main part of the cycle.
-    """
-
-    def schedule(self, n, **kwargs):
-        cycle = self.find_in_interval(n)
-        n = n - self.cum_cycles[cycle]
-        if self.verbosity_interval > 0:
-            if n % self.verbosity_interval == 0:
-                log.info(f"current step: {n}, recent lr-multiplier: {self.last_f}, current cycle {cycle}")
-
-        if n < self.lr_warm_up_steps[cycle]:
-            f = (self.f_max[cycle] - self.f_start[cycle]) / self.lr_warm_up_steps[cycle] * n + self.f_start[cycle]
-            self.last_f = f
-            return f
-        else:
-            f = self.f_min[cycle] + (self.f_max[cycle] - self.f_min[cycle]) * (self.cycle_lengths[cycle] - n) / (
-                self.cycle_lengths[cycle] - self.lr_warm_up_steps[cycle]
-            )
-            self.last_f = f
-            return f

lyra_2/_ext/imaginaire/lazy_config/__init__.py

@@ -1,60 +0,0 @@
-# Copyright (c) Facebook, Inc. and its affiliates.
-import os
-
-from omegaconf import DictConfig, OmegaConf
-
-from lyra_2._ext.imaginaire.lazy_config.instantiate import instantiate
-from lyra_2._ext.imaginaire.lazy_config.lazy import LazyCall, LazyConfig
-from lyra_2._ext.imaginaire.lazy_config.omegaconf_patch import to_object
-
-OmegaConf.to_object = to_object
-
-PLACEHOLDER = None
-LazyDict = DictConfig
-
-__all__ = ["instantiate", "LazyCall", "LazyConfig", "PLACEHOLDER", "LazyDict"]
-
-
-DOC_BUILDING = os.getenv("_DOC_BUILDING", False)  # set in docs/conf.py
-
-
-def fixup_module_metadata(module_name, namespace, keys=None):
-    """
-    Fix the __qualname__ of module members to be their exported api name, so
-    when they are referenced in docs, sphinx can find them. Reference:
-    https://github.com/python-trio/trio/blob/6754c74eacfad9cc5c92d5c24727a2f3b620624e/trio/_util.py#L216-L241
-    """
-    if not DOC_BUILDING:
-        return
-    seen_ids = set()
-
-    def fix_one(qualname, name, obj):
-        # avoid infinite recursion (relevant when using
-        # typing.Generic, for example)
-        if id(obj) in seen_ids:
-            return
-        seen_ids.add(id(obj))
-
-        mod = getattr(obj, "__module__", None)
-        if mod is not None and (mod.startswith(module_name) or mod.startswith("fvcore.")):
-            obj.__module__ = module_name
-            # Modules, unlike everything else in Python, put fully-qualitied
-            # names into their __name__ attribute. We check for "." to avoid
-            # rewriting these.
-            if hasattr(obj, "__name__") and "." not in obj.__name__:
-                obj.__name__ = name
-                obj.__qualname__ = qualname
-            if isinstance(obj, type):
-                for attr_name, attr_value in obj.__dict__.items():
-                    fix_one(objname + "." + attr_name, attr_name, attr_value)
-
-    if keys is None:
-        keys = namespace.keys()
-    for objname in keys:
-        if not objname.startswith("_"):
-            obj = namespace[objname]
-            fix_one(objname, objname, obj)
-
-
-fixup_module_metadata(__name__, globals(), __all__)
-del fixup_module_metadata

lyra_2/_ext/imaginaire/lazy_config/file_io.py

@@ -1,24 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from iopath.common.file_io import HTTPURLHandler, OneDrivePathHandler, PathHandler
-from iopath.common.file_io import PathManager as PathManagerBase
-
-__all__ = ["PathManager", "PathHandler"]
-
-
-PathManager = PathManagerBase()
-PathManager.register_handler(HTTPURLHandler())
-PathManager.register_handler(OneDrivePathHandler())

lyra_2/_ext/imaginaire/lazy_config/instantiate.py

@@ -1,120 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import collections.abc as abc
-import dataclasses
-import logging
-from typing import Any
-
-import attrs
-
-from lyra_2._ext.imaginaire.lazy_config.registry import _convert_target_to_string, locate
-
-__all__ = ["dump_dataclass", "instantiate"]
-
-
-def is_dataclass_or_attrs(target):
-    return dataclasses.is_dataclass(target) or attrs.has(target)
-
-
-def dump_dataclass(obj: Any):
-    """
-    Dump a dataclass recursively into a dict that can be later instantiated.
-
-    Args:
-        obj: a dataclass object
-
-    Returns:
-        dict
-    """
-    assert dataclasses.is_dataclass(obj) and not isinstance(obj, type), (
-        "dump_dataclass() requires an instance of a dataclass."
-    )
-    ret = {"_target_": _convert_target_to_string(type(obj))}
-    for f in dataclasses.fields(obj):
-        v = getattr(obj, f.name)
-        if dataclasses.is_dataclass(v):
-            v = dump_dataclass(v)
-        if isinstance(v, (list, tuple)):
-            v = [dump_dataclass(x) if dataclasses.is_dataclass(x) else x for x in v]
-        ret[f.name] = v
-    return ret
-
-
-def instantiate(cfg, *args, **kwargs):
-    """
-    Recursively instantiate objects defined in dictionaries by
-    "_target_" and arguments.
-
-    Args:
-        cfg: a dict-like object with "_target_" that defines the caller, and
-            other keys that define the arguments
-        args: Optional positional parameters pass-through.
-        kwargs: Optional named parameters pass-through.
-
-    Returns:
-        object instantiated by cfg
-    """
-    from omegaconf import DictConfig, ListConfig, OmegaConf
-
-    if isinstance(cfg, ListConfig):
-        lst = [instantiate(x) for x in cfg]
-        return ListConfig(lst, flags={"allow_objects": True})
-    if isinstance(cfg, list):
-        # Specialize for list, because many classes take
-        # list[objects] as arguments, such as ResNet, DatasetMapper
-        return [instantiate(x) for x in cfg]
-
-    # If input is a DictConfig backed by dataclasses (i.e. omegaconf's structured config),
-    # instantiate it to the actual dataclass.
-    if isinstance(cfg, DictConfig) and is_dataclass_or_attrs(cfg._metadata.object_type):
-        return OmegaConf.to_object(cfg)
-
-    if isinstance(cfg, abc.Mapping) and "_target_" in cfg:
-        # conceptually equivalent to hydra.utils.instantiate(cfg) with _convert_=all,
-        # but faster: https://github.com/facebookresearch/hydra/issues/1200
-        is_recursive = getattr(cfg, "_recursive_", True)
-        if is_recursive:
-            cfg = {k: instantiate(v) for k, v in cfg.items()}
-        else:
-            cfg = {k: v for k, v in cfg.items()}
-        # pop the _recursive_ key to avoid passing it as a parameter
-        if "_recursive_" in cfg:
-            cfg.pop("_recursive_")
-        cls = cfg.pop("_target_")
-        cls = instantiate(cls)
-
-        if isinstance(cls, str):
-            cls_name = cls
-            cls = locate(cls_name)
-            assert cls is not None, cls_name
-        else:
-            try:
-                cls_name = cls.__module__ + "." + cls.__qualname__
-            except Exception:
-                # target could be anything, so the above could fail
-                cls_name = str(cls)
-        assert callable(cls), f"_target_ {cls} does not define a callable object"
-        try:
-            # override config with kwargs
-            instantiate_kwargs = {}
-            instantiate_kwargs.update(cfg)
-            instantiate_kwargs.update(kwargs)
-            return cls(*args, **instantiate_kwargs)
-        except TypeError:
-            logger = logging.getLogger(__name__)
-            logger.error(f"Error when instantiating {cls_name}!")
-            raise
-    return cfg  # return as-is if don't know what to do

lyra_2/_ext/imaginaire/lazy_config/lazy.py

@@ -1,430 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import ast
-import builtins
-import collections.abc as abc
-import importlib
-import inspect
-import logging
-import os
-import pickle
-import uuid
-from collections import OrderedDict
-from contextlib import contextmanager
-from copy import deepcopy
-from dataclasses import is_dataclass
-from typing import Any, Dict, List, Tuple, Union
-
-import attrs
-import yaml
-from omegaconf import DictConfig, ListConfig, OmegaConf
-
-try:
-    import dill as dill_pickle
-except ImportError:
-    dill_pickle = None
-
-try:
-    import cloudpickle
-except ImportError:
-    cloudpickle = None
-
-from lyra_2._ext.imaginaire.lazy_config.file_io import PathManager
-from lyra_2._ext.imaginaire.lazy_config.registry import _convert_target_to_string
-
-__all__ = ["LazyCall", "LazyConfig"]
-
-
-def sort_dict(d: Dict[str, Any]) -> OrderedDict[str, Any]:
-    return OrderedDict(sorted(d.items(), key=lambda x: x[0]))
-
-
-def dict_representer(dumper: yaml.Dumper, data: OrderedDict[str, Any]) -> yaml.nodes.MappingNode:
-    return dumper.represent_mapping("tag:yaml.org,2002:map", data.items())
-
-
-def sort_recursive(obj: Union[Dict[str, Any], List[Any], Any]) -> Union[OrderedDict[str, Any], List[Any], Any]:
-    if isinstance(obj, dict):
-        return sort_dict({k: sort_recursive(v) for k, v in obj.items()})
-    elif isinstance(obj, list):
-        return [sort_recursive(item) for item in obj]
-    return obj
-
-
-yaml.add_representer(OrderedDict, dict_representer)
-
-OmegaConf.register_new_resolver("add", lambda *vals: sum(vals))
-OmegaConf.register_new_resolver("subtract", lambda *vals: vals[0] - sum(vals[1:]))
-
-
-def get_default_params(cls_or_func):
-    if callable(cls_or_func):
-        # inspect signature for function
-        signature = inspect.signature(cls_or_func)
-    else:
-        # inspect signature for class
-        signature = inspect.signature(cls_or_func.__init__)
-    params = signature.parameters
-    default_params = {
-        name: param.default for name, param in params.items() if param.default is not inspect.Parameter.empty
-    }
-    return default_params
-
-
-class LazyCall:
-    """
-    Wrap a callable so that when it's called, the call will not be executed,
-    but returns a dict that describes the call.
-
-    LazyCall object has to be called with only keyword arguments. Positional
-    arguments are not yet supported.
-
-    Examples:
-    ::
-        from detectron2.config import instantiate, LazyCall
-
-        layer_cfg = LazyCall(nn.Conv2d)(in_channels=32, out_channels=32)
-        layer_cfg.out_channels = 64   # can edit it afterwards
-        layer = instantiate(layer_cfg)
-    """
-
-    def __init__(self, target):
-        if not (callable(target) or isinstance(target, (str, abc.Mapping))):
-            raise TypeError(f"target of LazyCall must be a callable or defines a callable! Got {target}")
-        self._target = target
-
-    def __call__(self, **kwargs):
-        if is_dataclass(self._target) or attrs.has(self._target):
-            # omegaconf object cannot hold dataclass type
-            # https://github.com/omry/omegaconf/issues/784
-            target = _convert_target_to_string(self._target)
-        else:
-            target = self._target
-        kwargs["_target_"] = target
-
-        _final_params = get_default_params(self._target)
-        _final_params.update(kwargs)
-
-        return DictConfig(content=_final_params, flags={"allow_objects": True})
-
-
-def _visit_dict_config(cfg, func):
-    """
-    Apply func recursively to all DictConfig in cfg.
-    """
-    if isinstance(cfg, DictConfig):
-        func(cfg)
-        for v in cfg.values():
-            _visit_dict_config(v, func)
-    elif isinstance(cfg, ListConfig):
-        for v in cfg:
-            _visit_dict_config(v, func)
-
-
-def _validate_py_syntax(filename):
-    # see also https://github.com/open-mmlab/mmcv/blob/master/mmcv/utils/config.py
-    with PathManager.open(filename, "r") as f:
-        content = f.read()
-    try:
-        ast.parse(content)
-    except SyntaxError as e:
-        raise SyntaxError(f"Config file {filename} has syntax error!") from e
-
-
-def _cast_to_config(obj):
-    # if given a dict, return DictConfig instead
-    if isinstance(obj, dict):
-        return DictConfig(obj, flags={"allow_objects": True})
-    return obj
-
-
-_CFG_PACKAGE_NAME = "detectron2._cfg_loader"
-"""
-A namespace to put all imported config into.
-"""
-
-
-def _random_package_name(filename):
-    # generate a random package name when loading config files
-    return _CFG_PACKAGE_NAME + str(uuid.uuid4())[:4] + "." + os.path.basename(filename)
-
-
-@contextmanager
-def _patch_import():
-    """
-    Enhance relative import statements in config files, so that they:
-    1. locate files purely based on relative location, regardless of packages.
-       e.g. you can import file without having __init__
-    2. do not cache modules globally; modifications of module states has no side effect
-    3. support other storage system through PathManager, so config files can be in the cloud
-    4. imported dict are turned into omegaconf.DictConfig automatically
-    """
-    old_import = builtins.__import__
-
-    def find_relative_file(original_file, relative_import_path, level):
-        # NOTE: "from . import x" is not handled. Because then it's unclear
-        # if such import should produce `x` as a python module or DictConfig.
-        # This can be discussed further if needed.
-        relative_import_err = """
-Relative import of directories is not allowed within config files.
-Within a config file, relative import can only import other config files.
-""".replace("\n", " ")
-        if not len(relative_import_path):
-            raise ImportError(relative_import_err)
-
-        cur_file = os.path.dirname(original_file)
-        for _ in range(level - 1):
-            cur_file = os.path.dirname(cur_file)
-        cur_name = relative_import_path.lstrip(".")
-        for part in cur_name.split("."):
-            cur_file = os.path.join(cur_file, part)
-        if not cur_file.endswith(".py"):
-            cur_file += ".py"
-        if not PathManager.isfile(cur_file):
-            cur_file_no_suffix = cur_file[: -len(".py")]
-            if PathManager.isdir(cur_file_no_suffix):
-                raise ImportError(f"Cannot import from {cur_file_no_suffix}." + relative_import_err)
-            else:
-                raise ImportError(
-                    f"Cannot import name {relative_import_path} from {original_file}: {cur_file} does not exist."
-                )
-        return cur_file
-
-    def new_import(name, globals=None, locals=None, fromlist=(), level=0):
-        if (
-            # Only deal with relative imports inside config files
-            level != 0 and globals is not None and (globals.get("__package__", "") or "").startswith(_CFG_PACKAGE_NAME)
-        ):
-            cur_file = find_relative_file(globals["__file__"], name, level)
-            _validate_py_syntax(cur_file)
-            spec = importlib.machinery.ModuleSpec(_random_package_name(cur_file), None, origin=cur_file)
-            module = importlib.util.module_from_spec(spec)
-            module.__file__ = cur_file
-            with PathManager.open(cur_file) as f:
-                content = f.read()
-            exec(compile(content, cur_file, "exec"), module.__dict__)
-            for name in fromlist:  # turn imported dict into DictConfig automatically
-                val = _cast_to_config(module.__dict__[name])
-                module.__dict__[name] = val
-            return module
-        return old_import(name, globals, locals, fromlist=fromlist, level=level)
-
-    builtins.__import__ = new_import
-    yield new_import
-    builtins.__import__ = old_import
-
-
-class LazyConfig:
-    """
-    Provide methods to save, load, and overrides an omegaconf config object
-    which may contain definition of lazily-constructed objects.
-    """
-
-    @staticmethod
-    def load_rel(filename: str, keys: Union[None, str, Tuple[str, ...]] = None):
-        """
-        Similar to :meth:`load()`, but load path relative to the caller's
-        source file.
-
-        This has the same functionality as a relative import, except that this method
-        accepts filename as a string, so more characters are allowed in the filename.
-        """
-        caller_frame = inspect.stack()[1]
-        caller_fname = caller_frame[0].f_code.co_filename
-        assert caller_fname != "<string>", "load_rel Unable to find caller"
-        caller_dir = os.path.dirname(caller_fname)
-        filename = os.path.join(caller_dir, filename)
-        return LazyConfig.load(filename, keys)
-
-    @staticmethod
-    def load(filename: str, keys: Union[None, str, Tuple[str, ...]] = None):
-        """
-        Load a config file.
-
-        Args:
-            filename: absolute path or relative path w.r.t. the current working directory
-            keys: keys to load and return. If not given, return all keys
-                (whose values are config objects) in a dict.
-        """
-        has_keys = keys is not None
-        filename = filename.replace("/./", "/")  # redundant
-        if os.path.splitext(filename)[1] not in [".py", ".yaml", ".yml"]:
-            raise ValueError(f"Config file {filename} has to be a python or yaml file.")
-        if filename.endswith(".py"):
-            _validate_py_syntax(filename)
-
-            with _patch_import():
-                # Record the filename
-                module_namespace = {
-                    "__file__": filename,
-                    "__package__": _random_package_name(filename),
-                }
-                with PathManager.open(filename) as f:
-                    content = f.read()
-                # Compile first with filename to:
-                # 1. make filename appears in stacktrace
-                # 2. make load_rel able to find its parent's (possibly remote) location
-                exec(compile(content, filename, "exec"), module_namespace)
-
-            ret = module_namespace
-        else:
-            with PathManager.open(filename) as f:
-                obj = yaml.unsafe_load(f)
-            ret = OmegaConf.create(obj, flags={"allow_objects": True})
-
-        if has_keys:
-            if isinstance(keys, str):
-                return _cast_to_config(ret[keys])
-            else:
-                return tuple(_cast_to_config(ret[a]) for a in keys)
-        else:
-            if filename.endswith(".py"):
-                # when not specified, only load those that are config objects
-                ret = DictConfig(
-                    {
-                        name: _cast_to_config(value)
-                        for name, value in ret.items()
-                        if isinstance(value, (DictConfig, ListConfig, dict)) and not name.startswith("_")
-                    },
-                    flags={"allow_objects": True},
-                )
-            return ret
-
-    @staticmethod
-    def save_pkl(cfg, filename: str) -> str:
-        """
-        Saves a Config object to a file using pickle serialization. This method is typically used
-        when the configuration object contains complex objects, such as lambdas, that are not supported by
-        simpler serialization methods like YAML. The function attempts to create a deep copy of the configuration
-        object before serialization to ensure that the original object remains unmodified.
-
-        Args:
-            cfg: A Config object to be serialized and saved.
-            filename: The path and name of the file where the configuration should be saved. The function
-                      assumes the file extension indicates a pickle format (e.g., .pkl).
-
-        Returns:
-            str: The filename to which the configuration was saved. This can be used to verify the file location
-                 or log the outcome.
-
-        Notes:
-            - The function logs a warning if the configuration is successfully saved using pickle.
-            - If saving fails, an error is logged with the exception details.
-        """
-        logger = logging.getLogger(__name__)
-        try:
-            cfg = deepcopy(cfg)
-        except Exception:
-            pass
-
-        try:
-            with PathManager.open(filename, "wb") as f:
-                pickle.dump(cfg, f)
-            logger.warning(f"Config is saved using pickle at {filename}.")
-        except Exception as e:
-            logger.error(f"Failed to save config to {filename}: {e}. Trying dill or cloudpickle instead")
-            if dill_pickle:
-                try:
-                    with PathManager.open(filename, "wb") as f:
-                        pickle.dump(dill_pickle.dumps(cfg, recurse=True), f)
-                        logger.warning(f"Config is saved using dill at {filename}.")
-                except Exception as e:
-                    logger.error(f"Failed to save config to {filename}: {e}.")
-                    if cloudpickle:
-                        try:
-                            with PathManager.open(filename, "wb") as f:
-                                pickle.dump(cloudpickle.dumps(cfg), f)
-                            logger.warning(f"Config is saved using cloudpickle at {filename}.")
-                        except Exception as e:
-                            logger.error(f"Failed to save config to {filename}: {e}.")
-                    else:
-                        logger.error("cloudpickle is not available. Cannot save the config.")
-                        raise e
-
-        return filename
-
-    @staticmethod
-    def save_yaml(cfg, filename: str) -> str:
-        """
-        Saves a Config object to a file using YAML serialization. This method is beneficial when the configuration object's content needs to be human-readable and easily editable. YAML is suitable for configurations that do not contain complex types like lambdas, which must be handled differently. The function converts unserializable items to strings before saving to ensure compatibility with YAML serialization.
-
-        Args:
-            cfg: A Config object to be serialized and saved. It handles both DictConfig and ListConfig types.
-            filename: The path and name of the file where the configuration should be saved. The function does not require a specific file extension but typically uses '.yaml'.
-
-        Returns:
-            str: The filename to which the configuration was saved. This can be used to verify the file location or log the outcome.
-
-        Notes:
-            - The function logs a warning if the configuration is successfully saved using YAML.
-            - If saving fails, an error is logged with the exception details.
-        """
-        logger = logging.getLogger(__name__)
-        try:
-            cfg = deepcopy(cfg)
-        except Exception:
-            pass
-
-        # Define a function to check if an item is serializable to YAML
-        def is_serializable(item):
-            try:
-                OmegaConf.to_yaml(item)
-                return True
-            except Exception as e:
-                return False
-
-        # Function to convert unserializable items to strings
-        def serialize_config(config):
-            if isinstance(config, DictConfig):
-                for key, value in config.items():
-                    if isinstance(value, (DictConfig, ListConfig)):
-                        try:
-                            if "_target_" in value:
-                                default_params = get_default_params(value["_target_"])
-                                for default_key, default_v in default_params.items():
-                                    if default_key not in value:
-                                        value[default_key] = default_v
-                        except Exception as e:
-                            logger.error(f"Failed to add default argument values: {e}")
-
-                        serialize_config(value)
-                    else:
-                        if not is_serializable(value) and value is not None:
-                            config[key] = str(value)
-            elif isinstance(config, ListConfig):
-                for i, item in enumerate(config):
-                    if isinstance(item, (DictConfig, ListConfig)):
-                        serialize_config(item)
-                    else:
-                        if not is_serializable(item) and item is not None:
-                            config[i] = str(item)
-            else:
-                raise NotImplementedError("Input config must be a DictConfig or ListConfig.")
-            return config
-
-        # Convert Config object to a DictConfig object.
-        config_dict = attrs.asdict(cfg)
-        config_omegaconf = DictConfig(content=config_dict, flags={"allow_objects": True})
-
-        # Serialize the DictConfig object by converting non-serializable objects to strings.
-        config_omegaconf = serialize_config(config_omegaconf)
-
-        config_dict: Dict[str, Any] = OmegaConf.to_container(config_omegaconf, resolve=True)
-        sorted_config: OrderedDict[str, Any] = sort_recursive(config_dict)
-        with open(filename, "w") as f:
-            yaml.dump(sorted_config, f, default_flow_style=False)
-        logger.warning(f"Config is saved using omegaconf at {filename}.")
-        return filename

lyra_2/_ext/imaginaire/lazy_config/omegaconf_patch.py

@@ -1,65 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from typing import Any, Dict, List, Union
-
-from omegaconf import OmegaConf
-from omegaconf.base import DictKeyType, SCMode
-from omegaconf.dictconfig import DictConfig  # pragma: no cover
-
-
-def to_object(cfg: Any) -> Union[Dict[DictKeyType, Any], List[Any], None, str, Any]:
-    """
-    Converts an OmegaConf configuration object to a native Python container (dict or list), unless
-    the configuration is specifically created by LazyCall, in which case the original configuration
-    is returned directly.
-
-    This function serves as a modification of the original `to_object` method from OmegaConf,
-    preventing DictConfig objects created by LazyCall from being automatically converted to Python
-    dictionaries. This ensures that configurations meant to be lazily evaluated retain their intended
-    structure and behavior.
-
-    Differences from OmegaConf's original `to_object`:
-    - Adds a check at the beginning to return the configuration unchanged if it is created by LazyCall.
-
-    Reference:
-    - Original OmegaConf `to_object` method: https://github.com/omry/omegaconf/blob/master/omegaconf/omegaconf.py#L595
-
-    Args:
-        cfg (Any): The OmegaConf configuration object to convert.
-
-    Returns:
-        Union[Dict[DictKeyType, Any], List[Any], None, str, Any]: The converted Python container if
-        `cfg` is not a LazyCall created configuration, otherwise the unchanged `cfg`.
-
-    Examples:
-        >>> cfg = DictConfig({"key": "value", "_target_": "Model"})
-        >>> to_object(cfg)
-        DictConfig({"key": "value", "_target_": "Model"})
-
-        >>> cfg = DictConfig({"list": [1, 2, 3]})
-        >>> to_object(cfg)
-        {'list': [1, 2, 3]}
-    """
-    if isinstance(cfg, DictConfig) and "_target_" in cfg.keys():
-        return cfg
-
-    return OmegaConf.to_container(
-        cfg=cfg,
-        resolve=True,
-        throw_on_missing=True,
-        enum_to_str=False,
-        structured_config_mode=SCMode.INSTANTIATE,
-    )

lyra_2/_ext/imaginaire/lazy_config/registry.py

@@ -1,74 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import pydoc
-from typing import Any
-
-from fvcore.common.registry import Registry  # for backward compatibility.
-
-"""
-``Registry`` and `locate` provide ways to map a string (typically found
-in config files) to callable objects.
-"""
-
-__all__ = ["Registry", "locate"]
-
-
-def _convert_target_to_string(t: Any) -> str:
-    """
-    Inverse of ``locate()``.
-
-    Args:
-        t: any object with ``__module__`` and ``__qualname__``
-    """
-    module, qualname = t.__module__, t.__qualname__
-
-    # Compress the path to this object, e.g. ``module.submodule._impl.class``
-    # may become ``module.submodule.class``, if the later also resolves to the same
-    # object. This simplifies the string, and also is less affected by moving the
-    # class implementation.
-    module_parts = module.split(".")
-    for k in range(1, len(module_parts)):
-        prefix = ".".join(module_parts[:k])
-        candidate = f"{prefix}.{qualname}"
-        try:
-            if locate(candidate) is t:
-                return candidate
-        except ImportError:
-            pass
-    return f"{module}.{qualname}"
-
-
-def locate(name: str) -> Any:
-    """
-    Locate and return an object ``x`` using an input string ``{x.__module__}.{x.__qualname__}``,
-    such as "module.submodule.class_name".
-
-    Raise Exception if it cannot be found.
-    """
-    obj = pydoc.locate(name)
-
-    # Some cases (e.g. torch.optim.sgd.SGD) not handled correctly
-    # by pydoc.locate. Try a private function from hydra.
-    if obj is None:
-        try:
-            # from hydra.utils import get_method - will print many errors
-            from hydra.utils import _locate
-        except ImportError as e:
-            raise ImportError(f"Cannot dynamically locate object {name}!") from e
-        else:
-            obj = _locate(name)  # it raises if fails
-
-    return obj

lyra_2/_ext/imaginaire/model.py

@@ -1,129 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from typing import Any
-
-import torch
-
-from lyra_2._ext.imaginaire.lazy_config import LazyDict, instantiate
-
-
-class ImaginaireModel(torch.nn.Module):
-    """The base model class of Imaginaire. It is inherited from torch.nn.Module.
-
-    All models in Imaginaire should inherit ImaginaireModel. It should include the implementions for all the
-    computation graphs. All inheriting child classes should implement the following methods:
-    - training_step(): The training step of the model, including the loss computation.
-    - validation_step(): The validation step of the model, including the loss computation.
-    - forward(): The computation graph for model inference.
-    The following methods have default implementations in ImaginaireModel:
-    - init_optimizer_scheduler(): Creates the optimizer and scheduler for the model.
-    """
-
-    def __init__(self) -> None:
-        super().__init__()
-
-    def init_optimizer_scheduler(
-        self, optimizer_config: LazyDict, scheduler_config: LazyDict
-    ) -> tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LRScheduler]:
-        """Creates the optimizer and scheduler for the model.
-
-        Args:
-            config_model (ModelConfig): The config object for the model.
-
-        Returns:
-            optimizer (torch.optim.Optimizer): The model optimizer.
-            scheduler (torch.optim.lr_scheduler.LRScheduler): The optimization scheduler.
-        """
-        optimizer_config.params = self.parameters()
-        optimizer = instantiate(optimizer_config)
-        scheduler_config.optimizer = optimizer
-        scheduler = instantiate(scheduler_config)
-        return optimizer, scheduler
-
-    def training_step(
-        self, data_batch: dict[str, torch.Tensor], iteration: int
-    ) -> tuple[dict[str, torch.Tensor], torch.Tensor]:
-        """The training step of the model, including the loss computation.
-
-        Args:
-            data (dict[str, torch.Tensor]): Data batch (dictionary of tensors).
-            iteration (int): Current iteration number.
-
-        Returns:
-            output_batch (dict[str, torch.Tensor]): Auxiliary model output from the training batch.
-            loss (torch.Tensor): The total loss for backprop (weighted sum of various losses).
-        """
-        raise NotImplementedError
-
-    @torch.no_grad()
-    def validation_step(
-        self, data_batch: dict[str, torch.Tensor], iteration: int
-    ) -> tuple[dict[str, torch.Tensor], torch.Tensor]:
-        """The validation step of the model, including the loss computation.
-
-        Args:
-            data (dict[str, torch.Tensor]): Data batch (dictionary of tensors).
-            iteration (int): Current iteration number.
-
-        Returns:
-            output_batch (dict[str, torch.Tensor]): Auxiliary model output from the validation batch.
-            loss (torch.Tensor): The total loss (weighted sum of various losses).
-        """
-        raise NotImplementedError
-
-    @torch.inference_mode()
-    def forward(self, *args: Any, **kwargs: Any) -> Any:
-        """The computation graph for model inference.
-
-        Args:
-            *args: Whatever you decide to pass into the forward method.
-            **kwargs: Keyword arguments are also possible.
-
-        Return:
-            Your model's output.
-        """
-        raise NotImplementedError
-
-    def on_train_start(self, memory_format: torch.memory_format = torch.preserve_format) -> None:
-        """The model preparation before the training is launched
-
-        Args:
-            memory_format (torch.memory_format): Memory format of the model.
-        """
-        pass
-
-    def on_before_zero_grad(
-        self, optimizer: torch.optim.Optimizer, scheduler: torch.optim.lr_scheduler.LRScheduler, iteration: int
-    ) -> None:
-        """Hook before zero_grad() is called.
-
-        Args:
-            optimizer (torch.optim.Optimizer): The model optimizer.
-            scheduler (torch.optim.lr_scheduler.LRScheduler): The optimization scheduler.
-            iteration (int): Current iteration number.
-        """
-        pass
-
-    def on_after_backward(self, iteration: int = 0) -> None:
-        """Hook after loss.backward() is called.
-
-        This method is called immediately after the backward pass, allowing for custom operations
-        or modifications to be performed on the gradients before the optimizer step.
-
-        Args:
-            iteration (int): Current iteration number.
-        """
-        pass

lyra_2/_ext/imaginaire/trainer.py

@@ -1,379 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import functools
-import inspect
-import os
-import signal
-
-import torch
-import torch.distributed as dist
-import torch.utils.data
-
-from lyra_2._ext.imaginaire.utils.context_managers import distributed_init
-from lyra_2._ext.imaginaire.utils.profiling import maybe_enable_memory_snapshot, maybe_enable_profiling
-
-try:
-    from megatron.core import parallel_state
-
-    USE_MEGATRON = True
-except ImportError:
-    USE_MEGATRON = False
-    print("Megatron-core is not installed.")
-
-
-from lyra_2._ext.imaginaire.lazy_config import LazyConfig, instantiate
-from lyra_2._ext.imaginaire.model import ImaginaireModel
-from lyra_2._ext.imaginaire.utils import callback, distributed, ema, log, misc
-from lyra_2._ext.imaginaire.utils.checkpointer import Checkpointer
-from lyra_2._ext.imaginaire.utils.misc import StragglerDetectorV2
-
-try:
-    from lyra_2._src.callbacks.smart_stop import TimeoutException
-except ImportError:
-    # Define a dummy exception if smart_stop is not available
-    class TimeoutException(Exception):
-        pass
-
-
-class ImaginaireTrainer:
-    """The base trainer class of Imaginaire.
-
-    All trainers in Imaginaire should inherit ImaginaireTrainer. It contains the basic functionality for model training
-    (particularly suited for large-scale training), including data parallel (DDP/FSDP), model weight average (EMA),
-    mixed-precision training (fp16/bf16).
-
-    Attributes:
-        checkpointer (Checkpointer): checkpointer object to save/load model weights and optimizer states.
-        training_timer (misc.Timer): Timer object to time code blocks and functions.
-    """
-
-    def __init__(self, config):
-        """Constructor of the trainer.
-
-        Args:
-            config (Config): The config object for the Imaginaire codebase.
-        """
-        super().__init__()
-        self.config = config
-        # Set up the distributed computing environment.
-        with distributed_init():
-            distributed.init()
-            # Set up parallel states.
-            if hasattr(config.model, "context_parallel_size"):
-                if config.model_parallel.context_parallel_size > 1:
-                    raise ValueError(
-                        "Both config.model.context_parallel_size and config.model_parallel.context_parallel_size are set. "
-                        "config.model.context_parallel_size is deprecated. Please only set config.model_parallel.context_parallel_size."
-                    )
-                else:
-                    log.critical(
-                        "Using deprecated config.model.context_parallel_size. Please use config.model_parallel.context_parallel_size instead."
-                    )
-                    config.model_parallel.context_parallel_size = config.model.context_parallel_size
-            if USE_MEGATRON:
-                if (
-                    "create_gloo_process_groups"
-                    in inspect.signature(parallel_state.initialize_model_parallel).parameters
-                ):
-                    parallel_state.initialize_model_parallel(
-                        pipeline_model_parallel_size=config.model_parallel.pipeline_model_parallel_size,
-                        tensor_model_parallel_size=config.model_parallel.tensor_model_parallel_size,
-                        context_parallel_size=config.model_parallel.context_parallel_size,
-                        create_gloo_process_groups=False,
-                    )
-                else:
-                    parallel_state.initialize_model_parallel(
-                        pipeline_model_parallel_size=config.model_parallel.pipeline_model_parallel_size,
-                        tensor_model_parallel_size=config.model_parallel.tensor_model_parallel_size,
-                        context_parallel_size=config.model_parallel.context_parallel_size,
-                    )
-                # `config.model_parallel.sequence_parallel` is a bool that indicates whether to use sequence parallelism.
-                # It is not part of the original `parallel_state` API, so we need to set it manually.
-                parallel_state.sequence_parallel = config.model_parallel.sequence_parallel
-                if parallel_state.sequence_parallel:
-                    os.environ["CUDA_DEVICE_MAX_CONNECTIONS"] = "1"
-
-        # Create the local job directory, save the config file, and pipe to a local log.
-        if distributed.is_rank0():
-            os.makedirs(config.job.path_local, exist_ok=True)
-            # Save the config as .pkl for reproducibility.
-            LazyConfig.save_pkl(config, f"{config.job.path_local}/config.pkl")
-            # Save the config as .yaml for reading or parsing experiment hyperparameters.
-            LazyConfig.save_yaml(config, f"{config.job.path_local}/config.yaml")
-        dist.barrier()
-        log.init_loguru_file(f"{config.job.path_local}/stdout.log")
-        if distributed.is_rank0():
-            # Print important environment variables and the effective config.
-            log.info("Config:\n" + config.pretty_print(use_color=True))
-        misc.print_environ_variables(["HF_HOME", "IMAGINAIRE_OUTPUT_ROOT"])
-        # Set the random seed. If multi-GPU, different ranks are set with different seeds.
-        misc.set_random_seed(seed=config.trainer.seed, by_rank=True)
-        # Initialize cuDNN.
-        torch.backends.cudnn.deterministic = config.trainer.cudnn.deterministic
-        torch.backends.cudnn.benchmark = config.trainer.cudnn.benchmark
-        # Floating-point precision settings.
-        torch.backends.cudnn.allow_tf32 = torch.backends.cuda.matmul.allow_tf32 = True
-        # Initialize the callback functions.
-        self.callbacks = callback.CallBackGroup(config=config, trainer=self)
-        # Initialize the model checkpointer.
-        if config.checkpoint.type is None:
-            self.checkpointer = Checkpointer(config.checkpoint, config.job, callbacks=self.callbacks)
-        else:
-            self.checkpointer: Checkpointer = instantiate(
-                config.checkpoint.type, config.checkpoint, config.job, callbacks=self.callbacks
-            )
-        # Initialize the timer for speed benchmarking.
-        self.training_timer = misc.TrainingTimer()
-        # Initialize Straggler Detection
-        self.straggler_detector = StragglerDetectorV2(
-            enabled=self.config.trainer.straggler_detection.enabled,
-            report_freq=self.config.trainer.straggler_detection.report_freq,
-            profile_freq=self.config.trainer.straggler_detection.profile_freq,
-            max_diff=self.config.trainer.straggler_detection.max_diff,
-            raise_error=self.config.trainer.straggler_detection.raise_error,
-        )
-        self.straggler_detector.initialize()
-        # Send a TimeoutError if a training step takes over timeout_period seconds.
-        signal.signal(signal.SIGALRM, functools.partial(misc.timeout_handler, config.trainer.timeout_period))  # type: ignore
-
-    def train(
-        self,
-        model: ImaginaireModel,
-        dataloader_train: torch.utils.data.DataLoader,
-        dataloader_val: torch.utils.data.DataLoader,
-    ) -> None:
-        """The training function.
-
-        Args:
-            model (ImaginaireModel): The PyTorch model.
-            dataloader_train (torch.utils.data.DataLoader): The training data loader.
-            dataloader_val (torch.utils.data.DataLoader): The validation data loader.
-        """
-        # Leaving this for backward compability for now, but we can think about moving this to model.on_train_start for all models.
-        model = model.to("cuda", memory_format=self.config.trainer.memory_format)  # type: ignore
-        model.on_train_start(self.config.trainer.memory_format)
-
-        # Initialize the optimizer, scheduler, and grad_scaler.
-        self.callbacks.on_optimizer_init_start()
-        optimizer, scheduler = model.init_optimizer_scheduler(self.config.optimizer, self.config.scheduler)
-        grad_scaler = torch.amp.GradScaler("cuda", **self.config.trainer.grad_scaler_args)
-        self.callbacks.on_optimizer_init_end()
-        # Load the model checkpoint and get the starting iteration number.
-        iteration = self.checkpointer.load(model, optimizer, scheduler, grad_scaler)
-        grad_accum_iter = 0
-        log.info(f"Distributed parallelism mode: {self.config.trainer.distributed_parallelism}")
-        if self.config.trainer.distributed_parallelism == "ddp":
-            # Create a DDP model wrapper.
-            model_ddp = distributed.parallel_model_wrapper(self.config.trainer.ddp, model)
-        elif self.config.trainer.distributed_parallelism == "fsdp":
-            model_ddp = model
-        else:
-            raise ValueError(f"Unknown distributed parallelism mode: {self.config.trainer.distributed_parallelism}")
-
-        log.info("Starting training...")
-        self.callbacks.on_train_start(model, iteration=iteration)
-        # Initial validation.
-        if self.config.trainer.run_validation and iteration == 0:
-            self.validate(model, dataloader_val, iteration=iteration)
-        _end_training = False
-        _smart_stop_triggered = False
-        _oom_triggered = False
-        with (
-            maybe_enable_profiling(self.config, global_step=iteration) as torch_profiler,
-            maybe_enable_memory_snapshot(self.config, global_step=iteration) as memory_profiler,
-        ):
-            while True:
-                dataloader_train_iter = iter(dataloader_train)
-                while True:
-                    self.callbacks.on_before_dataloading(iteration)
-                    try:
-                        with (
-                            self.training_timer("dataloader_train"),
-                            self.straggler_detector.profile_section(
-                                "dataloading",
-                                self.config.trainer.straggler_detection.analyze_dataloading,
-                                profile_cuda=False,
-                            ),
-                        ):
-                            data_batch = next(dataloader_train_iter)
-                    except StopIteration:
-                        break
-                    finally:
-                        self.callbacks.on_after_dataloading(iteration)
-                    # If max_iter is reached, exit the training loop.
-                    if iteration >= self.config.trainer.max_iter:
-                        _end_training = True
-                        break
-                    # Move all tensors in the data batch to GPU device.
-                    data_batch = misc.to(data_batch, device="cuda")
-                    # The actual training step.
-                    self.callbacks.on_training_step_start(model, data_batch, iteration=iteration)
-                    self.callbacks.on_training_step_batch_start(model, data_batch, iteration=iteration)
-                    if not model.training:
-                        model_ddp.train()
-                    assert model_ddp.training, "model_ddp is not in training mode."
-                    assert model.training, "model is not in training mode."
-                    try:
-                        output_batch, loss, grad_accum_iter = self.training_step(
-                            model_ddp,
-                            optimizer,
-                            scheduler,
-                            grad_scaler,
-                            data_batch,
-                            iteration=iteration,
-                            grad_accum_iter=grad_accum_iter,
-                        )
-                    except torch.OutOfMemoryError as e:
-                        # CUDA OOM error - save checkpoint and exit gracefully
-                        _oom_triggered = True
-                        log.error(f"CUDA Out of Memory error caught: {e}")
-                        log.info("Saving checkpoint due to CUDA OOM error...")
-                        self.checkpointer.save(model, optimizer, scheduler, grad_scaler, iteration=iteration)
-                        log.success("Checkpoint saved successfully. Exiting training gracefully due to OOM.")
-                        _end_training = True
-                        break
-                    self.callbacks.on_training_step_batch_end(
-                        model, data_batch, output_batch, loss, iteration=iteration
-                    )
-                    # If the gradients are still being accumulated, continue to load the next training batch.
-                    if grad_accum_iter != 0:
-                        continue
-                    # Do the following when an actual optimizer (update) step has been made.
-                    iteration += 1
-                    # Save checkpoint.
-                    if iteration % self.config.checkpoint.save_iter == 0:
-                        self.checkpointer.save(model, optimizer, scheduler, grad_scaler, iteration=iteration)
-                    # Call on_training_step_end with SmartStop exception handling
-                    try:
-                        self.callbacks.on_training_step_end(model, data_batch, output_batch, loss, iteration=iteration)
-                    except TimeoutException as e:
-                        # Smart stop triggered - save checkpoint and exit gracefully
-                        _smart_stop_triggered = True
-                        log.warning(f"SmartStop exception caught: {e}")
-                        log.info("Saving checkpoint due to time limit...")
-                        self.checkpointer.save(model, optimizer, scheduler, grad_scaler, iteration=iteration)
-                        log.success("Checkpoint saved successfully. Exiting training gracefully.")
-                        _end_training = True
-                        break
-                    # Validation.
-                    if self.config.trainer.run_validation and iteration % self.config.trainer.validation_iter == 0:
-                        self.validate(model, dataloader_val, iteration=iteration)
-                    # This iteration is successful; reset the timeout signal.
-                    signal.alarm(self.config.trainer.timeout_period)
-                    self.straggler_detector.generate_report(iteration)
-                    if torch_profiler:
-                        torch_profiler.step()
-                    if memory_profiler:
-                        memory_profiler.step()
-                if _end_training:
-                    break
-        log.success("Done with training.")
-        if iteration % self.config.checkpoint.save_iter != 0 and not _smart_stop_triggered and not _oom_triggered:
-            self.checkpointer.save(model, optimizer, scheduler, grad_scaler, iteration=iteration)
-        self.callbacks.on_train_end(model, iteration=iteration)
-        self.checkpointer.finalize()
-        distributed.barrier()
-        self.callbacks.on_app_end()
-
-    def training_step(
-        self,
-        model_ddp: torch.nn.Module | distributed.DistributedDataParallel,
-        optimizer: torch.optim.Optimizer,
-        scheduler: torch.optim.lr_scheduler.LRScheduler,
-        grad_scaler: torch.amp.GradScaler,
-        data: dict[str, torch.Tensor],
-        iteration: int = 0,
-        grad_accum_iter: int = 0,
-    ) -> tuple[dict[str, torch.Tensor], torch.Tensor, int]:
-        """The training step.
-
-        Args:
-            model_ddp (torch.nn.Module | distributed.DistributedDataParallel): The model with a DDP wrapper or, the bare
-              module, depending on whether distributed training is enabled or not.
-            optimizer (torch.optim.Optimizer): The model optimizer.
-            scheduler (torch.optim.lr_scheduler.LRScheduler): The optimization scheduler.
-            grad_scaler (torch.amp.GradScaler): The gradient scaler (for mixed precision training).
-            data (dict[str, torch.Tensor]): Data batch (dictionary of tensors).
-            iteration (int): Current iteration number.
-            grad_accum_iter (int): Number of gradient accumulation iterations.
-
-        Returns:
-            output (dict[str, torch.Tensor]): The model output from the training data batch (dictionary of tensors).
-            loss (torch.Tensor): The total loss of the training data batch.
-        """
-        # Only let DDP sync gradient at the last iteration of the gradient accumulation window
-        with distributed.ddp_sync_grad(model_ddp, grad_accum_iter == self.config.trainer.grad_accum_iter - 1):
-            self.callbacks.on_before_forward(iteration=iteration)
-            with self.training_timer("forward"):
-                with self.straggler_detector.profile_section(
-                    "fwd", self.config.trainer.straggler_detection.analyze_forward
-                ):
-                    output_batch, loss = model_ddp.training_step(data, iteration)
-            self.callbacks.on_after_forward(iteration=iteration)
-            self.callbacks.on_before_backward(model_ddp, loss, iteration=iteration)
-            with self.training_timer("backward"):
-                with self.straggler_detector.profile_section(
-                    "bwd", self.config.trainer.straggler_detection.analyze_backward
-                ):
-                    loss_scaled = grad_scaler.scale(loss / self.config.trainer.grad_accum_iter)
-                    loss_scaled.backward()
-                    if self.config.trainer.distributed_parallelism == "ddp":
-                        model_ddp.module.on_after_backward()
-                    else:
-                        model_ddp.on_after_backward()
-            self.callbacks.on_after_backward(model_ddp, iteration=iteration)
-        grad_accum_iter += 1
-        if grad_accum_iter == self.config.trainer.grad_accum_iter:
-            with self.training_timer("optimizer_step"):
-                with self.straggler_detector.profile_section(
-                    "opt", self.config.trainer.straggler_detection.analyze_optimizer
-                ):
-                    self.callbacks.on_before_optimizer_step(
-                        model_ddp, optimizer, scheduler, grad_scaler, iteration=iteration
-                    )
-                    grad_scaler.step(optimizer)
-                    grad_scaler.update()
-                    scheduler.step()
-                    self.callbacks.on_before_zero_grad(model_ddp, optimizer, scheduler, iteration=iteration)
-                    if self.config.trainer.distributed_parallelism == "ddp":
-                        model_ddp.module.on_before_zero_grad(optimizer, scheduler, iteration=iteration)
-                    else:
-                        model_ddp.on_before_zero_grad(optimizer, scheduler, iteration=iteration)
-                    optimizer.zero_grad(set_to_none=True)
-            grad_accum_iter = 0
-        return output_batch, loss, grad_accum_iter
-
-    @torch.no_grad()
-    def validate(self, model: ImaginaireModel, dataloader_val: torch.utils.data.DataLoader, iteration: int = 0) -> None:
-        """Validate on the full validation dataset.
-
-        Args:
-            model (ImaginaireModel): The PyTorch model.
-            dataloader_val (torch.utils.data.DataLoader): The validation data loader.
-            iteration (int): Current iteration number.
-        """
-        self.callbacks.on_validation_start(model, dataloader_val, iteration=iteration)
-        model.eval()
-        # Evaluate on the full validation set.
-        with ema.ema_scope(model, enabled=model.config.ema.enabled):
-            for val_iter, data_batch in enumerate(dataloader_val):
-                if self.config.trainer.max_val_iter is not None and val_iter >= self.config.trainer.max_val_iter:
-                    break
-                data_batch = misc.to(data_batch, device="cuda")
-                self.callbacks.on_validation_step_start(model, data_batch, iteration=iteration)
-                output_batch, loss = model.validation_step(data_batch, iteration)
-                self.callbacks.on_validation_step_end(model, data_batch, output_batch, loss, iteration=iteration)
-        self.callbacks.on_validation_end(model, iteration=iteration)

lyra_2/_ext/imaginaire/types/denoise_prediction.py

@@ -1,28 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from __future__ import annotations
-
-from dataclasses import dataclass
-from typing import Optional
-
-import torch
-
-
-@dataclass
-class DenoisePrediction:
-    x0: torch.Tensor  # clean data prediction
-    eps: Optional[torch.Tensor] = None  # noise prediction
-    logvar: Optional[torch.Tensor] = None  # log variance of noise prediction, can be used a confidence / uncertainty

lyra_2/_ext/imaginaire/utils/__init__.py

@@ -1,14 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.

lyra_2/_ext/imaginaire/utils/callback.py

@@ -1,524 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from __future__ import annotations
-
-import sys
-import time
-import warnings
-from typing import TYPE_CHECKING, Any, Callable, Optional
-
-import omegaconf
-import torch
-import torch.distributed as dist
-import torch.utils.data
-import tqdm
-from lyra_2._ext.imaginaire.lazy_config import instantiate
-from lyra_2._ext.imaginaire.utils import distributed, log, misc
-from lyra_2._ext.imaginaire.utils.misc import get_local_tensor_if_DTensor
-
-try:
-    from megatron.core import parallel_state
-except ImportError:
-    parallel_state = None
-    print("Megatron-core is not installed.")
-
-
-if TYPE_CHECKING:
-    from lyra_2._ext.imaginaire.config import Config
-    from lyra_2._ext.imaginaire.model import ImaginaireModel
-    from lyra_2._ext.imaginaire.trainer import ImaginaireTrainer
-
-
-class CallBackGroup:
-    """A class for hosting a collection of callback objects.
-
-    It is used to execute callback functions of multiple callback objects with the same method name.
-    When callbackgroup.func(args) is executed, internally it loops through the objects in self._callbacks and runs
-    self._callbacks[0].func(args), self._callbacks[1].func(args), etc. The method name and arguments should match.
-
-    Attributes:
-        _callbacks (list[Callback]): List of callback objects.
-    """
-
-    def __init__(self, config: Config, trainer: ImaginaireTrainer) -> None:
-        """Initializes the list of callback objects.
-
-        Args:
-            config (Config): The config object for the Imaginaire codebase.
-            trainer (ImaginaireTrainer): The main trainer.
-        """
-        self._callbacks = []
-        callback_configs = config.trainer.callbacks
-        if callback_configs:
-            if isinstance(callback_configs, list) or isinstance(callback_configs, omegaconf.listconfig.ListConfig):
-                warnings.warn(
-                    "The 'config.trainer.callbacks' parameter should be a dict instead of a list. "
-                    "Please update your code",
-                    DeprecationWarning,
-                    stacklevel=2,
-                )
-                callback_configs = {f"callback_{i}": v for i, v in enumerate(callback_configs)}
-            for callback_name, current_callback_cfg in callback_configs.items():
-                if "_target_" not in current_callback_cfg:
-                    log.critical(
-                        f"Callback {callback_name} is missing the '_target_' field. \n SKip {current_callback_cfg}"
-                    )
-                    continue
-                log.info(f"Instantiating callback {callback_name}: {current_callback_cfg}")
-                _callback = instantiate(current_callback_cfg)
-                assert isinstance(_callback, Callback), f"{current_callback_cfg} is not a valid callback."
-                _callback.config = config
-                _callback.trainer = trainer
-                self._callbacks.append(_callback)
-
-    def __getattr__(self, method_name: str) -> Callable:
-        """Loops through the callback objects to call the corresponding callback function.
-
-        Args:
-            method_name (str): Callback method name.
-        """
-
-        def multi_callback_wrapper(*args, **kwargs) -> None:
-            for callback in self._callbacks:
-                assert hasattr(callback, method_name)
-                method = getattr(callback, method_name)
-                assert callable(method)
-                _ = method(*args, **kwargs)
-
-        return multi_callback_wrapper
-
-
-class Callback:
-    """The base class for all callbacks.
-
-    All callbacks should inherit from this class and adhere to the established method names and signatures.
-    """
-
-    def __init__(self, config: Optional["Config"] = None, trainer: Optional["ImaginaireTrainer"] = None):
-        """Initializes a Callback object.
-
-        Args:
-            config (Optional[Config]): The configuration object for the Imaginaire codebase, if available.
-            trainer (Optional[ImaginaireTrainer]): The main trainer handling the training loop, if available.
-
-        Notes:
-            The config and trainer parameters are optional to maintain backward compatibility.
-            In future releases, these parameters will be removed. Upon using these parameters, a deprecation
-            warning will be issued.
-
-        """
-        if config is not None or trainer is not None:
-            warnings.warn(
-                "The 'config' and 'trainer' parameters are deprecated and will be removed in a future release. "
-                "Please update your code to create Callback instances without these parameters.",
-                DeprecationWarning,
-                stacklevel=2,
-            )
-        del config, trainer
-
-    def on_train_start(self, model: ImaginaireModel, iteration: int = 0) -> None:
-        pass
-
-    def on_training_step_start(self, model: ImaginaireModel, data: dict[str, torch.Tensor], iteration: int = 0) -> None:
-        """
-        Called before the training step, for each batch. This is paired with on_training_step_end() but note that
-        when using gradient accumulation, while on_training_step_end() is only called when the optimizer is updated,
-        this function is called for every batch.
-        Use on_training_step_batch_start and on_training_step_batch_end if you need callbacks that are called
-        for every batch, albeit with the same iteration number.
-        FIXME - should this either be deprecated, or called only when a new training step is started after having updated
-        the optimizer?
-        """
-        pass
-
-    def on_training_step_batch_start(
-        self, model: ImaginaireModel, data: dict[str, torch.Tensor], iteration: int = 0
-    ) -> None:
-        """
-        Called before the training step, for each batch, similarly to on_training_step_start(). This function is paired with
-        on_training_step_batch_end(), and both functions are called for every batch even when using gradient accumulation.
-        Note that the iteration is only updated when the optimizer is updated, and therefore it may be the same for multiple invocations.
-        """
-        pass
-
-    def on_before_forward(self, iteration: int = 0) -> None:
-        pass
-
-    def on_after_forward(self, iteration: int = 0) -> None:
-        pass
-
-    def on_before_backward(
-        self, model_ddp: distributed.DistributedDataParallel, loss: torch.Tensor, iteration: int = 0
-    ) -> None:
-        pass
-
-    def on_after_backward(self, model_ddp: distributed.DistributedDataParallel, iteration: int = 0) -> None:
-        pass
-
-    def on_before_dataloading(self, iteration: int = 0) -> None:
-        pass
-
-    def on_after_dataloading(self, iteration: int = 0) -> None:
-        pass
-
-    def on_optimizer_init_start(self) -> None:
-        pass
-
-    def on_optimizer_init_end(self) -> None:
-        pass
-
-    def on_before_optimizer_step(
-        self,
-        model_ddp: distributed.DistributedDataParallel,
-        optimizer: torch.optim.Optimizer,
-        scheduler: torch.optim.lr_scheduler.LRScheduler,
-        grad_scaler: torch.amp.GradScaler,
-        iteration: int = 0,
-    ) -> None:
-        pass
-
-    def on_before_zero_grad(
-        self,
-        model_ddp: distributed.DistributedDataParallel,
-        optimizer: torch.optim.Optimizer,
-        scheduler: torch.optim.lr_scheduler.LRScheduler,
-        iteration: int = 0,
-    ) -> None:
-        pass
-
-    def on_training_step_batch_end(
-        self,
-        model: ImaginaireModel,
-        data_batch: dict[str, torch.Tensor],
-        output_batch: dict[str, torch.Tensor],
-        loss: torch.Tensor,
-        iteration: int = 0,
-    ) -> None:
-        """
-        Called at the end of a training step for every batch even when using gradient accumulation.
-        This is paired with on_training_step_batch_start(). Note that the iteration is only updated when the optimizer is updated,
-        and therefore it may be the same for multiple batches.
-        """
-        pass
-
-    def on_training_step_end(
-        self,
-        model: ImaginaireModel,
-        data_batch: dict[str, torch.Tensor],
-        output_batch: dict[str, torch.Tensor],
-        loss: torch.Tensor,
-        iteration: int = 0,
-    ) -> None:
-        """
-        Called at the end of a training step, but note that when using gradient accumulation, this is only called
-        when the optimizer is updated, and the iteration incremented, whereas on_training_step_start is called every time.
-        Use on_training_step_batch_start and on_training_step_batch_end if you need callbacks that are called
-        for every batch.
-        """
-        pass
-
-    def on_validation_start(
-        self, model: ImaginaireModel, dataloader_val: torch.utils.data.DataLoader, iteration: int = 0
-    ) -> None:
-        pass
-
-    def on_validation_step_start(
-        self, model: ImaginaireModel, data: dict[str, torch.Tensor], iteration: int = 0
-    ) -> None:
-        pass
-
-    def on_validation_step_end(
-        self,
-        model: ImaginaireModel,
-        data_batch: dict[str, torch.Tensor],
-        output_batch: dict[str, torch.Tensor],
-        loss: torch.Tensor,
-        iteration: int = 0,
-    ) -> None:
-        pass
-
-    def on_validation_end(self, model: ImaginaireModel, iteration: int = 0) -> None:
-        pass
-
-    def on_load_checkpoint_start(self, model: ImaginaireModel) -> None:
-        pass
-
-    def on_load_checkpoint_end(
-        self, model: ImaginaireModel, iteration: int = 0, checkpoint_path: Optional[str] = None
-    ) -> None:
-        pass
-
-    def on_load_checkpoint(self, model: ImaginaireModel, state_dict: dict[Any]) -> None:
-        """
-        Called when checkpoint loading is about to start, but after on_save_checkpoint_start().
-        FIXME - why do we need this callback, can't we just use on_save_checkpoint_start()?
-        """
-        pass
-
-    def on_save_checkpoint_start(self, model: ImaginaireModel, iteration: int = 0) -> None:
-        """
-        Called when checkpoint saving is about to start.
-        """
-        pass
-
-    def on_save_checkpoint_end(self, model: ImaginaireModel, iteration: int = 0) -> None:
-        """
-        Called when the synchronous part of checkpointing is finished, this function can be used
-        along with on_save_checkpoint_start() to measure the exposed (synchronous) checkpoint time.
-        Note that for asynchronous checkpoint, the checkpoint may still be ongoing, so this function
-        does not mean the checkpoint is finished for the asynchronous case, use on_save_checkpoint_success()
-        for that.
-        """
-        pass
-
-    def on_save_checkpoint_success(self, iteration: int = 0, elapsed_time: float = 0) -> None:
-        """
-        Called when checkpoint saving is fully finished, and succeeded. Not called if checkpoint failed.
-        For synchronous checkpoint, it is called at the same time as on_save_checkpoint_end(), but for asynchronous
-        checkpoint, it is called after the asynchronous part has also finished. For checkpointers with out-of-process
-        checkpointing, this function is called as soon as the notification is received from the checkpointer process,
-        which may not be immediately after the checkpoint has completed but later on. Therefore, if you need to measure
-        the full checkpoint duration for the asynchronous part, use the elapsed_time parameter, do not measure it directly
-        as this would be a significant overestimate.
-        """
-        pass
-
-    def on_save_checkpoint(self, model: ImaginaireModel, state_dict: dict[Any]) -> None:
-        pass
-
-    def on_train_end(self, model: ImaginaireModel, iteration: int = 0) -> None:
-        pass
-
-    def on_app_end(self) -> None:
-        pass
-
-
-class EMAModelCallback(Callback):
-    """The callback class for tracking EMA model weights."""
-
-    def on_train_start(self, model: ImaginaireModel, iteration: int = 0) -> None:
-        # Set up the EMA model weight tracker.
-        if model.config.ema.enabled:
-            assert hasattr(model, "ema"), "EMA should be initialized from ImaginaireModel"
-            # EMA model must be kept in FP32 precision.
-            model.ema = model.ema.to(dtype=torch.float32)
-        else:
-            assert not hasattr(model, "ema"), "There should be no EMA initialized."
-
-    def on_training_step_end(
-        self,
-        model: ImaginaireModel,
-        data_batch: dict[str, torch.Tensor],
-        output_batch: dict[str, torch.Tensor],
-        loss: torch.Tensor,
-        iteration: int = 0,
-    ) -> None:
-        # Update the EMA model with the new regular weights.
-        if model.config.ema.enabled:
-            model.ema.update_average(model, iteration)
-
-
-class ProgressBarCallback(Callback):
-    """The callback class for visualizing the training/validation progress bar in the console."""
-
-    @distributed.rank0_only
-    def on_train_start(self, model: ImaginaireModel, iteration: int = 0) -> None:
-        self.train_pbar = tqdm.trange(self.config.trainer.max_iter, initial=iteration, desc="Training")
-
-    @distributed.rank0_only
-    def on_training_step_end(
-        self,
-        model: ImaginaireModel,
-        data_batch: dict[str, torch.Tensor],
-        output_batch: dict[str, torch.Tensor],
-        loss: torch.Tensor,
-        iteration: int = 0,
-    ) -> None:
-        self.train_pbar.update()
-
-    @distributed.rank0_only
-    def on_validation_start(
-        self, model: ImaginaireModel, dataloader_val: torch.utils.data.DataLoader, iteration: int = 0
-    ) -> None:
-        if self.config.trainer.max_val_iter is not None:
-            num_iter = self.config.trainer.max_val_iter
-        else:
-            num_iter = len(dataloader_val)
-        assert num_iter is not None and num_iter > 0, f"Invalid number of validation iterations: {num_iter}"
-        self.val_pbar = tqdm.trange(num_iter, desc="Validating", position=1, leave=False)
-
-    @distributed.rank0_only
-    def on_validation_step_end(
-        self,
-        model: ImaginaireModel,
-        data_batch: dict[str, torch.Tensor],
-        output_batch: dict[str, torch.Tensor],
-        loss: torch.Tensor,
-        iteration: int = 0,
-    ) -> None:
-        self.val_pbar.update()
-
-    @distributed.rank0_only
-    def on_validation_end(self, model: ImaginaireModel, iteration: int = 0) -> None:
-        self.val_pbar.close()
-
-    @distributed.rank0_only
-    def on_train_end(self, model: ImaginaireModel, iteration: int = 0) -> None:
-        self.trainer.checkpointer.finalize()
-        self.train_pbar.close()
-
-
-class IterationLoggerCallback(Callback):
-    """The callback class for visualizing the training/validation progress bar in the console."""
-
-    @distributed.rank0_only
-    def on_train_start(self, model: ImaginaireModel, iteration: int = 0) -> None:
-        # self.train_pbar = tqdm.trange(self.config.trainer.max_iter, initial=iteration, desc="Training")
-        self.start_iteration_time = time.time()
-        self.elapsed_iteration_time = 0
-
-    @distributed.rank0_only
-    def on_training_step_start(self, model: ImaginaireModel, data: dict[str, torch.Tensor], iteration: int = 0) -> None:
-        self.start_iteration_time = time.time()
-
-    @distributed.rank0_only
-    def on_training_step_end(
-        self,
-        model: ImaginaireModel,
-        data_batch: dict[str, torch.Tensor],
-        output_batch: dict[str, torch.Tensor],
-        loss: torch.Tensor,
-        iteration: int = 0,
-    ) -> None:
-        # but this is only called when the optimizer is updated, so it's only the time for the last batch.
-        self.elapsed_iteration_time += time.time() - self.start_iteration_time
-
-        if iteration % self.config.trainer.logging_iter == 0:
-            avg_time = self.elapsed_iteration_time / self.config.trainer.logging_iter
-            log.info(f"Iteration: {iteration}, average iter time: {avg_time:2f}, total loss {loss.item():4f}")
-
-            self.elapsed_iteration_time = 0
-
-
-class LowPrecisionCallback(Callback):
-    """The callback class handling low precision training"""
-
-    def __init__(self, config: Config, trainer: ImaginaireTrainer, update_iter: int):
-        self.update_iter = update_iter
-
-    def on_train_start(self, model: ImaginaireModel, iteration: int = 0) -> None:
-        if model.precision == torch.float32:
-            log.critical("Using fp32. We should disable master weights update.")
-            self.update_iter = sys.maxsize
-        else:
-            assert model.precision in [
-                torch.bfloat16,
-                torch.float16,
-                torch.half,
-            ], "LowPrecisionCallback must use a low precision dtype."
-        self.precision_type = model.precision
-
-    def on_training_step_start(self, model: ImaginaireModel, data: dict[str, torch.Tensor], iteration: int = 0) -> None:
-        for k, v in data.items():
-            if isinstance(v, torch.Tensor) and torch.is_floating_point(data[k]):
-                data[k] = v.to(dtype=self.precision_type)
-
-    def on_validation_step_start(
-        self, model: ImaginaireModel, data: dict[str, torch.Tensor], iteration: int = 0
-    ) -> None:
-        for k, v in data.items():
-            if isinstance(v, torch.Tensor) and torch.is_floating_point(data[k]):
-                data[k] = v.to(dtype=self.precision_type)
-
-    def on_before_zero_grad(
-        self,
-        model_ddp: distributed.DistributedDataParallel,
-        optimizer: torch.optim.Optimizer,
-        scheduler: torch.optim.lr_scheduler.LRScheduler,
-        iteration: int = 0,
-    ) -> None:
-        if iteration % self.update_iter == 0:
-            if getattr(optimizer, "master_weights", False):
-                params, master_params = [], []
-                for group, group_master in zip(optimizer.param_groups, optimizer.param_groups_master):
-                    for p, p_master in zip(group["params"], group_master["params"]):
-                        params.append(get_local_tensor_if_DTensor(p.data))
-                        master_params.append(p_master.data)
-                torch._foreach_copy_(params, master_params)
-
-
-class NVTXCallback(Callback):
-    """The callback for creating NVTX ranges"""
-
-    def __init__(
-        self,
-        synchronize: bool = False,
-        config: Optional["Config"] = None,
-        trainer: Optional["ImaginaireTrainer"] = None,
-    ):
-        super().__init__(config, trainer)
-        self.synchronize = synchronize
-
-    def on_before_forward(self, iteration: int = 0) -> None:
-        if self.synchronize:
-            torch.cuda.synchronize()
-        torch.cuda.nvtx.range_push("forward")
-
-    def on_after_forward(self, iteration: int = 0) -> None:
-        if self.synchronize:
-            torch.cuda.synchronize()
-        torch.cuda.nvtx.range_pop()
-
-    def on_before_backward(
-        self, model_ddp: distributed.DistributedDataParallel, loss: torch.Tensor, iteration: int = 0
-    ) -> None:
-        if self.synchronize:
-            torch.cuda.synchronize()
-        torch.cuda.nvtx.range_push("backward")
-
-    def on_after_backward(self, model_ddp: distributed.DistributedDataParallel, iteration: int = 0) -> None:
-        if self.synchronize:
-            torch.cuda.synchronize()
-        torch.cuda.nvtx.range_pop()
-
-    def on_before_optimizer_step(
-        self,
-        model_ddp: distributed.DistributedDataParallel,
-        optimizer: torch.optim.Optimizer,
-        scheduler: torch.optim.lr_scheduler.LRScheduler,
-        grad_scaler: torch.amp.GradScaler,
-        iteration: int = 0,
-    ) -> None:
-        if self.synchronize:
-            torch.cuda.synchronize()
-        torch.cuda.nvtx.range_push("optimizer_step")
-
-    def on_before_zero_grad(
-        self,
-        model_ddp: distributed.DistributedDataParallel,
-        optimizer: torch.optim.Optimizer,
-        scheduler: torch.optim.lr_scheduler.LRScheduler,
-        iteration: int = 0,
-    ) -> None:
-        if self.synchronize:
-            torch.cuda.synchronize()
-        torch.cuda.nvtx.range_pop()
-
-    def on_before_dataloading(self, iteration: int = 0) -> None:
-        torch.cuda.nvtx.range_push("dataloading")
-
-    def on_after_dataloading(self, iteration: int = 0) -> None:
-        torch.cuda.nvtx.range_pop()

lyra_2/_ext/imaginaire/utils/checkpointer.py

@@ -1,372 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from __future__ import annotations
-
-import os
-import threading
-from typing import TYPE_CHECKING, List, NamedTuple, Tuple
-
-import torch
-
-from lyra_2._ext.imaginaire.model import ImaginaireModel
-from lyra_2._ext.imaginaire.utils import callback, distributed, log, misc, object_store
-
-if TYPE_CHECKING:
-    from lyra_2._ext.imaginaire.config import CheckpointConfig, JobConfig
-
-TORCH_VERSION: Tuple[int, ...] = tuple(int(x) for x in torch.__version__.split(".")[:2])
-if TORCH_VERSION >= (1, 11):
-    from torch.ao import quantization
-    from torch.ao.quantization import FakeQuantizeBase, ObserverBase
-elif (
-    TORCH_VERSION >= (1, 8)
-    and hasattr(torch.quantization, "FakeQuantizeBase")
-    and hasattr(torch.quantization, "ObserverBase")
-):
-    from torch import quantization
-    from torch.quantization import FakeQuantizeBase, ObserverBase
-
-
-class _IncompatibleKeys(
-    NamedTuple(
-        "IncompatibleKeys",
-        [
-            ("missing_keys", List[str]),
-            ("unexpected_keys", List[str]),
-            ("incorrect_shapes", List[Tuple[str, Tuple[int], Tuple[int]]]),
-        ],
-    )
-):
-    pass
-
-
-class Checkpointer:
-    """The checkpointer class. Supports checkpoint saving/loading to both local disk or object store."""
-
-    def __init__(self, config_checkpoint: CheckpointConfig, config_job: JobConfig, callbacks: callback.CallBackGroup):
-        """Constructor of the checkpointer.
-
-        Args:
-            config_checkpoint (CheckpointConfig): The config object for the checkpointer.
-        """
-        # Set the callback functions.
-        self.callbacks = callbacks
-
-        self.checkpoint_dir_local = f"{config_job.path_local}/checkpoints"
-        self.checkpoint_dir_object_store = f"{config_job.path}/checkpoints"
-        self.save_to_object_store = config_checkpoint.save_to_object_store.enabled
-        self.load_from_object_store = config_checkpoint.load_from_object_store.enabled
-        self.strict_resume = config_checkpoint.strict_resume
-        self.load_path = config_checkpoint.load_path or None
-        self.load_training_state = config_checkpoint.load_training_state
-        self.only_load_scheduler_state = config_checkpoint.only_load_scheduler_state
-        self.save_thread = None
-        # Create the object store client interface.
-        if self.save_to_object_store:
-            self.object_store_saver = object_store.ObjectStore(config_checkpoint.save_to_object_store)
-        if self.load_from_object_store:
-            self.object_store_loader = object_store.ObjectStore(config_checkpoint.load_from_object_store)
-
-    def save(
-        self,
-        model: ImaginaireModel,
-        optimizer: torch.optim.Optimizer,
-        scheduler: torch.optim.lr_scheduler.LRScheduler,
-        grad_scaler: torch.amp.GradScaler,
-        iteration: int,
-    ) -> None:
-        """Save network weights, optimizer parameters, scheduler parameters to a checkpoint.
-
-        Args:
-            model (ImaginaireModel): The PyTorch model.
-            optimizer (torch.optim.Optimizer): The model optimizer.
-            scheduler (torch.optim.lr_scheduler.LRScheduler): The optimization scheduler.
-            grad_scaler (torch.amp.GradScaler): The gradient scaler (for mixed precision training).
-            iteration (int): Current iteration number.
-        """
-        self.callbacks.on_save_checkpoint_start(model, iteration)
-
-        checkpoint_file = f"iter_{iteration:09}.pt"
-
-        if distributed.get_rank() == 0:
-            state_dict = dict(
-                model=model.state_dict(),
-                optimizer=optimizer.state_dict(),
-                scheduler=scheduler.state_dict(),
-                grad_scaler=grad_scaler.state_dict(),
-                iteration=iteration,
-            )
-            state_dict = misc.to(state_dict, device="cpu")
-            self.callbacks.on_save_checkpoint(model, state_dict=state_dict)
-            # Wait for previous saver thread to end.
-            if self.save_thread:
-                self.save_thread.join()
-            # Run the checkpoint saver in a separate thread.
-            self.save_thread = threading.Thread(
-                target=self._save_worker_object_store if self.save_to_object_store else self._save_worker_local,
-                daemon=False,
-                args=(state_dict, checkpoint_file, distributed.get_rank()),
-            )
-            self.save_thread.start()
-
-        # Note: Checkpoints are saved on a separate thread and this callback is not accurate.
-        # Please check logs from on_save_checkpoint_success() for better accuracy
-        self.callbacks.on_save_checkpoint_end(model=None, iteration=iteration)
-
-    @misc.timer("checkpoint saving (local)")
-    def _save_worker_local(self, state_dict: dict[str, torch.Tensor], checkpoint_file: str, rank: int = 0) -> None:
-        """Worker to save checkpoint to local disk, spawned with a child thread (runs in parallel with the training).
-
-        Args:
-            state_dict (dict[str, torch.Tensor]): The state dict of the model/optimizer/scheduler.
-            checkpoint_file (str): The file name of the model checkpoint.
-            rank (int): GPU device (default: 0).
-        """
-        checkpoint_path = os.path.join(self.checkpoint_dir_local, checkpoint_file)
-        os.makedirs(self.checkpoint_dir_local, exist_ok=True)
-        try:
-            torch.save(state_dict, checkpoint_path)
-            if rank == 0:
-                self._write_latest_checkpoint_file(checkpoint_file)
-            log.success(f"Saved checkpoint (local): {checkpoint_path}")
-            iteration = int(checkpoint_file.replace("iter_", "").replace(".pt", ""))
-            self.callbacks.on_save_checkpoint_success(iteration=iteration)
-        except Exception as e:  # noqa: BLE001
-            log.exception(f"Checkpoint failed to save (local): {e}")
-
-    @misc.timer("checkpoint saving (object store)")
-    def _save_worker_object_store(
-        self, state_dict: dict[str, torch.Tensor], checkpoint_file: str, rank: int = 0
-    ) -> None:
-        """Worker to upload checkpoint to object store, spawned with a child thread (in parallel with the training).
-
-        Args:
-            state_dict (dict[str, torch.Tensor]): The state dict of the model/optimizer/scheduler.
-            checkpoint_file (str): The file name of the model checkpoint.
-            rank (int): GPU device (default: 0).
-        """
-        checkpoint_path = os.path.join(self.checkpoint_dir_object_store, checkpoint_file)
-        try:
-            self.object_store_saver.save_object(state_dict, key=checkpoint_path, type="torch")
-            if rank == 0:
-                self._write_latest_checkpoint_file(checkpoint_file)
-            log.success(f"Saved checkpoint (object store): {checkpoint_path}")
-            iteration = int(checkpoint_file.replace("iter_", "").replace(".pt", ""))
-            self.callbacks.on_save_checkpoint_success(iteration=iteration)
-        except Exception as e:  # noqa: BLE001
-            log.exception(f"Checkpoint failed to upload (object store): {e}")
-
-    @misc.timer("checkpoint loading")
-    def load(
-        self,
-        model: ImaginaireModel,
-        optimizer: torch.optim.Optimizer | None = None,
-        scheduler: torch.optim.lr_scheduler.LRScheduler | None = None,
-        grad_scaler: torch.amp.GradScaler | None = None,
-    ) -> int:
-        """Load network weights and optimizer states from a checkpoint in a single process.
-
-        The priority of the checkpoint loading logic is:
-        1. Attempt to resume training if possible by looking for latest_checkpoint.txt under the same name.
-        2. If no latest checkpoint were found, it loads the model weights specified by config_checkpoint.path.
-           - This is typically used for inference mode.
-           - If config_checkpoint.load_optimizer_state is True, then also load the optimizer and scheduler states.
-        3. If none of the above, randomly initialize the model parameters and train from scratch.
-
-        Args:
-            model (ImaginaireModel): The PyTorch model.
-            optimizer (torch.optim.Optimizer | None): The model optimizer (default: None).
-            scheduler (torch.optim.lr_scheduler.LRScheduler | None): The optimization scheduler (default: None).
-            grad_scaler (torch.amp.GradScaler | None): The gradient scaler (for mixed precision training).
-
-        Returns:
-            iteration (int): the iteration number to start/resume from.
-        """
-        self.callbacks.on_load_checkpoint_start(model)
-
-        latest_checkpoint_file = self._read_latest_checkpoint_file()
-        if latest_checkpoint_file is not None:
-            # 1. Resume training from latest_checkpoint.txt under the same name.
-            checkpoint_dir = (
-                self.checkpoint_dir_object_store if self.load_from_object_store else self.checkpoint_dir_local
-            )
-            checkpoint_path = os.path.join(checkpoint_dir, latest_checkpoint_file)
-            resume = True
-            only_resume_scheduler = True
-        else:
-            if self.load_path:
-                # 2. Load the module weights specified by config_checkpoint.path.
-                checkpoint_path = self.load_path
-                resume = self.load_training_state
-                only_resume_scheduler = self.only_load_scheduler_state
-            else:
-                # 3. Randomly initialize the model parameters and train from scratch.
-                checkpoint_path = None
-                resume = False
-                only_resume_scheduler = False
-        # Load checkpoint.
-        if checkpoint_path is not None:
-            self._check_checkpoint_exists(checkpoint_path)
-            if self.load_from_object_store:
-                log.info(f"Loading checkpoint (object store): {checkpoint_path}")
-                state_dict = self.object_store_loader.load_object(key=checkpoint_path, type="torch", max_attempts=20)
-                log.success(f"Complete loading checkpoint (object store): {checkpoint_path}")
-            else:
-                log.info(f"Loading checkpoint (local): {checkpoint_path}")
-                state_dict = torch.load(checkpoint_path, map_location=lambda storage, loc: storage)
-                log.success(f"Complete loading checkpoint (local): {checkpoint_path}")
-            self.callbacks.on_load_checkpoint(model, state_dict=state_dict)
-            # Load the state dicts.
-            log.info("- Loading the model...")
-            model.load_state_dict(state_dict["model"], strict=self.strict_resume)
-            if resume or only_resume_scheduler:
-                iteration = state_dict["iteration"]
-                assert scheduler
-                log.info("- Loading the scheduler...")
-                scheduler.load_state_dict(state_dict["scheduler"])
-                scheduler.last_epoch = iteration
-            else:
-                iteration = 0
-            if resume:
-                assert optimizer
-                log.info("- Loading the optimizer...")
-                optimizer.load_state_dict(state_dict["optimizer"])
-                log.info("- Loading the gradient scaler...")
-                grad_scaler.load_state_dict(state_dict["grad_scaler"])
-                log.success(f"Done with loading the checkpoint (iteration {iteration}).")
-            else:
-                log.success("Done with loading the checkpoint.")
-        else:
-            # Checkpoint not found and not specified. We will train everything from scratch.
-            iteration = 0
-            log.info("Training from scratch.")
-        torch.cuda.empty_cache()
-
-        self.callbacks.on_load_checkpoint_end(model, iteration=iteration, checkpoint_path=checkpoint_path)
-
-        return iteration
-
-    def _read_latest_checkpoint_file(self) -> str | None:
-        """Get the file name of the latest saved checkpoint. If it doesn't exist, return None.
-
-        Returns:
-            checkpoint_file (str | None): file name of the latest saved checkpoint.
-        """
-        checkpoint_file = None
-        if self.load_from_object_store:
-            latest_path = os.path.join(self.checkpoint_dir_object_store, "latest_checkpoint.txt")
-            if self.object_store_loader.object_exists(key=latest_path):
-                checkpoint_file = self.object_store_loader.load_object(key=latest_path, type="text").strip()
-        else:
-            latest_path = os.path.join(self.checkpoint_dir_local, "latest_checkpoint.txt")
-            if os.path.isfile(latest_path):
-                checkpoint_file = open(latest_path).read().strip()
-        return checkpoint_file
-
-    def _write_latest_checkpoint_file(self, checkpoint_file: str) -> None:
-        """Track the file name of the latest saved checkpoint.
-
-        Args:
-            checkpoint_file (str): file name of the latest saved checkpoint.
-        """
-        content = f"{checkpoint_file}\n"
-        if self.save_to_object_store:
-            latest_path = os.path.join(self.checkpoint_dir_object_store, "latest_checkpoint.txt")
-            self.object_store_saver.save_object(content, key=latest_path, type="text")
-        else:
-            latest_path = os.path.join(self.checkpoint_dir_local, "latest_checkpoint.txt")
-            with open(latest_path, "w") as file:
-                file.write(content)
-
-    def _check_checkpoint_exists(self, checkpoint_path: str) -> None:
-        """If the file checkpoint_path does not exist, raise an error.
-
-        Args:
-            checkpoint_path (str): full path to the checkpoint.
-        """
-        if self.load_from_object_store:
-            if not self.object_store_loader.object_exists(key=checkpoint_path):
-                raise FileNotFoundError(f"File not found (object store): {checkpoint_path}")
-        else:
-            if not os.path.exists(checkpoint_path):
-                raise FileNotFoundError(f"File not found (local): {checkpoint_path}")
-
-    def finalize(self) -> None:
-        """Finalize the checkpointer."""
-        if self.save_thread:
-            self.save_thread.join()
-
-
-# https://github.com/facebookresearch/fvcore/blob/9d683aae73fb899dd35d6cf6720e5ef567761c57/fvcore/common/checkpoint.py
-def non_strict_load_model(model: torch.nn.Module, checkpoint_state_dict: dict) -> _IncompatibleKeys:
-    # workaround https://github.com/pytorch/pytorch/issues/24139
-    model_state_dict = model.state_dict()
-    incorrect_shapes = []
-    for k in list(checkpoint_state_dict.keys()):
-        if k in model_state_dict:
-            if "_extra_state" in k:  # Key introduced by TransformerEngine for FP8
-                log.warning(f"Skipping key {k} introduced by TransformerEngine for FP8 in the checkpoint.")
-                continue
-            model_param = model_state_dict[k]
-            # Allow mismatch for uninitialized parameters
-            if TORCH_VERSION >= (1, 8) and isinstance(model_param, torch.nn.parameter.UninitializedParameter):
-                continue
-            if not isinstance(model_param, torch.Tensor):
-                raise ValueError(
-                    f"Find non-tensor parameter {k} in the model. type: {type(model_param)} {type(checkpoint_state_dict[k])}, please check if this key is safe to skip or not."
-                )
-
-            shape_model = tuple(model_param.shape)
-            shape_checkpoint = tuple(checkpoint_state_dict[k].shape)
-            if shape_model != shape_checkpoint:
-                has_observer_base_classes = (
-                    TORCH_VERSION >= (1, 8)
-                    and hasattr(quantization, "ObserverBase")
-                    and hasattr(quantization, "FakeQuantizeBase")
-                )
-                if has_observer_base_classes:
-                    # Handle the special case of quantization per channel observers,
-                    # where buffer shape mismatches are expected.
-                    def _get_module_for_key(model: torch.nn.Module, key: str) -> torch.nn.Module:
-                        # foo.bar.param_or_buffer_name -> [foo, bar]
-                        key_parts = key.split(".")[:-1]
-                        cur_module = model
-                        for key_part in key_parts:
-                            cur_module = getattr(cur_module, key_part)
-                        return cur_module
-
-                    cls_to_skip = (
-                        ObserverBase,
-                        FakeQuantizeBase,
-                    )
-                    target_module = _get_module_for_key(model, k)
-                    if isinstance(target_module, cls_to_skip):
-                        # Do not remove modules with expected shape mismatches
-                        # them from the state_dict loading. They have special logic
-                        # in _load_from_state_dict to handle the mismatches.
-                        continue
-
-                incorrect_shapes.append((k, shape_checkpoint, shape_model))
-                checkpoint_state_dict.pop(k)
-    incompatible = model.load_state_dict(checkpoint_state_dict, strict=False)
-    # Remove keys with "_extra_state" suffix, which are non-parameter items introduced by TransformerEngine for FP8 handling
-    missing_keys = [k for k in incompatible.missing_keys if "_extra_state" not in k]
-    unexpected_keys = [k for k in incompatible.unexpected_keys if "_extra_state" not in k]
-    return _IncompatibleKeys(
-        missing_keys=missing_keys,
-        unexpected_keys=unexpected_keys,
-        incorrect_shapes=incorrect_shapes,
-    )

lyra_2/_ext/imaginaire/utils/config_helper.py

@@ -1,214 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import importlib
-import os
-import pkgutil
-import sys
-from dataclasses import fields as dataclass_fields
-from dataclasses import is_dataclass
-from typing import Any, Dict, Optional
-
-import attr
-import attrs
-from hydra import compose, initialize
-from hydra.core.config_store import ConfigStore
-from hydra.core.global_hydra import GlobalHydra
-from omegaconf import DictConfig, OmegaConf
-
-from lyra_2._ext.imaginaire.config import Config
-from lyra_2._ext.imaginaire.utils import log
-
-
-def is_attrs_or_dataclass(obj) -> bool:
-    """
-    Check if the object is an instance of an attrs class or a dataclass.
-
-    Args:
-        obj: The object to check.
-
-    Returns:
-        bool: True if the object is an instance of an attrs class or a dataclass, False otherwise.
-    """
-    return is_dataclass(obj) or attr.has(type(obj))
-
-
-def get_fields(obj):
-    """
-    Get the fields of an attrs class or a dataclass.
-
-    Args:
-        obj: The object to get fields from. Must be an instance of an attrs class or a dataclass.
-
-    Returns:
-        list: A list of field names.
-
-    Raises:
-        ValueError: If the object is neither an attrs class nor a dataclass.
-    """
-    if is_dataclass(obj):
-        return [field.name for field in dataclass_fields(obj)]
-    elif attr.has(type(obj)):
-        return [field.name for field in attr.fields(type(obj))]
-    else:
-        raise ValueError("The object is neither an attrs class nor a dataclass.")
-
-
-def override(config: Config, overrides: Optional[list[str]] = None) -> Config:
-    """
-    :param config: the instance of class `Config` (usually from `make_config`)
-    :param overrides: list of overrides for config
-    :return: the composed instance of class `Config`
-    """
-    # Store the class of the config for reconstruction after overriding.
-    # config_class = type(config)
-
-    # Convert Config object to a DictConfig object
-    config_dict = attrs.asdict(config)
-    config_omegaconf = DictConfig(content=config_dict, flags={"allow_objects": True})
-    # Enforce "--" separator between the script arguments and overriding configs.
-    if overrides:
-        if overrides[0] != "--":
-            raise ValueError(
-                f'Hydra config overrides must be separated with a "--" token. but got overrides={overrides}, and overrides[0]={overrides[0]}'
-            )
-        overrides = overrides[1:]
-    # Use Hydra to handle overrides
-    cs = ConfigStore.instance()
-    cs.store(name="config", node=config_omegaconf)
-    if not GlobalHydra().is_initialized():
-        with initialize(version_base=None):
-            config_omegaconf = compose(config_name="config", overrides=overrides)
-            OmegaConf.resolve(config_omegaconf)
-    else:
-        config_omegaconf = compose(config_name="config", overrides=overrides)
-        OmegaConf.resolve(config_omegaconf)
-
-    def config_from_dict(ref_instance: Any, kwargs: Any) -> Any:
-        """
-        Construct an instance of the same type as ref_instance using the provided dictionary or data or unstructured data
-
-        Args:
-            ref_instance: The reference instance to determine the type and fields when needed
-            kwargs: A dictionary of keyword arguments to use for constructing the new instance or primitive data or unstructured data
-
-        Returns:
-            Any: A new instance of the same type as ref_instance constructed using the provided kwargs or the primitive data or unstructured data
-
-        Raises:
-            AssertionError: If the fields do not match or if extra keys are found.
-            Exception: If there is an error constructing the new instance.
-        """
-        is_type = is_attrs_or_dataclass(ref_instance)
-        if not is_type:
-            return kwargs
-        else:
-            ref_fields = set(get_fields(ref_instance))
-            assert isinstance(kwargs, dict) or isinstance(kwargs, DictConfig), (
-                "kwargs must be a dictionary or a DictConfig"
-            )
-            keys = set(kwargs.keys())
-
-            # ref_fields must equal to or include all keys
-            extra_keys = keys - ref_fields
-            assert ref_fields == keys or keys.issubset(ref_fields), (
-                f"Fields mismatch: {ref_fields} != {keys}. Extra keys found: {extra_keys} \n \t when constructing {type(ref_instance)} with {keys}"
-            )
-
-            resolved_kwargs: Dict[str, Any] = {}
-            for f in keys:
-                resolved_kwargs[f] = config_from_dict(getattr(ref_instance, f), kwargs[f])
-            try:
-                new_instance = type(ref_instance)(**resolved_kwargs)
-            except Exception as e:
-                log.error(f"Error when constructing {type(ref_instance)} with {resolved_kwargs}")
-                log.error(e)
-                raise e
-            return new_instance
-
-    config = config_from_dict(config, config_omegaconf)
-
-    return config
-
-
-def get_config_module(config_file: str) -> str:
-    if not config_file.endswith(".py"):
-        log.error("Config file cannot be specified as module.")
-        log.error("Please provide the path to the Python config file (relative to the Imaginaire4 root).")
-    assert os.path.isfile(config_file), f"Imaginaire4 config file ({config_file}) not found."
-    # Convert to importable module format.
-    config_module = config_file.replace("/", ".").replace(".py", "")
-    return config_module
-
-
-def import_module(full_module_name: str, reload: bool = False):
-    """
-    Import a module by name.
-
-    Args:
-        full_module_name: The fully qualified name of the module to import.
-        reload: If True, reload the module if it's already imported.
-    """
-    if full_module_name in sys.modules and reload:
-        importlib.reload(sys.modules[full_module_name])
-    else:
-        importlib.import_module(full_module_name)
-
-
-def import_all_modules_from_package(package_path: str, reload: bool = False, skip_underscore: bool = True) -> None:
-    """
-    Import all modules from the specified package path recursively.
-
-    This function is typically used in conjunction with Hydra to ensure that all modules
-    within a specified package are imported, which is necessary for registering configurations.
-
-    Example usage:
-    ```python
-    import_all_modules_from_package("projects.cosmos.diffusion.v1.config.experiment", reload=True, skip_underscore=False)
-    ```
-
-    Args:
-        package_path (str): The dotted path to the package from which to import all modules.
-        reload (bool): Flag to determine whether to reload modules if they're already imported.
-        skip_underscore (bool): If True, skips importing modules that start with an underscore.
-    """
-    log.info(f"{'Reloading' if reload else 'Importing'} all modules from package {package_path}")
-    package = importlib.import_module(package_path)
-    package_directory = package.__path__
-
-    def import_modules_recursively(directory: str, prefix: str) -> None:
-        """
-        Recursively imports or reloads all modules in the given directory.
-
-        Args:
-            directory (str): The file system path to the current package directory.
-            prefix (str): The module prefix (e.g., 'projects.cosmos.diffusion.v1.config').
-        """
-        for _, module_name, is_pkg in pkgutil.iter_modules([directory]):
-            if skip_underscore and module_name.startswith("_"):
-                log.debug(f"Skipping module {module_name} as it starts with an underscore")
-                continue
-
-            full_module_name = f"{prefix}.{module_name}"
-            log.debug(f"{'Reloading' if reload else 'Importing'} module {full_module_name}")
-
-            import_module(full_module_name, reload=reload)
-
-            if is_pkg:
-                sub_package_directory = os.path.join(directory, module_name)
-                import_modules_recursively(sub_package_directory, full_module_name)
-
-    for directory in package_directory:
-        import_modules_recursively(directory, package_path)

lyra_2/_ext/imaginaire/utils/context_managers.py

@@ -1,55 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from contextlib import ExitStack, contextmanager
-from typing import Generator
-
-from lyra_2._ext.imaginaire.utils.misc import timer
-
-
-@contextmanager
-def data_loader_init() -> Generator[None, None, None]:
-    """
-    Wrap the data loader initialization with multiple context managers used for telemetry and one logger.
-    """
-    contexts = [
-        timer("init_data_loader"),
-    ]
-    with ExitStack() as stack:
-        yield [stack.enter_context(cm) for cm in contexts]
-
-
-@contextmanager
-def model_init(set_barrier: bool = False) -> Generator[None, None, None]:
-    """
-    Wrap the instantiation of the model with multiple context managers used for telemetry and one logger.
-    """
-    contexts = [
-        timer("init_model"),
-    ]
-    with ExitStack() as stack:
-        yield [stack.enter_context(cm) for cm in contexts]
-
-
-@contextmanager
-def distributed_init() -> Generator[None, None, None]:
-    """
-    Wrap the distributed initialization, used for telemetry and timers
-    """
-    contexts = [
-        timer("init_distributed"),
-    ]
-    with ExitStack() as stack:
-        yield [stack.enter_context(cm) for cm in contexts]

lyra_2/_ext/imaginaire/utils/count_params.py

@@ -1,29 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from torch import nn
-
-
-def disabled_train(self, mode: bool = True):
-    """Overwrite model.train with this function to make sure train/eval mode
-    does not change anymore."""
-    return self
-
-
-def count_params(model: nn.Module, verbose=False) -> int:
-    total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
-    if verbose:
-        print(f"{model.__class__.__name__} has {total_params * 1.0e-6:.2f} M params.")
-    return total_params

lyra_2/_ext/imaginaire/utils/device.py

@@ -1,114 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import gc
-import math
-import os
-
-import pynvml
-from loguru import logger as logging
-
-
-def get_gpu_architecture():
-    """
-    Retrieves the GPU architecture of the available GPUs.
-
-    Returns:
-        str: The GPU architecture, which can be "H100", "A100", or "Other".
-    """
-    try:
-        pynvml.nvmlInit()
-        device_count = pynvml.nvmlDeviceGetCount()
-        for i in range(device_count):
-            handle = pynvml.nvmlDeviceGetHandleByIndex(i)
-            model_name = pynvml.nvmlDeviceGetName(handle)
-            if isinstance(model_name, bytes):
-                model_name = model_name.decode("utf-8")
-            print(f"GPU {i}: Model: {model_name}")
-
-            # Check for specific models like H100 or A100
-            if "H100" in model_name or "H200" in model_name:
-                return "H100"
-            elif "A100" in model_name:
-                return "A100"
-            elif "L40S" in model_name:
-                return "L40S"
-            elif "B200" in model_name:
-                return "B200"
-    except pynvml.NVMLError as error:
-        print(f"Failed to get GPU info: {error}")
-    finally:
-        pynvml.nvmlShutdown()
-
-    # return "Other" incase of non hopper/ampere or error
-    return "Other"
-
-
-class GPUArchitectureNotSupported(Exception):
-    """
-    Custom exception raised when the expected GPU architecture is not supported.
-    """
-
-    pass
-
-
-def print_gpu_mem(str=None):
-    try:
-        pynvml.nvmlInit()
-        meminfo = pynvml.nvmlDeviceGetMemoryInfo(pynvml.nvmlDeviceGetHandleByIndex(0))
-        logging.info(
-            f"{str}: {meminfo.used / 1024 / 1024}/{meminfo.total / 1024 / 1024}MiB used ({meminfo.free / 1024 / 1024}MiB free)"
-        )
-    except pynvml.NVMLError as error:
-        print(f"Failed to get GPU memory info: {error}")
-
-
-def force_gc():
-    print_gpu_mem()
-    print("gc()")
-    gc.collect()
-    print_gpu_mem()
-    print("empty cuda cache")
-    # print(torch.cuda.memory_summary())
-    print_gpu_mem()
-
-
-def gpu0_has_80gb_or_less():
-    try:
-        pynvml.nvmlInit()
-        meminfo = pynvml.nvmlDeviceGetMemoryInfo(pynvml.nvmlDeviceGetHandleByIndex(0))
-        return meminfo.total / 1024 / 1024 / 1024 <= 80
-    except pynvml.NVMLError as error:
-        print(f"Failed to get GPU memory info: {error}")
-
-
-class Device:
-    _nvml_affinity_elements = math.ceil(os.cpu_count() / 64)  # type: ignore
-
-    def __init__(self, device_idx: int):
-        super().__init__()
-        self.handle = pynvml.nvmlDeviceGetHandleByIndex(device_idx)
-
-    def get_name(self) -> str:
-        return pynvml.nvmlDeviceGetName(self.handle)
-
-    def get_cpu_affinity(self) -> list[int]:
-        affinity_string = ""
-        for j in pynvml.nvmlDeviceGetCpuAffinity(self.handle, Device._nvml_affinity_elements):
-            # assume nvml returns list of 64 bit ints
-            affinity_string = "{:064b}".format(j) + affinity_string
-        affinity_list = [int(x) for x in affinity_string]
-        affinity_list.reverse()  # so core 0 is in 0th element of list
-        return [i for i, e in enumerate(affinity_list) if e != 0]

lyra_2/_ext/imaginaire/utils/distributed.py

@@ -1,443 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from __future__ import annotations
-
-import collections
-import collections.abc
-import ctypes
-import functools
-import os
-from contextlib import contextmanager
-from datetime import timedelta
-from typing import TYPE_CHECKING, Any, Callable, Container, Optional
-
-import pynvml
-import torch
-import torch.distributed as dist
-from torch.distributed import get_process_group_ranks
-
-from lyra_2._ext.imaginaire.utils.device import Device
-
-if dist.is_available():
-    from torch.distributed.distributed_c10d import _get_default_group
-    from torch.distributed.utils import _sync_module_states, _verify_param_shape_across_processes
-
-from lyra_2._ext.imaginaire.utils import log
-
-if TYPE_CHECKING:
-    from lyra_2._ext.imaginaire.config import DDPConfig
-
-try:
-    from megatron.core import parallel_state
-except ImportError:
-    print("Megatron-core is not installed.")
-
-
-def init() -> int | None:
-    """Initialize distributed training."""
-    if dist.is_initialized():
-        return torch.cuda.current_device()
-
-    # Set GPU affinity.
-    pynvml.nvmlInit()
-    local_rank = int(os.getenv("LOCAL_RANK", 0))
-    try:
-        device = Device(local_rank)
-        os.sched_setaffinity(0, device.get_cpu_affinity())
-    except Exception as e:
-        log.warning(f"Failed to set device affinity: {e}")
-    # Set up NCCL communication.
-    os.environ["TORCH_NCCL_BLOCKING_WAIT"] = "0"
-    os.environ["TORCH_NCCL_ASYNC_ERROR_HANDLING"] = "1"
-    if dist.is_available():
-        torch.cuda.set_device(local_rank)
-        # Get the timeout value from environment variable
-        timeout_seconds = os.getenv("TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC", 1800)
-        # Convert the timeout to an integer (if it isn't already) and then to a timedelta
-        timeout_timedelta = timedelta(seconds=int(timeout_seconds))
-        dist.init_process_group(backend="nccl", init_method="env://", timeout=timeout_timedelta)
-        log.info(
-            f"Initialized distributed training with local rank {local_rank} with timeout {timeout_seconds}",
-            rank0_only=False,
-        )
-    # Increase the L2 fetch granularity for faster speed.
-    _libcudart = ctypes.CDLL("libcudart.so")
-    # Set device limit on the current device.
-    p_value = ctypes.cast((ctypes.c_int * 1)(), ctypes.POINTER(ctypes.c_int))
-    _libcudart.cudaDeviceSetLimit(ctypes.c_int(0x05), ctypes.c_int(128))
-    _libcudart.cudaDeviceGetLimit(p_value, ctypes.c_int(0x05))
-    log.info(f"Training with {get_world_size()} GPUs.")
-
-
-def get_rank(group: Optional[dist.ProcessGroup] = None) -> int:
-    """Get the rank (GPU device) of the worker.
-
-    Returns:
-        rank (int): The rank of the worker.
-    """
-    rank = 0
-    if dist.is_available() and dist.is_initialized():
-        rank = dist.get_rank(group)
-    return rank
-
-
-def get_world_size(group: Optional[dist.ProcessGroup] = None) -> int:
-    """Get world size. How many GPUs are available in this job.
-
-    Returns:
-        world_size (int): The total number of GPUs available in this job.
-    """
-    world_size = 1
-    if dist.is_available() and dist.is_initialized():
-        world_size = dist.get_world_size(group)
-    return world_size
-
-
-def is_rank0() -> bool:
-    """Check if current process is the master GPU.
-
-    Returns:
-        (bool): True if this function is called from the master GPU, else False.
-    """
-    return get_rank() == 0
-
-
-def is_local_rank0() -> bool:
-    """Check if current process is the local master GPU in the current node.
-
-    Returns:
-        (bool): True if this function is called from the local master GPU, else False.
-    """
-    return torch.cuda.current_device() == 0
-
-
-def rank0_only(func: Callable) -> Callable:
-    """Apply this function only to the master GPU.
-
-    Example usage:
-        @rank0_only
-        def func(x):
-            return x + 3
-
-    Args:
-        func (Callable): a function.
-
-    Returns:
-        (Callable): A function wrapper executing the function only on the master GPU.
-    """
-
-    @functools.wraps(func)
-    def wrapper(*args, **kwargs):  # noqa: ANN202
-        if is_rank0():
-            return func(*args, **kwargs)
-        else:
-            return None
-
-    return wrapper
-
-
-def barrier() -> None:
-    """Barrier for all GPUs."""
-    if dist.is_available() and dist.is_initialized():
-        dist.barrier()
-
-
-def rank0_first(func: Callable) -> Callable:
-    """run the function on rank 0 first, then on other ranks."""
-
-    @functools.wraps(func)
-    def wrapper(*args, **kwargs):  # noqa: ANN202
-        if is_rank0():
-            result = func(*args, **kwargs)
-        barrier()
-        if not is_rank0():
-            result = func(*args, **kwargs)
-        return result
-
-    return wrapper
-
-
-def parallel_model_wrapper(config_ddp: DDPConfig, model: torch.nn.Module) -> torch.nn.Module | DistributedDataParallel:
-    """Wraps the model to enable data parallalism for training across multiple GPU devices.
-
-    Args:
-        config_ddp (DDPConfig): The data parallel config.
-        model (torch.nn.Module): The PyTorch module.
-
-    Returns:
-        model (torch.nn.Module | DistributedDataParallel): The data parallel model wrapper
-            if distributed environment is available, otherwise return the original model.
-    """
-    if dist.is_available() and dist.is_initialized():
-        local_rank = int(os.getenv("LOCAL_RANK", 0))
-        try:
-            ddp_group = parallel_state.get_data_parallel_group(with_context_parallel=True)
-        except Exception as e:
-            log.info(e)
-            log.info("parallel_state not initialized, treating all GPUs equally for DDP")
-            ddp_group = None
-
-        model = DistributedDataParallel(
-            model,
-            device_ids=[local_rank],
-            output_device=local_rank,
-            find_unused_parameters=config_ddp.find_unused_parameters,
-            static_graph=config_ddp.static_graph,
-            broadcast_buffers=config_ddp.broadcast_buffers,
-            process_group=ddp_group,
-        )
-    return model
-
-
-class DistributedDataParallel(torch.nn.parallel.DistributedDataParallel):
-    """This extends torch.nn.parallel.DistributedDataParallel with .training_step().
-
-    This borrows the concept of `forward-redirection` from Pytorch lightning. It wraps an ImaginaireModel such that
-    model.training_step() would be executed when calling self.training_step(), while preserving the behavior of calling
-    model() for Pytorch modules. Internally, this is a double rerouting mechanism (training_step -> forward ->
-    training_step), allowing us to preserve the function names and signatures.
-    """
-
-    def __init__(self, model: torch.nn.Module, *args, **kwargs):
-        super().__init__(model, *args, **kwargs)
-        self.show_sync_grad_static_graph_warning = True
-
-    def training_step(self, *args, **kwargs) -> Any:
-        # Cache the original model.forward() method.
-        original_forward = self.module.forward
-
-        def wrapped_training_step(*_args, **_kwargs):  # noqa: ANN202
-            # Unpatch immediately before calling training_step() because itself may want to call the real forward.
-            self.module.forward = original_forward
-            # The actual .training_step().
-            return self.module.training_step(*_args, **_kwargs)
-
-        # Patch the original_module's forward so we can redirect the arguments back to the real method.
-        self.module.forward = wrapped_training_step
-        # Call self, which implicitly calls self.forward() --> model.forward(), which is now model.training_step().
-        # Without calling self.forward() or model.forward() explciitly, implicit hooks are also executed.
-        return self(*args, **kwargs)
-
-
-@contextmanager
-def ddp_sync_grad(model, enabled):
-    r"""
-    Context manager to enable/disable gradient synchronizations across DDP processes for DDP model.
-    Modified from:
-    https://pytorch.org/docs/stable/_modules/torch/nn/parallel/distributed.html#DistributedDataParallel.no_sync
-    Note that this is incompatible with static_graph=True and will be an no-op if static_graph=True.
-
-    Within this context, gradients will be accumulated on module
-    variables, which will later be synchronized in the first
-    forward-backward pass exiting the context.
-
-    .. warning::
-        The forward pass should be included inside the context manager, or
-        else gradients will still be synchronized.
-    """
-    assert isinstance(model, torch.nn.Module)
-    if isinstance(model, DistributedDataParallel):
-        old_require_backward_grad_sync = model.require_backward_grad_sync
-        if model.static_graph and model.require_backward_grad_sync != enabled:
-            if model.show_sync_grad_static_graph_warning:
-                log.warning("DDP static_graph=True is incompatible with sync_grad(). Performance will be reduced.")
-                model.show_sync_grad_static_graph_warning = False
-        else:
-            model.require_backward_grad_sync = enabled
-    try:
-        yield
-    finally:
-        if isinstance(model, DistributedDataParallel):
-            model.require_backward_grad_sync = old_require_backward_grad_sync
-
-
-def collate_batches(data_batches: list[dict[str, torch.Tensor]]) -> torch.Tensor | dict[str, torch.Tensor]:
-    """Aggregate the list of data batches from all devices and process the results.
-
-    This is used for gathering validation data batches with lyra_2._ext.imaginaire.utils.dataloader.DistributedEvalSampler.
-    It will return the data/output of the entire validation set in its original index order. The sizes of data_batches
-    in different ranks may differ by 1 (if dataset size is not evenly divisible), in which case a dummy sample will be
-    created before calling dis.all_gather().
-
-    Args:
-        data_batches (list[dict[str, torch.Tensor]]): List of tensors or (hierarchical) dictionary where
-            leaf entries are tensors.
-
-    Returns:
-        data_gather (torch.Tensor | dict[str, torch.Tensor]): tensors or (hierarchical) dictionary where
-            leaf entries are concatenated tensors.
-    """
-    if isinstance(data_batches[0], torch.Tensor):
-        # Concatenate the local data batches.
-        data_concat = torch.cat(data_batches, dim=0)  # type: ignore
-        # Get the largest number of local samples from all ranks to determine whether to dummy-pad on this rank.
-        max_num_local_samples = torch.tensor(len(data_concat), device="cuda")
-        dist.all_reduce(max_num_local_samples, op=dist.ReduceOp.MAX)
-        if len(data_concat) < max_num_local_samples:
-            assert len(data_concat) + 1 == max_num_local_samples
-            dummy = torch.empty_like(data_concat[:1])
-            data_concat = torch.cat([data_concat, dummy], dim=0)
-            dummy_count = torch.tensor(1, device="cuda")
-        else:
-            dummy_count = torch.tensor(0, device="cuda")
-        # Get all concatenated batches from all ranks and concatenate again.
-        dist.all_reduce(dummy_count, op=dist.ReduceOp.SUM)
-        data_concat = all_gather_tensor(data_concat.contiguous())
-        data_collate = torch.stack(data_concat, dim=1).flatten(start_dim=0, end_dim=1)
-        # Remove the dummy samples.
-        if dummy_count > 0:
-            data_collate = data_collate[:-dummy_count]
-    elif isinstance(data_batches[0], collections.abc.Mapping):
-        data_collate = dict()
-        for key in data_batches[0].keys():
-            data_collate[key] = collate_batches([data[key] for data in data_batches])  # type: ignore
-    else:
-        raise TypeError
-    return data_collate
-
-
-@torch.no_grad()
-def all_gather_tensor(tensor: torch.Tensor) -> list[torch.Tensor]:
-    """Gather the corresponding tensor from all GPU devices to a list.
-
-    Args:
-        tensor (torch.Tensor): Pytorch tensor.
-
-    Returns:
-        tensor_list (list[torch.Tensor]): A list of Pytorch tensors gathered from all GPU devices.
-    """
-    tensor_list = [torch.zeros_like(tensor) for _ in range(get_world_size())]
-    dist.all_gather(tensor_list, tensor)
-    return tensor_list
-
-
-def broadcast(tensor, src, group=None, async_op=False):
-    world_size = get_world_size()
-    if world_size < 2:
-        return tensor
-    dist.broadcast(tensor, src=src, group=group, async_op=async_op)
-
-
-def dist_reduce_tensor(tensor, rank=0, reduce="mean"):
-    r"""Reduce to rank 0"""
-    world_size = get_world_size()
-    if world_size < 2:
-        return tensor
-    with torch.no_grad():
-        dist.reduce(tensor, dst=rank)
-        if get_rank() == rank:
-            if reduce == "mean":
-                tensor /= world_size
-            elif reduce == "sum":
-                pass
-            else:
-                raise NotImplementedError
-    return tensor
-
-
-def sync_model_states(
-    model: torch.nn.Module,
-    process_group: Optional[dist.ProcessGroup] = None,
-    src: int = 0,
-    params_and_buffers_to_ignore: Optional[Container[str]] = None,
-    broadcast_buffers: bool = True,
-):
-    """
-    Modify based on DDP source code
-    Synchronizes the parameters and buffers of a model across different processes in a distributed setting.
-
-    This function ensures that all processes in the specified process group have the same initial parameters and
-    buffers from the source rank, typically rank 0. It is useful when different processes start with different model
-    states and a synchronization is required to ensure consistency across all ranks.
-
-    Args:
-        model (nn.Module): The model whose parameters and buffers are to be synchronized.
-        process_group (dist.ProcessGroup, optional): The process group for communication. If None,
-            the default group is used. Defaults to None.
-        src (int, optional): The source rank from which parameters and buffers will be broadcasted.
-            Defaults to 0.
-        params_and_buffers_to_ignore (Optional[Container[str]], optional): A container of parameter and buffer
-            names to exclude from synchronization. Defaults to None, which means all parameters and buffers are
-            included.
-        broadcast_buffers (bool, optional): Whether to broadcast buffers or not. Defaults to True.
-
-    Side Effects:
-        This function modifies the state of the model in-place to synchronize it with the source rank's model state.
-
-    Raises:
-        RuntimeError: If the shapes of parameters across processes do not match, a runtime error will be raised.
-
-    Examples:
-        >>> # downloading duplicated model weights from s3 in each rank and save network bandwidth
-        >>> # useful and save our time when model weights are huge
-        >>> if dist.get_rank == 0:
-        >>>     model.load_state_dict(network_bound_weights_download_fn(s3_weights_path))
-        >>> dist.barrir()
-        >>> sync_model_states(model) # sync rank0 weights to other ranks
-    """
-    if not dist.is_available() or not dist.is_initialized():
-        return
-    if process_group is None:
-        process_group = _get_default_group()
-    if not params_and_buffers_to_ignore:
-        params_and_buffers_to_ignore = set()
-
-    log.info(
-        f"Synchronizing model states from rank {src} to all ranks in process group {get_process_group_ranks(process_group)}."
-    )
-
-    # Build tuple of (module, parameter) for all parameters that require grads.
-    modules_and_parameters = [
-        (module, parameter)
-        for module_name, module in model.named_modules()
-        for parameter in [
-            param
-            # Note that we access module.named_parameters instead of
-            # parameters(module). parameters(module) is only needed in the
-            # single-process multi device case, where it accesses replicated
-            # parameters through _former_parameters.
-            for param_name, param in module.named_parameters(recurse=False)
-            if f"{module_name}.{param_name}" not in params_and_buffers_to_ignore
-            # if param.requires_grad
-            # and f"{module_name}.{param_name}" not in params_and_buffers_to_ignore
-        ]
-    ]
-
-    # Deduplicate any parameters that might be shared across child modules.
-    memo = set()
-    modules_and_parameters = [
-        # "p not in memo" is the deduplication check.
-        # "not memo.add(p)" is always True, and it's only there to cause "add(p)" if needed.
-        (m, p)
-        for m, p in modules_and_parameters
-        if p not in memo and not memo.add(p)  # type: ignore[func-returns-value]
-    ]
-
-    # Build list of parameters.
-    parameters = [parameter for _, parameter in modules_and_parameters]
-    if len(parameters) == 0:
-        return
-
-    _verify_param_shape_across_processes(process_group, parameters)
-
-    _sync_module_states(
-        module=model,
-        process_group=process_group,
-        broadcast_bucket_size=int(250 * 1024 * 1024),
-        src=src,
-        params_and_buffers_to_ignore=params_and_buffers_to_ignore,
-        broadcast_buffers=broadcast_buffers,
-    )

lyra_2/_ext/imaginaire/utils/easy_io/__init__.py

@@ -1,14 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.

lyra_2/_ext/imaginaire/utils/easy_io/backends/__init__.py

@@ -1,34 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from lyra_2._ext.imaginaire.utils.easy_io.backends.base_backend import BaseStorageBackend
-from lyra_2._ext.imaginaire.utils.easy_io.backends.boto3_backend import Boto3Backend
-from lyra_2._ext.imaginaire.utils.easy_io.backends.http_backend import HTTPBackend
-from lyra_2._ext.imaginaire.utils.easy_io.backends.local_backend import LocalBackend
-from lyra_2._ext.imaginaire.utils.easy_io.backends.registry_utils import (
-    backends,
-    prefix_to_backends,
-    register_backend,
-)
-
-__all__ = [
-    "BaseStorageBackend",
-    "LocalBackend",
-    "HTTPBackend",
-    "Boto3Backend",
-    "register_backend",
-    "backends",
-    "prefix_to_backends",
-]

lyra_2/_ext/imaginaire/utils/easy_io/backends/auto_auth.py

@@ -1,64 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import contextlib
-import json
-from typing import Any, Optional
-
-from lyra_2._ext.imaginaire.utils import log
-from lyra_2._ext.imaginaire.utils.env_parsers.cred_env_parser import CRED_ENVS, CRED_ENVS_DICT
-
-DEPLOYMENT_ENVS = ["prod", "dev", "stg"]
-
-
-# context manger to open a file or read from env variable
-@contextlib.contextmanager
-def open_auth(s3_credential_path: Optional[Any], mode: str):
-    if not s3_credential_path:
-        log.info(f"No credential file provided {s3_credential_path}.")
-        yield None
-        return
-
-    name = s3_credential_path.split("/")[-1].split(".")[0]
-    if not name:
-        raise ValueError(f"Could not parse into env var: {s3_credential_path}")
-    cred_env_name = f"PROD_{name.upper()}"
-
-    if CRED_ENVS.APP_ENV in DEPLOYMENT_ENVS and cred_env_name in CRED_ENVS_DICT:
-        object_storage_config = get_creds_from_env(cred_env_name)
-        log.info(f"using ENV vars for {cred_env_name}")
-        yield object_storage_config
-    else:
-        log.info(f"using credential file: {s3_credential_path}")
-        with open(s3_credential_path, mode) as f:
-            yield f
-
-
-def get_creds_from_env(cred_env_name: str) -> dict[str, str]:
-    try:
-        object_storage_config = CRED_ENVS_DICT[cred_env_name]
-    except KeyError:
-        raise ValueError(f"Could not find {cred_env_name} in CRED_ENVS")
-    empty_args = {key.upper() for key in object_storage_config if object_storage_config[key] == ""}
-    if empty_args:
-        raise ValueError(f"Some required environment variable(s) were not provided for {cred_env_name}", empty_args)
-    return object_storage_config
-
-
-def json_load_auth(f):
-    if CRED_ENVS.APP_ENV in DEPLOYMENT_ENVS:
-        return f if f else {}
-    else:
-        return json.load(f)

lyra_2/_ext/imaginaire/utils/easy_io/backends/base_backend.py

@@ -1,60 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import os
-import os.path as osp
-from abc import ABCMeta, abstractmethod
-
-
-def mkdir_or_exist(dir_name, mode=0o777):
-    if dir_name == "":
-        return
-    dir_name = osp.expanduser(dir_name)
-    os.makedirs(dir_name, mode=mode, exist_ok=True)
-
-
-def has_method(obj, method):
-    return hasattr(obj, method) and callable(getattr(obj, method))
-
-
-class BaseStorageBackend(metaclass=ABCMeta):
-    """Abstract class of storage backends.
-
-    All backends need to implement two apis: :meth:`get()` and
-    :meth:`get_text()`.
-
-    - :meth:`get()` reads the file as a byte stream.
-    - :meth:`get_text()` reads the file as texts.
-    """
-
-    # a flag to indicate whether the backend can create a symlink for a file
-    # This attribute will be deprecated in future.
-    _allow_symlink = False
-
-    @property
-    def allow_symlink(self):
-        return self._allow_symlink
-
-    @property
-    def name(self):
-        return self.__class__.__name__
-
-    @abstractmethod
-    def get(self, filepath):
-        pass
-
-    @abstractmethod
-    def get_text(self, filepath):
-        pass

lyra_2/_ext/imaginaire/utils/easy_io/backends/boto3_backend.py

@@ -1,841 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import io
-import os
-import re
-import tempfile
-from contextlib import contextmanager
-from pathlib import Path
-from shutil import SameFileError
-from typing import Generator, Iterator, Optional, Tuple, Union
-
-from lyra_2._ext.imaginaire.utils import log
-from lyra_2._ext.imaginaire.utils.easy_io.backends.base_backend import (
-    BaseStorageBackend,
-    has_method,
-    mkdir_or_exist,
-)
-from lyra_2._ext.imaginaire.utils.easy_io.backends.boto3_client import Boto3Client
-
-
-class Boto3Backend(BaseStorageBackend):
-    """boto3 storage backend (for internal usage).
-
-    Boto3Backend supports reading and writing data to multiple clusters.
-    If the file path contains the cluster name, Boto3Backend will read data
-    from specified cluster or write data to it. Otherwise, Boto3Backend will
-    access the default cluster.
-
-    Args:
-        path_mapping (dict, optional): Path mapping dict from local path to
-            Boto3 path. When ``path_mapping={'src': 'dst'}``, ``src`` in
-            ``filepath`` will be replaced by ``dst``. Defaults to None.
-        s3_credential_path (str, optional): Config path of Boto3 client. Default: None.
-            `New in version 0.3.3`.
-
-    Examples:
-        >>> backend = Boto3Backend()
-        >>> filepath1 = 's3://path/of/file'
-        >>> filepath2 = 'cluster-name:s3://path/of/file'
-        >>> backend.get(filepath1)  # get data from default cluster
-        >>> client.get(filepath2)  # get data from 'cluster-name' cluster
-    """
-
-    def __init__(
-        self,
-        s3_credential_path: str = "",
-        path_mapping: Optional[dict] = None,
-    ):
-        self._client = Boto3Client(s3_credential_path=s3_credential_path)
-        assert isinstance(path_mapping, dict) or path_mapping is None
-        self.path_mapping = path_mapping
-        if path_mapping:
-            for k, v in path_mapping.items():
-                log.critical(f"Path mapping: {k} -> {v}", rank0_only=False)
-
-    def _map_path(self, filepath: Union[str, Path]) -> str:
-        """Map ``filepath`` to a string path whose prefix will be replaced by
-        :attr:`self.path_mapping`.
-
-        Args:
-            filepath (str or Path): Path to be mapped.
-        """
-        filepath = str(filepath)
-        if self.path_mapping is not None:
-            for k, v in self.path_mapping.items():
-                filepath = filepath.replace(k, v, 1)
-        return filepath
-
-    def _format_path(self, filepath: str) -> str:
-        """Convert a ``filepath`` to standard format of s3 oss.
-
-        If the ``filepath`` is concatenated by ``os.path.join``, in a Windows
-        environment, the ``filepath`` will be the format of
-        's3://bucket_name\\image.jpg'. By invoking :meth:`_format_path`, the
-        above ``filepath`` will be converted to 's3://bucket_name/image.jpg'.
-
-        Args:
-            filepath (str): Path to be formatted.
-        """
-        return re.sub(r"\\+", "/", filepath)
-
-    def _replace_prefix(self, filepath: Union[str, Path]) -> str:
-        filepath = str(filepath)
-        return filepath
-        # return filepath.replace('s3://', 's3://')
-
-    def get(self, filepath: Union[str, Path]) -> bytes:
-        """Read bytes from a given ``filepath`` with 'rb' mode.
-
-        Args:
-            filepath (str or Path): Path to read data.
-
-        Returns:
-            bytes: Return bytes read from filepath.
-
-        Examples:
-            >>> backend = Boto3Backend()
-            >>> filepath = 's3://path/of/file'
-            >>> backend.get(filepath)
-            b'hello world'
-        """
-        filepath = self._map_path(filepath)
-        filepath = self._format_path(filepath)
-        filepath = self._replace_prefix(filepath)
-        value = self._client.get(filepath)
-        return value
-
-    def get_text(
-        self,
-        filepath: Union[str, Path],
-        encoding: str = "utf-8",
-    ) -> str:
-        """Read text from a given ``filepath`` with 'r' mode.
-
-        Args:
-            filepath (str or Path): Path to read data.
-            encoding (str): The encoding format used to open the ``filepath``.
-                Defaults to 'utf-8'.
-
-        Returns:
-            str: Expected text reading from ``filepath``.
-
-        Examples:
-            >>> backend = Boto3Backend()
-            >>> filepath = 's3://path/of/file'
-            >>> backend.get_text(filepath)
-            'hello world'
-        """
-        return str(self.get(filepath), encoding=encoding)
-
-    def put(self, obj: Union[bytes, io.BytesIO], filepath: Union[str, Path]) -> None:
-        """Write bytes to a given ``filepath``.
-
-        Args:
-            obj (bytes): Data to be saved.
-            filepath (str or Path): Path to write data.
-
-        Examples:
-            >>> backend = Boto3Backend()
-            >>> filepath = 's3://path/of/file'
-            >>> backend.put(b'hello world', filepath)
-        """
-        filepath = self._map_path(filepath)
-        filepath = self._format_path(filepath)
-        filepath = self._replace_prefix(filepath)
-        self._client.put(obj, filepath)
-
-    def fast_put(self, obj: Union[bytes, io.BytesIO], filepath: Union[str, Path], num_processes: int = 32) -> None:
-        """Write bytes to a given ``filepath`` with multiple processes and async"""
-        assert num_processes > 1
-        filepath = self._map_path(filepath)
-        filepath = self._format_path(filepath)
-        filepath = self._replace_prefix(filepath)
-        self._client.fast_put(obj, filepath, num_processes=num_processes)
-
-    def put_text(
-        self,
-        obj: str,
-        filepath: Union[str, Path],
-        encoding: str = "utf-8",
-    ) -> None:
-        """Write text to a given ``filepath``.
-
-        Args:
-            obj (str): Data to be written.
-            filepath (str or Path): Path to write data.
-            encoding (str): The encoding format used to encode the ``obj``.
-                Defaults to 'utf-8'.
-
-        Examples:
-            >>> backend = Boto3Backend()
-            >>> filepath = 's3://path/of/file'
-            >>> backend.put_text('hello world', filepath)
-        """
-        self.put(bytes(obj, encoding=encoding), filepath)
-
-    def exists(self, filepath: Union[str, Path]) -> bool:
-        """Check whether a file path exists.
-
-        Args:
-            filepath (str or Path): Path to be checked whether exists.
-
-        Returns:
-            bool: Return ``True`` if ``filepath`` exists, ``False`` otherwise.
-
-        Examples:
-            >>> backend = Boto3Backend()
-            >>> filepath = 's3://path/of/file'
-            >>> backend.exists(filepath)
-            True
-        """
-        if not (has_method(self._client, "contains") and has_method(self._client, "isdir")):
-            raise NotImplementedError(
-                "Current version of Boto3 Python SDK has not supported "
-                "the `contains` and `isdir` methods, please use a higher"
-                "version or dev branch instead."
-            )
-
-        filepath = self._map_path(filepath)
-        filepath = self._format_path(filepath)
-        filepath = self._replace_prefix(filepath)
-        return self._client.contains(filepath) or self._client.isdir(filepath)
-
-    def isdir(self, filepath: Union[str, Path]) -> bool:
-        """Check whether a file path is a directory.
-
-        Args:
-            filepath (str or Path): Path to be checked whether it is a
-                directory.
-
-        Returns:
-            bool: Return ``True`` if ``filepath`` points to a directory,
-            ``False`` otherwise.
-
-        Examples:
-            >>> backend = Boto3Backend()
-            >>> filepath = 's3://path/of/dir'
-            >>> backend.isdir(filepath)
-            True
-        """
-        if not has_method(self._client, "isdir"):
-            raise NotImplementedError(
-                "Current version of Boto3 Python SDK has not supported "
-                "the `isdir` method, please use a higher version or dev"
-                " branch instead."
-            )
-
-        filepath = self._map_path(filepath)
-        filepath = self._format_path(filepath)
-        filepath = self._replace_prefix(filepath)
-        return self._client.isdir(filepath)
-
-    def isfile(self, filepath: Union[str, Path]) -> bool:
-        """Check whether a file path is a file.
-
-        Args:
-            filepath (str or Path): Path to be checked whether it is a file.
-
-        Returns:
-            bool: Return ``True`` if ``filepath`` points to a file, ``False``
-            otherwise.
-
-        Examples:
-            >>> backend = Boto3Backend()
-            >>> filepath = 's3://path/of/file'
-            >>> backend.isfile(filepath)
-            True
-        """
-        if not has_method(self._client, "contains"):
-            raise NotImplementedError(
-                "Current version of Boto3 Python SDK has not supported "
-                "the `contains` method, please use a higher version or "
-                "dev branch instead."
-            )
-
-        filepath = self._map_path(filepath)
-        filepath = self._format_path(filepath)
-        filepath = self._replace_prefix(filepath)
-        return self._client.contains(filepath)
-
-    def join_path(
-        self,
-        filepath: Union[str, Path],
-        *filepaths: Union[str, Path],
-    ) -> str:
-        r"""Concatenate all file paths.
-
-        Join one or more filepath components intelligently. The return value
-        is the concatenation of filepath and any members of \*filepaths.
-
-        Args:
-            filepath (str or Path): Path to be concatenated.
-
-        Returns:
-            str: The result after concatenation.
-
-        Examples:
-            >>> backend = Boto3Backend()
-            >>> filepath = 's3://path/of/file'
-            >>> backend.join_path(filepath, 'another/path')
-            's3://path/of/file/another/path'
-            >>> backend.join_path(filepath, '/another/path')
-            's3://path/of/file/another/path'
-        """
-        filepath = self._format_path(self._map_path(filepath))
-        if filepath.endswith("/"):
-            filepath = filepath[:-1]
-        formatted_paths = [filepath]
-        for path in filepaths:
-            formatted_path = self._format_path(self._map_path(path))
-            formatted_paths.append(formatted_path.lstrip("/"))
-
-        return "/".join(formatted_paths)
-
-    @contextmanager
-    def get_local_path(
-        self,
-        filepath: Union[str, Path],
-    ) -> Generator[Union[str, Path], None, None]:
-        """Download a file from ``filepath`` to a local temporary directory,
-        and return the temporary path.
-
-        ``get_local_path`` is decorated by :meth:`contxtlib.contextmanager`. It
-        can be called with ``with`` statement, and when exists from the
-        ``with`` statement, the temporary path will be released.
-
-        Args:
-            filepath (str or Path): Download a file from ``filepath``.
-
-        Yields:
-            Iterable[str]: Only yield one temporary path.
-
-        Examples:
-            >>> backend = Boto3Backend()
-            >>> # After existing from the ``with`` clause,
-            >>> # the path will be removed
-            >>> filepath = 's3://path/of/file'
-            >>> with backend.get_local_path(filepath) as path:
-            ...     # do something here
-        """
-        assert self.isfile(filepath)
-        try:
-            f = tempfile.NamedTemporaryFile(delete=False)
-            f.write(self.get(filepath))
-            f.close()
-            yield f.name
-        finally:
-            os.remove(f.name)
-
-    def copyfile(
-        self,
-        src: Union[str, Path],
-        dst: Union[str, Path],
-    ) -> str:
-        """Copy a file src to dst and return the destination file.
-
-        src and dst should have the same prefix. If dst specifies a directory,
-        the file will be copied into dst using the base filename from src. If
-        dst specifies a file that already exists, it will be replaced.
-
-        Args:
-            src (str or Path): A file to be copied.
-            dst (str or Path): Copy file to dst.
-
-        Returns:
-            str: The destination file.
-
-        Raises:
-            SameFileError: If src and dst are the same file, a SameFileError
-                will be raised.
-
-        Examples:
-            >>> backend = Boto3Backend()
-            >>> # dst is a file
-            >>> src = 's3://path/of/file'
-            >>> dst = 's3://path/of/file1'
-            >>> backend.copyfile(src, dst)
-            's3://path/of/file1'
-
-            >>> # dst is a directory
-            >>> dst = 's3://path/of/dir'
-            >>> backend.copyfile(src, dst)
-            's3://path/of/dir/file'
-        """
-        src = self._format_path(self._map_path(src))
-        dst = self._format_path(self._map_path(dst))
-        if self.isdir(dst):
-            dst = self.join_path(dst, src.split("/")[-1])
-
-        if src == dst:
-            raise SameFileError("src and dst should not be same")
-
-        self.put(self.get(src), dst)
-        return dst
-
-    def copytree(
-        self,
-        src: Union[str, Path],
-        dst: Union[str, Path],
-    ) -> str:
-        """Recursively copy an entire directory tree rooted at src to a
-        directory named dst and return the destination directory.
-
-        src and dst should have the same prefix.
-
-        Args:
-            src (str or Path): A directory to be copied.
-            dst (str or Path): Copy directory to dst.
-            backend_args (dict, optional): Arguments to instantiate the
-                prefix of uri corresponding backend. Defaults to None.
-
-        Returns:
-            str: The destination directory.
-
-        Raises:
-            FileExistsError: If dst had already existed, a FileExistsError will
-                be raised.
-
-        Examples:
-            >>> backend = Boto3Backend()
-            >>> src = 's3://path/of/dir'
-            >>> dst = 's3://path/of/dir1'
-            >>> backend.copytree(src, dst)
-            's3://path/of/dir1'
-        """
-        src = self._format_path(self._map_path(src))
-        dst = self._format_path(self._map_path(dst))
-
-        if self.exists(dst):
-            raise FileExistsError("dst should not exist")
-
-        for path in self.list_dir_or_file(src, list_dir=False, recursive=True):
-            src_path = self.join_path(src, path)
-            dst_path = self.join_path(dst, path)
-            self.put(self.get(src_path), dst_path)
-
-        return dst
-
-    def copyfile_from_local(
-        self,
-        src: Union[str, Path],
-        dst: Union[str, Path],
-    ) -> str:
-        """Upload a local file src to dst and return the destination file.
-
-        Args:
-            src (str or Path): A local file to be copied.
-            dst (str or Path): Copy file to dst.
-            backend_args (dict, optional): Arguments to instantiate the
-                prefix of uri corresponding backend. Defaults to None.
-
-        Returns:
-            str: If dst specifies a directory, the file will be copied into dst
-            using the base filename from src.
-
-        Examples:
-            >>> backend = Boto3Backend()
-            >>> # dst is a file
-            >>> src = 'path/of/your/file'
-            >>> dst = 's3://path/of/file1'
-            >>> backend.copyfile_from_local(src, dst)
-            's3://path/of/file1'
-
-            >>> # dst is a directory
-            >>> dst = 's3://path/of/dir'
-            >>> backend.copyfile_from_local(src, dst)
-            's3://path/of/dir/file'
-        """
-        dst = self._format_path(self._map_path(dst))
-        if self.isdir(dst):
-            dst = self.join_path(dst, os.path.basename(src))
-
-        with open(src, "rb") as f:
-            self.put(f.read(), dst)
-
-        return dst
-
-    def copytree_from_local(
-        self,
-        src: Union[str, Path],
-        dst: Union[str, Path],
-    ) -> str:
-        """Recursively copy an entire directory tree rooted at src to a
-        directory named dst and return the destination directory.
-
-        Args:
-            src (str or Path): A local directory to be copied.
-            dst (str or Path): Copy directory to dst.
-
-        Returns:
-            str: The destination directory.
-
-        Raises:
-            FileExistsError: If dst had already existed, a FileExistsError will
-                be raised.
-
-        Examples:
-            >>> backend = Boto3Backend()
-            >>> src = 'path/of/your/dir'
-            >>> dst = 's3://path/of/dir1'
-            >>> backend.copytree_from_local(src, dst)
-            's3://path/of/dir1'
-        """
-        dst = self._format_path(self._map_path(dst))
-        if self.exists(dst):
-            raise FileExistsError("dst should not exist")
-
-        src = str(src)
-
-        for cur_dir, _, files in os.walk(src):
-            for f in files:
-                src_path = os.path.join(cur_dir, f)
-                dst_path = self.join_path(dst, src_path.replace(src, ""))
-                self.copyfile_from_local(src_path, dst_path)
-
-        return dst
-
-    def copyfile_to_local(
-        self,
-        src: Union[str, Path],
-        dst: Union[str, Path],
-        dst_type: str,  # Choose from ["file", "dir"]
-    ) -> Union[str, Path]:
-        """Copy the file src to local dst and return the destination file.
-
-        If dst specifies a directory, the file will be copied into dst using
-        the base filename from src. If dst specifies a file that already
-        exists, it will be replaced.
-
-        Args:
-            src (str or Path): A file to be copied.
-            dst (str or Path): Copy file to to local dst.
-
-        Returns:
-            str: If dst specifies a directory, the file will be copied into dst
-            using the base filename from src.
-
-        Examples:
-            >>> backend = Boto3Backend()
-            >>> # dst is a file
-            >>> src = 's3://path/of/file'
-            >>> dst = 'path/of/your/file'
-            >>> backend.copyfile_to_local(src, dst)
-            'path/of/your/file'
-
-            >>> # dst is a directory
-            >>> dst = 'path/of/your/dir'
-            >>> backend.copyfile_to_local(src, dst)
-            'path/of/your/dir/file'
-        """
-        assert dst_type in ["file", "dir"]
-        # There is no good way to detect whether dst is a directory or a file, so we make dst_type required
-        if dst_type == "dir":
-            basename = os.path.basename(src)
-            if isinstance(dst, str):
-                dst = os.path.join(dst, basename)
-            else:
-                assert isinstance(dst, Path)
-                dst = dst / basename
-
-        # Create parent directory if it doesn't exist
-        parent_dir = os.path.dirname(dst)
-        os.makedirs(parent_dir, exist_ok=True)
-
-        try:
-            with open(dst, "wb") as f:
-                data = self.get(src)
-                f.write(data)
-        except Exception as e:
-            log.error(f"Failed to write file: {e}")
-            raise
-
-        return dst
-
-    def copytree_to_local(
-        self,
-        src: Union[str, Path],
-        dst: Union[str, Path],
-    ) -> Union[str, Path]:
-        """Recursively copy an entire directory tree rooted at src to a local
-        directory named dst and return the destination directory.
-
-        Args:
-            src (str or Path): A directory to be copied.
-            dst (str or Path): Copy directory to local dst.
-            backend_args (dict, optional): Arguments to instantiate the
-                prefix of uri corresponding backend. Defaults to None.
-
-        Returns:
-            str: The destination directory.
-
-        Examples:
-            >>> backend = Boto3Backend()
-            >>> src = 's3://path/of/dir'
-            >>> dst = 'path/of/your/dir'
-            >>> backend.copytree_to_local(src, dst)
-            'path/of/your/dir'
-        """
-        for path in self.list_dir_or_file(src, list_dir=False, recursive=True):
-            dst_path = os.path.join(dst, path)
-            mkdir_or_exist(os.path.dirname(dst_path))
-            with open(dst_path, "wb") as f:
-                f.write(self.get(self.join_path(src, path)))
-
-        return dst
-
-    def remove(self, filepath: Union[str, Path]) -> None:
-        """Remove a file.
-
-        Args:
-            filepath (str or Path): Path to be removed.
-
-        Raises:
-            FileNotFoundError: If filepath does not exist, an FileNotFoundError
-                will be raised.
-            IsADirectoryError: If filepath is a directory, an IsADirectoryError
-                will be raised.
-
-        Examples:
-            >>> backend = Boto3Backend()
-            >>> filepath = 's3://path/of/file'
-            >>> backend.remove(filepath)
-        """
-        if not has_method(self._client, "delete"):
-            raise NotImplementedError(
-                "Current version of Boto3 Python SDK has not supported "
-                "the `delete` method, please use a higher version or dev "
-                "branch instead."
-            )
-
-        if not self.exists(filepath):
-            raise FileNotFoundError(f"filepath {filepath} does not exist")
-
-        if self.isdir(filepath):
-            raise IsADirectoryError("filepath should be a file")
-
-        filepath = self._map_path(filepath)
-        filepath = self._format_path(filepath)
-        filepath = self._replace_prefix(filepath)
-        self._client.delete(filepath)
-
-    def rmtree(self, dir_path: Union[str, Path]) -> None:
-        """Recursively delete a directory tree.
-
-        Args:
-            dir_path (str or Path): A directory to be removed.
-
-        Examples:
-            >>> backend = Boto3Backend()
-            >>> dir_path = 's3://path/of/dir'
-            >>> backend.rmtree(dir_path)
-        """
-        for path in self.list_dir_or_file(dir_path, list_dir=False, recursive=True):
-            filepath = self.join_path(dir_path, path)
-            self.remove(filepath)
-
-    def copy_if_symlink_fails(
-        self,
-        src: Union[str, Path],
-        dst: Union[str, Path],
-    ) -> bool:
-        """Create a symbolic link pointing to src named dst.
-
-        Directly copy src to dst because PetrelBacekend does not support create
-        a symbolic link.
-
-        Args:
-            src (str or Path): A file or directory to be copied.
-            dst (str or Path): Copy a file or directory to dst.
-            backend_args (dict, optional): Arguments to instantiate the
-                prefix of uri corresponding backend. Defaults to None.
-
-        Returns:
-            bool: Return False because Boto3Backend does not support create
-            a symbolic link.
-
-        Examples:
-            >>> backend = Boto3Backend()
-            >>> src = 's3://path/of/file'
-            >>> dst = 's3://path/of/your/file'
-            >>> backend.copy_if_symlink_fails(src, dst)
-            False
-            >>> src = 's3://path/of/dir'
-            >>> dst = 's3://path/of/your/dir'
-            >>> backend.copy_if_symlink_fails(src, dst)
-            False
-        """
-        if self.isfile(src):
-            self.copyfile(src, dst)
-        else:
-            self.copytree(src, dst)
-        return False
-
-    def list_dir(self, dir_path: Union[str, Path]):
-        """List all folders in an S3 bucket with a given prefix.
-
-        Args:
-            dir_path (str | Path): Path of the directory.
-
-        Examples:
-            >>> backend = Boto3Backend()
-            >>> dir_path = 's3://path/of/dir'
-            >>> backend.list_dir(dir_path)
-        """
-        dir_path = self._map_path(dir_path)
-        dir_path = self._format_path(dir_path)
-        dir_path = self._replace_prefix(dir_path)
-        return self._client.ls_dir(dir_path)
-
-    def list_dir_or_file(  # pylint: disable=too-many-arguments
-        self,
-        dir_path: Union[str, Path],
-        list_dir: bool = True,
-        list_file: bool = True,
-        suffix: Optional[Union[str, Tuple[str]]] = None,
-        recursive: bool = False,
-    ) -> Iterator[str]:
-        """Scan a directory to find the interested directories or files in
-        arbitrary order.
-
-        Note:
-            Boto3 has no concept of directories but it simulates the directory
-            hierarchy in the filesystem through public prefixes. In addition,
-            if the returned path ends with '/', it means the path is a public
-            prefix which is a logical directory.
-
-        Note:
-            :meth:`list_dir_or_file` returns the path relative to ``dir_path``.
-            In addition, the returned path of directory will not contains the
-            suffix '/' which is consistent with other backends.
-
-        Args:
-            dir_path (str | Path): Path of the directory.
-            list_dir (bool): List the directories. Defaults to True.
-            list_file (bool): List the path of files. Defaults to True.
-            suffix (str or tuple[str], optional):  File suffix
-                that we are interested in. Defaults to None.
-            recursive (bool): If set to True, recursively scan the
-                directory. Defaults to False.
-
-        Yields:
-            Iterable[str]: A relative path to ``dir_path``.
-
-        Examples:
-            >>> backend = Boto3Backend()
-            >>> dir_path = 's3://path/of/dir'
-            >>> # list those files and directories in current directory
-            >>> for file_path in backend.list_dir_or_file(dir_path):
-            ...     print(file_path)
-            >>> # only list files
-            >>> for file_path in backend.list_dir_or_file(dir_path, list_dir=False):
-            ...     print(file_path)
-            >>> # only list directories
-            >>> for file_path in backend.list_dir_or_file(dir_path, list_file=False):
-            ...     print(file_path)
-            >>> # only list files ending with specified suffixes
-            >>> for file_path in backend.list_dir_or_file(dir_path, suffix='.txt'):
-            ...     print(file_path)
-            >>> # list all files and directory recursively
-            >>> for file_path in backend.list_dir_or_file(dir_path, recursive=True):
-            ...     print(file_path)
-        """  # noqa: E501
-        if not has_method(self._client, "list"):
-            raise NotImplementedError(
-                "Current version of Boto3 Python SDK has not supported "
-                "the `list` method, please use a higher version or dev"
-                " branch instead."
-            )
-
-        dir_path = self._map_path(dir_path)
-        dir_path = self._format_path(dir_path)
-        dir_path = self._replace_prefix(dir_path)
-        if list_dir and suffix is not None:
-            raise TypeError("`list_dir` should be False when `suffix` is not None")
-
-        if list_dir and not list_file and not recursive:
-            raise TypeError(
-                "Please use `list_dir` instead of `list_dir_or_file` when you only want to list the first level directories."
-            )
-
-        if (suffix is not None) and not isinstance(suffix, (str, tuple)):
-            raise TypeError("`suffix` must be a string or tuple of strings")
-
-        # Boto3's simulated directory hierarchy assumes that directory paths
-        # should end with `/`
-        if not dir_path.endswith("/"):
-            dir_path += "/"
-
-        root = dir_path
-
-        def _list_dir_or_file(dir_path, list_dir, list_file, suffix, recursive):
-            # Keep track of directories we've already yielded to avoid duplicates
-            yielded_dirs = set() if list_dir else None
-
-            for path in self._client.list(dir_path):
-                # All paths returned by S3 list are file paths, never directory paths
-                absolute_path = self.join_path(dir_path, path)
-                rel_path = absolute_path[len(root) :]
-
-                # If we want directories, extract directory prefixes from file paths
-                # boto3 client actually never return dir, it only return file paths
-                if list_dir and "/" in rel_path:
-                    if not recursive:
-                        # Non-recursive: only yield immediate child directory (first level)
-                        first_slash_pos = rel_path.find("/")
-                        immediate_child_dir = rel_path[:first_slash_pos]
-
-                        if immediate_child_dir not in yielded_dirs:
-                            yielded_dirs.add(immediate_child_dir)
-                            yield immediate_child_dir
-                    else:
-                        # Recursive: yield all directory levels
-                        path_parts = rel_path.split("/")[:-1]  # Exclude filename
-                        current_dir = ""
-                        for part in path_parts:
-                            if current_dir:
-                                current_dir += "/" + part
-                            else:
-                                current_dir = part
-
-                            if current_dir not in yielded_dirs:
-                                yielded_dirs.add(current_dir)
-                                yield current_dir
-
-                # Handle file listing
-                if (suffix is None or rel_path.endswith(suffix)) and list_file:
-                    yield rel_path
-
-        return _list_dir_or_file(dir_path, list_dir, list_file, suffix, recursive)
-
-    def generate_presigned_url(self, url: str, client_method: str = "get_object", expires_in: int = 3600) -> str:
-        """Generate the presigned url of video stream which can be passed to
-        mmcv.VideoReader. Now only work on Boto3 backend.
-
-        Note:
-            Now only work on Boto3 backend.
-
-        Args:
-            url (str): Url of video stream.
-            client_method (str): Method of client, 'get_object' or
-                'put_object'. Default: 'get_object'.
-            expires_in (int): expires, in seconds. Default: 3600.
-
-        Returns:
-            str: Generated presigned url.
-        """
-        raise NotImplementedError("generate_presigned_url is not supported in Boto3Backend")
-        return self._client.generate_presigned_url(url, client_method, expires_in)

lyra_2/_ext/imaginaire/utils/easy_io/backends/boto3_client.py

@@ -1,565 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import asyncio
-import concurrent.futures
-import io
-import os
-import time
-from math import ceil
-from multiprocessing import shared_memory
-from typing import Any, Dict, Generator, List, Tuple
-
-import boto3
-import numpy as np
-from botocore.config import Config as S3Config
-from botocore.exceptions import ClientError
-
-import lyra_2._ext.imaginaire.utils.easy_io.backends.auto_auth as auto
-from lyra_2._ext.imaginaire.utils import log
-from lyra_2._ext.imaginaire.utils.env_parsers.cred_env_parser import CRED_ENVS
-
-try:
-    import aioboto3
-    import aioboto3.session
-    from aiobotocore.config import AioConfig
-    from aiobotocore.session import AioSession
-except ImportError:
-    aioboto3 = None
-    AioSession = None
-
-MAX_RETRIES = 5
-RETRY_DELAY = 1  # seconds
-
-
-async def upload_single_part_async(
-    s3: AioSession, bucket: str, key: str, part_number: int, data: bytes, upload_id: str
-) -> Dict[str, Any]:
-    """
-    Uploads a single part of a file asynchronously to S3.
-
-    Args:
-        s3 (S3): The S3 client.
-        bucket (str): The S3 bucket name.
-        key (str): The S3 key (file path).
-        part_number (int): The part number of the upload.
-        data (bytes): The data to upload.
-        upload_id (str): The upload ID for the multipart upload.
-
-    Returns:
-        Dict[str, Any]: A dictionary containing the part number and ETag.
-    """
-    for attempt in range(MAX_RETRIES):
-        try:
-            response = await s3.upload_part(
-                Bucket=bucket, Key=key, PartNumber=part_number, UploadId=upload_id, Body=data
-            )
-            return {"PartNumber": part_number, "ETag": response["ETag"]}
-        except (ClientError, asyncio.TimeoutError, Exception) as e:
-            log.warning(f"Attempt {attempt + 1} failed for part {part_number}: {str(e)}", rank0_only=False)
-            if attempt < MAX_RETRIES - 1:
-                await asyncio.sleep(RETRY_DELAY * (2**attempt))  # Exponential backoff
-            else:
-                log.error(f"Failed to upload part {part_number} after {MAX_RETRIES} attempts", rank0_only=False)
-                raise
-
-
-async def upload_parts_async(
-    part_size: int,
-    part_numbers: range,
-    upload_id: str,
-    data: bytes,
-    bucket: str,
-    key: str,
-    client_config: Dict[str, Any],
-) -> List[Dict[str, Any]]:
-    """
-    Uploads multiple parts of a file asynchronously to S3.
-
-    Args:
-        part_size (int): The size of each part in bytes.
-        part_numbers (range): The range of part numbers to upload.
-        upload_id (str): The upload ID for the multipart upload.
-        data (bytes): The data to upload.
-        bucket (str): The S3 bucket name.
-        key (str): The S3 key (file path).
-        client_config (Dict[str, Any]): The S3 client configuration.
-
-    Returns:
-        List[Dict[str, Any]]: A list of dictionaries containing part numbers and ETags.
-    """
-    session = aioboto3.Session()
-    config = AioConfig(retries={"max_attempts": 3, "mode": "adaptive"}, connect_timeout=5, read_timeout=10)
-    start_idx = part_numbers[0]
-    async with session.client("s3", config=config, **client_config) as s3:
-        tasks = []
-        for part_number in part_numbers:
-            start = (part_number - start_idx) * part_size
-            end = min(start + part_size, len(data))
-            part_data = data[start:end]
-            tasks.append(upload_single_part_async(s3, bucket, key, part_number + 1, part_data, upload_id))
-
-        results = await asyncio.gather(*tasks, return_exceptions=True)
-
-        successful_parts = []
-        failed_parts = []
-        for part_number, result in enumerate(results, start=start_idx + 1):
-            if isinstance(result, Exception):
-                failed_parts.append(part_number)
-            else:
-                successful_parts.append(result)
-
-        if failed_parts:
-            log.error(f"Failed to upload parts: {failed_parts}", rank0_only=False)
-            raise Exception(f"Failed to upload {len(failed_parts)} parts")
-
-        successful_parts.sort(key=lambda part: part["PartNumber"])
-        return successful_parts
-
-
-def upload_parts_to_s3(args: Tuple[range, str, int, bytes, str, str, Dict[str, Any]]) -> List[Dict[str, Any]]:
-    """
-    Uploads parts of a file to S3 using a new event loop.
-
-    Args:
-        args (Tuple[range, str, int, bytes, str, str, Dict[str, Any]]): The arguments for uploading parts, including:
-            part_numbers (range): The range of part numbers to upload.
-            upload_id (str): The upload ID for the multipart upload.
-            part_size (int): The size of each part in bytes.
-            data (bytes): The data to upload.
-            bucket (str): The S3 bucket name.
-            key (str): The S3 key (file path).
-            client_config (Dict[str, Any]): The S3 client configuration.
-
-    Returns:
-        List[Dict[str, Any]]: A list of dictionaries containing part numbers and ETags.
-    """
-    part_numbers, upload_id, part_size, data, bucket, key, client_config = args
-    loop = asyncio.new_event_loop()
-    asyncio.set_event_loop(loop)
-    parts = loop.run_until_complete(
-        upload_parts_async(part_size, part_numbers, upload_id, data, bucket, key, client_config)
-    )
-    loop.close()
-    return parts
-
-
-async def download_single_part_async(
-    s3, bucket: str, key: str, part_number: int, start: int, end: int, shm_name: str, part_size: int
-) -> None:
-    """
-    Downloads a single part of a file asynchronously and writes it to shared memory.
-
-    Args:
-        s3 (S3): The S3 client.
-        bucket (str): The S3 bucket name.
-        key (str): The S3 key (file path).
-        part_number (int): The part number.
-        start (int): The start byte of the part.
-        end (int): The end byte of the part.
-        shm_name (str): The name of the shared memory block.
-        part_size (int): The size of each part in bytes.
-    """
-    for attempt in range(MAX_RETRIES):
-        try:
-            range_header = f"bytes={start}-{end}"
-            response = await s3.get_object(Bucket=bucket, Key=key, Range=range_header)
-            data = await response["Body"].read()
-
-            shm = shared_memory.SharedMemory(name=shm_name)
-            offset = part_number * part_size
-            shm.buf[offset : offset + len(data)] = data
-            shm.close()
-            return
-        except (ClientError, asyncio.TimeoutError, Exception) as e:
-            log.warning(f"Attempt {attempt + 1} failed for part {part_number}: {str(e)}", rank0_only=False)
-            if attempt < MAX_RETRIES - 1:
-                await asyncio.sleep(RETRY_DELAY * (2**attempt))  # Exponential backoff
-            else:
-                log.error(f"Failed to download part {part_number} after {MAX_RETRIES} attempts", rank0_only=False)
-                raise
-
-
-async def download_parts_async(
-    part_size: int, part_numbers: range, bucket: str, key: str, client_config: Dict[str, Any], shm_name: str
-) -> None:
-    """
-    Downloads multiple parts of a file asynchronously and writes them to shared memory.
-
-    Args:
-        part_size (int): The size of each part in bytes.
-        part_numbers (range): The range of part numbers to download.
-        bucket (str): The S3 bucket name.
-        key (str): The S3 key (file path).
-        client_config (Dict[str, Any]): The S3 client configuration.
-        shm_name (str): The name of the shared memory block.
-    """
-    session = aioboto3.Session()
-    config = AioConfig(retries={"max_attempts": 5, "mode": "adaptive"}, connect_timeout=10, read_timeout=30)
-    async with session.client("s3", config=config, **client_config) as s3:
-        tasks = [
-            download_single_part_async(
-                s3,
-                bucket,
-                key,
-                part_number,
-                part_number * part_size,
-                (part_number + 1) * part_size - 1,
-                shm_name,
-                part_size,
-            )
-            for part_number in part_numbers
-        ]
-        results = await asyncio.gather(*tasks, return_exceptions=True)
-        failed_parts = [part for part, result in zip(part_numbers, results) if isinstance(result, Exception)]
-
-        if failed_parts:
-            log.error(f"Failed to download parts: {failed_parts}", rank0_only=False)
-            raise Exception(f"Failed to download {len(failed_parts)} parts")
-
-
-def download_parts_to_s3(args: Tuple[range, int, str, str, Dict[str, Any], str]) -> bytes:
-    """
-    Downloads parts of a file using a new event loop.
-
-    Args:
-        args (Tuple[range, int, str, str, Dict[str, Any]]): The arguments for downloading parts, including:
-            part_numbers (range): The range of part numbers to download.
-            part_size (int): The size of each part in bytes.
-            bucket (str): The S3 bucket name.
-            key (str): The S3 key (file path).
-            client_config (Dict[str, Any]): The S3 client configuration.
-
-    Returns:
-        bytes: The combined file data from all downloaded parts.
-    """
-    part_numbers, part_size, bucket, key, client_config, shm_name = args
-    loop = asyncio.new_event_loop()
-    asyncio.set_event_loop(loop)
-    loop.run_until_complete(download_parts_async(part_size, part_numbers, bucket, key, client_config, shm_name))
-    loop.close()
-
-
-class Boto3Client:
-    def __init__(
-        self,
-        s3_credential_path: str,
-        max_attempt: int = 3,
-    ):
-        self.max_attempt = max_attempt
-        assert s3_credential_path, "s3_credential_path is required"
-        assert os.path.exists(s3_credential_path) or CRED_ENVS.APP_ENV in [
-            "prod",
-            "dev",
-            "stg",
-        ], f"Credential file not found: {s3_credential_path}"
-        with auto.open_auth(s3_credential_path, "r") as f:
-            conf = auto.json_load_auth(f)
-
-        s3_config = S3Config(
-            signature_version="s3v4",
-            s3={"addressing_style": "virtual"},
-            response_checksum_validation="when_required",
-            request_checksum_calculation="when_required",
-        )
-
-        self._client = boto3.client("s3", **conf, config=s3_config)
-        self._s3_cred_info = conf
-        self._mc_kv_store = None
-
-    def get(self, filepath):
-        filepath = self._check_path(filepath)
-
-        if self._mc_kv_store and self._mc_kv_store.available:
-            if self._mc_kv_store.has(filepath):
-                return self._mc_kv_store.get(filepath)
-
-        attempt = 0
-        while attempt < self.max_attempt:
-            try:
-                buffer = io.BytesIO()
-                self._client.download_fileobj(
-                    Bucket=filepath.split("/")[0],
-                    Key="/".join(filepath.split("/")[1:]),
-                    Fileobj=buffer,
-                )
-                buffer.seek(0)
-                if self._mc_kv_store and self._mc_kv_store.available:
-                    self._mc_kv_store.put(filepath, buffer.read())
-
-                return buffer.read()
-            except Exception as e:
-                attempt += 1
-                log.error(f"Got an exception: attempt={attempt} - {e} - {filepath}", rank0_only=False)
-
-        raise ConnectionError("Unable to read {} from. {} attempts tried.".format(filepath, attempt))
-
-    def _get_file_size(self, bucket, key, max_retries=10):
-        retries = 0
-        while retries < max_retries:
-            try:
-                # Try to get the file size
-                file_size = self._client.head_object(Bucket=bucket, Key=key)["ContentLength"]
-                return file_size  # Return file size if successful
-            except ClientError as e:
-                retries += 1
-                log.error(f"Attempt {retries} failed for s3://{bucket}/{key}: {e}", rank0_only=False)
-                if retries >= max_retries:
-                    raise  # Re-raise the exception after max retries
-                time.sleep(2)  # Wait for 2 seconds before retrying
-            except Exception as e:
-                retries += 1
-                log.error(
-                    f"Attempt {retries} failed for s3://{bucket}/{key}: due to an unexpected error: {e}",
-                    rank0_only=False,
-                )
-                if retries >= max_retries:
-                    raise  # Re-raise the exception after max retries
-                time.sleep(2)  # Wait for 2 seconds before retrying
-
-    def put(self, obj, filepath):
-        filepath = self._check_path(filepath)
-        bucket_name = filepath.split("/")[0]
-        key = "/".join(filepath.split("/")[1:])
-        attempt = 0
-        while attempt < self.max_attempt:
-            try:
-                # If obj is a string path to a local file, use upload_file instead
-                if isinstance(obj, str) and os.path.isfile(obj):
-                    self._client.upload_file(Filename=obj, Bucket=bucket_name, Key=key)
-                    return
-                if isinstance(obj, io.BytesIO):
-                    obj.seek(0)
-                    self._client.upload_fileobj(obj, Bucket=bucket_name, Key=key)
-                    return
-                if isinstance(obj, bytes):
-                    self._client.put_object(Body=obj, Bucket=bucket_name, Key=key)
-                    return
-                else:
-                    raise ValueError("Unsupported object type for upload")
-            except ClientError as e:
-                attempt += 1
-                log.error(f"Got an exception: attempt={attempt} - {e} - {filepath}", rank0_only=False)
-
-        raise ConnectionError("Unable to write {} to. {} attempts tried.".format(filepath, attempt))
-
-    def fast_put(self, obj, filepath, num_processes: int = 32):
-        assert aioboto3 is not None, "aioboto3 is required for fast_put"
-        original_filepath = filepath
-        filepath = self._check_path(filepath)
-        bucket = filepath.split("/")[0]
-        key = "/".join(filepath.split("/")[1:])
-        part_size = 16 * 1024 * 1024  # 16 MB part size
-
-        if isinstance(obj, bytes):
-            data = obj
-        elif isinstance(obj, str) and os.path.isfile(obj):
-            with open(obj, "rb") as f:
-                data = f.read()
-        elif isinstance(obj, io.BytesIO):
-            obj.seek(0)
-            data = obj.read()
-        else:
-            raise ValueError("Unsupported object type for upload")
-
-        file_size = len(data)
-        if file_size <= part_size * num_processes:
-            return self.put(data, original_filepath)
-        num_parts = ceil(file_size / part_size)
-        upload_id = self._client.create_multipart_upload(Bucket=bucket, Key=key)["UploadId"]
-
-        part_numbers = np.array_split(np.arange(num_parts), num_processes)
-
-        with concurrent.futures.ProcessPoolExecutor(max_workers=num_processes) as executor:
-            args = []
-            for i in range(num_processes):
-                cur_parts = part_numbers[i].tolist()
-                cur_data = data[cur_parts[0] * part_size : min(cur_parts[-1] * part_size + part_size, file_size)]
-                args.append((cur_parts, upload_id, part_size, cur_data, bucket, key, self._s3_cred_info))
-            results = executor.map(upload_parts_to_s3, args)
-            parts = []
-            for result in results:
-                parts.extend(result)
-
-        parts = sorted(parts, key=lambda part: part["PartNumber"])
-        self._client.complete_multipart_upload(
-            Bucket=bucket, Key=key, UploadId=upload_id, MultipartUpload={"Parts": parts}
-        )
-
-    def contains(self, filepath: str, max_retries=10) -> bool:
-        """
-        Checks if the specified object exists in the S3 bucket with retry logic for errors.
-
-        Args:
-            filepath (str): The s3 path of the file to check, must start with "s3://".
-
-        Returns:
-            bool: True if the object exists in the S3 bucket, False otherwise.
-
-        Raises:
-            ClientError: If an error response other than "404 Not Found" is returned from the S3 service.
-        """
-        filepath = self._check_path(filepath)
-        bucket = filepath.split("/")[0]
-        key = "/".join(filepath.split("/")[1:])
-
-        retries = 0
-        while retries < max_retries:
-            try:
-                # Try to check if the object exists
-                self._client.head_object(Bucket=bucket, Key=key)
-                return True  # Object exists
-            except ClientError as e:
-                if e.response["Error"]["Code"] == "404":
-                    return False  # Object does not exist
-                else:
-                    retries += 1
-                    print(f"Attempt {retries} failed with error: {e}")
-                    if retries >= max_retries:
-                        raise  # Re-raise the exception if max retries are reached
-                    time.sleep(2)  # Wait for 2 seconds before retrying
-            except Exception as e:
-                retries += 1
-                print(f"Attempt {retries} failed due to an unexpected error: {e}")
-                if retries >= max_retries:
-                    raise  # Re-raise the exception if max retries are reached
-                time.sleep(2)  # Wait for 2 seconds before retrying
-
-    def isdir(self, filepath: str, max_retries=10) -> bool:
-        """
-        Determines if the specified path corresponds to a directory in S3 with retry logic.
-
-        A directory in S3 is implied if there are any objects stored with the given prefix,
-        which means this function checks for the existence of any objects at or under the specified path.
-
-        Args:
-            filepath (str): The s3 path to check, must start with "s3://".
-
-        Returns:
-            bool: True if the specified path corresponds to a directory in S3, False otherwise.
-                Directories in S3 are not physical entities but are implied by object keys.
-
-        Raises:
-            ClientError: An error from the S3 API that isn't related to the absence of the directory
-                        (logged but not raised further).
-        """
-        filepath = self._check_path(filepath)
-        if not filepath.endswith("/"):
-            filepath += "/"
-
-        bucket = filepath.split("/")[0]
-        prefix = "/".join(filepath.split("/")[1:])
-
-        retries = 0
-        while retries < max_retries:
-            try:
-                # Try to check if any objects exist with the given prefix (i.e., directory in S3)
-                resp = self._client.list_objects_v2(Bucket=bucket, Prefix=prefix, Delimiter="/", MaxKeys=1)
-                # Check if any content or prefixes exist under the given path
-                return "CommonPrefixes" in resp or "Contents" in resp
-            except ClientError as e:
-                retries += 1
-                log.error(f"Attempt {retries} failed: {e}", rank0_only=False)
-                if retries >= max_retries:
-                    return False  # Return False if maximum retries are reached
-                time.sleep(2)  # Wait for 2 seconds before retrying
-            except Exception as e:
-                retries += 1
-                log.error(f"Attempt {retries} failed due to an unexpected error: {e}", rank0_only=False)
-                if retries >= max_retries:
-                    return False  # Return False if maximum retries are reached
-                time.sleep(2)  # Wait for 2 seconds before retrying
-
-    def delete(self, filepath):
-        filepath = self._check_path(filepath)
-        self._client.delete_object(Bucket=filepath.split("/")[0], Key="/".join(filepath.split("/")[1:]))
-
-    def ls_dir(self, filepath: str) -> Generator[str, None, None]:
-        """
-        List all folders in an S3 bucket with a given prefix.
-
-        Args:
-            filepath (str): The S3 path of the folder to list.
-
-        Yields:
-            str: The keys of the folders in the S3 bucket.
-        """
-        filepath = self._check_path(filepath)
-        bucket = filepath.split("/")[0]
-        prefix = "/".join(filepath.split("/")[1:])
-        continuation_token = None
-        if prefix and not prefix.endswith("/"):
-            prefix += "/"
-
-        while True:
-            if continuation_token:
-                resp = self._client.list_objects_v2(
-                    Bucket=bucket, Prefix=prefix, Delimiter="/", ContinuationToken=continuation_token
-                )
-            else:
-                resp = self._client.list_objects_v2(Bucket=bucket, Prefix=prefix, Delimiter="/")
-
-            if "CommonPrefixes" in resp:
-                for item in resp["CommonPrefixes"]:
-                    yield item["Prefix"][len(prefix) :]
-
-            # Check if there are more keys to retrieve
-            if resp.get("IsTruncated"):  # If IsTruncated is True, there are more keys
-                continuation_token = resp.get("NextContinuationToken")
-            else:
-                break
-
-    def list(self, filepath: str, exclude_prefix: str = None) -> Generator[str, None, None]:
-        """
-        List all keys in an S3 bucket with a given prefix, excluding files that start with
-        specified prefix.
-
-        Args:
-            filepath (str): The S3 path of the file to list.
-            exclude_prefix (str): Files starting with this prefix will be excluded from results.
-                                Defaults to "real".
-
-        Yields:
-            str: The keys of the files in the S3 bucket that don't start with exclude_prefix.
-        """
-        filepath = self._check_path(filepath)
-        bucket = filepath.split("/")[0]
-        prefix = "/".join(filepath.split("/")[1:])
-
-        continuation_token = None
-
-        while True:
-            if continuation_token:
-                resp = self._client.list_objects_v2(Bucket=bucket, Prefix=prefix, ContinuationToken=continuation_token)
-            else:
-                resp = self._client.list_objects_v2(Bucket=bucket, Prefix=prefix)
-
-            if "Contents" in resp:
-                for item in resp["Contents"]:
-                    key = item["Key"][len(prefix) :]
-                    # Skip files that start with the excluded prefix
-                    if exclude_prefix is None or not key.startswith(exclude_prefix):
-                        yield key
-
-            # Check if there are more keys to retrieve
-            if resp.get("IsTruncated"):  # If IsTruncated is True, there are more keys
-                continuation_token = resp.get("NextContinuationToken")
-            else:
-                break
-
-    def _check_path(self, filepath: str):
-        assert filepath.startswith("s3://")
-        filepath = filepath[5:]
-        return filepath

lyra_2/_ext/imaginaire/utils/easy_io/backends/http_backend.py

@@ -1,91 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import os
-import tempfile
-from contextlib import contextmanager
-from pathlib import Path
-from typing import Generator, Union
-from urllib.request import urlopen
-
-from lyra_2._ext.imaginaire.utils.easy_io.backends.base_backend import BaseStorageBackend
-
-
-class HTTPBackend(BaseStorageBackend):
-    """HTTP and HTTPS storage bachend."""
-
-    def get(self, filepath: str) -> bytes:
-        """Read bytes from a given ``filepath``.
-
-        Args:
-            filepath (str): Path to read data.
-
-        Returns:
-            bytes: Expected bytes object.
-
-        Examples:
-            >>> backend = HTTPBackend()
-            >>> backend.get('http://path/of/file')
-            b'hello world'
-        """
-        return urlopen(filepath).read()
-
-    def get_text(self, filepath, encoding="utf-8") -> str:
-        """Read text from a given ``filepath``.
-
-        Args:
-            filepath (str): Path to read data.
-            encoding (str): The encoding format used to open the ``filepath``.
-                Defaults to 'utf-8'.
-
-        Returns:
-            str: Expected text reading from ``filepath``.
-
-        Examples:
-            >>> backend = HTTPBackend()
-            >>> backend.get_text('http://path/of/file')
-            'hello world'
-        """
-        return urlopen(filepath).read().decode(encoding)
-
-    @contextmanager
-    def get_local_path(self, filepath: str) -> Generator[Union[str, Path], None, None]:
-        """Download a file from ``filepath`` to a local temporary directory,
-        and return the temporary path.
-
-        ``get_local_path`` is decorated by :meth:`contxtlib.contextmanager`. It
-        can be called with ``with`` statement, and when exists from the
-        ``with`` statement, the temporary path will be released.
-
-        Args:
-            filepath (str): Download a file from ``filepath``.
-
-        Yields:
-            Iterable[str]: Only yield one temporary path.
-
-        Examples:
-            >>> backend = HTTPBackend()
-            >>> # After existing from the ``with`` clause,
-            >>> # the path will be removed
-            >>> with backend.get_local_path('http://path/of/file') as path:
-            ...     # do something here
-        """
-        try:
-            f = tempfile.NamedTemporaryFile(delete=False)
-            f.write(self.get(filepath))
-            f.close()
-            yield f.name
-        finally:
-            os.remove(f.name)

lyra_2/_ext/imaginaire/utils/easy_io/backends/local_backend.py

@@ -1,550 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import io
-import os
-import os.path as osp
-import shutil
-from contextlib import contextmanager
-from pathlib import Path
-from typing import Generator, Iterator, Optional, Tuple, Union
-
-from lyra_2._ext.imaginaire.utils.easy_io.backends.base_backend import BaseStorageBackend, mkdir_or_exist
-
-
-class LocalBackend(BaseStorageBackend):
-    """Raw local storage backend."""
-
-    _allow_symlink = True
-
-    def get(self, filepath: Union[str, Path]) -> bytes:
-        """Read bytes from a given ``filepath`` with 'rb' mode.
-
-        Args:
-            filepath (str or Path): Path to read data.
-
-        Returns:
-            bytes: Expected bytes object.
-
-        Examples:
-            >>> backend = LocalBackend()
-            >>> filepath = '/path/of/file'
-            >>> backend.get(filepath)
-            b'hello world'
-        """
-        with open(filepath, "rb") as f:
-            value = f.read()
-        return value
-
-    def get_text(self, filepath: Union[str, Path], encoding: str = "utf-8") -> str:
-        """Read text from a given ``filepath`` with 'r' mode.
-
-        Args:
-            filepath (str or Path): Path to read data.
-            encoding (str): The encoding format used to open the ``filepath``.
-                Defaults to 'utf-8'.
-
-        Returns:
-            str: Expected text reading from ``filepath``.
-
-        Examples:
-            >>> backend = LocalBackend()
-            >>> filepath = '/path/of/file'
-            >>> backend.get_text(filepath)
-            'hello world'
-        """
-        with open(filepath, encoding=encoding) as f:
-            text = f.read()
-        return text
-
-    def put(self, obj: Union[bytes, io.BytesIO], filepath: Union[str, Path]) -> None:
-        """Write bytes to a given ``filepath`` with 'wb' mode.
-
-        Note:
-            ``put`` will create a directory if the directory of
-            ``filepath`` does not exist.
-
-        Args:
-            obj (bytes): Data to be written.
-            filepath (str or Path): Path to write data.
-
-        Examples:
-            >>> backend = LocalBackend()
-            >>> filepath = '/path/of/file'
-            >>> backend.put(b'hello world', filepath)
-        """
-        mkdir_or_exist(osp.dirname(filepath))
-        if isinstance(obj, io.BytesIO):
-            obj.seek(0)
-            obj = obj.getvalue()
-        with open(filepath, "wb") as f:
-            f.write(obj)
-
-    def put_text(self, obj: str, filepath: Union[str, Path], encoding: str = "utf-8") -> None:
-        """Write text to a given ``filepath`` with 'w' mode.
-
-        Note:
-            ``put_text`` will create a directory if the directory of
-            ``filepath`` does not exist.
-
-        Args:
-            obj (str): Data to be written.
-            filepath (str or Path): Path to write data.
-            encoding (str): The encoding format used to open the ``filepath``.
-                Defaults to 'utf-8'.
-
-        Examples:
-            >>> backend = LocalBackend()
-            >>> filepath = '/path/of/file'
-            >>> backend.put_text('hello world', filepath)
-        """
-        mkdir_or_exist(osp.dirname(filepath))
-        with open(filepath, "w", encoding=encoding) as f:
-            f.write(obj)
-
-    def exists(self, filepath: Union[str, Path]) -> bool:
-        """Check whether a file path exists.
-
-        Args:
-            filepath (str or Path): Path to be checked whether exists.
-
-        Returns:
-            bool: Return ``True`` if ``filepath`` exists, ``False`` otherwise.
-
-        Examples:
-            >>> backend = LocalBackend()
-            >>> filepath = '/path/of/file'
-            >>> backend.exists(filepath)
-            True
-        """
-        return osp.exists(filepath)
-
-    def isdir(self, filepath: Union[str, Path]) -> bool:
-        """Check whether a file path is a directory.
-
-        Args:
-            filepath (str or Path): Path to be checked whether it is a
-                directory.
-
-        Returns:
-            bool: Return ``True`` if ``filepath`` points to a directory,
-            ``False`` otherwise.
-
-        Examples:
-            >>> backend = LocalBackend()
-            >>> filepath = '/path/of/dir'
-            >>> backend.isdir(filepath)
-            True
-        """
-        return osp.isdir(filepath)
-
-    def isfile(self, filepath: Union[str, Path]) -> bool:
-        """Check whether a file path is a file.
-
-        Args:
-            filepath (str or Path): Path to be checked whether it is a file.
-
-        Returns:
-            bool: Return ``True`` if ``filepath`` points to a file, ``False``
-            otherwise.
-
-        Examples:
-            >>> backend = LocalBackend()
-            >>> filepath = '/path/of/file'
-            >>> backend.isfile(filepath)
-            True
-        """
-        return osp.isfile(filepath)
-
-    def join_path(self, filepath: Union[str, Path], *filepaths: Union[str, Path]) -> str:
-        r"""Concatenate all file paths.
-
-        Join one or more filepath components intelligently. The return value
-        is the concatenation of filepath and any members of \*filepaths.
-
-        Args:
-            filepath (str or Path): Path to be concatenated.
-
-        Returns:
-            str: The result of concatenation.
-
-        Examples:
-            >>> backend = LocalBackend()
-            >>> filepath1 = '/path/of/dir1'
-            >>> filepath2 = 'dir2'
-            >>> filepath3 = 'path/of/file'
-            >>> backend.join_path(filepath1, filepath2, filepath3)
-            '/path/of/dir/dir2/path/of/file'
-        """
-        return osp.join(filepath, *filepaths)
-
-    @contextmanager
-    def get_local_path(
-        self,
-        filepath: Union[str, Path],
-    ) -> Generator[Union[str, Path], None, None]:
-        """Only for unified API and do nothing.
-
-        Args:
-            filepath (str or Path): Path to be read data.
-            backend_args (dict, optional): Arguments to instantiate the
-                corresponding backend. Defaults to None.
-
-        Examples:
-            >>> backend = LocalBackend()
-            >>> with backend.get_local_path('s3://bucket/abc.jpg') as path:
-            ...     # do something here
-        """
-        yield filepath
-
-    def copyfile(
-        self,
-        src: Union[str, Path],
-        dst: Union[str, Path],
-    ) -> str:
-        """Copy a file src to dst and return the destination file.
-
-        src and dst should have the same prefix. If dst specifies a directory,
-        the file will be copied into dst using the base filename from src. If
-        dst specifies a file that already exists, it will be replaced.
-
-        Args:
-            src (str or Path): A file to be copied.
-            dst (str or Path): Copy file to dst.
-
-        Returns:
-            str: The destination file.
-
-        Raises:
-            SameFileError: If src and dst are the same file, a SameFileError
-                will be raised.
-
-        Examples:
-            >>> backend = LocalBackend()
-            >>> # dst is a file
-            >>> src = '/path/of/file'
-            >>> dst = '/path1/of/file1'
-            >>> # src will be copied to '/path1/of/file1'
-            >>> backend.copyfile(src, dst)
-            '/path1/of/file1'
-
-            >>> # dst is a directory
-            >>> dst = '/path1/of/dir'
-            >>> # src will be copied to '/path1/of/dir/file'
-            >>> backend.copyfile(src, dst)
-            '/path1/of/dir/file'
-        """
-        return shutil.copy(src, dst)
-
-    def copytree(
-        self,
-        src: Union[str, Path],
-        dst: Union[str, Path],
-    ) -> str:
-        """Recursively copy an entire directory tree rooted at src to a
-        directory named dst and return the destination directory.
-
-        src and dst should have the same prefix and dst must not already exist.
-
-
-        Args:
-            src (str or Path): A directory to be copied.
-            dst (str or Path): Copy directory to dst.
-
-        Returns:
-            str: The destination directory.
-
-        Raises:
-            FileExistsError: If dst had already existed, a FileExistsError will
-                be raised.
-
-        Examples:
-            >>> backend = LocalBackend()
-            >>> src = '/path/of/dir1'
-            >>> dst = '/path/of/dir2'
-            >>> backend.copytree(src, dst)
-            '/path/of/dir2'
-        """
-        return shutil.copytree(src, dst)
-
-    def copyfile_from_local(
-        self,
-        src: Union[str, Path],
-        dst: Union[str, Path],
-    ) -> str:
-        """Copy a local file src to dst and return the destination file. Same
-        as :meth:`copyfile`.
-
-        Args:
-            src (str or Path): A local file to be copied.
-            dst (str or Path): Copy file to dst.
-
-        Returns:
-            str: If dst specifies a directory, the file will be copied into dst
-            using the base filename from src.
-
-        Raises:
-            SameFileError: If src and dst are the same file, a SameFileError
-                will be raised.
-
-        Examples:
-            >>> backend = LocalBackend()
-            >>> # dst is a file
-            >>> src = '/path/of/file'
-            >>> dst = '/path1/of/file1'
-            >>> # src will be copied to '/path1/of/file1'
-            >>> backend.copyfile_from_local(src, dst)
-            '/path1/of/file1'
-
-            >>> # dst is a directory
-            >>> dst = '/path1/of/dir'
-            >>> # src will be copied to
-            >>> backend.copyfile_from_local(src, dst)
-            '/path1/of/dir/file'
-        """
-        return self.copyfile(src, dst)
-
-    def copytree_from_local(
-        self,
-        src: Union[str, Path],
-        dst: Union[str, Path],
-    ) -> str:
-        """Recursively copy an entire directory tree rooted at src to a
-        directory named dst and return the destination directory. Same as
-        :meth:`copytree`.
-
-        Args:
-            src (str or Path): A local directory to be copied.
-            dst (str or Path): Copy directory to dst.
-
-        Returns:
-            str: The destination directory.
-
-        Examples:
-            >>> backend = LocalBackend()
-            >>> src = '/path/of/dir1'
-            >>> dst = '/path/of/dir2'
-            >>> backend.copytree_from_local(src, dst)
-            '/path/of/dir2'
-        """
-        return self.copytree(src, dst)
-
-    def copyfile_to_local(
-        self,
-        src: Union[str, Path],
-        dst: Union[str, Path],
-        dst_type: Optional[str] = None,
-    ) -> str:
-        """Copy the file src to local dst and return the destination file. Same
-        as :meth:`copyfile`.
-
-        If dst specifies a directory, the file will be copied into dst using
-        the base filename from src. If dst specifies a file that already
-        exists, it will be replaced.
-
-        Args:
-            src (str or Path): A file to be copied.
-            dst (str or Path): Copy file to to local dst.
-
-        Returns:
-            str: If dst specifies a directory, the file will be copied into dst
-            using the base filename from src.
-
-        Examples:
-            >>> backend = LocalBackend()
-            >>> # dst is a file
-            >>> src = '/path/of/file'
-            >>> dst = '/path1/of/file1'
-            >>> # src will be copied to '/path1/of/file1'
-            >>> backend.copyfile_to_local(src, dst)
-            '/path1/of/file1'
-
-            >>> # dst is a directory
-            >>> dst = '/path1/of/dir'
-            >>> # src will be copied to
-            >>> backend.copyfile_to_local(src, dst)
-            '/path1/of/dir/file'
-        """
-        return self.copyfile(src, dst)
-
-    def copytree_to_local(
-        self,
-        src: Union[str, Path],
-        dst: Union[str, Path],
-    ) -> str:
-        """Recursively copy an entire directory tree rooted at src to a local
-        directory named dst and return the destination directory.
-
-        Args:
-            src (str or Path): A directory to be copied.
-            dst (str or Path): Copy directory to local dst.
-            backend_args (dict, optional): Arguments to instantiate the
-                prefix of uri corresponding backend. Defaults to None.
-
-        Returns:
-            str: The destination directory.
-
-        Examples:
-            >>> backend = LocalBackend()
-            >>> src = '/path/of/dir1'
-            >>> dst = '/path/of/dir2'
-            >>> backend.copytree_from_local(src, dst)
-            '/path/of/dir2'
-        """
-        return self.copytree(src, dst)
-
-    def remove(self, filepath: Union[str, Path]) -> None:
-        """Remove a file.
-
-        Args:
-            filepath (str or Path): Path to be removed.
-
-        Raises:
-            IsADirectoryError: If filepath is a directory, an IsADirectoryError
-                will be raised.
-            FileNotFoundError: If filepath does not exist, an FileNotFoundError
-                will be raised.
-
-        Examples:
-            >>> backend = LocalBackend()
-            >>> filepath = '/path/of/file'
-            >>> backend.remove(filepath)
-        """
-        if not self.exists(filepath):
-            raise FileNotFoundError(f"filepath {filepath} does not exist")
-
-        if self.isdir(filepath):
-            raise IsADirectoryError("filepath should be a file")
-
-        os.remove(filepath)
-
-    def rmtree(self, dir_path: Union[str, Path]) -> None:
-        """Recursively delete a directory tree.
-
-        Args:
-            dir_path (str or Path): A directory to be removed.
-
-        Examples:
-            >>> dir_path = '/path/of/dir'
-            >>> backend.rmtree(dir_path)
-        """
-        shutil.rmtree(dir_path)
-
-    def copy_if_symlink_fails(
-        self,
-        src: Union[str, Path],
-        dst: Union[str, Path],
-    ) -> bool:
-        """Create a symbolic link pointing to src named dst.
-
-        If failed to create a symbolic link pointing to src, directly copy src
-        to dst instead.
-
-        Args:
-            src (str or Path): Create a symbolic link pointing to src.
-            dst (str or Path): Create a symbolic link named dst.
-
-        Returns:
-            bool: Return True if successfully create a symbolic link pointing
-            to src. Otherwise, return False.
-
-        Examples:
-            >>> backend = LocalBackend()
-            >>> src = '/path/of/file'
-            >>> dst = '/path1/of/file1'
-            >>> backend.copy_if_symlink_fails(src, dst)
-            True
-            >>> src = '/path/of/dir'
-            >>> dst = '/path1/of/dir1'
-            >>> backend.copy_if_symlink_fails(src, dst)
-            True
-        """
-        try:
-            os.symlink(src, dst)
-            return True
-        except Exception:
-            if self.isfile(src):
-                self.copyfile(src, dst)
-            else:
-                self.copytree(src, dst)
-            return False
-
-    def list_dir_or_file(
-        self,
-        dir_path: Union[str, Path],
-        list_dir: bool = True,
-        list_file: bool = True,
-        suffix: Optional[Union[str, Tuple[str]]] = None,
-        recursive: bool = False,
-    ) -> Iterator[str]:
-        """Scan a directory to find the interested directories or files in
-        arbitrary order.
-
-        Note:
-            :meth:`list_dir_or_file` returns the path relative to ``dir_path``.
-
-        Args:
-            dir_path (str or Path): Path of the directory.
-            list_dir (bool): List the directories. Defaults to True.
-            list_file (bool): List the path of files. Defaults to True.
-            suffix (str or tuple[str], optional): File suffix that we are
-                interested in. Defaults to None.
-            recursive (bool): If set to True, recursively scan the directory.
-                Defaults to False.
-
-        Yields:
-            Iterable[str]: A relative path to ``dir_path``.
-
-        Examples:
-            >>> backend = LocalBackend()
-            >>> dir_path = '/path/of/dir'
-            >>> # list those files and directories in current directory
-            >>> for file_path in backend.list_dir_or_file(dir_path):
-            ...     print(file_path)
-            >>> # only list files
-            >>> for file_path in backend.list_dir_or_file(dir_path, list_dir=False):
-            ...     print(file_path)
-            >>> # only list directories
-            >>> for file_path in backend.list_dir_or_file(dir_path, list_file=False):
-            ...     print(file_path)
-            >>> # only list files ending with specified suffixes
-            >>> for file_path in backend.list_dir_or_file(dir_path, suffix='.txt'):
-            ...     print(file_path)
-            >>> # list all files and directory recursively
-            >>> for file_path in backend.list_dir_or_file(dir_path, recursive=True):
-            ...     print(file_path)
-        """  # noqa: E501
-        if list_dir and suffix is not None:
-            raise TypeError("`suffix` should be None when `list_dir` is True")
-
-        if (suffix is not None) and not isinstance(suffix, (str, tuple)):
-            raise TypeError("`suffix` must be a string or tuple of strings")
-
-        root = dir_path
-
-        def _list_dir_or_file(dir_path, list_dir, list_file, suffix, recursive):
-            for entry in os.scandir(dir_path):
-                if not entry.name.startswith(".") and entry.is_file():
-                    rel_path = osp.relpath(entry.path, root)
-                    if (suffix is None or rel_path.endswith(suffix)) and list_file:
-                        yield rel_path
-                elif osp.isdir(entry.path):
-                    if list_dir:
-                        rel_dir = osp.relpath(entry.path, root)
-                        yield rel_dir
-                    if recursive:
-                        yield from _list_dir_or_file(entry.path, list_dir, list_file, suffix, recursive)
-
-        return _list_dir_or_file(dir_path, list_dir, list_file, suffix, recursive)

lyra_2/_ext/imaginaire/utils/easy_io/backends/registry_utils.py

@@ -1,130 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import inspect
-from typing import Optional, Type, Union
-
-from lyra_2._ext.imaginaire.utils.easy_io.backends.base_backend import BaseStorageBackend
-from lyra_2._ext.imaginaire.utils.easy_io.backends.boto3_backend import Boto3Backend
-from lyra_2._ext.imaginaire.utils.easy_io.backends.http_backend import HTTPBackend
-from lyra_2._ext.imaginaire.utils.easy_io.backends.local_backend import LocalBackend
-
-backends: dict = {}
-prefix_to_backends: dict = {}
-
-
-def _register_backend(
-    name: str,
-    backend: Type[BaseStorageBackend],
-    force: bool = False,
-    prefixes: Union[str, list, tuple, None] = None,
-):
-    """Register a backend.
-
-    Args:
-        name (str): The name of the registered backend.
-        backend (BaseStorageBackend): The backend class to be registered,
-            which must be a subclass of :class:`BaseStorageBackend`.
-        force (bool): Whether to override the backend if the name has already
-            been registered. Defaults to False.
-        prefixes (str or list[str] or tuple[str], optional): The prefix
-            of the registered storage backend. Defaults to None.
-    """
-    global backends, prefix_to_backends
-
-    if not isinstance(name, str):
-        raise TypeError(f"the backend name should be a string, but got {type(name)}")
-
-    if not inspect.isclass(backend):
-        raise TypeError(f"backend should be a class, but got {type(backend)}")
-    if not issubclass(backend, BaseStorageBackend):
-        raise TypeError(f"backend {backend} is not a subclass of BaseStorageBackend")
-
-    if name in backends and not force:
-        raise ValueError(
-            f'{name} is already registered as a storage backend, add "force=True" if you want to override it'
-        )
-    backends[name] = backend
-
-    if prefixes is not None:
-        if isinstance(prefixes, str):
-            prefixes = [prefixes]
-        else:
-            assert isinstance(prefixes, (list, tuple))
-
-        for prefix in prefixes:
-            if prefix in prefix_to_backends and not force:
-                raise ValueError(
-                    f'{prefix} is already registered as a storage backend, add "force=True" if you want to override it'
-                )
-
-            prefix_to_backends[prefix] = backend
-
-
-def register_backend(
-    name: str,
-    backend: Optional[Type[BaseStorageBackend]] = None,
-    force: bool = False,
-    prefixes: Union[str, list, tuple, None] = None,
-):
-    """Register a backend.
-
-    Args:
-        name (str): The name of the registered backend.
-        backend (class, optional): The backend class to be registered,
-            which must be a subclass of :class:`BaseStorageBackend`.
-            When this method is used as a decorator, backend is None.
-            Defaults to None.
-        force (bool): Whether to override the backend if the name has already
-            been registered. Defaults to False.
-        prefixes (str or list[str] or tuple[str], optional): The prefix
-            of the registered storage backend. Defaults to None.
-
-    This method can be used as a normal method or a decorator.
-
-    Examples:
-
-        >>> class NewBackend(BaseStorageBackend):
-        ...     def get(self, filepath):
-        ...         return filepath
-        ...
-        ...     def get_text(self, filepath):
-        ...         return filepath
-        >>> register_backend('new', NewBackend)
-
-        >>> @register_backend('new')
-        ... class NewBackend(BaseStorageBackend):
-        ...     def get(self, filepath):
-        ...         return filepath
-        ...
-        ...     def get_text(self, filepath):
-        ...         return filepath
-    """
-    if backend is not None:
-        _register_backend(name, backend, force=force, prefixes=prefixes)
-        return
-
-    def _register(backend_cls):
-        _register_backend(name, backend_cls, force=force, prefixes=prefixes)
-        return backend_cls
-
-    return _register
-
-
-register_backend("local", LocalBackend, prefixes="")
-# To avoid breaking backward Compatibility, 's3' is also used as a
-# prefix for Boto3Backend
-register_backend("s3", Boto3Backend, prefixes=["s3"])
-register_backend("http", HTTPBackend, prefixes=["http", "https"])

lyra_2/_ext/imaginaire/utils/easy_io/easy_io.py

@@ -1,1085 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import json
-import warnings
-from contextlib import contextmanager
-from io import BytesIO, StringIO
-from pathlib import Path
-from typing import IO, Any, Generator, Iterator, Optional, Tuple, Union
-
-from lyra_2._ext.imaginaire.utils.easy_io.backends import backends, prefix_to_backends
-from lyra_2._ext.imaginaire.utils.easy_io.file_client import FileClient
-from lyra_2._ext.imaginaire.utils.easy_io.handlers import file_handlers
-
-backend_instances: dict = {}
-
-
-def is_filepath(filepath):
-    return isinstance(filepath, (str, Path))
-
-
-def _parse_uri_prefix(uri: Union[str, Path]) -> str:
-    """Parse the prefix of uri.
-
-    Args:
-        uri (str or Path): Uri to be parsed that contains the file prefix.
-
-    Examples:
-        >>> _parse_uri_prefix('/home/path/of/your/file')
-        ''
-        >>> _parse_uri_prefix('s3://path/of/your/file')
-        's3'
-        >>> _parse_uri_prefix('clusterName:s3://path/of/your/file')
-        's3'
-
-    Returns:
-        str: Return the prefix of uri if the uri contains '://'. Otherwise,
-        return ''.
-    """
-    assert is_filepath(uri)
-    uri = str(uri)
-    # if uri does not contains '://', the uri will be handled by
-    # LocalBackend by default
-    if "://" not in uri:
-        return ""
-    else:
-        prefix, _ = uri.split("://")
-        # In the case of Boto3Backend, the prefix may contain the cluster
-        # name like clusterName:s3://path/of/your/file
-        if ":" in prefix:
-            _, prefix = prefix.split(":")
-        return prefix
-
-
-def _get_file_backend(prefix: str, backend_args: dict):
-    """Return a file backend based on the prefix or backend_args.
-
-    Args:
-        prefix (str): Prefix of uri.
-        backend_args (dict): Arguments to instantiate the corresponding
-            backend.
-    """
-    # backend name has a higher priority
-    if "backend" in backend_args:
-        # backend_args should not be modified
-        backend_args_bak = backend_args.copy()
-        backend_name = backend_args_bak.pop("backend")
-        backend = backends[backend_name](**backend_args_bak)
-    else:
-        backend = prefix_to_backends[prefix](**backend_args)
-    return backend
-
-
-def set_s3_backend(
-    key: str = "s3:{}",
-    backend_args: Optional[dict] = None,
-):
-    """register s3 backend.
-
-    Args:
-        key str: The key to register the s3 backend. Defaults to s3.
-        backend_args (dict, optional): Arguments to instantiate the
-            corresponding backend. Defaults to None.
-    """
-    global backend_instances
-    if backend_args is None:
-        backend_args = {}
-    backend = _get_file_backend(key, backend_args)
-    backend_instances[key] = backend
-    return backend
-
-
-def get_file_backend(
-    uri: Union[str, Path, None] = None,
-    *,
-    backend_args: Optional[dict] = None,
-    enable_singleton: bool = False,
-    backend_key: Optional[str] = None,
-):
-    """Return a file backend based on the prefix of uri or backend_args.
-
-    Args:
-        uri (str or Path): Uri to be parsed that contains the file prefix.
-        backend_args (dict, optional): Arguments to instantiate the
-            corresponding backend. Defaults to None.
-        enable_singleton (bool): Whether to enable the singleton pattern.
-            If it is True, the backend created will be reused if the
-            signature is same with the previous one. Defaults to False.
-        backend_key: str: The key to register the backend. Defaults to None.
-
-    Returns:
-        BaseStorageBackend: Instantiated Backend object.
-
-    Examples:
-        >>> # get file backend based on the prefix of uri
-        >>> uri = 's3://path/of/your/file'
-        >>> backend = get_file_backend(uri)
-        >>> # get file backend based on the backend_args
-        >>> backend = get_file_backend(backend_args={'backend': 's3'})
-        >>> # backend name has a higher priority if 'backend' in backend_args
-        >>> backend = get_file_backend(uri, backend_args={'backend': 's3'})
-    """
-    global backend_instances
-    if backend_key is not None:
-        if backend_key in backend_instances:
-            return backend_instances[backend_key]
-
-    if backend_args is None:
-        backend_args = {}
-
-    if uri is None and "backend" not in backend_args and backend_key is None:
-        raise ValueError('uri should not be None when "backend" does not exist in backend_args and backend_key is None')
-
-    if uri is not None:
-        prefix = _parse_uri_prefix(uri)
-    else:
-        prefix = ""
-
-    if enable_singleton:
-        unique_key = f"{prefix}:{json.dumps(backend_args)}"
-        if unique_key in backend_instances:
-            return backend_instances[unique_key]
-
-        backend = _get_file_backend(prefix, backend_args)
-        backend_instances[unique_key] = backend
-        if backend_key is not None:
-            backend_instances[backend_key] = backend
-        return backend
-    else:
-        backend = _get_file_backend(prefix, backend_args)
-        return backend
-
-
-def get(
-    filepath: Union[str, Path],
-    backend_args: Optional[dict] = None,
-    backend_key: Optional[str] = None,
-) -> bytes:
-    """Read bytes from a given ``filepath`` with 'rb' mode.
-
-    Args:
-        filepath (str or Path): Path to read data.
-        backend_args (dict, optional): Arguments to instantiate the
-            corresponding backend. Defaults to None.
-        backend_key (str, optional): The key to get the backend from register.
-
-    Returns:
-        bytes: Expected bytes object.
-
-    Examples:
-        >>> filepath = '/path/of/file'
-        >>> get(filepath)
-        b'hello world'
-    """
-    backend = get_file_backend(
-        filepath,
-        backend_args=backend_args,
-        enable_singleton=True,
-        backend_key=backend_key,
-    )
-    return backend.get(filepath)
-
-
-def get_text(
-    filepath: Union[str, Path],
-    encoding="utf-8",
-    backend_args: Optional[dict] = None,
-    backend_key: Optional[str] = None,
-) -> str:
-    """Read text from a given ``filepath`` with 'r' mode.
-
-    Args:
-        filepath (str or Path): Path to read data.
-        encoding (str): The encoding format used to open the ``filepath``.
-            Defaults to 'utf-8'.
-        backend_args (dict, optional): Arguments to instantiate the
-            corresponding backend. Defaults to None.
-        backend_key (str, optional): The key to get the backend from register.
-
-    Returns:
-        str: Expected text reading from ``filepath``.
-
-    Examples:
-        >>> filepath = '/path/of/file'
-        >>> get_text(filepath)
-        'hello world'
-    """
-    backend = get_file_backend(
-        filepath,
-        backend_args=backend_args,
-        enable_singleton=True,
-        backend_key=backend_key,
-    )
-    return backend.get_text(filepath, encoding)
-
-
-def put(
-    obj: bytes,
-    filepath: Union[str, Path],
-    backend_args: Optional[dict] = None,
-    backend_key: Optional[str] = None,
-) -> None:
-    """Write bytes to a given ``filepath`` with 'wb' mode.
-
-    Note:
-        ``put`` should create a directory if the directory of
-        ``filepath`` does not exist.
-
-    Args:
-        obj (bytes): Data to be written.
-        filepath (str or Path): Path to write data.
-        backend_args (dict, optional): Arguments to instantiate the
-            corresponding backend. Defaults to None.
-        backend_key (str, optional): The key to get the backend from register.
-
-    Examples:
-        >>> filepath = '/path/of/file'
-        >>> put(b'hello world', filepath)
-    """
-    backend = get_file_backend(
-        filepath,
-        backend_args=backend_args,
-        enable_singleton=True,
-        backend_key=backend_key,
-    )
-    backend.put(obj, filepath)
-
-
-def put_text(
-    obj: str,
-    filepath: Union[str, Path],
-    backend_args: Optional[dict] = None,
-    backend_key: Optional[str] = None,
-) -> None:
-    """Write text to a given ``filepath`` with 'w' mode.
-
-    Note:
-        ``put_text`` should create a directory if the directory of
-        ``filepath`` does not exist.
-
-    Args:
-        obj (str): Data to be written.
-        filepath (str or Path): Path to write data.
-        encoding (str, optional): The encoding format used to open the
-            ``filepath``. Defaults to 'utf-8'.
-        backend_args (dict, optional): Arguments to instantiate the
-            corresponding backend. Defaults to None.
-        backend_key (str, optional): The key to get the backend from register.
-
-    Examples:
-        >>> filepath = '/path/of/file'
-        >>> put_text('hello world', filepath)
-    """
-    backend = get_file_backend(
-        filepath,
-        backend_args=backend_args,
-        enable_singleton=True,
-        backend_key=backend_key,
-    )
-    backend.put_text(obj, filepath)
-
-
-def exists(
-    filepath: Union[str, Path],
-    backend_args: Optional[dict] = None,
-    backend_key: Optional[str] = None,
-) -> bool:
-    """Check whether a file path exists.
-
-    Args:
-        filepath (str or Path): Path to be checked whether exists.
-        backend_args (dict, optional): Arguments to instantiate the
-            corresponding backend. Defaults to None.
-        backend_key (str, optional): The key to get the backend from register.
-
-    Returns:
-        bool: Return ``True`` if ``filepath`` exists, ``False`` otherwise.
-
-    Examples:
-        >>> filepath = '/path/of/file'
-        >>> exists(filepath)
-        True
-    """
-    backend = get_file_backend(
-        filepath,
-        backend_args=backend_args,
-        enable_singleton=True,
-        backend_key=backend_key,
-    )
-    return backend.exists(filepath)
-
-
-def isdir(
-    filepath: Union[str, Path],
-    backend_args: Optional[dict] = None,
-    backend_key: Optional[str] = None,
-) -> bool:
-    """Check whether a file path is a directory.
-
-    Args:
-        filepath (str or Path): Path to be checked whether it is a
-            directory.
-        backend_args (dict, optional): Arguments to instantiate the
-            corresponding backend. Defaults to None.
-        backend_key (str, optional): The key to get the backend from register.
-
-    Returns:
-        bool: Return ``True`` if ``filepath`` points to a directory,
-        ``False`` otherwise.
-
-    Examples:
-        >>> filepath = '/path/of/dir'
-        >>> isdir(filepath)
-        True
-    """
-    backend = get_file_backend(
-        filepath,
-        backend_args=backend_args,
-        enable_singleton=True,
-        backend_key=backend_key,
-    )
-    return backend.isdir(filepath)
-
-
-def isfile(
-    filepath: Union[str, Path],
-    backend_args: Optional[dict] = None,
-    backend_key: Optional[str] = None,
-) -> bool:
-    """Check whether a file path is a file.
-
-    Args:
-        filepath (str or Path): Path to be checked whether it is a file.
-        backend_args (dict, optional): Arguments to instantiate the
-            corresponding backend. Defaults to None.
-        backend_key (str, optional): The key to get the backend from register.
-
-    Returns:
-        bool: Return ``True`` if ``filepath`` points to a file, ``False``
-        otherwise.
-
-    Examples:
-        >>> filepath = '/path/of/file'
-        >>> isfile(filepath)
-        True
-    """
-    backend = get_file_backend(
-        filepath,
-        backend_args=backend_args,
-        enable_singleton=True,
-        backend_key=backend_key,
-    )
-    return backend.isfile(filepath)
-
-
-def join_path(
-    filepath: Union[str, Path],
-    *filepaths: Union[str, Path],
-    backend_args: Optional[dict] = None,
-    backend_key: Optional[str] = None,
-) -> Union[str, Path]:
-    r"""Concatenate all file paths.
-
-    Join one or more filepath components intelligently. The return value
-    is the concatenation of filepath and any members of \*filepaths.
-
-    Args:
-        filepath (str or Path): Path to be concatenated.
-        *filepaths (str or Path): Other paths to be concatenated.
-        backend_args (dict, optional): Arguments to instantiate the
-            corresponding backend. Defaults to None.
-        backend_key (str, optional): The key to get the backend from register.
-
-    Returns:
-        str: The result of concatenation.
-
-    Examples:
-        >>> filepath1 = '/path/of/dir1'
-        >>> filepath2 = 'dir2'
-        >>> filepath3 = 'path/of/file'
-        >>> join_path(filepath1, filepath2, filepath3)
-        '/path/of/dir/dir2/path/of/file'
-    """
-    backend = get_file_backend(
-        filepath,
-        backend_args=backend_args,
-        enable_singleton=True,
-        backend_key=backend_key,
-    )
-    return backend.join_path(filepath, *filepaths)
-
-
-@contextmanager
-def get_local_path(
-    filepath: Union[str, Path],
-    backend_args: Optional[dict] = None,
-    backend_key: Optional[str] = None,
-) -> Generator[Union[str, Path], None, None]:
-    """Download data from ``filepath`` and write the data to local path.
-
-    ``get_local_path`` is decorated by :meth:`contxtlib.contextmanager`. It
-    can be called with ``with`` statement, and when exists from the
-    ``with`` statement, the temporary path will be released.
-
-    Note:
-        If the ``filepath`` is a local path, just return itself and it will
-        not be released (removed).
-
-    Args:
-        filepath (str or Path): Path to be read data.
-        backend_args (dict, optional): Arguments to instantiate the
-            corresponding backend. Defaults to None.
-
-    Yields:
-        Iterable[str]: Only yield one path.
-
-    Examples:
-        >>> with get_local_path('s3://bucket/abc.jpg') as path:
-        ...     # do something here
-    """
-    backend = get_file_backend(
-        filepath,
-        backend_args=backend_args,
-        enable_singleton=True,
-        backend_key=backend_key,
-    )
-    with backend.get_local_path(str(filepath)) as local_path:
-        yield local_path
-
-
-def copyfile(
-    src: Union[str, Path],
-    dst: Union[str, Path],
-    backend_args: Optional[dict] = None,
-    backend_key: Optional[str] = None,
-) -> Union[str, Path]:
-    """Copy a file src to dst and return the destination file.
-
-    src and dst should have the same prefix. If dst specifies a directory,
-    the file will be copied into dst using the base filename from src. If
-    dst specifies a file that already exists, it will be replaced.
-
-    Args:
-        src (str or Path): A file to be copied.
-        dst (str or Path): Copy file to dst.
-        backend_args (dict, optional): Arguments to instantiate the
-            corresponding backend. Defaults to None.
-
-    Returns:
-        str: The destination file.
-
-    Raises:
-        SameFileError: If src and dst are the same file, a SameFileError will
-            be raised.
-
-    Examples:
-        >>> # dst is a file
-        >>> src = '/path/of/file'
-        >>> dst = '/path1/of/file1'
-        >>> # src will be copied to '/path1/of/file1'
-        >>> copyfile(src, dst)
-        '/path1/of/file1'
-
-        >>> # dst is a directory
-        >>> dst = '/path1/of/dir'
-        >>> # src will be copied to '/path1/of/dir/file'
-        >>> copyfile(src, dst)
-        '/path1/of/dir/file'
-    """
-    backend = get_file_backend(src, backend_args=backend_args, enable_singleton=True, backend_key=backend_key)
-    return backend.copyfile(src, dst)
-
-
-def copytree(
-    src: Union[str, Path],
-    dst: Union[str, Path],
-    backend_args: Optional[dict] = None,
-    backend_key: Optional[str] = None,
-) -> Union[str, Path]:
-    """Recursively copy an entire directory tree rooted at src to a directory
-    named dst and return the destination directory.
-
-    src and dst should have the same prefix and dst must not already exist.
-
-    Args:
-        src (str or Path): A directory to be copied.
-        dst (str or Path): Copy directory to dst.
-        backend_args (dict, optional): Arguments to instantiate the
-            corresponding backend. Defaults to None.
-        backend_key (str, optional): The key to get the backend from register.
-
-    Returns:
-        str: The destination directory.
-
-    Raises:
-        FileExistsError: If dst had already existed, a FileExistsError will be
-            raised.
-
-    Examples:
-        >>> src = '/path/of/dir1'
-        >>> dst = '/path/of/dir2'
-        >>> copytree(src, dst)
-        '/path/of/dir2'
-    """
-    backend = get_file_backend(src, backend_args=backend_args, enable_singleton=True, backend_key=backend_key)
-    return backend.copytree(src, dst)
-
-
-def copyfile_from_local(
-    src: Union[str, Path],
-    dst: Union[str, Path],
-    backend_args: Optional[dict] = None,
-    backend_key: Optional[str] = None,
-) -> Union[str, Path]:
-    """Copy a local file src to dst and return the destination file.
-
-    Note:
-        If the backend is the instance of LocalBackend, it does the same
-        thing with :func:`copyfile`.
-
-    Args:
-        src (str or Path): A local file to be copied.
-        dst (str or Path): Copy file to dst.
-        backend_args (dict, optional): Arguments to instantiate the
-            corresponding backend. Defaults to None.
-
-    Returns:
-        str: If dst specifies a directory, the file will be copied into dst
-        using the base filename from src.
-
-    Examples:
-        >>> # dst is a file
-        >>> src = '/path/of/file'
-        >>> dst = 's3://openmmlab/mmengine/file1'
-        >>> # src will be copied to 's3://openmmlab/mmengine/file1'
-        >>> copyfile_from_local(src, dst)
-        s3://openmmlab/mmengine/file1
-
-        >>> # dst is a directory
-        >>> dst = 's3://openmmlab/mmengine'
-        >>> # src will be copied to 's3://openmmlab/mmengine/file''
-        >>> copyfile_from_local(src, dst)
-        's3://openmmlab/mmengine/file'
-    """
-    backend = get_file_backend(dst, backend_args=backend_args, enable_singleton=True, backend_key=backend_key)
-    return backend.copyfile_from_local(src, dst)
-
-
-def copytree_from_local(
-    src: Union[str, Path],
-    dst: Union[str, Path],
-    backend_args: Optional[dict] = None,
-    backend_key: Optional[str] = None,
-) -> Union[str, Path]:
-    """Recursively copy an entire directory tree rooted at src to a directory
-    named dst and return the destination directory.
-
-    Note:
-        If the backend is the instance of LocalBackend, it does the same
-        thing with :func:`copytree`.
-
-    Args:
-        src (str or Path): A local directory to be copied.
-        dst (str or Path): Copy directory to dst.
-        backend_args (dict, optional): Arguments to instantiate the
-            corresponding backend. Defaults to None.
-
-    Returns:
-        str: The destination directory.
-
-    Examples:
-        >>> src = '/path/of/dir'
-        >>> dst = 's3://openmmlab/mmengine/dir'
-        >>> copyfile_from_local(src, dst)
-        's3://openmmlab/mmengine/dir'
-    """
-    backend = get_file_backend(dst, backend_args=backend_args, enable_singleton=True, backend_key=backend_key)
-    return backend.copytree_from_local(src, dst)
-
-
-def copyfile_to_local(
-    src: Union[str, Path],
-    dst: Union[str, Path],
-    dst_type: str,  # Choose from ["file", "dir"]
-    backend_args: Optional[dict] = None,
-    backend_key: Optional[str] = None,
-) -> Union[str, Path]:
-    """Copy the file src to local dst and return the destination file.
-
-    If dst specifies a directory, the file will be copied into dst using
-    the base filename from src. If dst specifies a file that already
-    exists, it will be replaced.
-
-    Note:
-        If the backend is the instance of LocalBackend, it does the same
-        thing with :func:`copyfile`.
-
-    Args:
-        src (str or Path): A file to be copied.
-        dst (str or Path): Copy file to to local dst.
-        backend_args (dict, optional): Arguments to instantiate the
-            corresponding backend. Defaults to None.
-
-    Returns:
-        str: If dst specifies a directory, the file will be copied into dst
-        using the base filename from src.
-
-    Examples:
-        >>> # dst is a file
-        >>> src = 's3://openmmlab/mmengine/file'
-        >>> dst = '/path/of/file'
-        >>> # src will be copied to '/path/of/file'
-        >>> copyfile_to_local(src, dst)
-        '/path/of/file'
-
-        >>> # dst is a directory
-        >>> dst = '/path/of/dir'
-        >>> # src will be copied to '/path/of/dir/file'
-        >>> copyfile_to_local(src, dst)
-        '/path/of/dir/file'
-    """
-    assert dst_type in ["file", "dir"]
-    Path(dst).parent.mkdir(parents=True, exist_ok=True)
-    backend = get_file_backend(src, backend_args=backend_args, enable_singleton=True, backend_key=backend_key)
-    return backend.copyfile_to_local(src, dst, dst_type=dst_type)
-
-
-def copytree_to_local(
-    src: Union[str, Path],
-    dst: Union[str, Path],
-    backend_args: Optional[dict] = None,
-    backend_key: Optional[str] = None,
-) -> Union[str, Path]:
-    """Recursively copy an entire directory tree rooted at src to a local
-    directory named dst and return the destination directory.
-
-    Note:
-        If the backend is the instance of LocalBackend, it does the same
-        thing with :func:`copytree`.
-
-    Args:
-        src (str or Path): A directory to be copied.
-        dst (str or Path): Copy directory to local dst.
-        backend_args (dict, optional): Arguments to instantiate the
-            corresponding backend. Defaults to None.
-
-    Returns:
-        str: The destination directory.
-
-    Examples:
-        >>> src = 's3://openmmlab/mmengine/dir'
-        >>> dst = '/path/of/dir'
-        >>> copytree_to_local(src, dst)
-        '/path/of/dir'
-    """
-    Path(dst).parent.mkdir(parents=True, exist_ok=True)
-    backend = get_file_backend(dst, backend_args=backend_args, enable_singleton=True, backend_key=backend_key)
-    return backend.copytree_to_local(src, dst)
-
-
-def remove(
-    filepath: Union[str, Path],
-    backend_args: Optional[dict] = None,
-    backend_key: Optional[str] = None,
-) -> None:
-    """Remove a file.
-
-    Args:
-        filepath (str, Path): Path to be removed.
-        backend_args (dict, optional): Arguments to instantiate the
-            corresponding backend. Defaults to None.
-
-    Raises:
-        FileNotFoundError: If filepath does not exist, an FileNotFoundError
-            will be raised.
-        IsADirectoryError: If filepath is a directory, an IsADirectoryError
-            will be raised.
-
-    Examples:
-        >>> filepath = '/path/of/file'
-        >>> remove(filepath)
-    """
-    backend = get_file_backend(
-        filepath,
-        backend_args=backend_args,
-        enable_singleton=True,
-        backend_key=backend_key,
-    )
-    backend.remove(filepath)
-
-
-def rmtree(
-    dir_path: Union[str, Path],
-    backend_args: Optional[dict] = None,
-    backend_key: Optional[str] = None,
-) -> None:
-    """Recursively delete a directory tree.
-
-    Args:
-        dir_path (str or Path): A directory to be removed.
-        backend_args (dict, optional): Arguments to instantiate the
-            corresponding backend. Defaults to None.
-
-    Examples:
-        >>> dir_path = '/path/of/dir'
-        >>> rmtree(dir_path)
-    """
-    backend = get_file_backend(
-        dir_path,
-        backend_args=backend_args,
-        enable_singleton=True,
-        backend_key=backend_key,
-    )
-    backend.rmtree(dir_path)
-
-
-def copy_if_symlink_fails(
-    src: Union[str, Path],
-    dst: Union[str, Path],
-    backend_args: Optional[dict] = None,
-    backend_key: Optional[str] = None,
-) -> bool:
-    """Create a symbolic link pointing to src named dst.
-
-    If failed to create a symbolic link pointing to src, directory copy src to
-    dst instead.
-
-    Args:
-        src (str or Path): Create a symbolic link pointing to src.
-        dst (str or Path): Create a symbolic link named dst.
-        backend_args (dict, optional): Arguments to instantiate the
-            corresponding backend. Defaults to None.
-
-    Returns:
-        bool: Return True if successfully create a symbolic link pointing to
-        src. Otherwise, return False.
-
-    Examples:
-        >>> src = '/path/of/file'
-        >>> dst = '/path1/of/file1'
-        >>> copy_if_symlink_fails(src, dst)
-        True
-        >>> src = '/path/of/dir'
-        >>> dst = '/path1/of/dir1'
-        >>> copy_if_symlink_fails(src, dst)
-        True
-    """
-    backend = get_file_backend(src, backend_args=backend_args, enable_singleton=True, backend_key=backend_key)
-    return backend.copy_if_symlink_fails(src, dst)
-
-
-def list_dir(
-    dir_path: Union[str, Path],
-    backend_args: Optional[dict] = None,
-    backend_key: Optional[str] = None,
-):
-    """List all folders in an S3 bucket with a given prefix.
-
-    Args:
-        dir_path (str | Path): Path of the directory.
-
-    Examples:
-        >>> dir_path = '/path/of/dir'
-        >>> for file_path in list_dir(dir_path):
-        ...     print(file_path)
-    """
-    if not dir_path.endswith("/"):
-        dir_path += "/"
-    backend = get_file_backend(
-        dir_path,
-        backend_args=backend_args,
-        enable_singleton=True,
-        backend_key=backend_key,
-    )
-
-    return backend.list_dir(dir_path)
-
-
-def list_dir_or_file(
-    dir_path: Union[str, Path],
-    list_dir: bool = True,
-    list_file: bool = True,
-    suffix: Optional[Union[str, Tuple[str]]] = None,
-    recursive: bool = False,
-    backend_args: Optional[dict] = None,
-    backend_key: Optional[str] = None,
-) -> Iterator[str]:
-    """Scan a directory to find the interested directories or files in
-    arbitrary order.
-
-    Note:
-        :meth:`list_dir_or_file` returns the path relative to ``dir_path``.
-
-    Args:
-        dir_path (str or Path): Path of the directory.
-        list_dir (bool): List the directories. Defaults to True.
-        list_file (bool): List the path of files. Defaults to True.
-        suffix (str or tuple[str], optional): File suffix that we are
-            interested in. Defaults to None.
-        recursive (bool): If set to True, recursively scan the directory.
-            Defaults to False.
-        backend_args (dict, optional): Arguments to instantiate the
-            corresponding backend. Defaults to None.
-
-    Yields:
-        Iterable[str]: A relative path to ``dir_path``.
-
-    Examples:
-        >>> dir_path = '/path/of/dir'
-        >>> for file_path in list_dir_or_file(dir_path):
-        ...     print(file_path)
-        >>> # list those files and directories in current directory
-        >>> for file_path in list_dir_or_file(dir_path):
-        ...     print(file_path)
-        >>> # only list files
-        >>> for file_path in list_dir_or_file(dir_path, list_dir=False):
-        ...     print(file_path)
-        >>> # only list directories
-        >>> for file_path in list_dir_or_file(dir_path, list_file=False):
-        ...     print(file_path)
-        >>> # only list files ending with specified suffixes
-        >>> for file_path in list_dir_or_file(dir_path, suffix='.txt'):
-        ...     print(file_path)
-        >>> # list all files and directory recursively
-        >>> for file_path in list_dir_or_file(dir_path, recursive=True):
-        ...     print(file_path)
-    """
-    backend = get_file_backend(
-        dir_path,
-        backend_args=backend_args,
-        enable_singleton=True,
-        backend_key=backend_key,
-    )
-    yield from backend.list_dir_or_file(dir_path, list_dir, list_file, suffix, recursive)
-
-
-def generate_presigned_url(
-    url: str,
-    client_method: str = "get_object",
-    expires_in: int = 3600,
-    backend_args: Optional[dict] = None,
-    backend_key: Optional[str] = None,
-) -> str:
-    """Generate the presigned url of video stream which can be passed to
-    mmcv.VideoReader. Now only work on s3 backend.
-
-    Note:
-        Now only work on s3 backend.
-
-    Args:
-        url (str): Url of video stream.
-        client_method (str): Method of client, 'get_object' or
-            'put_object'. Defaults to 'get_object'.
-        expires_in (int): expires, in seconds. Defaults to 3600.
-        backend_args (dict, optional): Arguments to instantiate the
-            corresponding backend. Defaults to None.
-
-    Returns:
-        str: Generated presigned url.
-    """
-    backend = get_file_backend(url, backend_args=backend_args, enable_singleton=True, backend_key=backend_key)
-    return backend.generate_presigned_url(url, client_method, expires_in)
-
-
-def load(
-    file: Union[str, Path, IO[Any]],
-    file_format: Optional[str] = None,
-    file_client_args: Optional[dict] = None,
-    fast_backend: bool = False,
-    backend_args: Optional[dict] = None,
-    backend_key: Optional[str] = None,
-    **kwargs,
-):
-    """Load data from json/yaml/pickle files.
-
-    This method provides a unified api for loading data from serialized files.
-
-    ``load`` supports loading data from serialized files those can be storaged
-    in different backends.
-
-    Args:
-        file (str or :obj:`Path` or file-like object): Filename or a file-like
-            object.
-        file_format (str, optional): If not specified, the file format will be
-            inferred from the file extension, otherwise use the specified one.
-            Currently supported formats include "json", "yaml/yml" and
-            "pickle/pkl".
-        file_client_args (dict, optional): Arguments to instantiate a
-            FileClient. See :class:`mmengine.fileio.FileClient` for details.
-            Defaults to None. It will be deprecated in future. Please use
-            ``backend_args`` instead.
-        fast_backend: bool: Whether to use multiprocess. Defaults to False.
-        backend_args (dict, optional): Arguments to instantiate the
-            prefix of uri corresponding backend. Defaults to None.
-            New in v0.2.0.
-
-    Examples:
-        >>> load('/path/of/your/file')  # file is storaged in disk
-        >>> load('https://path/of/your/file')  # file is storaged in Internet
-        >>> load('s3://path/of/your/file')  # file is storaged in s3
-
-    Returns:
-        The content from the file.
-    """
-    if isinstance(file, Path):
-        file = str(file)
-    if file_format is None and isinstance(file, str):
-        file_format = file.split(".")[-1]
-    # convert file_format to lower case
-    file_format = file_format.lower()
-    if file_format not in file_handlers:
-        raise TypeError(f"Unsupported format: {file_format}")
-
-    if file_client_args is not None:
-        warnings.warn(
-            '"file_client_args" will be deprecated in future. Please use "backend_args" instead',
-            DeprecationWarning,
-        )
-        if backend_args is not None:
-            raise ValueError('"file_client_args and "backend_args" cannot be set at the same time.')
-
-    handler = file_handlers[file_format]
-    if isinstance(file, str):
-        if file_client_args is not None:
-            file_client = FileClient.infer_client(file_client_args, file)
-            file_backend = file_client
-        else:
-            file_backend = get_file_backend(
-                file,
-                backend_args=backend_args,
-                backend_key=backend_key,
-                enable_singleton=True,
-            )
-
-        if handler.str_like:
-            with StringIO(file_backend.get_text(file)) as f:
-                obj = handler.load_from_fileobj(f, **kwargs)
-        else:
-            if fast_backend:
-                if hasattr(file_backend, "fast_get"):
-                    with BytesIO(file_backend.fast_get(file)) as f:
-                        obj = handler.load_from_fileobj(f, **kwargs)
-                else:
-                    warnings.warn(
-                        f"fast_backend is not supported by the backend, type {type(file_backend)} fallback to normal get"
-                    )
-                    with BytesIO(file_backend.get(file)) as f:
-                        obj = handler.load_from_fileobj(f, **kwargs)
-            else:
-                with BytesIO(file_backend.get(file)) as f:
-                    obj = handler.load_from_fileobj(f, **kwargs)
-    elif hasattr(file, "read"):
-        obj = handler.load_from_fileobj(file, **kwargs)
-    else:
-        raise TypeError('"file" must be a filepath str or a file-object')
-    return obj
-
-
-def dump(
-    obj: Any,
-    file: Union[str, Path, IO[Any], None] = None,
-    file_format: Optional[str] = None,
-    file_client_args: Optional[dict] = None,
-    fast_backend: bool = False,
-    backend_args: Optional[dict] = None,
-    backend_key: Optional[str] = None,
-    **kwargs,
-):
-    """Dump data to json/yaml/pickle strings or files.
-
-    This method provides a unified api for dumping data as strings or to files,
-    and also supports custom arguments for each file format.
-
-    ``dump`` supports dumping data as strings or to files which is saved to
-    different backends.
-
-    Args:
-        obj (any): The python object to be dumped.
-        file (str or :obj:`Path` or file-like object, optional): If not
-            specified, then the object is dumped to a str, otherwise to a file
-            specified by the filename or file-like object.
-        file_format (str, optional): Same as :func:`load`.
-        file_client_args (dict, optional): Arguments to instantiate a
-            FileClient. See :class:`mmengine.fileio.FileClient` for details.
-            Defaults to None. It will be deprecated in future. Please use
-            ``backend_args`` instead.
-        fast_backend: bool: Whether to use multiprocess. Defaults to False.
-        backend_args (dict, optional): Arguments to instantiate the
-            prefix of uri corresponding backend. Defaults to None.
-            New in v0.2.0.
-        backend_key: str: The key to register the backend. Defaults to None.
-
-    Examples:
-        >>> dump('hello world', '/path/of/your/file')  # disk
-        >>> dump('hello world', 's3://path/of/your/file')  # ceph or s3
-
-    Returns:
-        bool: True for success, False otherwise.
-    """
-    if isinstance(file, Path):
-        file = str(file)
-    if file_format is None:
-        if isinstance(file, str):
-            file_format = file.split(".")[-1]
-        elif file is None:
-            raise ValueError("file_format must be specified since file is None")
-    # convert file_format to lower case
-    file_format = file_format.lower()
-    if file_format not in file_handlers:
-        raise TypeError(f"Unsupported format: {file_format}")
-
-    if file_client_args is not None:
-        warnings.warn(
-            '"file_client_args" will be deprecated in future. Please use "backend_args" instead',
-            DeprecationWarning,
-        )
-        if backend_args is not None:
-            raise ValueError('"file_client_args" and "backend_args" cannot be set at the same time.')
-
-    handler = file_handlers[file_format]
-    if file is None:
-        return handler.dump_to_str(obj, **kwargs)
-    elif isinstance(file, str):
-        if file_client_args is not None:
-            file_client = FileClient.infer_client(file_client_args, file)
-            file_backend = file_client
-        else:
-            file_backend = get_file_backend(
-                file,
-                backend_args=backend_args,
-                backend_key=backend_key,
-                enable_singleton=True,
-            )
-
-        if handler.str_like:
-            with StringIO() as f:
-                handler.dump_to_fileobj(obj, f, **kwargs)
-                file_backend.put_text(f.getvalue(), file)
-        else:
-            with BytesIO() as f:
-                handler.dump_to_fileobj(obj, f, **kwargs)
-                if fast_backend:
-                    if hasattr(file_backend, "fast_put"):
-                        file_backend.fast_put(f, file)
-                    else:
-                        warnings.warn("fast_backend is not supported by the backend, fallback to normal put")
-                        file_backend.put(f, file)
-                else:
-                    file_backend.put(f, file)
-    elif hasattr(file, "write"):
-        handler.dump_to_fileobj(obj, file, **kwargs)
-    else:
-        raise TypeError('"file" must be a filename str or a file-object')

lyra_2/_ext/imaginaire/utils/easy_io/file_client.py

@@ -1,458 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import inspect
-from contextlib import contextmanager
-from pathlib import Path
-from typing import Any, Generator, Iterator, Optional, Tuple, Union
-
-from lyra_2._ext.imaginaire.utils.easy_io.backends import (
-    BaseStorageBackend,
-    Boto3Backend,
-    HTTPBackend,
-    LocalBackend,
-)
-
-
-def is_filepath(filepath):
-    return isinstance(filepath, (str, Path))
-
-
-class HardDiskBackend(LocalBackend):
-    """Raw hard disks storage backend."""
-
-    @property
-    def name(self):
-        return self.__class__.__name__
-
-
-class FileClient:
-    """A general file client to access files in different backends.
-
-    The client loads a file or text in a specified backend from its path
-    and returns it as a binary or text file. There are two ways to choose a
-    backend, the name of backend and the prefix of path. Although both of them
-    can be used to choose a storage backend, ``backend`` has a higher priority
-    that is if they are all set, the storage backend will be chosen by the
-    backend argument. If they are all `None`, the disk backend will be chosen.
-    Note that It can also register other backend accessor with a given name,
-    prefixes, and backend class. In addition, We use the singleton pattern to
-    avoid repeated object creation. If the arguments are the same, the same
-    object will be returned.
-
-    Warning:
-        `FileClient` will be deprecated in future. Please use io functions
-        in https://mmengine.readthedocs.io/en/latest/api/fileio.html#file-io
-
-    Args:
-        backend (str, optional): The storage backend type. Options are "disk",
-            "memcached", "lmdb", "http" and "s3". Defaults to None.
-        prefix (str, optional): The prefix of the registered storage backend.
-            Options are "s3", "http", "https". Defaults to None.
-
-    Examples:
-        >>> # only set backend
-        >>> file_client = FileClient(backend='s3')
-        >>> # only set prefix
-        >>> file_client = FileClient(prefix='s3')
-        >>> # set both backend and prefix but use backend to choose client
-        >>> file_client = FileClient(backend='s3', prefix='s3')
-        >>> # if the arguments are the same, the same object is returned
-        >>> file_client1 = FileClient(backend='s3')
-        >>> file_client1 is file_client
-        True
-
-    Attributes:
-        client (:obj:`BaseStorageBackend`): The backend object.
-    """
-
-    _backends = {
-        "disk": HardDiskBackend,
-        "s3": Boto3Backend,
-        "http": HTTPBackend,
-    }
-
-    _prefix_to_backends: dict = {
-        "s3": Boto3Backend,
-        "http": HTTPBackend,
-        "https": HTTPBackend,
-    }
-
-    _instances: dict = {}
-
-    client: Any
-
-    def __new__(cls, backend=None, prefix=None, **kwargs):
-        if backend is None and prefix is None:
-            backend = "disk"
-        if backend is not None and backend not in cls._backends:
-            raise ValueError(
-                f"Backend {backend} is not supported. Currently supported ones are {list(cls._backends.keys())}"
-            )
-        if prefix is not None and prefix not in cls._prefix_to_backends:
-            raise ValueError(
-                f"prefix {prefix} is not supported. Currently supported ones are {list(cls._prefix_to_backends.keys())}"
-            )
-
-        # concatenate the arguments to a unique key for determining whether
-        # objects with the same arguments were created
-        arg_key = f"{backend}:{prefix}"
-        for key, value in kwargs.items():
-            arg_key += f":{key}:{value}"
-
-        # if a backend was overridden, it will create a new object
-        if arg_key in cls._instances:
-            _instance = cls._instances[arg_key]
-        else:
-            # create a new object and put it to _instance
-            _instance = super().__new__(cls)
-            if backend is not None:
-                _instance.client = cls._backends[backend](**kwargs)
-            else:
-                _instance.client = cls._prefix_to_backends[prefix](**kwargs)
-
-            cls._instances[arg_key] = _instance
-
-        return _instance
-
-    @property
-    def name(self):
-        return self.client.name
-
-    @property
-    def allow_symlink(self):
-        return self.client.allow_symlink
-
-    @staticmethod
-    def parse_uri_prefix(uri: Union[str, Path]) -> Optional[str]:
-        """Parse the prefix of a uri.
-
-        Args:
-            uri (str | Path): Uri to be parsed that contains the file prefix.
-
-        Examples:
-            >>> FileClient.parse_uri_prefix('s3://path/of/your/file')
-            's3'
-
-        Returns:
-            str | None: Return the prefix of uri if the uri contains '://' else
-            ``None``.
-        """
-        assert is_filepath(uri)
-        uri = str(uri)
-        if "://" not in uri:
-            return None
-        else:
-            prefix, _ = uri.split("://")
-            # In the case of Boto3Backend, the prefix may contains the cluster
-            # name like clusterName:s3
-            if ":" in prefix:
-                _, prefix = prefix.split(":")
-            return prefix
-
-    @classmethod
-    def infer_client(
-        cls,
-        file_client_args: Optional[dict] = None,
-        uri: Optional[Union[str, Path]] = None,
-    ) -> "FileClient":
-        """Infer a suitable file client based on the URI and arguments.
-
-        Args:
-            file_client_args (dict, optional): Arguments to instantiate a
-                FileClient. Defaults to None.
-            uri (str | Path, optional): Uri to be parsed that contains the file
-                prefix. Defaults to None.
-
-        Examples:
-            >>> uri = 's3://path/of/your/file'
-            >>> file_client = FileClient.infer_client(uri=uri)
-            >>> file_client_args = {'backend': 's3'}
-            >>> file_client = FileClient.infer_client(file_client_args)
-
-        Returns:
-            FileClient: Instantiated FileClient object.
-        """
-        assert file_client_args is not None or uri is not None
-        if file_client_args is None:
-            file_prefix = cls.parse_uri_prefix(uri)  # type: ignore
-            return cls(prefix=file_prefix)
-        else:
-            return cls(**file_client_args)
-
-    @classmethod
-    def _register_backend(cls, name, backend, force=False, prefixes=None):
-        if not isinstance(name, str):
-            raise TypeError(f"the backend name should be a string, but got {type(name)}")
-        if not inspect.isclass(backend):
-            raise TypeError(f"backend should be a class but got {type(backend)}")
-        if not issubclass(backend, BaseStorageBackend):
-            raise TypeError(f"backend {backend} is not a subclass of BaseStorageBackend")
-        if not force and name in cls._backends:
-            raise KeyError(
-                f'{name} is already registered as a storage backend, add "force=True" if you want to override it'
-            )
-
-        if name in cls._backends and force:
-            for arg_key, instance in list(cls._instances.items()):
-                if isinstance(instance.client, cls._backends[name]):
-                    cls._instances.pop(arg_key)
-        cls._backends[name] = backend
-
-        if prefixes is not None:
-            if isinstance(prefixes, str):
-                prefixes = [prefixes]
-            else:
-                assert isinstance(prefixes, (list, tuple))
-            for prefix in prefixes:
-                if prefix not in cls._prefix_to_backends:
-                    cls._prefix_to_backends[prefix] = backend
-                elif (prefix in cls._prefix_to_backends) and force:
-                    overridden_backend = cls._prefix_to_backends[prefix]
-                    for arg_key, instance in list(cls._instances.items()):
-                        if isinstance(instance.client, overridden_backend):
-                            cls._instances.pop(arg_key)
-                else:
-                    raise KeyError(
-                        f"{prefix} is already registered as a storage backend,"
-                        ' add "force=True" if you want to override it'
-                    )
-
-    @classmethod
-    def register_backend(cls, name, backend=None, force=False, prefixes=None):
-        """Register a backend to FileClient.
-
-        This method can be used as a normal class method or a decorator.
-
-        .. code-block:: python
-
-            class NewBackend(BaseStorageBackend):
-
-                def get(self, filepath):
-                    return filepath
-
-                def get_text(self, filepath):
-                    return filepath
-
-            FileClient.register_backend('new', NewBackend)
-
-        or
-
-        .. code-block:: python
-
-            @FileClient.register_backend('new')
-            class NewBackend(BaseStorageBackend):
-
-                def get(self, filepath):
-                    return filepath
-
-                def get_text(self, filepath):
-                    return filepath
-
-        Args:
-            name (str): The name of the registered backend.
-            backend (class, optional): The backend class to be registered,
-                which must be a subclass of :class:`BaseStorageBackend`.
-                When this method is used as a decorator, backend is None.
-                Defaults to None.
-            force (bool, optional): Whether to override the backend if the name
-                has already been registered. Defaults to False.
-            prefixes (str or list[str] or tuple[str], optional): The prefixes
-                of the registered storage backend. Defaults to None.
-                `New in version 1.3.15.`
-        """
-        if backend is not None:
-            cls._register_backend(name, backend, force=force, prefixes=prefixes)
-            return
-
-        def _register(backend_cls):
-            cls._register_backend(name, backend_cls, force=force, prefixes=prefixes)
-            return backend_cls
-
-        return _register
-
-    def get(self, filepath: Union[str, Path]) -> Union[bytes, memoryview]:
-        """Read data from a given ``filepath`` with 'rb' mode.
-
-        Note:
-            There are two types of return values for ``get``, one is ``bytes``
-            and the other is ``memoryview``. The advantage of using memoryview
-            is that you can avoid copying, and if you want to convert it to
-            ``bytes``, you can use ``.tobytes()``.
-
-        Args:
-            filepath (str or Path): Path to read data.
-
-        Returns:
-            bytes | memoryview: Expected bytes object or a memory view of the
-            bytes object.
-        """
-        return self.client.get(filepath)
-
-    def get_text(self, filepath: Union[str, Path], encoding="utf-8") -> str:
-        """Read data from a given ``filepath`` with 'r' mode.
-
-        Args:
-            filepath (str or Path): Path to read data.
-            encoding (str): The encoding format used to open the ``filepath``.
-                Defaults to 'utf-8'.
-
-        Returns:
-            str: Expected text reading from ``filepath``.
-        """
-        return self.client.get_text(filepath, encoding)
-
-    def put(self, obj: bytes, filepath: Union[str, Path]) -> None:
-        """Write data to a given ``filepath`` with 'wb' mode.
-
-        Note:
-            ``put`` should create a directory if the directory of ``filepath``
-            does not exist.
-
-        Args:
-            obj (bytes): Data to be written.
-            filepath (str or Path): Path to write data.
-        """
-        self.client.put(obj, filepath)
-
-    def put_text(self, obj: str, filepath: Union[str, Path]) -> None:
-        """Write data to a given ``filepath`` with 'w' mode.
-
-        Note:
-            ``put_text`` should create a directory if the directory of
-            ``filepath`` does not exist.
-
-        Args:
-            obj (str): Data to be written.
-            filepath (str or Path): Path to write data.
-            encoding (str, optional): The encoding format used to open the
-                `filepath`. Defaults to 'utf-8'.
-        """
-        self.client.put_text(obj, filepath)
-
-    def remove(self, filepath: Union[str, Path]) -> None:
-        """Remove a file.
-
-        Args:
-            filepath (str, Path): Path to be removed.
-        """
-        self.client.remove(filepath)
-
-    def exists(self, filepath: Union[str, Path]) -> bool:
-        """Check whether a file path exists.
-
-        Args:
-            filepath (str or Path): Path to be checked whether exists.
-
-        Returns:
-            bool: Return ``True`` if ``filepath`` exists, ``False`` otherwise.
-        """
-        return self.client.exists(filepath)
-
-    def isdir(self, filepath: Union[str, Path]) -> bool:
-        """Check whether a file path is a directory.
-
-        Args:
-            filepath (str or Path): Path to be checked whether it is a
-                directory.
-
-        Returns:
-            bool: Return ``True`` if ``filepath`` points to a directory,
-            ``False`` otherwise.
-        """
-        return self.client.isdir(filepath)
-
-    def isfile(self, filepath: Union[str, Path]) -> bool:
-        """Check whether a file path is a file.
-
-        Args:
-            filepath (str or Path): Path to be checked whether it is a file.
-
-        Returns:
-            bool: Return ``True`` if ``filepath`` points to a file, ``False``
-            otherwise.
-        """
-        return self.client.isfile(filepath)
-
-    def join_path(self, filepath: Union[str, Path], *filepaths: Union[str, Path]) -> str:
-        r"""Concatenate all file paths.
-
-        Join one or more filepath components intelligently. The return value
-        is the concatenation of filepath and any members of \*filepaths.
-
-        Args:
-            filepath (str or Path): Path to be concatenated.
-
-        Returns:
-            str: The result of concatenation.
-        """
-        return self.client.join_path(filepath, *filepaths)
-
-    @contextmanager
-    def get_local_path(self, filepath: Union[str, Path]) -> Generator[Union[str, Path], None, None]:
-        """Download data from ``filepath`` and write the data to local path.
-
-        ``get_local_path`` is decorated by :meth:`contxtlib.contextmanager`. It
-        can be called with ``with`` statement, and when exists from the
-        ``with`` statement, the temporary path will be released.
-
-        Note:
-            If the ``filepath`` is a local path, just return itself.
-
-        .. warning::
-            ``get_local_path`` is an experimental interface that may change in
-            the future.
-
-        Args:
-            filepath (str or Path): Path to be read data.
-
-        Examples:
-            >>> file_client = FileClient(prefix='s3')
-            >>> with file_client.get_local_path('s3://bucket/abc.jpg') as path:
-            ...     # do something here
-
-        Yields:
-            Iterable[str]: Only yield one path.
-        """
-        with self.client.get_local_path(str(filepath)) as local_path:
-            yield local_path
-
-    def list_dir_or_file(  # pylint: disable=too-many-arguments
-        self,
-        dir_path: Union[str, Path],
-        list_dir: bool = True,
-        list_file: bool = True,
-        suffix: Optional[Union[str, Tuple[str]]] = None,
-        recursive: bool = False,
-    ) -> Iterator[str]:
-        """Scan a directory to find the interested directories or files in
-        arbitrary order.
-
-        Note:
-            :meth:`list_dir_or_file` returns the path relative to ``dir_path``.
-
-        Args:
-            dir_path (str | Path): Path of the directory.
-            list_dir (bool): List the directories. Defaults to True.
-            list_file (bool): List the path of files. Defaults to True.
-            suffix (str or tuple[str], optional):  File suffix
-                that we are interested in. Defaults to None.
-            recursive (bool): If set to True, recursively scan the
-                directory. Defaults to False.
-
-        Yields:
-            Iterable[str]: A relative path to ``dir_path``.
-        """
-        yield from self.client.list_dir_or_file(dir_path, list_dir, list_file, suffix, recursive)

lyra_2/_ext/imaginaire/utils/easy_io/handlers/__init__.py

@@ -1,29 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from lyra_2._ext.imaginaire.utils.easy_io.handlers.base import BaseFileHandler
-from lyra_2._ext.imaginaire.utils.easy_io.handlers.json_handler import JsonHandler
-from lyra_2._ext.imaginaire.utils.easy_io.handlers.pickle_handler import PickleHandler
-from lyra_2._ext.imaginaire.utils.easy_io.handlers.registry_utils import file_handlers, register_handler
-from lyra_2._ext.imaginaire.utils.easy_io.handlers.yaml_handler import YamlHandler
-
-__all__ = [
-    "BaseFileHandler",
-    "JsonHandler",
-    "PickleHandler",
-    "YamlHandler",
-    "register_handler",
-    "file_handlers",
-]

lyra_2/_ext/imaginaire/utils/easy_io/handlers/base.py

@@ -1,44 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from abc import ABCMeta, abstractmethod
-
-
-class BaseFileHandler(metaclass=ABCMeta):
-    # `str_like` is a flag to indicate whether the type of file object is
-    # str-like object or bytes-like object. Pickle only processes bytes-like
-    # objects but json only processes str-like object. If it is str-like
-    # object, `StringIO` will be used to process the buffer.
-    str_like = True
-
-    @abstractmethod
-    def load_from_fileobj(self, file, **kwargs):
-        pass
-
-    @abstractmethod
-    def dump_to_fileobj(self, obj, file, **kwargs):
-        pass
-
-    @abstractmethod
-    def dump_to_str(self, obj, **kwargs):
-        pass
-
-    def load_from_path(self, filepath, mode="r", **kwargs):
-        with open(filepath, mode) as f:
-            return self.load_from_fileobj(f, **kwargs)
-
-    def dump_to_path(self, obj, filepath, mode="w", **kwargs):
-        with open(filepath, mode) as f:
-            self.dump_to_fileobj(obj, f, **kwargs)

lyra_2/_ext/imaginaire/utils/easy_io/handlers/byte_handler.py

@@ -1,39 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from typing import IO
-
-from lyra_2._ext.imaginaire.utils.easy_io.handlers.base import BaseFileHandler
-
-
-class ByteHandler(BaseFileHandler):
-    str_like = False
-
-    def load_from_fileobj(self, file: IO[bytes], **kwargs):
-        file.seek(0)
-        # extra all bytes and return
-        return file.read()
-
-    def dump_to_fileobj(
-        self,
-        obj: bytes,
-        file: IO[bytes],
-        **kwargs,
-    ):
-        # write all bytes to file
-        file.write(obj)
-
-    def dump_to_str(self, obj, **kwargs):
-        raise NotImplementedError

lyra_2/_ext/imaginaire/utils/easy_io/handlers/csv_handler.py

@@ -1,42 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import csv
-from io import StringIO
-
-from lyra_2._ext.imaginaire.utils.easy_io.handlers.base import BaseFileHandler
-
-
-class CsvHandler(BaseFileHandler):
-    def load_from_fileobj(self, file, **kwargs):
-        del kwargs
-        reader = csv.reader(file)
-        return list(reader)
-
-    def dump_to_fileobj(self, obj, file, **kwargs):
-        del kwargs
-        writer = csv.writer(file)
-        if not all(isinstance(row, list) for row in obj):
-            raise ValueError("Each row must be a list")
-        writer.writerows(obj)
-
-    def dump_to_str(self, obj, **kwargs):
-        del kwargs
-        output = StringIO()
-        writer = csv.writer(output)
-        if not all(isinstance(row, list) for row in obj):
-            raise ValueError("Each row must be a list")
-        writer.writerows(obj)
-        return output.getvalue()

lyra_2/_ext/imaginaire/utils/easy_io/handlers/gzip_handler.py

@@ -1,33 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import gzip
-import pickle
-from io import BytesIO
-from typing import Any
-
-from lyra_2._ext.imaginaire.utils.easy_io.handlers.pickle_handler import PickleHandler
-
-
-class GzipHandler(PickleHandler):
-    str_like = False
-
-    def load_from_fileobj(self, file: BytesIO, **kwargs):
-        with gzip.GzipFile(fileobj=file, mode="rb") as f:
-            return pickle.load(f)
-
-    def dump_to_fileobj(self, obj: Any, file: BytesIO, **kwargs):
-        with gzip.GzipFile(fileobj=file, mode="wb") as f:
-            pickle.dump(obj, f)

lyra_2/_ext/imaginaire/utils/easy_io/handlers/imageio_video_handler.py

@@ -1,168 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from typing import IO, Any, Dict, Tuple
-
-import imageio
-import imageio.v3 as iio_v3
-import numpy as np
-import torch
-
-from lyra_2._ext.imaginaire.utils import log
-from lyra_2._ext.imaginaire.utils.easy_io.handlers.base import BaseFileHandler
-
-
-class ImageioVideoHandler(BaseFileHandler):
-    str_like = False
-
-    def load_from_fileobj(
-        self, file: IO[bytes], format: str = "mp4", mode: str = "rgb", **kwargs
-    ) -> Tuple[np.ndarray, Dict[str, Any]]:
-        """
-        Load video from a file-like object using imageio.v3 with specified format and color mode.
-
-        Parameters:
-            file (IO[bytes]): A file-like object containing video data.
-            format (str): Format of the video file (default 'mp4').
-            mode (str): Color mode of the video, 'rgb' or 'gray' (default 'rgb').
-
-        Returns:
-            tuple: A tuple containing an array of video frames and metadata about the video.
-        """
-        file.seek(0)
-
-        # The plugin argument in v3 replaces the format argument in v2
-        plugin = kwargs.pop("plugin", "pyav")
-
-        # Load all frames at once using v3 API
-        video_frames = iio_v3.imread(file, plugin=plugin, **kwargs)
-
-        # Handle grayscale conversion if needed
-        if mode == "gray":
-            import cv2
-
-            if len(video_frames.shape) == 4:  # (frames, height, width, channels)
-                gray_frames = []
-                for frame in video_frames:
-                    gray_frame = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY)
-                    gray_frame = np.expand_dims(gray_frame, axis=2)  # Keep dimensions consistent
-                    gray_frames.append(gray_frame)
-                video_frames = np.array(gray_frames)
-
-        # Extract metadata
-        # Note: iio_v3.imread doesn't return metadata directly like v2 did
-        # We need to extract it separately
-        file.seek(0)
-        metadata = self._extract_metadata(file, plugin=plugin)
-
-        return video_frames, metadata
-
-    def _extract_metadata(self, file: IO[bytes], plugin: str = "pyav") -> Dict[str, Any]:
-        """
-        Extract metadata from a video file.
-
-        Parameters:
-            file (IO[bytes]): File-like object containing video data.
-            plugin (str): Plugin to use for reading.
-
-        Returns:
-            dict: Video metadata.
-        """
-        try:
-            # Create a generator to read frames and metadata
-            metadata = iio_v3.immeta(file, plugin=plugin)
-
-            # Add some standard fields similar to v2 metadata format
-            if "fps" not in metadata and "duration" in metadata:
-                # Read the first frame to get shape information
-                file.seek(0)
-                first_frame = iio_v3.imread(file, plugin=plugin, index=0)
-                metadata["size"] = first_frame.shape[1::-1]  # (width, height)
-                metadata["source_size"] = metadata["size"]
-
-                # Create a consistent metadata structure with v2
-                metadata["plugin"] = plugin
-                if "codec" not in metadata:
-                    metadata["codec"] = "unknown"
-                if "pix_fmt" not in metadata:
-                    metadata["pix_fmt"] = "unknown"
-
-                # Calculate nframes if possible
-                if "fps" in metadata and "duration" in metadata:
-                    metadata["nframes"] = int(metadata["fps"] * metadata["duration"])
-                else:
-                    metadata["nframes"] = float("inf")
-
-            return metadata
-
-        except Exception as e:
-            # Fallback to basic metadata
-            return {
-                "plugin": plugin,
-                "nframes": float("inf"),
-                "codec": "unknown",
-                "fps": 30.0,  # Default values
-                "duration": 0,
-                "size": (0, 0),
-            }
-
-    def dump_to_fileobj(
-        self,
-        obj: np.ndarray | torch.Tensor,
-        file: IO[bytes],
-        format: str = "mp4",  # pylint: disable=redefined-builtin
-        fps: int = 17,
-        quality: int = 5,
-        ffmpeg_params=None,
-        **kwargs,
-    ):
-        """
-        Save an array of video frames to a file-like object using imageio.
-
-        Parameters:
-            obj (Union[np.ndarray, torch.Tensor]): An array of frames to be saved as video.
-            file (IO[bytes]): A file-like object to which the video data will be written.
-            format (str): Format of the video file (default 'mp4').
-            fps (int): Frames per second of the output video (default 17).
-            quality (int): Quality of the video (0-10, default 5).
-            ffmpeg_params (list): Additional parameters to pass to ffmpeg.
-
-        """
-        if isinstance(obj, torch.Tensor):
-            assert obj.dtype == torch.uint8, "Tensor must be of type uint8"
-            obj = obj.cpu().numpy()
-        h, w = obj.shape[1:-1]
-
-        # Default ffmpeg params that ensure width and height are set
-        default_ffmpeg_params = ["-s", f"{w}x{h}"]
-
-        # Use provided ffmpeg_params if any, otherwise use defaults
-        final_ffmpeg_params = ffmpeg_params if ffmpeg_params is not None else default_ffmpeg_params
-
-        mimsave_kwargs = {
-            "fps": fps,
-            "quality": quality,
-            "macro_block_size": 1,
-            "ffmpeg_params": final_ffmpeg_params,
-            "output_params": ["-f", "mp4"],
-        }
-        # Update with any other kwargs
-        mimsave_kwargs.update(kwargs)
-        log.debug(f"mimsave_kwargs: {mimsave_kwargs}")
-
-        imageio.mimsave(file, obj, format, **mimsave_kwargs)
-
-    def dump_to_str(self, obj, **kwargs):
-        raise NotImplementedError

lyra_2/_ext/imaginaire/utils/easy_io/handlers/json_handler.py

@@ -1,49 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import json
-
-import numpy as np
-
-from lyra_2._ext.imaginaire.utils.easy_io.handlers.base import BaseFileHandler
-
-
-def set_default(obj):
-    """Set default json values for non-serializable values.
-
-    It helps convert ``set``, ``range`` and ``np.ndarray`` data types to list.
-    It also converts ``np.generic`` (including ``np.int32``, ``np.float32``,
-    etc.) into plain numbers of plain python built-in types.
-    """
-    if isinstance(obj, (set, range)):
-        return list(obj)
-    elif isinstance(obj, np.ndarray):
-        return obj.tolist()
-    elif isinstance(obj, np.generic):
-        return obj.item()
-    raise TypeError(f"{type(obj)} is unsupported for json dump")
-
-
-class JsonHandler(BaseFileHandler):
-    def load_from_fileobj(self, file):
-        return json.load(file)
-
-    def dump_to_fileobj(self, obj, file, **kwargs):
-        kwargs.setdefault("default", set_default)
-        json.dump(obj, file, **kwargs)
-
-    def dump_to_str(self, obj, **kwargs):
-        kwargs.setdefault("default", set_default)
-        return json.dumps(obj, **kwargs)

lyra_2/_ext/imaginaire/utils/easy_io/handlers/jsonl_handler.py

@@ -1,80 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import json
-from typing import IO
-
-import numpy as np
-
-from lyra_2._ext.imaginaire.utils.easy_io.handlers.base import BaseFileHandler
-
-
-def set_default(obj):
-    """Set default json values for non-serializable values.
-
-    It helps convert ``set``, ``range`` and ``np.ndarray`` data types to list.
-    It also converts ``np.generic`` (including ``np.int32``, ``np.float32``,
-    etc.) into plain numbers of plain python built-in types.
-    """
-    if isinstance(obj, (set, range)):
-        return list(obj)
-    elif isinstance(obj, np.ndarray):
-        return obj.tolist()
-    elif isinstance(obj, np.generic):
-        return obj.item()
-    raise TypeError(f"{type(obj)} is unsupported for json dump")
-
-
-class JsonlHandler(BaseFileHandler):
-    """Handler for JSON lines (JSONL) files."""
-
-    def load_from_fileobj(self, file: IO[bytes]):
-        """Load JSON objects from a newline-delimited JSON (JSONL) file object.
-
-        Returns:
-            A list of Python objects loaded from each JSON line.
-        """
-        data = []
-        for line in file:
-            line = line.strip()
-            if not line:
-                continue  # skip empty lines if any
-            data.append(json.loads(line))
-        return data
-
-    def dump_to_fileobj(self, obj: IO[bytes], file, **kwargs):
-        """Dump a list of objects to a newline-delimited JSON (JSONL) file object.
-
-        Args:
-            obj: A list (or iterable) of objects to dump line by line.
-        """
-        kwargs.setdefault("default", set_default)
-        for item in obj:
-            file.write(json.dumps(item, **kwargs) + "\n")
-
-    def dump_to_str(self, obj, **kwargs):
-        """Dump a list of objects to a newline-delimited JSON (JSONL) string."""
-        kwargs.setdefault("default", set_default)
-        lines = [json.dumps(item, **kwargs) for item in obj]
-        return "\n".join(lines)
-
-
-if __name__ == "__main__":
-    from lyra_2._ext.imaginaire.utils.easy_io import easy_io
-
-    easy_io.dump([1, 2, 3], "test.jsonl", file_format="jsonl")
-    print(easy_io.load("test.jsonl"))
-    easy_io.dump([{"key1": 1, "key2": 2}, {"key1": 3, "key2": 4}], "test.jsonl", file_format="jsonl")
-    print(easy_io.load("test.jsonl"))

lyra_2/_ext/imaginaire/utils/easy_io/handlers/np_handler.py

@@ -1,89 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from io import BytesIO
-from typing import IO, Any
-
-import numpy as np
-
-from lyra_2._ext.imaginaire.utils.easy_io.handlers.base import BaseFileHandler
-
-
-class NumpyHandler(BaseFileHandler):
-    str_like = False
-
-    def load_from_fileobj(self, file: IO[bytes], **kwargs) -> Any:
-        """
-        Load a NumPy array from a file-like object.
-
-        Parameters:
-            file (IO[bytes]): The file-like object containing the NumPy array data.
-            **kwargs: Additional keyword arguments passed to `np.load`.
-
-        Returns:
-            numpy.ndarray: The loaded NumPy array.
-        """
-        return np.load(file, **kwargs)
-
-    def load_from_path(self, filepath: str, **kwargs) -> Any:
-        """
-        Load a NumPy array from a file path.
-
-        Parameters:
-            filepath (str): The path to the file to load.
-            **kwargs: Additional keyword arguments passed to `np.load`.
-
-        Returns:
-            numpy.ndarray: The loaded NumPy array.
-        """
-        return super().load_from_path(filepath, mode="rb", **kwargs)
-
-    def dump_to_str(self, obj: np.ndarray, **kwargs) -> str:
-        """
-        Serialize a NumPy array to a string in binary format.
-
-        Parameters:
-            obj (np.ndarray): The NumPy array to serialize.
-            **kwargs: Additional keyword arguments passed to `np.save`.
-
-        Returns:
-            str: The serialized NumPy array as a string.
-        """
-        with BytesIO() as f:
-            np.save(f, obj, **kwargs)
-            return f.getvalue()
-
-    def dump_to_fileobj(self, obj: np.ndarray, file: IO[bytes], **kwargs):
-        """
-        Dump a NumPy array to a file-like object.
-
-        Parameters:
-            obj (np.ndarray): The NumPy array to dump.
-            file (IO[bytes]): The file-like object to which the array is dumped.
-            **kwargs: Additional keyword arguments passed to `np.save`.
-        """
-        np.save(file, obj, **kwargs)
-
-    def dump_to_path(self, obj: np.ndarray, filepath: str, **kwargs):
-        """
-        Dump a NumPy array to a file path.
-
-        Parameters:
-            obj (np.ndarray): The NumPy array to dump.
-            filepath (str): The file path where the array should be saved.
-            **kwargs: Additional keyword arguments passed to `np.save`.
-        """
-        with open(filepath, "wb") as f:
-            np.save(f, obj, **kwargs)

lyra_2/_ext/imaginaire/utils/easy_io/handlers/pandas_handler.py

@@ -1,31 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import pandas as pd
-
-from lyra_2._ext.imaginaire.utils.easy_io.handlers.base import BaseFileHandler  # isort:skip
-
-
-class PandasHandler(BaseFileHandler):
-    str_like = False
-
-    def load_from_fileobj(self, file, **kwargs):
-        return pd.read_csv(file, **kwargs)
-
-    def dump_to_fileobj(self, obj, file, **kwargs):
-        obj.to_csv(file, **kwargs)
-
-    def dump_to_str(self, obj, **kwargs):
-        raise NotImplementedError("PandasHandler does not support dumping to str")