nshtrainer 1.2.0__py3-none-any.whl → 1.3.0__py3-none-any.whl

nshtrainer/_directory.py

@@ -19,20 +19,8 @@ class DirectoryConfig(C.Config):
     This isn't specific to the run; it is the parent directory of all runs.
     """
 
-    log: Path | None = None
-    """Base directory for all experiment tracking (e.g., WandB, Tensorboard, etc.) files. If None, will use nshtrainer/{id}/log/."""
-
-    stdio: Path | None = None
-    """stdout/stderr log directory to use for the trainer. If None, will use nshtrainer/{id}/stdio/."""
-
-    checkpoint: Path | None = None
-    """Checkpoint directory to use for the trainer. If None, will use nshtrainer/{id}/checkpoint/."""
-
-    activation: Path | None = None
-    """Activation directory to use for the trainer. If None, will use nshtrainer/{id}/activation/."""
-
-    profile: Path | None = None
-    """Directory to save profiling information to. If None, will use nshtrainer/{id}/profile/."""
+    logdir_basename: str = "nshtrainer"
+    """Base name for the log directory."""
 
     setup_callback: DirectorySetupCallbackConfig = DirectorySetupCallbackConfig()
     """Configuration for the directory setup PyTorch Lightning callback."""
@@ -41,11 +29,11 @@ class DirectoryConfig(C.Config):
         if (project_root_dir := self.project_root) is None:
             project_root_dir = Path.cwd()
 
-        # The default base dir is $CWD/nshtrainer/{id}/
-        base_dir = project_root_dir / "nshtrainer"
+        # The default base dir is $CWD/{logdir_basename}/{id}/
+        base_dir = project_root_dir / self.logdir_basename
         base_dir.mkdir(exist_ok=True)
 
-        # Add a .gitignore file to the nshtrainer directory
+        # Add a .gitignore file to the {logdir_basename} directory
         #   which will ignore all files except for the .gitignore file itself
         gitignore_path = base_dir / ".gitignore"
         if not gitignore_path.exists():
@@ -57,13 +45,8 @@ class DirectoryConfig(C.Config):
 
         return base_dir
 
-    def resolve_subdirectory(
-        self,
-        run_id: str,
-        # subdirectory: Literal["log", "stdio", "checkpoint", "activation", "profile"],
-        subdirectory: str,
-    ) -> Path:
-        # The subdir will be $CWD/nshtrainer/{id}/{log, stdio, checkpoint, activation}/
+    def resolve_subdirectory(self, run_id: str, subdirectory: str) -> Path:
+        # The subdir will be $CWD/{logdir_basename}/{id}/{log, stdio, checkpoint, activation}/
         if (subdir := getattr(self, subdirectory, None)) is not None:
             assert isinstance(subdir, Path), (
                 f"Expected a Path for {subdirectory}, got {type(subdir)}"
@@ -79,7 +62,7 @@ class DirectoryConfig(C.Config):
         if (log_dir := logger.log_dir) is not None:
             return log_dir
 
-        # Save to nshtrainer/{id}/log/{logger name}
+        # Save to {logdir_basename}/{id}/log/{logger name}
         log_dir = self.resolve_subdirectory(run_id, "log")
         log_dir = log_dir / logger.resolve_logger_dirname()
         # ^ NOTE: Logger must have a `name` attribute, as this is

nshtrainer/callbacks/directory_setup.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import logging
-import os
 from pathlib import Path
 from typing import Literal
 
@@ -27,6 +26,7 @@ class DirectorySetupCallbackConfig(CallbackConfigBase):
     def __bool__(self):
         return self.enabled
 
+    @override
     def create_callbacks(self, trainer_config):
         if not self:
             return
@@ -35,21 +35,28 @@ class DirectorySetupCallbackConfig(CallbackConfigBase):
 
 
 def _create_symlink_to_nshrunner(base_dir: Path):
-    # Resolve the current nshrunner session directory
-    if not (session_dir := os.environ.get("NSHRUNNER_SESSION_DIR")):
-        log.warning("NSHRUNNER_SESSION_DIR is not set. Skipping symlink creation.")
+    try:
+        import nshrunner as nr
+    except ImportError:
+        log.info("nshrunner is not installed. Skipping symlink creation to nshrunner.")
+        return
+
+    # Check if we are in a nshrunner session
+    if (session := nr.Session.from_current_session()) is None:
+        log.info("No current nshrunner session found. Skipping symlink creation.")
         return
-    session_dir = Path(session_dir)
+
+    session_dir = session.session_dir
     if not session_dir.exists() or not session_dir.is_dir():
         log.warning(
-            f"NSHRUNNER_SESSION_DIR is not a valid directory: {session_dir}. "
+            f"nshrunner's session_dir is not a valid directory: {session_dir}. "
             "Skipping symlink creation."
         )
         return
 
     # Create the symlink
     symlink_path = base_dir / "nshrunner"
-    if symlink_path.exists():
+    if symlink_path.exists(follow_symlinks=False):
         # If it already points to the correct directory, we're done
         if symlink_path.resolve() == session_dir.resolve():
             return
@@ -61,7 +68,7 @@ def _create_symlink_to_nshrunner(base_dir: Path):
         )
         symlink_path.unlink()
 
-    symlink_path.symlink_to(session_dir)
+    symlink_path.symlink_to(session_dir, target_is_directory=True)
 
 
 class DirectorySetupCallback(NTCallbackBase):

nshtrainer/callbacks/distributed_prediction_writer.py

@@ -4,15 +4,19 @@ import functools
 import logging
 from collections.abc import Iterator, Sequence
 from pathlib import Path
-from typing import Any, ClassVar, Literal, overload
+from typing import TYPE_CHECKING, ClassVar, Generic, Literal, cast, overload
 
 import torch
 from lightning.fabric.utilities.apply_func import move_data_to_device
 from lightning.pytorch.callbacks import BasePredictionWriter
-from typing_extensions import final, override
+from typing_extensions import TypeVar, final, override
 
 from .base import CallbackConfigBase, CallbackMetadataConfig, callback_registry
 
+if TYPE_CHECKING:
+    from ..model.base import IndividualSample
+
+
 log = logging.getLogger(__name__)
 
 
@@ -130,7 +134,15 @@ class DistributedPredictionWriter(BasePredictionWriter):
                 save(sample, output_dir / f"{sample['index']}.pt")
 
 
-class DistributedPredictionReader(Sequence[tuple[Any, Any]]):
+SampleT = TypeVar(
+    "SampleT",
+    bound="IndividualSample",
+    default="IndividualSample",
+    infer_variance=True,
+)
+
+
+class DistributedPredictionReader(Sequence[SampleT], Generic[SampleT]):
     def __init__(self, output_dir: Path):
         self.output_dir = output_dir
 
@@ -139,15 +151,13 @@ class DistributedPredictionReader(Sequence[tuple[Any, Any]]):
         return len(list(self.output_dir.glob("*.pt")))
 
     @overload
-    def __getitem__(self, index: int) -> tuple[Any, Any]: ...
+    def __getitem__(self, index: int) -> SampleT: ...
 
     @overload
-    def __getitem__(self, index: slice) -> list[tuple[Any, Any]]: ...
+    def __getitem__(self, index: slice) -> list[SampleT]: ...
 
     @override
-    def __getitem__(
-        self, index: int | slice
-    ) -> tuple[Any, Any] | list[tuple[Any, Any]]:
+    def __getitem__(self, index: int | slice) -> SampleT | list[SampleT]:
         if isinstance(index, slice):
             # Handle slice indexing
             indices = range(*index.indices(len(self)))
@@ -157,10 +167,11 @@ class DistributedPredictionReader(Sequence[tuple[Any, Any]]):
         path = self.output_dir / f"{index}.pt"
         if not path.exists():
             raise FileNotFoundError(f"File {path} does not exist.")
-        sample = torch.load(path)
-        return sample["batch"], sample["prediction"]
+
+        sample = cast(SampleT, torch.load(path))
+        return sample
 
     @override
-    def __iter__(self) -> Iterator[tuple[Any, Any]]:
+    def __iter__(self) -> Iterator[SampleT]:
         for i in range(len(self)):
             yield self[i]

nshtrainer/configs/trainer/__init__.py

@@ -22,9 +22,6 @@ from nshtrainer.trainer._config import (
     DebugFlagCallbackConfig as DebugFlagCallbackConfig,
 )
 from nshtrainer.trainer._config import DirectoryConfig as DirectoryConfig
-from nshtrainer.trainer._config import (
-    DistributedPredictionWriterConfig as DistributedPredictionWriterConfig,
-)
 from nshtrainer.trainer._config import (
     EarlyStoppingCallbackConfig as EarlyStoppingCallbackConfig,
 )
@@ -126,6 +123,9 @@ from nshtrainer.trainer.plugin.precision import (
 )
 from nshtrainer.trainer.plugin.precision import XLAPluginConfig as XLAPluginConfig
 from nshtrainer.trainer.trainer import AcceleratorConfigBase as AcceleratorConfigBase
+from nshtrainer.trainer.trainer import (
+    DistributedPredictionWriterConfig as DistributedPredictionWriterConfig,
+)
 from nshtrainer.trainer.trainer import StrategyConfigBase as StrategyConfigBase
 
 from . import _config as _config

nshtrainer/configs/trainer/_config/__init__.py

@@ -18,9 +18,6 @@ from nshtrainer.trainer._config import (
     DebugFlagCallbackConfig as DebugFlagCallbackConfig,
 )
 from nshtrainer.trainer._config import DirectoryConfig as DirectoryConfig
-from nshtrainer.trainer._config import (
-    DistributedPredictionWriterConfig as DistributedPredictionWriterConfig,
-)
 from nshtrainer.trainer._config import (
     EarlyStoppingCallbackConfig as EarlyStoppingCallbackConfig,
 )
@@ -73,7 +70,6 @@ __all__ = [
     "CheckpointSavingConfig",
     "DebugFlagCallbackConfig",
     "DirectoryConfig",
-    "DistributedPredictionWriterConfig",
     "EarlyStoppingCallbackConfig",
     "EnvironmentConfig",
     "GradientClippingConfig",

nshtrainer/configs/trainer/trainer/__init__.py

@@ -3,12 +3,16 @@ from __future__ import annotations
 __codegen__ = True
 
 from nshtrainer.trainer.trainer import AcceleratorConfigBase as AcceleratorConfigBase
+from nshtrainer.trainer.trainer import (
+    DistributedPredictionWriterConfig as DistributedPredictionWriterConfig,
+)
 from nshtrainer.trainer.trainer import EnvironmentConfig as EnvironmentConfig
 from nshtrainer.trainer.trainer import StrategyConfigBase as StrategyConfigBase
 from nshtrainer.trainer.trainer import TrainerConfig as TrainerConfig
 
 __all__ = [
     "AcceleratorConfigBase",
+    "DistributedPredictionWriterConfig",
     "EnvironmentConfig",
     "StrategyConfigBase",
     "TrainerConfig",

nshtrainer/trainer/_config.py

@@ -31,7 +31,6 @@ from .._hf_hub import HuggingFaceHubConfig
 from ..callbacks import (
     BestCheckpointCallbackConfig,
     CallbackConfig,
-    DistributedPredictionWriterConfig,
     EarlyStoppingCallbackConfig,
     LastCheckpointCallbackConfig,
     NormLoggingCallbackConfig,
@@ -702,14 +701,6 @@ class TrainerConfig(C.Config):
     auto_validate_metrics: MetricValidationCallbackConfig | None = None
     """If enabled, will automatically validate the metrics before starting the training routine."""
 
-    distributed_predict: DistributedPredictionWriterConfig | None = (
-        DistributedPredictionWriterConfig()
-    )
-    """If enabled, will use a custom BasePredictionWriter callback to automatically
-    handle distributed prediction. This is useful for running prediction on multiple GPUs
-    seamlessly.
-    """
-
     lightning_kwargs: LightningTrainerKwargs = LightningTrainerKwargs()
     """
     Additional keyword arguments to pass to the Lightning `pl.Trainer` constructor.
@@ -761,6 +752,7 @@ class TrainerConfig(C.Config):
         )
 
     def _nshtrainer_all_callback_configs(self) -> Iterable[CallbackConfigBase | None]:
+        yield self.directory.setup_callback
         yield self.early_stopping
         yield self.checkpoint_saving
         yield self.lr_monitor
@@ -777,7 +769,6 @@ class TrainerConfig(C.Config):
         yield self.reduce_lr_on_plateau_sanity_checking
         yield self.auto_set_debug_flag
         yield self.auto_validate_metrics
-        yield self.distributed_predict
         yield from self.callbacks
 
     def _nshtrainer_all_logger_configs(self) -> Iterable[LoggerConfigBase | None]:

nshtrainer/trainer/_distributed_prediction_result.py

@@ -0,0 +1,80 @@
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from pathlib import Path
+
+log = logging.getLogger(__name__)
+
+
+@dataclass
+class DistributedPredictionResult:
+    """Represents the results of a distributed prediction run.
+
+    This dataclass provides easy access to both raw and processed prediction data.
+    """
+
+    root_dir: Path
+    """Root directory where predictions are stored."""
+
+    @property
+    def raw_dir(self) -> Path:
+        """Directory containing raw prediction data."""
+        return self.root_dir / "raw"
+
+    @property
+    def processed_dir(self) -> Path:
+        """Directory containing processed prediction data."""
+        return self.root_dir / "processed"
+
+    def get_raw_predictions(self, dataloader_idx: int = 0) -> Path:
+        """Get the directory containing raw predictions for a specific dataloader.
+
+        Args:
+            dataloader_idx: Index of the dataloader
+
+        Returns:
+            Path to the raw predictions directory for the specified dataloader
+        """
+        raw_loader_dir = self.raw_dir / f"dataloader_{dataloader_idx}"
+        if not raw_loader_dir.exists():
+            log.warning(f"Raw predictions directory {raw_loader_dir} does not exist.")
+        return raw_loader_dir
+
+    def get_processed_reader(self, dataloader_idx: int = 0):
+        """Get a reader for processed predictions from a specific dataloader.
+
+        Args:
+            dataloader_idx: Index of the dataloader
+
+        Returns:
+            A DistributedPredictionReader for the processed predictions, or None if no data exists
+        """
+        from ..callbacks.distributed_prediction_writer import (
+            DistributedPredictionReader,
+        )
+
+        processed_loader_dir = self.processed_dir / f"dataloader_{dataloader_idx}"
+        if not processed_loader_dir.exists():
+            log.warning(
+                f"Processed predictions directory {processed_loader_dir} does not exist."
+            )
+            return None
+
+        return DistributedPredictionReader(processed_loader_dir)
+
+    @classmethod
+    def load(cls, path: Path | str):
+        """Load prediction results from a directory.
+
+        Args:
+            path: Path to the predictions directory
+
+        Returns:
+            A DistributedPredictionResult instance
+        """
+        path = Path(path)
+        if not path.exists():
+            raise FileNotFoundError(f"Predictions directory {path} does not exist.")
+
+        return cls(root_dir=path)

nshtrainer/trainer/trainer.py

@@ -4,7 +4,7 @@ import logging
 import os
 from collections.abc import Callable, Mapping, Sequence
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, cast
+from typing import TYPE_CHECKING, Any, cast, overload
 
 import torch
 from lightning.fabric.plugins.environments.lsf import LSFEnvironment
@@ -24,9 +24,14 @@ from typing_extensions import Never, Unpack, assert_never, deprecated, override
 
 from .._checkpoint.metadata import write_checkpoint_metadata
 from ..callbacks.base import resolve_all_callbacks
+from ..callbacks.distributed_prediction_writer import (
+    DistributedPredictionWriter,
+    DistributedPredictionWriterConfig,
+)
 from ..util._environment_info import EnvironmentConfig
 from ..util.bf16 import is_bf16_supported_no_emulation
 from ._config import LightningTrainerKwargs, TrainerConfig
+from ._distributed_prediction_result import DistributedPredictionResult
 from ._log_hparams import patch_log_hparams_function
 from ._runtime_callback import RuntimeTrackerCallback, Stage
 from .accelerator import AcceleratorConfigBase
@@ -537,13 +542,66 @@ class Trainer(LightningTrainer):
         )
         return cls(hparams)
 
+    @overload
     def distributed_predict(
         self,
         model: LightningModule | None = None,
         dataloaders: EVAL_DATALOADERS | LightningDataModule | None = None,
         datamodule: LightningDataModule | None = None,
         ckpt_path: str | Path | None = None,
-    ):
+        *,
+        config: DistributedPredictionWriterConfig,
+    ) -> DistributedPredictionResult: ...
+
+    @overload
+    def distributed_predict(
+        self,
+        model: LightningModule | None = None,
+        dataloaders: EVAL_DATALOADERS | LightningDataModule | None = None,
+        datamodule: LightningDataModule | None = None,
+        ckpt_path: str | Path | None = None,
+        *,
+        dirpath: Path | None = None,
+        move_to_cpu_on_save: bool = True,
+        save_raw: bool = True,
+        save_processed: bool = True,
+    ) -> DistributedPredictionResult: ...
+
+    def distributed_predict(
+        self,
+        model: LightningModule | None = None,
+        dataloaders: EVAL_DATALOADERS | LightningDataModule | None = None,
+        datamodule: LightningDataModule | None = None,
+        ckpt_path: str | Path | None = None,
+        *,
+        config: DistributedPredictionWriterConfig | None = None,
+        dirpath: Path | None = None,
+        move_to_cpu_on_save: bool = True,
+        save_raw: bool = True,
+        save_processed: bool = True,
+    ) -> DistributedPredictionResult:
+        if config is None:
+            config = DistributedPredictionWriterConfig(
+                dirpath=dirpath,
+                move_to_cpu_on_save=move_to_cpu_on_save,
+                save_raw=save_raw,
+                save_processed=save_processed,
+            )
+
+        # Remove any DistributedPredictionWriter callbacks that are already set
+        # and add the new one.
+        callbacks = self.callbacks.copy()
+        callbacks = [
+            callback
+            for callback in callbacks
+            if not isinstance(callback, DistributedPredictionWriter)
+        ]
+        writer_callbacks = list(config.create_callbacks(self.hparams))
+        assert len(writer_callbacks) == 1
+        callback = writer_callbacks[0]
+        callbacks.append(callback)
+        self.callbacks = self._callback_connector._reorder_callbacks(callbacks)
+
         self.predict(
             model,
             dataloaders,
@@ -551,3 +609,9 @@ class Trainer(LightningTrainer):
             return_predictions=False,
             ckpt_path=ckpt_path,
         )
+
+        # Wait for all processes to finish
+        self.strategy.barrier("Trainer.distributed_predict")
+
+        # Return an object that contains information about the predictions
+        return DistributedPredictionResult(root_dir=callback.output_dir)

{nshtrainer-1.2.0.dist-info → nshtrainer-1.3.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: nshtrainer
-Version: 1.2.0
+Version: 1.3.0
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com

{nshtrainer-1.2.0.dist-info → nshtrainer-1.3.0.dist-info}/RECORD

@@ -3,7 +3,7 @@ nshtrainer/__init__.py,sha256=VcqBfL8RgCcZDaY645nxeDmOspqerx4x46wggCMnS0E,692
 nshtrainer/_callback.py,sha256=ZDppiJ4d65tRXTEWYPZLH_F1xFizdz1pkWJe_sQ5uII,12564
 nshtrainer/_checkpoint/metadata.py,sha256=Hh5a7OkdknUEbkEwX6vS88-XLEeuVDoR6a3en2uLzQE,5597
 nshtrainer/_checkpoint/saver.py,sha256=utcrYKSosd04N9m2GIylufO5DO05D90qVU3mvadfApU,1658
-nshtrainer/_directory.py,sha256=SuXJe9xJXZkDXWWfeOS9rEDz6vZUA6mpnEdkAW0ZQnY,3193
+nshtrainer/_directory.py,sha256=RAG8e0y3VZwGIyy_D-GXgDMK5OvitQU6qEWxHTpWEeY,2490
 nshtrainer/_experimental/__init__.py,sha256=U4S_2y3zgLZVfMenHRaJFBW8yqh2mUBuI291LGQVOJ8,35
 nshtrainer/_hf_hub.py,sha256=4OsCbIITnZk_YLyoMrVyZ0SIN04FBxlC0ig2Et8UAdo,14287
 nshtrainer/callbacks/__init__.py,sha256=m6eJuprZfBELuKpngKXre33B9yPXkG7jlKVmI-0yXRQ,4000
@@ -15,8 +15,8 @@ nshtrainer/callbacks/checkpoint/best_checkpoint.py,sha256=aCs3E1eucfDlUeW2Iq_Ke7
 nshtrainer/callbacks/checkpoint/last_checkpoint.py,sha256=vn-as3ex7kaTRcKsIurVtM6kUSHYNwHJeYG82j2dMcc,3554
 nshtrainer/callbacks/checkpoint/on_exception_checkpoint.py,sha256=nljzETqkHwA-4g8mxaeFK5HxA8My0dlIPzIUscSMWyk,3525
 nshtrainer/callbacks/debug_flag.py,sha256=96fuP0C7C6dSs1GiMeUYzzs0X3Q4Pjt9JVWg3b75fU4,1748
-nshtrainer/callbacks/directory_setup.py,sha256=wPas_Ren8ANejogmIdKhqqgj4ulxz9AS_8xVIAfRXa0,2565
-nshtrainer/callbacks/distributed_prediction_writer.py,sha256=OSh2C6XF7Nki4eFByNVhwlt69izkxnlmfPx54w4rvBo,5274
+nshtrainer/callbacks/directory_setup.py,sha256=Ln6f0tCgoBscHeigIAWtCCoAmuWB-kPyaf7SylU7MYo,2773
+nshtrainer/callbacks/distributed_prediction_writer.py,sha256=PvxV9E9lHT-NQ-h1ld7WugajqiFyFXECsreUt3e7pxk,5440
 nshtrainer/callbacks/early_stopping.py,sha256=rC_qYKCQWjRQJFo0ky46uG0aDJdYP8vsSlKunk0bUVI,4765
 nshtrainer/callbacks/ema.py,sha256=dBFiUXG0xmyCw8-ayuSzJMKqSbepl6Ii5VIbhFlT5ug,12255
 nshtrainer/callbacks/finite_checks.py,sha256=3lZ3kEIjmYQfqTF0DcrgZ9_98ZLQhQj8usH7SgWst3o,2185
@@ -85,8 +85,8 @@ nshtrainer/configs/profiler/_base/__init__.py,sha256=ekYfPg-VDhCAFM5nJka2TxUYdRD
 nshtrainer/configs/profiler/advanced/__init__.py,sha256=-ThpUat16Ij_0avkMUVVA8wCWDG_q_tM7KQofnWQCtg,308
 nshtrainer/configs/profiler/pytorch/__init__.py,sha256=soAU1s2_Pa1na4gW8CK-iysJBO5M_7YeZC2_x40iEdg,294
 nshtrainer/configs/profiler/simple/__init__.py,sha256=3Wb11lPuFuyasq8xS1CZ4WLuBCLS_nVSQGVllvOOi0Y,289
-nshtrainer/configs/trainer/__init__.py,sha256=PF9rYuVpk0IuhjcxS_hmBTT6A0oq7AWZDcx0Gfqi7MM,8040
-nshtrainer/configs/trainer/_config/__init__.py,sha256=5B8pjyNHfyFJ6p8dD5VSHD1tw2CcZ87Eq2C_Req3t60,3977
+nshtrainer/configs/trainer/__init__.py,sha256=YLlDOUYDp_qURHhcmhCxTcY6K5AbmoTxdzBPB9SEZII,8040
+nshtrainer/configs/trainer/_config/__init__.py,sha256=6DXdtP-uH11TopQ7kzId9fco-wVkD7ZfevbBqDpN6TE,3817
 nshtrainer/configs/trainer/accelerator/__init__.py,sha256=3H6R3wlwbKL1TzDqGCChZk78-BcE2czLouo7Djiq3nA,898
 nshtrainer/configs/trainer/plugin/__init__.py,sha256=NkHQxMPkrtTtdIAO4dQUE9SWEcHRDB0yUXLkTjnl4dA,3332
 nshtrainer/configs/trainer/plugin/base/__init__.py,sha256=slW5z1FZw2qICXO9l9DnLIDB1Yl7KOcxPEZkyYIHrp4,276
@@ -95,7 +95,7 @@ nshtrainer/configs/trainer/plugin/io/__init__.py,sha256=AtGUuE0M16dTpX0q9NqvJiE4
 nshtrainer/configs/trainer/plugin/layer_sync/__init__.py,sha256=SYDZk2M6sgpt4sEuoURuS8EKYmaqGcvYxETE9jvTrEE,431
 nshtrainer/configs/trainer/plugin/precision/__init__.py,sha256=szlqSfK2XuWdkf72LQzQFv3SlWfKFdRUpBEYIxQ3TPs,1507
 nshtrainer/configs/trainer/strategy/__init__.py,sha256=50whNloJVBq_bdbLaPQnPBTeS1Rcs8MwxTCYBj1kKa4,273
-nshtrainer/configs/trainer/trainer/__init__.py,sha256=DDuBRx0kVNMW0z_sqKTUt8-Ql7bOpargi4KcHHvDu_c,486
+nshtrainer/configs/trainer/trainer/__init__.py,sha256=gOyfE4LlKP-pDJB_ILf79--GztnkF_QmEcexHgqGxOI,646
 nshtrainer/configs/util/__init__.py,sha256=qXittS7f7MyaqJnjvFLKnKsyb6bXTD3dEV16jXVDaH4,2104
 nshtrainer/configs/util/_environment_info/__init__.py,sha256=eB4E0Ck7XCeSC5gbUdA5thd7TXnjGCL0t8GZIFj7uCI,1644
 nshtrainer/configs/util/config/__init__.py,sha256=nEFiDG3-dvvTytYn1tEkPFzp7fgaGRp2j7toSN7yRGs,501
@@ -135,7 +135,8 @@ nshtrainer/profiler/advanced.py,sha256=XrM3FX0ThCv5UwUrrH0l4Ow4LGAtpiBww2N8QAU5N
 nshtrainer/profiler/pytorch.py,sha256=8K37XvPnCApUpIK8tA2zNMFIaIiTLSoxKQoiyCPBm1Q,2757
 nshtrainer/profiler/simple.py,sha256=PimjqcU-JuS-8C0ZGHAdwCxgNLij4x0FH6WXsjBQzZs,1005
 nshtrainer/trainer/__init__.py,sha256=fQ7gQRlGWX-90TYT0rttkQyvXDCzo7DAvJgr-jX1zsY,316
-nshtrainer/trainer/_config.py,sha256=tdWAYh-KGXBpgdY8fwvOejjRZN-AS2Ze0f_9s2VEuZ0,33556
+nshtrainer/trainer/_config.py,sha256=Lt9tuzxgVzVnyEFz61xbaPudfsXbKYUphOg-qMDHO8g,33203
+nshtrainer/trainer/_distributed_prediction_result.py,sha256=bQw8Z6PT694UUf-zQPkech6CxyUSy8bAIexfSfPej0U,2507
 nshtrainer/trainer/_log_hparams.py,sha256=XH2lZ4U_3AZBhOt91ocsEhdL_NRz35oWvqLCUFDohUs,2389
 nshtrainer/trainer/_runtime_callback.py,sha256=6F2Gq27Q8OFfN3RtdNC6QRA8ac0LC1hh4DUE3V5WgbI,4217
 nshtrainer/trainer/accelerator.py,sha256=Bqq-ry7DeCY4zw9_zBvTZiijpA-uUHrDjtbLV652m4M,2415
@@ -147,7 +148,7 @@ nshtrainer/trainer/plugin/layer_sync.py,sha256=-BbEyWZ063O7tZme7Gdu1lVxK6p1NeuLc
 nshtrainer/trainer/plugin/precision.py,sha256=7lf7KZd_yFyPmhLApjEIv0pkoDB5zdxi-7in0wRj3z8,5436
 nshtrainer/trainer/signal_connector.py,sha256=GhfGcSzfaTNhnj2QFkBDq5aT7FqbLMA7eC8SYQs8_8w,10828
 nshtrainer/trainer/strategy.py,sha256=VPTn5z3zvXTydY8IJchjhjcOfpvtoejnvUkq5E4WTus,1368
-nshtrainer/trainer/trainer.py,sha256=smoN61iixWYDWGFvxrt8VwryZVy_NzqqjUcgOid0gRA,21696
+nshtrainer/trainer/trainer.py,sha256=6oky6E8cjGqUNzJGyyTO551pE9A6YueOv5oxg1fZVR0,24129
 nshtrainer/util/_environment_info.py,sha256=MT8mBe6ZolRfKiwU-les1P-lPNPqXpHQcfADrh_A3uY,24629
 nshtrainer/util/bf16.py,sha256=9QhHZCkYSfYpIcxwAMoXyuh2yTSHBzT-EdLQB297jEs,762
 nshtrainer/util/config/__init__.py,sha256=Z39JJufSb61Lhn2GfVcv3eFW_eorOrN9-9llDWlnZZM,272
@@ -159,6 +160,6 @@ nshtrainer/util/seed.py,sha256=diMV8iwBKN7Xxt5pELmui-gyqyT80_CZzomrWhNss0k,316
 nshtrainer/util/slurm.py,sha256=HflkP5iI_r4UHMyPjw9R4dD5AHsJUpcfJw5PLvGYBRM,1603
 nshtrainer/util/typed.py,sha256=Xt5fUU6zwLKSTLUdenovnKK0N8qUq89Kddz2_XeykVQ,164
 nshtrainer/util/typing_utils.py,sha256=MjY-CUX9R5Tzat-BlFnQjwl1PQ_W2yZQoXhkYHlJ_VA,442
-nshtrainer-1.2.0.dist-info/METADATA,sha256=HkNLruaJJuf3ijnGe7NqNd9emBR6QHMRh2-taC5wTrU,960
-nshtrainer-1.2.0.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
-nshtrainer-1.2.0.dist-info/RECORD,,
+nshtrainer-1.3.0.dist-info/METADATA,sha256=M84AwXCuoJp21_m2IQKYDC-SFWDAdhOy-2fDL1jk9Lw,960
+nshtrainer-1.3.0.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
+nshtrainer-1.3.0.dist-info/RECORD,,