`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.