@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 7a6c8c3fa1
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 05:30 UTC

[Chore] Separate out `vllm.utils.network_utils` (#27164)

Signed-off-by: iAmir97 <Amir.balwel@embeddedllm.com> Co-authored-by: iAmir97 <Amir.balwel@embeddedllm.com>

iAmir97 committed Oct 19, 2025 at 17:06 UTC 7a6c8c3fa150957d154ea729490d38dc6f023229
37 files changed +516 -487
examples/offline_inference/data_parallel.py
+1 -1
@@ -33,7 +33,7 @@ import os
33 from time import sleep
34
35 from vllm import LLM, SamplingParams
36 -from vllm.utils import get_open_port
36 +from vllm.utils.network_utils import get_open_port
37
38
39 def parse_args():
examples/offline_inference/rlhf.py
+1 -1
@@ -38,7 +38,7 @@ from rlhf_utils import stateless_init_process_group
38 from transformers import AutoModelForCausalLM
39
40 from vllm import LLM, SamplingParams
41 -from vllm.utils import get_ip, get_open_port
41 +from vllm.utils.network_utils import get_ip, get_open_port
42
43
44 class MyLLM(LLM):
tests/distributed/test_multi_node_assignment.py
+1 -1
@@ -19,7 +19,7 @@ from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
19 from vllm import initialize_ray_cluster
20 from vllm.config import ParallelConfig
21 from vllm.executor.ray_utils import _wait_until_pg_removed
22 -from vllm.utils import get_ip
22 +from vllm.utils.network_utils import get_ip
23
24 VLLM_MULTI_NODE = os.getenv("VLLM_MULTI_NODE", "0") == "1"
25
tests/distributed/test_node_count.py
+1 -1
@@ -7,7 +7,7 @@ import torch.distributed as dist
7
8 from vllm.distributed.parallel_state import _node_count
9 from vllm.distributed.utils import StatelessProcessGroup
10 -from vllm.utils import get_ip, get_open_port
10 +from vllm.utils.network_utils import get_ip, get_open_port
11
12 if __name__ == "__main__":
13 dist.init_process_group(backend="gloo")
tests/distributed/test_same_node.py
+1 -1
@@ -7,7 +7,7 @@ import torch.distributed as dist
7
8 from vllm.distributed.parallel_state import in_the_same_node_as
9 from vllm.distributed.utils import StatelessProcessGroup
10 -from vllm.utils import get_ip, get_open_port
10 +from vllm.utils.network_utils import get_ip, get_open_port
11
12 if __name__ == "__main__":
13 dist.init_process_group(backend="gloo")
tests/distributed/test_shm_broadcast.py
+2 -1
@@ -10,7 +10,8 @@ import torch.distributed as dist
10
11 from vllm.distributed.device_communicators.shm_broadcast import MessageQueue
12 from vllm.distributed.utils import StatelessProcessGroup
13 -from vllm.utils import get_open_port, update_environment_variables
13 +from vllm.utils import update_environment_variables
14 +from vllm.utils.network_utils import get_open_port
15
16
17 def get_arrays(n: int, seed: int = 0) -> list[np.ndarray]:
tests/distributed/test_utils.py
+2 -4
@@ -10,10 +10,8 @@ import torch
10 import vllm.envs as envs
11 from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
12 from vllm.distributed.utils import StatelessProcessGroup
13 -from vllm.utils import (
14 - get_open_port,
15 - update_environment_variables,
16 -)
13 +from vllm.utils import update_environment_variables
14 +from vllm.utils.network_utils import get_open_port
15 from vllm.utils.torch_utils import cuda_device_count_stateless
16
17 from ..utils import multi_gpu_test
tests/entrypoints/openai/test_shutdown.py
+1 -1
@@ -9,7 +9,7 @@ import time
9 import openai
10 import pytest
11
12 -from ...utils import get_open_port
12 +from vllm.utils.network_utils import get_open_port
13
14 MODEL_NAME = "hmellor/tiny-random-LlamaForCausalLM"
15
tests/kernels/moe/modular_kernel_tools/parallel_utils.py
+1 -1
@@ -12,7 +12,7 @@ from typing_extensions import ParamSpec
12
13 from vllm.config import VllmConfig, set_current_vllm_config
14 from vllm.distributed import init_distributed_environment, initialize_model_parallel
15 -from vllm.utils import get_open_port
15 +from vllm.utils.network_utils import get_open_port
16
17 ## Parallel Processes Utils
18
tests/kernels/moe/parallel_utils.py
+2 -1
@@ -15,7 +15,8 @@ from torch.distributed import ProcessGroup
15 from torch.multiprocessing import spawn # pyright: ignore[reportPrivateImportUsage]
16 from typing_extensions import ParamSpec
17
18 -from vllm.utils import get_open_port, has_deep_ep
18 +from vllm.utils import has_deep_ep
19 +from vllm.utils.network_utils import get_open_port
20
21 if has_deep_ep():
22 from vllm.model_executor.layers.fused_moe.deepep_ht_prepare_finalize import (
tests/model_executor/model_loader/tensorizer_loader/conftest.py
+1 -1
@@ -8,7 +8,7 @@ from vllm import LLM, EngineArgs
8 from vllm.distributed import cleanup_dist_env_and_memory
9 from vllm.model_executor.model_loader import tensorizer as tensorizer_mod
10 from vllm.model_executor.model_loader.tensorizer import TensorizerConfig
11 -from vllm.utils import get_distributed_init_method, get_ip, get_open_port
11 +from vllm.utils.network_utils import get_distributed_init_method, get_ip, get_open_port
12 from vllm.v1.executor.abstract import UniProcExecutor
13 from vllm.v1.worker.worker_base import WorkerWrapperBase
14
tests/models/test_vision.py
+2 -1
@@ -19,7 +19,8 @@ from vllm.model_executor.models.vision import (
19 run_dp_sharded_vision_model,
20 )
21 from vllm.platforms import current_platform
22 -from vllm.utils import get_open_port, update_environment_variables
22 +from vllm.utils import update_environment_variables
23 +from vllm.utils.network_utils import get_open_port
24
25 pytestmark = pytest.mark.cpu_test
26
tests/utils.py
+1 -1
@@ -46,9 +46,9 @@ from vllm.platforms import current_platform
46 from vllm.transformers_utils.tokenizer import get_tokenizer
47 from vllm.utils import (
48 FlexibleArgumentParser,
49 - get_open_port,
49 )
50 from vllm.utils.mem_constants import GB_bytes
51 +from vllm.utils.network_utils import get_open_port
52 from vllm.utils.torch_utils import cuda_device_count_stateless
53
54 if current_platform.is_rocm():
tests/utils_/test_network_utils.py new
+126
@@ -0,0 +1,126 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import socket
4 +
5 +import pytest
6 +import zmq
7 +
8 +from vllm.utils.network_utils import (
9 + get_open_port,
10 + get_tcp_uri,
11 + join_host_port,
12 + make_zmq_path,
13 + make_zmq_socket,
14 + split_host_port,
15 + split_zmq_path,
16 +)
17 +
18 +
19 +def test_get_open_port(monkeypatch: pytest.MonkeyPatch):
20 + with monkeypatch.context() as m:
21 + m.setenv("VLLM_PORT", "5678")
22 + # make sure we can get multiple ports, even if the env var is set
23 + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s1:
24 + s1.bind(("localhost", get_open_port()))
25 + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s2:
26 + s2.bind(("localhost", get_open_port()))
27 + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s3:
28 + s3.bind(("localhost", get_open_port()))
29 +
30 +
31 +@pytest.mark.parametrize(
32 + "path,expected",
33 + [
34 + ("ipc://some_path", ("ipc", "some_path", "")),
35 + ("tcp://127.0.0.1:5555", ("tcp", "127.0.0.1", "5555")),
36 + ("tcp://[::1]:5555", ("tcp", "::1", "5555")), # IPv6 address
37 + ("inproc://some_identifier", ("inproc", "some_identifier", "")),
38 + ],
39 +)
40 +def test_split_zmq_path(path, expected):
41 + assert split_zmq_path(path) == expected
42 +
43 +
44 +@pytest.mark.parametrize(
45 + "invalid_path",
46 + [
47 + "invalid_path", # Missing scheme
48 + "tcp://127.0.0.1", # Missing port
49 + "tcp://[::1]", # Missing port for IPv6
50 + "tcp://:5555", # Missing host
51 + ],
52 +)
53 +def test_split_zmq_path_invalid(invalid_path):
54 + with pytest.raises(ValueError):
55 + split_zmq_path(invalid_path)
56 +
57 +
58 +def test_make_zmq_socket_ipv6():
59 + # Check if IPv6 is supported by trying to create an IPv6 socket
60 + try:
61 + sock = socket.socket(socket.AF_INET6, socket.SOCK_STREAM)
62 + sock.close()
63 + except OSError:
64 + pytest.skip("IPv6 is not supported on this system")
65 +
66 + ctx = zmq.Context()
67 + ipv6_path = "tcp://[::]:5555" # IPv6 loopback address
68 + socket_type = zmq.REP # Example socket type
69 +
70 + # Create the socket
71 + zsock: zmq.Socket = make_zmq_socket(ctx, ipv6_path, socket_type)
72 +
73 + # Verify that the IPV6 option is set
74 + assert zsock.getsockopt(zmq.IPV6) == 1, (
75 + "IPV6 option should be enabled for IPv6 addresses"
76 + )
77 +
78 + # Clean up
79 + zsock.close()
80 + ctx.term()
81 +
82 +
83 +def test_make_zmq_path():
84 + assert make_zmq_path("tcp", "127.0.0.1", "5555") == "tcp://127.0.0.1:5555"
85 + assert make_zmq_path("tcp", "::1", "5555") == "tcp://[::1]:5555"
86 +
87 +
88 +def test_get_tcp_uri():
89 + assert get_tcp_uri("127.0.0.1", 5555) == "tcp://127.0.0.1:5555"
90 + assert get_tcp_uri("::1", 5555) == "tcp://[::1]:5555"
91 +
92 +
93 +def test_split_host_port():
94 + # valid ipv4
95 + assert split_host_port("127.0.0.1:5555") == ("127.0.0.1", 5555)
96 + # invalid ipv4
97 + with pytest.raises(ValueError):
98 + # multi colon
99 + assert split_host_port("127.0.0.1::5555")
100 + with pytest.raises(ValueError):
101 + # tailing colon
102 + assert split_host_port("127.0.0.1:5555:")
103 + with pytest.raises(ValueError):
104 + # no colon
105 + assert split_host_port("127.0.0.15555")
106 + with pytest.raises(ValueError):
107 + # none int port
108 + assert split_host_port("127.0.0.1:5555a")
109 +
110 + # valid ipv6
111 + assert split_host_port("[::1]:5555") == ("::1", 5555)
112 + # invalid ipv6
113 + with pytest.raises(ValueError):
114 + # multi colon
115 + assert split_host_port("[::1]::5555")
116 + with pytest.raises(IndexError):
117 + # no colon
118 + assert split_host_port("[::1]5555")
119 + with pytest.raises(ValueError):
120 + # none int port
121 + assert split_host_port("[::1]:5555a")
122 +
123 +
124 +def test_join_host_port():
125 + assert join_host_port("127.0.0.1", 5555) == "127.0.0.1:5555"
126 + assert join_host_port("::1", 5555) == "[::1]:5555"
tests/utils_/test_utils.py
-119
@@ -6,7 +6,6 @@ import hashlib
6 import json
7 import os
8 import pickle
9 -import socket
9 import tempfile
10 from pathlib import Path
11 from unittest.mock import patch
@@ -14,7 +13,6 @@ from unittest.mock import patch
13 import pytest
14 import torch
15 import yaml
17 -import zmq
16 from transformers import AutoTokenizer
17 from vllm_test_utils.monitor import monitor
18
@@ -24,13 +22,6 @@ from vllm.transformers_utils.detokenizer_utils import convert_ids_list_to_tokens
22 from vllm.utils import (
23 FlexibleArgumentParser,
24 bind_kv_cache,
27 - get_open_port,
28 - get_tcp_uri,
29 - join_host_port,
30 - make_zmq_path,
31 - make_zmq_socket,
32 - split_host_port,
33 - split_zmq_path,
25 unique_filepath,
26 )
27 from vllm.utils.hashing import sha256
@@ -43,18 +34,6 @@ from vllm.utils.mem_utils import MemorySnapshot, memory_profiling
34 from ..utils import create_new_process_for_each_test, flat_product
35
36
46 -def test_get_open_port(monkeypatch: pytest.MonkeyPatch):
47 - with monkeypatch.context() as m:
48 - m.setenv("VLLM_PORT", "5678")
49 - # make sure we can get multiple ports, even if the env var is set
50 - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s1:
51 - s1.bind(("localhost", get_open_port()))
52 - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s2:
53 - s2.bind(("localhost", get_open_port()))
54 - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s3:
55 - s3.bind(("localhost", get_open_port()))
56 -
57 -
37 # Tests for FlexibleArgumentParser
38 @pytest.fixture
39 def parser():
@@ -573,104 +552,6 @@ def test_sha256(input: tuple):
552 assert digest != sha256(input + (1,))
553
554
576 -@pytest.mark.parametrize(
577 - "path,expected",
578 - [
579 - ("ipc://some_path", ("ipc", "some_path", "")),
580 - ("tcp://127.0.0.1:5555", ("tcp", "127.0.0.1", "5555")),
581 - ("tcp://[::1]:5555", ("tcp", "::1", "5555")), # IPv6 address
582 - ("inproc://some_identifier", ("inproc", "some_identifier", "")),
583 - ],
584 -)
585 -def test_split_zmq_path(path, expected):
586 - assert split_zmq_path(path) == expected
587 -
588 -
589 -@pytest.mark.parametrize(
590 - "invalid_path",
591 - [
592 - "invalid_path", # Missing scheme
593 - "tcp://127.0.0.1", # Missing port
594 - "tcp://[::1]", # Missing port for IPv6
595 - "tcp://:5555", # Missing host
596 - ],
597 -)
598 -def test_split_zmq_path_invalid(invalid_path):
599 - with pytest.raises(ValueError):
600 - split_zmq_path(invalid_path)
601 -
602 -
603 -def test_make_zmq_socket_ipv6():
604 - # Check if IPv6 is supported by trying to create an IPv6 socket
605 - try:
606 - sock = socket.socket(socket.AF_INET6, socket.SOCK_STREAM)
607 - sock.close()
608 - except socket.error:
609 - pytest.skip("IPv6 is not supported on this system")
610 -
611 - ctx = zmq.Context()
612 - ipv6_path = "tcp://[::]:5555" # IPv6 loopback address
613 - socket_type = zmq.REP # Example socket type
614 -
615 - # Create the socket
616 - zsock: zmq.Socket = make_zmq_socket(ctx, ipv6_path, socket_type)
617 -
618 - # Verify that the IPV6 option is set
619 - assert zsock.getsockopt(zmq.IPV6) == 1, (
620 - "IPV6 option should be enabled for IPv6 addresses"
621 - )
622 -
623 - # Clean up
624 - zsock.close()
625 - ctx.term()
626 -
627 -
628 -def test_make_zmq_path():
629 - assert make_zmq_path("tcp", "127.0.0.1", "5555") == "tcp://127.0.0.1:5555"
630 - assert make_zmq_path("tcp", "::1", "5555") == "tcp://[::1]:5555"
631 -
632 -
633 -def test_get_tcp_uri():
634 - assert get_tcp_uri("127.0.0.1", 5555) == "tcp://127.0.0.1:5555"
635 - assert get_tcp_uri("::1", 5555) == "tcp://[::1]:5555"
636 -
637 -
638 -def test_split_host_port():
639 - # valid ipv4
640 - assert split_host_port("127.0.0.1:5555") == ("127.0.0.1", 5555)
641 - # invalid ipv4
642 - with pytest.raises(ValueError):
643 - # multi colon
644 - assert split_host_port("127.0.0.1::5555")
645 - with pytest.raises(ValueError):
646 - # tailing colon
647 - assert split_host_port("127.0.0.1:5555:")
648 - with pytest.raises(ValueError):
649 - # no colon
650 - assert split_host_port("127.0.0.15555")
651 - with pytest.raises(ValueError):
652 - # none int port
653 - assert split_host_port("127.0.0.1:5555a")
654 -
655 - # valid ipv6
656 - assert split_host_port("[::1]:5555") == ("::1", 5555)
657 - # invalid ipv6
658 - with pytest.raises(ValueError):
659 - # multi colon
660 - assert split_host_port("[::1]::5555")
661 - with pytest.raises(IndexError):
662 - # no colon
663 - assert split_host_port("[::1]5555")
664 - with pytest.raises(ValueError):
665 - # none int port
666 - assert split_host_port("[::1]:5555a")
667 -
668 -
669 -def test_join_host_port():
670 - assert join_host_port("127.0.0.1", 5555) == "127.0.0.1:5555"
671 - assert join_host_port("::1", 5555) == "[::1]:5555"
672 -
673 -
555 def test_convert_ids_list_to_tokens():
556 tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-1.5B-Instruct")
557 token_ids = tokenizer.encode("Hello, world!")
vllm/config/parallel.py
+1 -1
@@ -18,7 +18,7 @@ from vllm.model_executor.layers.batch_invariant import (
18 vllm_is_batch_invariant,
19 )
20 from vllm.platforms import current_platform
21 -from vllm.utils import get_open_ports_list
21 +from vllm.utils.network_utils import get_open_ports_list
22 from vllm.utils.torch_utils import cuda_device_count_stateless
23
24 if TYPE_CHECKING:
vllm/distributed/device_communicators/shm_broadcast.py
+1 -1
@@ -27,7 +27,7 @@ from zmq import ( # type: ignore
27 import vllm.envs as envs
28 from vllm.distributed.utils import StatelessProcessGroup, sched_yield
29 from vllm.logger import init_logger
30 -from vllm.utils import (
30 +from vllm.utils.network_utils import (
31 get_ip,
32 get_open_port,
33 get_open_zmq_ipc_path,
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+1 -1
@@ -40,7 +40,7 @@ from vllm.distributed.utils import divide
40 from vllm.forward_context import ForwardContext
41 from vllm.logger import init_logger
42 from vllm.platforms import current_platform
43 -from vllm.utils import make_zmq_path, make_zmq_socket
43 +from vllm.utils.network_utils import make_zmq_path, make_zmq_socket
44 from vllm.v1.attention.backends.utils import get_kv_cache_layout
45 from vllm.v1.core.sched.output import SchedulerOutput
46
vllm/distributed/kv_transfer/kv_connector/v1/p2p/p2p_nccl_engine.py
+1 -1
@@ -25,7 +25,7 @@ from vllm.distributed.device_communicators.pynccl_wrapper import (
25 from vllm.distributed.kv_transfer.kv_connector.v1.p2p.tensor_memory_pool import ( # noqa: E501
26 TensorMemoryPool,
27 )
28 -from vllm.utils import get_ip
28 +from vllm.utils.network_utils import get_ip
29 from vllm.utils.torch_utils import current_stream
30
31 logger = logging.getLogger(__name__)
vllm/distributed/kv_transfer/kv_pipe/mooncake_pipe.py
+1 -1
@@ -15,7 +15,7 @@ from safetensors.torch import save as safetensors_save
15 from vllm.config.kv_transfer import KVTransferConfig
16 from vllm.distributed.kv_transfer.kv_pipe.base import KVPipeBase
17 from vllm.logger import init_logger
18 -from vllm.utils import join_host_port, make_zmq_path, split_host_port
18 +from vllm.utils.network_utils import join_host_port, make_zmq_path, split_host_port
19
20 logger = init_logger(__name__)
21 NONE_INT = -150886311
vllm/distributed/parallel_state.py
+1 -3
@@ -49,10 +49,8 @@ from vllm.distributed.device_communicators.base_device_communicator import (
49 )
50 from vllm.distributed.utils import StatelessProcessGroup
51 from vllm.logger import init_logger
52 -from vllm.utils import (
53 - get_distributed_init_method,
54 -)
52 from vllm.utils.import_utils import resolve_obj_by_qualname
53 +from vllm.utils.network_utils import get_distributed_init_method
54 from vllm.utils.torch_utils import (
55 direct_register_custom_op,
56 supports_custom_op,
vllm/distributed/utils.py
+1 -1
@@ -29,7 +29,7 @@ from torch.distributed.rendezvous import rendezvous
29
30 import vllm.envs as envs
31 from vllm.logger import init_logger
32 -from vllm.utils import get_tcp_uri
32 +from vllm.utils.network_utils import get_tcp_uri
33 from vllm.utils.torch_utils import is_torch_equal_or_newer
34
35 logger = init_logger(__name__)
vllm/engine/arg_utils.py
+2 -1
@@ -81,8 +81,9 @@ from vllm.transformers_utils.config import (
81 maybe_override_with_speculators,
82 )
83 from vllm.transformers_utils.utils import check_gguf_file
84 -from vllm.utils import FlexibleArgumentParser, get_ip, is_in_ray_actor
84 +from vllm.utils import FlexibleArgumentParser, is_in_ray_actor
85 from vllm.utils.mem_constants import GiB_bytes
86 +from vllm.utils.network_utils import get_ip
87 from vllm.v1.sample.logits_processor import LogitsProcessor
88
89 if TYPE_CHECKING:
vllm/entrypoints/cli/serve.py
+1 -1
@@ -21,9 +21,9 @@ from vllm.usage.usage_lib import UsageContext
21 from vllm.utils import (
22 FlexibleArgumentParser,
23 decorate_logs,
24 - get_tcp_uri,
24 set_process_title,
25 )
26 +from vllm.utils.network_utils import get_tcp_uri
27 from vllm.v1.engine.core import EngineCoreProc
28 from vllm.v1.engine.utils import CoreEngineProcManager, launch_core_engines
29 from vllm.v1.executor.abstract import Executor
vllm/entrypoints/launcher.py
+1 -1
@@ -18,7 +18,7 @@ from vllm.entrypoints.constants import (
18 )
19 from vllm.entrypoints.ssl import SSLCertRefresher
20 from vllm.logger import init_logger
21 -from vllm.utils import find_process_using_port
21 +from vllm.utils.network_utils import find_process_using_port
22 from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
23
24 logger = init_logger(__name__)
vllm/entrypoints/openai/api_server.py
+1 -1
@@ -115,9 +115,9 @@ from vllm.utils import (
115 Device,
116 FlexibleArgumentParser,
117 decorate_logs,
118 - is_valid_ipv6_address,
118 set_ulimit,
119 )
120 +from vllm.utils.network_utils import is_valid_ipv6_address
121 from vllm.v1.engine.exceptions import EngineDeadError
122 from vllm.v1.metrics.prometheus import get_prometheus_registry
123 from vllm.version import __version__ as VLLM_VERSION
vllm/executor/ray_distributed_executor.py
+2 -2
@@ -19,12 +19,12 @@ from vllm.logger import init_logger
19 from vllm.platforms import current_platform
20 from vllm.ray.ray_env import get_env_vars_to_copy
21 from vllm.sequence import ExecuteModelRequest
22 -from vllm.utils import (
22 +from vllm.utils.asyncio import make_async
23 +from vllm.utils.network_utils import (
24 get_distributed_init_method,
25 get_ip,
26 get_open_port,
27 )
27 -from vllm.utils.asyncio import make_async
28 from vllm.v1.outputs import SamplerOutput
29
30 if ray is not None:
vllm/executor/ray_utils.py
+1 -1
@@ -15,7 +15,7 @@ from vllm.executor.msgspec_utils import decode_hook, encode_hook
15 from vllm.logger import init_logger
16 from vllm.platforms import current_platform
17 from vllm.sequence import ExecuteModelRequest, IntermediateTensors
18 -from vllm.utils import get_ip
18 +from vllm.utils.network_utils import get_ip
19 from vllm.v1.outputs import AsyncModelRunnerOutput
20 from vllm.v1.worker.worker_base import WorkerWrapperBase
21
vllm/executor/uniproc_executor.py
+2 -1
@@ -13,7 +13,8 @@ import torch.distributed as dist
13 import vllm.envs as envs
14 from vllm.executor.executor_base import ExecutorBase
15 from vllm.logger import init_logger
16 -from vllm.utils import get_distributed_init_method, get_ip, get_open_port, run_method
16 +from vllm.utils import run_method
17 +from vllm.utils.network_utils import get_distributed_init_method, get_ip, get_open_port
18 from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType
19 from vllm.v1.outputs import AsyncModelRunnerOutput
20 from vllm.v1.worker.worker_base import WorkerWrapperBase
vllm/utils/__init__.py
+12 -322
@@ -7,12 +7,10 @@ import enum
7 import getpass
8 import importlib
9 import inspect
10 -import ipaddress
10 import json
11 import multiprocessing
12 import os
13 import signal
15 -import socket
14 import subprocess
15 import sys
16 import tempfile
@@ -33,15 +31,12 @@ from argparse import (
31 from collections import defaultdict
32 from collections.abc import (
33 Callable,
36 - Iterator,
34 Sequence,
35 )
36 from concurrent.futures.process import ProcessPoolExecutor
37 from functools import cache, partial, wraps
38 from pathlib import Path
39 from typing import TYPE_CHECKING, Any, TextIO, TypeVar
43 -from urllib.parse import urlparse
44 -from uuid import uuid4
40
41 import cloudpickle
42 import psutil
@@ -49,34 +44,36 @@ import regex as re
44 import setproctitle
45 import torch
46 import yaml
52 -import zmq
53 -import zmq.asyncio
47
48 import vllm.envs as envs
49 from vllm.logger import enable_trace_function_call, init_logger
50 from vllm.ray.lazy_utils import is_in_ray_actor
51
59 -_DEPRECATED_PROFILING = {"cprofile", "cprofile_context"}
52 +_DEPRECATED_MAPPINGS = {
53 + "cprofile": "profiling",
54 + "cprofile_context": "profiling",
55 + "get_open_port": "network_utils",
56 +}
57
58
59 def __getattr__(name: str) -> Any: # noqa: D401 - short deprecation docstring
63 - """Module-level getattr to handle deprecated profiling utilities."""
64 - if name in _DEPRECATED_PROFILING:
60 + """Module-level getattr to handle deprecated utilities."""
61 + if name in _DEPRECATED_MAPPINGS:
62 + submodule_name = _DEPRECATED_MAPPINGS[name]
63 warnings.warn(
64 f"vllm.utils.{name} is deprecated and will be removed in a future version. "
67 - f"Use vllm.utils.profiling.{name} instead.",
65 + f"Use vllm.utils.{submodule_name}.{name} instead.",
66 DeprecationWarning,
67 stacklevel=2,
68 )
71 - import vllm.utils.profiling as _prof
72 -
73 - return getattr(_prof, name)
69 + module = __import__(f"vllm.utils.{submodule_name}", fromlist=[submodule_name])
70 + return getattr(module, name)
71 raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
72
73
74 def __dir__() -> list[str]:
75 # expose deprecated names in dir() for better UX/tab-completion
79 - return sorted(list(globals().keys()) + list(_DEPRECATED_PROFILING))
76 + return sorted(list(globals().keys()) + list(_DEPRECATED_MAPPINGS.keys()))
77
78
79 if TYPE_CHECKING:
@@ -149,197 +146,6 @@ def random_uuid() -> str:
146 return str(uuid.uuid4().hex)
147
148
152 -def close_sockets(sockets: Sequence[zmq.Socket | zmq.asyncio.Socket]):
153 - for sock in sockets:
154 - if sock is not None:
155 - sock.close(linger=0)
156 -
157 -
158 -def get_ip() -> str:
159 - host_ip = envs.VLLM_HOST_IP
160 - if "HOST_IP" in os.environ and "VLLM_HOST_IP" not in os.environ:
161 - logger.warning(
162 - "The environment variable HOST_IP is deprecated and ignored, as"
163 - " it is often used by Docker and other software to"
164 - " interact with the container's network stack. Please "
165 - "use VLLM_HOST_IP instead to set the IP address for vLLM processes"
166 - " to communicate with each other."
167 - )
168 - if host_ip:
169 - return host_ip
170 -
171 - # IP is not set, try to get it from the network interface
172 -
173 - # try ipv4
174 - s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
175 - try:
176 - s.connect(("8.8.8.8", 80)) # Doesn't need to be reachable
177 - return s.getsockname()[0]
178 - except Exception:
179 - pass
180 -
181 - # try ipv6
182 - try:
183 - s = socket.socket(socket.AF_INET6, socket.SOCK_DGRAM)
184 - # Google's public DNS server, see
185 - # https://developers.google.com/speed/public-dns/docs/using#addresses
186 - s.connect(("2001:4860:4860::8888", 80)) # Doesn't need to be reachable
187 - return s.getsockname()[0]
188 - except Exception:
189 - pass
190 -
191 - warnings.warn(
192 - "Failed to get the IP address, using 0.0.0.0 by default."
193 - "The value can be set by the environment variable"
194 - " VLLM_HOST_IP or HOST_IP.",
195 - stacklevel=2,
196 - )
197 - return "0.0.0.0"
198 -
199 -
200 -def test_loopback_bind(address, family):
201 - try:
202 - s = socket.socket(family, socket.SOCK_DGRAM)
203 - s.bind((address, 0)) # Port 0 = auto assign
204 - s.close()
205 - return True
206 - except OSError:
207 - return False
208 -
209 -
210 -def get_loopback_ip() -> str:
211 - loopback_ip = envs.VLLM_LOOPBACK_IP
212 - if loopback_ip:
213 - return loopback_ip
214 -
215 - # VLLM_LOOPBACK_IP is not set, try to get it based on network interface
216 -
217 - if test_loopback_bind("127.0.0.1", socket.AF_INET):
218 - return "127.0.0.1"
219 - elif test_loopback_bind("::1", socket.AF_INET6):
220 - return "::1"
221 - else:
222 - raise RuntimeError(
223 - "Neither 127.0.0.1 nor ::1 are bound to a local interface. "
224 - "Set the VLLM_LOOPBACK_IP environment variable explicitly."
225 - )
226 -
227 -
228 -def is_valid_ipv6_address(address: str) -> bool:
229 - try:
230 - ipaddress.IPv6Address(address)
231 - return True
232 - except ValueError:
233 - return False
234 -
235 -
236 -def split_host_port(host_port: str) -> tuple[str, int]:
237 - # ipv6
238 - if host_port.startswith("["):
239 - host, port = host_port.rsplit("]", 1)
240 - host = host[1:]
241 - port = port.split(":")[1]
242 - return host, int(port)
243 - else:
244 - host, port = host_port.split(":")
245 - return host, int(port)
246 -
247 -
248 -def join_host_port(host: str, port: int) -> str:
249 - if is_valid_ipv6_address(host):
250 - return f"[{host}]:{port}"
251 - else:
252 - return f"{host}:{port}"
253 -
254 -
255 -def get_distributed_init_method(ip: str, port: int) -> str:
256 - return get_tcp_uri(ip, port)
257 -
258 -
259 -def get_tcp_uri(ip: str, port: int) -> str:
260 - if is_valid_ipv6_address(ip):
261 - return f"tcp://[{ip}]:{port}"
262 - else:
263 - return f"tcp://{ip}:{port}"
264 -
265 -
266 -def get_open_zmq_ipc_path() -> str:
267 - base_rpc_path = envs.VLLM_RPC_BASE_PATH
268 - return f"ipc://{base_rpc_path}/{uuid4()}"
269 -
270 -
271 -def get_open_zmq_inproc_path() -> str:
272 - return f"inproc://{uuid4()}"
273 -
274 -
275 -def get_open_port() -> int:
276 - """
277 - Get an open port for the vLLM process to listen on.
278 - An edge case to handle, is when we run data parallel,
279 - we need to avoid ports that are potentially used by
280 - the data parallel master process.
281 - Right now we reserve 10 ports for the data parallel master
282 - process. Currently it uses 2 ports.
283 - """
284 - if "VLLM_DP_MASTER_PORT" in os.environ:
285 - dp_master_port = envs.VLLM_DP_MASTER_PORT
286 - reserved_port_range = range(dp_master_port, dp_master_port + 10)
287 - while True:
288 - candidate_port = _get_open_port()
289 - if candidate_port not in reserved_port_range:
290 - return candidate_port
291 - return _get_open_port()
292 -
293 -
294 -def get_open_ports_list(count: int = 5) -> list[int]:
295 - """Get a list of open ports."""
296 - ports = set[int]()
297 - while len(ports) < count:
298 - ports.add(get_open_port())
299 - return list(ports)
300 -
301 -
302 -def _get_open_port() -> int:
303 - port = envs.VLLM_PORT
304 - if port is not None:
305 - while True:
306 - try:
307 - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
308 - s.bind(("", port))
309 - return port
310 - except OSError:
311 - port += 1 # Increment port number if already in use
312 - logger.info("Port %d is already in use, trying port %d", port - 1, port)
313 - # try ipv4
314 - try:
315 - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
316 - s.bind(("", 0))
317 - return s.getsockname()[1]
318 - except OSError:
319 - # try ipv6
320 - with socket.socket(socket.AF_INET6, socket.SOCK_STREAM) as s:
321 - s.bind(("", 0))
322 - return s.getsockname()[1]
323 -
324 -
325 -def find_process_using_port(port: int) -> psutil.Process | None:
326 - # TODO: We can not check for running processes with network
327 - # port on macOS. Therefore, we can not have a full graceful shutdown
328 - # of vLLM. For now, let's not look for processes in this case.
329 - # Ref: https://www.florianreinhard.de/accessdenied-in-psutil/
330 - if sys.platform.startswith("darwin"):
331 - return None
332 -
333 - our_pid = os.getpid()
334 - for conn in psutil.net_connections():
335 - if conn.laddr.port == port and (conn.pid is not None and conn.pid != our_pid):
336 - try:
337 - return psutil.Process(conn.pid)
338 - except psutil.NoSuchProcess:
339 - return None
340 - return None
341 -
342 -
149 def update_environment_variables(envs: dict[str, str]):
150 for k, v in envs.items():
151 if k in os.environ and os.environ[k] != v:
@@ -1119,122 +925,6 @@ def get_exception_traceback():
925 return err_str
926
927
1122 -def split_zmq_path(path: str) -> tuple[str, str, str]:
1123 - """Split a zmq path into its parts."""
1124 - parsed = urlparse(path)
1125 - if not parsed.scheme:
1126 - raise ValueError(f"Invalid zmq path: {path}")
1127 -
1128 - scheme = parsed.scheme
1129 - host = parsed.hostname or ""
1130 - port = str(parsed.port or "")
1131 -
1132 - if scheme == "tcp" and not all((host, port)):
1133 - # The host and port fields are required for tcp
1134 - raise ValueError(f"Invalid zmq path: {path}")
1135 -
1136 - if scheme != "tcp" and port:
1137 - # port only makes sense with tcp
1138 - raise ValueError(f"Invalid zmq path: {path}")
1139 -
1140 - return scheme, host, port
1141 -
1142 -
1143 -def make_zmq_path(scheme: str, host: str, port: int | None = None) -> str:
1144 - """Make a ZMQ path from its parts.
1145 -
1146 - Args:
1147 - scheme: The ZMQ transport scheme (e.g. tcp, ipc, inproc).
1148 - host: The host - can be an IPv4 address, IPv6 address, or hostname.
1149 - port: Optional port number, only used for TCP sockets.
1150 -
1151 - Returns:
1152 - A properly formatted ZMQ path string.
1153 - """
1154 - if port is None:
1155 - return f"{scheme}://{host}"
1156 - if is_valid_ipv6_address(host):
1157 - return f"{scheme}://[{host}]:{port}"
1158 - return f"{scheme}://{host}:{port}"
1159 -
1160 -
1161 -# Adapted from: https://github.com/sgl-project/sglang/blob/v0.4.1/python/sglang/srt/utils.py#L783 # noqa: E501
1162 -def make_zmq_socket(
1163 - ctx: zmq.asyncio.Context | zmq.Context, # type: ignore[name-defined]
1164 - path: str,
1165 - socket_type: Any,
1166 - bind: bool | None = None,
1167 - identity: bytes | None = None,
1168 - linger: int | None = None,
1169 -) -> zmq.Socket | zmq.asyncio.Socket: # type: ignore[name-defined]
1170 - """Make a ZMQ socket with the proper bind/connect semantics."""
1171 -
1172 - mem = psutil.virtual_memory()
1173 - socket = ctx.socket(socket_type)
1174 -
1175 - # Calculate buffer size based on system memory
1176 - total_mem = mem.total / 1024**3
1177 - available_mem = mem.available / 1024**3
1178 - # For systems with substantial memory (>32GB total, >16GB available):
1179 - # - Set a large 0.5GB buffer to improve throughput
1180 - # For systems with less memory:
1181 - # - Use system default (-1) to avoid excessive memory consumption
1182 - buf_size = int(0.5 * 1024**3) if total_mem > 32 and available_mem > 16 else -1
1183 -
1184 - if bind is None:
1185 - bind = socket_type not in (zmq.PUSH, zmq.SUB, zmq.XSUB)
1186 -
1187 - if socket_type in (zmq.PULL, zmq.DEALER, zmq.ROUTER):
1188 - socket.setsockopt(zmq.RCVHWM, 0)
1189 - socket.setsockopt(zmq.RCVBUF, buf_size)
1190 -
1191 - if socket_type in (zmq.PUSH, zmq.DEALER, zmq.ROUTER):
1192 - socket.setsockopt(zmq.SNDHWM, 0)
1193 - socket.setsockopt(zmq.SNDBUF, buf_size)
1194 -
1195 - if identity is not None:
1196 - socket.setsockopt(zmq.IDENTITY, identity)
1197 -
1198 - if linger is not None:
1199 - socket.setsockopt(zmq.LINGER, linger)
1200 -
1201 - if socket_type == zmq.XPUB:
1202 - socket.setsockopt(zmq.XPUB_VERBOSE, True)
1203 -
1204 - # Determine if the path is a TCP socket with an IPv6 address.
1205 - # Enable IPv6 on the zmq socket if so.
1206 - scheme, host, _ = split_zmq_path(path)
1207 - if scheme == "tcp" and is_valid_ipv6_address(host):
1208 - socket.setsockopt(zmq.IPV6, 1)
1209 -
1210 - if bind:
1211 - socket.bind(path)
1212 - else:
1213 - socket.connect(path)
1214 -
1215 - return socket
1216 -
1217 -
1218 -@contextlib.contextmanager
1219 -def zmq_socket_ctx(
1220 - path: str,
1221 - socket_type: Any,
1222 - bind: bool | None = None,
1223 - linger: int = 0,
1224 - identity: bytes | None = None,
1225 -) -> Iterator[zmq.Socket]:
1226 - """Context manager for a ZMQ socket"""
1227 -
1228 - ctx = zmq.Context() # type: ignore[attr-defined]
1229 - try:
1230 - yield make_zmq_socket(ctx, path, socket_type, bind=bind, identity=identity)
1231 - except KeyboardInterrupt:
1232 - logger.debug("Got Keyboard Interrupt.")
1233 -
1234 - finally:
1235 - ctx.destroy(linger=linger)
1236 -
1237 -
928 def _maybe_force_spawn():
929 """Check if we need to force the use of the `spawn` multiprocessing start
930 method.
vllm/utils/network_utils.py new
+331
@@ -0,0 +1,331 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import contextlib
4 +import ipaddress
5 +import os
6 +import socket
7 +import sys
8 +import warnings
9 +from collections.abc import (
10 + Iterator,
11 + Sequence,
12 +)
13 +from typing import Any
14 +from urllib.parse import urlparse
15 +from uuid import uuid4
16 +
17 +import psutil
18 +import zmq
19 +import zmq.asyncio
20 +
21 +import vllm.envs as envs
22 +from vllm.logger import init_logger
23 +
24 +logger = init_logger(__name__)
25 +
26 +
27 +def close_sockets(sockets: Sequence[zmq.Socket | zmq.asyncio.Socket]):
28 + for sock in sockets:
29 + if sock is not None:
30 + sock.close(linger=0)
31 +
32 +
33 +def get_ip() -> str:
34 + host_ip = envs.VLLM_HOST_IP
35 + if "HOST_IP" in os.environ and "VLLM_HOST_IP" not in os.environ:
36 + logger.warning(
37 + "The environment variable HOST_IP is deprecated and ignored, as"
38 + " it is often used by Docker and other software to"
39 + " interact with the container's network stack. Please "
40 + "use VLLM_HOST_IP instead to set the IP address for vLLM processes"
41 + " to communicate with each other."
42 + )
43 + if host_ip:
44 + return host_ip
45 +
46 + # IP is not set, try to get it from the network interface
47 +
48 + # try ipv4
49 + try:
50 + with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:
51 + s.connect(("8.8.8.8", 80)) # Doesn't need to be reachable
52 + return s.getsockname()[0]
53 + except Exception:
54 + pass
55 +
56 + # try ipv6
57 + try:
58 + with socket.socket(socket.AF_INET6, socket.SOCK_DGRAM) as s:
59 + # Google's public DNS server, see
60 + # https://developers.google.com/speed/public-dns/docs/using#addresses
61 + s.connect(("2001:4860:4860::8888", 80)) # Doesn't need to be reachable
62 + return s.getsockname()[0]
63 + except Exception:
64 + pass
65 +
66 + warnings.warn(
67 + "Failed to get the IP address, using 0.0.0.0 by default."
68 + "The value can be set by the environment variable"
69 + " VLLM_HOST_IP or HOST_IP.",
70 + stacklevel=2,
71 + )
72 + return "0.0.0.0"
73 +
74 +
75 +def test_loopback_bind(address, family):
76 + try:
77 + s = socket.socket(family, socket.SOCK_DGRAM)
78 + s.bind((address, 0)) # Port 0 = auto assign
79 + s.close()
80 + return True
81 + except OSError:
82 + return False
83 +
84 +
85 +def get_loopback_ip() -> str:
86 + loopback_ip = envs.VLLM_LOOPBACK_IP
87 + if loopback_ip:
88 + return loopback_ip
89 +
90 + # VLLM_LOOPBACK_IP is not set, try to get it based on network interface
91 +
92 + if test_loopback_bind("127.0.0.1", socket.AF_INET):
93 + return "127.0.0.1"
94 + elif test_loopback_bind("::1", socket.AF_INET6):
95 + return "::1"
96 + else:
97 + raise RuntimeError(
98 + "Neither 127.0.0.1 nor ::1 are bound to a local interface. "
99 + "Set the VLLM_LOOPBACK_IP environment variable explicitly."
100 + )
101 +
102 +
103 +def is_valid_ipv6_address(address: str) -> bool:
104 + try:
105 + ipaddress.IPv6Address(address)
106 + return True
107 + except ValueError:
108 + return False
109 +
110 +
111 +def split_host_port(host_port: str) -> tuple[str, int]:
112 + # ipv6
113 + if host_port.startswith("["):
114 + host, port = host_port.rsplit("]", 1)
115 + host = host[1:]
116 + port = port.split(":")[1]
117 + return host, int(port)
118 + else:
119 + host, port = host_port.split(":")
120 + return host, int(port)
121 +
122 +
123 +def join_host_port(host: str, port: int) -> str:
124 + if is_valid_ipv6_address(host):
125 + return f"[{host}]:{port}"
126 + else:
127 + return f"{host}:{port}"
128 +
129 +
130 +def get_distributed_init_method(ip: str, port: int) -> str:
131 + return get_tcp_uri(ip, port)
132 +
133 +
134 +def get_tcp_uri(ip: str, port: int) -> str:
135 + if is_valid_ipv6_address(ip):
136 + return f"tcp://[{ip}]:{port}"
137 + else:
138 + return f"tcp://{ip}:{port}"
139 +
140 +
141 +def get_open_zmq_ipc_path() -> str:
142 + base_rpc_path = envs.VLLM_RPC_BASE_PATH
143 + return f"ipc://{base_rpc_path}/{uuid4()}"
144 +
145 +
146 +def get_open_zmq_inproc_path() -> str:
147 + return f"inproc://{uuid4()}"
148 +
149 +
150 +def get_open_port() -> int:
151 + """
152 + Get an open port for the vLLM process to listen on.
153 + An edge case to handle, is when we run data parallel,
154 + we need to avoid ports that are potentially used by
155 + the data parallel master process.
156 + Right now we reserve 10 ports for the data parallel master
157 + process. Currently it uses 2 ports.
158 + """
159 + if "VLLM_DP_MASTER_PORT" in os.environ:
160 + dp_master_port = envs.VLLM_DP_MASTER_PORT
161 + reserved_port_range = range(dp_master_port, dp_master_port + 10)
162 + while True:
163 + candidate_port = _get_open_port()
164 + if candidate_port not in reserved_port_range:
165 + return candidate_port
166 + return _get_open_port()
167 +
168 +
169 +def get_open_ports_list(count: int = 5) -> list[int]:
170 + """Get a list of open ports."""
171 + ports = set[int]()
172 + while len(ports) < count:
173 + ports.add(get_open_port())
174 + return list(ports)
175 +
176 +
177 +def _get_open_port() -> int:
178 + port = envs.VLLM_PORT
179 + if port is not None:
180 + while True:
181 + try:
182 + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
183 + s.bind(("", port))
184 + return port
185 + except OSError:
186 + port += 1 # Increment port number if already in use
187 + logger.info("Port %d is already in use, trying port %d", port - 1, port)
188 + # try ipv4
189 + try:
190 + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
191 + s.bind(("", 0))
192 + return s.getsockname()[1]
193 + except OSError:
194 + # try ipv6
195 + with socket.socket(socket.AF_INET6, socket.SOCK_STREAM) as s:
196 + s.bind(("", 0))
197 + return s.getsockname()[1]
198 +
199 +
200 +def find_process_using_port(port: int) -> psutil.Process | None:
201 + # TODO: We can not check for running processes with network
202 + # port on macOS. Therefore, we can not have a full graceful shutdown
203 + # of vLLM. For now, let's not look for processes in this case.
204 + # Ref: https://www.florianreinhard.de/accessdenied-in-psutil/
205 + if sys.platform.startswith("darwin"):
206 + return None
207 +
208 + our_pid = os.getpid()
209 + for conn in psutil.net_connections():
210 + if conn.laddr.port == port and (conn.pid is not None and conn.pid != our_pid):
211 + try:
212 + return psutil.Process(conn.pid)
213 + except psutil.NoSuchProcess:
214 + return None
215 + return None
216 +
217 +
218 +def split_zmq_path(path: str) -> tuple[str, str, str]:
219 + """Split a zmq path into its parts."""
220 + parsed = urlparse(path)
221 + if not parsed.scheme:
222 + raise ValueError(f"Invalid zmq path: {path}")
223 +
224 + scheme = parsed.scheme
225 + host = parsed.hostname or ""
226 + port = str(parsed.port or "")
227 +
228 + if scheme == "tcp" and not all((host, port)):
229 + # The host and port fields are required for tcp
230 + raise ValueError(f"Invalid zmq path: {path}")
231 +
232 + if scheme != "tcp" and port:
233 + # port only makes sense with tcp
234 + raise ValueError(f"Invalid zmq path: {path}")
235 +
236 + return scheme, host, port
237 +
238 +
239 +def make_zmq_path(scheme: str, host: str, port: int | None = None) -> str:
240 + """Make a ZMQ path from its parts.
241 +
242 + Args:
243 + scheme: The ZMQ transport scheme (e.g. tcp, ipc, inproc).
244 + host: The host - can be an IPv4 address, IPv6 address, or hostname.
245 + port: Optional port number, only used for TCP sockets.
246 +
247 + Returns:
248 + A properly formatted ZMQ path string.
249 + """
250 + if port is None:
251 + return f"{scheme}://{host}"
252 + if is_valid_ipv6_address(host):
253 + return f"{scheme}://[{host}]:{port}"
254 + return f"{scheme}://{host}:{port}"
255 +
256 +
257 +# Adapted from: https://github.com/sgl-project/sglang/blob/v0.4.1/python/sglang/srt/utils.py#L783 # noqa: E501
258 +def make_zmq_socket(
259 + ctx: zmq.asyncio.Context | zmq.Context, # type: ignore[name-defined]
260 + path: str,
261 + socket_type: Any,
262 + bind: bool | None = None,
263 + identity: bytes | None = None,
264 + linger: int | None = None,
265 +) -> zmq.Socket | zmq.asyncio.Socket: # type: ignore[name-defined]
266 + """Make a ZMQ socket with the proper bind/connect semantics."""
267 +
268 + mem = psutil.virtual_memory()
269 + socket = ctx.socket(socket_type)
270 +
271 + # Calculate buffer size based on system memory
272 + total_mem = mem.total / 1024**3
273 + available_mem = mem.available / 1024**3
274 + # For systems with substantial memory (>32GB total, >16GB available):
275 + # - Set a large 0.5GB buffer to improve throughput
276 + # For systems with less memory:
277 + # - Use system default (-1) to avoid excessive memory consumption
278 + buf_size = int(0.5 * 1024**3) if total_mem > 32 and available_mem > 16 else -1
279 +
280 + if bind is None:
281 + bind = socket_type not in (zmq.PUSH, zmq.SUB, zmq.XSUB)
282 +
283 + if socket_type in (zmq.PULL, zmq.DEALER, zmq.ROUTER):
284 + socket.setsockopt(zmq.RCVHWM, 0)
285 + socket.setsockopt(zmq.RCVBUF, buf_size)
286 +
287 + if socket_type in (zmq.PUSH, zmq.DEALER, zmq.ROUTER):
288 + socket.setsockopt(zmq.SNDHWM, 0)
289 + socket.setsockopt(zmq.SNDBUF, buf_size)
290 +
291 + if identity is not None:
292 + socket.setsockopt(zmq.IDENTITY, identity)
293 +
294 + if linger is not None:
295 + socket.setsockopt(zmq.LINGER, linger)
296 +
297 + if socket_type == zmq.XPUB:
298 + socket.setsockopt(zmq.XPUB_VERBOSE, True)
299 +
300 + # Determine if the path is a TCP socket with an IPv6 address.
301 + # Enable IPv6 on the zmq socket if so.
302 + scheme, host, _ = split_zmq_path(path)
303 + if scheme == "tcp" and is_valid_ipv6_address(host):
304 + socket.setsockopt(zmq.IPV6, 1)
305 +
306 + if bind:
307 + socket.bind(path)
308 + else:
309 + socket.connect(path)
310 +
311 + return socket
312 +
313 +
314 +@contextlib.contextmanager
315 +def zmq_socket_ctx(
316 + path: str,
317 + socket_type: Any,
318 + bind: bool | None = None,
319 + linger: int = 0,
320 + identity: bytes | None = None,
321 +) -> Iterator[zmq.Socket]:
322 + """Context manager for a ZMQ socket"""
323 +
324 + ctx = zmq.Context() # type: ignore[attr-defined]
325 + try:
326 + yield make_zmq_socket(ctx, path, socket_type, bind=bind, identity=identity)
327 + except KeyboardInterrupt:
328 + logger.debug("Got Keyboard Interrupt.")
329 +
330 + finally:
331 + ctx.destroy(linger=linger)
vllm/v1/engine/coordinator.py
+2 -1
@@ -10,7 +10,8 @@ import zmq
10
11 from vllm.config import ParallelConfig
12 from vllm.logger import init_logger
13 -from vllm.utils import get_mp_context, make_zmq_socket, set_process_title
13 +from vllm.utils import get_mp_context, set_process_title
14 +from vllm.utils.network_utils import make_zmq_socket
15 from vllm.v1.engine import EngineCoreOutputs, EngineCoreRequestType
16 from vllm.v1.serial_utils import MsgpackDecoder
17 from vllm.v1.utils import get_engine_client_zmq_addr, shutdown
vllm/v1/engine/core.py
+1 -1
@@ -30,12 +30,12 @@ from vllm.tasks import POOLING_TASKS, SupportedTask
30 from vllm.transformers_utils.config import maybe_register_config_serialize_by_value
31 from vllm.utils import (
32 decorate_logs,
33 - make_zmq_socket,
33 set_process_title,
34 )
35 from vllm.utils.gc_utils import maybe_attach_gc_debug_callback
36 from vllm.utils.hashing import get_hash_fn_by_name
37 from vllm.utils.import_utils import resolve_obj_by_qualname
38 +from vllm.utils.network_utils import make_zmq_socket
39 from vllm.v1.core.kv_cache_utils import (
40 BlockHash,
41 generate_scheduler_kv_cache_config,
vllm/v1/engine/core_client.py
+2 -2
@@ -23,13 +23,13 @@ from vllm.config import VllmConfig
23 from vllm.logger import init_logger
24 from vllm.lora.request import LoRARequest
25 from vllm.tasks import SupportedTask
26 -from vllm.utils import (
26 +from vllm.utils.asyncio import in_loop
27 +from vllm.utils.network_utils import (
28 close_sockets,
29 get_open_port,
30 get_open_zmq_inproc_path,
31 make_zmq_socket,
32 )
32 -from vllm.utils.asyncio import in_loop
33 from vllm.v1.engine import (
34 EngineCoreOutputs,
35 EngineCoreRequest,
vllm/v1/engine/utils.py
+2 -1
@@ -20,7 +20,8 @@ from vllm.config import CacheConfig, ParallelConfig, VllmConfig
20 from vllm.logger import init_logger
21 from vllm.platforms import current_platform
22 from vllm.ray.ray_env import get_env_vars_to_copy
23 -from vllm.utils import get_mp_context, get_open_zmq_ipc_path, zmq_socket_ctx
23 +from vllm.utils import get_mp_context
24 +from vllm.utils.network_utils import get_open_zmq_ipc_path, zmq_socket_ctx
25 from vllm.v1.engine.coordinator import DPCoordinator
26 from vllm.v1.executor.abstract import Executor
27 from vllm.v1.utils import get_engine_client_zmq_addr, shutdown
vllm/v1/executor/multiproc_executor.py
+4 -2
@@ -38,11 +38,13 @@ from vllm.logger import init_logger
38 from vllm.utils import (
39 _maybe_force_spawn,
40 decorate_logs,
41 + get_mp_context,
42 + set_process_title,
43 +)
44 +from vllm.utils.network_utils import (
45 get_distributed_init_method,
46 get_loopback_ip,
43 - get_mp_context,
47 get_open_port,
45 - set_process_title,
48 )
49 from vllm.v1.core.sched.output import SchedulerOutput
50 from vllm.v1.executor.abstract import Executor, FailureCallback
vllm/v1/utils.py
+2 -6
@@ -25,12 +25,8 @@ from torch.autograd.profiler import record_function
25 import vllm.envs as envs
26 from vllm.logger import init_logger
27 from vllm.usage.usage_lib import UsageContext, is_usage_stats_enabled, usage_message
28 -from vllm.utils import (
29 - get_open_port,
30 - get_open_zmq_ipc_path,
31 - get_tcp_uri,
32 - kill_process_tree,
33 -)
28 +from vllm.utils import kill_process_tree
29 +from vllm.utils.network_utils import get_open_port, get_open_zmq_ipc_path, get_tcp_uri
30
31 if TYPE_CHECKING:
32 import numpy as np