581 lines
19 KiB
Python
581 lines
19 KiB
Python
"""Tests for the remote graph analysis boundary."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from api_intercept_test_support import * # noqa: F401,F403
|
|
|
|
def test_remote_partition_preserves_dag_around_ssh_only_llm(
|
|
api_intercept_module: Any,
|
|
remote_graph_analysis_module: Any,
|
|
) -> None:
|
|
"""Provider boundaries must not be undone by a coarse fanout cycle."""
|
|
prompt = {
|
|
"1": {"class_type": "RemoteImage", "inputs": {}},
|
|
"2": {
|
|
"class_type": "RemoteImageConsumer",
|
|
"inputs": {"image": ["1", 0]},
|
|
},
|
|
"3": {
|
|
"class_type": "ModalLLM",
|
|
"inputs": {
|
|
"image": ["2", 0],
|
|
"model_profile": "huihui-qwen3.8-27b-abliterated-q2-k-gguf",
|
|
},
|
|
},
|
|
"4": {
|
|
"class_type": "RemoteImageConsumer",
|
|
"inputs": {"image": ["1", 0], "prompt": ["3", 0]},
|
|
},
|
|
}
|
|
fake_nodes_module = SimpleNamespace(
|
|
NODE_CLASS_MAPPINGS={
|
|
"RemoteImage": _FakeRemoteImageNode,
|
|
"RemoteImageConsumer": _FakeRemoteImageConsumerNode,
|
|
"ModalLLM": _FakeTextNode,
|
|
},
|
|
NODE_DISPLAY_NAME_MAPPINGS={},
|
|
)
|
|
|
|
component_groups = remote_graph_analysis_module._remote_component_partition_groups(
|
|
prompt,
|
|
set(prompt),
|
|
remote_graph_analysis_module._build_consumer_map(prompt),
|
|
fake_nodes_module,
|
|
)
|
|
components = remote_graph_analysis_module._component_topological_order(
|
|
prompt,
|
|
component_groups,
|
|
)
|
|
|
|
assert components == [["1", "2"], ["3"], ["4"]]
|
|
assert remote_graph_analysis_module._component_execution_stages(
|
|
prompt,
|
|
component_groups,
|
|
) == [["1"], ["3"], ["4"]]
|
|
|
|
def test_remote_partition_replicates_non_transportable_fanout_around_ssh_llm(
|
|
api_intercept_module: Any,
|
|
component_planning_module: Any,
|
|
execution_scheduling_module: Any,
|
|
settings_module: Any,
|
|
sync_engine_module: Any,
|
|
monkeypatch: Any,
|
|
tmp_path: Path,
|
|
) -> None:
|
|
"""A safe shared runtime producer should be rebuilt after an SSH-only phase."""
|
|
prompt = {
|
|
"1": {"class_type": "VAELoader", "inputs": {}},
|
|
"2": {
|
|
"class_type": "VAEDecode",
|
|
"inputs": {"vae": ["1", 0]},
|
|
},
|
|
"3": {
|
|
"class_type": "ModalLLM",
|
|
"inputs": {
|
|
"image": ["2", 0],
|
|
"model_profile": "huihui-qwen3.8-27b-abliterated-q2-k-gguf",
|
|
},
|
|
},
|
|
"4": {
|
|
"class_type": "VAEDecode",
|
|
"inputs": {"vae": ["1", 0], "prompt": ["3", 0]},
|
|
},
|
|
}
|
|
pristine_prompt = json.loads(json.dumps(prompt))
|
|
fake_nodes_module = SimpleNamespace(
|
|
NODE_CLASS_MAPPINGS={
|
|
"VAELoader": _FakeVAELoaderNode,
|
|
"VAEDecode": _FakeVAEDecodeNode,
|
|
"ModalLLM": _FakeTextNode,
|
|
},
|
|
NODE_DISPLAY_NAME_MAPPINGS={},
|
|
)
|
|
remote_node_ids = set(prompt)
|
|
component_plans = component_planning_module._build_component_plans(
|
|
prompt,
|
|
remote_node_ids,
|
|
fake_nodes_module,
|
|
)
|
|
component_planning_module.validate_remote_component_transport_compatibility(
|
|
prompt,
|
|
component_plans,
|
|
fake_nodes_module,
|
|
)
|
|
assert [component.representative_node_id for component in component_plans] == [
|
|
"1",
|
|
"3",
|
|
"4",
|
|
]
|
|
replica_node_ids = {
|
|
node_id
|
|
for node_id in prompt
|
|
if node_id.startswith(component_planning_module._REMOTE_REPLICA_NODE_PREFIX)
|
|
}
|
|
assert len(replica_node_ids) == 1
|
|
replica_node_id = next(iter(replica_node_ids))
|
|
assert prompt["4"]["inputs"]["vae"] == [replica_node_id, 0]
|
|
assert prompt[replica_node_id] == prompt["1"]
|
|
assert replica_node_id in component_plans[2].node_ids
|
|
assert component_plans[0].boundary_inputs == []
|
|
assert component_plans[0].boundary_outputs[0].io_type == "IMAGE"
|
|
assert component_plans[2].boundary_inputs[0].io_type == "STRING"
|
|
assert component_plans[2].boundary_outputs == []
|
|
required_provider = execution_scheduling_module._component_required_provider(
|
|
component_plans[1],
|
|
prompt,
|
|
{
|
|
"huihui-qwen3.8-27b-abliterated-q2-k-gguf": SimpleNamespace(
|
|
backend="llama_cpp_server"
|
|
)
|
|
},
|
|
)
|
|
assert required_provider.value == "ssh_docker"
|
|
|
|
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="local",
|
|
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=tmp_path / "custom_nodes",
|
|
)
|
|
settings.custom_nodes_dir.mkdir()
|
|
sync_engine = sync_engine_module.ModalAssetSyncEngine.from_environment(settings)
|
|
|
|
def assign_to_modal(*, components: list[Any], **_kwargs: Any) -> dict[str, Any]:
|
|
"""Keep the rewrite test independent of real provider availability."""
|
|
return {
|
|
component.representative_node_id: api_intercept_module.ExecutionAssignment(
|
|
environment_id="modal:H200",
|
|
provider=api_intercept_module.ExecutionProvider.MODAL,
|
|
predicted_cost_usd=0.0,
|
|
predicted_completion_seconds=1.0,
|
|
)
|
|
for component in components
|
|
}
|
|
|
|
monkeypatch.setattr(
|
|
execution_scheduling_module,
|
|
"_plan_component_execution_assignments",
|
|
assign_to_modal,
|
|
)
|
|
rewritten_prompt, summary = api_intercept_module.rewrite_prompt_for_modal(
|
|
prompt=pristine_prompt,
|
|
workflow={
|
|
"nodes": [
|
|
{"id": node_id, "properties": {"is_modal_remote": True}}
|
|
for node_id in range(1, 5)
|
|
]
|
|
},
|
|
sync_engine=sync_engine,
|
|
settings=settings,
|
|
nodes_module=fake_nodes_module,
|
|
)
|
|
|
|
assert not any(
|
|
node_id.startswith(component_planning_module._REMOTE_REPLICA_NODE_PREFIX)
|
|
for node_id in rewritten_prompt
|
|
)
|
|
assert summary.component_node_ids_by_representative["4"] == ["4"]
|
|
assert not any(
|
|
node_id.startswith(component_planning_module._REMOTE_REPLICA_NODE_PREFIX)
|
|
for node_id in summary.rewritten_node_id_map
|
|
)
|
|
downstream_payload = rewritten_prompt["4"]["inputs"]["original_node_data"]
|
|
replica_payload_node_ids = {
|
|
node_id
|
|
for node_id in downstream_payload["subgraph_prompt"]
|
|
if node_id.startswith(component_planning_module._REMOTE_REPLICA_NODE_PREFIX)
|
|
}
|
|
assert len(replica_payload_node_ids) == 1
|
|
downstream_vae_input = downstream_payload["subgraph_prompt"]["4"]["inputs"][
|
|
"vae"
|
|
]
|
|
assert downstream_vae_input[0] in replica_payload_node_ids
|
|
|
|
def test_remote_partition_replicates_linked_model_loader_closure(
|
|
api_intercept_module: Any,
|
|
component_planning_module: Any,
|
|
) -> None:
|
|
"""A downstream provider phase should rebuild a linked loader chain, not sample."""
|
|
prompt = {
|
|
"1": {
|
|
"class_type": "UNETLoader",
|
|
"inputs": {"unet_name": "model.safetensors"},
|
|
},
|
|
"2": {
|
|
"class_type": "LoraLoaderModelOnly",
|
|
"inputs": {"model": ["1", 0], "lora_name": "adapter.safetensors"},
|
|
},
|
|
"3": {
|
|
"class_type": "RemoteSampler",
|
|
"inputs": {"model": ["2", 0]},
|
|
},
|
|
"4": {
|
|
"class_type": "ModalLLM",
|
|
"inputs": {
|
|
"image": ["3", 0],
|
|
"model_profile": "huihui-qwen3.8-27b-abliterated-q2-k-gguf",
|
|
},
|
|
},
|
|
"5": {
|
|
"class_type": "RemoteSampler",
|
|
"inputs": {"model": ["2", 0], "prompt": ["4", 0]},
|
|
},
|
|
}
|
|
fake_nodes_module = SimpleNamespace(
|
|
NODE_CLASS_MAPPINGS={
|
|
"UNETLoader": _FakeRemoteModelNode,
|
|
"LoraLoaderModelOnly": _FakeRemoteModelNode,
|
|
"RemoteSampler": _FakeRemoteSamplerNode,
|
|
"ModalLLM": _FakeTextNode,
|
|
},
|
|
NODE_DISPLAY_NAME_MAPPINGS={},
|
|
)
|
|
|
|
component_plans = component_planning_module._build_component_plans(
|
|
prompt,
|
|
set(prompt),
|
|
fake_nodes_module,
|
|
)
|
|
|
|
assert [component.representative_node_id for component in component_plans] == [
|
|
"1",
|
|
"4",
|
|
"5",
|
|
]
|
|
replica_node_ids = sorted(
|
|
node_id
|
|
for node_id in prompt
|
|
if node_id.startswith(component_planning_module._REMOTE_REPLICA_NODE_PREFIX)
|
|
)
|
|
assert len(replica_node_ids) == 2
|
|
replica_loader_id = next(
|
|
node_id
|
|
for node_id in replica_node_ids
|
|
if prompt[node_id]["class_type"] == "UNETLoader"
|
|
)
|
|
replica_lora_id = next(
|
|
node_id
|
|
for node_id in replica_node_ids
|
|
if prompt[node_id]["class_type"] == "LoraLoaderModelOnly"
|
|
)
|
|
assert prompt[replica_lora_id]["inputs"]["model"] == [replica_loader_id, 0]
|
|
assert prompt["5"]["inputs"]["model"] == [replica_lora_id, 0]
|
|
assert replica_node_ids == sorted(
|
|
set(component_plans[2].node_ids) - {"5"}
|
|
)
|
|
component_planning_module.validate_remote_component_transport_compatibility(
|
|
prompt,
|
|
component_plans,
|
|
fake_nodes_module,
|
|
)
|
|
|
|
def test_cross_provider_boundary_uses_transport_instead_of_remote_session(
|
|
api_intercept_module: Any,
|
|
component_planning_module: Any,
|
|
execution_environments_module: Any,
|
|
) -> None:
|
|
"""Session-backed references must never cross provider storage boundaries."""
|
|
source = api_intercept_module.LinkedOutputRef("1", 0)
|
|
boundary_output = api_intercept_module.BoundaryOutputSpec(
|
|
proxy_output_name="remote_output_0",
|
|
source=source,
|
|
io_type="IMAGE",
|
|
is_list=False,
|
|
)
|
|
producer = api_intercept_module.RemoteComponentPlan(
|
|
node_ids=["1"],
|
|
representative_node_id="1",
|
|
boundary_inputs=[],
|
|
boundary_outputs=[boundary_output],
|
|
execute_node_ids=["1"],
|
|
contains_output_node=False,
|
|
)
|
|
consumer = api_intercept_module.RemoteComponentPlan(
|
|
node_ids=["2"],
|
|
representative_node_id="2",
|
|
boundary_inputs=[
|
|
api_intercept_module.BoundaryInputSpec(
|
|
proxy_input_name="remote_input_0",
|
|
source=source,
|
|
io_type="IMAGE",
|
|
targets=[api_intercept_module.InputTarget("2", "image")],
|
|
)
|
|
],
|
|
boundary_outputs=[],
|
|
execute_node_ids=["2"],
|
|
contains_output_node=False,
|
|
)
|
|
assignment_type = execution_environments_module.ExecutionAssignment
|
|
provider_type = execution_environments_module.ExecutionProvider
|
|
|
|
session_component_ids = (
|
|
component_planning_module._mark_remote_to_remote_session_boundaries(
|
|
{
|
|
"1": {"class_type": "RemoteImage", "inputs": {}},
|
|
"2": {
|
|
"class_type": "RemoteImageConsumer",
|
|
"inputs": {"image": ["1", 0]},
|
|
},
|
|
},
|
|
[producer, consumer],
|
|
SimpleNamespace(
|
|
NODE_CLASS_MAPPINGS={
|
|
"RemoteImage": _FakeRemoteImageNode,
|
|
"RemoteImageConsumer": _FakeRemoteImageConsumerNode,
|
|
}
|
|
),
|
|
{
|
|
"1": assignment_type(
|
|
"modal:H200",
|
|
provider_type.MODAL,
|
|
None,
|
|
0.0,
|
|
),
|
|
"2": assignment_type(
|
|
"lambda",
|
|
provider_type.SSH_DOCKER,
|
|
0.0,
|
|
0.0,
|
|
),
|
|
},
|
|
)
|
|
)
|
|
|
|
assert session_component_ids == set()
|
|
assert boundary_output.session_output is False
|
|
assert boundary_output.session_consumer_node_ids == []
|
|
|
|
def test_non_modal_boundary_with_local_preview_uses_transport(
|
|
api_intercept_module: Any,
|
|
component_planning_module: Any,
|
|
execution_environments_module: Any,
|
|
) -> None:
|
|
"""A Vast bridge with a local consumer must not require Modal shared storage."""
|
|
source = api_intercept_module.LinkedOutputRef("1", 0)
|
|
boundary_output = api_intercept_module.BoundaryOutputSpec(
|
|
proxy_output_name="remote_output_0",
|
|
source=source,
|
|
io_type="IMAGE",
|
|
is_list=False,
|
|
)
|
|
producer = api_intercept_module.RemoteComponentPlan(
|
|
node_ids=["1"],
|
|
representative_node_id="1",
|
|
boundary_inputs=[],
|
|
boundary_outputs=[boundary_output],
|
|
execute_node_ids=["1"],
|
|
contains_output_node=False,
|
|
)
|
|
consumer = api_intercept_module.RemoteComponentPlan(
|
|
node_ids=["2"],
|
|
representative_node_id="2",
|
|
boundary_inputs=[
|
|
api_intercept_module.BoundaryInputSpec(
|
|
proxy_input_name="remote_input_0",
|
|
source=source,
|
|
io_type="IMAGE",
|
|
targets=[api_intercept_module.InputTarget("2", "image")],
|
|
)
|
|
],
|
|
boundary_outputs=[],
|
|
execute_node_ids=["2"],
|
|
contains_output_node=False,
|
|
)
|
|
assignment_type = execution_environments_module.ExecutionAssignment
|
|
provider_type = execution_environments_module.ExecutionProvider
|
|
vast_assignment = assignment_type(
|
|
"vast:profile:1234",
|
|
provider_type.VAST,
|
|
0.0,
|
|
0.0,
|
|
)
|
|
|
|
session_component_ids = (
|
|
component_planning_module._mark_remote_to_remote_session_boundaries(
|
|
{
|
|
"1": {"class_type": "RemoteImage", "inputs": {}},
|
|
"2": {
|
|
"class_type": "RemoteImageConsumer",
|
|
"inputs": {"image": ["1", 0]},
|
|
},
|
|
"3": {
|
|
"class_type": "PreviewImage",
|
|
"inputs": {"images": ["1", 0]},
|
|
},
|
|
},
|
|
[producer, consumer],
|
|
SimpleNamespace(
|
|
NODE_CLASS_MAPPINGS={
|
|
"RemoteImage": _FakeRemoteImageNode,
|
|
"RemoteImageConsumer": _FakeRemoteImageConsumerNode,
|
|
"PreviewImage": _FakePreviewImageNode,
|
|
}
|
|
),
|
|
{"1": vast_assignment, "2": vast_assignment},
|
|
)
|
|
)
|
|
|
|
assert session_component_ids == set()
|
|
assert boundary_output.session_output is False
|
|
assert boundary_output.session_consumer_node_ids == []
|
|
assert boundary_output.local_materializer_node_id is None
|
|
|
|
def test_transportable_list_boundary_preserves_scheduler_items(
|
|
api_intercept_module: Any,
|
|
component_planning_module: Any,
|
|
execution_environments_module: Any,
|
|
) -> None:
|
|
"""Keep a same-host list output in ComfyUI instead of one bridge token."""
|
|
source = api_intercept_module.LinkedOutputRef("1", 0)
|
|
boundary_output = api_intercept_module.BoundaryOutputSpec(
|
|
proxy_output_name="seed_list",
|
|
source=source,
|
|
io_type="INT",
|
|
is_list=True,
|
|
)
|
|
producer = api_intercept_module.RemoteComponentPlan(
|
|
node_ids=["1"],
|
|
representative_node_id="1",
|
|
boundary_inputs=[],
|
|
boundary_outputs=[boundary_output],
|
|
execute_node_ids=["1"],
|
|
contains_output_node=False,
|
|
)
|
|
consumer = api_intercept_module.RemoteComponentPlan(
|
|
node_ids=["2"],
|
|
representative_node_id="2",
|
|
boundary_inputs=[
|
|
api_intercept_module.BoundaryInputSpec(
|
|
proxy_input_name="seed",
|
|
source=source,
|
|
io_type="INT",
|
|
targets=[api_intercept_module.InputTarget("2", "seed")],
|
|
)
|
|
],
|
|
boundary_outputs=[],
|
|
execute_node_ids=["2"],
|
|
contains_output_node=False,
|
|
)
|
|
assignment_type = execution_environments_module.ExecutionAssignment
|
|
provider_type = execution_environments_module.ExecutionProvider
|
|
lambda_assignment = assignment_type(
|
|
"lambda",
|
|
provider_type.SSH_DOCKER,
|
|
0.0,
|
|
0.0,
|
|
)
|
|
|
|
session_component_ids = (
|
|
component_planning_module._mark_remote_to_remote_session_boundaries(
|
|
{
|
|
"1": {"class_type": "NextSeeds", "inputs": {}},
|
|
"2": {
|
|
"class_type": "NextSeeds",
|
|
"inputs": {"seed": ["1", 0]},
|
|
},
|
|
},
|
|
[producer, consumer],
|
|
SimpleNamespace(NODE_CLASS_MAPPINGS={}),
|
|
{"1": lambda_assignment, "2": lambda_assignment},
|
|
)
|
|
)
|
|
|
|
assert session_component_ids == set()
|
|
assert boundary_output.session_output is False
|
|
assert boundary_output.session_consumer_node_ids == []
|
|
|
|
def test_split_phase_order_accounts_for_local_feedback_dependencies(
|
|
api_intercept_module: Any,
|
|
component_planning_module: Any,
|
|
) -> None:
|
|
"""Split phase ordering should treat local re-entry paths as real dependencies."""
|
|
prompt = {
|
|
"3": {
|
|
"class_type": "RemoteSampler",
|
|
"inputs": {"model": ["14", 0]},
|
|
},
|
|
"11": {
|
|
"class_type": "VAEDecode",
|
|
"inputs": {"samples": ["3", 0]},
|
|
},
|
|
"14": {
|
|
"class_type": "RemoteModel",
|
|
"inputs": {},
|
|
},
|
|
"191": {
|
|
"class_type": "RemoteSampler",
|
|
"inputs": {"model": ["14", 0], "conditioning": ["358", 0]},
|
|
},
|
|
"357": {
|
|
"class_type": "BetterGrok",
|
|
"inputs": {"prompt_images": ["11", 0]},
|
|
},
|
|
"358": {
|
|
"class_type": "RemoteTextEncode",
|
|
"inputs": {"text": ["357", 1]},
|
|
},
|
|
}
|
|
component_prompt = {
|
|
"3": prompt["3"],
|
|
"14": prompt["14"],
|
|
"191": prompt["191"],
|
|
"358": prompt["358"],
|
|
}
|
|
|
|
ordered_execute_node_ids = (
|
|
component_planning_module._order_execute_node_ids_for_transportable_splits(
|
|
prompt=prompt,
|
|
component_prompt=component_prompt,
|
|
component_node_ids={"3", "14", "191", "358"},
|
|
execute_node_ids=["191", "3"],
|
|
)
|
|
)
|
|
|
|
assert ordered_execute_node_ids == ["3", "191"]
|
|
|
|
def test_boundary_source_signature_changes_with_upstream_prompt_structure(
|
|
api_intercept_module: Any,
|
|
prompt_payload_metadata_module: Any,
|
|
) -> None:
|
|
"""Non-transportable boundary provenance should change when the upstream prompt changes."""
|
|
source = api_intercept_module.LinkedOutputRef(node_id="2", output_index=0)
|
|
base_prompt = {
|
|
"1": {
|
|
"class_type": "CheckpointLoader",
|
|
"inputs": {"ckpt_name": "base.safetensors"},
|
|
},
|
|
"2": {
|
|
"class_type": "LoraLoader",
|
|
"inputs": {"model": ["1", 0], "strength_model": 0.8},
|
|
},
|
|
}
|
|
changed_prompt = {
|
|
"1": {
|
|
"class_type": "CheckpointLoader",
|
|
"inputs": {"ckpt_name": "base.safetensors"},
|
|
},
|
|
"2": {
|
|
"class_type": "LoraLoader",
|
|
"inputs": {"model": ["1", 0], "strength_model": 0.5},
|
|
},
|
|
}
|
|
|
|
first_signature = prompt_payload_metadata_module._boundary_source_signature(base_prompt, source)
|
|
second_signature = prompt_payload_metadata_module._boundary_source_signature(base_prompt, source)
|
|
changed_signature = prompt_payload_metadata_module._boundary_source_signature(changed_prompt, source)
|
|
|
|
assert first_signature == second_signature
|
|
assert changed_signature != first_signature
|
|
|