From 546a97ef691f242c899a5e0906d0e75f42694e95 Mon Sep 17 00:00:00 2001
From: Wenwei Zhang <40779233+ZwwWayne@users.noreply.github.com>
Date: Tue, 21 May 2024 01:45:06 +0800
Subject: [PATCH] [Misc]: allow user to specify port in distributed setting
 (#4914)

---
 vllm/envs.py  | 7 +++++++
 vllm/utils.py | 3 +++
 2 files changed, 10 insertions(+)

diff --git a/vllm/envs.py b/vllm/envs.py
index 68d8a074..56ff79e0 100644
--- a/vllm/envs.py
+++ b/vllm/envs.py
@@ -3,6 +3,7 @@ from typing import TYPE_CHECKING, Any, Callable, Dict, Optional
 
 if TYPE_CHECKING:
     VLLM_HOST_IP: str = ""
+    VLLM_PORT: Optional[int] = None
     VLLM_USE_MODELSCOPE: bool = False
     VLLM_INSTANCE_ID: Optional[str] = None
     VLLM_NCCL_SO_PATH: Optional[str] = None
@@ -96,6 +97,12 @@ environment_variables: Dict[str, Callable[[], Any]] = {
     'VLLM_HOST_IP':
     lambda: os.getenv('VLLM_HOST_IP', "") or os.getenv("HOST_IP", ""),
 
+    # used in distributed environment to manually set the communication port
+    # '0' is used to make mypy happy
+    'VLLM_PORT':
+    lambda: int(os.getenv('VLLM_PORT', '0'))
+    if 'VLLM_PORT' in os.environ else None,
+
     # If true, will load models from ModelScope instead of Hugging Face Hub.
     # note that the value is true or false, not numbers
     "VLLM_USE_MODELSCOPE":
diff --git a/vllm/utils.py b/vllm/utils.py
index f0e71f5e..552b43e7 100644
--- a/vllm/utils.py
+++ b/vllm/utils.py
@@ -282,6 +282,9 @@ def get_distributed_init_method(ip: str, port: int) -> str:
 
 
 def get_open_port() -> int:
+    port = envs.VLLM_PORT
+    if port is not None:
+        return port
     # try ipv4
     try:
         with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: