Merge branch 'master' into model_config

Author: oyangz

CHANGELOG.md

@@ -1,5 +1,17 @@
 # Changelog
 
+## v2.129.0 (2023-01-19)
+
+### Features
+
+ * add p2 deprecation for PT>=1.13
+ * TF2.11 Update to PySDK
+
+### Bug Fixes and Other Changes
+
+ * Improve Pipeline integ tests and fix resource leak
+ * Update TF version to 2.8.4
+
 ## v2.128.0 (2023-01-10)
 
 ### Features

VERSION

@@ -1 +1 @@
-2.128.1.dev0
+2.129.1.dev0

src/sagemaker/fw_utils.py

@@ -422,7 +422,20 @@ def tar_and_upload_dir(
     script_name = script if directory else os.path.basename(script)
     dependencies = dependencies or []
     key = "%s/sourcedir.tar.gz" % s3_key_prefix
-    tmp = tempfile.mkdtemp()
+    if (
+        settings is not None
+        and settings.local_download_dir is not None
+        and not (
+            os.path.exists(settings.local_download_dir)
+            and os.path.isdir(settings.local_download_dir)
+        )
+    ):
+        raise ValueError(
+            "Inputted directory for storing newly generated temporary directory does "
+            f"not exist: '{settings.local_download_dir}'"
+        )
+    local_download_dir = None if settings is None else settings.local_download_dir
+    tmp = tempfile.mkdtemp(dir=local_download_dir)
     encrypt_artifact = True if settings is None else settings.encrypt_repacked_artifacts
 
     try:

src/sagemaker/session_settings.py

@@ -18,17 +18,25 @@
 class SessionSettings(object):
     """Optional container class for settings to apply to a SageMaker session."""
 
-    def __init__(self, encrypt_repacked_artifacts=True) -> None:
+    def __init__(self, encrypt_repacked_artifacts=True, local_download_dir=None) -> None:
         """Initialize the ``SessionSettings`` of a SageMaker ``Session``.
 
         Args:
             encrypt_repacked_artifacts (bool): Flag to indicate whether to encrypt the artifacts
                 at rest in S3 using the default AWS managed KMS key for S3 when a custom KMS key
                 is not provided (Default: True).
+            local_download_dir (str): Optional. A path specifying the local directory
+                for downloading artifacts. (Default: None).
         """
         self._encrypt_repacked_artifacts = encrypt_repacked_artifacts
+        self._local_download_dir = local_download_dir
 
     @property
     def encrypt_repacked_artifacts(self) -> bool:
         """Return True if repacked artifacts at rest in S3 should be encrypted by default."""
         return self._encrypt_repacked_artifacts
+
+    @property
+    def local_download_dir(self) -> str:
+        """Return path specifying the local directory for downloading artifacts."""
+        return self._local_download_dir

src/sagemaker/utils.py

@@ -358,7 +358,7 @@ def create_tar_file(source_files, target=None):
 
 
 @contextlib.contextmanager
-def _tmpdir(suffix="", prefix="tmp"):
+def _tmpdir(suffix="", prefix="tmp", directory=None):
     """Create a temporary directory with a context manager.
 
     The file is deleted when the context exits.
@@ -369,11 +369,18 @@ def _tmpdir(suffix="", prefix="tmp"):
             suffix, otherwise there will be no suffix.
         prefix (str): If prefix is specified, the file name will begin with that
             prefix; otherwise, a default prefix is used.
+        directory (str): If a directory is specified, the file will be downloaded
+            in this directory; otherwise, a default directory is used.
 
     Returns:
         str: path to the directory
     """
-    tmp = tempfile.mkdtemp(suffix=suffix, prefix=prefix, dir=None)
+    if directory is not None and not (os.path.exists(directory) and os.path.isdir(directory)):
+        raise ValueError(
+            "Inputted directory for storing newly generated temporary "
+            f"directory does not exist: '{directory}'"
+        )
+    tmp = tempfile.mkdtemp(suffix=suffix, prefix=prefix, dir=directory)
     yield tmp
     shutil.rmtree(tmp)
 
@@ -427,7 +434,13 @@ def repack_model(
     """
     dependencies = dependencies or []
 
-    with _tmpdir() as tmp:
+    local_download_dir = (
+        None
+        if sagemaker_session.settings is None
+        or sagemaker_session.settings.local_download_dir is None
+        else sagemaker_session.settings.local_download_dir
+    )
+    with _tmpdir(directory=local_download_dir) as tmp:
         model_dir = _extract_model(model_uri, sagemaker_session, tmp)
 
         _create_or_update_code_dir(

tests/integ/sagemaker/workflow/helpers.py

@@ -0,0 +1,24 @@
+# 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.
+from __future__ import absolute_import
+
+from botocore.exceptions import WaiterError
+
+from sagemaker.workflow.pipeline import _PipelineExecution
+
+
+def wait_pipeline_execution(execution: _PipelineExecution, delay: int = 30, max_attempts: int = 60):
+    try:
+        execution.wait(delay=delay, max_attempts=max_attempts)
+    except WaiterError:
+        pass

tests/integ/sagemaker/workflow/test_automl_steps.py

@@ -16,8 +16,8 @@
 import re
 
 import pytest
-from botocore.exceptions import WaiterError
 
+from tests.integ.sagemaker.workflow.helpers import wait_pipeline_execution
 from sagemaker.workflow import ParameterString
 from sagemaker.workflow.automl_step import AutoMLStep
 from sagemaker.automl.automl import AutoML, AutoMLInput
@@ -133,10 +133,7 @@ def test_automl_step(pipeline_session, role, pipeline_name):
     try:
         _ = pipeline.create(role)
         execution = pipeline.start(parameters={})
-        try:
-            execution.wait(delay=30, max_attempts=60)
-        except WaiterError:
-            pass
+        wait_pipeline_execution(execution=execution)
 
         execution_steps = execution.list_steps()
         has_automl_job = False

tests/integ/sagemaker/workflow/test_clarify_check_steps.py

@@ -17,8 +17,8 @@
 import os
 
 import pytest
-from botocore.exceptions import WaiterError
 
+from tests.integ.sagemaker.workflow.helpers import wait_pipeline_execution
 from sagemaker.clarify import (
     BiasConfig,
     DataConfig,
@@ -182,10 +182,7 @@ def test_one_step_data_bias_pipeline_happycase(
 
             assert response["PipelineArn"] == create_arn
 
-            try:
-                execution.wait(delay=30, max_attempts=60)
-            except WaiterError:
-                pass
+            wait_pipeline_execution(execution=execution)
             execution_steps = execution.list_steps()
 
             assert len(execution_steps) == 1
@@ -272,10 +269,7 @@ def test_one_step_data_bias_pipeline_constraint_violation(
 
             assert response["PipelineArn"] == create_arn
 
-            try:
-                execution.wait(delay=30, max_attempts=60)
-            except WaiterError:
-                pass
+            wait_pipeline_execution(execution=execution)
             execution_steps = execution.list_steps()
 
             assert len(execution_steps) == 1

tests/integ/sagemaker/workflow/test_experiment.py

@@ -17,7 +17,7 @@
 
 import pytest
 
-from botocore.exceptions import WaiterError
+from tests.integ.sagemaker.workflow.helpers import wait_pipeline_execution
 from sagemaker.processing import ProcessingInput
 from sagemaker.session import get_execution_role
 from sagemaker.sklearn.processing import SKLearnProcessor
@@ -120,10 +120,7 @@ def test_pipeline_execution_with_default_experiment_config(
         pipeline.create(role)
         execution = pipeline.start(parameters={})
 
-        try:
-            execution.wait(delay=30, max_attempts=3)
-        except WaiterError:
-            pass
+        wait_pipeline_execution(execution=execution, max_attempts=3)
         execution_steps = execution.list_steps()
         assert len(execution_steps) == 1
         assert execution_steps[0]["StepName"] == "sklearn-process"
@@ -195,10 +192,7 @@ def test_pipeline_execution_with_custom_experiment_config(
         pipeline.create(role)
         execution = pipeline.start(parameters={})
 
-        try:
-            execution.wait(delay=30, max_attempts=3)
-        except WaiterError:
-            pass
+        wait_pipeline_execution(execution=execution, max_attempts=3)
         execution_steps = execution.list_steps()
         assert len(execution_steps) == 1
         assert execution_steps[0]["StepName"] == "sklearn-process"

tests/integ/sagemaker/workflow/test_fail_steps.py

@@ -13,8 +13,8 @@
 from __future__ import absolute_import
 
 import pytest
-from botocore.exceptions import WaiterError
 
+from tests.integ.sagemaker.workflow.helpers import wait_pipeline_execution
 from sagemaker import get_execution_role, utils
 from sagemaker.workflow.condition_step import ConditionStep
 from sagemaker.workflow.conditions import ConditionEquals
@@ -62,10 +62,7 @@ def test_two_step_fail_pipeline_with_str_err_msg(sagemaker_session, role, pipeli
         response = execution.describe()
         assert response["PipelineArn"] == pipeline_arn
 
-        try:
-            execution.wait(delay=30, max_attempts=60)
-        except WaiterError:
-            pass
+        wait_pipeline_execution(execution=execution)
         execution_steps = execution.list_steps()
 
         assert len(execution_steps) == 2
@@ -130,10 +127,7 @@ def test_two_step_fail_pipeline_with_parameter_err_msg(sagemaker_session, role,
         response = execution.describe()
         assert response["PipelineArn"] == pipeline_arn
 
-        try:
-            execution.wait(delay=30, max_attempts=60)
-        except WaiterError:
-            pass
+        wait_pipeline_execution(execution=execution)
         execution_steps = execution.list_steps()
 
         assert len(execution_steps) == 2
@@ -196,10 +190,7 @@ def test_two_step_fail_pipeline_with_join_fn(sagemaker_session, role, pipeline_n
         response = execution.describe()
         assert response["PipelineArn"] == pipeline_arn
 
-        try:
-            execution.wait(delay=30, max_attempts=60)
-        except WaiterError:
-            pass
+        wait_pipeline_execution(execution=execution)
         execution_steps = execution.list_steps()
 
         assert len(execution_steps) == 2
@@ -257,10 +248,7 @@ def test_two_step_fail_pipeline_with_no_err_msg(sagemaker_session, role, pipelin
         response = execution.describe()
         assert response["PipelineArn"] == pipeline_arn
 
-        try:
-            execution.wait(delay=30, max_attempts=60)
-        except WaiterError:
-            pass
+        wait_pipeline_execution(execution=execution)
         execution_steps = execution.list_steps()
 
         assert len(execution_steps) == 2

tests/integ/sagemaker/workflow/test_model_create_and_registration.py

@@ -23,9 +23,9 @@
 import re
 
 import pytest
-from botocore.exceptions import WaiterError
 
 import tests
+from tests.integ.sagemaker.workflow.helpers import wait_pipeline_execution
 from sagemaker.tensorflow import TensorFlow, TensorFlowModel
 from tests.integ.retry import retries
 from sagemaker.drift_check_baselines import DriftCheckBaselines
@@ -180,12 +180,14 @@ def test_conditional_pytorch_training_model_registration(
         )
 
         execution = pipeline.start(parameters={})
+        wait_pipeline_execution(execution=execution)
         assert re.match(
             rf"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}/execution/",
             execution.arn,
         )
 
         execution = pipeline.start(parameters={"GoodEnoughInput": 0})
+        wait_pipeline_execution(execution=execution)
         assert re.match(
             rf"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}/execution/",
             execution.arn,
@@ -259,12 +261,14 @@ def test_mxnet_model_registration(
         )
 
         execution = pipeline.start(parameters={})
+        wait_pipeline_execution(execution=execution)
         assert re.match(
             rf"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}/execution/",
             execution.arn,
         )
 
         execution = pipeline.start()
+        wait_pipeline_execution(execution=execution)
         assert re.match(
             rf"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}/execution/",
             execution.arn,
@@ -470,12 +474,14 @@ def test_sklearn_xgboost_sip_model_registration(
         )
 
         execution = pipeline.start(parameters={})
+        wait_pipeline_execution(execution=execution)
         assert re.match(
             rf"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}/execution/",
             execution.arn,
         )
 
         execution = pipeline.start()
+        wait_pipeline_execution(execution=execution)
         assert re.match(
             rf"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}/execution/",
             execution.arn,
@@ -656,10 +662,7 @@ def test_model_registration_with_drift_check_baselines(
 
             assert response["PipelineArn"] == create_arn
 
-            try:
-                execution.wait(delay=30, max_attempts=60)
-            except WaiterError:
-                pass
+            wait_pipeline_execution(execution=execution)
             execution_steps = execution.list_steps()
 
             assert len(execution_steps) == 1
@@ -797,12 +800,14 @@ def test_model_registration_with_model_repack(
         )
 
         execution = pipeline.start(parameters={})
+        wait_pipeline_execution(execution=execution)
         assert re.match(
             rf"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}/execution/",
             execution.arn,
         )
 
         execution = pipeline.start(parameters={"GoodEnoughInput": 0})
+        wait_pipeline_execution(execution=execution)
         assert re.match(
             rf"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}/execution/",
             execution.arn,
@@ -889,10 +894,7 @@ def test_model_registration_with_tensorflow_model_with_pipeline_model(
             rf"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}/execution/",
             execution.arn,
         )
-        try:
-            execution.wait(delay=30, max_attempts=60)
-        except WaiterError:
-            pass
+        wait_pipeline_execution(execution=execution)
         execution_steps = execution.list_steps()
 
         for step in execution_steps: