Skip to content

perf(imports): resolve top-level and rails exports lazily - #2256

Draft
Pouyanpi wants to merge 5 commits into
developfrom
pouyanpi/fix-package-wide-eager-init
Draft

perf(imports): resolve top-level and rails exports lazily#2256
Pouyanpi wants to merge 5 commits into
developfrom
pouyanpi/fix-package-wide-eager-init

Conversation

@Pouyanpi

@Pouyanpi Pouyanpi commented Aug 5, 2026

Copy link
Copy Markdown
Collaborator

Stacked PR. Base branch: refactor/iorails-actions-decouple-colang (#2241).
Review this PR only against that base. #2241 decouples llm_call and the rail
actions from Colang; this PR makes the package entry points lazy so those
decoupled paths are actually reachable without booting the Colang runtime.

Description

import nemoguardrails eagerly imported RailsConfig, LLMRails, and
Guardrails at package-import time. Because those pull in the Colang parsers,
the tracing/OpenTelemetry stack, aiohttp, jinja2, and more, a bare
import nemoguardrails booted ~870 modules (~0.5s) even for callers that
only needed a type or a submodule.

This PR resolves the public names of nemoguardrails and nemoguardrails.rails
lazily via PEP 562 module __getattr__, pointing each name at its narrowest
module (RailsConfig -> rails.llm.config, LLMRails -> rails.llm.llmrails,
Guardrails -> guardrails.guardrails). The heavy runtime is loaded only on first
use of a heavy symbol.

Making nemoguardrails.rails lazy has a second effect: submodules such as
nemoguardrails.rails.llm.options can now initialize the rails package without
dragging in LLMRails and its Colang runtime. That removes a circular import
(base_guardrails <-> rails.llm.llmrails) which previously made isolated
rail-action imports fail, and removes the need for an explicit
nemoguardrails.rails preload workaround in the top-level __getattr__.

Performance

Measured in fresh interpreters (median of 7). "Colang" is the number of
nemoguardrails.colang.* modules resident in sys.modules after the import;
"Modules" is total sys.modules. Absolute times vary with machine load; the
module counts do not.

Import Before (time) After (time) Before (Colang / modules) After (Colang / modules)
import nemoguardrails ~534 ms ~30 ms 32 / 874 0 / 199
from nemoguardrails import ChatMessage ~583 ms ~34 ms 32 / 874 0 / 207
from nemoguardrails import RailsConfig ~574 ms ~248 ms 32 / 874 0 / 455
import nemoguardrails.llm.call ~596 ms ~168 ms 32 / 874 0 / 354
import nemoguardrails.library.content_safety.actions ~577 ms ~293 ms 32 / 874 0 / 489
from nemoguardrails import LLMRails ~582 ms ~514 ms 32 / 874 32 / 784

Takeaways:

  • Import-only / lightweight-type consumers (import nemoguardrails,
    ChatMessage) drop from ~0.5s and ~870 modules to ~30 ms and ~200 modules,
    with zero Colang modules loaded.
  • Config-only, llm_call, and built-in rail-action paths load zero
    Colang modules (the invariant refactor(iorails): Decouple llm_call and rail actions from unrelated Colang modules #2241 established, now reachable end-to-end
    through the real package entry points).
  • LLMRails still loads the full Colang runtime, by design: it is the
    legacy Colang engine. It is even slightly lighter than before (784 vs 874
    modules) because it no longer eagerly loads the Guardrails/tracing branch.

Behavior preserved

  • __all__, dir(nemoguardrails), and unknown-attribute AttributeError.
  • Type-checker and IDE symbols, via if TYPE_CHECKING: re-exports.
  • The NEMO_GUARDRAILS_IORAILS_ENGINE alias of top-level LLMRails to
    Guardrails.
  • importlib.reload(nemoguardrails) re-evaluates that env alias: resolution is
    intentionally not cached into module globals, so a reload re-reads the
    environment rather than returning a stale binding.

Reviewer notes

  • cli/chat.py (1 line): stricter type resolution through the new
    TYPE_CHECKING re-exports lets ty infer chunk: str in the streaming loop,
    which makes an existing cast(str, chunk) redundant and trips the ty gate.
    The value is always str there, so the redundant cast is removed (the cast
    import is still used elsewhere in the file).
  • tests/llm/test_call_import_graph.py (from refactor(iorails): Decouple llm_call and rail actions from unrelated Colang modules #2241): its docstrings state
    that the runtime property "could never be observed" because
    nemoguardrails/__init__ "imports the world." That premise no longer holds
    after this PR. Those static import-graph tests still pass and remain valid;
    the new tests/test_lazy_imports.py is the runtime complement. The stale
    wording was left untouched here to keep the stack diff clean and can be
    refreshed when the stack settles.

Verification

  • make test (full pytest.ini testpaths): 6299 passed, 178 skipped.
  • pre-commit run --files ... on every changed file: all hooks pass, including
    ruff, ruff format, license insertion, and ty.
  • New tests/test_lazy_imports.py runs each assertion in an isolated
    subprocess so Colang loaded by one case (for example, resolving LLMRails)
    cannot leak into another and mask a regression. It covers: Colang-free
    import nemoguardrails, ChatMessage, RailsConfig, a built-in rail action,
    and nemoguardrails.llm.call; LLMRails loading the legacy runtime;
    Guardrails resolving; the env-var alias and its re-evaluation on reload; and
    the dir/__all__/invalid-attribute contracts.
  • No live LLM or provider calls added.

AI Assistance

  • No AI tools were used.
  • AI tools were used; a human reviewed and can explain every change (tool: Claude Code).

Checklist

  • I've read the CONTRIBUTING guidelines.
  • This PR links to a triaged issue assigned to me.
  • My PR title follows the project commit convention.
  • I've updated the documentation if applicable.
  • I've added tests if applicable.
  • I've noted any verification beyond CI and any checks I couldn't run.
  • I did not update generated changelog files manually.
  • I addressed all CodeRabbit, Greptile, and other review comments, or replied with why no change is needed.
  • @mentions of the person or team responsible for reviewing proposed changes.

tgasser-nv and others added 5 commits July 31, 2026 16:39
…ing Colang at import-time. Decouples actions from Colang
Resolve the public names exported from `nemoguardrails` and
`nemoguardrails.rails` lazily via PEP 562 `__getattr__` instead of importing
them eagerly at package import time.

A bare `import nemoguardrails` previously booted ~700 modules (Colang parsers,
tracing/OpenTelemetry, aiohttp, jinja2, ...) because the top-level package
eagerly imported `RailsConfig`, `LLMRails`, and `Guardrails`. Lazy resolution
defers that cost to first use of a heavy symbol, so lightweight consumers
(version probes, type imports, submodule imports, built-in rail actions) no
longer pay it.

Making `nemoguardrails.rails` lazy also lets `rails.llm.options` initialize the
package without pulling in the full runtime, which removes the circular import
that made isolated rail-action imports fail, and removes the need for the
explicit `nemoguardrails.rails` preload workaround.

Behavior preserved: `__all__`, `dir()`, type-checker exports (via
`TYPE_CHECKING`), unknown-attribute `AttributeError`, and the
`NEMO_GUARDRAILS_IORAILS_ENGINE` alias of `LLMRails` to `Guardrails`.
Resolution is intentionally not cached into module globals so that
`importlib.reload(nemoguardrails)` re-reads the env alias.

Also drops a now-redundant `cast(str, chunk)` in cli/chat.py that stricter
type resolution exposes.

Signed-off-by: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com>
@github-actions github-actions Bot added status: needs triage New issues that have not yet been reviewed or categorized. size: M labels Aug 5, 2026
@codecov

codecov Bot commented Aug 5, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 72.00000% with 7 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
nemoguardrails/rails/__init__.py 50.00% 5 Missing ⚠️
nemoguardrails/__init__.py 92.85% 1 Missing ⚠️
nemoguardrails/cli/chat.py 0.00% 1 Missing ⚠️

📢 Thoughts on this report? Let us know!

@tgasser-nv
tgasser-nv force-pushed the refactor/iorails-actions-decouple-colang branch from ff09623 to 428d84d Compare August 6, 2026 15:16
Base automatically changed from refactor/iorails-actions-decouple-colang to develop August 6, 2026 15:24
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

size: M status: needs triage New issues that have not yet been reviewed or categorized.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants