aws / aws/sagemaker-python-sdk
`prepare_for_smd()` missing return value causes `CustomOrchestrator` container health check failure
- Lingua principale
- Python
- Stelle
- 2.3k
- Fork
- 1.3k
- Merge medio
- 1g 22h
- PR unite (30g)
- 35
Descrizione
**PySDK Version**
- [ ] PySDK V2 (2.x)
- [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](https://aws.amazon.com/blogs/machine-learning/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`:
```python
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**:
```python
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](https://github.com/aws/sagemaker-python-sdk/blob/44410e2c55202f80d980e1dfd4533e66b9d44495/sagemaker-serve/src/sagemaker/serve/model_builder_servers.py#L747):
```python
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**
```python
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
│ ❱ 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.
Guida per i contributori
Apri la guida per i contributori
Direzione di ricerca
Inizia con sagemaker/serve/model_server/smd/prepare.py e model_builder_servers.py intorno alla riga 747, quindi confronta le versioni SDK e container di check_integrity.py. Esegui la riproduzione fornita di build e deployment di CustomOrchestrator e verifica che il valore di integrità generato raggiunga il container, così che l’health check abbia esito positivo e il componente di inferenza passi a InService.
Scritto dal modello di indicizzazione a partire dal testo della issue.
Valutazione
- Stack tecnologico
- aws, python
- Ambito
- backend-api-design, machine-learning
- Tipo di issue
- Bug
- Difficoltà
- 4/5
- Tempo stimato
- 3-5 giorni
- Stato di attività
- Attiva
- Chiarezza
- Abbastanza chiara
- Idoneità per principianti
- 65/100