aws / aws/sagemaker-python-sdk

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

Ouverte
#6,200 1 commentaire 0 réactions 0 personnes assignées Voir sur GitHub
component: model builder type: bug
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

Recevez les nouvelles issues par e-mail

Un résumé court des issues GitHub adaptées aux débutants.