feat: Support prefetching of specific envs

[hemildesai]

Useful to prefetch specific venvs

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
• Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Summary by CodeRabbit

• New Features

  • Added ability to selectively prefetch virtual environments with optional filters targeting specific actors by name.
  • Enabled command-line support for specifying prefetch filters.

• Tests

  • Added comprehensive test coverage for prefetching functionality with various filtering scenarios.

_{✏️ Tip: You can customize this high-level summary in your review settings.}

[coderabbitai]

📝 Walkthrough

Walkthrough

The prefetch_venvs utility function has been enhanced with optional filtering capabilities. The function now accepts a filters parameter to selectively prefetch virtual environments based on actor FQN patterns, includes CLI support via argparse, adds runtime logging, and removes unused frozen environment wrapper code.

Changes

Cohort / File(s)	Summary
Core Functionality Enhancement nemo_rl/utils/prefetch_venvs.py	Updated prefetch_venvs() signature to accept optional filters parameter; added argparse CLI support for filtered prefetching; integrated runtime logging of provided filters; removed unused frozen environment wrapper script creation code; filtering logic applied before venv processing with error tolerance.
Test Suite tests/unit/utils/test_prefetch_venvs.py	New comprehensive test suite with mock registry fixture and 9 test cases covering no-filter prefetching, single/multiple substring filters (OR logic), no-match scenarios, system Python actor exclusion, partial FQN matching, empty filter lists, error continuation behavior, and case-sensitive filtering.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Pre-merge checks and finishing touches

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Test Results For Major Changes	⚠️ Warning	PR description lacks test results or testing information despite comprehensive unit test suite being included in the codebase.	Update PR description to document the 9 unit tests added, their coverage areas, and test execution results.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title 'feat: Support prefetching of specific envs' accurately describes the main change—adding the ability to selectively prefetch virtual environments based on filters.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.

✨ 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 hemil/update-prefetch

---

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 (3)

nemo_rl/utils/prefetch_venvs.py (3)

23-30: Consider adding explicit type information to the docstring.

The docstring is clear but could be more explicit about the parameter type for better documentation clarity.

🔎 Suggested enhancement

 def prefetch_venvs(filters=None):
     """Prefetch all virtual environments that will be used by workers.
 
     Args:
-        filters: List of strings to match against actor FQNs. If provided, only
-                actors whose FQN contains at least one of the filter strings will
-                be prefetched. If None, all venvs are prefetched.
+        filters (Optional[List[str]]): List of strings to match against actor FQNs. 
+            If provided, only actors whose FQN contains at least one of the filter 
+            strings will be prefetched. If None, all venvs are prefetched.
     """

---

31-65: Consider using logger instead of print() for consistency.

The function uses print() statements throughout, while nemo_rl/utils/venvs.py uses logger.info() for similar output. Using a logger would provide better control over output levels, formatting, and integration with logging infrastructure.

🔎 Example refactor using logger

Add logger import at the top of the file:

import logging

logger = logging.getLogger(__name__)

Then replace print statements with logger calls:

-    print("Prefetching virtual environments...")
+    logger.info("Prefetching virtual environments...")
     if filters:
-        print(f"Filtering for: {filters}")
+        logger.info(f"Filtering for: {filters}")

Apply similar changes to other print statements in the function.

---

68-93: Simplify the conditional on line 93.

The conditional args.filters if args.filters else None is redundant since an empty list is falsy. The expression can be simplified.

🔎 Proposed simplification

     args = parser.parse_args()
 
-    prefetch_venvs(filters=args.filters if args.filters else None)
+    prefetch_venvs(filters=args.filters or None)

📜 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 669e70c and e2bb023.

📒 Files selected for processing (2)

• nemo_rl/utils/prefetch_venvs.py
• tests/unit/utils/test_prefetch_venvs.py

🧰 Additional context used

📓 Path-based instructions (4)

**/*.py

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

**/*.py: Conform code to Python 3.12+
Indent code with 4 spaces. Do not use tabs
Use snake_case for file names
Use PascalCase for class names
Use snake_case for function and method names
Use snake_case for local variables
Prefix variable names that start with a number with 'k' (e.g., k_99th_percentile)
Use upper snake_case with 'G' prefix for global variables (e.g., G_MY_GLOBAL)
Use upper snake_case for constants
Avoid shadowing variables declared in an outer scope
Initialize all externally visible members of a class in the constructor
Prefer docstrings over comments for interfaces that may be used outside a file
Reserve comments for code within a function or interfaces that are local to a file
If a piece of code is commented out, include a comment describing its usage and why it's commented out. Remove debug comments before merging
Use Google style docstrings for classes and functions in Python, which can be parsed by Sphinx
Avoid using reflection when functionality can be easily achieved without reflection
When using try-except blocks, limit the except clause to the smallest set of specific errors possible
When using try-except blocks for duck-typing, keep the body of the try as small as possible and use the else block for logic
YAML is the single source of truth for configuration defaults. Do not set non-None defaults in code for configuration values
For required configuration attributes, access config directly and expect presence (e.g., policy_cfg['precision']) without hidden defaults
Use typing.NotRequired to mark optional attributes in TypedDict for configuration
When adding a new config key to a TypedDict subclass, document the key's purpose, valid values/types, and recommended default, and reflect the default in exemplar YAMLs under examples/configs/*.yaml
Follow the Google Python Style Guide for Python code

Files:

• nemo_rl/utils/prefetch_venvs.py
• tests/unit/utils/test_prefetch_venvs.py

nemo_rl/**/*.py

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

For any source file under nemo_rl/*.py that defines a class or function decorated with @ray.remote, add a coverage pragma (# pragma: no cover) because these run in separate Ray processes

Files:

• nemo_rl/utils/prefetch_venvs.py

!(**/tests/**|**/test_*.py|**/test_*.sh)

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

Add the NVIDIA copyright header to all Python files and shell scripts (excluding tests). The header should include the current year

Files:

• nemo_rl/utils/prefetch_venvs.py
• tests/unit/utils/test_prefetch_venvs.py

**/*.{py,sh}

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

The NVIDIA copyright header should appear at the top of all Python files and shell scripts (excluding tests)

Files:

• nemo_rl/utils/prefetch_venvs.py
• tests/unit/utils/test_prefetch_venvs.py

🧬 Code graph analysis (2)

nemo_rl/utils/prefetch_venvs.py (1)

nemo_rl/utils/venvs.py (1)

• create_local_venv (34-102)

tests/unit/utils/test_prefetch_venvs.py (1)

nemo_rl/utils/prefetch_venvs.py (1)

• prefetch_venvs (23-65)

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

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

🔇 Additional comments (5)

nemo_rl/utils/prefetch_venvs.py (3)

14-14: LGTM: Import added for CLI support.

The argparse import is properly placed and necessary for the new CLI functionality.

---

38-40: LGTM: Filter logic correctly implements OR semantics.

The filtering logic properly checks if any filter string is a substring of the actor FQN, implementing the expected OR behavior for multiple filters.

---

57-63: LGTM: Appropriate error handling for prefetch utility.

The broad exception handling is appropriate here since prefetching is error-tolerant by design. The explicit comment and the test coverage (test_prefetch_venvs_continues_on_error) confirm this is the intended behavior.

tests/unit/utils/test_prefetch_venvs.py (2)

19-28: LGTM: Well-designed test fixture.

The mock registry provides good coverage of different actor types and executable configurations, enabling comprehensive testing of the filtering logic.

---

31-240: LGTM: Comprehensive test coverage.

The test suite thoroughly covers all filtering scenarios including:

• Base case with no filters
• Single and multiple filter strings with OR logic
• Non-matching filters
• System Python exclusion
• Partial substring matching
• Empty filter list behavior
• Error tolerance and continuation
• Case sensitivity

The tests are well-structured, use appropriate mocking, and have clear assertions validating the expected behavior.