_handler.py 12.4 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
    container_name='sb-workspace',
100
101
    docker_username=None,
    docker_password=None,
102
    no_docker=False,
103
    no_image_pull=False,
104
105
106
107
    host_file=None,
    host_list=None,
    host_username=None,
    host_password=None,
Yifan Xiong's avatar
Yifan Xiong committed
108
    private_key=None,
109
    output_dir=None,
Yifan Xiong's avatar
Yifan Xiong committed
110
111
    config_file=None,
    config_override=None
112
):
Yifan Xiong's avatar
Yifan Xiong committed
113
    """Process runner related arguments.
114
115

    Args:
Yifan Xiong's avatar
Yifan Xiong committed
116
        docker_image (str, optional): Docker image URI. Defaults to superbench/superbench:latest.
117
        container_name (str, optional): Docker container name. Defaults to sb-workspace.
118
119
        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.
120
        no_docker (bool, optional): Run on host directly without Docker. Defaults to False.
121
        no_image_pull (bool, optional): Skip pull and use local Docker image. Defaults to False.
122
123
124
125
126
        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.
127
        output_dir (str, optional): Path to output directory. Defaults to None.
Yifan Xiong's avatar
Yifan Xiong committed
128
129
130
131
132
133
134
135
136
        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.
137
138
139
140

    Raises:
        CLIError: If input arguments are invalid.
    """
141
142
    if bool(docker_username) != bool(docker_password):
        raise CLIError('Must specify both docker_username and docker_password if authentication is needed.')
143
144
    if not container_name or not container_name.strip():
        raise CLIError('container_name cannot be empty.')
145
146
    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
147
148
    host_file = check_argument_file('host_file', host_file)
    private_key = check_argument_file('private_key', private_key)
149

Yifan Xiong's avatar
Yifan Xiong committed
150
151
152
153
    # Docker config
    docker_config = OmegaConf.create(
        {
            'image': docker_image,
154
            'container_name': container_name.strip(),
Yifan Xiong's avatar
Yifan Xiong committed
155
156
157
            'username': docker_username,
            'password': docker_password,
            'registry': split_docker_domain(docker_image)[0],
158
            'skip': no_docker,
159
            'pull': not no_image_pull,
Yifan Xiong's avatar
Yifan Xiong committed
160
161
162
163
164
165
166
167
168
169
170
171
        }
    )
    # 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,
        }
    )
172

173
174
175
176
177
    sb_config, sb_output_dir = process_config_arguments(
        config_file=config_file,
        config_override=config_override,
        output_dir=output_dir,
    )
178

179
    return docker_config, ansible_config, sb_config, sb_output_dir
Yifan Xiong's avatar
Yifan Xiong committed
180
181
182


def version_command_handler():
183
    """Print the current SuperBench tool version in "{last tag}+g{git hash}.d{date}" format.
Yifan Xiong's avatar
Yifan Xiong committed
184
185
186
187

    Returns:
        str: current SuperBench tool version.
    """
188
189
190
191
    try:
        return version('superbench')
    except PackageNotFoundError:
        return superbench.__version__
Yifan Xiong's avatar
Yifan Xiong committed
192
193


194
def exec_command_handler(config_file=None, config_override=None, output_dir=None):
195
196
197
198
199
200
    """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.
201
        output_dir (str, optional): Path to output directory. Defaults to None.
202
203
204
205

    Raises:
        CLIError: If input arguments are invalid.
    """
206
207
208
209
210
    sb_config, sb_output_dir = process_config_arguments(
        config_file=config_file,
        config_override=config_override,
        output_dir=output_dir,
    )
211

212
    executor = SuperBenchExecutor(sb_config, sb_output_dir)
Yifan Xiong's avatar
Yifan Xiong committed
213
    executor.exec()
214
215


Yifan Xiong's avatar
Yifan Xiong committed
216
217
def deploy_command_handler(
    docker_image='superbench/superbench',
218
    container_name='sb-workspace',
Yifan Xiong's avatar
Yifan Xiong committed
219
220
    docker_username=None,
    docker_password=None,
221
    no_image_pull=False,
Yifan Xiong's avatar
Yifan Xiong committed
222
223
224
225
    host_file=None,
    host_list=None,
    host_username=None,
    host_password=None,
226
    output_dir=None,
Yifan Xiong's avatar
Yifan Xiong committed
227
228
229
230
231
232
233
234
235
236
237
238
    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.
239
        container_name (str, optional): Docker container name. Defaults to sb-workspace.
Yifan Xiong's avatar
Yifan Xiong committed
240
241
        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.
242
        no_image_pull (bool, optional): Skip pull and use local Docker image. Defaults to False.
Yifan Xiong's avatar
Yifan Xiong committed
243
244
245
246
        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.
247
        output_dir (str, optional): Path to output directory. Defaults to None.
Yifan Xiong's avatar
Yifan Xiong committed
248
249
250
251
252
        private_key (str, optional): Path to private key if needed. Defaults to None.

    Raises:
        CLIError: If input arguments are invalid.
    """
253
    docker_config, ansible_config, sb_config, sb_output_dir = process_runner_arguments(
Yifan Xiong's avatar
Yifan Xiong committed
254
        docker_image=docker_image,
255
        container_name=container_name,
Yifan Xiong's avatar
Yifan Xiong committed
256
257
        docker_username=docker_username,
        docker_password=docker_password,
258
        no_docker=False,
259
        no_image_pull=no_image_pull,
Yifan Xiong's avatar
Yifan Xiong committed
260
261
262
263
        host_file=host_file,
        host_list=host_list,
        host_username=host_username,
        host_password=host_password,
264
        output_dir=output_dir,
Yifan Xiong's avatar
Yifan Xiong committed
265
266
267
        private_key=private_key,
    )

268
    runner = SuperBenchRunner(sb_config, docker_config, ansible_config, sb_output_dir)
269
    runner.deploy()
270
271
    if runner.get_failure_count() != 0:
        sys.exit(runner.get_failure_count())
272
273
274


def run_command_handler(
Yifan Xiong's avatar
Yifan Xiong committed
275
    docker_image='superbench/superbench',
276
    container_name='sb-workspace',
277
278
    docker_username=None,
    docker_password=None,
279
    no_docker=False,
280
281
282
283
    host_file=None,
    host_list=None,
    host_username=None,
    host_password=None,
284
    output_dir=None,
285
286
    private_key=None,
    config_file=None,
287
288
    config_override=None,
    get_info=False,
289
290
291
292
293
294
):
    """Run the SuperBench benchmarks distributedly.

    Run all benchmarks on given nodes.

    Args:
Yifan Xiong's avatar
Yifan Xiong committed
295
        docker_image (str, optional): Docker image URI. Defaults to superbench/superbench:latest.
296
        container_name (str, optional): Docker container name. Defaults to sb-workspace.
297
298
        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.
299
        no_docker (bool, optional): Run on host directly without Docker. Defaults to False.
300
301
302
303
        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.
304
        output_dir (str, optional): Path to output directory. Defaults to None.
305
306
307
308
        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.
309
        get_info (bool, optional): Collect node system info. Defaults to False.
310
311
312
313

    Raises:
        CLIError: If input arguments are invalid.
    """
314
    docker_config, ansible_config, sb_config, sb_output_dir = process_runner_arguments(
Yifan Xiong's avatar
Yifan Xiong committed
315
        docker_image=docker_image,
316
        container_name=container_name,
Yifan Xiong's avatar
Yifan Xiong committed
317
318
        docker_username=docker_username,
        docker_password=docker_password,
319
        no_docker=no_docker,
320
        no_image_pull=False,
Yifan Xiong's avatar
Yifan Xiong committed
321
322
323
324
        host_file=host_file,
        host_list=host_list,
        host_username=host_username,
        host_password=host_password,
325
        output_dir=output_dir,
Yifan Xiong's avatar
Yifan Xiong committed
326
327
328
329
        private_key=private_key,
        config_file=config_file,
        config_override=config_override,
    )
330

331
    runner = SuperBenchRunner(sb_config, docker_config, ansible_config, sb_output_dir)
332

333
    runner.run()
334
335
336
    if get_info:
        runner.run_sys_info()

337
338
    if runner.get_failure_count() != 0:
        sys.exit(runner.get_failure_count())