def test_model_package_enable_network_isolation_with_no_product_id(sagemaker_session):
sagemaker_session.sagemaker_client.describe_model_package = Mock(
return_value=DESCRIBE_MODEL_PACKAGE_RESPONSE)
model_package = ModelPackage(role='role', model_package_arn='my-model-package',
sagemaker_session=sagemaker_session)
assert model_package.enable_network_isolation() is False
def test_model_package_create_transformer(sagemaker_session):
model_package = ModelPackage(role="role",
model_package_arn="my-model-package",
sagemaker_session=sagemaker_session)
model_package.name = "auto-generated-model"
transformer = model_package.transformer(instance_count=1,
instance_type="ml.m4.xlarge",
env={"test": True})
assert isinstance(transformer, sagemaker.transformer.Transformer)
assert transformer.model_name == "auto-generated-model"
assert transformer.instance_type == "ml.m4.xlarge"
assert transformer.env == {"test": True}
def test_create_sagemaker_model_generates_model_name(name_from_base,
sagemaker_session):
model_package_name = "my-model-package"
model_package = ModelPackage(role="role",
model_package_arn=model_package_name,
sagemaker_session=sagemaker_session)
model_package._create_sagemaker_model()
name_from_base.assert_called_with(model_package_name)
assert name_from_base.return_value == model_package.name
def test_model_package_create_transformer(sagemaker_session):
sagemaker_session.sagemaker_client.describe_model_package = Mock(
return_value=DESCRIBE_MODEL_PACKAGE_RESPONSE)
model_package = ModelPackage(role='role', model_package_arn='my-model-package',
sagemaker_session=sagemaker_session)
model_package.name = 'auto-generated-model'
transformer = model_package.transformer(instance_count=1, instance_type='ml.m4.xlarge',
env={'test': True})
assert isinstance(transformer, sagemaker.transformer.Transformer)
assert transformer.model_name == 'auto-generated-model'
assert transformer.instance_type == 'ml.m4.xlarge'
assert transformer.env == {'test': True}
def test_create_sagemaker_model_generates_model_name_each_time(
name_from_base, sagemaker_session):
model_package_name = "my-model-package"
model_package = ModelPackage(role="role",
model_package_arn=model_package_name,
sagemaker_session=sagemaker_session)
model_package._create_sagemaker_model()
model_package._create_sagemaker_model()
name_from_base.assert_called_with(model_package_name)
assert 2 == name_from_base.call_count
def test_model_package_enable_network_isolation_with_product_id(sagemaker_session):
model_package_response = copy.deepcopy(DESCRIBE_MODEL_PACKAGE_RESPONSE)
model_package_response['InferenceSpecification']['Containers'].append(
{
'Image': '1.dkr.ecr.us-east-2.amazonaws.com/some-container:latest',
'ModelDataUrl': 's3://bucket/output/model.tar.gz',
'ProductId': 'some-product-id'
}
)
sagemaker_session.sagemaker_client.describe_model_package = Mock(
return_value=model_package_response)
model_package = ModelPackage(role='role', model_package_arn='my-model-package',
sagemaker_session=sagemaker_session)
assert model_package.enable_network_isolation() is True
def test_model_package_enable_network_isolation_with_product_id(sagemaker_session):
model_package_response = copy.deepcopy(DESCRIBE_MODEL_PACKAGE_RESPONSE)
model_package_response["InferenceSpecification"]["Containers"].append(
{
"Image": "1.dkr.ecr.us-east-2.amazonaws.com/some-container:latest",
"ModelDataUrl": "s3://bucket/output/model.tar.gz",
"ProductId": "some-product-id",
}
)
sagemaker_session.sagemaker_client.describe_model_package = Mock(
return_value=model_package_response
)
model_package = ModelPackage(
role="role", model_package_arn="my-model-package", sagemaker_session=sagemaker_session
)
assert model_package.enable_network_isolation() is True
def _deploy_static_endpoint(execution_arn, sagemaker_session):
try:
model_package_arn = get_model_package_arn_from_static_pipeline(
execution_arn, sagemaker_session)
model_package = ModelPackage(
role=get_execution_role(sagemaker_session),
model_package_arn=model_package_arn,
sagemaker_session=sagemaker_session,
)
model_package.deploy(1,
"ml.t2.medium",
endpoint_name=STATIC_ENDPOINT_NAME)
time.sleep(120)
except ClientError as e:
if e.response["Error"]["Code"] == "ValidationException":
print(
f"Endpoint {STATIC_ENDPOINT_NAME} already exists. Continuing.")
pass
else:
raise (e)
def test_model_package_create_transformer_with_product_id(sagemaker_session):
model_package_response = copy.deepcopy(DESCRIBE_MODEL_PACKAGE_RESPONSE)
model_package_response['InferenceSpecification']['Containers'].append(
{
'Image': '1.dkr.ecr.us-east-2.amazonaws.com/some-container:latest',
'ModelDataUrl': 's3://bucket/output/model.tar.gz',
'ProductId': 'some-product-id'
}
)
sagemaker_session.sagemaker_client.describe_model_package = Mock(
return_value=model_package_response)
model_package = ModelPackage(role='role', model_package_arn='my-model-package',
sagemaker_session=sagemaker_session)
model_package.name = 'auto-generated-model'
transformer = model_package.transformer(instance_count=1, instance_type='ml.m4.xlarge',
env={'test': True})
assert isinstance(transformer, sagemaker.transformer.Transformer)
assert transformer.model_name == 'auto-generated-model'
assert transformer.instance_type == 'ml.m4.xlarge'
assert transformer.env is None
def test_model_package_create_transformer_with_product_id(sagemaker_session):
model_package_response = copy.deepcopy(DESCRIBE_MODEL_PACKAGE_RESPONSE)
model_package_response["InferenceSpecification"]["Containers"].append({
"Image":
"1.dkr.ecr.us-east-2.amazonaws.com/some-container:latest",
"ModelDataUrl":
"s3://bucket/output/model.tar.gz",
"ProductId":
"some-product-id",
})
sagemaker_session.sagemaker_client.describe_model_package = Mock(
return_value=model_package_response)
model_package = ModelPackage(role="role",
model_package_arn="my-model-package",
sagemaker_session=sagemaker_session)
model_package.name = "auto-generated-model"
transformer = model_package.transformer(instance_count=1,
instance_type="ml.m4.xlarge",
env={"test": True})
assert isinstance(transformer, sagemaker.transformer.Transformer)
assert transformer.model_name == "auto-generated-model"
assert transformer.instance_type == "ml.m4.xlarge"
assert transformer.env is None
def test_create_sagemaker_model_uses_model_name(name_from_base,
sagemaker_session):
model_name = "my-model"
model_package_name = "my-model-package"
model_package = ModelPackage(
role="role",
name=model_name,
model_package_arn=model_package_name,
sagemaker_session=sagemaker_session,
)
model_package._create_sagemaker_model()
assert model_name == model_package.name
name_from_base.assert_not_called()
sagemaker_session.create_model.assert_called_with(
model_name,
"role",
{"ModelPackageName": model_package_name},
vpc_config=None,
enable_network_isolation=False,
)
def test_model_package_enable_network_isolation_with_no_product_id(
sagemaker_session):
model_package = ModelPackage(role="role",
model_package_arn="my-model-package",
sagemaker_session=sagemaker_session)
assert model_package.enable_network_isolation() is False