pytrilogy 0.3.138__cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

trilogy/scripts/parallel_execution.py

@@ -0,0 +1,483 @@
+import threading
+from dataclasses import dataclass
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Callable, Protocol
+
+import networkx as nx
+
+from trilogy import Executor
+from trilogy.scripts.dependency import (
+    DependencyResolver,
+    DependencyStrategy,
+    ScriptNode,
+    create_script_nodes,
+)
+
+
+@dataclass
+class ExecutionResult:
+    """Result of executing a single script."""
+
+    node: ScriptNode
+    success: bool
+    error: Exception | None = None
+    duration: float = 0.0  # seconds
+
+
+@dataclass
+class ParallelExecutionSummary:
+    """Summary of a parallel execution run."""
+
+    total_scripts: int
+    successful: int
+    failed: int
+    total_duration: float
+    results: list[ExecutionResult]
+
+    @property
+    def all_succeeded(self) -> bool:
+        return self.failed == 0
+
+
+class ExecutionStrategy(Protocol):
+    """Protocol for execution traversal strategies."""
+
+    def execute(
+        self,
+        graph: nx.DiGraph,
+        resolver: DependencyResolver,
+        max_workers: int,
+        executor_factory: Callable[[ScriptNode], Any],
+        execution_fn: Callable[[Any, ScriptNode], None],
+        on_script_start: Callable[[ScriptNode], None] | None = None,
+        on_script_complete: Callable[[ExecutionResult], None] | None = None,
+    ) -> list[ExecutionResult]:
+        """
+        Execute scripts according to the strategy.
+
+        Args:
+            graph: The dependency graph (edges point from deps to dependents).
+            max_workers: Maximum parallel workers.
+            executor_factory: Factory to create executor for each script.
+            execution_fn: Function to execute a script.
+
+        Returns:
+            List of ExecutionResult for all scripts.
+        """
+        ...
+
+
+# Type aliases for cleaner signatures
+CompletedSet = set[ScriptNode]
+FailedSet = set[ScriptNode]
+InProgressSet = set[ScriptNode]
+ResultsList = list[ExecutionResult]
+RemainingDepsDict = dict[ScriptNode, int]
+ReadyList = list[ScriptNode]
+OnCompleteCallback = Callable[[ExecutionResult], None] | None
+
+
+def _propagate_failure(
+    failed_node: ScriptNode,
+    graph: nx.DiGraph,
+    completed: CompletedSet,
+    in_progress: InProgressSet,
+    results: ResultsList,
+    failed: FailedSet,
+    on_script_complete: OnCompleteCallback,
+) -> None:
+    """
+    Recursively mark all *unstarted* dependents of a failed node as failed and skipped.
+    """
+    for dependent in graph.successors(failed_node):
+        if dependent not in completed and dependent not in in_progress:
+            skip_result = ExecutionResult(
+                node=dependent,
+                success=False,
+                error=RuntimeError("Skipped due to failed dependency"),
+                duration=0.0,
+            )
+            results.append(skip_result)
+            completed.add(dependent)
+            failed.add(dependent)
+            if on_script_complete:
+                on_script_complete(skip_result)
+            _propagate_failure(
+                dependent,
+                graph,
+                completed,
+                in_progress,
+                results,
+                failed,
+                on_script_complete,
+            )
+
+
+def _get_next_ready(ready: ReadyList) -> ScriptNode | None:
+    """Get next ready node from the queue."""
+    if ready:
+        return ready.pop(0)
+    return None
+
+
+def _mark_node_complete(
+    node: ScriptNode,
+    success: bool,
+    graph: nx.DiGraph,
+    completed: CompletedSet,
+    failed: FailedSet,
+    in_progress: InProgressSet,
+    remaining_deps: RemainingDepsDict,
+    ready: ReadyList,
+    results: ResultsList,
+    on_script_complete: OnCompleteCallback,
+) -> None:
+    """
+    Mark a node as complete, update dependent counts, and add newly ready/skipped nodes.
+    """
+    in_progress.discard(node)
+    completed.add(node)
+    if not success:
+        failed.add(node)
+
+    # Update dependents
+    for dependent in graph.successors(node):
+        if dependent in completed or dependent in in_progress:
+            continue
+
+        if success:
+            remaining_deps[dependent] -= 1
+            if remaining_deps[dependent] == 0:
+                # Check if any dependency failed before running
+                deps = set(graph.predecessors(dependent))
+                if deps & failed:
+                    # Skip this node - dependency failed
+                    skip_result = ExecutionResult(
+                        node=dependent,
+                        success=False,
+                        error=RuntimeError("Skipped due to failed dependency"),
+                        duration=0.0,
+                    )
+                    results.append(skip_result)
+                    completed.add(dependent)
+                    failed.add(dependent)
+                    if on_script_complete:
+                        on_script_complete(skip_result)
+                    # Recursively mark dependents as failed
+                    _propagate_failure(
+                        dependent,
+                        graph,
+                        completed,
+                        in_progress,
+                        results,
+                        failed,
+                        on_script_complete,
+                    )
+                else:
+                    ready.append(dependent)
+        else:
+            # Current node failed - mark this dependent as skipped
+            if dependent not in failed:
+                skip_result = ExecutionResult(
+                    node=dependent,
+                    success=False,
+                    error=RuntimeError("Skipped due to failed dependency"),
+                    duration=0.0,
+                )
+                results.append(skip_result)
+                completed.add(dependent)
+                failed.add(dependent)
+                if on_script_complete:
+                    on_script_complete(skip_result)
+                # Recursively mark dependents as failed
+                _propagate_failure(
+                    dependent,
+                    graph,
+                    completed,
+                    in_progress,
+                    results,
+                    failed,
+                    on_script_complete,
+                )
+
+
+def _is_execution_done(completed: CompletedSet, total_count: int) -> bool:
+    """Check if all nodes have been processed."""
+    return len(completed) >= total_count
+
+
+def _execute_single(
+    node: ScriptNode,
+    executor_factory: Callable[[ScriptNode], Executor],
+    execution_fn: Callable[[Any, ScriptNode], None],
+) -> ExecutionResult:
+    """Execute a single script and return the result."""
+    start_time = datetime.now()
+    executor = None
+    try:
+        executor = executor_factory(node)
+        execution_fn(executor, node)
+
+        duration = (datetime.now() - start_time).total_seconds()
+        if executor:
+            executor.close()
+        return ExecutionResult(
+            node=node,
+            success=True,
+            error=None,
+            duration=duration,
+        )
+    except Exception as e:
+        duration = (datetime.now() - start_time).total_seconds()
+        if executor:
+            executor.close()  # Ensure executor is closed even on failure
+        return ExecutionResult(
+            node=node,
+            success=False,
+            error=e,
+            duration=duration,
+        )
+
+
+def _create_worker(
+    graph: nx.DiGraph,
+    lock: threading.Lock,
+    work_available: threading.Condition,
+    completed: CompletedSet,
+    failed: FailedSet,
+    in_progress: InProgressSet,
+    remaining_deps: RemainingDepsDict,
+    ready: ReadyList,
+    results: ResultsList,
+    total_count: int,
+    executor_factory: Callable[[ScriptNode], Any],
+    execution_fn: Callable[[Any, ScriptNode], None],
+    on_script_start: Callable[[ScriptNode], None] | None,
+    on_script_complete: OnCompleteCallback,
+) -> Callable[[], None]:
+    """
+    Create a worker function for thread execution to process the dependency graph.
+    """
+
+    def worker() -> None:
+        while True:
+            node = None
+
+            with work_available:
+                # Wait for work or global completion
+                while not ready and not _is_execution_done(completed, total_count):
+                    work_available.wait()
+
+                if _is_execution_done(completed, total_count):
+                    return
+
+                node = _get_next_ready(ready)
+                if node is None:
+                    # Should be impossible if total_count check is correct, but handles race condition safety
+                    continue
+
+                in_progress.add(node)
+
+            # Execute outside the lock
+            if node is not None:
+                if on_script_start:
+                    on_script_start(node)
+                result = _execute_single(node, executor_factory, execution_fn)
+
+                # Use the lock for state updates and notification
+                with lock:
+                    results.append(result)
+
+                    if on_script_complete:
+                        on_script_complete(result)
+
+                    _mark_node_complete(
+                        node,
+                        result.success,
+                        graph,
+                        completed,
+                        failed,
+                        in_progress,
+                        remaining_deps,
+                        ready,
+                        results,
+                        on_script_complete,
+                    )
+                    work_available.notify_all()  # Notify other workers of new ready/completed state
+
+    return worker
+
+
+class EagerBFSStrategy:
+    """
+    Eager Breadth-First Search (BFS) execution strategy.
+
+    Scripts execute as soon as all their dependencies complete, maximizing parallelism.
+    Uses a thread pool coordinated by locks and condition variables.
+    """
+
+    def execute(
+        self,
+        graph: nx.DiGraph,
+        resolver: DependencyResolver,
+        max_workers: int,
+        executor_factory: Callable[[ScriptNode], Any],
+        execution_fn: Callable[[Any, ScriptNode], None],
+        on_script_start: Callable[[ScriptNode], None] | None = None,
+        on_script_complete: Callable[[ExecutionResult], None] | None = None,
+    ) -> list[ExecutionResult]:
+        """Execute scripts eagerly as dependencies complete."""
+        if not graph.nodes():
+            return []
+
+        lock = threading.Lock()
+        work_available = threading.Condition(lock)
+
+        # Track state
+        completed: CompletedSet = set()
+        failed: FailedSet = set()
+        in_progress: InProgressSet = set()
+        results: ResultsList = []
+
+        # Calculate in-degrees (number of incomplete dependencies)
+        remaining_deps: RemainingDepsDict = {
+            node: graph.in_degree(node) for node in graph.nodes()
+        }
+
+        # Ready queue - nodes with all dependencies satisfied initially (in-degree 0)
+        ready: ReadyList = [node for node in graph.nodes() if remaining_deps[node] == 0]
+
+        total_count = len(graph.nodes())
+
+        # Create the worker function
+        worker = _create_worker(
+            graph=graph,
+            lock=lock,
+            work_available=work_available,
+            completed=completed,
+            failed=failed,
+            in_progress=in_progress,
+            remaining_deps=remaining_deps,
+            ready=ready,
+            results=results,
+            total_count=total_count,
+            executor_factory=executor_factory,
+            execution_fn=execution_fn,
+            on_script_start=on_script_start,
+            on_script_complete=on_script_complete,
+        )
+
+        # Start worker threads
+        workers = min(max_workers, total_count)
+        threads: list[threading.Thread] = []
+        for _ in range(workers):
+            t = threading.Thread(target=worker, daemon=True)
+            t.start()
+            threads.append(t)
+
+        # Wake up any waiting workers if we have initial work
+        with work_available:
+            work_available.notify_all()
+
+        # Wait for all threads to complete
+        for t in threads:
+            t.join()
+
+        return results
+
+
+class ParallelExecutor:
+    """
+    Executes scripts in parallel while respecting dependencies.
+
+    Uses an Eager BFS traversal by default, running scripts as soon as their
+    dependencies complete.
+    """
+
+    def __init__(
+        self,
+        max_workers: int = 5,
+        dependency_strategy: DependencyStrategy | None = None,
+        execution_strategy: ExecutionStrategy | None = None,
+    ):
+        """
+        Initialize the parallel executor.
+
+        Args:
+            max_workers: Maximum number of parallel workers.
+            dependency_strategy: Strategy for resolving dependencies.
+            execution_strategy: Strategy for traversing the graph during execution.
+        """
+        self.max_workers = max_workers
+        # Resolver finds dependencies and builds the graph
+        self.resolver = DependencyResolver(strategy=dependency_strategy)
+        # Execution strategy determines how the graph is traversed and executed
+        self.execution_strategy = execution_strategy or EagerBFSStrategy()
+
+    def execute(
+        self,
+        root: Path,
+        executor_factory: Callable[[ScriptNode], Any],
+        execution_fn: Callable[[Any, ScriptNode], None],
+        on_script_start: Callable[[ScriptNode], None] | None = None,
+        on_script_complete: Callable[[ExecutionResult], None] | None = None,
+    ) -> ParallelExecutionSummary:
+        """
+        Execute scripts in parallel respecting dependencies.
+
+        Args:
+            root: Root path (folder or single file) to find scripts.
+            executor_factory: Factory function to create an executor for a script.
+            execution_fn: Function that executes a script given (executor, node).
+
+        Returns:
+            ParallelExecutionSummary with all results.
+        """
+        start_time = datetime.now()
+
+        # Build dependency graph
+        if root.is_dir():
+            graph = self.resolver.build_folder_graph(root)
+            nodes = list(graph.nodes())
+        else:
+            nodes = create_script_nodes([root])
+            graph = self.resolver.build_graph(nodes)
+
+        # Total count of nodes for summary/completion check
+        total_scripts = len(nodes)
+
+        # Execute using the configured strategy
+        results = self.execution_strategy.execute(
+            graph=graph,
+            resolver=self.resolver,
+            max_workers=self.max_workers,
+            executor_factory=executor_factory,
+            execution_fn=execution_fn,
+            on_script_start=on_script_start,
+            on_script_complete=on_script_complete,
+        )
+
+        total_duration = (datetime.now() - start_time).total_seconds()
+        successful = sum(1 for r in results if r.success)
+
+        return ParallelExecutionSummary(
+            total_scripts=total_scripts,
+            successful=successful,
+            failed=total_scripts - successful,
+            total_duration=total_duration,
+            results=results,
+        )
+
+    def get_folder_execution_plan(self, folder: Path) -> nx.DiGraph:
+        """
+        Get the execution plan (dependency graph) for all scripts in a folder.
+        """
+        return self.resolver.build_folder_graph(folder)
+
+    def get_execution_plan(self, files: list[Path]) -> nx.DiGraph:
+        """
+        Get the execution plan (dependency graph) without executing.
+        """
+        nodes = create_script_nodes(files)
+        return self.resolver.build_graph(nodes)

trilogy/scripts/single_execution.py

@@ -0,0 +1,131 @@
+import traceback
+from datetime import datetime
+from typing import Any, Union
+
+from trilogy import Executor
+from trilogy.core.statements.execute import PROCESSED_STATEMENT_TYPES
+from trilogy.scripts.display import (
+    FETCH_LIMIT,
+    ResultSet,
+    create_progress_context,
+    print_error,
+    print_info,
+    print_results_table,
+    show_statement_result,
+    show_statement_type,
+)
+
+
+def get_statement_type(statement: PROCESSED_STATEMENT_TYPES) -> str:
+    """Get the type/class name of a statement."""
+    return type(statement).__name__
+
+
+def execute_single_statement(
+    exec: Executor,
+    query: PROCESSED_STATEMENT_TYPES,
+    idx: int,
+    total_queries: int,
+    use_progress=False,
+) -> tuple[bool, ResultSet | None, Any, Union[Exception, None]]:
+    """Execute a single statement and handle results/errors consistently."""
+    # Log the statement type before execution
+    statement_type = get_statement_type(query)
+    if not use_progress:  # Only show type when not using progress bar
+        show_statement_type(idx, total_queries, statement_type)
+
+    start_time = datetime.now()
+
+    try:
+        raw_results = exec.execute_statement(query)
+        results = (
+            ResultSet(
+                rows=raw_results.fetchmany(FETCH_LIMIT + 1), columns=raw_results.keys()
+            )
+            if raw_results
+            else None
+        )
+        duration = datetime.now() - start_time
+
+        if not use_progress:
+            show_statement_result(idx, total_queries, duration, bool(results))
+
+        return True, results, duration, None
+
+    except Exception as e:
+        duration = datetime.now() - start_time
+
+        if not use_progress:
+            show_statement_result(idx, total_queries, duration, False, str(e), type(e))
+
+        return False, None, duration, e
+
+
+def execute_queries_with_progress(
+    exec: Executor, queries: list[PROCESSED_STATEMENT_TYPES]
+) -> Exception | None:
+    """Execute queries with Rich progress bar. Returns True if all succeeded, False if any failed."""
+    progress = create_progress_context()
+    results_to_print = []
+    exception = None
+
+    with progress:
+        task = progress.add_task("Executing statements...", total=len(queries))
+
+        for idx, query in enumerate(queries):
+            statement_type = get_statement_type(query)
+            progress.update(
+                task, description=f"Statement {idx+1}/{len(queries)} ({statement_type})"
+            )
+
+            success, results, duration, error = execute_single_statement(
+                exec, query, idx, len(queries), use_progress=True
+            )
+
+            if not success:
+                exception = error
+
+            # Store results for printing after progress is done
+            results_to_print.append(
+                (idx, len(queries), duration, success, results, error)
+            )
+            progress.advance(task)
+            if exception:
+                break
+
+    # Print all results after progress bar is finished
+    for idx, total_queries, duration, success, results, error in results_to_print:
+        if error:
+            show_statement_result(
+                idx, total_queries, duration, False, str(error), type(error)
+            )
+            print_error(f"Full traceback:\n{traceback.format_exc()}")
+        else:
+            show_statement_result(idx, total_queries, duration, bool(results))
+            if results and not error:
+                print_results_table(results)
+
+    return exception
+
+
+def execute_queries_simple(
+    exec: Executor, queries: list[PROCESSED_STATEMENT_TYPES]
+) -> Exception | None:
+    """Execute queries with simple output. Returns True if all succeeded, False if any failed."""
+    exception = None
+
+    for idx, query in enumerate(queries):
+        if len(queries) > 1:
+            print_info(f"Executing statement {idx+1} of {len(queries)}...")
+
+        success, results, duration, error = execute_single_statement(
+            exec, query, idx, len(queries), use_progress=False
+        )
+
+        if not success:
+            exception = error
+
+        if results and not error:
+            print_results_table(results)
+
+    return exception