cli.py 4.51 KB
Newer Older
Dion Häfner's avatar
Dion Häfner committed
1
import functools
2 3
import sys
import time
Dion Häfner's avatar
Dion Häfner committed
4 5 6

import click

Dion Häfner's avatar
Dion Häfner committed
7
from veros.settings import SETTINGS
Dion Häfner's avatar
Dion Häfner committed
8

9 10
BACKENDS = ['numpy', 'bohrium']
LOGLEVELS = ['trace', 'debug', 'info', 'warning', 'error', 'critical']
11

Dion Häfner's avatar
Dion Häfner committed
12 13

class VerosSetting(click.ParamType):
14
    name = 'setting'
Dion Häfner's avatar
Dion Häfner committed
15 16 17 18 19 20 21
    current_key = None

    def convert(self, value, param, ctx):
        assert param.nargs == 2

        if self.current_key is None:
            if value not in SETTINGS:
22
                self.fail('Unknown setting %s' % value)
Dion Häfner's avatar
Dion Häfner committed
23 24 25 26 27 28 29 30 31 32 33 34 35 36
            self.current_key = value
            return value

        assert self.current_key in SETTINGS
        setting = SETTINGS[self.current_key]
        self.current_key = None

        if setting.type is bool:
            return click.BOOL(value)

        return setting.type(value)


def cli(run):
37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67
    """Decorator that wraps the decorated function with the Veros setup command line interface.

    Example:

        >>> @veros.tools.cli.cli()
        >>> def run_setup(override):
        ...     sim = MyVerosSetup(override=override)
        ...     sim.run()
        ...
        >>> if __name__ == '__main__':
        ...     run_setup()

    This script then automatically supports settings to be specified from the command line::

        $ python my_setup.py --help
        Usage: my_setup.py [OPTIONS]

        Options:
        -b, --backend [numpy|bohrium]   Backend to use for computations (default:
                                        numpy)
        -v, --loglevel [trace|debug|info|warning|error|critical]
                                        Log level used for output (default: info)
        -s, --override SETTING VALUE    Override default setting, may be specified
                                        multiple times
        -p, --profile-mode              Write a performance profile for debugging
                                        (default: false)
        -n, --num-proc INTEGER...       Number of processes in x and y dimension
                                        (requires execution via mpirun)
        --help                          Show this message and exit.

    """
68 69 70 71 72 73
    @click.command('veros-run')
    @click.option('-b', '--backend', default='numpy', type=click.Choice(BACKENDS),
                  help='Backend to use for computations (default: numpy)', envvar='VEROS_BACKEND')
    @click.option('-v', '--loglevel', default='info', type=click.Choice(LOGLEVELS),
                  help='Log level used for output (default: info)', envvar='VEROS_LOGLEVEL')
    @click.option('-s', '--override', nargs=2, multiple=True, metavar='SETTING VALUE',
Dion Häfner's avatar
Dion Häfner committed
74
                  type=VerosSetting(), default=tuple(),
75 76 77 78
                  help='Override default setting, may be specified multiple times')
    @click.option('-p', '--profile-mode', is_flag=True, default=False, type=click.BOOL, envvar='VEROS_PROFILE',
                  help='Write a performance profile for debugging (default: false)')
    @click.option('-n', '--num-proc', nargs=2, default=[1, 1], type=click.INT,
Dion Häfner's avatar
cleanup  
Dion Häfner committed
79 80 81
                  help='Number of processes in x and y dimension')
    @click.option('--slave', default=False, is_flag=True, hidden=True,
                  help='Indicates that this process is an MPI worker (for internal use)')
Dion Häfner's avatar
Dion Häfner committed
82
    @functools.wraps(run)
83 84 85 86 87 88 89 90 91 92 93 94 95 96 97
    def wrapped(*args, slave, **kwargs):
        from veros import runtime_settings, runtime_state

        total_proc = kwargs['num_proc'][0] * kwargs['num_proc'][1]

        if total_proc > 1 and runtime_state.proc_num == 1 and not slave:
            from mpi4py import MPI

            comm = MPI.COMM_SELF.Spawn(
                sys.executable,
                args=['-m', 'mpi4py'] + list(sys.argv) + ['--slave'],
                maxprocs=total_proc
            )

            futures = [comm.irecv(source=p) for p in range(total_proc)]
Dion Häfner's avatar
cleanup  
Dion Häfner committed
98 99 100 101 102 103 104 105 106
            while True:
                done, success = zip(*(f.test() for f in futures))

                if any(s is False for s in success):
                    raise RuntimeError('An MPI worker encountered an error')

                if all(done):
                    break

107 108 109
                time.sleep(0.1)

            return
Dion Häfner's avatar
Dion Häfner committed
110

111
        kwargs['override'] = dict(kwargs['override'])
Dion Häfner's avatar
Dion Häfner committed
112

113
        for setting in ('backend', 'profile_mode', 'num_proc', 'loglevel'):
Dion Häfner's avatar
Dion Häfner committed
114 115
            setattr(runtime_settings, setting, kwargs.pop(setting))

116 117
        try:
            run(*args, **kwargs)
Dion Häfner's avatar
cleanup  
Dion Häfner committed
118 119 120 121 122
        except:  # noqa: E722
            status = False
            raise
        else:
            status = True
123 124
        finally:
            if slave:
Dion Häfner's avatar
cleanup  
Dion Häfner committed
125
                runtime_settings.mpi_comm.Get_parent().send(status, dest=0)
Dion Häfner's avatar
Dion Häfner committed
126 127

    return wrapped