22
22
23
23
from sagemaker .processing import ProcessingInput , ProcessingOutput , ScriptProcessor , Processor
24
24
from sagemaker .sklearn .processing import SKLearnProcessor
25
+ from sagemaker .utils import sts_regional_endpoint
25
26
from tests .integ import DATA_DIR
26
27
from tests .integ .kms_utils import get_or_create_kms_key
27
28
28
29
ROLE = "SageMakerRole"
29
30
DEFAULT_REGION = "us-west-2"
30
- CUSTOM_BUCKET_PATH = "sagemaker-custom-bucket"
31
+ CUSTOM_BUCKET_PATH_PREFIX = "sagemaker-custom-bucket"
31
32
32
33
33
34
@pytest .fixture (scope = "module" )
@@ -49,11 +50,17 @@ def sagemaker_session_with_custom_bucket(
49
50
else None
50
51
)
51
52
53
+ region = boto_session .region_name
54
+ account = boto_session .client (
55
+ "sts" , region_name = region , endpoint_url = sts_regional_endpoint (region )
56
+ ).get_caller_identity ()["Account" ]
57
+ custom_default_bucket = "{}-{}-{}" .format (CUSTOM_BUCKET_PATH_PREFIX , region , account )
58
+
52
59
return Session (
53
60
boto_session = boto_session ,
54
61
sagemaker_client = sagemaker_client ,
55
62
sagemaker_runtime_client = runtime_client ,
56
- default_bucket = CUSTOM_BUCKET_PATH ,
63
+ default_bucket = custom_default_bucket ,
57
64
)
58
65
59
66
@@ -255,10 +262,10 @@ def test_sklearn_with_custom_default_bucket(
255
262
job_description = sklearn_processor .latest_job .describe ()
256
263
257
264
assert job_description ["ProcessingInputs" ][0 ]["InputName" ] == "dummy_input"
258
- assert CUSTOM_BUCKET_PATH in job_description ["ProcessingInputs" ][0 ]["S3Input" ]["S3Uri" ]
265
+ assert CUSTOM_BUCKET_PATH_PREFIX in job_description ["ProcessingInputs" ][0 ]["S3Input" ]["S3Uri" ]
259
266
260
267
assert job_description ["ProcessingInputs" ][1 ]["InputName" ] == "code"
261
- assert CUSTOM_BUCKET_PATH in job_description ["ProcessingInputs" ][1 ]["S3Input" ]["S3Uri" ]
268
+ assert CUSTOM_BUCKET_PATH_PREFIX in job_description ["ProcessingInputs" ][1 ]["S3Input" ]["S3Uri" ]
262
269
263
270
assert job_description ["ProcessingJobName" ].startswith ("test-sklearn-with-customizations" )
264
271
@@ -564,7 +571,7 @@ def test_processor_with_custom_bucket(
564
571
job_description = processor .latest_job .describe ()
565
572
566
573
assert job_description ["ProcessingInputs" ][0 ]["InputName" ] == "code"
567
- assert CUSTOM_BUCKET_PATH in job_description ["ProcessingInputs" ][0 ]["S3Input" ]["S3Uri" ]
574
+ assert CUSTOM_BUCKET_PATH_PREFIX in job_description ["ProcessingInputs" ][0 ]["S3Input" ]["S3Uri" ]
568
575
569
576
assert job_description ["ProcessingJobName" ].startswith ("test-processor" )
570
577
0 commit comments