From 231215c141f332b52d46275b9a0020da0dbe4c06 Mon Sep 17 00:00:00 2001 From: qidewenwhen Date: Tue, 19 Mar 2024 17:18:42 -0700 Subject: [PATCH] fix: Create workflow module scoped sagemaker_session to resolve test race condition --- tests/integ/sagemaker/workflow/conftest.py | 39 ++++++++++++++++++++++ 1 file changed, 39 insertions(+) create mode 100644 tests/integ/sagemaker/workflow/conftest.py diff --git a/tests/integ/sagemaker/workflow/conftest.py b/tests/integ/sagemaker/workflow/conftest.py new file mode 100644 index 0000000000..7448459245 --- /dev/null +++ b/tests/integ/sagemaker/workflow/conftest.py @@ -0,0 +1,39 @@ +# 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 + +import pytest +from botocore.config import Config + +from conftest import CUSTOM_S3_OBJECT_KEY_PREFIX +from sagemaker import Session + + +@pytest.fixture(scope="module") +def sagemaker_session( + sagemaker_client_config, + boto_session, +): + sagemaker_client_config.setdefault("config", Config(retries=dict(max_attempts=10))) + sagemaker_client = ( + boto_session.client("sagemaker", **sagemaker_client_config) + if sagemaker_client_config + else None + ) + + return Session( + boto_session=boto_session, + sagemaker_client=sagemaker_client, + sagemaker_config={}, + default_bucket_prefix=CUSTOM_S3_OBJECT_KEY_PREFIX, + )