def test_create_sagemaker_model_accelerator_type(prepare_container_def, sagemaker_session):
model = Model(MODEL_IMAGE, MODEL_DATA, name=MODEL_NAME, sagemaker_session=sagemaker_session)
accelerator_type = "ml.eia.medium"
model._create_sagemaker_model(INSTANCE_TYPE, accelerator_type=accelerator_type)
prepare_container_def.assert_called_with(INSTANCE_TYPE, accelerator_type=accelerator_type)
def test_register_calls_model_package_args(get_model_package_args,
sagemaker_session):
source_dir = "s3://blah/blah/blah"
t = Model(
entry_point=ENTRY_POINT_INFERENCE,
role=ROLE,
sagemaker_session=sagemaker_session,
source_dir=source_dir,
image_uri=IMAGE_URI,
model_data=MODEL_DATA,
)
t.register(
SUPPORTED_CONTENT_TYPES,
SUPPORTED_RESPONSE_MIME_TYPES,
SUPPORTED_REALTIME_INFERENCE_INSTANCE_TYPES,
SUPPORTED_BATCH_TRANSFORM_INSTANCE_TYPES,
marketplace_cert=True,
description=MODEL_DESCRIPTION,
model_package_name=MODEL_NAME,
validation_specification=VALIDATION_SPECIFICATION,
)
# check that the kwarg validation_specification was passed to the internal method 'get_model_package_args'
assert (
"validation_specification"
in get_model_package_args.call_args_list[0][1]
), "validation_specification kwarg was not passed to get_model_package_args"
# check that the kwarg validation_specification is identical to the one passed into the method 'register'
assert (
VALIDATION_SPECIFICATION == get_model_package_args.call_args_list[0][1]
["validation_specification"]
), """ValidationSpecification from model.register method is not identical to validation_spec from
def _test_hub_model(sagemaker_session, framework_version, ecr_image, instance_type, model_dir, accelerator_type=None):
endpoint_name = sagemaker.utils.unique_name_from_base("sagemaker-huggingface-serving-hub-model")
env = {
"HF_MODEL_ID": "sshleifer/tiny-distilbert-base-uncased-finetuned-sst-2-english",
"HF_TASK": "text-classification",
}
hf_model = Model(
env=env,
role="SageMakerRole",
image_uri=ecr_image,
sagemaker_session=sagemaker_session,
predictor_cls=Predictor,
)
with timeout_and_delete_endpoint(endpoint_name, sagemaker_session, minutes=30):
predictor = hf_model.deploy(
initial_instance_count=1,
instance_type=instance_type,
endpoint_name=endpoint_name,
)
data = {
"inputs": "Camera - You are awarded a SiPix Digital Camera! call 09061221066 fromm landline. Delivery within 28 days."
}
predictor.serializer = JSONSerializer()
predictor.deserializer = JSONDeserializer()
output = predictor.predict(data)
assert "score" in output[0]
def test_model_create_transformer_network_isolation(sagemaker_session):
model = Model(
MODEL_IMAGE, MODEL_DATA, sagemaker_session=sagemaker_session, enable_network_isolation=True
)
transformer = model.transformer(1, "ml.m4.xlarge", env={"should_be": "overwritten"})
assert transformer.env is None
def test_create_sagemaker_model_instance_type(prepare_container_def,
sagemaker_session):
model = Model(MODEL_DATA, MODEL_IMAGE, sagemaker_session=sagemaker_session)
model._create_sagemaker_model(INSTANCE_TYPE)
prepare_container_def.assert_called_with(INSTANCE_TYPE,
accelerator_type=None)
def test_create_sagemaker_model_generates_model_name(base_name_from_image,
name_from_base,
prepare_container_def,
sagemaker_session):
container_def = {
"Image": MODEL_IMAGE,
"Environment": {},
"ModelDataUrl": MODEL_DATA
}
prepare_container_def.return_value = container_def
model = Model(
MODEL_IMAGE,
MODEL_DATA,
sagemaker_session=sagemaker_session,
)
model._create_sagemaker_model(INSTANCE_TYPE)
base_name_from_image.assert_called_with(MODEL_IMAGE)
name_from_base.assert_called_with(base_name_from_image.return_value)
sagemaker_session.create_model.assert_called_with(
MODEL_NAME,
None,
container_def,
vpc_config=None,
enable_network_isolation=False,
tags=None,
)
def test_git_support_succeed_model_class(tar_and_upload_dir, git_clone_repo,
sagemaker_session):
git_clone_repo.side_effect = lambda gitconfig, entrypoint, sourcedir, dependency: {
"entry_point": "entry_point",
"source_dir": "/tmp/repo_dir/source_dir",
"dependencies": ["/tmp/repo_dir/foo", "/tmp/repo_dir/bar"],
}
entry_point = "entry_point"
source_dir = "source_dir"
dependencies = ["foo", "bar"]
git_config = {"repo": GIT_REPO, "branch": BRANCH, "commit": COMMIT}
model = Model(
sagemaker_session=sagemaker_session,
entry_point=entry_point,
source_dir=source_dir,
dependencies=dependencies,
git_config=git_config,
image_uri=IMAGE_URI,
)
model.prepare_container_def(instance_type=INSTANCE_TYPE)
git_clone_repo.assert_called_with(git_config, entry_point, source_dir,
dependencies)
assert model.entry_point == "entry_point"
assert model.source_dir == "/tmp/repo_dir/source_dir"
assert model.dependencies == ["/tmp/repo_dir/foo", "/tmp/repo_dir/bar"]
def test_deploy_data_capture_config(production_variant, name_from_base,
sagemaker_session):
model = Model(MODEL_IMAGE,
MODEL_DATA,
role=ROLE,
name=MODEL_NAME,
sagemaker_session=sagemaker_session)
data_capture_config = Mock()
data_capture_config_dict = {"EnableCapture": True}
data_capture_config._to_request_dict.return_value = data_capture_config_dict
model.deploy(
instance_type=INSTANCE_TYPE,
initial_instance_count=INSTANCE_COUNT,
data_capture_config=data_capture_config,
)
data_capture_config._to_request_dict.assert_called_with()
sagemaker_session.endpoint_from_production_variants.assert_called_with(
name=ENDPOINT_NAME,
production_variants=[BASE_PRODUCTION_VARIANT],
tags=None,
kms_key=None,
wait=True,
data_capture_config_dict=data_capture_config_dict,
async_inference_config_dict=None,
)
def test_deploy_no_role(sagemaker_session):
model = Model(MODEL_IMAGE, MODEL_DATA, sagemaker_session=sagemaker_session)
with pytest.raises(ValueError,
match="Role can not be null for deploying a model"):
model.deploy(instance_type=INSTANCE_TYPE,
initial_instance_count=INSTANCE_COUNT)
def test_create_sagemaker_model_optional_model_params(base_name_from_image,
name_from_base,
prepare_container_def,
sagemaker_session):
container_def = {
"Image": MODEL_IMAGE,
"Environment": {},
"ModelDataUrl": MODEL_DATA
}
prepare_container_def.return_value = container_def
vpc_config = {"Subnets": ["123"], "SecurityGroupIds": ["456", "789"]}
model = Model(
MODEL_IMAGE,
MODEL_DATA,
name=MODEL_NAME,
role=ROLE,
vpc_config=vpc_config,
enable_network_isolation=True,
sagemaker_session=sagemaker_session,
)
model._create_sagemaker_model(INSTANCE_TYPE)
base_name_from_image.assert_not_called()
name_from_base.assert_not_called()
sagemaker_session.create_model.assert_called_with(
MODEL_NAME,
ROLE,
container_def,
vpc_config=vpc_config,
enable_network_isolation=True,
tags=None,
)
def test_deploy_predictor_cls(production_variant, sagemaker_session):
model = Model(
MODEL_IMAGE,
MODEL_DATA,
role=ROLE,
name=MODEL_NAME,
predictor_cls=sagemaker.predictor.Predictor,
sagemaker_session=sagemaker_session,
)
endpoint_name = "foo"
predictor = model.deploy(
instance_type=INSTANCE_TYPE,
initial_instance_count=INSTANCE_COUNT,
endpoint_name=endpoint_name,
)
assert isinstance(predictor, sagemaker.predictor.Predictor)
assert predictor.endpoint_name == endpoint_name
assert predictor.sagemaker_session == sagemaker_session
endpoint_name_async = "foo-async"
predictor_async = model.deploy(
instance_type=INSTANCE_TYPE,
initial_instance_count=INSTANCE_COUNT,
endpoint_name=endpoint_name_async,
async_inference_config=AsyncInferenceConfig(),
)
assert isinstance(predictor_async,
sagemaker.predictor_async.AsyncPredictor)
assert predictor_async.name == model.name
assert predictor_async.endpoint_name == endpoint_name_async
assert predictor_async.sagemaker_session == sagemaker_session
def test_deploy_async_inference(production_variant, name_from_base,
sagemaker_session):
model = Model(MODEL_IMAGE,
MODEL_DATA,
role=ROLE,
name=MODEL_NAME,
sagemaker_session=sagemaker_session)
async_inference_config = AsyncInferenceConfig(output_path="s3://some-path")
async_inference_config_dict = {
"OutputConfig": {
"S3OutputPath": "s3://some-path",
},
}
model.deploy(
instance_type=INSTANCE_TYPE,
initial_instance_count=INSTANCE_COUNT,
async_inference_config=async_inference_config,
)
sagemaker_session.endpoint_from_production_variants.assert_called_with(
name=ENDPOINT_NAME,
production_variants=[BASE_PRODUCTION_VARIANT],
tags=None,
kms_key=None,
wait=True,
data_capture_config_dict=None,
async_inference_config_dict=async_inference_config_dict,
)
def test_model_create_transformer(create_sagemaker_model, sagemaker_session):
model_name = "auto-generated-model"
model = Model(MODEL_IMAGE,
MODEL_DATA,
name=model_name,
sagemaker_session=sagemaker_session)
instance_type = "ml.m4.xlarge"
transformer = model.transformer(instance_count=1,
instance_type=instance_type)
create_sagemaker_model.assert_called_with(instance_type, tags=None)
assert isinstance(transformer, sagemaker.transformer.Transformer)
assert transformer.model_name == model_name
assert transformer.instance_type == instance_type
assert transformer.instance_count == 1
assert transformer.sagemaker_session == sagemaker_session
assert transformer.base_transform_job_name == model_name
assert transformer.strategy is None
assert transformer.env is None
assert transformer.output_path is None
assert transformer.output_kms_key is None
assert transformer.accept is None
assert transformer.assemble_with is None
assert transformer.volume_kms_key is None
assert transformer.max_concurrent_transforms is None
assert transformer.max_payload is None
assert transformer.tags is None
def endpoint_name(sagemaker_session):
endpoint_name = unique_name_from_base("model-quality-monitor-integ")
xgb_model_data = sagemaker_session.upload_data(
path=os.path.join(XGBOOST_DATA_PATH, "xgb_model.tar.gz"),
key_prefix="integ-test-data/xgboost/model",
)
xgb_image = image_uris.retrieve("xgboost",
sagemaker_session.boto_region_name,
version="1",
image_scope="inference")
with tests.integ.timeout.timeout_and_delete_endpoint_by_name(
endpoint_name=endpoint_name,
sagemaker_session=sagemaker_session,
hours=2):
xgb_model = Model(
model_data=xgb_model_data,
image_uri=xgb_image,
name=endpoint_name, # model name
role=ROLE,
sagemaker_session=sagemaker_session,
)
xgb_model.deploy(
INSTANCE_COUNT,
INSTANCE_TYPE,
endpoint_name=endpoint_name,
data_capture_config=DataCaptureConfig(
True, sagemaker_session=sagemaker_session),
)
yield endpoint_name
def test_create_sagemaker_model(name_from_image, prepare_container_def,
sagemaker_session):
name_from_image.return_value = MODEL_NAME
container_def = {
"Image": MODEL_IMAGE,
"Environment": {},
"ModelDataUrl": MODEL_DATA
}
prepare_container_def.return_value = container_def
model = Model(MODEL_DATA, MODEL_IMAGE, sagemaker_session=sagemaker_session)
model._create_sagemaker_model(INSTANCE_TYPE)
prepare_container_def.assert_called_with(INSTANCE_TYPE,
accelerator_type=None)
name_from_image.assert_called_with(MODEL_IMAGE)
sagemaker_session.create_model.assert_called_with(
MODEL_NAME,
None,
container_def,
vpc_config=None,
enable_network_isolation=False,
tags=None)
def test_deploy_accelerator_type(production_variant, create_sagemaker_model,
sagemaker_session):
model = Model(MODEL_DATA,
MODEL_IMAGE,
role=ROLE,
name=MODEL_NAME,
sagemaker_session=sagemaker_session)
production_variant_result = copy.deepcopy(BASE_PRODUCTION_VARIANT)
production_variant_result["AcceleratorType"] = ACCELERATOR_TYPE
production_variant.return_value = production_variant_result
model.deploy(
instance_type=INSTANCE_TYPE,
initial_instance_count=INSTANCE_COUNT,
accelerator_type=ACCELERATOR_TYPE,
)
create_sagemaker_model.assert_called_with(INSTANCE_TYPE, ACCELERATOR_TYPE,
None)
production_variant.assert_called_with(MODEL_NAME,
INSTANCE_TYPE,
INSTANCE_COUNT,
accelerator_type=ACCELERATOR_TYPE)
sagemaker_session.endpoint_from_production_variants.assert_called_with(
name=MODEL_NAME,
production_variants=[production_variant_result],
tags=None,
kms_key=None,
wait=True,
data_capture_config_dict=None,
)
def test_deploy_update_endpoint(sagemaker_session):
model = Model(MODEL_DATA,
MODEL_IMAGE,
role=ROLE,
sagemaker_session=sagemaker_session)
model.deploy(instance_type=INSTANCE_TYPE,
initial_instance_count=INSTANCE_COUNT,
update_endpoint=True)
sagemaker_session.create_endpoint_config.assert_called_with(
name=model.name,
model_name=model.name,
initial_instance_count=INSTANCE_COUNT,
instance_type=INSTANCE_TYPE,
accelerator_type=None,
tags=None,
kms_key=None,
data_capture_config_dict=None,
)
config_name = sagemaker_session.create_endpoint_config(
name=model.name,
model_name=model.name,
initial_instance_count=INSTANCE_COUNT,
instance_type=INSTANCE_TYPE,
accelerator_type=ACCELERATOR_TYPE,
)
sagemaker_session.update_endpoint.assert_called_with(model.name,
config_name,
wait=True)
sagemaker_session.create_endpoint.assert_not_called()
def test_delete_model_no_name(sagemaker_session):
model = Model(MODEL_IMAGE, MODEL_DATA, sagemaker_session=sagemaker_session)
with pytest.raises(
ValueError, match="The SageMaker model must be created first before attempting to delete."
):
model.delete_model()
sagemaker_session.delete_model.assert_not_called()
def test_model_create_transformer_base_name(sagemaker_session):
model = Model(MODEL_IMAGE, MODEL_DATA, sagemaker_session=sagemaker_session)
base_name = "foo"
model._base_name = base_name
transformer = model.transformer(1, "ml.m4.xlarge")
assert base_name == transformer.base_transform_job_name
def test_deploy_wrong_serverless_config(sagemaker_session):
model = Model(MODEL_IMAGE, MODEL_DATA, role=ROLE)
with pytest.raises(
ValueError,
match=
"serverless_inference_config needs to be a ServerlessInferenceConfig object",
):
model.deploy(serverless_inference_config={})
def test_delete_model(sagemaker_session):
model = Model(MODEL_IMAGE,
MODEL_DATA,
name=MODEL_NAME,
sagemaker_session=sagemaker_session)
model.delete_model()
sagemaker_session.delete_model.assert_called_with(model.name)
def _create_model(sagemaker_session=None):
model = Model(MODEL_IMAGE,
MODEL_DATA,
role="role",
sagemaker_session=sagemaker_session)
model._compilation_job_name = "compilation-test-name"
model._is_compiled_model = True
return model
def test_prepare_container_def_with_model_data_and_env():
env = {"FOO": "BAR"}
model = Model(MODEL_IMAGE, MODEL_DATA, env=env)
expected = {"Image": MODEL_IMAGE, "Environment": env, "ModelDataUrl": MODEL_DATA}
container_def = model.prepare_container_def(INSTANCE_TYPE, "ml.eia.medium")
assert expected == container_def
container_def = model.prepare_container_def()
assert expected == container_def
def test_prepare_container_def_with_image_config():
image_config = {"RepositoryAccessMode": "Vpc"}
model = Model(MODEL_IMAGE, image_config=image_config)
expected = {
"Image": MODEL_IMAGE,
"ImageConfig": {"RepositoryAccessMode": "Vpc"},
"Environment": {},
}
container_def = model.prepare_container_def()
assert expected == container_def
def test_create_sagemaker_model_tags(prepare_container_def, sagemaker_session):
container_def = {"Image": MODEL_IMAGE, "Environment": {}, "ModelDataUrl": MODEL_DATA}
prepare_container_def.return_value = container_def
model = Model(MODEL_IMAGE, MODEL_DATA, name=MODEL_NAME, sagemaker_session=sagemaker_session)
tags = {"Key": "foo", "Value": "bar"}
model._create_sagemaker_model(INSTANCE_TYPE, tags=tags)
sagemaker_session.create_model.assert_called_with(
MODEL_NAME, None, container_def, vpc_config=None, enable_network_isolation=False, tags=tags
)
def test_jumpstart_inference_model_class(setup):
model_id, model_version = "catboost-classification-model", "1.0.0"
instance_type, instance_count = "ml.m5.xlarge", 1
print("Starting inference...")
image_uri = image_uris.retrieve(
region=None,
framework=None,
image_scope="inference",
model_id=model_id,
model_version=model_version,
instance_type=instance_type,
)
script_uri = script_uris.retrieve(model_id=model_id,
model_version=model_version,
script_scope="inference")
model_uri = model_uris.retrieve(model_id=model_id,
model_version=model_version,
model_scope="inference")
model = Model(
image_uri=image_uri,
model_data=model_uri,
source_dir=script_uri,
entry_point=INFERENCE_ENTRY_POINT_SCRIPT_NAME,
role=get_sm_session().get_caller_identity_arn(),
sagemaker_session=get_sm_session(),
enable_network_isolation=True,
)
model.deploy(
initial_instance_count=instance_count,
instance_type=instance_type,
tags=[{
"Key": JUMPSTART_TAG,
"Value": os.environ[ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID]
}],
)
endpoint_invoker = EndpointInvoker(endpoint_name=model.endpoint_name, )
download_inference_assets()
ground_truth_label, features = get_tabular_data(
InferenceTabularDataname.MULTICLASS)
response = endpoint_invoker.invoke_tabular_endpoint(features)
assert response is not None
def test_prepare_container_def_with_model_src_s3_returns_correct_url(
sagemaker_session):
model = Model(
entry_point=ENTRY_POINT_INFERENCE,
role=ROLE,
sagemaker_session=sagemaker_session,
source_dir=SCRIPT_URI,
image_uri=MODEL_IMAGE,
model_data=Properties("Steps.MyStep"),
)
container_def = model.prepare_container_def(INSTANCE_TYPE, "ml.eia.medium")
assert container_def["Environment"][
"SAGEMAKER_SUBMIT_DIRECTORY"] == SCRIPT_URI
def test_compile_validates_model_data():
model = Model(MODEL_IMAGE)
with pytest.raises(ValueError) as e:
model.compile(
target_instance_family="ml_c4",
input_shape={"data": [1, 3, 1024, 1024]},
output_path="s3://output",
role="role",
framework="tensorflow",
job_name="compile-model",
)
assert "You must provide an S3 path to the compressed model artifacts." in str(e)
def test_deploy(name_from_base, prepare_container_def, production_variant,
sagemaker_session):
production_variant.return_value = BASE_PRODUCTION_VARIANT
container_def = {
"Image": MODEL_IMAGE,
"Environment": {},
"ModelDataUrl": MODEL_DATA
}
prepare_container_def.return_value = container_def
model = Model(MODEL_IMAGE,
MODEL_DATA,
role=ROLE,
sagemaker_session=sagemaker_session)
model.deploy(instance_type=INSTANCE_TYPE,
initial_instance_count=INSTANCE_COUNT)
name_from_base.assert_called_with(MODEL_IMAGE)
assert 2 == name_from_base.call_count
prepare_container_def.assert_called_with(INSTANCE_TYPE,
accelerator_type=None,
serverless_inference_config=None)
production_variant.assert_called_with(
MODEL_NAME,
INSTANCE_TYPE,
INSTANCE_COUNT,
accelerator_type=None,
serverless_inference_config=None,
)
sagemaker_session.create_model.assert_called_with(
MODEL_NAME,
ROLE,
container_def,
vpc_config=None,
enable_network_isolation=False,
tags=None)
sagemaker_session.endpoint_from_production_variants.assert_called_with(
name=MODEL_NAME,
production_variants=[BASE_PRODUCTION_VARIANT],
tags=None,
kms_key=None,
wait=True,
data_capture_config_dict=None,
async_inference_config_dict=None,
)
def test_create_sagemaker_model_creates_correct_session(local_session, session):
model = Model(MODEL_IMAGE, MODEL_DATA)
model._create_sagemaker_model("local")
assert model.sagemaker_session == local_session.return_value
model = Model(MODEL_IMAGE, MODEL_DATA)
model._create_sagemaker_model("ml.m5.xlarge")
assert model.sagemaker_session == session.return_value