fix: construct default session when not present (#463)

Author: metrizable

src/sagemaker/workflow/_utils.py

@@ -215,6 +215,8 @@ def __init__(
         transform_instances,
         model_package_group_name=None,
         image_uri=None,
+        model_metrics=None,
+        approval_status="PendingManualApproval",
         compile_model_family=None,
         **kwargs,
     ):
@@ -236,6 +238,9 @@ def __init__(
                 versioned (default: None).
             image_uri (str): The container image uri for Model Package, if not specified,
                 Estimator's training container image will be used (default: None).
+            model_metrics (ModelMetrics): ModelMetrics object (default: None).
+            approval_status (str): Model Approval Status, values can be "Approved", "Rejected",
+                or "PendingManualApproval" (default: "PendingManualApproval").
             compile_model_family (str): Instance family for compiled model, if specified, a compiled
                 model will be used (default: None).
             **kwargs: additional arguments to `create_model`.
@@ -249,6 +254,8 @@ def __init__(
         self.transform_instances = transform_instances
         self.model_package_group_name = model_package_group_name
         self.image_uri = image_uri
+        self.model_metrics = model_metrics
+        self.approval_status = approval_status
         self.compile_model_family = compile_model_family
         self.kwargs = kwargs
 
@@ -300,6 +307,8 @@ def arguments(self) -> RequestType:
             inference_instances=self.inference_instances,
             transform_instances=self.transform_instances,
             model_package_group_name=self.model_package_group_name,
+            model_metrics=self.model_metrics,
+            approval_status=self.approval_status,
         )
         request_dict = model.sagemaker_session._get_create_model_package_request(
             **model_package_args
@@ -309,8 +318,6 @@ def arguments(self) -> RequestType:
             request_dict.pop("CertifyForMarketplace")
         if "Description" in request_dict:
             request_dict.pop("Description")
-        if "ModelApprovalStatus" in request_dict:
-            request_dict.pop("ModelApprovalStatus")
 
         return request_dict
 

src/sagemaker/workflow/pipeline.py

@@ -49,13 +49,15 @@ class Pipeline(Entity):
             `else_steps` of any `ConditionStep`. In particular, any steps that are within the
             `if_steps` or `else_steps` of a `ConditionStep` cannot be listed in the steps of a
             pipeline.
-
+        sagemaker_session (sagemaker.session.Session): Session object which manages interactions
+            with Amazon SageMaker APIs and any other AWS services needed. If not specified, the
+            pipeline creates one using the default AWS configuration chain.
     """
 
     name: str = attr.ib(factory=str)
     parameters: Sequence[Parameter] = attr.ib(factory=list)
     steps: Sequence[Union[Step, StepCollection]] = attr.ib(factory=list)
-    sagemaker_session: Session = attr.ib(default=Session)
+    sagemaker_session: Session = attr.ib(factory=Session)
 
     _version: str = "2020-12-01"
     _metadata: Dict[str, Any] = dict()
@@ -180,7 +182,7 @@ def start(
         self,
         parameters: Dict[str, Any] = None,
         execution_description: str = None,
-    ) -> Dict[str, Any]:
+    ):
         """Starts a Pipeline execution in the Workflow service.
 
         Args:
@@ -189,7 +191,7 @@ def start(
             execution_description (str): Description of the execution.
 
         Returns:
-            Response dict from service.
+            A `_PipelineExecution` instance, if successful.
         """
         exists = True
         try:
@@ -286,10 +288,17 @@ def update_args(args: Dict[str, Any], **kwargs):
 
 @attr.s
 class _PipelineExecution:
-    """Internal class for encapsulating pipeline execution instances"""
+    """Internal class for encapsulating pipeline execution instances.
+
+    Attributes:
+        arn (str): The arn of the pipeline exeuction.
+        sagemaker_session (sagemaker.session.Session): Session object which manages interactions
+            with Amazon SageMaker APIs and any other AWS services needed. If not specified, the
+            pipeline creates one using the default AWS configuration chain.
+    """
 
     arn: str = attr.ib()
-    sagemaker_session: Session = attr.ib(default=Session)
+    sagemaker_session: Session = attr.ib(factory=Session)
 
     def stop(self):
         """Stops a pipeline execution."""

src/sagemaker/workflow/properties.py

@@ -13,7 +13,7 @@
 """The properties definitions for workflow."""
 from __future__ import absolute_import
 
-from typing import Union
+from typing import Dict, Union
 
 import attr
 
@@ -79,7 +79,7 @@ def __init__(self, path: str, shape_name: str = None):
             shape_name (str): botocore sagemaker service model shape name.
         """
         super(PropertiesList, self).__init__(path, shape_name)
-        self._items = dict()
+        self._items: Dict[Union[int, str], Properties] = dict()
 
     def __getitem__(self, item: Union[int, str]):
         """Populate the indexing item with a Property, working for both list and dictionary.
@@ -114,7 +114,7 @@ class PropertyFile(Expression):
     path: str = attr.ib()
 
     @property
-    def expr(self):
+    def expr(self) -> Dict[str, str]:
         """The expression dict for a `PropertyFile`."""
         return {
             "PropertyFileName": self.name,

src/sagemaker/workflow/steps.py

@@ -226,7 +226,7 @@ def __init__(
         self,
         name: str,
         transformer: Transformer,
-        inputs: TransformInput = None,
+        inputs: TransformInput,
     ):
         """Constructs a TrainingStep, given an `Transformer` instance.
 
@@ -237,7 +237,6 @@ def __init__(
             name (str): The name of the transform step.
             transformer (Transformer): A `sagemaker.transformer.Transformer` instance.
             inputs (TransformInput): A `sagemaker.inputs.TransformInput` instance.
-                Defaults to `None`.
         """
         super(TransformStep, self).__init__(name, StepTypeEnum.TRANSFORM)
         self.transformer = transformer

tests/integ/test_artifact_analytics.py

@@ -45,6 +45,7 @@ def generate_artifacts(sagemaker_session):
 
 
 @pytest.mark.canary_quick
+@pytest.mark.skip("Failing as restricted to the SageMaker/Pipeline runtimes")
 def test_artifact_analytics(sagemaker_session):
     with generate_artifacts(sagemaker_session):
         analytics = ArtifactAnalytics(

tests/integ/test_workflow.py

@@ -21,6 +21,7 @@
 import pytest
 
 from botocore.config import Config
+from botocore.exceptions import WaiterError
 from sagemaker.inputs import TrainingInput
 from sagemaker.processing import ProcessingInput, ProcessingOutput
 from sagemaker.pytorch.estimator import PyTorch
@@ -263,6 +264,10 @@ def test_one_step_sklearn_processing_pipeline(
         response = execution.describe()
         assert response["PipelineArn"] == create_arn
 
+        try:
+            execution.wait(delay=5, max_attempts=6)
+        except WaiterError:
+            pass
         execution_steps = execution.list_steps()
         assert len(execution_steps) == 1
         assert execution_steps[0]["StepName"] == "sklearn-process"

tests/unit/sagemaker/workflow/__init__.py

@@ -0,0 +1,34 @@
+# Copyright 2020 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.
+# language governing permissions and limitations under the License.
+"""Helper methods for testing."""
+from __future__ import absolute_import
+
+
+def ordered(obj):
+    """Helper function for dict comparison.
+
+    Recursively orders a json-like dict or list of dicts.
+
+    Args:
+        obj: either a list or a dict
+
+    Returns:
+        either a sorted list of elements or sorted list of tuples
+    """
+    if isinstance(obj, dict):
+        return sorted((k, ordered(v)) for k, v in obj.items())
+    if isinstance(obj, list):
+        return sorted(ordered(x) for x in obj)
+    else:
+        return obj

tests/unit/sagemaker/workflow/helpers.py

@@ -28,6 +28,7 @@
     Step,
     StepTypeEnum,
 )
+from tests.unit.sagemaker.workflow.helpers import ordered
 
 
 class CustomStep(Step):
@@ -49,16 +50,6 @@ def properties(self):
         return self._properties
 
 
-def ordered(obj):
-    """Helper function for dict comparison"""
-    if isinstance(obj, dict):
-        return sorted((k, ordered(v)) for k, v in obj.items())
-    if isinstance(obj, list):
-        return sorted(ordered(x) for x in obj)
-    else:
-        return obj
-
-
 @pytest.fixture
 def role_arn():
     return "arn:role"
@@ -211,11 +202,16 @@ def test_pipeline_basic():
     )
 
 
-def test_pipeline_two_step():
+def test_pipeline_two_step(sagemaker_session_mock):
     parameter = ParameterString("MyStr")
     step1 = CustomStep(name="MyStep1", input_data=parameter)
     step2 = CustomStep(name="MyStep2", input_data=step1.properties.S3Uri)
-    pipeline = Pipeline(name="MyPipeline", parameters=[parameter], steps=[step1, step2])
+    pipeline = Pipeline(
+        name="MyPipeline",
+        parameters=[parameter],
+        steps=[step1, step2],
+        sagemaker_session=sagemaker_session_mock,
+    )
     assert pipeline.to_request() == {
         "Version": "2020-12-01",
         "Metadata": {},

tests/unit/sagemaker/workflow/test_pipeline.py

@@ -32,6 +32,7 @@
     StepCollection,
     RegisterModel,
 )
+from tests.unit.sagemaker.workflow.helpers import ordered
 
 REGION = "us-west-2"
 BUCKET = "my-bucket"
@@ -124,20 +125,23 @@ def test_register_model(estimator):
         inference_instances=["inference_instance"],
         transform_instances=["transform_instance"],
     )
-    assert register_model.request_dicts() == [
-        {
-            "Name": "RegisterModelStep",
-            "Type": "RegisterModel",
-            "Arguments": {
-                "InferenceSpecification": {
-                    "Containers": [
-                        {"Image": "fakeimage", "ModelDataUrl": "s3://my-bucket/model.tar.gz"}
-                    ],
-                    "SupportedContentTypes": ["content_type"],
-                    "SupportedRealtimeInferenceInstanceTypes": ["inference_instance"],
-                    "SupportedResponseMIMETypes": ["response_type"],
-                    "SupportedTransformInstanceTypes": ["transform_instance"],
+    assert ordered(register_model.request_dicts()) == ordered(
+        [
+            {
+                "Name": "RegisterModelStep",
+                "Type": "RegisterModel",
+                "Arguments": {
+                    "InferenceSpecification": {
+                        "Containers": [
+                            {"Image": "fakeimage", "ModelDataUrl": "s3://my-bucket/model.tar.gz"}
+                        ],
+                        "SupportedContentTypes": ["content_type"],
+                        "SupportedRealtimeInferenceInstanceTypes": ["inference_instance"],
+                        "SupportedResponseMIMETypes": ["response_type"],
+                        "SupportedTransformInstanceTypes": ["transform_instance"],
+                    },
+                    "ModelApprovalStatus": "PendingManualApproval",
                 },
             },
-        },
-    ]
+        ]
+    )