|
import threading |
|
from typing import MutableMapping, NamedTuple |
|
|
|
import wandb |
|
|
|
|
|
class FileStats(NamedTuple): |
|
deduped: bool |
|
total: int |
|
uploaded: int |
|
failed: bool |
|
artifact_file: bool |
|
|
|
|
|
class Summary(NamedTuple): |
|
uploaded_bytes: int |
|
total_bytes: int |
|
deduped_bytes: int |
|
|
|
|
|
class FileCountsByCategory(NamedTuple): |
|
artifact: int |
|
wandb: int |
|
media: int |
|
other: int |
|
|
|
|
|
class Stats: |
|
def __init__(self) -> None: |
|
self._stats: MutableMapping[str, FileStats] = {} |
|
self._lock = threading.Lock() |
|
|
|
def init_file( |
|
self, save_name: str, size: int, is_artifact_file: bool = False |
|
) -> None: |
|
with self._lock: |
|
self._stats[save_name] = FileStats( |
|
deduped=False, |
|
total=size, |
|
uploaded=0, |
|
failed=False, |
|
artifact_file=is_artifact_file, |
|
) |
|
|
|
def set_file_deduped(self, save_name: str) -> None: |
|
with self._lock: |
|
orig = self._stats[save_name] |
|
self._stats[save_name] = orig._replace( |
|
deduped=True, |
|
uploaded=orig.total, |
|
) |
|
|
|
def update_uploaded_file(self, save_name: str, total_uploaded: int) -> None: |
|
with self._lock: |
|
self._stats[save_name] = self._stats[save_name]._replace( |
|
uploaded=total_uploaded, |
|
) |
|
|
|
def update_failed_file(self, save_name: str) -> None: |
|
with self._lock: |
|
self._stats[save_name] = self._stats[save_name]._replace( |
|
uploaded=0, |
|
failed=True, |
|
) |
|
|
|
def summary(self) -> Summary: |
|
|
|
|
|
with self._lock: |
|
stats = list(self._stats.values()) |
|
return Summary( |
|
uploaded_bytes=sum(f.uploaded for f in stats), |
|
total_bytes=sum(f.total for f in stats), |
|
deduped_bytes=sum(f.total for f in stats if f.deduped), |
|
) |
|
|
|
def file_counts_by_category(self) -> FileCountsByCategory: |
|
artifact_files = 0 |
|
wandb_files = 0 |
|
media_files = 0 |
|
other_files = 0 |
|
|
|
|
|
with self._lock: |
|
file_stats = list(self._stats.items()) |
|
for save_name, stats in file_stats: |
|
if stats.artifact_file: |
|
artifact_files += 1 |
|
elif wandb.wandb_lib.filenames.is_wandb_file(save_name): |
|
wandb_files += 1 |
|
elif save_name.startswith("media"): |
|
media_files += 1 |
|
else: |
|
other_files += 1 |
|
return FileCountsByCategory( |
|
artifact=artifact_files, |
|
wandb=wandb_files, |
|
media=media_files, |
|
other=other_files, |
|
) |
|
|