fix: order of populating container list (#2523)

* Fixing the order of populating container list

* pylint error fix

Co-authored-by: Shreya Pandit <[email protected]>
Co-authored-by: Ahsan Khan <[email protected]>

Author: 3 people

src/sagemaker/workflow/step_collections.py

@@ -136,10 +136,8 @@ def __init__(
         elif model is not None:
             if isinstance(model, PipelineModel):
                 self.model_list = model.models
-                self.container_def_list = model.pipeline_container_def(inference_instances[0])
             elif isinstance(model, Model):
                 self.model_list = [model]
-                self.container_def_list = [model.prepare_container_def(inference_instances[0])]
 
             for model_entity in self.model_list:
                 if estimator is not None:
@@ -154,10 +152,10 @@ def __init__(
                     source_dir = model_entity.source_dir
                     dependencies = model_entity.dependencies
                     kwargs = dict(**kwargs, output_kms_key=model_entity.model_kms_key)
-                    name = model_entity.name or model_entity._framework_name
+                    model_name = model_entity.name or model_entity._framework_name
 
                     repack_model_step = _RepackModelStep(
-                        name=f"{name}RepackModel",
+                        name=f"{model_name}RepackModel",
                         depends_on=depends_on,
                         sagemaker_session=sagemaker_session,
                         role=role,
@@ -171,10 +169,14 @@ def __init__(
                     model_entity.model_data = (
                         repack_model_step.properties.ModelArtifacts.S3ModelArtifacts
                     )
-
                     # remove kwargs consumed by model repacking step
                     kwargs.pop("output_kms_key", None)
 
+            if isinstance(model, PipelineModel):
+                self.container_def_list = model.pipeline_container_def(inference_instances[0])
+            elif isinstance(model, Model):
+                self.container_def_list = [model.prepare_container_def(inference_instances[0])]
+
         register_model_step = _RegisterModelStep(
             name=name,
             estimator=estimator,