Merge branch 'master' into docstring-fix

Author: jerrypeng7773

src/sagemaker/lambda_helper.py

@@ -15,6 +15,7 @@
 
 from io import BytesIO
 import zipfile
+import time
 from botocore.exceptions import ClientError
 from sagemaker.session import Session
 
@@ -134,32 +135,35 @@ def update(self):
         Returns: boto3 response from Lambda's update_function method.
         """
         lambda_client = _get_lambda_client(self.session)
-
-        if self.script is not None:
-            try:
-                response = lambda_client.update_function_code(
-                    FunctionName=self.function_name, ZipFile=_zip_lambda_code(self.script)
-                )
-                return response
-            except ClientError as e:
-                error = e.response["Error"]
-                raise ValueError(error)
-        else:
+        retry_attempts = 7
+        for i in range(retry_attempts):
             try:
-                response = lambda_client.update_function_code(
-                    FunctionName=(self.function_name or self.function_arn),
-                    S3Bucket=self.s3_bucket,
-                    S3Key=_upload_to_s3(
-                        s3_client=_get_s3_client(self.session),
-                        function_name=self.function_name,
-                        zipped_code_dir=self.zipped_code_dir,
-                        s3_bucket=self.s3_bucket,
-                    ),
-                )
+                if self.script is not None:
+                    response = lambda_client.update_function_code(
+                        FunctionName=self.function_name, ZipFile=_zip_lambda_code(self.script)
+                    )
+                else:
+                    response = lambda_client.update_function_code(
+                        FunctionName=(self.function_name or self.function_arn),
+                        S3Bucket=self.s3_bucket,
+                        S3Key=_upload_to_s3(
+                            s3_client=_get_s3_client(self.session),
+                            function_name=self.function_name,
+                            zipped_code_dir=self.zipped_code_dir,
+                            s3_bucket=self.s3_bucket,
+                        ),
+                    )
                 return response
             except ClientError as e:
                 error = e.response["Error"]
-                raise ValueError(error)
+                code = error["Code"]
+                if code == "ResourceConflictException":
+                    if i == retry_attempts - 1:
+                        raise ValueError(error)
+                    # max wait time = 2**0 + 2**1 + .. + 2**6 = 127 seconds
+                    time.sleep(2**i)
+                else:
+                    raise ValueError(error)
 
     def upsert(self):
         """Method to create a lambda function or update it if it already exists

src/sagemaker/workflow/_utils.py

@@ -17,7 +17,7 @@
 import shutil
 import tarfile
 import tempfile
-from typing import List, Union, Optional
+from typing import List, Union, Optional, TYPE_CHECKING
 from sagemaker import image_uris
 from sagemaker.inputs import TrainingInput
 from sagemaker.estimator import EstimatorBase
@@ -34,6 +34,9 @@
 from sagemaker.utils import _save_model, download_file_from_url
 from sagemaker.workflow.retry import RetryPolicy
 
+if TYPE_CHECKING:
+    from sagemaker.workflow.step_collections import StepCollection
+
 FRAMEWORK_VERSION = "0.23-1"
 INSTANCE_TYPE = "ml.m5.large"
 REPACK_SCRIPT = "_repack_model.py"
@@ -57,7 +60,7 @@ def __init__(
         description: str = None,
         source_dir: str = None,
         dependencies: List = None,
-        depends_on: Union[List[str], List[Step]] = None,
+        depends_on: Optional[List[Union[str, Step, "StepCollection"]]] = None,
         retry_policies: List[RetryPolicy] = None,
         subnets=None,
         security_group_ids=None,
@@ -124,8 +127,9 @@ def __init__(
                         >>>     |------ virtual-env
 
                     This is not supported with "local code" in Local Mode.
-            depends_on (List[str] or List[Step]): A list of step names or instances
-                    this step depends on (default: None).
+            depends_on (List[Union[str, Step, StepCollection]]): The list of `Step`/`StepCollection`
+                names or `Step` instances or `StepCollection` instances that the current `Step`
+                depends on (default: None).
             retry_policies (List[RetryPolicy]): The list of retry policies for the current step
                 (default: None).
             subnets (list[str]): List of subnet ids. If not specified, the re-packing
@@ -274,7 +278,7 @@ def __init__(
         compile_model_family=None,
         display_name: str = None,
         description=None,
-        depends_on: Optional[Union[List[str], List[Step]]] = None,
+        depends_on: Optional[List[Union[str, Step, "StepCollection"]]] = None,
         retry_policies: Optional[List[RetryPolicy]] = None,
         tags=None,
         container_def_list=None,
@@ -311,8 +315,9 @@ def __init__(
                 if specified, a compiled model will be used (default: None).
             display_name (str): The display name of this `_RegisterModelStep` step (default: None).
             description (str): Model Package description (default: None).
-            depends_on (List[str] or List[Step]): A list of step names or instances
-                this step depends on (default: None).
+            depends_on (List[Union[str, Step, StepCollection]]): The list of `Step`/`StepCollection`
+                names or `Step` instances or `StepCollection` instances that the current `Step`
+                depends on (default: None).
             retry_policies (List[RetryPolicy]): The list of retry policies for the current step
                 (default: None).
             tags (List[dict[str, str]]): A list of dictionaries containing key-value pairs used to

src/sagemaker/workflow/callback_step.py

@@ -13,7 +13,7 @@
 """The step definitions for workflow."""
 from __future__ import absolute_import
 
-from typing import List, Dict, Union
+from typing import List, Dict, Union, Optional
 from enum import Enum
 
 import attr
@@ -27,6 +27,7 @@
 from sagemaker.workflow.entities import (
     DefaultEnumMeta,
 )
+from sagemaker.workflow.step_collections import StepCollection
 from sagemaker.workflow.steps import Step, StepTypeEnum, CacheConfig
 
 
@@ -86,7 +87,7 @@ def __init__(
         display_name: str = None,
         description: str = None,
         cache_config: CacheConfig = None,
-        depends_on: Union[List[str], List[Step]] = None,
+        depends_on: Optional[List[Union[str, Step, StepCollection]]] = None,
     ):
         """Constructs a CallbackStep.
 
@@ -99,8 +100,9 @@ def __init__(
             display_name (str): The display name of the callback step.
             description (str): The description of the callback step.
             cache_config (CacheConfig):  A `sagemaker.workflow.steps.CacheConfig` instance.
-            depends_on (List[str] or List[Step]): A list of step names or step instances
-                this `sagemaker.workflow.steps.CallbackStep` depends on
+            depends_on (List[Union[str, Step, StepCollection]]): A list of `Step`/`StepCollection`
+                names or `Step` instances or `StepCollection` instances that this `CallbackStep`
+                depends on.
         """
         super(CallbackStep, self).__init__(
             name, display_name, description, StepTypeEnum.CALLBACK, depends_on

src/sagemaker/workflow/clarify_check_step.py

@@ -18,7 +18,7 @@
 import os
 import tempfile
 from abc import ABC
-from typing import List, Union
+from typing import List, Union, Optional
 
 import attr
 
@@ -40,6 +40,7 @@
 from sagemaker.workflow import PipelineNonPrimitiveInputTypes, is_pipeline_variable
 from sagemaker.workflow.entities import RequestType
 from sagemaker.workflow.properties import Properties
+from sagemaker.workflow.step_collections import StepCollection
 from sagemaker.workflow.steps import Step, StepTypeEnum, CacheConfig
 from sagemaker.workflow.check_job_config import CheckJobConfig
 
@@ -158,7 +159,7 @@ def __init__(
         display_name: str = None,
         description: str = None,
         cache_config: CacheConfig = None,
-        depends_on: Union[List[str], List[Step]] = None,
+        depends_on: Optional[List[Union[str, Step, StepCollection]]] = None,
     ):
         """Constructs a ClarifyCheckStep.
 
@@ -180,8 +181,9 @@ def __init__(
             description (str): The description of the ClarifyCheckStep step (default: None).
             cache_config (CacheConfig):  A `sagemaker.workflow.steps.CacheConfig` instance
                 (default: None).
-            depends_on (List[str] or List[Step]): A list of step names or step instances
-                this `sagemaker.workflow.steps.ClarifyCheckStep` depends on (default: None).
+            depends_on (List[Union[str, Step, StepCollection]]): A list of `Step`/`StepCollection`
+                names or `Step` instances or `StepCollection` instances that this `ClarifyCheckStep`
+                depends on (default: None).
         """
         if (
             not isinstance(clarify_check_config, DataBiasCheckConfig)

src/sagemaker/workflow/condition_step.py

@@ -13,17 +13,17 @@
 """The step definitions for workflow."""
 from __future__ import absolute_import
 
-from typing import List, Union
+from typing import List, Union, Optional
 
 import attr
 
 from sagemaker.deprecations import deprecated_class
 from sagemaker.workflow.conditions import Condition
+from sagemaker.workflow.step_collections import StepCollection
 from sagemaker.workflow.steps import (
     Step,
     StepTypeEnum,
 )
-from sagemaker.workflow.step_collections import StepCollection
 from sagemaker.workflow.utilities import list_to_request
 from sagemaker.workflow.entities import (
     RequestType,
@@ -41,7 +41,7 @@ class ConditionStep(Step):
     def __init__(
         self,
         name: str,
-        depends_on: Union[List[str], List[Step]] = None,
+        depends_on: Optional[List[Union[str, Step, StepCollection]]] = None,
         display_name: str = None,
         description: str = None,
         conditions: List[Condition] = None,
@@ -56,6 +56,9 @@ def __init__(
 
         Args:
             name (str): The name of the condition step.
+            depends_on (List[Union[str, Step, StepCollection]]): The list of `Step`/StepCollection`
+                names or `Step` instances or `StepCollection` instances that the current `Step`
+                depends on.
             display_name (str): The display name of the condition step.
             description (str): The description of the condition step.
             conditions (List[Condition]): A list of `sagemaker.workflow.conditions.Condition`

src/sagemaker/workflow/emr_step.py

@@ -13,14 +13,15 @@
 """The step definitions for workflow."""
 from __future__ import absolute_import
 
-from typing import List
+from typing import List, Union, Optional
 
 from sagemaker.workflow.entities import (
     RequestType,
 )
 from sagemaker.workflow.properties import (
     Properties,
 )
+from sagemaker.workflow.step_collections import StepCollection
 from sagemaker.workflow.steps import Step, StepTypeEnum, CacheConfig
 
 
@@ -70,7 +71,7 @@ def __init__(
         description: str,
         cluster_id: str,
         step_config: EMRStepConfig,
-        depends_on: List[str] = None,
+        depends_on: Optional[List[Union[str, Step, StepCollection]]] = None,
         cache_config: CacheConfig = None,
     ):
         """Constructs a EMRStep.
@@ -81,8 +82,9 @@ def __init__(
             description(str): The description of the EMR step.
             cluster_id(str): The ID of the running EMR cluster.
             step_config(EMRStepConfig): One StepConfig to be executed by the job flow.
-            depends_on(List[str]):
-                A list of step names this `sagemaker.workflow.steps.EMRStep` depends on
+            depends_on (List[Union[str, Step, StepCollection]]): A list of `Step`/`StepCollection`
+                names or `Step` instances or `StepCollection` instances that this `EMRStep`
+                depends on.
             cache_config(CacheConfig):  A `sagemaker.workflow.steps.CacheConfig` instance.
 
         """

src/sagemaker/workflow/fail_step.py

@@ -13,12 +13,13 @@
 """The `Step` definitions for SageMaker Pipelines Workflows."""
 from __future__ import absolute_import
 
-from typing import List, Union
+from typing import List, Union, Optional
 
 from sagemaker.workflow import PipelineNonPrimitiveInputTypes
 from sagemaker.workflow.entities import (
     RequestType,
 )
+from sagemaker.workflow.step_collections import StepCollection
 from sagemaker.workflow.steps import Step, StepTypeEnum
 
 
@@ -31,7 +32,7 @@ def __init__(
         error_message: Union[str, PipelineNonPrimitiveInputTypes] = None,
         display_name: str = None,
         description: str = None,
-        depends_on: Union[List[str], List[Step]] = None,
+        depends_on: Optional[List[Union[str, Step, StepCollection]]] = None,
     ):
         """Constructs a `FailStep`.
 
@@ -45,8 +46,9 @@ def __init__(
             display_name (str): The display name of the `FailStep`.
                 The display name provides better UI readability. (default: None).
             description (str): The description of the `FailStep` (default: None).
-            depends_on (List[str] or List[Step]): A list of `Step` names or `Step` instances
-                that this `FailStep` depends on.
+            depends_on (List[Union[str, Step, StepCollection]]): A list of `Step`/`StepCollection`
+                names or `Step` instances or `StepCollection` instances that this `FailStep`
+                depends on.
                 If a listed `Step` name does not exist, an error is returned (default: None).
         """
         super(FailStep, self).__init__(

src/sagemaker/workflow/lambda_step.py

@@ -13,7 +13,7 @@
 """The step definitions for workflow."""
 from __future__ import absolute_import
 
-from typing import List, Dict
+from typing import List, Dict, Optional, Union
 from enum import Enum
 
 import attr
@@ -27,6 +27,7 @@
 from sagemaker.workflow.entities import (
     DefaultEnumMeta,
 )
+from sagemaker.workflow.step_collections import StepCollection
 from sagemaker.workflow.steps import Step, StepTypeEnum, CacheConfig
 from sagemaker.lambda_helper import Lambda
 
@@ -87,7 +88,7 @@ def __init__(
         inputs: dict = None,
         outputs: List[LambdaOutput] = None,
         cache_config: CacheConfig = None,
-        depends_on: List[str] = None,
+        depends_on: Optional[List[Union[str, Step, StepCollection]]] = None,
     ):
         """Constructs a LambdaStep.
 
@@ -102,8 +103,9 @@ def __init__(
                 to the lambda function.
             outputs (List[LambdaOutput]): List of outputs from the lambda function.
             cache_config (CacheConfig):  A `sagemaker.workflow.steps.CacheConfig` instance.
-            depends_on (List[str]): A list of step names this `sagemaker.workflow.steps.LambdaStep`
-                depends on
+            depends_on (List[Union[str, Step, StepCollection]]): A list of `Step`/`StepCollection`
+                names or `Step` instances or `StepCollection` instances that this `LambdaStep`
+                depends on.
         """
         super(LambdaStep, self).__init__(
             name, display_name, description, StepTypeEnum.LAMBDA, depends_on

src/sagemaker/workflow/model_step.py

@@ -40,7 +40,7 @@ def __init__(
         self,
         name: str,
         step_args: _ModelStepArguments,
-        depends_on: Optional[Union[List[str], List[Step]]] = None,
+        depends_on: Optional[List[Union[str, Step, StepCollection]]] = None,
         retry_policies: Optional[Union[List[RetryPolicy], Dict[str, List[RetryPolicy]]]] = None,
         display_name: Optional[str] = None,
         description: Optional[str] = None,
@@ -51,8 +51,9 @@ def __init__(
             name (str): The name of the `ModelStep`. A name is required and must be
                 unique within a pipeline.
             step_args (_ModelStepArguments): The arguments for the `ModelStep` definition.
-            depends_on (List[str] or List[Step]): A list of `Step` names or `Step` instances
-                that this `ModelStep` depends on.
+            depends_on (List[Union[str, Step, StepCollection]]): A list of `Step`/`StepCollection`
+                names or `Step` instances or `StepCollection` instances that the first step,
+                in this `ModelStep` collection, depends on.
                 If a listed `Step` name does not exist, an error is returned (default: None).
             retry_policies (List[RetryPolicy] or Dict[str, List[RetryPolicy]]): The list of retry
                 policies for the `ModelStep` (default: None).