Merge branch 'dev' into feat/override-jumpstart-content-bucket

Author: evakravi

README.rst

@@ -90,6 +90,7 @@ SageMaker Python SDK is tested on:
 - Python 3.6
 - Python 3.7
 - Python 3.8
+- Python 3.9
 
 AWS Permissions
 ~~~~~~~~~~~~~~~

doc/workflows/pipelines/sagemaker.workflow.pipelines.rst

@@ -147,3 +147,5 @@ Steps
 .. autoclass:: sagemaker.workflow.clarify_check_step.ClarifyCheckConfig
 
 .. autoclass:: sagemaker.workflow.clarify_check_step.ClarifyCheckStep
+
+.. autoclass:: sagemaker.workflow.fail_step.FailStep

setup.py

@@ -33,7 +33,7 @@ def read_version():
 
 # Declare minimal set for installation
 required_packages = [
-    "attrs",
+    "attrs==20.3.0",
     "boto3>=1.20.21",
     "google-pasta",
     "numpy>=1.9.0",
@@ -49,36 +49,38 @@ def read_version():
 # Specific use case dependencies
 extras = {
     "local": [
-        "urllib3>=1.21.1,!=1.25,!=1.25.1",
-        "docker-compose>=1.25.2",
-        "docker==5.0.0",
-        "PyYAML>=5.3, <6",  # PyYAML version has to match docker-compose requirements
+        "urllib3==1.26.8",
+        "docker-compose==1.29.2",
+        "docker~=5.0.0",
+        "PyYAML==5.4.1",  # PyYAML version has to match docker-compose requirements
     ],
-    "scipy": ["scipy>=0.19.0"],
+    "scipy": ["scipy==1.5.4"],
 }
 # Meta dependency groups
 extras["all"] = [item for group in extras.values() for item in group]
 # Tests specific dependencies (do not need to be included in 'all')
 extras["test"] = (
     [
         extras["all"],
-        "tox",
-        "flake8",
-        "pytest<6.1.0",
-        "pytest-cov",
-        "pytest-rerunfailures",
-        "pytest-timeout",
+        "tox==3.24.5",
+        "flake8==4.0.1",
+        "pytest==6.0.2",
+        "pytest-cov==3.0.0",
+        "pytest-rerunfailures==10.2",
+        "pytest-timeout==2.1.0",
         "pytest-xdist==2.4.0",
-        "coverage<6.2",
-        "mock",
-        "contextlib2",
-        "awslogs",
-        "black",
+        "coverage>=5.2, <6.2",
+        "mock==4.0.3",
+        "contextlib2==21.6.0",
+        "awslogs==0.14.0",
+        "black==22.1.0",
         "stopit==1.1.2",
-        "apache-airflow==1.10.11",
-        "fabric>=2.0",
-        "requests>=2.20.0, <3",
-        "sagemaker-experiments",
+        "apache-airflow==2.2.3",
+        "apache-airflow-providers-amazon==3.0.0",
+        "attrs==20.3.0",
+        "fabric==2.6.0",
+        "requests==2.27.1",
+        "sagemaker-experiments==0.1.35",
     ],
 )
 

src/sagemaker/model.py

@@ -23,7 +23,6 @@
 import sagemaker
 from sagemaker import (
     fw_utils,
-    image_uris,
     local,
     s3,
     session,
@@ -657,34 +656,6 @@ def _compilation_job_config(
             "job_name": job_name,
         }
 
-    def _compilation_image_uri(self, region, target_instance_type, framework, framework_version):
-        """Retrieve the Neo or Inferentia image URI.
-
-        Args:
-            region (str): The AWS region.
-            target_instance_type (str): Identifies the device on which you want to run
-                your model after compilation, for example: ml_c5. For valid values, see
-                https://docs.aws.amazon.com/sagemaker/latest/dg/API_OutputConfig.html.
-            framework (str): The framework name.
-            framework_version (str): The framework version.
-        """
-        framework_prefix = ""
-        framework_suffix = ""
-
-        if framework == "xgboost":
-            framework_suffix = "-neo"
-        elif target_instance_type.startswith("ml_inf"):
-            framework_prefix = "inferentia-"
-        else:
-            framework_prefix = "neo-"
-
-        return image_uris.retrieve(
-            "{}{}{}".format(framework_prefix, framework, framework_suffix),
-            region,
-            instance_type=target_instance_type,
-            version=framework_version,
-        )
-
     def package_for_edge(
         self,
         output_path,
@@ -849,12 +820,7 @@ def compile(
             if target_instance_family == "ml_eia2":
                 pass
             elif target_instance_family.startswith("ml_"):
-                self.image_uri = self._compilation_image_uri(
-                    self.sagemaker_session.boto_region_name,
-                    target_instance_family,
-                    framework,
-                    framework_version,
-                )
+                self.image_uri = job_status.get("InferenceImage", None)
                 self._is_compiled_model = True
             else:
                 LOGGER.warning(

src/sagemaker/workflow/fail_step.py

@@ -0,0 +1,71 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""The `Step` definitions for SageMaker Pipelines Workflows."""
+from __future__ import absolute_import
+
+from typing import List, Union
+
+from sagemaker.workflow import PipelineNonPrimitiveInputTypes
+from sagemaker.workflow.entities import (
+    RequestType,
+)
+from sagemaker.workflow.steps import Step, StepTypeEnum
+
+
+class FailStep(Step):
+    """`FailStep` for SageMaker Pipelines Workflows."""
+
+    def __init__(
+        self,
+        name: str,
+        error_message: Union[str, PipelineNonPrimitiveInputTypes] = None,
+        display_name: str = None,
+        description: str = None,
+        depends_on: Union[List[str], List[Step]] = None,
+    ):
+        """Constructs a `FailStep`.
+
+        Args:
+            name (str): The name of the `FailStep`. A name is required and must be
+                unique within a pipeline.
+            error_message (str or PipelineNonPrimitiveInputTypes):
+                An error message defined by the user.
+                Once the `FailStep` is reached, the execution fails and the
+                error message is set as the failure reason (default: None).
+            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.
+                If a listed `Step` name does not exist, an error is returned (default: None).
+        """
+        super(FailStep, self).__init__(
+            name, display_name, description, StepTypeEnum.FAIL, depends_on
+        )
+        self.error_message = error_message if error_message is not None else ""
+
+    @property
+    def arguments(self) -> RequestType:
+        """The arguments dictionary that is used to define the `FailStep`."""
+        return dict(ErrorMessage=self.error_message)
+
+    @property
+    def properties(self):
+        """A `Properties` object is not available for the `FailStep`.
+
+        Executing a `FailStep` will terminate the pipeline.
+        `FailStep` properties should not be referenced.
+        """
+        raise RuntimeError(
+            "FailStep is a terminal step and the Properties object is not available for it."
+        )