aws / aws/sagemaker-python-sdk
`prepare_for_smd()` missing return value causes `CustomOrchestrator` container health check failure
- Langage dominant
- Python
- Étoiles
- 2.3k
- Forks
- 1.3k
- Merge moyen
- 1 j 22 h
- PR mergées (30 j)
- 35
Description
**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.
Guide de contribution
Ouvrir le guide de contribution
Piste de recherche
Commencez par sagemaker/serve/model_server/smd/prepare.py et model_builder_servers.py autour de la ligne 747, puis comparez les versions SDK et conteneur de check_integrity.py. Exécutez la reproduction fournie du build et du déploiement de CustomOrchestrator, et vérifiez que la valeur d’intégrité générée atteint le conteneur afin que le health check réussisse et que le composant d’inférence passe à InService.
Rédigé par le modèle d'indexation à partir du texte de l'issue.
Évaluation
- Stack technique
- aws, python
- Domaine
- backend-api-design, machine-learning
- Type d'issue
- Bug
- Difficulté
- 4/5
- Temps estimé
- 3-5 jours
- Activité
- Active
- Clarté
- Plutôt claire
- Accessibilité débutants
- 65/100