ModelBuilder.build() ignores the configured session when it reloads a model
Maintainers usually reply within 2 days
A pull request for this has already been merged.
- #6487 by @mohamedzeidan2021 — merged
Assessment
- Difficulty
- 2/5
- Estimated time
- 1-3 hours
- Newbie friendliness
- 72/100
- Issue type
- Bug
- Clarity
- Clearly specified
- Activity status
- Active
- Tech stack
- aws, python
- Domain
- cloud, machine-learning
Research direction
The bug is in the final return of ModelBuilder.build() in sagemaker/serve/model_builder.py, where Model.get() is called without the configured session. Read how self.sagemaker_session is set in that class, then check whether Model.get() accepts a session argument. Add a unit test that uses different default and target boto sessions and asserts that Model.get() receives the target session after CreateModel. Done when that test passes and the existing ModelBuilder tests still pass.
Written by the indexing model from the issue text.
Description
PySDK Version
- PySDK V2 (2.x)
- PySDK V3 (3.x)
Describe the bug
ModelBuilder.build() creates a SageMaker model with its configured
sagemaker_session. It then reloads the model with Model.get() without
passing that session.
If the configured session and default boto session use different AWS accounts,
the create request succeeds in the configured account. The subsequent describe
request uses the default account and fails because the model does not exist
there.
The relevant call in sagemaker/serve/model_builder.py is:
return Model.get(model_name=self.model_name, region=self.region)
To reproduce
Configure two AWS CLI profiles that use different AWS accounts. The target
profile must have SageMaker permissions and permission to pass the execution
role.
Set these environment variables:
export TARGET_PROFILE="<profile-for-target-account>"
export DEFAULT_PROFILE="<profile-for-different-account>"
export TARGET_ROLE_ARN="arn:aws:iam::<target-account-id>:role/<sagemaker-execution-role>"
export TARGET_BUCKET="<existing-bucket-in-target-account>"
export IMAGE_URI="<inference-image-uri-accessible-from-target-account>"
export AWS_REGION="us-east-2"
Run this script:
import os
import uuid
import boto3
from botocore.exceptions import ClientError
from sagemaker.core.helper.session_helper import Session
from sagemaker.core.resources import Model
from sagemaker.core.shapes import ContainerDefinition
from sagemaker.serve import ModelBuilder
region = os.environ["AWS_REGION"]
target_profile = os.environ["TARGET_PROFILE"]
default_profile = os.environ["DEFAULT_PROFILE"]
role_arn = os.environ["TARGET_ROLE_ARN"]
bucket = os.environ["TARGET_BUCKET"]
image_uri = os.environ["IMAGE_URI"]
model_name = f"modelbuilder-session-repro-{uuid.uuid4().hex[:8]}"
# Make the process default session use a different account.
boto3.setup_default_session(profile_name=default_profile, region_name=region)
target_boto_session = boto3.Session(
profile_name=target_profile,
region_name=region,
)
target_sagemaker_session = Session(
boto_session=target_boto_session,
default_bucket=bucket,
)
target_account = target_boto_session.client("sts").get_caller_identity()["Account"]
default_account = boto3.client("sts", region_name=region).get_caller_identity()["Account"]
print(f"Target account: {target_account}")
print(f"Default account: {default_account}")
source_model = Model(
model_name=f"{model_name}-source",
primary_container=ContainerDefinition(image=image_uri),
execution_role_arn=role_arn,
)
builder = ModelBuilder(
model=[source_model],
role_arn=role_arn,
sagemaker_session=target_sagemaker_session,
instance_type="ml.m5.large",
)
try:
builder.build(model_name=model_name)
finally:
# The create request succeeds before build() raises.
try:
target_sagemaker_session.sagemaker_client.delete_model(
ModelName=model_name,
)
except ClientError as error:
if error.response["Error"]["Code"] != "ValidationException":
raise
The script raises an error similar to:
botocore.exceptions.ClientError: An error occurred (ValidationException)
when calling the DescribeModel operation:
Could not find model "modelbuilder-session-repro-1234abcd".
Before cleanup, DescribeModel succeeds through target_sagemaker_session.
The model is present in the target account.
Expected behavior
ModelBuilder.build() must use the configured session for all operations.
The method must return the created Model resource from the target account.
The reload call can pass the boto session:
return Model.get(
model_name=self.model_name,
session=self.sagemaker_session.boto_session,
region=self.region,
)
Suggested patch:
diff --git a/src/sagemaker/serve/model_builder.py b/src/sagemaker/serve/model_builder.py
--- a/src/sagemaker/serve/model_builder.py
+++ b/src/sagemaker/serve/model_builder.py
@@
- return Model.get(model_name=self.model_name, region=self.region)
+ return Model.get(
+ model_name=self.model_name,
+ session=self.sagemaker_session.boto_session,
+ region=self.region,
+ )
The regression test must configure different default and target boto sessions.
It must make sure that Model.get() receives the target boto session after
CreateModel.
Screenshots or logs
Target account: <target-account-id>
Default account: <different-account-id>
Creating model with name: modelbuilder-session-repro-1234abcd
Model has been created: modelbuilder-session-repro-1234abcd
ValidationException: Could not find model
"modelbuilder-session-repro-1234abcd".
CloudTrail shows CreateModel in the target account. The failing
DescribeModel request uses the default account.
System information
- SageMaker Python SDK version: 3.24.0
- Framework name (eg. PyTorch) or algorithm (eg. KMeans): Not framework-specific
- Framework version: N/A
- Python version: 3.12.0
- CPU or GPU: CPU, Apple Silicon (arm64)
- Custom Docker image (Y/N): N
Additional context
The failure is hidden when the default boto session uses the same account as
the configured SageMaker session.
This also affects multi-container inference pipelines. CreateModel completes
in the target account, but endpoint creation never starts because the model
reload fails first.
- Dominant language
- Python
- Stars
- 2.3k
- Forks
- 1.3k
- Avg merge
- 2d 20h
- Merged PRs (30d)
- 85
Getting set up
- No Dockerfile or Docker Compose file
- No pull request template
- Read the contributing guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
More from aws/sagemaker-python-sdk
-
Difficulty 2/5 1-3 hours Newbie friendliness 72/100
aws/sagemaker-python-sdk#6495 ·
Maintainers usually reply within 2 days
-
[Bug] ModelTrainer with no input channels emits InputDataConfig: [], which CreatePipeline rejects (min=1) — v2 omitted the keyPossibly taken @sagemaker-bot claimed this 17 days ago. Open
Difficulty 2/5 1-3 hours Newbie friendliness 76/100
aws/sagemaker-python-sdk#6156 · 2 comments ·
Maintainers usually reply within 2 days
-
sagemaker-train should depend on mlflow-skinny, following sagemaker-mlflow 0.5.0Possibly taken @mohamedzeidan2021 claimed this 18 days ago. Open
Difficulty 2/5 Half a day Newbie friendliness 72/100
aws/sagemaker-python-sdk#6152 ·
Maintainers usually reply within 2 days
-
[Bug] ModelTrainer drops TrainingJobName for PipelineSession, breaking use_custom_job_prefix on TrainingStepPossibly taken @mohamedzeidan2021 claimed this 22 days ago. Opencomponent: pipelines type: bug
Difficulty 2/5 1-3 hours Newbie friendliness 72/100
aws/sagemaker-python-sdk#5776 · 1 comment · 1 reaction ·
Maintainers usually reply within 2 days
-
`ModelTrainer` bug with method to load hyperparameters from file for Amazon Nova RecipePossibly taken @sagemaker-bot claimed this 17 days ago. Open
Difficulty 2/5 1-3 hours Newbie friendliness 82/100
aws/sagemaker-python-sdk#5770 ·
Maintainers usually reply within 2 days
All issues in aws/sagemaker-python-sdk
Similar issues
-
bot-triaged module: LrScheduler module: optimizer triaged
Difficulty 2/5 1-3 hours Newbie friendliness 70/100
Maintainers usually reply within 1 day
-
documentation glm5 llm-model
Difficulty 2/5 1-3 hours Newbie friendliness 78/100
vllm-project/vllm-ascend#19062 ·
Maintainers usually reply within 1 day
-
author-page correction
Difficulty 2/5 1-3 hours Newbie friendliness 68/100
acl-org/acl-anthology#10496 · 1 comment ·
Maintainers usually reply within 1 day
-
Issue Type: Bug Report
Difficulty 2/5 1-3 hours Newbie friendliness 82/100
celery/celery#10875 · 1 comment ·
Maintainers usually reply within 1 day
-
enhancement good first issue
Difficulty 2/5 1-3 hours Newbie friendliness 76/100
SasanLabs/VulnerableApp#836 ·
Maintainers usually reply within 1 day