addressed PR comments for UT's and added default value None to missing parameter

Author: rahven14

src/sagemaker/pipeline.py

@@ -84,7 +84,7 @@ def __init__(
         self.enable_network_isolation = enable_network_isolation
         self.endpoint_name = None
 
-    def pipeline_container_def(self, instance_type):
+    def pipeline_container_def(self, instance_type=None):
         """The pipeline definition for deploying this model.
 
         This is the dict created by ``sagemaker.pipeline_container_def()``.
@@ -321,7 +321,10 @@ def register(
             )
         else:
             container_def = [
-                {"Image": image_uri or model.image_uri, "ModelDataUrl": model.model_data}
+                {
+                    "Image": image_uri or model.image_uri,
+                    "ModelDataUrl": model.model_data,
+                }
                 for model in self.models
             ]
 

tests/unit/sagemaker/workflow/test_pipeline_session.py

@@ -124,7 +124,9 @@ def test_pipeline_session_context_for_model_step(pipeline_session_mock):
     assert len(register_step_args.need_runtime_repack) == 0
 
 
-def test_pipeline_session_context_for_model_step_without_instance_types(pipeline_session_mock):
+def test_pipeline_session_context_for_model_step_without_instance_types(
+    pipeline_session_mock,
+):
     model = Model(
         name="MyModel",
         image_uri="fakeimage",
@@ -134,27 +136,38 @@ def test_pipeline_session_context_for_model_step_without_instance_types(pipeline
         source_dir=f"{DATA_DIR}",
         role=_ROLE,
     )
-    # CreateModelStep requires runtime repack
-    create_step_args = model.create(
-        instance_type="c4.4xlarge",
-        accelerator_type="ml.eia1.medium",
-    )
-    # The context should be cleaned up before return
-    assert pipeline_session_mock.context is None
-    assert create_step_args.create_model_request
-    assert not create_step_args.create_model_package_request
-    assert len(create_step_args.need_runtime_repack) == 1
 
-    # _RegisterModelStep does not require runtime repack
-    model.entry_point = None
-    model.source_dir = None
     register_step_args = model.register(
         content_types=["text/csv"],
         response_types=["text/csv"],
         model_package_group_name="MyModelPackageGroup",
     )
-    # The context should be cleaned up before return
-    assert not pipeline_session_mock.context
-    assert not register_step_args.create_model_request
-    assert register_step_args.create_model_package_request
-    assert len(register_step_args.need_runtime_repack) == 0
+
+    expected_output = {
+        "ModelPackageGroupName": "MyModelPackageGroup",
+        "InferenceSpecification": {
+            "Containers": [
+                {
+                    "Image": "fakeimage",
+                    "Environment": {
+                        "SAGEMAKER_PROGRAM": "dummy_script.py",
+                        "SAGEMAKER_SUBMIT_DIRECTORY": "/opt/ml/model/code",
+                        "SAGEMAKER_CONTAINER_LOG_LEVEL": "20",
+                        "SAGEMAKER_REGION": "us-west-2",
+                    },
+                    "ModelDataUrl": ParameterString(
+                        name="ModelData",
+                        default_value="s3://my-bucket/file",
+                    ),
+                }
+            ],
+            "SupportedContentTypes": ["text/csv"],
+            "SupportedResponseMIMETypes": ["text/csv"],
+            "SupportedRealtimeInferenceInstanceTypes": None,
+            "SupportedTransformInstanceTypes": None,
+        },
+        "CertifyForMarketplace": False,
+        "ModelApprovalStatus": "PendingManualApproval",
+    }
+
+    assert register_step_args.create_model_package_request == expected_output