Repository navigation
fix: can't find transformers_modules error for moonlight - #1124
Conversation
WalkthroughAdds patch_transformers_module_dir to venv utilities and calls it when constructing worker runtime environments to insert the Transformers modules directory into PYTHONPATH after setting virtual-environment variables. No public API signatures changed. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant Caller
participant WorkerGroups
participant venvs as venvs.patch_transformers_module_dir
participant Env as runtime_env["env_vars"]
participant Worker
Caller->>WorkerGroups: create_workers_from_bundle_indices(...)
WorkerGroups->>Env: set VIRTUAL_ENV, UV_PROJECT_ENVIRONMENT
WorkerGroups->>venvs: patch_transformers_module_dir(Env)
venvs-->>WorkerGroups: Env updated (PYTHONPATH includes modules dir)
WorkerGroups->>Worker: spawn worker with runtime_env
Worker-->>WorkerGroups: worker ready
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Poem
Pre-merge checks and finishing touches✅ Passed checks (3 passed)
✨ Finishing touches
🧪 Generate unit tests
Tip 👮 Agentic pre-merge checks are now available in preview!Pro plan users can now enable pre-merge checks in their settings to enforce checklists before merging PRs.
Please see the documentation for more information. Example: reviews:
pre_merge_checks:
custom_checks:
- name: "Undocumented Breaking Changes"
mode: "warning"
instructions: |
Pass/fail criteria: All breaking changes to public APIs, CLI flags, environment variables, configuration keys, database schemas, or HTTP/GraphQL endpoints must be documented in the "Breaking Change" section of the PR description and in CHANGELOG.md. Exclude purely internal or private changes (e.g., code not exported from package entry points or explicitly marked as internal).Please share your feedback with us on this Discord post. Comment |
|
@yfw to review |
There was a problem hiding this comment.
Actionable comments posted: 2
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
nemo_rl/distributed/worker_groups.py(2 hunks)nemo_rl/utils/venvs.py(1 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
nemo_rl/distributed/worker_groups.py (1)
nemo_rl/utils/venvs.py (2)
create_local_venv_on_each_node(152-189)patch_transformers_module_dir(193-203)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: CI quality check
🔇 Additional comments (1)
nemo_rl/distributed/worker_groups.py (1)
30-30: Import looks fine.No coupling concerns; colocating venv helpers here is reasonable.
ℹ️ File Consistency CheckCheck based on commit: d32f4e6 (PR #1124 from This is a test comment This check ensures that related file implementations remain synchronized across the codebase. If you believe this warning is incorrect or the files should intentionally differ, please add a comment explaining the reasoning. |
ℹ️ File Consistency CheckCheck based on commit: 6dda8a3 (PR #1124 from This is a test comment This check ensures that related file implementations remain synchronized across the codebase. If you believe this warning is incorrect or the files should intentionally differ, please add a comment explaining the reasoning. |
ℹ️ File Consistency CheckCheck based on commit: 20e99a1 (PR #1124 from This is a test comment This check ensures that related file implementations remain synchronized across the codebase. If you believe this warning is incorrect or the files should intentionally differ, please add a comment explaining the reasoning. |
ℹ️ File Consistency CheckCheck based on commit: dac601a (PR #1124 from This is a test comment This check ensures that related file implementations remain synchronized across the codebase. If you believe this warning is incorrect or the files should intentionally differ, please add a comment explaining the reasoning. |
There was a problem hiding this comment.
Actionable comments posted: 0
♻️ Duplicate comments (2)
nemo_rl/utils/venvs.py (2)
208-213: Useos.pathsepand de-duplicatePYTHONPATHentries.Hardcoding ":" is not cross‑platform and duplicate entries bloat envs. The diff above addresses both by using
os.pathsepand checking membership before prepending.
202-207: Avoid driver-sidetransformersimport and hard assert; compute HF modules path robustly, no-op if absent.Importing
transformershere can fail on the driver (only installed inside worker venvs) and the hard assert bricks clusters before any modules are downloaded. Prefer a resilient lookup (env ->huggingface_hub.constants-> conservative defaults), use a no-op with a debug log if missing, and add a return type. Also normalize/expand paths.Apply this diff:
-def patch_transformers_module_dir(env_vars: dict[str, str]): - """Updates the PYTHONPATH in the given env_vars to include downloaded Transformers modules. +def patch_transformers_module_dir(env_vars: dict[str, str]) -> dict[str, str]: + """Prepend the Hugging Face dynamic modules cache to PYTHONPATH so imports like + `transformers_modules.*` resolve at runtime. No-op if the cache dir is absent. - Some models, such as `moonshotai/Moonlight-16B-A3B`, may raise - `ModuleNotFoundError: No module named 'transformers_modules'` during - tokenizer deserialization when passed to a `Policy` object. This patch - prevents that error. + Some models (e.g., Moonlight) raise `ModuleNotFoundError: transformers_modules` + during tokenizer deserialization on Ray workers. This ensures the modules cache + is discoverable without requiring `transformers` on the driver. - Assumption: the cache directory is shared across virtual environments. + Assumption: the cache directory is shared across virtual environments. """ - from transformers.utils.hub import TRANSFORMERS_CACHE - - module_dir = os.path.join(TRANSFORMERS_CACHE, "..", "modules") - assert os.path.exists(module_dir), ( - f"TRANSFORMERS_CACHE is not set or does not exist: {module_dir}" - ) - if "PYTHONPATH" not in env_vars: - env_vars["PYTHONPATH"] = module_dir - else: - env_vars["PYTHONPATH"] = f"{module_dir}:{env_vars['PYTHONPATH']}" - - return env_vars + # 1) Explicit override if provided + module_dir: str | None = os.environ.get("HF_MODULES_CACHE") + if module_dir: + module_dir = os.path.expanduser(module_dir) + else: + # 2) Prefer lightweight HF Hub constants (no heavy import on the driver) + try: + from huggingface_hub.constants import HF_MODULES_CACHE as _HF_MODULES_CACHE # type: ignore + module_dir = os.path.expanduser(_HF_MODULES_CACHE) + except Exception: + # 3) Last-resort default (avoid importing transformers on the driver) + module_dir = os.path.join(os.path.expanduser("~/.cache/huggingface"), "modules") + + module_dir = os.path.normpath(module_dir) + if not os.path.isdir(module_dir): + logger.debug("HF modules cache not found at %s; skipping PYTHONPATH patch.", module_dir) + return env_vars + + existing = env_vars.get("PYTHONPATH", "") + parts = [p for p in existing.split(os.pathsep) if p] if existing else [] + if module_dir not in parts: + env_vars["PYTHONPATH"] = os.pathsep.join([module_dir] + parts) if parts else module_dir + return env_vars
🧹 Nitpick comments (2)
nemo_rl/utils/venvs.py (2)
192-213: Optional: consider lightweighthuggingface_hubfallback check.If you prefer keeping a fallback that derives the path from
transformerswhen available, we can add a nested try importingTRANSFORMERS_CACHEand compute../modules. Omitted in the diff to avoid heavy imports on drivers.If you want that variant, I can provide an alternative diff that adds the guarded
transformersfallback.
192-201: Docstring: clarify behavior and call-site assumptions.Add one sentence to the function docstring stating it’s a safe no‑op when the Transformers modules cache is absent and that it should be called only after runtime_env["env_vars"] is fully composed (e.g., after VIRTUAL_ENV / UV_PROJECT_ENVIRONMENT are set). Call site confirmed: nemo_rl/distributed/worker_groups.py invokes patch_transformers_module_dir(runtime_env["env_vars"]) after setting those env vars (line 516).
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
nemo_rl/distributed/worker_groups.py(2 hunks)nemo_rl/utils/venvs.py(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- nemo_rl/distributed/worker_groups.py
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (2)
- GitHub Check: Lint check
- GitHub Check: Post submodule check comment / Comment on PR
🔇 Additional comments (1)
nemo_rl/utils/venvs.py (1)
192-213: Ignore HF_MODULES_CACHE — it's not a public constant; keep using TRANSFORMERS_CACHE.huggingface_hub.constants does not export HF_MODULES_CACHE; the current approach (from transformers.utils.hub import TRANSFORMERS_CACHE and deriving the modules dir via TRANSFORMERS_CACHE/../modules) is correct.
Likely an incorrect or invalid review comment.
|
@joyang-nv @terrykong I think this PR has an issue. When CI ran in the queue, it failed with this in the functional tests https://github.com/NVIDIA-NeMo/RL/actions/runs/17822565027/job/50687190084 I see. Thanks for info. I will fix it. |
Signed-off-by: Jonas Yang <joyang@nvidia.com>
8b3437b to
3fef216
Compare
|
@terrykong can you approve this once you saw CI pass? And not sure whether still need @chtruong814 to help merge? |
Signed-off-by: Jonas Yang <joyang@nvidia.com> Signed-off-by: Jonas Yang CN <joyang@nvidia.com> Signed-off-by: Terry Kong <terryk@nvidia.com> Co-authored-by: Terry Kong <terrycurtiskong@gmail.com> Co-authored-by: Terry Kong <terryk@nvidia.com>
…#1124) Signed-off-by: Jonas Yang <joyang@nvidia.com> Signed-off-by: Jonas Yang CN <joyang@nvidia.com> Signed-off-by: Terry Kong <terryk@nvidia.com> Co-authored-by: Terry Kong <terrycurtiskong@gmail.com> Co-authored-by: Terry Kong <terryk@nvidia.com>
…#1124) Signed-off-by: Jonas Yang <joyang@nvidia.com> Signed-off-by: Jonas Yang CN <joyang@nvidia.com> Signed-off-by: Terry Kong <terryk@nvidia.com> Co-authored-by: Terry Kong <terrycurtiskong@gmail.com> Co-authored-by: Terry Kong <terryk@nvidia.com> Signed-off-by: yuanhangs <yuanhangs@nvidia.com>
What does this PR do ?
When you download some HF models with modules like moonlight, venvs might has assertion which doesn't have transformers_modules module. This is due to downloaded modules path is not added to
PYTHONPATH.Issues
List issues that this PR closes (syntax):
Usage
# Add a code snippet demonstrating how to use thisBefore your PR is "Ready for review"
Pre checks:
Additional Information
Summary by CodeRabbit
Bug Fixes
Chores