Exemplo n.º 1
0
    def __init__(self,
                 name,
                 codebase,
                 safe_mode=False,
                 workbase=GFDL_WORK,
                 database=GFDL_DATA):
        super(Experiment, self).__init__()
        self.name = name
        self.codebase = codebase
        self.safe_mode = safe_mode

        # set the default locations of working directory,
        # executable directory, restart file storage, and
        # output data directory.
        self.workdir = P(workbase, 'experiment', self.name)
        self.rundir = P(
            self.workdir,
            'run')  # temporary area an individual run will be performed
        self.datadir = P(
            database,
            self.name)  # where run data will be moved to upon completion
        self.restartdir = P(self.datadir,
                            'restarts')  # where restarts will be stored
        self.template_dir = P(_module_directory, 'templates')

        self.env_source = get_env_file()

        self.templates = Environment(
            loader=FileSystemLoader(self.template_dir))

        self.diag_table = DiagTable()
        self.field_table_file = P(self.codebase.srcdir, 'extra', 'model',
                                  self.codebase.name, 'field_table')
        self.inputfiles = []

        self.namelist = Namelist()
Exemplo n.º 2
0
class Experiment(Logger, EventEmitter):
    """A basic GFDL experiment"""

    RESOLUTIONS = {
        'T170': {
            'lon_max': 512,
            'lat_max': 256,
            'num_fourier': 170,
            'num_spherical': 171
        },
        'T85': {
            'lon_max': 256,
            'lat_max': 128,
            'num_fourier': 85,
            'num_spherical': 86
        },
        'T42': {
            'lon_max': 128,
            'lat_max': 64,
            'num_fourier': 42,
            'num_spherical': 43,
        },
        'T21': {
            'lon_max': 64,
            'lat_max': 32,
            'num_fourier': 21,
            'num_spherical': 22,
        },
    }

    runfmt = 'run%04d'
    restartfmt = 'res%04d.tar.gz'

    def __init__(self,
                 name,
                 codebase,
                 safe_mode=False,
                 workbase=GFDL_WORK,
                 database=GFDL_DATA):
        super(Experiment, self).__init__()
        self.name = name
        self.codebase = codebase
        self.safe_mode = safe_mode

        # set the default locations of working directory,
        # executable directory, restart file storage, and
        # output data directory.
        self.workdir = P(workbase, 'experiment', self.name)
        self.rundir = P(
            self.workdir,
            'run')  # temporary area an individual run will be performed
        self.datadir = P(
            database,
            self.name)  # where run data will be moved to upon completion
        self.restartdir = P(self.datadir,
                            'restarts')  # where restarts will be stored
        self.template_dir = P(_module_directory, 'templates')

        self.env_source = get_env_file()

        self.templates = Environment(
            loader=FileSystemLoader(self.template_dir))

        self.diag_table = DiagTable()
        self.field_table_file = P(self.codebase.srcdir, 'extra', 'model',
                                  self.codebase.name, 'field_table')
        self.inputfiles = []

        self.namelist = Namelist()

    @destructive
    def rm_workdir(self):
        try:
            sh.rm(['-r', self.workdir])
        except sh.ErrorReturnCode:
            self.log.warning(
                'Tried to remove working directory but it doesnt exist')

    @destructive
    def rm_datadir(self):
        try:
            sh.rm(['-r', self.datadir])
        except sh.ErrorReturnCode:
            self.log.warning(
                'Tried to remove data directory but it doesnt exist')

    @destructive
    @useworkdir
    def clear_workdir(self):
        self.rm_workdir()
        mkdir(self.workdir)
        self.log.info('Emptied working directory %r' % self.workdir)

    @destructive
    @useworkdir
    def clear_rundir(self):
        #sh.cd(self.workdir)
        try:
            sh.rm(['-r', self.rundir])
        except sh.ErrorReturnCode:
            self.log.warning(
                'Tried to remove run directory but it doesnt exist')
        mkdir(self.rundir)
        self.log.info('Emptied run directory %r' % self.rundir)

    def get_restart_file(self, i):
        return P(self.restartdir, self.restartfmt % i)

    def get_outputdir(self, run):
        return P(self.datadir, self.runfmt % run)

    def set_resolution(self, res, num_levels=None):
        """Set the resolution of the model, based on the standard triangular
        truncations of the spectral core.  For example,
            exp.set_resolution('T85', 25)
        creates a spectral core with enough modes to natively correspond to
        a 256x128 lon-lat resolution."""
        delta = self.RESOLUTIONS[res]
        if num_levels is not None:
            delta['num_levels'] = num_levels
        self.update_namelist({'spectral_dynamics_nml': delta})

    def update_namelist(self, new_vals):
        """Update the namelist sections, overwriting existing values."""
        for sec in new_vals:
            if sec not in self.namelist:
                self.namelist[sec] = {}
            nml = self.namelist[sec]
            nml.update(new_vals[sec])

    def write_namelist(self, outdir):
        namelist_file = P(outdir, 'input.nml')
        self.log.info('Writing namelist to %r' % namelist_file)
        self.namelist.write(namelist_file)

    def write_diag_table(self, outdir):
        outfile = P(outdir, 'diag_table')
        self.log.info('Writing diag_table to %r' % outfile)
        if self.diag_table.is_valid():
            if self.diag_table.calendar is None:
                # diagnose the calendar from the namelist
                cal = self.get_calendar()
                self.diag_table.calendar = cal
            self.diag_table.write(outfile)
        else:
            self.log.error(
                "No output files defined in the DiagTable. Stopping.")
            raise ValueError()

    def write_field_table(self, outdir):
        self.log.info('Writing field_table to %r' % P(outdir, 'field_table'))
        sh.cp(self.field_table_file, P(outdir, 'field_table'))

    def log_output(self, outputstring):
        line = outputstring.strip()
        if 'warning' in line.lower():
            self.log.warn(line)
        else:
            self.log.debug(line)
        #return clean_log_debug(outputstring)

    def delete_restart(self, run):
        resfile = self.get_restart_file(run)
        if os.path.isfile(resfile):
            sh.rm(resfile)
            self.log.info('Deleted restart file %s' % resfile)

    def get_calendar(self):
        """Get the value of 'main_nml/calendar.
        Returns a string name of calendar, or None if not set in namelist.'"""
        if 'main_nml' in self.namelist:
            return self.namelist['main_nml'].get('calendar')
        else:
            return None

    @destructive
    @useworkdir
    def run(self,
            i,
            restart_file=None,
            use_restart=True,
            multi_node=False,
            num_cores=8,
            overwrite_data=False,
            save_run=False,
            run_idb=False,
            nice_score=0):
        """Run the model.
            `num_cores`: Number of mpi cores to distribute over.
            `restart_file` (optional): A path to a valid restart archive.  If None and `use_restart=True`,
                                       restart file (i-1) will be used.
            `save_run`:  If True, copy the entire working directory over to GFDL_DATA
                         so that the run can rerun without the python script.
                         (This uses a lot of data storage!)

        """

        self.clear_rundir()

        indir = P(self.rundir, 'INPUT')
        outdir = P(self.datadir, self.runfmt % i)
        resdir = P(self.rundir, 'RESTART')

        if os.path.isdir(outdir):
            if overwrite_data:
                self.log.warning(
                    'Data for run %d already exists and overwrite_data is True. Overwriting.'
                    % i)
                sh.rm('-r', outdir)
            else:
                self.log.warn(
                    'Data for run %d already exists but overwrite_data is False. Stopping.'
                    % i)
                return False

        # make the output run folder and copy over the input files
        mkdir([indir, resdir, self.restartdir])

        self.codebase.write_source_control_status(
            P(self.rundir, 'git_hash_used.txt'))
        self.write_namelist(self.rundir)
        self.write_field_table(self.rundir)
        self.write_diag_table(self.rundir)

        for filename in self.inputfiles:
            sh.cp([filename, P(indir, os.path.split(filename)[1])])

        mpirun_opts = ''

        if multi_node:
            mpirun_opts += ' -bootstrap pbsdsh -f $PBS_NODEFILE'

        if use_restart:
            if not restart_file:
                # get the restart from previous iteration
                restart_file = self.get_restart_file(i - 1)
            if not os.path.isfile(restart_file):
                self.log.error('Restart file not found, expecting file %r' %
                               restart_file)
                raise IOError('Restart file not found, expecting file %r' %
                              restart_file)
            else:
                self.log.info('Using restart file %r' % restart_file)

            self.extract_restart_archive(restart_file, indir)
        else:
            self.log.info('Running without restart file')
            restart_file = None

        vars = {
            'rundir': self.rundir,
            'execdir': self.codebase.builddir,
            'executable': self.codebase.executable_name,
            'env_source': self.env_source,
            'mpirun_opts': mpirun_opts,
            'num_cores': num_cores,
            'run_idb': run_idb,
            'nice_score': nice_score
        }

        runscript = self.templates.get_template('run.sh')

        # employ the template to create a runscript
        t = runscript.stream(**vars).dump(P(self.rundir, 'run.sh'))

        def _outhandler(line):
            handled = self.emit('run:output', self, line)
            if not handled:  # only log the output when no event handler is used
                self.log_output(line)

        self.emit('run:ready', self, i)
        self.log.info("Beginning run %d" % i)
        try:
            #for line in sh.bash(P(self.rundir, 'run.sh'), _iter=True, _err_to_out=True):
            proc = sh.bash(P(self.rundir, 'run.sh'),
                           _bg=True,
                           _out=_outhandler,
                           _err_to_out=True)
            self.log.info('process running as {}'.format(proc.process.pid))
            proc.wait()
            completed = True
        except KeyboardInterrupt as e:
            self.log.error("Manual interrupt, killing process.")
            proc.process.terminate()
            proc.wait()
            #log.info("Cleaning run directory.")
            #self.clear_rundir()
            raise e
        except sh.ErrorReturnCode as e:
            completed = False
            self.log.error("Run %d failed. See log for details." % i)
            self.log.error("Error: %r" % e)
            self.emit('run:failed', self)
            raise FailedRunError()

        self.emit('run:completed', self, i)
        self.log.info('Run %d complete' % i)
        mkdir(outdir)

        if num_cores > 1:
            # use postprocessing tool to combine the output from several cores
            codebase_combine_script = P(self.codebase.builddir,
                                        'mppnccombine_run.sh')
            if not os.path.exists(codebase_combine_script):
                self.log.warning(
                    'combine script does not exist in the commit you are running Isca from.  Falling back to using $GFDL_BASE mppnccombine_run.sh script'
                )
                sh.ln('-s',
                      P(GFDL_BASE, 'postprocessing', 'mppnccombine_run.sh'),
                      codebase_combine_script)
            combinetool = sh.Command(codebase_combine_script)
            for file in self.diag_table.files:
                netcdf_file = '%s.nc' % file
                filebase = P(self.rundir, netcdf_file)
                combinetool(self.codebase.builddir, filebase)
                # copy the combined netcdf file into the data archive directory
                sh.cp(filebase, P(outdir, netcdf_file))
                # remove all netcdf fragments from the run directory
                sh.rm(glob.glob(filebase + '*'))
                self.log.debug('%s combined and copied to data directory' %
                               netcdf_file)

            for restart in glob.glob(P(resdir, '*.res.nc.0000')):
                restartfile = restart.replace('.0000', '')
                combinetool(self.codebase.builddir, restartfile)
                sh.rm(glob.glob(restartfile + '.????'))
                self.log.debug("Restart file %s combined" % restartfile)

            self.emit('run:combined', self)

        # make the restart archive and delete the restart files
        self.make_restart_archive(self.get_restart_file(i), resdir)
        sh.rm('-r', resdir)

        if save_run:
            # copy the complete run directory to GFDL_DATA so that the run can
            # be recreated without the python script if required
            mkdir(resdir)
            sh.cp(['-a', self.rundir, outdir])
        else:
            # just save some useful diagnostic information
            self.write_namelist(outdir)
            self.write_field_table(outdir)
            self.write_diag_table(outdir)
            self.codebase.write_source_control_status(
                P(outdir, 'git_hash_used.txt'))

        self.clear_rundir()

        return True

    def make_restart_archive(self, archive_file, restart_directory):
        with tarfile.open(archive_file, 'w:gz') as tar:
            tar.add(restart_directory, arcname='.')
        self.log.info("Restart archive created at %s" % archive_file)

    def extract_restart_archive(self, archive_file, input_directory):
        with tarfile.open(archive_file, 'r:gz') as tar:
            tar.extractall(path=input_directory)
        self.log.info("Restart %s extracted to %s" %
                      (archive_file, input_directory))

    def derive(self, new_experiment_name):
        """Derive a new experiment based on this one."""
        new_exp = Experiment(new_experiment_name, self.codebase)
        new_exp.namelist = self.namelist.copy()
        new_exp.diag_table = self.diag_table.copy()
        new_exp.inputfiles = self.inputfiles[:]

        return new_exp