vllm_omni.engine.cfg_companion_tracker ¶

CFG companion request tracker for the Omni orchestrator.

Encapsulates all bookkeeping for Classifier-Free Guidance companion requests (parent/companion ID mapping, completion tracking, deferred forwarding, and cleanup).

logger `module-attribute` ¶

logger = logging.getLogger(__name__)

CfgCompanionTracker ¶

Manages CFG companion request lifecycle in the orchestrator scheduling loop.

abort_parents ¶

abort_parents(request_ids: list[str]) -> list[str]

all_companions_done ¶

all_companions_done(parent_id: str) -> bool

cleanup_parent ¶

cleanup_parent(parent_id: str) -> list[str]

defer_parent ¶

defer_parent(
    parent_id: str, engine_outputs: Any, stage_id: int
) -> None

Hold parent result while waiting for companions to finish.

Deferred parents are released when their companions' processed outputs arrive, and failed by the orchestrator when a companion errors or is aborted. TODO: a deadline for deferred parents and detection of a companion replica dying silently (no error output, no abort) still need a liveness owner in the orchestrator loop.

get_companion_outputs ¶

get_companion_outputs(parent_id: str) -> list[Any]

Return companion outputs (role-registration order) for bundling.

Non-destructive: outputs stay stashed until cleanup_parent so a re-submission of the same parent (e.g. a streaming terminal update) bundles the complete set again instead of an empty one.

get_companion_request_ids ¶

get_companion_request_ids(parent_id: str) -> dict[str, str]

Return {role: companion_request_id} for a parent.

get_parent_id ¶

get_parent_id(req_id: str) -> str | None

Return the parent request id for a companion, or None.

has_companions ¶

has_companions(parent_id: str) -> bool

is_companion ¶

is_companion(req_id: str) -> bool

is_companion_done ¶

is_companion_done(companion_id: str) -> bool

True once this companion's processed output was stashed and marked done.

on_companion_completed ¶

on_companion_completed(companion_id: str) -> str | None

Mark done. Returns parent_id only if parent is pending and all companions finished.

pop_pending_parent ¶

pop_pending_parent(parent_id: str) -> dict[str, Any] | None

register_companion ¶

register_companion(
    parent_id: str, role: str, companion_id: str
) -> None

register_parent ¶

register_parent(parent_id: str) -> None

set_companion_output ¶

set_companion_output(
    companion_id: str, output: Any
) -> None

Stash companion engine output for the parent to bundle at forward time.