Skip to content

chore: excessive jumpstart bucket logging #4053

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Aug 11, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 12 additions & 0 deletions src/sagemaker/jumpstart/accessors.py
Original file line number Diff line number Diff line change
Expand Up @@ -42,8 +42,20 @@ class JumpStartModelsAccessor(object):
_cache: Optional[cache.JumpStartModelsCache] = None
_curr_region = JUMPSTART_DEFAULT_REGION_NAME

_content_bucket: Optional[str] = None

_cache_kwargs: Dict[str, Any] = {}

@staticmethod
def set_jumpstart_content_bucket(content_bucket: str) -> None:
"""Sets JumpStart content bucket."""
JumpStartModelsAccessor._content_bucket = content_bucket

@staticmethod
def get_jumpstart_content_bucket() -> Optional[str]:
"""Returns JumpStart content bucket."""
return JumpStartModelsAccessor._content_bucket

@staticmethod
def _validate_and_mutate_region_cache_kwargs(
cache_kwargs: Optional[Dict[str, Any]] = None, region: Optional[str] = None
Expand Down
38 changes: 27 additions & 11 deletions src/sagemaker/jumpstart/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -72,21 +72,37 @@ def get_jumpstart_content_bucket(
RuntimeError: If JumpStart is not launched in ``region``.
"""

old_content_bucket: Optional[
str
] = accessors.JumpStartModelsAccessor.get_jumpstart_content_bucket()

info_logs: List[str] = []

bucket_to_return: Optional[str] = None
if (
constants.ENV_VARIABLE_JUMPSTART_CONTENT_BUCKET_OVERRIDE in os.environ
and len(os.environ[constants.ENV_VARIABLE_JUMPSTART_CONTENT_BUCKET_OVERRIDE]) > 0
):
bucket_override = os.environ[constants.ENV_VARIABLE_JUMPSTART_CONTENT_BUCKET_OVERRIDE]
constants.JUMPSTART_LOGGER.info("Using JumpStart bucket override: '%s'", bucket_override)
return bucket_override
try:
return constants.JUMPSTART_REGION_NAME_TO_LAUNCHED_REGION_DICT[region].content_bucket
except KeyError:
formatted_launched_regions_str = get_jumpstart_launched_regions_message()
raise ValueError(
f"Unable to get content bucket for JumpStart in {region} region. "
f"{formatted_launched_regions_str}"
)
bucket_to_return = os.environ[constants.ENV_VARIABLE_JUMPSTART_CONTENT_BUCKET_OVERRIDE]
info_logs.append(f"Using JumpStart bucket override: '{bucket_to_return}'")
else:
try:
bucket_to_return = constants.JUMPSTART_REGION_NAME_TO_LAUNCHED_REGION_DICT[
region
].content_bucket
except KeyError:
formatted_launched_regions_str = get_jumpstart_launched_regions_message()
raise ValueError(
f"Unable to get content bucket for JumpStart in {region} region. "
f"{formatted_launched_regions_str}"
)

accessors.JumpStartModelsAccessor.set_jumpstart_content_bucket(bucket_to_return)

if bucket_to_return != old_content_bucket:
for info_log in info_logs:
constants.JUMPSTART_LOGGER.info(info_log)
return bucket_to_return


def get_formatted_manifest(
Expand Down
5 changes: 1 addition & 4 deletions tests/unit/sagemaker/jumpstart/test_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -57,10 +57,7 @@ def test_get_jumpstart_content_bucket_override():
with patch("logging.Logger.info") as mocked_info_log:
random_region = "random_region"
assert "some-val" == utils.get_jumpstart_content_bucket(random_region)
mocked_info_log.assert_called_once_with(
"Using JumpStart bucket override: '%s'",
"some-val",
)
mocked_info_log.assert_called_once_with("Using JumpStart bucket override: 'some-val'")


def test_get_jumpstart_launched_regions_message():
Expand Down