chore: add unit tests

Author: evakravi

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

@@ -16,6 +16,7 @@
 from unittest import mock
 import unittest
 from inspect import signature
+from mock import Mock
 
 import pytest
 from sagemaker.async_inference.async_inference_config import AsyncInferenceConfig
@@ -30,12 +31,16 @@
 from sagemaker.jumpstart.artifacts.metric_definitions import (
     _retrieve_default_training_metric_definitions,
 )
-from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
+from sagemaker.jumpstart.constants import (
+    DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
+    JUMPSTART_DEFAULT_REGION_NAME,
+)
 from sagemaker.jumpstart.enums import JumpStartScriptScope, JumpStartTag
 
 from sagemaker.jumpstart.estimator import JumpStartEstimator
 
 from sagemaker.jumpstart.utils import get_jumpstart_content_bucket
+from sagemaker.session import Session
 from sagemaker.session_settings import SessionSettings
 from tests.integ.sagemaker.jumpstart.utils import get_training_dataset_for_model_and_version
 from sagemaker.model import Model
@@ -44,6 +49,7 @@
     get_special_model_spec,
     overwrite_dictionary,
 )
+import boto3
 
 
 execution_role = "fake role! do not use!"
@@ -1773,6 +1779,56 @@ def test_model_artifact_variant_estimator(
             ],
         )
 
+    @mock.patch("sagemaker.jumpstart.estimator.get_default_predictor")
+    @mock.patch("sagemaker.jumpstart.estimator.Estimator.__init__")
+    @mock.patch("sagemaker.jumpstart.estimator.Estimator.fit")
+    @mock.patch("sagemaker.jumpstart.estimator.Estimator.deploy")
+    @mock.patch("sagemaker.jumpstart.estimator.is_valid_model_id")
+    @mock.patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
+    @mock.patch("sagemaker.jumpstart.factory.model.JUMPSTART_DEFAULT_REGION_NAME", region)
+    def test_jumpstart_estimator_session(
+        self,
+        mock_get_model_specs: mock.Mock,
+        mock_is_valid_model_id: mock.Mock,
+        mock_deploy,
+        mock_fit,
+        mock_init,
+        get_default_predictor,
+    ):
+
+        mock_is_valid_model_id.return_value = True
+
+        model_id, _ = "js-trainable-model", "*"
+
+        mock_get_model_specs.side_effect = get_special_model_spec
+
+        region = "eu-west-1"  # some non-default region
+
+        if region == JUMPSTART_DEFAULT_REGION_NAME:
+            region = "us-west-2"
+
+        session = Session(boto_session=boto3.session.Session(region_name=region))
+
+        assert session.boto_region_name != JUMPSTART_DEFAULT_REGION_NAME
+
+        session.get_caller_identity_arn = Mock(return_value="blah")
+
+        estimator = JumpStartEstimator(model_id=model_id, sagemaker_session=session)
+        estimator.fit()
+
+        estimator.deploy()
+
+        assert len(mock_get_model_specs.call_args_list) > 1
+
+        regions = {call[1]["region"] for call in mock_get_model_specs.call_args_list}
+
+        assert len(regions) == 1
+        assert list(regions)[0] == region
+
+        s3_clients = {call[1]["s3_client"] for call in mock_get_model_specs.call_args_list}
+        assert len(s3_clients) == 1
+        assert list(s3_clients)[0] == session.s3_client
+
 
 def test_jumpstart_estimator_requires_model_id():
     with pytest.raises(ValueError):

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

@@ -15,18 +15,22 @@
 from typing import Optional, Set
 from unittest import mock
 import unittest
-from mock import MagicMock
+from mock import MagicMock, Mock
 import pytest
 from sagemaker.async_inference.async_inference_config import AsyncInferenceConfig
 from sagemaker.jumpstart.artifacts.environment_variables import (
     _retrieve_default_environment_variables,
 )
-from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
+from sagemaker.jumpstart.constants import (
+    DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
+    JUMPSTART_DEFAULT_REGION_NAME,
+)
 from sagemaker.jumpstart.enums import JumpStartScriptScope, JumpStartTag
 
 from sagemaker.jumpstart.model import JumpStartModel
 from sagemaker.model import Model
 from sagemaker.predictor import Predictor
+from sagemaker.session import Session
 from sagemaker.session_settings import SessionSettings
 from sagemaker.enums import EndpointType
 from sagemaker.compute_resource_requirements.resource_requirements import ResourceRequirements
@@ -36,6 +40,7 @@
     overwrite_dictionary,
     get_special_model_spec_for_inference_component_based_endpoint,
 )
+import boto3
 
 execution_role = "fake role! do not use!"
 region = "us-west-2"
@@ -1252,6 +1257,52 @@ def test_model_registry_accept_and_response_types(
             response_types=["application/json;verbose", "application/json"],
         )
 
+    @mock.patch("sagemaker.jumpstart.model.get_default_predictor")
+    @mock.patch("sagemaker.jumpstart.model.Model.__init__")
+    @mock.patch("sagemaker.jumpstart.model.Model.deploy")
+    @mock.patch("sagemaker.jumpstart.model.is_valid_model_id")
+    @mock.patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
+    @mock.patch("sagemaker.jumpstart.factory.model.JUMPSTART_DEFAULT_REGION_NAME", region)
+    def test_jumpstart_model_session(
+        self,
+        mock_get_model_specs: mock.Mock,
+        mock_is_valid_model_id: mock.Mock,
+        mock_deploy,
+        mock_init,
+        get_default_predictor,
+    ):
+
+        mock_is_valid_model_id.return_value = True
+
+        model_id, _ = "model_data_s3_prefix_model", "*"
+
+        mock_get_model_specs.side_effect = get_special_model_spec
+
+        region = "eu-west-1"  # some non-default region
+
+        if region == JUMPSTART_DEFAULT_REGION_NAME:
+            region = "us-west-2"
+
+        session = Session(boto_session=boto3.session.Session(region_name=region))
+
+        assert session.boto_region_name != JUMPSTART_DEFAULT_REGION_NAME
+
+        session.get_caller_identity_arn = Mock(return_value="blah")
+
+        model = JumpStartModel(model_id=model_id, sagemaker_session=session)
+        model.deploy()
+
+        assert len(mock_get_model_specs.call_args_list) > 1
+
+        regions = {call[1]["region"] for call in mock_get_model_specs.call_args_list}
+
+        assert len(regions) == 1
+        assert list(regions)[0] == region
+
+        s3_clients = {call[1]["s3_client"] for call in mock_get_model_specs.call_args_list}
+        assert len(s3_clients) == 1
+        assert list(s3_clients)[0] == session.s3_client
+
 
 def test_jumpstart_model_requires_model_id():
     with pytest.raises(ValueError):

tests/unit/sagemaker/jumpstart/test_cache.py

@@ -18,6 +18,7 @@
 from unittest.mock import Mock, call, mock_open
 from botocore.stub import Stubber
 import botocore
+import boto3
 
 from mock.mock import MagicMock
 import pytest
@@ -27,6 +28,7 @@
 from sagemaker.jumpstart.constants import (
     ENV_VARIABLE_JUMPSTART_MANIFEST_LOCAL_ROOT_DIR_OVERRIDE,
     ENV_VARIABLE_JUMPSTART_SPECS_LOCAL_ROOT_DIR_OVERRIDE,
+    JUMPSTART_DEFAULT_REGION_NAME,
 )
 from sagemaker.jumpstart.types import (
     JumpStartModelHeader,
@@ -854,3 +856,37 @@ def test_jumpstart_local_metadata_override_specs_not_exist_both_directories(
             ),
         ]
     )
+
+
+@pytest.mark.parametrize(
+    "s3_bucket_name, s3_client, region",
+    [
+        (
+            "jumpstart-cache-prod",
+            boto3.client("s3", region_name="blah-blah"),
+            JUMPSTART_DEFAULT_REGION_NAME,
+        ),
+        (
+            "jumpstart-cache-prod-us-west-2",
+            boto3.client("s3", region_name="us-west-2"),
+            "us-west-2",
+        ),
+        ("jumpstart-cache-prod", boto3.client("s3", region_name="us-east-2"), "us-east-2"),
+    ],
+)
+def test_get_region_fallback_success(s3_bucket_name, s3_client, region):
+    cache = JumpStartModelsCache()
+    assert region == cache._get_region_fallback(s3_bucket_name, s3_client)
+
+
+@pytest.mark.parametrize(
+    "s3_bucket_name, s3_client",
+    [
+        ("jumpstart-cache-prod-us-west-2", boto3.client("s3", region_name="us-east-2")),
+        ("jumpstart-cache-prod-us-west-2-us-east-2", boto3.client("s3", region_name="us-east-2")),
+    ],
+)
+def test_get_region_fallback_failure(s3_bucket_name, s3_client):
+    cache = JumpStartModelsCache()
+    with pytest.raises(ValueError):
+        cache._get_region_fallback(s3_bucket_name, s3_client)