mirror of
https://github.com/deepspeedai/DeepSpeed.git
synced 2025-10-20 15:33:51 +08:00
348 lines
10 KiB
Python
348 lines
10 KiB
Python
"""
|
|
Copyright 2021 The Microsoft DeepSpeed Team
|
|
Licensed under the MIT license.
|
|
|
|
Functionality of swapping optimizer tensors to/from (NVMe) storage devices.
|
|
"""
|
|
import os
|
|
import sys
|
|
import argparse
|
|
import json
|
|
import itertools
|
|
import subprocess
|
|
import shutil
|
|
|
|
from test_ds_aio_utils import refine_integer_value
|
|
from perf_sweep_utils import READ_OP_DESC, WRITE_OP_DESC, BENCH_LOG_DIR, \
|
|
READ_IO_DIR, WRITE_IO_DIR, READ_LOG_DIR, WRITE_LOG_DIR
|
|
from deepspeed.ops.op_builder import AsyncIOBuilder
|
|
|
|
OTHER_OPTIONS = '--handle'
|
|
PERF_SCRIPT = 'test_ds_aio.py'
|
|
DEFAULT_SWEEP_CONFIG = {
|
|
"block_size": ["128K", "256K"],
|
|
"queue_depth": [4, 16, 32],
|
|
"overlap_events": [True, False],
|
|
"io_parallel": [2, 8],
|
|
"single_submit": [False]
|
|
}
|
|
|
|
|
|
class Job(object):
|
|
|
|
def __init__(self, cmd_line, output_file=None, work_dir=None):
|
|
self.cmd_line = cmd_line
|
|
self.output_file = output_file
|
|
self.work_dir = work_dir
|
|
self.output_fd = None
|
|
|
|
def cmd(self):
|
|
return self.cmd_line
|
|
|
|
def get_stdout(self):
|
|
return self.output_fd
|
|
|
|
def get_stderr(self):
|
|
return self.output_fd
|
|
|
|
def get_cwd(self):
|
|
return self.work_dir
|
|
|
|
def open_output_file(self):
|
|
if self.output_file is not None:
|
|
self.output_fd = open(self.output_file, 'w')
|
|
|
|
def close_output_file(self):
|
|
if self.output_fd is not None:
|
|
self.output_fd.close()
|
|
self.output_fd = None
|
|
|
|
|
|
class SweepConfig(object):
|
|
|
|
def __init__(self, args):
|
|
self.nvme_dir = args.nvme_dir
|
|
self.io_size = args.io_size
|
|
self.search_space = get_sweep_config_dict(args.sweep_config)
|
|
self.read = not args.no_read
|
|
self.write = not args.no_write
|
|
self.flush_cache = not args.no_sudo
|
|
self.log_dir = args.log_dir
|
|
self.loops = args.loops
|
|
self.other_options = f'{OTHER_OPTIONS} --loops {args.loops}'
|
|
|
|
|
|
def parse_arguments():
|
|
parser = argparse.ArgumentParser()
|
|
|
|
parser.add_argument('--nvme_dir',
|
|
required=True,
|
|
type=str,
|
|
help='Directory in which to perform I/O tests. A writeable directory on a NVMe device.')
|
|
|
|
parser.add_argument('--sweep_config', type=str, default=None, help='Performance sweep configuration json file.')
|
|
|
|
parser.add_argument('--no_read', action='store_true', help='Disable read performance measurements.')
|
|
|
|
parser.add_argument('--no_write', action='store_true', help='Disable write performance measurements.')
|
|
|
|
parser.add_argument('--io_size',
|
|
type=str,
|
|
default="400M",
|
|
help='Number of I/O bytes to read/write for performance measurements.')
|
|
|
|
parser.add_argument(
|
|
'--no_sudo',
|
|
action='store_true',
|
|
help=
|
|
'Run without sudo access. Page cache will not be flushed and reported read speeds may be higher than actual.')
|
|
|
|
parser.add_argument(
|
|
'--log_dir',
|
|
type=str,
|
|
default=BENCH_LOG_DIR,
|
|
help=f'Output directory for performance log files. Default is {os.path.join(".", BENCH_LOG_DIR)}')
|
|
|
|
parser.add_argument('--loops', type=int, default=1, help='Count of operation repetitions')
|
|
|
|
args = parser.parse_args()
|
|
print(f'args = {args}')
|
|
|
|
return args
|
|
|
|
|
|
def dump_cmd_lines(cmd_lines):
|
|
print(f'cmd line count = {len(cmd_lines)}')
|
|
for i, cmd in enumerate(cmd_lines):
|
|
print(f'{i}: {cmd}')
|
|
|
|
|
|
def get_sweep_config_dict(sweep_config_json):
|
|
if sweep_config_json is None:
|
|
return DEFAULT_SWEEP_CONFIG
|
|
|
|
with open(sweep_config_json) as fp:
|
|
sweep_config = json.load(fp)
|
|
return sweep_config
|
|
|
|
|
|
def get_sweep_cmd_lines(sweep_config_dict):
|
|
|
|
def flatten_options(key, value_list):
|
|
flat_list = []
|
|
for v in value_list:
|
|
if not type(v) is bool:
|
|
flat_list.append(f'--{key} {v}')
|
|
elif v:
|
|
flat_list.append(f'--{key}')
|
|
else:
|
|
flat_list.append(' ')
|
|
|
|
return flat_list
|
|
|
|
flat_list = [flatten_options(key, value) for key, value in sweep_config_dict.items()]
|
|
cmd_list = list(itertools.product(*flat_list))
|
|
cmd_list = [list(cmd) for cmd in cmd_list]
|
|
#dump_cmd_lines(cmd_list)
|
|
return cmd_list
|
|
|
|
|
|
def run_job(job):
|
|
args = ' '.join(job.cmd())
|
|
print(f'args = {args}')
|
|
job.open_output_file()
|
|
proc = subprocess.run(args=args, shell=True, stdout=job.get_stdout(), stderr=job.get_stderr(), cwd=job.get_cwd())
|
|
job.close_output_file()
|
|
assert proc.returncode == 0, \
|
|
f"This command failed: {job.cmd()}"
|
|
|
|
|
|
def launch_sweep(sweep_jobs, sync_job, flush_cache_job):
|
|
for perf_job in sweep_jobs:
|
|
if flush_cache_job is not None:
|
|
run_job(sync_job)
|
|
run_job(flush_cache_job)
|
|
|
|
run_job(perf_job)
|
|
|
|
run_job(sync_job)
|
|
|
|
|
|
def create_cmd_tags(cmd_line):
|
|
tags = {}
|
|
for param_value in cmd_line:
|
|
fields = param_value.split()
|
|
if len(fields) == 1:
|
|
tags[fields[0]] = None
|
|
elif len(fields) == 2:
|
|
tags[fields[0]] = fields[1]
|
|
return tags
|
|
|
|
|
|
def get_log_file(io_op_desc, cmd_line):
|
|
QUEUE_DEPTH = "--queue_depth"
|
|
BLOCK_SIZE = "--block_size"
|
|
SINGLE_SUBMIT = "--single_submit"
|
|
OVERLAP_EVENTS = "--overlap_events"
|
|
THREAD_COUNT = "--threads"
|
|
IO_PARALLEL = "--io_parallel"
|
|
|
|
tag_map = {
|
|
QUEUE_DEPTH: "d",
|
|
BLOCK_SIZE: "bs",
|
|
SINGLE_SUBMIT: "single",
|
|
OVERLAP_EVENTS: "overlap",
|
|
THREAD_COUNT: "t",
|
|
IO_PARALLEL: "p"
|
|
}
|
|
|
|
tag_default = {
|
|
QUEUE_DEPTH: 1,
|
|
BLOCK_SIZE: "1M",
|
|
SINGLE_SUBMIT: "block",
|
|
OVERLAP_EVENTS: "sequential",
|
|
THREAD_COUNT: 1,
|
|
IO_PARALLEL: 1
|
|
}
|
|
|
|
def get_default_value(tag):
|
|
value = tag_default[tag]
|
|
if tag in [SINGLE_SUBMIT, OVERLAP_EVENTS]:
|
|
return value
|
|
return f'{tag_map[tag]}{value}'
|
|
|
|
def get_config_value(tag, value):
|
|
tag_key = tag_map[tag]
|
|
if value is None:
|
|
return tag_key
|
|
return f'{tag_key}{value}'
|
|
|
|
tag_list = [SINGLE_SUBMIT, OVERLAP_EVENTS, THREAD_COUNT, IO_PARALLEL, QUEUE_DEPTH, BLOCK_SIZE]
|
|
log_tags = [io_op_desc]
|
|
cmd_tags = create_cmd_tags(cmd_line)
|
|
for tag in tag_list:
|
|
if tag in cmd_tags:
|
|
log_tags.append(get_config_value(tag, cmd_tags[tag]))
|
|
else:
|
|
log_tags.append(get_default_value(tag))
|
|
|
|
log_file = '_'.join(log_tags)
|
|
log_file += '.txt'
|
|
return log_file
|
|
|
|
|
|
def create_perf_jobs(io_op_desc, log_dir, cmd_lines):
|
|
py_cmd = ['python', os.path.join(script_path(), PERF_SCRIPT)]
|
|
|
|
perf_jobs = []
|
|
for cmd in cmd_lines:
|
|
log_file = os.path.join(log_dir, get_log_file(io_op_desc, cmd))
|
|
job = Job(cmd_line=py_cmd + cmd, output_file=log_file)
|
|
perf_jobs.append(job)
|
|
|
|
return perf_jobs
|
|
|
|
|
|
def script_path():
|
|
return os.path.dirname(os.path.realpath(sys.argv[0]))
|
|
|
|
|
|
def async_io_setup():
|
|
return AsyncIOBuilder().is_compatible()
|
|
|
|
|
|
def get_block_size_and_count(io_bytes):
|
|
block_size = 1
|
|
block_count = io_bytes
|
|
bytes_in_KB = 1024
|
|
|
|
while block_count % bytes_in_KB == 0:
|
|
block_size *= bytes_in_KB
|
|
block_count /= bytes_in_KB
|
|
|
|
return int(block_size), int(block_count)
|
|
|
|
|
|
def create_read_file(sweep_config):
|
|
read_folder = os.path.join(sweep_config.nvme_dir, f'{READ_IO_DIR}')
|
|
os.makedirs(read_folder, exist_ok=True)
|
|
read_file_name = os.path.join(read_folder, f'random_{sweep_config.io_size}B.pt')
|
|
block_size, block_count = get_block_size_and_count(refine_integer_value(sweep_config.io_size))
|
|
dd_job = Job(cmd_line=[f'dd if=/dev/urandom of={read_file_name} bs={block_size} count={block_count}'])
|
|
print(f'[Start] Create read file of {sweep_config.io_size} bytes by running {dd_job.cmd()} ....')
|
|
run_job(dd_job)
|
|
print(f'[Done] Create read file of {sweep_config.io_size} bytes by running {dd_job.cmd()} ....')
|
|
return read_folder, read_file_name
|
|
|
|
|
|
def remove_folder(folder):
|
|
assert os.path.isdir(folder), f"Error: cannot remove {folder} - folder not found"
|
|
shutil.rmtree(folder)
|
|
|
|
|
|
def run_read_sweep(sweep_config, flush_cache_job, sync_job, cmd_lines):
|
|
read_folder, read_file_name = create_read_file(sweep_config)
|
|
read_option = f'--read_file {read_file_name}'
|
|
read_cmd_lines = [[f'{read_option} {sweep_config.other_options}'] + cmd for cmd in cmd_lines]
|
|
#dump_cmd_lines(read_cmd_lines)
|
|
|
|
log_folder = os.path.join(sweep_config.log_dir, f'{READ_LOG_DIR}')
|
|
os.makedirs(log_folder, exist_ok=True)
|
|
|
|
perf_jobs = create_perf_jobs(io_op_desc=READ_OP_DESC, log_dir=log_folder, cmd_lines=read_cmd_lines)
|
|
|
|
launch_sweep(sweep_jobs=perf_jobs, sync_job=sync_job, flush_cache_job=flush_cache_job)
|
|
|
|
remove_folder(read_folder)
|
|
|
|
|
|
def run_write_sweep(sweep_config, flush_cache_job, sync_job, cmd_lines):
|
|
write_folder = os.path.join(sweep_config.nvme_dir, f'{WRITE_IO_DIR}')
|
|
os.makedirs(write_folder, exist_ok=True)
|
|
write_file_name = os.path.join(write_folder, f'random_{sweep_config.io_size}B.pt')
|
|
write_option = f'--write_size {sweep_config.io_size} --write_file {write_file_name}'
|
|
write_cmd_lines = [[f'{write_option} {sweep_config.other_options}'] + cmd for cmd in cmd_lines]
|
|
#dump_cmd_lines(write_cmd_lines)
|
|
|
|
log_folder = os.path.join(sweep_config.log_dir, f'{WRITE_LOG_DIR}')
|
|
os.makedirs(log_folder, exist_ok=True)
|
|
|
|
perf_jobs = create_perf_jobs(io_op_desc=WRITE_OP_DESC, log_dir=log_folder, cmd_lines=write_cmd_lines)
|
|
|
|
launch_sweep(sweep_jobs=perf_jobs, sync_job=sync_job, flush_cache_job=flush_cache_job)
|
|
|
|
remove_folder(write_folder)
|
|
|
|
|
|
def main():
|
|
print("Running performance sweep of deepspeed nvme library")
|
|
|
|
if not async_io_setup():
|
|
error_msg = """
|
|
Failing because environment is not properly configured for deepspeed async i/o module.
|
|
Possible fix: apt install libaio-dev.
|
|
"""
|
|
print(error_msg)
|
|
quit()
|
|
|
|
args = parse_arguments()
|
|
sweep_config = SweepConfig(args)
|
|
cmd_lines = get_sweep_cmd_lines(sweep_config.search_space)
|
|
|
|
if sweep_config.flush_cache:
|
|
flush_cache_job = Job(cmd_line=['sudo', 'bash -c', "'echo 1 > /proc/sys/vm/drop_caches'"])
|
|
else:
|
|
flush_cache_job = None
|
|
|
|
sync_job = Job(cmd_line=['sync'])
|
|
|
|
if sweep_config.read:
|
|
run_read_sweep(sweep_config, flush_cache_job, sync_job, cmd_lines)
|
|
|
|
if sweep_config.write:
|
|
run_write_sweep(sweep_config, flush_cache_job, sync_job, cmd_lines)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|