aws / aws/sagemaker-python-sdk
`prepare_for_smd()` missing return value causes `CustomOrchestrator` container health check failure
- 主要言語
- Python
- スター
- 2.3k
- フォーク
- 1.3k
- 平均マージ
- 1日 22時間
- マージ済み PR(30日)
- 35
説明
**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.
コントリビューションガイド
調査の方向性
sagemaker/serve/model_server/smd/prepare.py と model_builder_servers.py の 747 行目付近から始め、SDK 版とコンテナ版の check_integrity.py を比較します。提供されている CustomOrchestrator の build と deployment の再現手順を実行し、生成された整合性値がコンテナに到達して health check が成功し、推論コンポーネントが InService になることを確認します。
索引モデルが issue の本文から書いたものです。
評価
- 技術スタック
- aws, python
- 領域
- backend-api-design, machine-learning
- issue の種類
- バグ
- 難易度
- 4/5
- 見積もり時間
- 3〜5日
- 活発さ
- 活発
- 明瞭さ
- おおむね明確
- 初心者へのやさしさ
- 65/100