test(recipes): validate all recipe factories

yaoyu-33 · yaoyu-33 · commit 6fe7e49f8294 · 2026-07-16T16:50:51.000-07:00
Signed-off-by: yaoyu-33 &lt;yaoyu.094@gmail.com&gt;
diff --git a/tests/unit_tests/recipes/recipe_test_utils.py b/tests/unit_tests/recipes/recipe_test_utils.py
@@ -17,8 +17,12 @@
 from __future__ import annotations
 
 import importlib
+import inspect
+import pkgutil
+import socket
+import sys
 from collections.abc import Callable, Iterator
-from types import ModuleType
+from types import ModuleType, SimpleNamespace
 
 import pytest
 
@@ -59,3 +63,170 @@ def patch_recipe_module_global(
         if hasattr(owner, name) and id(owner) not in patched_ids:
             monkeypatch.setattr(owner, name, value)
             patched_ids.add(id(owner))
+
+
+def discover_recipe_factories(
+    package: ModuleType,
+    *,
+    exclude_module_prefixes: tuple[str, ...] = (),
+) -> tuple[Callable[..., object], ...]:
+    """Import and return every public recipe factory defined below ``package``.
+
+    Compatibility modules often re-export the same callable under a second
+    name. Restricting discovery to functions owned by each leaf module keeps
+    the construction matrix focused on unique implementations.
+    """
+    package_path = getattr(package, "__path__", None)
+    if package_path is None:
+        raise ValueError(f"{package.__name__} is not a package")
+
+    factories: dict[tuple[str, str], Callable[..., object]] = {}
+    prefix = f"{package.__name__}."
+    for module_info in pkgutil.walk_packages(package_path, prefix=prefix):
+        if module_info.name.startswith(exclude_module_prefixes):
+            continue
+
+        module = importlib.import_module(module_info.name)
+        for name, value in vars(module).items():
+            if name.startswith("_") or not name.endswith("_config") or not inspect.isfunction(value):
+                continue
+            if value.__module__ != module.__name__:
+                continue
+            factories[(value.__module__, value.__qualname__)] = value
+
+    return tuple(factories[key] for key in sorted(factories))
+
+
+def recipe_factory_key(factory: Callable[..., object]) -> tuple[str, str]:
+    """Return a stable identity for a recipe factory or compatibility alias."""
+    return factory.__module__, factory.__qualname__
+
+
+def recipe_factory_id(factory: Callable[..., object]) -> str:
+    """Return a readable pytest parameter ID for a recipe factory."""
+    module, name = recipe_factory_key(factory)
+    return f"{module}:{name}"
+
+
+def exported_recipe_factory_keys(module: ModuleType) -> set[tuple[str, str]]:
+    """Return identities of public ``*_config`` callables exported by a module."""
+    return {
+        recipe_factory_key(value)
+        for name, value in vars(module).items()
+        if not name.startswith("_") and name.endswith("_config") and callable(value)
+    }
+
+
+class _OfflineModelProvider:
+    """Mutable provider with the fields HF-backed recipes read before writing."""
+
+    def __init__(self) -> None:
+        self.apply_rope_fusion = False
+        self.context_parallel_size = 1
+        self.cross_entropy_fusion_impl = "native"
+        self.csa_compress_ratios = [0] * 33
+        self.dsa_indexer_skip_topk_offset = 0
+        self.dsa_indexer_topk_freq = 1
+        self.experimental_attention_variant = "dsa"
+        self.make_vocab_size_divisible_by = 128
+        self.moe_flex_dispatcher_backend = None
+        self.mtp_num_layers = 1
+        self.num_layers = 32
+        self.num_moe_experts = 8
+        self.rotary_base = 10000.0
+        self.rotary_scaling_factor = 1
+        self.seq_length = 4096
+        self.tensor_model_parallel_size = 1
+        self.use_te_rng_tracker = False
+        self.use_transformer_engine_op_fuser = False
+        self.vocab_size = 256000
+        self.yarn_original_max_position_embeddings = 32768
+
+    def finalize(self) -> None:
+        """Match the model-provider interface without initializing a model."""
+
+
+class _OfflineAutoBridge:
+    """Build a local provider without reading a Hugging Face configuration."""
+
+    @classmethod
+    def from_hf_pretrained(cls, *args: object, **kwargs: object) -> "_OfflineAutoBridge":
+        del args, kwargs
+        return cls()
+
+    def to_megatron_provider(self, *args: object, **kwargs: object) -> _OfflineModelProvider:
+        del args, kwargs
+        return _OfflineModelProvider()
+
+
+class _OfflineTokenizer:
+    """Small tokenizer stand-in used by Gemma vocabulary adjustment helpers."""
+
+    def __len__(self) -> int:
+        return 256000
+
+
+class _OfflinePreTrainedFlux:
+    """Expose the local config surface consumed by ``FluxBridge``."""
+
+    def __init__(self, model_name_or_path: object, **kwargs: object) -> None:
+        del kwargs
+        self.model_name_or_path = str(model_name_or_path)
+        self.config = SimpleNamespace(
+            num_attention_heads=24,
+            attention_head_dim=128,
+            in_channels=64,
+            patch_size=1,
+            num_layers=19,
+            num_single_layers=38,
+            joint_attention_dim=4096,
+            pooled_projection_dim=768,
+            guidance_embeds=True,
+            axes_dims_rope=[16, 56, 56],
+            ffn_dim=12288,
+        )
+
+
+def patch_recipe_construction_dependencies(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Keep exhaustive recipe construction deterministic, CPU-only, and offline."""
+    monkeypatch.setenv("HF_DATASETS_OFFLINE", "1")
+    monkeypatch.setenv("HF_HUB_OFFLINE", "1")
+    monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1")
+
+    recipe_prefixes = ("megatron.bridge.recipes.", "megatron.bridge.perf_recipes.")
+
+    def skip_flex_dispatcher_hardware_probe(*args: object, **kwargs: object) -> None:
+        del args, kwargs
+
+    for module_name, module in tuple(sys.modules.items()):
+        if module is None or not module_name.startswith(recipe_prefixes):
+            continue
+        if hasattr(module, "AutoBridge"):
+            monkeypatch.setattr(module, "AutoBridge", _OfflineAutoBridge)
+        if hasattr(module, "apply_flex_dispatcher_backend"):
+            monkeypatch.setattr(module, "apply_flex_dispatcher_backend", skip_flex_dispatcher_hardware_probe)
+
+    flux_recipe_module = importlib.import_module("megatron.bridge.recipes.flux.h100.flux")
+    monkeypatch.setattr(flux_recipe_module, "PreTrainedFlux", _OfflinePreTrainedFlux)
+
+    deepseek_v4_recipe_module = importlib.import_module("megatron.bridge.recipes.deepseek.h100.deepseek_v4")
+    monkeypatch.setattr(deepseek_v4_recipe_module, "deepseek_v4_supports_blackwell_fused_kernels", lambda: False)
+
+    from transformers import AutoTokenizer
+
+    def load_offline_tokenizer(*args: object, **kwargs: object) -> _OfflineTokenizer:
+        del args, kwargs
+        return _OfflineTokenizer()
+
+    monkeypatch.setattr(
+        AutoTokenizer,
+        "from_pretrained",
+        staticmethod(load_offline_tokenizer),
+    )
+
+    def reject_network(*args: object, **kwargs: object) -> None:
+        del args, kwargs
+        raise AssertionError("recipe construction attempted network access")
+
+    monkeypatch.setattr(socket, "create_connection", reject_network)
+    monkeypatch.setattr(socket.socket, "connect", reject_network)
diff --git a/tests/unit_tests/recipes/test_all_perf_recipe_factories.py b/tests/unit_tests/recipes/test_all_perf_recipe_factories.py
@@ -0,0 +1,86 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Offline construction tests for every public performance recipe factory."""
+
+import importlib
+import inspect
+import pkgutil
+from collections.abc import Callable
+
+import pytest
+
+from megatron.bridge.training.config import ConfigContainer
+from tests.unit_tests.recipes.recipe_test_utils import (
+    discover_recipe_factories,
+    exported_recipe_factory_keys,
+    patch_recipe_construction_dependencies,
+    recipe_factory_id,
+    recipe_factory_key,
+)
+
+
+pytestmark = pytest.mark.unit
+
+_PERF_RECIPES_PACKAGE = importlib.import_module("megatron.bridge.perf_recipes")
+_PERF_RECIPE_FACTORIES = discover_recipe_factories(_PERF_RECIPES_PACKAGE)
+_PERF_RECIPE_FAMILY_PACKAGES = tuple(
+    importlib.import_module(module_info.name)
+    for module_info in pkgutil.iter_modules(
+        _PERF_RECIPES_PACKAGE.__path__,
+        prefix=f"{_PERF_RECIPES_PACKAGE.__name__}.",
+    )
+    if module_info.ispkg
+)
+_PERF_RECIPE_FAMILY_PACKAGES_BY_NAME = {package.__name__: package for package in _PERF_RECIPE_FAMILY_PACKAGES}
+
+
+@pytest.fixture(autouse=True)
+def _keep_recipe_construction_offline(monkeypatch: pytest.MonkeyPatch) -> None:
+    patch_recipe_construction_dependencies(monkeypatch)
+
+
+def test_all_perf_recipe_factories_are_exported() -> None:
+    """Every canonical factory remains available from its public family package."""
+    canonical = {recipe_factory_key(factory) for factory in _PERF_RECIPE_FACTORIES}
+    exported = set().union(*(exported_recipe_factory_keys(package) for package in _PERF_RECIPE_FAMILY_PACKAGES))
+    assert exported == canonical
+    for factory in _PERF_RECIPE_FACTORIES:
+        family_name = factory.__module__.split(".", maxsplit=4)[3]
+        family_package = _PERF_RECIPE_FAMILY_PACKAGES_BY_NAME[f"{_PERF_RECIPES_PACKAGE.__name__}.{family_name}"]
+        assert getattr(family_package, factory.__name__) is factory
+
+
+@pytest.mark.parametrize("recipe_factory", _PERF_RECIPE_FACTORIES, ids=recipe_factory_id)
+def test_perf_recipe_factory_builds_config(recipe_factory: Callable[..., object]) -> None:
+    """Every performance recipe can be built without GPU or network access."""
+    inspect.signature(recipe_factory).bind()
+
+    cfg = recipe_factory()
+
+    assert isinstance(cfg, ConfigContainer)
+    for section in (
+        "model",
+        "train",
+        "optimizer",
+        "scheduler",
+        "dataset",
+        "logger",
+        "tokenizer",
+        "checkpoint",
+        "rng",
+        "ddp",
+        "dist",
+    ):
+        assert getattr(cfg, section) is not None
diff --git a/tests/unit_tests/recipes/test_all_recipe_factories.py b/tests/unit_tests/recipes/test_all_recipe_factories.py
@@ -0,0 +1,97 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Offline construction tests for every public library recipe factory."""
+
+import importlib
+import inspect
+from collections.abc import Callable
+
+import pytest
+
+from megatron.bridge.training.config import ConfigContainer
+from tests.unit_tests.recipes.recipe_test_utils import (
+    discover_recipe_factories,
+    exported_recipe_factory_keys,
+    patch_recipe_construction_dependencies,
+    recipe_factory_id,
+    recipe_factory_key,
+)
+
+
+pytestmark = pytest.mark.unit
+
+_RECIPES_PACKAGE = importlib.import_module("megatron.bridge.recipes")
+_RECIPE_FACTORIES = discover_recipe_factories(
+    _RECIPES_PACKAGE,
+    exclude_module_prefixes=("megatron.bridge.recipes.utils",),
+)
+_UNSUPPORTED_FACTORY_KEYS = {
+    (
+        "megatron.bridge.recipes.qwen.h100.qwen3_next",
+        "qwen3_next_80b_a3b_peft_1gpu_h100_bf16_config",
+    ),
+}
+_RUNNABLE_RECIPE_FACTORIES = tuple(
+    factory for factory in _RECIPE_FACTORIES if recipe_factory_key(factory) not in _UNSUPPORTED_FACTORY_KEYS
+)
+_UNSUPPORTED_RECIPE_FACTORIES = tuple(
+    factory for factory in _RECIPE_FACTORIES if recipe_factory_key(factory) in _UNSUPPORTED_FACTORY_KEYS
+)
+
+
+@pytest.fixture(autouse=True)
+def _keep_recipe_construction_offline(monkeypatch: pytest.MonkeyPatch) -> None:
+    patch_recipe_construction_dependencies(monkeypatch)
+
+
+def test_all_recipe_factories_are_exported() -> None:
+    """Every canonical factory remains available from the public recipe package."""
+    canonical = {recipe_factory_key(factory) for factory in _RECIPE_FACTORIES}
+    assert exported_recipe_factory_keys(_RECIPES_PACKAGE) == canonical
+    for factory in _RECIPE_FACTORIES:
+        assert getattr(_RECIPES_PACKAGE, factory.__name__) is factory
+
+
+@pytest.mark.parametrize("recipe_factory", _RUNNABLE_RECIPE_FACTORIES, ids=recipe_factory_id)
+def test_recipe_factory_builds_config(recipe_factory: Callable[..., object]) -> None:
+    """Every supported recipe can be called with defaults without GPU or network access."""
+    inspect.signature(recipe_factory).bind()
+
+    cfg = recipe_factory()
+
+    assert isinstance(cfg, ConfigContainer)
+    for section in (
+        "model",
+        "train",
+        "optimizer",
+        "scheduler",
+        "dataset",
+        "logger",
+        "tokenizer",
+        "checkpoint",
+        "rng",
+        "ddp",
+        "dist",
+    ):
+        assert getattr(cfg, section) is not None
+
+
+@pytest.mark.parametrize("recipe_factory", _UNSUPPORTED_RECIPE_FACTORIES, ids=recipe_factory_id)
+def test_intentionally_unsupported_recipe_has_clear_error(recipe_factory: Callable[..., object]) -> None:
+    """Keep the documented Qwen3-Next PEFT limitation explicit in the exhaustive matrix."""
+    inspect.signature(recipe_factory).bind()
+
+    with pytest.raises(NotImplementedError, match="PEFT is not currently supported for Qwen3-Next"):
+        recipe_factory()
