perf(imports): stop loading optional integrations to import grpo

[tianyi-zhang-02]

What does this PR do ?

Stops import nemo_rl.algorithms.grpo from loading wandb, mlflow, swanlab, matplotlib, fastapi and uvicorn. None are needed unless the corresponding feature is used, but every CLI entrypoint, Ray actor and unit-test collection paid for all of them.

Warm cache, median of 9 runs:

	before	after
import nemo_rl.algorithms.grpo	9.04 s	3.77 s (−58%)
modules loaded	7150	5256 (−26%)

Afterwards wandb, mlflow, fastapi and uvicorn are absent from sys.modules. The seconds are from one laptop; the module count is deterministic and the number I would trust.

What moved, in descending order of cost:

• utils/logger.py — mlflow, swanlab, wandb, matplotlib.pyplot become lazy module proxies. This is most of the win (nemo_rl.utils.logger alone: 5597 → 3409 modules).
• experience/{metric_utils,rollout_manager,rollouts}.py — from wandb import Table/Histogram, one call site each.
• utils/memory_tracker.py — from ray.scripts.scripts import memory_summary, used only in get_snapshot_str().
• generation/vllm/{vllm_worker_async,vllm_sparse_refit}.py — fastapi/uvicorn behind TYPE_CHECKING.

Issues

Supersedes and closes #3437, which tried a piece of this and does not actually work — applying it alone leaves the module count at exactly 7150. vllm_sparse_refit.py:29-31 imports fastapi at module scope as well, and the vLLM workers are not the real source anyway:

algorithms/grpo.py:127     -> utils/memory_tracker.py:19
                           -> from ray.scripts.scripts import memory_summary
                              (Ray's CLI entrypoint, pulls in the dashboard stack)

Usage

No config or API change. nemo_rl.utils.logger.wandb and friends resolve exactly as before.

Before your PR is "Ready for review"

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests? (none — see below)
• Did you run the unit tests and functional tests locally? (unit only)
• Did you add or update any necessary documentation? (n/a)

Additional Information

Why a proxy in logger.py rather than function-scope imports. tests/unit/utils/test_logger.py has ~100 patch("nemo_rl.utils.logger.<name>") targets. Moving those imports into methods removes the module attribute and breaks 28 tests — I tried that first. Lazy module-level names mean neither the 15 call sites nor a single test changed. plt.Figure annotations are quoted so importing the module does not touch the proxy.

No new tests. Behaviour is unchanged, and asserting "module X is not in sys.modules" is easy to write but easy to make flaky. Happy to add a subprocess test pinning whichever of the four you want guarded.

Verification. 195 unit tests across tests/unit/utils/test_logger.py, tests/unit/experience/ and tests/unit/algorithms/test_utils.py pass before and after, no test modified. ruff check and ruff format --check clean; none of the seven files are in pyrefly.toml. I have not run the GPU or functional suites.

cc. @RayenTian @yuki-97

[copy-pr-bot]

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.

[tianyi-zhang-02]

Force-pushed after rebuilding this on current main. Two things were wrong with the previous version, both worth recording:

It silently reverted ~115 lines of the vLLM 0.25 migration. I had taken vllm_worker_async.py wholesale from an older branch, so the diff carried that branch's version of the whole file rather than just the import change. Against current main that reads as a revert: OnlineRenderer back to OpenAIServingRender, the parser= argument back to the tool_parser=/reasoning_parser= pair, the spec_lookahead handling deleted, and — worst — the val_temperature/val_top_p validation-sampling branch replaced by a bare assert request.temperature == generation_config["temperature"], which would have failed every validation rollout. It merged cleanly because main had not touched those exact lines, so nothing would have flagged it.

Request was left as an unquoted annotation on two nested async defs in vllm_sparse_refit.py after moving the import under TYPE_CHECKING. Annotations on a nested def evaluate when the enclosing function runs, so that is a NameError at request-handling time, not a typing nit. ruff check --select TC catches it; the repo's ruff config selects only D,F, so CI would not have.

The branch now differs from main by 11 lines in vllm_worker_async.py — only the import deferral and two quoted annotations.

Re-measured on the rebuilt branch, warm cache, median of 7:

	before	after
import nemo_rl.algorithms.grpo	7.71 s	3.42 s (−56%)
modules	7098	5203 (−27%)

wandb, mlflow, fastapi and uvicorn are absent from sys.modules afterwards. 196 tests across tests/unit/utils/test_logger.py and tests/unit/experience/ pass identically before and after, no test modified. ruff check, ruff format --check and ruff check --select TC all clean.