Skip to content

chore: telemetry for deployment configs #4806

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.

Sign up for GitHub

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 16 commits into from
Aug 5, 2024
Merged

chore: telemetry for deployment configs #4806

Show file tree
Hide file tree
Changes from 6 commits
Commits
Show all changes
16 commits
Select commit Hold shift + click to select a range
78874b5
chore: telemetry for deployment configs
evakravi Jul 26, 2024
49a1c82
chore: minor fixes
evakravi Jul 26, 2024
3fc60f2
chore: address minor issues
evakravi Jul 26, 2024
df711e3
fix: flake8
evakravi Jul 26, 2024
8f886fe
fix: model type for estimator
evakravi Jul 26, 2024
086b920
Merge branch 'master' into chore/telemetry-for-deployment-configs
evakravi Jul 26, 2024
35b85c2
Merge branch 'master' into chore/telemetry-for-deployment-configs
evakravi Jul 27, 2024
5150ef2
Merge branch 'master' into chore/telemetry-for-deployment-configs
evakravi Jul 30, 2024
e4fa7b0
chore: add ranking name argument to get_top_ranked_config_name
evakravi Jul 30, 2024
e34936b
Merge branch 'master' into chore/telemetry-for-deployment-configs
evakravi Jul 31, 2024
12f8f95
chore: use named args
evakravi Jul 31, 2024
3534b79
fix: remove tuple from model type
evakravi Aug 5, 2024
9d635a4
Merge remote-tracking branch 'origin' into chore/telemetry-for-deploy…
evakravi Aug 5, 2024
3852722
Merge branch 'master' into chore/telemetry-for-deployment-configs
evakravi Aug 5, 2024
fdaa7fd
chore: add comment explaining test
evakravi Aug 5, 2024
f490e46
Merge branch 'master' into chore/telemetry-for-deployment-configs
evakravi Aug 5, 2024
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
2 changes: 2 additions & 0 deletions src/sagemaker/jumpstart/exceptions.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -150,6 +150,7 @@ def __init__(
model. (Default: None).

"""
version = version or "*"
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

sanity check: this would change the behavior on line 157 right? Is that intentional?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This was intentional since version was being supplied as None, not being omitted.
I'm going to try removing this tho

if message:
self.message = message
else:
Expand Down Expand Up @@ -198,6 +199,7 @@ def __init__(
version: Optional[str] = None,
message: Optional[str] = None,
):
version = version or "*"
if message:
self.message = message
else:
Expand Down
32 changes: 18 additions & 14 deletions src/sagemaker/jumpstart/factory/estimator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,7 @@
add_jumpstart_model_info_tags,
get_eula_message,
get_default_jumpstart_session_with_user_agent_suffix,
get_top_ranked_config_name,
update_dict_if_key_not_present,
resolve_estimator_sagemaker_config_field,
verify_model_region_and_return_specs,
Expand Down Expand Up @@ -204,7 +205,9 @@ def get_init_kwargs(

estimator_init_kwargs = _add_model_version_to_kwargs(estimator_init_kwargs)
estimator_init_kwargs = _add_vulnerable_and_deprecated_status_to_kwargs(estimator_init_kwargs)
estimator_init_kwargs = _add_sagemaker_session_to_kwargs(estimator_init_kwargs)
estimator_init_kwargs = _add_sagemaker_session_with_custom_user_agent_to_kwargs(
estimator_init_kwargs
)
estimator_init_kwargs = _add_region_to_kwargs(estimator_init_kwargs)
estimator_init_kwargs = _add_instance_type_and_count_to_kwargs(estimator_init_kwargs)
estimator_init_kwargs = _add_image_uri_to_kwargs(estimator_init_kwargs)
Expand Down Expand Up @@ -438,12 +441,17 @@ def _add_region_to_kwargs(kwargs: JumpStartKwargs) -> JumpStartKwargs:
return kwargs


def _add_sagemaker_session_to_kwargs(kwargs: JumpStartKwargs) -> JumpStartKwargs:
def _add_sagemaker_session_with_custom_user_agent_to_kwargs(
kwargs: JumpStartKwargs,
) -> JumpStartKwargs:
"""Sets session in kwargs based on default or override, returns full kwargs."""
kwargs.sagemaker_session = (
kwargs.sagemaker_session
or get_default_jumpstart_session_with_user_agent_suffix(
kwargs.model_id, kwargs.model_version, kwargs.hub_arn
model_id=kwargs.model_id,
model_version=kwargs.model_version,
config_name=None,
is_hub_content=kwargs.hub_arn is not None,
)
)
return kwargs
Expand Down Expand Up @@ -903,20 +911,16 @@ def _add_config_name_to_kwargs(
) -> JumpStartEstimatorInitKwargs:
"""Sets tags in kwargs based on default or override, returns full kwargs."""

specs = verify_model_region_and_return_specs(
kwargs.config_name = kwargs.config_name or get_top_ranked_config_name(
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for refactoring this!

region=kwargs.region,
model_id=kwargs.model_id,
version=kwargs.model_version,
model_version=kwargs.model_version,
sagemaker_session=kwargs.sagemaker_session,
scope=JumpStartScriptScope.TRAINING,
region=kwargs.region,
tolerate_vulnerable_model=kwargs.tolerate_vulnerable_model,
model_type=kwargs.model_type[0],
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

sanity-check: any risk of getting an IndexError here?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actually, i did this cause the typing said it was a tuple, but looking at other places of the code, we do kwargs.model_type.

@malav-shastri -- any idea what to do here?

tolerate_deprecated_model=kwargs.tolerate_deprecated_model,
sagemaker_session=kwargs.sagemaker_session,
config_name=kwargs.config_name,
tolerate_vulnerable_model=kwargs.tolerate_vulnerable_model,
hub_arn=kwargs.hub_arn,
)

if specs.training_configs and specs.training_configs.get_top_config_from_ranking():
kwargs.config_name = (
kwargs.config_name or specs.training_configs.get_top_config_from_ranking().config_name
)

return kwargs
131 changes: 83 additions & 48 deletions src/sagemaker/jumpstart/factory/model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@
)
from sagemaker.jumpstart.artifacts.resource_names import _retrieve_resource_name_base
from sagemaker.jumpstart.constants import (
DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
INFERENCE_ENTRY_POINT_SCRIPT_NAME,
JUMPSTART_DEFAULT_REGION_NAME,
JUMPSTART_LOGGER,
Expand All @@ -54,6 +55,7 @@
add_jumpstart_model_info_tags,
get_default_jumpstart_session_with_user_agent_suffix,
get_neo_content_bucket,
get_top_ranked_config_name,
update_dict_if_key_not_present,
resolve_model_sagemaker_config_field,
verify_model_region_and_return_specs,
Expand Down Expand Up @@ -155,15 +157,15 @@ def _add_region_to_kwargs(kwargs: JumpStartModelInitKwargs) -> JumpStartModelIni
return kwargs


def _add_sagemaker_session_to_kwargs(
def _add_sagemaker_session_with_custom_user_agent_to_kwargs(
kwargs: Union[JumpStartModelInitKwargs, JumpStartModelDeployKwargs]
) -> JumpStartModelInitKwargs:
"""Sets session in kwargs based on default or override, returns full kwargs."""

kwargs.sagemaker_session = (
kwargs.sagemaker_session
or get_default_jumpstart_session_with_user_agent_suffix(
kwargs.model_id, kwargs.model_version, kwargs.hub_arn
kwargs.model_id, kwargs.model_version, kwargs.config_name, kwargs.hub_arn
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

consider using named args here to avoid the risk of mixing them up.

)
)

Expand Down Expand Up @@ -244,6 +246,32 @@ def _add_instance_type_to_kwargs(
kwargs.instance_type,
)

specs = verify_model_region_and_return_specs(
model_id=kwargs.model_id,
version=kwargs.model_version,
scope=JumpStartScriptScope.INFERENCE,
region=kwargs.region,
tolerate_vulnerable_model=kwargs.tolerate_vulnerable_model,
tolerate_deprecated_model=kwargs.tolerate_deprecated_model,
sagemaker_session=kwargs.sagemaker_session,
model_type=kwargs.model_type,
config_name=kwargs.config_name,
)

if specs.inference_configs and kwargs.config_name not in specs.inference_configs.configs:
return kwargs

resolved_config = (
specs.inference_configs.configs[kwargs.config_name].resolved_config
if specs.inference_configs
else None
)
if resolved_config is None:
return kwargs
supported_instance_types = resolved_config.get("supported_inference_instance_types", [])
if kwargs.instance_type not in supported_instance_types:
JUMPSTART_LOGGER.warning("Overriding instance type to %s", kwargs.instance_type)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit (I know it isn't from this PR): could we use f-string please, also update to Overriding instance type to %s even though it is not in the list of supported instances.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

i think we did this cause there's some linters that complain


return kwargs


Expand Down Expand Up @@ -662,33 +690,25 @@ def _add_config_name_to_init_kwargs(kwargs: JumpStartModelInitKwargs) -> JumpSta
ValueError: If the instance_type is not supported with the current config.
"""

specs = verify_model_region_and_return_specs(
# we need to create a default JS session (without custom user agent)
# in order to retrieve config name info
temp_session = kwargs.sagemaker_session or DEFAULT_JUMPSTART_SAGEMAKER_SESSION

kwargs.config_name = kwargs.config_name or get_top_ranked_config_name(
region=kwargs.region,
model_id=kwargs.model_id,
version=kwargs.model_version,
model_version=kwargs.model_version,
sagemaker_session=temp_session,
scope=JumpStartScriptScope.INFERENCE,
region=kwargs.region,
tolerate_vulnerable_model=kwargs.tolerate_vulnerable_model,
tolerate_deprecated_model=kwargs.tolerate_deprecated_model,
sagemaker_session=kwargs.sagemaker_session,
model_type=kwargs.model_type,
config_name=kwargs.config_name,
tolerate_deprecated_model=kwargs.tolerate_deprecated_model,
tolerate_vulnerable_model=kwargs.tolerate_vulnerable_model,
hub_arn=kwargs.hub_arn,
)
if specs.inference_configs:
default_config_name = specs.inference_configs.get_top_config_from_ranking().config_name
kwargs.config_name = kwargs.config_name or default_config_name

if not kwargs.config_name:
return kwargs

if kwargs.config_name not in set(specs.inference_configs.configs.keys()):
raise ValueError(
f"Config {kwargs.config_name} is not supported for model {kwargs.model_id}."
)
if kwargs.config_name is None:
return kwargs

resolved_config = specs.inference_configs.configs[kwargs.config_name].resolved_config
supported_instance_types = resolved_config.get("supported_inference_instance_types", [])
if kwargs.instance_type not in supported_instance_types:
JUMPSTART_LOGGER.warning("Overriding instance type to %s", kwargs.instance_type)
return kwargs


Expand Down Expand Up @@ -740,27 +760,41 @@ def _add_config_name_to_deploy_kwargs(
ValueError: If the instance_type is not supported with the current config.
"""

specs = verify_model_region_and_return_specs(
model_id=kwargs.model_id,
version=kwargs.model_version,
scope=JumpStartScriptScope.INFERENCE,
region=kwargs.region,
tolerate_vulnerable_model=kwargs.tolerate_vulnerable_model,
tolerate_deprecated_model=kwargs.tolerate_deprecated_model,
sagemaker_session=kwargs.sagemaker_session,
model_type=kwargs.model_type,
config_name=kwargs.config_name,
)
# we need to create a default JS session (without custom user agent)
# in order to retrieve config name info
temp_session = kwargs.sagemaker_session or DEFAULT_JUMPSTART_SAGEMAKER_SESSION
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not sure if I follow this, if kwargs.sagemaker_session is None, don't we fallback to the default session here?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yes, but the problem is we need a session in order to derive what the config name is.

For the session to get the config name, we use the DEFAULT_JUMPSTART_SAGEMAKER_SESSION without any special user agent, unless a custom session is provided


if training_config_name:
kwargs.config_name = _select_inference_config_from_training_config(

specs = verify_model_region_and_return_specs(
model_id=kwargs.model_id,
version=kwargs.model_version,
scope=JumpStartScriptScope.INFERENCE,
region=kwargs.region,
tolerate_vulnerable_model=kwargs.tolerate_vulnerable_model,
tolerate_deprecated_model=kwargs.tolerate_deprecated_model,
sagemaker_session=temp_session,
model_type=kwargs.model_type,
config_name=kwargs.config_name,
)
default_config_name = _select_inference_config_from_training_config(
specs=specs, training_config_name=training_config_name
)
return kwargs

if specs.inference_configs:
default_config_name = specs.inference_configs.get_top_config_from_ranking().config_name
kwargs.config_name = kwargs.config_name or default_config_name
else:
default_config_name = get_top_ranked_config_name(
region=kwargs.region,
model_id=kwargs.model_id,
model_version=kwargs.model_version,
sagemaker_session=temp_session,
scope=JumpStartScriptScope.INFERENCE,
model_type=kwargs.model_type,
tolerate_deprecated_model=kwargs.tolerate_deprecated_model,
tolerate_vulnerable_model=kwargs.tolerate_vulnerable_model,
hub_arn=kwargs.hub_arn,
)

kwargs.config_name = kwargs.config_name or default_config_name

return kwargs

Expand Down Expand Up @@ -839,15 +873,15 @@ def get_deploy_kwargs(
routing_config=routing_config,
)

deploy_kwargs = _add_sagemaker_session_to_kwargs(kwargs=deploy_kwargs)
deploy_kwargs = _add_config_name_to_deploy_kwargs(
kwargs=deploy_kwargs, training_config_name=training_config_name
)

deploy_kwargs = _add_model_version_to_kwargs(kwargs=deploy_kwargs)

deploy_kwargs = _add_endpoint_name_to_kwargs(kwargs=deploy_kwargs)
deploy_kwargs = _add_sagemaker_session_with_custom_user_agent_to_kwargs(kwargs=deploy_kwargs)

deploy_kwargs = _add_config_name_to_deploy_kwargs(
kwargs=deploy_kwargs, training_config_name=training_config_name
)
deploy_kwargs = _add_endpoint_name_to_kwargs(kwargs=deploy_kwargs)

deploy_kwargs = _add_instance_type_to_kwargs(kwargs=deploy_kwargs)

Expand Down Expand Up @@ -1030,11 +1064,14 @@ def get_init_kwargs(
)

model_init_kwargs = _add_vulnerable_and_deprecated_status_to_kwargs(kwargs=model_init_kwargs)
model_init_kwargs = _add_model_version_to_kwargs(kwargs=model_init_kwargs)
model_init_kwargs = _add_config_name_to_init_kwargs(kwargs=model_init_kwargs)

model_init_kwargs = _add_sagemaker_session_to_kwargs(kwargs=model_init_kwargs)
model_init_kwargs = _add_sagemaker_session_with_custom_user_agent_to_kwargs(
kwargs=model_init_kwargs
)
model_init_kwargs = _add_region_to_kwargs(kwargs=model_init_kwargs)

model_init_kwargs = _add_model_version_to_kwargs(kwargs=model_init_kwargs)
model_init_kwargs = _add_model_name_to_kwargs(kwargs=model_init_kwargs)

model_init_kwargs = _add_instance_type_to_kwargs(
Expand Down Expand Up @@ -1062,8 +1099,6 @@ def get_init_kwargs(

model_init_kwargs = _add_resources_to_kwargs(kwargs=model_init_kwargs)

model_init_kwargs = _add_config_name_to_init_kwargs(kwargs=model_init_kwargs)

model_init_kwargs = _add_additional_model_data_sources_to_kwargs(kwargs=model_init_kwargs)

return model_init_kwargs
55 changes: 53 additions & 2 deletions src/sagemaker/jumpstart/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1109,11 +1109,15 @@ def get_jumpstart_configs(


def get_jumpstart_user_agent_extra_suffix(
model_id: Optional[str], model_version: Optional[str], is_hub_content: Optional[bool]
model_id: Optional[str],
model_version: Optional[str],
config_name: Optional[str],
is_hub_content: Optional[bool],
) -> str:
"""Returns the model-specific user agent string to be added to requests."""
sagemaker_python_sdk_headers = get_user_agent_extra_suffix()
jumpstart_specific_suffix = f"md/js_model_id#{model_id} md/js_model_ver#{model_version}"
config_specific_suffix = f"md/js_config#{config_name}"
hub_specific_suffix = f"md/js_is_hub_content#{is_hub_content}"

if os.getenv(constants.ENV_VARIABLE_DISABLE_JUMPSTART_TELEMETRY, None):
Expand All @@ -1128,19 +1132,66 @@ def get_jumpstart_user_agent_extra_suffix(
else:
headers = f"{sagemaker_python_sdk_headers} {jumpstart_specific_suffix}"

if config_name:
headers = f"{headers} {config_specific_suffix}"

return headers


def get_top_ranked_config_name(
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

non-blocking: are we assuming overall ranking in this method? Wouldn't we need to be provided which ranking array to use?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That's an excellent point. How about I make it an optional argument, the ranking name, with the default value being overall?

region: str,
model_id: str,
model_version: str,
sagemaker_session: Optional[Session] = constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
scope: enums.JumpStartScriptScope = enums.JumpStartScriptScope.INFERENCE,
model_type: enums.JumpStartModelType = enums.JumpStartModelType.OPEN_WEIGHTS,
tolerate_deprecated_model: bool = False,
tolerate_vulnerable_model: bool = False,
hub_arn: Optional[str] = None,
) -> Optional[str]:
"""Returns the top ranked config name for the given model ID and region.

Raises:
ValueError: If the script scope is not supported by JumpStart.
"""
model_specs = verify_model_region_and_return_specs(
model_id=model_id,
version=model_version,
scope=scope,
region=region,
hub_arn=hub_arn,
tolerate_vulnerable_model=tolerate_vulnerable_model,
tolerate_deprecated_model=tolerate_deprecated_model,
sagemaker_session=sagemaker_session,
model_type=model_type,
)

if scope == enums.JumpStartScriptScope.INFERENCE:
return (
model_specs.inference_configs.get_top_config_from_ranking().config_name
if model_specs.inference_configs
else None
)
if scope == enums.JumpStartScriptScope.TRAINING:
return (
model_specs.training_configs.get_top_config_from_ranking().config_name
if model_specs.training_configs
else None
)
raise ValueError(f"Unsupported script scope: {scope}.")


def get_default_jumpstart_session_with_user_agent_suffix(
model_id: Optional[str] = None,
model_version: Optional[str] = None,
config_name: Optional[str] = None,
is_hub_content: Optional[bool] = False,
) -> Session:
"""Returns default JumpStart SageMaker Session with model-specific user agent suffix."""
botocore_session = botocore.session.get_session()
botocore_config = botocore.config.Config(
user_agent_extra=get_jumpstart_user_agent_extra_suffix(
model_id, model_version, is_hub_content
model_id, model_version, config_name, is_hub_content
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

consider using named args here.

),
)
botocore_session.set_default_client_config(botocore_config)
Expand Down
Loading
Loading