Spaces:
Configuration error
Configuration error
| """Parallel execution runner for executing BuildTasks concurrently in stages (PH2-022).""" | |
| from __future__ import annotations | |
| import logging | |
| import os | |
| import time | |
| from concurrent.futures import ThreadPoolExecutor, as_completed, Future | |
| from typing import Any, Dict, Generator, List, Optional | |
| from core.analysis_registry import AnalysisRegistry, AnalysisNode | |
| from core.incremental_build_planner import BuildTask | |
| from models.build_event import ( | |
| BuildEvent, | |
| TaskQueued, | |
| TaskStarted, | |
| TaskCompleted, | |
| TaskSkipped, | |
| TaskFailed, | |
| StageCompleted, | |
| CacheHitEvent, | |
| LoadEvent, | |
| CacheMissEvent, | |
| ProgressEvent, | |
| SaveEvent, | |
| BuildTimeEvent, | |
| ErrorEvent, | |
| ) | |
| logger = logging.getLogger(__name__) | |
| def execute_single_task( | |
| task: BuildTask, | |
| node: AnalysisNode, | |
| repo_name: str, | |
| repo_path: Optional[str], | |
| files: Optional[List[Dict[str, str]]], | |
| context: Any, | |
| force_rebuild: bool, | |
| build_id: str, | |
| service_resolver: Optional[Any] = None, | |
| snapshot_store: Optional[Any] = None, | |
| ) -> List[BuildEvent]: | |
| """Execute a single build task and capture all events generated.""" | |
| from core.logging_context import build_id_var, repository_var, analysis_var | |
| token_build = build_id_var.set(build_id) | |
| token_repo = repository_var.set(repo_name) | |
| token_analysis = analysis_var.set(node.name) | |
| events: List[BuildEvent] = [] | |
| try: | |
| if task.mode == "SKIP": | |
| events.append(TaskSkipped(repo_name, node.name)) | |
| events.append(CacheHitEvent(repo_name, node.name)) | |
| events.append(LoadEvent(repo_name, node.name)) | |
| events.append(BuildTimeEvent(repo_name, node.name, 0.0)) | |
| return events | |
| events.append(TaskStarted(repo_name, node.name)) | |
| events.append(CacheMissEvent(repo_name, node.name)) | |
| start_time = time.time() | |
| try: | |
| # Resolve service singleton via resolver or direct instantiation | |
| service_instance = None | |
| if service_resolver is not None: | |
| service_instance = service_resolver(node.service_class) | |
| elif node.service_class is not None: | |
| try: | |
| service_instance = node.service_class() | |
| except Exception: | |
| service_instance = None | |
| if service_instance is not None: | |
| has_new_api = hasattr(service_instance, "build_full") | |
| if node.name == "Symbol Index": | |
| if has_new_api: | |
| if task.mode == "FULL": | |
| service_instance.build_full( | |
| repo_name, repo_path=repo_path, files=files | |
| ) | |
| else: | |
| service_instance.build_partial( | |
| repo_name, | |
| task.changed_files, | |
| repo_path=repo_path, | |
| files=files, | |
| ) | |
| else: | |
| service_instance.build( | |
| repo_name, repo_path=repo_path, files=files | |
| ) | |
| elif node.name == "Dependency Graph": | |
| if has_new_api: | |
| if task.mode == "FULL": | |
| service_instance.build_full( | |
| repo_name, repo_path=repo_path, files=files | |
| ) | |
| else: | |
| service_instance.build_partial( | |
| repo_name, | |
| task.changed_files, | |
| repo_path=repo_path, | |
| files=files, | |
| ) | |
| else: | |
| service_instance.build( | |
| repo_name, | |
| repo_path=repo_path, | |
| files=files, | |
| force_rebuild=force_rebuild, | |
| ) | |
| elif node.name == "Call Graph": | |
| if has_new_api: | |
| if task.mode == "FULL": | |
| gen = service_instance.build_full( | |
| repo_name, context=context, files=files | |
| ) | |
| else: | |
| gen = service_instance.build_partial( | |
| repo_name, | |
| task.changed_files, | |
| context=context, | |
| files=files, | |
| ) | |
| for progress in gen: | |
| events.append( | |
| ProgressEvent( | |
| repo_name, node.name, progress.get("message", "") | |
| ) | |
| ) | |
| else: | |
| service_instance.build(repo_name, context=context) | |
| elif node.name == "Git History": | |
| if has_new_api: | |
| if task.mode == "FULL": | |
| gen = service_instance.build_full( | |
| repo_name, since_days=30, context=context | |
| ) | |
| else: | |
| gen = service_instance.build_partial( | |
| repo_name, | |
| task.changed_files, | |
| since_days=30, | |
| context=context, | |
| ) | |
| else: | |
| gen = service_instance.build(repo_name, since_days=30) | |
| for progress in gen: | |
| events.append( | |
| ProgressEvent( | |
| repo_name, node.name, progress.get("message", "") | |
| ) | |
| ) | |
| elif node.name == "API Surface": | |
| if has_new_api: | |
| if task.mode == "FULL": | |
| gen = service_instance.build_full( | |
| repo_name, context=context, files=files | |
| ) | |
| else: | |
| gen = service_instance.build_partial( | |
| repo_name, | |
| task.changed_files, | |
| context=context, | |
| files=files, | |
| ) | |
| for progress in gen: | |
| pass | |
| else: | |
| service_instance.build(repo_name, context=context) | |
| else: | |
| # Custom/mock service execution fallback | |
| if has_new_api: | |
| if task.mode == "FULL": | |
| service_instance.build_full( | |
| repo_name, repo_path=repo_path, files=files | |
| ) | |
| else: | |
| service_instance.build_partial( | |
| repo_name, | |
| task.changed_files, | |
| repo_path=repo_path, | |
| files=files, | |
| ) | |
| elif hasattr(service_instance, "build"): | |
| service_instance.build( | |
| repo_name, repo_path=repo_path, files=files | |
| ) | |
| else: | |
| # Placeholder for future analyses (service_instance is None) | |
| if snapshot_store is not None: | |
| store = snapshot_store | |
| else: | |
| from storage.snapshot_store import JsonSnapshotStore | |
| store = JsonSnapshotStore() | |
| key = node.outputs[0] if node.outputs else node.name.lower() | |
| store.save( | |
| repo_name, | |
| key, | |
| {"status": "placeholder", "_schema_version": node.schema_version}, | |
| ) | |
| events.append(SaveEvent(repo_name, node.name)) | |
| duration = (time.time() - start_time) * 1000 | |
| events.append(TaskCompleted(repo_name, node.name, duration)) | |
| events.append(BuildTimeEvent(repo_name, node.name, duration)) | |
| from core.metrics import metrics_registry | |
| metrics_registry.record_task_duration( | |
| repo_name, node.name, duration / 1000.0 | |
| ) | |
| except Exception as exc: | |
| duration = (time.time() - start_time) * 1000 | |
| events.append(TaskFailed(repo_name, node.name, str(exc))) | |
| events.append(ErrorEvent(repo_name, str(exc), node=node.name)) | |
| from core.metrics import metrics_registry | |
| metrics_registry.record_task_duration( | |
| repo_name, node.name, duration / 1000.0 | |
| ) | |
| raise | |
| finally: | |
| build_id_var.reset(token_build) | |
| repository_var.reset(token_repo) | |
| analysis_var.reset(token_analysis) | |
| return events | |
| class ParallelExecutionRunner: | |
| """Orchestrates stage-based parallel build execution using ThreadPoolExecutor.""" | |
| def __init__( | |
| self, | |
| repo_name: str, | |
| repo_path: Optional[str], | |
| files: Optional[List[Dict[str, str]]], | |
| context: Any, | |
| registry: AnalysisRegistry, | |
| max_workers: Optional[int] = None, | |
| force_rebuild: bool = False, | |
| build_id: Optional[str] = None, | |
| service_resolver: Optional[Any] = None, | |
| snapshot_store: Optional[Any] = None, | |
| ) -> None: | |
| import uuid | |
| self.repo_name = repo_name | |
| self.repo_path = repo_path | |
| self.files = files | |
| self.context = context | |
| self.registry = registry | |
| self.force_rebuild = force_rebuild | |
| self.build_id = build_id or str(uuid.uuid4()) | |
| self.service_resolver = service_resolver | |
| self.snapshot_store = snapshot_store | |
| # Default workers count: max(1, CPU cores - 1) | |
| if max_workers is None: | |
| max_workers = max(1, (os.cpu_count() or 4) - 1) | |
| self.max_workers = max_workers | |
| def run_stages( | |
| self, stages: List[List[BuildTask]] | |
| ) -> Generator[BuildEvent, None, None]: | |
| """Runs the planned stages sequentially, executing tasks within each stage in parallel.""" | |
| # Yield queued events for all tasks first | |
| for stage in stages: | |
| for task in stage: | |
| yield TaskQueued(self.repo_name, task.analysis) | |
| with ThreadPoolExecutor(max_workers=self.max_workers) as executor: | |
| for stage_idx, stage in enumerate(stages): | |
| tasks_in_stage = [task.analysis for task in stage] | |
| logger.info( | |
| "Running stage %d containing tasks: %s", | |
| stage_idx, | |
| tasks_in_stage, | |
| ) | |
| # Map task name to AnalysisNode | |
| node_map = {node.name: node for node in self.registry.get_build_order()} | |
| futures: Dict[Future[List[BuildEvent]], BuildTask] = {} | |
| for task in stage: | |
| node = node_map.get(task.analysis) | |
| if node is None: | |
| continue | |
| future = executor.submit( | |
| execute_single_task, | |
| task, | |
| node, | |
| self.repo_name, | |
| self.repo_path, | |
| self.files, | |
| self.context, | |
| self.force_rebuild, | |
| self.build_id, | |
| self.service_resolver, | |
| self.snapshot_store, | |
| ) | |
| futures[future] = task | |
| # Monitor futures as they complete | |
| failed = False | |
| exception_to_raise = None | |
| completed_task_events: Dict[str, List[BuildEvent]] = {} | |
| try: | |
| for future in as_completed(futures): | |
| task = futures[future] | |
| try: | |
| # future.result() propagates exceptions raised by workers | |
| task_events = future.result() | |
| completed_task_events[task.analysis] = task_events | |
| except Exception as exc: | |
| failed = True | |
| exception_to_raise = exc | |
| logger.error( | |
| "Task %s failed: %s. Cancelling remaining tasks in stage %d.", | |
| task.analysis, | |
| exc, | |
| stage_idx, | |
| ) | |
| # Cancel all other pending futures in this stage | |
| for pending_future in futures: | |
| if pending_future != future: | |
| pending_future.cancel() | |
| break | |
| except Exception as exc: | |
| failed = True | |
| exception_to_raise = exc | |
| if failed: | |
| if exception_to_raise: | |
| raise exception_to_raise | |
| else: | |
| raise RuntimeError( | |
| f"Stage {stage_idx} failed due to worker error." | |
| ) | |
| # Stage completed successfully: yield events alphabetically | |
| stage_tasks_sorted = sorted(stage, key=lambda t: t.analysis) | |
| for task in stage_tasks_sorted: | |
| task_events = completed_task_events.get(task.analysis, []) | |
| for event in task_events: | |
| yield event | |
| # Yield StageCompleted event | |
| yield StageCompleted(self.repo_name, stage_idx, tasks_in_stage) | |