cp: fix: Fix the logger error in non-colocated sync-grpo code path (1355) into r0.4.0

[chtruong814]

beep boop [🤖]: Hi @youngeunkwon0405 👋,

we've cherry picked #1355 into  for you! 🚀

Please review and approve this cherry pick by your convenience!

Summary by CodeRabbit

• Bug Fixes
  • Fixed metrics logging behavior to prevent potential errors in performance and training logs. Adjusted how certain metrics are handled during the logging process to ensure compatibility with external monitoring systems.

[coderabbitai]

📝 Walkthrough

Walkthrough

The PR relocates the deletion of per_worker_token_counts from the GRPO algorithm's metric logging in grpo.py to a centralized cleanup step within the print_performance_metrics utility function in utils.py, ensuring consistent metric cleanup before logging.

Changes

Cohort / File(s)	Summary
Metric cleanup relocation nemo_rl/algorithms/grpo.py, nemo_rl/algorithms/utils.py	Removed conditional deletion of per_worker_token_counts from grpo.py; added centralized cleanup step in print_performance_metrics after FLOPS logging to prevent wandb logging errors.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

Possibly related PRs

• NVIDIA-NeMo/RL#1183: Modifies performance-metrics logging flow and per_worker_token_counts handling in the same functions and utilities.
• NVIDIA-NeMo/RL#1355: Directly related metric cleanup refactoring, moving per_worker_token_counts deletion from GRPO training code to centralized utility function.

Suggested labels

CI:L1, r0.4.0

Suggested reviewers

• terrykong

Pre-merge checks and finishing touches

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Title Check	⚠️ Warning	The pull request title includes cherry-pick metadata and branch information rather than clearly and concisely summarizing the actual change to fix a logging error in the sync-grpo code path.	Rewrite the title as a short, single sentence describing the fix, for example “Fix logger error in non-colocated sync-GRPO code path,” and remove cherry-pick and branch details from the title.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Test Results For Major Changes	✅ Passed	This PR backports a logging bugfix that only adjusts the handling of a metrics field without introducing new features, breaking changes, algorithmic modifications, or performance-critical alterations, so it is considered minor and the absence of explicit test results or performance benchmarks aligns with the guidelines.

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment
• Commit unit tests in branch cherry-pick-1355-r0.4.0

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 0

🧹 Nitpick comments (1)

nemo_rl/algorithms/utils.py (1)

556-564: Use .pop() for more idiomatic cleanup.

The cleanup logic is well-placed and correctly prevents wandb logging errors. However, the pattern can be simplified using .pop() with a default value.

Apply this diff:

-    # Clean up metrics to avoid wandb logging errors
-    # Dict structures cannot be logged to wandb
-    if "per_worker_token_counts" in metrics:
-        del metrics["per_worker_token_counts"]
+    # Clean up metrics to avoid wandb logging errors
+    # Dict structures cannot be logged to wandb
+    metrics.pop("per_worker_token_counts", None)

This change makes the code more concise and eliminates the need for an existence check.

Based on static analysis hints.

📜 Review details

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 1b81f38 and 0443968.

📒 Files selected for processing (2)

• nemo_rl/algorithms/grpo.py (0 hunks)
• nemo_rl/algorithms/utils.py (1 hunks)

💤 Files with no reviewable changes (1)

• nemo_rl/algorithms/grpo.py

🧰 Additional context used

📓 Path-based instructions (2)

**/*.py

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

**/*.py: Follow the Google Python Style Guide for all Python code
Target Python 3.12+ for all Python code in NeMo-RL
Indent Python code with 4 spaces; do not use tabs
Python filenames should be snake_case (e.g., some_file.py)
Class names should be PascalCase
Function and method names should be snake_case
Local variable names should be snake_case; if starting with a number, prefix with k (e.g., k_99th_percentile)
Global variables should be UPPER_SNAKE_CASE and prefixed with G_ (e.g., G_MY_GLOBAL)
Constants should be UPPER_SNAKE_CASE
Avoid shadowing variables declared in an outer scope
Initialize all externally visible members of a class in the constructor
For public interfaces used outside a file, prefer docstrings over comments
Use comments mainly for code within a function or interfaces local to a file
Commented-out code must include a nearby comment explaining usage and why it is commented out; otherwise remove before merging
Use Google-style docstrings for classes and functions (Sphinx-parseable)
Avoid using reflection when functionality can be easily achieved without it
Limit except clauses to the smallest specific set of exceptions possible
For duck-typing via try/except, keep the try body minimal and use else for main logic
Add the NVIDIA copyright header (with current year) at the top of all Python files, excluding tests/ and test-only scripts

Files:

• nemo_rl/algorithms/utils.py

nemo_rl/**/*.py

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

nemo_rl/**/*.py: Do not set non-None configuration defaults in code; YAML is the single source of truth for defaults
Access required config attributes directly (e.g., policy_cfg["precision"]) and assume presence; do not introduce hidden defaults
Express configuration optionality via TypedDict using typing.NotRequired
When adding a new config key to a TypedDict subclass, document the key’s purpose, valid values/types, and recommended default in code
For any class or function decorated with @ray.remote, add '# pragma: no cover' on the class/def line (and on remote functions)

Files:

• nemo_rl/algorithms/utils.py

🪛 Ruff (0.14.0)

nemo_rl/algorithms/utils.py

563-563: Use pop instead of key in dict followed by del dict[key]

Replace if statement with .pop(..., None)

(RUF051)

⏰ 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). (5)

• GitHub Check: Lint check
• GitHub Check: Lint check
• GitHub Check: Lint check
• GitHub Check: Post submodule check comment / Comment on PR
• GitHub Check: Post automodel integration comment / Comment on PR