perf: streamline modal transport and sync

This commit is contained in:
Ken Simpson
2026-07-14 19:53:08 -07:00
parent 5a587d4ef3
commit c6af844b19
16 changed files with 533 additions and 24 deletions
+5
View File
@@ -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. |
+25 -1
View File
@@ -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
)
+76 -7
View File
@@ -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:
+1
View File
@@ -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.
+6
View File
@@ -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]
+2
View File
@@ -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,
}
+212 -12
View File
@@ -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]:
+6
View File
@@ -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,
)
+64 -4
View File
@@ -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,
*,
+2
View File
@@ -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)
+2
View File
@@ -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 {
+31
View File
@@ -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:
+1
View File
@@ -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",
}
+42
View File
@@ -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")
+2
View File
@@ -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(
+56
View File
@@ -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,