fix: SDK Defaults Config - Handle config injection for None Sessions (#3878)

Co-authored-by: Ruban Hussain <[email protected]>
Co-authored-by: Kalyani Nikure <[email protected]>

Author: 3 people

src/sagemaker/jumpstart/utils.py

@@ -38,6 +38,7 @@
     JumpStartVersionedModelId,
 )
 from sagemaker.session import Session
+from sagemaker.config import load_sagemaker_config
 from sagemaker.utils import resolve_value_from_config
 from sagemaker.workflow import is_pipeline_variable
 
@@ -455,6 +456,9 @@ def resolve_model_intelligent_default_field(
     over intelligent defaults. For all other fields, intelligent defaults takes precedence
     over the JumpStart default fields.
     """
+    # In case, sagemaker_session is None, get sagemaker_config from load_sagemaker_config()
+    # to resolve value from config for the respective field_name parameter
+    _sagemaker_config = load_sagemaker_config() if (sagemaker_session is None) else None
 
     # We allow customers to define a role which takes precedence
     # over intelligent defaults
@@ -464,6 +468,7 @@ def resolve_model_intelligent_default_field(
             config_path=MODEL_EXECUTION_ROLE_ARN_PATH,
             default_value=default_value or sagemaker_session.get_caller_identity_arn(),
             sagemaker_session=sagemaker_session,
+            sagemaker_config=_sagemaker_config,
         )
 
     # JumpStart Models have certain default field values. We want
@@ -474,6 +479,7 @@ def resolve_model_intelligent_default_field(
             config_path=MODEL_ENABLE_NETWORK_ISOLATION_PATH,
             sagemaker_session=sagemaker_session,
             default_value=default_value,
+            sagemaker_config=_sagemaker_config,
         )
         return resolved_val if resolved_val is not None else field_val
 
@@ -494,6 +500,11 @@ def resolve_estimator_intelligent_default_field(
     over the JumpStart default fields.
     """
 
+    # Workaround for config injection if sagemaker_session is None, since in
+    # that case sagemaker_session will not be initialized until
+    # `_init_sagemaker_session_if_does_not_exist` is called later
+    _sagemaker_config = load_sagemaker_config() if (sagemaker_session is None) else None
+
     # We allow customers to define a role which takes precedence
     # over intelligent defaults
     if field_name == "role":
@@ -502,6 +513,7 @@ def resolve_estimator_intelligent_default_field(
             config_path=TRAINING_JOB_ROLE_ARN_PATH,
             default_value=default_value or sagemaker_session.get_caller_identity_arn(),
             sagemaker_session=sagemaker_session,
+            sagemaker_config=_sagemaker_config,
         )
 
     # JumpStart Estimators have certain default field values. We want
@@ -513,6 +525,7 @@ def resolve_estimator_intelligent_default_field(
             config_path=TRAINING_JOB_ENABLE_NETWORK_ISOLATION_PATH,
             sagemaker_session=sagemaker_session,
             default_value=default_value,
+            sagemaker_config=_sagemaker_config,
         )
         return resolved_val if resolved_val is not None else field_val
 
@@ -523,6 +536,7 @@ def resolve_estimator_intelligent_default_field(
             config_path=TRAINING_JOB_INTER_CONTAINER_ENCRYPTION_PATH,
             sagemaker_session=sagemaker_session,
             default_value=default_value,
+            sagemaker_config=_sagemaker_config,
         )
         return resolved_val if resolved_val is not None else field_val
 

src/sagemaker/model.py

@@ -41,6 +41,7 @@
     MODEL_EXECUTION_ROLE_ARN_PATH,
     MODEL_PRIMARY_CONTAINER_ENVIRONMENT_PATH,
     ENDPOINT_CONFIG_ASYNC_KMS_KEY_ID_PATH,
+    load_sagemaker_config,
 )
 from sagemaker.session import Session
 from sagemaker.model_metrics import ModelMetrics
@@ -313,13 +314,25 @@ def __init__(
         self.name = name
         self._base_name = None
         self.sagemaker_session = sagemaker_session
+
+        # Workaround for config injection if sagemaker_session is None, since in
+        # that case sagemaker_session will not be initialized until
+        # `_init_sagemaker_session_if_does_not_exist` is called later
+        self._sagemaker_config = (
+            load_sagemaker_config() if (self.sagemaker_session is None) else None
+        )
+
         self.role = resolve_value_from_config(
             role,
             MODEL_EXECUTION_ROLE_ARN_PATH,
             sagemaker_session=self.sagemaker_session,
+            sagemaker_config=self._sagemaker_config,
         )
         self.vpc_config = resolve_value_from_config(
-            vpc_config, MODEL_VPC_CONFIG_PATH, sagemaker_session=self.sagemaker_session
+            vpc_config,
+            MODEL_VPC_CONFIG_PATH,
+            sagemaker_session=self.sagemaker_session,
+            sagemaker_config=self._sagemaker_config,
         )
         self.endpoint_name = None
         self._is_compiled_model = False
@@ -330,16 +343,17 @@ def __init__(
         self._enable_network_isolation = resolve_value_from_config(
             enable_network_isolation,
             MODEL_ENABLE_NETWORK_ISOLATION_PATH,
+            default_value=False,
             sagemaker_session=self.sagemaker_session,
+            sagemaker_config=self._sagemaker_config,
         )
         self.env = resolve_value_from_config(
             env,
             MODEL_PRIMARY_CONTAINER_ENVIRONMENT_PATH,
             default_value={},
             sagemaker_session=self.sagemaker_session,
+            sagemaker_config=self._sagemaker_config,
         )
-        if self._enable_network_isolation is None:
-            self._enable_network_isolation = False
         self.model_kms_key = model_kms_key
         self.image_config = image_config
         self.entry_point = entry_point
@@ -542,9 +556,9 @@ def _init_sagemaker_session_if_does_not_exist(self, instance_type=None):
             return
 
         if instance_type in ("local", "local_gpu"):
-            self.sagemaker_session = local.LocalSession()
+            self.sagemaker_session = local.LocalSession(sagemaker_config=self._sagemaker_config)
         else:
-            self.sagemaker_session = session.Session()
+            self.sagemaker_session = session.Session(sagemaker_config=self._sagemaker_config)
 
     def prepare_container_def(
         self,

src/sagemaker/pipeline.py

@@ -22,6 +22,7 @@
     MODEL_VPC_CONFIG_PATH,
     MODEL_ENABLE_NETWORK_ISOLATION_PATH,
     MODEL_EXECUTION_ROLE_ARN_PATH,
+    load_sagemaker_config,
 )
 from sagemaker.drift_check_baselines import DriftCheckBaselines
 from sagemaker.metadata_properties import MetadataProperties
@@ -90,19 +91,33 @@ def __init__(
         self.models = models
         self.predictor_cls = predictor_cls
         self.name = name
-        self.sagemaker_session = sagemaker_session
         self.endpoint_name = None
+        self.sagemaker_session = sagemaker_session
+
+        # In case, sagemaker_session is None, get sagemaker_config from load_sagemaker_config()
+        # to resolve value from config for the respective parameter
+        self._sagemaker_config = (
+            load_sagemaker_config() if (self.sagemaker_session is None) else None
+        )
+
         self.role = resolve_value_from_config(
-            role, MODEL_EXECUTION_ROLE_ARN_PATH, sagemaker_session=self.sagemaker_session
+            role,
+            MODEL_EXECUTION_ROLE_ARN_PATH,
+            sagemaker_session=self.sagemaker_session,
+            sagemaker_config=self._sagemaker_config,
         )
         self.vpc_config = resolve_value_from_config(
-            vpc_config, MODEL_VPC_CONFIG_PATH, sagemaker_session=self.sagemaker_session
+            vpc_config,
+            MODEL_VPC_CONFIG_PATH,
+            sagemaker_session=self.sagemaker_session,
+            sagemaker_config=self._sagemaker_config,
         )
         self.enable_network_isolation = resolve_value_from_config(
             direct_input=enable_network_isolation,
             config_path=MODEL_ENABLE_NETWORK_ISOLATION_PATH,
             default_value=False,
             sagemaker_session=self.sagemaker_session,
+            sagemaker_config=self._sagemaker_config,
         )
 
         if not self.role:
@@ -214,7 +229,7 @@ def deploy(
             is not None. Otherwise, return None.
         """
         if not self.sagemaker_session:
-            self.sagemaker_session = Session()
+            self.sagemaker_session = Session(sagemaker_config=self._sagemaker_config)
 
         containers = self.pipeline_container_def(instance_type)
 
@@ -303,7 +318,7 @@ def _create_sagemaker_pipeline_model(self, instance_type):
                 support or not.
         """
         if not self.sagemaker_session:
-            self.sagemaker_session = Session()
+            self.sagemaker_session = Session(sagemaker_config=self._sagemaker_config)
 
         containers = self.pipeline_container_def(instance_type)
 

src/sagemaker/utils.py

@@ -1046,6 +1046,7 @@ def resolve_value_from_config(
     config_path: str = None,
     default_value=None,
     sagemaker_session=None,
+    sagemaker_config: dict = None,
 ):
     """Decides which value for the caller to use.
 
@@ -1059,19 +1060,30 @@ def resolve_value_from_config(
 
     Args:
         direct_input: The value that the caller of this method starts with. Usually this is an
-        input to the caller's class or method.
+            input to the caller's class or method.
         config_path (str): A string denoting the path used to lookup the value in the
-        sagemaker config.
+            sagemaker config.
         default_value: The value used if not present elsewhere.
         sagemaker_session (sagemaker.session.Session): A SageMaker Session object, used for
-        SageMaker interactions (default: None).
+            SageMaker interactions (default: None).
+        sagemaker_config (dict): The sdk defaults config that is normally accessed through a
+            Session object by doing `session.sagemaker_config`. (default: None) This parameter will
+            be checked for the config value if (and only if) sagemaker_session is None. This
+            parameter exists for the rare cases where the user provided no Session but a default
+            Session cannot be initialized before config injection is needed. In that case,
+            the config dictionary may be loaded and passed here before a default Session object
+            is created.
 
     Returns:
         The value that should be used by the caller
     """
 
     config_value = (
-        get_sagemaker_config_value(sagemaker_session, config_path) if config_path else None
+        get_sagemaker_config_value(
+            sagemaker_session, config_path, sagemaker_config=sagemaker_config
+        )
+        if config_path
+        else None
     )
     _log_sagemaker_config_single_substitution(direct_input, config_value, config_path)
 
@@ -1084,23 +1096,33 @@ def resolve_value_from_config(
     return default_value
 
 
-def get_sagemaker_config_value(sagemaker_session, key):
+def get_sagemaker_config_value(sagemaker_session, key, sagemaker_config: dict = None):
     """Returns the value that corresponds to the provided key from the configuration file.
 
     Args:
         key: Key Path of the config file entry.
         sagemaker_session (sagemaker.session.Session): A SageMaker Session object, used for
-        SageMaker interactions.
+            SageMaker interactions.
+        sagemaker_config (dict): The sdk defaults config that is normally accessed through a
+            Session object by doing `session.sagemaker_config`. (default: None) This parameter will
+            be checked for the config value if (and only if) sagemaker_session is None. This
+            parameter exists for the rare cases where no Session provided but a default Session
+            cannot be initialized before config injection is needed. In that case, the config
+            dictionary may be loaded and passed here before a default Session object is created.
 
     Returns:
         object: The corresponding default value in the configuration file.
     """
-    if not sagemaker_session:
+    if sagemaker_session:
+        config_to_check = sagemaker_session.sagemaker_config
+    else:
+        config_to_check = sagemaker_config
+
+    if not config_to_check:
         return None
 
-    if sagemaker_session.sagemaker_config:
-        validate_sagemaker_config(sagemaker_session.sagemaker_config)
-    config_value = get_config_value(key, sagemaker_session.sagemaker_config)
+    validate_sagemaker_config(config_to_check)
+    config_value = get_config_value(key, config_to_check)
     # Copy the value so any modifications to the output will not modify the source config
     return copy.deepcopy(config_value)
 

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

@@ -57,7 +57,7 @@
 metadata_inference_role = "th1234567iant role"
 
 
-def config_value_impl(sagemaker_session: Session, config_path: str):
+def config_value_impl(sagemaker_session: Session, config_path: str, sagemaker_config: dict):
     if config_path == TRAINING_JOB_ENABLE_NETWORK_ISOLATION_PATH:
         return config_enable_network_isolation
 

tests/unit/sagemaker/jumpstart/model/test_intelligent_defaults.py

@@ -42,7 +42,7 @@
 metadata_enable_network_isolation = random.choice([True, False])
 
 
-def config_value_impl(sagemaker_session: Session, config_path: str):
+def config_value_impl(sagemaker_session: Session, config_path: str, sagemaker_config: dict):
     if config_path == MODEL_EXECUTION_ROLE_ARN_PATH:
         return config_role
 

tests/unit/test_utils.py

@@ -35,6 +35,7 @@
 from sagemaker.utils import (
     retry_with_backoff,
     check_and_get_run_experiment_config,
+    get_sagemaker_config_value,
     resolve_value_from_config,
     resolve_class_attribute_from_config,
     resolve_nested_dict_value_from_config,
@@ -1190,6 +1191,10 @@ def test_resolve_value_from_config():
     sagemaker_session.sagemaker_config.update(
         {"SageMaker": {"EndpointConfig": {"KmsKeyId": "CONFIG_VALUE"}}}
     )
+    sagemaker_config = {
+        "SchemaVersion": "1.0",
+        "SageMaker": {"EndpointConfig": {"KmsKeyId": "CONFIG_VALUE"}},
+    }
 
     # direct_input should be respected
     assert (
@@ -1223,6 +1228,13 @@ def test_resolve_value_from_config():
 
     assert resolve_value_from_config(None, None, None, sagemaker_session) is None
 
+    # Config value from sagemaker_config should be returned
+    # if no direct_input and sagemaker_session is None
+    assert (
+        resolve_value_from_config(None, config_key_path, None, None, sagemaker_config)
+        == "CONFIG_VALUE"
+    )
+
     # Different falsy direct_inputs
     assert resolve_value_from_config("", config_key_path, None, sagemaker_session) == ""
 
@@ -1239,6 +1251,58 @@ def test_resolve_value_from_config():
     mock_info_logger.reset_mock()
 
 
+def test_get_sagemaker_config_value():
+    mock_config_logger = Mock()
+
+    mock_info_logger = Mock()
+    mock_config_logger.info = mock_info_logger
+    # using a shorter name for inside the test
+    sagemaker_session = MagicMock()
+    sagemaker_session.sagemaker_config = {"SchemaVersion": "1.0"}
+    config_key_path = "SageMaker.EndpointConfig.KmsKeyId"
+    sagemaker_session.sagemaker_config.update(
+        {"SageMaker": {"EndpointConfig": {"KmsKeyId": "CONFIG_VALUE"}}}
+    )
+    sagemaker_config = {
+        "SchemaVersion": "1.0",
+        "SageMaker": {"EndpointConfig": {"KmsKeyId": "CONFIG_VALUE"}},
+    }
+
+    # Tests that the function returns the correct value when the key exists in the sagemaker_session configuration.
+    assert (
+        get_sagemaker_config_value(
+            sagemaker_session=sagemaker_session, key=config_key_path, sagemaker_config=None
+        )
+        == "CONFIG_VALUE"
+    )
+
+    # Tests that the function correctly uses the sagemaker_config to get value for the requested
+    # config_key_path when sagemaker_session is None.
+    assert (
+        get_sagemaker_config_value(
+            sagemaker_session=None, key=config_key_path, sagemaker_config=sagemaker_config
+        )
+        == "CONFIG_VALUE"
+    )
+
+    # Tests that the function returns None when the key does not exist in the configuration.
+    invalid_key = "inavlid_key"
+    assert (
+        get_sagemaker_config_value(
+            sagemaker_session=sagemaker_session, key=invalid_key, sagemaker_config=sagemaker_config
+        )
+        is None
+    )
+
+    # Tests that the function returns None when sagemaker_session and sagemaker_config are None.
+    assert (
+        get_sagemaker_config_value(
+            sagemaker_session=None, key=config_key_path, sagemaker_config=None
+        )
+        is None
+    )
+
+
 @patch("jsonschema.validate")
 @pytest.mark.parametrize(
     "existing_value, config_value, default_value",