aws / aws/sagemaker-python-sdk
SageMakerClient CustomSession uses default values instead of passed user session
- Dominant language
- Python
- Stars
- 2.3k
- Forks
- 1.3k
- Avg merge
- 2d 2h
- Merged PRs (30d)
- 32
Description
**PySDK Version**
- [ ] PySDK V2 (2.x)
- [x] PySDK V3 (3.x)
**Describe the bug**
SageMakerClient builds the sagemaker client from botocore.session.get_session(), ignoring the provided user session resulting in cross-account CreateTrainingJob failures.
In sagemaker-core 2.16.0, SageMakerClient.__init__ ([sagemaker/core/utils/utils.py, ~L395–406](https://github.com/aws/sagemaker-python-sdk/blob/v3.15.1/sagemaker-core/src/sagemaker/core/utils/utils.py#L380-L408)) creates self.sagemaker_client from a fresh default botocore session (botocore.session.get_session() → custom_session) rather than the session argument, under a # TODO: Remove post-launch custom service-model loader. The sagemaker-runtime/featurestore/metrics clients correctly use session. Because the class is a SingletonMeta singleton, the first (default-credential) client is reused process-wide ignoring passed in user sessions.
Impact: callers that pass an explicit boto3 session (e.g. ModelTrainer(sagemaker_session=…) → TrainingJob.create(session=…)) have the control-plane call issued under the ambient default AWS profile instead. When the execution role is in a different account than the default profile, CreateTrainingJob fails with `ValidationException: RoleArn: Cross-account pass role is not allowed`.
**To reproduce**
If a user has multiple AWS profiles (all for different accounts) set up in ~/.aws/config and the default in ~/.aws/credentials is NOT the same account as the intended profile that is used to create the SageMaker Session then CreateTrainingJob will fail with `ValidationException: RoleArn: Cross-account pass role is not allowed.`
```python
from sagemaker.core.helper.session_helper import Session
from sagemaker.train import ModelTrainer
from sagemaker.train.configs import SourceCode, Compute, InputData
import boto3
sess = Session(boto_session=boto3.Session(profile_name="acct-A", region_name="us-east-2"))
ModelTrainer(
sagemaker_session=sess, # account A
role="arn:aws:iam:::role/", # account A
training_image=".dkr.ecr.us-east-2.amazonaws.com/img:latest",
compute=Compute(instance_type="ml.m5.xlarge", instance_count=1),
source_code=SourceCode(source_dir="src", entry_script="train.py"),
).train(input_data_config=[InputData(channel_name="train", data_source="s3://bucket/train/")])
# -> ValidationException: RoleArn: Cross-account pass role is not allowed
# (only when your *default* AWS profile is a different account than A)
```
**Expected behavior**
A training job to be created and executed in the specified account. This is a migration from v2 to v3 and the v2 implementation of this works without issue creating training jobs in the targeted accounts based on the profile passed to the Session.
**Screenshots or logs**
If applicable, add screenshots or logs to help explain your problem.
**System information**
A description of your system. Please provide:
- **SageMaker Python SDK version**: 3.15.1
- **Framework name (eg. PyTorch) or algorithm (eg. KMeans)**: SKLearn
- **Framework version**: 1.4-2
- **Python version**: py3
- **CPU or GPU**: CPU
- **Custom Docker image (Y/N)**: N
**Additional context**
This issue also seems to persist to the latest version of 3.16 since the same custom_session is also used to initiate the sagemaker_client.
Contributor guide
Assessment
This issue has not been assessed yet.