diff --git a/fastvideo/dataset/__init__.py b/fastvideo/dataset/__init__.py index b581b9219..a5f1feb72 100644 --- a/fastvideo/dataset/__init__.py +++ b/fastvideo/dataset/__init__.py @@ -2,17 +2,10 @@ from torchvision import transforms from torchvision.transforms import Lambda -from fastvideo.dataset.parquet_dataset_map_style import ( - build_parquet_map_style_dataloader) -from fastvideo.dataset.parquet_dataset_streaming_style import ( - LatentsParquetStreamingDataset, - build_parquet_streaming_style_dataloader, -) -from fastvideo.dataset.ltx2_precomputed_dataset import ( - build_ltx2_precomputed_dataloader, LTX2PrecomputedDataset) +from fastvideo.dataset.parquet_dataset_map_style import build_parquet_map_style_dataloader +from fastvideo.dataset.ltx2_precomputed_dataset import build_ltx2_precomputed_dataloader, LTX2PrecomputedDataset from fastvideo.dataset.preprocessing_datasets import VideoCaptionMergedDataset, TextDataset -from fastvideo.dataset.transform import (CenterCropResizeVideo, Normalize255, - TemporalRandomCrop) +from fastvideo.dataset.transform import CenterCropResizeVideo, Normalize255, TemporalRandomCrop from fastvideo.dataset.validation_dataset import ValidationDataset @@ -25,27 +18,31 @@ def getdataset(args) -> VideoCaptionMergedDataset: resize = [ CenterCropResizeVideo((args.max_height, args.max_width)), ] - transform = transforms.Compose([ - # Normalize255(), - *resize, - ]) - transform_topcrop = transforms.Compose([ - Normalize255(), - *resize_topcrop, - norm_fun, - ]) - return VideoCaptionMergedDataset(data_merge_path=args.data_merge_path, - args=args, - transform=transform, - temporal_sample=temporal_sample, - transform_topcrop=transform_topcrop, - seed=args.seed) - + transform = transforms.Compose( + [ + # Normalize255(), + *resize, + ] + ) + transform_topcrop = transforms.Compose( + [ + Normalize255(), + *resize_topcrop, + norm_fun, + ] + ) + return VideoCaptionMergedDataset( + data_merge_path=args.data_merge_path, + args=args, + transform=transform, + temporal_sample=temporal_sample, + transform_topcrop=transform_topcrop, + seed=args.seed, + ) + def gettextdataset(args) -> TextDataset: - return TextDataset(data_merge_path=args.data_merge_path, - args=args, - seed=args.seed) + return TextDataset(data_merge_path=args.data_merge_path, args=args, seed=args.seed) __all__ = [ @@ -58,3 +55,14 @@ __all__ = [ "VideoCaptionMergedDataset", "TextDataset", ] + + +def __getattr__(name: str): + if name in { + "LatentsParquetStreamingDataset", + "build_parquet_streaming_style_dataloader", + }: + from fastvideo.dataset import parquet_dataset_streaming_style + + return getattr(parquet_dataset_streaming_style, name) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/fastvideo/dataset/benchmarks/benchmark_parquet_streaming_style.py b/fastvideo/dataset/benchmarks/benchmark_parquet_streaming_style.py index 8c2dd4e2c..fabc7bc33 100644 --- a/fastvideo/dataset/benchmarks/benchmark_parquet_streaming_style.py +++ b/fastvideo/dataset/benchmarks/benchmark_parquet_streaming_style.py @@ -15,7 +15,8 @@ import time from fastvideo.dataset.dataloader.schema import pyarrow_schema_t2v from fastvideo.dataset.parquet_dataset_streaming_style import ( - LatentsParquetStreamingDataset, ) + LatentsParquetStreamingDataset, +) def main() -> None: @@ -26,8 +27,7 @@ def main() -> None: parser.add_argument("--read-batch-size", type=int, default=4) parser.add_argument("--batches", type=int, default=4) args = parser.parse_args() - manifest_path = args.manifest_path or os.path.join( - tempfile.gettempdir(), "fastvideo-streaming-benchmark.json") + manifest_path = args.manifest_path or os.path.join(tempfile.gettempdir(), "fastvideo-streaming-benchmark.json") dataset = LatentsParquetStreamingDataset( args.data_path, args.batch_size, @@ -44,13 +44,15 @@ def main() -> None: if index + 1 >= args.batches: break elapsed = time.perf_counter() - started - print({ - "batches": index + 1, - "rows": rows, - "seconds": elapsed, - "rows_per_second": rows / elapsed, - "projected_columns": len(pyarrow_schema_t2v.names), - }) + print( + { + "batches": index + 1, + "rows": rows, + "seconds": elapsed, + "rows_per_second": rows / elapsed, + "projected_columns": len(pyarrow_schema_t2v.names), + } + ) if __name__ == "__main__": diff --git a/fastvideo/dataset/parquet_dataset_streaming_style.py b/fastvideo/dataset/parquet_dataset_streaming_style.py index 1f4641f5c..3fe820852 100644 --- a/fastvideo/dataset/parquet_dataset_streaming_style.py +++ b/fastvideo/dataset/parquet_dataset_streaming_style.py @@ -8,14 +8,12 @@ the requested T2V columns, and consumes each assigned row group sequentially. from __future__ import annotations -import fcntl import hashlib import json import os import random import tempfile from collections.abc import Iterator -from pathlib import Path from typing import Any import pyarrow as pa @@ -25,13 +23,13 @@ from torch.utils.data import IterableDataset, get_worker_info from torchdata.stateful_dataloader import StatefulDataLoader from fastvideo.dataset.utils import collate_rows_from_parquet_schema -from fastvideo.distributed import (get_sp_world_size, get_world_rank, - get_world_size) +from fastvideo.distributed import get_sp_world_size, get_world_rank, get_world_size from fastvideo.logger import init_logger logger = init_logger(__name__) -_MANIFEST_VERSION = 1 +_MANIFEST_VERSION = 2 +_STATE_VERSION = 2 def _barrier() -> None: @@ -47,16 +45,54 @@ def _assert_user_owned_manifest(dataset_root: str, manifest_path: str) -> None: manifest_parent = _real(os.path.dirname(manifest_path) or ".") if os.path.commonpath([dataset_root, manifest_parent]) == dataset_root: raise ValueError( - "streaming_manifest_path must be outside the read-only dataset " - f"root ({dataset_root}); got {manifest_path}") + f"streaming_manifest_path must be outside the read-only dataset root ({dataset_root}); got {manifest_path}" + ) def _manifest_fingerprint(row_groups: list[dict[str, Any]]) -> str: - payload = json.dumps(row_groups, sort_keys=True, - separators=(",", ":")).encode("utf-8") + payload = json.dumps(row_groups, sort_keys=True, separators=(",", ":")).encode("utf-8") return hashlib.sha256(payload).hexdigest() +def _manifest_file_path(dataset_root: str, relative_path: str) -> str: + file_path = _real(os.path.join(dataset_root, relative_path)) + if os.path.commonpath([dataset_root, file_path]) != dataset_root: + raise ValueError(f"Manifest file escapes dataset root: {relative_path!r}") + return file_path + + +def _manifest_matches_source( + dataset_root: str, + manifest: dict[str, Any], +) -> bool: + """Validate cached metadata without reopening every Parquet footer.""" + row_groups = manifest.get("row_groups") + if not isinstance(row_groups, list) or not row_groups: + return False + expected_sample_start = 0 + stats: dict[str, os.stat_result] = {} + try: + for item in row_groups: + relative_path = str(item["file"]) + rows = int(item["rows"]) + row_group = int(item["row_group"]) + sample_start = int(item["sample_start"]) + if rows <= 0 or row_group < 0 or sample_start != expected_sample_start: + return False + expected_sample_start += rows + if relative_path not in stats: + file_path = _manifest_file_path(dataset_root, relative_path) + stats[relative_path] = os.stat(file_path) + stat = stats[relative_path] + if int(item["file_size"]) != stat.st_size or int(item["file_mtime_ns"]) != stat.st_mtime_ns: + return False + except (KeyError, OSError, TypeError, ValueError): + return False + return int(manifest.get("total_rows", -1)) == expected_sample_start and manifest.get( + "fingerprint" + ) == _manifest_fingerprint(row_groups) + + def _scan_manifest(dataset_root: str, columns: tuple[str, ...]) -> dict[str, Any]: parquet_files: list[str] = [] for root, _, files in os.walk(dataset_root): @@ -65,27 +101,29 @@ def _scan_manifest(dataset_root: str, columns: tuple[str, ...]) -> dict[str, Any parquet_files.append(_real(os.path.join(root, name))) parquet_files.sort() if not parquet_files: - raise FileNotFoundError( - f"No parquet files found under dataset path: {dataset_root}") + raise FileNotFoundError(f"No parquet files found under dataset path: {dataset_root}") row_groups: list[dict[str, Any]] = [] sample_start = 0 for file_path in parquet_files: + file_stat = os.stat(file_path) parquet_file = pq.ParquetFile(file_path) missing = sorted(set(columns) - set(parquet_file.schema_arrow.names)) if missing: - raise ValueError( - f"Parquet file {file_path} is missing projected T2V columns: " - f"{missing}") + raise ValueError(f"Parquet file {file_path} is missing projected T2V columns: {missing}") relative_path = os.path.relpath(file_path, dataset_root) for row_group_index in range(parquet_file.num_row_groups): rows = parquet_file.metadata.row_group(row_group_index).num_rows - row_groups.append({ - "file": relative_path, - "row_group": row_group_index, - "rows": rows, - "sample_start": sample_start, - }) + row_groups.append( + { + "file": relative_path, + "row_group": row_group_index, + "rows": rows, + "sample_start": sample_start, + "file_size": file_stat.st_size, + "file_mtime_ns": file_stat.st_mtime_ns, + } + ) sample_start += rows return { @@ -101,8 +139,7 @@ def _scan_manifest(dataset_root: str, columns: tuple[str, ...]) -> dict[str, Any def _write_json_atomic(path: str, payload: dict[str, Any]) -> None: parent = os.path.dirname(path) os.makedirs(parent, exist_ok=True) - fd, temporary_path = tempfile.mkstemp( - dir=parent, prefix=f".{os.path.basename(path)}.", suffix=".tmp") + fd, temporary_path = tempfile.mkstemp(dir=parent, prefix=f".{os.path.basename(path)}.", suffix=".tmp") try: with os.fdopen(fd, "w", encoding="utf-8") as handle: json.dump(payload, handle, sort_keys=True) @@ -115,8 +152,12 @@ def _write_json_atomic(path: str, payload: dict[str, Any]) -> None: def load_or_create_streaming_manifest( - dataset_root: str, manifest_path: str, - columns: tuple[str, ...]) -> dict[str, Any]: + dataset_root: str, manifest_path: str, columns: tuple[str, ...] +) -> dict[str, Any]: + if os.name != "posix": + raise RuntimeError("The streaming Parquet loader currently requires POSIX file locks") + import fcntl + dataset_root = _real(dataset_root) manifest_path = _real(manifest_path) _assert_user_owned_manifest(dataset_root, manifest_path) @@ -129,24 +170,22 @@ def load_or_create_streaming_manifest( if os.path.exists(manifest_path): with open(manifest_path, encoding="utf-8") as handle: candidate = json.load(handle) - if (candidate.get("version") == _MANIFEST_VERSION - and candidate.get("dataset_root") == dataset_root - and candidate.get("columns") == list(columns) - and candidate.get("fingerprint") == - _manifest_fingerprint(candidate.get("row_groups", []))): + if ( + candidate.get("version") == _MANIFEST_VERSION + and candidate.get("dataset_root") == dataset_root + and candidate.get("columns") == list(columns) + and _manifest_matches_source(dataset_root, candidate) + ): manifest = candidate if manifest is None: - logger.info("Building safe JSON Parquet manifest at %s", - manifest_path) + logger.info("Building safe JSON Parquet manifest at %s", manifest_path) manifest = _scan_manifest(dataset_root, columns) _write_json_atomic(manifest_path, manifest) fcntl.flock(lock_handle.fileno(), fcntl.LOCK_UN) return manifest -def _contiguous_row_group_shards( - row_groups: list[dict[str, Any]], num_shards: int -) -> list[list[dict[str, Any]]]: +def _contiguous_row_group_shards(row_groups: list[dict[str, Any]], num_shards: int) -> list[list[dict[str, Any]]]: """Split ordered row groups into row-balanced contiguous shards.""" if num_shards < 1: raise ValueError("num_shards must be positive") @@ -161,8 +200,7 @@ def _contiguous_row_group_shards( item = row_groups[cursor] item_rows = int(item["rows"]) items_left = len(row_groups) - cursor - if (shard_rows and shard_rows + item_rows > target - and items_left >= shards_left - 1): + if shard_rows and shard_rows + item_rows > target and items_left >= shards_left - 1: break shards[shard_index].append(item) shard_rows += item_rows @@ -205,8 +243,8 @@ class LatentsParquetStreamingDataset(IterableDataset): raise ValueError("batch_size and read_batch_size must be positive") if not manifest_path: raise ValueError( - "streaming_manifest_path is required; place it in user-owned " - "storage, never inside the shared dataset") + "streaming_manifest_path is required; place it in user-owned storage, never inside the shared dataset" + ) self.path = _real(path) self.batch_size = int(batch_size) @@ -236,8 +274,7 @@ class LatentsParquetStreamingDataset(IterableDataset): self.sp_group_index = self.global_rank // self.sp_world_size if self.global_rank == 0: - load_or_create_streaming_manifest( - self.path, self.manifest_path, self.columns) + load_or_create_streaming_manifest(self.path, self.manifest_path, self.columns) _barrier() with open(self.manifest_path, encoding="utf-8") as handle: self.manifest = json.load(handle) @@ -248,14 +285,11 @@ class LatentsParquetStreamingDataset(IterableDataset): self.manifest_fingerprint = str(self.manifest["fingerprint"]) num_shards = self.num_sp_groups * self.num_workers - self.shards = _contiguous_row_group_shards( - self.manifest["row_groups"], num_shards) - shard_rows = [sum(int(item["rows"]) for item in shard) - for shard in self.shards] + self.shards = _contiguous_row_group_shards(self.manifest["row_groups"], num_shards) + shard_rows = [sum(int(item["rows"]) for item in shard) for shard in self.shards] self.samples_per_worker = min(shard_rows) // self.batch_size * self.batch_size if self.samples_per_worker == 0: - raise ValueError( - "Dataset is too small for the requested DP/SP/worker layout") + raise ValueError("Dataset is too small for the requested DP/SP/worker layout") self._epoch = 0 self._row_group_position = 0 @@ -264,8 +298,10 @@ class LatentsParquetStreamingDataset(IterableDataset): self._loaded_state: dict[str, Any] | None = None logger.info( "Streaming parquet: %d rows, %d row groups, %d samples per worker", - int(self.manifest["total_rows"]), len(self.manifest["row_groups"]), - self.samples_per_worker) + int(self.manifest["total_rows"]), + len(self.manifest["row_groups"]), + self.samples_per_worker, + ) def _worker_id(self) -> int: info = get_worker_info() @@ -285,14 +321,12 @@ class LatentsParquetStreamingDataset(IterableDataset): if self._loaded_state is not None: saved_worker = int(self._loaded_state.get("worker_id", worker_id)) if saved_worker != worker_id: - raise ValueError( - f"Dataloader state belongs to worker {saved_worker}, not {worker_id}") + raise ValueError(f"Dataloader state belongs to worker {saved_worker}, not {worker_id}") self._loaded_state = None shard = self._worker_shard(worker_id) rows: list[dict[str, Any]] = [] - while (self._yielded_samples < self.samples_per_worker - and self._row_group_position < len(shard)): + while self._yielded_samples < self.samples_per_worker and self._row_group_position < len(shard): item = shard[self._row_group_position] file_path = os.path.join(self.path, item["file"]) parquet_file = pq.ParquetFile(file_path) @@ -310,8 +344,7 @@ class LatentsParquetStreamingDataset(IterableDataset): consumed_in_group = record_end continue start = max(self._row_offset - consumed_in_group, 0) - for local_index, row in enumerate( - record_batch.slice(start).to_pylist(), start=start): + for local_index, row in enumerate(record_batch.slice(start).to_pylist(), start=start): global_offset = consumed_in_group + local_index row["_sample_index"] = int(item["sample_start"]) + global_offset rows.append(row) @@ -337,8 +370,7 @@ class LatentsParquetStreamingDataset(IterableDataset): self._row_offset = 0 if self._yielded_samples != self.samples_per_worker: - raise RuntimeError( - "Streaming shard ended before producing the balanced sample count") + raise RuntimeError("Streaming shard ended before producing the balanced sample count") self._epoch += 1 self._row_group_position = 0 self._row_offset = 0 @@ -349,8 +381,9 @@ class LatentsParquetStreamingDataset(IterableDataset): def state_dict(self) -> dict[str, Any]: return { - "version": 1, + "version": _STATE_VERSION, "manifest_fingerprint": self.manifest_fingerprint, + "topology": self._resume_topology(), "worker_id": self._worker_id(), "epoch": self._epoch, "row_group_position": self._row_group_position, @@ -359,28 +392,47 @@ class LatentsParquetStreamingDataset(IterableDataset): } def load_state_dict(self, state_dict: dict[str, Any]) -> None: + if int(state_dict.get("version", -1)) != _STATE_VERSION: + raise ValueError("Cannot resume: unsupported streaming dataset state") if state_dict.get("manifest_fingerprint") != self.manifest_fingerprint: raise ValueError("Cannot resume: streaming dataset manifest changed") + if state_dict.get("topology") != self._resume_topology(): + raise ValueError("Cannot resume: streaming loader topology or sampling config changed") self._epoch = int(state_dict["epoch"]) self._row_group_position = int(state_dict["row_group_position"]) self._row_offset = int(state_dict["row_offset"]) self._yielded_samples = int(state_dict["yielded_samples"]) self._loaded_state = dict(state_dict) - def get_validation_negative_prompt( - self) -> tuple[Any, Any, str]: + def _resume_topology(self) -> dict[str, Any]: + return { + "global_rank": self.global_rank, + "world_size": self.world_size, + "sp_world_size": self.sp_world_size, + "sp_group_index": self.sp_group_index, + "num_workers": self.num_workers, + "batch_size": self.batch_size, + "read_batch_size": self.read_batch_size, + "seed": self.seed, + "shuffle_row_groups": self.shuffle_row_groups, + "samples_per_worker": self.samples_per_worker, + } + + def get_validation_negative_prompt(self) -> tuple[Any, Any, str]: first = self.manifest["row_groups"][0] parquet_file = pq.ParquetFile(os.path.join(self.path, first["file"])) - first_batch = next(parquet_file.iter_batches( - batch_size=1, - row_groups=[int(first["row_group"])], - columns=list(self.columns), - use_threads=False, - )) + first_batch = next( + parquet_file.iter_batches( + batch_size=1, + row_groups=[int(first["row_group"])], + columns=list(self.columns), + use_threads=False, + ) + ) row = first_batch.to_pylist()[0] batch = collate_rows_from_parquet_schema( - [row], self.parquet_schema, self.text_padding_length, cfg_rate=0.0, - seed=self.seed) + [row], self.parquet_schema, self.text_padding_length, cfg_rate=0.0, seed=self.seed + ) embedding = batch["text_embedding"] mask = batch["text_attention_mask"] prompt = batch["info_list"][0]["prompt"] diff --git a/fastvideo/tests/dataset/test_parquet_dataset_streaming_style.py b/fastvideo/tests/dataset/test_parquet_dataset_streaming_style.py index 38846a1af..5faf764ec 100644 --- a/fastvideo/tests/dataset/test_parquet_dataset_streaming_style.py +++ b/fastvideo/tests/dataset/test_parquet_dataset_streaming_style.py @@ -41,15 +41,16 @@ def parquet_root(tmp_path: Path) -> Path: root = tmp_path / "shared-read-only-dataset" root.mkdir() schema = pyarrow_schema_t2v.append(pa.field("unused_large_column", pa.binary())) - pq.write_table(pa.Table.from_pylist([_row(i) for i in range(12)], schema=schema), - root / "part-0.parquet", row_group_size=3) - pq.write_table(pa.Table.from_pylist([_row(i) for i in range(12, 24)], schema=schema), - root / "part-1.parquet", row_group_size=3) + pq.write_table( + pa.Table.from_pylist([_row(i) for i in range(12)], schema=schema), root / "part-0.parquet", row_group_size=3 + ) + pq.write_table( + pa.Table.from_pylist([_row(i) for i in range(12, 24)], schema=schema), root / "part-1.parquet", row_group_size=3 + ) return root -def _patch_dist(monkeypatch: pytest.MonkeyPatch, rank: int = 0, - world_size: int = 1, sp_size: int = 1) -> None: +def _patch_dist(monkeypatch: pytest.MonkeyPatch, rank: int = 0, world_size: int = 1, sp_size: int = 1) -> None: monkeypatch.setattr(streaming.dist, "is_initialized", lambda: True) monkeypatch.setattr(streaming, "get_world_rank", lambda: rank) monkeypatch.setattr(streaming, "get_world_size", lambda: world_size) @@ -72,7 +73,8 @@ def _dataset(parquet_root: Path, cache_root: Path, **kwargs): def test_workers_zero_projects_columns_and_writes_json_manifest( - parquet_root: Path, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + parquet_root: Path, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: _patch_dist(monkeypatch) recorded_columns = [] original = pq.ParquetFile.iter_batches @@ -95,7 +97,8 @@ def test_workers_zero_projects_columns_and_writes_json_manifest( def test_state_dict_resumes_at_exact_next_batch( - parquet_root: Path, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + parquet_root: Path, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: _patch_dist(monkeypatch) first_dataset = _dataset(parquet_root, tmp_path / "owned-cache") iterator = iter(first_dataset) @@ -110,7 +113,8 @@ def test_state_dict_resumes_at_exact_next_batch( def test_odd_row_group_tails_are_carried_into_next_group( - parquet_root: Path, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + parquet_root: Path, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: _patch_dist(monkeypatch) dataset = _dataset(parquet_root, tmp_path / "owned-cache") ids = [item["id"] for batch in dataset for item in batch["info_list"]] @@ -118,14 +122,13 @@ def test_odd_row_group_tails_are_carried_into_next_group( def test_dp_sp_shards_are_identical_within_sp_and_disjoint_across_dp( - parquet_root: Path, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + parquet_root: Path, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: ids_by_rank = {} for rank in range(4): _patch_dist(monkeypatch, rank=rank, world_size=4, sp_size=2) dataset = _dataset(parquet_root, tmp_path / "owned-cache") - ids_by_rank[rank] = { - item["id"] for batch in dataset for item in batch["info_list"] - } + ids_by_rank[rank] = {item["id"] for batch in dataset for item in batch["info_list"]} assert ids_by_rank[0] == ids_by_rank[1] assert ids_by_rank[2] == ids_by_rank[3] assert ids_by_rank[0].isdisjoint(ids_by_rank[2]) @@ -133,27 +136,26 @@ def test_dp_sp_shards_are_identical_within_sp_and_disjoint_across_dp( def test_equal_number_of_row_groups_and_dp_shards_stays_nonempty( - parquet_root: Path, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + parquet_root: Path, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: ids_by_rank = [] for rank in range(4): _patch_dist(monkeypatch, rank=rank, world_size=4, sp_size=1) dataset = _dataset(parquet_root, tmp_path / "owned-cache") - ids_by_rank.append({ - item["id"] for batch in dataset for item in batch["info_list"] - }) + ids_by_rank.append({item["id"] for batch in dataset for item in batch["info_list"]}) assert all(ids for ids in ids_by_rank) assert len(set().union(*ids_by_rank)) == 24 -def test_manifest_must_not_be_written_inside_dataset( - parquet_root: Path, monkeypatch: pytest.MonkeyPatch) -> None: +def test_manifest_must_not_be_written_inside_dataset(parquet_root: Path, monkeypatch: pytest.MonkeyPatch) -> None: _patch_dist(monkeypatch) with pytest.raises(ValueError, match="outside the read-only dataset"): _dataset(parquet_root, parquet_root / "cache") def test_stateful_dataloader_restores_dataset_cursor( - parquet_root: Path, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + parquet_root: Path, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: _, loader = streaming.build_parquet_streaming_style_dataloader( path=str(parquet_root), batch_size=2, @@ -182,3 +184,30 @@ def test_stateful_dataloader_restores_dataset_cursor( resumed.load_state_dict(state) actual = next(iter(resumed))["info_list"] assert [item["id"] for item in actual] == [item["id"] for item in expected] + + +def test_resume_rejects_changed_topology(parquet_root: Path, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + _patch_dist(monkeypatch) + dataset = _dataset(parquet_root, tmp_path / "owned-cache") + next(iter(dataset)) + state = dataset.state_dict() + state["topology"]["batch_size"] = 99 + + resumed = _dataset(parquet_root, tmp_path / "owned-cache") + with pytest.raises(ValueError, match="topology or sampling config changed"): + resumed.load_state_dict(state) + + +def test_manifest_is_rebuilt_when_source_file_changes( + parquet_root: Path, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + _patch_dist(monkeypatch) + dataset = _dataset(parquet_root, tmp_path / "owned-cache") + old_fingerprint = dataset.manifest_fingerprint + source_file = parquet_root / "part-1.parquet" + table = pq.read_table(source_file) + pq.write_table(table, source_file, row_group_size=2) + + refreshed = _dataset(parquet_root, tmp_path / "owned-cache") + assert refreshed.manifest_fingerprint != old_fingerprint + assert refreshed.manifest["total_rows"] == 24 diff --git a/fastvideo/train/utils/dataloader.py b/fastvideo/train/utils/dataloader.py index 583797cf8..fd16ba339 100644 --- a/fastvideo/train/utils/dataloader.py +++ b/fastvideo/train/utils/dataloader.py @@ -6,7 +6,8 @@ from typing import Any, TYPE_CHECKING if TYPE_CHECKING: from fastvideo.train.utils.training_config import ( - DataConfig, ) + DataConfig, + ) def build_parquet_t2v_train_dataloader( @@ -17,11 +18,12 @@ def build_parquet_t2v_train_dataloader( ) -> Any: """Build a parquet dataloader for T2V-style datasets.""" - from fastvideo.dataset import (build_parquet_map_style_dataloader, - build_parquet_streaming_style_dataloader) - dataloader_type = str(getattr(data_config, "dataloader_type", "map")).strip().lower() if dataloader_type == "streaming": + from fastvideo.dataset.parquet_dataset_streaming_style import ( + build_parquet_streaming_style_dataloader, + ) + _dataset, dataloader = build_parquet_streaming_style_dataloader( data_config.data_path, data_config.train_batch_size, @@ -37,11 +39,11 @@ def build_parquet_t2v_train_dataloader( ) return dataloader if dataloader_type != "map": - raise ValueError( - f"Unsupported T2V dataloader_type: {dataloader_type!r}; " - "expected 'map' or 'streaming'") + raise ValueError(f"Unsupported T2V dataloader_type: {dataloader_type!r}; expected 'map' or 'streaming'") - _dataset, dataloader = (build_parquet_map_style_dataloader( + from fastvideo.dataset import build_parquet_map_style_dataloader + + _dataset, dataloader = build_parquet_map_style_dataloader( data_config.data_path, data_config.train_batch_size, num_data_workers=(data_config.dataloader_num_workers), @@ -50,7 +52,7 @@ def build_parquet_t2v_train_dataloader( drop_last=True, text_padding_length=int(text_len), seed=int(data_config.seed or 0), - )) + ) return dataloader @@ -62,9 +64,10 @@ def build_parquet_matrixgame2_train_dataloader( """Build a parquet dataloader for Matrix-Game 2.0 datasets.""" from fastvideo.dataset import ( - build_parquet_map_style_dataloader, ) + build_parquet_map_style_dataloader, + ) - _dataset, dataloader = (build_parquet_map_style_dataloader( + _dataset, dataloader = build_parquet_map_style_dataloader( data_config.data_path, data_config.train_batch_size, num_data_workers=(data_config.dataloader_num_workers), @@ -73,5 +76,5 @@ def build_parquet_matrixgame2_train_dataloader( drop_last=True, text_padding_length=512, seed=int(data_config.seed or 0), - )) + ) return dataloader