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):