fix(engine): dedup ZMQ ports across stages to stop OmniMasterServer bind races

Each (stage_id, replica_id) route allocated its 3 handshake/input/output
ports via a separate get_open_ports_list(count=3) call. That helper only
guarantees uniqueness *within* a single call, so across routes the OS could
hand back a just-closed ephemeral port. Two stages then drew the same port
and the second engine subprocess to bind() it died with
"zmq.error.ZMQError: Address already in use", surfacing as
"Orchestrator initialization failed" / "Stage N replica M exited before API
server became ready". This flaked multi-stage startup (e.g. Qwen3-Omni
thinker/talker/code2wav) — one CI shard passed while a parallel shard failed
on the same commit (build 2630, dev/vllm-align).

OmniMasterServer now tracks every port it hands out in self._allocated_ports
(seeded with the registration port and the coordinator ROUTER port) and a new
_alloc_unique_ports() redraws any collision with a bounded retry budget,
raising loudly on port exhaustion instead of spinning. Adds _port_from_zmq_address()
helper and two CPU-only regression tests.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

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

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.