fix: failing config tests

Author: evakravi

src/sagemaker/jumpstart/factory/estimator.py

@@ -207,7 +207,9 @@ def get_init_kwargs(
         enable_session_tag_chaining=enable_session_tag_chaining,
     )
 
-    estimator_init_kwargs = _set_temp_sagemaker_session_if_not_set(kwargs=estimator_init_kwargs)
+    estimator_init_kwargs, orig_session = _set_temp_sagemaker_session_if_not_set(
+        kwargs=estimator_init_kwargs
+    )
     estimator_init_kwargs.specs = verify_model_region_and_return_specs(
         **get_model_info_default_kwargs(
             estimator_init_kwargs, include_model_version=False, include_tolerate_flags=False
@@ -223,7 +225,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_with_custom_user_agent_to_kwargs(
-        estimator_init_kwargs
+        estimator_init_kwargs, orig_session
     )
     estimator_init_kwargs = _add_region_to_kwargs(estimator_init_kwargs)
     estimator_init_kwargs = _add_instance_type_and_count_to_kwargs(estimator_init_kwargs)
@@ -280,7 +282,7 @@ def get_fit_kwargs(
         config_name=config_name,
     )
 
-    estimator_fit_kwargs = _set_temp_sagemaker_session_if_not_set(kwargs=estimator_fit_kwargs)
+    estimator_fit_kwargs, _ = _set_temp_sagemaker_session_if_not_set(kwargs=estimator_fit_kwargs)
     estimator_fit_kwargs.specs = verify_model_region_and_return_specs(
         **get_model_info_default_kwargs(
             estimator_fit_kwargs, include_model_version=False, include_tolerate_flags=False
@@ -472,17 +474,14 @@ def _add_region_to_kwargs(kwargs: JumpStartKwargs) -> JumpStartKwargs:
 
 
 def _add_sagemaker_session_with_custom_user_agent_to_kwargs(
-    kwargs: JumpStartKwargs,
+    kwargs: JumpStartKwargs, orig_session: Optional[Session]
 ) -> 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(
-            model_id=kwargs.model_id,
-            model_version=kwargs.model_version,
-            config_name=None,
-            is_hub_content=kwargs.hub_arn is not None,
-        )
+    kwargs.sagemaker_session = orig_session or get_default_jumpstart_session_with_user_agent_suffix(
+        model_id=kwargs.model_id,
+        model_version=kwargs.model_version,
+        config_name=None,
+        is_hub_content=kwargs.hub_arn is not None,
     )
     return kwargs
 

src/sagemaker/jumpstart/factory/model.py

@@ -161,18 +161,16 @@ def _add_region_to_kwargs(kwargs: JumpStartModelInitKwargs) -> JumpStartModelIni
 
 
 def _add_sagemaker_session_with_custom_user_agent_to_kwargs(
-    kwargs: Union[JumpStartModelInitKwargs, JumpStartModelDeployKwargs]
+    kwargs: Union[JumpStartModelInitKwargs, JumpStartModelDeployKwargs],
+    orig_session: Optional[Session],
 ) -> 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(
-            model_id=kwargs.model_id,
-            model_version=kwargs.model_version,
-            config_name=kwargs.config_name,
-            is_hub_content=kwargs.hub_arn is not None,
-        )
+    kwargs.sagemaker_session = orig_session or get_default_jumpstart_session_with_user_agent_suffix(
+        model_id=kwargs.model_id,
+        model_version=kwargs.model_version,
+        config_name=kwargs.config_name,
+        is_hub_content=kwargs.hub_arn is not None,
     )
 
     return kwargs
@@ -686,7 +684,7 @@ def get_deploy_kwargs(
         config_name=config_name,
         routing_config=routing_config,
     )
-    deploy_kwargs = _set_temp_sagemaker_session_if_not_set(kwargs=deploy_kwargs)
+    deploy_kwargs, orig_session = _set_temp_sagemaker_session_if_not_set(kwargs=deploy_kwargs)
     deploy_kwargs.specs = verify_model_region_and_return_specs(
         **get_model_info_default_kwargs(
             deploy_kwargs, include_model_version=False, include_tolerate_flags=False
@@ -705,7 +703,9 @@ def get_deploy_kwargs(
 
     deploy_kwargs = _add_model_version_to_kwargs(kwargs=deploy_kwargs)
 
-    deploy_kwargs = _add_sagemaker_session_with_custom_user_agent_to_kwargs(kwargs=deploy_kwargs)
+    deploy_kwargs = _add_sagemaker_session_with_custom_user_agent_to_kwargs(
+        kwargs=deploy_kwargs, orig_session=orig_session
+    )
 
     deploy_kwargs = _add_endpoint_name_to_kwargs(kwargs=deploy_kwargs)
 
@@ -890,7 +890,9 @@ def get_init_kwargs(
         config_name=config_name,
         additional_model_data_sources=additional_model_data_sources,
     )
-    model_init_kwargs = _set_temp_sagemaker_session_if_not_set(kwargs=model_init_kwargs)
+    model_init_kwargs, orig_session = _set_temp_sagemaker_session_if_not_set(
+        kwargs=model_init_kwargs
+    )
     model_init_kwargs.specs = verify_model_region_and_return_specs(
         **get_model_info_default_kwargs(
             model_init_kwargs, include_model_version=False, include_tolerate_flags=False
@@ -908,7 +910,7 @@ def get_init_kwargs(
     model_init_kwargs = _add_config_name_to_init_kwargs(kwargs=model_init_kwargs)
 
     model_init_kwargs = _add_sagemaker_session_with_custom_user_agent_to_kwargs(
-        kwargs=model_init_kwargs
+        kwargs=model_init_kwargs, orig_session=orig_session
     )
     model_init_kwargs = _add_region_to_kwargs(kwargs=model_init_kwargs)
 

src/sagemaker/jumpstart/factory/utils.py

@@ -13,7 +13,7 @@
 """This module stores JumpStart factory utilities."""
 
 from __future__ import absolute_import
-from typing import Union
+from typing import Tuple, Union
 
 from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
 from sagemaker.jumpstart.types import (
@@ -23,6 +23,7 @@
     JumpStartModelDeployKwargs,
     JumpStartModelInitKwargs,
 )
+from sagemaker.session import Session
 
 KwargsType = Union[
     JumpStartModelDeployKwargs,
@@ -65,13 +66,14 @@ def get_model_info_default_kwargs(
     return kwargs_dict
 
 
-def _set_temp_sagemaker_session_if_not_set(kwargs: KwargsType) -> KwargsType:
-    """Sets a temporary sagemaker session if one is not set.
+def _set_temp_sagemaker_session_if_not_set(kwargs: KwargsType) -> Tuple[KwargsType, Session]:
+    """Sets a temporary sagemaker session if one is not set, and returns original session.
 
     We need to create a default JS session (without custom user agent)
     in order to retrieve config name info.
     """
 
+    orig_session = kwargs.sagemaker_session
     if kwargs.sagemaker_session is None:
         kwargs.sagemaker_session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION
-    return kwargs
+    return kwargs, orig_session