Skip to content

refactor(components): enforce public imports in ci - #3794

Open
akoumpa wants to merge 2 commits into
NVIDIA-NeMo:mainfrom
akoumpa:akoumparouli/enforce-component-exports
Open

refactor(components): enforce public imports in ci#3794
akoumpa wants to merge 2 commits into
NVIDIA-NeMo:mainfrom
akoumpa:akoumparouli/enforce-component-exports

Conversation

@akoumpa

@akoumpa akoumpa commented Sep 2, 2026

Copy link
Copy Markdown
Contributor

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 under nemo_automodel.

Changelog

  • Add a custom Import Linter component-interface contract that scans all runtime modules under nemo_automodel.
  • Reject private component-submodule imports and component-root names absent from __all__.
  • Migrate existing recipe, framework, model, CLI, and cross-component imports to their public component interfaces.
  • Add lazy package exports to preserve optional-dependency and import-cycle behavior.
  • Add focused contract tests and update repository policy documentation.

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.
  • Focused and upstream-overlap pytest selection — 450 passed.
  • 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:

  • Make sure you read and followed Contributor guidelines
  • Did you write any new necessary tests?
  • Did you add or update any necessary documentation?

Additional Information

  • Draft while the full CI matrix runs.

@copy-pr-bot

copy-pr-bot Bot commented Sep 2, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@akoumpa

akoumpa commented Sep 2, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 2015c92

@github-actions

github-actions Bot commented Sep 2, 2026

Copy link
Copy Markdown
Contributor

@akoumpa

akoumpa commented Sep 2, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test b2f7098

Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>

@jgerh jgerh left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Completed tech pubs review of markdown files and provided a few copyedits and suggested text revisions.

Comment thread docs/about/key-features.mdx Outdated
Comment thread docs/about/key-features.mdx Outdated
Comment thread docs/about/key-features.mdx Outdated
Comment thread docs/about/key-features.mdx Outdated
Comment thread docs/about/key-features.mdx Outdated
Comment thread docs/repository-structure.mdx Outdated
Comment thread docs/repository-structure.mdx Outdated
Comment thread docs/repository-structure.mdx Outdated
Comment thread docs/repository-structure.mdx Outdated
Comment thread docs/repository-structure.mdx Outdated

@yuhezhang-ai yuhezhang-ai left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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):

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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.

Comment thread docs/about/key-features.mdx Outdated
| `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.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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"),

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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>
@akoumpa

akoumpa commented Sep 4, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test b3490ee

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants