Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
test_ds_aio.py 2.86 KB
Newer Older
aiss's avatar
push Deepspeed 0.6.3 rocm version  
aiss committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
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 
"""
Copyright 2020 The Microsoft DeepSpeed Team
Licensed under the MIT license.

Functionality of swapping optimizer tensors to/from (NVMe) storage devices.
"""

import os
import torch
import argparse
import time
import sys
from multiprocessing import Pool
import multiprocessing as mp
from ds_aio_basic import aio_basic_multiprocessing
from ds_aio_handle import aio_handle_multiprocessing
from test_ds_aio_utils import refine_args


def parse_arguments():
    parser = argparse.ArgumentParser()

    parser.add_argument('--read_file', type=str, default=None, help='Read file.')

    parser.add_argument('--write_file', type=str, default=None, help='Write file.')

    parser.add_argument('--write_size',
                        type=str,
                        default=None,
                        help='Number of bytes to write.')

    parser.add_argument('--block_size', type=str, default='1M', help='I/O block size.')

    parser.add_argument('--queue_depth', type=int, default=32, help='I/O queue depth.')

    parser.add_argument('--threads',
                        type=int,
                        default=1,
                        help='Thread parallelism count.')

    parser.add_argument(
        '--single_submit',
        action='store_true',
        help=
        'Submit I/O requests in singles (default is submit queue_depth amount at once.).'
    )

    parser.add_argument('--overlap_events',
                        action='store_true',
                        help='Overlap I/O submission and completion requests.')

    parser.add_argument('--validate',
                        action='store_true',
                        help='Perform validation in library.')

    parser.add_argument('--handle', action='store_true', help='Use AIO handle.')

    parser.add_argument('--loops',
                        type=int,
                        default=1,
                        help='Count of operation repetitions')

    parser.add_argument('--io_parallel',
                        type=int,
                        default=None,
                        help='Per iop parallelism')

    parser.add_argument('--gpu', action='store_true', help='Use GPU memory')

    args = parser.parse_args()
    print(f'args = {args}')
    return args


def validate_args(args):
    if args.read_file and not os.path.isfile(args.read_file):
        print(f'args validation error: {args.read_file} not found')
        return False

    return True


def main():
    print(f'Testing deepspeed_aio python frontend')

    args = parse_arguments()
    refine_args(args)
    if not validate_args(args):
        quit()

    mp.set_start_method('spawn')
    multiprocess_function = aio_handle_multiprocessing if args.handle else aio_basic_multiprocessing
    if args.read_file:
        multiprocess_function(args, True)

    if args.write_file:
        multiprocess_function(args, False)


if __name__ == "__main__":
    main()