nshtrainer 1.2.1__tar.gz → 1.3.1__tar.gz

{nshtrainer-1.2.1 → nshtrainer-1.3.1}/PKG-INFO

@@ -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 → nshtrainer-1.3.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "nshtrainer"
-version = "1.2.1"
+version = "1.3.1"
 description = ""
 authors = [{ name = "Nima Shoghi", email = "nimashoghi@gmail.com" }]
 requires-python = ">=3.10,<4.0"

{nshtrainer-1.2.1 → nshtrainer-1.3.1}/src/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-1.2.1 → nshtrainer-1.3.1}/src/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-1.2.1 → nshtrainer-1.3.1}/src/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-1.2.1 → nshtrainer-1.3.1}/src/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-1.2.1 → nshtrainer-1.3.1}/src/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-1.2.1 → nshtrainer-1.3.1}/src/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-1.3.1/src/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-1.2.1 → nshtrainer-1.3.1}/src/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)