def test_byo_framework_model_config(sagemaker_session):
byo_model = model.FrameworkModel(
model_data="{{ model_data }}",
image="{{ image }}",
role="{{ role }}",
entry_point="{{ entry_point }}",
source_dir="{{ source_dir }}",
env={"{{ key }}": "{{ value }}"},
name="model",
sagemaker_session=sagemaker_session,
)
config = airflow.model_config(instance_type="ml.c4.xlarge",
model=byo_model)
expected_config = {
"ModelName": "model",
"PrimaryContainer": {
"Image": "{{ image }}",
"Environment": {
"{{ key }}": "{{ value }}",
"SAGEMAKER_PROGRAM": "{{ entry_point }}",
"SAGEMAKER_SUBMIT_DIRECTORY":
"s3://output/model/source/sourcedir.tar.gz",
"SAGEMAKER_ENABLE_CLOUDWATCH_METRICS": "false",
"SAGEMAKER_CONTAINER_LOG_LEVEL": "20",
"SAGEMAKER_REGION": "us-west-2",
},
"ModelDataUrl": "{{ model_data }}",
},
"ExecutionRoleArn": "{{ role }}",
"S3Operations": {
"S3Upload": [{
"Path": "{{ source_dir }}",
"Bucket": "output",
"Key": "model/source/sourcedir.tar.gz",
"Tar": True,
}]
},
}
assert config == expected_config
def test_byo_framework_model_config(sagemaker_session):
byo_model = model.FrameworkModel(model_data="{{ model_data }}",
image="{{ image }}",
role="{{ role }}",
entry_point="{{ entry_point }}",
source_dir="{{ source_dir }}",
env={"{{ key }}": "{{ value }}"},
name='model',
sagemaker_session=sagemaker_session)
config = airflow.model_config(instance_type='ml.c4.xlarge',
model=byo_model)
expected_config = {
'ModelName': 'model',
'PrimaryContainer': {
'Image': '{{ image }}',
'Environment': {
'{{ key }}': '{{ value }}',
'SAGEMAKER_PROGRAM': '{{ entry_point }}',
'SAGEMAKER_SUBMIT_DIRECTORY':
's3://output/model/source/sourcedir.tar.gz',
'SAGEMAKER_ENABLE_CLOUDWATCH_METRICS': 'false',
'SAGEMAKER_CONTAINER_LOG_LEVEL': '20',
'SAGEMAKER_REGION': 'us-west-2'
},
'ModelDataUrl': '{{ model_data }}'
},
'ExecutionRoleArn': '{{ role }}',
'S3Operations': {
'S3Upload': [{
'Path': '{{ source_dir }}',
'Bucket': 'output',
'Key': 'model/source/sourcedir.tar.gz',
'Tar': True
}]
}
}
assert config == expected_config