diff --git a/remote/modal_app.py b/remote/modal_app.py index 679d1c0..8d4b44b 100644 --- a/remote/modal_app.py +++ b/remote/modal_app.py @@ -2632,6 +2632,39 @@ def _lookup_modal_interrupt_store() -> Any | None: return interrupt_store +def _remote_interrupt_flag_value() -> dict[str, float]: + """Return the shared Modal interrupt-store value for one cancellation request.""" + return {"requested_at": time.time()} + + +def _write_remote_interrupt_flag(interrupt_store: Any, prompt_id: str, component_id: str) -> None: + """Write one remote cancellation request with the blocking Modal Dict API.""" + interrupt_store.put( + _remote_interrupt_flag_key(prompt_id, component_id), + _remote_interrupt_flag_value(), + ) + + +async def _write_remote_interrupt_flag_async( + interrupt_store: Any, + prompt_id: str, + component_id: str, +) -> None: + """Write one remote cancellation request without blocking the async caller.""" + put_method = getattr(interrupt_store, "put", None) + put_async = getattr(put_method, "aio", None) + if callable(put_async): + result = put_async( + _remote_interrupt_flag_key(prompt_id, component_id), + _remote_interrupt_flag_value(), + ) + if inspect.isawaitable(result): + await result + return + + await asyncio.to_thread(_write_remote_interrupt_flag, interrupt_store, prompt_id, component_id) + + def _request_remote_interrupt(payload: dict[str, Any]) -> bool: """Write one remote cancellation request into the shared Modal interrupt store.""" _abandon_local_modal_workflow_gate(payload, "local interrupt requested") @@ -2640,10 +2673,24 @@ def _request_remote_interrupt(payload: dict[str, Any]) -> bool: return False prompt_id, component_id = _remote_interrupt_key(payload) - interrupt_store.put( - _remote_interrupt_flag_key(prompt_id, component_id), - {"requested_at": time.time()}, + _write_remote_interrupt_flag(interrupt_store, prompt_id, component_id) + logger.info( + "Propagated local interrupt to Modal prompt=%s component=%s through shared control state.", + prompt_id, + component_id, ) + return True + + +async def _request_remote_interrupt_async(payload: dict[str, Any]) -> bool: + """Write one remote cancellation request into the shared Modal interrupt store asynchronously.""" + _abandon_local_modal_workflow_gate(payload, "local interrupt requested") + interrupt_store = await asyncio.to_thread(_lookup_modal_interrupt_store) + if interrupt_store is None: + return False + + prompt_id, component_id = _remote_interrupt_key(payload) + await _write_remote_interrupt_flag_async(interrupt_store, prompt_id, component_id) logger.info( "Propagated local interrupt to Modal prompt=%s component=%s through shared control state.", prompt_id, @@ -2723,6 +2770,32 @@ def request_remote_modal_prompt_interrupt(prompt_id: str) -> bool: return True +async def request_remote_modal_prompt_interrupt_async(prompt_id: str) -> bool: + """Request cancellation for every active Modal invocation belonging to one prompt asynchronously.""" + normalized_prompt_id = str(prompt_id) + _abandon_local_modal_workflow_gate( + {"prompt_id": normalized_prompt_id}, + "prompt-level interrupt requested", + ) + with _ACTIVE_REMOTE_INVOCATIONS_LOCK: + invocations = list(_ACTIVE_REMOTE_INVOCATIONS_BY_PROMPT.get(normalized_prompt_id, {}).values()) + if not invocations: + return False + + logger.info( + "Requesting async remote Modal cancellation for prompt=%s across %d active component(s).", + normalized_prompt_id, + len(invocations), + ) + for invocation in invocations: + if invocation.cancellation_event is not None: + invocation.cancellation_event.set() + await _request_remote_interrupt_async( + {"prompt_id": invocation.prompt_id, "component_id": invocation.component_id} + ) + return True + + def _sync_local_interrupt_to_cancellation_event( payload: dict[str, Any], cancellation_event: threading.Event | None, @@ -2816,6 +2889,38 @@ def _handle_modal_wait_cancellation( return interrupt_sent, cancellation_started_at +async def _handle_modal_wait_cancellation_async( + payload: dict[str, Any], + cancellation_event: threading.Event, + *, + interrupt_sent: bool, + cancellation_started_at: float | None, +) -> tuple[bool, float | None]: + """Propagate and bound local waiting after cancellation during an async Modal wait phase.""" + if not _sync_local_interrupt_to_cancellation_event(payload, cancellation_event): + return interrupt_sent, cancellation_started_at + + if not interrupt_sent: + await _request_remote_interrupt_async(payload) + return True, time.monotonic() + + if cancellation_started_at is None: + return interrupt_sent, time.monotonic() + + grace_seconds = max(0.0, get_settings().remote_cancel_grace_seconds) + if time.monotonic() - cancellation_started_at >= grace_seconds: + logger.info( + "Modal component=%s did not reach a cancellable remote call within %.3fs of local interrupt; releasing the local prompt while remote cancellation continues.", + payload.get("component_id"), + grace_seconds, + ) + raise ModalRemoteInvocationError( + "Remote Modal call did not reach a cancellable remote phase after local interrupt propagation." + ) + + return interrupt_sent, cancellation_started_at + + def _invoke_remote_call_with_interrupts( *, payload: dict[str, Any], @@ -4291,19 +4396,26 @@ async def _invoke_bound_remote_engine_async( cancellation_event, ) wrapped_future = asyncio.wrap_future(future) + interrupt_sent = False + cancellation_started_at: float | None = None try: while True: try: response = await asyncio.wait_for(asyncio.shield(wrapped_future), timeout=0.1) break except asyncio.TimeoutError: - _sync_local_interrupt_to_cancellation_event(payload, cancellation_event) + interrupt_sent, cancellation_started_at = await _handle_modal_wait_cancellation_async( + payload, + cancellation_event, + interrupt_sent=interrupt_sent, + cancellation_started_at=cancellation_started_at, + ) continue except asyncio.CancelledError: cancellation_event.set() prompt_id = str(payload.get("prompt_id")) if payload.get("prompt_id") is not None else None if prompt_id is not None: - request_remote_modal_prompt_interrupt(prompt_id) + await request_remote_modal_prompt_interrupt_async(prompt_id) raise except Exception: if cancellation_event.is_set() or _local_processing_interrupted(): @@ -5497,7 +5609,7 @@ async def invoke_remote_engine_async(payload: dict[str, Any], kwargs_payload: by response = await asyncio.wait_for(asyncio.shield(wrapped_future), timeout=0.1) break except asyncio.TimeoutError: - interrupt_sent, cancellation_started_at = _handle_modal_wait_cancellation( + interrupt_sent, cancellation_started_at = await _handle_modal_wait_cancellation_async( payload, cancellation_event, interrupt_sent=interrupt_sent, @@ -5508,7 +5620,7 @@ async def invoke_remote_engine_async(payload: dict[str, Any], kwargs_payload: by cancellation_event.set() prompt_id = str(payload.get("prompt_id")) if payload.get("prompt_id") is not None else None if prompt_id is not None: - request_remote_modal_prompt_interrupt(prompt_id) + await request_remote_modal_prompt_interrupt_async(prompt_id) raise except Exception: if cancellation_event.is_set() or _local_processing_interrupted(): diff --git a/tests/test_modal_executor_node.py b/tests/test_modal_executor_node.py index 0794ed1..5d37c60 100644 --- a/tests/test_modal_executor_node.py +++ b/tests/test_modal_executor_node.py @@ -5803,6 +5803,80 @@ def test_remote_modal_interrupt_callback_writes_shared_control_flag( assert isinstance(interrupt_value["requested_at"], float) +def test_request_remote_interrupt_async_uses_shared_control_async_put( + remote_modal_app_module: Any, + monkeypatch: Any, +) -> None: + """Async local cancellation should use Modal Dict.put.aio instead of blocking put.""" + + class FakePut: + """Modal Dict put double that records sync and async writes separately.""" + + def __init__(self) -> None: + """Initialize captured writes.""" + self.sync_calls: list[tuple[str, Any]] = [] + self.async_calls: list[tuple[str, Any]] = [] + + def __call__(self, key: str, value: Any, *, skip_if_exists: bool = False) -> bool: + """Record one blocking write.""" + del skip_if_exists + self.sync_calls.append((key, value)) + return True + + async def aio(self, key: str, value: Any, *, skip_if_exists: bool = False) -> bool: + """Record one async write.""" + del skip_if_exists + self.async_calls.append((key, value)) + return True + + class FakeInterruptStore: + """Simple Modal Dict double exposing put.aio.""" + + def __init__(self) -> None: + """Initialize the fake async put handle.""" + self.put = FakePut() + + interrupt_store = FakeInterruptStore() + + class FakeModalDict: + """Minimal modal.Dict shim that returns the fake interrupt store.""" + + @staticmethod + def from_name( + name: str, + *, + environment_name: str | None = None, + create_if_missing: bool = False, + client: Any | None = None, + ) -> FakeInterruptStore: + return interrupt_store + + monkeypatch.setattr( + remote_modal_app_module, + "modal", + types.SimpleNamespace(Dict=FakeModalDict), + ) + monkeypatch.setenv("COMFY_MODAL_INTERRUPT_DICT_NAME", "shared-interrupts") + remote_modal_app_module.get_settings.cache_clear() + remote_modal_app_module._MODAL_INTERRUPT_DICTS.clear() + try: + wrote_interrupt = asyncio.run( + remote_modal_app_module._request_remote_interrupt_async( + {"prompt_id": "prompt-1", "component_id": "component-2"} + ) + ) + finally: + remote_modal_app_module.get_settings.cache_clear() + remote_modal_app_module._MODAL_INTERRUPT_DICTS.clear() + + assert wrote_interrupt is True + assert interrupt_store.put.sync_calls == [] + assert len(interrupt_store.put.async_calls) == 1 + interrupt_key, interrupt_value = interrupt_store.put.async_calls[0] + assert interrupt_key == "prompt-1:component-2" + assert isinstance(interrupt_value["requested_at"], float) + + def test_request_remote_modal_prompt_interrupt_cancels_active_components( remote_modal_app_module: Any, ) -> None: