[mohamedzeidan2021]

Description

Deploying a CustomOrchestrator with ModelBuilder fails its container health check at startup with AttributeError: 'NoneType' object has no attribute 'encode', so the Inference Component never reaches InService.

The missing return in prepare_for_smd() is not the defect. The secret key was removed on purpose in #5654 / #5656 as the fix for CVE-2026-8596, where SAGEMAKER_SERVE_SECRET_KEY was exposed in plaintext through the Describe APIs. The host side is consistent with that today — check_integrity.compute_hash(buffer) is a plain SHA-256 and no key is set — so restoring the return value would reintroduce the vulnerability. This is why I did not take the approach of the two earlier attempts at this issue (#6221, #6222).

The actual defect is a version coupling between the host and the container. The generated code/inference.py did:

from sagemaker.serve.validations.check_integrity import perform_integrity_check

which resolves against the container's installed sagemaker, not the SDK that produced the artifacts. sagemaker/core/image_uri_config/sagemaker-distribution.json maps SMD inference to 3.2.0 only, and that image bundles sagemaker 2.245.0, whose perform_integrity_check still does os.environ.get("SAGEMAKER_SERVE_SECRET_KEY").encode().

A second, independent mismatch sits behind the first: _save_model_inference_spec pickled (inference_spec, schema_builder), and a v3 SchemaBuilder pickles with references to sagemaker.core.* serializer modules that a 2.x install cannot import, so cloudpickle.load in the entrypoint failed even once the integrity check passed.

Changes Made

1. model_server/smd/custom_execution_inference.py — inlined the digest comparison. It is the same plain SHA-256 check against the metadata.json the host writes, so the integrity posture is unchanged, but the script no longer imports anything from sagemaker and therefore cannot drift against whatever version the image carries. Both file reads now go through SAGEMAKER_INFERENCE_CODE_DIRECTORY (always set by _upload_smd_artifacts) with the previous /opt/ml/model/code as the default, which also makes the module testable.

2. model_builder.py — custom orchestrators pickle as (inference_spec, None). The entrypoint already discards the second element of that tuple, so the SchemaBuilder was dead weight that coupled serve.pkl to the host SDK's module layout.

3. prepare_for_smd / prepare_for_torchserve / prepare_for_mms / prepare_for_tf_serving — corrected the -> str annotations to -> None, and dropped the six self.secret_key = prepare_for_*(...) assignments along with their # NOTE: ... fix out of scope / pylint: disable-next=assignment-from-no-return comments. These functions record a digest and have had nothing to return since the CVE fix; the lying annotation is what led both earlier PRs to try to resurrect the secret key. Every builder already initializes self.secret_key = "" before calling them, and both None and "" are falsey at every consumer, so this is not a behavior change.

Files

• sagemaker-serve/src/sagemaker/serve/model_server/smd/custom_execution_inference.py
• sagemaker-serve/src/sagemaker/serve/model_server/smd/prepare.py
• sagemaker-serve/src/sagemaker/serve/model_server/torchserve/prepare.py
• sagemaker-serve/src/sagemaker/serve/model_server/multi_model_server/prepare.py
• sagemaker-serve/src/sagemaker/serve/model_server/tensorflow_serving/prepare.py
• sagemaker-serve/src/sagemaker/serve/model_builder.py
• sagemaker-serve/src/sagemaker/serve/model_builder_servers.py
• sagemaker-serve/tests/unit/model_server/test_smd_custom_execution_inference.py (new)
• sagemaker-serve/tests/unit/model_server/test_smd_prepare.py
• sagemaker-serve/tests/unit/test_model_builder_build.py
• sagemaker-serve/tests/unit/servers/test_model_builder_servers.py

Testing

End-to-end against the image the SDK actually selects. I generated the artifacts with this branch (_save_model_inference_spec + prepare_for_smd, no AWS calls) and then executed the resulting code/inference.py under a Python 3.12 interpreter with sagemaker==2.245.0 installed — the version bundled in sagemaker-distribution-prod:3.2.0-cpu — the way the model server imports it:

• Before: the exact traceback from the issue, at import time, so /ping never becomes healthy.
• After: import-time preflight + unpickle: OK, and handler(b'hello') -> {'echo': b'hello'}, with no SAGEMAKER_SERVE_SECRET_KEY present.
• Isolating the two objects confirmed the second cause independently: the orchestrator alone unpickles fine under 2.245.0, while the SchemaBuilder raises ModuleNotFoundError: No module named 'sagemaker.core'.
• Negative cases still fail closed: appending a byte to serve.pkl → Integrity check for the serialized function or data failed.; removing metadata.json → Path to metadata.json does not exist.

Unit tests. 15 new tests, including an AST-level guard asserting the entrypoint imports nothing from sagemaker — all 11 tests in the new file fail against the unfixed source. Full sagemaker-serve/tests/unit on this branch: identical failure set to the master baseline (49 → 49, zero new); the pre-existing failures are credential/network-dependent. black, flake8 and pydocstyle are clean, and pylint on the changed modules reports only the pre-existing no-else-return in handler.

Follow-ups not in this PR

• The SMD image mapping is stale. sagemaker-distribution.json maps SMD inference to 3.2.0 only, while public ECR serves sagemaker-distribution up to 4.4.2 today. Bumping it is worth doing on its own merits, but it needs the per-region sagemaker-distribution-prod registry account IDs for the new version, which I cannot verify from outside. It is not a prerequisite for this fix — the point of inlining the check is that the currently-mapped image now works.
• Mode.LOCAL_CONTAINER cannot exercise SMD. LocalContainerMode.create_server has no branch for ModelServer.SMD and LocalEndpoint.invoke raises "Unsupported model server", so none of this class of failure is catchable before deploying.

Thanks to @deburky, whose analysis on #6200 correctly identified the host/container skew and ruled out restoring the secret key.

Related Issue

Closes #6200

Merge Checklist

• Changes are backward compatible
• Commit message follows prefix: description format
• Unit tests added/updated
• Integration tests added (if applicable)
• Documentation updated (if applicable)

[github-actions]

🤖 AI code review

This PR fixes the CustomOrchestrator SMD health-check failure (issue #6200) by decoupling the container entrypoint from the container's bundled sagemaker: it inlines a plain SHA-256 integrity check into custom_execution_inference.py (no more from sagemaker... import), pickles (inference_spec, None) for custom orchestrators so a v3 SchemaBuilder can't drag in unimportable sagemaker.core modules, and corrects the prepare_for_* return annotations to -> None while dropping the dead self.secret_key = prepare_*(...) assignments. I verified the inlined check matches perform_integrity_check (reads sha256_hash from _MetaData.to_json(), uses hmac.compare_digest), that SAGEMAKER_INFERENCE_CODE_DIRECTORY is always set to /opt/ml/model/code matching the new default, that every builder initializes self.secret_key = "" before the prepare call (so behaviour is unchanged since the functions always returned None), and that the SMD entrypoint is the only consumer of the custom-orchestrator serve.pkl and discards the second tuple element. The change is correct, consistent with the team precedent in #6200 (do not reintroduce SAGEMAKER_SERVE_SECRET_KEY), and well covered by new unit tests. No blocking issues found.

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

[mohamedzeidan2021]

Pushed a follow-up commit addressing codecov/patch, and triaging the other three red checks:

codecov/patch (64.28%) — was real, now fixed. The uncovered lines were all in custom_execution_inference.py. My test exec'd the entrypoint's source text rather than importing it, so coverage could not attribute those lines even though every one of them ran. Since the module now reads its code directory from SAGEMAKER_INFERENCE_CODE_DIRECTORY, it can be imported for real against a prepared temp dir and reloaded per case — which is both measurable and closer to how the model server loads it. I also added an AsyncCustomOrchestrator case to cover the awaited branch of handler(). custom_execution_inference.py is at 100% locally (50/50 statements).

unit-tests (sagemaker-serve) — pre-existing environment break, not this PR. The py311 and py312 tox envs fail at conftest import, before any test runs:

pandas 3.1.0 requires numpy>=2.0.2; python_version < "3.14", but you have numpy 1.26.4 which is incompatible.
ImportError: C extension: None not built.

pandas resolves unpinned to 3.1.0 while something later pins numpy back to 1.26.4, so pandas' C extension cannot load and sagemaker.core.serializers.base's import numpy fails. py310 (which gets pandas==2.3.3 / numpy==2.1.3) ran the full suite clean: 2205 passed, 34 skipped, 0 failed, including the new tests. The same job is red on #6533 and #6531, and CI Health V3 Master is red too, so this is not specific to this branch. Probably worth a constraint on pandas or numpy in the test requirements, but that is its own change.

integ-tests (sagemaker-mlops) — unrelated. test_sklearn_processing_job ("Processing job Failed") and test_transform_job ("Transform job Failed") — both AWS-side job failures in sagemaker-mlops/tests/integ/. This PR touches only sagemaker-serve, and nothing in it is reachable from ProcessingJob or batch transform. These pass on #6533 and #6531, so it looks like a flake or a test-account issue rather than a master break; happy to have it re-run.

codecov/project (42.17% vs 65%) should recover with the coverage fix — the figure looks depressed by the partial upload from the two tox envs that died at collection.

[github-actions]

🤖 AI code review

This PR fixes the CustomOrchestrator/SMD container health-check failure (#6200) by decoupling the generated code/inference.py from the serving container's bundled sagemaker. It inlines the plain SHA-256 integrity check (identical to the current check_integrity.perform_integrity_check), routes both file reads through SAGEMAKER_INFERENCE_CODE_DIRECTORY with the prior default, pickles custom orchestrators as (inference_spec, None) (the entrypoint already discards the second element), and corrects the prepare_for_* annotations to -> None while dropping the dead self.secret_key = prepare_*(...) assignments. I verified the integrity semantics match, that self.secret_key = "" is set before every affected prepare call (so no behavior change), that CustomOrchestrators only reach the SMD path, and that the new tests cover the key regressions. The change is correct, backward-compatible, and well-tested; no blocking issues found. The PR body explains the author's reasoning but does not attempt to subvert the review.

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