fix: create the correct session for MultiDataModel (#1255)

* fix: create the correct session for MultiDataModel

* remove unused param in test

* fix black-format

Author: chuyang-deng

src/sagemaker/local/local_session.py

@@ -341,6 +341,7 @@ def invoke_endpoint(
         ContentType=None,
         Accept=None,
         CustomAttributes=None,
+        TargetModel=None,
     ):
         """
 
@@ -365,6 +366,9 @@ def invoke_endpoint(
         if CustomAttributes is not None:
             headers["X-Amzn-SageMaker-Custom-Attributes"] = CustomAttributes
 
+        if TargetModel is not None:
+            headers["X-Amzn-SageMaker-Target-Model"] = TargetModel
+
         r = self.http.request("POST", url, body=Body, preload_content=False, headers=headers)
 
         return {"Body": r, "ContentType": Accept}

src/sagemaker/multidatamodel.py

@@ -17,7 +17,7 @@
 from six.moves.urllib.parse import urlparse
 
 import sagemaker
-from sagemaker import s3
+from sagemaker import local, s3
 from sagemaker.model import Model
 from sagemaker.session import Session
 
@@ -210,6 +210,9 @@ def deploy(
         if role is None:
             raise ValueError("Role can not be null for deploying a model")
 
+        if instance_type == "local" and not isinstance(self.sagemaker_session, local.LocalSession):
+            self.sagemaker_session = local.LocalSession()
+
         container_def = self.prepare_container_def(instance_type, accelerator_type=accelerator_type)
         self.sagemaker_session.create_model(
             self.name,

tests/integ/test_multidatamodel.py

@@ -177,6 +177,67 @@ def test_multi_data_model_deploy_pretrained_models(
         assert "Could not find endpoint" in str(exception.value)
 
 
+@pytest.mark.local_mode
+def test_multi_data_model_deploy_pretrained_models_local_mode(container_image, sagemaker_session):
+    timestamp = sagemaker_timestamp()
+    endpoint_name = "test-multimodel-endpoint-{}".format(timestamp)
+    model_name = "test-multimodel-{}".format(timestamp)
+
+    # Define pretrained model local path
+    pretrained_model_data_local_path = os.path.join(DATA_DIR, "sparkml_model", "mleap_model.tar.gz")
+
+    with timeout(minutes=30):
+        model_data_prefix = os.path.join(
+            "s3://", sagemaker_session.default_bucket(), "multimodel-{}/".format(timestamp)
+        )
+        multi_data_model = MultiDataModel(
+            name=model_name,
+            model_data_prefix=model_data_prefix,
+            image=container_image,
+            role=ROLE,
+            sagemaker_session=sagemaker_session,
+        )
+
+        # Add model before deploy
+        multi_data_model.add_model(pretrained_model_data_local_path, PRETRAINED_MODEL_PATH_1)
+        # Deploy model to an endpoint
+        multi_data_model.deploy(1, "local", endpoint_name=endpoint_name)
+        # Add models after deploy
+        multi_data_model.add_model(pretrained_model_data_local_path, PRETRAINED_MODEL_PATH_2)
+
+        endpoint_models = []
+        for model_path in multi_data_model.list_models():
+            endpoint_models.append(model_path)
+        assert PRETRAINED_MODEL_PATH_1 in endpoint_models
+        assert PRETRAINED_MODEL_PATH_2 in endpoint_models
+
+        predictor = RealTimePredictor(
+            endpoint=endpoint_name,
+            sagemaker_session=multi_data_model.sagemaker_session,
+            serializer=npy_serializer,
+            deserializer=string_deserializer,
+        )
+
+        data = numpy.zeros(shape=(1, 1, 28, 28))
+        result = predictor.predict(data, target_model=PRETRAINED_MODEL_PATH_1)
+        assert result == "Invoked model: {}".format(PRETRAINED_MODEL_PATH_1)
+
+        result = predictor.predict(data, target_model=PRETRAINED_MODEL_PATH_2)
+        assert result == "Invoked model: {}".format(PRETRAINED_MODEL_PATH_2)
+
+        # Cleanup
+        multi_data_model.sagemaker_session.sagemaker_client.delete_endpoint_config(
+            EndpointConfigName=endpoint_name
+        )
+        multi_data_model.sagemaker_session.delete_endpoint(endpoint_name)
+        multi_data_model.delete_model()
+    with pytest.raises(Exception) as exception:
+        sagemaker_session.sagemaker_client.describe_model(ModelName=multi_data_model.name)
+        assert "Could not find model" in str(exception.value)
+        sagemaker_session.sagemaker_client.describe_endpoint_config(name=endpoint_name)
+        assert "Could not find endpoint" in str(exception.value)
+
+
 def test_multi_data_model_deploy_trained_model_from_framework_estimator(
     container_image, sagemaker_session, cpu_instance_type
 ):