diff --git a/tests/distributed/omni_connectors/test_shm_ring.py b/tests/distributed/omni_connectors/test_shm_ring.py index c2ecc3dec..e5e5e5938 100644 --- a/tests/distributed/omni_connectors/test_shm_ring.py +++ b/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)) diff --git a/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py b/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py index 9fc3cc42d..10c1c5de4 100644 --- a/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py +++ b/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,