fix: avoid blocking modal interrupt writes in async paths
This commit is contained in:
Ken Simpson
+119
-7
@@ -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():
|
||||
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user