nshtrainer 1.2.1__py3-none-any.whl → 1.3.1__py3-none-any.whl

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/__init__.py

@@ -1,7 +1,13 @@
 from __future__ import annotations
 
 from ..callbacks import callback_registry as callback_registry
+from ..callbacks.distributed_prediction_writer import (
+    DistributedPredictionReader as DistributedPredictionReader,
+)
 from ._config import TrainerConfig as TrainerConfig
+from ._distributed_prediction_result import (
+    DistributedPredictionResult as DistributedPredictionResult,
+)
 from .accelerator import accelerator_registry as accelerator_registry
 from .plugin import plugin_registry as plugin_registry
 from .trainer import Trainer as Trainer

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.
@@ -778,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.1.dist-info → nshtrainer-1.3.1.dist-info}/METADATA

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

{nshtrainer-1.2.1.dist-info → nshtrainer-1.3.1.dist-info}/RECORD

@@ -16,7 +16,7 @@ nshtrainer/callbacks/checkpoint/last_checkpoint.py,sha256=vn-as3ex7kaTRcKsIurVtM
 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=Ln6f0tCgoBscHeigIAWtCCoAmuWB-kPyaf7SylU7MYo,2773
-nshtrainer/callbacks/distributed_prediction_writer.py,sha256=OSh2C6XF7Nki4eFByNVhwlt69izkxnlmfPx54w4rvBo,5274
+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
@@ -134,8 +134,9 @@ nshtrainer/profiler/_base.py,sha256=kFcSVn9gJuMwgDxbfyHh46CmEAIPZjxw3yjPbKgzvwA,
 nshtrainer/profiler/advanced.py,sha256=XrM3FX0ThCv5UwUrrH0l4Ow4LGAtpiBww2N8QAU5NOQ,1160
 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=EaAZCajuvDoVAn79zGebUgg0ijz-i3eLhwxYq8oTNe8,33600
+nshtrainer/trainer/__init__.py,sha256=jRaHdaFK8wxNrN1bleT9cf29iZahL_-XkWo5TWz2CmA,550
+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.1.dist-info/METADATA,sha256=MrF6xvpgRjy2HsFwCIU9YZXMYcOb9ItAXY8m66P2CoQ,960
-nshtrainer-1.2.1.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
-nshtrainer-1.2.1.dist-info/RECORD,,
+nshtrainer-1.3.1.dist-info/METADATA,sha256=RCFzQ6YlNZmaYUMcLR4RMotPI3X3QXFwI6MWyN5nkjE,960
+nshtrainer-1.3.1.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
+nshtrainer-1.3.1.dist-info/RECORD,,