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