Skip to content

fix(serve): make torch an optional dependency - #6166

Open
goelakash wants to merge 1 commit into
aws:masterfrom
goelakash:fix-optional-torch-serve
Open

fix(serve): make torch an optional dependency#6166
goelakash wants to merge 1 commit into
aws:masterfrom
goelakash:fix-optional-torch-serve

Conversation

@goelakash

Copy link
Copy Markdown
Contributor

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.

DEFAULT_SERIALIZERS_BY_FRAMEWORK instantiated TorchTensorSerializer at
module scope, so importing sagemaker.serve required torch. Store the
classes instead and instantiate on lookup, and move torch to an extra.

Fixes aws#5531
@goelakash
goelakash force-pushed the fix-optional-torch-serve branch from 1725779 to 38b75f5 Compare August 7, 2026 03:32
@goelakash
goelakash temporarily deployed to manual-approval August 7, 2026 03:32 — with GitHub Actions Inactive
@goelakash
goelakash temporarily deployed to manual-approval August 7, 2026 03:32 — with GitHub Actions Inactive
@goelakash
goelakash marked this pull request as ready for review August 7, 2026 03:33
@goelakash
goelakash temporarily deployed to manual-approval August 7, 2026 03:33 — with GitHub Actions Inactive
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Make torch an optional dependency / extra

1 participant