Conversation
Signed-off-by: Yuki Huang <yukih@nvidia.com>
|
/ok to test c8cf4c1 |
ℹ️ File Consistency CheckCheck based on commit: c8cf4c1 (PR #2094 from ✅ DTensor Policy Worker Synchronization CheckBoth DTensor policy worker files were modified in this PR:
Please ensure that the changes are consistent between both files where applicable. 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. |
📝 WalkthroughWalkthroughThe changes introduce environment-specific actor class identifier support in worker group management, add a new method for executing named methods across all workers with shared input data, refactor worker class implementations to separate Ray remote constraints from concrete logic, and expand documentation and examples for research templates with new worker extension patterns. Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes 🚥 Pre-merge checks | ✅ 2 | ❌ 2❌ Failed checks (1 warning, 1 inconclusive)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches
🧪 Generate unit tests (beta)
Tip Try Coding Plans. Let us write the prompt for your AI agent so you can ship faster (with fewer bugs). Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
research/template_project/single_update.py (1)
86-92:⚠️ Potential issue | 🟠 MajorUse an extension FQN that is importable under the documented invocation.
Line 91 hard-codes
research.template_project.template_project.worker_extension.DTensorPolicyWorkerV2Extension, but this is inconsistent with the script's own imports (line 34 usesfrom template_project.data_utils import ...) and the documented invocation method (uv run single_update.pyfrom theresearch/template_projectdirectory). When the script is run as documented,template_project.*is directly importable, butresearch.*is not in the module search path. Ray worker processes will fail to load the extension withModuleNotFoundError: No module named 'research'unless the repo root is manually added toPYTHONPATH.Suggested fix
- worker_extension_cls="research.template_project.template_project.worker_extension.DTensorPolicyWorkerV2Extension", + worker_extension_cls="template_project.worker_extension.DTensorPolicyWorkerV2Extension",🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@research/template_project/single_update.py` around lines 86 - 92, The hard-coded extension FQN in the Policy constructor is not importable when running from the documented working directory; change the worker_extension_cls value passed to Policy (the Policy(...) instantiation) from "research.template_project.template_project.worker_extension.DTensorPolicyWorkerV2Extension" to the package-relative import used elsewhere, e.g. "template_project.worker_extension.DTensorPolicyWorkerV2Extension", so the worker extension can be imported when running `uv run single_update.py`.
🧹 Nitpick comments (1)
nemo_rl/models/policy/lm_policy.py (1)
136-142: Fail fast on invalid extension classes.This currently prints an advisory but defers all validation to Ray startup. Importing the extension class up front and checking
issubclass(...)against the resolved base worker would turn typos and incompatible extensions into a much clearer local error.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@nemo_rl/models/policy/lm_policy.py` around lines 136 - 142, When a worker_extension_cls is provided, validate it immediately by checking issubclass(worker_extension_cls, worker_builder_cls) and raise a clear TypeError if the check fails instead of merely printing; on success set worker_builder_cls_for_env = worker_builder_cls and then override worker_builder_cls = worker_extension_cls so downstream code uses the validated extension. Ensure the error message names the offending worker_extension_cls and the expected base worker_builder_cls to aid debugging.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py`:
- Around line 1129-1133: The coverage pragma is placed on the `@ray.remote`
decorator closing line instead of the class definition; move the "# pragma: no
cover" comment to the class line for DTensorPolicyWorkerV2 so the class
declaration reads with the pragma (i.e., ensure the coverage comment is attached
to the "class DTensorPolicyWorkerV2(DTensorPolicyWorkerV2Impl):" line), leaving
the
`@ray.remote`(runtime_env=get_runtime_env_for_policy_worker("dtensor_policy_worker_v2"))
decorator unchanged.
In `@research/README.md`:
- Around line 28-30: The documentation currently swaps the repository references
for example config paths; update the README sentence that mentions
`research/template_project/configs` and `examples/configs` so they point to the
correct repository examples: `research/template_project/configs` is the research
template example and `examples/configs` is the core repository example—edit the
text near the "Configuration" section to swap the descriptions so each path is
labeled with its correct source.
---
Outside diff comments:
In `@research/template_project/single_update.py`:
- Around line 86-92: The hard-coded extension FQN in the Policy constructor is
not importable when running from the documented working directory; change the
worker_extension_cls value passed to Policy (the Policy(...) instantiation) from
"research.template_project.template_project.worker_extension.DTensorPolicyWorkerV2Extension"
to the package-relative import used elsewhere, e.g.
"template_project.worker_extension.DTensorPolicyWorkerV2Extension", so the
worker extension can be imported when running `uv run single_update.py`.
---
Nitpick comments:
In `@nemo_rl/models/policy/lm_policy.py`:
- Around line 136-142: When a worker_extension_cls is provided, validate it
immediately by checking issubclass(worker_extension_cls, worker_builder_cls) and
raise a clear TypeError if the check fails instead of merely printing; on
success set worker_builder_cls_for_env = worker_builder_cls and then override
worker_builder_cls = worker_extension_cls so downstream code uses the validated
extension. Ensure the error message names the offending worker_extension_cls and
the expected base worker_builder_cls to aid debugging.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 973e0925-c2cd-4c67-9924-08830faf7698
📒 Files selected for processing (9)
nemo_rl/distributed/worker_groups.pynemo_rl/models/policy/lm_policy.pynemo_rl/models/policy/workers/dtensor_policy_worker.pynemo_rl/models/policy/workers/dtensor_policy_worker_v2.pynemo_rl/models/policy/workers/megatron_policy_worker.pyresearch/README.mdresearch/template_project/README.mdresearch/template_project/single_update.pyresearch/template_project/template_project/worker_extension.py
ℹ️ File Consistency CheckCheck based on commit: cd05258 (PR #2094 from ✅ DTensor Policy Worker Synchronization CheckBoth DTensor policy worker files were modified in this PR:
Please ensure that the changes are consistent between both files where applicable. 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. |
|
/ok to test cd05258 |
Signed-off-by: Yuki Huang <yukih@nvidia.com>
Signed-off-by: Yuki Huang <yukih@nvidia.com>
Signed-off-by: Yuki Huang <yukih@nvidia.com>
yuki-97
force-pushed
the
yukih/research-template
branch
from
cd05258 to
4361031
Compare
ℹ️ File Consistency CheckCheck based on commit: 4361031 (PR #2094 from ✅ DTensor Policy Worker Synchronization CheckBoth DTensor policy worker files were modified in this PR:
Please ensure that the changes are consistent between both files where applicable. 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. |
|
/ok to test 4361031 |
Signed-off-by: Yuki Huang <yukih@nvidia.com>
ℹ️ File Consistency CheckCheck based on commit: ade0c15 (PR #2094 from ✅ DTensor Policy Worker Synchronization CheckBoth DTensor policy worker files were modified in this PR:
Please ensure that the changes are consistent between both files where applicable. 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. |
1 participant
Support research project to add custom functions to policy workers by using extension class.
Summary by CodeRabbit
Release Notes
New Features
worker_extension_clsparameter in Policy initializationDocumentation