def device_host_file_size_matches(dhf,
total_bytes,
device_chunk_overhead=0,
serialized_chunk_overhead=1024):
byte_sum = dhf.device_buffer.fast.total_weight
# `dhf.host_buffer.fast` is only available when Worker's `memory_limit != 0`
if hasattr(dhf.host_buffer, "fast"):
byte_sum += dhf.host_buffer.fast.total_weight
else:
byte_sum += sum([sizeof(b) for b in dhf.host_buffer.values()])
# `dhf.disk` is only available when Worker's `memory_limit != 0`
if dhf.disk is not None:
file_path = [
os.path.join(dhf.disk.directory, safe_key(k))
for k in dhf.disk.keys()
]
file_size = [os.path.getsize(f) for f in file_path]
byte_sum += sum(file_size)
# Allow up to chunk_overhead bytes overhead per chunk
device_overhead = len(dhf.device) * device_chunk_overhead
host_overhead = len(dhf.host) * serialized_chunk_overhead
disk_overhead = (len(dhf.disk) *
serialized_chunk_overhead if dhf.disk is not None else 0)
return (byte_sum >= total_bytes and byte_sum <=
total_bytes + device_overhead + host_overhead + disk_overhead)
def assert_device_host_file_size(dhf,
total_bytes,
device_chunk_overhead=0,
serialized_chunk_overhead=1024):
byte_sum = dhf.device_buffer.fast.total_weight + dhf.host_buffer.fast.total_weight
file_path = [
os.path.join(dhf.disk.directory, safe_key(k)) for k in dhf.disk.keys()
]
file_size = [os.path.getsize(f) for f in file_path]
byte_sum += sum(file_size)
# Allow up to chunk_overhead bytes overhead per chunk on disk
device_overhead = len(dhf.device) * device_chunk_overhead
host_overhead = len(dhf.host) * serialized_chunk_overhead
disk_overhead = len(dhf.disk) * serialized_chunk_overhead
assert (byte_sum >= total_bytes and byte_sum <=
total_bytes + device_overhead + host_overhead + disk_overhead)