[goelakash]

Fixes #5531

sagemaker-serve lists torch>=2.0.0 as a required dependency, so any install of sagemaker, sagemaker-serve, or sagemaker-mlops pulls torch (and on Linux, the CUDA/cuDNN/NCCL stack) even for API-only use. sagemaker-core already treats torch as an extra.

The blocker was DEFAULT_SERIALIZERS_BY_FRAMEWORK in serve/constants.py, which instantiated TorchTensorSerializer() at module scope. That runs from torch import Tensor on any import sagemaker.serve. Every other torch reference in the package is already lazy.

Changes

• Store serializer/deserializer classes in DEFAULT_SERIALIZERS_BY_FRAMEWORK; instantiate at lookup in _fetch_serializer_and_deserializer_for_framework.
• Move torch>=2.0.0 to a torch extra, matching sagemaker-core.
• Update the two docstring examples and the affected unit tests.
• Add tests/unit/test_optional_torch_dependency.py, mirroring the existing sagemaker-core subprocess pattern.

Installing with the torch extra is unchanged. TorchTensorSerializer() still raises the same ImportError if torch is missing when actually used.

Testing

• tests/unit in a venv with and without torch: no new failures vs. master (53 pre-existing failures in both).
• New tests fail on master and pass with this change.
• Verified in a container with no torch installed: from sagemaker.serve import ModelBuilder and import sagemaker.mlops both succeed.
• Image size for all four packages on Amazon Linux 2023: 1.45 GB -> 536 MB.

[mohamedzeidan2021]

Thanks for this @goelakash, and apologies for how long it sat. Doing a proper deep dive on it now — here's where things landed.

Your diagnosis was right, and it was the correct read of the problem: DEFAULT_SERIALIZERS_BY_FRAMEWORK instantiating TorchTensorSerializer() at module scope was the actual blocker, and every other torch reference in the package was already lazy.

While this PR was waiting, #6512 landed (Oct 21) and independently fixed the same root cause with a wider scope — it also made onnx, onnxruntime, tritonclient[http] and mlflow optional, and used a _LazySerializerRegistry(Mapping) so that DEFAULT_SERIALIZERS_BY_FRAMEWORK[framework] keeps returning instances. That avoids the API change of handing callers classes, which is what forced the churn in test_constants.py and test_model_builder_utils_new.py here. That's purely an accident of timing and review latency on our side, not a judgment on this PR — and it's the reason all four of the current conflicts are in regions where #6512's version is the broader fix.

So rather than ask you to resolve conflicts into a change that would mostly be reverted, I've opened #6531 to carry forward the parts of this PR that #6512 did not cover, with you credited as co-author on the commit:

1. TorchTensorSerializer.__init__ still imports torch. The lazy registry only defers the failure to _fetch_serializer_and_deserializer_for_framework, so a torch-free client deploying a prebuilt PyTorch container (framework detected from the image URI, no tensor ever handled) still broke. This was the most valuable catch in your PR and fix: make sagemaker-serve test tooling and serving backends optional #6512 missed it.
2. TorchTensorTranslator's bare import torch — your friendly ImportError naming the extra.
3. pip install sagemaker[torch] — your root pyproject.toml passthrough extra, extended to onnx/triton/mlflow for symmetry.
4. Plus one defect in fix: make sagemaker-serve test tooling and serving backends optional #6512 itself that your PR's structure would have avoided: the registry inherited Mapping.__contains__, which answers via __getitem__, so framework in DEFAULT_SERIALIZERS_BY_FRAMEWORK instantiated the pair as a side effect and raised ImportError out of a membership test.

One change I made to your approach on (1): instead of the structural check

type(data).__module__.split(".")[0] == "torch" and callable(getattr(data, "detach", None)) ...

#6531 uses sys.modules.get("torch") and then a real isinstance. A torch.Tensor can't exist unless torch is already imported, so that's exactly as strict as the original isinstance — no widening to objects that merely look like tensors — while still never triggering an import.

Closing this in favor of #6531. Thank you for the diagnosis, the measurements (1.45 GB → 536 MB is what made the case concrete), and for your patience with the delay.