def test_deploy_amazon_alg_model_config(sagemaker_session):
pca_model = pca.PCAModel(
model_data="{{ model_data }}",
role="{{ role }}",
sagemaker_session=sagemaker_session)
config = airflow.deploy_config(pca_model,
initial_instance_count="{{ instance_count }}",
instance_type='ml.c4.xlarge')
expected_config = {
'Model': {
'ModelName': "pca-%s" % TIME_STAMP,
'PrimaryContainer': {
'Image': '174872318107.dkr.ecr.us-west-2.amazonaws.com/pca:1',
'Environment': {},
'ModelDataUrl': '{{ model_data }}'},
'ExecutionRoleArn': '{{ role }}'},
'EndpointConfig': {
'EndpointConfigName': "pca-%s" % TIME_STAMP,
'ProductionVariants': [{
'InstanceType': 'ml.c4.xlarge',
'InitialInstanceCount': '{{ instance_count }}',
'ModelName': "pca-%s" % TIME_STAMP,
'VariantName': 'AllTraffic',
'InitialVariantWeight': 1
}]
},
'Endpoint': {
'EndpointName': "pca-%s" % TIME_STAMP,
'EndpointConfigName': "pca-%s" % TIME_STAMP
}
}
assert config == expected_config
def test_deploy_framework_model_config(sagemaker_session):
chainer_model = chainer.ChainerModel(
model_data="{{ model_data }}",
role="{{ role }}",
entry_point="{{ entry_point }}",
source_dir="{{ source_dir }}",
image=None,
py_version='py3',
framework_version='5.0.0',
model_server_workers="{{ model_server_worker }}",
sagemaker_session=sagemaker_session)
config = airflow.deploy_config(chainer_model,
initial_instance_count="{{ instance_count }}",
instance_type="ml.m4.xlarge")
expected_config = {
'Model': {
'ModelName': "sagemaker-chainer-%s" % TIME_STAMP,
'PrimaryContainer': {
'Image': '520713654638.dkr.ecr.us-west-2.amazonaws.com/sagemaker-chainer:5.0.0-cpu-py3',
'Environment': {
'SAGEMAKER_PROGRAM': '{{ entry_point }}',
'SAGEMAKER_SUBMIT_DIRECTORY': "s3://output/sagemaker-chainer-%s/source/sourcedir.tar.gz"
% TIME_STAMP,
'SAGEMAKER_ENABLE_CLOUDWATCH_METRICS': 'false',
'SAGEMAKER_CONTAINER_LOG_LEVEL': '20',
'SAGEMAKER_REGION': 'us-west-2',
'SAGEMAKER_MODEL_SERVER_WORKERS': '{{ model_server_worker }}'
},
'ModelDataUrl': '{{ model_data }}'},
'ExecutionRoleArn': '{{ role }}'
},
'EndpointConfig': {
'EndpointConfigName': "sagemaker-chainer-%s" % TIME_STAMP,
'ProductionVariants': [{
'InstanceType': 'ml.m4.xlarge',
'InitialInstanceCount': '{{ instance_count }}',
'ModelName': "sagemaker-chainer-%s" % TIME_STAMP,
'VariantName': 'AllTraffic',
'InitialVariantWeight': 1
}]
},
'Endpoint': {
'EndpointName': "sagemaker-chainer-%s" % TIME_STAMP,
'EndpointConfigName': "sagemaker-chainer-%s" % TIME_STAMP
},
'S3Operations': {
'S3Upload': [{
'Path': '{{ source_dir }}',
'Bucket': 'output',
'Key': "sagemaker-chainer-%s/source/sourcedir.tar.gz" % TIME_STAMP,
'Tar': True
}]
}
}
assert config == expected_config
def test_deploy_amazon_alg_model_config(sagemaker_session):
pca_model = pca.PCAModel(model_data="{{ model_data }}",
role="{{ role }}",
sagemaker_session=sagemaker_session)
config = airflow.deploy_config(
pca_model,
initial_instance_count="{{ instance_count }}",
instance_type="ml.c4.xlarge")
expected_config = {
"Model": {
"ModelName": "pca-%s" % TIME_STAMP,
"PrimaryContainer": {
"Image": "174872318107.dkr.ecr.us-west-2.amazonaws.com/pca:1",
"Environment": {},
"ModelDataUrl": "{{ model_data }}",
},
"ExecutionRoleArn": "{{ role }}",
},
"EndpointConfig": {
"EndpointConfigName":
"pca-%s" % TIME_STAMP,
"ProductionVariants": [{
"InstanceType": "ml.c4.xlarge",
"InitialInstanceCount": "{{ instance_count }}",
"ModelName": "pca-%s" % TIME_STAMP,
"VariantName": "AllTraffic",
"InitialVariantWeight": 1,
}],
},
"Endpoint": {
"EndpointName": "pca-%s" % TIME_STAMP,
"EndpointConfigName": "pca-%s" % TIME_STAMP,
},
}
assert config == expected_config
def test_deploy_framework_model_config(sagemaker_session):
chainer_model = chainer.ChainerModel(
model_data="{{ model_data }}",
role="{{ role }}",
entry_point="{{ entry_point }}",
source_dir="{{ source_dir }}",
image=None,
py_version="py3",
framework_version="5.0.0",
model_server_workers="{{ model_server_worker }}",
sagemaker_session=sagemaker_session,
)
config = airflow.deploy_config(
chainer_model,
initial_instance_count="{{ instance_count }}",
instance_type="ml.m4.xlarge")
expected_config = {
"Model": {
"ModelName": "sagemaker-chainer-%s" % TIME_STAMP,
"PrimaryContainer": {
"Image":
"520713654638.dkr.ecr.us-west-2.amazonaws.com/sagemaker-chainer:5.0.0-cpu-py3",
"Environment": {
"SAGEMAKER_PROGRAM":
"{{ entry_point }}",
"SAGEMAKER_SUBMIT_DIRECTORY":
"s3://output/sagemaker-chainer-%s/source/sourcedir.tar.gz"
% TIME_STAMP,
"SAGEMAKER_ENABLE_CLOUDWATCH_METRICS":
"false",
"SAGEMAKER_CONTAINER_LOG_LEVEL":
"20",
"SAGEMAKER_REGION":
"us-west-2",
"SAGEMAKER_MODEL_SERVER_WORKERS":
"{{ model_server_worker }}",
},
"ModelDataUrl": "{{ model_data }}",
},
"ExecutionRoleArn": "{{ role }}",
},
"EndpointConfig": {
"EndpointConfigName":
"sagemaker-chainer-%s" % TIME_STAMP,
"ProductionVariants": [{
"InstanceType": "ml.m4.xlarge",
"InitialInstanceCount": "{{ instance_count }}",
"ModelName": "sagemaker-chainer-%s" % TIME_STAMP,
"VariantName": "AllTraffic",
"InitialVariantWeight": 1,
}],
},
"Endpoint": {
"EndpointName": "sagemaker-chainer-%s" % TIME_STAMP,
"EndpointConfigName": "sagemaker-chainer-%s" % TIME_STAMP,
},
"S3Operations": {
"S3Upload": [{
"Path":
"{{ source_dir }}",
"Bucket":
"output",
"Key":
"sagemaker-chainer-%s/source/sourcedir.tar.gz" % TIME_STAMP,
"Tar":
True,
}]
},
}
assert config == expected_config
