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_amazon_alg_model_config(sagemaker_session):
pca_model = pca.PCAModel(model_data="{{ model_data }}",
role="{{ role }}",
sagemaker_session=sagemaker_session)
config = airflow.model_config(instance_type="ml.c4.xlarge",
model=pca_model)
expected_config = {
"ModelName": "pca-%s" % TIME_STAMP,
"PrimaryContainer": {
"Image": "174872318107.dkr.ecr.us-west-2.amazonaws.com/pca:1",
"Environment": {},
"ModelDataUrl": "{{ model_data }}",
},
"ExecutionRoleArn": "{{ role }}",
}
assert config == expected_config
def test_amazon_alg_model_config(sagemaker_session):
pca_model = pca.PCAModel(model_data="{{ model_data }}",
role="{{ role }}",
sagemaker_session=sagemaker_session)
config = airflow.model_config(instance_type='ml.c4.xlarge',
model=pca_model)
expected_config = {
'ModelName': "pca-%s" % TIME_STAMP,
'PrimaryContainer': {
'Image': '174872318107.dkr.ecr.us-west-2.amazonaws.com/pca:1',
'Environment': {},
'ModelDataUrl': '{{ model_data }}'
},
'ExecutionRoleArn': '{{ role }}'
}
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
