예제 #1
0
def parse_args():
    parser = argparse.ArgumentParser()
    env_vars = docker_rf_utils.get_env_vars('CURW_')

    def check_key(k, d_val):
        if k in env_vars and not env_vars[k]:
            return env_vars[k]
        else:
            return d_val

    parser.add_argument('-run_id',
                        default=check_key('run_id',
                                          docker_rf_utils.id_generator()))
    parser.add_argument('-db_config',
                        default=check_key(
                            'db_config',
                            docker_rf_utils.get_base64_encoded_str('{}')))
    parser.add_argument('-wrf_config',
                        default=check_key(
                            'wrf_config',
                            docker_rf_utils.get_base64_encoded_str('{}')))
    parser.add_argument('-overwrite', default=check_key('overwrite', 'False'))
    parser.add_argument('-data_type', default=check_key('data_type', 'data'))
    parser.add_argument('-procedures',
                        default=check_key('procedures', str(sys.maxsize)))

    return parser.parse_args()
def initialize_config(config_str, configs_prefix='wrf_config_', **context):
    config = json.loads(config_str)

    for k, v in context['templates_dict'].items():
        if k.startswith(configs_prefix):
            config[k.replace(configs_prefix, '')] = v

    config_str = json.dumps(config, sort_keys=True)
    logging.info('Initialized wrf_config:\n' + config_str)

    b64_config = docker_rf_utils.get_base64_encoded_str(config_str)
    logging.info('Returned xcom:\n' + b64_config)
    return b64_config
clean_up = PythonOperator(
    task_id='cleanup-config',
    python_callable=clean_up_wrf_run,
    provide_context=True,
    op_args=['init-config'],
    dag=dag,
)

wps = CurwDockerOperator(
    task_id='wps',
    image=wrf_image,
    docker_url=docker_url,
    command=airflow_docker_utils.get_docker_cmd(
        '{{ task_instance.xcom_pull(task_ids=\'gen-run-id\') }}',
        '{{ task_instance.xcom_pull(task_ids=\'init-config\') }}', 'wps',
        docker_rf_utils.get_base64_encoded_str(airflow_vars[nl_wps_key]),
        docker_rf_utils.get_base64_encoded_str(airflow_vars[nl_inp_keys[0]]),
        gcs_volumes),
    cpus=1,
    volumes=docker_volumes,
    auto_remove=True,
    privileged=True,
    dag=dag,
    pool=wrf_pool,
)

wrf = CurwDockerOperator(
    task_id='wrf',
    image=wrf_image,
    docker_url=docker_url,
    command=airflow_docker_utils.get_docker_cmd(
clean_up = PythonOperator(
    task_id='cleanup-config',
    python_callable=clean_up_wrf_run,
    provide_context=True,
    op_args=['init-config'],
    dag=dag,
)

wps = CurwDockerOperator(
    task_id='wps',
    image=wrf_image,
    docker_url=docker_url,
    command=airflow_docker_utils.get_docker_cmd(
        '{{ task_instance.xcom_pull(task_ids=\'gen-run-id\') }}',
        '{{ task_instance.xcom_pull(task_ids=\'init-config\') }}', 'wps',
        docker_rf_utils.get_base64_encoded_str(airflow_vars[nl_wps_key]),
        docker_rf_utils.get_base64_encoded_str(airflow_vars[nl_inp_keys[0]]),
        gcs_volumes),
    cpus=1,
    volumes=docker_volumes,
    auto_remove=True,
    privileged=True,
    dag=dag,
    pool=wrf_pool,
)

generate_run_id >> init_config >> wps

select_wrf = DummyOperator(task_id='select-wrf', dag=dag)

for i in range(parallel_runs):