-
Notifications
You must be signed in to change notification settings - Fork 220
[train][Test] Clean up http server setup in unit test #886
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Conversation
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Code Review
This pull request refactors the HTTP server setup logic in the unit tests by extracting a set_up_http_server() helper function. This is a good improvement that reduces code duplication and improves maintainability.
My review includes two main points:
- A potential
NameErrorin thefinallyblocks of the refactored tests. The server thread and port variables are used for cleanup but might not be initialized if an error occurs earlier in thetryblock. I've suggested a pattern to make the cleanup more robust. - An improvement to the port finding logic within the new helper function to use a more idiomatic approach.
| shutdown_server(host=SERVER_HOST, port=server_port, max_wait_seconds=5) | ||
| if server_thread.is_alive(): | ||
| server_thread.join(timeout=5) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The variables server_port and server_thread are defined inside the try block. If an exception occurs before their assignment (e.g., during init_inference_engines), they will not be defined, and this finally block will raise a NameError, preventing proper cleanup. To make the cleanup robust, you should initialize these variables to None before the try block and check for their existence here.
This same issue exists in other refactored tests in this file: test_http_endpoint_completions_routing_and_batching, test_http_endpoint_openai_api_with_weight_sync, test_http_endpoint_with_remote_servers, and test_http_endpoint_error_handling.
if server_port is not None:
shutdown_server(host=SERVER_HOST, port=server_port, max_wait_seconds=5)
if server_thread is not None and server_thread.is_alive():
server_thread.join(timeout=5)| def _find_available_port(host: str, start_port: int, max_attempts: int = 100) -> int: | ||
| """Find an available port starting from start_port.""" | ||
| import socket | ||
|
|
||
| for port in range(start_port, start_port + max_attempts): | ||
| try: | ||
| with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: | ||
| s.bind((host, port)) | ||
| return port | ||
| except OSError: | ||
| continue | ||
| raise RuntimeError(f"Could not find available port in range {start_port}-{start_port + max_attempts}") | ||
|
|
||
| # Find an available port | ||
| server_port = _find_available_port(SERVER_HOST, SERVER_PORT_START) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This implementation for finding an available port can be simplified and made more idiomatic by binding to port 0, which lets the operating system choose an available ephemeral port. This avoids manually iterating through a range of ports.
Note that both the current implementation and the suggested one have a time-of-check-to-time-of-use (TOCTOU) race condition, as the port is released after being found and before the server binds to it. A more robust solution would involve passing the bound socket to the server, but that would require changes to the serve function and is likely out of scope for this PR.
def _find_available_port(host: str) -> int:
"""Find an available port by binding to port 0."""
import socket
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.bind((host, 0))
return s.getsockname()[1]
# Find an available port
server_port = _find_available_port(SERVER_HOST)
Remove redundant code by extracting out
set_up_http_server()inskyrl-train/tests/gpu/gpu_ci/test_inference_engine_client_http_endpoint.py