[Bugfix] Claim fallback stage segments before tracking readers

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

tests/distributed/omni_connectors/test_shm_ring.py

@@ -153,6 +153,28 @@ def test_native_scalar_restore_preserves_nested_markers_and_input_containers():
     assert isinstance(decoded["mixed"][2][0], dict)
 
 
+def test_fallback_miss_does_not_register_an_unclaimed_allocation(edge, monkeypatch, caplog):
+    sender, receiver = edge
+    key = "fallback-claim-" + uuid.uuid4().hex
+    _, _, metadata = sender.put("0", "1", key, torch.arange(4096))
+    assert "shm" in metadata
+    original = shared_memory.SharedMemory
+    calls = []
+
+    def tracked_open(*args, **kwargs):
+        calls.append((args, kwargs))
+        return original(*args, **kwargs)
+
+    monkeypatch.setattr(shared_memory, "SharedMemory", tracked_open)
+    with open(f"/dev/shm/shm_{key}_lockfile.lock", "rb+") as lock:
+        fcntl.flock(lock, fcntl.LOCK_EX)
+        assert receiver.get("0", "1", key) is None
+        assert not calls
+    assert torch.equal(receiver.get("0", "1", key)[0], torch.arange(4096))
+    assert receiver.get("0", "1", key, metadata) is None
+    assert not any(record.levelname == "ERROR" for record in caplog.records)
+
+
 def test_tensor_frames_preserve_shared_output_reconstruction(edge):
     sender, receiver = edge
     completion = asdict(CompletionOutput(index=0, text="", token_ids=[1], cumulative_logprob=0.0, logprobs=None))

vllm_omni/distributed/omni_connectors/connectors/shm_connector.py

@@ -272,7 +272,7 @@ class SharedMemoryConnector(OmniConnectorBase):
             obj = self.deserialize_obj(data_bytes)
             result = (obj, int(shm_handle.get("size", 0)))
             return result
-        except BlockingIOError:
+        except (BlockingIOError, FileNotFoundError):
             return None
         except Exception as e:
             logger.error(f"SharedMemoryConnector shm get failed for req : {e}")
@@ -286,13 +286,15 @@ class SharedMemoryConnector(OmniConnectorBase):
 
     def _get_by_key(self, get_key: str) -> tuple[Any, int] | None:
         """Read a SHM segment addressed purely by *get_key*."""
-        shm = None
         try:
-            shm = shm_pkg.SharedMemory(name=get_key)
-            if shm is None or shm.size == 0:
+            # Inspect size without registering a reader before it owns the
+            # delivery lock. Competing readers must not re-register an already
+            # consumed allocation with the application's resource tracker.
+            size = os.stat(f"/dev/shm/{get_key}").st_size
+            if size == 0:
                 return None
             lock_file = f"/dev/shm/shm_{get_key}_lockfile.lock"
-            shm_handle = {"name": get_key, "size": shm.size}
+            shm_handle = {"name": get_key, "size": size}
             result = self._get_data_with_lock(lock_file, shm_handle)
             if result is not None:
                 with self._pending_keys_lock:
@@ -300,20 +302,9 @@ class SharedMemoryConnector(OmniConnectorBase):
             return result
         except FileNotFoundError:
             return None
-        except ValueError as e:
-            # A receiver can observe a newly-created POSIX SHM object before
-            # the writer has finished sizing it. Treat that as "not ready yet"
-            # so async polling can retry without a traceback.
-            if "empty file" in str(e):
-                return None
-            logger.debug("_get_by_key: unexpected error reading SHM segment %s", get_key, exc_info=True)
-            return None
         except Exception:
             logger.debug("_get_by_key: unexpected error reading SHM segment %s", get_key, exc_info=True)
             return None
-        finally:
-            if shm:
-                shm.close()
 
     def get(
         self,