517 lines
21 KiB
Python
517 lines
21 KiB
Python
"""Remote prompt analysis, preparation cancellation, and queue dispatch routes."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import threading
|
|
import time
|
|
import uuid
|
|
from dataclasses import dataclass, field
|
|
from typing import Any, Mapping
|
|
|
|
from aiohttp import web
|
|
|
|
if __package__:
|
|
from .modal_executor_node import MODAL_PROMPT_ID_EXTRA_PNGINFO_KEY
|
|
from .queue_bridge import (
|
|
_cancel_remote_preparation,
|
|
_clear_remote_preparation,
|
|
_queue_prompt_json,
|
|
_queued_ssh_environment_ids,
|
|
_set_remote_preparation,
|
|
)
|
|
from .remote_graph_analysis import (
|
|
analyze_remote_node_selection,
|
|
requested_remote_node_ids,
|
|
)
|
|
from .remote_plan_types import ModalPromptValidationError, RewriteSummary
|
|
from .route_context import RouteContext
|
|
from .settings import ModalSyncSettings, modal_gpu_from_workflow, settings_for_modal_gpu
|
|
from .sync_engine import (
|
|
SyncCancelledError,
|
|
begin_r2_writeback_prompt,
|
|
finish_r2_writeback_prompt,
|
|
)
|
|
else: # pragma: no cover - flat import inside the Modal container.
|
|
from modal_executor_node import MODAL_PROMPT_ID_EXTRA_PNGINFO_KEY
|
|
from queue_bridge import (
|
|
_cancel_remote_preparation,
|
|
_clear_remote_preparation,
|
|
_queue_prompt_json,
|
|
_queued_ssh_environment_ids,
|
|
_set_remote_preparation,
|
|
)
|
|
from remote_graph_analysis import (
|
|
analyze_remote_node_selection,
|
|
requested_remote_node_ids,
|
|
)
|
|
from remote_plan_types import ModalPromptValidationError, RewriteSummary
|
|
from route_context import RouteContext
|
|
from settings import ModalSyncSettings, modal_gpu_from_workflow, settings_for_modal_gpu
|
|
from sync_engine import (
|
|
SyncCancelledError,
|
|
begin_r2_writeback_prompt,
|
|
finish_r2_writeback_prompt,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class _QueueRequestState:
|
|
"""Track one queue request across validation, rewrite, and dispatch stages."""
|
|
|
|
json_data: dict[str, Any] | None = None
|
|
workflow: dict[str, Any] | None = None
|
|
remote_node_ids: list[str] = field(default_factory=list)
|
|
request_modal_gpu: str | None = None
|
|
summary: RewriteSummary = field(default_factory=RewriteSummary)
|
|
preparation_prompt_id: str | None = None
|
|
preparation_cancellation: threading.Event = field(default_factory=threading.Event)
|
|
configurator_node_id: str | None = None
|
|
r2_writeback_prompt_id: str | None = None
|
|
prompt_queued: bool = False
|
|
request_settings: ModalSyncSettings | None = None
|
|
client_id: str | None = None
|
|
prompt_id: str | None = None
|
|
status_modal_gpu: str | None = None
|
|
selected_modal_gpus: list[str] = field(default_factory=list)
|
|
|
|
|
|
async def _initialize_queue_request(
|
|
request: web.Request,
|
|
ctx: RouteContext,
|
|
state: _QueueRequestState,
|
|
) -> None:
|
|
"""Parse request JSON and resolve workflow-scoped selection metadata."""
|
|
json_data = await request.json()
|
|
json_data.setdefault("prompt_id", str(uuid.uuid4()))
|
|
json_data.setdefault("extra_data", {})
|
|
json_data["extra_data"]["prompt_id"] = json_data["prompt_id"]
|
|
if json_data.get("client_id") is not None:
|
|
json_data["extra_data"]["client_id"] = json_data["client_id"]
|
|
state.json_data = json_data
|
|
state.client_id = str(json_data.get("client_id")) if json_data.get("client_id") else None
|
|
state.prompt_id = str(json_data.get("prompt_id")) if json_data.get("prompt_id") else None
|
|
extra_pnginfo = (json_data.get("extra_data") or {}).get("extra_pnginfo") or {}
|
|
state.workflow = extra_pnginfo.get("workflow")
|
|
state.remote_node_ids = sorted(
|
|
requested_remote_node_ids(
|
|
prompt=json_data.get("prompt", {}),
|
|
workflow=state.workflow,
|
|
settings=ctx.settings,
|
|
)
|
|
)
|
|
state.configurator_node_id = ctx.configurator_node_id(json_data.get("prompt", {}))
|
|
|
|
|
|
def _reserve_queue_preparation(
|
|
prompt_server: Any,
|
|
state: _QueueRequestState,
|
|
) -> None:
|
|
"""Reserve R2 and queue-visible preparation state for one remote prompt."""
|
|
if state.prompt_id is None or state.json_data is None:
|
|
return
|
|
begin_r2_writeback_prompt(state.prompt_id)
|
|
state.r2_writeback_prompt_id = state.prompt_id
|
|
extra_pnginfo = (state.json_data.get("extra_data") or {}).get("extra_pnginfo") or {}
|
|
extra_pnginfo[MODAL_PROMPT_ID_EXTRA_PNGINFO_KEY] = state.prompt_id
|
|
state.json_data["extra_data"]["extra_pnginfo"] = extra_pnginfo
|
|
logger.debug(
|
|
"Attached prompt-scoped Modal execution metadata prompt_id=%s.",
|
|
state.prompt_id,
|
|
)
|
|
if _set_remote_preparation(
|
|
prompt_server,
|
|
prompt_id=state.prompt_id,
|
|
prompt=state.json_data.get("prompt", {}),
|
|
extra_data=state.json_data.get("extra_data", {}),
|
|
cancellation_event=state.preparation_cancellation,
|
|
):
|
|
state.preparation_prompt_id = state.prompt_id
|
|
|
|
|
|
def _resolve_queue_settings(ctx: RouteContext, state: _QueueRequestState) -> None:
|
|
"""Resolve workflow GPU settings and status identity for one request."""
|
|
assert state.json_data is not None
|
|
try:
|
|
state.request_settings = settings_for_modal_gpu(
|
|
ctx.settings,
|
|
modal_gpu_from_workflow(state.workflow, ctx.settings.modal_gpu),
|
|
)
|
|
except ValueError as exc:
|
|
raise ModalPromptValidationError(str(exc)) from exc
|
|
logger.info(
|
|
"Resolved workflow Modal GPU selection gpu=%s prompt_id=%s.",
|
|
state.request_settings.modal_gpu,
|
|
state.prompt_id,
|
|
)
|
|
state.status_modal_gpu = (
|
|
None
|
|
if ctx.prompt_uses_configurator(state.json_data.get("prompt", {}))
|
|
else state.request_settings.modal_gpu
|
|
)
|
|
state.request_modal_gpu = state.status_modal_gpu
|
|
|
|
|
|
def _emit_setup_status(
|
|
prompt_server: Any,
|
|
ctx: RouteContext,
|
|
state: _QueueRequestState,
|
|
message: str,
|
|
current: int | None = None,
|
|
total: int | None = None,
|
|
*,
|
|
environment_id: str | None = None,
|
|
) -> None:
|
|
"""Publish one cancellable queue-time setup event."""
|
|
if state.preparation_cancellation.is_set():
|
|
raise SyncCancelledError("Remote workflow preparation was cancelled.")
|
|
ctx.emit_status(
|
|
prompt_server=prompt_server,
|
|
phase="setup",
|
|
client_id=state.client_id,
|
|
prompt_id=state.prompt_id,
|
|
node_ids=state.remote_node_ids,
|
|
configurator_node_id=state.configurator_node_id,
|
|
modal_gpu=state.status_modal_gpu,
|
|
component_node_ids_by_representative=(
|
|
state.summary.component_node_ids_by_representative or None
|
|
),
|
|
status_message=message,
|
|
status_current=current,
|
|
status_total=total,
|
|
**(
|
|
{"execution_environment_id": environment_id}
|
|
if environment_id is not None
|
|
else {}
|
|
),
|
|
)
|
|
|
|
|
|
def _emit_execution_plan(
|
|
prompt_server: Any,
|
|
ctx: RouteContext,
|
|
state: _QueueRequestState,
|
|
assignments: dict[str, dict[str, Any]],
|
|
configurations: list[dict[str, Any]],
|
|
) -> None:
|
|
"""Publish scheduler choices before capacity acquisition can block."""
|
|
component_nodes = {
|
|
component_id: list(assignment.get("node_ids", []))
|
|
for component_id, assignment in assignments.items()
|
|
}
|
|
ctx.emit_status(
|
|
prompt_server=prompt_server,
|
|
phase="setup",
|
|
client_id=state.client_id,
|
|
prompt_id=state.prompt_id,
|
|
node_ids=state.remote_node_ids,
|
|
configurator_node_id=state.configurator_node_id,
|
|
modal_gpu=state.status_modal_gpu,
|
|
component_node_ids_by_representative=component_nodes,
|
|
status_message="Remote execution plan ready",
|
|
remote_execution_assignments=assignments,
|
|
remote_execution_configurations=configurations,
|
|
)
|
|
|
|
|
|
def _store_rewrite_metadata(ctx: RouteContext, state: _QueueRequestState) -> None:
|
|
"""Attach rewrite and execution-plan metadata to the queued request."""
|
|
assert state.json_data is not None
|
|
assert state.request_settings is not None
|
|
modal_extra = state.json_data.setdefault("extra_data", {}).setdefault("modal", {})
|
|
modal_extra.update(
|
|
{
|
|
"gpu": state.request_settings.modal_gpu,
|
|
"remote_node_ids": state.summary.remote_node_ids,
|
|
"remote_component_ids": state.summary.remote_component_ids,
|
|
"component_dependency_ids_by_representative": state.summary.component_dependency_ids_by_representative,
|
|
"component_execution_stages": state.summary.component_execution_stages,
|
|
"mapped_component_ids": state.summary.mapped_component_ids,
|
|
"estimated_max_parallel_requests": state.summary.estimated_max_parallel_requests,
|
|
"max_parallel_requests_upper_bound": state.summary.max_parallel_requests_upper_bound,
|
|
"synced_assets": [asset.remote_path for asset in state.summary.synced_assets],
|
|
}
|
|
)
|
|
state.json_data["extra_data"]["remote_execution"] = {
|
|
"assignments": ctx.execution_assignments_payload(
|
|
state.summary,
|
|
state.request_settings,
|
|
),
|
|
"configurations": list(state.summary.remote_configurations),
|
|
}
|
|
if state.summary.custom_nodes_bundle is not None:
|
|
modal_extra["custom_nodes_bundle"] = state.summary.custom_nodes_bundle.remote_path
|
|
|
|
|
|
async def _rewrite_remote_prompt(
|
|
prompt_server: Any,
|
|
ctx: RouteContext,
|
|
state: _QueueRequestState,
|
|
) -> None:
|
|
"""Analyze, prepare capacity, rewrite the prompt, and attach dispatch metadata."""
|
|
assert state.json_data is not None
|
|
assert state.request_settings is not None
|
|
_emit_setup_status(prompt_server, ctx, state, "Preparing remote workflow")
|
|
rewrite_started_at = time.perf_counter()
|
|
occupied_environment_ids = _queued_ssh_environment_ids(
|
|
prompt_server,
|
|
excluding_prompt_id=state.prompt_id,
|
|
)
|
|
if occupied_environment_ids:
|
|
logger.info(
|
|
"Queued workflow may reuse SSH capacity after earlier prompts finish environments=%s.",
|
|
sorted(occupied_environment_ids),
|
|
)
|
|
rewritten_prompt, state.summary = await ctx.rewrite_prompt(
|
|
prompt=state.json_data["prompt"],
|
|
workflow=state.workflow,
|
|
sync_engine=ctx.sync_engine,
|
|
settings=state.request_settings,
|
|
extra_data=state.json_data.get("extra_data"),
|
|
status_callback=lambda message, current=None, total=None: _emit_setup_status(
|
|
prompt_server, ctx, state, message, current, total
|
|
),
|
|
environment_status_callback=lambda environment_id, message, current=None, total=None: _emit_setup_status(
|
|
prompt_server,
|
|
ctx,
|
|
state,
|
|
message,
|
|
current,
|
|
total,
|
|
environment_id=environment_id,
|
|
),
|
|
plan_callback=lambda assignments, configurations: _emit_execution_plan(
|
|
prompt_server, ctx, state, assignments, configurations
|
|
),
|
|
cancellation_check=state.preparation_cancellation.is_set,
|
|
occupied_environment_ids=occupied_environment_ids,
|
|
)
|
|
state.selected_modal_gpus = ctx.selected_modal_gpus(
|
|
state.summary,
|
|
state.request_settings.modal_gpu,
|
|
)
|
|
state.status_modal_gpu = (
|
|
state.selected_modal_gpus[0] if len(state.selected_modal_gpus) == 1 else None
|
|
)
|
|
state.request_modal_gpu = state.status_modal_gpu
|
|
logger.info(
|
|
"Modal prompt rewrite finished in %.3fs for %d remote nodes across %d components.",
|
|
time.perf_counter() - rewrite_started_at,
|
|
len(state.summary.remote_node_ids),
|
|
len(state.summary.remote_component_ids),
|
|
)
|
|
state.remote_node_ids = list(state.summary.remote_node_ids)
|
|
state.json_data["prompt"] = rewritten_prompt
|
|
if state.json_data.get("partial_execution_targets"):
|
|
rewritten_targets = {
|
|
state.summary.rewritten_node_id_map.get(str(target), str(target))
|
|
for target in state.json_data["partial_execution_targets"]
|
|
}
|
|
state.json_data["partial_execution_targets"] = sorted(rewritten_targets)
|
|
_store_rewrite_metadata(ctx, state)
|
|
_emit_setup_status(prompt_server, ctx, state, "Submitting remote workflow")
|
|
|
|
|
|
def _modal_queue_response_payload(
|
|
ctx: RouteContext,
|
|
state: _QueueRequestState,
|
|
) -> dict[str, Any] | None:
|
|
"""Build the Modal-specific successful queue response payload."""
|
|
if not state.summary.remote_node_ids:
|
|
return None
|
|
assert state.request_settings is not None
|
|
return {
|
|
"modal_gpu": state.status_modal_gpu,
|
|
"remote_execution_configurator_node_id": state.configurator_node_id,
|
|
"remote_execution_modal_gpus": state.selected_modal_gpus,
|
|
"modal_remote_node_ids": list(state.summary.remote_node_ids),
|
|
"modal_sandwiched_local_node_ids": list(state.summary.sandwiched_local_node_ids),
|
|
"modal_parallel_local_branch_node_ids": list(
|
|
state.summary.parallel_local_branch_node_ids
|
|
),
|
|
"remote_execution_assignments": ctx.execution_assignments_payload(
|
|
state.summary,
|
|
state.request_settings,
|
|
),
|
|
"remote_execution_configurations": list(state.summary.remote_configurations),
|
|
"modal_components": [
|
|
{
|
|
"representative_node_id": representative_node_id,
|
|
"node_ids": list(component_node_ids),
|
|
}
|
|
for representative_node_id, component_node_ids in sorted(
|
|
state.summary.component_node_ids_by_representative.items()
|
|
)
|
|
],
|
|
}
|
|
|
|
|
|
def _queue_error_response(
|
|
prompt_server: Any,
|
|
ctx: RouteContext,
|
|
state: _QueueRequestState,
|
|
exc: BaseException,
|
|
*,
|
|
phase: str,
|
|
status: int,
|
|
) -> web.Response:
|
|
"""Emit request-scoped UI failure state and return a JSON error response."""
|
|
if state.json_data is not None:
|
|
kwargs = {
|
|
"status_message": str(exc)
|
|
if phase == "execution_interrupted"
|
|
else None,
|
|
"error_message": str(exc) if phase == "error" else None,
|
|
}
|
|
ctx.emit_status(
|
|
prompt_server=prompt_server,
|
|
phase=phase,
|
|
client_id=state.client_id,
|
|
prompt_id=state.prompt_id,
|
|
node_ids=state.remote_node_ids,
|
|
configurator_node_id=state.configurator_node_id,
|
|
modal_gpu=state.request_modal_gpu,
|
|
**{key: value for key, value in kwargs.items() if value is not None},
|
|
)
|
|
return web.json_response(
|
|
{
|
|
"error": str(exc),
|
|
"node_errors": [],
|
|
**({"cancelled": True} if phase == "execution_interrupted" else {}),
|
|
},
|
|
status=status,
|
|
)
|
|
|
|
|
|
async def _handle_modal_queue_prompt(
|
|
request: web.Request,
|
|
prompt_server: Any,
|
|
ctx: RouteContext,
|
|
) -> web.Response:
|
|
"""Run the validate, analyze, plan, rewrite, and dispatch queue stages."""
|
|
logger.info("Received Modal queue request.")
|
|
state = _QueueRequestState()
|
|
request_started_at = time.perf_counter()
|
|
try:
|
|
await _initialize_queue_request(request, ctx, state)
|
|
assert state.json_data is not None
|
|
if "prompt" in state.json_data and not state.remote_node_ids:
|
|
logger.info(
|
|
"No workflow nodes are marked for Modal execution; forwarding prompt without Modal status or rewrite."
|
|
)
|
|
return await _queue_prompt_json(prompt_server, state.json_data)
|
|
_reserve_queue_preparation(prompt_server, state)
|
|
_resolve_queue_settings(ctx, state)
|
|
if "prompt" in state.json_data:
|
|
await _rewrite_remote_prompt(prompt_server, ctx, state)
|
|
response = await _queue_prompt_json(
|
|
prompt_server,
|
|
state.json_data,
|
|
modal_response_payload=_modal_queue_response_payload(ctx, state),
|
|
)
|
|
state.prompt_queued = response.status < 400
|
|
return response
|
|
except SyncCancelledError as exc:
|
|
logger.info("Remote workflow preparation cancelled: %s", exc)
|
|
return _queue_error_response(
|
|
prompt_server, ctx, state, exc, phase="execution_interrupted", status=409
|
|
)
|
|
except FileNotFoundError as exc:
|
|
logger.exception("Modal asset sync failed.")
|
|
return _queue_error_response(
|
|
prompt_server, ctx, state, exc, phase="error", status=400
|
|
)
|
|
except ModalPromptValidationError as exc:
|
|
logger.exception("Modal prompt validation failed.")
|
|
return _queue_error_response(
|
|
prompt_server, ctx, state, exc, phase="error", status=400
|
|
)
|
|
except Exception as exc:
|
|
logger.exception("Modal queue handler failed.")
|
|
return _queue_error_response(
|
|
prompt_server, ctx, state, exc, phase="error", status=500
|
|
)
|
|
finally:
|
|
logger.info(
|
|
"Modal queue request completed in %.3fs.",
|
|
time.perf_counter() - request_started_at,
|
|
)
|
|
if state.preparation_prompt_id is not None:
|
|
_clear_remote_preparation(prompt_server, state.preparation_prompt_id)
|
|
if state.r2_writeback_prompt_id is not None and not state.prompt_queued:
|
|
finish_r2_writeback_prompt(state.r2_writeback_prompt_id)
|
|
|
|
|
|
def register_queue_routes(prompt_server: Any, ctx: RouteContext) -> None:
|
|
"""Register analysis, cancellation, and prompt queue routes."""
|
|
@prompt_server.routes.post(ctx.cancel_preparation_route_path)
|
|
async def cancel_remote_preparation(request: web.Request) -> web.Response:
|
|
"""Cancel one prompt while it is still preparing remote execution."""
|
|
payload = await request.json()
|
|
prompt_id = str(payload.get("prompt_id") or "").strip()
|
|
cancelled = _cancel_remote_preparation(prompt_server, prompt_id)
|
|
return web.json_response({"cancelled": cancelled, "prompt_id": prompt_id})
|
|
|
|
@prompt_server.routes.post(ctx.analysis_route_path)
|
|
async def modal_analyze_remote_nodes(request: web.Request) -> web.Response:
|
|
"""Analyze which workflow nodes should be marked remote for the current graph."""
|
|
logger.info("Received Modal remote-node analysis request.")
|
|
try:
|
|
request_started_at = time.perf_counter()
|
|
json_data = await request.json()
|
|
prompt = json_data.get("prompt")
|
|
if not isinstance(prompt, dict):
|
|
raise ValueError("Modal remote-node analysis requires a 'prompt' object.")
|
|
workflow = json_data.get("workflow")
|
|
seed_node_ids = json_data.get("seed_node_ids") or []
|
|
if not isinstance(seed_node_ids, list):
|
|
raise ValueError(
|
|
"Modal remote-node analysis requires 'seed_node_ids' to be a list."
|
|
)
|
|
analysis = analyze_remote_node_selection(
|
|
prompt=prompt,
|
|
workflow=workflow if isinstance(workflow, dict) else None,
|
|
seed_workflow_node_paths=[str(node_id) for node_id in seed_node_ids],
|
|
settings=ctx.settings,
|
|
)
|
|
logger.info(
|
|
"Modal remote-node analysis finished in %.3fs with %d requested nodes and %d additions.",
|
|
time.perf_counter() - request_started_at,
|
|
len(analysis.requested_workflow_node_paths),
|
|
len(analysis.added_workflow_node_paths),
|
|
)
|
|
return web.json_response(
|
|
{
|
|
"requested_node_ids": analysis.requested_node_ids,
|
|
"requested_workflow_node_paths": analysis.requested_workflow_node_paths,
|
|
"current_remote_node_ids": analysis.current_remote_node_ids,
|
|
"current_remote_workflow_node_paths": analysis.current_remote_workflow_node_paths,
|
|
"resolved_remote_node_ids": analysis.resolved_remote_node_ids,
|
|
"resolved_workflow_node_paths": analysis.resolved_workflow_node_paths,
|
|
"added_node_ids": analysis.added_node_ids,
|
|
"added_workflow_node_paths": analysis.added_workflow_node_paths,
|
|
"sandwiched_local_node_ids": analysis.sandwiched_local_node_ids,
|
|
"reasons": [
|
|
{
|
|
"node_id": reason.node_id,
|
|
"class_type": reason.class_type,
|
|
"required_by_node_id": reason.required_by_node_id,
|
|
"required_by_class_type": reason.required_by_class_type,
|
|
"output_index": reason.output_index,
|
|
"io_type": reason.io_type,
|
|
}
|
|
for reason in analysis.reasons
|
|
],
|
|
}
|
|
)
|
|
except (TypeError, ValueError) as exc:
|
|
logger.warning("Modal remote-node analysis request was invalid: %s", exc)
|
|
return web.json_response({"error": str(exc), "node_errors": []}, status=400)
|
|
|
|
@prompt_server.routes.post(ctx.settings.route_path)
|
|
async def modal_queue_prompt(request: web.Request) -> web.Response:
|
|
"""Handle prompt queue requests that include remote execution markers."""
|
|
return await _handle_modal_queue_prompt(request, prompt_server, ctx)
|