aws / aws/sagemaker-python-sdk

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

オープン
#6,200 コメント 1 件 リアクション 0 件 担当者 0 名 GitHub で見る
component: model builder type: bug
主要言語
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

新しい issue をメールで受け取る

初心者向けの GitHub issue を短くまとめたダイジェスト。