Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
ds_aio_handle.py 4.96 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
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176 
"""
Copyright 2020 The Microsoft DeepSpeed Team
Licensed under the MIT license.

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

import torch
import os
import time
from multiprocessing import Pool, Barrier
from deepspeed.ops.aio import AsyncIOBuilder
from test_ds_aio_utils import report_results, task_log, task_barrier


def pre_handle(args, tid, read_op):
    io_string = "Read" if read_op else "Write"
    num_bytes = os.path.getsize(args.read_file) if read_op else args.write_size
    file = args.read_file if read_op else f'{args.write_file}.{tid}'

    task_log(tid, f'Allocate tensor of size {num_bytes} bytes')
    if args.gpu:
        buffer = torch.empty(num_bytes, dtype=torch.uint8, device='cuda')
    else:
        buffer = torch.empty(num_bytes, dtype=torch.uint8, device='cpu').pin_memory()
    task_log(
        tid,
        f'{io_string} file {file} of size {num_bytes} bytes from buffer on device {buffer.device}'
    )

    io_parallel = args.io_parallel if args.io_parallel else 1
    handle = AsyncIOBuilder().load().aio_handle(args.block_size,
                                                args.queue_depth,
                                                args.single_submit,
                                                args.overlap_events,
                                                io_parallel)
    task_log(tid, f'created deepspeed aio handle')

    ctxt = {}
    ctxt['file'] = file
    ctxt['num_bytes'] = num_bytes
    ctxt['handle'] = handle
    ctxt['buffer'] = buffer
    ctxt['elapsed_sec'] = 0

    return ctxt


def pre_handle_read(pool_params):
    args, tid = pool_params
    ctxt = pre_handle(args, tid, True)
    return ctxt


def pre_handle_write(pool_params):
    args, tid = pool_params
    ctxt = pre_handle(args, tid, False)
    return ctxt


def post_handle(pool_params):
    _, _, ctxt = pool_params
    ctxt["buffer"].detach()
    ctxt["buffer"] = None
    return ctxt


def main_parallel_read(pool_params):
    args, tid, ctxt = pool_params
    handle = ctxt['handle']

    start_time = time.time()
    ret = handle.pread(ctxt['buffer'], ctxt['file'], args.validate, True)
    assert ret != -1
    handle.wait()
    end_time = time.time()
    ctxt['elapsed_sec'] += end_time - start_time

    return ctxt


def main_parallel_write(pool_params):
    args, tid, ctxt = pool_params
    handle = ctxt['handle']
    start_time = time.time()
    ret = handle.pwrite(ctxt['buffer'], ctxt['file'], args.validate, True)
    assert ret != -1
    handle.wait()
    end_time = time.time()
    ctxt['elapsed_sec'] += end_time - start_time

    return ctxt


def main_handle_read(pool_parms):
    args, tid, ctxt = pool_parms
    handle = ctxt['handle']

    start_time = time.time()
    ret = handle.read(ctxt['buffer'], ctxt['file'], args.validate)
    assert ret != -1
    end_time = time.time()
    ctxt['elapsed_sec'] += end_time - start_time

    return ctxt


def main_handle_write(pool_parms):
    args, tid, ctxt = pool_parms
    handle = ctxt['handle']
    start_time = time.time()
    ret = handle.write(ctxt['buffer'], ctxt['file'], args.validate)
    assert ret != -1
    end_time = time.time()
    ctxt['elapsed_sec'] += end_time - start_time

    return ctxt


def get_schedule(args, read_op):
    schedule = {}
    if read_op:
        schedule['pre'] = pre_handle_read
        schedule['post'] = post_handle
        schedule['main'] = main_parallel_read if args.io_parallel else main_handle_read
    else:
        schedule['pre'] = pre_handle_write
        schedule['post'] = post_handle
        schedule['main'] = main_parallel_write if args.io_parallel else main_handle_write

    return schedule


def _aio_handle_tasklet(pool_params):
    args, tid, read_op = pool_params

    # Create schedule
    schedule = get_schedule(args, read_op)
    task_log(tid, f'schedule = {schedule}')
    task_barrier(aio_barrier, args.threads)

    # Run pre task
    task_log(tid, f'running pre-task')
    ctxt = schedule["pre"]((args, tid))
    task_barrier(aio_barrier, args.threads)

    # Run main tasks in a loop
    ctxt["main_task_sec"] = 0
    for i in range(args.loops):
        task_log(tid, f'running main task {i}')
        start_time = time.time()
        ctxt = schedule["main"]((args, tid, ctxt))
        task_barrier(aio_barrier, args.threads)
        stop_time = time.time()
        ctxt["main_task_sec"] += stop_time - start_time

    # Run post task
    task_log(tid, f'running post-task')
    ctxt = schedule["post"]((args, tid, ctxt))
    task_barrier(aio_barrier, args.threads)

    return ctxt["main_task_sec"], ctxt["elapsed_sec"], ctxt["num_bytes"] * args.loops


def _init_tasklet(b):
    global aio_barrier
    aio_barrier = b


def aio_handle_multiprocessing(args, read_op):
    b = Barrier(args.threads)
    pool_params = [(args, p, read_op) for p in range(args.threads)]
    with Pool(processes=args.threads, initializer=_init_tasklet, initargs=(b, )) as p:
        pool_results = p.map(_aio_handle_tasklet, pool_params)

    report_results(args, read_op, pool_results)