wehub-resource-sync
78 lines
2.4 KiB
Python
78 lines
2.4 KiB
Python
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import List
|
|
from datasets import Dataset, concatenate_datasets
|
|
import logging
|
|
import shutil
|
|
|
|
logger = logging.getLogger("raft")
|
|
|
|
@dataclass
|
|
class Checkpoint:
|
|
path: Path
|
|
num: int
|
|
|
|
def load(self) -> Dataset:
|
|
return Dataset.load_from_disk(self.path)
|
|
|
|
def __lt__(self, other: 'Checkpoint') -> bool:
|
|
return self.num < other.num
|
|
|
|
def __eq__(self, other: 'Checkpoint') -> bool:
|
|
return self.num == other.num
|
|
|
|
def __hash__(self) -> int:
|
|
return hash(self.num)
|
|
|
|
class Checkpointing:
|
|
|
|
def __init__(self, checkpoints_dir: Path) -> None:
|
|
self.checkpoints_dir = checkpoints_dir
|
|
|
|
def missing_checkpoints(self, num) -> List[int]:
|
|
return [n for n in range(0, num) if not (self.checkpoints_dir / f"checkpoint-{n}").exists()]
|
|
|
|
def save_checkpoint(self, ds: Dataset, num: int):
|
|
checkpoint_path = self.checkpoints_dir / ("checkpoint-" + str(num))
|
|
ds.save_to_disk(checkpoint_path)
|
|
|
|
def load_checkpoint(self, num: int):
|
|
checkpoint_path = self.checkpoints_dir / ("checkpoint-" + str(num))
|
|
if checkpoint_path.exists():
|
|
return Dataset.load_from_disk(checkpoint_path)
|
|
return None
|
|
|
|
def get_checkpoints(self) -> List[Checkpoint]:
|
|
checkpoints = []
|
|
if not self.checkpoints_dir.exists():
|
|
return checkpoints
|
|
for dir_path in self.checkpoints_dir.iterdir():
|
|
if dir_path.is_dir() and dir_path.name.startswith("checkpoint-"):
|
|
num = int(dir_path.name.split("-")[1])
|
|
checkpoints.append(Checkpoint(dir_path, num))
|
|
return checkpoints
|
|
|
|
def has_checkpoints(self) -> bool:
|
|
return len(self.get_checkpoints()) > 0
|
|
|
|
def collect_checkpoints(self) -> Dataset:
|
|
ds_list = list([checkpoint.load() for checkpoint in self.get_checkpoints()])
|
|
ds = concatenate_datasets(ds_list)
|
|
return ds
|
|
|
|
def delete_checkpoints(self):
|
|
shutil.rmtree(self.checkpoints_dir)
|
|
|
|
def checkpointed(checkpointing: Checkpointing):
|
|
def wrapped(func):
|
|
def wrapper(chunk_id, *args, **kwargs):
|
|
ds = checkpointing.load_checkpoint(chunk_id)
|
|
if ds:
|
|
return ds
|
|
ds = func(chunk_id=chunk_id, *args, **kwargs)
|
|
if ds.num_rows > 0:
|
|
checkpointing.save_checkpoint(ds, chunk_id)
|
|
return ds
|
|
return wrapper
|
|
return wrapped
|