633 lines
26 KiB
Python
633 lines
26 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
from __future__ import annotations
|
|
|
|
import threading
|
|
import time
|
|
from collections import defaultdict
|
|
from contextlib import AbstractContextManager, nullcontext
|
|
from dataclasses import dataclass, field
|
|
from queue import Full, Queue
|
|
from types import SimpleNamespace
|
|
from typing import Any, cast
|
|
|
|
import torch
|
|
from vllm.logger import init_logger
|
|
|
|
from vllm_omni.data_entry_keys import OmniPayload, flatten_payload, unflatten_payload
|
|
from vllm_omni.worker.omni_connector_model_runner_mixin import (
|
|
OmniConnectorModelRunnerMixin,
|
|
should_accumulate_full_payload_output,
|
|
)
|
|
from vllm_omni.worker.payload_span import get_tensor_span, merge_tensor_spans
|
|
from vllm_omni.worker_v2.delivery import (
|
|
DeliveryTimeoutError,
|
|
OmniDeliveryManager,
|
|
)
|
|
|
|
|
|
@dataclass
|
|
class _NativeRequestState:
|
|
request_id: str
|
|
external_req_id: str
|
|
prompt_token_ids: list[int]
|
|
additional_information: Any = None
|
|
sampling_params: Any = None
|
|
num_computed_tokens: int = 0
|
|
resumable: bool = False
|
|
output_token_ids: list[int] = field(default_factory=list)
|
|
finished: bool = False
|
|
|
|
def snapshot(self, *, include_token_history: bool) -> SimpleNamespace:
|
|
prompt = list(self.prompt_token_ids) if include_token_history else []
|
|
output = list(self.output_token_ids) if include_token_history else []
|
|
finished = self.finished
|
|
request = SimpleNamespace(
|
|
request_id=self.request_id,
|
|
req_id=self.request_id,
|
|
external_req_id=self.external_req_id,
|
|
prompt_token_ids=prompt,
|
|
output_token_ids=output,
|
|
all_token_ids=prompt + output,
|
|
output_token_count=len(self.output_token_ids),
|
|
last_output_token_id=self.output_token_ids[-1] if self.output_token_ids else None,
|
|
additional_information=self.additional_information,
|
|
sampling_params=self.sampling_params,
|
|
num_computed_tokens=self.num_computed_tokens,
|
|
resumable=self.resumable,
|
|
)
|
|
request.is_finished = lambda: finished
|
|
return request
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class _NativeOutputBatch:
|
|
req_ids: list[str]
|
|
inter_stage_outputs: list[Any | None] | None
|
|
sampled_token_ids: list[list[int]] | None
|
|
|
|
|
|
_NATIVE_OUTPUT_STOP = object()
|
|
_NATIVE_OUTPUT_QUEUE_DEPTH = 8
|
|
_DEFAULT_DELIVERY_TIMEOUT_S = 30.0
|
|
_DEFAULT_SHUTDOWN_TIMEOUT_S = 5.0
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
class OmniRunnerDataPlane(OmniConnectorModelRunnerMixin):
|
|
"""MRv2-owned stage transport and request-side payload state."""
|
|
|
|
def __init__(self, vllm_config: Any, model_config: Any) -> None:
|
|
self._native_requests: dict[str, _NativeRequestState] = {}
|
|
self._native_output_lock = threading.RLock()
|
|
self._native_send_lock = threading.Lock()
|
|
self._native_outputs_in_flight: dict[str, int] = defaultdict(int)
|
|
self._native_terminal_pending: set[str] = set()
|
|
# Stage payload builders read the stage config through
|
|
# ``_get_model_config()`` (e.g. Qwen3-Omni's Thinker reads
|
|
# ``hf_config.talker_config.accept_hidden_layer``).
|
|
self.vllm_config = vllm_config
|
|
self.model_config = model_config
|
|
self.init_omni_connectors(model_config=model_config)
|
|
self._delivery_manager = OmniDeliveryManager(
|
|
delivery_timeout_s=self._connector_delivery_timeout(),
|
|
shutdown_timeout_s=_DEFAULT_SHUTDOWN_TIMEOUT_S,
|
|
)
|
|
self._can_send = (
|
|
self._custom_process_func is not None
|
|
if self._async_chunk
|
|
else should_accumulate_full_payload_output(model_config, self._custom_process_func)
|
|
)
|
|
self._start_output_worker(max_pending_batches=_NATIVE_OUTPUT_QUEUE_DEPTH)
|
|
|
|
def _connector_delivery_timeout(self) -> float:
|
|
config = getattr(getattr(self, "_omni_connector", None), "config", None)
|
|
extra = config.get("extra", {}) if isinstance(config, dict) else {}
|
|
value = extra.get("delivery_timeout_s", _DEFAULT_DELIVERY_TIMEOUT_S)
|
|
try:
|
|
timeout = float(value)
|
|
except (TypeError, ValueError) as error:
|
|
raise ValueError(f"invalid connector delivery_timeout_s: {value!r}") from error
|
|
if timeout <= 0:
|
|
raise ValueError("connector delivery_timeout_s must be positive")
|
|
return timeout
|
|
|
|
def _create_send_completion(self, *, request_id: str, put_key: str) -> Any:
|
|
return self._delivery_manager.create_ticket(
|
|
request_id=request_id,
|
|
put_key=put_key,
|
|
)
|
|
|
|
def _start_output_worker(self, *, max_pending_batches: int) -> None:
|
|
if max_pending_batches <= 0:
|
|
raise ValueError("max_pending_batches must be positive")
|
|
self._native_output_queue: Queue[_NativeOutputBatch | object] = Queue(maxsize=max_pending_batches)
|
|
self._native_output_error_lock = threading.Lock()
|
|
self._native_output_error: BaseException | None = None
|
|
self._connector_send_error_sink = self._record_output_error
|
|
self._native_output_closed = False
|
|
stage_id = getattr(self, "_stage_id", "unknown")
|
|
self._native_output_worker: threading.Thread | None = threading.Thread(
|
|
target=self._output_worker_loop,
|
|
name=f"omni-native-output-{stage_id}",
|
|
daemon=True,
|
|
)
|
|
self._native_output_worker.start()
|
|
|
|
def _record_output_error(self, error: BaseException) -> None:
|
|
first_error = False
|
|
with self._native_output_error_lock:
|
|
if self._native_output_error is None:
|
|
self._native_output_error = error
|
|
first_error = True
|
|
if first_error:
|
|
manager = getattr(self, "_delivery_manager", None)
|
|
if manager is not None:
|
|
manager.quarantine(error)
|
|
|
|
def _raise_output_error(self) -> None:
|
|
with self._native_output_error_lock:
|
|
error = self._native_output_error
|
|
if error is not None:
|
|
raise RuntimeError(f"native output worker failed: {error}") from error
|
|
|
|
def _output_worker_loop(self) -> None:
|
|
while True:
|
|
batch = self._native_output_queue.get()
|
|
try:
|
|
if batch is _NATIVE_OUTPUT_STOP:
|
|
return
|
|
if self._native_output_error is not None:
|
|
continue
|
|
assert isinstance(batch, _NativeOutputBatch)
|
|
self.complete_outputs(
|
|
req_ids=batch.req_ids,
|
|
inter_stage_outputs=batch.inter_stage_outputs,
|
|
sampled_token_ids=batch.sampled_token_ids,
|
|
)
|
|
except BaseException as error:
|
|
self._record_output_error(error)
|
|
finally:
|
|
self._native_output_queue.task_done()
|
|
|
|
def enqueue_outputs(
|
|
self,
|
|
*,
|
|
req_ids: list[str],
|
|
inter_stage_outputs: list[Any | None] | None,
|
|
sampled_token_ids: list[list[int]] | None,
|
|
) -> None:
|
|
"""Transfer one completed model batch to the ordered output worker."""
|
|
self._raise_output_error()
|
|
if self._native_output_closed:
|
|
raise RuntimeError("native output worker is closed")
|
|
batch = _NativeOutputBatch(
|
|
req_ids=req_ids,
|
|
inter_stage_outputs=inter_stage_outputs,
|
|
sampled_token_ids=sampled_token_ids,
|
|
)
|
|
while True:
|
|
try:
|
|
self._native_output_queue.put(batch, timeout=0.1)
|
|
break
|
|
except Full:
|
|
self._raise_output_error()
|
|
self._raise_output_error()
|
|
|
|
def drain_outputs(self) -> None:
|
|
queue = self._native_output_queue
|
|
pending_batches = max(1, queue.qsize() + 1)
|
|
timeout_s = self._delivery_manager.delivery_timeout_s * pending_batches
|
|
deadline = time.monotonic() + timeout_s
|
|
while queue.unfinished_tasks:
|
|
self._raise_output_error()
|
|
remaining = deadline - time.monotonic()
|
|
if remaining <= 0:
|
|
error = DeliveryTimeoutError(
|
|
"native output queue drain timed out after "
|
|
f"{timeout_s:.3f}s with {queue.unfinished_tasks} "
|
|
"unfinished batch(es)"
|
|
)
|
|
self._record_output_error(error)
|
|
self._raise_output_error()
|
|
with queue.all_tasks_done:
|
|
if queue.unfinished_tasks:
|
|
queue.all_tasks_done.wait(timeout=min(0.05, remaining))
|
|
self._raise_output_error()
|
|
|
|
def _stop_output_worker(self) -> None:
|
|
worker = getattr(self, "_native_output_worker", None)
|
|
if worker is None:
|
|
return
|
|
self._native_output_closed = True
|
|
self._delivery_manager.shutdown(RuntimeError("MRv2 output worker shutdown"))
|
|
deadline = time.monotonic() + self._delivery_manager.shutdown_timeout_s
|
|
while worker.is_alive():
|
|
remaining = deadline - time.monotonic()
|
|
if remaining <= 0:
|
|
break
|
|
try:
|
|
self._native_output_queue.put(
|
|
_NATIVE_OUTPUT_STOP,
|
|
timeout=min(0.05, remaining),
|
|
)
|
|
break
|
|
except Full:
|
|
continue
|
|
worker.join(timeout=max(0.0, deadline - time.monotonic()))
|
|
if worker.is_alive():
|
|
raise RuntimeError("MRv2 output worker did not stop within its shutdown deadline")
|
|
self._native_output_worker = None
|
|
|
|
def register_request(self, request_data: Any) -> None:
|
|
self._raise_output_error()
|
|
req_id = str(request_data.req_id)
|
|
external_req_id = str(getattr(request_data, "external_req_id", None) or req_id)
|
|
with self._native_output_lock:
|
|
self._native_outputs_in_flight.pop(req_id, None)
|
|
self._native_terminal_pending.discard(req_id)
|
|
self._native_requests[req_id] = _NativeRequestState(
|
|
request_id=req_id,
|
|
external_req_id=external_req_id,
|
|
prompt_token_ids=list(getattr(request_data, "prompt_token_ids", None) or []),
|
|
additional_information=getattr(request_data, "additional_information", None),
|
|
sampling_params=getattr(request_data, "sampling_params", None),
|
|
num_computed_tokens=int(getattr(request_data, "num_computed_tokens", 0) or 0),
|
|
resumable=bool(getattr(request_data, "resumable", False)),
|
|
)
|
|
|
|
def reserve_outputs(self, req_ids: list[str]) -> None:
|
|
"""Keep request state alive until deferred runner outputs are consumed."""
|
|
self._raise_output_error()
|
|
with self._native_output_lock:
|
|
for req_id in dict.fromkeys(req_ids):
|
|
if req_id in self._native_requests:
|
|
self._native_outputs_in_flight[req_id] += 1
|
|
|
|
def _ready_terminal_requests(self) -> set[str]:
|
|
return {
|
|
req_id for req_id in self._native_terminal_pending if self._native_outputs_in_flight.get(req_id, 0) == 0
|
|
}
|
|
|
|
def _emit_ready_terminals(self) -> int:
|
|
ready = self._ready_terminal_requests()
|
|
if not ready:
|
|
return 0
|
|
emitted = self.emit_chunks(
|
|
req_ids=[],
|
|
inter_stage_outputs=None,
|
|
sampled_token_ids=None,
|
|
terminal_req_ids=ready,
|
|
)
|
|
self._native_terminal_pending.difference_update(ready)
|
|
for req_id in ready:
|
|
self._native_outputs_in_flight.pop(req_id, None)
|
|
return emitted
|
|
|
|
def request_terminal(self, req_ids: set[str]) -> int:
|
|
"""Emit terminal markers only after all deferred outputs are enqueued."""
|
|
self._raise_output_error()
|
|
with self._native_output_lock:
|
|
self._native_terminal_pending.update(req_id for req_id in req_ids if req_id in self._native_requests)
|
|
return self._emit_ready_terminals()
|
|
|
|
def abort_requests(self, req_ids: set[str]) -> int:
|
|
"""Cancel deferred outputs and terminate each live request once."""
|
|
with self._native_output_lock:
|
|
active_req_ids = {req_id for req_id in req_ids if req_id in self._native_requests}
|
|
# Receivers are registered before requests acquire model slots.
|
|
# Cancelling during that interval must still stop recv polling.
|
|
for req_id in req_ids - active_req_ids:
|
|
self.cleanup_finished_request(req_id)
|
|
if not active_req_ids:
|
|
return 0
|
|
for req_id in active_req_ids:
|
|
self._native_outputs_in_flight.pop(req_id, None)
|
|
self._native_terminal_pending.discard(req_id)
|
|
# Cancellation must not flush a partially accumulated utterance.
|
|
self._pending_full_payload_send.pop(req_id, None)
|
|
return self.emit_chunks(
|
|
req_ids=[],
|
|
inter_stage_outputs=None,
|
|
sampled_token_ids=None,
|
|
terminal_req_ids=active_req_ids,
|
|
)
|
|
|
|
def complete_outputs(
|
|
self,
|
|
*,
|
|
req_ids: list[str],
|
|
inter_stage_outputs: list[Any | None] | None,
|
|
sampled_token_ids: list[list[int]] | None,
|
|
) -> int:
|
|
"""Commit one deferred output batch, then release its lifecycle holds."""
|
|
send_lock_held = False
|
|
with self._native_output_lock:
|
|
entries, finished_req_ids = self._build_chunk_entries(
|
|
req_ids=req_ids,
|
|
inter_stage_outputs=inter_stage_outputs,
|
|
sampled_token_ids=sampled_token_ids,
|
|
terminal_req_ids=set(),
|
|
)
|
|
# Claim connector ordering before publishing the updated request
|
|
# state. request_terminal() can still register a pending terminal
|
|
# while the enqueue runs, but abort cannot overtake this output.
|
|
if entries:
|
|
self._native_send_lock.acquire()
|
|
send_lock_held = True
|
|
try:
|
|
emitted = self._send_chunk_entries(
|
|
entries,
|
|
send_lock_held=send_lock_held,
|
|
)
|
|
finally:
|
|
if send_lock_held:
|
|
self._native_send_lock.release()
|
|
with self._native_output_lock:
|
|
self._cleanup_finished_entries(finished_req_ids)
|
|
for req_id in dict.fromkeys(req_ids):
|
|
in_flight = self._native_outputs_in_flight.get(req_id, 0)
|
|
if in_flight > 1:
|
|
self._native_outputs_in_flight[req_id] = in_flight - 1
|
|
elif in_flight == 1:
|
|
self._native_outputs_in_flight.pop(req_id, None)
|
|
return emitted + self._emit_ready_terminals()
|
|
|
|
def register_receivers(self, handles: list[Any]) -> None:
|
|
for handle in handles:
|
|
self.register_chunk_recv(handle)
|
|
|
|
@classmethod
|
|
def _copy_payload_structure(cls, value: Any) -> Any:
|
|
if isinstance(value, dict):
|
|
return {key: cls._copy_payload_structure(item) for key, item in value.items()}
|
|
if isinstance(value, list):
|
|
return list(value)
|
|
return value
|
|
|
|
def _accumulate_payload(self, req_id: str, payload_data: OmniPayload) -> OmniPayload:
|
|
"""Accumulate chunk payloads (concat tensors, extend lists)."""
|
|
if req_id not in self._send_side_request_payload:
|
|
self._send_side_request_payload[req_id] = dict(payload_data)
|
|
return cast(OmniPayload, dict(self._send_side_request_payload[req_id]))
|
|
|
|
origin = self._send_side_request_payload[req_id]
|
|
merged = dict(origin)
|
|
raw_ok = payload_data.get("meta", {}).get("override_keys", []) if isinstance(payload_data, dict) else []
|
|
override_keys = {tuple(k) if isinstance(k, list) else k for k in raw_ok}
|
|
|
|
for key, value in payload_data.items():
|
|
if isinstance(value, dict):
|
|
origin_sub = origin.get(key)
|
|
merged_sub = dict(origin_sub) if isinstance(origin_sub, dict) else {}
|
|
span_handled: set[str] = set()
|
|
if key == "embed" and isinstance(origin_sub, dict):
|
|
for tk, sk, ek in (("decode", "decode_token_start", "decode_token_end"),):
|
|
if tk not in value or (key, tk) in override_keys:
|
|
continue
|
|
span = merge_tensor_spans(
|
|
get_tensor_span(origin_sub, tensor_key=tk, start_key=sk, end_key=ek),
|
|
get_tensor_span(value, tensor_key=tk, start_key=sk, end_key=ek),
|
|
)
|
|
if span is None:
|
|
continue
|
|
t, s, e = span
|
|
merged_sub[tk] = t
|
|
merged_sub[sk] = s
|
|
merged_sub[ek] = e
|
|
span_handled |= {tk, sk, ek}
|
|
for qual, qval in value.items():
|
|
if qual in span_handled:
|
|
continue
|
|
if key == "meta":
|
|
# Per-chunk state (``finished``, the V1 sender's
|
|
# ``is_segment_finished`` flag, ...), not rows.
|
|
merged_sub[qual] = qval
|
|
continue
|
|
if (key, qual) in override_keys:
|
|
merged_sub[qual] = qval
|
|
continue
|
|
osv = merged_sub.get(qual)
|
|
if isinstance(qval, torch.Tensor) and isinstance(osv, torch.Tensor):
|
|
merged_sub[qual] = torch.cat([osv, qval], dim=0)
|
|
elif isinstance(qval, list) and isinstance(osv, list):
|
|
merged_sub[qual] = osv + qval
|
|
else:
|
|
merged_sub[qual] = qval
|
|
merged[key] = merged_sub
|
|
else:
|
|
if key in override_keys:
|
|
merged[key] = value
|
|
continue
|
|
ov = origin.get(key)
|
|
if isinstance(value, torch.Tensor) and isinstance(ov, torch.Tensor):
|
|
merged[key] = torch.cat([ov, value], dim=0)
|
|
elif isinstance(value, list) and isinstance(ov, list):
|
|
merged[key] = ov + value
|
|
else:
|
|
merged[key] = value
|
|
|
|
self._send_side_request_payload[req_id] = merged
|
|
return cast(OmniPayload, dict(merged))
|
|
|
|
def pop_local_stage_payload(self, req_id: str) -> Any:
|
|
"""Hand one accumulated delta to the model and acknowledge its rows.
|
|
|
|
Connector accumulation and ``model_intermediate_buffer`` are separate
|
|
ownership domains. Once decode rows are handed to the model, remove
|
|
them from connector accumulation so the next chunk cannot replay the
|
|
same rows into Talker's ``cached_decode``.
|
|
"""
|
|
with self._lock:
|
|
payload = self._local_stage_payload_cache.pop(req_id, None)
|
|
if payload is None:
|
|
return None
|
|
delivered = self._copy_payload_structure(payload)
|
|
external_req_id = self._request_ids_mapping.get(req_id, req_id)
|
|
accumulated = self._send_side_request_payload.get(external_req_id)
|
|
if isinstance(accumulated, dict):
|
|
embed = accumulated.get("embed")
|
|
if isinstance(embed, dict):
|
|
embed.pop("decode", None)
|
|
embed.pop("decode_token_start", None)
|
|
embed.pop("decode_token_end", None)
|
|
hidden_states = accumulated.get("hidden_states")
|
|
if isinstance(hidden_states, dict):
|
|
hidden_states.pop("output", None)
|
|
ids = accumulated.get("ids")
|
|
if isinstance(ids, dict):
|
|
ids.pop("output", None)
|
|
return delivered
|
|
|
|
def emit_chunks(
|
|
self,
|
|
*,
|
|
req_ids: list[str],
|
|
inter_stage_outputs: list[Any | None] | None,
|
|
sampled_token_ids: list[list[int]] | None,
|
|
terminal_req_ids: set[str],
|
|
) -> int:
|
|
with self._native_output_lock:
|
|
entries, finished_req_ids = self._build_chunk_entries(
|
|
req_ids=req_ids,
|
|
inter_stage_outputs=inter_stage_outputs,
|
|
sampled_token_ids=sampled_token_ids,
|
|
terminal_req_ids=terminal_req_ids,
|
|
)
|
|
emitted = self._send_chunk_entries(entries)
|
|
self._cleanup_finished_entries(finished_req_ids)
|
|
return emitted
|
|
|
|
def _build_chunk_entries(
|
|
self,
|
|
*,
|
|
req_ids: list[str],
|
|
inter_stage_outputs: list[Any | None] | None,
|
|
sampled_token_ids: list[list[int]] | None,
|
|
terminal_req_ids: set[str],
|
|
) -> tuple[list[tuple[Any, Any | None]], list[str]]:
|
|
if not getattr(self, "_can_send", True):
|
|
receive_finished_ids = [req_id for req_id in sorted(terminal_req_ids) if req_id in self._native_requests]
|
|
return [], receive_finished_ids
|
|
payload_by_req = {
|
|
req_id: inter_stage_outputs[index]
|
|
for index, req_id in enumerate(req_ids)
|
|
if inter_stage_outputs is not None and index < len(inter_stage_outputs)
|
|
}
|
|
sampled_by_req = {
|
|
req_id: sampled_token_ids[index]
|
|
for index, req_id in enumerate(req_ids)
|
|
if sampled_token_ids is not None and index < len(sampled_token_ids)
|
|
}
|
|
entries: list[tuple[Any, Any | None]] = []
|
|
finished_req_ids: list[str] = []
|
|
for req_id in dict.fromkeys([*req_ids, *sorted(terminal_req_ids)]):
|
|
state = self._native_requests.get(req_id)
|
|
if state is None:
|
|
continue
|
|
state.output_token_ids.extend(int(token_id) for token_id in sampled_by_req.get(req_id, []))
|
|
state.finished = req_id in terminal_req_ids
|
|
payload = payload_by_req.get(req_id)
|
|
if not self._async_chunk:
|
|
# Full-payload consumers execute once, after the producer has
|
|
# finished. Reuse V1's concat/replace contract for flat outputs
|
|
# and retain the final token history for codec-row cropping.
|
|
if isinstance(payload, dict):
|
|
self.accumulate_full_payload_output(req_id, flatten_payload(payload), state)
|
|
if not state.finished:
|
|
continue
|
|
accumulated = self._pending_full_payload_send.pop(req_id, None)
|
|
payload = self._materialize_full_payload_entry(accumulated)[0] if accumulated is not None else None
|
|
if self._async_chunk and isinstance(payload, dict):
|
|
payload = unflatten_payload(payload)
|
|
if payload is None and not state.finished:
|
|
continue
|
|
include_token_history = bool(
|
|
not self._async_chunk or state.resumable or self._put_req_chunk.get(state.external_req_id, 0) == 0
|
|
)
|
|
entries.append(
|
|
(
|
|
state.snapshot(include_token_history=include_token_history),
|
|
payload,
|
|
)
|
|
)
|
|
if state.finished:
|
|
finished_req_ids.append(req_id)
|
|
return entries, finished_req_ids
|
|
|
|
def _send_chunk_entries(
|
|
self,
|
|
entries: list[tuple[Any, Any | None]],
|
|
*,
|
|
send_lock_held: bool = False,
|
|
) -> int:
|
|
if not entries:
|
|
return 0
|
|
context: AbstractContextManager
|
|
if send_lock_held:
|
|
context = nullcontext()
|
|
else:
|
|
context = self._native_send_lock
|
|
with context:
|
|
return self.send_chunks(entries, propagate_errors=True, wait_for_delivery=True)
|
|
|
|
def _cleanup_finished_entries(self, finished_req_ids: list[str]) -> None:
|
|
for req_id in finished_req_ids:
|
|
self._native_requests.pop(req_id, None)
|
|
self.cleanup_finished_request(req_id)
|
|
|
|
def shutdown_omni_connectors(self) -> None:
|
|
"""Bound MRv2 shutdown even if a connector call cannot be cancelled.
|
|
|
|
Connector ``close()`` runs concurrently with the I/O threads so a
|
|
backend that supports cancellation can release a blocked ``put()``.
|
|
Backends that do not support it are quarantined and left only on
|
|
daemon threads after the bounded deadline; V1 keeps its existing
|
|
shutdown implementation in the shared mixin.
|
|
"""
|
|
timeout_s = self._delivery_manager.shutdown_timeout_s
|
|
deadline = time.monotonic() + timeout_s
|
|
self._delivery_manager.shutdown(RuntimeError("MRv2 connector shutdown"))
|
|
self._stop_event.set()
|
|
self._work_available.set()
|
|
self._save_work_available.set()
|
|
|
|
close_errors: list[BaseException] = []
|
|
connector = getattr(self, "_omni_connector", None)
|
|
|
|
def close_connector() -> None:
|
|
if connector is None:
|
|
return
|
|
try:
|
|
connector.close()
|
|
except BaseException as error:
|
|
close_errors.append(error)
|
|
|
|
close_thread = threading.Thread(
|
|
target=close_connector,
|
|
name=f"omni-native-connector-close-{getattr(self, '_stage_id', 'unknown')}",
|
|
daemon=True,
|
|
)
|
|
close_thread.start()
|
|
|
|
live_threads: list[str] = []
|
|
for name, thread in (
|
|
("recv", getattr(self, "_recv_thread", None)),
|
|
("save", getattr(self, "_save_thread", None)),
|
|
("connector-close", close_thread),
|
|
):
|
|
if thread is None:
|
|
continue
|
|
thread.join(timeout=max(0.0, deadline - time.monotonic()))
|
|
if thread.is_alive():
|
|
live_threads.append(name)
|
|
|
|
if close_errors:
|
|
raise RuntimeError("MRv2 connector close failed") from close_errors[0]
|
|
if live_threads:
|
|
raise RuntimeError(
|
|
f"MRv2 connector shutdown exceeded {timeout_s:.3f}s; blocked thread(s): {', '.join(live_threads)}"
|
|
)
|
|
|
|
def close(self) -> None:
|
|
error: BaseException | None = None
|
|
try:
|
|
self.drain_outputs()
|
|
except BaseException as exc:
|
|
error = exc
|
|
finally:
|
|
self._delivery_manager.shutdown(error or RuntimeError("MRv2 data plane close"))
|
|
try:
|
|
self._stop_output_worker()
|
|
except BaseException as exc:
|
|
if error is None:
|
|
error = exc
|
|
try:
|
|
self.shutdown_omni_connectors()
|
|
except BaseException as exc:
|
|
if error is None:
|
|
error = exc
|
|
if error is not None:
|
|
raise error
|