fix(vllm): support IPv6 communicator hosts

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]]

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