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