|
40 | 40 | ) |
41 | 41 |
|
42 | 42 | PYPROJECT = Path(__file__).resolve().parents[2] / "pyproject.toml" |
| 43 | +ROOT_PYPROJECT = Path(__file__).resolve().parents[3] / "pyproject.toml" |
43 | 44 | OPTIONAL_BACKENDS = {"torch", "onnx", "onnxruntime", "tritonclient", "mlflow"} |
44 | 45 | # Modules that must not be imported just by importing the package. |
45 | 46 | MUST_NOT_IMPORT = sorted(OPTIONAL_BACKENDS | {"pytest"}) |
@@ -82,6 +83,20 @@ def test_test_extra_still_installs_the_backends(self, project): |
82 | 83 |
|
83 | 84 | assert OPTIONAL_BACKENDS <= names |
84 | 85 |
|
| 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 | + |
85 | 100 |
|
86 | 101 | class TestNoEagerBackendImport: |
87 | 102 | def test_importing_sagemaker_serve_needs_no_optional_backend(self): |
@@ -117,6 +132,99 @@ def find_spec(self, name, path=None, target=None): |
117 | 132 | assert result.returncode == 0, result.stderr |
118 | 133 |
|
119 | 134 |
|
| 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 | + |
120 | 228 | class TestLazySerializerRegistry: |
121 | 229 | def test_pairs_are_not_built_until_requested(self): |
122 | 230 | built = [] |
@@ -148,3 +256,19 @@ def test_registry_still_behaves_like_the_mapping_it_replaced(self): |
148 | 256 |
|
149 | 257 | assert serializer.CONTENT_TYPE == "application/x-npy" |
150 | 258 | 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