feat: support pipeline versioning

For commit: aws/sagemaker-python-sdk-staging@9bfe85a

Author: zhaoqizqwang

sagemaker-mlops/src/sagemaker/mlops/workflow/pipeline.py

@@ -130,6 +130,15 @@ def __init__(
             self.sagemaker_session.boto_session.client("scheduler"),
         )
 
+    @property
+    def latest_pipeline_version_id(self):
+        """Retrieves the latest version id of this pipeline"""
+        summaries = self.list_pipeline_versions(max_results=1)["PipelineVersionSummaries"]
+        if not summaries:
+            return None
+        else:
+            return summaries[0].get("PipelineVersionId")
+
     def create(
         self,
         role_arn: str = None,
@@ -219,15 +228,22 @@ def _create_args(
         )
         return kwargs
 
-    def describe(self) -> Dict[str, Any]:
+    def describe(self, pipeline_version_id: int = None) -> Dict[str, Any]:
         """Describes a Pipeline in the Workflow service.
 
+        Args:
+            pipeline_version_id (Optional[str]): version ID of the pipeline to describe.
+
         Returns:
             Response dict from the service. See `boto3 client documentation
             <https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/\
 sagemaker.html#SageMaker.Client.describe_pipeline>`_
         """
-        return self.sagemaker_session.sagemaker_client.describe_pipeline(PipelineName=self.name)
+        kwargs = dict(PipelineName=self.name)
+        if pipeline_version_id:
+            kwargs["PipelineVersionId"] = pipeline_version_id
+
+        return self.sagemaker_session.sagemaker_client.describe_pipeline(**kwargs)
 
     def update(
         self,
@@ -337,6 +353,7 @@ def start(
         execution_description: str = None,
         parallelism_config: ParallelismConfiguration = None,
         selective_execution_config: SelectiveExecutionConfig = None,
+        pipeline_version_id: int = None,
     ):
         """Starts a Pipeline execution in the Workflow service.
 
@@ -350,6 +367,8 @@ def start(
                 over the parallelism configuration of the parent pipeline.
             selective_execution_config (Optional[SelectiveExecutionConfig]): The configuration for
                 selective step execution.
+            pipeline_version_id (Optional[str]): version ID of the pipeline to start the execution from. If not
+                specified, uses the latest version ID.
 
         Returns:
             A `_PipelineExecution` instance, if successful.
@@ -371,6 +390,7 @@ def start(
             PipelineExecutionDisplayName=execution_display_name,
             ParallelismConfiguration=parallelism_config,
             SelectiveExecutionConfig=selective_execution_config,
+            PipelineVersionId=pipeline_version_id,
         )
         if self.sagemaker_session.local_mode:
             update_args(kwargs, PipelineParameters=parameters)
@@ -466,6 +486,32 @@ def list_executions(
             if key in response
         }
 
+    def list_pipeline_versions(
+        self, sort_order: str = None, max_results: int = None, next_token: str = None
+    ) -> str:
+        """Lists a pipeline's versions.
+
+        Args:
+            sort_order (str): The sort order for results (Ascending/Descending).
+            max_results (int): The maximum number of pipeline executions to return in the response.
+            next_token (str):  If the result of the previous `ListPipelineExecutions` request was
+                truncated, the response includes a `NextToken`. To retrieve the next set of pipeline
+                executions, use the token in the next request.
+
+        Returns:
+            List of Pipeline Version Summaries. See
+            boto3 client list_pipeline_versions
+            https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker/client/list_pipeline_versions.html#
+        """
+        kwargs = dict(PipelineName=self.name)
+        update_args(
+            kwargs,
+            SortOrder=sort_order,
+            NextToken=next_token,
+            MaxResults=max_results,
+        )
+        return self.sagemaker_session.sagemaker_client.list_pipeline_versions(**kwargs)
+
     def _get_latest_execution_arn(self):
         """Retrieves the latest execution of this pipeline"""
         response = self.list_executions(

sagemaker-mlops/tests/unit/workflow/test_pipeline.py

@@ -417,6 +417,62 @@ def test_pipeline_graph_iteration(mock_step):
     assert len(steps) == 1
 
 
+
+
+
+def test_generate_step_map_duplicate_names():
+    from sagemaker.mlops.workflow.pipeline import _generate_step_map
+    
+    step1 = Mock(spec=Step)
+    step1.name = "duplicate"
+    step2 = Mock(spec=Step)
+    step2.name = "duplicate"
+    
+    step_map = {}
+    with pytest.raises(ValueError, match="duplicate names"):
+        _generate_step_map([step1, step2], step_map)
+
+
+def test_pipeline_latest_version_id(mock_session, mock_step):
+    pipeline = Pipeline(name="test-pipeline", steps=[mock_step], sagemaker_session=mock_session)
+    mock_session.sagemaker_client.list_pipeline_versions.return_value = {
+        "PipelineVersionSummaries": [{"PipelineVersionId": 123}]
+    }
+    assert pipeline.latest_pipeline_version_id == 123
+
+
+def test_pipeline_latest_version_id_none(mock_session, mock_step):
+    pipeline = Pipeline(name="test-pipeline", steps=[mock_step], sagemaker_session=mock_session)
+    mock_session.sagemaker_client.list_pipeline_versions.return_value = {
+        "PipelineVersionSummaries": []
+    }
+    assert pipeline.latest_pipeline_version_id is None
+
+
+def test_pipeline_describe_with_version_id(mock_session, mock_step):
+    pipeline = Pipeline(name="test-pipeline", steps=[mock_step], sagemaker_session=mock_session)
+    pipeline.describe(pipeline_version_id=123)
+    mock_session.sagemaker_client.describe_pipeline.assert_called_once_with(
+        PipelineName="test-pipeline", PipelineVersionId=123
+    )
+
+
+def test_pipeline_start_with_version_id(mock_session, mock_step):
+    pipeline = Pipeline(name="test-pipeline", steps=[mock_step], sagemaker_session=mock_session)
+    mock_session.sagemaker_client.start_pipeline_execution.return_value = {"PipelineExecutionArn": "arn"}
+    pipeline.start(pipeline_version_id=123)
+    call_kwargs = mock_session.sagemaker_client.start_pipeline_execution.call_args[1]
+    assert call_kwargs["PipelineVersionId"] == 123
+
+
+def test_pipeline_list_versions(mock_session, mock_step):
+    pipeline = Pipeline(name="test-pipeline", steps=[mock_step], sagemaker_session=mock_session)
+    pipeline.list_pipeline_versions(sort_order="Descending", max_results=10)
+    mock_session.sagemaker_client.list_pipeline_versions.assert_called_once_with(
+        PipelineName="test-pipeline", SortOrder="Descending", MaxResults=10
+    )
+
+
 def test_pipeline_execution_result_waiter_error(mock_session):
     from sagemaker.mlops.workflow.pipeline import _PipelineExecution
     from botocore.exceptions import WaiterError