Extract mapped remote execution
This commit is contained in:
Ken Simpson
@@ -516,7 +516,7 @@ COMFYUI_ROOT=/tmp/comfyui-modal-test/ComfyUI \
|
||||
|
||||
For full local development including the remote and Apple-local extras, use `uv sync --extra remote --extra local-apple --group test`.
|
||||
|
||||
The prompt interception layer is organized as a compatibility aggregator over focused modules: planning values, graph analysis, component planning, prompt rewriting, execution scheduling, queue lifecycle, diagnostics, and HTTP route domains each have a dedicated owner. Route handlers share an immutable `RouteContext`, while `api_intercept.py` retains only registration and compatibility re-exports. Host-side remote execution follows the same ownership model: deployment loading, runtime compatibility checks, and auto-deploy state live in `remote/modal_deployment.py`; proactive warmup, snapshot profiles, and local-gap keepalive state live in `remote/modal_warmup.py`; and adjacent LLM-staging and mapped-execution callbacks are injected by the `remote/modal_app.py` orchestrator. Blocking Modal executors are constructed lazily on first use instead of during import. Root modules support both package imports used by ComfyUI and flat imports used by the Modal container; `tests/test_import_shapes.py` guards both layouts and the stable `comfyui_modal_sync_cloud` surface.
|
||||
The prompt interception layer is organized as a compatibility aggregator over focused modules: planning values, graph analysis, component planning, prompt rewriting, execution scheduling, queue lifecycle, diagnostics, and HTTP route domains each have a dedicated owner. Route handlers share an immutable `RouteContext`, while `api_intercept.py` retains only registration and compatibility re-exports. Host-side remote execution follows the same ownership model: deployment loading, runtime compatibility checks, and auto-deploy state live in `remote/modal_deployment.py`; proactive warmup, snapshot profiles, and local-gap keepalive state live in `remote/modal_warmup.py`; and mapped or implicit-batch lane scheduling lives in `remote/mapped_execution.py`. Adjacent invocation and LLM-staging callbacks are injected by the `remote/modal_app.py` orchestrator. Blocking Modal executors are constructed lazily on first use instead of during import. Root modules support both package imports used by ComfyUI and flat imports used by the Modal container; `tests/test_import_shapes.py` guards both layouts and the stable `comfyui_modal_sync_cloud` surface.
|
||||
|
||||
### Vast.ai API Simulator
|
||||
|
||||
|
||||
@@ -659,3 +659,4 @@
|
||||
- A moved handler resolves imported functions in its new module, so tests must patch the new owner rather than a compatibility re-export. Inject the queue rewrite and status callbacks through the context where importing the interceptor would create a cycle; this preserves registration-time monkeypatching without coupling route modules back to the aggregator.
|
||||
- Deployment extraction has the same dependency-direction constraint as route extraction. Keep Modal SDK guards, app-loading locks, version caches, and auto-deploy coordination in the deployment module, then inject warmup, snapshot, and LLM-staging hooks from the host orchestrator. A lazy upward import avoids an import-time cycle but still obscures ownership and makes test patches resolve in surprising modules.
|
||||
- Importing a host orchestration module should not allocate thread pools before any workflow needs them. Move warmup and keepalive executors with their state owner, guard each lazy constructor with a lock, and keep the blocking invocation executor lazy in the residual orchestrator; regression tests can then patch the actual executor owner without retaining compatibility aliases to mutable state.
|
||||
- Mapped execution mixes pure payload partitioning with calls back into the host invocation core. Move its lane registry and scheduling algorithms together, but inject only the true upward edges (ordinary invocation, affinity-bound invocation, and interrupt handling). Direct dependencies on serialization, local subgraph execution, deployment lookup, and UI progress remain imports whose tests patch the new mapped-execution owner.
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
+51
-1220
File diff suppressed because it is too large
Load Diff
@@ -485,6 +485,12 @@ def modal_warmup_module(extension_package: object) -> object:
|
||||
return importlib.import_module(f"{PACKAGE_NAME}.remote.modal_warmup")
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def mapped_execution_module(extension_package: object) -> object:
|
||||
"""Return mapped and implicit-batch execution helpers."""
|
||||
return importlib.import_module(f"{PACKAGE_NAME}.remote.mapped_execution")
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def serialization_module(extension_package: object) -> object:
|
||||
"""Return the serialization helpers module."""
|
||||
|
||||
@@ -8441,13 +8441,14 @@ def test_emit_local_remote_dispatch_status_marks_component_starting(
|
||||
|
||||
def test_emit_local_mapped_lane_progress_start_marks_lane_as_setup_only(
|
||||
remote_modal_app_module: Any,
|
||||
mapped_execution_module: Any,
|
||||
monkeypatch: Any,
|
||||
) -> None:
|
||||
"""Provisioning a mapped lane should emit a dedicated setup-only lane progress event."""
|
||||
emitted_progress: list[dict[str, Any]] = []
|
||||
|
||||
monkeypatch.setattr(
|
||||
remote_modal_app_module,
|
||||
mapped_execution_module,
|
||||
"_emit_local_modal_progress",
|
||||
lambda **kwargs: emitted_progress.append(kwargs),
|
||||
)
|
||||
@@ -10262,6 +10263,7 @@ def test_invoke_remote_engine_payload_releases_local_prompt_after_cancel_grace(
|
||||
|
||||
def test_invoke_mapped_remote_engine_async_runs_explicit_mapped_phase_items(
|
||||
remote_modal_app_module: Any,
|
||||
mapped_execution_module: Any,
|
||||
serialization_module: Any,
|
||||
monkeypatch: Any,
|
||||
) -> None:
|
||||
@@ -10280,12 +10282,12 @@ def test_invoke_mapped_remote_engine_async_runs_explicit_mapped_phase_items(
|
||||
return (f"done:{hydrated_inputs['remote_input_0']}",)
|
||||
|
||||
monkeypatch.setattr(
|
||||
remote_modal_app_module,
|
||||
mapped_execution_module,
|
||||
"_execute_subgraph_prompt",
|
||||
fake_execute_subgraph_prompt,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
remote_modal_app_module,
|
||||
mapped_execution_module,
|
||||
"_emit_local_modal_progress",
|
||||
lambda **kwargs: progress_updates.append(kwargs),
|
||||
)
|
||||
@@ -10594,6 +10596,7 @@ def test_implicit_mapping_propagates_scheduler_list_through_downstream_component
|
||||
|
||||
def test_invoke_mapped_remote_engine_async_splits_int_inputs_for_direct_targets(
|
||||
remote_modal_app_module: Any,
|
||||
mapped_execution_module: Any,
|
||||
serialization_module: Any,
|
||||
monkeypatch: Any,
|
||||
) -> None:
|
||||
@@ -10609,7 +10612,7 @@ def test_invoke_mapped_remote_engine_async_splits_int_inputs_for_direct_targets(
|
||||
return (f"seed:{hydrated_inputs['remote_input_0']}",)
|
||||
|
||||
monkeypatch.setattr(
|
||||
remote_modal_app_module,
|
||||
mapped_execution_module,
|
||||
"_execute_subgraph_prompt",
|
||||
fake_execute_subgraph_prompt,
|
||||
)
|
||||
@@ -10692,6 +10695,7 @@ def test_invoke_mapped_remote_engine_async_splits_int_inputs_for_direct_targets(
|
||||
|
||||
def test_invoke_mapped_remote_engine_async_executes_static_branch_once(
|
||||
remote_modal_app_module: Any,
|
||||
mapped_execution_module: Any,
|
||||
serialization_module: Any,
|
||||
monkeypatch: Any,
|
||||
) -> None:
|
||||
@@ -10723,7 +10727,7 @@ def test_invoke_mapped_remote_engine_async_executes_static_branch_once(
|
||||
return (f"mapped:{hydrated_inputs['remote_input_1']}",)
|
||||
|
||||
monkeypatch.setattr(
|
||||
remote_modal_app_module,
|
||||
mapped_execution_module,
|
||||
"_execute_subgraph_prompt",
|
||||
fake_execute_subgraph_prompt,
|
||||
)
|
||||
@@ -12678,6 +12682,7 @@ def test_implicitly_mapped_subgraph_shared_model_keeps_unbatched_sampler_single_
|
||||
|
||||
def test_implicitly_mapped_subgraph_seeds_remote_lanes_before_item_dispatch(
|
||||
remote_modal_app_module: Any,
|
||||
mapped_execution_module: Any,
|
||||
serialization_module: Any,
|
||||
session_state_module: Any,
|
||||
monkeypatch: Any,
|
||||
@@ -12688,10 +12693,11 @@ def test_implicitly_mapped_subgraph_seeds_remote_lanes_before_item_dispatch(
|
||||
|
||||
monkeypatch.setenv("COMFY_MODAL_EXECUTION_MODE", "remote")
|
||||
monkeypatch.setattr(remote_modal_app_module, "modal", object())
|
||||
monkeypatch.setattr(mapped_execution_module, "modal", object())
|
||||
monkeypatch.setattr(remote_modal_app_module, "_mapped_execution_parallelism", lambda total_items: 2)
|
||||
monkeypatch.setattr(remote_modal_app_module, "ensure_remote_warm_capacity", lambda *args, **kwargs: 0)
|
||||
monkeypatch.setattr(
|
||||
remote_modal_app_module,
|
||||
mapped_execution_module,
|
||||
"_lookup_deployed_remote_engine",
|
||||
lambda payload, affinity_key_override=None: (
|
||||
f"engine:{affinity_key_override or remote_modal_app_module._remote_worker_affinity_key(payload)}"
|
||||
@@ -12729,7 +12735,7 @@ def test_implicitly_mapped_subgraph_seeds_remote_lanes_before_item_dispatch(
|
||||
fake_invoke_bound_remote_engine_async,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
remote_modal_app_module,
|
||||
mapped_execution_module,
|
||||
"_emit_local_mapped_lane_progress_start",
|
||||
lambda payload, lane_index, item_index=None: observed_lane_setup_starts.append(
|
||||
(str(payload["component_id"]), int(lane_index))
|
||||
@@ -12847,6 +12853,7 @@ def test_implicitly_mapped_subgraph_seeds_remote_lanes_before_item_dispatch(
|
||||
|
||||
def test_implicitly_mapped_subgraph_allows_ready_lanes_to_start_before_slowest_seed_finishes(
|
||||
remote_modal_app_module: Any,
|
||||
mapped_execution_module: Any,
|
||||
serialization_module: Any,
|
||||
session_state_module: Any,
|
||||
monkeypatch: Any,
|
||||
@@ -12856,10 +12863,11 @@ def test_implicitly_mapped_subgraph_allows_ready_lanes_to_start_before_slowest_s
|
||||
|
||||
monkeypatch.setenv("COMFY_MODAL_EXECUTION_MODE", "remote")
|
||||
monkeypatch.setattr(remote_modal_app_module, "modal", object())
|
||||
monkeypatch.setattr(mapped_execution_module, "modal", object())
|
||||
monkeypatch.setattr(remote_modal_app_module, "_mapped_execution_parallelism", lambda total_items: 2)
|
||||
monkeypatch.setattr(remote_modal_app_module, "ensure_remote_warm_capacity", lambda *args, **kwargs: 0)
|
||||
monkeypatch.setattr(
|
||||
remote_modal_app_module,
|
||||
mapped_execution_module,
|
||||
"_lookup_deployed_remote_engine",
|
||||
lambda payload, affinity_key_override=None: (
|
||||
f"engine:{affinity_key_override or remote_modal_app_module._remote_worker_affinity_key(payload)}"
|
||||
@@ -12989,6 +12997,7 @@ def test_implicitly_mapped_subgraph_allows_ready_lanes_to_start_before_slowest_s
|
||||
|
||||
def test_implicitly_mapped_subgraph_returns_once_all_items_finish_without_waiting_for_unused_lane(
|
||||
remote_modal_app_module: Any,
|
||||
mapped_execution_module: Any,
|
||||
serialization_module: Any,
|
||||
session_state_module: Any,
|
||||
monkeypatch: Any,
|
||||
@@ -12998,10 +13007,11 @@ def test_implicitly_mapped_subgraph_returns_once_all_items_finish_without_waitin
|
||||
|
||||
monkeypatch.setenv("COMFY_MODAL_EXECUTION_MODE", "remote")
|
||||
monkeypatch.setattr(remote_modal_app_module, "modal", object())
|
||||
monkeypatch.setattr(mapped_execution_module, "modal", object())
|
||||
monkeypatch.setattr(remote_modal_app_module, "_mapped_execution_parallelism", lambda total_items: 2)
|
||||
monkeypatch.setattr(remote_modal_app_module, "ensure_remote_warm_capacity", lambda *args, **kwargs: 0)
|
||||
monkeypatch.setattr(
|
||||
remote_modal_app_module,
|
||||
mapped_execution_module,
|
||||
"_lookup_deployed_remote_engine",
|
||||
lambda payload, affinity_key_override=None: (
|
||||
f"engine:{affinity_key_override or remote_modal_app_module._remote_worker_affinity_key(payload)}"
|
||||
@@ -13133,6 +13143,7 @@ def test_implicitly_mapped_subgraph_returns_once_all_items_finish_without_waitin
|
||||
|
||||
def test_implicitly_mapped_subgraph_ignores_late_lane_failure_after_all_items_complete(
|
||||
remote_modal_app_module: Any,
|
||||
mapped_execution_module: Any,
|
||||
serialization_module: Any,
|
||||
session_state_module: Any,
|
||||
monkeypatch: Any,
|
||||
@@ -13142,10 +13153,11 @@ def test_implicitly_mapped_subgraph_ignores_late_lane_failure_after_all_items_co
|
||||
|
||||
monkeypatch.setenv("COMFY_MODAL_EXECUTION_MODE", "remote")
|
||||
monkeypatch.setattr(remote_modal_app_module, "modal", object())
|
||||
monkeypatch.setattr(mapped_execution_module, "modal", object())
|
||||
monkeypatch.setattr(remote_modal_app_module, "_mapped_execution_parallelism", lambda total_items: 2)
|
||||
monkeypatch.setattr(remote_modal_app_module, "ensure_remote_warm_capacity", lambda *args, **kwargs: 0)
|
||||
monkeypatch.setattr(
|
||||
remote_modal_app_module,
|
||||
mapped_execution_module,
|
||||
"_lookup_deployed_remote_engine",
|
||||
lambda payload, affinity_key_override=None: (
|
||||
f"engine:{affinity_key_override or remote_modal_app_module._remote_worker_affinity_key(payload)}"
|
||||
@@ -13271,6 +13283,7 @@ def test_implicitly_mapped_subgraph_ignores_late_lane_failure_after_all_items_co
|
||||
|
||||
def test_implicitly_mapped_subgraph_skips_outer_fanout_for_input_is_list_targets(
|
||||
remote_modal_app_module: Any,
|
||||
mapped_execution_module: Any,
|
||||
serialization_module: Any,
|
||||
monkeypatch: Any,
|
||||
) -> None:
|
||||
@@ -13288,7 +13301,7 @@ def test_implicitly_mapped_subgraph_skips_outer_fanout_for_input_is_list_targets
|
||||
}
|
||||
},
|
||||
)()
|
||||
monkeypatch.setattr(remote_modal_app_module, "_load_nodes_module", lambda: fake_nodes_module)
|
||||
monkeypatch.setattr(mapped_execution_module, "_load_nodes_module", lambda: fake_nodes_module)
|
||||
|
||||
async def fake_invoke_remote_engine_async(
|
||||
payload: dict[str, Any],
|
||||
@@ -13359,6 +13372,7 @@ def test_implicitly_mapped_subgraph_skips_outer_fanout_for_input_is_list_targets
|
||||
|
||||
def test_invoke_remote_engine_async_bypasses_implicit_mapping_once_for_input_is_list_targets(
|
||||
remote_modal_app_module: Any,
|
||||
mapped_execution_module: Any,
|
||||
serialization_module: Any,
|
||||
monkeypatch: Any,
|
||||
) -> None:
|
||||
@@ -13376,7 +13390,7 @@ def test_invoke_remote_engine_async_bypasses_implicit_mapping_once_for_input_is_
|
||||
}
|
||||
},
|
||||
)()
|
||||
monkeypatch.setattr(remote_modal_app_module, "_load_nodes_module", lambda: fake_nodes_module)
|
||||
monkeypatch.setattr(mapped_execution_module, "_load_nodes_module", lambda: fake_nodes_module)
|
||||
|
||||
def fake_invoke_remote_engine(
|
||||
payload: dict[str, Any],
|
||||
|
||||
@@ -57,6 +57,9 @@ def _load_modal_app_module() -> types.ModuleType:
|
||||
def test_mapped_stream_progress_preserves_real_node_id(monkeypatch: Any) -> None:
|
||||
"""Mapped lane progress should target the real executing node, not the representative."""
|
||||
modal_app = _load_modal_app_module()
|
||||
mapped_execution = sys.modules[
|
||||
f"{TEST_PACKAGE_NAME}.remote.mapped_execution"
|
||||
]
|
||||
emitted_progress: list[dict[str, Any]] = []
|
||||
|
||||
def capture_progress(**kwargs: Any) -> None:
|
||||
@@ -64,6 +67,9 @@ def test_mapped_stream_progress_preserves_real_node_id(monkeypatch: Any) -> None
|
||||
emitted_progress.append(kwargs)
|
||||
|
||||
monkeypatch.setattr(modal_app, "_emit_local_modal_progress", capture_progress)
|
||||
monkeypatch.setattr(
|
||||
mapped_execution, "_emit_local_modal_progress", capture_progress
|
||||
)
|
||||
|
||||
payload = {
|
||||
"prompt_id": "prompt-1",
|
||||
@@ -91,8 +97,8 @@ def test_mapped_stream_progress_preserves_real_node_id(monkeypatch: Any) -> None
|
||||
]
|
||||
)
|
||||
|
||||
with modal_app._MAPPED_PROGRESS_NODE_IDS_LOCK:
|
||||
modal_app._MAPPED_PROGRESS_NODE_IDS.clear()
|
||||
with mapped_execution._MAPPED_PROGRESS_NODE_IDS_LOCK:
|
||||
mapped_execution._MAPPED_PROGRESS_NODE_IDS.clear()
|
||||
|
||||
modal_app._consume_remote_payload_stream(payload, stream_events)
|
||||
|
||||
@@ -111,8 +117,8 @@ def test_mapped_stream_progress_preserves_real_node_id(monkeypatch: Any) -> None
|
||||
"aggregate_only": False,
|
||||
}
|
||||
]
|
||||
with modal_app._MAPPED_PROGRESS_NODE_IDS_LOCK:
|
||||
assert modal_app._MAPPED_PROGRESS_NODE_IDS[
|
||||
with mapped_execution._MAPPED_PROGRESS_NODE_IDS_LOCK:
|
||||
assert mapped_execution._MAPPED_PROGRESS_NODE_IDS[
|
||||
("prompt-1", "component-1", "7")
|
||||
] == "node-b"
|
||||
|
||||
|
||||
Reference in New Issue
Block a user