fix: better error handling and message in refit #1477
Conversation
Signed-off-by: Zhiyu Li <[email protected]>
📝 WalkthroughWalkthroughThe PR enhances error handling in IPC communication by replacing silent failures with explicit exception raising. It adds traceback imports and modifies two functions to catch specific ZMQ exceptions and re-raise them as RuntimeError or TimeoutError with detailed diagnostic information, improving visibility into failures. Changes
Estimated code review effort🎯 2 (Simple) | ⏱️ ~10 minutes
Pre-merge checks and finishing touches❌ Failed checks (1 inconclusive)
✅ Passed checks (2 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
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: 3
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
nemo_rl/models/generation/vllm/vllm_backend.py(2 hunks)nemo_rl/models/policy/utils.py(2 hunks)
🧰 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/models/generation/vllm/vllm_backend.pynemo_rl/models/policy/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/models/generation/vllm/vllm_backend.pynemo_rl/models/policy/utils.py
🪛 Ruff (0.14.3)
nemo_rl/models/generation/vllm/vllm_backend.py
164-167: Within an except clause, raise exceptions with raise ... from err or raise ... from None to distinguish them from errors in exception handling
(B904)
164-167: Avoid specifying long messages outside the exception class
(TRY003)
nemo_rl/models/policy/utils.py
487-491: Within an except clause, raise exceptions with raise ... from err or raise ... from None to distinguish them from errors in exception handling
(B904)
487-491: Avoid specifying long messages outside the exception class
(TRY003)
493-498: Within an except clause, raise exceptions with raise ... from err or raise ... from None to distinguish them from errors in exception handling
(B904)
493-498: Avoid specifying long messages outside the exception class
(TRY003)
⏰ 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 automodel integration comment / Comment on PR
- GitHub Check: Post submodule check comment / Comment on PR
🔇 Additional comments (2)
nemo_rl/models/generation/vllm/vllm_backend.py (1)
15-15: LGTM!The
tracebackimport is necessary for the enhanced error reporting in the exception handler.nemo_rl/models/policy/utils.py (1)
18-18: LGTM!The
tracebackandzmqimports are necessary for the enhanced ZMQ error handling.Also applies to: 23-23
| except Exception as e: | ||
| print( | ||
| f"Error in VllmInternalWorkerExtension.update_weights_via_ipc_zmq: {e}" | ||
| raise RuntimeError( | ||
| f"Error in VllmInternalWorkerExtension.update_weights_via_ipc_zmq: {e}.\n" | ||
| f"{traceback.format_exc()}" | ||
| ) | ||
| return False |
There was a problem hiding this comment.
Remove unreachable code and preserve exception chain.
Two issues need to be addressed:
- Line 168 is unreachable code - the
return Falsestatement can never execute after theraiseon line 164. - The exception chain should be preserved using
raise ... from eto maintain the original exception context for debugging.
Apply this diff to fix both issues:
except Exception as e:
- raise RuntimeError(
- f"Error in VllmInternalWorkerExtension.update_weights_via_ipc_zmq: {e}.\n"
- f"{traceback.format_exc()}"
- )
- return False
+ raise RuntimeError(
+ f"Error in VllmInternalWorkerExtension.update_weights_via_ipc_zmq: {e}.\n"
+ f"{traceback.format_exc()}"
+ ) from e🧰 Tools
🪛 Ruff (0.14.3)
163-163: Do not catch blind exception: Exception
(BLE001)
164-167: Within an except clause, raise exceptions with raise ... from err or raise ... from None to distinguish them from errors in exception handling
(B904)
164-167: Avoid specifying long messages outside the exception class
(TRY003)
🤖 Prompt for AI Agents
In nemo_rl/models/generation/vllm/vllm_backend.py around lines 163 to 168,
remove the unreachable "return False" after the raise and preserve the original
exception chain by re-raising the RuntimeError from the caught exception (use
"raise RuntimeError(...) from e"); ensure you only raise and do not leave any
code after the raise.
| except zmq.Again: | ||
| timeout_ms = zmq_socket.getsockopt(zmq.RCVTIMEO) | ||
| raise TimeoutError( | ||
| f"{worker_name} (rank {rank}): ZMQ communication timeout after {timeout_ms}ms in policy worker side. " | ||
| f"The generation worker may be dead or unresponsive. " | ||
| f"This typically indicates the generation worker has crashed or is not responding to weight streaming." | ||
| ) |
There was a problem hiding this comment.
Preserve exception chain for better debugging.
The exception handling logic is correct, but the exception chain should be preserved using raise ... from to maintain the original exception context.
Apply this diff:
except zmq.Again:
timeout_ms = zmq_socket.getsockopt(zmq.RCVTIMEO)
raise TimeoutError(
f"{worker_name} (rank {rank}): ZMQ communication timeout after {timeout_ms}ms in policy worker side. "
f"The generation worker may be dead or unresponsive. "
f"This typically indicates the generation worker has crashed or is not responding to weight streaming."
- )
+ ) from NoneNote: Using from None here is appropriate because we're converting a low-level ZMQ timeout to a more semantic TimeoutError, and the original zmq.Again exception doesn't add additional context beyond what's already captured in the message.
🧰 Tools
🪛 Ruff (0.14.3)
487-491: Within an except clause, raise exceptions with raise ... from err or raise ... from None to distinguish them from errors in exception handling
(B904)
487-491: Avoid specifying long messages outside the exception class
(TRY003)
🤖 Prompt for AI Agents
nemo_rl/models/policy/utils.py around lines 485-491: the except block catches
zmq.Again and raises a new TimeoutError but does not preserve the exception
chain; modify the raise to use "raise TimeoutError(... ) from None" so the new
semantic TimeoutError replaces the low-level ZMQ exception (per reviewer note
that from None is appropriate).
| except zmq.ZMQError as e: | ||
| raise RuntimeError( | ||
| f"{worker_name} (rank {rank}): ZMQ error during weight streaming: {e} (errno: {e.errno}). " | ||
| f"Error details: {e.strerror}. " | ||
| f"This may indicate network issues or the peer process has terminated unexpectedly.\n" | ||
| f"{traceback.format_exc()}" | ||
| ) |
There was a problem hiding this comment.
Preserve exception chain for better debugging.
The exception handling provides excellent diagnostic information, but the exception chain should be preserved using raise ... from e to maintain the original ZMQ exception context.
Apply this diff:
except zmq.ZMQError as e:
raise RuntimeError(
f"{worker_name} (rank {rank}): ZMQ error during weight streaming: {e} (errno: {e.errno}). "
f"Error details: {e.strerror}. "
f"This may indicate network issues or the peer process has terminated unexpectedly.\n"
f"{traceback.format_exc()}"
- )
+ ) from e📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| except zmq.ZMQError as e: | |
| raise RuntimeError( | |
| f"{worker_name} (rank {rank}): ZMQ error during weight streaming: {e} (errno: {e.errno}). " | |
| f"Error details: {e.strerror}. " | |
| f"This may indicate network issues or the peer process has terminated unexpectedly.\n" | |
| f"{traceback.format_exc()}" | |
| ) | |
| except zmq.ZMQError as e: | |
| raise RuntimeError( | |
| f"{worker_name} (rank {rank}): ZMQ error during weight streaming: {e} (errno: {e.errno}). " | |
| f"Error details: {e.strerror}. " | |
| f"This may indicate network issues or the peer process has terminated unexpectedly.\n" | |
| f"{traceback.format_exc()}" | |
| ) from e |
🧰 Tools
🪛 Ruff (0.14.3)
493-498: Within an except clause, raise exceptions with raise ... from err or raise ... from None to distinguish them from errors in exception handling
(B904)
493-498: Avoid specifying long messages outside the exception class
(TRY003)
🤖 Prompt for AI Agents
In nemo_rl/models/policy/utils.py around lines 492 to 498, the RuntimeError
raised on catching zmq.ZMQError should preserve the original exception chain;
modify the raise to use "raise RuntimeError(... ) from e" so the ZMQError is
attached as the __cause__, keeping the existing detailed message intact.
What does this PR do ?
Add a one line overview of what this PR aims to accomplish.
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