fix: jumpstart async inference config predictor support (#3970)

Author: evakravi

src/sagemaker/jumpstart/estimator.py

@@ -973,7 +973,7 @@ def deploy(
         )
 
         # If no predictor class was passed, add defaults to predictor
-        if self.orig_predictor_cls is None:
+        if self.orig_predictor_cls is None and async_inference_config is None:
             return get_default_predictor(
                 predictor=predictor,
                 model_id=self.model_id,

src/sagemaker/jumpstart/model.py

@@ -432,7 +432,7 @@ def deploy(
         predictor = super(JumpStartModel, self).deploy(**deploy_kwargs.to_kwargs_dict())
 
         # If no predictor class was passed, add defaults to predictor
-        if self.orig_predictor_cls is None:
+        if self.orig_predictor_cls is None and async_inference_config is None:
             return get_default_predictor(
                 predictor=predictor,
                 model_id=self.model_id,

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

@@ -18,6 +18,7 @@
 from inspect import signature
 
 import pytest
+from sagemaker.async_inference.async_inference_config import AsyncInferenceConfig
 
 from sagemaker.debugger.profiler_config import ProfilerConfig
 from sagemaker.estimator import Estimator
@@ -640,6 +641,56 @@ def test_no_predictor_returns_default_predictor(
         self.assertEqual(type(predictor), Predictor)
         self.assertEqual(predictor, default_predictor_with_presets)
 
+    @mock.patch("sagemaker.jumpstart.estimator.get_default_predictor")
+    @mock.patch("sagemaker.jumpstart.estimator.is_valid_model_id")
+    @mock.patch("sagemaker.jumpstart.factory.model.Session")
+    @mock.patch("sagemaker.jumpstart.factory.estimator.Session")
+    @mock.patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
+    @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.factory.estimator.JUMPSTART_DEFAULT_REGION_NAME", region)
+    @mock.patch("sagemaker.jumpstart.factory.model.JUMPSTART_DEFAULT_REGION_NAME", region)
+    def test_no_predictor_yes_async_inference_config(
+        self,
+        mock_estimator_deploy: mock.Mock,
+        mock_estimator_fit: mock.Mock,
+        mock_estimator_init: mock.Mock,
+        mock_get_model_specs: mock.Mock,
+        mock_session_estimator: mock.Mock,
+        mock_session_model: mock.Mock,
+        mock_is_valid_model_id: mock.Mock,
+        mock_get_default_predictor: mock.Mock,
+    ):
+        mock_estimator_deploy.return_value = default_predictor
+
+        mock_get_default_predictor.return_value = default_predictor_with_presets
+
+        mock_is_valid_model_id.return_value = True
+
+        model_id, _ = "js-trainable-model-prepacked", "*"
+
+        mock_get_model_specs.side_effect = get_special_model_spec
+
+        mock_session_estimator.return_value = sagemaker_session
+        mock_session_model.return_value = sagemaker_session
+
+        estimator = JumpStartEstimator(
+            model_id=model_id,
+        )
+
+        channels = {
+            "training": f"s3://{get_jumpstart_content_bucket(region)}/"
+            f"some-training-dataset-doesn't-matter",
+        }
+
+        estimator.fit(channels)
+
+        predictor = estimator.deploy(async_inference_config=AsyncInferenceConfig())
+
+        mock_get_default_predictor.assert_not_called()
+        self.assertEqual(type(predictor), Predictor)
+
     @mock.patch("sagemaker.jumpstart.estimator.get_default_predictor")
     @mock.patch("sagemaker.jumpstart.estimator.is_valid_model_id")
     @mock.patch("sagemaker.jumpstart.factory.model.Session")

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

@@ -16,6 +16,7 @@
 from unittest import mock
 import unittest
 import pytest
+from sagemaker.async_inference.async_inference_config import AsyncInferenceConfig
 from sagemaker.jumpstart.enums import JumpStartScriptScope
 
 from sagemaker.jumpstart.model import JumpStartModel
@@ -417,6 +418,42 @@ def test_no_predictor_returns_default_predictor(
         self.assertEqual(type(predictor), Predictor)
         self.assertEqual(predictor, default_predictor_with_presets)
 
+    @mock.patch("sagemaker.jumpstart.model.get_default_predictor")
+    @mock.patch("sagemaker.jumpstart.model.is_valid_model_id")
+    @mock.patch("sagemaker.jumpstart.factory.model.Session")
+    @mock.patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
+    @mock.patch("sagemaker.jumpstart.model.Model.__init__")
+    @mock.patch("sagemaker.jumpstart.model.Model.deploy")
+    @mock.patch("sagemaker.jumpstart.factory.model.JUMPSTART_DEFAULT_REGION_NAME", region)
+    def test_no_predictor_yes_async_inference_config(
+        self,
+        mock_model_deploy: mock.Mock,
+        mock_model_init: mock.Mock,
+        mock_get_model_specs: mock.Mock,
+        mock_session: mock.Mock,
+        mock_is_valid_model_id: mock.Mock,
+        mock_get_default_predictor: mock.Mock,
+    ):
+        mock_get_default_predictor.return_value = default_predictor_with_presets
+
+        mock_model_deploy.return_value = default_predictor
+
+        mock_is_valid_model_id.return_value = True
+
+        model_id, _ = "js-model-class-model-prepacked", "*"
+
+        mock_get_model_specs.side_effect = get_special_model_spec
+
+        mock_session.return_value = sagemaker_session
+
+        model = JumpStartModel(
+            model_id=model_id,
+        )
+
+        model.deploy(async_inference_config=AsyncInferenceConfig())
+
+        mock_get_default_predictor.assert_not_called()
+
     @mock.patch("sagemaker.jumpstart.model.get_default_predictor")
     @mock.patch("sagemaker.jumpstart.model.is_valid_model_id")
     @mock.patch("sagemaker.jumpstart.factory.model.Session")