diff --git a/tests/test_vllm_client_server.py b/tests/test_vllm_client_server.py index 534cfcf9550..c0b82a602e7 100644 --- a/tests/test_vllm_client_server.py +++ b/tests/test_vllm_client_server.py @@ -14,13 +14,18 @@ import os import subprocess +from unittest.mock import patch from types import SimpleNamespace import pytest from transformers import AutoModelForCausalLM, AutoProcessor, AutoTokenizer from transformers.testing_utils import torch_device -from trl.generation.vllm_client import VLLMClient +from trl.generation.vllm_client import ( + VLLMClient, + _format_http_host, + _resolve_communicator_host, +) from trl.generation.vllm_generation import extract_logprobs from trl.import_utils import is_vllm_available from trl.scripts.vllm_serve import chunk_list @@ -39,6 +44,22 @@ from vllm import LLM, SamplingParams +class TestVLLMClientAddressing(TrlTestCase): + def test_communicator_host_strips_ipv6_brackets(self): + assert _resolve_communicator_host("[2001:db8::1]") == "2001:db8::1" + assert _resolve_communicator_host("2001:db8::1") == "2001:db8::1" + + @patch("trl.generation.vllm_client.socket.gethostbyname", return_value="127.0.0.1") + def test_communicator_host_resolves_hostname(self, gethostbyname): + assert _resolve_communicator_host("localhost") == "127.0.0.1" + gethostbyname.assert_called_once_with("localhost") + + def test_http_host_brackets_only_ipv6_literals(self): + assert _format_http_host("[2001:db8::1]") == "[2001:db8::1]" + assert _format_http_host("2001:db8::1") == "[2001:db8::1]" + assert _format_http_host("127.0.0.1") == "127.0.0.1" + assert _format_http_host("localhost") == "localhost" + class TestChunkList(TrlTestCase): def test_even_split(self): assert chunk_list([1, 2, 3, 4, 5, 6], 2) == [[1, 2, 3], [4, 5, 6]] diff --git a/trl/generation/vllm_client.py b/trl/generation/vllm_client.py index 73f08f451da..ac305a19b41 100644 --- a/trl/generation/vllm_client.py +++ b/trl/generation/vllm_client.py @@ -48,6 +48,31 @@ logger = logging.getLogger(__name__) +def _strip_ipv6_brackets(host: str) -> str: + """Return an IPv6 literal without URL-only brackets.""" + if host.startswith("[") and host.endswith("]"): + return host[1:-1] + return host + + +def _resolve_communicator_host(host: str) -> str: + """Return the TCPStore/NCCL host while preserving legacy hostname resolution.""" + host = _strip_ipv6_brackets(host) + for family in (socket.AF_INET, socket.AF_INET6): + try: + socket.inet_pton(family, host) + return host + except OSError: + pass + return socket.gethostbyname(host) + + +def _format_http_host(host: str) -> str: + """Bracket an IPv6 literal when embedding it in an HTTP URL.""" + host = _strip_ipv6_brackets(host) + return f"[{host}]" if ":" in host else host + + def pil_to_base64(image): buffer = BytesIO() image.save(buffer, format="PNG") @@ -155,13 +180,13 @@ def __init__( if base_url is not None: # Parse the base_url to extract host and port parsed_url = urlparse(base_url) - self.host = socket.gethostbyname(parsed_url.hostname) + self.host = _resolve_communicator_host(parsed_url.hostname) scheme = parsed_url.scheme or "http" self.base_url = f"{scheme}://{parsed_url.netloc}{parsed_url.path}" else: - self.host = host + self.host = _resolve_communicator_host(host) self.server_port = server_port - self.base_url = f"http://{self.host}:{self.server_port}" + self.base_url = f"http://{_format_http_host(self.host)}:{self.server_port}" self.group_port = group_port self.check_server(connection_timeout) # check server and fail after timeout @@ -193,7 +218,9 @@ def check_server(self, total_timeout: float = 0.0, retry_interval: float = 2.0): else: if response.status_code == 200: if "X-Forwarded-For" in response.headers: - self.host = response.headers["X-Forwarded-For"] + self.host = _resolve_communicator_host( + response.headers["X-Forwarded-For"] + ) logger.info("Server is up!") return None