def test_prepare_container_def(tfo, time, sagemaker_session):
framework_model = DummyFrameworkModel(sagemaker_session)
sparkml_model = SparkMLModel(
model_data=MODEL_DATA_2,
role=ROLE,
sagemaker_session=sagemaker_session,
env={"SAGEMAKER_DEFAULT_INVOCATIONS_ACCEPT": "text/csv"},
)
model = PipelineModel(
models=[framework_model, sparkml_model], role=ROLE, sagemaker_session=sagemaker_session
)
assert model.pipeline_container_def(INSTANCE_TYPE) == [
{
"Environment": {
"SAGEMAKER_PROGRAM": "blah.py",
"SAGEMAKER_SUBMIT_DIRECTORY": "s3://mybucket/mi-1-2017-10-10-14-14-15/sourcedir.tar.gz",
"SAGEMAKER_CONTAINER_LOG_LEVEL": "20",
"SAGEMAKER_REGION": "us-west-2",
},
"Image": "mi-1",
"ModelDataUrl": "s3://bucket/model_1.tar.gz",
},
{
"Environment": {"SAGEMAKER_DEFAULT_INVOCATIONS_ACCEPT": "text/csv"},
"Image": "246618743249.dkr.ecr.us-west-2.amazonaws.com"
+ "/sagemaker-sparkml-serving:2.2",
"ModelDataUrl": "s3://bucket/model_2.tar.gz",
},
]
def test_prepare_container_def(tfo, time, sagemaker_session):
framework_model = DummyFrameworkModel(sagemaker_session)
sparkml_model = SparkMLModel(
model_data=MODEL_DATA_2,
role=ROLE,
sagemaker_session=sagemaker_session,
env={'SAGEMAKER_DEFAULT_INVOCATIONS_ACCEPT': 'text/csv'})
model = PipelineModel(models=[framework_model, sparkml_model],
role=ROLE,
sagemaker_session=sagemaker_session)
assert model.pipeline_container_def(INSTANCE_TYPE) == [{
'Environment': {
'SAGEMAKER_PROGRAM': 'blah.py',
'SAGEMAKER_SUBMIT_DIRECTORY':
's3://mybucket/mi-1-2017-10-10-14-14-15/sourcedir.tar.gz',
'SAGEMAKER_CONTAINER_LOG_LEVEL': '20',
'SAGEMAKER_REGION': 'us-west-2',
'SAGEMAKER_ENABLE_CLOUDWATCH_METRICS': 'false'
},
'Image':
'mi-1',
'ModelDataUrl':
's3://bucket/model_1.tar.gz'
}, {
'Environment': {
'SAGEMAKER_DEFAULT_INVOCATIONS_ACCEPT': 'text/csv'
},
'Image':
'246618743249.dkr.ecr.us-west-2.amazonaws.com' +
'/sagemaker-sparkml-serving:2.2',
'ModelDataUrl':
's3://bucket/model_2.tar.gz'
}]
