fix: close sockets on bind failure, fix __exit__ traceback in trainers#1032
fix: close sockets on bind failure, fix __exit__ traceback in trainers#1032mango766 wants to merge 1 commit intoinclusionAI:mainfrom
Conversation
- is_port_free(): use try/finally so sockets are closed even when bind() raises OSError - RLTrainer/SFTTrainer __exit__: return False instead of raise exc_value to preserve the original traceback
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request addresses two distinct issues to enhance the robustness and debuggability of the system. It resolves a socket leak by guaranteeing socket closure in network utility functions and improves error reporting in trainer classes by ensuring that original exception tracebacks are preserved during exit handling. Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Changelog
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
There was a problem hiding this comment.
Code Review
This pull request introduces two important fixes. First, it correctly handles socket cleanup in is_port_free by using try...finally blocks, preventing resource leaks on bind failures. Second, it improves exception handling in the __exit__ methods of RLTrainer and SFTTrainer by returning False instead of re-raising the exception, which preserves the original traceback for better debugging. The changes are correct and improve the robustness of the code. I've added one suggestion to further refactor the socket handling logic to be more idiomatic.
| sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) | ||
| try: | ||
| sock.bind(("", port)) | ||
| sock.close() | ||
| except OSError: | ||
| return False | ||
| finally: | ||
| sock.close() | ||
|
|
||
| # Check UDP | ||
| sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) | ||
| try: | ||
| sock.bind(("", port)) | ||
| sock.close() | ||
| return True | ||
| except OSError: | ||
| return False | ||
| finally: | ||
| sock.close() |
There was a problem hiding this comment.
While using try...finally correctly fixes the socket leak, this can be made more idiomatic and concise by using a with statement. The socket object can be used as a context manager, which will automatically handle closing the socket, even if an exception occurs.
| sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) | |
| try: | |
| sock.bind(("", port)) | |
| sock.close() | |
| except OSError: | |
| return False | |
| finally: | |
| sock.close() | |
| # Check UDP | |
| sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) | |
| try: | |
| sock.bind(("", port)) | |
| sock.close() | |
| return True | |
| except OSError: | |
| return False | |
| finally: | |
| sock.close() | |
| try: | |
| with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: | |
| sock.bind(("", port)) | |
| except OSError: | |
| return False | |
| # Check UDP | |
| try: | |
| with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as sock: | |
| sock.bind(("", port)) | |
| return True | |
| except OSError: | |
| return False |
Fixes #1031
Two small fixes:
1. Socket leak in
is_port_free()When
sock.bind()fails, the socket was never closed. Switched totry/finallyso it's always cleaned up.2. Broken traceback in trainer
__exit__RLTrainerandSFTTrainerboth hadraise exc_valuein__exit__, which replaces the original traceback with one pointing at the raise statement. ReturningFalselets Python re-raise naturally with the full traceback.