Ken SimpsonandClaude Opus 5.5
dba3a6abd9
* Expect resolved STRING type for ModalMapInput boundaries
69a33e3 made _remote_output_io_type resolve wildcard "*" outputs through
their linked input, so a ModalMapInput boundary fed by a STRING list now
reports STRING, matching the mapped_input type the test already expects.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Add Remote Progressive Preview node
A local PreviewImage-style node that shows one empty slot per mapped item
when remote work starts and fills each slot as soon as that item's images
stream back, instead of waiting for ComfyUI to deliver the whole result.
- Planning records direct consumers as progressive_preview_target_node_ids
and exempts them from local bridge materializers.
- The remote shim streams boundary outputs for those targets; mapped item
payloads carry map_item_count and dispatch emits a start event.
- The host saves streamed images to temp, records them per prompt/node, and
sends modal_progressive_preview websocket events.
- The node reuses streamed slots on final execution and only materializes
bridge refs that were never streamed.
- web/progressive_preview.js renders the slot grid and hands off to the
native preview on the final executed event.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Execute every mapped item instead of reusing shared-session outputs
Mapped items share one remote session. The restored-session short-circuit,
meant for identical per-lane seed phases, looked outputs up by node id
only, so after the first item on a container every later item skipped
sampling and returned that earlier item's image under its own bridge key.
Payloads carrying map_item_index now always execute.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Keep finished progressive previews handed off to the native preview
Hide the panel with an inline display style so the frontend's DOM-widget
wrapper cannot override the hidden attribute, and ignore replayed events
for a prompt whose final images were already delivered.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---------
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
666 lines
24 KiB
Python
666 lines
24 KiB
Python
"""Local UI forwarding for streamed Modal payload events."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
import time
|
|
from dataclasses import dataclass, field
|
|
from typing import Any, Iterator
|
|
|
|
from ..serialization import (
|
|
coerce_serialized_node_outputs,
|
|
deserialize_value,
|
|
)
|
|
from .local_ui_events import (
|
|
_emit_local_executed_output,
|
|
_emit_local_modal_progress,
|
|
_emit_local_modal_status,
|
|
_emit_local_preview_boundary_output,
|
|
_emit_local_preview_image,
|
|
_emit_local_progressive_preview_slot,
|
|
_progress_stream_event_metadata,
|
|
_remote_execution_destination,
|
|
_remote_execution_identity,
|
|
_remote_prompt_ancestor_node_ids,
|
|
_should_forward_suppressed_stream_event,
|
|
)
|
|
from .mapped_execution import _remember_mapped_lane_node_id
|
|
from .modal_container_logs import (
|
|
_coerce_modal_task_id,
|
|
_is_remote_container_log_stream_enabled,
|
|
_release_remote_container_log_stream,
|
|
_retain_remote_container_log_stream,
|
|
)
|
|
from .modal_deployment import ModalRemoteInvocationError
|
|
from .modal_warmup import _schedule_speculative_affinity_prewarm
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def _close_remote_payload_stream(stream_events: Iterator[dict[str, Any]]) -> None:
|
|
"""Best-effort close one streamed Modal iterator after the terminal result arrives."""
|
|
close_callable = getattr(stream_events, "close", None)
|
|
if not callable(close_callable):
|
|
return
|
|
try:
|
|
close_result = close_callable()
|
|
except (AttributeError, RuntimeError, TypeError, ValueError) as exc:
|
|
logger.debug("Ignoring remote payload stream close failure: %s", exc)
|
|
return
|
|
if not asyncio.iscoroutine(close_result):
|
|
return
|
|
try:
|
|
asyncio.run(close_result)
|
|
except (RuntimeError, ValueError) as exc:
|
|
logger.debug("Ignoring async remote payload stream close failure: %s", exc)
|
|
|
|
|
|
|
|
|
|
|
|
@dataclass
|
|
class _PayloadStreamState:
|
|
"""Track local forwarding state for one streamed remote invocation."""
|
|
|
|
payload: dict[str, Any]
|
|
prompt_id: str | None
|
|
client_id: str | None
|
|
node_ids: list[str]
|
|
modal_gpu: str | None
|
|
suppress_status_stream: bool
|
|
component_id: str
|
|
invocation_id: str
|
|
stream_started_at: float
|
|
previous_event_at: float
|
|
result_payload: bytes | bytearray | None = None
|
|
suppressed_progress_metadata: dict[str, dict[str, str | None]] = field(
|
|
default_factory=dict
|
|
)
|
|
active_remote_task_id: str | None = None
|
|
active_remote_log_task_id: str | None = None
|
|
should_close_stream: bool = False
|
|
event_count: int = 0
|
|
progress_event_count: int = 0
|
|
speculative_prewarm_checked: bool = False
|
|
|
|
|
|
def _new_payload_stream_state(payload: dict[str, Any]) -> _PayloadStreamState:
|
|
"""Build initial stream-forwarding state from one remote payload."""
|
|
prompt_id = (
|
|
str(payload["prompt_id"]) if payload.get("prompt_id") is not None else None
|
|
)
|
|
extra_data = payload.get("extra_data") or {}
|
|
started_at = time.monotonic()
|
|
return _PayloadStreamState(
|
|
payload=payload,
|
|
prompt_id=prompt_id,
|
|
client_id=(
|
|
str(extra_data["client_id"])
|
|
if extra_data.get("client_id") is not None
|
|
else None
|
|
),
|
|
node_ids=[str(node_id) for node_id in payload.get("component_node_ids", [])],
|
|
modal_gpu=(
|
|
str(payload["modal_gpu"]) if payload.get("modal_gpu") is not None else None
|
|
),
|
|
suppress_status_stream=bool(payload.get("suppress_status_stream")),
|
|
component_id=str(payload.get("component_id") or "payload"),
|
|
invocation_id=str(payload.get("invocation_id") or "none"),
|
|
stream_started_at=started_at,
|
|
previous_event_at=started_at,
|
|
)
|
|
|
|
|
|
def _handle_remote_log_event(
|
|
state: _PayloadStreamState,
|
|
stream_event: dict[str, Any],
|
|
) -> None:
|
|
"""Retain remote task identity and optionally start its container log stream."""
|
|
task_id = _coerce_modal_task_id(stream_event.get("task_id"))
|
|
if task_id is not None and state.active_remote_task_id is None:
|
|
state.active_remote_task_id = task_id
|
|
if (
|
|
task_id is not None
|
|
and state.active_remote_log_task_id is None
|
|
and _is_remote_container_log_stream_enabled()
|
|
):
|
|
state.active_remote_log_task_id = _retain_remote_container_log_stream(task_id)
|
|
|
|
|
|
def _mapped_progress_lane_id(
|
|
payload: dict[str, Any],
|
|
stream_event: dict[str, Any],
|
|
) -> str | None:
|
|
"""Return the explicit event lane or the payload's mapped lane fallback."""
|
|
if stream_event.get("lane_id") is not None:
|
|
return str(stream_event["lane_id"])
|
|
if payload.get("mapped_progress_lane_id") is not None:
|
|
return str(payload["mapped_progress_lane_id"])
|
|
return None
|
|
|
|
|
|
def _node_progress_kwargs(
|
|
state: _PayloadStreamState,
|
|
stream_event: dict[str, Any],
|
|
*,
|
|
reported_node_id: str,
|
|
display_node_id: str,
|
|
real_node_id: str | None,
|
|
lane_id: str | None,
|
|
aggregate_only: bool,
|
|
) -> dict[str, Any]:
|
|
"""Build local progress emitter arguments from one remote node event."""
|
|
kwargs: dict[str, Any] = {
|
|
"prompt_id": state.prompt_id,
|
|
"client_id": state.client_id,
|
|
"node_id": reported_node_id,
|
|
"value": float(stream_event.get("value", 0.0)),
|
|
"max_value": float(stream_event.get("max", 1.0)),
|
|
"display_node_id": display_node_id,
|
|
"real_node_id": real_node_id,
|
|
"lane_id": lane_id,
|
|
"clear": bool(stream_event.get("clear", False)),
|
|
"item_index": (
|
|
int(stream_event["item_index"])
|
|
if stream_event.get("item_index") is not None
|
|
else (
|
|
int(state.payload["map_item_index"])
|
|
if state.payload.get("map_item_index") is not None
|
|
else None
|
|
)
|
|
),
|
|
"aggregate_only": aggregate_only,
|
|
}
|
|
kwargs.update(
|
|
_remote_execution_identity(state.payload, state.active_remote_task_id)
|
|
)
|
|
for field_name in ("stage", "message", "unit"):
|
|
if stream_event.get(field_name) is not None:
|
|
kwargs[field_name] = str(stream_event[field_name])
|
|
if stream_event.get("indeterminate") is not None:
|
|
kwargs["indeterminate"] = bool(stream_event["indeterminate"])
|
|
for field_name in (
|
|
"elapsed_seconds",
|
|
"time_to_first_token_seconds",
|
|
"tokens_per_second",
|
|
):
|
|
if stream_event.get(field_name) is not None:
|
|
kwargs[field_name] = float(stream_event[field_name])
|
|
completed_ancestors = _remote_prompt_ancestor_node_ids(
|
|
state.payload, real_node_id or display_node_id
|
|
)
|
|
if completed_ancestors:
|
|
kwargs["completed_ancestor_node_ids"] = completed_ancestors
|
|
return kwargs
|
|
|
|
|
|
def _handle_node_progress_event(
|
|
state: _PayloadStreamState,
|
|
stream_event: dict[str, Any],
|
|
) -> None:
|
|
"""Filter and forward one remote per-node progress event."""
|
|
metadata = _progress_stream_event_metadata(stream_event)
|
|
filter_node_id = metadata["filter_node_id"] if metadata is not None else None
|
|
lane_id = _mapped_progress_lane_id(state.payload, stream_event)
|
|
aggregate_only = bool(stream_event.get("aggregate_only", False))
|
|
if (
|
|
state.suppress_status_stream
|
|
and lane_id is None
|
|
and not aggregate_only
|
|
and not _should_forward_suppressed_stream_event(
|
|
state.payload, filter_node_id
|
|
)
|
|
):
|
|
logger.debug(
|
|
"Suppressing streamed Modal node progress for component=%s node_id=%s real_node_id=%s because it does not belong to this mapped/static payload.",
|
|
state.payload.get("component_id"),
|
|
stream_event.get("node_id"),
|
|
stream_event.get("real_node_id"),
|
|
)
|
|
return
|
|
reported_node_id = metadata["node_id"] if metadata is not None else None
|
|
if reported_node_id is None:
|
|
return
|
|
display_node_id = (
|
|
metadata["display_node_id"] if metadata is not None else str(reported_node_id)
|
|
)
|
|
real_node_id = metadata["real_node_id"] if metadata is not None else None
|
|
progress_node_id = real_node_id or display_node_id
|
|
if lane_id is not None:
|
|
_remember_mapped_lane_node_id(state.payload, lane_id, progress_node_id)
|
|
elif state.suppress_status_stream and not aggregate_only and metadata is not None:
|
|
state.suppressed_progress_metadata[str(metadata["filter_node_id"])] = {
|
|
"node_id": str(reported_node_id),
|
|
"display_node_id": display_node_id,
|
|
"real_node_id": real_node_id,
|
|
}
|
|
logger.debug(
|
|
"Forwarding streamed Modal node progress for component=%s node_id=%s real_node_id=%s value=%s max=%s lane_id=%s.",
|
|
state.payload.get("component_id"),
|
|
reported_node_id,
|
|
real_node_id,
|
|
stream_event.get("value"),
|
|
stream_event.get("max"),
|
|
lane_id,
|
|
)
|
|
_emit_local_modal_progress(
|
|
**_node_progress_kwargs(
|
|
state,
|
|
stream_event,
|
|
reported_node_id=str(reported_node_id),
|
|
display_node_id=display_node_id,
|
|
real_node_id=real_node_id,
|
|
lane_id=lane_id,
|
|
aggregate_only=aggregate_only,
|
|
)
|
|
)
|
|
|
|
|
|
def _handle_executed_event(
|
|
state: _PayloadStreamState,
|
|
stream_event: dict[str, Any],
|
|
) -> None:
|
|
"""Forward one remote executed-output event when it belongs to this payload."""
|
|
reported_node_id = stream_event.get("node_id")
|
|
if reported_node_id is None:
|
|
return
|
|
if not _should_forward_suppressed_stream_event(
|
|
state.payload, reported_node_id
|
|
):
|
|
logger.debug(
|
|
"Suppressing streamed Modal executed output for component=%s node_id=%s because it does not belong to this mapped/static payload.",
|
|
state.payload.get("component_id"),
|
|
reported_node_id,
|
|
)
|
|
return
|
|
logger.debug(
|
|
"Forwarding streamed Modal executed output for component=%s node_id=%s.",
|
|
state.payload.get("component_id"),
|
|
reported_node_id,
|
|
)
|
|
_emit_local_executed_output(
|
|
prompt_id=state.prompt_id,
|
|
client_id=state.client_id,
|
|
node_id=str(reported_node_id),
|
|
display_node_id=(
|
|
str(stream_event["display_node_id"])
|
|
if stream_event.get("display_node_id") is not None
|
|
else None
|
|
),
|
|
output_payload=deserialize_value(stream_event.get("output")),
|
|
)
|
|
|
|
|
|
def _handle_preview_event(
|
|
state: _PayloadStreamState,
|
|
stream_event: dict[str, Any],
|
|
) -> None:
|
|
"""Forward one streamed preview image when it belongs to this payload."""
|
|
reported_node_id = stream_event.get("node_id")
|
|
image_bytes = deserialize_value(stream_event.get("image_bytes"))
|
|
if reported_node_id is None or not isinstance(image_bytes, bytes):
|
|
return
|
|
if not _should_forward_suppressed_stream_event(
|
|
state.payload, reported_node_id
|
|
):
|
|
logger.debug(
|
|
"Suppressing streamed Modal preview image for component=%s node_id=%s because it does not belong to this mapped/static payload.",
|
|
state.payload.get("component_id"),
|
|
reported_node_id,
|
|
)
|
|
return
|
|
logger.debug(
|
|
"Forwarding streamed Modal preview image for component=%s node_id=%s.",
|
|
state.payload.get("component_id"),
|
|
reported_node_id,
|
|
)
|
|
_emit_local_preview_image(
|
|
prompt_id=state.prompt_id,
|
|
client_id=state.client_id,
|
|
node_id=str(reported_node_id),
|
|
display_node_id=(
|
|
str(stream_event["display_node_id"])
|
|
if stream_event.get("display_node_id") is not None
|
|
else None
|
|
),
|
|
parent_node_id=(
|
|
str(stream_event["parent_node_id"])
|
|
if stream_event.get("parent_node_id") is not None
|
|
else None
|
|
),
|
|
real_node_id=(
|
|
str(stream_event["real_node_id"])
|
|
if stream_event.get("real_node_id") is not None
|
|
else None
|
|
),
|
|
image_type=str(stream_event.get("image_type", "PNG")),
|
|
image_bytes=image_bytes,
|
|
max_size=(
|
|
int(stream_event["max_size"])
|
|
if stream_event.get("max_size") is not None
|
|
else None
|
|
),
|
|
)
|
|
|
|
|
|
def _handle_boundary_output_event(
|
|
state: _PayloadStreamState,
|
|
stream_event: dict[str, Any],
|
|
) -> None:
|
|
"""Forward one streamed boundary value to its local preview targets."""
|
|
target_node_ids = _stream_event_node_ids(stream_event, "preview_target_node_ids")
|
|
progressive_node_ids = _stream_event_node_ids(
|
|
stream_event, "progressive_preview_target_node_ids"
|
|
)
|
|
if not target_node_ids and not progressive_node_ids:
|
|
return
|
|
logger.debug(
|
|
"Forwarding streamed Modal boundary output previews for component=%s source_node=%s targets=%s progressive_targets=%s.",
|
|
state.payload.get("component_id"),
|
|
stream_event.get("node_id"),
|
|
target_node_ids,
|
|
progressive_node_ids,
|
|
)
|
|
image_value = deserialize_value(stream_event.get("value"))
|
|
if progressive_node_ids:
|
|
_emit_local_progressive_preview_slot(
|
|
prompt_id=state.prompt_id,
|
|
client_id=state.client_id,
|
|
node_ids=progressive_node_ids,
|
|
slot_index=_optional_payload_int(state.payload, "map_item_index") or 0,
|
|
slot_count=_optional_payload_int(state.payload, "map_item_count"),
|
|
image_value=image_value,
|
|
)
|
|
if target_node_ids:
|
|
_emit_local_preview_boundary_output(
|
|
prompt_id=state.prompt_id,
|
|
client_id=state.client_id,
|
|
preview_target_node_ids=target_node_ids,
|
|
image_value=image_value,
|
|
)
|
|
|
|
|
|
def _stream_event_node_ids(stream_event: dict[str, Any], key: str) -> list[str]:
|
|
"""Return the non-empty node ids listed under one stream event key."""
|
|
return [str(node_id) for node_id in stream_event.get(key, []) or [] if str(node_id)]
|
|
|
|
|
|
def _optional_payload_int(payload: dict[str, Any], key: str) -> int | None:
|
|
"""Return one optional integer payload field."""
|
|
value = payload.get(key)
|
|
return int(value) if value is not None else None
|
|
|
|
|
|
def _handle_cached_node_event(
|
|
state: _PayloadStreamState,
|
|
stream_event: dict[str, Any],
|
|
) -> None:
|
|
"""Forward one cached-node marker through the local progress channel."""
|
|
metadata = _progress_stream_event_metadata(stream_event)
|
|
filter_node_id = metadata["filter_node_id"] if metadata is not None else None
|
|
if state.suppress_status_stream and not _should_forward_suppressed_stream_event(
|
|
state.payload, filter_node_id
|
|
):
|
|
logger.debug(
|
|
"Suppressing streamed Modal cached-node marker for component=%s node_id=%s real_node_id=%s because it does not belong to this mapped/static payload.",
|
|
state.payload.get("component_id"),
|
|
stream_event.get("node_id"),
|
|
stream_event.get("real_node_id"),
|
|
)
|
|
return
|
|
reported_node_id = metadata["node_id"] if metadata is not None else None
|
|
if reported_node_id is None:
|
|
return
|
|
display_node_id = (
|
|
metadata["display_node_id"] if metadata is not None else str(reported_node_id)
|
|
)
|
|
real_node_id = metadata["real_node_id"] if metadata is not None else None
|
|
logger.debug(
|
|
"Forwarding streamed Modal cached-node marker for component=%s node_id=%s real_node_id=%s.",
|
|
state.payload.get("component_id"),
|
|
reported_node_id,
|
|
real_node_id,
|
|
)
|
|
_emit_local_modal_progress(
|
|
prompt_id=state.prompt_id,
|
|
client_id=state.client_id,
|
|
node_id=str(reported_node_id),
|
|
value=0.0,
|
|
max_value=1.0,
|
|
display_node_id=display_node_id,
|
|
real_node_id=real_node_id,
|
|
cached_hit=True,
|
|
**_remote_execution_identity(state.payload, state.active_remote_task_id),
|
|
)
|
|
|
|
|
|
def _clear_suppressed_progress(state: _PayloadStreamState) -> None:
|
|
"""Clear retained node progress when a suppressed stream reaches a terminal phase."""
|
|
for metadata in state.suppressed_progress_metadata.values():
|
|
_emit_local_modal_progress(
|
|
prompt_id=state.prompt_id,
|
|
client_id=state.client_id,
|
|
node_id=str(metadata["node_id"]),
|
|
value=0.0,
|
|
max_value=1.0,
|
|
display_node_id=metadata["display_node_id"],
|
|
real_node_id=metadata["real_node_id"],
|
|
clear=True,
|
|
)
|
|
state.suppressed_progress_metadata.clear()
|
|
|
|
|
|
def _handle_status_event(
|
|
state: _PayloadStreamState,
|
|
stream_event: dict[str, Any],
|
|
) -> None:
|
|
"""Forward one component-level remote execution status event."""
|
|
remote_phase = str(stream_event.get("phase", "executing"))
|
|
logger.info(
|
|
"Forwarding streamed Modal progress for component=%s phase=%s active_node_id=%s.",
|
|
state.payload.get("component_id"),
|
|
remote_phase,
|
|
stream_event.get("active_node_id"),
|
|
)
|
|
if state.suppress_status_stream:
|
|
if remote_phase in {
|
|
"execution_success",
|
|
"execution_error",
|
|
"execution_interrupted",
|
|
}:
|
|
_clear_suppressed_progress(state)
|
|
return
|
|
identity = _remote_execution_identity(
|
|
state.payload, state.active_remote_task_id
|
|
)
|
|
if remote_phase == "execution_success":
|
|
_emit_local_modal_status(
|
|
prompt_id=state.prompt_id,
|
|
client_id=state.client_id,
|
|
phase="finalizing",
|
|
node_ids=state.node_ids,
|
|
modal_gpu=state.modal_gpu,
|
|
status_message=(
|
|
"Receiving remote outputs from "
|
|
f"{_remote_execution_destination(state.payload)}"
|
|
),
|
|
**identity,
|
|
)
|
|
return
|
|
active_node_id = (
|
|
str(stream_event["active_node_id"])
|
|
if stream_event.get("active_node_id") is not None
|
|
else None
|
|
)
|
|
_emit_local_modal_status(
|
|
prompt_id=state.prompt_id,
|
|
client_id=state.client_id,
|
|
phase=remote_phase,
|
|
node_ids=state.node_ids,
|
|
modal_gpu=state.modal_gpu,
|
|
active_node_id=active_node_id,
|
|
completed_ancestor_node_ids=(
|
|
_remote_prompt_ancestor_node_ids(state.payload, active_node_id) or None
|
|
),
|
|
active_node_class_type=(
|
|
str(stream_event["active_node_class_type"])
|
|
if stream_event.get("active_node_class_type") is not None
|
|
else None
|
|
),
|
|
active_node_role=(
|
|
str(stream_event["active_node_role"])
|
|
if stream_event.get("active_node_role") is not None
|
|
else None
|
|
),
|
|
**identity,
|
|
)
|
|
|
|
|
|
def _handle_progress_stream_event(
|
|
state: _PayloadStreamState,
|
|
stream_event: dict[str, Any],
|
|
) -> None:
|
|
"""Dispatch one progress envelope to its event-specific handler."""
|
|
event_type = str(stream_event.get("event_type", ""))
|
|
handlers = {
|
|
"node_progress": _handle_node_progress_event,
|
|
"executed": _handle_executed_event,
|
|
"preview": _handle_preview_event,
|
|
"boundary_output": _handle_boundary_output_event,
|
|
"node_cached": _handle_cached_node_event,
|
|
}
|
|
handler = handlers.get(event_type)
|
|
if handler is not None:
|
|
handler(state, stream_event)
|
|
return
|
|
_handle_status_event(state, stream_event)
|
|
|
|
|
|
def _handle_result_event(
|
|
state: _PayloadStreamState,
|
|
stream_event: dict[str, Any],
|
|
*,
|
|
event_received_at: float,
|
|
seconds_since_previous_event: float,
|
|
) -> None:
|
|
"""Validate and retain the final transport-safe result payload."""
|
|
candidate_outputs = stream_event.get("outputs")
|
|
candidate_bytes = (
|
|
len(candidate_outputs)
|
|
if isinstance(candidate_outputs, bytes | bytearray)
|
|
else -1
|
|
)
|
|
logger.info(
|
|
"Received streamed Modal result component=%s prompt_id=%s invocation_id=%s "
|
|
"result_bytes=%d stream_elapsed_seconds=%.3f seconds_since_previous_event=%.3f "
|
|
"event_count=%d progress_event_count=%d.",
|
|
state.component_id,
|
|
state.prompt_id or "none",
|
|
state.invocation_id,
|
|
candidate_bytes,
|
|
event_received_at - state.stream_started_at,
|
|
seconds_since_previous_event,
|
|
state.event_count,
|
|
state.progress_event_count,
|
|
)
|
|
try:
|
|
state.result_payload = coerce_serialized_node_outputs(candidate_outputs)
|
|
except TypeError as exc:
|
|
raise ModalRemoteInvocationError(
|
|
"Modal streamed payload result did not include transport-safe outputs."
|
|
) from exc
|
|
state.should_close_stream = True
|
|
|
|
|
|
def _close_consumed_payload_stream(
|
|
state: _PayloadStreamState,
|
|
stream_events: Iterator[dict[str, Any]],
|
|
) -> None:
|
|
"""Close a completed result stream and release any retained log stream."""
|
|
if state.should_close_stream:
|
|
close_started_at = time.monotonic()
|
|
logger.info(
|
|
"Starting local Modal result stream close component=%s prompt_id=%s invocation_id=%s.",
|
|
state.component_id,
|
|
state.prompt_id or "none",
|
|
state.invocation_id,
|
|
)
|
|
_close_remote_payload_stream(stream_events)
|
|
logger.info(
|
|
"Finished local Modal result stream close in %.3fs component=%s prompt_id=%s invocation_id=%s.",
|
|
time.monotonic() - close_started_at,
|
|
state.component_id,
|
|
state.prompt_id or "none",
|
|
state.invocation_id,
|
|
)
|
|
if state.active_remote_log_task_id is not None:
|
|
_release_remote_container_log_stream(state.active_remote_log_task_id)
|
|
|
|
|
|
def _consume_remote_payload_stream(
|
|
payload: dict[str, Any],
|
|
stream_events: Iterator[dict[str, Any]],
|
|
) -> bytes:
|
|
"""Forward remote progress events into the local UI and return final bytes."""
|
|
state = _new_payload_stream_state(payload)
|
|
logger.info(
|
|
"Starting local Modal stream consumption component=%s prompt_id=%s invocation_id=%s.",
|
|
state.component_id,
|
|
state.prompt_id or "none",
|
|
state.invocation_id,
|
|
)
|
|
try:
|
|
for stream_event in stream_events:
|
|
event_received_at = time.monotonic()
|
|
seconds_since_previous_event = (
|
|
event_received_at - state.previous_event_at
|
|
)
|
|
state.previous_event_at = event_received_at
|
|
state.event_count += 1
|
|
if not state.speculative_prewarm_checked:
|
|
state.speculative_prewarm_checked = True
|
|
_schedule_speculative_affinity_prewarm(
|
|
payload, reason="current_remote_stream_started"
|
|
)
|
|
event_kind = str(stream_event.get("kind", ""))
|
|
if event_kind == "remote_logs":
|
|
_handle_remote_log_event(state, stream_event)
|
|
continue
|
|
if event_kind == "progress":
|
|
state.progress_event_count += 1
|
|
_handle_progress_stream_event(state, stream_event)
|
|
continue
|
|
if event_kind == "result":
|
|
_handle_result_event(
|
|
state,
|
|
stream_event,
|
|
event_received_at=event_received_at,
|
|
seconds_since_previous_event=seconds_since_previous_event,
|
|
)
|
|
break
|
|
logger.debug(
|
|
"Ignoring unexpected streamed Modal event kind=%s for component=%s.",
|
|
event_kind,
|
|
payload.get("component_id"),
|
|
)
|
|
finally:
|
|
_close_consumed_payload_stream(state, stream_events)
|
|
|
|
if state.result_payload is None:
|
|
raise ModalRemoteInvocationError(
|
|
f"Modal streamed payload for component={payload.get('component_id')!r} did not yield a final result."
|
|
)
|
|
logger.info(
|
|
"Finished local Modal stream consumption in %.3fs component=%s prompt_id=%s invocation_id=%s "
|
|
"result_bytes=%d event_count=%d progress_event_count=%d.",
|
|
time.monotonic() - state.stream_started_at,
|
|
state.component_id,
|
|
state.prompt_id or "none",
|
|
state.invocation_id,
|
|
len(state.result_payload),
|
|
state.event_count,
|
|
state.progress_event_count,
|
|
)
|
|
return bytes(state.result_payload)
|