fix(serve): make torch an optional dependency - #6166
Open
goelakash wants to merge 1 commit into
Open
Conversation
goelakash
had a problem deploying
to
manual-approval
August 7, 2026 03:29 — with 
GitHub Actions
Error
GitHub Actions
Error
goelakash
had a problem deploying
to
manual-approval
August 7, 2026 03:29 — with
GitHub Actions
Error
GitHub Actions
Error
goelakash
had a problem deploying
to
manual-approval
August 7, 2026 03:29 — with
GitHub Actions
Error
GitHub Actions
Error
goelakash
force-pushed
the
fix-optional-torch-serve
branch
from
August 7, 2026 03:32
1725779 to
38b75f5
Compare
goelakash
temporarily deployed
to
manual-approval
August 7, 2026 03:32 — with
GitHub Actions
Inactive
GitHub Actions
Inactive
goelakash
had a problem deploying
to
manual-approval
August 7, 2026 03:32 — with
GitHub Actions
Error
GitHub Actions
Error
goelakash
temporarily deployed
to
manual-approval
August 7, 2026 03:32 — with
GitHub Actions
Inactive
GitHub Actions
Inactive
goelakash
marked this pull request as ready for review
August 7, 2026 03:33
goelakash
temporarily deployed
to
manual-approval
August 7, 2026 03:33 — with
GitHub Actions
Inactive
GitHub Actions
Inactive
Installing sagemaker-serve (or the umbrella sagemaker, which depends on it) pulled torch, and on Linux the ~2.9 GB CUDA/cuDNN/NCCL closure, even for API-only use. sagemaker-core already declared torch an extra; sagemaker-serve declared it required, so serve users paid regardless. The import-time blocker was DEFAULT_SERIALIZERS_BY_FRAMEWORK in serve/constants.py, which instantiated TorchTensorSerializer() at module scope, running `from torch import Tensor` on any `import sagemaker.serve`. Changes: - Store serializer/deserializer classes (not instances) in DEFAULT_SERIALIZERS_BY_FRAMEWORK; instantiate on lookup. Removes the import-time torch dependency. - Move torch>=2.0.0 to a `torch` extra in sagemaker-serve, matching sagemaker-core. Add a `torch` extra to the umbrella sagemaker package (sagemaker-serve[torch]) so `pip install sagemaker[torch]` works. - Duck-type the tensor check in TorchTensorSerializer: serialization only needs the object's own detach()/numpy(), so it no longer imports torch at all and works whenever the caller holds a tensor. - Make the in-process model server's torch import lazy with a CPU fallback, so importing it no longer requires torch. - Fix the error messages that told serve users to install the wrong package: TorchTensorDeserializer, the Triton translator, and the ONNX export path now name sagemaker-serve[torch]. - Update and add tests asserting the whole surface imports and serializes without torch, and that the paths that genuinely construct torch objects (deserializer, Triton translator) still raise a clear error naming the extra. Fixes aws#5531 Breaking change: users who installed plain sagemaker-serve and relied on a torch code path (deserializing to tensor/pt, or ONNX/Triton export of a PyTorch model) must now install sagemaker-serve[torch]. Both paths inherently require a torch object the caller supplies, so torch is already present in practice; a GitHub-wide search found no external callers of these paths that do not already import torch. --- X-AI-Prompt: address the reviewer concern on whether making torch optional breaks existing customers; find ways to reduce the torch dependency and update pyproject and docs; consolidate all changes in the goelakash93 fork for a possible major version bump. X-AI-Tool: Claude Code
goelakash
force-pushed
the
fix-optional-torch-serve
branch
from
August 7, 2026 20:49
38b75f5 to
299438e
Compare
goelakash
requested a deployment
to
manual-approval
August 7, 2026 20:49 — with
GitHub Actions
Waiting
GitHub Actions
Waiting
goelakash
requested a deployment
to
manual-approval
August 7, 2026 20:49 — with
GitHub Actions
Waiting
GitHub Actions
Waiting
goelakash
requested a deployment
to
manual-approval
August 7, 2026 20:49 — with
GitHub Actions
Waiting
GitHub Actions
Waiting
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #5531
sagemaker-serveliststorch>=2.0.0as a required dependency, so any install ofsagemaker,sagemaker-serve, orsagemaker-mlopspulls torch (and on Linux, the CUDA/cuDNN/NCCL stack) even for API-only use.sagemaker-corealready treats torch as an extra.The blocker was
DEFAULT_SERIALIZERS_BY_FRAMEWORKinserve/constants.py, which instantiatedTorchTensorSerializer()at module scope. That runsfrom torch import Tensoron anyimport sagemaker.serve. Every other torch reference in the package is already lazy.Changes
DEFAULT_SERIALIZERS_BY_FRAMEWORK; instantiate at lookup in_fetch_serializer_and_deserializer_for_framework.torch>=2.0.0to atorchextra, matchingsagemaker-core.tests/unit/test_optional_torch_dependency.py, mirroring the existingsagemaker-coresubprocess pattern.Installing with the
torchextra is unchanged.TorchTensorSerializer()still raises the sameImportErrorif torch is missing when actually used.Testing
tests/unitin a venv with and without torch: no new failures vs.master(53 pre-existing failures in both).masterand pass with this change.from sagemaker.serve import ModelBuilderandimport sagemaker.mlopsboth succeed.