"""Parallel runtime engine for executing workflows with concurrent node execution.
This module provides a parallel execution engine for Kailash workflows,
specifically designed to run independent nodes concurrently for maximum performance.
"""
import asyncio
import contextvars
import logging
import time
from collections import deque
from datetime import UTC, datetime
from typing import Any
from kailash.nodes.base_async import AsyncNode
from kailash.runtime._time_limits import (
_TimeLimitClassifier,
_validate_limits,
arm_time_limits_async,
)
from kailash.runtime.cancellation import CancellationToken
from kailash.sdk_exceptions import (
HardTimeLimitExceeded,
RuntimeExecutionError,
SoftTimeLimitExceeded,
WorkflowCancelledError,
WorkflowExecutionError,
WorkflowValidationError,
)
from kailash.tracking import TaskManager, TaskStatus
from kailash.tracking.metrics_collector import MetricsCollector
from kailash.tracking.models import TaskMetrics
from kailash.workflow.graph import Workflow
logger = logging.getLogger(__name__)
[docs]
class ParallelRuntime:
"""Parallel execution engine for workflows.
This runtime provides true concurrent execution of independent nodes in a workflow,
allowing for maximum performance with both synchronous and asynchronous nodes.
Key features:
- Concurrent execution of independent nodes
- Dynamic scheduling based on dependency resolution
- Support for both sync and async nodes
- Configurable parallelism limits
- Detailed execution metrics and visualization
Usage:
runtime = ParallelRuntime(max_workers=8)
results, run_id = await runtime.execute(workflow, parameters={...})
"""
[docs]
def __init__(self, max_workers: int = 8, debug: bool = False):
"""Initialize the parallel runtime.
Args:
max_workers: Maximum number of concurrent node executions
debug: Whether to enable debug logging
"""
self.max_workers = max_workers
self.debug = debug
self.logger = logger
if debug:
self.logger.setLevel(logging.DEBUG)
else:
self.logger.setLevel(logging.INFO)
self.semaphore: asyncio.Semaphore | None = (
None # Will be initialized during execution
)
[docs]
async def execute(
self,
workflow: Workflow,
task_manager: TaskManager | None = None,
parameters: dict[str, dict[str, Any]] | None = None,
*,
soft_time_limit: float | None = None,
time_limit: float | None = None,
**kwargs: Any,
) -> tuple[dict[str, Any], str | None]:
"""Execute a workflow with parallel node execution.
Args:
workflow: Workflow to execute
task_manager: Optional task manager for tracking
parameters: Optional parameter overrides per node
soft_time_limit: Optional advisory deadline in seconds (#912
Shard 1 slot; enforcement lands Shard 2).
time_limit: Optional unconditional kill deadline in seconds.
**kwargs: Forward-compatibility kwargs for additive #912 Shard 1
contract.
Returns:
Tuple of (results dict, run_id)
Raises:
RuntimeExecutionError: If execution fails
WorkflowValidationError: If workflow is invalid
"""
# #912 Shard 1: validate typed time-limit kwargs at the entry point.
_validate_limits(soft_time_limit, time_limit)
if not workflow:
raise RuntimeExecutionError("No workflow provided")
run_id = None
start_time = time.time()
# #912 Shard 6: arm asyncio-task-based deadlines around the
# parallel execution path (mirrors the AsyncLocalRuntime
# wiring). When at least one limit is set, we layer a fresh
# cancellation token; the soft timer cancels it; the post-
# completion poll raises the typed deadline exception per
# Shard 2 invariant 5.
_has_time_limit = soft_time_limit is not None or time_limit is not None
cancellable = None
_attempt_token: CancellationToken | None = None
if _has_time_limit:
_attempt_token = CancellationToken()
cancellable = arm_time_limits_async(
_attempt_token,
soft_time_limit=soft_time_limit,
time_limit=time_limit,
)
try:
try:
# Validate workflow
workflow.validate(runtime_parameters=parameters)
# Initialize semaphore for concurrent execution control
self.semaphore = asyncio.Semaphore(self.max_workers)
# Initialize tracking
if task_manager:
try:
run_id = task_manager.create_run(
workflow_name=workflow.name,
metadata={
"parameters": parameters,
"debug": self.debug,
"runtime": "parallel",
"max_workers": self.max_workers,
},
)
except Exception as e:
self.logger.warning(f"Failed to create task run: {e}")
# Continue without tracking
# Execute workflow with parallel node execution
results = await self._execute_workflow_parallel(
workflow=workflow,
task_manager=task_manager,
run_id=run_id,
parameters=parameters or {},
)
# Mark run as completed
if task_manager and run_id:
try:
end_time = time.time()
execution_time = end_time - start_time
task_manager.update_run_status(run_id, "completed")
except Exception as e:
self.logger.warning(f"Failed to update run status: {e}")
# #912 Shard 6: post-completion poll for hard-deadline-
# fired-after-success / soft-fired-but-completed.
if cancellable is not None:
if cancellable.hard_deadline_reached:
raise HardTimeLimitExceeded(
f"workflow exceeded hard time limit "
f"(time_limit={cancellable.time_limit}s + "
f"grace_seconds={cancellable.grace_seconds}s)"
)
if (
_attempt_token is not None
and _attempt_token.is_cancelled
and cancellable.soft_time_limit is not None
):
raise SoftTimeLimitExceeded(
f"workflow exceeded soft time limit "
f"(soft_time_limit={cancellable.soft_time_limit}s)"
)
return results, run_id
except (SoftTimeLimitExceeded, HardTimeLimitExceeded):
# Typed deadline exceptions MUST propagate untouched
# above the broad except below.
if task_manager and run_id:
try:
task_manager.update_run_status(
run_id, "failed", error="Time limit exceeded"
)
except Exception:
pass
raise
except WorkflowCancelledError as cancel_exc:
# Classify token-cancellation into the typed deadline
# subclass when timers were armed.
if cancellable is not None:
classified = _TimeLimitClassifier(cancellable).classify(cancel_exc)
if classified is not cancel_exc:
raise classified from cancel_exc
raise
except WorkflowValidationError:
# Re-raise validation errors as-is
if task_manager and run_id:
try:
task_manager.update_run_status(
run_id, "failed", error="Validation failed"
)
except Exception:
pass
raise
except Exception as e:
# Mark run as failed
if task_manager and run_id:
try:
task_manager.update_run_status(run_id, "failed", error=str(e))
except Exception:
pass
# Wrap other errors in RuntimeExecutionError
raise RuntimeExecutionError(
f"Parallel workflow execution failed: {type(e).__name__}: {e}"
) from e
finally:
# Always release timer tasks even on the no-limit path.
if cancellable is not None:
cancellable.disarm()
async def _execute_workflow_parallel(
self,
workflow: Workflow,
task_manager: TaskManager | None,
run_id: str | None,
parameters: dict[str, dict[str, Any]],
) -> dict[str, Any]:
"""Execute the workflow nodes in parallel where possible.
This method uses a dynamic scheduling approach to run independent nodes
concurrently while respecting dependencies.
Args:
workflow: Workflow to execute
task_manager: Task manager for tracking
run_id: Run ID for tracking
parameters: Parameter overrides
Returns:
Dictionary of node results
Raises:
WorkflowExecutionError: If execution fails
"""
# Initialize result storage and tracking
results = {}
node_outputs = {}
node_tasks = {}
failed_nodes = set()
# Calculate initial dependencies for each node
dependencies = {
node: set(workflow.graph.predecessors(node))
for node in workflow.graph.nodes()
}
ready_nodes = deque([node for node, deps in dependencies.items() if not deps])
pending_nodes = set(workflow.graph.nodes()) - set(ready_nodes)
self.logger.info(
f"Starting parallel execution with {len(ready_nodes)} initially ready nodes"
)
# Process nodes until all are complete
while ready_nodes or pending_nodes or node_tasks:
# Schedule ready nodes up to max_workers limit
while ready_nodes and len(node_tasks) < self.max_workers:
node_id = ready_nodes.popleft()
# Skip if node already failed
if node_id in failed_nodes:
continue
# Create and start task for this node
task = asyncio.create_task(
self._execute_node(
workflow=workflow,
node_id=node_id,
node_outputs=node_outputs,
parameters=parameters.get(node_id, {}),
task_manager=task_manager,
run_id=run_id,
)
)
node_tasks[node_id] = task
self.logger.debug(f"Scheduled node {node_id} for execution")
# Wait for any node to complete if we have active tasks
if node_tasks:
# Wait for the first task to complete
done, _ = await asyncio.wait(
node_tasks.values(), return_when=asyncio.FIRST_COMPLETED
)
# Process completed nodes
for task in done:
# Find the node_id for this task
completed_node_id = next(
node_id
for node_id, node_task in node_tasks.items()
if node_task == task
)
# Remove from active tasks
node_tasks.pop(completed_node_id)
try:
# Get result and add to outputs
node_result, success = task.result()
results[completed_node_id] = node_result
if success:
node_outputs[completed_node_id] = node_result
self.logger.info(
f"Node {completed_node_id} completed successfully"
)
# Update dependent nodes
for dependent in workflow.graph.successors(
completed_node_id
):
if dependent in pending_nodes:
dependencies[dependent].remove(completed_node_id)
# If all dependencies are satisfied, mark as ready
if not dependencies[dependent]:
ready_nodes.append(dependent)
pending_nodes.remove(dependent)
self.logger.debug(
f"Node {dependent} is now ready"
)
else:
# Node failed, mark it and check if we should continue
failed_nodes.add(completed_node_id)
self.logger.error(f"Node {completed_node_id} failed")
# Determine if we should stop execution
if self._should_stop_on_error(workflow, completed_node_id):
error_msg = f"Node '{completed_node_id}' failed"
raise WorkflowExecutionError(error_msg)
# Update dependent nodes to also mark as failed
self._mark_dependent_nodes_as_failed(
workflow,
completed_node_id,
failed_nodes,
pending_nodes,
ready_nodes,
)
except Exception as e:
# Handle unexpected task exceptions
failed_nodes.add(completed_node_id)
self.logger.error(
f"Unexpected error in node {completed_node_id}: {e}"
)
# Determine if we should stop execution
if self._should_stop_on_error(workflow, completed_node_id):
error_msg = f"Node '{completed_node_id}' failed with unexpected error: {e}"
raise WorkflowExecutionError(error_msg) from e
# Mark dependents as failed
self._mark_dependent_nodes_as_failed(
workflow,
completed_node_id,
failed_nodes,
pending_nodes,
ready_nodes,
)
else:
# No active tasks but we still have pending nodes - this indicates a deadlock
if pending_nodes:
remaining = list(pending_nodes)
raise WorkflowExecutionError(
f"Deadlock detected. Nodes waiting for dependencies: {remaining}"
)
# No tasks and no pending nodes means we're done
break
self.logger.info(
f"Parallel execution complete. Succeeded: {len(results) - len(failed_nodes)}, Failed: {len(failed_nodes)}"
)
return results
async def _execute_node(
self,
workflow: Workflow,
node_id: str,
node_outputs: dict[str, dict[str, Any]],
parameters: dict[str, Any],
task_manager: TaskManager | None,
run_id: str | None,
) -> tuple[dict[str, Any], bool]:
"""Execute a single node asynchronously.
Args:
workflow: The workflow being executed
node_id: ID of the node to execute
node_outputs: Dictionary of outputs from previously executed nodes
parameters: Parameter overrides for this node
task_manager: Task manager for tracking
run_id: Run ID for tracking
Returns:
Tuple of (node_result, success)
Note:
This method never raises exceptions - it returns success=False instead
to allow the caller to handle failures appropriately.
"""
# Get node instance
node_instance = workflow._node_instances.get(node_id)
if not node_instance:
self.logger.error(f"Node instance '{node_id}' not found in workflow")
return {"error": "Node instance not found"}, False
# Start task tracking
task = None
try:
if task_manager and run_id:
task = task_manager.create_task(
run_id=run_id,
node_id=node_id,
node_type=node_instance.__class__.__name__,
started_at=datetime.now(UTC),
)
except Exception as e:
self.logger.warning(f"Failed to create task for node '{node_id}': {e}")
try:
# Limit concurrent execution
assert self.semaphore is not None, "Semaphore not initialized"
async with self.semaphore:
# Update task status
if task:
task.update_status(TaskStatus.RUNNING)
# Prepare inputs
inputs = self._prepare_node_inputs(
workflow=workflow,
node_id=node_id,
node_instance=node_instance,
node_outputs=node_outputs,
parameters=parameters,
)
if self.debug:
self.logger.debug(f"Node {node_id} inputs: {inputs}")
# Execute node with metrics collection
collector = MetricsCollector()
if isinstance(node_instance, AsyncNode):
# Use async execution for AsyncNode
outputs, performance_metrics = await collector.collect_async(
node_instance.execute_async(**inputs), node_id=node_id
)
else:
# Use sync execution in an executor for regular Node
loop = asyncio.get_running_loop()
async def execute_with_metrics():
with collector.collect(node_id=node_id) as context:
# Propagate the caller's contextvars across the
# thread-pool boundary so a ContextVar set before
# execution is visible inside node.run() (#1200).
ctx = contextvars.copy_context()
result = await loop.run_in_executor(
None,
lambda: ctx.run(
lambda: node_instance.execute(**inputs)
),
)
return result, context.result()
outputs, performance_metrics = await execute_with_metrics()
# Update task status with enhanced metrics
if task:
task.update_status(
TaskStatus.COMPLETED,
result=outputs,
ended_at=datetime.now(UTC),
metadata={"execution_time": performance_metrics.duration},
)
# Convert and save performance metrics
if task_manager:
task_metrics_data = performance_metrics.to_task_metrics()
task_metrics = TaskMetrics(**task_metrics_data)
task_manager.update_task_metrics(task.task_id, task_metrics)
self.logger.info(
f"Node {node_id} completed successfully in {performance_metrics.duration:.3f}s"
)
return outputs, True
except Exception as e:
self.logger.error(f"Node {node_id} failed: {e}", exc_info=self.debug)
# Update task status
if task:
task.update_status(
TaskStatus.FAILED, error=str(e), ended_at=datetime.now(UTC)
)
# Return error result
error_result = {
"error": str(e),
"error_type": type(e).__name__,
"failed": True,
}
return error_result, False
def _prepare_node_inputs(
self,
workflow: Workflow,
node_id: str,
node_instance: Any,
node_outputs: dict[str, dict[str, Any]],
parameters: dict[str, Any],
) -> dict[str, Any]:
"""Prepare inputs for a node execution.
Args:
workflow: The workflow being executed
node_id: Current node ID
node_instance: Current node instance
node_outputs: Outputs from previously executed nodes
parameters: Parameter overrides
Returns:
Dictionary of inputs for the node
Raises:
WorkflowExecutionError: If input preparation fails
"""
inputs = {}
# Start with node configuration
inputs.update(node_instance.config)
# Add connected inputs from other nodes
for edge in workflow.graph.in_edges(node_id, data=True):
source_node_id = edge[0]
mapping = edge[2].get("mapping", {})
if source_node_id in node_outputs:
source_outputs = node_outputs[source_node_id]
# Check if the source node failed
if isinstance(source_outputs, dict) and source_outputs.get("failed"):
raise WorkflowExecutionError(
f"Cannot use outputs from failed node '{source_node_id}'"
)
for source_key, target_key in mapping.items():
if source_key in source_outputs:
inputs[target_key] = source_outputs[source_key]
else:
self.logger.warning(
f"Source output '{source_key}' not found in node '{source_node_id}'. "
f"Available outputs: {list(source_outputs.keys())}"
)
# Apply parameter overrides
inputs.update(parameters)
return inputs
def _should_stop_on_error(self, workflow: Workflow, node_id: str) -> bool:
"""Determine if execution should stop when a node fails.
Args:
workflow: The workflow being executed
node_id: Failed node ID
Returns:
Whether to stop execution
"""
# Check if any downstream nodes depend on this node
has_dependents = workflow.graph.out_degree(node_id) > 0
# For now, stop if the failed node has dependents
# Future: implement configurable error handling policies
return has_dependents
def _mark_dependent_nodes_as_failed(
self,
workflow: Workflow,
failed_node: str,
failed_nodes: set[str],
pending_nodes: set[str],
ready_nodes: deque[str],
) -> None:
"""Mark all dependent nodes as failed.
Args:
workflow: The workflow being executed
failed_node: The node that failed
failed_nodes: Set to track failed nodes
pending_nodes: Set of nodes waiting for dependencies
ready_nodes: Queue of nodes ready to execute
"""
# Get all descendants of the failed node
descendants = workflow.graph.descendants(failed_node)
# Mark all descendants as failed
for node in descendants:
failed_nodes.add(node)
# Remove from pending or ready as appropriate
if node in pending_nodes:
pending_nodes.remove(node)
# Need to handle as list comprehension since deque doesn't support
# efficient removal of arbitrary elements
if node in ready_nodes:
ready_nodes_list = list(ready_nodes)
ready_nodes_list.remove(node)
ready_nodes.clear()
ready_nodes.extend(ready_nodes_list)
self.logger.debug(
f"Marked {len(descendants)} dependent nodes as failed due to failure of node {failed_node}"
)