fix(serve): make torch an optional dependency

Installing sagemaker-serve (or the umbrella sagemaker, which depends on
it) pulled torch, and on Linux the ~2.9 GB CUDA/cuDNN/NCCL closure, even
for API-only use. sagemaker-core already declared torch an extra;
sagemaker-serve declared it required, so serve users paid regardless.

The import-time blocker was DEFAULT_SERIALIZERS_BY_FRAMEWORK in
serve/constants.py, which instantiated TorchTensorSerializer() at module
scope, running `from torch import Tensor` on any `import sagemaker.serve`.

Changes:

- Store serializer/deserializer classes (not instances) in
  DEFAULT_SERIALIZERS_BY_FRAMEWORK; instantiate on lookup. Removes the
  import-time torch dependency.
- Move torch>=2.0.0 to a `torch` extra in sagemaker-serve, matching
  sagemaker-core. Add a `torch` extra to the umbrella sagemaker package
  (sagemaker-serve[torch]) so `pip install sagemaker[torch]` works.
- Duck-type the tensor check in TorchTensorSerializer: serialization only
  needs the object's own detach()/numpy(), so it no longer imports torch
  at all and works whenever the caller holds a tensor.
- Make the in-process model server's torch import lazy with a CPU
  fallback, so importing it no longer requires torch.
- Fix the error messages that told serve users to install the wrong
  package: TorchTensorDeserializer, the Triton translator, and the ONNX
  export path now name sagemaker-serve[torch].
- Update and add tests asserting the whole surface imports and serializes
  without torch, and that the paths that genuinely construct torch
  objects (deserializer, Triton translator) still raise a clear error
  naming the extra.

Fixes #5531

Breaking change: users who installed plain sagemaker-serve and relied on
a torch code path (deserializing to tensor/pt, or ONNX/Triton export of a
PyTorch model) must now install sagemaker-serve[torch]. Both paths
inherently require a torch object the caller supplies, so torch is
already present in practice; a GitHub-wide search found no external
callers of these paths that do not already import torch.

---

X-AI-Prompt: address the reviewer concern on whether making torch optional breaks existing customers; find ways to reduce the torch dependency and update pyproject and docs; consolidate all changes in the goelakash93 fork for a possible major version bump.
X-AI-Tool: Claude Code

Author: goelakash

pyproject.toml

@@ -42,6 +42,7 @@ dependencies = [
 train = ["sagemaker-train"]
 serve = ["sagemaker-serve"]
 mlops = ["sagemaker-mlops"]
+torch = ["sagemaker-serve[torch]"]
 all = ["sagemaker-train", "sagemaker-serve", "sagemaker-mlops"]
 
 [project.urls]

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

@@ -368,7 +368,8 @@ def __init__(self, accept="tensor/pt"):
         except ImportError as e:
             raise ImportError(
                 "Unable to import torch. Please install torch to use TorchTensorDeserializer: "
-                "pip install 'sagemaker-core[torch]'"
+                "pip install 'sagemaker-core[torch]' "
+                "(or 'sagemaker-serve[torch]' if you installed sagemaker-serve)"
             ) from e
 
     def deserialize(self, stream, content_type="tensor/pt"):

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

@@ -443,17 +443,22 @@ class TorchTensorSerializer(SimpleBaseSerializer):
 
     def __init__(self, content_type="tensor/pt"):
         super(TorchTensorSerializer, self).__init__(content_type=content_type)
-        try:
-            from torch import Tensor
+        self.numpy_serializer = NumpySerializer()
 
-            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
+    @staticmethod
+    def _is_torch_tensor(data):
+        """Recognize a torch.Tensor without importing torch.
 
-        self.numpy_serializer = NumpySerializer()
+        Serialization only needs the tensor's own detach()/numpy() methods, so
+        the type is identified structurally. A caller holding a real tensor
+        already has torch installed; importing it here would add nothing but a
+        multi-hundred-megabyte dependency for everyone else.
+        """
+        return (
+            type(data).__module__.split(".")[0] == "torch"
+            and callable(getattr(data, "detach", None))
+            and callable(getattr(data, "numpy", None))
+        )
 
     def serialize(self, data):
         """Serialize torch.Tensor to a buffer.
@@ -464,7 +469,7 @@ def serialize(self, data):
         Returns:
             raw-bytes: The data serialized as raw-bytes from the input.
         """
-        if isinstance(data, self.torch_tensor):
+        if self._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

@@ -121,16 +121,30 @@ def test_deserializer_module_imports_without_torch():
     )
 
 
-def test_torch_tensor_serializer_raises_import_error_without_torch():
-    """Verify TorchTensorSerializer raises ImportError when torch is not installed."""
+def test_torch_tensor_serializer_works_without_torch():
+    """Serializing a tensor needs only its own detach()/numpy(), not the torch import,
+    so TorchTensorSerializer must construct and serialize without torch installed."""
+    import numpy as np
+
     import sagemaker.core.serializers.base as ser_module
 
     saved = {}
     try:
         saved = _block_torch()
 
-        with pytest.raises(ImportError, match="Unable to import torch"):
-            ser_module.TorchTensorSerializer()
+        class FakeTensor:
+            __module__ = "torch"
+
+            def detach(self):
+                return self
+
+            def numpy(self):
+                return np.array([1, 2, 3])
+
+        serializer = ser_module.TorchTensorSerializer()
+        assert serializer.serialize(FakeTensor())
+        with pytest.raises(ValueError, match="not a torch.Tensor"):
+            serializer.serialize([1, 2, 3])
     finally:
         _restore_torch(saved)
 

sagemaker-serve/pyproject.toml

@@ -34,11 +34,13 @@ dependencies = [
     "psutil",
     "tritonclient[http]",
     "onnx",
-    "onnxruntime",
-    "torch>=2.0.0"
+    "onnxruntime"
 ]
 
 [project.optional-dependencies]
+torch = [
+    "torch>=2.0.0",
+]
 test = [
     "pytest",
     "pytest-cov",

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

@@ -20,17 +20,18 @@
 
 Example:
     Using Framework enum::
-    
+
         from sagemaker.serve.constants import Framework, DEFAULT_SERIALIZERS_BY_FRAMEWORK
-        
+
         # Get serializers for PyTorch
-        serializer, deserializer = DEFAULT_SERIALIZERS_BY_FRAMEWORK[Framework.PYTORCH]
+        serializer_cls, deserializer_cls = DEFAULT_SERIALIZERS_BY_FRAMEWORK[Framework.PYTORCH]
+        serializer, deserializer = serializer_cls(), deserializer_cls()
 """
 from __future__ import absolute_import, annotations
 
 # Standard library imports
 from enum import Enum
-from typing import Dict, Set, Tuple
+from typing import Callable, Dict, Set, Tuple
 
 # SageMaker imports
 from sagemaker.serve.mode.function_pointers import Mode
@@ -96,7 +97,8 @@ class Framework(Enum):
         Using framework enum::
         
             if detected_framework == Framework.PYTORCH:
-                serializer, deserializer = DEFAULT_SERIALIZERS_BY_FRAMEWORK[Framework.PYTORCH]
+                ser_cls, deser_cls = DEFAULT_SERIALIZERS_BY_FRAMEWORK[Framework.PYTORCH]
+                serializer, deserializer = ser_cls(), deser_cls()
     """
     XGBOOST = "XGBoost"
     LDA = "LDA"
@@ -116,18 +118,18 @@ class Framework(Enum):
 # Framework Serialization Mapping
 # ========================================
 
-DEFAULT_SERIALIZERS_BY_FRAMEWORK: Dict[Framework, Tuple] = {
-    Framework.XGBOOST: (LibSVMSerializer(), CSVDeserializer()),
-    Framework.LDA: (RecordSerializer(), RecordDeserializer()),
-    Framework.PYTORCH: (TorchTensorSerializer(), JSONDeserializer()),
-    Framework.TENSORFLOW: (NumpySerializer(), JSONDeserializer()),
-    Framework.MXNET: (RecordSerializer(), JSONDeserializer()),
-    Framework.CHAINER: (NumpySerializer(), JSONDeserializer()),
-    Framework.SKLEARN: (NumpySerializer(), NumpyDeserializer()),
-    Framework.HUGGINGFACE: (JSONSerializer(), JSONDeserializer()),
-    Framework.DJL: (JSONSerializer(), JSONDeserializer()),
-    Framework.SPARKML: (NumpySerializer(), JSONDeserializer()),
-    Framework.NTM: (RecordSerializer(), JSONDeserializer()),
-    Framework.SMD: (JSONSerializer(), JSONDeserializer()),
+DEFAULT_SERIALIZERS_BY_FRAMEWORK: Dict[Framework, Tuple[Callable, Callable]] = {
+    Framework.XGBOOST: (LibSVMSerializer, CSVDeserializer),
+    Framework.LDA: (RecordSerializer, RecordDeserializer),
+    Framework.PYTORCH: (TorchTensorSerializer, JSONDeserializer),
+    Framework.TENSORFLOW: (NumpySerializer, JSONDeserializer),
+    Framework.MXNET: (RecordSerializer, JSONDeserializer),
+    Framework.CHAINER: (NumpySerializer, JSONDeserializer),
+    Framework.SKLEARN: (NumpySerializer, NumpyDeserializer),
+    Framework.HUGGINGFACE: (JSONSerializer, JSONDeserializer),
+    Framework.DJL: (JSONSerializer, JSONDeserializer),
+    Framework.SPARKML: (NumpySerializer, JSONDeserializer),
+    Framework.NTM: (RecordSerializer, JSONDeserializer),
+    Framework.SMD: (JSONSerializer, JSONDeserializer),
 }
 

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(
+                "Unable to import torch. Translating to torch.Tensor requires torch: "
+                "pip install 'sagemaker-serve[torch]'"
+            ) from e
 
         self.convert_from_numpy = torch.from_numpy  # pylint: disable=E1101
         self.CONTENT_TYPE = "tensor/pt"

sagemaker-serve/src/sagemaker/serve/model_builder_utils.py

@@ -1288,7 +1288,8 @@ def _fetch_serializer_and_deserializer_for_framework(self, framework: str) -> Tu
         """
         framework_enum = self._normalize_framework_to_enum(framework)
         if framework_enum and framework_enum in DEFAULT_SERIALIZERS_BY_FRAMEWORK:
-            return DEFAULT_SERIALIZERS_BY_FRAMEWORK[framework_enum]
+            serializer_cls, deserializer_cls = DEFAULT_SERIALIZERS_BY_FRAMEWORK[framework_enum]
+            return serializer_cls(), deserializer_cls()
         return NumpySerializer(), JSONDeserializer()
 
     def _normalize_framework_to_enum(
@@ -3036,9 +3037,10 @@ def _export_pytorch_to_onnx(
 
         except ModuleNotFoundError:
             raise ImportError(
-                "Launching Triton with ModelBuilder for a PyTorch model requires onnx module "
-                "but it was not found in your environment. "
-                "Checkout the instructions on the installation page of its repo: "
+                "Launching Triton with ModelBuilder for a PyTorch model requires the torch and "
+                "onnx modules but one of them was not found in your environment. "
+                "Install torch with: pip install 'sagemaker-serve[torch]'. "
+                "For onnx, check the instructions on the installation page of its repo: "
                 "https://onnxruntime.ai/docs/install/ "
                 "And follow the ones that match your environment. "
                 "Please note that you may need to restart your runtime after installation."

sagemaker-serve/src/sagemaker/serve/model_server/in_process_model_server/app.py

@@ -6,7 +6,6 @@
 import io
 import logging
 import threading
-import torch
 from typing import Optional, Type
 
 from sagemaker.serve.spec.inference_spec import InferenceSpec
@@ -62,7 +61,12 @@ def __init__(
                         "Unable to import transformers, check if transformers is installed."
                     )
 
-                device = 0 if torch.cuda.is_available() else -1
+                try:
+                    import torch
+
+                    device = 0 if torch.cuda.is_available() else -1
+                except ImportError:
+                    device = -1
 
                 self._load_model = pipeline(task, model=self.model, device=device)
             except Exception:

sagemaker-serve/tests/unit/test_constants.py

@@ -56,24 +56,24 @@ class TestDefaultSerializersByFramework(unittest.TestCase):
     def test_all_frameworks_have_serializers(self):
         for framework in Framework:
             self.assertIn(framework, DEFAULT_SERIALIZERS_BY_FRAMEWORK)
-            serializer, deserializer = DEFAULT_SERIALIZERS_BY_FRAMEWORK[framework]
-            self.assertIsNotNone(serializer)
-            self.assertIsNotNone(deserializer)
+            serializer_cls, deserializer_cls = DEFAULT_SERIALIZERS_BY_FRAMEWORK[framework]
+            self.assertIsNotNone(serializer_cls)
+            self.assertIsNotNone(deserializer_cls)
 
     def test_pytorch_serializers(self):
-        serializer, deserializer = DEFAULT_SERIALIZERS_BY_FRAMEWORK[Framework.PYTORCH]
-        self.assertEqual(serializer.__class__.__name__, "TorchTensorSerializer")
-        self.assertEqual(deserializer.__class__.__name__, "JSONDeserializer")
+        serializer_cls, deserializer_cls = DEFAULT_SERIALIZERS_BY_FRAMEWORK[Framework.PYTORCH]
+        self.assertEqual(serializer_cls.__name__, "TorchTensorSerializer")
+        self.assertEqual(deserializer_cls.__name__, "JSONDeserializer")
 
     def test_tensorflow_serializers(self):
-        serializer, deserializer = DEFAULT_SERIALIZERS_BY_FRAMEWORK[Framework.TENSORFLOW]
-        self.assertEqual(serializer.__class__.__name__, "NumpySerializer")
-        self.assertEqual(deserializer.__class__.__name__, "JSONDeserializer")
+        serializer_cls, deserializer_cls = DEFAULT_SERIALIZERS_BY_FRAMEWORK[Framework.TENSORFLOW]
+        self.assertEqual(serializer_cls.__name__, "NumpySerializer")
+        self.assertEqual(deserializer_cls.__name__, "JSONDeserializer")
 
     def test_sklearn_serializers(self):
-        serializer, deserializer = DEFAULT_SERIALIZERS_BY_FRAMEWORK[Framework.SKLEARN]
-        self.assertEqual(serializer.__class__.__name__, "NumpySerializer")
-        self.assertEqual(deserializer.__class__.__name__, "NumpyDeserializer")
+        serializer_cls, deserializer_cls = DEFAULT_SERIALIZERS_BY_FRAMEWORK[Framework.SKLEARN]
+        self.assertEqual(serializer_cls.__name__, "NumpySerializer")
+        self.assertEqual(deserializer_cls.__name__, "NumpyDeserializer")
 
 
 if __name__ == "__main__":