[Misc]: allow user to specify port in distributed setting (#4914)

This commit is contained in:
Wenwei Zhang 2024-05-21 01:45:06 +08:00 committed by GitHub
parent da5a0b539d
commit 546a97ef69
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 10 additions and 0 deletions

View File

@ -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":

View File

@ -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: