| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import ctypes |
| 5 | import gc |
| 6 | import struct |
| 7 | import threading |
| 8 | import weakref |
| 9 | from concurrent.futures import Future, ThreadPoolExecutor |
| 10 | from dataclasses import replace |
| 11 | from types import SimpleNamespace |
| 12 | from unittest.mock import MagicMock, patch |
| 13 | |
| 14 | import pytest |
| 15 | import torch |
| 16 | |
| 17 | from tests.v1.ec_connector.unit.utils import create_ec_vllm_config |
| 18 | from vllm.config import MultiModalConfig |
| 19 | from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorRole |
| 20 | from vllm.distributed.ec_transfer.ec_connector.mooncake.config import MooncakeECConfig |
| 21 | from vllm.distributed.ec_transfer.ec_connector.mooncake.metadata import ( |
| 22 | ECMooncakeConnectorMetadata, |
| 23 | ECMooncakePushSpec, |
| 24 | ) |
| 25 | from vllm.distributed.ec_transfer.ec_connector.mooncake_ec_connector import ( |
| 26 | ECMooncakeConnector, |
| 27 | ) |
| 28 | from vllm.distributed.ec_transfer.ec_connector.mooncake_store_embedding import ( |
| 29 | backend as store_backend, |
| 30 | ) |
| 31 | from vllm.distributed.ec_transfer.ec_connector.mooncake_store_embedding import ( |
| 32 | store_client, |
| 33 | ) |
| 34 | from vllm.distributed.ec_transfer.ec_connector.mooncake_store_embedding.data import ( |
| 35 | TensorSpec, |
| 36 | build_embedding_namespace, |
| 37 | make_embedding_key, |
| 38 | ) |
| 39 | |
| 40 | pytestmark = pytest.mark.cpu_test |
| 41 | KEY = ( |
| 42 | "test@embedding@model:model@revision:revision" |
| 43 | "@encoder:encoder@dtype:torch.float32@protocol:v3" |
| 44 | ) |
| 45 | SPEC = TensorSpec((2, 2), "torch.float32", 16) |
| 46 | |
| 47 | |
| 48 | def test_embedding_key_combines_namespace_and_identifier(): |
| 49 | config = create_ec_vllm_config(ec_role="ec_producer") |
| 50 | config.ec_transfer_config.ec_connector_extra_config.update( |
| 51 | embedding_cache_prefix="team@cache", embedding_model_identity="shared/model" |
| 52 | ) |
| 53 | config.model_config.revision = "rev:1" |
| 54 | config.model_config.dtype = torch.float32 |
| 55 | config.model_config.multimodal_config = None |
| 56 | assert make_embedding_key(build_embedding_namespace(config), "image@1") == ( |
| 57 | "team@cache@embedding@model:shared/model@revision:rev:1" |
| 58 | "@encoder:encoder:default@dtype:torch.float32@protocol:v3@id:image@1" |
| 59 | ) |
| 60 | |
| 61 | |
| 62 | def test_shared_reuse_requires_content_stable_identifiers(): |
| 63 | """Encoder-only cache-off rendering uses process-local counter identifiers.""" |
| 64 | config = create_ec_vllm_config(ec_role="ec_producer") |
| 65 | config.ec_transfer_config.ec_connector = "ECMooncakeConnector" |
| 66 | config.ec_transfer_config.ec_connector_extra_config["cross_encoder_cache"] = True |
| 67 | config.model_config.multimodal_config = MultiModalConfig(mm_processor_cache_gb=0) |
| 68 | config.use_v2_model_runner = True |
| 69 | config.lora_config = None |
| 70 | with pytest.raises(ValueError, match="mm_processor_cache_gb"): |
| 71 | MooncakeECConfig.from_vllm_config(config) |
| 72 | |
| 73 | |
| 74 | class MemoryStore: |
| 75 | """Only the native boundary is faked; codec, resolver and publisher are real.""" |
| 76 | |
| 77 | def __init__(self): |
| 78 | self.objects: dict[str, bytes] = {} |
| 79 | self.registered: dict[int, int] = {} |
| 80 | |
| 81 | def close(self): |
| 82 | return 0 |
| 83 | |
| 84 | def batch_is_exist(self, keys): |
| 85 | return [int(k in self.objects) for k in keys] |
| 86 | |
| 87 | def register_buffer(self, addr, size): |
| 88 | self.registered[addr] = size |
| 89 | return 0 |
| 90 | |
| 91 | def unregister_buffer(self, addr): |
| 92 | del self.registered[addr] |
| 93 | return 0 |
| 94 | |
| 95 | def batch_put_from_multi_buffers(self, keys, pointers, sizes, config): |
| 96 | [key], [ptrs], [lengths] = keys, pointers, sizes |
| 97 | self.objects[key] = b"".join( |
| 98 | ctypes.string_at(p, n) for p, n in zip(ptrs, lengths, strict=True) |
| 99 | ) |
| 100 | return [0] |
| 101 | |
| 102 | def batch_get_into(self, keys, pointers, sizes): |
| 103 | results = [] |
| 104 | for key, ptr, size in zip(keys, pointers, sizes, strict=True): |
| 105 | assert any( |
| 106 | base <= ptr and ptr + size <= base + capacity |
| 107 | for base, capacity in self.registered.items() |
| 108 | ) |
| 109 | obj = self.objects.get(key) |
| 110 | if obj is None: |
| 111 | results.append(-704) |
| 112 | elif len(obj) > size: |
| 113 | results.append(-600) |
| 114 | else: |
| 115 | ctypes.memmove(ptr, obj, len(obj)) |
| 116 | results.append(len(obj)) |
| 117 | return results |
| 118 | |
| 119 | |
| 120 | @pytest.fixture |
| 121 | def backend(): |
| 122 | instance = store_backend.MooncakeEmbeddingStoreBackend( |
| 123 | store_client.MooncakeEmbeddingStoreClient( |
| 124 | MemoryStore(), read_buffer_bytes=1024 |
| 125 | ), |
| 126 | KEY, |
| 127 | max_pending_items=32, |
| 128 | max_pending_bytes=2 * 1024**3, |
| 129 | ) |
| 130 | yield instance |
| 131 | # Some tests deliberately poison I/O; always stop the Python executor. |
| 132 | instance._executor.shutdown(wait=True) |
| 133 | if not instance.store_client._poisoned and ( |
| 134 | instance.store_client._read_buffer is not None |
| 135 | or instance.store_client._put_header is not None |
| 136 | ): |
| 137 | instance.store_client.close() |
| 138 | |
| 139 | |
| 140 | def test_tensor_round_trip(backend): |
| 141 | client, native = backend.store_client, backend.store_client.store |
| 142 | native.register_buffer = MagicMock(wraps=native.register_buffer) |
| 143 | header_ptr = None |
| 144 | for index, dtype in enumerate((torch.float16, torch.bfloat16, torch.float32)): |
| 145 | tensor = torch.arange(6, dtype=dtype).reshape(index + 1, -1) |
| 146 | key = make_embedding_key(KEY, str(dtype)) |
| 147 | client.put_tensor(key, tensor) |
| 148 | pointer = ctypes.addressof(client._put_header) |
| 149 | header_ptr = header_ptr or pointer |
| 150 | assert pointer == header_ptr |
| 151 | expected = TensorSpec(tuple(tensor.shape), str(dtype), tensor.nbytes) |
| 152 | with torch.device("meta"): |
| 153 | loaded = client.load_tensors({key: expected}, "cpu")[key] |
| 154 | assert client._read_buffer.device.type == "cpu" |
| 155 | assert loaded.dtype == dtype and torch.equal(loaded, tensor) |
| 156 | assert len(native.objects[key]) == 24 + tensor.nbytes |
| 157 | assert ( |
| 158 | sum(c.args[0] == header_ptr for c in native.register_buffer.call_args_list) == 1 |
| 159 | ) |
| 160 | assert len(native.registered) == 2 |
| 161 | |
| 162 | |
| 163 | @pytest.mark.parametrize("capacity", [80, 120]) |
| 164 | def test_batch_reads_reuse_registration_without_aliasing_outputs(backend, capacity): |
| 165 | client, native = backend.store_client, backend.store_client.store |
| 166 | client._read_buffer_bytes = capacity |
| 167 | tensors = {key: torch.full((2, 2), float(i)) for i, key in enumerate("abc")} |
| 168 | for key, tensor in tensors.items(): |
| 169 | client.put_tensor(key, tensor) |
| 170 | native.register_buffer = MagicMock(wraps=native.register_buffer) |
| 171 | native.unregister_buffer = MagicMock(wraps=native.unregister_buffer) |
| 172 | native.batch_get_into = MagicMock(wraps=native.batch_get_into) |
| 173 | expected = dict.fromkeys(tensors, SPEC) |
| 174 | first = client.load_tensors(expected, "cpu") |
| 175 | for key, tensor in tensors.items(): |
| 176 | native.objects[key] = native.objects[key][:24] + struct.pack( |
| 177 | "<4f", *[tensor[0, 0].item() + 10] * 4 |
| 178 | ) |
| 179 | second = client.load_tensors(expected, "cpu") |
| 180 | assert all(torch.equal(first[key], tensor) for key, tensor in tensors.items()) |
| 181 | assert all(torch.equal(second[key], tensor + 10) for key, tensor in tensors.items()) |
| 182 | assert native.batch_get_into.call_count == (4 if capacity == 80 else 2) |
| 183 | native.register_buffer.assert_called_once() |
| 184 | native.unregister_buffer.assert_not_called() |
| 185 | client.close() |
| 186 | assert native.unregister_buffer.call_count == 2 |
| 187 | assert not native.registered |
| 188 | |
| 189 | |
| 190 | @pytest.mark.parametrize("mode", ["miss", "oversized"]) |
| 191 | def test_unloadable_inputs_do_not_allocate_staging(backend, mode): |
| 192 | client, native = backend.store_client, backend.store_client.store |
| 193 | client._read_buffer_bytes = 39 |
| 194 | if mode == "oversized": |
| 195 | client.put_tensor(make_embedding_key(KEY, "a"), torch.ones(2, 2)) |
| 196 | native.register_buffer = MagicMock(wraps=native.register_buffer) |
| 197 | assert not backend.resolve_inputs({"a": SPEC}, {}, "cpu") |
| 198 | assert client._read_buffer is None |
| 199 | if mode == "miss": |
| 200 | assert client._put_header is None |
| 201 | client.close() |
| 202 | native.register_buffer.assert_not_called() |
| 203 | |
| 204 | |
| 205 | @pytest.mark.parametrize( |
| 206 | "results", [[40, -800], [40], [40, True], [[40], [40]], None, [40, 41]] |
| 207 | ) |
| 208 | def test_unsafe_batch_result_retains_whole_staging(backend, results): |
| 209 | client, native = backend.store_client, backend.store_client.store |
| 210 | native.batch_get_into = MagicMock(return_value=results) |
| 211 | native.unregister_buffer = MagicMock(wraps=native.unregister_buffer) |
| 212 | for _ in range(2): |
| 213 | with pytest.raises(store_client.EmbeddingStoreError, match="unconfirmed"): |
| 214 | client.load_tensors(dict.fromkeys("ab", SPEC), "cpu") |
| 215 | native.batch_get_into.assert_called_once() |
| 216 | with pytest.raises(store_client.EmbeddingStoreError, match="unconfirmed"): |
| 217 | client.close() |
| 218 | native.unregister_buffer.assert_not_called() |
| 219 | assert client._read_buffer is client._unsafe_owners[0] |
| 220 | |
| 221 | |
| 222 | @pytest.mark.parametrize("failure", ["allocation", "copy", "record", "completion"]) |
| 223 | def test_copy_failure_preserves_staging_lifetime(backend, failure): |
| 224 | """Emulate CUDA allocation/event boundaries without requiring a GPU.""" |
| 225 | client = backend.store_client |
| 226 | for key in "ab": |
| 227 | client.put_tensor(key, torch.ones(2, 2)) |
| 228 | allocate, copy = torch.empty, torch.Tensor.copy_ |
| 229 | targets: list[torch.Tensor] = [] |
| 230 | copies: list[torch.Tensor] = [] |
| 231 | ready = MagicMock() |
| 232 | if failure == "record": |
| 233 | ready.record.side_effect = RuntimeError("record failure") |
| 234 | elif failure == "completion": |
| 235 | ready.synchronize.side_effect = RuntimeError("copy failure") |
| 236 | |
| 237 | def allocate_on_cpu(*args, **kwargs): |
| 238 | is_target = kwargs.get("device") == torch.device("cuda") |
| 239 | if is_target and targets and failure == "allocation": |
| 240 | raise torch.OutOfMemoryError("allocation failure") |
| 241 | kwargs.update(device="cpu", pin_memory=False) |
| 242 | tensor = allocate(*args, **kwargs) |
| 243 | if is_target: |
| 244 | targets.append(tensor) |
| 245 | return tensor |
| 246 | |
| 247 | def record_copy(target, source, **kwargs): |
| 248 | if copies and failure == "copy": |
| 249 | raise RuntimeError("copy failure") |
| 250 | copies.append(target) |
| 251 | return copy(target, source, **kwargs) |
| 252 | |
| 253 | with ( |
| 254 | patch.object(torch, "empty", side_effect=allocate_on_cpu), |
| 255 | patch.object(torch, "Event", return_value=ready), |
| 256 | patch.object(torch.accelerator, "current_stream"), |
| 257 | patch.object(torch.Tensor, "copy_", new=record_copy), |
| 258 | pytest.raises(RuntimeError, match="allocation failure|unconfirmed"), |
| 259 | ): |
| 260 | client.load_tensors(dict.fromkeys("ab", SPEC), "cuda") |
| 261 | if failure == "allocation": |
| 262 | assert not copies and not client._poisoned |
| 263 | else: |
| 264 | assert len(copies) == (1 if failure == "copy" else 2) |
| 265 | assert client._poisoned |
| 266 | assert any(owner is client._read_buffer for owner in client._unsafe_owners) |
| 267 | assert any(owner is ready for owner in client._unsafe_owners) |
| 268 | assert all( |
| 269 | any(owner is target for owner in client._unsafe_owners) |
| 270 | for target in targets |
| 271 | ) |
| 272 | with pytest.raises(store_client.EmbeddingStoreError, match="unconfirmed"): |
| 273 | client.close() |
| 274 | |
| 275 | |
| 276 | def test_non_matrix_output_is_rejected_before_registration(backend): |
| 277 | client = backend.store_client |
| 278 | client.store.register_buffer = MagicMock() |
| 279 | with pytest.raises(store_client.EmbeddingStoreOperationError, match="2D"): |
| 280 | client.put_tensor(make_embedding_key(KEY, "a"), torch.ones(2, 2, 2)) |
| 281 | client.store.register_buffer.assert_not_called() |
| 282 | |
| 283 | |
| 284 | def make_worker(backend): |
| 285 | from vllm.distributed.ec_transfer.ec_connector.mooncake.worker import ( |
| 286 | ECMooncakeWorker, |
| 287 | ) |
| 288 | |
| 289 | worker = object.__new__(ECMooncakeWorker) |
| 290 | worker.is_producer = True |
| 291 | worker.is_consumer = False |
| 292 | worker._output_store = backend |
| 293 | worker._buffer_device = "cpu" |
| 294 | worker._producer_pushes = MagicMock() |
| 295 | worker._producer_pushes.reserve.return_value = (MagicMock(), True) |
| 296 | worker._producer_pushes.poll.return_value = [] |
| 297 | worker._producer_pushes.cancel_requests.return_value = [] |
| 298 | worker._producer_pushes.pending = False |
| 299 | worker._control_executor = MagicMock() |
| 300 | worker._consumer_memory = MagicMock() |
| 301 | worker._consumer_memory.drain_reclaimed.return_value = set() |
| 302 | worker._completed_loads = set() |
| 303 | worker._failed_loads = set() |
| 304 | worker._failed_saves = set() |
| 305 | worker._push_ready = threading.Event() |
| 306 | worker._flush_pending_pushes = MagicMock() |
| 307 | worker._bind_push_source = MagicMock() |
| 308 | connector = object.__new__(ECMooncakeConnector) |
| 309 | connector._worker = worker |
| 310 | connector._role = ECConnectorRole.WORKER |
| 311 | connector._is_producer = True |
| 312 | connector._is_consumer = False |
| 313 | connector._connector_metadata = None |
| 314 | return connector |
| 315 | |
| 316 | |
| 317 | @pytest.mark.parametrize( |
| 318 | "mode", |
| 319 | [ |
| 320 | "miss", |
| 321 | "legacy", |
| 322 | "hit", |
| 323 | "partial", |
| 324 | "get-failure", |
| 325 | "short-data", |
| 326 | "lease-expired", |
| 327 | "registration-failure", |
| 328 | ], |
| 329 | ) |
| 330 | def test_reuse_pipeline(backend, mode): |
| 331 | from vllm.multimodal.inputs import ( |
| 332 | MultiModalFeatureSpec, |
| 333 | MultiModalKwargsItem, |
| 334 | PlaceholderRange, |
| 335 | ) |
| 336 | from vllm.v1.worker.gpu.ec_connector import ActiveECConnector |
| 337 | from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache |
| 338 | from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner |
| 339 | from vllm.v1.worker.gpu.model_states.interface import ModelState |
| 340 | |
| 341 | native = backend.store_client.store |
| 342 | values = {"a": 1.0, "b": 2.0, "c": 3.0} |
| 343 | if mode == "legacy": |
| 344 | old_namespace = KEY.replace("@protocol:v3", "@protocol:v2") |
| 345 | for key in "abc": |
| 346 | native.objects[make_embedding_key(old_namespace, key)] = b"old format" |
| 347 | if mode == "partial": |
| 348 | hits = {"a", "c"} |
| 349 | elif mode not in ("miss", "legacy"): |
| 350 | hits = set("abc") |
| 351 | else: |
| 352 | hits = set() |
| 353 | for key in hits: |
| 354 | backend.store_client.put_tensor( |
| 355 | make_embedding_key(KEY, key), torch.full((2, 2), values[key]) |
| 356 | ) |
| 357 | if mode == "short-data": |
| 358 | key = make_embedding_key(KEY, "b") |
| 359 | native.objects[key] = native.objects[key][:-1] |
| 360 | hits.remove("b") |
| 361 | if mode in ("get-failure", "lease-expired"): |
| 362 | read = native.batch_get_into |
| 363 | |
| 364 | def failed_read(keys, pointers, sizes): |
| 365 | results = read(keys, pointers, sizes) |
| 366 | results[keys.index(make_embedding_key(KEY, "b"))] = ( |
| 367 | -707 if mode == "lease-expired" else -704 |
| 368 | ) |
| 369 | return results |
| 370 | |
| 371 | native.batch_get_into = failed_read |
| 372 | hits.remove("b") |
| 373 | if mode == "registration-failure": |
| 374 | native.register_buffer = MagicMock(return_value=-500) |
| 375 | hits.clear() |
| 376 | connector = make_worker(backend) |
| 377 | metadata = ECMooncakeConnectorMetadata( |
| 378 | store_candidates=dict.fromkeys("abc", SPEC), |
| 379 | pushes=[ |
| 380 | ECMooncakePushSpec( |
| 381 | mm_hash=k, |
| 382 | shape=(2, 2), |
| 383 | dtype="float32", |
| 384 | nbytes=16, |
| 385 | consumer_zmq="ipc:///tmp/test", |
| 386 | transfer_id=k, |
| 387 | request_id="req", |
| 388 | ) |
| 389 | for k in "abc" |
| 390 | ], |
| 391 | ) |
| 392 | cache = EncoderCache() |
| 393 | cache.mm_features["req"] = [ |
| 394 | MultiModalFeatureSpec( |
| 395 | data=MultiModalKwargsItem({}), |
| 396 | modality="image", |
| 397 | identifier=k, |
| 398 | mm_position=PlaceholderRange(offset=i * 2, length=2), |
| 399 | ) |
| 400 | for i, k in enumerate("abc") |
| 401 | ] |
| 402 | runner = EncoderRunner( |
| 403 | model=None, |
| 404 | max_num_tokens=8, |
| 405 | hidden_size=2, |
| 406 | encoder_cache=cache, |
| 407 | dtype=torch.float32, |
| 408 | device=torch.device("cpu"), |
| 409 | ) |
| 410 | runner.enable_timing = False |
| 411 | missing = [k for k in "abc" if k not in hits] |
| 412 | runner.execute_mm_encoder = MagicMock( |
| 413 | return_value=[torch.full((2, 2), values[k] + 10) for k in missing] |
| 414 | ) |
| 415 | state = SimpleNamespace(encoder_runner=runner, encoder_cache=cache) |
| 416 | with patch( |
| 417 | "vllm.v1.worker.gpu.ec_connector.get_ec_transfer", return_value=connector |
| 418 | ): |
| 419 | adapter = ActiveECConnector(SimpleNamespace(), cache.encoder_outputs) |
| 420 | with adapter.maybe_get_output( |
| 421 | SimpleNamespace(ec_connector_metadata=metadata, finished_req_ids=set()) |
| 422 | ) as output: |
| 423 | assert runner.prepare_mm_inputs({"req": [0, 1, 2]})[0] == missing |
| 424 | ModelState.execute_mm_encoder(state, {"req": [0, 1, 2]}) |
| 425 | assert runner.execute_mm_encoder.call_count == bool(missing) |
| 426 | assert output.ec_connector_worker_meta.pending_saves is False |
| 427 | assert all( |
| 428 | torch.equal( |
| 429 | cache.encoder_outputs[k], |
| 430 | torch.full((2, 2), values[k] + (0 if k in hits else 10)), |
| 431 | ) |
| 432 | for k in "abc" |
| 433 | ) |
| 434 | assert { |
| 435 | c.args[1] for c in connector._worker._bind_push_source.call_args_list |
| 436 | } == set("abc") |
| 437 | backend.shutdown() |
| 438 | assert not native.registered |
| 439 | if mode in ("miss", "legacy"): |
| 440 | other = store_backend.MooncakeEmbeddingStoreBackend( |
| 441 | store_client.MooncakeEmbeddingStoreClient(native, read_buffer_bytes=1024), |
| 442 | KEY, |
| 443 | max_pending_items=32, |
| 444 | max_pending_bytes=2 * 1024**3, |
| 445 | ) |
| 446 | try: |
| 447 | loaded: dict[str, torch.Tensor] = {} |
| 448 | assert other.resolve_inputs( |
| 449 | dict.fromkeys("abc", SPEC), |
| 450 | loaded, |
| 451 | "cpu", |
| 452 | ) == set("abc") |
| 453 | assert all(torch.equal(loaded[k], cache.encoder_outputs[k]) for k in "abc") |
| 454 | finally: |
| 455 | other.shutdown() |
| 456 | |
| 457 | |
| 458 | @pytest.mark.parametrize( |
| 459 | "mismatch", |
| 460 | ["revision", "shape", "dtype", "magic", "short-header", "short-data", "long-data"], |
| 461 | ) |
| 462 | def test_incompatible_output_is_a_miss(backend, mismatch): |
| 463 | key = make_embedding_key(KEY, "a") |
| 464 | dtype = torch.float16 if mismatch == "dtype" else torch.float32 |
| 465 | backend.store_client.put_tensor(key, torch.ones(2, 2, dtype=dtype)) |
| 466 | if mismatch == "revision": |
| 467 | backend.namespace = KEY.replace("@revision:revision", "@revision:other") |
| 468 | objects = backend.store_client.store.objects |
| 469 | if mismatch == "magic": |
| 470 | objects[key] = b"\0" * 4 + objects[key][4:] |
| 471 | elif mismatch == "short-header": |
| 472 | objects[key] = objects[key][:23] |
| 473 | elif mismatch == "short-data": |
| 474 | objects[key] = objects[key][:-1] |
| 475 | elif mismatch == "long-data": |
| 476 | objects[key] += b"extra" |
| 477 | expected = replace(SPEC, shape=(4, 1)) if mismatch == "shape" else SPEC |
| 478 | outputs: dict[str, torch.Tensor] = {} |
| 479 | assert not backend.resolve_inputs({"a": expected}, outputs, "cpu") |
| 480 | assert not outputs |
| 481 | tensor = torch.ones(2, 2) |
| 482 | backend.store_client.put_tensor( |
| 483 | make_embedding_key(backend.namespace, "valid"), tensor |
| 484 | ) |
| 485 | assert backend.resolve_inputs({"valid": SPEC}, outputs, "cpu") == {"valid"} |
| 486 | assert torch.equal(outputs["valid"], tensor) |
| 487 | assert len(backend.store_client.store.registered) == 2 |
| 488 | |
| 489 | |
| 490 | def test_partial_registration_rejection_releases_publication(backend): |
| 491 | """A rejected header registration must not prevent a later publication.""" |
| 492 | native = backend.store_client.store |
| 493 | register = native.register_buffer |
| 494 | put = native.batch_put_from_multi_buffers |
| 495 | native.register_buffer = lambda addr, size: ( |
| 496 | -500 if size == 24 else register(addr, size) |
| 497 | ) |
| 498 | native.batch_put_from_multi_buffers = MagicMock() |
| 499 | tensor = torch.ones(2, 2) |
| 500 | owner = weakref.ref(tensor) |
| 501 | assert backend._enqueue_save("a", tensor) |
| 502 | storage_owner = weakref.ref(tensor.untyped_storage()) |
| 503 | future = backend._pending["a"][0] |
| 504 | with pytest.raises(store_client.EmbeddingStoreOperationError, match="register"): |
| 505 | future.result(timeout=5) |
| 506 | del tensor |
| 507 | gc.collect() |
| 508 | assert owner() is None and not native.registered |
| 509 | native.batch_put_from_multi_buffers.assert_not_called() |
| 510 | backend.resolve_inputs({}, {}, "cpu") |
| 511 | assert backend._pending_bytes == 0 |
| 512 | gc.collect() |
| 513 | assert storage_owner() is None # The failed Future is still alive. |
| 514 | assert backend.store_client._put_header is None |
| 515 | native.register_buffer = register |
| 516 | native.batch_put_from_multi_buffers = put |
| 517 | backend.store_client.put_tensor("retry", torch.ones(2, 2)) |
| 518 | assert len(native.registered) == 1 |
| 519 | |
| 520 | |
| 521 | @pytest.mark.parametrize( |
| 522 | "operation", |
| 523 | ["read-unregister", "header-unregister", "put", "rejected-put", "unregister"], |
| 524 | ) |
| 525 | def test_native_io_owners(backend, operation): |
| 526 | """Uncertain I/O retains owners; completed rejection frees them even via Future.""" |
| 527 | native, client = backend.store_client.store, backend.store_client |
| 528 | put = native.batch_put_from_multi_buffers |
| 529 | backend.max_pending_bytes = SPEC.nbytes |
| 530 | tensor = torch.ones(2, 2) |
| 531 | owner = weakref.ref(tensor) |
| 532 | rejected = operation == "rejected-put" |
| 533 | if operation in ("read-unregister", "header-unregister"): |
| 534 | key = make_embedding_key(KEY, "a") |
| 535 | client.put_tensor(key, tensor) |
| 536 | if operation == "read-unregister": |
| 537 | client.load_tensors({key: SPEC}, "cpu") |
| 538 | native.unregister_buffer = MagicMock(return_value=-500) |
| 539 | with pytest.raises(store_client.EmbeddingStoreError, match="unregister"): |
| 540 | client.close() |
| 541 | owner = weakref.ref( |
| 542 | client._read_buffer |
| 543 | if operation == "read-unregister" |
| 544 | else client._put_header |
| 545 | ) |
| 546 | else: |
| 547 | if operation == "unregister": |
| 548 | native.unregister_buffer = MagicMock(return_value=-500) |
| 549 | else: |
| 550 | native.batch_put_from_multi_buffers = MagicMock( |
| 551 | return_value=[-1700 if rejected else -800] |
| 552 | ) |
| 553 | assert backend._enqueue_save("a", tensor) |
| 554 | future = backend._pending["a"][0] |
| 555 | with pytest.raises(store_client.EmbeddingStoreError): |
| 556 | future.result(timeout=5) |
| 557 | del tensor |
| 558 | gc.collect() |
| 559 | if rejected: |
| 560 | assert owner() is None # A live Future must not retain the tensor. |
| 561 | tensor = torch.ones(4, 4) # Too large for admission after reaping. |
| 562 | caller_owner = weakref.ref(tensor) |
| 563 | backend._step_candidates = {"next"} |
| 564 | backend.save_output("next", tensor) |
| 565 | del tensor |
| 566 | gc.collect() |
| 567 | assert caller_owner() is None |
| 568 | assert backend._pending_bytes == 0 and len(native.registered) == 1 |
| 569 | header = client._put_header |
| 570 | native.batch_put_from_multi_buffers = put |
| 571 | client.put_tensor("retry", torch.ones(1, 4)) |
| 572 | assert client._put_header is header |
| 573 | else: |
| 574 | assert owner() is not None and native.registered |
| 575 | header = bytes(client._put_header) |
| 576 | with pytest.raises(store_client.EmbeddingStoreError, match="unconfirmed"): |
| 577 | client.put_tensor("retry", torch.ones(1, 4)) |
| 578 | assert bytes(client._put_header) == header |
| 579 | with pytest.raises(store_client.EmbeddingStoreError, match="unconfirmed"): |
| 580 | client.close() |
| 581 | |
| 582 | |
| 583 | def test_publication_budget_charges_and_releases_backing_storage(backend): |
| 584 | tensor = torch.arange(16, dtype=torch.float32)[:4].view(2, 2) |
| 585 | owner = weakref.ref(tensor) |
| 586 | future: Future[None] = Future() |
| 587 | backend.max_pending_items = 1 |
| 588 | backend.max_pending_bytes = tensor.nbytes |
| 589 | with patch.object(backend._executor, "submit", return_value=future): |
| 590 | assert not backend._enqueue_save("a", tensor) |
| 591 | backend.max_pending_bytes = 64 |
| 592 | assert backend._enqueue_save("a", tensor) |
| 593 | storage_owner = weakref.ref(tensor.untyped_storage()) |
| 594 | assert not backend._enqueue_save("b", tensor) |
| 595 | del tensor |
| 596 | gc.collect() |
| 597 | backend.reap() # The final step-end poll precedes publication completion. |
| 598 | assert owner() is not None |
| 599 | backend._save(make_embedding_key(KEY, "a"), backend._pending["a"][1]) |
| 600 | future.set_result(None) |
| 601 | assert owner() is None and storage_owner() is not None |
| 602 | |
| 603 | def load_after_reaping(*args): |
| 604 | assert not backend._pending and backend._pending_bytes == 0 |
| 605 | assert storage_owner() is None |
| 606 | return {} |
| 607 | |
| 608 | with patch.object(backend.store_client, "load_tensors", load_after_reaping): |
| 609 | assert not backend.resolve_inputs({"b": SPEC}, {}, "cpu") |
| 610 | |
| 611 | |
| 612 | @pytest.mark.parametrize("rejected", [False, True]) |
| 613 | def test_shared_storage_budget_is_released_after_last_publication(backend, rejected): |
| 614 | batch = torch.arange(16, dtype=torch.float32).view(4, 4) |
| 615 | first: Future[None] = Future() |
| 616 | last: Future[None] = Future() |
| 617 | separate: Future[None] = Future() |
| 618 | backend.max_pending_items = 2 |
| 619 | backend.max_pending_bytes = batch.nbytes |
| 620 | with patch.object(backend._executor, "submit", side_effect=[first, last, separate]): |
| 621 | assert backend._enqueue_save("a", batch[:2]) |
| 622 | assert backend._enqueue_save("b", batch[2:]) |
| 623 | assert backend._pending_bytes == batch.nbytes |
| 624 | assert not backend._enqueue_save("a", batch[:2]) |
| 625 | assert not backend._enqueue_save("extra", batch[:1]) |
| 626 | if rejected: |
| 627 | first.set_exception(store_client.EmbeddingStoreOperationError("rejected")) |
| 628 | else: |
| 629 | first.set_result(None) |
| 630 | backend.reap() |
| 631 | assert backend._pending_bytes == batch.nbytes |
| 632 | assert not backend._enqueue_save("c", torch.ones(2, 2)) |
| 633 | last.set_result(None) |
| 634 | backend.reap() |
| 635 | assert backend._pending_bytes == 0 |
| 636 | assert backend._enqueue_save("c", torch.ones(2, 2)) |
| 637 | assert backend._pending_bytes == 16 |
| 638 | separate.set_result(None) |
| 639 | backend.reap() |
| 640 | assert backend._pending_bytes == 0 |
| 641 | |
| 642 | |
| 643 | @pytest.mark.parametrize("shared_storage", [False, True]) |
| 644 | def test_completion_race_preserves_failure_and_reclaims_other_items( |
| 645 | backend, shared_storage |
| 646 | ): |
| 647 | class RacingFuture(Future[None]): |
| 648 | def done(self): |
| 649 | observed = super().done() |
| 650 | if not observed: |
| 651 | self.set_exception(store_client.EmbeddingStoreError("uncertain")) |
| 652 | return observed |
| 653 | |
| 654 | racing = RacingFuture() |
| 655 | other: Future[None] = Future() |
| 656 | tensor = torch.ones(2, 2) |
| 657 | other_tensor = tensor[:1] if shared_storage else torch.ones(2, 2) |
| 658 | with patch.object(backend._executor, "submit", side_effect=[racing, other]): |
| 659 | assert backend._enqueue_save("race", tensor) |
| 660 | assert backend._enqueue_save("other", other_tensor) |
| 661 | backend.reap() |
| 662 | assert backend._pending_bytes == (16 if shared_storage else 32) |
| 663 | other.set_result(None) |
| 664 | with pytest.raises(store_client.EmbeddingStoreError): |
| 665 | backend.reap() |
| 666 | assert backend._pending_bytes == 16 |
| 667 | with pytest.raises(store_client.EmbeddingStoreError): |
| 668 | backend._enqueue_save("later", torch.ones(2, 2)) |
| 669 | |
| 670 | |
| 671 | def test_readiness_and_shutdown_do_not_block_p2p_completion(backend): |
| 672 | ready, waiting, entered = [threading.Event() for _ in range(3)] |
| 673 | |
| 674 | original = backend._executor.shutdown |
| 675 | |
| 676 | def closing(*args, **kwargs): |
| 677 | entered.set() |
| 678 | return original(*args, **kwargs) |
| 679 | |
| 680 | native = backend.store_client.store |
| 681 | |
| 682 | def close_store(): |
| 683 | assert ready.is_set() and not native.registered |
| 684 | assert backend._pending_bytes == 0 |
| 685 | return 0 |
| 686 | |
| 687 | native.close = MagicMock(side_effect=close_store) |
| 688 | put = native.batch_put_from_multi_buffers |
| 689 | |
| 690 | def waiting_put(*args): |
| 691 | waiting.set() |
| 692 | assert ready.wait(5) |
| 693 | assert backend.store_client._put_header is not None and native.registered |
| 694 | return put(*args) |
| 695 | |
| 696 | native.batch_put_from_multi_buffers = waiting_put |
| 697 | connector = make_worker(backend) |
| 698 | connector.bind_connector_metadata( |
| 699 | ECMooncakeConnectorMetadata(store_candidates={"a": SPEC}) |
| 700 | ) |
| 701 | cache: dict[str, torch.Tensor] = {} |
| 702 | with ( |
| 703 | patch.object( |
| 704 | store_backend, "_record_tensor_ready_event", return_value=MagicMock() |
| 705 | ), |
| 706 | patch.object(backend._executor, "shutdown", side_effect=closing), |
| 707 | ThreadPoolExecutor(max_workers=1) as closer, |
| 708 | ): |
| 709 | connector.start_save_caches(encoder_cache=cache) |
| 710 | cache["a"] = torch.arange(4, dtype=torch.float32).view(2, 2) |
| 711 | connector.save_caches(cache, "a") |
| 712 | assert connector.build_connector_worker_meta().pending_saves is False |
| 713 | assert waiting.wait(5) and not backend.store_client.store.objects |
| 714 | close = closer.submit(backend.shutdown) |
| 715 | try: |
| 716 | assert entered.wait(5) and not close.done() |
| 717 | assert not backend._enqueue_save("late", torch.ones(2, 2)) |
| 718 | finally: |
| 719 | ready.set() |
| 720 | close.result(timeout=5) |
| 721 | native.close.assert_called_once_with() |
| 722 | assert not backend.store_client.store.registered and backend._pending_bytes == 0 |
| 723 | |
| 724 | assert backend.store_client.store.objects[make_embedding_key(KEY, "a")][ |
| 725 | 24: |
| 726 | ] == struct.pack("<4f", 0, 1, 2, 3) |