Updating the input to RegisterModelStep to Pipeline Model

Author: sreedes

src/sagemaker/session.py

@@ -2522,7 +2522,8 @@ def _create_model_request(
         if isinstance(container_definition, list):
             request["Containers"] = container_definition
         else:
-            request["PrimaryContainer"] = container_definition
+            request["Containers"] = [container_definition]
+
         if tags:
             request["Tags"] = tags
 

src/sagemaker/workflow/step_collections.py

@@ -68,7 +68,7 @@ def __init__(
         compile_model_family=None,
         description=None,
         tags=None,
-        models=None,
+        pipeline_model=None,
         **kwargs,
     ):
         """Construct steps `_RepackModelStep` and `_RegisterModelStep` based on the estimator.
@@ -100,7 +100,8 @@ def __init__(
                 that tags will only be applied to newly created model package groups; if the
                 name of an existing group is passed to "model_package_group_name",
                 tags will not be applied.
-            models (list): A list of models.
+            pipeline_model (object): A PipelineModel object that comprises a list of models
+                which gets executed as a serial inference pipeline.
             **kwargs: additional arguments to `create_model`.
         """
         steps: List[Step] = []
@@ -132,14 +133,14 @@ def __init__(
         kwargs.pop("dependencies", None)
         kwargs.pop("output_kms_key", None)
 
-        if models is not None:
-            for model in models:
+        if pipeline_model is not None:
+            for model in pipeline_model.models:
                 if estimator is not None:
                     sagemaker_session = estimator.sagemaker_session
                     role = estimator.role
                 else:
                     sagemaker_session = model.sagemaker_session
-                    role = model.role
+                    role = pipeline_model.role
                 if hasattr(model, "entry_point"):
                     repack_model = True
                     entry_point = model.entry_point
@@ -174,7 +175,7 @@ def __init__(
             compile_model_family=compile_model_family,
             description=description,
             tags=tags,
-            model_list=models,
+            model_list=pipeline_model.models,
             **kwargs,
         )
         if not repack_model:

tests/integ/test_workflow.py

@@ -36,6 +36,7 @@
 from sagemaker import image_uris
 from sagemaker.inputs import CreateModelInput, TrainingInput
 from sagemaker.model import Model
+from sagemaker import PipelineModel
 from sagemaker.processing import ProcessingInput, ProcessingOutput, FeatureStoreOutput
 from sagemaker.pytorch.estimator import PyTorch
 from sagemaker.tuner import HyperparameterTuner, IntegerParameter
@@ -1018,10 +1019,12 @@ def test_mxnet_model_registration(
         py_version="py3",
         sagemaker_session=sagemaker_session,
     )
+    
+    pipeline_model = PipelineModel([model],role)
 
     step_register = RegisterModel(
         name="mxnet-register-model",
-        models=[model],
+        pipeline_model=pipeline_model,
         content_types=["*"],
         response_types=["*"],
         inference_instances=["ml.m5.xlarge"],
@@ -1202,10 +1205,12 @@ def test_sklearn_xgboost_sip_model_registration(
         role=role,
         sagemaker_session=sagemaker_session,
     )
+    
+    pipeline_model = PipelineModel([xgboost_model, sklearn_model],role)
 
     step_register = RegisterModel(
         name="AbaloneRegisterModel",
-        models=[xgboost_model, sklearn_model],
+        pipeline_model=pipeline_model,
         content_types=["application/json"],
         response_types=["application/json"],
         inference_instances=["ml.t2.medium", "ml.m5.xlarge"],

tests/unit/sagemaker/workflow/test_step_collections.py

@@ -27,6 +27,7 @@
     PropertyMock,
 )
 
+from sagemaker import PipelineModel
 from sagemaker.estimator import Estimator
 from sagemaker.model import Model
 from sagemaker.tensorflow import TensorFlow
@@ -279,6 +280,8 @@ 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(
         name="RegisterModelStep",
@@ -291,7 +294,7 @@ def test_register_model_sip(estimator, model_metrics):
         model_metrics=model_metrics,
         approval_status="Approved",
         description="description",
-        models=model_list,
+        pipeline_model=pipeline_model,
         depends_on=["TestStep"],
     )
     assert ordered(register_model.request_dicts()) == ordered(