Method test_offset_write

tests/unit/ops/aio/test_aio.py:334–381

(self, tmpdir, file_partitions, use_cuda_pinned_tensor, use_fd)

    @pytest.mark.parametrize('use_fd', [False, True])
    def test_offset_write(self, tmpdir, file_partitions, use_cuda_pinned_tensor, use_fd):

        _skip_for_invalid_environment(use_cuda_pinned_tensor=use_cuda_pinned_tensor)
        ref_file = _get_file_path(tmpdir, '_py_random')
        aio_file = _get_file_path(tmpdir, '_aio_random')
        partition_unit_size = BLOCK_SIZE
        file_size = sum(file_partitions) * partition_unit_size

        h = AsyncIOBuilder().load().aio_handle(BLOCK_SIZE, QUEUE_DEPTH, True, True, IO_PARALLEL)

        if use_cuda_pinned_tensor:
            data_buffer = torch.ByteTensor(list(os.urandom(file_size))).pin_memory()
        else:
            data_buffer = h.new_cpu_locked_tensor(file_size, torch.empty(0, dtype=torch.uint8))

        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(data_buffer, 0, file_offsets[i], file_partitions[i] * partition_unit_size)

            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:
                os.path.isfile(aio_fd)
                os.close(aio_fd)

            filecmp.clear_cache()
            assert filecmp.cmp(ref_file, aio_file, shallow=False)

        ref_fd.close()

        if not use_cuda_pinned_tensor:
            h.free_cpu_locked_tensor(data_buffer)

    def test_offset_read(self, tmpdir, file_partitions, use_cuda_pinned_tensor):