Populating the container definition during Model Registration

Author: sreedes

src/sagemaker/model.py

@@ -178,21 +178,23 @@ def register(
         """
         if self.model_data is None:
             raise ValueError("SageMaker Model Package cannot be created without model data.")
+        if image_uri is not None:
+            self.image_uri = image_uri
+        container_def = self.prepare_container_def()
 
         model_pkg_args = sagemaker.get_model_package_args(
             content_types,
             response_types,
             inference_instances,
             transform_instances,
-            model_package_name,
-            model_package_group_name,
-            self.model_data,
-            image_uri or self.image_uri,
-            model_metrics,
-            metadata_properties,
-            marketplace_cert,
-            approval_status,
-            description,
+            model_package_name=model_package_name,
+            model_package_group_name=model_package_group_name,
+            model_metrics=model_metrics,
+            metadata_properties=metadata_properties,
+            marketplace_cert=marketplace_cert,
+            approval_status=approval_status,
+            description=description,
+            container_def_list=[container_def],
             drift_check_baselines=drift_check_baselines,
         )
         model_package = self.sagemaker_session.create_model_package_from_containers(

tests/unit/test_estimator.py

@@ -3069,6 +3069,7 @@ def test_register_inference_image(sagemaker_session):
         "containers": [
             {
                 "Image": inference_image,
+                "Environment": {},
                 "ModelDataUrl": estimator.model_data,
             }
         ],