fix: close the remaining gaps in optional torch support

pyproject.toml

@@ -45,6 +45,15 @@ serve = ["sagemaker-serve"]
 mlops = ["sagemaker-mlops"]
 all = ["sagemaker-train", "sagemaker-serve", "sagemaker-mlops"]
 
+# sagemaker-serve's framework backends are optional, so pass its extras through
+# to let callers write 'pip install sagemaker[torch]' instead of having to name
+# the subpackage. 'all' is deliberately left as the subpackage list: installing
+# every backend is 'pip install sagemaker-serve[all]'.
+torch = ["sagemaker-serve[torch]"]
+onnx = ["sagemaker-serve[onnx]"]
+triton = ["sagemaker-serve[triton]"]
+mlflow = ["sagemaker-serve[mlflow]"]
+
 [project.urls]
 Homepage = "https://github.com/aws/sagemaker-python-sdk"
 

sagemaker-core/src/sagemaker/core/serializers/base.py

@@ -19,6 +19,7 @@
 import csv
 import io
 import json
+import sys
 import numpy as np
 from pandas import DataFrame
 from six import with_metaclass
@@ -433,6 +434,18 @@ def serialize(self, data):
         raise ValueError(f"Object of type {type(data)} is not String serializable.")
 
 
+def _is_torch_tensor(data):
+    """Return whether ``data`` is a ``torch.Tensor``, without importing torch.
+
+    A ``torch.Tensor`` instance cannot exist unless torch has already been
+    imported, so a missing ``sys.modules`` entry is proof that ``data`` is not
+    one. This keeps the check exact while letting callers that never pass a
+    tensor avoid pulling in an optional multi-hundred-megabyte dependency.
+    """
+    torch = sys.modules.get("torch")
+    return torch is not None and isinstance(data, torch.Tensor)
+
+
 class TorchTensorSerializer(SimpleBaseSerializer):
     """Serialize torch.Tensor to a buffer by converting tensor to numpy and call NumpySerializer.
 
@@ -444,17 +457,24 @@ class TorchTensorSerializer(SimpleBaseSerializer):
 
     def __init__(self, content_type="tensor/pt"):
         super(TorchTensorSerializer, self).__init__(content_type=content_type)
+        self.numpy_serializer = NumpySerializer()
+
+    @property
+    def torch_tensor(self):
+        """The ``torch.Tensor`` type, imported on first access.
+
+        torch is an optional dependency (``sagemaker-core[torch]``), so it is not
+        imported until a caller actually needs the type.
+        """
         try:
             from torch import Tensor
-
-            self.torch_tensor = Tensor
         except ImportError as e:
             raise ImportError(
                 "Unable to import torch. Please install torch to use TorchTensorSerializer: "
                 "pip install 'sagemaker-core[torch]'"
             ) from e
 
-        self.numpy_serializer = NumpySerializer()
+        return Tensor
 
     def serialize(self, data):
         """Serialize torch.Tensor to a buffer.
@@ -465,7 +485,7 @@ def serialize(self, data):
         Returns:
             raw-bytes: The data serialized as raw-bytes from the input.
         """
-        if isinstance(data, self.torch_tensor):
+        if _is_torch_tensor(data):
             try:
                 return self.numpy_serializer.serialize(data.detach().numpy())
             except Exception as e:

sagemaker-core/tests/unit/test_optional_torch_dependency.py

@@ -122,16 +122,39 @@ def test_deserializer_module_imports_without_torch():
     ), f"Subprocess failed:\nstdout: {result.stdout}\nstderr: {result.stderr}"
 
 
-def test_torch_tensor_serializer_raises_import_error_without_torch():
-    """Verify TorchTensorSerializer raises ImportError when torch is not installed."""
+def test_torch_tensor_serializer_constructs_without_torch():
+    """Verify TorchTensorSerializer can be built and rejects non-tensors without torch.
+
+    Construction must not import torch, so that a framework default serializer can
+    be selected on a torch-free install. Nothing is lost: an input that is actually
+    a torch.Tensor can only exist where torch is installed.
+    """
+    import sagemaker.core.serializers.base as ser_module
+
+    saved = {}
+    try:
+        saved = _block_torch()
+
+        serializer = ser_module.TorchTensorSerializer()
+
+        with pytest.raises(ValueError, match="not a torch.Tensor"):
+            serializer.serialize([1.0, 2.0, 3.0])
+    finally:
+        _restore_torch(saved)
+
+
+def test_torch_tensor_serializer_type_access_raises_import_error_without_torch():
+    """Verify the torch.Tensor type itself is still reported as missing."""
     import sagemaker.core.serializers.base as ser_module
 
     saved = {}
     try:
         saved = _block_torch()
 
+        serializer = ser_module.TorchTensorSerializer()
+
         with pytest.raises(ImportError, match="Unable to import torch"):
-            ser_module.TorchTensorSerializer()
+            serializer.torch_tensor  # pylint: disable=pointless-statement
     finally:
         _restore_torch(saved)
 

sagemaker-serve/src/sagemaker/serve/constants.py

@@ -144,6 +144,15 @@ def __getitem__(self, framework: Framework) -> Tuple:
             self._instances[framework] = (serializer_cls(), deserializer_cls())
         return self._instances[framework]
 
+    def __contains__(self, framework) -> bool:
+        """Return whether the framework has a default pair, without building it.
+
+        ``Mapping.__contains__`` would answer this by calling ``__getitem__``,
+        which defeats the laziness and surfaces an optional backend's
+        ``ImportError`` out of a membership test.
+        """
+        return framework in self._factories
+
     def __iter__(self):
         """Iterate over the frameworks that have a default pair."""
         return iter(self._factories)

sagemaker-serve/src/sagemaker/serve/marshalling/triton_translator.py

@@ -14,7 +14,13 @@ class TorchTensorTranslator:
     """Translate torch.Tensor from and to numpy.ndarray"""
 
     def __init__(self) -> None:
-        import torch
+        try:
+            import torch
+        except ImportError as e:
+            raise ImportError(
+                "Translating to torch.Tensor requires the torch module but it was not found "
+                "in your environment. Install it with 'pip install sagemaker-serve[torch]'."
+            ) from e
 
         self.convert_from_numpy = torch.from_numpy  # pylint: disable=E1101
         self.CONTENT_TYPE = "tensor/pt"

sagemaker-serve/tests/unit/test_optional_backend_dependencies.py

@@ -40,6 +40,7 @@
 )
 
 PYPROJECT = Path(__file__).resolve().parents[2] / "pyproject.toml"
+ROOT_PYPROJECT = Path(__file__).resolve().parents[3] / "pyproject.toml"
 OPTIONAL_BACKENDS = {"torch", "onnx", "onnxruntime", "tritonclient", "mlflow"}
 # Modules that must not be imported just by importing the package.
 MUST_NOT_IMPORT = sorted(OPTIONAL_BACKENDS | {"pytest"})
@@ -82,6 +83,20 @@ def test_test_extra_still_installs_the_backends(self, project):
 
         assert OPTIONAL_BACKENDS <= names
 
+    @pytest.mark.parametrize("extra", ["torch", "onnx", "triton", "mlflow"])
+    def test_top_level_package_passes_the_backend_extras_through(self, extra):
+        """``pip install sagemaker[torch]`` must work, not just ``sagemaker-serve[torch]``."""
+        if tomllib is None:
+            pytest.skip("tomllib requires Python 3.11+")
+        if not ROOT_PYPROJECT.exists():
+            pytest.skip("top-level pyproject.toml is not available outside the repo")
+
+        with open(ROOT_PYPROJECT, "rb") as handle:
+            root = tomllib.load(handle)["project"]
+
+        assert root["name"] == "sagemaker"
+        assert root["optional-dependencies"][extra] == [f"sagemaker-serve[{extra}]"]
+
 
 class TestNoEagerBackendImport:
     def test_importing_sagemaker_serve_needs_no_optional_backend(self):
@@ -117,6 +132,99 @@ def find_spec(self, name, path=None, target=None):
         assert result.returncode == 0, result.stderr
 
 
+class TestPyTorchDefaultsWithoutTorch:
+    def test_pytorch_default_pair_is_selectable_without_torch(self):
+        """Deploying a prebuilt PyTorch container needs no local torch: the framework is
+        detected from the image URI, so picking its default pair must not require the
+        backend. Subprocess so blocking torch cannot leak into the session."""
+        script = textwrap.dedent("""
+            import sys
+
+            class Blocker:
+                def find_spec(self, name, path=None, target=None):
+                    if name.split(".")[0] == "torch":
+                        raise ImportError("torch is blocked for this test")
+                    return None
+
+            sys.meta_path.insert(0, Blocker())
+
+            from sagemaker.serve.constants import DEFAULT_SERIALIZERS_BY_FRAMEWORK, Framework
+
+            serializer, deserializer = DEFAULT_SERIALIZERS_BY_FRAMEWORK[Framework.PYTORCH]
+
+            assert type(serializer).__name__ == "TorchTensorSerializer"
+            assert serializer.CONTENT_TYPE == "tensor/pt"
+            assert "torch" not in sys.modules
+            """)
+
+        result = subprocess.run(
+            [sys.executable, "-c", script], capture_output=True, text=True, check=False
+        )
+
+        assert result.returncode == 0, result.stderr
+
+    def test_serializer_rejects_non_tensors_without_torch(self):
+        """The no-torch path must stay a type error, not an ImportError."""
+        script = textwrap.dedent("""
+            import sys
+
+            class Blocker:
+                def find_spec(self, name, path=None, target=None):
+                    if name.split(".")[0] == "torch":
+                        raise ImportError("torch is blocked for this test")
+                    return None
+
+            sys.meta_path.insert(0, Blocker())
+
+            from sagemaker.core.serializers import TorchTensorSerializer
+
+            try:
+                TorchTensorSerializer().serialize([1.0, 2.0, 3.0])
+            except ValueError as e:
+                assert "not a torch.Tensor" in str(e), str(e)
+            else:
+                raise AssertionError("expected ValueError")
+            """)
+
+        result = subprocess.run(
+            [sys.executable, "-c", script], capture_output=True, text=True, check=False
+        )
+
+        assert result.returncode == 0, result.stderr
+
+
+class TestMissingBackendErrors:
+    def test_triton_translator_names_the_extra(self):
+        """A bare ``import torch`` here surfaced a raw ModuleNotFoundError that left the
+        caller no idea which extra to install. Subprocess so blocking torch is contained."""
+        script = textwrap.dedent("""
+            import sys
+
+            class Blocker:
+                def find_spec(self, name, path=None, target=None):
+                    if name.split(".")[0] == "torch":
+                        raise ImportError("torch is blocked for this test")
+                    return None
+
+            sys.meta_path.insert(0, Blocker())
+
+            from sagemaker.serve.marshalling.triton_translator import TorchTensorTranslator
+
+            try:
+                TorchTensorTranslator()
+            except ImportError as e:
+                assert "sagemaker-serve[torch]" in str(e), str(e)
+            else:
+                raise AssertionError("expected ImportError")
+            """)
+
+        result = subprocess.run(
+            [sys.executable, "-c", script], capture_output=True, text=True, check=False
+        )
+
+        assert result.returncode == 0, result.stderr
+
+
 class TestLazySerializerRegistry:
     def test_pairs_are_not_built_until_requested(self):
         built = []
@@ -148,3 +256,19 @@ def test_registry_still_behaves_like_the_mapping_it_replaced(self):
 
         assert serializer.CONTENT_TYPE == "application/x-npy"
         assert deserializer is not None
+
+    def test_membership_does_not_build_the_pair(self):
+        """``Mapping.__contains__`` answers via ``__getitem__``, which would both
+        defeat the laziness and raise a backend's ImportError out of an ``in`` test."""
+        built = []
+
+        class Boom:
+            def __init__(self):
+                built.append("built")
+                raise ImportError("backend missing")
+
+        registry = _LazySerializerRegistry({Framework.PYTORCH: (Boom, Boom)})
+
+        assert Framework.PYTORCH in registry
+        assert Framework.SKLEARN not in registry
+        assert built == []