[mohamedzeidan2021]

Follow-up to #6512, which made the sagemaker-serve serving backends extras.

Credit where it's due: all four of these gaps were diagnosed first by @goelakash in #6166, which has been open since August. That PR was overtaken by #6512 and now conflicts with it in four places where #6512's version is the broader fix (_LazySerializerRegistry keeps DEFAULT_SERIALIZERS_BY_FRAMEWORK[x] returning instances, so no call sites change; it also covers onnx, onnxruntime, tritonclient and mlflow). This PR carries forward the part of #6166 that #6512 did not cover, rebased onto master. @goelakash is credited as co-author on the commit.

What was still broken

Selecting the PyTorch default serializer still required torch. TorchTensorSerializer.__init__ imported torch, so _LazySerializerRegistry 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 handles a tensor, so it hit an ImportError for a dependency it had no use for — exactly the API-only case #5531 is about.

The type check now consults sys.modules:

def _is_torch_tensor(data):
    torch = sys.modules.get("torch")
    return torch is not None and isinstance(data, torch.Tensor)

A torch.Tensor instance cannot exist unless torch has already been imported, so this is exactly as strict as the isinstance check it replaces — no duck-typing, no widening — while never triggering an import. torch_tensor becomes a property, preserving 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 is implemented as try: self[key] except KeyError: return False. So framework_enum in DEFAULT_SERIALIZERS_BY_FRAMEWORK instantiated the pair as a side effect, defeating the laziness, and for PYTORCH without torch raised ImportError out of a membership test instead of 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 (torch, onnx, triton, mlflow) through. all is deliberately unchanged, so no existing install starts pulling torch.

Testing

• All 10 new and changed assertions fail on master and pass here.
• sagemaker-core and sagemaker-serve unit tests matching serial|deserial|torch|translator|constants: 374 passed. The only failures are 5 pre-existing TestTensorflowTensorTranslator errors (No module named 'tensorflow'), identical on master.
• Verified end to end with torch blocked by a meta_path hook: _fetch_serializer_and_deserializer_for_framework("pytorch") returns (TorchTensorSerializer, JSONDeserializer) with "torch" not in sys.modules.
• black==26.3.1, flake8, pylint (10.00/10) and pydocstyle are clean on every changed file.

Open question for reviewers

sagemaker[all] is left as the three subpackages, so it stays torch-free. The conventional reading of all is "all optional functionality", which would make it sagemaker-serve[all] — but that silently re-adds the full backend set for anyone already using sagemaker[all]. Flagging rather than deciding; happy to change it.

[github-actions]

🤖 AI code review

This PR closes remaining gaps in optional torch support: TorchTensorSerializer no longer imports torch at construction (the type check now consults sys.modules and torch_tensor becomes a lazy property), _LazySerializerRegistry gets a __contains__ that doesn't build pairs, TorchTensorTranslator raises a clearer ImportError naming the extra, and the top-level sagemaker package forwards sagemaker-serve's backend extras. The change is well-targeted, consistent with the established optional-dependency approach (PR #5713/#5714), and comes with thorough subprocess-based tests that block torch via a meta_path hook. Overall it looks correct; only a minor consistency note below.

_{Reviewed commit c04b8a4. Automated review; verify before acting.}

[github-actions]

🟡 Low: _is_torch_tensor here uses the sys.modules approach, while sagemaker.serve.builder.schema_builder._is_torch_tensor (already in the tree) still does from torch import Tensor inside a try/except. The two helpers now have the same name but different behaviour (one can trigger a torch import, the other never does). Worth a brief comment or aligning the serve-side helper so future readers don't assume they're equivalent.