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_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.model_config(instance_type="ml.c4.xlarge", model=chainer_model) expected_config = { "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 }}", "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_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.model_config(instance_type='ml.c4.xlarge', model=chainer_model) expected_config = { 'ModelName': "sagemaker-chainer-{{ execution_date.strftime('%Y-%m-%d-%H-%M-%S') }}", '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-" "{{ execution_date.strftime('%Y-%m-%d-%H-%M-%S') }}" "/source/sourcedir.tar.gz", '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 }}', 'S3Operations': { 'S3Upload': [{ 'Path': '{{ source_dir }}', 'Bucket': 'output', 'Key': "sagemaker-chainer-{{ execution_date.strftime('%Y-%m-%d-%H-%M-%S') }}/source/sourcedir.tar.gz", 'Tar': True }] } } assert config == expected_config