fix(vllm): preserve hostname resolution for communicator

Author: yikun-c

trl/generation/vllm_client.py

@@ -47,16 +47,28 @@
 logger = logging.getLogger(__name__)
 
 
-def _normalize_communicator_host(host: str) -> str:
-    """Return the bare host literal required by TCPStore/NCCL communicators."""
+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 = _normalize_communicator_host(host)
+    host = _strip_ipv6_brackets(host)
     return f"[{host}]" if ":" in host else host
 
 
@@ -167,11 +179,11 @@ def __init__(
         if base_url is not None:
             # Parse the base_url to extract host and port
             parsed_url = urlparse(base_url)
-            self.host = _normalize_communicator_host(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 = _normalize_communicator_host(host)
+            self.host = _resolve_communicator_host(host)
             self.server_port = server_port
             self.base_url = f"http://{_format_http_host(self.host)}:{self.server_port}"
         self.group_port = group_port