| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Session objects for the ECCPUConnector. |
| 4 | |
| 5 | Layer hierarchy: connector → session → transport (connection) |
| 6 | |
| 7 | Two long-lived sessions: |
| 8 | |
| 9 | ProducerSession — owns ZmqServerTransport. The connector calls poll_step() |
| 10 | once per engine step, which drains the transport without blocking and |
| 11 | feeds the raw bytes to poll(). poll() decodes XferReqs, grants or NACKs |
| 12 | them, pins source blocks, sends XferAck replies, drains NIXL completion |
| 13 | notifications, and sweeps expired xfers. |
| 14 | |
| 15 | ConsumerSession — owns one ZmqClientConnection and the NIXL agent |
| 16 | registration for that peer. The connector calls poll(messages, now) |
| 17 | once per engine step with the raw bytes collected for this peer. |
| 18 | poll() decodes XferAcks, dispatches them to ConsumerXfer objects, |
| 19 | and advances all active xfer state machines. |
| 20 | |
| 21 | Internal data classes: |
| 22 | |
| 23 | ProducerXfer — one pinned read grant, keyed by session_id. |
| 24 | ConsumerXfer — state machine for one NIXL READ, keyed by mm_hash. |
| 25 | """ |
| 26 | |
| 27 | import enum |
| 28 | import logging |
| 29 | import time |
| 30 | import uuid |
| 31 | from dataclasses import dataclass, field |
| 32 | from typing import Any |
| 33 | |
| 34 | import msgspec |
| 35 | |
| 36 | from vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq import ( |
| 37 | ZmqClientConnection, |
| 38 | ZmqClientTransport, |
| 39 | ZmqServerTransport, |
| 40 | ) |
| 41 | from vllm.distributed.ec_transfer.ec_connector.cpu.data.base import ( |
| 42 | DataTransport, |
| 43 | ) |
| 44 | from vllm.distributed.ec_transfer.ec_connector.cpu.protocol import ( |
| 45 | EC_CONNECTOR_VERSION, |
| 46 | EXPECTED_NACKS, |
| 47 | RETRYABLE_NACKS, |
| 48 | XferAck, |
| 49 | XferReq, |
| 50 | XferStatus, |
| 51 | ) |
| 52 | from vllm.distributed.ec_transfer.ec_connector.cpu.scheduler.embedding_cache import ( |
| 53 | EmbeddingCache, |
| 54 | ) |
| 55 | from vllm.distributed.ec_transfer.ec_connector.cpu.utils import PeerAddr |
| 56 | from vllm.logger import init_logger |
| 57 | |
| 58 | logger = init_logger(__name__) |
| 59 | |
| 60 | PRODUCER_PIN_LEASE_S = 30.0 |
| 61 | _CONSUMER_READ_TIMEOUT_S = 20.0 |
| 62 | _CONSUMER_QUARANTINE_TIMEOUT_S = 60.0 |
| 63 | |
| 64 | |
| 65 | # ── XferState ───────────────────────────────────────────────────────────────── |
| 66 | |
| 67 | |
| 68 | class XferState(enum.Enum): |
| 69 | """Lifecycle states of a single ConsumerXfer.""" |
| 70 | |
| 71 | WAITING_ACK = "waiting_ack" # XferReq sent, awaiting XferAck |
| 72 | READING = "reading" # NIXL READ in flight |
| 73 | DONE = "done" # READ completed; promote to local cache |
| 74 | ACK_TIMEOUT = "ack_timeout" # XferAck never arrived; free blocks + retry |
| 75 | READ_FAILED = "read_failed" # Unexpected NIXL state; free blocks + tombstone |
| 76 | QUARANTINED = "quarantined" # Read timed out, NIXL unabortable; keep blocks |
| 77 | SETTLED = "settled" # Quarantined DMA done; safe to free blocks |
| 78 | |
| 79 | |
| 80 | # ── ProducerXfer ────────────────────────────────────────────────────────────── |
| 81 | |
| 82 | |
| 83 | @dataclass |
| 84 | class ProducerXfer: |
| 85 | """One pinned read grant on the producer side. |
| 86 | |
| 87 | Pure data — lifecycle operations (pin/unpin) are performed by ProducerSession. |
| 88 | Keyed by (consumer_session_id, mm_hash) in ProducerSession._active_xfers. |
| 89 | """ |
| 90 | |
| 91 | mm_hash: str |
| 92 | block_indices: list[int] |
| 93 | deadline: float |
| 94 | |
| 95 | def is_expired(self) -> bool: |
| 96 | return time.monotonic() > self.deadline |
| 97 | |
| 98 | |
| 99 | # ── ConsumerXfer ────────────────────────────────────────────────────────────── |
| 100 | |
| 101 | |
| 102 | @dataclass |
| 103 | class ConsumerXfer: |
| 104 | """State machine for one NIXL READ on the consumer side. |
| 105 | |
| 106 | transfer_handle is the in-flight NIXL handle returned by post_read() — |
| 107 | used for check_xfer_state() / release_xfer_handle(). The dlist handle |
| 108 | for the remote peer lives inside DataTransport, not here. |
| 109 | """ |
| 110 | |
| 111 | mm_hash: str |
| 112 | block_indices: list[int] |
| 113 | addr: PeerAddr |
| 114 | deadline: float |
| 115 | data: DataTransport |
| 116 | consumer_session_id: str # session identity; combined with mm_hash for notif_msg |
| 117 | |
| 118 | transfer_handle: Any | None = field(default=None) |
| 119 | _quarantined: bool = field(default=False, init=False) |
| 120 | |
| 121 | def handle_ack(self, ack: XferAck, agent_name: str) -> bool: |
| 122 | """Issue the NIXL READ on an OK XferAck. Returns False on any NACK.""" |
| 123 | if ack.status != XferStatus.OK: |
| 124 | return False |
| 125 | # notif_msg encodes both session_id and mm_hash so the producer can |
| 126 | # do a direct (session_id, mm_hash) → ProducerXfer lookup on completion. |
| 127 | notif_msg = f"{self.consumer_session_id}:{self.mm_hash}".encode() |
| 128 | self.transfer_handle = self.data.post_read( |
| 129 | self.block_indices, |
| 130 | agent_name, |
| 131 | ack.src_block_indices, |
| 132 | notif_msg=notif_msg, |
| 133 | ) |
| 134 | self.deadline = time.monotonic() + _CONSUMER_READ_TIMEOUT_S |
| 135 | logger.debug( |
| 136 | "EC consumer: NIXL READ posted mm_hash=%s agent=%s local_blocks=%d " |
| 137 | "remote_blocks=%d", |
| 138 | self.mm_hash, |
| 139 | agent_name, |
| 140 | len(self.block_indices), |
| 141 | len(ack.src_block_indices), |
| 142 | ) |
| 143 | return True |
| 144 | |
| 145 | def poll(self, now: float) -> XferState: |
| 146 | if self._quarantined: |
| 147 | return self._poll_quarantined(now) |
| 148 | if self.transfer_handle is None: |
| 149 | return ( |
| 150 | XferState.ACK_TIMEOUT if now > self.deadline else XferState.WAITING_ACK |
| 151 | ) |
| 152 | return self._poll_reading(now) |
| 153 | |
| 154 | def _poll_reading(self, now: float) -> XferState: |
| 155 | assert self.transfer_handle is not None |
| 156 | try: |
| 157 | state = self.data.check_xfer_state(self.transfer_handle) |
| 158 | except Exception: |
| 159 | logger.exception( |
| 160 | "EC Connector: check_xfer_state failed for mm_hash=%s", self.mm_hash |
| 161 | ) |
| 162 | self.data.release_xfer_handle(self.transfer_handle) |
| 163 | self.transfer_handle = None |
| 164 | return XferState.READ_FAILED |
| 165 | if state == "DONE": |
| 166 | self.data.release_xfer_handle(self.transfer_handle) |
| 167 | self.transfer_handle = None |
| 168 | logger.debug("EC consumer: NIXL READ completed mm_hash=%s", self.mm_hash) |
| 169 | return XferState.DONE |
| 170 | if state == "PROC": |
| 171 | if now > self.deadline: |
| 172 | self._quarantined = True |
| 173 | self.deadline = now + _CONSUMER_QUARANTINE_TIMEOUT_S |
| 174 | logger.warning( |
| 175 | "EC Connector: READ for mm_hash=%s from %s:%d timed out; " |
| 176 | "quarantining", |
| 177 | self.mm_hash, |
| 178 | self.addr[0], |
| 179 | self.addr[1], |
| 180 | ) |
| 181 | return XferState.QUARANTINED |
| 182 | return XferState.READING |
| 183 | logger.warning( |
| 184 | "EC Connector: READ for mm_hash=%s unexpected NIXL state %r", |
| 185 | self.mm_hash, |
| 186 | state, |
| 187 | ) |
| 188 | self.data.release_xfer_handle(self.transfer_handle) |
| 189 | self.transfer_handle = None |
| 190 | return XferState.READ_FAILED |
| 191 | |
| 192 | def _poll_quarantined(self, now: float) -> XferState: |
| 193 | if self.transfer_handle is None: |
| 194 | return XferState.SETTLED |
| 195 | # If quarantine has been too long, assume NIXL is stuck and evict blocks. |
| 196 | if now > self.deadline: |
| 197 | self.data.release_xfer_handle(self.transfer_handle) |
| 198 | self.transfer_handle = None |
| 199 | logger.warning( |
| 200 | "EC Connector: quarantined mm_hash=%s from %s:%d timed out; " |
| 201 | "evicting blocks", |
| 202 | self.mm_hash, |
| 203 | self.addr[0], |
| 204 | self.addr[1], |
| 205 | ) |
| 206 | return XferState.READ_FAILED |
| 207 | try: |
| 208 | state = self.data.check_xfer_state(self.transfer_handle) |
| 209 | except Exception: |
| 210 | logger.exception( |
| 211 | "EC Connector: check_xfer_state failed for quarantined mm_hash=%s", |
| 212 | self.mm_hash, |
| 213 | ) |
| 214 | state = None |
| 215 | if state == "PROC": |
| 216 | return XferState.QUARANTINED |
| 217 | self.data.release_xfer_handle(self.transfer_handle) |
| 218 | self.transfer_handle = None |
| 219 | logger.debug( |
| 220 | "EC Connector: quarantined mm_hash=%s settled (state=%s)", |
| 221 | self.mm_hash, |
| 222 | state, |
| 223 | ) |
| 224 | return XferState.SETTLED |
| 225 | |
| 226 | def cancel(self) -> None: |
| 227 | """Cancel a WAITING_ACK xfer (peer down). No tombstone — caller allows retry.""" |
| 228 | assert self.transfer_handle is None, "cancel() requires WAITING_ACK state" |
| 229 | |
| 230 | def release(self) -> None: |
| 231 | """Force-release the transfer handle. Shutdown only.""" |
| 232 | if self.transfer_handle is not None: |
| 233 | self.data.release_xfer_handle(self.transfer_handle) |
| 234 | self.transfer_handle = None |
| 235 | |
| 236 | |
| 237 | # ── ProducerSession ─────────────────────────────────────────────────────────── |
| 238 | |
| 239 | |
| 240 | class ProducerSession: |
| 241 | """Long-lived producer session. Owns ZmqServerTransport. |
| 242 | |
| 243 | poll_step() runs once per engine step: drains the transport without |
| 244 | blocking, decodes XferReqs, grants or NACKs each one, drains NIXL completion |
| 245 | notifications, and sweeps expired xfers. |
| 246 | """ |
| 247 | |
| 248 | def __init__( |
| 249 | self, |
| 250 | transport: ZmqServerTransport, |
| 251 | data: DataTransport, |
| 252 | cache: EmbeddingCache, |
| 253 | compat_hash: str, |
| 254 | ) -> None: |
| 255 | self._transport = transport |
| 256 | self._data = data |
| 257 | self._cache = cache # EmbeddingCache — single block owner |
| 258 | self._compat_hash = compat_hash |
| 259 | |
| 260 | # Key: "{consumer_session_id}:{mm_hash}" — matches the NIXL notif_msg exactly, |
| 261 | # so completion notification lookup is a direct dict pop with no parsing. |
| 262 | self._active_xfers: dict[str, ProducerXfer] = {} |
| 263 | |
| 264 | self._req_decoder = msgspec.msgpack.Decoder(XferReq) |
| 265 | self._req_encoder = msgspec.msgpack.Encoder() |
| 266 | |
| 267 | def close(self) -> None: |
| 268 | for xfer in self._active_xfers.values(): |
| 269 | self._cache.unpin(xfer.mm_hash) # release any still-held pins |
| 270 | self._active_xfers.clear() |
| 271 | self._transport.close() |
| 272 | |
| 273 | def poll_step(self) -> None: |
| 274 | """Drain the transport without blocking and process what arrived.""" |
| 275 | self.poll(self._transport.poll(timeout_ms=0)) |
| 276 | |
| 277 | def poll(self, messages: list[tuple[bytes, bytes]]) -> None: |
| 278 | """Decode and process inbound XferReqs, drain NIXL notifs, sweep timeouts. |
| 279 | |
| 280 | messages: (identity, payload) pairs from ZmqServerTransport.poll(). |
| 281 | """ |
| 282 | for identity, payload in messages: |
| 283 | try: |
| 284 | req = self._req_decoder.decode(payload) |
| 285 | except (msgspec.DecodeError, msgspec.ValidationError): |
| 286 | logger.warning("EC Connector: dropped malformed XferReq") |
| 287 | continue |
| 288 | ack = self._grant_or_nack(req) |
| 289 | self._transport.send(identity, self._req_encoder.encode(ack)) |
| 290 | |
| 291 | self._drain_notifs() |
| 292 | self._sweep_timeouts() |
| 293 | |
| 294 | def _grant_or_nack(self, req: XferReq) -> XferAck: |
| 295 | if req.connector_version != EC_CONNECTOR_VERSION: |
| 296 | logger.warning( |
| 297 | "EC Connector: incompatible version for mm_hash=%s", req.mm_hash |
| 298 | ) |
| 299 | return XferAck(mm_hash=req.mm_hash, status=XferStatus.NACK_VERSION) |
| 300 | if req.compatibility_hash != self._compat_hash: |
| 301 | logger.warning( |
| 302 | "EC Connector: incompatible compat hash for mm_hash=%s", req.mm_hash |
| 303 | ) |
| 304 | return XferAck(mm_hash=req.mm_hash, status=XferStatus.NACK_INCOMPAT) |
| 305 | if not req.session_id: |
| 306 | # Grants are keyed by (session_id, mm_hash), so an omitted id from |
| 307 | # one peer aliases another's: the first completion would unpin |
| 308 | # blocks the other peer is still reading. |
| 309 | logger.warning( |
| 310 | "EC Connector: XferReq without session_id for mm_hash=%s", req.mm_hash |
| 311 | ) |
| 312 | return XferAck(mm_hash=req.mm_hash, status=XferStatus.NACK_INTERNAL) |
| 313 | entry = self._cache.pin_if_ready(req.mm_hash) |
| 314 | if entry is None or not entry.ready: |
| 315 | status = ( |
| 316 | XferStatus.NACK_MISSING if entry is None else XferStatus.NACK_NOT_READY |
| 317 | ) |
| 318 | logger.log( |
| 319 | logging.INFO if status in EXPECTED_NACKS else logging.WARNING, |
| 320 | "EC Connector: mm_hash=%s not serveable in local cache (%s); NACKing", |
| 321 | req.mm_hash, |
| 322 | status.name, |
| 323 | ) |
| 324 | return XferAck(mm_hash=req.mm_hash, status=status) |
| 325 | block_indices = list(entry.block_ids) |
| 326 | key = f"{req.session_id}:{req.mm_hash}" |
| 327 | if key in self._active_xfers: |
| 328 | # This peer is re-requesting a grant it still holds, having timed |
| 329 | # out waiting. Overwriting it drops the pin taken for it, which |
| 330 | # neither a completion nor the sweep would ever release. |
| 331 | self._cache.unpin(req.mm_hash) |
| 332 | self._active_xfers[key] = ProducerXfer( |
| 333 | mm_hash=req.mm_hash, |
| 334 | block_indices=block_indices, |
| 335 | deadline=time.monotonic() + PRODUCER_PIN_LEASE_S, |
| 336 | ) |
| 337 | logger.debug( |
| 338 | "EC producer: serving NIXL READ request mm_hash=%s key=%s blocks=%d", |
| 339 | req.mm_hash, |
| 340 | key, |
| 341 | len(block_indices), |
| 342 | ) |
| 343 | return XferAck( |
| 344 | mm_hash=req.mm_hash, |
| 345 | status=XferStatus.OK, |
| 346 | session_id=req.session_id, # echo consumer's session_id back |
| 347 | src_block_indices=block_indices, |
| 348 | agent_metadata=self._data.get_agent_metadata(), |
| 349 | mem_descriptor=self._data.get_mem_descriptor(), |
| 350 | ) |
| 351 | |
| 352 | def _drain_notifs(self) -> None: |
| 353 | try: |
| 354 | notifs = self._data.get_new_notifs() |
| 355 | except Exception: |
| 356 | logger.exception("EC Connector: get_new_notifs failed") |
| 357 | return |
| 358 | for msgs in notifs.values(): |
| 359 | for msg in msgs: |
| 360 | try: |
| 361 | key = msg.decode("utf-8") # "{session_id}:{mm_hash}" |
| 362 | except Exception: |
| 363 | logger.exception("EC Connector: Dropped malformed NIXL notif") |
| 364 | continue |
| 365 | xfer = self._active_xfers.pop(key, None) |
| 366 | if xfer is not None: |
| 367 | self._cache.unpin(xfer.mm_hash) |
| 368 | logger.debug( |
| 369 | "EC producer: NIXL READ completed mm_hash=%s key=%s", |
| 370 | xfer.mm_hash, |
| 371 | key, |
| 372 | ) |
| 373 | |
| 374 | def _sweep_timeouts(self) -> None: |
| 375 | for key, xfer in list(self._active_xfers.items()): |
| 376 | if xfer.is_expired(): |
| 377 | logger.warning( |
| 378 | "EC producer: consumer did not complete NIXL READ " |
| 379 | "mm_hash=%s key=%s within the pin lease; releasing " |
| 380 | "pinned blocks", |
| 381 | xfer.mm_hash, |
| 382 | key, |
| 383 | ) |
| 384 | self._cache.unpin(xfer.mm_hash) |
| 385 | del self._active_xfers[key] |
| 386 | |
| 387 | |
| 388 | # ── ConsumerSession ─────────────────────────────────────────────────────────── |
| 389 | |
| 390 | |
| 391 | @dataclass |
| 392 | class ConsumerSessionResults: |
| 393 | """Results from ConsumerSession.take_results().""" |
| 394 | |
| 395 | completed: set[str] # promote to local cache |
| 396 | tombstoned: set[str] # free blocks + prevent retry this step |
| 397 | quarantined: set[str] # timed-out, DMA still running: prevent retry |
| 398 | # but DO NOT free blocks — session holds them |
| 399 | retryable: set[str] # free blocks, allow another attempt (not ready, peer down) |
| 400 | settled: list[tuple[str, list[int]]] # (mm_hash, block_indices) quarantine cleared |
| 401 | |
| 402 | |
| 403 | class ConsumerSession: |
| 404 | """Long-lived per-producer session on the consumer side. |
| 405 | |
| 406 | Owns the ZMQ DEALER connection and the NIXL agent registration for one |
| 407 | producer peer. The connector calls poll(messages, now) once per engine |
| 408 | step with the raw bytes collected from this peer's DEALER socket. |
| 409 | |
| 410 | poll() decodes XferAcks, dispatches them to the right ConsumerXfer, |
| 411 | advances all active xfer state machines, and drains quarantined xfers. |
| 412 | Results accumulate in internal sets; call take_results() to collect them. |
| 413 | """ |
| 414 | |
| 415 | def __init__( |
| 416 | self, |
| 417 | addr: PeerAddr, |
| 418 | zmq_conn: ZmqClientConnection, |
| 419 | transport: ZmqClientTransport, |
| 420 | data: DataTransport, |
| 421 | compat_hash: str, |
| 422 | ) -> None: |
| 423 | self._addr = addr |
| 424 | self._zmq = zmq_conn |
| 425 | self._transport = transport |
| 426 | self._data = data |
| 427 | self._compat_hash = compat_hash |
| 428 | self._session_id = str( |
| 429 | uuid.uuid4() |
| 430 | ) # stable identity for (session_id:mm_hash) keys |
| 431 | |
| 432 | self._nixl_agent_name: str | None = None |
| 433 | self._nixl_metadata_bytes: bytes | None = None |
| 434 | |
| 435 | self._xfers: dict[str, ConsumerXfer] = {} # mm_hash → xfer |
| 436 | self._quarantined: list[ConsumerXfer] = [] # DMA pending, to poll |
| 437 | self._newly_quarantined: set[str] = set() # quarantined since last take_results |
| 438 | |
| 439 | self._completed: set[str] = set() |
| 440 | self._tombstoned: set[str] = set() |
| 441 | self._retryable: set[str] = set() |
| 442 | self._settled: list[tuple[str, list[int]]] = [] |
| 443 | |
| 444 | self._encoder = msgspec.msgpack.Encoder() |
| 445 | self._decoder = msgspec.msgpack.Decoder(XferAck) |
| 446 | |
| 447 | @property |
| 448 | def addr(self) -> PeerAddr: |
| 449 | return self._addr |
| 450 | |
| 451 | def start_xfer( |
| 452 | self, mm_hash: str, block_indices: list[int], deadline: float |
| 453 | ) -> None: |
| 454 | """Create a ConsumerXfer and send the XferReq over the connection.""" |
| 455 | xfer = ConsumerXfer( |
| 456 | mm_hash=mm_hash, |
| 457 | block_indices=block_indices, |
| 458 | addr=self._addr, |
| 459 | deadline=deadline, |
| 460 | data=self._data, |
| 461 | consumer_session_id=self._session_id, |
| 462 | ) |
| 463 | self._zmq.send( |
| 464 | self._encoder.encode( |
| 465 | XferReq( |
| 466 | mm_hash=mm_hash, |
| 467 | compatibility_hash=self._compat_hash, |
| 468 | session_id=self._session_id, |
| 469 | ) |
| 470 | ) |
| 471 | ) |
| 472 | self._xfers[mm_hash] = xfer |
| 473 | logger.debug( |
| 474 | "EC Connector: XferReq sent mm_hash=%s to %s:%d", |
| 475 | mm_hash, |
| 476 | self._addr[0], |
| 477 | self._addr[1], |
| 478 | ) |
| 479 | |
| 480 | def poll(self, messages: list[bytes], now: float) -> None: |
| 481 | """Decode inbound XferAcks, dispatch to xfers, and advance all state machines. |
| 482 | |
| 483 | messages: raw bytes from ZmqClientTransport.poll() for this peer. |
| 484 | Results accumulate until take_results() is called. |
| 485 | """ |
| 486 | for raw in messages: |
| 487 | try: |
| 488 | ack = self._decoder.decode(raw) |
| 489 | except (msgspec.DecodeError, msgspec.ValidationError): |
| 490 | logger.warning( |
| 491 | "EC Connector: dropped malformed XferAck from %s:%d", |
| 492 | self._addr[0], |
| 493 | self._addr[1], |
| 494 | ) |
| 495 | continue |
| 496 | self._dispatch_ack(ack) |
| 497 | |
| 498 | for mm_hash, xfer in list(self._xfers.items()): |
| 499 | state = xfer.poll(now) |
| 500 | if state == XferState.DONE: |
| 501 | self._completed.add(mm_hash) |
| 502 | del self._xfers[mm_hash] |
| 503 | elif state == XferState.ACK_TIMEOUT: |
| 504 | # The scheduler bounds retries with a per-request deadline. |
| 505 | logger.warning( |
| 506 | "EC consumer: no XferAck for mm_hash=%s from %s:%d in time; " |
| 507 | "will re-request", |
| 508 | mm_hash, |
| 509 | self._addr[0], |
| 510 | self._addr[1], |
| 511 | ) |
| 512 | self._retryable.add(mm_hash) |
| 513 | del self._xfers[mm_hash] |
| 514 | elif state == XferState.READ_FAILED: |
| 515 | self._tombstoned.add(mm_hash) |
| 516 | del self._xfers[mm_hash] |
| 517 | elif state == XferState.QUARANTINED: |
| 518 | # DMA may still be writing — caller must NOT free blocks. |
| 519 | self._quarantined.append(xfer) |
| 520 | self._newly_quarantined.add(mm_hash) |
| 521 | del self._xfers[mm_hash] |
| 522 | |
| 523 | self._drain_quarantine(now) |
| 524 | |
| 525 | def _dispatch_ack(self, ack: XferAck) -> None: |
| 526 | xfer = self._xfers.get(ack.mm_hash) |
| 527 | if xfer is None or xfer.transfer_handle is not None: |
| 528 | return # unknown mm_hash or duplicate ack |
| 529 | |
| 530 | if ack.status != XferStatus.OK: |
| 531 | logger.log( |
| 532 | logging.INFO if ack.status in EXPECTED_NACKS else logging.WARNING, |
| 533 | "EC Connector: NACK %s from %s:%d for mm_hash=%s", |
| 534 | ack.status.name, |
| 535 | self._addr[0], |
| 536 | self._addr[1], |
| 537 | ack.mm_hash, |
| 538 | ) |
| 539 | del self._xfers[ack.mm_hash] |
| 540 | # The producer announced this encoding and its save is still in |
| 541 | # flight, so another attempt on a later step can be granted. Every |
| 542 | # other NACK is terminal for this mm_hash. |
| 543 | if ack.status in RETRYABLE_NACKS: |
| 544 | self._retryable.add(ack.mm_hash) |
| 545 | else: |
| 546 | self._tombstoned.add(ack.mm_hash) |
| 547 | return |
| 548 | |
| 549 | try: |
| 550 | agent_name = self._ensure_registered(ack.agent_metadata, ack.mem_descriptor) |
| 551 | except Exception: |
| 552 | logger.exception( |
| 553 | "EC Connector: failed to register peer %s:%d for mm_hash=%s", |
| 554 | self._addr[0], |
| 555 | self._addr[1], |
| 556 | ack.mm_hash, |
| 557 | ) |
| 558 | del self._xfers[ack.mm_hash] |
| 559 | self._tombstoned.add(ack.mm_hash) |
| 560 | return |
| 561 | |
| 562 | if not xfer.handle_ack(ack, agent_name): |
| 563 | del self._xfers[ack.mm_hash] |
| 564 | self._tombstoned.add(ack.mm_hash) |
| 565 | |
| 566 | def _ensure_registered(self, metadata: bytes, mem_descriptor: bytes) -> str: |
| 567 | """Register or re-register the peer in DataTransport. Returns agent_name.""" |
| 568 | if self._nixl_agent_name is not None and self._nixl_metadata_bytes == metadata: |
| 569 | return self._nixl_agent_name |
| 570 | if self._nixl_agent_name is not None: |
| 571 | logger.info( |
| 572 | "EC Connector: producer %s:%d restarted; re-registering", |
| 573 | self._addr[0], |
| 574 | self._addr[1], |
| 575 | ) |
| 576 | self._data.remove_remote_peer(self._nixl_agent_name) |
| 577 | agent_name = self._data.add_remote_peer(metadata, mem_descriptor) |
| 578 | self._nixl_agent_name = agent_name |
| 579 | self._nixl_metadata_bytes = metadata |
| 580 | logger.debug( |
| 581 | "EC Connector: registered peer %s:%d agent=%s", |
| 582 | self._addr[0], |
| 583 | self._addr[1], |
| 584 | agent_name, |
| 585 | ) |
| 586 | return agent_name |
| 587 | |
| 588 | def _drain_quarantine(self, now: float) -> None: |
| 589 | still_pending: list[ConsumerXfer] = [] |
| 590 | for xfer in list(self._quarantined): |
| 591 | state = xfer.poll(now) |
| 592 | if state == XferState.SETTLED: |
| 593 | self._settled.append((xfer.mm_hash, xfer.block_indices)) |
| 594 | else: |
| 595 | still_pending.append(xfer) |
| 596 | self._quarantined = still_pending |
| 597 | |
| 598 | def on_peer_down(self) -> None: |
| 599 | """WAITING_ACK xfers → cancel (retry allowed); READING xfers → quarantine.""" |
| 600 | for mm_hash, xfer in list(self._xfers.items()): |
| 601 | if xfer.transfer_handle is None: |
| 602 | xfer.cancel() |
| 603 | self._retryable.add(mm_hash) |
| 604 | else: |
| 605 | self._quarantined.append(xfer) |
| 606 | self._newly_quarantined.add(mm_hash) |
| 607 | self._xfers.clear() |
| 608 | |
| 609 | def take_results(self) -> ConsumerSessionResults: |
| 610 | """Return and clear all accumulated results since the last call.""" |
| 611 | results = ConsumerSessionResults( |
| 612 | completed=self._completed, |
| 613 | tombstoned=self._tombstoned, |
| 614 | quarantined=self._newly_quarantined, |
| 615 | retryable=self._retryable, |
| 616 | settled=self._settled, |
| 617 | ) |
| 618 | self._completed = set() |
| 619 | self._tombstoned = set() |
| 620 | self._newly_quarantined = set() |
| 621 | self._retryable = set() |
| 622 | self._settled = [] |
| 623 | return results |
| 624 | |
| 625 | def has_in_flight(self) -> bool: |
| 626 | return bool(self._xfers) or bool(self._quarantined) |
| 627 | |
| 628 | def close(self) -> None: |
| 629 | """Release all transfer handles, deregister NIXL peer, close connection. |
| 630 | |
| 631 | Connection teardown is routed through the transport so the pooled |
| 632 | connection is dropped in lock-step with the session — a session and its |
| 633 | DEALER connection have a 1:1 lifecycle, so a later connect() to the same |
| 634 | peer builds a fresh socket instead of handing back this dead one. |
| 635 | """ |
| 636 | for xfer in self._xfers.values(): |
| 637 | xfer.release() |
| 638 | self._xfers.clear() |
| 639 | for xfer in self._quarantined: |
| 640 | xfer.release() |
| 641 | self._quarantined.clear() |
| 642 | if self._nixl_agent_name is not None: |
| 643 | self._data.remove_remote_peer(self._nixl_agent_name) |
| 644 | self._transport.remove(self._addr) |