MCPcopy
hub / github.com/deepspeedai/DeepSpeed / test_offset_read

Method test_offset_read

tests/unit/ops/aio/test_aio.py:383–413  ·  view source on GitHub ↗
(self, tmpdir, file_partitions, use_cuda_pinned_tensor)

Source from the content-addressed store, hash-verified

381 h.free_cpu_locked_tensor(data_buffer)
382
383 def test_offset_read(self, tmpdir, file_partitions, use_cuda_pinned_tensor):
384
385 _skip_for_invalid_environment(use_cuda_pinned_tensor=use_cuda_pinned_tensor)
386 partition_unit_size = BLOCK_SIZE
387 file_size = sum(file_partitions) * partition_unit_size
388 ref_file, _ = _do_ref_write(tmpdir, 0, file_size)
389 h = AsyncIOBuilder().load().aio_handle(BLOCK_SIZE, QUEUE_DEPTH, True, True, IO_PARALLEL)
390
391 if use_cuda_pinned_tensor:
392 data_buffer = torch.zeros(file_size, dtype=torch.uint8, device='cpu').pin_memory()
393 else:
394 data_buffer = h.new_cpu_locked_tensor(file_size, torch.empty(0, dtype=torch.uint8))
395
396 file_offsets = []
397 next_offset = 0
398 for i in range(len(file_partitions)):
399 file_offsets.append(next_offset)
400 next_offset += file_partitions[i] * partition_unit_size
401
402 with open(ref_file, 'rb') as ref_fd:
403 for i in range(len(file_partitions)):
404 ref_fd.seek(file_offsets[i])
405 bytes_to_read = file_partitions[i] * partition_unit_size
406 ref_buf = list(ref_fd.read(bytes_to_read))
407
408 dst_tensor = torch.narrow(data_buffer, 0, 0, bytes_to_read)
409 assert 1 == h.sync_pread(dst_tensor, ref_file, file_offsets[i])
410 assert dst_tensor.tolist() == ref_buf
411
412 if not use_cuda_pinned_tensor:
413 h.free_cpu_locked_tensor(data_buffer)

Callers

nothing calls this directly

Calls 7

AsyncIOBuilderClass · 0.90
appendMethod · 0.80
_do_ref_writeFunction · 0.70
loadMethod · 0.45
pin_memoryMethod · 0.45
readMethod · 0.45

Tested by

no test coverage detected