`prepare_for_smd()` missing return value causes `CustomOrchestrator` container health check failure

[lopezfelipe]

PySDK Version

• PySDK V2 (2.x)
• PySDK V3 (3.x)

Describe the bug

When deploying a CustomOrchestrator as an Inference Component — as documented in the Build and deploy AI inference workflows with new enhancements to the Amazon SageMaker Python SDK blog post — the container fails its health check during startup with:

AttributeError: 'NoneType' object has no attribute 'encode'

The error occurs in the container's _pickle_file_integrity_check() at line 26 of check_integrity.py:

actual_hash_value = compute_hash(buffer=buffer, secret_key=secret_key)

where secret_key = os.environ.get("SAGEMAKER_SERVE_SECRET_KEY") is None.

Root cause:

prepare_for_smd() in model_server/smd/prepare.py computes the hash and writes metadata.json, but has no return statement:

def prepare_for_smd(model_path, shared_libs, dependencies, inference_spec=None) -> str:
 ...
 hash_value = compute_hash(buffer=buffer)
 with open(str(code_dir.joinpath("metadata.json")), "wb") as metadata:
 metadata.write(_MetaData(hash_value).to_json())
 # ← No return statement — returns None implicitly

In model_builder_servers.py L747:

self.secret_key = prepare_for_smd(...) # = None

Since self.secret_key is None, the SAGEMAKER_SERVE_SECRET_KEY environment variable is never set on the deployed Model/IC. At runtime, the container's integrity check reads the env var, gets None, and crashes.

Additionally, there is a version mismatch between:

• The SDK's local check_integrity.py (uses plain SHA-256, no secret key)
• The container image's bundled check_integrity.py (uses HMAC with a secret key)

To reproduce

from sagemaker.serve.model_builder import ModelBuilder, SchemaBuilder
from sagemaker.serve.spec.inference_base import CustomOrchestrator
from sagemaker.core.inference_config import ResourceRequirements
from sagemaker.core.helper.session_helper import Session, get_execution_role

class MyOrchestrator(CustomOrchestrator):
 def __init__(self, endpoint_name, component_names):
 super().__init__()
 self.endpoint_name = endpoint_name
 self.component_names = component_names

 def handle(self, data, context=None):
 import json
 response = self.client.invoke_endpoint(
 EndpointName=self.endpoint_name,
 InferenceComponentName=self.component_names[0],
 Body=data if isinstance(data, (str, bytes)) else json.dumps(data),
 ContentType="application/json"
 )
 return json.loads(response["Body"].read())

role = get_execution_role()
sess = Session()

orchestrator = ModelBuilder(
 inference_spec=MyOrchestrator(
 endpoint_name="my-existing-endpoint",
 component_names=["base-ic", "adapter-ic"],
 ),
 dependencies={"auto": False, "custom": ["cloudpickle"]},
 sagemaker_session=sess,
 role_arn=role,
 schema_builder=SchemaBuilder(sample_input="Test", sample_output={"generated_text": "test"}),
)

# Workaround for missing constructor fields (separate issue)
orchestrator.resource_requirements = ResourceRequirements(
 requests={"memory": 4096, "num_accelerators": 1, "copies": 1, "num_cpus": 2}
)
orchestrator.inference_component_name = "my-orchestrator-ic"

orchestrator.build()

# Verify secret_key is None after build:
print(f"secret_key: {orchestrator.secret_key}") # Prints: secret_key: None

# Deploy via boto3 (workaround for separate _deploy_for_ic bug):
orchestrator_model_name = orchestrator.built_model.model_name
sm_client = sess.sagemaker_client
sm_client.create_inference_component(
 InferenceComponentName="my-orchestrator-ic",
 EndpointName="my-existing-endpoint",
 VariantName="AllTraffic",
 Specification={
 "ModelName": orchestrator_model_name,
 "ComputeResourceRequirements": {
 "NumberOfAcceleratorDevicesRequired": 1,
 "MinMemoryRequiredInMb": 4096,
 "NumberOfCpuCoresRequired": 2,
 },
 "StartupParameters": {
 "ModelDataDownloadTimeoutInSeconds": 300,
 "ContainerStartupHealthCheckTimeoutInSeconds": 300,
 }
 },
 RuntimeConfig={"CopyCount": 1}
)
# IC fails health check — see CloudWatch logs below

Expected behavior

The CustomOrchestrator IC should pass its container health check and become InService. The SAGEMAKER_SERVE_SECRET_KEY should be correctly generated during build() and propagated to the container environment.

Screenshots or logs

CloudWatch logs from the IC's container (/aws/sagemaker/InferenceComponents/my-orchestrator-ic):

/opt/ml/model/code/inference.py:60 in <module>
│ ❱ 60 _run_preflight_diagnostics()

/opt/ml/model/code/inference.py:38 in _run_preflight_diagnostics
│ ❱ 38 │ _pickle_file_integrity_check()

/opt/ml/model/code/inference.py:57 in _pickle_file_integrity_check
│ ❱ 57 │ perform_integrity_check(buffer=buffer, metadata_path=metadata_path)

/opt/conda/lib/python3.12/site-packages/sagemaker/serve/validations/check_integrity.py:26 in perform_integrity_check
│ ❱ 26 │ actual_hash_value = compute_hash(buffer=buffer, secret_key=secret_key)

AttributeError: 'NoneType' object has no attribute 'encode'

The container then fails the ping health check and the IC never reaches InService.

System information

• SageMaker Python SDK version: sagemaker-serve 1.20.0 (SDK V3)
• Framework name: SageMaker Distribution (SMD) container for CustomOrchestrator
• Framework version: sagemaker-distribution-prod:3.2.0-cpu
• Python version: 3.12
• CPU or GPU: GPU (ml.g6.12xlarge endpoint)
• Custom Docker image (Y/N): N

Additional context

There appear to be two sub-issues:

1. prepare_for_smd()** has no return statement** — it should return the computed hash (or a generated secret key) so that self.secret_key is set to a real value in model_builder_servers.py L747.
2. Version mismatch between SDK and container — The SDK's local check_integrity.py uses plain SHA-256 (hashlib.sha256(buffer).hexdigest()), but the container image (sagemaker-distribution-prod:3.2.0-cpu) still has an older version that uses HMAC with a secret key (hmac.new(secret_key.encode(), msg=buffer, digestmod=hashlib.sha256)). These need to be aligned.