Write model-specific feature dictionaries as resumable TensorDict shards.
Pipelines publish a dict[str, Tensor] with a common batch dimension; this writer handles buffering, per-rank shard names, metadata, and resume discovery without knowing what any feature means.
Source code in fastvideo/dataset/v2a_feature_cache.py
| def __init__(
self,
output_root: str | Path,
*,
rank: int,
samples_per_shard: int,
) -> None:
if samples_per_shard <= 0:
raise ValueError("samples_per_shard must be positive")
self.output_root = Path(output_root).expanduser().resolve()
self.worker_root = self.output_root / f"worker_{rank:05d}"
self.worker_root.mkdir(parents=True, exist_ok=True)
self.samples_per_shard = int(samples_per_shard)
self._features: dict[str, list[torch.Tensor]] = defaultdict(list)
self._metadata: list[dict[str, Any]] = []
self.existing_ids = self._read_existing_ids(self.output_root)
existing_indices = []
for path in self.worker_root.glob("shard_*"):
try:
existing_indices.append(int(path.name.split("_")[-1]))
except ValueError:
continue
self._next_shard = max(existing_indices, default=-1) + 1
|