chore: import upstream snapshot with attribution

python/ray/train/torch/torch_checkpoint.py

@@ -0,0 +1,196 @@
+import os
+import tempfile
+import warnings
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Dict, Optional
+
+import torch
+
+from ray.air._internal.torch_utils import (
+    consume_prefix_in_state_dict_if_present_not_in_place,
+    load_torch_model,
+)
+from ray.train._internal.framework_checkpoint import FrameworkCheckpoint
+from ray.util.annotations import PublicAPI
+
+if TYPE_CHECKING:
+    from ray.data.preprocessor import Preprocessor
+
+ENCODED_DATA_KEY = "torch_encoded_data"
+
+
+@PublicAPI(stability="beta")
+class TorchCheckpoint(FrameworkCheckpoint):
+    """A :class:`~ray.train.Checkpoint` with Torch-specific functionality."""
+
+    MODEL_FILENAME = "model.pt"
+
+    @classmethod
+    def from_state_dict(
+        cls,
+        state_dict: Dict[str, Any],
+        *,
+        preprocessor: Optional["Preprocessor"] = None,
+    ) -> "TorchCheckpoint":
+        """Create a :class:`~ray.train.Checkpoint` that stores a model state dictionary.
+
+        .. tip::
+
+            This is the recommended method for creating
+            :class:`TorchCheckpoints<TorchCheckpoint>`.
+
+        Args:
+            state_dict: The model state dictionary to store in the checkpoint.
+            preprocessor: A fitted preprocessor to be applied before inference.
+
+        Returns:
+            A :class:`TorchCheckpoint` containing the specified state dictionary.
+
+        Examples:
+
+            .. testcode::
+
+                import torch
+                import torch.nn as nn
+                from ray.train.torch import TorchCheckpoint
+
+                # Set manual seed
+                torch.manual_seed(42)
+
+                # Function to create a NN model
+                def create_model() -> nn.Module:
+                    model = nn.Sequential(nn.Linear(1, 10),
+                            nn.ReLU(),
+                            nn.Linear(10,1))
+                    return model
+
+                # Create a TorchCheckpoint from our model's state_dict
+                model = create_model()
+                checkpoint = TorchCheckpoint.from_state_dict(model.state_dict())
+
+                # Now load the model from the TorchCheckpoint by providing the
+                # model architecture
+                model_from_chkpt = checkpoint.get_model(create_model())
+
+                # Assert they have the same state dict
+                assert str(model.state_dict()) == str(model_from_chkpt.state_dict())
+                print("worked")
+
+            .. testoutput::
+                :hide:
+
+                ...
+        """
+        tempdir = tempfile.mkdtemp()
+
+        model_path = Path(tempdir, cls.MODEL_FILENAME).as_posix()
+        stripped_state_dict = consume_prefix_in_state_dict_if_present_not_in_place(
+            state_dict, "module."
+        )
+        torch.save(stripped_state_dict, model_path)
+
+        checkpoint = cls.from_directory(tempdir)
+        if preprocessor:
+            checkpoint.set_preprocessor(preprocessor)
+        return checkpoint
+
+    @classmethod
+    def from_model(
+        cls,
+        model: torch.nn.Module,
+        *,
+        preprocessor: Optional["Preprocessor"] = None,
+    ) -> "TorchCheckpoint":
+        """Create a :class:`~ray.train.Checkpoint` that stores a Torch model.
+
+        .. note::
+
+            PyTorch recommends storing state dictionaries. To create a
+            :class:`TorchCheckpoint` from a state dictionary, call
+            :meth:`~ray.train.torch.TorchCheckpoint.from_state_dict`. To learn more
+            about state dictionaries, read
+            `Saving and Loading Models <https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict>`_. # noqa: E501
+
+        Args:
+            model: The Torch model to store in the checkpoint.
+            preprocessor: A fitted preprocessor to be applied before inference.
+
+        Returns:
+            A :class:`TorchCheckpoint` containing the specified model.
+
+        Examples:
+
+            .. testcode::
+
+                from ray.train.torch import TorchCheckpoint
+                import torch
+
+                # Create model identity and send a random tensor to it
+                model = torch.nn.Identity()
+                input = torch.randn(2, 2)
+                output = model(input)
+
+                # Create a checkpoint
+                checkpoint = TorchCheckpoint.from_model(model)
+                print(checkpoint)
+
+            .. testoutput::
+                :hide:
+
+                ...
+        """
+        tempdir = tempfile.mkdtemp()
+
+        model_path = Path(tempdir, cls.MODEL_FILENAME).as_posix()
+        torch.save(model, model_path)
+
+        checkpoint = cls.from_directory(tempdir)
+        if preprocessor:
+            checkpoint.set_preprocessor(preprocessor)
+        return checkpoint
+
+    def get_model(self, model: Optional[torch.nn.Module] = None) -> torch.nn.Module:
+        """Retrieve the model stored in this checkpoint.
+
+        .. warning::
+
+            The checkpoint path must point to a **trusted** source.
+            Checkpoints created with
+            :meth:`~ray.train.torch.TorchCheckpoint.from_model` store the entire
+            ``nn.Module`` via pickle serialization. Loading such a checkpoint from an
+            untrusted path (shared storage, downloaded artifact, checkpoint produced by
+            a different party) is equivalent to executing arbitrary Python code. Prefer
+            checkpoints created with
+            :meth:`~ray.train.torch.TorchCheckpoint.from_state_dict`, which stores
+            only model weights and is safe to load from untrusted sources.
+
+        Args:
+            model: If the checkpoint contains a model state dict, and not
+                the model itself, then the state dict will be loaded to this
+                ``model``. Otherwise, the model will be discarded.
+
+        Returns:
+            The loaded ``torch.nn.Module``.
+        """
+        with self.as_directory() as tempdir:
+            model_path = Path(tempdir, self.MODEL_FILENAME).as_posix()
+            if not os.path.exists(model_path):
+                raise RuntimeError(
+                    "`model.pt` not found within this checkpoint. Make sure you "
+                    "created this `TorchCheckpoint` from one of its public "
+                    "constructors (`from_state_dict` or `from_model`)."
+                )
+            model_or_state_dict = torch.load(
+                model_path, map_location="cpu", weights_only=False
+            )
+
+        if isinstance(model_or_state_dict, torch.nn.Module):
+            if model:
+                warnings.warn(
+                    "TorchCheckpoint already contains all information needed. "
+                    "Discarding provided `model` argument."
+                )
+        model = load_torch_model(
+            saved_model=model_or_state_dict, model_definition=model
+        )
+        return model