[Core] Cache host ring discovery without blocking empty polls

Signed-off-by: Sy03 <1370724210@qq.com>

docs/design/feature/stage_host_transport.md

@@ -58,6 +58,9 @@ controlled comparisons or deployments with older receivers. Nonzero capacities
 must be between 4 KiB and 64 MiB. All ring endpoints must support this protocol.
 The launcher supplies a common deployment scope; discovery scans only that
 scope's directory. Receive attempts do not wait for another process's file lock.
+An unchanged empty channel skips its frame lock. Channel registration and
+retirement advance a serialized discovery stamp, so empty polling does not
+rescan the directory and new producers remain discoverable without a timer.
 Ring endpoints must share a Linux PID namespace. Deployments that share SHM
 across different PID namespaces must use `host_ring_bytes: 0`.
 

tests/distributed/omni_connectors/test_shm_ring.py

@@ -263,6 +263,75 @@ def test_flat_scope_and_metadata_work_across_different_parents(monkeypatch):
         receiver.close()
 
 
+def test_empty_poll_reuses_discovery_and_never_waits_for_a_ring_lock(edge, monkeypatch):
+    sender, receiver = edge
+    sender.put("0", "1", "prime-poll", torch.arange(3))
+    assert receiver.get("0", "1", "prime-poll") is not None
+
+    def unexpected_work(*args, **kwargs):
+        pytest.fail("an empty, unchanged channel must not rescan discovery or acquire its frame lock")
+
+    monkeypatch.setattr("vllm_omni.distributed.omni_connectors.connectors.shm_ring.glob.iglob", unexpected_work)
+    monkeypatch.setattr(fcntl, "flock", unexpected_work)
+    for index in range(128):
+        assert receiver.get("0", "1", f"missing-poll-{index}") is None
+    monkeypatch.undo()
+    sender.put("0", "1", "new-poll", torch.arange(5))
+    assert receiver.get("0", "1", "new-poll")[0].tolist() == list(range(5))
+
+
+def test_new_producer_is_discovered_after_a_miss_even_when_the_clock_does_not_advance(edge, monkeypatch):
+    sender, receiver = edge
+    sender.put("0", "1", "first-producer", torch.arange(3))
+    assert receiver.get("0", "1", "first-producer") is not None
+    assert receiver.get("0", "1", "second-producer") is None
+    old_version = receiver._host_ring._discovery_versions["0", "1"]
+    monkeypatch.setattr("vllm_omni.distributed.omni_connectors.connectors.shm_ring.time.time_ns", lambda: 0)
+    with SharedMemoryConnector(sender.config) as second:
+        second.put("0", "1", "second-producer", torch.arange(5))
+        assert receiver.get("0", "1", "second-producer")[0].tolist() == list(range(5))
+        new_version = receiver._host_ring._discovery_versions["0", "1"]
+        assert new_version[2] == old_version[2] + 1
+
+
+def test_one_receiver_discovers_each_input_edge_at_the_same_registry_version(edge):
+    sender, receiver = edge
+    with SharedMemoryConnector({**sender.config, "stage_id": 2}) as second:
+        sender.put("0", "1", "fanin-first", torch.arange(3))
+        second.put("2", "1", "fanin-second", torch.arange(5))
+        assert receiver.get("0", "1", "fanin-first")[0].tolist() == list(range(3))
+        assert receiver.get("2", "1", "fanin-second")[0].tolist() == list(range(5))
+
+
+def test_discovery_attaches_all_producers_even_when_the_first_one_has_the_requested_key(edge, monkeypatch):
+    sender, receiver = edge
+    with SharedMemoryConnector(sender.config) as second:
+        first_metadata = sender.put("0", "1", "multi-first", torch.arange(3))[2]
+        second_metadata = second.put("0", "1", "multi-second", torch.arange(5))[2]
+        paths = [
+            f"/dev/shm/{receiver._host_ring.directory}/{metadata['host_ring']['name']}"
+            for metadata in (first_metadata, second_metadata)
+        ]
+        monkeypatch.setattr(
+            "vllm_omni.distributed.omni_connectors.connectors.shm_ring.glob.iglob", lambda _: iter(paths)
+        )
+        assert receiver.get("0", "1", "multi-first")[0].tolist() == list(range(3))
+        assert receiver.get("0", "1", "multi-second")[0].tolist() == list(range(5))
+
+
+def test_closed_channel_releases_reader_cache_while_the_producer_process_is_still_alive(edge):
+    sender, receiver = edge
+    sender.put("0", "1", "close-cache", torch.arange(3))
+    assert receiver.get("0", "1", "close-cache") is not None
+    channel = next(iter(receiver._host_ring.readers.values()))
+    fd = channel.fd
+    sender.close()
+    assert receiver.get("0", "1", "not-published-after-close") is None
+    assert not receiver._host_ring.readers
+    with pytest.raises(OSError):
+        os.fstat(fd)
+
+
 def test_disabled_ring_and_unsupported_tree_keep_legacy_serializer(edge):
     sender, receiver = edge
     with SharedMemoryConnector({"host_ring_bytes": 0}) as legacy:
@@ -307,7 +376,8 @@ def _locked_writer(name, ready, release):
         os.close(fd)
 
 
-def test_deadline_receive_never_waits_for_another_process_ring_lock(edge):
+@pytest.mark.parametrize("use_metadata", [False, True])
+def test_deadline_receive_never_waits_for_another_process_ring_lock(edge, use_metadata):
     sender, receiver = edge
     _, _, metadata = sender.put("0", "1", "locked", torch.arange(3))
     ctx = multiprocessing.get_context("spawn")
@@ -317,7 +387,10 @@ def test_deadline_receive_never_waits_for_another_process_ring_lock(edge):
     try:
         assert ready.wait(15)
         start = time.monotonic()
-        assert receiver.get_with_deadline("0", "1", "locked", metadata, deadline=start + 1) is None
+        assert (
+            receiver.get_with_deadline("0", "1", "locked", metadata if use_metadata else None, deadline=start + 1)
+            is None
+        )
         assert time.monotonic() - start < 1
     finally:
         release.set()
@@ -326,7 +399,7 @@ def test_deadline_receive_never_waits_for_another_process_ring_lock(edge):
             writer.kill()
             writer.join()
     assert writer.exitcode == 0
-    assert receiver.get("0", "1", "locked", metadata)[0].tolist() == [0, 1, 2]
+    assert receiver.get("0", "1", "locked", metadata if use_metadata else None)[0].tolist() == [0, 1, 2]
 
 
 def _competing_reader(scope, barrier, results):

vllm_omni/distributed/omni_connectors/connectors/shm_ring.py

@@ -18,6 +18,7 @@ import os
 import select
 import struct
 import threading
+import time
 import uuid
 from collections import deque
 from collections.abc import Iterator
@@ -50,6 +51,25 @@ def _process_start(pid: int) -> int:
         return int(file.read().rsplit(")", 1)[1].split()[19])
 
 
+@contextmanager
+def _registry_update(directory: str) -> Iterator[int]:
+    """Publish a strictly changing discovery stamp at channel birth/retirement.
+
+    Kernel directory mtimes can coincide within a clock tick. Serialize marker
+    updates and preserve the previous stamp before the filesystem changes it.
+    Receivers can then cache discovery without a timer or a per-message scan.
+    """
+    fd = os.open(directory, os.O_RDONLY | os.O_DIRECTORY)
+    try:
+        fcntl.flock(fd, fcntl.LOCK_EX)
+        previous = os.fstat(fd).st_mtime_ns
+        yield fd
+        stamp = max(previous + 1, time.time_ns())
+        os.utime(fd, ns=(stamp, stamp))
+    finally:
+        os.close(fd)
+
+
 @dataclass
 class _Channel:
     name: str
@@ -209,7 +229,8 @@ class _Channel:
                 pass
         if self.registry is not None:
             try:
-                os.unlink(self.registry)
+                with _registry_update(os.path.dirname(self.registry)) as directory_fd:
+                    os.unlink(os.path.basename(self.registry), dir_fd=directory_fd)
             except FileNotFoundError:
                 pass
 
@@ -227,6 +248,7 @@ class HostRingTransport:
         self.producers: dict[tuple[str, str], _Channel] = {}
         self._disabled_edges: set[tuple[str, str]] = set()
         self.readers: dict[str, _Channel] = {}
+        self._discovery_versions: dict[tuple[str, str], tuple[int, int, int]] = {}
         self._lock = threading.RLock()
         self._closed = False
 
@@ -254,8 +276,9 @@ class HostRingTransport:
                     channel.registry = f"/dev/shm/{self.directory}/{name}"
                     # Discovery is confined to this deployment directory. The
                     # marker appears only after the allocation is ready.
-                    marker = os.open(channel.registry, os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0o600)
-                    os.close(marker)
+                    with _registry_update(os.path.dirname(channel.registry)) as directory_fd:
+                        marker = os.open(name, os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0o600, dir_fd=directory_fd)
+                        os.close(marker)
                 except OSError as error:
                     if channel is not None:
                         channel.close()
@@ -318,6 +341,12 @@ class HostRingTransport:
         if not channel.producer_alive():
             self.readers.pop(channel.name, None)
             channel.close()
+            self._discovery_versions.clear()
+            return None
+        # This optimistic check only rejects a miss; successful reads still
+        # take the shared lock. An unchanged publication cursor cannot have
+        # introduced a new key since this receiver's last locked scan.
+        if key not in channel.index and struct.unpack_from("<Q", channel.mapping, 16)[0] == channel.cursor:
             return None
         try:
             fcntl.flock(channel.fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
@@ -364,9 +393,22 @@ class HostRingTransport:
                     result = self._receive(channel, key, None)
                     if result is not None:
                         return result
+            directory = f"/dev/shm/{self.directory}"
+            try:
+                status = os.stat(directory)
+            except FileNotFoundError:
+                self._discovery_versions.clear()
+                self._reap_readers()
+                return None
+            version = (status.st_dev, status.st_ino, status.st_mtime_ns)
+            edge = (source, destination)
+            if version == self._discovery_versions.get(edge):
+                return None
             self._reap_readers()
-            # Re-discover only on a miss, retaining first-message delivery and
-            # support for newly started producer processes without a poll timer.
+            # A new producer advances the stamp before publishing any frames,
+            # retaining first-message discovery without scanning every miss.
+            discovered_result = None
+            discovery_complete = True
             for path in glob.iglob(f"/dev/shm/{self.directory}/{edge_prefix}*"):
                 name = os.path.basename(path)
                 if name in self.readers:
@@ -377,16 +419,25 @@ class HostRingTransport:
                     # A published marker whose allocation or producer is gone
                     # cannot become valid again: allocation names are unique.
                     try:
-                        os.unlink(path)
+                        with _registry_update(directory) as directory_fd:
+                            os.unlink(name, dir_fd=directory_fd)
                     except FileNotFoundError:
                         pass
                     continue
                 except (BlockingIOError, ValueError):
+                    # Attachment may race a writer holding the frame lock.
+                    # Retry discovery after it releases, even at this stamp.
+                    discovery_complete = False
                     continue
-                result = self._receive(channel, key, None)
-                if result is not None:
-                    return result
-            return None
+                if discovered_result is None:
+                    discovered_result = self._receive(channel, key, None)
+            # Attach every producer before caching this edge's stamp. Returning
+            # as soon as one frame is found would hide its unvisited peers.
+            if discovery_complete:
+                self._discovery_versions[edge] = version
+            else:
+                self._discovery_versions.pop(edge, None)
+            return discovered_result
 
     def cancel(self, key: str | None = None, prefix: str | None = None) -> int:
         with self._lock:
@@ -424,6 +475,7 @@ class HostRingTransport:
                 channel.close()
             self.producers.clear()
             self.readers.clear()
+            self._discovery_versions.clear()
             try:
                 os.rmdir(f"/dev/shm/{self.directory}")
             except (FileNotFoundError, OSError):