Skip to content

Commit 6a5d943

Browse files
mohamedzeidan2021Mohamed Zeidangoelakash
authored
fix: close the remaining gaps in optional torch support (#6531)
#6512 made the sagemaker-serve serving backends extras, but four gaps kept torch from being genuinely optional (#5531). All four were diagnosed first in #6166. Selecting the PyTorch default serializer still required torch. ``TorchTensorSerializer.__init__`` imported torch, so the lazy registry only deferred the failure to ``_fetch_serializer_and_deserializer_for_framework``. A client deploying a prebuilt PyTorch container detects the framework from the image URI and never touches a tensor, so it hit an ImportError for a dependency it had no use for. The type check now consults ``sys.modules``: a torch.Tensor cannot exist unless torch has already been imported, so the check stays exact without triggering an import. ``torch_tensor`` becomes a property, keeping the attribute and its ImportError for callers that want the type itself. Membership tests built the pair they were asking about. The registry inherited ``Mapping.__contains__``, which answers by calling ``__getitem__`` - defeating the laziness and raising a backend's ImportError out of an ``in`` test rather than returning True. ``TorchTensorTranslator`` raised a bare ModuleNotFoundError, leaving the caller to guess which extra to install. ``pip install sagemaker[torch]`` did not work. The top-level package now passes sagemaker-serve's backend extras through. ``all`` is unchanged, so no existing install starts pulling torch. Co-authored-by: Mohamed Zeidan <zeidmo@amazon.com> Co-authored-by: goelakash <goelakash93@gmail.com>
1 parent bdb9eee commit 6a5d943

6 files changed

Lines changed: 199 additions & 8 deletions

File tree

‎pyproject.toml‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,15 @@ serve = ["sagemaker-serve"]
4545
mlops = ["sagemaker-mlops"]
4646
all = ["sagemaker-train", "sagemaker-serve", "sagemaker-mlops"]
4747

48+
# sagemaker-serve's framework backends are optional, so pass its extras through
49+
# to let callers write 'pip install sagemaker[torch]' instead of having to name
50+
# the subpackage. 'all' is deliberately left as the subpackage list: installing
51+
# every backend is 'pip install sagemaker-serve[all]'.
52+
torch = ["sagemaker-serve[torch]"]
53+
onnx = ["sagemaker-serve[onnx]"]
54+
triton = ["sagemaker-serve[triton]"]
55+
mlflow = ["sagemaker-serve[mlflow]"]
56+
4857
[project.urls]
4958
Homepage = "https://github.com/aws/sagemaker-python-sdk"
5059

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

Lines changed: 24 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import csv
2020
import io
2121
import json
22+
import sys
2223
import numpy as np
2324
from pandas import DataFrame
2425
from six import with_metaclass
@@ -433,6 +434,18 @@ def serialize(self, data):
433434
raise ValueError(f"Object of type {type(data)} is not String serializable.")
434435

435436

437+
def _is_torch_tensor(data):
438+
"""Return whether ``data`` is a ``torch.Tensor``, without importing torch.
439+
440+
A ``torch.Tensor`` instance cannot exist unless torch has already been
441+
imported, so a missing ``sys.modules`` entry is proof that ``data`` is not
442+
one. This keeps the check exact while letting callers that never pass a
443+
tensor avoid pulling in an optional multi-hundred-megabyte dependency.
444+
"""
445+
torch = sys.modules.get("torch")
446+
return torch is not None and isinstance(data, torch.Tensor)
447+
448+
436449
class TorchTensorSerializer(SimpleBaseSerializer):
437450
"""Serialize torch.Tensor to a buffer by converting tensor to numpy and call NumpySerializer.
438451
@@ -444,17 +457,24 @@ class TorchTensorSerializer(SimpleBaseSerializer):
444457

445458
def __init__(self, content_type="tensor/pt"):
446459
super(TorchTensorSerializer, self).__init__(content_type=content_type)
460+
self.numpy_serializer = NumpySerializer()
461+
462+
@property
463+
def torch_tensor(self):
464+
"""The ``torch.Tensor`` type, imported on first access.
465+
466+
torch is an optional dependency (``sagemaker-core[torch]``), so it is not
467+
imported until a caller actually needs the type.
468+
"""
447469
try:
448470
from torch import Tensor
449-
450-
self.torch_tensor = Tensor
451471
except ImportError as e:
452472
raise ImportError(
453473
"Unable to import torch. Please install torch to use TorchTensorSerializer: "
454474
"pip install 'sagemaker-core[torch]'"
455475
) from e
456476

457-
self.numpy_serializer = NumpySerializer()
477+
return Tensor
458478

459479
def serialize(self, data):
460480
"""Serialize torch.Tensor to a buffer.
@@ -465,7 +485,7 @@ def serialize(self, data):
465485
Returns:
466486
raw-bytes: The data serialized as raw-bytes from the input.
467487
"""
468-
if isinstance(data, self.torch_tensor):
488+
if _is_torch_tensor(data):
469489
try:
470490
return self.numpy_serializer.serialize(data.detach().numpy())
471491
except Exception as e:

‎sagemaker-core/tests/unit/test_optional_torch_dependency.py‎

Lines changed: 26 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -122,16 +122,39 @@ def test_deserializer_module_imports_without_torch():
122122
), f"Subprocess failed:\nstdout: {result.stdout}\nstderr: {result.stderr}"
123123

124124

125-
def test_torch_tensor_serializer_raises_import_error_without_torch():
126-
"""Verify TorchTensorSerializer raises ImportError when torch is not installed."""
125+
def test_torch_tensor_serializer_constructs_without_torch():
126+
"""Verify TorchTensorSerializer can be built and rejects non-tensors without torch.
127+
128+
Construction must not import torch, so that a framework default serializer can
129+
be selected on a torch-free install. Nothing is lost: an input that is actually
130+
a torch.Tensor can only exist where torch is installed.
131+
"""
132+
import sagemaker.core.serializers.base as ser_module
133+
134+
saved = {}
135+
try:
136+
saved = _block_torch()
137+
138+
serializer = ser_module.TorchTensorSerializer()
139+
140+
with pytest.raises(ValueError, match="not a torch.Tensor"):
141+
serializer.serialize([1.0, 2.0, 3.0])
142+
finally:
143+
_restore_torch(saved)
144+
145+
146+
def test_torch_tensor_serializer_type_access_raises_import_error_without_torch():
147+
"""Verify the torch.Tensor type itself is still reported as missing."""
127148
import sagemaker.core.serializers.base as ser_module
128149

129150
saved = {}
130151
try:
131152
saved = _block_torch()
132153

154+
serializer = ser_module.TorchTensorSerializer()
155+
133156
with pytest.raises(ImportError, match="Unable to import torch"):
134-
ser_module.TorchTensorSerializer()
157+
serializer.torch_tensor # pylint: disable=pointless-statement
135158
finally:
136159
_restore_torch(saved)
137160

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

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -144,6 +144,15 @@ def __getitem__(self, framework: Framework) -> Tuple:
144144
self._instances[framework] = (serializer_cls(), deserializer_cls())
145145
return self._instances[framework]
146146

147+
def __contains__(self, framework) -> bool:
148+
"""Return whether the framework has a default pair, without building it.
149+
150+
``Mapping.__contains__`` would answer this by calling ``__getitem__``,
151+
which defeats the laziness and surfaces an optional backend's
152+
``ImportError`` out of a membership test.
153+
"""
154+
return framework in self._factories
155+
147156
def __iter__(self):
148157
"""Iterate over the frameworks that have a default pair."""
149158
return iter(self._factories)

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

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,13 @@ class TorchTensorTranslator:
1414
"""Translate torch.Tensor from and to numpy.ndarray"""
1515

1616
def __init__(self) -> None:
17-
import torch
17+
try:
18+
import torch
19+
except ImportError as e:
20+
raise ImportError(
21+
"Translating to torch.Tensor requires the torch module but it was not found "
22+
"in your environment. Install it with 'pip install sagemaker-serve[torch]'."
23+
) from e
1824

1925
self.convert_from_numpy = torch.from_numpy # pylint: disable=E1101
2026
self.CONTENT_TYPE = "tensor/pt"

‎sagemaker-serve/tests/unit/test_optional_backend_dependencies.py‎

Lines changed: 124 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@
4040
)
4141

4242
PYPROJECT = Path(__file__).resolve().parents[2] / "pyproject.toml"
43+
ROOT_PYPROJECT = Path(__file__).resolve().parents[3] / "pyproject.toml"
4344
OPTIONAL_BACKENDS = {"torch", "onnx", "onnxruntime", "tritonclient", "mlflow"}
4445
# Modules that must not be imported just by importing the package.
4546
MUST_NOT_IMPORT = sorted(OPTIONAL_BACKENDS | {"pytest"})
@@ -82,6 +83,20 @@ def test_test_extra_still_installs_the_backends(self, project):
8283

8384
assert OPTIONAL_BACKENDS <= names
8485

86+
@pytest.mark.parametrize("extra", ["torch", "onnx", "triton", "mlflow"])
87+
def test_top_level_package_passes_the_backend_extras_through(self, extra):
88+
"""``pip install sagemaker[torch]`` must work, not just ``sagemaker-serve[torch]``."""
89+
if tomllib is None:
90+
pytest.skip("tomllib requires Python 3.11+")
91+
if not ROOT_PYPROJECT.exists():
92+
pytest.skip("top-level pyproject.toml is not available outside the repo")
93+
94+
with open(ROOT_PYPROJECT, "rb") as handle:
95+
root = tomllib.load(handle)["project"]
96+
97+
assert root["name"] == "sagemaker"
98+
assert root["optional-dependencies"][extra] == [f"sagemaker-serve[{extra}]"]
99+
85100

86101
class TestNoEagerBackendImport:
87102
def test_importing_sagemaker_serve_needs_no_optional_backend(self):
@@ -117,6 +132,99 @@ def find_spec(self, name, path=None, target=None):
117132
assert result.returncode == 0, result.stderr
118133

119134

135+
class TestPyTorchDefaultsWithoutTorch:
136+
def test_pytorch_default_pair_is_selectable_without_torch(self):
137+
"""Deploying a prebuilt PyTorch container needs no local torch: the framework is
138+
detected from the image URI, so picking its default pair must not require the
139+
backend. Subprocess so blocking torch cannot leak into the session."""
140+
script = textwrap.dedent("""
141+
import sys
142+
143+
class Blocker:
144+
def find_spec(self, name, path=None, target=None):
145+
if name.split(".")[0] == "torch":
146+
raise ImportError("torch is blocked for this test")
147+
return None
148+
149+
sys.meta_path.insert(0, Blocker())
150+
151+
from sagemaker.serve.constants import DEFAULT_SERIALIZERS_BY_FRAMEWORK, Framework
152+
153+
serializer, deserializer = DEFAULT_SERIALIZERS_BY_FRAMEWORK[Framework.PYTORCH]
154+
155+
assert type(serializer).__name__ == "TorchTensorSerializer"
156+
assert serializer.CONTENT_TYPE == "tensor/pt"
157+
assert "torch" not in sys.modules
158+
""")
159+
160+
result = subprocess.run(
161+
[sys.executable, "-c", script], capture_output=True, text=True, check=False
162+
)
163+
164+
assert result.returncode == 0, result.stderr
165+
166+
def test_serializer_rejects_non_tensors_without_torch(self):
167+
"""The no-torch path must stay a type error, not an ImportError."""
168+
script = textwrap.dedent("""
169+
import sys
170+
171+
class Blocker:
172+
def find_spec(self, name, path=None, target=None):
173+
if name.split(".")[0] == "torch":
174+
raise ImportError("torch is blocked for this test")
175+
return None
176+
177+
sys.meta_path.insert(0, Blocker())
178+
179+
from sagemaker.core.serializers import TorchTensorSerializer
180+
181+
try:
182+
TorchTensorSerializer().serialize([1.0, 2.0, 3.0])
183+
except ValueError as e:
184+
assert "not a torch.Tensor" in str(e), str(e)
185+
else:
186+
raise AssertionError("expected ValueError")
187+
""")
188+
189+
result = subprocess.run(
190+
[sys.executable, "-c", script], capture_output=True, text=True, check=False
191+
)
192+
193+
assert result.returncode == 0, result.stderr
194+
195+
196+
class TestMissingBackendErrors:
197+
def test_triton_translator_names_the_extra(self):
198+
"""A bare ``import torch`` here surfaced a raw ModuleNotFoundError that left the
199+
caller no idea which extra to install. Subprocess so blocking torch is contained."""
200+
script = textwrap.dedent("""
201+
import sys
202+
203+
class Blocker:
204+
def find_spec(self, name, path=None, target=None):
205+
if name.split(".")[0] == "torch":
206+
raise ImportError("torch is blocked for this test")
207+
return None
208+
209+
sys.meta_path.insert(0, Blocker())
210+
211+
from sagemaker.serve.marshalling.triton_translator import TorchTensorTranslator
212+
213+
try:
214+
TorchTensorTranslator()
215+
except ImportError as e:
216+
assert "sagemaker-serve[torch]" in str(e), str(e)
217+
else:
218+
raise AssertionError("expected ImportError")
219+
""")
220+
221+
result = subprocess.run(
222+
[sys.executable, "-c", script], capture_output=True, text=True, check=False
223+
)
224+
225+
assert result.returncode == 0, result.stderr
226+
227+
120228
class TestLazySerializerRegistry:
121229
def test_pairs_are_not_built_until_requested(self):
122230
built = []
@@ -148,3 +256,19 @@ def test_registry_still_behaves_like_the_mapping_it_replaced(self):
148256

149257
assert serializer.CONTENT_TYPE == "application/x-npy"
150258
assert deserializer is not None
259+
260+
def test_membership_does_not_build_the_pair(self):
261+
"""``Mapping.__contains__`` answers via ``__getitem__``, which would both
262+
defeat the laziness and raise a backend's ImportError out of an ``in`` test."""
263+
built = []
264+
265+
class Boom:
266+
def __init__(self):
267+
built.append("built")
268+
raise ImportError("backend missing")
269+
270+
registry = _LazySerializerRegistry({Framework.PYTORCH: (Boom, Boom)})
271+
272+
assert Framework.PYTORCH in registry
273+
assert Framework.SKLEARN not in registry
274+
assert built == []

0 commit comments

Comments
 (0)