diff --git a/README.md b/README.md index c0af74c..18eebc4 100644 --- a/README.md +++ b/README.md @@ -156,6 +156,8 @@ Split proxies and mapped phases use prompt-scoped remote sessions for live non-t Remote subgraph runs can persist transport-safe node outputs into a shared Modal `Dict` using ComfyUI input-signature semantics. The cache skips non-serializable outputs and entries above the configured size cap. +Tensor and byte boundary values use a versioned binary envelope, so safetensors bytes cross Modal directly instead of expanding through base64 JSON. Readers remain compatible with legacy JSON payloads during deployment replacement. Streamed progress uses a bounded queue that coalesces stale progress when a consumer falls behind while preserving result, error, and completion events. + Cancellation uses a shared Modal `Dict` control store plus local polling of ComfyUI cancellation. Remote workers retire themselves after poisoned CUDA/runtime failures or stuck cancellation, while deterministic prompt and custom-node errors preserve healthy warm workers. ## Asset And Custom Node Sync @@ -171,6 +173,8 @@ Absolute paths and model names resolvable through ComfyUI `folder_paths` work. A In remote mode, assets and custom-node archives are uploaded into the configured Modal volume. In local mode, the default backend is a local mirror used for development and tests. +Asset paths are planned once per queued prompt. Repeated references across nodes or remote components share one hash, sync-index lookup, and upload decision while every component still receives the same content-addressed remote path and reload metadata. + Custom-node sync is enabled by default in remote mode and disabled by default in local mode. When enabled, Modal-Sync packages `custom_nodes/` as a whole-tree manifest plus content-addressed archives for each top-level custom-node package. Unchanged package digests are reused through a Modal `Dict` sync index instead of probing the volume for many marker files. When a synced top-level custom-node package has a `requirements.txt`, those requirements are folded into the Modal image build. `-r other-file.txt` includes are followed relative to the declaring package; pip option and constraint lines are ignored. @@ -236,6 +240,7 @@ Boolean values accept `1`, `true`, `yes`, `on`, `0`, `false`, `no`, and `off`. | `COMFY_MODAL_MAX_INFLIGHT_CALLS` | `4` | Maximum local Modal calls dispatched at once; mapped fan-out is clamped to this budget. | | `COMFY_MODAL_EXECUTION_TIMEOUT_SECONDS` | `3600` | Maximum runtime for one Modal workflow call. | | `COMFY_MODAL_STARTUP_TIMEOUT_SECONDS` | `900` | Maximum Modal container startup and snapshot-restore time. | +| `COMFY_MODAL_STREAM_EVENT_QUEUE_MAXSIZE` | `256` | Maximum buffered remote progress/result envelopes; stale progress is coalesced when full. | | `COMFY_MODAL_ENABLE_PROACTIVE_WARMUP` | `true` | Start background warmup from runtime parallelism signals such as mapped fan-out. | | `COMFY_MODAL_ENABLE_LOADER_PREWARM` | `true` | During warmup, execute synthetic loader prompts for root literal model-loader nodes. | | `COMFY_MODAL_PROACTIVE_WARMUP_HEAD_START_SECONDS` | `2.0` | Bounded wait for exact mapped warmup slots before lane seeding starts. | diff --git a/api_intercept.py b/api_intercept.py index 5936434..61ec3ce 100644 --- a/api_intercept.py +++ b/api_intercept.py @@ -24,7 +24,13 @@ from .modal_executor_node import ( ) from .session_state import RemoteSessionHandle from .settings import ModalSyncSettings, get_settings -from .sync_engine import ModalAssetSyncEngine, ModalVolumeBackend, SyncedAsset, modal +from .sync_engine import ( + AssetSyncRequestCache, + ModalAssetSyncEngine, + ModalVolumeBackend, + SyncedAsset, + modal, +) logger = logging.getLogger(__name__) @@ -2578,6 +2584,7 @@ def _sync_component_prompt_inputs( component: RemoteComponentPlan, rewritten_prompt: dict[str, Any], sync_engine: ModalAssetSyncEngine, + request_cache: AssetSyncRequestCache, status_callback: Any | None = None, ) -> tuple[dict[str, Any], list[SyncedAsset]]: """Build a synced prompt payload for one remote component.""" @@ -2593,6 +2600,7 @@ def _sync_component_prompt_inputs( synced_inputs, node_assets = sync_engine.sync_prompt_inputs( copy.deepcopy(prompt_node.get("inputs", {})), status_callback=status_callback, + request_cache=request_cache, ) synced_assets.extend(node_assets) component_prompt[node_id] = { @@ -2614,6 +2622,14 @@ def _sync_component_prompt_inputs( return component_prompt, synced_assets +def _deduplicate_synced_assets(synced_assets: list[SyncedAsset]) -> list[SyncedAsset]: + """Return one request summary record per content-addressed remote path.""" + unique_assets_by_remote_path: dict[str, SyncedAsset] = {} + for synced_asset in synced_assets: + unique_assets_by_remote_path.setdefault(synced_asset.remote_path, synced_asset) + return list(unique_assets_by_remote_path.values()) + + def _build_component_payload( component: RemoteComponentPlan, component_prompt: dict[str, Any], @@ -3554,17 +3570,25 @@ def rewrite_prompt_for_modal( synced_component_prompts: dict[str, dict[str, Any]] = {} synced_assets_by_component_id: dict[str, list[SyncedAsset]] = {} + request_asset_cache = resolved_sync_engine.create_request_asset_cache( + rewritten_prompt[node_id].get("inputs", {}) + for component in components + for node_id in component.node_ids + ) for component in components: component_prompt, synced_assets = _sync_component_prompt_inputs( component=component, rewritten_prompt=rewritten_prompt, sync_engine=resolved_sync_engine, + request_cache=request_asset_cache, status_callback=status_callback, ) synced_component_prompts[component.representative_node_id] = component_prompt synced_assets_by_component_id[component.representative_node_id] = list(synced_assets) summary.synced_assets.extend(synced_assets) + summary.synced_assets = _deduplicate_synced_assets(summary.synced_assets) + requires_volume_reload = any(asset.uploaded for asset in summary.synced_assets) or ( summary.custom_nodes_bundle is not None and summary.custom_nodes_bundle.uploaded ) diff --git a/comfyui_modal_sync_cloud.py b/comfyui_modal_sync_cloud.py index 4565ae8..843c69a 100644 --- a/comfyui_modal_sync_cloud.py +++ b/comfyui_modal_sync_cloud.py @@ -4543,6 +4543,59 @@ def _execute_mapped_subgraph_payload( ) +class _BoundedStreamEventBuffer: + """Bound progress memory while preserving terminal stream events.""" + + def __init__(self, maxsize: int) -> None: + """Initialize a bounded event queue and close signal.""" + self._queue: queue.Queue[tuple[str, Any]] = queue.Queue(maxsize=max(4, maxsize)) + self._closed = threading.Event() + self._dropped_progress_events = 0 + self._dropped_lock = threading.Lock() + + @property + def dropped_progress_events(self) -> int: + """Return how many stale progress events were coalesced away.""" + with self._dropped_lock: + return self._dropped_progress_events + + def publish_progress(self, payload: Any) -> None: + """Publish the newest progress event without exceeding the queue bound.""" + event = ("progress", payload) + while not self._closed.is_set(): + try: + self._queue.put_nowait(event) + return + except queue.Full: + try: + discarded_event = self._queue.get_nowait() + except queue.Empty: + continue + if discarded_event[0] != "progress": + self._queue.put_nowait(discarded_event) + return + with self._dropped_lock: + self._dropped_progress_events += 1 + + def publish_terminal(self, event_kind: str, payload: Any) -> bool: + """Publish a result, error, or completion unless the consumer closed.""" + while not self._closed.is_set(): + try: + self._queue.put((event_kind, payload), timeout=0.1) + return True + except queue.Full: + continue + return False + + def get(self) -> tuple[str, Any]: + """Wait for and return the next buffered stream event.""" + return self._queue.get() + + def close(self) -> None: + """Release any producer waiting after the consumer stops.""" + self._closed.set() + + def _stream_remote_payload_events( payload: dict[str, Any], kwargs_payload: bytes | bytearray | str | dict[str, Any], @@ -4551,12 +4604,14 @@ def _stream_remote_payload_events( interrupt_flag_key: str | None = None, ) -> Iterator[dict[str, Any]]: """Yield progress and result events for one remote payload execution.""" - event_queue: queue.Queue[tuple[str, Any]] = queue.Queue() + event_buffer = _BoundedStreamEventBuffer( + get_settings().stream_event_queue_maxsize + ) task_id = os.getenv("MODAL_TASK_ID") def publish_status(progress_state: dict[str, Any]) -> None: """Queue a progress envelope for the remote caller.""" - event_queue.put(("progress", serialize_mapping(progress_state))) + event_buffer.publish_progress(serialize_mapping(progress_state)) def execute_payload() -> None: """Run the payload in a worker thread and enqueue the terminal outcome.""" @@ -4599,11 +4654,11 @@ def _stream_remote_payload_events( execute_node_kwargs["interrupt_flag_key"] = interrupt_flag_key outputs = execute_node_locally(payload, kwargs_payload, **execute_node_kwargs) except Exception as exc: # pragma: no cover - exercised through generator consumer tests. - event_queue.put(("error", exc)) + event_buffer.publish_terminal("error", exc) else: - event_queue.put(("result", outputs)) + event_buffer.publish_terminal("result", outputs) finally: - event_queue.put(("done", None)) + event_buffer.publish_terminal("done", None) worker_thread = threading.Thread( target=execute_payload, @@ -4615,7 +4670,7 @@ def _stream_remote_payload_events( worker_thread.start() try: while True: - event_kind, event_payload = event_queue.get() + event_kind, event_payload = event_buffer.get() if event_kind == "progress": yield {"kind": "progress", **event_payload} continue @@ -4627,7 +4682,14 @@ def _stream_remote_payload_events( if event_kind == "done": return finally: + event_buffer.close() worker_thread.join(timeout=1.0) + if event_buffer.dropped_progress_events: + logger.info( + "Coalesced %d stale remote progress event(s) for component=%s to keep the stream buffer bounded.", + event_buffer.dropped_progress_events, + payload.get("component_id"), + ) def _should_ignore_repo_path(path: Path) -> bool: @@ -5154,7 +5216,14 @@ if modal is not None: # pragma: no branch - remote entrypoint configuration. ) image = ( modal.Image.debian_slim(python_version=REMOTE_PYTHON_VERSION) - .env({"COMFY_MODAL_RUNTIME_FINGERPRINT": runtime_identity.fingerprint}) + .env( + { + "COMFY_MODAL_RUNTIME_FINGERPRINT": runtime_identity.fingerprint, + "COMFY_MODAL_STREAM_EVENT_QUEUE_MAXSIZE": str( + settings.stream_event_queue_maxsize + ), + } + ) .pip_install(*_comfyui_runtime_packages()) ) if custom_node_packages: diff --git a/docs/LEARNINGS.md b/docs/LEARNINGS.md index ea484c6..8b2c965 100644 --- a/docs/LEARNINGS.md +++ b/docs/LEARNINGS.md @@ -294,3 +294,4 @@ - Modal app deletion is not a complete reset for Modal-Sync. Persistent state also lives in named Modal Dicts for interrupts, node output cache, durable session bridges, sync index, and snapshot profiles, plus the shared Modal Volume for uploaded assets and custom-node bundles. Expose Dict deletion separately from Volume deletion because clearing caches is a debugging reset, while deleting the Volume forces a full reupload. - Local CPU count is not a safe default for Modal GPU fan-out. Remote calls need an explicit in-flight budget shared by ordinary component dispatch, mapped lanes, and proactive warmup; Modal execution and container startup also need separate deadlines. Queue-time hashing, packaging, and upload work must run off ComfyUI's event loop, and deterministic node failures should preserve a healthy warm worker while poisoned CUDA/runtime failures retire it. - A protocol integer only detects explicit wire-contract changes; it cannot distinguish two deployments built from different source, dependency, or runtime settings. The client and worker need the same deterministic fingerprint over remote Python source, bundled ComfyUI source, custom-node requirements, exact package pins, and deployment-shaping options. Cache compatibility by that fingerprint and bake it into a fixed Python image so a same-protocol stale worker is replaced before it can run payloads. +- Base64 is a costly default for tensor-heavy RPC payloads. A 12 MiB tensor payload measured about 16 MiB and 66.6 ms to encode through safetensors plus base64 JSON, while a versioned raw-attachment envelope stayed near 12 MiB and encoded in 31.8 ms (14.5 ms decode) in the same local check. Keep legacy JSON reads for rolling compatibility, deduplicate asset decisions across the whole queued prompt, and bound streamed progress memory by coalescing stale progress without dropping terminal events. diff --git a/pyproject.toml b/pyproject.toml index 7090ba6..da9a729 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,12 +20,18 @@ Repository = "https://github.com/ttulttul/ComfyUI-Modal" [dependency-groups] test = [ "aiohttp", + "einops", "numpy", "pillow", + "psutil", "pytest", + "pyyaml", "safetensors", + "scipy", "torch", + "torchvision", "tqdm", + "transformers", ] [tool.comfy] diff --git a/runtime_environment.py b/runtime_environment.py index 6732edb..a9cc83b 100644 --- a/runtime_environment.py +++ b/runtime_environment.py @@ -81,6 +81,7 @@ class RemoteRuntimeSettings(Protocol): session_bridge_dict_name: str snapshot_profile_dict_name: str startup_timeout_seconds: int + stream_event_queue_maxsize: int sync_index_dict_name: str volume_name: str @@ -241,6 +242,7 @@ def _runtime_options(settings: RemoteRuntimeSettings) -> dict[str, Any]: "session_bridge_dict_name": settings.session_bridge_dict_name, "snapshot_profile_dict_name": settings.snapshot_profile_dict_name, "startup_timeout_seconds": settings.startup_timeout_seconds, + "stream_event_queue_maxsize": settings.stream_event_queue_maxsize, "sync_index_dict_name": settings.sync_index_dict_name, "volume_name": settings.volume_name, } diff --git a/serialization.py b/serialization.py index fc19941..687102c 100644 --- a/serialization.py +++ b/serialization.py @@ -6,6 +6,7 @@ import base64 import copy import json import logging +import struct from collections.abc import Mapping, Sequence from typing import Any @@ -18,6 +19,10 @@ _BYTES_KIND = "bytes" _TUPLE_KIND = "tuple" _MAPPED_OUTPUT_KIND = "mapped_output" _BATCHABLE_TENSOR_IO_TYPES = frozenset({"IMAGE", "MASK", "NOISE", "SIGMAS"}) +_BINARY_ENVELOPE_MAGIC = b"CMODALB1" +_BINARY_ENVELOPE_VERSION = 1 +_BINARY_HEADER_LENGTH_BYTES = 8 +_MAX_BINARY_HEADER_BYTES = 16 * 1024 * 1024 class MappedOutputValue(list[Any]): @@ -89,6 +94,23 @@ def _deserialize_tensor(payload: Mapping[str, Any]) -> Any: return tensor_map[_VALUE_KEY] +def _serialize_tensor_bytes(value: Any) -> bytes: + """Serialize a tensor directly to safetensors bytes without base64 expansion.""" + from safetensors.torch import save + + torch = _import_torch() + if not isinstance(value, torch.Tensor): + raise TypeError("Expected a torch.Tensor payload.") + return save({_VALUE_KEY: value.detach().contiguous()}) + + +def _deserialize_tensor_bytes(payload: bytes) -> Any: + """Deserialize direct safetensors bytes back into a tensor.""" + from safetensors.torch import load + + return load(payload)[_VALUE_KEY] + + def serialize_value(value: Any) -> Any: """Convert a Python value into a JSON-safe execution payload.""" if _is_scalar(value): @@ -169,30 +191,206 @@ def deserialize_value(payload: Any) -> Any: return {str(key): deserialize_value(value) for key, value in payload.items()} +def _serialize_transport_value(value: Any, attachments: list[bytes]) -> Any: + """Convert one value into JSON metadata plus raw binary attachments.""" + if _is_scalar(value): + return value + + if isinstance(value, MappedOutputValue): + return { + _KIND_KEY: _MAPPED_OUTPUT_KIND, + "items": [_serialize_transport_value(item, attachments) for item in value.items], + "io_type": value.io_type, + "is_list": value.is_list, + } + + try: + torch = _import_torch() + except ModuleNotFoundError: + torch = None + + if torch is not None and isinstance(value, torch.Tensor): + attachment_index = len(attachments) + attachments.append(_serialize_tensor_bytes(value)) + return {_KIND_KEY: _TENSOR_KIND, "attachment": attachment_index} + + if isinstance(value, bytes): + attachment_index = len(attachments) + attachments.append(value) + return {_KIND_KEY: _BYTES_KIND, "attachment": attachment_index} + + if isinstance(value, tuple): + return { + _KIND_KEY: _TUPLE_KIND, + "items": [_serialize_transport_value(item, attachments) for item in value], + } + + if isinstance(value, list): + return [_serialize_transport_value(item, attachments) for item in value] + + if isinstance(value, Mapping): + return { + str(key): _serialize_transport_value(item, attachments) + for key, item in value.items() + } + + raise TypeError( + "ComfyUI-Modal can only serialize JSON-compatible values, bytes, " + "and torch tensors. Unsupported value type: " + f"{type(value)!r}" + ) + + +def _serialize_transport_payload(value: Any, *, sort_keys: bool) -> bytes: + """Serialize one transport payload, using raw attachments when needed.""" + attachments: list[bytes] = [] + metadata = _serialize_transport_value(value, attachments) + if not attachments: + return json.dumps(metadata, sort_keys=sort_keys).encode("utf-8") + + header = json.dumps( + { + "version": _BINARY_ENVELOPE_VERSION, + "payload": metadata, + "attachment_lengths": [len(attachment) for attachment in attachments], + }, + sort_keys=sort_keys, + separators=(",", ":"), + ).encode("utf-8") + return b"".join( + ( + _BINARY_ENVELOPE_MAGIC, + struct.pack(">Q", len(header)), + header, + *attachments, + ) + ) + + +def _deserialize_transport_value(payload: Any, attachments: tuple[bytes, ...]) -> Any: + """Reconstruct one value from metadata and binary attachments.""" + if _is_scalar(payload): + return payload + if isinstance(payload, list): + return [_deserialize_transport_value(item, attachments) for item in payload] + if not isinstance(payload, Mapping): + raise TypeError(f"Unsupported payload type: {type(payload)!r}") + + kind = payload.get(_KIND_KEY) + attachment_index = payload.get("attachment") + if kind in {_TENSOR_KIND, _BYTES_KIND} and attachment_index is not None: + if isinstance(attachment_index, bool) or not isinstance(attachment_index, int): + raise ValueError("Binary attachment indices must be integers.") + if attachment_index < 0 or attachment_index >= len(attachments): + raise ValueError(f"Binary attachment index {attachment_index} is out of range.") + attachment = attachments[attachment_index] + if kind == _TENSOR_KIND: + return _deserialize_tensor_bytes(attachment) + return attachment + if kind == _TENSOR_KIND: + return _deserialize_tensor(payload) + if kind == _BYTES_KIND: + encoded = payload["payload"] + return base64.b64decode(encoded.encode("ascii")) + if kind == _TUPLE_KIND: + return tuple( + _deserialize_transport_value(item, attachments) + for item in payload["items"] + ) + if kind == _MAPPED_OUTPUT_KIND: + return MappedOutputValue( + items=tuple( + _deserialize_transport_value(item, attachments) + for item in payload["items"] + ), + io_type=str(payload.get("io_type", "*")), + is_list=bool(payload.get("is_list", False)), + ) + return { + str(key): _deserialize_transport_value(value, attachments) + for key, value in payload.items() + } + + +def _deserialize_binary_envelope(payload: bytes) -> tuple[Any, tuple[bytes, ...]]: + """Parse one validated binary transport envelope.""" + length_offset = len(_BINARY_ENVELOPE_MAGIC) + header_offset = length_offset + _BINARY_HEADER_LENGTH_BYTES + if len(payload) < header_offset: + raise ValueError("Binary transport envelope is truncated before its header.") + header_length = struct.unpack(">Q", payload[length_offset:header_offset])[0] + if header_length > _MAX_BINARY_HEADER_BYTES: + raise ValueError( + f"Binary transport header exceeds {_MAX_BINARY_HEADER_BYTES} bytes." + ) + attachment_offset = header_offset + header_length + if attachment_offset > len(payload): + raise ValueError("Binary transport envelope is truncated inside its header.") + try: + header = json.loads(payload[header_offset:attachment_offset].decode("utf-8")) + except (UnicodeDecodeError, json.JSONDecodeError) as exc: + raise ValueError("Binary transport envelope header is invalid JSON.") from exc + if not isinstance(header, Mapping): + raise ValueError("Binary transport envelope header must be a mapping.") + if header.get("version") != _BINARY_ENVELOPE_VERSION: + raise ValueError(f"Unsupported binary transport version {header.get('version')!r}.") + attachment_lengths = header.get("attachment_lengths") + if not isinstance(attachment_lengths, list): + raise ValueError("Binary transport envelope is missing attachment lengths.") + + attachments: list[bytes] = [] + next_offset = attachment_offset + for attachment_length in attachment_lengths: + if ( + isinstance(attachment_length, bool) + or not isinstance(attachment_length, int) + or attachment_length < 0 + ): + raise ValueError("Binary attachment lengths must be non-negative integers.") + next_attachment_offset = next_offset + attachment_length + if next_attachment_offset > len(payload): + raise ValueError("Binary transport envelope is truncated inside an attachment.") + attachments.append(bytes(memoryview(payload)[next_offset:next_attachment_offset])) + next_offset = next_attachment_offset + if next_offset != len(payload): + raise ValueError("Binary transport envelope contains trailing bytes.") + return header.get("payload"), tuple(attachments) + + +def _deserialize_transport_payload(payload: bytes | bytearray | str) -> tuple[Any, tuple[bytes, ...]]: + """Decode either the binary transport or the legacy JSON representation.""" + if isinstance(payload, str): + return json.loads(payload), () + normalized_payload = bytes(payload) + if normalized_payload.startswith(_BINARY_ENVELOPE_MAGIC): + return _deserialize_binary_envelope(normalized_payload) + return json.loads(normalized_payload.decode("utf-8")), () + + def serialize_node_inputs(inputs: Mapping[str, Any]) -> bytes: """Serialize node keyword arguments into transport bytes.""" - payload = serialize_mapping(inputs) - return json.dumps(payload, sort_keys=True).encode("utf-8") + return _serialize_transport_payload(dict(inputs), sort_keys=True) def deserialize_node_inputs(payload: bytes | bytearray | str | Mapping[str, Any]) -> dict[str, Any]: """Deserialize node keyword arguments from transport bytes.""" if isinstance(payload, Mapping): raw_payload = dict(payload) + attachments: tuple[bytes, ...] = () else: - if isinstance(payload, (bytes, bytearray)): - payload = payload.decode("utf-8") - raw_payload = json.loads(payload) + raw_payload, attachments = _deserialize_transport_payload(payload) if not isinstance(raw_payload, Mapping): raise TypeError("Serialized node inputs must decode to a mapping.") - return {str(key): deserialize_value(value) for key, value in raw_payload.items()} + return { + str(key): _deserialize_transport_value(value, attachments) + for key, value in raw_payload.items() + } def serialize_node_outputs(outputs: Sequence[Any]) -> bytes: """Serialize node outputs into transport bytes.""" - payload = [serialize_value(value) for value in outputs] - return json.dumps(payload).encode("utf-8") + return _serialize_transport_payload(list(outputs), sort_keys=False) def coerce_serialized_node_outputs(outputs: bytes | bytearray | str | Sequence[Any] | Any) -> bytes: @@ -212,14 +410,16 @@ def deserialize_node_outputs(payload: bytes | bytearray | str | Sequence[Any]) - """Deserialize node outputs from transport bytes.""" if isinstance(payload, Sequence) and not isinstance(payload, (bytes, bytearray, str)): raw_payload = list(payload) + attachments: tuple[bytes, ...] = () else: - if isinstance(payload, (bytes, bytearray)): - payload = payload.decode("utf-8") - raw_payload = json.loads(payload) + raw_payload, attachments = _deserialize_transport_payload(payload) if not isinstance(raw_payload, list): raise TypeError("Serialized node outputs must decode to a list.") - return tuple(deserialize_value(value) for value in raw_payload) + return tuple( + _deserialize_transport_value(value, attachments) + for value in raw_payload + ) def _split_tensor_batch(value: Any) -> list[Any]: diff --git a/settings.py b/settings.py index 0c67866..2ec8df0 100644 --- a/settings.py +++ b/settings.py @@ -48,6 +48,7 @@ _SETTINGS_ENV_KEYS = ( "COMFY_MODAL_PROACTIVE_WARMUP_HEAD_START_SECONDS", "COMFY_MODAL_REMOTE_CANCEL_GRACE_SECONDS", "COMFY_MODAL_REMOTE_CANCEL_RESTART_SECONDS", + "COMFY_MODAL_STREAM_EVENT_QUEUE_MAXSIZE", "COMFY_MODAL_STREAM_REMOTE_CONTAINER_LOGS", ) @@ -91,6 +92,7 @@ class ModalSyncSettings: proactive_warmup_head_start_seconds: float = 2.0 remote_cancel_grace_seconds: float = 2.0 remote_cancel_restart_seconds: float = 1.0 + stream_event_queue_maxsize: int = 256 stream_remote_container_logs: bool = False @@ -298,6 +300,10 @@ def _get_settings_cached( "COMFY_MODAL_REMOTE_CANCEL_RESTART_SECONDS", 1.0, ), + stream_event_queue_maxsize=max( + 4, + _read_int_env("COMFY_MODAL_STREAM_EVENT_QUEUE_MAXSIZE", 256), + ), stream_remote_container_logs=_read_bool_env("COMFY_MODAL_STREAM_REMOTE_CONTAINER_LOGS") or False, ) diff --git a/sync_engine.py b/sync_engine.py index a67bee5..b2e5f94 100644 --- a/sync_engine.py +++ b/sync_engine.py @@ -14,7 +14,7 @@ import zipfile from concurrent.futures import Future, ThreadPoolExecutor from dataclasses import dataclass, field from pathlib import Path -from typing import Any, Callable, Protocol +from typing import Any, Callable, Iterable, Protocol from .settings import ModalSyncSettings, get_settings @@ -106,6 +106,43 @@ class SyncedAsset: uploaded: bool +@dataclass +class AssetSyncRequestCache: + """Deduplicate asset hashing, index lookup, and upload within one queued prompt.""" + + planned_paths: tuple[Path, ...] + _assets_by_path: dict[Path, SyncedAsset] = field(default_factory=dict) + _positions_by_path: dict[Path, int] = field(init=False, default_factory=dict) + + def __post_init__(self) -> None: + """Index stable request progress positions once.""" + self._positions_by_path = { + path.resolve(): index + for index, path in enumerate(self.planned_paths, start=1) + } + + def get(self, local_path: Path) -> SyncedAsset | None: + """Return a previously synced asset for this request when available.""" + return self._assets_by_path.get(local_path.resolve()) + + def put(self, synced_asset: SyncedAsset) -> None: + """Remember one request-scoped sync result by its resolved local path.""" + self._assets_by_path[synced_asset.local_path.resolve()] = synced_asset + + def progress(self, local_path: Path) -> tuple[int, int]: + """Return the stable one-based progress position for a planned asset.""" + resolved_path = local_path.resolve() + return self._positions_by_path[resolved_path], len(self.planned_paths) + + def synced_assets(self) -> tuple[SyncedAsset, ...]: + """Return unique sync results in request planning order.""" + return tuple( + self._assets_by_path[path] + for path in self.planned_paths + if path in self._assets_by_path + ) + + @dataclass(frozen=True) class _CustomNodesArchiveSpec: """Deterministic archive spec for one top-level custom_nodes payload slice.""" @@ -485,6 +522,7 @@ class ModalAssetSyncEngine: inputs: dict[str, Any], *, status_callback: SyncStatusCallback | None = None, + request_cache: AssetSyncRequestCache | None = None, ) -> tuple[dict[str, Any], list[SyncedAsset]]: """Rewrite file-like prompt inputs to mirrored storage paths.""" synced_assets: list[SyncedAsset] = [] @@ -498,13 +536,24 @@ class ModalAssetSyncEngine: if isinstance(value, str): maybe_path = self._resolve_model_path(value) if maybe_path is not None: - syncable_asset_index += 1 + cached_asset = request_cache.get(maybe_path) if request_cache is not None else None + if cached_asset is not None: + synced_assets.append(cached_asset) + return cached_asset.remote_path + if request_cache is None: + syncable_asset_index += 1 + item_index = syncable_asset_index + total_items = len(syncable_asset_paths) + else: + item_index, total_items = request_cache.progress(maybe_path) synced_asset = self.sync_file( maybe_path, status_callback=status_callback, - item_index=syncable_asset_index, - total_items=len(syncable_asset_paths), + item_index=item_index, + total_items=total_items, ) + if request_cache is not None: + request_cache.put(synced_asset) synced_assets.append(synced_asset) return synced_asset.remote_path return value @@ -522,6 +571,17 @@ class ModalAssetSyncEngine: ) return rewritten_inputs, synced_assets + def create_request_asset_cache( + self, + prompt_input_values: Iterable[Any], + ) -> AssetSyncRequestCache: + """Plan unique syncable assets for one queued prompt in stable encounter order.""" + unique_paths: dict[Path, None] = {} + for prompt_input_value in prompt_input_values: + for local_path in self._collect_syncable_asset_paths(prompt_input_value): + unique_paths.setdefault(local_path.resolve(), None) + return AssetSyncRequestCache(planned_paths=tuple(unique_paths)) + def sync_custom_nodes_directory( self, *, diff --git a/tests/conftest.py b/tests/conftest.py index 479ac9a..e633aab 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -162,7 +162,9 @@ def reset_modal_environment(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> monkeypatch.delenv("COMFY_MODAL_STARTUP_TIMEOUT_SECONDS", raising=False) monkeypatch.delenv("COMFY_MODAL_ENABLE_PROACTIVE_WARMUP", raising=False) monkeypatch.delenv("COMFY_MODAL_REMOTE_CANCEL_GRACE_SECONDS", raising=False) + monkeypatch.delenv("COMFY_MODAL_REMOTE_CANCEL_RESTART_SECONDS", raising=False) monkeypatch.delenv("COMFY_MODAL_SESSION_BRIDGE_DICT_NAME", raising=False) + monkeypatch.delenv("COMFY_MODAL_STREAM_EVENT_QUEUE_MAXSIZE", raising=False) monkeypatch.delenv("COMFY_MODAL_STREAM_REMOTE_CONTAINER_LOGS", raising=False) monkeypatch.delenv("COMFY_MODAL_SCALEDOWN_WINDOW", raising=False) monkeypatch.delenv("COMFY_MODAL_MIN_CONTAINERS", raising=False) diff --git a/tests/test_api_intercept.py b/tests/test_api_intercept.py index 0ac3060..84b44fd 100644 --- a/tests/test_api_intercept.py +++ b/tests/test_api_intercept.py @@ -1754,8 +1754,10 @@ def test_rewrite_uses_one_request_wide_volume_reload_marker_across_components( component: Any, rewritten_prompt: dict[str, Any], sync_engine: Any, + request_cache: Any, status_callback: Any = None, ) -> tuple[dict[str, Any], list[Any]]: + del sync_engine, request_cache, status_callback if component.representative_node_id == "1": return {"1": rewritten_prompt["1"]}, [] return { diff --git a/tests/test_modal_executor_node.py b/tests/test_modal_executor_node.py index b008b0e..dccdefe 100644 --- a/tests/test_modal_executor_node.py +++ b/tests/test_modal_executor_node.py @@ -1664,6 +1664,37 @@ def test_modal_cloud_streams_tensor_safe_progress_and_result_events( assert torch.equal(decoded_outputs[0], tensor) +def test_modal_cloud_stream_buffer_coalesces_progress_and_preserves_terminal_events( + modal_cloud_module: Any, +) -> None: + """A slow stream consumer should have bounded progress memory without losing results.""" + event_buffer = modal_cloud_module._BoundedStreamEventBuffer(maxsize=4) + for value in range(100): + event_buffer.publish_progress({"value": value}) + + terminal_thread = threading.Thread( + target=lambda: ( + event_buffer.publish_terminal("result", b"outputs"), + event_buffer.publish_terminal("done", None), + ) + ) + terminal_thread.start() + observed_events: list[tuple[str, Any]] = [] + while True: + event = event_buffer.get() + observed_events.append(event) + if event[0] == "done": + break + terminal_thread.join(timeout=1.0) + event_buffer.close() + + progress_values = [payload["value"] for kind, payload in observed_events if kind == "progress"] + assert progress_values == [96, 97, 98, 99] + assert observed_events[-2:] == [("result", b"outputs"), ("done", None)] + assert event_buffer.dropped_progress_events == 96 + assert not terminal_thread.is_alive() + + def test_modal_cloud_only_reloads_volume_for_requests_with_new_uploads( modal_cloud_module: Any, ) -> None: diff --git a/tests/test_runtime_environment.py b/tests/test_runtime_environment.py index ac7475f..871c31d 100644 --- a/tests/test_runtime_environment.py +++ b/tests/test_runtime_environment.py @@ -25,6 +25,7 @@ def _runtime_settings(**overrides: Any) -> SimpleNamespace: "session_bridge_dict_name": "comfy-modal-sync-session-bridges", "snapshot_profile_dict_name": "comfy-modal-sync-snapshot-profiles", "startup_timeout_seconds": 900, + "stream_event_queue_maxsize": 256, "sync_index_dict_name": "comfy-modal-sync-sync-index", "volume_name": "comfy-universal-storage", } diff --git a/tests/test_serialization.py b/tests/test_serialization.py index 2624a47..991c4c1 100644 --- a/tests/test_serialization.py +++ b/tests/test_serialization.py @@ -2,6 +2,7 @@ from __future__ import annotations +import json from typing import Any import pytest @@ -37,6 +38,47 @@ def test_tensor_round_trip(serialization_module: Any) -> None: assert torch.equal(decoded[0], tensor) +def test_tensor_transport_uses_raw_binary_attachments(serialization_module: Any) -> None: + """Tensor transport should avoid base64's one-third payload expansion.""" + torch = pytest.importorskip("torch") + tensor = torch.arange(256 * 1024, dtype=torch.float32) + + encoded = serialization_module.serialize_node_inputs({"tensor": tensor}) + legacy_json = json.dumps( + serialization_module.serialize_mapping({"tensor": tensor}), + sort_keys=True, + ).encode("utf-8") + + assert encoded.startswith(serialization_module._BINARY_ENVELOPE_MAGIC) + assert len(encoded) < len(legacy_json) * 0.8 + assert torch.equal( + serialization_module.deserialize_node_inputs(encoded)["tensor"], + tensor, + ) + + +def test_binary_transport_keeps_legacy_json_compatibility(serialization_module: Any) -> None: + """New readers should continue to accept already-deployed JSON/base64 payloads.""" + torch = pytest.importorskip("torch") + tensor = torch.arange(8, dtype=torch.float32) + legacy_payload = json.dumps( + serialization_module.serialize_mapping({"tensor": tensor, "bytes": b"abc"}) + ).encode("utf-8") + + decoded = serialization_module.deserialize_node_inputs(legacy_payload) + + assert torch.equal(decoded["tensor"], tensor) + assert decoded["bytes"] == b"abc" + + +def test_binary_transport_rejects_truncated_attachments(serialization_module: Any) -> None: + """A partial Modal payload should fail explicitly instead of decoding corrupted data.""" + encoded = serialization_module.serialize_node_inputs({"bytes": b"abcdef"}) + + with pytest.raises(ValueError, match="truncated inside an attachment"): + serialization_module.deserialize_node_inputs(encoded[:-1]) + + def test_serialize_mapping_supports_nested_tensors(serialization_module: Any) -> None: """Transport mapping helpers should encode tensor values safely.""" torch = pytest.importorskip("torch") diff --git a/tests/test_settings.py b/tests/test_settings.py index c3901f1..44770a3 100644 --- a/tests/test_settings.py +++ b/tests/test_settings.py @@ -122,6 +122,7 @@ def test_settings_reads_remote_invocation_runtime_controls( monkeypatch.setenv("COMFY_MODAL_MAX_INFLIGHT_CALLS", "7") monkeypatch.setenv("COMFY_MODAL_EXECUTION_TIMEOUT_SECONDS", "1800") monkeypatch.setenv("COMFY_MODAL_STARTUP_TIMEOUT_SECONDS", "600") + monkeypatch.setenv("COMFY_MODAL_STREAM_EVENT_QUEUE_MAXSIZE", "32") settings_module.get_settings.cache_clear() try: settings = settings_module.get_settings() @@ -131,6 +132,7 @@ def test_settings_reads_remote_invocation_runtime_controls( assert settings.max_inflight_calls == 7 assert settings.execution_timeout_seconds == 1800 assert settings.startup_timeout_seconds == 600 + assert settings.stream_event_queue_maxsize == 32 def test_settings_reads_proactive_warmup_override( diff --git a/tests/test_sync_engine.py b/tests/test_sync_engine.py index be91620..2e9cbe9 100644 --- a/tests/test_sync_engine.py +++ b/tests/test_sync_engine.py @@ -51,6 +51,62 @@ def test_sync_file_deduplicates_by_hash( assert first.remote_path in sync_index_path.read_text(encoding="utf-8") +def test_request_asset_cache_syncs_repeated_prompt_asset_once( + settings_module: Any, + sync_engine_module: Any, + monkeypatch: Any, + tmp_path: Path, +) -> None: + """Repeated references across remote nodes should share one request sync decision.""" + monkeypatch.setattr(sync_engine_module, "modal", None) + asset_path = tmp_path / "model.safetensors" + asset_path.write_bytes(b"model-bytes") + settings = settings_module.ModalSyncSettings( + app_name="app", + auto_deploy=True, + allow_ephemeral_fallback=False, + enable_memory_snapshot=True, + enable_gpu_memory_snapshot=False, + execution_mode="local", + sync_custom_nodes=False, + volume_name="volume", + route_path="/modal/queue_prompt", + marker_property="is_modal_remote", + local_storage_root=tmp_path / "storage", + remote_storage_root="/storage", + custom_nodes_archive_name="custom_nodes_bundle.zip", + comfyui_root=None, + custom_nodes_dir=None, + ) + engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings) + sync_calls: list[Path] = [] + original_sync_file = engine.sync_file + + def recording_sync_file(local_path: Path, *args: Any, **kwargs: Any) -> Any: + """Record actual request sync decisions before delegating.""" + sync_calls.append(local_path.resolve()) + return original_sync_file(local_path, *args, **kwargs) + + monkeypatch.setattr(engine, "sync_file", recording_sync_file) + request_cache = engine.create_request_asset_cache( + ({"model": str(asset_path)}, {"second_model": str(asset_path)}) + ) + + first_inputs, first_assets = engine.sync_prompt_inputs( + {"model": str(asset_path)}, + request_cache=request_cache, + ) + second_inputs, second_assets = engine.sync_prompt_inputs( + {"second_model": str(asset_path)}, + request_cache=request_cache, + ) + + assert sync_calls == [asset_path.resolve()] + assert first_inputs["model"] == second_inputs["second_model"] + assert first_assets == second_assets + assert request_cache.synced_assets() == (first_assets[0],) + + def test_sync_file_emits_upload_status( settings_module: Any, sync_engine_module: Any,