param_runner 10.1 KB
Newer Older
PLian's avatar
PLian committed
1
#!/usr/bin/env python3
David Trudgian's avatar
David Trudgian committed
2
3

import logging
David Trudgian's avatar
David Trudgian committed
4
import os
PLian's avatar
PLian committed
5
import shutil
6
7
import sys
import subprocess
8
import argparse
9
import colorlog
10
import pathlib
11
12

parentdir = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
PLian's avatar
PLian committed
13
sys.path.insert(0, parentdir)
14

PLian's avatar
PLian committed
15
import param_runner
David Trudgian's avatar
David Trudgian committed
16
from param_runner import __version__, executors, optimizers, param
17

18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
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
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
handler = colorlog.StreamHandler()
handler.setFormatter(colorlog.ColoredFormatter(
    '%(log_color)s%(levelname)-8s %(message)s'))

logger = colorlog.getLogger()
logger.addHandler(handler)


class Main(object):
    def __init__(self):
        parser = argparse.ArgumentParser(
            usage='''param_runner <command> [<paramfile>]

    check  <paramfile>    Check if your yaml file is validate
    run    <paramfile>    Run the job on the local computer
    submit <paramfile>    Run the job on BioHPC clusters

    init <spearmint>      Install Spearmint and Python2 environment (required by spearmint)
    
    test                  Run the test cases
    examples              Show example files
    uninstall             Uninstall param_runner
''')
        parser.add_argument('command', help='Subcommand to run')
        parser.add_argument('-V', '--version', action='version', version=__version__, help='Show version')
        # parse_args defaults to [1:] for args, but you need to
        # exclude the rest of the args too, or validation will fail
        args = parser.parse_args(sys.argv[1:2])
        if not hasattr(self, args.command):
            print('Unrecognized command')
            parser.print_help()
            exit(1)
        # use dispatch pattern to invoke method with same name
        getattr(self, args.command)()

    def check(self):
        parser = argparse.ArgumentParser(usage='param_runner check [-h] [-v] <paramfile>',
            description='Check if your yaml file is validate')
        # prefixing the argument with -- means it's optional
        parser.add_argument('-v', '--verbose', action='store_true', help='Show debug information')
        parser.add_argument('paramfile', type=str, default='', help='Yaml file which contains your configurations')
        # now that we're inside a subcommand, ignore the first
        # TWO argvs, ie the command (git) and the subcommand (commit)
        args = parser.parse_args(sys.argv[2:])
        if args.verbose:
            logger.setLevel(logging.DEBUG)
        else:
            logger.setLevel(logging.INFO)

        self.__prepare_param(args.paramfile)

    def run(self):
        parser = argparse.ArgumentParser(usage='param_runner run [-h] [-v] <paramfile>',
            description='Run the job on the local computer')
        parser.add_argument('-v', '--verbose', action='store_true', help='Show debug information')
        parser.add_argument('paramfile', type=str, default='', help='Yaml file which contains your configurations')
        args = parser.parse_args(sys.argv[2:])
        if args.verbose:
            logger.setLevel(logging.DEBUG)
        else:
            logger.setLevel(logging.INFO)

        p = self.__prepare_param(args.paramfile)

        exe = executors.LocalExecutor(os.path.dirname((os.path.abspath(args.paramfile))), p)
        opt = optimizers.get_optimizer(p, exe)
        ret = opt.run()

        if ret:
            logger.info("Done.")

    def submit(self):
        parser = argparse.ArgumentParser(usage='param_runner submit [-h] [-v] <paramfile>',
            description='This command will generate a sbatch file and try to submit it to the BioHPC clusters')
        parser.add_argument('-v', '--verbose', action='store_true', help='Show debug information')
        parser.add_argument('paramfile', type=str, default='', help='Yaml file which contains your configurations')
        args = parser.parse_args(sys.argv[2:])
        if args.verbose:
            logger.setLevel(logging.DEBUG)
            print('Running param_runner --verbose submit %s' % args.paramfile)
        else:
            logger.setLevel(logging.INFO)
            print('Running param_runner submit %s' % args.paramfile)

        p = self.__prepare_param(args.paramfile)

        # prepare the sbatch file
        try:
            exe = executors.SlurmWrapperExecutor(os.path.dirname((os.path.abspath(args.paramfile))), p)
            opt = optimizers.get_optimizer(p, exe)
            ret = opt.run()
        except Exception as e:
            logger.error("Could not prepare the sbatch for submition.")
            logger.error(e)
            raise
David Trudgian's avatar
David Trudgian committed
113

114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
        # submit the sbatch file
        if ret:
            try:
                out = subprocess.check_output([shutil.which('sbatch'), exe.sbatch_file])
                logger.debug(out.decode())
            except Exception as e:
                logger.error("Could not submit the sbatch file %s." % exe.sbatch_file)
                logger.error(e)
            logger.info("Done.")

    def init(self):
        parser = argparse.ArgumentParser(usage='param_runner init [-h] [-v] spearmint',
            description='Install the computing environment. Currently, only spearmint needs to be installed'
                        ' before running.')
        parser.add_argument('-v', '--verbose', action='store_true', help='Show debug information')
        parser.add_argument('spearmint', type=str, default='spearmint',
                            help='Install Spearmint and Python2 environment (required by spearmint)')
        args = parser.parse_args(sys.argv[2:])
        if args.verbose:
            logger.setLevel(logging.DEBUG)
        else:
            logger.setLevel(logging.INFO)
David Trudgian's avatar
David Trudgian committed
136

PLian's avatar
PLian committed
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
        sp_install_path = os.path.join(os.path.dirname(param_runner.__file__), 'lib')
        logger.info("Will install Python 2.7 and Spearmint into {0}".format(sp_install_path))
        os.chdir(sp_install_path)
        try:
            shutil.rmtree(os.path.join(sp_install_path, 'spearmint', 'lib'))
        except:
            pass

        try:
            shutil.unpack_archive('spearmint.tgz', extract_dir=os.path.dirname(os.path.dirname(param_runner.__file__)))
            subprocess.check_call(['/bin/bash', 'setup_sp.sh', '--prefix={0}'.format(sp_install_path)])
        except Exception as e:
            logger.error('Failed to init spearmint.')
            logger.error(e)

152
153
154
155
156
157
158
159
160
161
    def test(self):
        parser = argparse.ArgumentParser(usage='param_runner test [-h] [-v]',
            description='Run the test cases')
        parser.add_argument('-v', '--verbose', action='store_true', help='Show debug information')
        args = parser.parse_args(sys.argv[2:])
        if args.verbose:
            logger.setLevel(logging.DEBUG)
        else:
            logger.setLevel(logging.INFO)

PLian's avatar
PLian committed
162
163
164
165
166
167
168
169
170
171
        install_path = os.path.dirname(param_runner.__file__)
        print("Testing the installation of param_runner with Pytest")
        os.chdir(install_path)

        try:
            subprocess.check_call([shutil.which('pytest'), 'test'])
        except Exception as e:
            logger.error('Failed to test param_runner.')
            logger.error(e)

172
173
174
175
176
177
178
179
180
181
    def uninstall(self):
        parser = argparse.ArgumentParser(usage='param_runner uninstall [-h] [-v]',
            description='Uninstall param_runner')
        parser.add_argument('-v', '--verbose', action='store_true', help='Show debug information')
        args = parser.parse_args(sys.argv[2:])
        if args.verbose:
            logger.setLevel(logging.DEBUG)
        else:
            logger.setLevel(logging.INFO)

PLian's avatar
PLian committed
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
        sp_install_path = os.path.join(os.path.dirname(param_runner.__file__), 'lib')
        print("Will uninstall param_runner")
        os.chdir(sp_install_path)

        try:
            shutil.rmtree(os.path.join(sp_install_path, 'spearmint'))
        except Exception as e:
            print(e)
            pass

        try:
            subprocess.check_call([shutil.which('pip'), 'uninstall', '-y', 'param_runner'])
        except Exception as e:
            logger.error('Failed to uninstall param_runner.')
            logger.error(e)

198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
    def examples(self):
        parser = argparse.ArgumentParser(usage='param_runner examples [-h] [-v]',
            description='Show example files')
        parser.add_argument('-v', '--verbose', action='store_true', help='Show debug information')
        args = parser.parse_args(sys.argv[2:])
        if args.verbose:
            logger.setLevel(logging.DEBUG)
        else:
            logger.setLevel(logging.INFO)

        examples_path = os.path.join(os.path.dirname(param_runner.__file__), 'examples')
        os.chdir(examples_path)
        # remove cache files
        [p.unlink() for p in pathlib.Path('.').rglob('*.py[co]')]
        [p.rmdir() for p in pathlib.Path('.').rglob('__pycache__')]
        # list the example files tree
        print("Example files for param_runner:\n")
        print(examples_path)
        for l in self.__tree(pathlib.Path.cwd()):
            print(l)

    def __prepare_param(self, param_file):
PLian's avatar
PLian committed
220
221
222
223
224
225
        print("param_runner - version %s" % __version__)
        print("-------------------------------")
        print("Parameter exploration runner for the BioHPC Cluster")
        print("P. Lian, UT Southwestern BioHPC")
        print("biohpc-help@utsouthwestern.edu")
        print("")
David Trudgian's avatar
David Trudgian committed
226
227

        try:
228
229
            p = param.ParamFile(param_file)
            p.load()
230
            return p
David Trudgian's avatar
David Trudgian committed
231
        except Exception as e:
PLian's avatar
PLian committed
232
233
            logger.info("Could not finish the task!\n")
            logger.error("[%s] %s\n" % (e.__class__.__name__, e))
234
            raise
David Trudgian's avatar
David Trudgian committed
235

236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
    def __tree(self, dir_path: pathlib.Path, prefix: str = ''):
        # prefix components:
        space = '    '
        branch = '│   '
        # pointers:
        tee = '├── '
        last = '└── '

        contents = list(dir_path.iterdir())
        # contents each get pointers that are ├── with a final └── :
        pointers = [tee] * (len(contents) - 1) + [last]
        for pointer, path in zip(pointers, contents):
            yield prefix + pointer + path.name
            if path.is_dir():  # extend the prefix and recurse:
                extension = branch if pointer == tee else space
                # i.e. space because last, └── , above so no more |
                yield from self.__tree(path, prefix=prefix + extension)

David Trudgian's avatar
David Trudgian committed
254
255

if __name__ == '__main__':
256
    Main()