class WrfTask(CurwTask):
def __init__(self, wrf_config_key='wrf_config', **kwargs):
self.wrf_config_key = wrf_config_key
super(WrfTask, self).__init__(**kwargs)
def set_config(self, **kwargs):
wrf_config_json = None
if 'ti' in kwargs:
wrf_config_json = kwargs['ti'].xcom_pull(task_ids=None,
key=self.wrf_config_key)
if wrf_config_json is not None:
logging.info('wrf_config from xcom using %s key: %s' %
(self.wrf_config_key, wrf_config_json))
self.config = WrfConfig(json.loads(wrf_config_json))
else:
try:
self.config = WrfConfig(
Variable.get(self.wrf_config_key, deserialize_json=True))
logging.info('wrf_config from variable using %s key: %s' %
(self.wrf_config_key, self.config.to_json_string))
except KeyError:
raise CurwAriflowTasksException('Unable to find WrfConfig')
def add_config_item(self, key, value):
self.config.set(key, value)
Variable.set(self.wrf_config_key, self.config.to_json_string())
def acquire_wrf_lock(wrf_config_key):
config = WrfConfig(Variable.get(wrf_config_key, deserialize_json=True))
file_path = os.path.join(utils.get_em_real_dir(config.get('wrf_home')),
'wrf.lock')
logging.info('acquiring lock %s' % file_path)
with open(file_path, 'w') as lock:
lock.write(config.to_string())
def get_wrf_config(**kwargs):
if 'ti' in kwargs:
wrf_config_json = kwargs['ti'].xcom_pull(task_ids=None,
key='wrf_config_json')
logging.info('wrf_config from xcom: ' + wrf_config_json)
wrf_config = WrfConfig(json.loads(wrf_config_json))
else:
try:
wrf_config = WrfConfig(
Variable.get('wrf_config', deserialize_json=True))
except KeyError:
raise RunUngribTaskException('Unable to find WrfConfig')
return wrf_config
def get_gfs_download_subdag(parent_dag_name,
child_dag_name,
args,
wrf_config_key='wrf_config',
test_mode=False):
dag_subdag = DAG(
dag_id='%s.%s' % (parent_dag_name, child_dag_name),
default_args=args,
schedule_interval=None,
)
try:
wrf_config = WrfConfig(
configs=Variable.get(wrf_config_key, deserialize_json=True))
except KeyError as e:
logging.error('Key error %s' % str(e))
return dag_subdag
period = wrf_config.get('period')
step = wrf_config.get('gfs_step')
gfs_dir = wrf_config.get('gfs_dir')
try:
gfs_date, gfs_cycle, start = utils.get_appropriate_gfs_inventory(
wrf_config)
logging.info('Gfs date %s, gfs cycle %s, start inventory %s' %
(gfs_date, gfs_cycle, str(start).zfill(3)))
except KeyError as e:
# raise WrfRunException(str(e))
logging.error(
'Unable to find the key: %s. Returining an empty subdag' % str(e))
return dag_subdag
gfs_clean = PythonOperator(
python_callable=gfs_cleanup,
task_id='%s-task-%s' % (child_dag_name, 'gfs_cleanup'),
op_args=[wrf_config.get('gfs_clean'), gfs_dir],
default_args=args,
dag=dag_subdag,
)
for i in range(0, period * 24 + 1, step):
logging.info('Adding %dth inventory as %d' % (int(start) + i, i))
t = PythonOperator(
python_callable=download_inventory.download_i_th_inventory,
task_id='%s-task-%s' % (child_dag_name, i),
op_args=[
int(start) + i,
wrf_config.get('gfs_url'),
wrf_config.get('gfs_inv'), gfs_date, gfs_cycle,
wrf_config.get('gfs_res'), gfs_dir,
wrf_config.get('nfs_dir'), test_mode
],
# provide_context=True,
default_args=args,
dag=dag_subdag,
)
gfs_clean >> t
return dag_subdag
def test_rainfall_extraction():
rf_task = RainfallExtraction()
rf_task.config = WrfConfig({
'wrf_home': '/home/curw/Desktop/temp',
'wrf_output_dir': '/home/curw/Desktop/temp',
'start_date': '2017-08-13_00:00'
})
rf_task.process()
def set_config(self, **kwargs):
wrf_config_json = None
if 'ti' in kwargs:
wrf_config_json = kwargs['ti'].xcom_pull(task_ids=None,
key=self.wrf_config_key)
if wrf_config_json is not None:
logging.info('wrf_config from xcom using %s key: %s' %
(self.wrf_config_key, wrf_config_json))
self.config = WrfConfig(json.loads(wrf_config_json))
else:
try:
self.config = WrfConfig(
Variable.get(self.wrf_config_key, deserialize_json=True))
logging.info('wrf_config from variable using %s key: %s' %
(self.wrf_config_key, self.config.to_json_string))
except KeyError:
raise CurwAriflowTasksException('Unable to find WrfConfig')
def test_rainfall_extraction2():
rf_task = RainfallExtractionD01()
rf_task.config = WrfConfig({
'wrf_home': '/home/nira/Desktop/temp',
'wrf_output_dir': '/home/nira/Desktop/temp',
'nfs_dir': '/home/nira/Desktop/temp',
'start_date': '2017-09-24_00:00'
})
rf_task.process()
class CurwWrfFileLockSensor(BaseSensorOperator):
@apply_defaults
def __init__(self, config_key, *args, **kwargs):
self.config_key = config_key
self.config = WrfConfig()
super(CurwWrfFileLockSensor, self).__init__(*args, **kwargs)
def poke(self, context):
file_path = os.path.join(
utils.get_em_real_dir(self.get_config().get('wrf_home')),
'wrf.lock')
logging.info('Poking %s' % file_path)
if os.path.exists(file_path) and os.path.isfile(file_path):
return False
else:
return True
def get_config(self):
if self.config.is_empty():
self.config = WrfConfig(
Variable.get(self.config_key, deserialize_json=True))
return self.config
def release_wrf_lock(wrf_config_key):
config = WrfConfig(Variable.get(wrf_config_key, deserialize_json=True))
file_path = os.path.join(utils.get_em_real_dir(config.get('wrf_home')),
'wrf.lock')
logging.info('releasing lock %s' % file_path)
os.remove(file_path)
def __init__(self, config=WrfConfig()):
self.config = config
def get_config(self):
if self.config.is_empty():
self.config = WrfConfig(
Variable.get(self.config_key, deserialize_json=True))
return self.config
def __init__(self, config_key, *args, **kwargs):
self.config_key = config_key
self.config = WrfConfig()
super(CurwWrfFileLockSensor, self).__init__(*args, **kwargs)