Add unit tests for train_rl.get_rollout_kwargs_for_parallelism

[igorts-git]

Description

This is part of the RL testing Fixit.
It adds a few simple unit tests to the train_rl.get_rollout_kwargs_for_parallelism function.

FIXES: b/490410802, b/490410719, b/490411080, b/490410543

Pair-programmed with Gemini-CLI.

Tests

This PR is a collection of new unit tests. No runtime functionality has been changed.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[github-actions]

🤖 Hi @igorts-git, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[github-actions]

🤖 Hi @igorts-git, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[github-actions]

## 📋 Review Summary

This PR successfully adds clear and targeted unit tests for the train_rl.get_rollout_kwargs_for_parallelism function, improving test coverage for the RL testing fixit. The tests are logically structured and effectively cover both successful auto-derivation cases and common error scenarios.

🔍 General Feedback

• Good Use of Mocks: The usage of SimpleNamespace to mock sampler_config is clean and minimizes setup overhead for the tests.
• Clear Assertions: Utilizing assertRaisesRegex for error testing effectively ensures that the correct validation branches are being triggered.
• Test Structure: The test cases are well-documented with descriptive docstrings and comments, making the intent of each test immediately understandable.

[xuefgu]

Thanks!