|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
import torch |
|
|
import os |
|
|
import time |
|
|
from .test_ds_aio_utils import task_log, create_filename, create_file, create_page_locked_tensor |
|
|
from .ds_aio_constants import * |
|
|
|
|
|
|
|
|
class TorchIO_Engine(object): |
|
|
|
|
|
def __init__(self, args, tid, read_op): |
|
|
self.ctxt = self._create_context(args, tid, read_op) |
|
|
self.zipfile_serialization = not args.torch_legacy_save |
|
|
|
|
|
def fini(self): |
|
|
self.ctxt[BUFFER].detach() |
|
|
self.ctxt[BUFFER] = None |
|
|
|
|
|
def read(self, args, tid): |
|
|
start_time = time.time() |
|
|
torch.load(f=self.ctxt[FILE], map_location=self.ctxt[BUFFER].device) |
|
|
end_time = time.time() |
|
|
self.ctxt[ELAPSED_SEC] += end_time - start_time |
|
|
|
|
|
def write(self, args, tid): |
|
|
|
|
|
if os.path.isfile(self.ctxt[FILE]): |
|
|
os.remove(self.ctxt[FILE]) |
|
|
|
|
|
start_time = time.time() |
|
|
torch.save(obj=self.ctxt[BUFFER], f=self.ctxt[FILE], _use_new_zipfile_serialization=self.zipfile_serialization) |
|
|
end_time = time.time() |
|
|
self.ctxt[ELAPSED_SEC] += end_time - start_time |
|
|
|
|
|
def _create_context(self, args, tid, read_op): |
|
|
io_string = "Read" if read_op else "Write" |
|
|
device_id, folder = args.mapping_list[tid] |
|
|
filename = create_filename(folder, args.read, args.io_size, tid) |
|
|
if args.read and not (os.path.isfile(filename) and os.path.getsize(filename) == args.io_size): |
|
|
create_file(filename, args.io_size) |
|
|
|
|
|
task_log(tid, f'Allocate tensor of size {args.io_size} bytes') |
|
|
|
|
|
if args.gpu: |
|
|
buffer = torch.randint(high=128, size=(args.io_size, ), dtype=torch.uint8, device=f'cuda:{device_id}') |
|
|
else: |
|
|
buffer = create_page_locked_tensor(args.io_size, True) |
|
|
|
|
|
task_log(tid, |
|
|
f'{io_string} file {filename} of size {args.io_size} bytes from buffer on device {buffer.device}', |
|
|
force=True) |
|
|
|
|
|
task_log(tid, f'created torch_io engine') |
|
|
|
|
|
ctxt = {} |
|
|
ctxt[FILE] = filename |
|
|
ctxt[NUM_BYTES] = args.io_size |
|
|
ctxt[BUFFER] = buffer |
|
|
ctxt[ELAPSED_SEC] = 0 |
|
|
return ctxt |
|
|
|
