chore: telemetry for deployment configs

Author: evakravi

src/sagemaker/jumpstart/exceptions.py

@@ -150,6 +150,7 @@ def __init__(
                 model. (Default: None).
 
         """
+        version = version or "*"
         if message:
             self.message = message
         else:
@@ -198,6 +199,7 @@ def __init__(
         version: Optional[str] = None,
         message: Optional[str] = None,
     ):
+        version = version or "*"
         if message:
             self.message = message
         else:

src/sagemaker/jumpstart/factory/estimator.py

@@ -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,
@@ -204,7 +205,7 @@ 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_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)
@@ -438,12 +439,15 @@ def _add_region_to_kwargs(kwargs: JumpStartKwargs) -> JumpStartKwargs:
     return kwargs
 
 
-def _add_sagemaker_session_to_kwargs(kwargs: JumpStartKwargs) -> JumpStartKwargs:
+def _add_sagemaker_session_with_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
@@ -903,20 +907,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(
+        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,
         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

src/sagemaker/jumpstart/factory/model.py

@@ -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,
@@ -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,
@@ -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
         )
     )
 
@@ -662,33 +664,47 @@ def _add_config_name_to_init_kwargs(kwargs: JumpStartModelInitKwargs) -> JumpSta
         ValueError: If the instance_type is not supported with the current config.
     """
 
+    # 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,
+        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,
+    )
+
+    if kwargs.config_name is None:
+        return kwargs
+
     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,
+        sagemaker_session=temp_session,
         model_type=kwargs.model_type,
         config_name=kwargs.config_name,
     )
-    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}."
-            )
 
-        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)
+    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)
     return kwargs
 
 
@@ -740,27 +756,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
 
     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
 
@@ -839,16 +869,16 @@ 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_sagemaker_session_with_custom_user_agent_to_kwargs(kwargs=deploy_kwargs)
 
     deploy_kwargs = _add_model_version_to_kwargs(kwargs=deploy_kwargs)
 
     deploy_kwargs = _add_endpoint_name_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_instance_type_to_kwargs(kwargs=deploy_kwargs)
 
     deploy_kwargs.initial_instance_count = initial_instance_count or 1
@@ -1030,11 +1060,14 @@ def get_init_kwargs(
     )
 
     model_init_kwargs = _add_vulnerable_and_deprecated_status_to_kwargs(kwargs=model_init_kwargs)
+    model_init_kwargs = _add_config_name_to_init_kwargs(kwargs=model_init_kwargs)
+    model_init_kwargs = _add_model_version_to_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(
@@ -1062,8 +1095,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

src/sagemaker/jumpstart/utils.py

@@ -1109,11 +1109,16 @@ 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}"
+    print(config_name)
     hub_specific_suffix = f"md/js_is_hub_content#{is_hub_content}"
 
     if os.getenv(constants.ENV_VARIABLE_DISABLE_JUMPSTART_TELEMETRY, None):
@@ -1128,19 +1133,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(
+    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
         ),
     )
     botocore_session.set_default_client_config(botocore_config)

tests/integ/sagemaker/jumpstart/model/test_jumpstart_model.py

@@ -396,3 +396,20 @@ def test_jumpstart_model_with_deployment_configs(setup):
     response = predictor.predict(payload, custom_attributes="accept_eula=true")
 
     assert response is not None
+
+
+def test_jumpstart_session_with_config_name():
+    model = JumpStartModel(model_id="meta-textgeneration-llama-2-7b", model_version="*")
+    assert model.config_name != None
+    session = model.sagemaker_session
+
+    with mock.patch("botocore.client.BaseClient._make_request") as mock_make_request:
+        try:
+            session.sagemaker_client.list_endpoints()
+        except Exception:
+            pass
+
+    assert (
+        "md/js_model_id#meta-textgeneration-llama-2-7b md/js_model_ver#* md/js_config#tgi"
+        in mock_make_request.call_args[0][1]["headers"]["User-Agent"]
+    )

tests/unit/sagemaker/jumpstart/estimator/test_estimator.py

@@ -1637,7 +1637,7 @@ def test_training_passes_role_to_deploy(
     @mock.patch("sagemaker.jumpstart.estimator.validate_model_id_and_get_type")
     @mock.patch(
         "sagemaker.jumpstart.factory.model.get_default_jumpstart_session_with_user_agent_suffix",
-        sagemaker_session,
+        lambda *largs, **kwargs: sagemaker_session,
     )
     @mock.patch(
         "sagemaker.jumpstart.factory.estimator.get_default_jumpstart_session_with_user_agent_suffix",