def setUp(self):
self.sagemaker = SageMakerTrainingOperator(
task_id='test_sagemaker_operator',
aws_conn_id='sagemaker_test_id',
config=create_training_params,
wait_for_completion=False,
check_interval=5)
class TestSageMakerTrainingOperator(unittest.TestCase):
def setUp(self):
configuration.load_test_config()
self.sagemaker = SageMakerTrainingOperator(
task_id='test_sagemaker_operator',
aws_conn_id='sagemaker_test_id',
config=create_training_params,
wait_for_completion=False,
check_interval=5)
def test_parse_config_integers(self):
self.sagemaker.parse_config_integers()
self.assertEqual(
self.sagemaker.config['ResourceConfig']['InstanceCount'],
int(self.sagemaker.config['ResourceConfig']['InstanceCount']))
self.assertEqual(
self.sagemaker.config['ResourceConfig']['VolumeSizeInGB'],
int(self.sagemaker.config['ResourceConfig']['VolumeSizeInGB']))
self.assertEqual(
self.sagemaker.config['StoppingCondition']['MaxRuntimeInSeconds'],
int(self.sagemaker.config['StoppingCondition']
['MaxRuntimeInSeconds']))
@mock.patch.object(SageMakerHook, 'get_conn')
@mock.patch.object(SageMakerHook, 'create_training_job')
def test_execute(self, mock_training, mock_client):
mock_training.return_value = {
'TrainingJobArn': 'testarn',
'ResponseMetadata': {
'HTTPStatusCode': 200
}
}
self.sagemaker.execute(None)
mock_training.assert_called_once_with(create_training_params,
wait_for_completion=False,
print_log=True,
check_interval=5,
max_ingestion_time=None)
@mock.patch.object(SageMakerHook, 'get_conn')
@mock.patch.object(SageMakerHook, 'create_training_job')
def test_execute_with_failure(self, mock_training, mock_client):
mock_training.return_value = {
'TrainingJobArn': 'testarn',
'ResponseMetadata': {
'HTTPStatusCode': 404
}
}
self.assertRaises(AirflowException, self.sagemaker.execute, None)
def setUp(self):
configuration.load_test_config()
self.sagemaker = SageMakerTrainingOperator(
task_id='test_sagemaker_operator',
aws_conn_id='sagemaker_test_id',
config=create_training_params,
wait_for_completion=False,
check_interval=5
)
class TestSageMakerTrainingOperator(unittest.TestCase):
def setUp(self):
configuration.load_test_config()
self.sagemaker = SageMakerTrainingOperator(
task_id='test_sagemaker_operator',
aws_conn_id='sagemaker_test_id',
config=create_training_params,
wait_for_completion=False,
check_interval=5
)
def test_parse_config_integers(self):
self.sagemaker.parse_config_integers()
self.assertEqual(self.sagemaker.config['ResourceConfig']['InstanceCount'],
int(self.sagemaker.config['ResourceConfig']['InstanceCount']))
self.assertEqual(self.sagemaker.config['ResourceConfig']['VolumeSizeInGB'],
int(self.sagemaker.config['ResourceConfig']['VolumeSizeInGB']))
self.assertEqual(self.sagemaker.config['StoppingCondition']['MaxRuntimeInSeconds'],
int(self.sagemaker.config['StoppingCondition']['MaxRuntimeInSeconds']))
@mock.patch.object(SageMakerHook, 'get_conn')
@mock.patch.object(SageMakerHook, 'create_training_job')
def test_execute(self, mock_training, mock_client):
mock_training.return_value = {'TrainingJobArn': 'testarn',
'ResponseMetadata':
{'HTTPStatusCode': 200}}
self.sagemaker.execute(None)
mock_training.assert_called_once_with(create_training_params,
wait_for_completion=False,
print_log=True,
check_interval=5,
max_ingestion_time=None
)
@mock.patch.object(SageMakerHook, 'get_conn')
@mock.patch.object(SageMakerHook, 'create_training_job')
def test_execute_with_failure(self, mock_training, mock_client):
mock_training.return_value = {'TrainingJobArn': 'testarn',
'ResponseMetadata':
{'HTTPStatusCode': 404}}
self.assertRaises(AirflowException, self.sagemaker.execute, None)
def test_pytorch_airflow_config_uploads_data_source_to_s3_when_inputs_not_provided(
sagemaker_session, cpu_instance_type):
with timeout(seconds=AIRFLOW_CONFIG_TIMEOUT_IN_SECONDS):
estimator = PyTorch(
entry_point=PYTORCH_MNIST_SCRIPT,
role=ROLE,
framework_version="1.1.0",
train_instance_count=2,
train_instance_type=cpu_instance_type,
hyperparameters={
"epochs": 6,
"backend": "gloo"
},
)
train_config = sm_airflow.training_config(estimator=estimator)
uploaded_s3_data = train_config["HyperParameters"][
"sagemaker_submit_directory"].strip('"')
transform_config = sm_airflow.transform_config_from_estimator(
estimator=estimator,
task_id="transform_config",
task_type="training",
instance_count=SINGLE_INSTANCE_COUNT,
instance_type=cpu_instance_type,
data=uploaded_s3_data,
content_type="text/csv",
)
default_args = {
"owner": "airflow",
"start_date": airflow.utils.dates.days_ago(2),
"provide_context": True,
}
dag = DAG("tensorflow_example",
default_args=default_args,
schedule_interval="@once")
train_op = SageMakerTrainingOperator(task_id="tf_training",
config=train_config,
wait_for_completion=True,
dag=dag)
transform_op = SageMakerTransformOperator(task_id="transform_operator",
config=transform_config,
wait_for_completion=True,
dag=dag)
transform_op.set_upstream(train_op)
_assert_that_s3_url_contains_data(sagemaker_session, uploaded_s3_data)
def _build_airflow_workflow(estimator,
instance_type,
inputs=None,
mini_batch_size=None):
training_config = sm_airflow.training_config(
estimator=estimator, inputs=inputs, mini_batch_size=mini_batch_size)
model = estimator.create_model()
assert model is not None
model_config = sm_airflow.model_config(instance_type, model)
assert model_config is not None
transform_config = sm_airflow.transform_config_from_estimator(
estimator=estimator,
task_id="transform_config",
task_type="training",
instance_count=SINGLE_INSTANCE_COUNT,
instance_type=estimator.train_instance_type,
data=inputs,
content_type="text/csv",
input_filter="$",
output_filter="$",
)
default_args = {
"owner": "airflow",
"start_date": airflow.utils.dates.days_ago(2),
"provide_context": True,
}
dag = DAG("tensorflow_example",
default_args=default_args,
schedule_interval="@once")
train_op = SageMakerTrainingOperator(task_id="tf_training",
config=training_config,
wait_for_completion=True,
dag=dag)
transform_op = SageMakerTransformOperator(task_id="transform_operator",
config=transform_config,
wait_for_completion=True,
dag=dag)
transform_op.set_upstream(train_op)
return training_config
# prepare the data for training
prepare_task = PythonOperator(task_id='preparing',
dag=dag,
provide_context=False,
python_callable=prepare.prepare,
op_kwargs=config["prepare_data"])
branching = BranchPythonOperator(task_id='branching',
dag=dag,
python_callable=lambda: "model_tuning"
if hpo_enabled else "model_training")
# launch sagemaker training job and wait until it completes
train_model_task = SageMakerTrainingOperator(task_id='model_training',
dag=dag,
config=train_config,
aws_conn_id='airflow-sagemaker',
wait_for_completion=True,
check_interval=30)
# launch sagemaker hyperparameter job and wait until it completes
tune_model_task = SageMakerTuningOperator(task_id='model_tuning',
dag=dag,
config=tuner_config,
aws_conn_id='airflow-sagemaker',
wait_for_completion=True,
check_interval=30)
# launch sagemaker batch transform job and wait until it completes
batch_transform_task = SageMakerTransformOperator(
task_id='predicting',
dag=dag,
task_id='sm_proc_job',
dag=dag,
provide_context=True,
python_callable=sm_proc_job.sm_proc_job,
op_kwargs={
'role': role,
'sess': sess,
'bucket': bucket,
'spark_repo_uri': config["processing_job"]['spark_repo_uri'],
'base_job_name': config["processing_job"]["base_job_name"]
})
# Train xgboost model task
train_model_task = SageMakerTrainingOperator(task_id='xgboost_model_training',
dag=dag,
config=train_config,
aws_conn_id='airflow-sagemaker',
wait_for_completion=True,
check_interval=30)
# Inference pipeline endpoint task
inference_pipeline_task = PythonOperator(
task_id='inference_pipeline',
dag=dag,
python_callable=inference_pipeline_ep.inference_pipeline_ep,
op_kwargs={
'role': role,
'sess': sess,
'spark_model_uri':
config['inference_pipeline']['inputs']['spark_model'],
'pipeline_model_name':
config['inference_pipeline']['pipeline_model_name'],
# define airflow DAG and tasks
# =============================================================================
# define airflow DAG
default_args = {
'owner': 'airflow',
'start_date': airflow.utils.dates.days_ago(2),
'provide_context': True
}
dag = DAG(dag_id='train_dkn',
default_args=default_args,
schedule_interval='@once')
train_op = SageMakerTrainingOperator(task_id='tf_training',
config=train_dkn_config,
wait_for_completion=True,
dag=dag)
task_def_op = PythonOperator(task_id='task_definition',
python_callable=task_def,
op_args=['gw1'],
provide_context=True,
dag=dag)
deploy_ecs_op = PythonOperator(task_id='run_task',
python_callable=deploy_model_service,
op_args=['gw1'],
provide_context=True,
dag=dag)
deploy_ecs_op.set_upstream(task_def_op)
# prepare the data for training
prepare_task = PythonOperator(task_id='preparing',
dag=dag,
provide_context=False,
python_callable=prepare.prepare,
op_kwargs=config["prepare_data"])
branching = BranchPythonOperator(task_id='branching',
dag=dag,
python_callable=lambda: "model_tuning"
if hpo_enabled else "model_training")
# launch sagemaker training job and wait until it completes
train_model_task = SageMakerTrainingOperator(task_id='model_training',
dag=dag,
config=train_config,
aws_conn_id='airflow-sagemaker',
wait_for_completion=True,
check_interval=30)
# launch sagemaker hyperparameter job and wait until it completes
tune_model_task = SageMakerTuningOperator(task_id='model_tuning',
dag=dag,
config=tuner_config,
aws_conn_id='airflow-sagemaker',
wait_for_completion=True,
check_interval=30)
# launch sagemaker batch transform job and wait until it completes
batch_transform_task = SageMakerTransformOperator(
task_id='predicting',
dag=dag,
sagemaker_train_model = SageMakerTrainingOperator(
task_id="sagemaker_train_model",
config={
"TrainingJobName":
"mnistclassifier-{{ execution_date.strftime('%Y-%m-%d-%H-%M-%S') }}",
"AlgorithmSpecification": {
"TrainingImage":
"438346466558.dkr.ecr.eu-west-1.amazonaws.com/kmeans:1",
"TrainingInputMode": "File",
},
"HyperParameters": {
"k": "10",
"feature_dim": "784"
},
"InputDataConfig": [{
"ChannelName": "train",
"DataSource": {
"S3DataSource": {
"S3DataType": "S3Prefix",
"S3Uri": "s3://your-bucket/mnist_data",
"S3DataDistributionType": "FullyReplicated",
}
},
}],
"OutputDataConfig": {
"S3OutputPath": "s3://your-bucket/mnistclassifier-output"
},
"ResourceConfig": {
"InstanceType": "ml.c4.xlarge",
"InstanceCount": 1,
"VolumeSizeInGB": 10,
},
"RoleArn": ("arn:aws:iam::297623009465:role/service-role/"
"AmazonSageMaker-ExecutionRole-20180905T153196"),
"StoppingCondition": {
"MaxRuntimeInSeconds": 24 * 60 * 60
},
},
wait_for_completion=True,
print_log=True,
check_interval=10,
dag=dag,
)
schedule_interval=None,
concurrency=1,
max_active_runs=1,
) as dag:
process_task = PythonOperator(
task_id="process",
dag=dag,
#provide_context=False,
python_callable=preprocess_glue,
)
train_task = SageMakerTrainingOperator(
task_id="train",
config=training_config,
aws_conn_id="airflow-sagemaker",
wait_for_completion=True,
check_interval=60, #check status of the job every minute
max_ingestion_time=
None, #allow training job to run as long as it needs, change for early stop
)
endpoint_deploy_task = SageMakerEndpointOperator(
task_id="endpoint-deploy",
config=endpoint_config,
aws_conn_id="sagemaker-airflow",
wait_for_completion=True,
check_interval=60, #check status of endpoint deployment every minute
max_ingestion_time=None,
operation=
'create', #change to update if you are updating rather than creating an endpoint
)