Select Git revision
scheduler.py
scheduler.py 11.18 KiB
#!/usr/bin/python3
"""
high level control script
submitting jobs into SLURM queue
"""
import os, sys, shutil
import datetime as dt
from slurmpy import Slurm
from config.cfg import exp, cluster
from scripts.utils import script_to_str, symlink
if __name__ == "__main__":
# necessary to find modules in folder, since SLURM runs the script elsewhere
sys.path.append(os.getcwd())
# allow scripts to access the configuration
symlink(cluster.scriptsdir+'/../config', cluster.scriptsdir+'/config')
log_dir = cluster.archivedir()+'/logs/'
slurm_scripts_dir = cluster.archivedir()+'/slurm-scripts/'
print('logging to', log_dir)
print('scripts, which are submitted to SLURM:', slurm_scripts_dir)
def my_Slurm(*args, cfg_update=dict(), **kwargs):
"""Shortcut to slurmpy's class; keep certain default kwargs
and only update some with kwarg `cfg_update`
see https://github.com/brentp/slurmpy
"""
return Slurm(*args, slurm_kwargs=dict(cluster.slurm_cfg, **cfg_update),
log_dir=log_dir, scripts_dir=slurm_scripts_dir, **kwargs)
class Cmdline(object):
def __init__(self, name, cfg_update):
self.name = name
def run(self, cmd, **kwargs):
print('running', self.name, 'without SLURM')
os.system(cmd)
def backup_scripts():
current = cluster.scriptsdir
main_a = cluster.archivedir()+'/DART-WRF/'
old_a = main_a+'/old/'
os.makedirs(cluster.archivedir(), exist_ok=True)
os.makedirs(main_a, exist_ok=True)
os.makedirs(old_a, exist_ok=True)
def func(a, b, method): # call method if not link or directory
if os.path.islink(a) or os.path.isdir(a):
pass
else:
method(a, b)
# archive existing files
for f in os.listdir(main_a):
func(os.path.join(main_a, f), old_a+'/'+f, shutil.move)
# reproducibility
for f in ['scheduler.py', 'config/clusters.py', 'config/cfg.py']:
fname = os.path.basename(f)
func(current+'/../'+f, main_a+'/'+fname, shutil.copy)
for f in os.listdir(current):
func(os.path.join(current, f), main_a+'/', shutil.copy)
def prepare_wrfinput():
"""Create WRF/run directories and wrfinput files
"""
s = my_Slurm("prep_wrfinput", cfg_update={"time": "5", "mail-type": "BEGIN"})
id = s.run(cluster.python+' '+cluster.scriptsdir+'/prepare_wrfinput.py')
cmd = """# run ideal.exe in parallel, then add geodata
export SLURM_STEP_GRES=none
for ((n=1; n<="""+str(exp.n_ens)+"""; n++))
do
rundir="""+cluster.userdir+'/run_WRF/'+exp.expname+"""/$n
echo $rundir
cd $rundir
mpirun -np 1 ./ideal.exe &
done
wait
for ((n=1; n<="""+str(exp.n_ens)+"""; n++))
do
rundir="""+cluster.userdir+'/run_WRF/'+exp.expname+"""/$n
mv $rundir/rsl.out.0000 $rundir/rsl.out.input
done
"""
s = my_Slurm("ideal", cfg_update={"ntasks": str(exp.n_ens), "nodes": "1",
"time": "10", "mem-per-cpu": "2G"})
id = s.run(cmd, depends_on=[id])
return id
def update_wrfinput_from_archive(valid_time, background_init_time, exppath, depends_on=None):
"""Given that directories with wrfinput files exist,
update these wrfinput files according to wrfout files
"""
s = my_Slurm("upd_wrfinput", cfg_update={"time": "5"})
# path of initial conditions, <iens> is replaced by member index
IC_path = exppath + background_init_time.strftime('/%Y-%m-%d_%H:%M/') \
+'*iens*/'+valid_time.strftime('/wrfout_d01_%Y-%m-%d_%H:%M:%S')
id = s.run(cluster.python+' '+cluster.scriptsdir+'/update_wrfinput_from_wrfout.py '
+IC_path, depends_on=[depends_on])
return id
def run_ENS(begin, end, depends_on=None, first_minute=True):
"""Run forecast for 1 minute, save output.
Then run whole timespan with 5 minutes interval.
"""
id = depends_on
if first_minute:
# first minute forecast (needed for validating an assimilation)
hist_interval = 1
radt = 1 # calc CFRAC also in first minute
begin_plus1 = begin+dt.timedelta(minutes=1)
s = my_Slurm("preWRF1", cfg_update=dict(time="2"))
id = s.run(' '.join([cluster.python,
cluster.scriptsdir+'/prepare_namelist.py',
begin.strftime('%Y-%m-%d_%H:%M'),
begin_plus1.strftime('%Y-%m-%d_%H:%M'),
str(hist_interval), str(radt),]),
depends_on=[id])
s = my_Slurm("runWRF1", cfg_update={"nodes": "1", "array": "1-"+str(exp.n_nodes),
"time": "2", "mem-per-cpu": "2G"})
cmd = script_to_str(cluster.run_WRF).replace('<expname>', exp.expname)
id = s.run(cmd, depends_on=[id])
# apply forward operator (DART filter without assimilation)
s = my_Slurm("fwOP-1m", cfg_update=dict(time="10", ntasks=48))
id = s.run(cluster.python+' '+cluster.scriptsdir+'/apply_obs_op_dart.py '
+ begin.strftime('%Y-%m-%d_%H:%M')+' '
+ begin_plus1.strftime('%Y-%m-%d_%H:%M'),
depends_on=[id])
# whole forecast timespan
hist_interval = 5
radt = 5
s = my_Slurm("preWRF2", cfg_update=dict(time="2"))
id = s.run(' '.join([cluster.python,
cluster.scriptsdir+'/prepare_namelist.py',
begin.strftime('%Y-%m-%d_%H:%M'),
end.strftime('%Y-%m-%d_%H:%M'),
str(hist_interval), str(radt),]),
depends_on=[id])
time_in_simulation_hours = (end-begin).total_seconds()/3600
runtime_wallclock_mins_expected = int(4+time_in_simulation_hours*9) # usually below 8 min/hour
s = my_Slurm("runWRF2", cfg_update={"nodes": "1", "array": "1-"+str(exp.n_nodes),
"time": str(runtime_wallclock_mins_expected), "mem-per-cpu": "2G"})
cmd = script_to_str(cluster.run_WRF).replace('<expname>', exp.expname)
id = s.run(cmd, depends_on=[id])
# not needed, since wrf.exe writes directly to archive folder
#s = my_Slurm("archiveWRF", cfg_update=dict(nodes="1", ntasks="1", time="10"))
#id3 = s.run(cluster.python+' '+cluster.scriptsdir+'/archive_wrf.py '
# + begin.strftime('%Y-%m-%d_%H:%M'), depends_on=[id2])
return id
def assimilate(assim_time, prior_init_time,
prior_path_exp=False, depends_on=None):
"""Creates observations from a nature run and assimilates them.
Args:
assim_time (dt.datetime): timestamp of prior wrfout files
prior_init_time (dt.datetime):
timestamp to find the directory where the prior wrfout files are
prior_path_exp (bool or str):
put a `str` to take the prior from a different experiment
if False: use `archivedir` (defined in config) to get prior state
if str: use this directory to get prior state
"""
if not prior_path_exp:
prior_path_exp = cluster.archivedir()
elif not isinstance(prior_path_exp, str):
raise TypeError('prior_path_exp either str or False, is '+str(type(prior_path_exp)))
# prepare state of nature run, from which observation is sampled
#s = my_Slurm("prepNature", cfg_update=dict(time="2"))
#id = s.run(cluster.python+' '+cluster.scriptsdir+'/prepare_nature.py '
# +time.strftime('%Y-%m-%d_%H:%M'), depends_on=[depends_on])
# prepare prior model state
s = my_Slurm("preAssim", cfg_update=dict(time="2"))
id = s.run(cluster.python+' '+cluster.scriptsdir+'/pre_assim.py '
+assim_time.strftime('%Y-%m-%d_%H:%M ')
+prior_init_time.strftime('%Y-%m-%d_%H:%M ')
+prior_path_exp, depends_on=[depends_on])
# prepare nature run, generate observations
s = my_Slurm("Assim", cfg_update={"nodes": "1", "ntasks": "96", "time": "30",
"mem": "300G", "ntasks-per-node": "96", "ntasks-per-core": "2"})
id = s.run(cluster.python+' '+cluster.scriptsdir+'/assim_synth_obs.py '
+time.strftime('%Y-%m-%d_%H:%M'), depends_on=[id])
# # actuall assimilation step
# s = my_Slurm("Assim", cfg_update=dict(nodes="1", ntasks="48", time="50", mem="200G"))
# cmd = 'cd '+cluster.dartrundir+'; mpirun -np 48 ./filter; rm obs_seq_all.out'
# id = s.run(cmd, depends_on=[id])
# s = my_Slurm("archiveAssim", cfg_update=dict(time="10"))
# id = s.run(cluster.python+' '+cluster.scriptsdir+'/archive_assim.py '
# + assim_time.strftime('%Y-%m-%d_%H:%M'), depends_on=[id])
s = my_Slurm("updateIC", cfg_update=dict(time="8"))
id = s.run(cluster.python+' '+cluster.scriptsdir+'/update_wrfinput_from_filteroutput.py '
+assim_time.strftime('%Y-%m-%d_%H:%M ')
+prior_init_time.strftime('%Y-%m-%d_%H:%M ')
+prior_path_exp, depends_on=[id])
return id
def create_satimages(init_time, depends_on=None):
s = my_Slurm("pRTTOV", cfg_update={"ntasks": "48", "time": "30"})
s.run(cluster.python+' /home/fs71386/lkugler/RTTOV-WRF/run_init.py '+cluster.archivedir()
+init_time.strftime('/%Y-%m-%d_%H:%M/'),
depends_on=[depends_on])
def mailme(depends_on=None):
if depends_on:
s = my_Slurm("AllFinished", cfg_update={"time": "1", "mail-type": "BEGIN"})
s.run('sleep 1', depends_on=[depends_on])
################################
if __name__ == "__main__":
print('starting osse')
timedelta_integrate = dt.timedelta(minutes=75)
timedelta_btw_assim = dt.timedelta(minutes=60)
backup_scripts()
id = None
start_from_existing_state = False
is_new_run = not start_from_existing_state
if is_new_run:
id = prepare_wrfinput() # create initial conditions
# spin up the ensemble
init_time = dt.datetime(2008, 7, 30, 6, 0)
integration_end_time = dt.datetime(2008, 7, 30, 9, 0)
id = run_ENS(begin=init_time,
end=integration_end_time,
first_minute=False,
depends_on=id)
time = integration_end_time
elif start_from_existing_state:
# get initial conditions from archive
init_time = dt.datetime(2008, 7, 30, 6)
integration_end_time = dt.datetime(2008, 7, 30, 10)
exppath_arch = '/gpfs/data/fs71386/lkugler/sim_archive/exp_v1.13_P0_ps-t2'
id = update_wrfinput_from_archive(integration_end_time, init_time, exppath_arch, depends_on=id)
# values for assimilation
assim_time = integration_end_time
prior_init_time = init_time
prior_path_exp = False #exppath_arch
while time <= dt.datetime(2008, 7, 30, 18):
id = assimilate(assim_time,
prior_init_time,
prior_path_exp=prior_path_exp,
depends_on=id)
prior_path_exp = False # use own exp path
# integration
this_forecast_init = assim_time # start integration from here
this_forecast_end = assim_time + timedelta_integrate
id = run_ENS(begin=this_forecast_init,
end=this_forecast_end,
depends_on=id)
create_satimages(this_forecast_init, depends_on=id)
# increment time
time += timedelta_btw_assim
# values for next iteration
assim_time = time
prior_init_time = time - timedelta_btw_assim
mailme(id)