-
Notifications
You must be signed in to change notification settings - Fork 1.2k
feat: support pipeline versioning #5248
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: master
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -125,6 +125,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") | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Is this always guaranteed to be the latest version when sort_order is None? There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Yes, the default sort order is descending |
||
|
||
def create( | ||
self, | ||
role_arn: str = None, | ||
|
@@ -166,7 +175,8 @@ def create( | |
kwargs, | ||
Tags=tags, | ||
) | ||
return self.sagemaker_session.sagemaker_client.create_pipeline(**kwargs) | ||
response = self.sagemaker_session.sagemaker_client.create_pipeline(**kwargs) | ||
return response | ||
Comment on lines
-169
to
+179
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Nit : Is this change needed ? There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Ah no. Would you like me to revert this? |
||
|
||
def _create_args( | ||
self, role_arn: str, description: str, parallelism_config: ParallelismConfiguration | ||
|
@@ -214,15 +224,21 @@ 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, | ||
|
@@ -257,7 +273,8 @@ def update( | |
return self.sagemaker_session.sagemaker_client.update_pipeline(self, description) | ||
|
||
kwargs = self._create_args(role_arn, description, parallelism_config) | ||
return self.sagemaker_session.sagemaker_client.update_pipeline(**kwargs) | ||
response = self.sagemaker_session.sagemaker_client.update_pipeline(**kwargs) | ||
return response | ||
|
||
def upsert( | ||
self, | ||
|
@@ -332,6 +349,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. | ||
|
||
|
@@ -345,6 +363,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. | ||
|
@@ -366,6 +386,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) | ||
|
@@ -461,6 +482,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( | ||
|
@@ -855,7 +902,7 @@ def describe(self): | |
sagemaker.html#SageMaker.Client.describe_pipeline_execution>`_. | ||
""" | ||
return self.sagemaker_session.sagemaker_client.describe_pipeline_execution( | ||
PipelineExecutionArn=self.arn, | ||
PipelineExecutionArn=self.arn | ||
) | ||
|
||
def list_steps(self): | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I think we had previously relaxed the requirement to enable airflow and resolve some dependency conflicts. Any reason why this is needed ?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
yes, we launched new apis UpdatePipelineVersion and ListPipelineVersions that were added to boto in version 1.39.5- boto/boto3@a1a8371#diff-2c623f3c6a917be56c59d43279244996836262cb1e12d9d0786c9c49eef6b43c
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
So currently we needed to downgrade to a lower version because airflow, which is another open source dependency depends on a lower boto version : #5245
If we do make this update , it would break customer experiences .