Jumpstart with Cross-Account Role Assumption Fails with GetObject Access Denial

[mencarellic]

Describe the bug
When using sagemaker.jumpstart.model method and an assumed role (cross-accounts) via Boto3, the call to JumpStartModel fails with a denial for a GetObject call. When running this without the AssumeRole, the same code works fine. However, due to some compliance requirements, I need to assume the role and run the code.

I can also create a model, endpoint, etc, using pure Boto3 however, I lose some of the abstraction, so I'd prefer to use the Sagemaker SDK.

To reproduce
Attempt to use the Jumpstart method while also using an assumed role. I am using this script:

import boto3
import sagemaker
from sagemaker.jumpstart.model import JumpStartModel

sagemaker_role = "arn:aws:iam::123456789012:role/sagemaker-role"
session_name = "AssumedRoleSession"
region = "us-west-2"
model_id = "huggingface-text2text-flan-t5-xxl-fp16"
model_version = "*"

sts_client = boto3.client("sts")
response = sts_client.assume_role(RoleArn=sagemaker_role, RoleSessionName=session_name)
assumed_session = boto3.Session(
 aws_access_key_id=response["Credentials"]["AccessKeyId"],
 aws_secret_access_key=response["Credentials"]["SecretAccessKey"],
 aws_session_token=response["Credentials"]["SessionToken"],
 region_name=region
)

sagemaker_session = sagemaker.Session(boto_session=assumed_session)
model = JumpStartModel(
 model_id=model_id,
 model_version=model_version,
 region=region,
 role=sagemaker_role,
 vpc_config={
 "SecurityGroupIds": [
 "sg-00000000000000000",
 ],
 "Subnets": [
 "subnet-00000000000000000",
 "subnet-11111111111111111",
 ],
 },
 sagemaker_session=sagemaker_session,
 enable_network_isolation=True,
)

model.deploy(
 initial_instance_count=1,
 instance_type='summary-v1-2023-08-03T00-00-00Z',
 endpoint_name='ml.g5.12xlarge',
)

Expected behavior
I expect a model, endpoint configuration, and endpoint to be created based on the parameters specified.

Screenshots or logs
Error:

Traceback (most recent call last):
 File "<stdin>", line 1, in <module>
 File "<stdin>", line 7, in create_jumpstart_model
 File "/usr/local/lib/python3.11/site-packages/sagemaker/jumpstart/model.py", line 266, in __init__
 if not _is_valid_model_id_hook():
 ^^^^^^^^^^^^^^^^^^^^^^^^^
 File "/usr/local/lib/python3.11/site-packages/sagemaker/jumpstart/model.py", line 259, in _is_valid_model_id_hook
 return is_valid_model_id(
 ^^^^^^^^^^^^^^^^^^
 File "/usr/local/lib/python3.11/site-packages/sagemaker/jumpstart/utils.py", line 592, in is_valid_model_id
 models_manifest_list = accessors.JumpStartModelsAccessor._get_manifest(region=region)
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 File "/usr/local/lib/python3.11/site-packages/sagemaker/jumpstart/accessors.py", line 97, in _get_manifest
 return JumpStartModelsAccessor._cache.get_manifest() # type: ignore
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 File "/usr/local/lib/python3.11/site-packages/sagemaker/jumpstart/cache.py", line 342, in get_manifest
 manifest_dict = self._s3_cache.get(
 ^^^^^^^^^^^^^^^^^^^
 File "/usr/local/lib/python3.11/site-packages/sagemaker/utilities/cache.py", line 103, in get
 self.put(key)
 File "/usr/local/lib/python3.11/site-packages/sagemaker/utilities/cache.py", line 126, in put
 value = self._retrieval_function( # type: ignore
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 File "/usr/local/lib/python3.11/site-packages/sagemaker/jumpstart/cache.py", line 323, in _retrieval_function
 formatted_body, etag = self._get_json_file(s3_key, file_type)
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 File "/usr/local/lib/python3.11/site-packages/sagemaker/jumpstart/cache.py", line 266, in _get_json_file
 file_content, etag = self._get_json_file_and_etag_from_s3(key)
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 File "/usr/local/lib/python3.11/site-packages/sagemaker/jumpstart/cache.py", line 243, in _get_json_file_and_etag_from_s3
 response = self._s3_client.get_object(Bucket=self.s3_bucket_name, Key=key)
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 File "/usr/local/lib/python3.11/site-packages/botocore/client.py", line 535, in _api_call
 return self._make_api_call(operation_name, kwargs)
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 File "/usr/local/lib/python3.11/site-packages/botocore/client.py", line 980, in _make_api_call
 raise error_class(parsed_response, operation_name)
botocore.exceptions.ClientError: An error occurred (AccessDenied) when calling the GetObject operation: Access Denied

If I create an S3 client with the same assumed session and try to fetch the manifest object (I'm assuming it is this one, maybe not: s3://jumpstart-cache-prod-us-west-2/models_manifest.json), I can do so:

>>> s3_client = assumed_session.client("s3", region_name=region)
>>> s3_bucket = 'jumpstart-cache-prod-us-west-2'
>>> s3_key = 'models_manifest.json'
>>> local_file_path = 'models_manifest.json'
>>> s3_client.download_file(s3_bucket, s3_key, local_file_path)
>>> with open(local_file_path, 'r') as file:
... first_10_lines = ''.join([next(file) for _ in range(10)])
...
>>> print(first_10_lines)
[
 {
 "model_id": "autogluon-classification-ensemble",
 "version": "1.1.1",
 "min_version": "2.103.0",
 "spec_key": "community_models/autogluon-classification-ensemble/specs_v1.1.1.json"
 },
 {
 "model_id": "autogluon-classification-ensemble",
 "version": "1.1.0",

System information
A description of your system. Please provide:

• SageMaker Python SDK version: 2.174.0
• Framework name (eg. PyTorch) or algorithm (eg. KMeans):
• Framework version:
• Python version: 3.11
• CPU or GPU: GPU
• Custom Docker image (Y/N): N

Additional context
Add any other context about the problem here.

[evakravi]

A PR has been raised with a fix. In the meantime, to get unblocked, you may do pip install git+https://github.com/evakravi/sagemaker-python-sdk@fix/js-cache-s3-client to use a working version of sagemaker.