merge fix

Author: mwfongAWS

setup.py

@@ -48,7 +48,7 @@ def read_requirements(filename):
 # Declare minimal set for installation
 required_packages = [
     "attrs>=20.3.0,<23",
-    "boto3>=1.26.28,<2.0",
+    "boto3>=1.26.131,<2.0",
     "cloudpickle==2.2.1",
     "google-pasta",
     "numpy>=1.9.0,<2.0",

src/sagemaker/remote_function/client.py

@@ -745,7 +745,7 @@ def map(self, func, *iterables):
         futures = map(self.submit, itertools.repeat(func), *iterables)
         return [future.result() for future in futures]
 
-    def shutdown(self, wait=True):
+    def shutdown(self):
         """Prevent more function executions to be submitted to this executor."""
         with self._state_condition:
             self._shutdown = True
@@ -756,15 +756,15 @@ def shutdown(self, wait=True):
             self._state_condition.notify_all()
 
         if self._workers is not None:
-            self._workers.shutdown(wait)
+            self._workers.shutdown(wait=True)
 
     def __enter__(self):
         """Create an executor instance and return it"""
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         """Make sure the executor instance is shutdown."""
-        self.shutdown(wait=False)
+        self.shutdown()
         return False
 
     @staticmethod

src/sagemaker/serverless/serverless_inference_config.py

@@ -12,10 +12,11 @@
 # language governing permissions and limitations under the License.
 """This module contains code related to the ServerlessInferenceConfig class.
 
-Codes are used for configuring async inference endpoint. Use it when deploying
+Codes are used for configuring serverless inference endpoint. Use it when deploying
 the model to the endpoints.
 """
 from __future__ import print_function, absolute_import
+from typing import Optional
 
 
 class ServerlessInferenceConfig(object):
@@ -29,6 +30,7 @@ def __init__(
         self,
         memory_size_in_mb: int = 2048,
         max_concurrency: int = 5,
+        provisioned_concurrency: Optional[int] = None,
     ):
         """Initialize a ServerlessInferenceConfig object for serverless inference configuration.
 
@@ -40,9 +42,13 @@ def __init__(
             max_concurrency (int): Optional. The maximum number of concurrent invocations
                 your serverless endpoint can process. If no value is provided, Amazon
                 SageMaker will choose the default value for you. (Default: 5)
+            provisioned_concurrency (int): Optional. The provisioned concurrency of your
+                serverless endpoint. If no value is provided, Amazon SageMaker will not
+                apply provisioned concucrrency to your Serverless endpoint. (Default: None)
         """
         self.memory_size_in_mb = memory_size_in_mb
         self.max_concurrency = max_concurrency
+        self.provisioned_concurrency = provisioned_concurrency
 
     def _to_request_dict(self):
         """Generates a request dictionary using the parameters provided to the class."""
@@ -51,4 +57,7 @@ def _to_request_dict(self):
             "MaxConcurrency": self.max_concurrency,
         }
 
+        if self.provisioned_concurrency is not None:
+            request_dict["ProvisionedConcurrency"] = self.provisioned_concurrency
+
         return request_dict

tests/integ/test_serverless_inference.py

@@ -44,14 +44,30 @@ def test_serverless_walkthrough(sagemaker_session, cpu_instance_type, training_s
         pca.extra_components = 5
         pca.fit(pca.record_set(training_set[0][:100]), job_name=job_name)
 
-    with timeout_and_delete_endpoint_by_name(job_name, sagemaker_session):
+    serverless_name = unique_name_from_base("pca-serverless")
+    with timeout_and_delete_endpoint_by_name(serverless_name, sagemaker_session):
 
         predictor_serverless = pca.deploy(
-            endpoint_name=job_name, serverless_inference_config=ServerlessInferenceConfig()
+            endpoint_name=serverless_name, serverless_inference_config=ServerlessInferenceConfig()
         )
 
         result = predictor_serverless.predict(training_set[0][:5])
 
         assert len(result) == 5
         for record in result:
             assert record.label["projection"] is not None
+
+    # Test out Serverless Provisioned Concurrency endpoint happy case
+    serverless_pc_name = unique_name_from_base("pca-serverless-pc")
+    with timeout_and_delete_endpoint_by_name(serverless_pc_name, sagemaker_session):
+
+        predictor_serverless_pc = pca.deploy(
+            endpoint_name=serverless_pc_name,
+            serverless_inference_config=ServerlessInferenceConfig(provisioned_concurrency=1),
+        )
+
+        result = predictor_serverless_pc.predict(training_set[0][:5])
+
+        assert len(result) == 5
+        for record in result:
+            assert record.label["projection"] is not None

tests/unit/sagemaker/remote_function/test_client.py

@@ -518,11 +518,6 @@ def test_executor_submit_happy_case(mock_start, mock_job_settings, parallelism):
         future_3 = e.submit(job_function, 9, 10, c=11, d=12)
         future_4 = e.submit(job_function, 13, 14, c=15, d=16)
 
-    future_1.wait()
-    future_2.wait()
-    future_3.wait()
-    future_4.wait()
-
     mock_start.assert_has_calls(
         [
             call(ANY, job_function, (1, 2), {"c": 3, "d": 4}, None),
@@ -531,6 +526,10 @@ def test_executor_submit_happy_case(mock_start, mock_job_settings, parallelism):
             call(ANY, job_function, (13, 14), {"c": 15, "d": 16}, None),
         ]
     )
+    mock_job_1.describe.assert_called()
+    mock_job_2.describe.assert_called()
+    mock_job_3.describe.assert_called()
+    mock_job_4.describe.assert_called()
 
     assert future_1.done()
     assert future_2.done()
@@ -555,15 +554,14 @@ def test_executor_submit_with_run(mock_start, mock_job_settings, run_obj):
             future_1 = e.submit(job_function, 1, 2, c=3, d=4)
             future_2 = e.submit(job_function, 5, 6, c=7, d=8)
 
-    future_1.wait()
-    future_2.wait()
-
     mock_start.assert_has_calls(
         [
             call(ANY, job_function, (1, 2), {"c": 3, "d": 4}, run_info),
             call(ANY, job_function, (5, 6), {"c": 7, "d": 8}, run_info),
         ]
     )
+    mock_job_1.describe.assert_called()
+    mock_job_2.describe.assert_called()
 
     assert future_1.done()
     assert future_2.done()
@@ -573,15 +571,14 @@ def test_executor_submit_with_run(mock_start, mock_job_settings, run_obj):
             future_3 = e.submit(job_function, 9, 10, c=11, d=12)
             future_4 = e.submit(job_function, 13, 14, c=15, d=16)
 
-    future_3.wait()
-    future_4.wait()
-
     mock_start.assert_has_calls(
         [
             call(ANY, job_function, (9, 10), {"c": 11, "d": 12}, run_info),
             call(ANY, job_function, (13, 14), {"c": 15, "d": 16}, run_info),
         ]
     )
+    mock_job_3.describe.assert_called()
+    mock_job_4.describe.assert_called()
 
     assert future_3.done()
     assert future_4.done()
@@ -633,7 +630,7 @@ def test_executor_fails_to_start_job(mock_start, *args):
 
     with pytest.raises(TypeError):
         future_1.result()
-    future_2.wait()
+    print(future_2._state)
     assert future_2.done()
 
 
@@ -698,8 +695,6 @@ def test_executor_describe_job_throttled_temporarily(mock_start, *args):
         # submit second job
         future_2 = e.submit(job_function, 5, 6, c=7, d=8)
 
-    future_1.wait()
-    future_2.wait()
     assert future_1.done()
     assert future_2.done()
 
@@ -719,9 +714,9 @@ def test_executor_describe_job_failed_permanently(mock_start, *args):
         future_2 = e.submit(job_function, 5, 6, c=7, d=8)
 
     with pytest.raises(RuntimeError):
-        future_1.result()
+        future_1.done()
     with pytest.raises(RuntimeError):
-        future_2.result()
+        future_2.done()
 
 
 @pytest.mark.parametrize(

tests/unit/sagemaker/serverless/test_serverless_inference_config.py

@@ -16,21 +16,54 @@
 
 DEFAULT_MEMORY_SIZE_IN_MB = 2048
 DEFAULT_MAX_CONCURRENCY = 5
+DEFAULT_PROVISIONED_CONCURRENCY = 5
 
 DEFAULT_REQUEST_DICT = {
     "MemorySizeInMB": DEFAULT_MEMORY_SIZE_IN_MB,
     "MaxConcurrency": DEFAULT_MAX_CONCURRENCY,
 }
 
+PROVISIONED_CONCURRENCY_REQUEST_DICT = {
+    "MemorySizeInMB": DEFAULT_MEMORY_SIZE_IN_MB,
+    "MaxConcurrency": DEFAULT_MAX_CONCURRENCY,
+    "ProvisionedConcurrency": DEFAULT_PROVISIONED_CONCURRENCY,
+}
+
 
 def test_init():
     serverless_inference_config = ServerlessInferenceConfig()
 
     assert serverless_inference_config.memory_size_in_mb == DEFAULT_MEMORY_SIZE_IN_MB
     assert serverless_inference_config.max_concurrency == DEFAULT_MAX_CONCURRENCY
 
+    serverless_provisioned_concurrency_inference_config = ServerlessInferenceConfig(
+        provisioned_concurrency=DEFAULT_PROVISIONED_CONCURRENCY
+    )
+
+    assert (
+        serverless_provisioned_concurrency_inference_config.memory_size_in_mb
+        == DEFAULT_MEMORY_SIZE_IN_MB
+    )
+    assert (
+        serverless_provisioned_concurrency_inference_config.max_concurrency
+        == DEFAULT_MAX_CONCURRENCY
+    )
+    assert (
+        serverless_provisioned_concurrency_inference_config.provisioned_concurrency
+        == DEFAULT_PROVISIONED_CONCURRENCY
+    )
+
 
 def test_to_request_dict():
     serverless_inference_config_dict = ServerlessInferenceConfig()._to_request_dict()
 
     assert serverless_inference_config_dict == DEFAULT_REQUEST_DICT
+
+    serverless_provisioned_concurrency_inference_config_dict = ServerlessInferenceConfig(
+        provisioned_concurrency=DEFAULT_PROVISIONED_CONCURRENCY
+    )._to_request_dict()
+
+    assert (
+        serverless_provisioned_concurrency_inference_config_dict
+        == PROVISIONED_CONCURRENCY_REQUEST_DICT
+    )