From b4a2ce93fb4af62df92be007ce11e83b6eb982e1 Mon Sep 17 00:00:00 2001 From: "svc-glamr@nvidia.com" Date: Tue, 3 Nov 2026 08:31:22 +0000 Subject: [PATCH 1/3] fix(multimodal): validate NIXL READ metadata before borrowing embedding buffers NixlReadEmbeddingReceiver.receive_embeddings borrowed a pooled descriptor before checking that the request described a usable NIXL READ. A malformed request (for example WRITE metadata, which passes RdmaMetadata validation but is rejected by begin_read) or a request larger than the pooled buffer removed the descriptor from the pool and left its _data_size changed. Check the READ contract (kind, exactly one remote descriptor, non-zero pointer, matching size, non-empty notification key, decodable agent metadata) before borrowing, as MmKwargsNixlReceiver does since #15482. Return the descriptor with its original size if building the view fails. Behavior once begin_read is called is unchanged. Refs DYN-5224 Signed-off-by: svc-glamr@nvidia.com --- .../common/multimodal/embedding_transfer.py | 42 ++++-- .../multimodal/test_embedding_transfer.py | 128 ++++++++++++++++++ 2 files changed, 162 insertions(+), 8 deletions(-) diff --git a/components/src/dynamo/common/multimodal/embedding_transfer.py b/components/src/dynamo/common/multimodal/embedding_transfer.py index e6768abf9a67..4e209d977626 100644 --- a/components/src/dynamo/common/multimodal/embedding_transfer.py +++ b/components/src/dynamo/common/multimodal/embedding_transfer.py @@ -900,6 +900,24 @@ async def receive_embeddings( readable_metadata = self._nixl_connect.RdmaMetadata.model_validate( request.serialized_request ) + tensor_size_bytes = embeddings_dtype.itemsize * math.prod(embeddings_shape) + # Schema validation alone accepts non-READ kinds, empty descriptors and + # defaults. Check the READ contract before borrowing a pooled buffer. + if ( + readable_metadata.operation_kind + != self._nixl_connect.OperationKind.READ.value + ): + raise ValueError("Embedding transfer metadata must be of kind READ") + remote_desc = readable_metadata.to_descriptors() + if isinstance(remote_desc, list): + raise ValueError("Embedding READ requires exactly one remote descriptor") + if remote_desc.ptr == 0: + raise ValueError("Remote descriptor pointer cannot be zero") + if remote_desc.size != tensor_size_bytes: + raise ValueError("Local and remote descriptors must be the same size") + if not readable_metadata.notification_key: + raise ValueError("READ notification_key must not be empty") + self._nixl_connect.Remote.decode_metadata(readable_metadata.nixl_metadata) original_descriptor_size = None if self.warmedup_descriptors.empty(): @@ -913,14 +931,22 @@ async def receive_embeddings( descriptor = self.warmedup_descriptors.get() # Slide view of pre-allocated tensor original_descriptor_size = descriptor._data_size - tensor_size_bytes = embeddings_dtype.itemsize * math.prod(embeddings_shape) - descriptor._data_size = tensor_size_bytes - assert descriptor._data_ref is not None - encodings_tensor = ( - descriptor._data_ref[:tensor_size_bytes] - .view(dtype=embeddings_dtype) - .view(embeddings_shape) - ) + # No transfer can start before begin_read, so a failure here (for + # example a request larger than the pooled buffer) returns the + # descriptor. Do not extend this past begin_read: NIXL may still + # write into the buffer once a read has been created. + try: + descriptor._data_size = tensor_size_bytes + assert descriptor._data_ref is not None + encodings_tensor = ( + descriptor._data_ref[:tensor_size_bytes] + .view(dtype=embeddings_dtype) + .view(embeddings_shape) + ) + except BaseException: + descriptor._data_size = original_descriptor_size + self.warmedup_descriptors.put(descriptor) + raise dynamic_descriptor = False with _nvtx.annotate("mm:nixl:begin_read", color="pink"): diff --git a/components/src/dynamo/common/tests/multimodal/test_embedding_transfer.py b/components/src/dynamo/common/tests/multimodal/test_embedding_transfer.py index 68d3ce26e08b..23cb2a1f7618 100644 --- a/components/src/dynamo/common/tests/multimodal/test_embedding_transfer.py +++ b/components/src/dynamo/common/tests/multimodal/test_embedding_transfer.py @@ -6,11 +6,14 @@ import asyncio import logging import time +import zlib from random import Random, randint +from unittest.mock import AsyncMock, MagicMock import pytest import torch +from dynamo import nixl_connect from dynamo.common.multimodal.embedding_transfer import ( LocalEmbeddingReceiver, LocalEmbeddingSender, @@ -19,6 +22,7 @@ NixlWriteEmbeddingReceiver, NixlWriteEmbeddingSender, RingBuffer, + TransferRequest, ) # GPU tier is set per-class/per-test below (gpu_0 for local/ring buffer, gpu_1 @@ -153,6 +157,130 @@ async def test_gpu_benchmark(self): await benchmark(sender, receiver, from_cuda=True) +@pytest.mark.gpu_0 # Native NIXL agent is mocked; buffers are CPU tensors +class TestNixlReadEmbeddingReceiverPreparation: + @pytest.fixture + def receiver(self, monkeypatch): + # Real receiver with a one-descriptor pool of registered CPU buffers. + # Only native NIXL agent calls are mocked. + native = MagicMock() + native.register_memory.side_effect = lambda *_: object() + # A read that slips past validation fails instead of polling forever. + native.check_xfer_state.return_value = "ERR" + api = MagicMock() + api.nixl_agent.return_value = native + monkeypatch.setattr(nixl_connect, "nixl_api", api) + receiver = NixlReadEmbeddingReceiver( + embedding_hidden_size=16, max_item_mm_token=1, max_items=1 + ) + descriptor = receiver.warmedup_descriptors.get_nowait() + receiver.warmedup_descriptors.put(descriptor) + receiver.connector.begin_read = AsyncMock(wraps=receiver.connector.begin_read) + try: + yield receiver, descriptor, native + finally: + if descriptor._connection is not None: + descriptor.deregister_with_connector(descriptor._connection) + + @staticmethod + def request(shape=(2, 2), size=16): + return TransferRequest( + embeddings_shape=list(shape), + embedding_dtype_str="torch.float32", + serialized_request={ + "descriptors": [{"ptr": 1, "size": size, "device": "cpu"}], + "operation_kind": nixl_connect.OperationKind.READ.value, + "notification_key": "embedding-0", + "nixl_metadata": zlib.compress(b"mock-agent").hex(), + }, + ) + + @staticmethod + def assert_pool_preserved(receiver, descriptor, native, original_size): + assert receiver.warmedup_descriptors.qsize() == 1 + pooled = receiver.warmedup_descriptors.get_nowait() + receiver.warmedup_descriptors.put(pooled) + assert pooled is descriptor + assert descriptor._data_size == original_size + assert receiver.inuse_descriptors == {} + receiver.connector.begin_read.assert_not_called() + native.initialize_xfer.assert_not_called() + native.add_remote_agent.assert_not_called() + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "invalid_field, invalid_value, error", + [ + ("operation_kind", nixl_connect.OperationKind.WRITE.value, "kind READ"), + ("operation_kind", None, "kind READ"), + ("descriptors", [], "at least one"), + ( + "descriptors", + [{"ptr": 1, "size": 8}, {"ptr": 2, "size": 8}], + "exactly one", + ), + ("descriptors", [{"ptr": 1, "size": 32}], "same size"), + ("descriptors", [{"size": 16}], "cannot be zero"), + ("notification_key", "", "notification_key"), + ("nixl_metadata", "", "cannot be empty"), + ("nixl_metadata", "not-hex", "fromhex"), + ("nixl_metadata", "00", "decompress"), + ], + ids=[ + "write", + "missing-kind", + "empty-descriptors", + "multiple-descriptors", + "size-mismatch", + "missing-pointer", + "empty-notification", + "empty-agent", + "invalid-hex", + "invalid-compression", + ], + ) + async def test_invalid_read_metadata_does_not_borrow_buffer( + self, receiver, invalid_field, invalid_value, error + ): + receiver, descriptor, native = receiver + original_size = descriptor._data_size + request = self.request() + if invalid_value is None: + request.serialized_request.pop(invalid_field) + else: + request.serialized_request[invalid_field] = invalid_value + # These all pass schema validation but cannot start a matching READ. + nixl_connect.RdmaMetadata.model_validate(request.serialized_request) + with pytest.raises((ValueError, zlib.error), match=error): + await receiver.receive_embeddings(request) + self.assert_pool_preserved(receiver, descriptor, native, original_size) + + @pytest.mark.asyncio + async def test_oversized_request_returns_borrowed_buffer(self, receiver): + receiver, descriptor, native = receiver + original_size = descriptor._data_size + # Valid READ metadata, but larger than the 16-byte pooled buffer, so + # building the view fails after the descriptor has been borrowed. + with pytest.raises(RuntimeError): + await receiver.receive_embeddings(self.request(shape=(4, 2), size=32)) + self.assert_pool_preserved(receiver, descriptor, native, original_size) + + @pytest.mark.asyncio + async def test_failed_read_keeps_buffer_out_of_pool(self, receiver): + receiver, descriptor, _ = receiver + # Once a read has been created, NIXL may still write into the buffer, + # so it must not be recycled even if the read fails. + read_op = MagicMock() + read_op.wait_for_completion = AsyncMock(side_effect=RuntimeError("failed")) + receiver.connector.begin_read = AsyncMock(return_value=read_op) + with pytest.raises(RuntimeError, match="failed"): + await receiver.receive_embeddings(self.request()) + receiver.connector.begin_read.assert_awaited_once() + assert receiver.connector.begin_read.await_args.args[1] is descriptor + assert receiver.warmedup_descriptors.qsize() == 0 + assert receiver.inuse_descriptors == {} + + @pytest.mark.gpu_0 # Echo tensor worker is CPU-only (no GPU required) class TestRingBuffer: def test_does_not_wrap_twice_over_live_buffers(self): From 481341ca937e65942334caab7a31509858c7e65b Mon Sep 17 00:00:00 2001 From: "svc-glamr@nvidia.com" Date: Tue, 3 Nov 2026 08:41:47 +0000 Subject: [PATCH 2/3] chore: remove factory working notes Signed-off-by: svc-glamr@nvidia.com --- .../src/dynamo/common/multimodal/embedding_transfer.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/components/src/dynamo/common/multimodal/embedding_transfer.py b/components/src/dynamo/common/multimodal/embedding_transfer.py index 4e209d977626..96830d2e2975 100644 --- a/components/src/dynamo/common/multimodal/embedding_transfer.py +++ b/components/src/dynamo/common/multimodal/embedding_transfer.py @@ -931,10 +931,8 @@ async def receive_embeddings( descriptor = self.warmedup_descriptors.get() # Slide view of pre-allocated tensor original_descriptor_size = descriptor._data_size - # No transfer can start before begin_read, so a failure here (for - # example a request larger than the pooled buffer) returns the - # descriptor. Do not extend this past begin_read: NIXL may still - # write into the buffer once a read has been created. + # Return the descriptor only on failures before begin_read: once a + # read exists, NIXL may still write into the buffer. try: descriptor._data_size = tensor_size_bytes assert descriptor._data_ref is not None From f6ba3a064ecab26355fa2273546bbea5928d5f82 Mon Sep 17 00:00:00 2001 From: "svc-glamr@nvidia.com" Date: Tue, 3 Nov 2026 09:01:44 +0000 Subject: [PATCH 3/3] fix(multimodal): reject empty decoded NIXL agent metadata before borrowing Signed-off-by: svc-glamr@nvidia.com --- .../src/dynamo/common/multimodal/embedding_transfer.py | 5 ++++- .../common/tests/multimodal/test_embedding_transfer.py | 2 ++ 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/components/src/dynamo/common/multimodal/embedding_transfer.py b/components/src/dynamo/common/multimodal/embedding_transfer.py index 96830d2e2975..5e23f2e49195 100644 --- a/components/src/dynamo/common/multimodal/embedding_transfer.py +++ b/components/src/dynamo/common/multimodal/embedding_transfer.py @@ -917,7 +917,10 @@ async def receive_embeddings( raise ValueError("Local and remote descriptors must be the same size") if not readable_metadata.notification_key: raise ValueError("READ notification_key must not be empty") - self._nixl_connect.Remote.decode_metadata(readable_metadata.nixl_metadata) + if not self._nixl_connect.Remote.decode_metadata( + readable_metadata.nixl_metadata + ): + raise ValueError("Decoded NIXL agent metadata cannot be empty") original_descriptor_size = None if self.warmedup_descriptors.empty(): diff --git a/components/src/dynamo/common/tests/multimodal/test_embedding_transfer.py b/components/src/dynamo/common/tests/multimodal/test_embedding_transfer.py index 23cb2a1f7618..a472f5c7b6b3 100644 --- a/components/src/dynamo/common/tests/multimodal/test_embedding_transfer.py +++ b/components/src/dynamo/common/tests/multimodal/test_embedding_transfer.py @@ -225,6 +225,7 @@ def assert_pool_preserved(receiver, descriptor, native, original_size): ("nixl_metadata", "", "cannot be empty"), ("nixl_metadata", "not-hex", "fromhex"), ("nixl_metadata", "00", "decompress"), + ("nixl_metadata", zlib.compress(b"").hex(), "Decoded NIXL agent"), ], ids=[ "write", @@ -237,6 +238,7 @@ def assert_pool_preserved(receiver, descriptor, native, original_size): "empty-agent", "invalid-hex", "invalid-compression", + "empty-decoded-agent", ], ) async def test_invalid_read_metadata_does_not_borrow_buffer(