Skip to content

Commit cd1bd9a

Browse files
committed
fix: Create workflow module scoped sagemaker_session to resolve test race condition
1 parent c8d1428 commit cd1bd9a

File tree

1 file changed

+40
-0
lines changed

1 file changed

+40
-0
lines changed
Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License"). You
4+
# may not use this file except in compliance with the License. A copy of
5+
# the License is located at
6+
#
7+
# http://aws.amazon.com/apache2.0/
8+
#
9+
# or in the "license" file accompanying this file. This file is
10+
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
11+
# ANY KIND, either express or implied. See the License for the specific
12+
# language governing permissions and limitations under the License.
13+
from __future__ import absolute_import
14+
15+
import pytest
16+
from botocore.config import Config
17+
18+
from sagemaker import Session
19+
20+
CUSTOM_S3_OBJECT_KEY_PREFIX = "session-default-prefix"
21+
22+
23+
@pytest.fixture(scope="module")
24+
def sagemaker_session(
25+
sagemaker_client_config,
26+
boto_session,
27+
):
28+
sagemaker_client_config.setdefault("config", Config(retries=dict(max_attempts=10)))
29+
sagemaker_client = (
30+
boto_session.client("sagemaker", **sagemaker_client_config)
31+
if sagemaker_client_config
32+
else None
33+
)
34+
35+
return Session(
36+
boto_session=boto_session,
37+
sagemaker_client=sagemaker_client,
38+
sagemaker_config={},
39+
default_bucket_prefix=CUSTOM_S3_OBJECT_KEY_PREFIX,
40+
)

0 commit comments

Comments
 (0)