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_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_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