diff --git a/pyproject.toml b/pyproject.toml index bde73bd982..e52df54d09 100644 --- a/pyproject.toml +++ b/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] diff --git a/sagemaker-core/src/sagemaker/core/deserializers/base.py b/sagemaker-core/src/sagemaker/core/deserializers/base.py index 03138ed577..6dbd5c8cef 100644 --- a/sagemaker-core/src/sagemaker/core/deserializers/base.py +++ b/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"): diff --git a/sagemaker-core/src/sagemaker/core/serializers/base.py b/sagemaker-core/src/sagemaker/core/serializers/base.py index 84b9832c63..e8fbb97285 100644 --- a/sagemaker-core/src/sagemaker/core/serializers/base.py +++ b/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: diff --git a/sagemaker-core/tests/unit/test_optional_torch_dependency.py b/sagemaker-core/tests/unit/test_optional_torch_dependency.py index 2b7efbc227..da68397c73 100644 --- a/sagemaker-core/tests/unit/test_optional_torch_dependency.py +++ b/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) diff --git a/sagemaker-serve/pyproject.toml b/sagemaker-serve/pyproject.toml index d98a84827f..2300a9967d 100644 --- a/sagemaker-serve/pyproject.toml +++ b/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", diff --git a/sagemaker-serve/src/sagemaker/serve/constants.py b/sagemaker-serve/src/sagemaker/serve/constants.py index 591f851a71..98e8bef1af 100644 --- a/sagemaker-serve/src/sagemaker/serve/constants.py +++ b/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), } diff --git a/sagemaker-serve/src/sagemaker/serve/marshalling/triton_translator.py b/sagemaker-serve/src/sagemaker/serve/marshalling/triton_translator.py index f7a941c2f2..05759129db 100644 --- a/sagemaker-serve/src/sagemaker/serve/marshalling/triton_translator.py +++ b/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" diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder_utils.py b/sagemaker-serve/src/sagemaker/serve/model_builder_utils.py index f4a90a2d05..21364a5ecb 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder_utils.py +++ b/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." diff --git a/sagemaker-serve/src/sagemaker/serve/model_server/in_process_model_server/app.py b/sagemaker-serve/src/sagemaker/serve/model_server/in_process_model_server/app.py index edf8b5748a..fabcb42bf7 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_server/in_process_model_server/app.py +++ b/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: diff --git a/sagemaker-serve/tests/unit/test_constants.py b/sagemaker-serve/tests/unit/test_constants.py index 43187a0f75..c69bfd57f8 100644 --- a/sagemaker-serve/tests/unit/test_constants.py +++ b/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__": diff --git a/sagemaker-serve/tests/unit/test_model_builder_utils_new.py b/sagemaker-serve/tests/unit/test_model_builder_utils_new.py index 9c8a59a2c9..53d64fdad8 100644 --- a/sagemaker-serve/tests/unit/test_model_builder_utils_new.py +++ b/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) diff --git a/sagemaker-serve/tests/unit/test_optional_torch_dependency.py b/sagemaker-serve/tests/unit/test_optional_torch_dependency.py new file mode 100644 index 0000000000..32e290f56d --- /dev/null +++ b/sagemaker-serve/tests/unit/test_optional_torch_dependency.py @@ -0,0 +1,163 @@ +# 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_tensor_serializer_works_without_torch(): + """Serializing a tensor needs only its own detach()/numpy(), not the torch import.""" + result = _run_without_torch( + """ + import numpy as np + + from sagemaker.core.serializers import TorchTensorSerializer + + class FakeTensor: + __module__ = "torch" + + def detach(self): + return self + + def numpy(self): + return np.array([1, 2, 3]) + + assert TorchTensorSerializer().serialize(FakeTensor()) + print("OK") + """ + ) + assert "OK" in result.stdout, result.stderr + + +def test_tensor_serializer_still_rejects_non_tensors(): + """The structural check must not widen to arbitrary objects.""" + result = _run_without_torch( + """ + from sagemaker.core.serializers import TorchTensorSerializer + + try: + TorchTensorSerializer().serialize([1, 2, 3]) + raise AssertionError("expected ValueError") + except ValueError: + print("OK") + """ + ) + assert "OK" in result.stdout, result.stderr + + +def test_in_process_server_imports_without_torch(): + """The in-process server module must not import torch at module scope.""" + result = _run_without_torch( + """ + from sagemaker.serve.model_server.in_process_model_server.app import InProcessServer + + assert InProcessServer is not None + print("OK") + """ + ) + assert "OK" in result.stdout, result.stderr + + +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_tensor_deserializer_names_the_extra(): + """Deserializing must construct a real tensor, so it still needs torch - and the + error must name an extra the caller can actually install.""" + result = _run_without_torch( + """ + from sagemaker.core.deserializers import TorchTensorDeserializer + + try: + TorchTensorDeserializer() + raise AssertionError("expected ImportError") + except ImportError as e: + assert "[torch]" in str(e) + print("OK") + """ + ) + assert "OK" in result.stdout, result.stderr + + +def test_triton_translator_names_the_extra(): + """Translating to torch.Tensor needs torch; the error must name the extra.""" + result = _run_without_torch( + """ + from sagemaker.serve.marshalling.triton_translator import TorchTensorTranslator + + try: + TorchTensorTranslator() + raise AssertionError("expected ImportError") + except ImportError as e: + assert "sagemaker-serve[torch]" in str(e) + print("OK") + """ + ) + assert "OK" in result.stdout, result.stderr