jamtur01's picture
Upload folder using huggingface_hub
9c6594c verified
# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0
# DeepSpeed Team
from .constants import *
BASE_STAT_KEYS = [
CLOSE_COUNT_KEY, FILENO_COUNT_KEY, FLUSH_COUNT_KEY, WRITE_COUNT_KEY, WRITE_BYTES_KEY, WRITE_SEC_KEY,
WRITE_SPEED_KEY
]
class BaseFileWriter(object):
def __init__(self, file_path):
self._file_path = file_path
self._stats = {k: 0 for k in BASE_STAT_KEYS}
def close(self):
pass
def fileno(self):
pass
def flush(self):
pass
def write(self, buffer):
pass
def file_path(self):
return self._file_path
def _incr_stats(self, key, incr=1):
self._stats[key] += incr
def _dump_state(self):
if self._stats[WRITE_SEC_KEY] > 0:
self._stats[WRITE_SPEED_KEY] = (self._stats[WRITE_BYTES_KEY] / self._stats[WRITE_SEC_KEY] / (1024**3))
state = self._stats
state[FILE_PATH_KEY] = self.file_path()
print(f'stats = {self._stats}')