fix: minor jumpstart logging improvements (#1366)

Author: evakravi

src/sagemaker/jumpstart/exceptions.py

@@ -42,8 +42,8 @@ def get_wildcard_model_version_msg(
 
     return (
         f"Using model '{model_id}' with wildcard version identifier '{wildcard_model_version}'. "
-        f"Please consider pinning to version '{full_model_version}' to "
-        f"ensure stable results. {_MAJOR_VERSION_WARNING_MSG}"
+        f"You can pin to version '{full_model_version}' "
+        f"for more stable results. {_MAJOR_VERSION_WARNING_MSG}"
     )
 
 
@@ -53,9 +53,9 @@ def get_old_model_version_msg(
     """Returns customer-facing message associated with using an old model version."""
 
     return (
-        f"Using model '{model_id}' with old version '{current_model_version}'. "
-        f"Please consider upgrading to version '{latest_model_version}'"
-        f". {_MAJOR_VERSION_WARNING_MSG}"
+        f"Using model '{model_id}' with version '{current_model_version}'. "
+        f"You can upgrade to version '{latest_model_version}' to get the latest model "
+        f"specifications. {_MAJOR_VERSION_WARNING_MSG}"
     )
 
 

src/sagemaker/jumpstart/factory/estimator.py

@@ -586,7 +586,7 @@ def _add_env_to_kwargs(
         )
 
     for key, value in extra_env_vars.items():
-        update_dict_if_key_not_present(
+        kwargs.environment = update_dict_if_key_not_present(
             kwargs.environment,
             key,
             value,

src/sagemaker/jumpstart/utils.py

@@ -755,13 +755,5 @@ def is_valid_model_id(
     if script == enums.JumpStartScriptScope.INFERENCE:
         return model_id in model_id_set
     if script == enums.JumpStartScriptScope.TRAINING:
-        return (
-            model_id in model_id_set
-            and accessors.JumpStartModelsAccessor.get_model_specs(
-                region=region,
-                model_id=model_id,
-                version=model_version,
-                s3_client=s3_client,
-            ).training_supported
-        )
+        return model_id in model_id_set
     raise ValueError(f"Unsupported script: {script}")

tests/integ/sagemaker/jumpstart/estimator/test_jumpstart_estimator.py

@@ -13,6 +13,7 @@
 from __future__ import absolute_import
 import os
 import time
+import mock
 
 import pytest
 from sagemaker.jumpstart.constants import JUMPSTART_DEFAULT_REGION_NAME
@@ -127,7 +128,8 @@ def test_gated_model_training(setup):
     assert response is not None
 
 
-def test_instatiating_estimator_not_too_slow(setup):
+@mock.patch("sagemaker.jumpstart.cache.JUMPSTART_LOGGER.warning")
+def test_instatiating_estimator(mock_warning_logger, setup):
 
     model_id = "xgboost-classification-model"
 
@@ -142,3 +144,5 @@ def test_instatiating_estimator_not_too_slow(setup):
     elapsed_time = time.perf_counter() - start_time
 
     assert elapsed_time <= MAX_INIT_TIME_SECONDS
+
+    mock_warning_logger.assert_called_once()

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

@@ -13,6 +13,7 @@
 from __future__ import absolute_import
 import os
 import time
+from unittest import mock
 
 import pytest
 
@@ -114,7 +115,8 @@ def test_model_package_arn_jumpstart_model(setup):
     assert response is not None
 
 
-def test_instatiating_model_not_too_slow(setup):
+@mock.patch("sagemaker.jumpstart.cache.JUMPSTART_LOGGER.warning")
+def test_instatiating_model(mock_warning_logger, setup):
 
     model_id = "catboost-regression-model"
 
@@ -130,6 +132,8 @@ def test_instatiating_model_not_too_slow(setup):
 
     assert elapsed_time <= MAX_INIT_TIME_SECONDS
 
+    mock_warning_logger.assert_called_once()
+
 
 def test_jumpstart_model_register(setup):
     model_id = "huggingface-txt2img-conflictx-complex-lineart"

tests/unit/sagemaker/jumpstart/constants.py

@@ -2777,7 +2777,8 @@
                 "g5": {
                     "regional_properties": {"image_uri": "$gpu_ecr_uri_1"},
                     "properties": {
-                        "gated_model_key_env_var_value": "meta-training/train-meta-textgeneration-llama-2-7b.tar.gz"
+                        "gated_model_key_env_var_value": "meta-training/train-meta-textgeneration-llama-2-7b.tar.gz",
+                        "environment_variables": {"SELF_DESTRUCT": "true"},
                     },
                 },
                 "local_gpu": {"regional_properties": {"image_uri": "$gpu_ecr_uri_1"}},

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

@@ -495,6 +495,7 @@ def test_gated_model_non_model_package_s3_uri(
             encrypt_inter_container_traffic=True,
             enable_network_isolation=True,
             environment={
+                "SELF_DESTRUCT": "true",
                 "accept_eula": True,
                 "SageMakerGatedModelS3Uri": "s3://top-secret-private-"
                 "models-bucket/meta-training/train-meta-textgeneration-llama-2-7b.tar.gz",

tests/unit/sagemaker/jumpstart/test_cache.py

@@ -587,8 +587,8 @@ def test_jumpstart_cache_makes_correct_s3_calls(
         )
         mocked_warning_log.assert_called_once_with(
             "Using model 'pytorch-ic-imagenet-inception-v3-classification-4' with wildcard "
-            "version identifier '*'. Please consider pinning to version '2.0.0' to "
-            "ensure stable results. Note that models may have different input/output "
+            "version identifier '*'. You can pin to version '2.0.0' for more "
+            "stable results. Note that models may have different input/output "
             "signatures after a major version upgrade."
         )
         mocked_warning_log.reset_mock()

tests/unit/sagemaker/jumpstart/test_exceptions.py

@@ -21,7 +21,7 @@
 def test_get_wildcard_model_version_msg():
     assert (
         "Using model 'mother_of_all_models' with wildcard version identifier '*'. "
-        "Please consider pinning to version '1.2.3' to ensure stable results. "
+        "You can pin to version '1.2.3' for more stable results. "
         "Note that models may have different input/output signatures after a "
         "major version upgrade."
         == get_wildcard_model_version_msg("mother_of_all_models", "*", "1.2.3")
@@ -30,8 +30,8 @@ def test_get_wildcard_model_version_msg():
 
 def test_get_old_model_version_msg():
     assert (
-        "Using model 'mother_of_all_models' with old version '1.0.0'. "
-        "Please consider upgrading to version '1.2.3'. Note that models "
-        "may have different input/output signatures after a major "
+        "Using model 'mother_of_all_models' with version '1.0.0'. "
+        "You can upgrade to version '1.2.3' to get the latest model specifications. "
+        "Note that models may have different input/output signatures after a major "
         "version upgrade." == get_old_model_version_msg("mother_of_all_models", "1.0.0", "1.2.3")
     )

tests/unit/sagemaker/jumpstart/test_utils.py

@@ -967,9 +967,9 @@ def test_jumpstart_old_model_spec(mock_get_manifest):
         )
 
         mocked_info_log.assert_called_once_with(
-            "Using model 'tensorflow-ic-imagenet-inception-v3-classification-4' with old version '1.0.0'. "
-            "Please consider upgrading to version '1.1.0'. Note that models may have different "
-            "input/output signatures after a major version upgrade."
+            "Using model 'tensorflow-ic-imagenet-inception-v3-classification-4' with version '1.0.0'. "
+            "You can upgrade to version '1.1.0' to get the latest model specifications. Note that models "
+            "may have different input/output signatures after a major version upgrade."
         )
 
         mocked_info_log.reset_mock()
@@ -1191,12 +1191,6 @@ def test_is_valid_model_id_true(
             mock_get_manifest.assert_called_once_with(
                 region=JUMPSTART_DEFAULT_REGION_NAME, s3_client=mock_s3_client_value
             )
-            mock_get_model_specs.assert_called_once_with(
-                region=JUMPSTART_DEFAULT_REGION_NAME,
-                model_id="bee",
-                version="*",
-                s3_client=mock_s3_client_value,
-            )
 
     @patch("sagemaker.jumpstart.utils.accessors.JumpStartModelsAccessor._get_manifest")
     @patch("sagemaker.jumpstart.utils.accessors.JumpStartModelsAccessor.get_model_specs")
@@ -1254,13 +1248,7 @@ def test_is_valid_model_id_false(self, mock_get_model_specs: Mock, mock_get_mani
             mock_get_model_specs.reset_mock()
 
             mock_get_model_specs.return_value = Mock(training_supported=False)
-            self.assertFalse(utils.is_valid_model_id("ay", script=JumpStartScriptScope.TRAINING))
+            self.assertTrue(utils.is_valid_model_id("ay", script=JumpStartScriptScope.TRAINING))
             mock_get_manifest.assert_called_once_with(
                 region=JUMPSTART_DEFAULT_REGION_NAME, s3_client=mock_s3_client_value
             )
-            mock_get_model_specs.assert_called_once_with(
-                region=JUMPSTART_DEFAULT_REGION_NAME,
-                model_id="ay",
-                version="*",
-                s3_client=mock_s3_client_value,
-            )