nodes/ttulttul-ComfyUI-Modal
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Cancel Vast preparation and bootstrap model downloads

Browse Source
This commit is contained in:
Ken Simpson
2026-08-24 15:35:58 -07:00
parent 4b4d3d894a
commit c6bbb3ec6c
15 changed files with 720 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+2
View File
@@ -173,6 +173,8 @@ Vast startup uses every available lifecycle field instead of relying only on `ac
Direct Vast SSH operations tolerate bounded connection-level closures, resets, and handshake interruptions with jittered exponential retries. Authentication rejection and deterministic remote-command failures remain immediate. Streamed uploads reopen the local source from byte zero on every attempt, and the remote atomic writer verifies the expected byte count before publishing, so a killed connection cannot leave a truncated file at the content-addressed destination.
Cancelling a workflow also cancels queue-time remote preparation. File hashing checks the prompt-scoped cancellation signal between chunks, and active Vast SSH uploads or Hugging Face downloads terminate their local SSH process instead of continuing behind ComfyUI's cancelling status. Vast Hugging Face acquisition runs from a small content-addressed tool bundle uploaded by the controller, so an older worker image cannot force a multi-gigabyte fallback upload merely because it lacks the newest materializer module. The launch environment does not overwrite the runtime fingerprint baked into the published worker image. An owned lease may be adopted across controller-only fingerprint drift when its configured worker image is unchanged and its worker protocol is still compatible, avoiding a duplicate rental; changing the configured image still requires matching capacity.
Set credentials and the worker image in the ComfyUI process environment. Neither value is written to workflows, the lease registry, API responses, or logs:
```bash
api_intercept.py
+93 -3
View File
@@ -64,6 +64,7 @@ from .sync_engine import (
MODEL_FILE_EXTENSIONS,
ModalAssetSyncEngine,
ModalVolumeBackend,
SyncCancelledError,
SyncedAsset,
modal,
resolve_model_path,
@@ -100,6 +101,9 @@ _MODAL_UI_EVENTS_LOCK = threading.Lock()
_MODAL_UI_EVENTS_BY_CLIENT: dict[str, deque[dict[str, Any]]] = {}
_MODAL_INTERRUPT_QUEUE_BRIDGE_ATTR = "__comfy_modal_interrupt_queue_bridge_installed"
_REMOTE_PREPARATION_PROMPTS_ATTR = "__comfy_modal_remote_preparation_prompts"
_REMOTE_PREPARATION_CANCELLATIONS_ATTR = (
"__comfy_modal_remote_preparation_cancellations"
)
_REMOTE_PREPARATION_LOCK_ATTR = "__comfy_modal_remote_preparation_lock"
_REMOTE_TOGGLE_WIDGET_NAME = "Run Remotely"
_LEGACY_REMOTE_TOGGLE_WIDGET_NAME = "Run on Modal"
@@ -7060,8 +7064,11 @@ def rewrite_prompt_for_modal(
extra_data: dict[str, Any] | None = None,
status_callback: SetupStatusCallback | None = None,
plan_callback: ExecutionPlanStatusCallback | None = None,
cancellation_check: Callable[[], bool] | None = None,
) -> tuple[dict[str, Any], RewriteSummary]:
"""Rewrite connected remote components into Modal proxy nodes."""
if cancellation_check is not None and cancellation_check():
raise SyncCancelledError("Remote workflow preparation was cancelled.")
resolved_settings = settings or get_settings()
remote_node_ids = requested_remote_node_ids(
prompt=prompt,
@@ -7255,9 +7262,9 @@ def rewrite_prompt_for_modal(
"Vast.ai assignment has no active controller service."
)
lease = vast_leases_by_environment[assignment.environment_id]
sync_engines_by_environment[assignment.environment_id] = (
vast_service.sync_engine(lease)
)
vast_sync_engine = vast_service.sync_engine(lease)
vast_sync_engine.cancellation_check = cancellation_check
sync_engines_by_environment[assignment.environment_id] = vast_sync_engine
continue
host = ssh_hosts_by_id.get(assignment.environment_id)
if host is None:
@@ -7595,6 +7602,7 @@ async def rewrite_prompt_for_modal_async(
extra_data: dict[str, Any] | None = None,
status_callback: SetupStatusCallback | None = None,
plan_callback: ExecutionPlanStatusCallback | None = None,
cancellation_check: Callable[[], bool] | None = None,
) -> tuple[dict[str, Any], RewriteSummary]:
"""Prepare one Modal prompt without blocking ComfyUI's event loop."""
return await asyncio.to_thread(
@@ -7607,6 +7615,7 @@ async def rewrite_prompt_for_modal_async(
extra_data=extra_data,
status_callback=status_callback,
plan_callback=plan_callback,
cancellation_check=cancellation_check,
)
@@ -7821,6 +7830,13 @@ def _delete_modal_volume_route_path(route_path: str) -> str:
return f"{route_path.rstrip('/')}/delete_volume"
def _cancel_preparation_route_path(route_path: str) -> str:
"""Return the sibling route used to cancel queue-time remote preparation."""
if route_path.endswith("/queue_prompt"):
return f"{route_path.removesuffix('/queue_prompt')}/cancel_preparation"
return f"{route_path.rstrip('/')}/cancel_preparation"
def _remote_environments_route_path(route_path: str) -> str:
"""Return the provider-neutral environment management route."""
del route_path
@@ -7995,8 +8011,14 @@ def _install_modal_interrupt_queue_bridge(prompt_server: Any) -> None:
return
preparation_prompts: dict[str, tuple[Any, ...]] = {}
preparation_cancellations: dict[str, threading.Event] = {}
preparation_lock = threading.RLock()
setattr(prompt_queue, _REMOTE_PREPARATION_PROMPTS_ATTR, preparation_prompts)
setattr(
prompt_queue,
_REMOTE_PREPARATION_CANCELLATIONS_ATTR,
preparation_cancellations,
)
setattr(prompt_queue, _REMOTE_PREPARATION_LOCK_ATTR, preparation_lock)
def preparation_items() -> list[tuple[Any, ...]]:
@@ -8092,6 +8114,7 @@ def _set_remote_preparation(
prompt_id: str,
prompt: Mapping[str, Any],
extra_data: Mapping[str, Any],
cancellation_event: threading.Event | None = None,
) -> bool:
"""Register one pre-queue remote workflow and publish its active queue state."""
prompt_queue = getattr(prompt_server, "prompt_queue", None)
@@ -8108,6 +8131,13 @@ def _set_remote_preparation(
[],
{},
)
cancellations = getattr(
prompt_queue,
_REMOTE_PREPARATION_CANCELLATIONS_ATTR,
None,
)
if isinstance(cancellations, dict) and cancellation_event is not None:
cancellations[prompt_id] = cancellation_event
queue_updated = getattr(prompt_server, "queue_updated", None)
if callable(queue_updated):
queue_updated()
@@ -8123,6 +8153,13 @@ def _clear_remote_preparation(prompt_server: Any, prompt_id: str) -> None:
return
with preparation_lock:
removed = preparations.pop(prompt_id, None)
cancellations = getattr(
prompt_queue,
_REMOTE_PREPARATION_CANCELLATIONS_ATTR,
None,
)
if isinstance(cancellations, dict):
cancellations.pop(prompt_id, None)
if removed is not None:
queue_updated = getattr(prompt_server, "queue_updated", None)
if callable(queue_updated):
@@ -8171,6 +8208,9 @@ def setup_modal_queue_route(
delete_volume_route_path = _delete_modal_volume_route_path(
resolved_settings.route_path
)
cancel_preparation_route_path = _cancel_preparation_route_path(
resolved_settings.route_path
)
remote_environments_route_path = _remote_environments_route_path(
resolved_settings.route_path
)
@@ -8190,6 +8230,28 @@ def setup_modal_queue_route(
vast_registry = _vast_registry(resolved_settings)
_install_modal_interrupt_queue_bridge(prompt_server)
@prompt_server.routes.post(cancel_preparation_route_path)
async def cancel_remote_preparation(request: web.Request) -> web.Response:
"""Cancel one prompt while it is still preparing remote execution."""
payload = await request.json()
prompt_id = str(payload.get("prompt_id") or "").strip()
prompt_queue = getattr(prompt_server, "prompt_queue", None)
cancellations = getattr(
prompt_queue,
_REMOTE_PREPARATION_CANCELLATIONS_ATTR,
None,
)
cancellation_event = (
cancellations.get(prompt_id)
if isinstance(cancellations, dict) and prompt_id
else None
)
if cancellation_event is None:
return web.json_response({"cancelled": False, "prompt_id": prompt_id})
cancellation_event.set()
logger.info("Cancelled remote preparation for prompt %s.", prompt_id)
return web.json_response({"cancelled": True, "prompt_id": prompt_id})
if hasattr(prompt_server.routes, "get"):
@prompt_server.routes.get(remote_environments_route_path)
@@ -8616,6 +8678,7 @@ def setup_modal_queue_route(
request_modal_gpu: str | None = None
summary = RewriteSummary()
preparation_prompt_id: str | None = None
preparation_cancellation = threading.Event()
configurator_node_id: str | None = None
try:
request_started_at = time.perf_counter()
@@ -8668,6 +8731,7 @@ def setup_modal_queue_route(
prompt_id=prompt_id,
prompt=json_data.get("prompt", {}),
extra_data=json_data.get("extra_data", {}),
cancellation_event=preparation_cancellation,
):
preparation_prompt_id = prompt_id
@@ -8698,6 +8762,10 @@ def setup_modal_queue_route(
total: int | None = None,
) -> None:
"""Forward one queue-time Modal setup update into the websocket stream."""
if preparation_cancellation.is_set():
raise SyncCancelledError(
"Remote workflow preparation was cancelled."
)
_emit_modal_status(
prompt_server=prompt_server,
phase="setup",
@@ -8748,6 +8816,7 @@ def setup_modal_queue_route(
extra_data=json_data.get("extra_data"),
status_callback=emit_setup_status,
plan_callback=emit_execution_plan,
cancellation_check=preparation_cancellation.is_set,
)
selected_modal_gpus = _selected_modal_gpus(
summary,
@@ -8862,6 +8931,27 @@ def setup_modal_queue_route(
time.perf_counter() - request_started_at,
)
return response
except SyncCancelledError as exc:
logger.info("Remote workflow preparation cancelled: %s", exc)
if json_data is not None:
_emit_modal_status(
prompt_server=prompt_server,
phase="execution_interrupted",
client_id=str(json_data.get("client_id"))
if json_data.get("client_id")
else None,
prompt_id=str(json_data.get("prompt_id"))
if json_data.get("prompt_id")
else None,
node_ids=remote_node_ids,
configurator_node_id=configurator_node_id,
modal_gpu=request_modal_gpu,
status_message=str(exc),
)
return web.json_response(
{"error": str(exc), "node_errors": [], "cancelled": True},
status=409,
)
except FileNotFoundError as exc:
logger.exception("Modal asset sync failed.")
if json_data is not None:
docs/LEARNINGS.md
+4
View File
@@ -540,3 +540,7 @@
- A websocket setup-error event can arrive before the failed queue request returns. If that event ends synthetic execution without forwarding its `error_message`, ComfyUI records a generic queue failure before the HTTP catch can preserve the backend detail. Use the same backend message for the synthetic error, global status, and configurator terminal state.
- Vast lifecycle fields are not individually authoritative: fall back from an absent `actual_status` to `cur_state` and `intended_status`. When all lifecycle fields, the provider message, and the SSH host remain absent continuously, use a shorter bounded grace period instead of presenting `unknown` for the entire readiness timeout; observable loading and download states must retain the normal window.
- Vast may close an individual SSH connection even after the same instance has passed worker readiness and immediately accepts a subsequent probe. Treat explicit connection closures, resets, and handshake failures as transient transport events with bounded jittered retries, while leaving authentication and remote-command failures immediate. A retried streamed upload must reopen its source and verify the received byte count before atomic publication; otherwise a severed stream can promote a truncated temporary file into a trusted content-addressed path.
- ComfyUI's `/interrupt` flag only reaches work already executing through its prompt worker. Remote planning and asset synchronization happen inside the queue HTTP request's `to_thread`, so cancellation needs its own prompt-scoped event, chunk-level hashing checks, and subprocess termination for active SSH transfers; changing the status pill alone does not stop network or disk work.
- A controller must not inject its expected runtime fingerprint into a rented worker's environment, because `runtime-info` then reports the expectation rather than the image's actual embedded identity. Preserve the fingerprint baked by the image builder and compare it to the controller expectation before trusting the worker.
- Queue-time tooling can evolve before a costly worker image is rebuilt. Ship small controller-owned helpers, such as the verified Hugging Face materializer, as deterministic content-addressed zipapps over SSH; this preserves current validation behavior without turning a missing Python module into a 50+ GB fallback upload.
- A full repository-source fingerprint changes for controller-only edits even when the configured Vast worker image and its wire protocol remain usable. Persist the image reference on new lease records, allow legacy records one compatibility adoption, and reuse only the same configured image with the exact worker protocol; otherwise a cancellation or UI fix can rent another expensive GPU without improving remote execution.
sync_engine.py
+87 -5
View File
@@ -22,6 +22,7 @@ from .settings import ModalSyncSettings, get_settings
logger = logging.getLogger(__name__)
SyncStatusCallback = Callable[[str, int | None, int | None], None]
CancellationCheck = Callable[[], bool]
_SYNC_EXTENSIONS = frozenset({".safetensors", ".ckpt", ".gguf", ".pt", ".vae"})
MODEL_FILE_EXTENSIONS = frozenset(
@@ -181,6 +182,39 @@ class HuggingFaceMaterializingBackend(Protocol):
"""Materialize one immutable file and report whether it succeeded."""
@runtime_checkable
class CancellableVolumeBackend(Protocol):
"""Optional backend capability for interruptible file uploads."""
def put_file_cancellable(
self,
local_path: Path,
remote_path: str,
*,
cancellation_check: CancellationCheck,
) -> None:
"""Upload a file while cooperatively observing cancellation."""
@runtime_checkable
class CancellableHuggingFaceMaterializingBackend(Protocol):
"""Optional backend capability for interruptible remote downloads."""
def materialize_huggingface_file_cancellable(
self,
source: HuggingFaceAssetSource,
remote_path: str,
*,
token: str | None,
cancellation_check: CancellationCheck,
) -> bool:
"""Materialize one immutable file while observing cancellation."""
class SyncCancelledError(RuntimeError):
"""Raised when queue-time asset preparation is cancelled."""
class SyncIndexBackend(Protocol):
"""Minimal metadata index interface used to deduplicate deterministic uploads."""
@@ -561,6 +595,7 @@ class ModalAssetSyncEngine:
sync_index: SyncIndexBackend | None = None
huggingface_asset_registry: HuggingFaceAssetRegistry | None = None
huggingface_asset_discovery: HuggingFaceAssetDiscovery | None = None
cancellation_check: CancellationCheck | None = None
_hash_cache: dict[str, dict[str, Any]] = field(init=False, default_factory=dict)
_path_resolution_cache: dict[str, str | None] = field(init=False, default_factory=dict)
_hash_cache_dirty: bool = field(init=False, default=False)
@@ -617,6 +652,7 @@ class ModalAssetSyncEngine:
total_items: int | None = None,
) -> SyncedAsset:
"""Sync a file into content-addressable remote storage."""
self._raise_if_cancelled()
source_path = local_path.expanduser().absolute()
if not source_path.is_file():
raise FileNotFoundError(f"Asset not found: {source_path}")
@@ -668,6 +704,7 @@ class ModalAssetSyncEngine:
request_cache: AssetSyncRequestCache | None = None,
) -> tuple[dict[str, Any], list[SyncedAsset]]:
"""Rewrite file-like prompt inputs to mirrored storage paths."""
self._raise_if_cancelled()
synced_assets: list[SyncedAsset] = []
sync_started_at = time.perf_counter()
logger.info("Scanning prompt inputs for syncable assets.")
@@ -676,6 +713,7 @@ class ModalAssetSyncEngine:
def rewrite(value: Any) -> Any:
nonlocal syncable_asset_index
self._raise_if_cancelled()
if isinstance(value, str):
maybe_path = self._resolve_model_path(value)
if maybe_path is not None:
@@ -898,6 +936,7 @@ class ModalAssetSyncEngine:
huggingface_source: HuggingFaceAssetSource | None = None,
) -> _ContentAddressedSyncResult:
"""Upload one deterministic file only when its digest is absent from the sync index."""
self._raise_if_cancelled()
existing_record = self._lookup_sync_record(sync_key)
if existing_record is not None:
indexed_remote_path = str(existing_record["remote_path"])
@@ -926,11 +965,32 @@ class ModalAssetSyncEngine:
status_current,
status_total,
)
if self.volume.materialize_huggingface_file(
huggingface_source,
remote_path,
token=os.getenv("HF_TOKEN") or os.getenv("HUGGING_FACE_HUB_TOKEN"),
token = os.getenv("HF_TOKEN") or os.getenv("HUGGING_FACE_HUB_TOKEN")
if (
self.cancellation_check is not None
and isinstance(
self.volume,
CancellableHuggingFaceMaterializingBackend,
)
):
try:
materialized = self.volume.materialize_huggingface_file_cancellable(
huggingface_source,
remote_path,
token=token,
cancellation_check=self.cancellation_check,
)
except InterruptedError as exc:
raise SyncCancelledError(
"Remote workflow preparation was cancelled."
) from exc
else:
materialized = self.volume.materialize_huggingface_file(
huggingface_source,
remote_path,
token=token,
)
if materialized:
self._store_sync_record(
sync_key=sync_key,
remote_path=remote_path,
@@ -955,7 +1015,22 @@ class ModalAssetSyncEngine:
status_current,
status_total,
)
self.volume.put_file(local_path, remote_path)
if self.cancellation_check is not None and isinstance(
self.volume,
CancellableVolumeBackend,
):
try:
self.volume.put_file_cancellable(
local_path,
remote_path,
cancellation_check=self.cancellation_check,
)
except InterruptedError as exc:
raise SyncCancelledError(
"Remote workflow preparation was cancelled."
) from exc
else:
self.volume.put_file(local_path, remote_path)
self._store_sync_record(
sync_key=sync_key,
remote_path=remote_path,
@@ -1292,6 +1367,7 @@ class ModalAssetSyncEngine:
def _hash_file(self, path: Path) -> str:
"""Compute the SHA256 digest for a file."""
self._raise_if_cancelled()
resolved_path = path.resolve()
stat_result = resolved_path.stat()
cache_key = str(resolved_path)
@@ -1307,6 +1383,7 @@ class ModalAssetSyncEngine:
digest = hashlib.sha256()
with resolved_path.open("rb") as handle:
for chunk in iter(lambda: handle.read(1024 * 1024), b""):
self._raise_if_cancelled()
digest.update(chunk)
sha256 = digest.hexdigest()
self._hash_cache[cache_key] = {
@@ -1318,6 +1395,11 @@ class ModalAssetSyncEngine:
self._mark_hash_cache_dirty()
return sha256
def _raise_if_cancelled(self) -> None:
"""Abort queue-time synchronization when its prompt was cancelled."""
if self.cancellation_check is not None and self.cancellation_check():
raise SyncCancelledError("Remote workflow preparation was cancelled.")
def _hash_directory(self, path: Path) -> str:
"""Compute a stable SHA256 digest for a directory tree."""
hash_started_at = time.perf_counter()
tests/test_api_intercept.py
+47
View File
@@ -5933,6 +5933,18 @@ def test_container_status_route_is_queue_route_sibling(api_intercept_module: Any
assert api_intercept_module._container_status_route_path("/modal/queue_prompt") == (
"/modal/container_status"
)
def test_cancel_preparation_route_is_queue_route_sibling(
api_intercept_module: Any,
) -> None:
"""Queue-time cancellation should use a stable sibling route."""
assert api_intercept_module._cancel_preparation_route_path(
"/modal/queue_prompt"
) == "/modal/cancel_preparation"
assert api_intercept_module._cancel_preparation_route_path(
"/custom/modal"
) == "/custom/modal/cancel_preparation"
assert api_intercept_module._container_status_route_path("/custom/modal") == (
"/custom/modal/container_status"
)
@@ -6243,6 +6255,41 @@ def test_remote_preparation_bridge_clears_failed_submission(
assert prompt_queue.get_current_queue_volatile() == ([], [])
def test_remote_preparation_bridge_tracks_prompt_cancellation(
api_intercept_module: Any,
) -> None:
"""Queue-time work should expose a prompt-scoped cancellation event."""
class FakePromptQueue:
"""Provide the queue method required by the bridge."""
def get_current_queue(self) -> tuple[list[Any], list[Any]]:
"""Return no native work."""
return [], []
prompt_queue = FakePromptQueue()
prompt_server = SimpleNamespace(prompt_queue=prompt_queue)
cancellation_event = api_intercept_module.threading.Event()
api_intercept_module._install_modal_interrupt_queue_bridge(prompt_server)
assert api_intercept_module._set_remote_preparation(
prompt_server,
prompt_id="prompt-cancel",
prompt={},
extra_data={},
cancellation_event=cancellation_event,
)
cancellations = getattr(
prompt_queue,
api_intercept_module._REMOTE_PREPARATION_CANCELLATIONS_ATTR,
)
assert cancellations["prompt-cancel"] is cancellation_event
api_intercept_module._clear_remote_preparation(prompt_server, "prompt-cancel")
assert "prompt-cancel" not in cancellations
def test_selected_vast_capacity_streams_setup_status(
api_intercept_module: Any,
vast_models_module: Any,
tests/test_sync_engine.py
+71
View File
@@ -10,6 +10,8 @@ import types
from typing import Any
import zipfile
import pytest
def test_sync_file_deduplicates_by_hash(
settings_module: Any,
@@ -1722,3 +1724,72 @@ def test_modal_volume_backend_applies_shared_rate_limit_backoff_across_calls(
assert backend._record_shared_rate_limit_backoff() == 0.25
assert backend.exists("/hashes/second.done") is True
assert 0.25 in sleep_calls
def test_cancellable_backend_interruption_becomes_sync_cancellation(
settings_module: Any,
sync_engine_module: Any,
tmp_path: Path,
) -> None:
"""An interrupted transport should stop queue-time sync without indexing it."""
asset_path = tmp_path / "model.safetensors"
asset_path.write_bytes(b"model")
class CancellableVolume:
"""Expose an upload that reports prompt cancellation."""
def exists(self, remote_path: str) -> bool:
"""Report empty storage."""
del remote_path
return False
def put_file(self, local_path: Path, remote_path: str) -> None:
"""Reject the non-cancellable path."""
del local_path, remote_path
raise AssertionError("cancellable upload path was not used")
def put_file_cancellable(
self,
local_path: Path,
remote_path: str,
*,
cancellation_check: Any,
) -> None:
"""Report that the active transport was interrupted."""
del local_path, remote_path
raise InterruptedError("cancelled")
def put_bytes(self, payload: bytes, remote_path: str) -> None:
"""Accept protocol-required byte writes."""
del payload, remote_path
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="vast",
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(
volume=CancellableVolume(),
settings=settings,
cancellation_check=lambda: False,
)
with pytest.raises(sync_engine_module.SyncCancelledError, match="cancelled"):
engine._sync_content_addressed_file(
local_path=asset_path,
remote_path="assets/model.safetensors",
sync_key="test-key",
source_description=str(asset_path),
)
tests/test_vast_leases.py
+46
View File
@@ -129,6 +129,52 @@ def test_manager_recovers_offer_race_then_reuses_persisted_lease(
asyncio.run(scenario())
def test_manager_adopts_legacy_lease_for_unchanged_worker_image(
tmp_path: Any,
vast_api_module: Any,
vast_leases_module: Any,
vast_models_module: Any,
vast_simulator_module: Any,
) -> None:
"""Controller fingerprint drift should not rent again for the same worker image."""
async def scenario() -> None:
state = vast_simulator_module.VastSimulatorState(polls_until_running=1)
async with _running_simulator(
vast_simulator_module.create_vast_simulator_app(state)
) as base_url:
registry = vast_leases_module.VastLeaseRegistry.for_user_directory(tmp_path)
client = vast_api_module.VastApiClient(state.api_key, base_url=base_url)
profile = _profile(vast_models_module)
first_manager = vast_leases_module.VastLeaseManager(
api_client=client,
registry=registry,
owner_id="comfy-owner",
runtime_fingerprint="a" * 64,
launch_spec_factory=_launch_factory(vast_models_module),
startup_timeout_seconds=2.0,
)
first = await first_manager.ensure_lease(profile)
second_manager = vast_leases_module.VastLeaseManager(
api_client=client,
registry=registry,
owner_id="comfy-owner",
runtime_fingerprint="b" * 64,
worker_image="ghcr.io/example/worker@sha256:same",
launch_spec_factory=_launch_factory(vast_models_module),
startup_timeout_seconds=2.0,
)
second = await second_manager.ensure_lease(profile)
assert second.instance_id == first.instance_id
assert second.runtime_fingerprint == "b" * 64
assert second.worker_image == "ghcr.io/example/worker@sha256:same"
assert len(state.instances) == 1
asyncio.run(scenario())
def test_active_work_suspends_idle_destruction_and_finish_resets_deadline(
tmp_path: Any,
vast_api_module: Any,
tests/test_vast_runtime.py
+27
View File
@@ -60,6 +60,8 @@ def test_launch_spec_starts_worker_and_watchdog_without_exporting_api_key(
assert "\nCOMFY_VAST_ENV\n" in payload["onstart"]
assert "\n COMFY_VAST_ENV\n" not in payload["onstart"]
assert payload["env"]["COMFY_MODAL_LLM_EXECUTION_TARGET"] == "vast"
assert "COMFY_MODAL_RUNTIME_FINGERPRINT" not in payload["env"]
assert "COMFY_MODAL_RUNTIME_FINGERPRINT" not in payload["onstart"]
def test_configuration_requires_explicit_published_image(
@@ -124,6 +126,31 @@ def test_runtime_manager_rejects_image_fingerprint_drift(
manager.ensure_worker()
def test_runtime_manager_accepts_same_protocol_fingerprint_drift(
vast_runtime_module: Any,
) -> None:
"""Controller-only source changes may reuse the same protocol-compatible image."""
manager = vast_runtime_module.VastRuntimeManager(
runner=FakeRunner(
[
{
"protocol_version": vast_runtime_module.REMOTE_PROTOCOL_VERSION,
"runtime_fingerprint": "b" * 64,
"worker_socket_ready": True,
}
]
),
configuration=vast_runtime_module.VastRuntimeConfiguration(
image="worker",
runtime_fingerprint="a" * 64,
),
)
info = manager.ensure_worker()
assert info["worker_socket_ready"] is True
def test_runtime_manager_logs_changed_ssh_readiness_diagnostics(
vast_runtime_module: Any,
caplog: Any,
tests/test_vast_ssh.py
+87 -2
View File
@@ -27,6 +27,7 @@ class FakeRunner:
input_file: Path | None = None,
timeout_seconds: float | None = None,
check: bool = True,
cancellation_check: Any | None = None,
) -> Any:
"""Record one command and infer test/atomic-write behavior."""
arguments = tuple(remote_argv)
@@ -37,6 +38,7 @@ class FakeRunner:
"input_file": input_file,
"timeout": timeout_seconds,
"check": check,
"cancellation_check": cancellation_check,
}
)
if arguments[:2] == ("test", "-f"):
@@ -47,7 +49,9 @@ class FakeRunner:
)
if arguments[:2] == ("python", "-c"):
self.existing.add(arguments[-1])
if arguments[:3] == ("python", "-m", "remote.huggingface_materializer"):
if arguments[:2] == ("python", "/storage/runtime-tools") or (
len(arguments) == 2 and arguments[1].endswith(".pyz")
):
request = __import__("json").loads(input_payload or b"{}")
self.existing.add(f"/storage/{request['remote_path']}")
return self.module.VastSshCommandResult(stdout=b"", stderr=b"", returncode=0)
@@ -344,13 +348,94 @@ def test_huggingface_materialization_keeps_token_out_of_command_arguments(
materialize_call = next(
call
for call in runner.calls
if call["argv"][:3] == ("python", "-m", "remote.huggingface_materializer")
if len(call["argv"]) == 2
and call["argv"][0] == "python"
and call["argv"][-1].endswith(".pyz")
)
assert "hf_private_token" not in " ".join(materialize_call["argv"])
assert b"hf_private_token" in materialize_call["input"]
assert backend.exists("assets/model.safetensors") is True
def test_materializer_bundle_is_self_contained_and_executable(
vast_ssh_module: Any,
tmp_path: Path,
) -> None:
"""The uploaded zipapp should load without worker-image package sources."""
bundle_path = tmp_path / "materializer.pyz"
bundle = vast_ssh_module._huggingface_materializer_bundle()
bundle_path.write_bytes(bundle)
assert vast_ssh_module._huggingface_materializer_bundle() == bundle
completed = subprocess.run(
[sys.executable, str(bundle_path)],
input=b"{}",
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
check=False,
)
assert completed.returncode != 0
assert b"No module named" not in completed.stderr
assert b"source must be an object" in completed.stderr
def test_runner_terminates_active_ssh_process_when_cancelled(
tmp_path: Path,
vast_ssh_module: Any,
monkeypatch: Any,
) -> None:
"""Prompt cancellation should terminate an in-flight OpenSSH subprocess."""
runner = vast_ssh_module.VastSshRunner(
vast_ssh_module.VastSshConnection(
host="ssh.example.invalid",
port=22,
known_hosts_path=(tmp_path / "known-hosts").resolve(),
)
)
class FakeProcess:
"""Represent a running SSH process that only exits after termination."""
returncode = -15
def __init__(self) -> None:
"""Track termination calls."""
self.terminated = False
def terminate(self) -> None:
"""Record graceful process termination."""
self.terminated = True
def wait(self, timeout: float | None = None) -> int:
"""Return the terminated status immediately."""
del timeout
return self.returncode
def kill(self) -> None:
"""Fail if graceful termination unexpectedly needs escalation."""
raise AssertionError("kill should not be needed")
def communicate(self, **kwargs: Any) -> tuple[bytes, bytes]:
"""Remain active until the cancellation check fires."""
del kwargs
raise subprocess.TimeoutExpired("ssh", 0.25)
process = FakeProcess()
monkeypatch.setattr(
vast_ssh_module.subprocess,
"Popen",
lambda *args, **kwargs: process,
)
checks = iter((False, True))
with pytest.raises(vast_ssh_module.VastSshCancelledError, match="cancelled"):
runner.run(("python", "worker.py"), cancellation_check=lambda: next(checks))
assert process.terminated is True
@pytest.mark.parametrize("path", ["../secret", "a/../../secret", ".", ""])
def test_volume_backend_rejects_path_traversal(
path: str,
tests/test_web_modal_toggle.py
+3 -1
View File
@@ -229,7 +229,9 @@ def test_cancel_click_shows_immediate_modal_cancelling_feedback() -> None:
source = _modal_toggle_source()
assert "function patchInterruptFeedback()" in source
assert "api.fetchApi = function modalFetchApi(resource, options) {" in source
assert "api.fetchApi = async function modalFetchApi(resource, options) {" in source
assert "MODAL_CANCEL_PREPARATION_ROUTE" in source
assert 'body: JSON.stringify({ prompt_id: promptId })' in source
assert "function promptIdsFromInterruptRequest(resource, options)" in source
assert 'route.includes("/interrupt")' in source
assert "function markPromptCancellationRequested(promptId)" in source
vast_leases.py
+32 -1
View File
@@ -79,6 +79,7 @@ class VastLeaseRecord:
active_invocations: int = 0
draining: bool = False
last_error: str | None = None
worker_image: str | None = None
def __post_init__(self) -> None:
"""Validate ownership, lifecycle, and scheduling fields."""
@@ -111,6 +112,8 @@ class VastLeaseRecord:
raise ValueError("Vast lease label does not match its owner identity.")
if self.ssh_port is not None and self.ssh_port <= 0:
raise ValueError("ssh_port must be positive when present.")
if self.worker_image is not None and not self.worker_image.strip():
raise ValueError("worker_image must be non-empty when present.")
@property
def environment_id(self) -> str:
@@ -162,6 +165,9 @@ class VastLeaseRecord:
active_invocations=int(payload.get("active_invocations", 0)),
draining=bool(payload.get("draining", False)),
last_error=(str(payload["last_error"]) if payload.get("last_error") else None),
worker_image=(
str(payload["worker_image"]) if payload.get("worker_image") else None
),
)
@@ -328,6 +334,7 @@ class VastLeaseManager:
],
startup_timeout_seconds: float = 900.0,
clock: Callable[[], float] = time.time,
worker_image: str | None = None,
) -> None:
"""Configure acquisition, persistence, and lifecycle dependencies."""
if not owner_id.strip():
@@ -343,6 +350,7 @@ class VastLeaseManager:
self.launch_spec_factory = launch_spec_factory
self.startup_timeout_seconds = startup_timeout_seconds
self.clock = clock
self.worker_image = worker_image.strip() if worker_image is not None else None
self._profile_locks_guard = threading.Lock()
self._profile_locks: dict[str, asyncio.Lock] = {}
@@ -389,7 +397,16 @@ class VastLeaseManager:
if lease.owner_id == self.owner_id
and lease.profile_id == slot_profile_id
and lease.profile_fingerprint == fingerprint
and lease.runtime_fingerprint == self.runtime_fingerprint
and (
(
self.worker_image is None
and lease.runtime_fingerprint == self.runtime_fingerprint
)
or (
self.worker_image is not None
and lease.worker_image in {None, self.worker_image}
)
)
and not lease.draining
]
for lease in sorted(candidates, key=lambda item: item.created_at_epoch):
@@ -406,6 +423,19 @@ class VastLeaseManager:
status_callback=status_callback,
)
refreshed = self._refresh_record(lease, instance)
if refreshed.runtime_fingerprint != self.runtime_fingerprint:
logger.info(
"Adopting compatible Vast lease instance_id=%d from runtime=%s "
"to runtime=%s for unchanged worker image.",
refreshed.instance_id,
refreshed.runtime_fingerprint[:12],
self.runtime_fingerprint[:12],
)
refreshed = replace(
refreshed,
runtime_fingerprint=self.runtime_fingerprint,
worker_image=self.worker_image,
)
self.registry.upsert(refreshed)
if refreshed.ready_for_work:
logger.info(
@@ -486,6 +516,7 @@ class VastLeaseManager:
profile_name=profile.profile_name,
profile_fingerprint=vast_profile_fingerprint(profile),
runtime_fingerprint=self.runtime_fingerprint,
worker_image=self.worker_image,
label=label,
actual_status="loading",
ssh_host=None,
vast_runtime.py
+16 -5
View File
@@ -13,11 +13,13 @@ from pathlib import PurePosixPath
from typing import Any, Callable, Mapping
if __package__:
from .remote_protocol import REMOTE_PROTOCOL_VERSION
from .remote.vast_watchdog import VastWatchdogSnapshot
from .vast_leases import VastLeaseRecord
from .vast_models import VastInstanceLaunchSpec, VastResourceProfile
from .vast_ssh import VastSshError, VastSshRunner, VastSshVolumeBackend
else: # pragma: no cover - direct debugging imports.
from remote_protocol import REMOTE_PROTOCOL_VERSION
from remote.vast_watchdog import VastWatchdogSnapshot
from vast_leases import VastLeaseRecord
from vast_models import VastInstanceLaunchSpec, VastResourceProfile
@@ -92,7 +94,6 @@ class VastRuntimeConfiguration:
environment = {
"COMFY_MODAL_REMOTE_WORKER": "1",
"COMFY_MODAL_LLM_EXECUTION_TARGET": "vast",
"COMFY_MODAL_RUNTIME_FINGERPRINT": self.runtime_fingerprint,
"COMFYUI_ROOT": str(self.remote_comfyui_root),
"COMFY_MODAL_COMFYUI_ROOT": str(self.remote_comfyui_root),
"COMFY_MODAL_LOCAL_STORAGE_ROOT": str(self.remote_storage_root),
@@ -159,10 +160,20 @@ class VastRuntimeManager:
else:
actual_fingerprint = str(info.get("runtime_fingerprint") or "")
if actual_fingerprint != self.configuration.runtime_fingerprint:
raise VastSshError(
"Vast worker image fingerprint mismatch: expected "
f"{self.configuration.runtime_fingerprint[:12]}, found "
f"{actual_fingerprint[:12] or 'empty'}."
actual_protocol = int(info.get("protocol_version") or 0)
if actual_protocol != REMOTE_PROTOCOL_VERSION:
raise VastSshError(
"Vast worker image fingerprint mismatch: expected "
f"{self.configuration.runtime_fingerprint[:12]}, found "
f"{actual_fingerprint[:12] or 'empty'}, with incompatible "
f"protocol {actual_protocol}."
)
logger.warning(
"Reusing Vast worker with compatible protocol despite source "
"fingerprint drift expected=%s actual=%s protocol=%d.",
self.configuration.runtime_fingerprint[:12],
actual_fingerprint[:12] or "empty",
actual_protocol,
)
if bool(info.get("worker_socket_ready")):
logger.info(
vast_service.py
+2 -1
View File
@@ -163,6 +163,7 @@ class VastService:
runtime_fingerprint=runtime_configuration.runtime_fingerprint,
launch_spec_factory=runtime_configuration.launch_spec,
startup_timeout_seconds=runtime_configuration.startup_timeout_seconds,
worker_image=runtime_configuration.image,
)
@classmethod
@@ -614,7 +615,7 @@ class VastService:
if lease.owner_id == self.settings.app_name
and lease.profile_id == profile.profile_id
and lease.profile_fingerprint == fingerprint
and lease.runtime_fingerprint == self.runtime_configuration.runtime_fingerprint
and lease.worker_image in {None, self.runtime_configuration.image}
and lease.ready_for_work
]
return min(matches, key=lambda lease: lease.created_at_epoch, default=None)
vast_ssh.py
+185 -1
View File
@@ -2,6 +2,7 @@
from __future__ import annotations
import hashlib
import json
import logging
import os
@@ -9,7 +10,9 @@ import random
import shlex
import subprocess
import time
import zipfile
from dataclasses import dataclass, field
from io import BytesIO
from pathlib import Path, PurePosixPath
from typing import BinaryIO, Callable, Sequence
@@ -76,6 +79,10 @@ class VastSshError(RuntimeError):
"""Raised when a direct Vast SSH or filesystem operation fails."""
class VastSshCancelledError(VastSshError):
"""Raised when the caller cancels an active direct SSH command."""
@dataclass(frozen=True)
class VastSshConnection:
"""Describe one Vast instance SSH endpoint without private credentials."""
@@ -223,6 +230,7 @@ class VastSshRunner:
input_file: Path | None = None,
timeout_seconds: float | None = None,
check: bool = True,
cancellation_check: Callable[[], bool] | None = None,
) -> VastSshCommandResult:
"""Execute one command with bounded transient transport retries."""
resolved_input_file = _resolved_input_file(input_payload, input_file)
@@ -233,6 +241,7 @@ class VastSshRunner:
input_payload=input_payload,
input_file=resolved_input_file,
timeout_seconds=timeout_seconds,
cancellation_check=cancellation_check,
)
if _is_transient_ssh_transport_failure(result):
if attempt >= self.retry_attempts:
@@ -253,6 +262,8 @@ class VastSshRunner:
result.stderr_text.strip() or result.stdout_text.strip(),
)
self.sleep(delay_seconds)
if cancellation_check is not None and cancellation_check():
raise VastSshCancelledError("Vast SSH command was cancelled.")
continue
logger.debug(
"Finished Vast SSH command host=%s command=%s returncode=%d "
@@ -275,6 +286,7 @@ class VastSshRunner:
input_payload: bytes | None,
input_file: Path | None,
timeout_seconds: float | None,
cancellation_check: Callable[[], bool] | None,
) -> VastSshCommandResult:
"""Execute one OpenSSH subprocess attempt, reopening streamed input files."""
try:
@@ -284,6 +296,7 @@ class VastSshRunner:
input_payload=input_payload,
input_handle=None,
timeout_seconds=timeout_seconds,
cancellation_check=cancellation_check,
)
else:
with input_file.open("rb") as input_handle:
@@ -292,6 +305,7 @@ class VastSshRunner:
input_payload=None,
input_handle=input_handle,
timeout_seconds=timeout_seconds,
cancellation_check=cancellation_check,
)
except FileNotFoundError as exc:
raise VastSshError(
@@ -338,6 +352,7 @@ class VastSshRunner:
input_payload: bytes | None,
input_handle: BinaryIO | None,
timeout_seconds: float | None,
cancellation_check: Callable[[], bool] | None,
) -> subprocess.CompletedProcess[bytes]:
"""Run OpenSSH with either a small byte payload or a streamed file handle."""
timeout = (
@@ -345,6 +360,14 @@ class VastSshRunner:
if timeout_seconds is None
else timeout_seconds
)
if cancellation_check is not None:
return self._run_cancellable_subprocess(
remote_argv,
input_payload=input_payload,
input_handle=input_handle,
timeout_seconds=timeout,
cancellation_check=cancellation_check,
)
if input_handle is not None:
return subprocess.run(
self.command(remote_argv),
@@ -363,6 +386,59 @@ class VastSshRunner:
timeout=timeout,
)
def _run_cancellable_subprocess(
self,
remote_argv: Sequence[str],
*,
input_payload: bytes | None,
input_handle: BinaryIO | None,
timeout_seconds: float,
cancellation_check: Callable[[], bool],
) -> subprocess.CompletedProcess[bytes]:
"""Run OpenSSH while promptly terminating it when cancellation is requested."""
if cancellation_check():
raise VastSshCancelledError("Vast SSH command was cancelled.")
process = subprocess.Popen(
self.command(remote_argv),
stdin=input_handle if input_handle is not None else subprocess.PIPE,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
)
started_at = time.monotonic()
payload = input_payload
while True:
if cancellation_check():
process.terminate()
try:
process.wait(timeout=5.0)
except subprocess.TimeoutExpired:
process.kill()
process.wait()
raise VastSshCancelledError("Vast SSH command was cancelled.")
remaining = timeout_seconds - (time.monotonic() - started_at)
if remaining <= 0:
process.kill()
stdout, stderr = process.communicate()
raise subprocess.TimeoutExpired(
self.command(remote_argv),
timeout_seconds,
output=stdout,
stderr=stderr,
)
try:
stdout, stderr = process.communicate(
input=payload,
timeout=min(0.25, remaining),
)
return subprocess.CompletedProcess(
args=self.command(remote_argv),
returncode=process.returncode,
stdout=stdout,
stderr=stderr,
)
except subprocess.TimeoutExpired:
payload = None
def popen(self, remote_argv: Sequence[str]) -> subprocess.Popen[bytes]:
"""Start one streaming direct SSH command."""
try:
@@ -387,6 +463,7 @@ class VastSshVolumeBackend:
runner: VastSshRunner
storage_root: PurePosixPath = DEFAULT_VAST_STORAGE_ROOT
_exists_cache: dict[str, bool] = field(default_factory=dict)
_materializer_path: PurePosixPath | None = field(init=False, default=None)
def __post_init__(self) -> None:
"""Require an absolute non-root storage location."""
@@ -429,6 +506,36 @@ class VastSshVolumeBackend:
)
self._exists_cache[normalized] = True
def put_file_cancellable(
self,
local_path: Path,
remote_path: str,
*,
cancellation_check: Callable[[], bool],
) -> None:
"""Stream one file while terminating OpenSSH promptly on cancellation."""
resolved_path = local_path.expanduser().resolve()
if not resolved_path.is_file():
raise FileNotFoundError(f"Asset not found: {resolved_path}")
size_bytes = resolved_path.stat().st_size
normalized = _validated_storage_path(remote_path)
try:
self.runner.run(
(
"python",
"-c",
_ATOMIC_STDIN_WRITER,
str(size_bytes),
str(self.storage_root / normalized),
),
input_file=resolved_path,
timeout_seconds=max(60.0, size_bytes / (4 * 1024 * 1024)),
cancellation_check=cancellation_check,
)
except VastSshCancelledError as exc:
raise InterruptedError("Vast file upload was cancelled.") from exc
self._exists_cache[normalized] = True
def put_bytes(self, payload: bytes, remote_path: str) -> None:
"""Atomically upload bytes under the storage root."""
self._put_payload(payload, remote_path)
@@ -452,6 +559,22 @@ class VastSshVolumeBackend:
token: str | None,
) -> bool:
"""Ask the Vast worker to fetch and verify one immutable Hugging Face file."""
return self.materialize_huggingface_file_cancellable(
source,
remote_path,
token=token,
cancellation_check=lambda: False,
)
def materialize_huggingface_file_cancellable(
self,
source: HuggingFaceAssetSource,
remote_path: str,
*,
token: str | None,
cancellation_check: Callable[[], bool],
) -> bool:
"""Run the current materializer bundle and allow prompt cancellation."""
normalized = _validated_storage_path(remote_path)
request_payload = json.dumps(
{
@@ -464,11 +587,17 @@ class VastSshVolumeBackend:
).encode("utf-8")
timeout_seconds = max(900.0, source.size_bytes / (2 * 1024 * 1024))
try:
materializer_path = self._ensure_materializer_bundle(
cancellation_check=cancellation_check,
)
self.runner.run(
("python", "-m", "remote.huggingface_materializer"),
("python", str(materializer_path)),
input_payload=request_payload,
timeout_seconds=timeout_seconds,
cancellation_check=cancellation_check,
)
except VastSshCancelledError as exc:
raise InterruptedError("Vast Hugging Face download was cancelled.") from exc
except VastSshError as exc:
logger.warning(
"Vast Hugging Face materialization failed for %s; falling back to SSH upload: %s",
@@ -479,6 +608,60 @@ class VastSshVolumeBackend:
self._exists_cache[normalized] = True
return True
def _ensure_materializer_bundle(
self,
*,
cancellation_check: Callable[[], bool],
) -> PurePosixPath:
"""Upload a tiny current-source zipapp independent of the worker image version."""
if self._materializer_path is not None:
return self._materializer_path
payload = _huggingface_materializer_bundle()
digest = hashlib.sha256(payload).hexdigest()
relative_path = f"runtime-tools/huggingface-materializer-{digest}.pyz"
normalized = _validated_storage_path(relative_path)
target = self.storage_root / normalized
if not self.exists(normalized):
self.runner.run(
(
"python",
"-c",
_ATOMIC_STDIN_WRITER,
str(len(payload)),
str(target),
),
input_payload=payload,
timeout_seconds=60.0,
cancellation_check=cancellation_check,
)
self._exists_cache[normalized] = True
self._materializer_path = target
return target
def _huggingface_materializer_bundle() -> bytes:
"""Build an executable zip containing the current verified-download implementation."""
repo_root = Path(__file__).resolve().parent
members = {
"__main__.py": (
b"from remote.huggingface_materializer import main\n"
b"raise SystemExit(main())\n"
),
"remote/__init__.py": b"",
"remote/huggingface_materializer.py": (
repo_root / "remote" / "huggingface_materializer.py"
).read_bytes(),
"huggingface_assets.py": (repo_root / "huggingface_assets.py").read_bytes(),
}
output = BytesIO()
with zipfile.ZipFile(output, "w", compression=zipfile.ZIP_DEFLATED) as archive:
for member_name, member_payload in members.items():
member = zipfile.ZipInfo(member_name, date_time=(1980, 1, 1, 0, 0, 0))
member.compress_type = zipfile.ZIP_DEFLATED
member.external_attr = 0o644 << 16
archive.writestr(member, member_payload)
return output.getvalue()
def vast_connection_from_lease(
*,
@@ -517,6 +700,7 @@ def _validated_storage_path(remote_path: str) -> str:
__all__ = [
"DEFAULT_VAST_STORAGE_ROOT",
"VastSshCommandResult",
"VastSshCancelledError",
"VastSshConnection",
"VastSshError",
"VastSshRunner",
web/modal_toggle.js
+18 -2
View File
@@ -11,6 +11,10 @@ const MODAL_PROGRESS_STATE_ROUTE = MODAL_ROUTE.replace(/\/queue_prompt$/, "/prog
const MODAL_CONTAINER_STATUS_ROUTE = MODAL_ROUTE.replace(/\/queue_prompt$/, "/container_status");
const MODAL_DELETE_CACHES_ROUTE = MODAL_ROUTE.replace(/\/queue_prompt$/, "/delete_caches");
const MODAL_DELETE_VOLUME_ROUTE = MODAL_ROUTE.replace(/\/queue_prompt$/, "/delete_volume");
const MODAL_CANCEL_PREPARATION_ROUTE = MODAL_ROUTE.replace(
/\/queue_prompt$/,
"/cancel_preparation",
);
const COMFY_QUEUE_ROUTE = "/queue";
const COMFY_HISTORY_ROUTE = "/history";
const INTERNAL_NODE_PREFIX = "ModalUniversalExecutor";
@@ -1592,12 +1596,24 @@ function patchInterruptFeedback() {
return;
}
const originalFetchApi = api.fetchApi;
api.fetchApi = function modalFetchApi(resource, options) {
api.fetchApi = async function modalFetchApi(resource, options) {
const promptIds = promptIdsFromInterruptRequest(resource, options);
for (const promptId of promptIds) {
markPromptCancellationRequested(promptId);
}
return originalFetchApi.apply(this, arguments);
if (promptIds.length > 0) {
await Promise.all(promptIds.map(async (promptId) => {
try {
await originalFetchApi.call(this, MODAL_CANCEL_PREPARATION_ROUTE, {
method: "POST",
body: JSON.stringify({ prompt_id: promptId }),
});
} catch (error) {
console.warn(`Unable to cancel remote preparation for prompt ${promptId}.`, error);
}
}));
}
return originalFetchApi.call(this, resource, options);
};
api.__modalInterruptFeedbackPatched = true;
}