diff --git a/pyproject.toml b/pyproject.toml index e99ce5c9f8..6df1184ac2 100644 --- a/pyproject.toml +++ b/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" diff --git a/sagemaker-core/src/sagemaker/core/serializers/base.py b/sagemaker-core/src/sagemaker/core/serializers/base.py index 9c01618889..381cd3cd8c 100644 --- a/sagemaker-core/src/sagemaker/core/serializers/base.py +++ b/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: diff --git a/sagemaker-core/tests/unit/test_optional_torch_dependency.py b/sagemaker-core/tests/unit/test_optional_torch_dependency.py index de084a41c3..ec34e8080e 100644 --- a/sagemaker-core/tests/unit/test_optional_torch_dependency.py +++ b/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) diff --git a/sagemaker-serve/src/sagemaker/serve/constants.py b/sagemaker-serve/src/sagemaker/serve/constants.py index a5a0c44ab6..1fd91f6339 100644 --- a/sagemaker-serve/src/sagemaker/serve/constants.py +++ b/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) diff --git a/sagemaker-serve/src/sagemaker/serve/marshalling/triton_translator.py b/sagemaker-serve/src/sagemaker/serve/marshalling/triton_translator.py index f7a941c2f2..ac6f43245e 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( + "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" diff --git a/sagemaker-serve/tests/unit/test_optional_backend_dependencies.py b/sagemaker-serve/tests/unit/test_optional_backend_dependencies.py index 3dff7a37e0..2c8c88d6ee 100644 --- a/sagemaker-serve/tests/unit/test_optional_backend_dependencies.py +++ b/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 == []