diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index ff467dc65..d5a50cf29 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/tests/engine/test_async_omni_engine_stage_init.py @@ -874,3 +874,54 @@ def test_build_engine_args_dict_uses_diffusion_attention_config_key(): assert "attention_config" not in engine_args_dict assert engine_args_dict["diffusion_attention_config"].default.backend == "FLASH_ATTN" + + +def test_omni_master_server_allocates_globally_unique_route_ports(monkeypatch): + """Regression: two stages must never draw the same ZMQ port. + + ``get_open_ports_list`` only dedups within a single call, so per-route + allocation used to let a later stage reuse an earlier stage's port. The + second engine to ``bind()`` then died with ``zmq.error.ZMQError: Address + already in use`` (flaky multi-stage startup, e.g. Qwen3-Omni thinker/talker/ + code2wav). ``OmniMasterServer`` now dedups every port it hands out. + """ + import itertools + + from vllm_omni.engine import stage_engine_startup as ses + + # A colliding prefix (repeats + the master port 9000) followed by an endless + # fresh stream, so the only way to succeed is to redraw the collisions. + supply = itertools.chain([9000, 9000, 9001, 9001, 9000], itertools.count(9002)) + + def fake_get_open_ports_list(count): + return [next(supply) for _ in range(count)] + + monkeypatch.setattr(ses, "get_open_ports_list", fake_get_open_ports_list) + + server = ses.OmniMasterServer( + master_address="127.0.0.1", + master_port=9000, # seed: a route must not reuse the registration port + stage_ids=[0, 1, 2], + ) + + ports = [] + for sid in (0, 1, 2): + alloc = server.get_allocation(sid) + for addr in ( + alloc.handshake_bind_address, + alloc.input_bind_address, + alloc.output_bind_address, + ): + ports.append(ses._port_from_zmq_address(addr)) + + assert len(ports) == len(set(ports)), f"duplicate route ports allocated: {ports}" + assert 9000 not in ports, "route reused the master registration port" + + +def test_port_from_zmq_address_parsing(): + from vllm_omni.engine.stage_engine_startup import _port_from_zmq_address + + assert _port_from_zmq_address("tcp://127.0.0.1:34277") == 34277 + assert _port_from_zmq_address(None) is None + assert _port_from_zmq_address("ipc:///tmp/sock") is None + assert _port_from_zmq_address("tcp://host:not-a-port") is None diff --git a/vllm_omni/engine/stage_engine_startup.py b/vllm_omni/engine/stage_engine_startup.py index 8f6c199ec..7a55bb93f 100644 --- a/vllm_omni/engine/stage_engine_startup.py +++ b/vllm_omni/engine/stage_engine_startup.py @@ -138,6 +138,21 @@ class StageReplicaResources: # --------------------------------------------------------------------------- +def _port_from_zmq_address(address: str | None) -> int | None: + """Extract the TCP port from a ``tcp://host:port`` ZMQ address, else ``None``. + + Non-TCP transports (``ipc://``, ``inproc://``) and unparsable values return + ``None`` so callers can skip seeding them into the port-dedup set. + """ + if not address: + return None + tail = address.rsplit(":", 1)[-1] + try: + return int(tail) + except ValueError: + return None + + class OmniMasterServer: """Registration server for single-stage engine startup.""" @@ -158,10 +173,25 @@ class OmniMasterServer: self._stage_configs: dict[StageRoute, Any] = {} self._stage_coordinator_addresses: dict[StageRoute, StageCoordinatorAddresses] = {} self._stage_config_events: dict[StageRoute, threading.Event] = {} + # Ports already handed out by *this* server. ``get_open_ports_list`` only + # guarantees uniqueness *within* a single call; a per-route call cannot + # see ports drawn for other routes, so two stages/replicas can draw the + # same ephemeral port and the second engine to ``bind()`` it dies with + # ``zmq.error.ZMQError: Address already in use``. Multi-stage models + # (e.g. Qwen3-Omni: thinker/talker/code2wav) allocate many routes and hit + # this intermittently. We dedup every allocated port against this set, + # seeded with the registration port (and coordinator ROUTER port) that + # are already bound on the same host. + self._allocated_ports: set[int] = set() + if master_port: + self._allocated_ports.add(int(master_port)) # Coordinator ROUTER address echoed back in every registration reply # so OmniCoordClientForStage knows where to connect from inside the # engine subprocess. self._coordinator_router_address = coordinator_router_address + coord_port = _port_from_zmq_address(coordinator_router_address) + if coord_port is not None: + self._allocated_ports.add(coord_port) # Fires only for *newly assigned* (auto-assigned) replicas, not for # head-side pre-allocated slots that already have head-side clients. self._on_register = on_register @@ -235,7 +265,7 @@ class OmniMasterServer: self._stage_config_events[route] = threading.Event() self._stage_coordinator_addresses[route] = StageCoordinatorAddresses() - hs_port, inp_port, out_port = get_open_ports_list(count=3) + hs_port, inp_port, out_port = self._alloc_unique_ports(3) alloc = StageAllocation( handshake_bind_address=f"tcp://{self._address}:{hs_port}", handshake_connect_address=f"tcp://{self._address}:{hs_port}", @@ -247,6 +277,32 @@ class OmniMasterServer: self._stage_routes[route] = alloc return alloc + def _alloc_unique_ports(self, count: int) -> list[int]: + """Return ``count`` open ports unique across every route this server owns. + + ``get_open_ports_list`` dedups only within its own call, so we redraw any + port that collides with one already recorded in ``self._allocated_ports`` + (registration/coordinator ports plus every previously-allocated route + port) and register the winners before returning. Callers already hold the + relevant allocation context (``__init__`` is single-threaded; + registration holds ``self._alloc_lock``). + """ + picked: list[int] = [] + # Bounded retry budget so a pathological host (near port exhaustion) + # fails loudly instead of spinning forever. + for _ in range(64): + for port in get_open_ports_list(count=count - len(picked)): + if port in self._allocated_ports: + continue + self._allocated_ports.add(port) + picked.append(port) + if len(picked) == count: + return picked + raise RuntimeError( + f"[OmniMasterServer] Could not allocate {count} unique open ports " + f"(host={self._address}); {len(self._allocated_ports)} already in use." + ) + def _next_free_replica_id(self, stage_id: int) -> int: """Return the next replica id to assign for an auto-assign registration.