chore: add unit tests

Author: evakravi

src/sagemaker/jumpstart/factory/estimator.py

@@ -56,7 +56,6 @@
     JUMPSTART_LOGGER,
     TRAINING_ENTRY_POINT_SCRIPT_NAME,
     SAGEMAKER_GATED_MODEL_S3_URI_TRAINING_ENV_VAR_KEY,
-    JUMPSTART_MODEL_HUB_NAME,
 )
 from sagemaker.jumpstart.enums import JumpStartScriptScope, JumpStartModelType
 from sagemaker.jumpstart.factory import model

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

@@ -0,0 +1,108 @@
+import pytest
+from unittest.mock import patch, Mock
+from sagemaker.jumpstart.factory.estimator import (
+    _add_model_uri_to_kwargs,
+    get_model_info_default_kwargs,
+)
+from sagemaker.jumpstart.types import JumpStartEstimatorInitKwargs
+from sagemaker.jumpstart.enums import JumpStartScriptScope
+
+
+class TestAddModelUriToKwargs:
+    @pytest.fixture
+    def mock_kwargs(self):
+        return JumpStartEstimatorInitKwargs(
+            model_id="test-model",
+            model_version="1.0.0",
+            instance_type="ml.m5.large",
+            model_uri=None,
+        )
+
+    @patch(
+        "sagemaker.jumpstart.factory.estimator._model_supports_training_model_uri",
+        return_value=True,
+    )
+    @patch("sagemaker.jumpstart.factory.estimator.model_uris.retrieve")
+    def test_add_model_uri_to_kwargs_default_uri(
+        self, mock_retrieve, mock_supports_training, mock_kwargs
+    ):
+        """Test adding default model URI when none is provided."""
+        default_uri = "s3://jumpstart-models/training/test-model/1.0.0"
+        mock_retrieve.return_value = default_uri
+
+        result = _add_model_uri_to_kwargs(mock_kwargs)
+
+        mock_supports_training.assert_called_once()
+        mock_retrieve.assert_called_once_with(
+            model_scope=JumpStartScriptScope.TRAINING,
+            instance_type=mock_kwargs.instance_type,
+            **get_model_info_default_kwargs(mock_kwargs),
+        )
+        assert result.model_uri == default_uri
+
+    @patch(
+        "sagemaker.jumpstart.factory.estimator._model_supports_training_model_uri",
+        return_value=True,
+    )
+    @patch(
+        "sagemaker.jumpstart.factory.estimator._model_supports_incremental_training",
+        return_value=True,
+    )
+    @patch("sagemaker.jumpstart.factory.estimator.model_uris.retrieve")
+    def test_add_model_uri_to_kwargs_custom_uri_with_incremental(
+        self, mock_retrieve, mock_supports_incremental, mock_supports_training, mock_kwargs
+    ):
+        """Test using custom model URI with incremental training support."""
+        default_uri = "s3://jumpstart-models/training/test-model/1.0.0"
+        custom_uri = "s3://custom-bucket/my-model"
+        mock_retrieve.return_value = default_uri
+        mock_kwargs.model_uri = custom_uri
+
+        result = _add_model_uri_to_kwargs(mock_kwargs)
+
+        mock_supports_training.assert_called_once()
+        mock_supports_incremental.assert_called_once()
+        assert result.model_uri == custom_uri
+
+    @patch(
+        "sagemaker.jumpstart.factory.estimator._model_supports_training_model_uri",
+        return_value=True,
+    )
+    @patch(
+        "sagemaker.jumpstart.factory.estimator._model_supports_incremental_training",
+        return_value=False,
+    )
+    @patch("sagemaker.jumpstart.factory.estimator.model_uris.retrieve")
+    @patch("sagemaker.jumpstart.factory.estimator.JUMPSTART_LOGGER.warning")
+    def test_add_model_uri_to_kwargs_custom_uri_without_incremental(
+        self,
+        mock_warning,
+        mock_retrieve,
+        mock_supports_incremental,
+        mock_supports_training,
+        mock_kwargs,
+    ):
+        """Test using custom model URI without incremental training support logs warning."""
+        default_uri = "s3://jumpstart-models/training/test-model/1.0.0"
+        custom_uri = "s3://custom-bucket/my-model"
+        mock_retrieve.return_value = default_uri
+        mock_kwargs.model_uri = custom_uri
+
+        result = _add_model_uri_to_kwargs(mock_kwargs)
+
+        mock_supports_training.assert_called_once()
+        mock_supports_incremental.assert_called_once()
+        mock_warning.assert_called_once()
+        assert "does not support incremental training" in mock_warning.call_args[0][0]
+        assert result.model_uri == custom_uri
+
+    @patch(
+        "sagemaker.jumpstart.factory.estimator._model_supports_training_model_uri",
+        return_value=False,
+    )
+    def test_add_model_uri_to_kwargs_no_training_support(self, mock_supports_training, mock_kwargs):
+        """Test when model doesn't support training model URI."""
+        result = _add_model_uri_to_kwargs(mock_kwargs)
+
+        mock_supports_training.assert_called_once()
+        assert result.model_uri is None

tests/unit/sagemaker/jumpstart/test_cache.py

@@ -1288,3 +1288,78 @@ def test_jumpstart_cache_handles_versioning_correctly_non_sem_ver(retrieval_func
     assert_key = JumpStartVersionedModelId("test-model", "abc")
 
     assert result == assert_key
+
+
+@patch("sagemaker.jumpstart.utils.get_region_fallback", lambda *args, **kwargs: "dummy-region")
+@patch(
+    "sagemaker.jumpstart.utils.get_jumpstart_content_bucket", lambda *args, **kwargs: "dummy-bucket"
+)
+def test_get_json_file_from_s3():
+    """Test _get_json_file retrieves from S3 in normal mode."""
+    cache = JumpStartModelsCache()
+    test_key = "test/file/path.json"
+    test_json_data = {"key": "value"}
+    test_etag = "test-etag-123"
+
+    with patch.object(
+        JumpStartModelsCache,
+        "_get_json_file_and_etag_from_s3",
+        return_value=(test_json_data, test_etag),
+    ) as mock_s3_get:
+        result, etag = cache._get_json_file(test_key, JumpStartS3FileType.OPEN_WEIGHT_MANIFEST)
+
+        mock_s3_get.assert_called_once_with(test_key)
+        assert result == test_json_data
+        assert etag == test_etag
+
+
+@patch("sagemaker.jumpstart.utils.get_region_fallback", lambda *args, **kwargs: "dummy-region")
+@patch(
+    "sagemaker.jumpstart.utils.get_jumpstart_content_bucket", lambda *args, **kwargs: "dummy-bucket"
+)
+def test_get_json_file_from_local_supported_type():
+    """Test _get_json_file retrieves from local override for supported file types."""
+    cache = JumpStartModelsCache()
+    test_key = "test/file/path.json"
+    test_json_data = {"key": "value"}
+
+    with (
+        patch.object(JumpStartModelsCache, "_is_local_metadata_mode", return_value=True),
+        patch.object(
+            JumpStartModelsCache, "_get_json_file_from_local_override", return_value=test_json_data
+        ) as mock_local_get,
+    ):
+        result, etag = cache._get_json_file(test_key, JumpStartS3FileType.OPEN_WEIGHT_MANIFEST)
+
+        mock_local_get.assert_called_once_with(test_key, JumpStartS3FileType.OPEN_WEIGHT_MANIFEST)
+        assert result == test_json_data
+        assert etag is None
+
+
+@patch("sagemaker.jumpstart.utils.get_region_fallback", lambda *args, **kwargs: "dummy-region")
+@patch(
+    "sagemaker.jumpstart.utils.get_jumpstart_content_bucket", lambda *args, **kwargs: "dummy-bucket"
+)
+def test_get_json_file_local_mode_unsupported_type():
+    """Test _get_json_file falls back to S3 for unsupported file types in local mode."""
+    cache = JumpStartModelsCache()
+    test_key = "test/file/path.json"
+    test_json_data = {"key": "value"}
+    test_etag = "test-etag-123"
+
+    with (
+        patch.object(JumpStartModelsCache, "_is_local_metadata_mode", return_value=True),
+        patch.object(
+            JumpStartModelsCache,
+            "_get_json_file_and_etag_from_s3",
+            return_value=(test_json_data, test_etag),
+        ) as mock_s3_get,
+        patch("sagemaker.jumpstart.cache.JUMPSTART_LOGGER.warning") as mock_warning,
+    ):
+        result, etag = cache._get_json_file(test_key, JumpStartS3FileType.PROPRIETARY_MANIFEST)
+
+        mock_s3_get.assert_called_once_with(test_key)
+        mock_warning.assert_called_once()
+        assert "not supported for local override" in mock_warning.call_args[0][0]
+        assert result == test_json_data
+        assert etag == test_etag

tests/unit/sagemaker/jumpstart/test_types.py

@@ -39,6 +39,8 @@
     INIT_KWARGS,
 )
 
+from unittest.mock import Mock
+
 INSTANCE_TYPE_VARIANT = JumpStartInstanceTypeVariants(
     {
         "regional_aliases": {
@@ -329,14 +331,66 @@ def test_jumpstart_model_header():
     assert header1 == header3
 
 
-def test_use_training_model_artifact():
-    specs1 = JumpStartModelSpecs(BASE_SPEC)
-    assert specs1.use_training_model_artifact()
-    specs1.gated_bucket = True
-    assert not specs1.use_training_model_artifact()
-    specs1.gated_bucket = False
-    specs1.training_model_package_artifact_uris = {"region1": "blah", "region2": "blah2"}
-    assert not specs1.use_training_model_artifact()
+class TestUseTrainingModelArtifact:
+    @pytest.fixture
+    def mock_specs(self):
+        specs = Mock(spec=JumpStartModelSpecs)
+        specs.training_instance_type_variants = Mock()
+        specs.supported_training_instance_types = ["ml.p3.2xlarge", "ml.g4dn.xlarge"]
+        specs.training_model_package_artifact_uris = {}
+        specs.training_artifact_key = None
+        return specs
+
+    def test_use_training_model_artifact_with_env_var(self, mock_specs):
+        """Test when instance type variants have env var values."""
+        mock_specs.training_instance_type_variants.get_instance_specific_gated_model_key_env_var_value.side_effect = [
+            "some-value",
+            None,
+        ]
+
+        result = JumpStartModelSpecs.use_training_model_artifact(mock_specs)
+
+        assert result is False
+        mock_specs.training_instance_type_variants.get_instance_specific_gated_model_key_env_var_value.assert_any_call(
+            "ml.p3.2xlarge"
+        )
+
+    def test_use_training_model_artifact_with_package_uris(self, mock_specs):
+        """Test when model has training package artifact URIs."""
+        mock_specs.training_instance_type_variants.get_instance_specific_gated_model_key_env_var_value.return_value = (
+            None
+        )
+        mock_specs.training_model_package_artifact_uris = {
+            "ml.p3.2xlarge": "arn:aws:sagemaker:ap-southeast-1:192199979996:model-package/llama2-13b-e155a2e0347b323fb882f1875851c5d3"
+        }
+
+        result = JumpStartModelSpecs.use_training_model_artifact(mock_specs)
+
+        assert result is False
+
+    def test_use_training_model_artifact_with_artifact_key(self, mock_specs):
+        """Test when model has training artifact key."""
+        mock_specs.training_instance_type_variants.get_instance_specific_gated_model_key_env_var_value.return_value = (
+            None
+        )
+        mock_specs.training_model_package_artifact_uris = {}
+        mock_specs.training_artifact_key = "some-key"
+
+        result = JumpStartModelSpecs.use_training_model_artifact(mock_specs)
+
+        assert result is True
+
+    def test_use_training_model_artifact_without_artifact_key(self, mock_specs):
+        """Test when model has no training artifact key."""
+        mock_specs.training_instance_type_variants.get_instance_specific_gated_model_key_env_var_value.return_value = (
+            None
+        )
+        mock_specs.training_model_package_artifact_uris = {}
+        mock_specs.training_artifact_key = None
+
+        result = JumpStartModelSpecs.use_training_model_artifact(mock_specs)
+
+        assert result is False
 
 
 def test_jumpstart_model_specs():