fix(serve): make torch an optional dependency

DEFAULT_SERIALIZERS_BY_FRAMEWORK instantiated TorchTensorSerializer at
module scope, so importing sagemaker.serve required torch. Store the
classes instead and instantiate on lookup, and move torch to an extra.

Fixes #5531

Author: goelakash

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,20 @@ 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()),
+# Values are classes, not instances: instantiating TorchTensorSerializer at
+# import time forces `import torch` on every consumer of this module.
+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/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(

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__":

sagemaker-serve/tests/unit/test_model_builder_utils_new.py

@@ -470,11 +470,14 @@ def test_fetch_serializer_for_known_framework(self, mock_default_serializers):
         """Test fetching serializer for known framework."""
         mock_serializer = Mock()
         mock_deserializer = Mock()
-        mock_default_serializers.__getitem__.return_value = (mock_serializer, mock_deserializer)
+        mock_default_serializers.__getitem__.return_value = (
+            Mock(return_value=mock_serializer),
+            Mock(return_value=mock_deserializer),
+        )
         mock_default_serializers.__contains__.return_value = True
-        
+
         serializer, deserializer = self.utils._fetch_serializer_and_deserializer_for_framework("pytorch")
-        
+
         self.assertEqual(serializer, mock_serializer)
         self.assertEqual(deserializer, mock_deserializer)
 

sagemaker-serve/tests/unit/test_optional_torch_dependency.py

@@ -0,0 +1,96 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""Tests to verify torch dependency is optional in sagemaker-serve.
+
+Runs in a subprocess so that blocking torch cannot affect other tests in
+the session.
+"""
+from __future__ import absolute_import
+
+import subprocess
+import sys
+import textwrap
+
+
+def _run_without_torch(body):
+    """Execute body in a subprocess where importing torch raises ImportError."""
+    script = textwrap.dedent(
+        """
+        import sys
+
+        class _TorchBlocker:
+            \"\"\"Meta path finder that makes `import torch` fail.\"\"\"
+
+            def find_spec(self, name, path=None, target=None):
+                if name == "torch" or name.startswith("torch."):
+                    raise ImportError("torch is blocked for this test")
+                return None
+
+        sys.meta_path.insert(0, _TorchBlocker())
+        for _name in [n for n in sys.modules if n == "torch" or n.startswith("torch.")]:
+            del sys.modules[_name]
+        """
+    ) + textwrap.dedent(body)
+    return subprocess.run(
+        [sys.executable, "-c", script], capture_output=True, text=True
+    )
+
+
+def test_constants_imports_without_torch():
+    """serve.constants must not instantiate TorchTensorSerializer at import time."""
+    result = _run_without_torch(
+        """
+        from sagemaker.serve.constants import (
+            Framework,
+            DEFAULT_SERIALIZERS_BY_FRAMEWORK,
+        )
+
+        assert Framework.PYTORCH in DEFAULT_SERIALIZERS_BY_FRAMEWORK
+        print("OK")
+        """
+    )
+    assert "OK" in result.stdout, result.stderr
+
+
+def test_model_builder_imports_without_torch():
+    """ModelBuilder must be importable for API-only use without torch installed."""
+    result = _run_without_torch(
+        """
+        from sagemaker.serve import ModelBuilder
+
+        assert ModelBuilder is not None
+        print("OK")
+        """
+    )
+    assert "OK" in result.stdout, result.stderr
+
+
+def test_pytorch_serializer_still_requires_torch():
+    """Resolving the PyTorch entry must still raise a clear error without torch."""
+    result = _run_without_torch(
+        """
+        from sagemaker.serve.constants import (
+            Framework,
+            DEFAULT_SERIALIZERS_BY_FRAMEWORK,
+        )
+
+        serializer_cls, _ = DEFAULT_SERIALIZERS_BY_FRAMEWORK[Framework.PYTORCH]
+        try:
+            serializer_cls()
+            raise AssertionError("expected ImportError")
+        except ImportError as e:
+            assert "torch" in str(e)
+            print("OK")
+        """
+    )
+    assert "OK" in result.stdout, result.stderr