Extract prompt affinity planning
This commit is contained in:
Ken Simpson
@@ -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
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
+24
-449
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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."""
|
||||
|
||||
@@ -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"},
|
||||
|
||||
@@ -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[
|
||||
|
||||
Reference in New Issue
Block a user