def test_deploy_config_from_framework_estimator(sagemaker_session):
mxnet_estimator = mxnet.MXNet(
entry_point="{{ entry_point }}",
source_dir="{{ source_dir }}",
py_version='py3',
framework_version='1.3.0',
role="{{ role }}",
train_instance_count=1,
train_instance_type='ml.m4.xlarge',
sagemaker_session=sagemaker_session,
base_job_name="{{ base_job_name }}",
hyperparameters={'batch_size': 100})
train_data = "{{ train_data }}"
# simulate training
airflow.training_config(mxnet_estimator, train_data)
config = airflow.deploy_config_from_estimator(estimator=mxnet_estimator,
task_id='task_id',
task_type='training',
initial_instance_count="{{ instance_count}}",
instance_type="ml.c4.large",
endpoint_name="mxnet-endpoint")
expected_config = {
'Model': {
'ModelName': "sagemaker-mxnet-%s" % TIME_STAMP,
'PrimaryContainer': {
'Image': '520713654638.dkr.ecr.us-west-2.amazonaws.com/sagemaker-mxnet:1.3.0-cpu-py3',
'Environment': {
'SAGEMAKER_PROGRAM': '{{ entry_point }}',
'SAGEMAKER_SUBMIT_DIRECTORY': "s3://output/{{ ti.xcom_pull(task_ids='task_id')['Training']"
"['TrainingJobName'] }}/source/sourcedir.tar.gz",
'SAGEMAKER_ENABLE_CLOUDWATCH_METRICS': 'false',
'SAGEMAKER_CONTAINER_LOG_LEVEL': '20',
'SAGEMAKER_REGION': 'us-west-2'},
'ModelDataUrl': "s3://output/{{ ti.xcom_pull(task_ids='task_id')['Training']['TrainingJobName'] }}"
"/output/model.tar.gz"
},
'ExecutionRoleArn': '{{ role }}'
},
'EndpointConfig': {
'EndpointConfigName': "sagemaker-mxnet-%s" % TIME_STAMP,
'ProductionVariants': [{
'InstanceType': 'ml.c4.large',
'InitialInstanceCount': '{{ instance_count}}',
'ModelName': "sagemaker-mxnet-%s" % TIME_STAMP,
'VariantName': 'AllTraffic',
'InitialVariantWeight': 1
}]
},
'Endpoint': {
'EndpointName': 'mxnet-endpoint',
'EndpointConfigName': "sagemaker-mxnet-%s" % TIME_STAMP
}
}
assert config == expected_config
def test_model_config_from_framework_estimator(sagemaker_session):
mxnet_estimator = mxnet.MXNet(
entry_point="{{ entry_point }}",
source_dir="{{ source_dir }}",
py_version="py3",
framework_version="1.3.0",
role="{{ role }}",
train_instance_count=1,
train_instance_type="ml.m4.xlarge",
sagemaker_session=sagemaker_session,
base_job_name="{{ base_job_name }}",
hyperparameters={"batch_size": 100},
)
data = "{{ training_data }}"
# simulate training
airflow.training_config(mxnet_estimator, data)
config = airflow.model_config_from_estimator(
instance_type="ml.c4.xlarge",
estimator=mxnet_estimator,
task_id="task_id",
task_type="training",
)
expected_config = {
"ModelName": "sagemaker-mxnet-%s" % TIME_STAMP,
"PrimaryContainer": {
"Image":
"520713654638.dkr.ecr.us-west-2.amazonaws.com/sagemaker-mxnet:1.3.0-cpu-py3",
"Environment": {
"SAGEMAKER_PROGRAM":
"{{ entry_point }}",
"SAGEMAKER_SUBMIT_DIRECTORY":
"s3://output/{{ ti.xcom_pull(task_ids='task_id')['Training']"
"['TrainingJobName'] }}/source/sourcedir.tar.gz",
"SAGEMAKER_ENABLE_CLOUDWATCH_METRICS":
"false",
"SAGEMAKER_CONTAINER_LOG_LEVEL":
"20",
"SAGEMAKER_REGION":
"us-west-2",
},
"ModelDataUrl":
"s3://output/{{ ti.xcom_pull(task_ids='task_id')['Training']['TrainingJobName'] }}"
"/output/model.tar.gz",
},
"ExecutionRoleArn": "{{ role }}",
}
assert config == expected_config
def test_model_config_from_framework_estimator(sagemaker_session):
mxnet_estimator = mxnet.MXNet(entry_point="{{ entry_point }}",
source_dir="{{ source_dir }}",
py_version='py3',
framework_version='1.3.0',
role="{{ role }}",
train_instance_count=1,
train_instance_type='ml.m4.xlarge',
sagemaker_session=sagemaker_session,
base_job_name="{{ base_job_name }}",
hyperparameters={'batch_size': 100})
data = "{{ training_data }}"
# simulate training
airflow.training_config(mxnet_estimator, data)
config = airflow.model_config_from_estimator(instance_type='ml.c4.xlarge',
estimator=mxnet_estimator)
expected_config = {
'ModelName':
"{{ base_job_name }}-{{ execution_date.strftime('%Y-%m-%d-%H-%M-%S') }}",
'PrimaryContainer': {
'Image':
'520713654638.dkr.ecr.us-west-2.amazonaws.com/sagemaker-mxnet:1.3.0-cpu-py3',
'Environment': {
'SAGEMAKER_PROGRAM':
'{{ entry_point }}',
'SAGEMAKER_SUBMIT_DIRECTORY':
"s3://output/{{ base_job_name }}-"
"{{ 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'
},
'ModelDataUrl':
"s3://output/{{ base_job_name }}-"
"{{ execution_date.strftime('%Y-%m-%d-%H-%M-%S') }}/output/model.tar.gz"
},
'ExecutionRoleArn': '{{ role }}'
}
assert config == expected_config
def test_framework_tuning_config(sagemaker_session):
mxnet_estimator = mxnet.MXNet(entry_point="{{ entry_point }}",
source_dir="{{ source_dir }}",
py_version='py3',
framework_version='1.3.0',
role="{{ role }}",
train_instance_count=1,
train_instance_type='ml.m4.xlarge',
sagemaker_session=sagemaker_session,
base_job_name="{{ base_job_name }}",
hyperparameters={'batch_size': 100})
hyperparameter_ranges = {
'optimizer': tuner.CategoricalParameter(['sgd', 'Adam']),
'learning_rate': tuner.ContinuousParameter(0.01, 0.2),
'num_epoch': tuner.IntegerParameter(10, 50)
}
objective_metric_name = 'Validation-accuracy'
metric_definitions = [{
'Name': 'Validation-accuracy',
'Regex': 'Validation-accuracy=([0-9\\.]+)'
}]
mxnet_tuner = tuner.HyperparameterTuner(
estimator=mxnet_estimator,
objective_metric_name=objective_metric_name,
hyperparameter_ranges=hyperparameter_ranges,
metric_definitions=metric_definitions,
strategy='Bayesian',
objective_type='Maximize',
max_jobs="{{ max_job }}",
max_parallel_jobs="{{ max_parallel_job }}",
tags=[{
'{{ key }}': '{{ value }}'
}],
base_tuning_job_name="{{ base_job_name }}")
data = "{{ training_data }}"
config = airflow.tuning_config(mxnet_tuner, data)
expected_config = {
'HyperParameterTuningJobName':
"{{ base_job_name }}-{{ execution_date.strftime('%y%m%d-%H%M') }}",
'HyperParameterTuningJobConfig': {
'Strategy': 'Bayesian',
'HyperParameterTuningJobObjective': {
'Type': 'Maximize',
'MetricName': 'Validation-accuracy'
},
'ResourceLimits': {
'MaxNumberOfTrainingJobs': '{{ max_job }}',
'MaxParallelTrainingJobs': '{{ max_parallel_job }}'
},
'ParameterRanges': {
'ContinuousParameterRanges': [{
'Name': 'learning_rate',
'MinValue': '0.01',
'MaxValue': '0.2'
}],
'CategoricalParameterRanges': [{
'Name': 'optimizer',
'Values': ['"sgd"', '"Adam"']
}],
'IntegerParameterRanges': [{
'Name': 'num_epoch',
'MinValue': '10',
'MaxValue': '50'
}]
}
},
'TrainingJobDefinition': {
'AlgorithmSpecification': {
'TrainingImage':
'520713654638.dkr.ecr.us-west-2.amazonaws.com/sagemaker-mxnet:1.3.0-cpu-py3',
'TrainingInputMode':
'File',
'MetricDefinitions': [{
'Name':
'Validation-accuracy',
'Regex':
'Validation-accuracy=([0-9\\.]+)'
}]
},
'OutputDataConfig': {
'S3OutputPath': 's3://output/'
},
'StoppingCondition': {
'MaxRuntimeInSeconds': 86400
},
'ResourceConfig': {
'InstanceCount': 1,
'InstanceType': 'ml.m4.xlarge',
'VolumeSizeInGB': 30
},
'RoleArn':
'{{ role }}',
'InputDataConfig': [{
'DataSource': {
'S3DataSource': {
'S3DataDistributionType': 'FullyReplicated',
'S3DataType': 'S3Prefix',
'S3Uri': '{{ training_data }}'
}
},
'ChannelName': 'training'
}],
'StaticHyperParameters': {
'batch_size':
'100',
'sagemaker_submit_directory':
'"s3://output/{{ base_job_name }}'
'-{{ execution_date.strftime(\'%Y-%m-%d-%H-%M-%S\') }}'
'/source/sourcedir.tar.gz"',
'sagemaker_program':
'"{{ entry_point }}"',
'sagemaker_enable_cloudwatch_metrics':
'false',
'sagemaker_container_log_level':
'20',
'sagemaker_job_name':
'"{{ base_job_name }}-'
'{{ execution_date.strftime(\'%Y-%m-%d-%H-%M-%S\') }}"',
'sagemaker_region':
'"us-west-2"'
}
},
'Tags': [{
'{{ key }}': '{{ value }}'
}],
'S3Operations': {
'S3Upload': [{
'Path': '{{ source_dir }}',
'Bucket': 'output',
'Key': "{{ base_job_name }}-"
"{{ execution_date.strftime('%Y-%m-%d-%H-%M-%S') }}/source/sourcedir.tar.gz",
'Tar': True
}]
}
}
assert config == expected_config
def test_framework_tuning_config(sagemaker_session):
mxnet_estimator = mxnet.MXNet(
entry_point="{{ entry_point }}",
source_dir="{{ source_dir }}",
py_version="py3",
framework_version="1.3.0",
role="{{ role }}",
train_instance_count=1,
train_instance_type="ml.m4.xlarge",
sagemaker_session=sagemaker_session,
base_job_name="{{ base_job_name }}",
hyperparameters={"batch_size": 100},
)
hyperparameter_ranges = {
"optimizer": tuner.CategoricalParameter(["sgd", "Adam"]),
"learning_rate": tuner.ContinuousParameter(0.01, 0.2),
"num_epoch": tuner.IntegerParameter(10, 50),
}
objective_metric_name = "Validation-accuracy"
metric_definitions = [{
"Name": "Validation-accuracy",
"Regex": "Validation-accuracy=([0-9\\.]+)"
}]
mxnet_tuner = tuner.HyperparameterTuner(
estimator=mxnet_estimator,
objective_metric_name=objective_metric_name,
hyperparameter_ranges=hyperparameter_ranges,
metric_definitions=metric_definitions,
strategy="Bayesian",
objective_type="Maximize",
max_jobs="{{ max_job }}",
max_parallel_jobs="{{ max_parallel_job }}",
tags=[{
"{{ key }}": "{{ value }}"
}],
base_tuning_job_name="{{ base_job_name }}",
)
data = "{{ training_data }}"
config = airflow.tuning_config(mxnet_tuner, data)
expected_config = {
"HyperParameterTuningJobName": "{{ base_job_name }}-%s" % TIME_STAMP,
"HyperParameterTuningJobConfig": {
"Strategy": "Bayesian",
"HyperParameterTuningJobObjective": {
"Type": "Maximize",
"MetricName": "Validation-accuracy",
},
"ResourceLimits": {
"MaxNumberOfTrainingJobs": "{{ max_job }}",
"MaxParallelTrainingJobs": "{{ max_parallel_job }}",
},
"ParameterRanges": {
"ContinuousParameterRanges": [{
"Name": "learning_rate",
"MinValue": "0.01",
"MaxValue": "0.2",
"ScalingType": "Auto",
}],
"CategoricalParameterRanges": [{
"Name": "optimizer",
"Values": ['"sgd"', '"Adam"']
}],
"IntegerParameterRanges": [{
"Name": "num_epoch",
"MinValue": "10",
"MaxValue": "50",
"ScalingType": "Auto"
}],
},
},
"TrainingJobDefinition": {
"AlgorithmSpecification": {
"TrainingImage":
"520713654638.dkr.ecr.us-west-2.amazonaws.com/sagemaker-mxnet:1.3.0-cpu-py3",
"TrainingInputMode":
"File",
"MetricDefinitions": [{
"Name":
"Validation-accuracy",
"Regex":
"Validation-accuracy=([0-9\\.]+)"
}],
},
"OutputDataConfig": {
"S3OutputPath": "s3://output/"
},
"StoppingCondition": {
"MaxRuntimeInSeconds": 86400
},
"ResourceConfig": {
"InstanceCount": 1,
"InstanceType": "ml.m4.xlarge",
"VolumeSizeInGB": 30,
},
"RoleArn":
"{{ role }}",
"InputDataConfig": [{
"DataSource": {
"S3DataSource": {
"S3DataDistributionType": "FullyReplicated",
"S3DataType": "S3Prefix",
"S3Uri": "{{ training_data }}",
}
},
"ChannelName": "training",
}],
"StaticHyperParameters": {
"batch_size":
"100",
"sagemaker_submit_directory":
'"s3://output/{{ base_job_name }}-%s/source/sourcedir.tar.gz"'
% TIME_STAMP,
"sagemaker_program":
'"{{ entry_point }}"',
"sagemaker_enable_cloudwatch_metrics":
"false",
"sagemaker_container_log_level":
"20",
"sagemaker_job_name":
'"{{ base_job_name }}-%s"' % TIME_STAMP,
"sagemaker_region":
'"us-west-2"',
},
},
"Tags": [{
"{{ key }}": "{{ value }}"
}],
"S3Operations": {
"S3Upload": [{
"Path":
"{{ source_dir }}",
"Bucket":
"output",
"Key":
"{{ base_job_name }}-%s/source/sourcedir.tar.gz" % TIME_STAMP,
"Tar":
True,
}]
},
}
assert config == expected_config
def test_transform_config_from_framework_estimator(sagemaker_session):
mxnet_estimator = mxnet.MXNet(entry_point="{{ entry_point }}",
source_dir="{{ source_dir }}",
py_version='py3',
framework_version='1.3.0',
role="{{ role }}",
train_instance_count=1,
train_instance_type='ml.m4.xlarge',
sagemaker_session=sagemaker_session,
base_job_name="{{ base_job_name }}",
hyperparameters={'batch_size': 100})
train_data = "{{ train_data }}"
transform_data = "{{ transform_data }}"
# simulate training
airflow.training_config(mxnet_estimator, train_data)
config = airflow.transform_config_from_estimator(
estimator=mxnet_estimator,
task_id='task_id',
task_type='training',
instance_count="{{ instance_count }}",
instance_type="ml.p2.xlarge",
data=transform_data)
expected_config = {
'Model': {
'ModelName': "sagemaker-mxnet-%s" % TIME_STAMP,
'PrimaryContainer': {
'Image':
'520713654638.dkr.ecr.us-west-2.amazonaws.com/sagemaker-mxnet:1.3.0-gpu-py3',
'Environment': {
'SAGEMAKER_PROGRAM':
'{{ entry_point }}',
'SAGEMAKER_SUBMIT_DIRECTORY':
"s3://output/{{ ti.xcom_pull(task_ids='task_id')"
"['Training']['TrainingJobName'] }}"
"/source/sourcedir.tar.gz",
'SAGEMAKER_ENABLE_CLOUDWATCH_METRICS':
'false',
'SAGEMAKER_CONTAINER_LOG_LEVEL':
'20',
'SAGEMAKER_REGION':
'us-west-2'
},
'ModelDataUrl':
"s3://output/{{ ti.xcom_pull(task_ids='task_id')['Training']['TrainingJobName'] }}"
"/output/model.tar.gz"
},
'ExecutionRoleArn': '{{ role }}'
},
'Transform': {
'TransformJobName': "{{ base_job_name }}-%s" % TIME_STAMP,
'ModelName': "sagemaker-mxnet-%s" % TIME_STAMP,
'TransformInput': {
'DataSource': {
'S3DataSource': {
'S3DataType': 'S3Prefix',
'S3Uri': '{{ transform_data }}'
}
}
},
'TransformOutput': {
'S3OutputPath':
"s3://output/{{ base_job_name }}-%s" % TIME_STAMP
},
'TransformResources': {
'InstanceCount': '{{ instance_count }}',
'InstanceType': 'ml.p2.xlarge'
},
'Environment': {}
}
}
assert config == expected_config
