From 7dbba8ceff13cb87b1c86e126fa3bef1e37a7a9e Mon Sep 17 00:00:00 2001 From: Ken Simpson Date: Sat, 29 Aug 2026 12:54:47 -0700 Subject: [PATCH] Extract prompt affinity planning --- .gitignore | 1 + README.md | 2 + docs/LEARNINGS.md | 4 + prompt_affinity_planning.py | 473 ++++++++++++++++++++++++++++++ prompt_rewrite.py | 473 ++---------------------------- tests/conftest.py | 6 + tests/test_api_intercept.py | 8 +- tests/test_modal_executor_node.py | 4 +- 8 files changed, 516 insertions(+), 455 deletions(-) create mode 100644 prompt_affinity_planning.py diff --git a/.gitignore b/.gitignore index 57a99c3..70c7ac5 100644 --- a/.gitignore +++ b/.gitignore @@ -62,6 +62,7 @@ !/output_artifacts.py !/prompt_payload_metadata.py !/prompt_diagnostics.py +!/prompt_affinity_planning.py !/prompt_rewrite.py !/proxy_node_factory.py !/proxy_payloads.py diff --git a/README.md b/README.md index 78d68ad..d69015f 100644 --- a/README.md +++ b/README.md @@ -520,6 +520,8 @@ The prompt interception layer is organized as a compatibility aggregator over fo Execution scheduling keeps provider-neutral orchestration in `execution_scheduling.py` while resource estimation and history signatures live in `execution_resource_estimates.py`, credential-safe plan rendering and its R2 usage cache live in `execution_plan_reporting.py`, and runtime backend construction plus assignment stamping live in `execution_assignment_runtime.py`. +Prompt rewriting keeps payload and proxy construction in `prompt_rewrite.py`; downstream frontier discovery, parallel local-branch release, local-gap keepalive, and speculative affinity prewarm decoration live in `prompt_affinity_planning.py`. + Resident LLM inference uses the same ownership boundaries: `llm_types.py` defines backend-neutral values, `llm_inputs.py` normalizes multimodal and file inputs, `vllm_instrumentation.py` owns process-wide vLLM execution policy and Triton telemetry, and `llm_backend_transformers.py`, `llm_backend_llamacpp.py`, and `llm_backend_vllm.py` isolate backend implementations. `modal_llm_runtime.py` remains the model-residency coordinator, while `local_llm_runtime.py` consumes shared types and input preparation without importing private helpers from the Modal backend. CPU-side staging adapts the general cross-process ownership lock in `snapshot_lease.py` to LLM-specific progress events. Asset synchronization is composed from focused collaborators: `sync_protocols.py` holds provider-neutral contracts and values, `sync_backends.py` owns local and Modal persistence adapters, `sync_hashing.py` owns the persistent content-hash cache, `sync_custom_nodes.py` builds and transfers custom-node archives and manifests, and `sync_r2_transfer.py` owns worker preflight plus idle-gated R2 read-through and write-back. `sync_engine.py` coordinates those collaborators and retains compatibility entrypoints. diff --git a/docs/LEARNINGS.md b/docs/LEARNINGS.md index 8830f3e..d773ddb 100644 --- a/docs/LEARNINGS.md +++ b/docs/LEARNINGS.md @@ -691,3 +691,7 @@ The entrypoint logger has a similar identity constraint: tests and operators ins ## Mutable caches move with their reporting owner The short-lived R2 bucket-usage cache is plan-reporting state, not scheduler state. Moving the cache, lock, credential lookup, and rendering helpers together into `execution_plan_reporting.py` prevents compatibility re-exports from creating a second mutable binding. Tests that clear the cache or replace credential clients must patch that owner directly; callers may continue to use the compatibility functions exported by `execution_scheduling.py`. + +## Affinity decoration is downstream graph planning + +Parallel local-branch release, local-gap keepalive, and speculative worker prewarm all traverse the rewritten proxy graph after payload construction. Keeping them together in `prompt_affinity_planning.py` gives that post-rewrite graph mutation one owner while leaving `prompt_rewrite.py` focused on constructing payloads and proxy nodes. diff --git a/prompt_affinity_planning.py b/prompt_affinity_planning.py new file mode 100644 index 0000000..3a6b3cb --- /dev/null +++ b/prompt_affinity_planning.py @@ -0,0 +1,473 @@ +"""Affinity, local-gap keepalive, and speculative prewarm prompt rewriting.""" + +from __future__ import annotations + +import copy +from collections import deque +import logging +from typing import Any, Mapping +import uuid + +if __package__: + from .component_planning import _component_ancestors_of_local_source + from .execution_environments import ExecutionProvider + from .modal_executor_node import ( + MODAL_PARALLEL_LOCAL_PASSTHROUGH_NODE_ID, + ensure_modal_parallel_local_passthrough_registered, + registered_proxy_execution_payload, + update_registered_proxy_payload_fields, + ) + from .remote_graph_analysis import ( + _build_consumer_map, + _downstream_node_ids_from_targets, + _is_link, + _node_output_refs, + ) + from .remote_plan_types import InputTarget, LinkedOutputRef +else: # pragma: no cover - flat import inside the Modal container. + from component_planning import _component_ancestors_of_local_source + from execution_environments import ExecutionProvider + from modal_executor_node import ( + MODAL_PARALLEL_LOCAL_PASSTHROUGH_NODE_ID, + ensure_modal_parallel_local_passthrough_registered, + registered_proxy_execution_payload, + update_registered_proxy_payload_fields, + ) + from remote_graph_analysis import ( + _build_consumer_map, + _downstream_node_ids_from_targets, + _is_link, + _node_output_refs, + ) + from remote_plan_types import InputTarget, LinkedOutputRef + +logger = logging.getLogger(__name__) + +_SPECULATIVE_PREWARM_TARGET_KEY = "speculative_remote_prewarm_target" +_SPECULATIVE_PREWARM_PAYLOAD_FIELDS = frozenset( + { + "component_id", + "prompt_id", + "modal_gpu", + "execution_provider", + "execution_environment_id", + "remote_worker_affinity_group", + "remote_local_gap_pool", + "subgraph_prompt", + "requires_volume_reload", + "volume_reload_marker", + "uploaded_volume_paths", + "custom_nodes_bundle", + "snapshot_profile_key", + } +) + + +def _nearest_downstream_remote_component_ids( + *, + rewritten_prompt: dict[str, Any], + consumers: dict[LinkedOutputRef, list[InputTarget]], + seed_targets: list[InputTarget], + remote_component_id_set: set[str], + nodes_module: Any, +) -> list[str]: + """Return the first remote proxies reached along every downstream branch.""" + nearest_remote_component_ids: set[str] = set() + visited_node_ids: set[str] = set() + pending_node_ids = deque(str(target.node_id) for target in seed_targets) + while pending_node_ids: + node_id = pending_node_ids.popleft() + if node_id in visited_node_ids or node_id not in rewritten_prompt: + continue + visited_node_ids.add(node_id) + if node_id in remote_component_id_set: + nearest_remote_component_ids.add(node_id) + continue + for output_ref in _node_output_refs( + rewritten_prompt, + node_id, + nodes_module, + ): + for downstream_target in consumers.get(output_ref, []): + pending_node_ids.append(str(downstream_target.node_id)) + return sorted(nearest_remote_component_ids) + + +def _parallelize_non_returning_local_branches( + *, + rewritten_prompt: dict[str, Any], + remote_component_ids: list[str], + nodes_module: Any, +) -> list[str]: + """Release local-only taps once the nearest remote continuations are in flight.""" + remote_component_id_set = set(remote_component_ids) + if len(remote_component_id_set) < 2: + return [] + + consumers = _build_consumer_map(rewritten_prompt) + ensure_modal_parallel_local_passthrough_registered(nodes_module) + passthrough_node_ids: list[str] = [] + dispatch_group_id = uuid.uuid4().hex + for component_id in remote_component_ids: + embedded_payload = ( + rewritten_prompt.get(component_id, {}) + .get("inputs", {}) + .get("original_node_data") + ) + if not isinstance(embedded_payload, Mapping): + continue + payload_prompt_id = embedded_payload.get("prompt_id") + if payload_prompt_id is not None and str(payload_prompt_id).strip(): + dispatch_group_id = str(payload_prompt_id) + break + + for source_component_id in remote_component_ids: + for source in _node_output_refs( + rewritten_prompt, + source_component_id, + nodes_module, + ): + output_consumers = list(consumers.get(source, [])) + if not output_consumers: + continue + downstream_remote_component_ids = _nearest_downstream_remote_component_ids( + rewritten_prompt=rewritten_prompt, + consumers=consumers, + seed_targets=output_consumers, + remote_component_id_set=remote_component_id_set, + nodes_module=nodes_module, + ) + downstream_remote_component_ids = [ + component_id + for component_id in downstream_remote_component_ids + if component_id != source_component_id + ] + if not downstream_remote_component_ids: + continue + + local_only_consumers: list[InputTarget] = [] + for target in output_consumers: + if target.node_id in remote_component_id_set: + continue + branch_node_ids = _downstream_node_ids_from_targets( + prompt=rewritten_prompt, + consumers=consumers, + seed_targets=[target], + nodes_module=nodes_module, + ) + if branch_node_ids & remote_component_id_set: + continue + local_only_consumers.append(target) + if not local_only_consumers: + continue + + passthrough_node_id = ( + f"__{MODAL_PARALLEL_LOCAL_PASSTHROUGH_NODE_ID}__" + f"{len(passthrough_node_ids)}" + ) + while passthrough_node_id in rewritten_prompt: + passthrough_node_id = f"{passthrough_node_id}_proxy" + passthrough_inputs: dict[str, Any] = { + "value": [source.node_id, source.output_index], + "dispatch_context": { + "dispatch_group_id": dispatch_group_id, + "component_ids": downstream_remote_component_ids, + }, + } + rewritten_prompt[passthrough_node_id] = { + "class_type": MODAL_PARALLEL_LOCAL_PASSTHROUGH_NODE_ID, + "inputs": passthrough_inputs, + "_meta": {"title": "Modal Parallel Local Branch"}, + } + for target in local_only_consumers: + rewritten_prompt[target.node_id]["inputs"][target.input_name] = [ + passthrough_node_id, + 0, + ] + for component_id in downstream_remote_component_ids: + prompt_inputs = rewritten_prompt[component_id].get("inputs", {}) + embedded_payload = prompt_inputs.get("original_node_data") + if not isinstance(embedded_payload, Mapping): + continue + prompt_inputs[ + "original_node_data" + ] = update_registered_proxy_payload_fields( + component_id, + embedded_payload, + { + "parallel_local_dispatch_group_id": dispatch_group_id, + "signal_parallel_local_dispatch": True, + }, + ) + passthrough_node_ids.append(passthrough_node_id) + logger.info( + "Parallelized local-only consumers %s of remote output %s:%d after downstream Modal dispatches %s via passthrough=%s.", + [target.node_id for target in local_only_consumers], + source.node_id, + source.output_index, + downstream_remote_component_ids, + passthrough_node_id, + ) + return passthrough_node_ids + + +def _configure_local_gap_keepalive_payloads( + *, + rewritten_prompt: dict[str, Any], + remote_component_ids: list[str], + sandwiched_local_node_ids: set[str], +) -> None: + """Mark only remote proxies that precede a local-to-remote execution gap.""" + if not sandwiched_local_node_ids: + return + + payloads_by_component_id: dict[str, Mapping[str, Any]] = {} + for component_id in remote_component_ids: + payload = ( + rewritten_prompt.get(component_id, {}) + .get("inputs", {}) + .get("original_node_data") + ) + if not isinstance(payload, dict): + continue + payloads_by_component_id[component_id] = payload + + keepalive_component_ids: set[str] = set() + continuation_component_ids: set[str] = set() + for local_node_id in sandwiched_local_node_ids: + upstream_component_ids: set[str] = set() + downstream_component_ids: set[str] = set() + for component_id in remote_component_ids: + if _component_ancestors_of_local_source( + prompt=rewritten_prompt, + source_node_id=local_node_id, + component_node_id_set={component_id}, + ): + upstream_component_ids.add(component_id) + if _component_ancestors_of_local_source( + prompt=rewritten_prompt, + source_node_id=component_id, + component_node_id_set={local_node_id}, + ): + downstream_component_ids.add(component_id) + + for upstream_component_id in upstream_component_ids: + upstream_payload = payloads_by_component_id.get(upstream_component_id) + if upstream_payload is None: + continue + upstream_environment_id = str( + upstream_payload.get("execution_environment_id") or "" + ) + if upstream_payload.get("execution_provider") != ExecutionProvider.MODAL.value: + continue + for downstream_component_id in downstream_component_ids: + downstream_payload = payloads_by_component_id.get( + downstream_component_id + ) + if ( + downstream_payload is None + or downstream_payload.get("execution_provider") + != ExecutionProvider.MODAL.value + or str(downstream_payload.get("execution_environment_id") or "") + != upstream_environment_id + ): + continue + keepalive_component_ids.add(upstream_component_id) + continuation_component_ids.add(downstream_component_id) + + for component_id, embedded_payload in payloads_by_component_id.items(): + if component_id not in keepalive_component_ids | continuation_component_ids: + continue + payload_fields: dict[str, Any] = {"remote_local_gap_pool": True} + if component_id in keepalive_component_ids: + payload_fields["keepalive_after_remote_component"] = True + if component_id in continuation_component_ids: + payload_fields["stop_local_gap_keepalive_before_remote_component"] = True + rewritten_prompt[component_id]["inputs"][ + "original_node_data" + ] = update_registered_proxy_payload_fields( + component_id, + embedded_payload, + payload_fields, + ) + logger.info( + "Configured local-gap Modal pool for components=%s keepalive_producers=%s continuations=%s.", + sorted(keepalive_component_ids | continuation_component_ids), + sorted(keepalive_component_ids), + sorted(continuation_component_ids), + ) + + +def _remote_proxy_payload( + rewritten_prompt: Mapping[str, Any], component_id: str +) -> Mapping[str, Any] | None: + """Return the registered execution payload embedded in one remote proxy.""" + prompt_node = rewritten_prompt.get(component_id) + if not isinstance(prompt_node, Mapping): + return None + inputs = prompt_node.get("inputs") + if not isinstance(inputs, Mapping): + return None + payload = inputs.get("original_node_data") + if not isinstance(payload, Mapping): + return None + return registered_proxy_execution_payload(component_id, payload) + + +def _remote_proxy_dependency_edges( + rewritten_prompt: Mapping[str, Any], + component_ids: set[str], +) -> dict[str, set[str]]: + """Return nearest remote dependencies while traversing intervening local nodes.""" + dependency_edges = {component_id: set() for component_id in component_ids} + for downstream_component_id in component_ids: + downstream_node = rewritten_prompt.get(downstream_component_id) + if not isinstance(downstream_node, Mapping): + continue + pending_node_ids = [ + str(input_value[0]) + for input_value in (downstream_node.get("inputs") or {}).values() + if _is_link(input_value) + ] + visited_node_ids: set[str] = set() + while pending_node_ids: + upstream_node_id = pending_node_ids.pop() + if upstream_node_id in visited_node_ids: + continue + visited_node_ids.add(upstream_node_id) + if upstream_node_id in component_ids: + dependency_edges[upstream_node_id].add(downstream_component_id) + continue + upstream_node = rewritten_prompt.get(upstream_node_id) + if not isinstance(upstream_node, Mapping): + continue + pending_node_ids.extend( + str(input_value[0]) + for input_value in (upstream_node.get("inputs") or {}).values() + if _is_link(input_value) + ) + return dependency_edges + + +def _component_descendant_distances( + component_id: str, + dependency_edges: Mapping[str, set[str]], +) -> dict[str, int]: + """Return shortest downstream distances from one remote component.""" + distances: dict[str, int] = {} + pending = deque( + (descendant_id, 1) + for descendant_id in dependency_edges.get(component_id, set()) + ) + while pending: + descendant_id, distance = pending.popleft() + previous_distance = distances.get(descendant_id) + if previous_distance is not None and previous_distance <= distance: + continue + distances[descendant_id] = distance + pending.extend( + (nested_descendant_id, distance + 1) + for nested_descendant_id in dependency_edges.get(descendant_id, set()) + ) + return distances + + +def _speculative_prewarm_target_payload( + payload: Mapping[str, Any], +) -> dict[str, Any]: + """Return the minimal side-effect-free payload needed to prepare one worker.""" + return { + field_name: copy.deepcopy(payload[field_name]) + for field_name in _SPECULATIVE_PREWARM_PAYLOAD_FIELDS + if field_name in payload + } + + +def _configure_speculative_affinity_prewarm_payloads( + *, + rewritten_prompt: dict[str, Any], + execution_stages: list[list[str]], +) -> None: + """Attach one next-affinity worker preparation target to each eligible proxy.""" + stage_index_by_component_id = { + component_id: stage_index + for stage_index, stage in enumerate(execution_stages) + for component_id in stage + } + payloads_by_component_id = { + component_id: payload + for component_id in stage_index_by_component_id + if (payload := _remote_proxy_payload(rewritten_prompt, component_id)) + is not None + } + dependency_edges = _remote_proxy_dependency_edges( + rewritten_prompt, + set(payloads_by_component_id), + ) + + configured_targets: dict[str, str] = {} + for component_id, payload in payloads_by_component_id.items(): + if payload.get("execution_provider") != ExecutionProvider.MODAL.value: + continue + execution_environment_id = str( + payload.get("execution_environment_id") or "" + ) + affinity_group = ( + str(payload.get("remote_worker_affinity_group") or "comfy").strip().lower() + ) + descendant_distances = _component_descendant_distances( + component_id, dependency_edges + ) + candidate_ids = sorted( + ( + descendant_id + for descendant_id in descendant_distances + if descendant_id in payloads_by_component_id + and payloads_by_component_id[descendant_id].get( + "execution_provider" + ) + == ExecutionProvider.MODAL.value + and str( + payloads_by_component_id[descendant_id].get( + "execution_environment_id" + ) + or "" + ) + == execution_environment_id + and str( + payloads_by_component_id[descendant_id].get( + "remote_worker_affinity_group" + ) + or "comfy" + ) + .strip() + .lower() + != affinity_group + ), + key=lambda descendant_id: ( + descendant_distances[descendant_id], + stage_index_by_component_id[descendant_id], + descendant_id, + ), + ) + if not candidate_ids: + continue + + target_component_id = candidate_ids[0] + target_payload = _speculative_prewarm_target_payload( + payloads_by_component_id[target_component_id] + ) + rewritten_prompt[component_id]["inputs"][ + "original_node_data" + ] = update_registered_proxy_payload_fields( + component_id, + payload, + {_SPECULATIVE_PREWARM_TARGET_KEY: target_payload}, + ) + configured_targets[component_id] = target_component_id + + if configured_targets: + logger.info( + "Configured one-step speculative Modal worker prewarm targets=%s.", + configured_targets, + ) diff --git a/prompt_rewrite.py b/prompt_rewrite.py index eb20ffa..e883f05 100644 --- a/prompt_rewrite.py +++ b/prompt_rewrite.py @@ -5,14 +5,12 @@ from __future__ import annotations import copy import logging import uuid -from collections import deque from dataclasses import dataclass from typing import Any, Mapping if __package__: from .component_planning import ( _boundary_output_payload, - _component_ancestors_of_local_source, _component_has_local_reentry_dependency, _component_has_parallel_local_remote_fanout, _component_upstream_closure, @@ -23,32 +21,36 @@ if __package__: _subgraph_topological_node_order, _subset_component_prompt, ) - from .execution_environments import ExecutionProvider from .modal_executor_node import ( MODAL_ARTIFACT_FINALIZER_MAX_COMPONENTS, MODAL_ARTIFACT_FINALIZER_NODE_ID, MODAL_COMPONENT_COMPLETION_OUTPUT_NAME, MODAL_LOCAL_BRIDGE_MATERIALIZER_NODE_ID, - MODAL_PARALLEL_LOCAL_PASSTHROUGH_NODE_ID, ensure_modal_artifact_finalizer_registered, ensure_modal_component_proxy_node_registered, ensure_modal_local_bridge_materializer_registered, - ensure_modal_parallel_local_passthrough_registered, register_cache_friendly_proxy_payload, register_modal_map_input_warmup_context, - registered_proxy_execution_payload, - update_registered_proxy_payload_fields, ) from .prompt_payload_metadata import ( _attach_snapshot_profile_key, _serialize_boundary_input_specs, ) + from .prompt_affinity_planning import ( + _SPECULATIVE_PREWARM_PAYLOAD_FIELDS, + _SPECULATIVE_PREWARM_TARGET_KEY, + _component_descendant_distances, + _configure_local_gap_keepalive_payloads, + _configure_speculative_affinity_prewarm_payloads, + _nearest_downstream_remote_component_ids, + _parallelize_non_returning_local_branches, + _remote_proxy_dependency_edges, + _remote_proxy_payload, + _speculative_prewarm_target_payload, + ) from .remote_graph_analysis import ( - _build_consumer_map, - _downstream_node_ids_from_targets, _is_link, _iter_payload_input_strings, - _node_output_refs, _normalize_output_metadata, _remote_output_io_type, _remote_output_is_list, @@ -67,7 +69,6 @@ if __package__: else: # pragma: no cover - flat import inside the Modal container. from component_planning import ( _boundary_output_payload, - _component_ancestors_of_local_source, _component_has_local_reentry_dependency, _component_has_parallel_local_remote_fanout, _component_upstream_closure, @@ -78,32 +79,36 @@ else: # pragma: no cover - flat import inside the Modal container. _subgraph_topological_node_order, _subset_component_prompt, ) - from execution_environments import ExecutionProvider from modal_executor_node import ( MODAL_ARTIFACT_FINALIZER_MAX_COMPONENTS, MODAL_ARTIFACT_FINALIZER_NODE_ID, MODAL_COMPONENT_COMPLETION_OUTPUT_NAME, MODAL_LOCAL_BRIDGE_MATERIALIZER_NODE_ID, - MODAL_PARALLEL_LOCAL_PASSTHROUGH_NODE_ID, ensure_modal_artifact_finalizer_registered, ensure_modal_component_proxy_node_registered, ensure_modal_local_bridge_materializer_registered, - ensure_modal_parallel_local_passthrough_registered, register_cache_friendly_proxy_payload, register_modal_map_input_warmup_context, - registered_proxy_execution_payload, - update_registered_proxy_payload_fields, ) from prompt_payload_metadata import ( _attach_snapshot_profile_key, _serialize_boundary_input_specs, ) + from prompt_affinity_planning import ( + _SPECULATIVE_PREWARM_PAYLOAD_FIELDS, + _SPECULATIVE_PREWARM_TARGET_KEY, + _component_descendant_distances, + _configure_local_gap_keepalive_payloads, + _configure_speculative_affinity_prewarm_payloads, + _nearest_downstream_remote_component_ids, + _parallelize_non_returning_local_branches, + _remote_proxy_dependency_edges, + _remote_proxy_payload, + _speculative_prewarm_target_payload, + ) from remote_graph_analysis import ( - _build_consumer_map, - _downstream_node_ids_from_targets, _is_link, _iter_payload_input_strings, - _node_output_refs, _normalize_output_metadata, _remote_output_io_type, _remote_output_is_list, @@ -122,26 +127,6 @@ else: # pragma: no cover - flat import inside the Modal container. logger = logging.getLogger(__name__) -_SPECULATIVE_PREWARM_TARGET_KEY = "speculative_remote_prewarm_target" -_SPECULATIVE_PREWARM_PAYLOAD_FIELDS = frozenset( - { - "component_id", - "prompt_id", - "modal_gpu", - "execution_provider", - "execution_environment_id", - "remote_worker_affinity_group", - "remote_local_gap_pool", - "subgraph_prompt", - "requires_volume_reload", - "volume_reload_marker", - "uploaded_volume_paths", - "custom_nodes_bundle", - "snapshot_profile_key", - } -) - - def _sync_component_prompt_inputs( component: RemoteComponentPlan, rewritten_prompt: dict[str, Any], @@ -1397,413 +1382,3 @@ def _attach_modal_artifact_finalizer( remote_component_ids, ) return finalizer_node_id - - -def _nearest_downstream_remote_component_ids( - *, - rewritten_prompt: dict[str, Any], - consumers: dict[LinkedOutputRef, list[InputTarget]], - seed_targets: list[InputTarget], - remote_component_id_set: set[str], - nodes_module: Any, -) -> list[str]: - """Return the first remote proxies reached along every downstream branch.""" - nearest_remote_component_ids: set[str] = set() - visited_node_ids: set[str] = set() - pending_node_ids = deque(str(target.node_id) for target in seed_targets) - while pending_node_ids: - node_id = pending_node_ids.popleft() - if node_id in visited_node_ids or node_id not in rewritten_prompt: - continue - visited_node_ids.add(node_id) - if node_id in remote_component_id_set: - nearest_remote_component_ids.add(node_id) - continue - for output_ref in _node_output_refs( - rewritten_prompt, - node_id, - nodes_module, - ): - for downstream_target in consumers.get(output_ref, []): - pending_node_ids.append(str(downstream_target.node_id)) - return sorted(nearest_remote_component_ids) - - -def _parallelize_non_returning_local_branches( - *, - rewritten_prompt: dict[str, Any], - remote_component_ids: list[str], - nodes_module: Any, -) -> list[str]: - """Release local-only taps once the nearest remote continuations are in flight.""" - remote_component_id_set = set(remote_component_ids) - if len(remote_component_id_set) < 2: - return [] - - consumers = _build_consumer_map(rewritten_prompt) - ensure_modal_parallel_local_passthrough_registered(nodes_module) - passthrough_node_ids: list[str] = [] - dispatch_group_id = uuid.uuid4().hex - for component_id in remote_component_ids: - embedded_payload = ( - rewritten_prompt.get(component_id, {}) - .get("inputs", {}) - .get("original_node_data") - ) - if not isinstance(embedded_payload, Mapping): - continue - payload_prompt_id = embedded_payload.get("prompt_id") - if payload_prompt_id is not None and str(payload_prompt_id).strip(): - dispatch_group_id = str(payload_prompt_id) - break - - for source_component_id in remote_component_ids: - for source in _node_output_refs( - rewritten_prompt, - source_component_id, - nodes_module, - ): - output_consumers = list(consumers.get(source, [])) - if not output_consumers: - continue - downstream_remote_component_ids = _nearest_downstream_remote_component_ids( - rewritten_prompt=rewritten_prompt, - consumers=consumers, - seed_targets=output_consumers, - remote_component_id_set=remote_component_id_set, - nodes_module=nodes_module, - ) - downstream_remote_component_ids = [ - component_id - for component_id in downstream_remote_component_ids - if component_id != source_component_id - ] - if not downstream_remote_component_ids: - continue - - local_only_consumers: list[InputTarget] = [] - for target in output_consumers: - if target.node_id in remote_component_id_set: - continue - branch_node_ids = _downstream_node_ids_from_targets( - prompt=rewritten_prompt, - consumers=consumers, - seed_targets=[target], - nodes_module=nodes_module, - ) - if branch_node_ids & remote_component_id_set: - continue - local_only_consumers.append(target) - if not local_only_consumers: - continue - - passthrough_node_id = ( - f"__{MODAL_PARALLEL_LOCAL_PASSTHROUGH_NODE_ID}__" - f"{len(passthrough_node_ids)}" - ) - while passthrough_node_id in rewritten_prompt: - passthrough_node_id = f"{passthrough_node_id}_proxy" - passthrough_inputs: dict[str, Any] = { - "value": [source.node_id, source.output_index], - "dispatch_context": { - "dispatch_group_id": dispatch_group_id, - "component_ids": downstream_remote_component_ids, - }, - } - rewritten_prompt[passthrough_node_id] = { - "class_type": MODAL_PARALLEL_LOCAL_PASSTHROUGH_NODE_ID, - "inputs": passthrough_inputs, - "_meta": {"title": "Modal Parallel Local Branch"}, - } - for target in local_only_consumers: - rewritten_prompt[target.node_id]["inputs"][target.input_name] = [ - passthrough_node_id, - 0, - ] - for component_id in downstream_remote_component_ids: - prompt_inputs = rewritten_prompt[component_id].get("inputs", {}) - embedded_payload = prompt_inputs.get("original_node_data") - if not isinstance(embedded_payload, Mapping): - continue - prompt_inputs[ - "original_node_data" - ] = update_registered_proxy_payload_fields( - component_id, - embedded_payload, - { - "parallel_local_dispatch_group_id": dispatch_group_id, - "signal_parallel_local_dispatch": True, - }, - ) - passthrough_node_ids.append(passthrough_node_id) - logger.info( - "Parallelized local-only consumers %s of remote output %s:%d after downstream Modal dispatches %s via passthrough=%s.", - [target.node_id for target in local_only_consumers], - source.node_id, - source.output_index, - downstream_remote_component_ids, - passthrough_node_id, - ) - return passthrough_node_ids - - -def _configure_local_gap_keepalive_payloads( - *, - rewritten_prompt: dict[str, Any], - remote_component_ids: list[str], - sandwiched_local_node_ids: set[str], -) -> None: - """Mark only remote proxies that precede a local-to-remote execution gap.""" - if not sandwiched_local_node_ids: - return - - payloads_by_component_id: dict[str, Mapping[str, Any]] = {} - for component_id in remote_component_ids: - payload = ( - rewritten_prompt.get(component_id, {}) - .get("inputs", {}) - .get("original_node_data") - ) - if not isinstance(payload, dict): - continue - payloads_by_component_id[component_id] = payload - - keepalive_component_ids: set[str] = set() - continuation_component_ids: set[str] = set() - for local_node_id in sandwiched_local_node_ids: - upstream_component_ids: set[str] = set() - downstream_component_ids: set[str] = set() - for component_id in remote_component_ids: - if _component_ancestors_of_local_source( - prompt=rewritten_prompt, - source_node_id=local_node_id, - component_node_id_set={component_id}, - ): - upstream_component_ids.add(component_id) - if _component_ancestors_of_local_source( - prompt=rewritten_prompt, - source_node_id=component_id, - component_node_id_set={local_node_id}, - ): - downstream_component_ids.add(component_id) - - for upstream_component_id in upstream_component_ids: - upstream_payload = payloads_by_component_id.get(upstream_component_id) - if upstream_payload is None: - continue - upstream_environment_id = str( - upstream_payload.get("execution_environment_id") or "" - ) - if upstream_payload.get("execution_provider") != ExecutionProvider.MODAL.value: - continue - for downstream_component_id in downstream_component_ids: - downstream_payload = payloads_by_component_id.get( - downstream_component_id - ) - if ( - downstream_payload is None - or downstream_payload.get("execution_provider") - != ExecutionProvider.MODAL.value - or str(downstream_payload.get("execution_environment_id") or "") - != upstream_environment_id - ): - continue - keepalive_component_ids.add(upstream_component_id) - continuation_component_ids.add(downstream_component_id) - - for component_id, embedded_payload in payloads_by_component_id.items(): - if component_id not in keepalive_component_ids | continuation_component_ids: - continue - payload_fields: dict[str, Any] = {"remote_local_gap_pool": True} - if component_id in keepalive_component_ids: - payload_fields["keepalive_after_remote_component"] = True - if component_id in continuation_component_ids: - payload_fields["stop_local_gap_keepalive_before_remote_component"] = True - rewritten_prompt[component_id]["inputs"][ - "original_node_data" - ] = update_registered_proxy_payload_fields( - component_id, - embedded_payload, - payload_fields, - ) - logger.info( - "Configured local-gap Modal pool for components=%s keepalive_producers=%s continuations=%s.", - sorted(keepalive_component_ids | continuation_component_ids), - sorted(keepalive_component_ids), - sorted(continuation_component_ids), - ) - - -def _remote_proxy_payload( - rewritten_prompt: Mapping[str, Any], component_id: str -) -> Mapping[str, Any] | None: - """Return the registered execution payload embedded in one remote proxy.""" - prompt_node = rewritten_prompt.get(component_id) - if not isinstance(prompt_node, Mapping): - return None - inputs = prompt_node.get("inputs") - if not isinstance(inputs, Mapping): - return None - payload = inputs.get("original_node_data") - if not isinstance(payload, Mapping): - return None - return registered_proxy_execution_payload(component_id, payload) - - -def _remote_proxy_dependency_edges( - rewritten_prompt: Mapping[str, Any], - component_ids: set[str], -) -> dict[str, set[str]]: - """Return nearest remote dependencies while traversing intervening local nodes.""" - dependency_edges = {component_id: set() for component_id in component_ids} - for downstream_component_id in component_ids: - downstream_node = rewritten_prompt.get(downstream_component_id) - if not isinstance(downstream_node, Mapping): - continue - pending_node_ids = [ - str(input_value[0]) - for input_value in (downstream_node.get("inputs") or {}).values() - if _is_link(input_value) - ] - visited_node_ids: set[str] = set() - while pending_node_ids: - upstream_node_id = pending_node_ids.pop() - if upstream_node_id in visited_node_ids: - continue - visited_node_ids.add(upstream_node_id) - if upstream_node_id in component_ids: - dependency_edges[upstream_node_id].add(downstream_component_id) - continue - upstream_node = rewritten_prompt.get(upstream_node_id) - if not isinstance(upstream_node, Mapping): - continue - pending_node_ids.extend( - str(input_value[0]) - for input_value in (upstream_node.get("inputs") or {}).values() - if _is_link(input_value) - ) - return dependency_edges - - -def _component_descendant_distances( - component_id: str, - dependency_edges: Mapping[str, set[str]], -) -> dict[str, int]: - """Return shortest downstream distances from one remote component.""" - distances: dict[str, int] = {} - pending = deque( - (descendant_id, 1) - for descendant_id in dependency_edges.get(component_id, set()) - ) - while pending: - descendant_id, distance = pending.popleft() - previous_distance = distances.get(descendant_id) - if previous_distance is not None and previous_distance <= distance: - continue - distances[descendant_id] = distance - pending.extend( - (nested_descendant_id, distance + 1) - for nested_descendant_id in dependency_edges.get(descendant_id, set()) - ) - return distances - - -def _speculative_prewarm_target_payload( - payload: Mapping[str, Any], -) -> dict[str, Any]: - """Return the minimal side-effect-free payload needed to prepare one worker.""" - return { - field_name: copy.deepcopy(payload[field_name]) - for field_name in _SPECULATIVE_PREWARM_PAYLOAD_FIELDS - if field_name in payload - } - - -def _configure_speculative_affinity_prewarm_payloads( - *, - rewritten_prompt: dict[str, Any], - execution_stages: list[list[str]], -) -> None: - """Attach one next-affinity worker preparation target to each eligible proxy.""" - stage_index_by_component_id = { - component_id: stage_index - for stage_index, stage in enumerate(execution_stages) - for component_id in stage - } - payloads_by_component_id = { - component_id: payload - for component_id in stage_index_by_component_id - if (payload := _remote_proxy_payload(rewritten_prompt, component_id)) - is not None - } - dependency_edges = _remote_proxy_dependency_edges( - rewritten_prompt, - set(payloads_by_component_id), - ) - - configured_targets: dict[str, str] = {} - for component_id, payload in payloads_by_component_id.items(): - if payload.get("execution_provider") != ExecutionProvider.MODAL.value: - continue - execution_environment_id = str( - payload.get("execution_environment_id") or "" - ) - affinity_group = ( - str(payload.get("remote_worker_affinity_group") or "comfy").strip().lower() - ) - descendant_distances = _component_descendant_distances( - component_id, dependency_edges - ) - candidate_ids = sorted( - ( - descendant_id - for descendant_id in descendant_distances - if descendant_id in payloads_by_component_id - and payloads_by_component_id[descendant_id].get( - "execution_provider" - ) - == ExecutionProvider.MODAL.value - and str( - payloads_by_component_id[descendant_id].get( - "execution_environment_id" - ) - or "" - ) - == execution_environment_id - and str( - payloads_by_component_id[descendant_id].get( - "remote_worker_affinity_group" - ) - or "comfy" - ) - .strip() - .lower() - != affinity_group - ), - key=lambda descendant_id: ( - descendant_distances[descendant_id], - stage_index_by_component_id[descendant_id], - descendant_id, - ), - ) - if not candidate_ids: - continue - - target_component_id = candidate_ids[0] - target_payload = _speculative_prewarm_target_payload( - payloads_by_component_id[target_component_id] - ) - rewritten_prompt[component_id]["inputs"][ - "original_node_data" - ] = update_registered_proxy_payload_fields( - component_id, - payload, - {_SPECULATIVE_PREWARM_TARGET_KEY: target_payload}, - ) - configured_targets[component_id] = target_component_id - - if configured_targets: - logger.info( - "Configured one-step speculative Modal worker prewarm targets=%s.", - configured_targets, - ) diff --git a/tests/conftest.py b/tests/conftest.py index 5b75bc9..908b702 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -345,6 +345,12 @@ def prompt_rewrite_module(extension_package: object) -> object: return importlib.import_module(f"{PACKAGE_NAME}.prompt_rewrite") +@pytest.fixture(scope="session") +def prompt_affinity_planning_module(extension_package: object) -> object: + """Return affinity, keepalive, and speculative prewarm rewrite helpers.""" + return importlib.import_module(f"{PACKAGE_NAME}.prompt_affinity_planning") + + @pytest.fixture(scope="session") def execution_scheduling_module(extension_package: object) -> object: """Return provider-neutral scheduling and capacity helpers.""" diff --git a/tests/test_api_intercept.py b/tests/test_api_intercept.py index a33d208..08f0283 100644 --- a/tests/test_api_intercept.py +++ b/tests/test_api_intercept.py @@ -4524,7 +4524,7 @@ def test_planner_resolved_llm_profile_is_attached_to_matching_payload( def test_planner_attaches_next_distinct_affinity_as_speculative_prewarm_target( api_intercept_module: Any, - prompt_rewrite_module: Any, + prompt_affinity_planning_module: Any, ) -> None: """Each proxy should prepare only its nearest reachable future worker group.""" rewritten_prompt = { @@ -4585,7 +4585,7 @@ def test_planner_attaches_next_distinct_affinity_as_speculative_prewarm_target( }, } - prompt_rewrite_module._configure_speculative_affinity_prewarm_payloads( + prompt_affinity_planning_module._configure_speculative_affinity_prewarm_payloads( rewritten_prompt=rewritten_prompt, execution_stages=[["spec-a", "spec-b"], ["spec-c"]], ) @@ -4612,7 +4612,7 @@ def test_planner_attaches_next_distinct_affinity_as_speculative_prewarm_target( def test_planner_does_not_bridge_local_gap_keepalive_across_providers( api_intercept_module: Any, - prompt_rewrite_module: Any, + prompt_affinity_planning_module: Any, ) -> None: """A Modal producer must not retain a slot for an SSH continuation.""" rewritten_prompt = { @@ -4643,7 +4643,7 @@ def test_planner_does_not_bridge_local_gap_keepalive_across_providers( }, } - prompt_rewrite_module._configure_local_gap_keepalive_payloads( + prompt_affinity_planning_module._configure_local_gap_keepalive_payloads( rewritten_prompt=rewritten_prompt, remote_component_ids=["modal-producer", "ssh-consumer"], sandwiched_local_node_ids={"local-gap"}, diff --git a/tests/test_modal_executor_node.py b/tests/test_modal_executor_node.py index 6dcfb00..81505b9 100644 --- a/tests/test_modal_executor_node.py +++ b/tests/test_modal_executor_node.py @@ -1433,7 +1433,7 @@ def test_local_bridge_materialization_downloads_object_backed_output( def test_parallel_local_dispatch_frontier_stops_at_nearest_remote_component( api_intercept_module: Any, - prompt_rewrite_module: Any, + prompt_affinity_planning_module: Any, remote_graph_analysis_module: Any, ) -> None: """A preview gate should not wait for remote descendants of a remote frontier.""" @@ -1461,7 +1461,7 @@ def test_parallel_local_dispatch_frontier_stops_at_nearest_remote_component( consumers = remote_graph_analysis_module._build_consumer_map(rewritten_prompt) nearest_component_ids = ( - prompt_rewrite_module._nearest_downstream_remote_component_ids( + prompt_affinity_planning_module._nearest_downstream_remote_component_ids( rewritten_prompt=rewritten_prompt, consumers=consumers, seed_targets=consumers[