Fix for local tests

Author: sreedes

src/sagemaker/session.py

@@ -2521,8 +2521,10 @@ def _create_model_request(
         request = {"ModelName": name, "ExecutionRoleArn": role}
         if isinstance(container_definition, list):
             request["Containers"] = container_definition
-        else:
+        elif container_definition["ModelPackageName"] is not None:
             request["Containers"] = [container_definition]
+        else:
+            request["PrimaryContainer"] = container_definition
 
         if tags:
             request["Tags"] = tags

src/sagemaker/workflow/step_collections.py

@@ -139,7 +139,7 @@ def __init__(
                     sagemaker_session = estimator.sagemaker_session
                     role = estimator.role
                 else:
-                    sagemaker_session = model.sagemaker_session
+                    sagemaker_session = pipeline_model.sagemaker_session or model.sagemaker_session
                     role = pipeline_model.role
                 if hasattr(model, "entry_point"):
                     repack_model = True

tests/integ/test_workflow.py

@@ -1019,8 +1019,7 @@ def test_mxnet_model_registration(
         py_version="py3",
         sagemaker_session=sagemaker_session,
     )
-    
-    pipeline_model = PipelineModel([model],role)
+    pipeline_model = PipelineModel([model], role, sagemaker_session)
 
     step_register = RegisterModel(
         name="mxnet-register-model",
@@ -1205,8 +1204,10 @@ def test_sklearn_xgboost_sip_model_registration(
         role=role,
         sagemaker_session=sagemaker_session,
     )
-    
-    pipeline_model = PipelineModel([xgboost_model, sklearn_model],role)
+
+    pipeline_model = PipelineModel(
+        [xgboost_model, sklearn_model], role, sagemaker_session=sagemaker_session
+    )
 
     step_register = RegisterModel(
         name="AbaloneRegisterModel",

tests/unit/sagemaker/workflow/test_step_collections.py

@@ -280,7 +280,7 @@ def test_register_model_sip(estimator, model_metrics):
         Model(image_uri="fakeimage1", model_data="Url1", env=[{"k1": "v1"}, {"k2": "v2"}]),
         Model(image_uri="fakeimage2", model_data="Url2", env=[{"k3": "v3"}, {"k4": "v4"}]),
     ]
-    
+
     pipeline_model = PipelineModel(model_list, ROLE)
 
     register_model = RegisterModel(