diff --git a/fastvideo/tests/train/utils/test_config.py b/fastvideo/tests/train/utils/test_config.py index a80cf0517..151a773b7 100644 --- a/fastvideo/tests/train/utils/test_config.py +++ b/fastvideo/tests/train/utils/test_config.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 """CPU-only unit tests for :func:`load_run_config`.""" + from __future__ import annotations from pathlib import Path @@ -27,8 +28,7 @@ def _minimal_yaml() -> dict[str, Any]: }, }, "method": { - "_target_": - "fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod", + "_target_": "fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod", }, "training": {}, } @@ -44,8 +44,7 @@ def test_minimal_yaml_loads_happy_path(tmp_path: Path) -> None: assert isinstance(cfg, RunConfig) assert isinstance(cfg.training, TrainingConfig) - assert cfg.method["_target_"] == ( - "fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod") + assert cfg.method["_target_"] == ("fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod") assert "student" in cfg.models assert cfg.callbacks == {} # raw retains the original YAML dict for downstream logging. @@ -64,6 +63,10 @@ def test_minimal_yaml_applies_all_defaults(tmp_path: Path) -> None: assert t.data.train_batch_size == 1 assert t.data.dataloader_num_workers == 0 + assert t.data.dataloader_type == "map" + assert t.data.streaming_manifest_path == "" + assert t.data.streaming_read_batch_size == 8 + assert t.data.streaming_shuffle_row_groups is True assert t.data.training_cfg_rate == 0.0 assert t.data.seed == 0 @@ -107,6 +110,10 @@ def test_full_yaml_populates_all_training_fields(tmp_path: Path) -> None: "data_path": "/some/path", "train_batch_size": 2, "dataloader_num_workers": 4, + "dataloader_type": "streaming", + "streaming_manifest_path": "/owned/index/manifest.json", + "streaming_read_batch_size": 2, + "streaming_shuffle_row_groups": False, "training_cfg_rate": 0.1, "seed": 42, "num_height": 256, @@ -136,9 +143,7 @@ def test_full_yaml_populates_all_training_fields(tmp_path: Path) -> None: "project_name": "myproj", "run_name": "myrun", }, - "vsa": { - "sparsity": 0.5 - }, + "vsa": {"sparsity": 0.5}, "model": { "weighting_scheme": "logit_normal", "logit_mean": 0.5, @@ -156,6 +161,10 @@ def test_full_yaml_populates_all_training_fields(tmp_path: Path) -> None: assert t.data.train_batch_size == 2 assert t.data.data_path == "/some/path" + assert t.data.dataloader_type == "streaming" + assert t.data.streaming_manifest_path == "/owned/index/manifest.json" + assert t.data.streaming_read_batch_size == 2 + assert t.data.streaming_shuffle_row_groups is False assert t.data.num_frames == 33 assert t.data.seed == 42 @@ -230,10 +239,13 @@ def test_missing_config_file_raises(tmp_path: Path) -> None: # --------------------------------------------------------------------------- -@pytest.mark.parametrize("betas_value, expected", [ - ([0.8, 0.9], (0.8, 0.9)), - ("0.9,0.999", (0.9, 0.999)), -]) +@pytest.mark.parametrize( + "betas_value, expected", + [ + ([0.8, 0.9], (0.8, 0.9)), + ("0.9,0.999", (0.9, 0.999)), + ], +) def test_betas_parses_list_and_string_forms( tmp_path: Path, betas_value: Any, @@ -347,8 +359,7 @@ def test_callbacks_passed_through_when_present(tmp_path: Path) -> None: data = _minimal_yaml() data["callbacks"] = { "grad_clip": { - "_target_": - "fastvideo.train.callbacks.grad_clip.GradNormClipCallback", + "_target_": "fastvideo.train.callbacks.grad_clip.GradNormClipCallback", "max_grad_norm": 1.0, }, } diff --git a/fastvideo/train/utils/config.py b/fastvideo/train/utils/config.py index fde664c9a..a81314fb3 100644 --- a/fastvideo/train/utils/config.py +++ b/fastvideo/train/utils/config.py @@ -54,7 +54,8 @@ class RunConfig: if dataclasses.is_dataclass(obj) and not isinstance(obj, type): return { f.name: _safe_asdict(getattr(obj, f.name)) - for f in dataclasses.fields(obj) if not callable(getattr(obj, f.name)) + for f in dataclasses.fields(obj) + if not callable(getattr(obj, f.name)) } if isinstance(obj, dict): return {k: _safe_asdict(v) for k, v in obj.items()} @@ -87,8 +88,7 @@ def _resolve_existing_file(path: str) -> str: def _require_mapping(raw: Any, *, where: str) -> dict[str, Any]: if not isinstance(raw, dict): - raise ValueError(f"Expected mapping at {where}, " - f"got {type(raw).__name__}") + raise ValueError(f"Expected mapping at {where}, got {type(raw).__name__}") return raw @@ -110,8 +110,7 @@ def get_optional_int(mapping: dict[str, Any], key: str, *, where: str) -> int | return int(raw) if isinstance(raw, str) and raw.strip(): return int(raw) - raise ValueError(f"Expected int at {where}, " - f"got {type(raw).__name__}") + raise ValueError(f"Expected int at {where}, got {type(raw).__name__}") def get_optional_float(mapping: dict[str, Any], key: str, *, where: str) -> float | None: @@ -124,8 +123,7 @@ def get_optional_float(mapping: dict[str, Any], key: str, *, where: str) -> floa return float(raw) if isinstance(raw, str) and raw.strip(): return float(raw) - raise ValueError(f"Expected float at {where}, " - f"got {type(raw).__name__}") + raise ValueError(f"Expected float at {where}, got {type(raw).__name__}") def parse_betas(raw: Any, *, where: str) -> tuple[float, float]: @@ -136,11 +134,9 @@ def parse_betas(raw: Any, *, where: str) -> tuple[float, float]: if isinstance(raw, str): parts = [p.strip() for p in raw.split(",") if p.strip()] if len(parts) != 2: - raise ValueError(f"Expected betas as 'b1,b2' at {where}, " - f"got {raw!r}") + raise ValueError(f"Expected betas as 'b1,b2' at {where}, got {raw!r}") return float(parts[0]), float(parts[1]) - raise ValueError(f"Expected betas as 'b1,b2' at {where}, " - f"got {type(raw).__name__}") + raise ValueError(f"Expected betas as 'b1,b2' at {where}, got {type(raw).__name__}") # ---- config convenience helpers ---- @@ -182,8 +178,7 @@ def require_non_negative_int( raise ValueError(f"Missing required key {loc!r}") val = get_optional_int(mapping, key, where=loc) if val is None or val < 0: - raise ValueError(f"{loc} must be a non-negative integer, " - f"got {raw!r}") + raise ValueError(f"{loc} must be a non-negative integer, got {raw!r}") return val @@ -203,8 +198,7 @@ def require_non_negative_float( raise ValueError(f"Missing required key {loc!r}") val = get_optional_float(mapping, key, where=loc) if val is None or val < 0.0: - raise ValueError(f"{loc} must be a non-negative float, " - f"got {raw!r}") + raise ValueError(f"{loc} must be a non-negative float, got {raw!r}") return val @@ -226,12 +220,10 @@ def require_choice( return default raise ValueError(f"Missing required key {loc!r}") if not isinstance(raw, str) or not raw.strip(): - raise ValueError(f"{loc} must be a non-empty string, " - f"got {type(raw).__name__}") + raise ValueError(f"{loc} must be a non-empty string, got {type(raw).__name__}") val = raw.strip().lower() if val not in choices: - raise ValueError(f"{loc} must be one of {sorted(choices)}, " - f"got {raw!r}") + raise ValueError(f"{loc} must be one of {sorted(choices)}, got {raw!r}") return val @@ -250,8 +242,7 @@ def require_bool( return default raise ValueError(f"Missing required key {loc!r}") if not isinstance(raw, bool): - raise ValueError(f"{loc} must be a bool, " - f"got {type(raw).__name__}") + raise ValueError(f"{loc} must be a bool, got {type(raw).__name__}") return raw @@ -359,9 +350,11 @@ def _build_training_config( preprocessed_data_type = str(da.get("preprocessed_data_type", "t2v") or "t2v").strip().lower() if preprocessed_data_type not in {"t2v", "text_only"}: - raise ValueError("training.data.preprocessed_data_type must be one of " - "{'t2v', 'text_only'}, got " - f"{preprocessed_data_type!r}") + raise ValueError( + "training.data.preprocessed_data_type must be one of " + "{'t2v', 'text_only'}, got " + f"{preprocessed_data_type!r}" + ) return TrainingConfig( distributed=DistributedConfig( @@ -377,6 +370,10 @@ def _build_training_config( preprocessed_data_type=preprocessed_data_type, train_batch_size=int(da.get("train_batch_size", 1) or 1), dataloader_num_workers=int(da.get("dataloader_num_workers", 0) or 0), + dataloader_type=str(da.get("dataloader_type", "map") or "map").strip().lower(), + streaming_manifest_path=str(da.get("streaming_manifest_path", "") or ""), + streaming_read_batch_size=int(da.get("streaming_read_batch_size", 8) or 8), + streaming_shuffle_row_groups=bool(da.get("streaming_shuffle_row_groups", True)), training_cfg_rate=float(da.get("training_cfg_rate", 0.0) or 0.0), seed=int(da.get("seed", 0) or 0), num_height=int(da.get("num_height", 0) or 0), @@ -427,7 +424,9 @@ def _build_training_config( ) -def _parse_cli_overrides(overrides: list[str], ) -> dict[str, Any]: +def _parse_cli_overrides( + overrides: list[str], +) -> dict[str, Any]: """Parse ``--dotted.key value`` CLI overrides. Returns a flat dict mapping dotted keys to parsed @@ -530,8 +529,7 @@ def load_run_config( role_str = _require_str(role, where="models.") model_cfg = _require_mapping(model_cfg_raw, where=f"models.{role_str}") if "_target_" not in model_cfg: - raise ValueError(f"models.{role_str} must have a " - "'_target_' key") + raise ValueError(f"models.{role_str} must have a '_target_' key") models[role_str] = dict(model_cfg) # --- method ---