refactor(components): enforce public imports in ci - #3794
Conversation
|
/ok to test 2015c92 |
|
🌿 Preview your docs: https://nvidia-preview-preview-b3490ee9d891.docs.buildwithfern.com/nemo/automodel |
2015c92 to
b2f7098
Compare
|
/ok to test b2f7098 |
Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>
jgerh
left a comment
There was a problem hiding this comment.
Completed tech pubs review of markdown files and provided a few copyedits and suggested text revisions.
yuhezhang-ai
left a comment
There was a problem hiding this comment.
Requesting changes for several correctness and maintainability issues in the new enforcement: stale mock targets make tests order-dependent, the default unit-test environment cannot collect the new contract tests, the inherited independence check is made vacuous, lazy exports erase static types, and the static all evaluator can approve names that do not exist at runtime. I left inline details. The public-API relabeling and documentation scope should also be resolved deliberately. The broader import migration otherwise appears runtime-compatible.
| return None | ||
|
|
||
| from nemo_automodel.components.datasets.utils import add_causal_masks_to_batch | ||
| from nemo_automodel.components.datasets import add_causal_masks_to_batch |
There was a problem hiding this comment.
[P1] Update the remaining patches to this package-level target. tests/unit_tests/recipes/test_train_ft.py:1909,1935,1958 still patches nemo_automodel.components.datasets.utils.add_causal_masks_to_batch. Because datasets.getattr caches the resolved object in globals(), that submodule mock can either be cached in the package namespace beyond teardown or be ignored if the package attribute was materialized earlier, making the tests order-dependent. Please patch nemo_automodel.components.datasets.add_causal_masks_to_batch instead. The same stale-target pattern remains in test_diffusion_train_metrics.py:299 and _diffusers/test_auto_diffusion_pipeline.py:1204.
| from unittest.mock import Mock | ||
|
|
||
| from tools.component_imports import ComponentImport, find_component_imports | ||
| from tools.import_linter_contracts import ComponentInterfaceContract |
There was a problem hiding this comment.
[P1] Keep default unit-test collection independent of lint-only dependencies. This unconditional import pulls in grimp and importlinter, but those packages are only in the linting dependency group while the default groups are build, docs, and test. A documented uv sync --frozen followed by pytest tests/unit_tests/ therefore fails during collection. Please either move the dependencies into the test group or isolate/conditionally import only the contract-specific tests; a module-level skip would also hide the scanner tests that do not need these packages.
| component_imports = find_component_imports(_PROJECT_ROOT, modules) | ||
|
|
||
| for component_import in component_imports: | ||
| self._remove_direct_import(graph, component_import) |
There was a problem hiding this comment.
[P1] Preserve the independence check on the original graph. This removes every scanner-discovered edge into a configured component before super().check() runs. Any chain between configured components must end on one of those edges, so the inherited invalid_chains check cannot detect it; consumer-to-component ignore_imports entries also become unmatched. For example, a loss -> moe -> distributed dependency that main rejects is kept here. Please run the independence check on the unmodified graph and combine it with the interface scan. If interface-only enforcement is intentional, derive directly from Contract and remove the dead chain/ignore machinery instead.
| __all__ += sorted(_LAZY_ATTRS.keys()) | ||
|
|
||
|
|
||
| def __getattr__(name: str) -> object: |
There was a problem hiding this comment.
[P1] Preserve static types for lazy exports. With getattr(name: str) -> object, imports of every name known only through _LAZY_ATTRS are typed as object; this pattern affects roughly 250 exports across the component packages. The strict directories stay green only because selected names are declared under TYPE_CHECKING or imported eagerly. Please provide TYPE_CHECKING declarations or stubs for every lazy export and add a parity check between those declarations and _LAZY_ATTRS.keys().
| if value is not _UNKNOWN: | ||
| exports.update(_string_set(value)) | ||
|
|
||
| for node in ast.walk(tree): |
There was a problem hiding this comment.
[P1] Evaluate all only at module level and in execution order. Walking the whole tree collects all.append/extend calls inside functions, unreachable branches, and statements that occur before a later reset, so the scanner can approve a symbol the package does not export at runtime. There is a related state bug above: all += ... updates exports but not values["all"], so a later all = sorted(all) drops the additions. Please process supported top-level statements in order, keep the tracked value synchronized, reject unsupported shapes with a targeted path/line error, and cover nested/conditional/reset/AugAssign cases.
| | `launcher/` | Interactive, SkyPilot, and NeMo-Run job launch; Slurm uses the root-level `slurm.sub` script | | ||
|
|
||
| Each component can be used independently and has no cross-module imports. | ||
| Every component consumer—including recipes and other components—imports symbols from the target component's package API. The package's `__all__` defines that public interface. |
There was a problem hiding this comment.
[P2] Scope this claim to what CI actually enforces. The contract config covers 11 of the 19 top-level component packages; _peft, attention, cuda_graphs, eval, models, moe, quantization, and speculative are omitted, and migrated code still imports private modules from some of them. Please either extend the configured scope or narrow this wording to the component packages currently covered.
| "Checkpointer": (".checkpointing", "Checkpointer"), | ||
| "StateDictAdapter": (".state_dict_adapter", "StateDictAdapter"), | ||
| "find_latest_checkpoint": (".utils", "find_latest_checkpoint"), | ||
| "get_checkpoint_tensor_dtypes": (".utils", "_get_checkpoint_tensor_dtypes"), |
There was a problem hiding this comment.
[P2] Avoid publishing private helpers by aliasing them. This exposes a public name backed by a private implementation, and the same pattern is repeated for a number of helpers that consumers then alias back to underscore-prefixed names. Since adding a name to all creates a long-lived API obligation, please either promote and rename the implementation consistently or route callers through an existing public operation instead of relabeling private internals solely for this migration.
Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>
|
/ok to test b3490ee |
What does this PR do ?
Enforces imports of configured Automodel components through package-level symbols declared in
__all__, including recipes and every other runtime consumer undernemo_automodel.Changelog
nemo_automodel.__all__.Validation
uvx --from import-linter==2.4.0 lint-imports --debug --verbose --no-cache— 1 contract kept, 0 broken.uvx --from "ruff~=0.12.0" ruff check .— passed.uvx --from "ruff~=0.12.0" ruff format --check .— 799 files already formatted.uvx --from bandit bandit -r app.py nemo_automodel examples scripts tools tutorials -t B614 -q— passed.git diff HEAD^ --check— passed.Before your PR is "Ready for review"
Pre checks:
Additional Information