260 lines
8.4 KiB
Python
260 lines
8.4 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
"""Model-neutral value types shared by the duplex engine components.
|
|
|
|
Immutable DTOs plus the ``DuplexStagePort`` base class that ``DuplexOrchestrator``
|
|
implements for the session manager/runner.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import base64
|
|
from abc import ABC, abstractmethod
|
|
from collections.abc import Iterable, Mapping
|
|
from dataclasses import dataclass, field
|
|
from enum import Enum
|
|
from types import MappingProxyType
|
|
|
|
import regex as re
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class DuplexFence:
|
|
"""Engine-internal session identity used for stage request ids and stale filtering."""
|
|
|
|
session_id: str
|
|
epoch: int = 0
|
|
turn_id: int = 0
|
|
|
|
|
|
class DuplexOutputAction(str, Enum):
|
|
DIRECT_RESPONSE = "direct_response"
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class DuplexAppendPlan:
|
|
prompt: dict[str, object]
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class DuplexOutputDecision:
|
|
action: DuplexOutputAction
|
|
metadata: Mapping[str, object] = field(default_factory=dict)
|
|
final_output_type: str = "text"
|
|
|
|
def __post_init__(self) -> None:
|
|
object.__setattr__(self, "metadata", MappingProxyType(dict(self.metadata)))
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class DuplexRequestIdentity:
|
|
session_id: str
|
|
fence: DuplexFence
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class DuplexStageRequestContext:
|
|
request_id: str
|
|
session_id: str
|
|
fence: DuplexFence
|
|
stage_id: int
|
|
final_stage_id: int
|
|
config_generation: int
|
|
sampling_params: tuple[object, ...]
|
|
session_config: Mapping[str, object] = field(default_factory=dict)
|
|
runtime_config: Mapping[str, object] = field(default_factory=dict)
|
|
|
|
def __post_init__(self) -> None:
|
|
object.__setattr__(self, "sampling_params", tuple(self.sampling_params))
|
|
object.__setattr__(self, "session_config", MappingProxyType(dict(self.session_config)))
|
|
object.__setattr__(self, "runtime_config", MappingProxyType(dict(self.runtime_config)))
|
|
|
|
@property
|
|
def stage_sampling_params(self) -> object:
|
|
return self.sampling_params[self.stage_id]
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class DuplexStageSubmission:
|
|
context: DuplexStageRequestContext
|
|
prompt: Mapping[str, object]
|
|
already_submitted: bool
|
|
# True: resume/update an existing stage0 id. False: open a new ephemeral id.
|
|
# Distinct from DuplexCapabilities.supports_core_resumable_request.
|
|
resumable: bool = True
|
|
|
|
def __post_init__(self) -> None:
|
|
object.__setattr__(self, "prompt", MappingProxyType(dict(self.prompt)))
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class DuplexStageSubmissionResult:
|
|
request_id: str
|
|
stage_id: int
|
|
replica_id: int
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class DuplexOutputContext:
|
|
identity: DuplexRequestIdentity
|
|
final_stage_id: int
|
|
segment_finished: bool
|
|
segment_token_ids: tuple[int, ...] = ()
|
|
segment_output_metadata: Mapping[str, object] = field(default_factory=dict)
|
|
|
|
def __post_init__(self) -> None:
|
|
object.__setattr__(self, "segment_token_ids", tuple(self.segment_token_ids))
|
|
object.__setattr__(
|
|
self,
|
|
"segment_output_metadata",
|
|
MappingProxyType(dict(self.segment_output_metadata)),
|
|
)
|
|
|
|
|
|
class DuplexStagePort(ABC):
|
|
"""Narrow stage-management surface the session runner/manager use (implemented by DuplexOrchestrator)."""
|
|
|
|
@property
|
|
@abstractmethod
|
|
def stage_count(self) -> int: ...
|
|
|
|
@abstractmethod
|
|
def sampling_defaults(self) -> tuple[object, ...]: ...
|
|
|
|
@abstractmethod
|
|
def ensure_request(self, context: DuplexStageRequestContext) -> None: ...
|
|
|
|
@abstractmethod
|
|
async def submit(self, submission: DuplexStageSubmission) -> DuplexStageSubmissionResult: ...
|
|
|
|
@abstractmethod
|
|
async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None: ...
|
|
|
|
@abstractmethod
|
|
async def abort_requests(self, request_ids: list[str]) -> None: ...
|
|
|
|
|
|
def duplex_data_plane_request_info(result: dict[str, object]) -> tuple[str | None, int | None]:
|
|
stage_results = result.get("stage_results")
|
|
if not isinstance(stage_results, list):
|
|
return None, None
|
|
for item in stage_results:
|
|
if not isinstance(item, dict):
|
|
continue
|
|
inner = item.get("result")
|
|
if not isinstance(inner, dict) or inner.get("data_plane_append") is not True:
|
|
continue
|
|
request_id = inner.get("request_id")
|
|
if not isinstance(request_id, str) or not request_id:
|
|
continue
|
|
response_stage_id = inner.get("response_stage_id")
|
|
return request_id, response_stage_id if isinstance(response_stage_id, int) else None
|
|
return None, None
|
|
|
|
|
|
def duplex_resource_request_id(fence: DuplexFence, role: str) -> str:
|
|
if not role or any(
|
|
character not in "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_-" for character in role
|
|
):
|
|
raise ValueError(f"invalid duplex resource role: {role!r}")
|
|
encoded_session_id = base64.urlsafe_b64encode(fence.session_id.encode("utf-8")).decode("ascii").rstrip("=")
|
|
return f"duplex-s.{encoded_session_id}.e.{fence.epoch}.r.{role}"
|
|
|
|
|
|
def duplex_ephemeral_stage_request_id(fence: DuplexFence, *, stage_id: int) -> str:
|
|
"""Turn-scoped Stage request id for non-resumable (ephemeral) duplex models."""
|
|
return duplex_resource_request_id(fence, f"stage{stage_id}-turn{fence.turn_id}")
|
|
|
|
|
|
_EPHEMERAL_TURN_IN_REQUEST_ID = re.compile(r"\.r\.stage\d+-turn(\d+)$")
|
|
|
|
|
|
def duplex_turn_id_from_request_id(request_id: str | None) -> int | None:
|
|
"""Parse ephemeral ``…r.stage{N}-turn{T}`` ids."""
|
|
if not isinstance(request_id, str):
|
|
return None
|
|
match = _EPHEMERAL_TURN_IN_REQUEST_ID.search(request_id)
|
|
return int(match.group(1)) if match else None
|
|
|
|
|
|
def duplex_same_turn_request_ids(request_id: str, candidate_ids: Iterable[str]) -> list[str]:
|
|
"""Other ephemeral stage ids from the same session, epoch, and turn.
|
|
|
|
``request_id`` itself is not included. Non-ephemeral ids are ignored.
|
|
"""
|
|
turn_id = duplex_turn_id_from_request_id(request_id)
|
|
if turn_id is None or ".r." not in request_id:
|
|
return []
|
|
prefix = request_id.rsplit(".r.", 1)[0] + ".r."
|
|
return [
|
|
candidate
|
|
for candidate in candidate_ids
|
|
if candidate != request_id
|
|
and isinstance(candidate, str)
|
|
and candidate.startswith(prefix)
|
|
and duplex_turn_id_from_request_id(candidate) == turn_id
|
|
]
|
|
|
|
|
|
def _duplex_resource_request_fields(request_id: str | None) -> tuple[str, str, str] | None:
|
|
"""Split ``duplex-s.<b64url>.e.<epoch>.r.<role>`` or return None."""
|
|
if not isinstance(request_id, str):
|
|
return None
|
|
parts = request_id.split(".")
|
|
if len(parts) != 6 or parts[0] != "duplex-s" or parts[2] != "e" or parts[4] != "r":
|
|
return None
|
|
return parts[1], parts[3], parts[5]
|
|
|
|
|
|
def duplex_session_id_from_request_id(request_id: str | None) -> str | None:
|
|
"""Decode the session id from ``duplex-s.<b64url>.e.<epoch>.r.<role>``."""
|
|
fields = _duplex_resource_request_fields(request_id)
|
|
if fields is None:
|
|
return None
|
|
encoded, _, _ = fields
|
|
pad = "=" * (-len(encoded) % 4)
|
|
try:
|
|
return base64.urlsafe_b64decode(encoded + pad).decode("utf-8")
|
|
except (ValueError, UnicodeDecodeError):
|
|
return None
|
|
|
|
|
|
def duplex_resource_request_belongs_to_session(request_id: str, session_id: str) -> bool:
|
|
"""Return whether a current-format resource request belongs to a session."""
|
|
fields = _duplex_resource_request_fields(request_id)
|
|
if fields is None:
|
|
return False
|
|
encoded, epoch, role = fields
|
|
try:
|
|
int(epoch)
|
|
except ValueError:
|
|
return False
|
|
if not role or any(
|
|
character not in "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_-" for character in role
|
|
):
|
|
return False
|
|
encoded_session_id = base64.urlsafe_b64encode(session_id.encode("utf-8")).decode("ascii").rstrip("=")
|
|
return encoded == encoded_session_id
|
|
|
|
|
|
__all__ = [
|
|
"DuplexFence",
|
|
"DuplexAppendPlan",
|
|
"DuplexOutputAction",
|
|
"DuplexOutputContext",
|
|
"DuplexOutputDecision",
|
|
"DuplexRequestIdentity",
|
|
"DuplexStagePort",
|
|
"DuplexStageRequestContext",
|
|
"DuplexStageSubmission",
|
|
"DuplexStageSubmissionResult",
|
|
"duplex_data_plane_request_info",
|
|
"duplex_ephemeral_stage_request_id",
|
|
"duplex_turn_id_from_request_id",
|
|
"duplex_resource_request_belongs_to_session",
|
|
"duplex_resource_request_id",
|
|
"duplex_same_turn_request_ids",
|
|
"duplex_session_id_from_request_id",
|
|
]
|