_handler.py 11.6 KB
Newer Older
1
2
3
4
5
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.

"""SuperBench CLI command handler."""

6
import sys
7
from pathlib import Path
8
from importlib_metadata import version, PackageNotFoundError
9
10

from knack.util import CLIError
11
from omegaconf import OmegaConf
12
13

import superbench
14
15
from superbench.runner import SuperBenchRunner
from superbench.executor import SuperBenchExecutor
16
from superbench.common.utils import create_sb_output_dir, get_sb_config
17
18
19
20
21
22
23
24
25


def check_argument_file(name, file):
    """Check file path in CLI arguments.

    Args:
        name (str): argument name.
        file (str): file path.

Yifan Xiong's avatar
Yifan Xiong committed
26
27
28
    Returns:
        str: Absolute file path if it exists.

29
30
31
    Raises:
        CLIError: If file does not exist.
    """
Yifan Xiong's avatar
Yifan Xiong committed
32
33
34
35
36
    if file:
        if not Path(file).exists():
            raise CLIError('{} {} does not exist.'.format(name, file))
        return str(Path(file).resolve())
    return file
37
38


Yifan Xiong's avatar
Yifan Xiong committed
39
40
41
42
43
44
45
def split_docker_domain(name):
    """Split Docker image name to domain and remainder part.

    Ported from https://github.com/distribution/distribution/blob/v2.7.1/reference/normalize.go#L62-L76.

    Args:
        name (str): Docker image name.
46
47

    Returns:
Yifan Xiong's avatar
Yifan Xiong committed
48
49
        str: Docker registry domain.
        str: Remainder part.
50
    """
Yifan Xiong's avatar
Yifan Xiong committed
51
52
    legacy_default_domain = 'index.docker.io'
    default_domain = 'docker.io'
53

Yifan Xiong's avatar
Yifan Xiong committed
54
55
56
57
58
59
60
61
62
63
64
    i = name.find('/')
    domain, remainder = '', ''
    if i == -1 or ('.' not in name[:i] and ':' not in name[:i] and name[:i] != 'localhost'):
        domain, remainder = default_domain, name
    else:
        domain, remainder = name[:i], name[i + 1:]
    if domain == legacy_default_domain:
        domain = default_domain
    if domain == default_domain and '/' not in remainder:
        remainder = 'library/{}'.format(remainder)
    return domain, remainder
65

Yifan Xiong's avatar
Yifan Xiong committed
66

67
def process_config_arguments(config_file=None, config_override=None, output_dir=None):
Yifan Xiong's avatar
Yifan Xiong committed
68
69
70
71
72
73
    """Process configuration arguments.

    Args:
        config_file (str, optional): Path to SuperBench config file. Defaults to None.
        config_override (str, optional): Extra arguments to override config_file,
            following [Hydra syntax](https://hydra.cc/docs/advanced/override_grammar/basic). Defaults to None.
74
        output_dir (str, optional): Path to output directory. Defaults to None.
Yifan Xiong's avatar
Yifan Xiong committed
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91

    Returns:
        DictConfig: SuperBench config object.
        str: Dir for output.

    Raises:
        CLIError: If input arguments are invalid.
    """
    config_file = check_argument_file('config_file', config_file)

    # SuperBench config
    sb_config = get_sb_config(config_file)
    if config_override:
        sb_config_from_override = OmegaConf.from_dotlist(config_override)
        sb_config = OmegaConf.merge(sb_config, sb_config_from_override)

    # Create output directory
92
    sb_output_dir = create_sb_output_dir(output_dir)
Yifan Xiong's avatar
Yifan Xiong committed
93

94
    return sb_config, sb_output_dir
Yifan Xiong's avatar
Yifan Xiong committed
95
96
97
98


def process_runner_arguments(
    docker_image='superbench/superbench',
99
100
    docker_username=None,
    docker_password=None,
101
    no_docker=False,
102
    no_image_pull=False,
103
104
105
106
    host_file=None,
    host_list=None,
    host_username=None,
    host_password=None,
Yifan Xiong's avatar
Yifan Xiong committed
107
    private_key=None,
108
    output_dir=None,
Yifan Xiong's avatar
Yifan Xiong committed
109
110
    config_file=None,
    config_override=None
111
):
Yifan Xiong's avatar
Yifan Xiong committed
112
    """Process runner related arguments.
113
114

    Args:
Yifan Xiong's avatar
Yifan Xiong committed
115
        docker_image (str, optional): Docker image URI. Defaults to superbench/superbench:latest.
116
117
        docker_username (str, optional): Docker registry username if authentication is needed. Defaults to None.
        docker_password (str, optional): Docker registry password if authentication is needed. Defaults to None.
118
        no_docker (bool, optional): Run on host directly without Docker. Defaults to False.
119
        no_image_pull (bool, optional): Skip pull and use local Docker image. Defaults to False.
120
121
122
123
124
        host_file (str, optional): Path to Ansible inventory host file. Defaults to None.
        host_list (str, optional): Comma separated host list. Defaults to None.
        host_username (str, optional): Host username if needed. Defaults to None.
        host_password (str, optional): Host password or key passphase if needed. Defaults to None.
        private_key (str, optional): Path to private key if needed. Defaults to None.
125
        output_dir (str, optional): Path to output directory. Defaults to None.
Yifan Xiong's avatar
Yifan Xiong committed
126
127
128
129
130
131
132
133
134
        config_file (str, optional): Path to SuperBench config file. Defaults to None.
        config_override (str, optional): Extra arguments to override config_file,
            following [Hydra syntax](https://hydra.cc/docs/advanced/override_grammar/basic). Defaults to None.

    Returns:
        DictConfig: SuperBench config object.
        DictConfig: Docker config object.
        DictConfig: Ansible config object.
        str: Dir for output.
135
136
137
138

    Raises:
        CLIError: If input arguments are invalid.
    """
139
140
    if bool(docker_username) != bool(docker_password):
        raise CLIError('Must specify both docker_username and docker_password if authentication is needed.')
141
142
    if not (host_file or host_list):
        raise CLIError('Must specify one of host_file or host_list.')
Yifan Xiong's avatar
Yifan Xiong committed
143
144
    host_file = check_argument_file('host_file', host_file)
    private_key = check_argument_file('private_key', private_key)
145

Yifan Xiong's avatar
Yifan Xiong committed
146
147
148
149
150
151
152
    # Docker config
    docker_config = OmegaConf.create(
        {
            'image': docker_image,
            'username': docker_username,
            'password': docker_password,
            'registry': split_docker_domain(docker_image)[0],
153
            'skip': no_docker,
154
            'pull': not no_image_pull,
Yifan Xiong's avatar
Yifan Xiong committed
155
156
157
158
159
160
161
162
163
164
165
166
        }
    )
    # Ansible config
    ansible_config = OmegaConf.create(
        {
            'host_file': host_file,
            'host_list': host_list,
            'host_username': host_username,
            'host_password': host_password,
            'private_key': private_key,
        }
    )
167

168
169
170
171
172
    sb_config, sb_output_dir = process_config_arguments(
        config_file=config_file,
        config_override=config_override,
        output_dir=output_dir,
    )
173

174
    return docker_config, ansible_config, sb_config, sb_output_dir
Yifan Xiong's avatar
Yifan Xiong committed
175
176
177


def version_command_handler():
178
    """Print the current SuperBench tool version in "{last tag}+g{git hash}.d{date}" format.
Yifan Xiong's avatar
Yifan Xiong committed
179
180
181
182

    Returns:
        str: current SuperBench tool version.
    """
183
184
185
186
    try:
        return version('superbench')
    except PackageNotFoundError:
        return superbench.__version__
Yifan Xiong's avatar
Yifan Xiong committed
187
188


189
def exec_command_handler(config_file=None, config_override=None, output_dir=None):
190
191
192
193
194
195
    """Run the SuperBench benchmarks locally.

    Args:
        config_file (str, optional): Path to SuperBench config file. Defaults to None.
        config_override (str, optional): Extra arguments to override config_file,
            following [Hydra syntax](https://hydra.cc/docs/advanced/override_grammar/basic). Defaults to None.
196
        output_dir (str, optional): Path to output directory. Defaults to None.
197
198
199
200

    Raises:
        CLIError: If input arguments are invalid.
    """
201
202
203
204
205
    sb_config, sb_output_dir = process_config_arguments(
        config_file=config_file,
        config_override=config_override,
        output_dir=output_dir,
    )
206

207
    executor = SuperBenchExecutor(sb_config, sb_output_dir)
Yifan Xiong's avatar
Yifan Xiong committed
208
    executor.exec()
209
210


Yifan Xiong's avatar
Yifan Xiong committed
211
212
213
214
def deploy_command_handler(
    docker_image='superbench/superbench',
    docker_username=None,
    docker_password=None,
215
    no_image_pull=False,
Yifan Xiong's avatar
Yifan Xiong committed
216
217
218
219
    host_file=None,
    host_list=None,
    host_username=None,
    host_password=None,
220
    output_dir=None,
Yifan Xiong's avatar
Yifan Xiong committed
221
222
223
224
225
226
227
228
229
230
231
232
233
234
    private_key=None
):
    """Deploy the SuperBench environments to all given nodes.

    Deploy SuperBench environments on all nodes, including:
    1. check drivers
    2. install required system dependencies
    3. install Docker and container runtime
    4. pull Docker image

    Args:
        docker_image (str, optional): Docker image URI. Defaults to superbench/superbench:latest.
        docker_username (str, optional): Docker registry username if authentication is needed. Defaults to None.
        docker_password (str, optional): Docker registry password if authentication is needed. Defaults to None.
235
        no_image_pull (bool, optional): Skip pull and use local Docker image. Defaults to False.
Yifan Xiong's avatar
Yifan Xiong committed
236
237
238
239
        host_file (str, optional): Path to Ansible inventory host file. Defaults to None.
        host_list (str, optional): Comma separated host list. Defaults to None.
        host_username (str, optional): Host username if needed. Defaults to None.
        host_password (str, optional): Host password or key passphase if needed. Defaults to None.
240
        output_dir (str, optional): Path to output directory. Defaults to None.
Yifan Xiong's avatar
Yifan Xiong committed
241
242
243
244
245
        private_key (str, optional): Path to private key if needed. Defaults to None.

    Raises:
        CLIError: If input arguments are invalid.
    """
246
    docker_config, ansible_config, sb_config, sb_output_dir = process_runner_arguments(
Yifan Xiong's avatar
Yifan Xiong committed
247
248
249
        docker_image=docker_image,
        docker_username=docker_username,
        docker_password=docker_password,
250
        no_docker=False,
251
        no_image_pull=no_image_pull,
Yifan Xiong's avatar
Yifan Xiong committed
252
253
254
255
        host_file=host_file,
        host_list=host_list,
        host_username=host_username,
        host_password=host_password,
256
        output_dir=output_dir,
Yifan Xiong's avatar
Yifan Xiong committed
257
258
259
        private_key=private_key,
    )

260
    runner = SuperBenchRunner(sb_config, docker_config, ansible_config, sb_output_dir)
261
    runner.deploy()
262
263
    if runner.get_failure_count() != 0:
        sys.exit(runner.get_failure_count())
264
265
266


def run_command_handler(
Yifan Xiong's avatar
Yifan Xiong committed
267
    docker_image='superbench/superbench',
268
269
    docker_username=None,
    docker_password=None,
270
    no_docker=False,
271
272
273
274
    host_file=None,
    host_list=None,
    host_username=None,
    host_password=None,
275
    output_dir=None,
276
277
278
279
280
281
282
283
284
    private_key=None,
    config_file=None,
    config_override=None
):
    """Run the SuperBench benchmarks distributedly.

    Run all benchmarks on given nodes.

    Args:
Yifan Xiong's avatar
Yifan Xiong committed
285
        docker_image (str, optional): Docker image URI. Defaults to superbench/superbench:latest.
286
287
        docker_username (str, optional): Docker registry username if authentication is needed. Defaults to None.
        docker_password (str, optional): Docker registry password if authentication is needed. Defaults to None.
288
        no_docker (bool, optional): Run on host directly without Docker. Defaults to False.
289
290
291
292
        host_file (str, optional): Path to Ansible inventory host file. Defaults to None.
        host_list (str, optional): Comma separated host list. Defaults to None.
        host_username (str, optional): Host username if needed. Defaults to None.
        host_password (str, optional): Host password or key passphase if needed. Defaults to None.
293
        output_dir (str, optional): Path to output directory. Defaults to None.
294
295
296
297
298
299
300
301
        private_key (str, optional): Path to private key if needed. Defaults to None.
        config_file (str, optional): Path to SuperBench config file. Defaults to None.
        config_override (str, optional): Extra arguments to override config_file,
            following [Hydra syntax](https://hydra.cc/docs/advanced/override_grammar/basic). Defaults to None.

    Raises:
        CLIError: If input arguments are invalid.
    """
302
    docker_config, ansible_config, sb_config, sb_output_dir = process_runner_arguments(
Yifan Xiong's avatar
Yifan Xiong committed
303
304
305
        docker_image=docker_image,
        docker_username=docker_username,
        docker_password=docker_password,
306
        no_docker=no_docker,
307
        no_image_pull=False,
Yifan Xiong's avatar
Yifan Xiong committed
308
309
310
311
        host_file=host_file,
        host_list=host_list,
        host_username=host_username,
        host_password=host_password,
312
        output_dir=output_dir,
Yifan Xiong's avatar
Yifan Xiong committed
313
314
315
316
        private_key=private_key,
        config_file=config_file,
        config_override=config_override,
    )
317

318
    runner = SuperBenchRunner(sb_config, docker_config, ansible_config, sb_output_dir)
319
    runner.run()
320
321
    if runner.get_failure_count() != 0:
        sys.exit(runner.get_failure_count())