320 lines
12 KiB
Python
320 lines
12 KiB
Python
"""Tests for the routes queue boundary."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from api_intercept_test_support import * # noqa: F401,F403
|
|
|
|
def test_queue_prompt_route_does_not_warm_modal_at_queue_time(
|
|
api_intercept_module: Any,
|
|
prompt_interception_module: Any,
|
|
queue_bridge_module: Any,
|
|
remote_modal_app_module: Any,
|
|
settings_module: Any,
|
|
monkeypatch: Any,
|
|
tmp_path: Path,
|
|
) -> None:
|
|
"""Accepting a queued prompt should not launch Modal warmup containers."""
|
|
|
|
class FakeRoutes:
|
|
"""Capture aiohttp route registrations."""
|
|
|
|
def __init__(self) -> None:
|
|
"""Initialize the route handler map."""
|
|
self.handlers: dict[str, Any] = {}
|
|
|
|
def post(self, path: str) -> Any:
|
|
"""Return a decorator that records one POST handler."""
|
|
|
|
def register(handler: Any) -> Any:
|
|
"""Store the decorated handler unchanged."""
|
|
self.handlers[path] = handler
|
|
return handler
|
|
|
|
return register
|
|
|
|
class FakePromptQueue:
|
|
"""Minimal prompt queue sink."""
|
|
|
|
def __init__(self) -> None:
|
|
"""Initialize captured queue items."""
|
|
self.items: list[tuple[Any, ...]] = []
|
|
|
|
def put(self, item: tuple[Any, ...]) -> None:
|
|
"""Record one queued prompt item."""
|
|
self.items.append(item)
|
|
|
|
class FakePromptServer:
|
|
"""Minimal PromptServer double with route registration."""
|
|
|
|
def __init__(self) -> None:
|
|
"""Initialize routing and queue state."""
|
|
self.number = 0
|
|
self.routes = FakeRoutes()
|
|
self.prompt_queue = FakePromptQueue()
|
|
|
|
def trigger_on_prompt(self, json_data: dict[str, Any]) -> dict[str, Any]:
|
|
"""Return the prompt unchanged."""
|
|
return json_data
|
|
|
|
class FakeRequest:
|
|
"""Minimal aiohttp request double."""
|
|
|
|
async def json(self) -> dict[str, Any]:
|
|
"""Return one Modal-marked prompt request."""
|
|
return {
|
|
"prompt_id": "prompt-queue-warmup",
|
|
"prompt": {"1": {"class_type": "RemoteImage", "inputs": {}}},
|
|
"extra_data": {
|
|
"extra_pnginfo": {
|
|
"workflow": {
|
|
"extra": {"comfy_modal": {"gpu": "B300"}},
|
|
"nodes": [
|
|
{
|
|
"id": 1,
|
|
"type": "RemoteImage",
|
|
"properties": {"is_modal_remote": True},
|
|
}
|
|
]
|
|
}
|
|
}
|
|
},
|
|
}
|
|
|
|
class FakeExecutionModule:
|
|
"""Minimal execution module exposing prompt validation."""
|
|
|
|
SENSITIVE_EXTRA_DATA_KEYS: tuple[str, ...] = ()
|
|
|
|
@staticmethod
|
|
async def validate_prompt(
|
|
prompt_id: str,
|
|
prompt: dict[str, Any],
|
|
partial_execution_targets: Any,
|
|
) -> tuple[bool, None, list[str], list[Any]]:
|
|
"""Accept the supplied prompt with one fake execution target."""
|
|
return True, None, ["1"], []
|
|
|
|
def fail_queue_time_warmup(*_args: Any, **_kwargs: Any) -> int:
|
|
"""Fail the test if queue handling tries to launch proactive warmup."""
|
|
raise AssertionError("queue route must not schedule Modal warmup")
|
|
|
|
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="remote",
|
|
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,
|
|
)
|
|
prompt_server = FakePromptServer()
|
|
summary = api_intercept_module.RewriteSummary(
|
|
remote_node_ids=["1"],
|
|
remote_component_ids=["1"],
|
|
component_node_ids_by_representative={"1": ["1"]},
|
|
sandwiched_local_node_ids=["4"],
|
|
component_execution_stages=[["1"]],
|
|
estimated_max_parallel_requests=1,
|
|
max_parallel_requests_upper_bound=1,
|
|
execution_assignments_by_representative={
|
|
"1": api_intercept_module.ExecutionAssignment(
|
|
environment_id="modal:B300",
|
|
provider=api_intercept_module.ExecutionProvider.MODAL,
|
|
predicted_cost_usd=0.01,
|
|
predicted_completion_seconds=10.0,
|
|
)
|
|
},
|
|
)
|
|
observed_rewrite_settings: list[Any] = []
|
|
|
|
def capture_rewrite_settings(**kwargs: Any) -> tuple[dict[str, Any], Any]:
|
|
"""Capture the workflow-derived settings passed into prompt rewriting."""
|
|
observed_rewrite_settings.append(kwargs["settings"])
|
|
return kwargs["prompt"], summary
|
|
|
|
monkeypatch.setattr(api_intercept_module, "_ROUTE_REGISTERED", False)
|
|
monkeypatch.setattr(
|
|
api_intercept_module,
|
|
"_get_server_module",
|
|
lambda: SimpleNamespace(PromptServer=SimpleNamespace(instance=prompt_server)),
|
|
)
|
|
monkeypatch.setattr(queue_bridge_module, "_get_execution_module", lambda: FakeExecutionModule)
|
|
monkeypatch.setattr(api_intercept_module, "_emit_modal_status", lambda **_kwargs: None)
|
|
monkeypatch.setattr(
|
|
prompt_interception_module,
|
|
"rewrite_prompt_for_modal",
|
|
capture_rewrite_settings,
|
|
)
|
|
monkeypatch.setattr(remote_modal_app_module, "ensure_remote_warm_capacity", fail_queue_time_warmup)
|
|
|
|
api_intercept_module.setup_modal_queue_route(
|
|
prompt_server=prompt_server,
|
|
sync_engine=object(),
|
|
settings=settings,
|
|
)
|
|
response = asyncio.run(prompt_server.routes.handlers["/modal/queue_prompt"](FakeRequest()))
|
|
|
|
response_payload = json.loads(response.text)
|
|
assert response_payload["prompt_id"] == "prompt-queue-warmup"
|
|
assert response_payload["modal_gpu"] == "B300"
|
|
assert response_payload["modal_remote_node_ids"] == ["1"]
|
|
assert response_payload["modal_sandwiched_local_node_ids"] == ["4"]
|
|
assert observed_rewrite_settings[0].modal_gpu == "B300"
|
|
queued_extra_data = prompt_server.prompt_queue.items[0][3]
|
|
assert queued_extra_data["modal"]["gpu"] == "B300"
|
|
assert queued_extra_data["extra_pnginfo"][
|
|
api_intercept_module.MODAL_PROMPT_ID_EXTRA_PNGINFO_KEY
|
|
] == "prompt-queue-warmup"
|
|
assert len(prompt_server.prompt_queue.items) == 1
|
|
|
|
def test_queue_prompt_route_without_remote_nodes_skips_modal_status_and_rewrite(
|
|
api_intercept_module: Any,
|
|
prompt_interception_module: Any,
|
|
queue_bridge_module: Any,
|
|
settings_module: Any,
|
|
monkeypatch: Any,
|
|
tmp_path: Path,
|
|
) -> None:
|
|
"""Prompts with no Modal-enabled workflow nodes should queue without Modal UI setup."""
|
|
|
|
class FakeRoutes:
|
|
"""Capture aiohttp route registrations."""
|
|
|
|
def __init__(self) -> None:
|
|
"""Initialize the route handler map."""
|
|
self.handlers: dict[str, Any] = {}
|
|
|
|
def post(self, path: str) -> Any:
|
|
"""Return a decorator that records one POST handler."""
|
|
|
|
def register(handler: Any) -> Any:
|
|
"""Store the decorated handler unchanged."""
|
|
self.handlers[path] = handler
|
|
return handler
|
|
|
|
return register
|
|
|
|
class FakePromptQueue:
|
|
"""Minimal prompt queue sink."""
|
|
|
|
def __init__(self) -> None:
|
|
"""Initialize captured queue items."""
|
|
self.items: list[tuple[Any, ...]] = []
|
|
|
|
def put(self, item: tuple[Any, ...]) -> None:
|
|
"""Record one queued prompt item."""
|
|
self.items.append(item)
|
|
|
|
class FakePromptServer:
|
|
"""Minimal PromptServer double with route registration."""
|
|
|
|
def __init__(self) -> None:
|
|
"""Initialize routing and queue state."""
|
|
self.number = 0
|
|
self.routes = FakeRoutes()
|
|
self.prompt_queue = FakePromptQueue()
|
|
|
|
def trigger_on_prompt(self, json_data: dict[str, Any]) -> dict[str, Any]:
|
|
"""Return the prompt unchanged."""
|
|
return json_data
|
|
|
|
class FakeRequest:
|
|
"""Minimal aiohttp request double."""
|
|
|
|
async def json(self) -> dict[str, Any]:
|
|
"""Return one ordinary prompt request."""
|
|
return {
|
|
"prompt_id": "prompt-no-modal",
|
|
"prompt": {"1": {"class_type": "LocalImage", "inputs": {}}},
|
|
"extra_data": {"extra_pnginfo": {"workflow": {"nodes": []}}},
|
|
}
|
|
|
|
class FakeExecutionModule:
|
|
"""Minimal execution module exposing prompt validation."""
|
|
|
|
SENSITIVE_EXTRA_DATA_KEYS: tuple[str, ...] = ()
|
|
|
|
@staticmethod
|
|
async def validate_prompt(
|
|
prompt_id: str,
|
|
prompt: dict[str, Any],
|
|
partial_execution_targets: Any,
|
|
) -> tuple[bool, None, list[str], list[Any]]:
|
|
"""Accept the supplied prompt with one fake execution target."""
|
|
return True, None, ["1"], []
|
|
|
|
def fail_modal_status(*_args: Any, **_kwargs: Any) -> None:
|
|
"""Fail if the no-remote fast path emits Modal UI state."""
|
|
raise AssertionError("no-remote prompts must not emit Modal status")
|
|
|
|
def fail_rewrite(*_args: Any, **_kwargs: Any) -> tuple[dict[str, Any], Any]:
|
|
"""Fail if the no-remote fast path enters Modal prompt rewriting."""
|
|
raise AssertionError("no-remote prompts must not be rewritten for Modal")
|
|
|
|
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="remote",
|
|
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,
|
|
)
|
|
prompt_server = FakePromptServer()
|
|
monkeypatch.setattr(api_intercept_module, "_ROUTE_REGISTERED", False)
|
|
monkeypatch.setattr(
|
|
api_intercept_module,
|
|
"_get_server_module",
|
|
lambda: SimpleNamespace(PromptServer=SimpleNamespace(instance=prompt_server)),
|
|
)
|
|
monkeypatch.setattr(queue_bridge_module, "_get_execution_module", lambda: FakeExecutionModule)
|
|
monkeypatch.setattr(api_intercept_module, "_emit_modal_status", fail_modal_status)
|
|
monkeypatch.setattr(
|
|
prompt_interception_module,
|
|
"rewrite_prompt_for_modal",
|
|
fail_rewrite,
|
|
)
|
|
|
|
api_intercept_module.setup_modal_queue_route(
|
|
prompt_server=prompt_server,
|
|
sync_engine=object(),
|
|
settings=settings,
|
|
)
|
|
response = asyncio.run(prompt_server.routes.handlers["/modal/queue_prompt"](FakeRequest()))
|
|
|
|
response_payload = json.loads(response.text)
|
|
assert response_payload["prompt_id"] == "prompt-no-modal"
|
|
assert "modal_remote_node_ids" not in response_payload
|
|
assert len(prompt_server.prompt_queue.items) == 1
|
|
|
|
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"
|
|
)
|
|
|