Skip to content
Snippets Groups Projects
Select Git revision
  • 1fe79b32b1e11c54339b2bf7402ecf0d6a08b99b
  • consistent_config default protected
2 results

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)