| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Scheduler-side planning and observation for Mooncake cache transfers. |
| 4 | |
| 5 | The Scheduler prepares Producer reservations, consumes Consumer readiness |
| 6 | events, emits per-step Worker metadata, and converts Worker results back into |
| 7 | request availability without owning tensor memory or running data transfers. |
| 8 | """ |
| 9 | |
| 10 | from __future__ import annotations |
| 11 | |
| 12 | import math |
| 13 | import time |
| 14 | from collections import OrderedDict |
| 15 | from concurrent.futures import Future, ThreadPoolExecutor |
| 16 | from typing import TYPE_CHECKING, Any |
| 17 | |
| 18 | import torch |
| 19 | |
| 20 | from vllm.distributed.ec_transfer.ec_connector.base import ( |
| 21 | ECConnectorMetadata, |
| 22 | ) |
| 23 | from vllm.distributed.ec_transfer.ec_connector.cpu.common import ( |
| 24 | _get_encoder_cache_hidden_dim, |
| 25 | ) |
| 26 | from vllm.distributed.ec_transfer.ec_connector.mooncake.config import ( |
| 27 | _RESERVATION_TTL_SECONDS, |
| 28 | MooncakeECConfig, |
| 29 | ) |
| 30 | from vllm.distributed.ec_transfer.ec_connector.mooncake.control import ( |
| 31 | ControlClient, |
| 32 | EventInbox, |
| 33 | make_cancel_request, |
| 34 | ) |
| 35 | from vllm.distributed.ec_transfer.ec_connector.mooncake.metadata import ( |
| 36 | ECMooncakeConnectorMetadata, |
| 37 | ECMooncakeLoadSpec, |
| 38 | ECMooncakePushSpec, |
| 39 | ECMooncakeWorkerMetadata, |
| 40 | ) |
| 41 | from vllm.distributed.ec_transfer.ec_connector.mooncake.state import ( |
| 42 | SchedulerTransferState, |
| 43 | SchedulerTransferTable, |
| 44 | ) |
| 45 | from vllm.distributed.ec_transfer.ec_connector.mooncake.transfer import ( |
| 46 | ensure_mooncake_available, |
| 47 | ) |
| 48 | from vllm.distributed.ec_transfer.ec_connector.mooncake_store_embedding.data import ( |
| 49 | TensorSpec, |
| 50 | ) |
| 51 | from vllm.distributed.ec_transfer.ec_connector.utils import ( |
| 52 | PlaceholderMetadataResolver, |
| 53 | collect_ec_item_metadata, |
| 54 | ) |
| 55 | from vllm.logger import init_logger |
| 56 | from vllm.v1.core.sched.output import SchedulerOutput |
| 57 | from vllm.v1.outputs import ECConnectorOutput |
| 58 | |
| 59 | if TYPE_CHECKING: |
| 60 | from vllm.config import VllmConfig |
| 61 | |
| 62 | logger = init_logger(__name__) |
| 63 | |
| 64 | # Unmatched `ec_items` entries are a protocol desync, not a per-request |
| 65 | # event: warn on the first few, then only once in a while. |
| 66 | _MAX_UNRESOLVED_TRANSFER_ID_WARNINGS = 5 |
| 67 | |
| 68 | _DRAIN_MIN_INTERVAL = 0.005 |
| 69 | _MAX_PENDING_EVENTS = 4096 |
| 70 | _MAX_TERMINAL_TRANSFER_RECORDS = 1 << 16 |
| 71 | _CANCEL_ATTEMPTS = 2 |
| 72 | _CONTROL_WORKERS = 8 |
| 73 | |
| 74 | |
| 75 | class ECMooncakeScheduler: |
| 76 | """Coordinate Mooncake transfers from the vLLM Scheduler process.""" |
| 77 | |
| 78 | def __init__(self, vllm_config: VllmConfig) -> None: |
| 79 | ensure_mooncake_available() |
| 80 | config = MooncakeECConfig.from_vllm_config(vllm_config) |
| 81 | |
| 82 | self._store_candidates: list[tuple[str, str, TensorSpec]] = [] |
| 83 | self._cross_encoder_cache = config.cross_encoder_cache |
| 84 | self._is_producer = config.is_producer |
| 85 | self._is_consumer = config.is_consumer |
| 86 | self._control_addr = config.control_addr |
| 87 | self._push_wait_timeout = config.push_wait_timeout_s |
| 88 | self._encoder_cache_hidden_dim = ( |
| 89 | _get_encoder_cache_hidden_dim(vllm_config) if config.is_producer else None |
| 90 | ) |
| 91 | self._model_config = vllm_config.model_config |
| 92 | self._control_client = ControlClient(config.control_timeout_ms) |
| 93 | self._control_executor = ThreadPoolExecutor( |
| 94 | max_workers=_CONTROL_WORKERS, |
| 95 | thread_name_prefix="ec-mooncake-control", |
| 96 | ) |
| 97 | self._event_inbox = EventInbox(self._control_client, self._control_executor) |
| 98 | |
| 99 | self._metadata_resolver = PlaceholderMetadataResolver(vllm_config.model_config) |
| 100 | self._unresolved_transfer_ids = 0 |
| 101 | self._drain_pending = True |
| 102 | self._drained_at = 0.0 |
| 103 | self._pending_cancels: dict[str, Future[Any]] = {} |
| 104 | self._transfers = SchedulerTransferTable( |
| 105 | config.pool_size, _RESERVATION_TTL_SECONDS |
| 106 | ) |
| 107 | self._scheduler_pending_work = False |
| 108 | self._pushes_to_prepare: dict[str, ECMooncakePushSpec] = {} |
| 109 | self._prepared_push_transfer_ids: set[str] = set() |
| 110 | self._event_ready_shards: OrderedDict[str, set[int]] = OrderedDict() |
| 111 | # Mirror of the engine's encoder cache, maintained from the alloc and |
| 112 | # free notifications the Scheduler already sends. Tracking it here |
| 113 | # keeps `ensure_cache_available` on the upstream two-argument shape. |
| 114 | self._local_cache: set[str] = set() |
| 115 | self._failed_saves: set[str] = set() |
| 116 | |
| 117 | def _cancel_remote(self, consumer_zmq: str, transfer_id: str) -> bool: |
| 118 | pending = None |
| 119 | for _ in range(_CANCEL_ATTEMPTS): |
| 120 | pending = self._control_client.discover_shards(consumer_zmq) |
| 121 | if pending is not None: |
| 122 | break |
| 123 | if pending is None: |
| 124 | raise RuntimeError( |
| 125 | f"Could not discover every EC consumer shard at {consumer_zmq}" |
| 126 | ) |
| 127 | |
| 128 | cancelled = False |
| 129 | error: BaseException | None = None |
| 130 | for _ in range(_CANCEL_ATTEMPTS): |
| 131 | failed = [] |
| 132 | for addr in pending: |
| 133 | try: |
| 134 | result = self._control_client.request( |
| 135 | addr, |
| 136 | make_cancel_request(transfer_id), |
| 137 | ) |
| 138 | except Exception as exc: |
| 139 | if error is None: |
| 140 | error = exc |
| 141 | failed.append(addr) |
| 142 | continue |
| 143 | cancelled |= isinstance(result, dict) and bool(result.get("cancelled")) |
| 144 | if not failed: |
| 145 | return cancelled |
| 146 | pending = failed |
| 147 | if error is not None: |
| 148 | raise error |
| 149 | return cancelled |
| 150 | |
| 151 | def _note_awaiting_push( |
| 152 | self, |
| 153 | mm_hash: str, |
| 154 | transfer_id: str, |
| 155 | request_id: str, |
| 156 | now: float, |
| 157 | ) -> None: |
| 158 | record = self._transfers.wait_for_event( |
| 159 | transfer_id, |
| 160 | request_id, |
| 161 | mm_hash, |
| 162 | now + self._push_wait_timeout, |
| 163 | ) |
| 164 | if record is None: |
| 165 | # The id names another encoding; the request is already failed. |
| 166 | return |
| 167 | if record.state is not SchedulerTransferState.WAITING_EVENT: |
| 168 | return |
| 169 | assert record.deadline is not None |
| 170 | if now < record.deadline: |
| 171 | return |
| 172 | elapsed = now - record.deadline + self._push_wait_timeout |
| 173 | self._transfers.mark_unavailable(transfer_id, now) |
| 174 | logger.warning( |
| 175 | "EC Mooncake waited %.1fs for a push of mm_hash=%s " |
| 176 | "(transfer_id=%s) that never arrived; " |
| 177 | "requests needing it fail with a retryable error.", |
| 178 | elapsed, |
| 179 | mm_hash, |
| 180 | transfer_id, |
| 181 | ) |
| 182 | |
| 183 | def take_unavailable_requests(self) -> set[str]: |
| 184 | failed, self._failed_saves = self._failed_saves, set() |
| 185 | return failed | self._transfers.take_unavailable_requests() |
| 186 | |
| 187 | def _poll_pending_cancels(self) -> None: |
| 188 | pending = {} |
| 189 | for transfer_id, future in self._pending_cancels.items(): |
| 190 | if not future.done(): |
| 191 | pending[transfer_id] = future |
| 192 | continue |
| 193 | try: |
| 194 | future.result() |
| 195 | except Exception: |
| 196 | logger.warning( |
| 197 | "EC Mooncake reservation cancellation failed", exc_info=True |
| 198 | ) |
| 199 | self._pending_cancels = pending |
| 200 | |
| 201 | def _note_shard_ready(self, data: dict[str, Any]) -> bool: |
| 202 | shard_count = self._event_inbox.shard_count |
| 203 | if shard_count <= 1: |
| 204 | return True |
| 205 | transfer_id = str(data["transfer_id"]) |
| 206 | record = self._transfers.get(transfer_id) |
| 207 | if ( |
| 208 | record is not None |
| 209 | and record.state is not SchedulerTransferState.WAITING_EVENT |
| 210 | ): |
| 211 | return False |
| 212 | shard = data.get("shard") |
| 213 | shards = self._event_ready_shards.setdefault(transfer_id, set()) |
| 214 | self._event_ready_shards.move_to_end(transfer_id) |
| 215 | shards.add(int(shard) if shard is not None else len(shards)) |
| 216 | if len(shards) < shard_count: |
| 217 | while len(self._event_ready_shards) > _MAX_PENDING_EVENTS: |
| 218 | self._event_ready_shards.popitem(last=False) |
| 219 | return False |
| 220 | self._event_ready_shards.pop(transfer_id, None) |
| 221 | return True |
| 222 | |
| 223 | def _forget_shard_readiness(self, transfer_id: str) -> None: |
| 224 | self._event_ready_shards.pop(transfer_id, None) |
| 225 | |
| 226 | def _store_pushed_spec(self, data: dict[str, Any]) -> bool: |
| 227 | transfer_id = str(data["transfer_id"]) |
| 228 | identifier = str(data["mm_hash"]) |
| 229 | _, accepted = self._transfers.observe_ready( |
| 230 | ECMooncakeLoadSpec( |
| 231 | mm_hash=identifier, |
| 232 | nbytes=int(data["nbytes"]), |
| 233 | shape=tuple(int(value) for value in data["shape"]), |
| 234 | dtype=str(data["dtype"]), |
| 235 | transfer_id=transfer_id, |
| 236 | ), |
| 237 | time.monotonic() + _RESERVATION_TTL_SECONDS, |
| 238 | ) |
| 239 | return accepted |
| 240 | |
| 241 | def _queue_cancel( |
| 242 | self, |
| 243 | transfer_id: str, |
| 244 | mm_hash: str = "", |
| 245 | request_id: str = "", |
| 246 | ) -> None: |
| 247 | if not self._transfers.cancel( |
| 248 | transfer_id, |
| 249 | time.monotonic(), |
| 250 | mm_hash=mm_hash, |
| 251 | request_id=request_id, |
| 252 | ): |
| 253 | return |
| 254 | self._forget_shard_readiness(transfer_id) |
| 255 | self._pending_cancels[transfer_id] = self._control_executor.submit( |
| 256 | self._cancel_remote, |
| 257 | self._control_addr, |
| 258 | transfer_id, |
| 259 | ) |
| 260 | |
| 261 | def _expire_transfers(self) -> None: |
| 262 | now = time.monotonic() |
| 263 | expired = self._transfers.expire(now, _MAX_TERMINAL_TRANSFER_RECORDS) |
| 264 | for record in expired: |
| 265 | self._queue_cancel(record.transfer_id) |
| 266 | |
| 267 | def _drain_push_notifications(self) -> None: |
| 268 | now = time.monotonic() |
| 269 | if not self._drain_pending and now - self._drained_at < _DRAIN_MIN_INTERVAL: |
| 270 | return |
| 271 | self._drain_pending = False |
| 272 | self._drained_at = now |
| 273 | self._poll_pending_cancels() |
| 274 | self._expire_transfers() |
| 275 | events = self._event_inbox.drain(self._control_addr) |
| 276 | for data in events: |
| 277 | try: |
| 278 | if not isinstance(data, dict): |
| 279 | raise TypeError("EC readiness event must be an object") |
| 280 | if not data.get("ready"): |
| 281 | continue |
| 282 | self._accept_ready_event(data) |
| 283 | except (KeyError, TypeError, ValueError): |
| 284 | # Readiness events cross a plain PULL socket, so a malformed |
| 285 | # one costs that event and not the engine. |
| 286 | logger.warning( |
| 287 | "Discarding a malformed EC readiness event.", exc_info=True |
| 288 | ) |
| 289 | |
| 290 | def _accept_ready_event(self, data: dict[str, Any]) -> None: |
| 291 | transfer_id = str(data["transfer_id"]) |
| 292 | record = self._transfers.get(transfer_id) |
| 293 | if record is not None and record.state in { |
| 294 | SchedulerTransferState.CANCELLED, |
| 295 | SchedulerTransferState.UNAVAILABLE, |
| 296 | SchedulerTransferState.EXPIRED, |
| 297 | SchedulerTransferState.FAILED, |
| 298 | }: |
| 299 | return |
| 300 | if not self._note_shard_ready(data): |
| 301 | return |
| 302 | self._store_pushed_spec(data) |
| 303 | |
| 304 | def has_cache_item(self, identifier: str) -> bool: |
| 305 | if not self._is_consumer: |
| 306 | return False |
| 307 | self._drain_push_notifications() |
| 308 | return self._transfers.has_state( |
| 309 | identifier, |
| 310 | ( |
| 311 | SchedulerTransferState.READY, |
| 312 | SchedulerTransferState.RESIDENT, |
| 313 | SchedulerTransferState.AVAILABLE, |
| 314 | ), |
| 315 | ) |
| 316 | |
| 317 | def _warn_unresolved_transfer_id( |
| 318 | self, request: Any, index: int, where: str |
| 319 | ) -> None: |
| 320 | """Report an `ec_items` entry that cannot be matched to a feature. |
| 321 | |
| 322 | A caller that cannot resolve a transfer id has to fall back to a |
| 323 | locally invented one or skip its bookkeeping entirely, and either way |
| 324 | the two sides of a transfer stop agreeing on its name. That desyncs |
| 325 | silently -- reservations are never released and only surface minutes |
| 326 | later as a full buffer pool -- so say it out loud the first time. |
| 327 | |
| 328 | `mm_hash` in `ec_items` must be a value some engine reported, never one |
| 329 | the caller derived itself: `mm_features[i].identifier` folds in the |
| 330 | engine's media_io_kwargs and mm_processor_kwargs, so a media uuid alone |
| 331 | no longer equals it. Omit the field to match by position instead. |
| 332 | """ |
| 333 | self._unresolved_transfer_ids += 1 |
| 334 | count = self._unresolved_transfer_ids |
| 335 | if count > _MAX_UNRESOLVED_TRANSFER_ID_WARNINGS and count % 1000: |
| 336 | return |
| 337 | params = getattr(request, "ec_transfer_params", None) |
| 338 | items = (params or {}).get("ec_items") or [] |
| 339 | logger.warning( |
| 340 | "EC Mooncake could not resolve a transfer id at %s for req=%s " |
| 341 | "index=%d: feature identifier=%s does not match any of the %d " |
| 342 | "ec_items %s (ec_transfer_params present=%s). Occurrence %d. " |
| 343 | "ec_items[].mm_hash must echo an engine-reported identifier, or " |
| 344 | "be omitted so the entry is matched by position.", |
| 345 | where, |
| 346 | request.request_id, |
| 347 | index, |
| 348 | request.mm_features[index].identifier[:16], |
| 349 | len(items), |
| 350 | [str(item.get("mm_hash"))[:16] for item in items[:4]], |
| 351 | params is not None, |
| 352 | count, |
| 353 | ) |
| 354 | |
| 355 | @staticmethod |
| 356 | def _request_transfer_id(request: Any, index: int) -> str | None: |
| 357 | params = getattr(request, "ec_transfer_params", None) or {} |
| 358 | items = params.get("ec_items") or [] |
| 359 | mm_hash = request.mm_features[index].identifier |
| 360 | if index < len(items): |
| 361 | item = items[index] |
| 362 | if item.get("mm_hash") in (None, mm_hash) and item.get("transfer_id"): |
| 363 | return str(item["transfer_id"]) |
| 364 | for item in items: |
| 365 | if item.get("mm_hash") == mm_hash and item.get("transfer_id"): |
| 366 | return str(item["transfer_id"]) |
| 367 | return None |
| 368 | |
| 369 | def ensure_cache_available(self, request: Any, num_computed_tokens: int) -> bool: |
| 370 | if self._is_producer: |
| 371 | for index, feature in enumerate(request.mm_features): |
| 372 | if ( |
| 373 | feature.mm_position.offset + feature.mm_position.length |
| 374 | > num_computed_tokens |
| 375 | ): |
| 376 | self._prepare_push_spec(request, index) |
| 377 | if not self._is_consumer: |
| 378 | return True |
| 379 | |
| 380 | self._drain_push_notifications() |
| 381 | # One timestamp for the whole decision, so the deadlines it hands out |
| 382 | # cannot disagree between features. |
| 383 | now = time.monotonic() |
| 384 | all_ready = True |
| 385 | for index, feature in enumerate(request.mm_features): |
| 386 | if ( |
| 387 | feature.mm_position.offset + feature.mm_position.length |
| 388 | <= num_computed_tokens |
| 389 | ): |
| 390 | continue |
| 391 | mm_hash = feature.identifier |
| 392 | transfer_id = self._request_transfer_id(request, index) |
| 393 | if transfer_id is not None: |
| 394 | self._transfers.touch_available( |
| 395 | transfer_id, now + _RESERVATION_TTL_SECONDS |
| 396 | ) |
| 397 | if mm_hash in self._local_cache: |
| 398 | continue |
| 399 | if self._transfers.has_state(mm_hash, (SchedulerTransferState.READY,)): |
| 400 | continue |
| 401 | if self._transfers.has_state(mm_hash, (SchedulerTransferState.LOADING,)): |
| 402 | all_ready = False |
| 403 | continue |
| 404 | record = self._transfers.begin_load( |
| 405 | mm_hash, |
| 406 | transfer_id, |
| 407 | request.request_id, |
| 408 | now + self._push_wait_timeout, |
| 409 | ) |
| 410 | if record is not None: |
| 411 | self._scheduler_pending_work = True |
| 412 | all_ready = False |
| 413 | else: |
| 414 | waiting_id = transfer_id or f"{request.request_id}:{index}" |
| 415 | self._note_awaiting_push(mm_hash, waiting_id, request.request_id, now) |
| 416 | all_ready = False |
| 417 | return all_ready |
| 418 | |
| 419 | def _prepare_push_spec(self, request: Any, index: int) -> None: |
| 420 | params = getattr(request, "ec_transfer_params", None) or {} |
| 421 | consumer_zmq = params.get("consumer_zmq") |
| 422 | mm_hash = request.mm_features[index].identifier |
| 423 | transfer_id = self._request_transfer_id(request, index) |
| 424 | if transfer_id is None: |
| 425 | # Still push: after the proxy has rewritten the item to embeds the |
| 426 | # consumer has no media left to fall back on, so a nameless push |
| 427 | # beats none. But the consumer knows this transfer by the id it |
| 428 | # sent, not by the one invented here, so its cancel will never |
| 429 | # reach the reservation this push is about to take. |
| 430 | if consumer_zmq: |
| 431 | self._warn_unresolved_transfer_id(request, index, "push prepare") |
| 432 | transfer_id = f"{request.request_id}:{index}" |
| 433 | if not consumer_zmq or transfer_id in self._prepared_push_transfer_ids: |
| 434 | return |
| 435 | output = self._encoder_output_spec(request, index) |
| 436 | self._pushes_to_prepare[transfer_id] = ECMooncakePushSpec( |
| 437 | mm_hash=mm_hash, |
| 438 | nbytes=output.nbytes, |
| 439 | shape=output.shape, |
| 440 | dtype=output.dtype.removeprefix("torch."), |
| 441 | consumer_zmq=str(consumer_zmq), |
| 442 | transfer_id=transfer_id, |
| 443 | request_id=request.request_id, |
| 444 | ) |
| 445 | self._prepared_push_transfer_ids.add(transfer_id) |
| 446 | |
| 447 | def _encoder_output_spec(self, request: Any, index: int) -> TensorSpec: |
| 448 | dtype = self._model_config.dtype |
| 449 | assert isinstance(dtype, torch.dtype) |
| 450 | assert self._encoder_cache_hidden_dim is not None |
| 451 | shape = (request.get_num_encoder_embeds(index), self._encoder_cache_hidden_dim) |
| 452 | return TensorSpec(shape, str(dtype), math.prod(shape) * dtype.itemsize) |
| 453 | |
| 454 | def update_state_after_alloc(self, request: Any, index: int) -> None: |
| 455 | self._local_cache.add(request.mm_features[index].identifier) |
| 456 | if self._is_producer: |
| 457 | self._prepare_push_spec(request, index) |
| 458 | if ( |
| 459 | self._cross_encoder_cache |
| 460 | and request.mm_features[index].modality == "image" |
| 461 | ): |
| 462 | self._store_candidates.append( |
| 463 | ( |
| 464 | request.request_id, |
| 465 | request.mm_features[index].identifier, |
| 466 | self._encoder_output_spec(request, index), |
| 467 | ) |
| 468 | ) |
| 469 | |
| 470 | def update_state_after_free(self, request: Any, index: int) -> None: |
| 471 | if not self._is_consumer: |
| 472 | return |
| 473 | transfer_id = self._request_transfer_id(request, index) |
| 474 | if transfer_id is None: |
| 475 | self._warn_unresolved_transfer_id(request, index, "encoder-cache free") |
| 476 | return |
| 477 | self._queue_cancel( |
| 478 | transfer_id, |
| 479 | mm_hash=request.mm_features[index].identifier, |
| 480 | request_id=request.request_id, |
| 481 | ) |
| 482 | |
| 483 | def build_connector_meta( |
| 484 | self, scheduler_output: SchedulerOutput |
| 485 | ) -> ECConnectorMetadata: |
| 486 | for mm_hash in scheduler_output.free_encoder_mm_hashes: |
| 487 | self._local_cache.discard(mm_hash) |
| 488 | self._transfers.release_ready(mm_hash, time.monotonic()) |
| 489 | for transfer_id in self._transfers.drain_orphaned(): |
| 490 | self._queue_cancel(transfer_id) |
| 491 | meta = ECMooncakeConnectorMetadata( |
| 492 | freed=scheduler_output.free_encoder_mm_hashes |
| 493 | ) |
| 494 | preempted = scheduler_output.preempted_req_ids or set() |
| 495 | meta.store_candidates = { |
| 496 | identifier: output |
| 497 | for request_id, identifier, output in self._store_candidates |
| 498 | if request_id not in preempted |
| 499 | } |
| 500 | self._store_candidates.clear() |
| 501 | for push_spec in self._pushes_to_prepare.values(): |
| 502 | meta.pushes.append(push_spec) |
| 503 | self._pushes_to_prepare.clear() |
| 504 | for record in self._transfers.take_loads_to_dispatch(): |
| 505 | assert record.spec is not None |
| 506 | meta.loads.append(record.spec) |
| 507 | self._poll_pending_cancels() |
| 508 | self._drain_pending = True |
| 509 | return meta |
| 510 | |
| 511 | def update_connector_output(self, connector_output: ECConnectorOutput) -> None: |
| 512 | meta = connector_output.ec_connector_worker_meta |
| 513 | if not isinstance(meta, ECMooncakeWorkerMetadata): |
| 514 | return |
| 515 | for mm_hash in meta.loaded: |
| 516 | self._transfers.complete_load(mm_hash) |
| 517 | for mm_hash in meta.failed_loads: |
| 518 | self._transfers.fail_load(mm_hash, time.monotonic()) |
| 519 | for mm_hash in meta.reclaimed: |
| 520 | self._transfers.reclaim(mm_hash, time.monotonic()) |
| 521 | self._scheduler_pending_work = meta.pending_saves |
| 522 | self._failed_saves.update(meta.failed_saves) |
| 523 | |
| 524 | def has_pending_push_work(self) -> bool: |
| 525 | return self._scheduler_pending_work |
| 526 | |
| 527 | def request_finished(self, request: Any) -> tuple[bool, dict[str, Any] | None]: |
| 528 | if self._is_consumer: |
| 529 | for index in range(len(request.mm_features)): |
| 530 | transfer_id = self._request_transfer_id(request, index) |
| 531 | if transfer_id is None: |
| 532 | self._warn_unresolved_transfer_id(request, index, "request finish") |
| 533 | continue |
| 534 | self._queue_cancel( |
| 535 | transfer_id, |
| 536 | mm_hash=request.mm_features[index].identifier, |
| 537 | request_id=request.request_id, |
| 538 | ) |
| 539 | if self._is_producer and self._prepared_push_transfer_ids: |
| 540 | for index in range(len(request.mm_features)): |
| 541 | transfer_id = self._request_transfer_id(request, index) |
| 542 | if transfer_id is None: |
| 543 | transfer_id = f"{request.request_id}:{index}" |
| 544 | self._prepared_push_transfer_ids.discard(transfer_id) |
| 545 | if not self._is_producer: |
| 546 | return False, None |
| 547 | |
| 548 | items = collect_ec_item_metadata(request.mm_features, self._metadata_resolver) |
| 549 | for index, feature in enumerate(request.mm_features): |
| 550 | transfer_id = self._request_transfer_id(request, index) |
| 551 | if transfer_id is not None: |
| 552 | items[feature.identifier]["transfer_id"] = transfer_id |
| 553 | |
| 554 | if not items: |
| 555 | return False, None |
| 556 | return False, items |
| 557 | |
| 558 | def close(self) -> None: |
| 559 | self._control_executor.shutdown(wait=True, cancel_futures=True) |
| 560 | self._control_client.close() |
| 561 | self._event_inbox.close() |