fix: pipelines local mode minor bug fixes

Co-authored-by: Namrata Madan <[email protected]>

Author: 2 people

src/sagemaker/local/entities.py

@@ -669,10 +669,8 @@ def start(self, **kwargs):
         execution = _LocalPipelineExecution(execution_id, self.pipeline, **kwargs)
 
         self._executions[execution_id] = execution
-        logger.info(
-            "Starting execution for pipeline %s. Execution ID is %s",
-            self.pipeline.name,
-            execution_id,
+        print(
+            f"Starting execution for pipeline {self.pipeline.name}. Execution ID is {execution_id}"
         )
         self.last_modified_time = datetime.datetime.now().timestamp()
 
@@ -690,6 +688,8 @@ def __init__(
         PipelineExecutionDescription=None,
         PipelineExecutionDisplayName=None,
     ):
+        from sagemaker.workflow.pipeline import PipelineGraph
+
         self.pipeline = pipeline
         self.pipeline_execution_name = execution_id
         self.pipeline_execution_description = PipelineExecutionDescription
@@ -699,7 +699,8 @@ def __init__(
         self.creation_time = datetime.datetime.now().timestamp()
         self.last_modified_time = self.creation_time
         self.step_execution = {}
-        self._initialize_step_execution(self.pipeline.steps)
+        self.pipeline_dag = PipelineGraph.from_pipeline(self.pipeline)
+        self._initialize_step_execution(self.pipeline_dag.step_map.values())
         self.pipeline_parameters = self._initialize_and_validate_parameters(PipelineParameters)
         self._blocked_steps = {}
 
@@ -732,37 +733,36 @@ def update_execution_success(self):
         """Mark execution as succeeded."""
         self.status = _LocalExecutionStatus.SUCCEEDED.value
         self.last_modified_time = datetime.datetime.now().timestamp()
-        logger.info("Pipeline execution %s SUCCEEDED", self.pipeline_execution_name)
+        print(f"Pipeline execution {self.pipeline_execution_name} SUCCEEDED")
 
     def update_execution_failure(self, step_name, failure_message):
         """Mark execution as failed."""
         self.status = _LocalExecutionStatus.FAILED.value
-        self.failure_reason = f"Step {step_name} failed with message: {failure_message}"
+        self.failure_reason = f"Step '{step_name}' failed with message: {failure_message}"
         self.last_modified_time = datetime.datetime.now().timestamp()
-        logger.info(
-            "Pipeline execution %s FAILED because step %s failed.",
-            self.pipeline_execution_name,
-            step_name,
+        print(
+            f"Pipeline execution {self.pipeline_execution_name} FAILED because step "
+            f"'{step_name}' failed."
         )
 
     def update_step_properties(self, step_name, step_properties):
         """Update pipeline step execution output properties."""
         self.step_execution.get(step_name).update_step_properties(step_properties)
-        logger.info("Pipeline step %s SUCCEEDED.", step_name)
+        print(f"Pipeline step '{step_name}' SUCCEEDED.")
 
     def update_step_failure(self, step_name, failure_message):
         """Mark step_name as failed."""
+        print(f"Pipeline step '{step_name}' FAILED. Failure message is: {failure_message}")
         self.step_execution.get(step_name).update_step_failure(failure_message)
-        logger.info("Pipeline step %s FAILED. Failure message is: %s", step_name, failure_message)
 
     def mark_step_executing(self, step_name):
         """Update pipelines step's status to EXECUTING and start_time to now."""
-        logger.info("Starting pipeline step: %s", step_name)
+        print(f"Starting pipeline step: '{step_name}'")
         self.step_execution.get(step_name).mark_step_executing()
 
     def _initialize_step_execution(self, steps):
         """Initialize step_execution dict."""
-        from sagemaker.workflow.steps import StepTypeEnum
+        from sagemaker.workflow.steps import StepTypeEnum, Step
 
         supported_steps_types = (
             StepTypeEnum.TRAINING,
@@ -774,16 +774,17 @@ def _initialize_step_execution(self, steps):
         )
 
         for step in steps:
-            if step.step_type not in supported_steps_types:
-                error_msg = self._construct_validation_exception_message(
-                    "Step type {} is not supported in local mode.".format(step.step_type.value)
+            if isinstance(step, Step):
+                if step.step_type not in supported_steps_types:
+                    error_msg = self._construct_validation_exception_message(
+                        "Step type {} is not supported in local mode.".format(step.step_type.value)
+                    )
+                    raise ClientError(error_msg, "start_pipeline_execution")
+                self.step_execution[step.name] = _LocalPipelineExecutionStep(
+                    step.name, step.step_type, step.description, step.display_name
                 )
-                raise ClientError(error_msg, "start_pipeline_execution")
-            self.step_execution[step.name] = _LocalPipelineExecutionStep(
-                step.name, step.step_type, step.description, step.display_name
-            )
-            if step.step_type == StepTypeEnum.CONDITION:
-                self._initialize_step_execution(step.if_steps + step.else_steps)
+                if step.step_type == StepTypeEnum.CONDITION:
+                    self._initialize_step_execution(step.if_steps + step.else_steps)
 
     def _initialize_and_validate_parameters(self, overridden_parameters):
         """Initialize and validate pipeline parameters."""

src/sagemaker/local/local_session.py

@@ -448,6 +448,7 @@ def update_pipeline(
             }
             raise ClientError(error_response, "update_pipeline")
         LocalSagemakerClient._pipelines[pipeline.name].pipeline_description = pipeline_description
+        LocalSagemakerClient._pipelines[pipeline.name].pipeline = pipeline
         LocalSagemakerClient._pipelines[
             pipeline.name
         ].last_modified_time = datetime.now().timestamp()

src/sagemaker/local/pipeline.py

@@ -17,11 +17,12 @@
 import json
 from copy import deepcopy
 from datetime import datetime
-from typing import Dict, List
+from typing import Dict, List, Union
 from botocore.exceptions import ClientError
 
 from sagemaker.workflow.conditions import ConditionTypeEnum
 from sagemaker.workflow.steps import StepTypeEnum, Step
+from sagemaker.workflow.step_collections import StepCollection
 from sagemaker.workflow.entities import PipelineVariable
 from sagemaker.workflow.parameters import Parameter
 from sagemaker.workflow.functions import Join, JsonGet, PropertyFile
@@ -256,8 +257,7 @@ def execute(self):
             return self.pipline_executor.local_sagemaker_client.describe_training_job(job_name)
         except Exception as e:  # pylint: disable=W0703
             self.pipline_executor.execution.update_step_failure(
-                self.step.name,
-                f"Error when executing step {self.step.name} of type {type(self.step)}: {e}",
+                self.step.name, f"{type(e).__name__}: {str(e)}"
             )
 
 
@@ -291,21 +291,18 @@ def execute(self):
 
         except Exception as e:  # pylint: disable=W0703
             self.pipline_executor.execution.update_step_failure(
-                self.step.name,
-                f"Error when executing step {self.step.name} of type {type(self.step)}: {e}",
+                self.step.name, f"{type(e).__name__}: {str(e)}"
             )
 
 
 class _ConditionStepExecutor(_StepExecutor):
     """Executor class to execute ConditionStep locally"""
 
     def execute(self):
-        def _block_all_downstream_steps(steps: List[Step]):
+        def _block_all_downstream_steps(steps: List[Union[Step, StepCollection]]):
             steps_to_block = set()
             for step in steps:
-                steps_to_block.update(
-                    self.pipline_executor.pipeline_dag.get_steps_in_sub_dag(step.name)
-                )
+                steps_to_block.update(self.pipline_executor.pipeline_dag.get_steps_in_sub_dag(step))
             self.pipline_executor._blocked_steps.update(steps_to_block)
 
         if_steps = self.step.if_steps
@@ -469,8 +466,7 @@ def execute(self):
             return self.pipline_executor.local_sagemaker_client.describe_transform_job(job_name)
         except Exception as e:  # pylint: disable=W0703
             self.pipline_executor.execution.update_step_failure(
-                self.step.name,
-                f"Error when executing step {self.step.name} of type {type(self.step)}: {e}",
+                self.step.name, f"{type(e).__name__}: {str(e)}"
             )
 
 
@@ -485,8 +481,7 @@ def execute(self):
             return self.pipline_executor.local_sagemaker_client.describe_model(model_name)
         except Exception as e:  # pylint: disable=W0703
             self.pipline_executor.execution.update_step_failure(
-                self.step.name,
-                f"Error when executing step {self.step.name} of type {type(self.step)}: {e}",
+                self.step.name, f"{type(e).__name__}: {str(e)}"
             )
 
 

src/sagemaker/workflow/pipeline.py

@@ -26,7 +26,6 @@
 from sagemaker import s3
 from sagemaker._studio import _append_project_tags
 from sagemaker.session import Session
-from sagemaker.local import LocalSession
 from sagemaker.workflow.callback_step import CallbackOutput, CallbackStep
 from sagemaker.workflow.lambda_step import LambdaOutput, LambdaStep
 from sagemaker.workflow.entities import (
@@ -162,9 +161,7 @@ def _create_args(
 
         # If pipeline definition is large, upload to S3 bucket and
         # provide PipelineDefinitionS3Location to request instead.
-        if len(pipeline_definition.encode("utf-8")) < 1024 * 100 or isinstance(
-            self.sagemaker_session, LocalSession
-        ):
+        if len(pipeline_definition.encode("utf-8")) < 1024 * 100:
             kwargs["PipelineDefinition"] = pipeline_definition
         else:
             desired_s3_uri = s3.s3_path_join(
@@ -660,19 +657,28 @@ def is_cyclic_helper(current_step):
                     return True
         return False
 
-    def get_steps_in_sub_dag(self, current_step: str, steps: Set[str] = None) -> Set[str]:
+    def get_steps_in_sub_dag(
+        self, current_step: Union[Step, StepCollection], sub_dag_steps: Set[str] = None
+    ) -> Set[str]:
         """Get names of all steps (including current step) in the sub dag of current step.
 
         Returns a set of step names in the sub dag.
         """
-        if steps is None:
-            steps = set()
-        if current_step not in self.adjacency_list:
-            raise ValueError("Step: %s does not exist in the pipeline." % current_step)
-        steps.add(current_step)
-        for step in self.adjacency_list[current_step]:
-            self.get_steps_in_sub_dag(step, steps)
-        return steps
+        if sub_dag_steps is None:
+            sub_dag_steps = set()
+
+        if isinstance(current_step, StepCollection):
+            current_steps = current_step.steps
+        else:
+            current_steps = [current_step]
+
+        for step in current_steps:
+            if step.name not in self.adjacency_list:
+                raise ValueError("Step: %s does not exist in the pipeline." % step.name)
+            sub_dag_steps.add(step.name)
+            for sub_step in self.adjacency_list[step.name]:
+                self.get_steps_in_sub_dag(self.step_map.get(sub_step), sub_dag_steps)
+        return sub_dag_steps
 
     def __iter__(self):
         """Perform topological sort traversal of the Pipeline Graph."""

tests/integ/test_local_mode.py

@@ -30,11 +30,11 @@
 
 from sagemaker.model import Model
 from sagemaker.transformer import Transformer
-from sagemaker.inputs import CreateModelInput
 from sagemaker.processing import ProcessingInput, ProcessingOutput, ScriptProcessor
 from sagemaker.sklearn.processing import SKLearnProcessor
 from sagemaker.workflow.pipeline import Pipeline
-from sagemaker.workflow.steps import TrainingStep, ProcessingStep, TransformStep, CreateModelStep
+from sagemaker.workflow.steps import TrainingStep, ProcessingStep, TransformStep
+from sagemaker.workflow.model_step import ModelStep
 from sagemaker.workflow.parameters import ParameterInteger
 from sagemaker.workflow.condition_step import ConditionStep
 from sagemaker.workflow.fail_step import FailStep
@@ -546,8 +546,8 @@ def test_local_pipeline_with_training_and_transform_steps(
     mxnet_training_latest_py_version,
     tmpdir,
 ):
-    instance_count = ParameterInteger(name="InstanceCountParam")
     session = LocalPipelineNoS3Session()
+    instance_count = ParameterInteger(name="InstanceCountParam")
     data_path = os.path.join(DATA_DIR, "mxnet_mnist")
     script_path = os.path.join(data_path, "check_env.py")
     output_path = "file://%s" % (str(tmpdir))
@@ -587,19 +587,12 @@ def test_local_pipeline_with_training_and_transform_steps(
     )
 
     # define create model step
-    inputs = CreateModelInput(
-        instance_type="local",
-        accelerator_type="local",
-    )
-    create_model_step = CreateModelStep(
-        name="mxnet_mnist_model",
-        model=model,
-        inputs=inputs,
-    )
+    model_step_args = model.create(instance_type="local", accelerator_type="local")
+    model_step = ModelStep(name="mxnet_mnist_model", step_args=model_step_args)
 
     # define transformer
     transformer = Transformer(
-        model_name=create_model_step.properties.ModelName,
+        model_name=model_step.properties.ModelName,
         instance_type="local",
         instance_count=instance_count,
         output_path=output_path,
@@ -619,7 +612,7 @@ def test_local_pipeline_with_training_and_transform_steps(
     pipeline = Pipeline(
         name="local_pipeline_training_transform",
         parameters=[instance_count],
-        steps=[training_step, create_model_step, transform_step],
+        steps=[training_step, model_step, transform_step],
         sagemaker_session=session,
     )
 

tests/unit/sagemaker/local/test_local_session.py

@@ -887,16 +887,25 @@ def test_create_describe_update_pipeline():
         steps=[CustomStep(name="MyStep", input_data=parameter)],
         sagemaker_session=LocalSession(),
     )
+    definition = pipeline.definition()
     pipeline.create("dummy-role", "pipeline-description")
 
     pipeline_describe_response1 = pipeline.describe()
     assert pipeline_describe_response1["PipelineArn"] == "MyPipeline"
-    assert pipeline_describe_response1["PipelineDefinition"] == pipeline.definition()
+    assert pipeline_describe_response1["PipelineDefinition"] == definition
     assert pipeline_describe_response1["PipelineDescription"] == "pipeline-description"
 
+    pipeline = Pipeline(
+        name="MyPipeline",
+        parameters=[parameter],
+        steps=[CustomStep(name="MyStepUpdated", input_data=parameter)],
+        sagemaker_session=LocalSession(),
+    )
+    updated_definition = pipeline.definition()
     pipeline.update("dummy-role", "pipeline-description-2")
     pipeline_describe_response2 = pipeline.describe()
     assert pipeline_describe_response2["PipelineDescription"] == "pipeline-description-2"
+    assert pipeline_describe_response2["PipelineDefinition"] == updated_definition
     assert (
         pipeline_describe_response2["CreationTime"]
         != pipeline_describe_response2["LastModifiedTime"]

tests/unit/sagemaker/workflow/test_pipeline_graph.py

@@ -61,7 +61,7 @@ def pipeline_graph_get_sub_dag(sagemaker_session_mock):
     step_f = CustomStep(name="stepF", depends_on=[step_c])
     step_g = CustomStep(name="stepG", depends_on=[step_e, step_d])
     step_h = CustomStep(name="stepH", depends_on=[step_g])
-    step_i = CustomStep(name="stepI", depends_on=[step_h])
+    step_i = CustomStepCollection(name="stepI", depends_on=[step_h])
     step_j = CustomStep(name="stepJ", depends_on=[step_h])
 
     pipeline = Pipeline(
@@ -312,7 +312,8 @@ def test_pipeline_graph_cyclic(sagemaker_session_mock):
                 "stepF",
                 "stepG",
                 "stepH",
-                "stepI",
+                "stepI-0",
+                "stepI-1",
                 "stepJ",
             },
         ),
@@ -326,22 +327,24 @@ def test_pipeline_graph_cyclic(sagemaker_session_mock):
                 "stepF",
                 "stepG",
                 "stepH",
-                "stepI",
+                "stepI-0",
+                "stepI-1",
                 "stepJ",
             },
         ),
-        ("stepC", {"stepC", "stepE", "stepF", "stepG", "stepH", "stepI", "stepJ"}),
-        ("stepD", {"stepD", "stepG", "stepH", "stepI", "stepJ"}),
-        ("stepE", {"stepE", "stepG", "stepH", "stepI", "stepJ"}),
+        ("stepC", {"stepC", "stepE", "stepF", "stepG", "stepH", "stepI-0", "stepI-1", "stepJ"}),
+        ("stepD", {"stepD", "stepG", "stepH", "stepI-0", "stepI-1", "stepJ"}),
+        ("stepE", {"stepE", "stepG", "stepH", "stepI-0", "stepI-1", "stepJ"}),
         ("stepF", {"stepF"}),
-        ("stepG", {"stepG", "stepH", "stepI", "stepJ"}),
-        ("stepH", {"stepH", "stepI", "stepJ"}),
-        ("stepI", {"stepI"}),
+        ("stepG", {"stepG", "stepH", "stepI-0", "stepI-1", "stepJ"}),
+        ("stepH", {"stepH", "stepI-0", "stepI-1", "stepJ"}),
+        ("stepI", {"stepI-0", "stepI-1"}),
         ("stepJ", {"stepJ"}),
     ],
 )
 def test_get_steps_in_sub_dag(pipeline_graph_get_sub_dag, step_name, expected_steps):
-    sub_steps = pipeline_graph_get_sub_dag.get_steps_in_sub_dag(step_name)
+    step = pipeline_graph_get_sub_dag.step_map.get(step_name)
+    sub_steps = pipeline_graph_get_sub_dag.get_steps_in_sub_dag(step)
     assert sub_steps == expected_steps