444 lines
15 KiB
Python
444 lines
15 KiB
Python
"""Tests for the prompt interception boundary."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from api_intercept_test_support import * # noqa: F401,F403
|
|
|
|
def test_queue_prompt_json_includes_resolved_modal_metadata(
|
|
api_intercept_module: Any,
|
|
queue_bridge_module: Any,
|
|
monkeypatch: Any,
|
|
) -> None:
|
|
"""Successful queue responses should include resolved remote nodes and component membership."""
|
|
|
|
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 for queue-response tests."""
|
|
|
|
def __init__(self) -> None:
|
|
"""Initialize queue state."""
|
|
self.number = 0
|
|
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 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"], []
|
|
|
|
monkeypatch.setattr(queue_bridge_module, "_get_execution_module", lambda: FakeExecutionModule)
|
|
prompt_server = FakePromptServer()
|
|
|
|
response = asyncio.run(
|
|
queue_bridge_module._queue_prompt_json(
|
|
prompt_server,
|
|
{
|
|
"prompt_id": "prompt-1",
|
|
"prompt": {"1": {"class_type": "Anything", "inputs": {}}},
|
|
"extra_data": {},
|
|
},
|
|
modal_response_payload={
|
|
"modal_remote_node_ids": ["1", "2"],
|
|
"modal_components": [
|
|
{
|
|
"representative_node_id": "1",
|
|
"node_ids": ["1", "2"],
|
|
}
|
|
],
|
|
},
|
|
)
|
|
)
|
|
|
|
response_payload = json.loads(response.text)
|
|
assert response_payload["prompt_id"] == "prompt-1"
|
|
assert response_payload["modal_remote_node_ids"] == ["1", "2"]
|
|
assert response_payload["modal_components"] == [
|
|
{
|
|
"representative_node_id": "1",
|
|
"node_ids": ["1", "2"],
|
|
}
|
|
]
|
|
|
|
def test_queue_prompt_json_logs_rewritten_modal_diagnostics_on_validation_failure(
|
|
api_intercept_module: Any,
|
|
queue_bridge_module: Any,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""Validation failures for Modal prompts should log the rewritten dependency graph."""
|
|
|
|
class FakePromptQueue:
|
|
"""Minimal prompt queue sink."""
|
|
|
|
def put(self, _item: tuple[Any, ...]) -> None:
|
|
"""Fail if an invalid prompt reaches the queue."""
|
|
raise AssertionError("invalid prompt must not be queued")
|
|
|
|
class FakePromptServer:
|
|
"""Minimal PromptServer double for validation-failure tests."""
|
|
|
|
number = 0
|
|
prompt_queue = FakePromptQueue()
|
|
|
|
def trigger_on_prompt(self, json_data: dict[str, Any]) -> dict[str, Any]:
|
|
"""Return the prompt unchanged."""
|
|
return json_data
|
|
|
|
class FakeExecutionModule:
|
|
"""Minimal execution module that rejects the prompt."""
|
|
|
|
SENSITIVE_EXTRA_DATA_KEYS: tuple[str, ...] = ()
|
|
|
|
@staticmethod
|
|
async def validate_prompt(
|
|
prompt_id: str,
|
|
prompt: dict[str, Any],
|
|
partial_execution_targets: Any,
|
|
) -> tuple[bool, dict[str, Any], list[str], dict[str, Any]]:
|
|
"""Reject the supplied prompt with a dependency-cycle shaped error."""
|
|
del prompt_id, prompt, partial_execution_targets
|
|
return (
|
|
False,
|
|
{
|
|
"type": "execution_error",
|
|
"message": "Dependency cycle detected",
|
|
"details": "",
|
|
"extra_info": {},
|
|
},
|
|
[],
|
|
{},
|
|
)
|
|
|
|
prompt = {
|
|
"1": {
|
|
"class_type": "ModalUniversalExecutor_a",
|
|
"inputs": {"remote_input_0": ["2", 0]},
|
|
},
|
|
"2": {
|
|
"class_type": "ModalUniversalExecutor_b",
|
|
"inputs": {"remote_input_0": ["1", 0]},
|
|
},
|
|
}
|
|
|
|
diagnostic_calls: list[dict[str, Any]] = []
|
|
|
|
def record_diagnostics(**kwargs: Any) -> None:
|
|
"""Record one rewritten-prompt diagnostics request."""
|
|
diagnostic_calls.append(dict(kwargs))
|
|
|
|
monkeypatch.setattr(queue_bridge_module, "_get_execution_module", lambda: FakeExecutionModule)
|
|
monkeypatch.setattr(
|
|
queue_bridge_module,
|
|
"_log_modal_rewritten_prompt_diagnostics",
|
|
record_diagnostics,
|
|
)
|
|
|
|
response = asyncio.run(
|
|
queue_bridge_module._queue_prompt_json(
|
|
FakePromptServer(),
|
|
{
|
|
"prompt_id": "prompt-cycle",
|
|
"prompt": prompt,
|
|
"extra_data": {
|
|
"modal": {
|
|
"remote_component_ids": ["1", "2"],
|
|
}
|
|
},
|
|
},
|
|
)
|
|
)
|
|
|
|
assert response.status == 400
|
|
assert diagnostic_calls == [
|
|
{
|
|
"prompt_id": "prompt-cycle",
|
|
"prompt": prompt,
|
|
"reason": "comfy_validation_failure",
|
|
"level": api_intercept_module.logging.WARNING,
|
|
}
|
|
]
|
|
|
|
def test_remote_execution_configurator_identity_is_preserved(
|
|
api_intercept_module: Any,
|
|
) -> None:
|
|
"""Queue-time UI events should address the exact serialized configurator node."""
|
|
prompt = {
|
|
"12": {"class_type": "KSampler", "inputs": {}},
|
|
"99": {
|
|
"class_type": "RemoteExecutionConfigurator",
|
|
"inputs": {"configurations.configuration_0": ["20", 0]},
|
|
},
|
|
}
|
|
|
|
assert (
|
|
api_intercept_module._remote_execution_configurator_node_id(prompt)
|
|
== "99"
|
|
)
|
|
assert api_intercept_module._remote_execution_configurator_node_id(
|
|
{
|
|
**prompt,
|
|
"100": {
|
|
"class_type": "RemoteExecutionConfigurator",
|
|
"inputs": {},
|
|
},
|
|
}
|
|
) is None
|
|
|
|
def test_modal_prompt_rewrite_keeps_event_loop_responsive(
|
|
api_intercept_module: Any,
|
|
prompt_interception_module: Any,
|
|
monkeypatch: Any,
|
|
) -> None:
|
|
"""Hashing and upload preparation should execute outside the ComfyUI event loop."""
|
|
rewrite_started = threading.Event()
|
|
release_rewrite = threading.Event()
|
|
|
|
def blocking_rewrite(**kwargs: Any) -> tuple[dict[str, Any], Any]:
|
|
"""Hold one fake rewrite until the async test proves the loop is responsive."""
|
|
rewrite_started.set()
|
|
assert release_rewrite.wait(timeout=1.0)
|
|
return kwargs["prompt"], api_intercept_module.RewriteSummary()
|
|
|
|
monkeypatch.setattr(
|
|
prompt_interception_module,
|
|
"rewrite_prompt_for_modal",
|
|
blocking_rewrite,
|
|
)
|
|
|
|
async def run_test() -> None:
|
|
"""Run the blocking rewrite and an independent event-loop callback together."""
|
|
rewrite_task = asyncio.create_task(
|
|
api_intercept_module.rewrite_prompt_for_modal_async(
|
|
prompt={"1": {"class_type": "RemoteImage", "inputs": {}}},
|
|
workflow=None,
|
|
)
|
|
)
|
|
deadline = asyncio.get_running_loop().time() + 1.0
|
|
while not rewrite_started.is_set():
|
|
if asyncio.get_running_loop().time() >= deadline:
|
|
raise AssertionError("background rewrite did not start")
|
|
await asyncio.sleep(0)
|
|
|
|
loop_progress: list[str] = []
|
|
asyncio.get_running_loop().call_soon(loop_progress.append, "responsive")
|
|
await asyncio.sleep(0)
|
|
assert loop_progress == ["responsive"]
|
|
assert not rewrite_task.done()
|
|
|
|
release_rewrite.set()
|
|
await asyncio.wait_for(rewrite_task, timeout=1.0)
|
|
|
|
asyncio.run(run_test())
|
|
|
|
def test_environment_setup_status_callback_preserves_environment_identity(
|
|
api_intercept_module: Any,
|
|
) -> None:
|
|
"""Environment setup updates should also retain prompt-wide progress."""
|
|
prompt_updates: list[tuple[str, int | None, int | None]] = []
|
|
environment_updates: list[tuple[str, str, int | None, int | None]] = []
|
|
callback = api_intercept_module._environment_setup_status_callback(
|
|
"vast:48602895",
|
|
lambda message, current, total: prompt_updates.append(
|
|
(message, current, total)
|
|
),
|
|
lambda environment_id, message, current, total: environment_updates.append(
|
|
(environment_id, message, current, total)
|
|
),
|
|
)
|
|
|
|
assert callback is not None
|
|
callback("Uploading asset", 3, 10)
|
|
|
|
assert prompt_updates == [("Uploading asset", 3, 10)]
|
|
assert environment_updates == [
|
|
("vast:48602895", "Uploading asset", 3, 10)
|
|
]
|
|
|
|
def test_remote_environment_assets_are_prepared_in_parallel(
|
|
api_intercept_module: Any,
|
|
) -> None:
|
|
"""Distinct environments should enter asset preparation concurrently."""
|
|
barrier = threading.Barrier(2)
|
|
state_lock = threading.Lock()
|
|
active_environment_count = 0
|
|
maximum_active_environment_count = 0
|
|
environment_events: list[tuple[str, str]] = []
|
|
|
|
class FakeSyncEngine:
|
|
"""Block custom-node setup until both environment workers are active."""
|
|
|
|
def __init__(self) -> None:
|
|
"""Record the environment-local preparation order."""
|
|
self.calls: list[str] = []
|
|
|
|
def preflight_r2_access(self, *, status_callback: Any) -> None:
|
|
"""Represent an environment without configured R2 backing."""
|
|
del status_callback
|
|
self.calls.append("r2_preflight")
|
|
|
|
def sync_custom_nodes_directory(self, *, status_callback: Any) -> None:
|
|
"""Prove the second environment starts before the first can finish."""
|
|
nonlocal active_environment_count, maximum_active_environment_count
|
|
self.calls.append("custom_nodes")
|
|
with state_lock:
|
|
active_environment_count += 1
|
|
maximum_active_environment_count = max(
|
|
maximum_active_environment_count,
|
|
active_environment_count,
|
|
)
|
|
status_callback("Uploading custom nodes", None, None)
|
|
try:
|
|
barrier.wait(timeout=2.0)
|
|
finally:
|
|
with state_lock:
|
|
active_environment_count -= 1
|
|
|
|
def create_request_asset_cache(self, values: Any) -> object:
|
|
"""Consume the environment-local input plan and return a sentinel cache."""
|
|
tuple(values)
|
|
self.calls.append("asset_plan")
|
|
return object()
|
|
|
|
def sync_prompt_inputs(
|
|
self,
|
|
inputs: dict[str, Any],
|
|
*,
|
|
status_callback: Any,
|
|
request_cache: object,
|
|
) -> tuple[dict[str, Any], list[Any]]:
|
|
"""Return one prepared prompt after publishing environment-local status."""
|
|
assert request_cache is not None
|
|
self.calls.append("prompt_assets")
|
|
status_callback("Downloading prompt asset", 1, 1)
|
|
return inputs, []
|
|
|
|
components = [
|
|
SimpleNamespace(representative_node_id="a", node_ids=["a"]),
|
|
SimpleNamespace(representative_node_id="b", node_ids=["b"]),
|
|
]
|
|
assignments = {
|
|
"a": SimpleNamespace(
|
|
environment_id="vast:big:1",
|
|
provider=api_intercept_module.ExecutionProvider.VAST,
|
|
),
|
|
"b": SimpleNamespace(
|
|
environment_id="lambda",
|
|
provider=api_intercept_module.ExecutionProvider.SSH_DOCKER,
|
|
),
|
|
}
|
|
vast_engine = FakeSyncEngine()
|
|
ssh_engine = FakeSyncEngine()
|
|
results = api_intercept_module._prepare_remote_environment_assets(
|
|
components=components,
|
|
assignments_by_component_id=assignments,
|
|
sync_engines_by_environment={
|
|
"vast:big:1": vast_engine,
|
|
"lambda": ssh_engine,
|
|
},
|
|
rewritten_prompt={
|
|
"a": {"class_type": "VAELoader", "inputs": {"vae": "a.safetensors"}},
|
|
"b": {"class_type": "VAELoader", "inputs": {"vae": "b.safetensors"}},
|
|
},
|
|
sync_custom_nodes=True,
|
|
status_callback=None,
|
|
environment_status_callback=(
|
|
lambda environment_id, message, _current, _total: environment_events.append(
|
|
(environment_id, message)
|
|
)
|
|
),
|
|
)
|
|
|
|
assert maximum_active_environment_count == 2
|
|
assert list(results) == ["vast:big:1", "lambda"]
|
|
assert results["vast:big:1"].component_prompts["a"]["a"]["inputs"] == {
|
|
"vae": "a.safetensors"
|
|
}
|
|
assert results["lambda"].component_prompts["b"]["b"]["inputs"] == {
|
|
"vae": "b.safetensors"
|
|
}
|
|
assert vast_engine.calls == [
|
|
"r2_preflight",
|
|
"custom_nodes",
|
|
"asset_plan",
|
|
"prompt_assets",
|
|
]
|
|
assert ssh_engine.calls == vast_engine.calls
|
|
for environment_id in results:
|
|
messages = [
|
|
message
|
|
for event_environment_id, message in environment_events
|
|
if event_environment_id == environment_id
|
|
]
|
|
expected_completion = (
|
|
"Ready for remote execution"
|
|
if environment_id == "vast:big:1"
|
|
else "Remote assets prepared; SSH runtime starts on dispatch"
|
|
)
|
|
assert messages == [
|
|
"Preparing remote assets",
|
|
"Uploading custom nodes",
|
|
"Downloading prompt asset",
|
|
expected_completion,
|
|
]
|
|
|
|
def test_remote_environment_asset_worker_failures_bubble_up(
|
|
api_intercept_module: Any,
|
|
) -> None:
|
|
"""A failed environment worker must fail queue preparation with its cause."""
|
|
|
|
class FailingSyncEngine:
|
|
"""Fail before prompt assets are considered prepared."""
|
|
|
|
def preflight_r2_access(self, *, status_callback: Any) -> None:
|
|
"""Represent an environment without configured R2 backing."""
|
|
del status_callback
|
|
|
|
def sync_custom_nodes_directory(self, *, status_callback: Any) -> None:
|
|
"""Raise the representative environment-specific setup failure."""
|
|
status_callback("Uploading custom nodes", None, None)
|
|
raise OSError("remote storage is unavailable")
|
|
|
|
with pytest.raises(OSError, match="remote storage is unavailable"):
|
|
api_intercept_module._prepare_remote_environment_assets(
|
|
components=[
|
|
SimpleNamespace(representative_node_id="a", node_ids=["a"])
|
|
],
|
|
assignments_by_component_id={
|
|
"a": SimpleNamespace(
|
|
environment_id="vast:broken:1",
|
|
provider=api_intercept_module.ExecutionProvider.VAST,
|
|
)
|
|
},
|
|
sync_engines_by_environment={
|
|
"vast:broken:1": FailingSyncEngine()
|
|
},
|
|
rewritten_prompt={
|
|
"a": {"class_type": "VAELoader", "inputs": {}}
|
|
},
|
|
sync_custom_nodes=True,
|
|
status_callback=None,
|
|
environment_status_callback=lambda *_args: None,
|
|
)
|
|
|