Description
Describe the bug
We operate two SageMaker pipelines: training and batch inference. In the training pipeline, we create a model using the CreateModelStep API, where we explicitly specify the image_uri for inference. The process completes successfully, creating a new model package version which is then registered with the Model Registry. Upon reviewing the inference specification of the generated model package, the image_uri correctly reflects the specified Docker image for inference.
In the inference pipeline, we import the model package version following the guidelines in the documentation [1] and a sample project [2]. However, the image_uri is not preserved and defaults back to the training image.
As a result, we are unable to run the Batch Transform step because the job incorrectly uses the training image instead of the specified inference image.
[1] https://docs.aws.amazon.com/sagemaker/latest/dg/sagemaker-mkt-model-pkg-model.html#sagemaker-mkt-model-pkg-model-sdk
[2] https://github.com/aws-samples/aws-enterprise-mlops-framework/blob/b6eea322b44b6d90a110ae91298308ba060f96d1/mlops-multi-account-cdk/mlops-sm-project-template/mlops_sm_project_template/templates/train_deploy_batch_inference_product/seed_code/build_app/ml_pipelines/inference/pipeline.py#L221
To reproduce
# Training pipeline
step_create_model = CreateModelStep(
name=f'model-{project_id}',
display_name="Create Model",
model=Model(
predictor_cls=AutoGluonTabularPredictor,
image_uri=inferenceImageUriParameter.to_string(),
model_data=step_train.properties.ModelArtifacts.S3ModelArtifacts,
sagemaker_session=pipeline_session,
role=role,
name=f'model-{project_id}'
),
inputs=CreateModelInput(
instance_type="ml.m5.xlarge",
),
depends_on=[step_train]
)
# Inference pipeline
model_package = ModelPackage(
role=role,
model_package_arn=model_package_arn, # e.g. arn:aws:sagemaker:eu-west-1:77777777777:model-package/mlops-example-1-p-jmxjq22cpa2r/21
sagemaker_session=pipeline_session,
)
step_create_model = ModelStep(
name="CreateModel",
display_name="Create Model",
step_args=model_package.create(
instance_type="ml.m5.xlarge"
),
depends_on=[step_load_data]
)Expected behavior
The image uri in the Inference specification is retained.
System information
- SageMaker Python SDK version: 2.232.2
- Framework name (eg. PyTorch) or algorithm (eg. KMeans): N/A
- Framework version: N/A
- Python version: 3.
- CPU or GPU: N/A
- Custom Docker image (Y/N): Y