|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
from .constants import * |
|
|
from .base_file_writer import BaseFileWriter |
|
|
from .utils import obj_serialization_details |
|
|
|
|
|
|
|
|
class MockFileWriter(BaseFileWriter): |
|
|
|
|
|
def __init__(self, file_path): |
|
|
super(MockFileWriter, self).__init__(file_path) |
|
|
self._fp = open(file_path, 'wb') |
|
|
self._stats[SAVE_STORAGE_KEY] = 0 |
|
|
self._stats[SAVE_STORAGE_BYTES_KEY] = 0 |
|
|
self._get_serialization_details = obj_serialization_details() |
|
|
|
|
|
def close(self): |
|
|
self._incr_stats(CLOSE_COUNT_KEY) |
|
|
self._fp.close() |
|
|
|
|
|
def fileno(self): |
|
|
self._incr_stats(FILENO_COUNT_KEY) |
|
|
return INVALID_FD |
|
|
|
|
|
def flush(self): |
|
|
self._incr_stats(FLUSH_COUNT_KEY) |
|
|
self._fp.flush() |
|
|
|
|
|
def write(self, buffer): |
|
|
return self._write(len(buffer)) |
|
|
|
|
|
def save_torch_storage_object_list(self, storage_obj_list, save_size): |
|
|
num_bytes = sum([self._save_torch_storage_object(obj, save_size) for obj in storage_obj_list]) |
|
|
return num_bytes |
|
|
|
|
|
def _save_torch_storage_object(self, storage_obj, save_size): |
|
|
details = self._get_serialization_details(storage_obj) |
|
|
self._incr_stats(SAVE_STORAGE_KEY) |
|
|
self._incr_stats(SAVE_STORAGE_BYTES_KEY, details.size) |
|
|
num_written_bytes = self._write(STORAGE_OBJ_SIZE) if save_size else 0 |
|
|
return num_written_bytes + self._write(details.size) |
|
|
|
|
|
def _write(self, num_bytes): |
|
|
self._incr_stats(WRITE_COUNT_KEY) |
|
|
self._incr_stats(WRITE_BYTES_KEY, num_bytes) |
|
|
return num_bytes |
|
|
|
