Merge branch 'master' into mwfongAWS-SM-sdk

Author: mwfongAWS

src/sagemaker/djl_inference/model.py

@@ -135,15 +135,17 @@ def _read_existing_serving_properties(directory: str):
     return properties
 
 
-def _get_model_config_properties_from_s3(model_s3_uri: str):
+def _get_model_config_properties_from_s3(model_s3_uri: str, sagemaker_session: Session):
     """Placeholder docstring"""
 
-    s3_files = s3.S3Downloader.list(model_s3_uri)
+    s3_files = s3.S3Downloader.list(model_s3_uri, sagemaker_session=sagemaker_session)
     model_config = None
     for config in defaults.VALID_MODEL_CONFIG_FILES:
         config_file = os.path.join(model_s3_uri, config)
         if config_file in s3_files:
-            model_config = json.loads(s3.S3Downloader.read_file(config_file))
+            model_config = json.loads(
+                s3.S3Downloader.read_file(config_file, sagemaker_session=sagemaker_session)
+            )
             break
     if not model_config:
         raise ValueError(
@@ -198,7 +200,8 @@ def __new__(
                 "containing folder"
             )
         if model_id.startswith("s3://"):
-            model_config = _get_model_config_properties_from_s3(model_id)
+            sagemaker_session = kwargs.get("sagemaker_session")
+            model_config = _get_model_config_properties_from_s3(model_id, sagemaker_session)
         else:
             model_config = _get_model_config_properties_from_hf(model_id)
         if model_config.get("_class_name") == "StableDiffusionPipeline":

src/sagemaker/model.py

@@ -574,6 +574,7 @@ def _upload_code(self, key_prefix: str, repack: bool = False) -> None:
                 script=self.entry_point,
                 directory=self.source_dir,
                 dependencies=self.dependencies,
+                kms_key=self.model_kms_key,
                 settings=self.sagemaker_session.settings,
             )
 

tests/conftest.py

@@ -52,6 +52,7 @@
     "me-south-1",
     "sa-east-1",
     "us-west-1",
+    "ap-south-1",  # no p3 availability
 ]
 
 NO_T2_REGIONS = ["eu-north-1", "ap-east-1", "me-south-1"]

tests/integ/test_inference_recommender.py

@@ -301,7 +301,6 @@ def test_default_right_size_and_deploy_registered_model_sklearn(
             predictor.delete_endpoint()
 
 
-@pytest.mark.skip(reason="This test is currently failing. Skipping until fixed")
 @pytest.mark.slow_test
 def test_default_right_size_and_deploy_unregistered_model_sklearn(
     default_right_sized_unregistered_model, sagemaker_session
@@ -346,7 +345,6 @@ def test_default_right_size_and_deploy_unregistered_base_model(
             predictor.delete_endpoint()
 
 
-@pytest.mark.skip(reason="This test is currently failing. Skipping until fixed")
 @pytest.mark.slow_test
 def test_advanced_right_size_and_deploy_unregistered_model_sklearn(
     advanced_right_sized_unregistered_model, sagemaker_session

tests/unit/test_djl_inference.py

@@ -174,6 +174,9 @@ def test_create_model_automatic_engine_selection(mock_s3_list, mock_read_file, s
             sagemaker_session=sagemaker_session,
             number_of_partitions=2,
         )
+        mock_s3_list.assert_any_call(
+            VALID_UNCOMPRESSED_MODEL_DATA, sagemaker_session=sagemaker_session
+        )
         if model_type == defaults.STABLE_DIFFUSION_MODEL_TYPE:
             assert ds_model.engine == DJLServingEngineEntryPointDefaults.STABLE_DIFFUSION
         else: