wehub-resource-sync
358 lines
12 KiB
Python
358 lines
12 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
import pytest
|
|
import os
|
|
import filecmp
|
|
import torch
|
|
import deepspeed
|
|
import deepspeed.comm as dist
|
|
from deepspeed.accelerator import get_accelerator
|
|
from deepspeed.ops.op_builder import GDSBuilder
|
|
from unit.common import DistributedTest
|
|
|
|
KILO_BYTE = 1024 * 256
|
|
BLOCK_SIZE = KILO_BYTE
|
|
QUEUE_DEPTH = 2
|
|
IO_SIZE = 4 * BLOCK_SIZE
|
|
IO_PARALLEL = 2
|
|
|
|
if not deepspeed.ops.__compatible_ops__[GDSBuilder.NAME]:
|
|
pytest.skip('Skip tests since gds is not compatible', allow_module_level=True)
|
|
|
|
|
|
def _get_local_rank():
|
|
if get_accelerator().is_available():
|
|
return dist.get_rank()
|
|
return 0
|
|
|
|
|
|
def _do_ref_write(tmpdir, index=0, file_size=IO_SIZE):
|
|
file_suffix = f'{_get_local_rank()}_{index}'
|
|
ref_file = os.path.join(tmpdir, f'_py_random_{file_suffix}.pt')
|
|
ref_buffer = os.urandom(file_size)
|
|
with open(ref_file, 'wb') as f:
|
|
f.write(ref_buffer)
|
|
|
|
return ref_file, ref_buffer
|
|
|
|
|
|
def _get_file_path(tmpdir, file_prefix, index=0):
|
|
file_suffix = f'{_get_local_rank()}_{index}'
|
|
return os.path.join(tmpdir, f'{file_prefix}_{file_suffix}.pt')
|
|
|
|
|
|
def _get_test_write_file(tmpdir, index):
|
|
file_suffix = f'{_get_local_rank()}_{index}'
|
|
return os.path.join(tmpdir, f'_gds_write_random_{file_suffix}.pt')
|
|
|
|
|
|
def _get_test_write_file_and_device_buffer(tmpdir, ref_buffer, gds_handle, index=0):
|
|
test_file = _get_test_write_file(tmpdir, index)
|
|
test_buffer = get_accelerator().ByteTensor(list(ref_buffer))
|
|
gds_handle.pin_device_tensor(test_buffer)
|
|
return test_file, test_buffer
|
|
|
|
|
|
def _validate_handle_state(handle, single_submit, overlap_events):
|
|
assert handle.get_single_submit() == single_submit
|
|
assert handle.get_overlap_events() == overlap_events
|
|
assert handle.get_intra_op_parallelism() == IO_PARALLEL
|
|
assert handle.get_block_size() == BLOCK_SIZE
|
|
assert handle.get_queue_depth() == QUEUE_DEPTH
|
|
|
|
|
|
@pytest.mark.parametrize("single_submit", [True, False])
|
|
@pytest.mark.parametrize("overlap_events", [True, False])
|
|
class TestRead(DistributedTest):
|
|
world_size = 1
|
|
reuse_dist_env = True
|
|
if not get_accelerator().is_available():
|
|
init_distributed = False
|
|
set_dist_env = False
|
|
|
|
def test_parallel_read(self, tmpdir, single_submit, overlap_events):
|
|
|
|
h = GDSBuilder().load().gds_handle(BLOCK_SIZE, QUEUE_DEPTH, single_submit, overlap_events, IO_PARALLEL)
|
|
|
|
gds_buffer = torch.empty(IO_SIZE, dtype=torch.uint8, device=get_accelerator().device_name())
|
|
h.pin_device_tensor(gds_buffer)
|
|
|
|
_validate_handle_state(h, single_submit, overlap_events)
|
|
|
|
ref_file, _ = _do_ref_write(tmpdir)
|
|
read_status = h.sync_pread(gds_buffer, ref_file, 0)
|
|
assert read_status == 1
|
|
|
|
with open(ref_file, 'rb') as f:
|
|
ref_buffer = list(f.read())
|
|
assert ref_buffer == gds_buffer.tolist()
|
|
|
|
h.unpin_device_tensor(gds_buffer)
|
|
|
|
def test_async_read(self, tmpdir, single_submit, overlap_events):
|
|
|
|
h = GDSBuilder().load().gds_handle(BLOCK_SIZE, QUEUE_DEPTH, single_submit, overlap_events, IO_PARALLEL)
|
|
|
|
gds_buffer = torch.empty(IO_SIZE, dtype=torch.uint8, device=get_accelerator().device_name())
|
|
h.pin_device_tensor(gds_buffer)
|
|
|
|
_validate_handle_state(h, single_submit, overlap_events)
|
|
|
|
ref_file, _ = _do_ref_write(tmpdir)
|
|
read_status = h.async_pread(gds_buffer, ref_file, 0)
|
|
assert read_status == 0
|
|
|
|
wait_status = h.wait()
|
|
assert wait_status == 1
|
|
|
|
with open(ref_file, 'rb') as f:
|
|
ref_buffer = list(f.read())
|
|
assert ref_buffer == gds_buffer.tolist()
|
|
|
|
h.unpin_device_tensor(gds_buffer)
|
|
|
|
|
|
@pytest.mark.parametrize("single_submit", [True, False])
|
|
@pytest.mark.parametrize("overlap_events", [True, False])
|
|
class TestWrite(DistributedTest):
|
|
world_size = 1
|
|
reuse_dist_env = True
|
|
if not get_accelerator().is_available():
|
|
init_distributed = False
|
|
set_dist_env = False
|
|
|
|
def test_parallel_write(self, tmpdir, single_submit, overlap_events):
|
|
|
|
ref_file, ref_buffer = _do_ref_write(tmpdir)
|
|
h = GDSBuilder().load().gds_handle(BLOCK_SIZE, QUEUE_DEPTH, single_submit, overlap_events, IO_PARALLEL)
|
|
|
|
gds_file, gds_buffer = _get_test_write_file_and_device_buffer(tmpdir, ref_buffer, h)
|
|
|
|
_validate_handle_state(h, single_submit, overlap_events)
|
|
|
|
write_status = h.sync_pwrite(gds_buffer, gds_file, 0)
|
|
assert write_status == 1
|
|
|
|
h.unpin_device_tensor(gds_buffer)
|
|
|
|
assert os.path.isfile(gds_file)
|
|
|
|
filecmp.clear_cache()
|
|
assert filecmp.cmp(ref_file, gds_file, shallow=False)
|
|
|
|
def test_async_write(self, tmpdir, single_submit, overlap_events):
|
|
ref_file, ref_buffer = _do_ref_write(tmpdir)
|
|
|
|
h = GDSBuilder().load().gds_handle(BLOCK_SIZE, QUEUE_DEPTH, single_submit, overlap_events, IO_PARALLEL)
|
|
gds_file, gds_buffer = _get_test_write_file_and_device_buffer(tmpdir, ref_buffer, h)
|
|
|
|
_validate_handle_state(h, single_submit, overlap_events)
|
|
|
|
write_status = h.async_pwrite(gds_buffer, gds_file, 0)
|
|
assert write_status == 0
|
|
|
|
wait_status = h.wait()
|
|
assert wait_status == 1
|
|
|
|
h.unpin_device_tensor(gds_buffer)
|
|
|
|
assert os.path.isfile(gds_file)
|
|
|
|
filecmp.clear_cache()
|
|
assert filecmp.cmp(ref_file, gds_file, shallow=False)
|
|
|
|
|
|
@pytest.mark.sequential
|
|
class TestAsyncQueue(DistributedTest):
|
|
world_size = 1
|
|
if not get_accelerator().is_available():
|
|
init_distributed = False
|
|
set_dist_env = False
|
|
|
|
@pytest.mark.parametrize("async_queue", [2, 3])
|
|
def test_read(self, tmpdir, async_queue):
|
|
|
|
ref_files = []
|
|
for i in range(async_queue):
|
|
f, _ = _do_ref_write(tmpdir, i)
|
|
ref_files.append(f)
|
|
|
|
single_submit = True
|
|
overlap_events = True
|
|
h = GDSBuilder().load().gds_handle(BLOCK_SIZE, QUEUE_DEPTH, single_submit, overlap_events, IO_PARALLEL)
|
|
|
|
gds_buffers = [
|
|
torch.empty(IO_SIZE, dtype=torch.uint8, device=get_accelerator().device_name()) for _ in range(async_queue)
|
|
]
|
|
for buf in gds_buffers:
|
|
h.pin_device_tensor(buf)
|
|
|
|
_validate_handle_state(h, single_submit, overlap_events)
|
|
|
|
for i in range(async_queue):
|
|
read_status = h.async_pread(gds_buffers[i], ref_files[i], 0)
|
|
assert read_status == 0
|
|
|
|
wait_status = h.wait()
|
|
assert wait_status == async_queue
|
|
|
|
for i in range(async_queue):
|
|
with open(ref_files[i], 'rb') as f:
|
|
ref_buffer = list(f.read())
|
|
assert ref_buffer == gds_buffers[i].tolist()
|
|
|
|
for t in gds_buffers:
|
|
h.unpin_device_tensor(t)
|
|
|
|
@pytest.mark.parametrize("async_queue", [2, 3])
|
|
def test_write(self, tmpdir, async_queue):
|
|
ref_files = []
|
|
ref_buffers = []
|
|
for i in range(async_queue):
|
|
f, buf = _do_ref_write(tmpdir, i)
|
|
ref_files.append(f)
|
|
ref_buffers.append(buf)
|
|
|
|
single_submit = True
|
|
overlap_events = True
|
|
h = GDSBuilder().load().gds_handle(BLOCK_SIZE, QUEUE_DEPTH, single_submit, overlap_events, IO_PARALLEL)
|
|
|
|
gds_files = []
|
|
gds_buffers = []
|
|
for i in range(async_queue):
|
|
f, buf = _get_test_write_file_and_device_buffer(tmpdir, ref_buffers[i], h, i)
|
|
gds_files.append(f)
|
|
gds_buffers.append(buf)
|
|
|
|
_validate_handle_state(h, single_submit, overlap_events)
|
|
|
|
for i in range(async_queue):
|
|
read_status = h.async_pwrite(gds_buffers[i], gds_files[i], 0)
|
|
assert read_status == 0
|
|
|
|
wait_status = h.wait()
|
|
assert wait_status == async_queue
|
|
|
|
for t in gds_buffers:
|
|
h.unpin_device_tensor(t)
|
|
|
|
for i in range(async_queue):
|
|
assert os.path.isfile(gds_files[i])
|
|
|
|
filecmp.clear_cache()
|
|
assert filecmp.cmp(ref_files[i], gds_files[i], shallow=False)
|
|
|
|
|
|
@pytest.mark.parametrize("use_new_api", [True, False])
|
|
class TestLockDeviceTensor(DistributedTest):
|
|
world_size = 2
|
|
reuse_dist_env = True
|
|
if not get_accelerator().is_available():
|
|
init_distributed = False
|
|
set_dist_env = False
|
|
|
|
def test_pin_device_tensor(self, use_new_api):
|
|
|
|
h = GDSBuilder().load().gds_handle(BLOCK_SIZE, QUEUE_DEPTH, True, True, IO_PARALLEL)
|
|
|
|
unpinned_buffer = torch.empty(IO_SIZE, dtype=torch.uint8, device=get_accelerator().device_name())
|
|
if use_new_api:
|
|
pinned_buffer = h.new_pinned_device_tensor(unpinned_buffer.numel(), unpinned_buffer)
|
|
else:
|
|
pinned_buffer = torch.empty_like(unpinned_buffer)
|
|
h.pin_device_tensor(pinned_buffer)
|
|
|
|
assert unpinned_buffer.device == pinned_buffer.device
|
|
assert unpinned_buffer.dtype == pinned_buffer.dtype
|
|
assert unpinned_buffer.numel() == pinned_buffer.numel()
|
|
|
|
if use_new_api:
|
|
h.free_pinned_device_tensor(pinned_buffer)
|
|
else:
|
|
h.unpin_device_tensor(pinned_buffer)
|
|
|
|
|
|
@pytest.mark.parametrize('file_partitions', [[1, 1, 1], [1, 1, 2], [1, 2, 1], [2, 1, 1]])
|
|
class TestAsyncFileOffset(DistributedTest):
|
|
world_size = 1
|
|
|
|
@pytest.mark.parametrize('use_fd', [False, True])
|
|
def test_offset_write(self, tmpdir, use_fd, file_partitions):
|
|
ref_file = _get_file_path(tmpdir, '_py_random')
|
|
aio_file = _get_file_path(tmpdir, '_aio_random')
|
|
partition_unit_size = IO_SIZE
|
|
file_size = sum(file_partitions) * partition_unit_size
|
|
|
|
h = GDSBuilder().load().gds_handle(BLOCK_SIZE, QUEUE_DEPTH, True, True, IO_PARALLEL)
|
|
|
|
gds_buffer = torch.empty(file_size, dtype=torch.uint8, device=get_accelerator().device_name())
|
|
h.pin_device_tensor(gds_buffer)
|
|
|
|
file_offsets = []
|
|
next_offset = 0
|
|
for i in range(len(file_partitions)):
|
|
file_offsets.append(next_offset)
|
|
next_offset += file_partitions[i] * partition_unit_size
|
|
|
|
ref_fd = open(ref_file, 'wb')
|
|
for i in range(len(file_partitions)):
|
|
src_buffer = torch.narrow(gds_buffer, 0, file_offsets[i],
|
|
file_partitions[i] * partition_unit_size).to(device='cpu')
|
|
|
|
ref_fd.write(src_buffer.numpy().tobytes())
|
|
ref_fd.flush()
|
|
|
|
if use_fd:
|
|
aio_fd = os.open(aio_file, flags=os.O_DIRECT | os.O_CREAT | os.O_WRONLY)
|
|
write_status = h.async_pwrite(buffer=src_buffer, fd=aio_fd, file_offset=file_offsets[i])
|
|
else:
|
|
write_status = h.async_pwrite(buffer=src_buffer, filename=aio_file, file_offset=file_offsets[i])
|
|
|
|
assert write_status == 0
|
|
wait_status = h.wait()
|
|
assert wait_status == 1
|
|
|
|
if use_fd:
|
|
assert os.path.isfile(aio_fd)
|
|
os.close(aio_fd)
|
|
|
|
filecmp.clear_cache()
|
|
assert filecmp.cmp(ref_file, aio_file, shallow=False)
|
|
|
|
ref_fd.close()
|
|
|
|
h.unpin_device_tensor(gds_buffer)
|
|
|
|
def test_offset_read(self, tmpdir, file_partitions):
|
|
partition_unit_size = BLOCK_SIZE
|
|
file_size = sum(file_partitions) * partition_unit_size
|
|
ref_file, _ = _do_ref_write(tmpdir, 0, file_size)
|
|
h = GDSBuilder().load().gds_handle(BLOCK_SIZE, QUEUE_DEPTH, True, True, IO_PARALLEL)
|
|
|
|
gds_buffer = torch.empty(file_size, dtype=torch.uint8, device=get_accelerator().device_name())
|
|
h.pin_device_tensor(gds_buffer)
|
|
|
|
file_offsets = []
|
|
next_offset = 0
|
|
for i in range(len(file_partitions)):
|
|
file_offsets.append(next_offset)
|
|
next_offset += file_partitions[i] * partition_unit_size
|
|
|
|
with open(ref_file, 'rb') as ref_fd:
|
|
for i in range(len(file_partitions)):
|
|
ref_fd.seek(file_offsets[i])
|
|
bytes_to_read = file_partitions[i] * partition_unit_size
|
|
ref_buf = list(ref_fd.read(bytes_to_read))
|
|
|
|
dst_tensor = torch.narrow(gds_buffer, 0, 0, bytes_to_read)
|
|
read_status = h.async_pread(dst_tensor, ref_file, file_offsets[i])
|
|
assert read_status == 0
|
|
wait_status = h.wait()
|
|
assert wait_status == 1
|
|
assert dst_tensor.tolist() == ref_buf
|
|
|
|
h.unpin_device_tensor(gds_buffer)
|