def test_transform_config_from_amazon_alg_estimator(sagemaker_session): knn_estimator = knn.KNN( role="{{ role }}", train_instance_count="{{ instance_count }}", train_instance_type="ml.m4.xlarge", k=16, sample_size=128, predictor_type="regressor", sagemaker_session=sagemaker_session, ) record = amazon_estimator.RecordSet("{{ record }}", 10000, 100, "S3Prefix") transform_data = "{{ transform_data }}" # simulate training airflow.training_config(knn_estimator, record, mini_batch_size=256) config = airflow.transform_config_from_estimator( estimator=knn_estimator, task_id="task_id", task_type="training", instance_count="{{ instance_count }}", instance_type="ml.p2.xlarge", data=transform_data, ) expected_config = { "Model": { "ModelName": "knn-%s" % TIME_STAMP, "PrimaryContainer": { "Image": "174872318107.dkr.ecr.us-west-2.amazonaws.com/knn:1", "Environment": {}, "ModelDataUrl": "s3://output/{{ ti.xcom_pull(task_ids='task_id')['Training']['TrainingJobName'] }}" "/output/model.tar.gz", }, "ExecutionRoleArn": "{{ role }}", }, "Transform": { "TransformJobName": "knn-%s" % TIME_STAMP, "ModelName": "knn-%s" % TIME_STAMP, "TransformInput": { "DataSource": { "S3DataSource": { "S3DataType": "S3Prefix", "S3Uri": "{{ transform_data }}" } } }, "TransformOutput": { "S3OutputPath": "s3://output/knn-%s" % TIME_STAMP }, "TransformResources": { "InstanceCount": "{{ instance_count }}", "InstanceType": "ml.p2.xlarge", }, }, } assert config == expected_config
def test_transform_config_from_amazon_alg_estimator(sagemaker_session): knn_estimator = knn.KNN(role="{{ role }}", train_instance_count="{{ instance_count }}", train_instance_type='ml.m4.xlarge', k=16, sample_size=128, predictor_type='regressor', sagemaker_session=sagemaker_session) record = amazon_estimator.RecordSet("{{ record }}", 10000, 100, 'S3Prefix') transform_data = "{{ transform_data }}" # simulate training airflow.training_config(knn_estimator, record, mini_batch_size=256) config = airflow.transform_config_from_estimator( estimator=knn_estimator, instance_count="{{ instance_count }}", instance_type="ml.p2.xlarge", data=transform_data) expected_config = { 'Model': { 'ModelName': "knn-{{ execution_date.strftime('%Y-%m-%d-%H-%M-%S') }}", 'PrimaryContainer': { 'Image': '174872318107.dkr.ecr.us-west-2.amazonaws.com/knn:1', 'Environment': {}, 'ModelDataUrl': "s3://output/knn-{{ execution_date.strftime('%Y-%m-%d-%H-%M-%S') }}" "/output/model.tar.gz" }, 'ExecutionRoleArn': '{{ role }}' }, 'Transform': { 'TransformJobName': "knn-{{ execution_date.strftime('%Y-%m-%d-%H-%M-%S') }}", 'ModelName': "knn-{{ execution_date.strftime('%Y-%m-%d-%H-%M-%S') }}", 'TransformInput': { 'DataSource': { 'S3DataSource': { 'S3DataType': 'S3Prefix', 'S3Uri': '{{ transform_data }}' } } }, 'TransformOutput': { 'S3OutputPath': "s3://output/knn-{{ execution_date.strftime('%Y-%m-%d-%H-%M-%S') }}" }, 'TransformResources': { 'InstanceCount': '{{ instance_count }}', 'InstanceType': 'ml.p2.xlarge' } } } assert config == expected_config
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
# train_config specifies SageMaker training configuration train_config = training_config(estimator=fm_estimator, inputs=config["train_model"]["inputs"]) # create tuner fm_tuner = HyperparameterTuner(estimator=fm_estimator, **config["tune_model"]["tuner_config"]) # create tuning config tuner_config = tuning_config(tuner=fm_tuner, inputs=config["tune_model"]["inputs"]) # create transform config transform_config = transform_config_from_estimator( estimator=fm_estimator, task_id="model_tuning" if hpo_enabled else "model_training", task_type="tuning" if hpo_enabled else "training", **config["batch_transform"]["transform_config"]) # ============================================================================= # define airflow DAG and tasks # ============================================================================= # define airflow DAG args = {'owner': 'airflow', 'start_date': airflow.utils.dates.days_ago(2)} dag = DAG( dag_id='sagemaker-ml-pipeline', default_args=args, schedule_interval=None,
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
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