diff --git a/tests/unit/torch/puzzletron/test_import_boundary.py b/tests/unit/torch/puzzletron/test_import_boundary.py index 7cf74bf5fb3..7ae4308b8d6 100644 --- a/tests/unit/torch/puzzletron/test_import_boundary.py +++ b/tests/unit/torch/puzzletron/test_import_boundary.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Regression tests for importing Puzzletron without optional AutoModel code.""" +"""Regression tests for Puzzletron's optional-dependency import boundaries.""" from __future__ import annotations @@ -25,69 +25,91 @@ import pytest REPOSITORY_ROOT = Path(__file__).resolve().parents[4] +PUZZLETRON_PACKAGE = "modelopt.torch.puzzletron" pytestmark = pytest.mark.skipif( sys.platform == "win32", reason="Puzzletron imports fcntl-backed runtime modules that are unavailable on Windows", ) -def _run_without_automodel(script: str) -> subprocess.CompletedProcess[str]: - environment = os.environ.copy() - # These probes validate fresh-interpreter import behavior. Coverage - # auto-start traces the entire child import graph and exceeds the unit - # suite's 60-second cap before the probe can return. - environment.pop("COVERAGE_PROCESS_START", None) - return subprocess.run( - [ - sys.executable, - "-c", - "import sys; sys.modules['nemo_automodel'] = None; " + script, - ], - cwd=REPOSITORY_ROOT, - env=environment, - capture_output=True, - text=True, - check=False, - ) - - def test_lightweight_puzzletron_import_does_not_require_automodel() -> None: - result = _run_without_automodel( - "from modelopt.torch.puzzletron.identity import stable_hash; " - "assert stable_hash({'ready': True}, prefix='ci')" + result = _run_fresh( + f"from {PUZZLETRON_PACKAGE}.identity import stable_hash; " + "assert stable_hash({'ready': True}, prefix='ci')", + unavailable_module="nemo_automodel", ) assert result.returncode == 0, result.stderr def test_resolved_setup_config_is_available_only_from_its_explicit_module() -> None: - result = _run_without_automodel( + result = _run_fresh( "import puzzletron_setup.v2 as setup_v2; " "from puzzletron_setup.v2.resolved import ResolvedCampaignConfig; " - "assert not hasattr(setup_v2, 'ResolvedCampaignConfig')" + "assert not hasattr(setup_v2, 'ResolvedCampaignConfig')", + unavailable_module="nemo_automodel", ) assert result.returncode == 0, result.stderr def test_resolved_setup_config_import_does_not_initialize_torch() -> None: - result = _run_without_automodel( + result = _run_fresh( "from puzzletron_setup.v2.resolved import ResolvedCampaignConfig; " "assert 'torch' not in sys.modules; " - "assert not any(name.startswith('modelopt.torch') for name in sys.modules)" + "assert not any(name.startswith('modelopt.torch') for name in sys.modules)", + unavailable_module="nemo_automodel", ) assert result.returncode == 0, result.stderr +def test_prepare_dataset_uses_fire_only_for_cli_execution() -> None: + result = _run_fresh( + "import importlib.util, pathlib, types; loader = importlib.util.spec_from_file_location; " + f"stub_names = ('datasets', 'numpy', '{PUZZLETRON_PACKAGE}.tools.logger'); " + "sys.modules.update({name: types.ModuleType(name) for name in stub_names}); " + "sys.modules[stub_names[-1]].mprint = lambda *args, **kwargs: None; " + f"module_name = '{PUZZLETRON_PACKAGE}.dataset.prepare_dataset'; " + "path = pathlib.Path(*module_name.split('.')).with_suffix('.py'); spec = loader(module_name, path); " + "module = importlib.util.module_from_spec(spec); spec.loader.exec_module(module); " + "fire_module = types.ModuleType('fire'); fire_module.Fire = lambda component: print(component.__name__); " + "sys.modules['fire'] = fire_module; spec = loader('__main__', path); " + "module = importlib.util.module_from_spec(spec); module.__package__ = module_name.rpartition('.')[0]; " + "module.__spec__ = None; spec.loader.exec_module(module)", + unavailable_module="fire", + ) + + assert result.returncode == 0, result.stderr + assert result.stdout.strip() == "process_and_save_dataset" + + def test_automodel_recipe_loader_reports_missing_dependency() -> None: - result = _run_without_automodel( - "from modelopt.torch.puzzletron.diagnostics.width_slice_equivalence " + result = _run_fresh( + f"from {PUZZLETRON_PACKAGE}.diagnostics.width_slice_equivalence " "import _replace_block_scoring_recipe; " "\ntry:\n _replace_block_scoring_recipe()" "\nexcept ImportError as error:" "\n assert 'requires a compatible NeMo AutoModel' in str(error)" - "\nelse:\n raise AssertionError('AutoModel-backed recipe unexpectedly loaded')" + "\nelse:\n raise AssertionError('AutoModel-backed recipe unexpectedly loaded')", + unavailable_module="nemo_automodel", ) assert result.returncode == 0, result.stderr + + +def _run_fresh( + script: str, *, unavailable_module: str | None = None +) -> subprocess.CompletedProcess[str]: + environment = os.environ.copy() + environment.pop("COVERAGE_PROCESS_START", None) + unavailable = f"sys.modules[{unavailable_module!r}] = None; " if unavailable_module else "" + prelude = "import sys; " + unavailable + return subprocess.run( + [sys.executable, "-c", prelude + script], + cwd=REPOSITORY_ROOT, + env=environment, + capture_output=True, + text=True, + check=False, + ) diff --git a/tests/unit/torch/puzzletron/test_prepare_dataset_imports.py b/tests/unit/torch/puzzletron/test_prepare_dataset_imports.py deleted file mode 100644 index 80544906ffa..00000000000 --- a/tests/unit/torch/puzzletron/test_prepare_dataset_imports.py +++ /dev/null @@ -1,23 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -import ast -from pathlib import Path - - -def test_prepare_dataset_imports_fire_only_for_cli_execution(): - source = Path("modelopt/torch/puzzletron/dataset/prepare_dataset.py").read_text() - tree = ast.parse(source) - - top_level_fire_imports = [ - node - for node in tree.body - if isinstance(node, (ast.Import, ast.ImportFrom)) - and ( - (isinstance(node, ast.Import) and any(alias.name == "fire" for alias in node.names)) - or (isinstance(node, ast.ImportFrom) and node.module == "fire") - ) - ] - - assert top_level_fire_imports == [] - assert "from fire import Fire" in source