Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
ds_aio_basic.py 4.28 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 
"""
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 test_ds_aio_utils import report_results, task_log, task_barrier
aiss's avatar
push dsv0.8.2 version  
aiss committed
13
14 
from deepspeed.accelerator import get_accelerator
from deepspeed.ops.op_builder import AsyncIOBuilder
aiss's avatar
push Deepspeed 0.6.3 rocm version  
aiss committed
15
16
17
18
19
20
21
22 


def pre_basic(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')
aiss's avatar
push dsv0.8.2 version  
aiss committed
23
24
25
26 
    buffer = get_accelerator().pin_memory(
        torch.empty(num_bytes,
                    dtype=torch.uint8,
                    device='cpu'))
aiss's avatar
push Deepspeed 0.6.3 rocm version  
aiss committed
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 
    task_log(
        tid,
        f'{io_string} file {file} of size {num_bytes} bytes from buffer on device {buffer.device}'
    )

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

    return ctxt


def pre_basic_read(pool_params):
    args, tid = pool_params
    ctxt = pre_basic(args, tid, True)
    return ctxt


def pre_basic_write(pool_params):
    args, tid = pool_params
    ctxt = pre_basic(args, tid, False)
    return ctxt


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


def main_basic_read(pool_params):
    args, tid, ctxt = pool_params
    start_time = time.time()
    AsyncIOBuilder().load().aio_read(ctxt['buffer'],
                                     ctxt['file'],
                                     args.block_size,
                                     args.queue_depth,
                                     args.single_submit,
                                     args.overlap_events,
                                     args.validate)
    end_time = time.time()
    ctxt['elapsed_sec'] += end_time - start_time

    return ctxt


def main_basic_write(pool_params):
    args, tid, ctxt = pool_params
    start_time = time.time()
    AsyncIOBuilder().load().aio_write(ctxt['buffer'],
                                      ctxt['file'],
                                      args.block_size,
                                      args.queue_depth,
                                      args.single_submit,
                                      args.overlap_events,
                                      args.validate)
    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_basic_read
        schedule['post'] = post_basic
        schedule['main'] = main_basic_read
    else:
        schedule['pre'] = pre_basic_write
        schedule['post'] = post_basic
        schedule['main'] = main_basic_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_basic_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)