def test_update_endpoint_no_args(name_from_base):
new_endpoint_config_name = "new-endpoint-config"
name_from_base.return_value = new_endpoint_config_name
sagemaker_session = empty_sagemaker_session()
existing_endpoint_config_name = "existing-endpoint-config"
predictor = Predictor(ENDPOINT, sagemaker_session=sagemaker_session)
predictor._endpoint_config_name = existing_endpoint_config_name
predictor.update_endpoint()
assert ["model-1", "model-2"] == predictor._model_names
assert new_endpoint_config_name == predictor._endpoint_config_name
name_from_base.assert_called_with(existing_endpoint_config_name)
sagemaker_session.create_endpoint_config_from_existing.assert_called_with(
existing_endpoint_config_name,
new_endpoint_config_name,
new_tags=None,
new_kms_key=None,
new_data_capture_config_dict=None,
new_production_variants=None,
)
sagemaker_session.update_endpoint.assert_called_with(
ENDPOINT, new_endpoint_config_name, wait=True)
def test_update_endpoint_no_one_default_model_name_with_instance_type_and_count(
):
sagemaker_session = empty_sagemaker_session()
predictor = Predictor(ENDPOINT, sagemaker_session=sagemaker_session)
with pytest.raises(ValueError) as exception:
predictor.update_endpoint(initial_instance_count=2,
instance_type="ml.c4.xlarge")
assert "Unable to choose a default model for a new EndpointConfig" in str(
exception.value)
def test_update_endpoint_no_instance_type_or_no_instance_count():
sagemaker_session = empty_sagemaker_session()
predictor = Predictor(ENDPOINT, sagemaker_session=sagemaker_session)
bad_args = ({
"instance_type": "ml.c4.xlarge"
}, {
"initial_instance_count": 2
})
for args in bad_args:
with pytest.raises(ValueError) as exception:
predictor.update_endpoint(**args)
expected_msg = "Missing initial_instance_count and/or instance_type."
assert expected_msg in str(exception.value)
def test_update_endpoint_all_args(name_from_base, production_variant):
new_endpoint_config_name = "new-endpoint-config"
name_from_base.return_value = new_endpoint_config_name
sagemaker_session = empty_sagemaker_session()
existing_endpoint_config_name = "existing-endpoint-config"
predictor = Predictor(ENDPOINT, sagemaker_session=sagemaker_session)
predictor._endpoint_config_name = existing_endpoint_config_name
new_instance_count = 2
new_instance_type = "ml.c4.xlarge"
new_accelerator_type = "ml.eia1.medium"
new_model_name = "new-model"
new_tags = {"Key": "foo", "Value": "bar"}
new_kms_key = "new-key"
new_data_capture_config_dict = {}
predictor.update_endpoint(
initial_instance_count=new_instance_count,
instance_type=new_instance_type,
accelerator_type=new_accelerator_type,
model_name=new_model_name,
tags=new_tags,
kms_key=new_kms_key,
data_capture_config_dict=new_data_capture_config_dict,
wait=False,
)
assert [new_model_name] == predictor._model_names
assert new_endpoint_config_name == predictor._endpoint_config_name
production_variant.assert_called_with(
new_model_name,
new_instance_type,
initial_instance_count=new_instance_count,
accelerator_type=new_accelerator_type,
)
sagemaker_session.create_endpoint_config_from_existing.assert_called_with(
existing_endpoint_config_name,
new_endpoint_config_name,
new_tags=new_tags,
new_kms_key=new_kms_key,
new_data_capture_config_dict=new_data_capture_config_dict,
new_production_variants=[production_variant.return_value],
)
sagemaker_session.update_endpoint.assert_called_with(
ENDPOINT, new_endpoint_config_name, wait=False)
def test_update_endpoint_instance_type_and_count(name_from_base,
production_variant):
new_endpoint_config_name = "new-endpoint-config"
name_from_base.return_value = new_endpoint_config_name
sagemaker_session = empty_sagemaker_session()
existing_endpoint_config_name = "existing-endpoint-config"
existing_model_name = "existing-model"
predictor = Predictor(ENDPOINT, sagemaker_session=sagemaker_session)
predictor._endpoint_config_name = existing_endpoint_config_name
predictor._model_names = [existing_model_name]
new_instance_count = 2
new_instance_type = "ml.c4.xlarge"
predictor.update_endpoint(
initial_instance_count=new_instance_count,
instance_type=new_instance_type,
)
assert [existing_model_name] == predictor._model_names
assert new_endpoint_config_name == predictor._endpoint_config_name
production_variant.assert_called_with(
existing_model_name,
new_instance_type,
initial_instance_count=new_instance_count,
accelerator_type=None,
)
sagemaker_session.create_endpoint_config_from_existing.assert_called_with(
existing_endpoint_config_name,
new_endpoint_config_name,
new_tags=None,
new_kms_key=None,
new_data_capture_config_dict=None,
new_production_variants=[production_variant.return_value],
)
sagemaker_session.update_endpoint.assert_called_with(
ENDPOINT, new_endpoint_config_name, wait=True)
def test_multi_data_model_deploy_pretrained_models_update_endpoint(
container_image, sagemaker_session, cpu_instance_type, alternative_cpu_instance_type
):
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_and_delete_endpoint_by_name(endpoint_name, sagemaker_session):
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_uri=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, cpu_instance_type, endpoint_name=endpoint_name)
# Add model after deploy
multi_data_model.add_model(pretrained_model_data_local_path, PRETRAINED_MODEL_PATH_2)
# List model assertions
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 = Predictor(
endpoint_name=endpoint_name,
sagemaker_session=sagemaker_session,
serializer=NumpySerializer(),
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)
endpoint_desc = sagemaker_session.sagemaker_client.describe_endpoint(
EndpointName=endpoint_name
)
old_config_name = endpoint_desc["EndpointConfigName"]
# Update endpoint
predictor.update_endpoint(
initial_instance_count=1, instance_type=alternative_cpu_instance_type
)
endpoint_desc = sagemaker_session.sagemaker_client.describe_endpoint(
EndpointName=endpoint_name
)
new_config_name = endpoint_desc["EndpointConfigName"]
new_config = sagemaker_session.sagemaker_client.describe_endpoint_config(
EndpointConfigName=new_config_name
)
assert old_config_name != new_config_name
assert new_config["ProductionVariants"][0]["InstanceType"] == alternative_cpu_instance_type
assert new_config["ProductionVariants"][0]["InitialInstanceCount"] == 1
# Cleanup
sagemaker_session.sagemaker_client.delete_endpoint_config(
EndpointConfigName=old_config_name
)
sagemaker_session.sagemaker_client.delete_endpoint_config(
EndpointConfigName=new_config_name
)
multi_data_model.delete_model()
with pytest.raises(Exception) as exception:
sagemaker_session.sagemaker_client.describe_model(ModelName=model_name)
assert "Could not find model" in str(exception.value)
sagemaker_session.sagemaker_client.describe_endpoint_config(name=old_config_name)
assert "Could not find endpoint" in str(exception.value)
sagemaker_session.sagemaker_client.describe_endpoint_config(name=new_config_name)
assert "Could not find endpoint" in str(exception.value)
