diff --git a/docs/execution-graph.md b/docs/execution-graph.md index 0adb9fdd..a3c5d399 100644 --- a/docs/execution-graph.md +++ b/docs/execution-graph.md @@ -575,6 +575,7 @@ metis_engine: - index - memory private_policy: + max_concurrency: 50 inputs: review: simple_llm_review finding_dedup: @@ -592,7 +593,13 @@ metis_engine: A handler receives validated inputs and a small context containing its granted capabilities together with the repository lookup contract, CodeGraph -materialize/load API, model runner, runtime limits, and callbacks. A node +materialize/load API, model runner, runtime limits, shared job scheduler, and +callbacks. `max_concurrency` limits that node's active jobs without reserving +threads; omitted values use the global `metis_engine.max_workers` limit. All +nodes on one engine share that worker pool, so their combined active work never +exceeds the global limit; larger node values are capped by it. A node submits +bounded work through `invocation.context.jobs.run(...)`; it does not construct +its own executor. A node accesses only the names in `invocation.context.capabilities`. A node declares `request: ReviewCommand` when it needs the review-stage request; the runner supplies that typed stage input without extra YAML. The context does not expose diff --git a/src/metis/configuration.py b/src/metis/configuration.py index 3dd62229..cd643a06 100644 --- a/src/metis/configuration.py +++ b/src/metis/configuration.py @@ -71,7 +71,12 @@ def load_runtime_config(config_path=None, enable_psql=False): engine_cfg = cfg.get("metis_engine", {}) runtime["max_token_length"] = engine_cfg.get("max_token_length", 100000) - runtime["max_workers"] = engine_cfg.get("max_workers", 8) + runtime["max_workers"] = _required_positive_int( + engine_cfg, + "max_workers", + section="metis_engine", + default=8, + ) runtime["embed_dim"] = engine_cfg.get("embed_dim", 1536) runtime["doc_chunk_size"] = engine_cfg.get("doc_chunk_size", 1024) runtime["doc_chunk_overlap"] = engine_cfg.get("doc_chunk_overlap", 200) @@ -195,8 +200,14 @@ def _positive_int(value: object, *, fallback: int) -> int: return parsed -def _required_positive_int(values: dict[str, Any], key: str, *, section: str) -> int: - value = values.get(key) +def _required_positive_int( + values: dict[str, Any], + key: str, + *, + section: str, + default: int | None = None, +) -> int: + value = values.get(key, default) if isinstance(value, bool) or not isinstance(value, int) or value <= 0: raise ValueError(f"{section}.{key} must be a positive integer") return value diff --git a/src/metis/engine/capabilities/navigation.py b/src/metis/engine/capabilities/navigation.py index 99831c84..492a7087 100644 --- a/src/metis/engine/capabilities/navigation.py +++ b/src/metis/engine/capabilities/navigation.py @@ -8,6 +8,7 @@ import re import shutil import subprocess +from threading import Lock from typing import Sequence from pydantic import BaseModel @@ -43,6 +44,9 @@ def __init__( self.max_chars = max_chars self._has_grep = shutil.which("grep") is not None self._has_find = shutil.which("find") is not None + # Serialize subprocess tools; add a measured FD-aware pool if tool + # latency becomes material. + self._subprocess_lock = Lock() def _resolve_path(self, raw_path: str) -> Path: return resolve_path_within_root(self.codebase_path, raw_path) @@ -53,13 +57,14 @@ def _run( *, ok_returncodes: tuple[int, ...] = (0,), ) -> str: - proc = subprocess.run( - list(argv), - cwd=str(self.codebase_path), - capture_output=True, - text=True, - timeout=self.timeout_seconds, - ) + with self._subprocess_lock: + proc = subprocess.run( + list(argv), + cwd=str(self.codebase_path), + capture_output=True, + text=True, + timeout=self.timeout_seconds, + ) stdout = (proc.stdout or "").strip() stderr = (proc.stderr or "").strip() if proc.returncode not in ok_returncodes: diff --git a/src/metis/engine/concurrency.py b/src/metis/engine/concurrency.py index 51efe6b0..2bd72bd1 100644 --- a/src/metis/engine/concurrency.py +++ b/src/metis/engine/concurrency.py @@ -6,29 +6,121 @@ import threading from collections.abc import Callable from collections.abc import Sequence +from concurrent.futures import Executor from concurrent.futures import FIRST_COMPLETED from concurrent.futures import Future from concurrent.futures import ThreadPoolExecutor from concurrent.futures import wait -from functools import partial +from typing import TYPE_CHECKING from metis import runlog from metis.usage import submit_with_current_context +if TYPE_CHECKING: + from metis.engine.execution.contracts import NodeJobs + + +class JobScheduler: + def __init__(self, max_workers: int) -> None: + self._max_workers = max_workers + self._executor = ThreadPoolExecutor( + max_workers=max_workers, + thread_name_prefix="metis", + ) + + def limit(self, max_concurrency: int) -> "NodeJobs": + return _NodeJobs( + self._executor, + self._max_workers, + ).limit(max_concurrency) + + def close(self) -> None: + self._executor.shutdown() + + +class _NodeJobs: + def __init__(self, executor: Executor, max_concurrency: int) -> None: + self._executor = executor + self._max_concurrency = max_concurrency + + def limit(self, max_concurrency: int) -> "_NodeJobs": + if max_concurrency < 1: + raise ValueError("max_concurrency must be positive") + return _NodeJobs( + self._executor, + min(max_concurrency, self._max_concurrency), + ) + + def run[JobT, ResultT]( + self, + jobs: Sequence[JobT], + worker: Callable[[JobT], ResultT], + *, + label: str | None, + result_key: Callable[[JobT], object], + on_complete: Callable[[JobT, int, int], None] | None = None, + ) -> list[ResultT]: + if not jobs: + return [] + + total = len(jobs) + worker_count = min(self._max_concurrency, total) + results: list[ResultT] = [] + + def invoke_worker(job: JobT) -> ResultT: + if label is None: + return worker(job) + key = result_key(job) + with runlog.span( + "task", + label, + {"kind": "concurrent_job", "key": key}, + ) as task_span: + runlog.bump("tasks") + result = worker(job) + task_span.end(attributes={"result": result}) + return result + + def collect( + job: JobT, + completed: int, + result_func: Callable[[], ResultT], + ) -> None: + results.append(result_func()) + if on_complete: + on_complete(job, completed, total) -def coerce_worker_count( - max_workers: int | str | None, - *, - default: int = 1, -) -> int: - value = default if max_workers is None or max_workers == "" else max_workers - return max(1, int(value)) - + job_iterator = iter(jobs) + futures: dict[Future[ResultT], JobT] = {} + completed = 0 + try: + for job in job_iterator: + futures[ + submit_with_current_context(self._executor, invoke_worker, job) + ] = job + if len(futures) >= worker_count: + break -def bounded_worker_count(max_workers: int | str | None, item_count: int) -> int: - if item_count <= 1: - return 1 - return min(coerce_worker_count(max_workers), item_count) + while futures: + completed_futures, _pending = wait( + futures, + return_when=FIRST_COMPLETED, + ) + for future in completed_futures: + completed += 1 + collect(futures.pop(future), completed, future.result) + for job in job_iterator: + futures[ + submit_with_current_context(self._executor, invoke_worker, job) + ] = job + if len(futures) >= worker_count: + break + except BaseException: + for future in futures: + future.cancel() + wait(futures) + raise + return results def serialized_progress_callback(callback): @@ -45,70 +137,3 @@ def _serialized(event): _serialized._metis_serialized_progress_callback = True return _serialized - - -def run_jobs[JobT, ResultT]( - jobs: Sequence[JobT], - worker: Callable[[JobT], ResultT], - *, - max_workers: int | str | None, - label: str, - result_key: Callable[[JobT], object], - on_complete: Callable[[JobT, int, int], None] | None = None, -) -> list[ResultT]: - if not jobs: - return [] - - total = len(jobs) - worker_count = bounded_worker_count(max_workers, total) - results: list[ResultT] = [] - - def invoke_worker(job: JobT) -> ResultT: - key = result_key(job) - with runlog.span( - "task", - label, - {"kind": "concurrent_job", "key": key}, - ) as task_span: - runlog.bump("tasks") - result = worker(job) - task_span.end(attributes={"result": result}) - return result - - def collect( - job: JobT, - completed: int, - result_func: Callable[[], ResultT], - ) -> None: - results.append(result_func()) - if on_complete: - on_complete(job, completed, total) - - if worker_count == 1: - for completed, job in enumerate(jobs, start=1): - collect(job, completed, partial(invoke_worker, job)) - return results - - with ThreadPoolExecutor(max_workers=worker_count) as executor: - job_iterator = iter(jobs) - futures: dict[Future[ResultT], JobT] = {} - completed = 0 - - for job in job_iterator: - futures[submit_with_current_context(executor, invoke_worker, job)] = job - if len(futures) >= worker_count: - break - - while futures: - completed_futures, _pending = wait( - futures, - return_when=FIRST_COMPLETED, - ) - for future in completed_futures: - completed += 1 - collect(futures.pop(future), completed, future.result) - for job in job_iterator: - futures[submit_with_current_context(executor, invoke_worker, job)] = job - if len(futures) >= worker_count: - break - return results diff --git a/src/metis/engine/core.py b/src/metis/engine/core.py index 4b2236c3..f7afe0e2 100644 --- a/src/metis/engine/core.py +++ b/src/metis/engine/core.py @@ -12,6 +12,7 @@ from metis import runlog from metis.chat_model_options import merge_chat_model_kwargs +from metis.configuration import _required_positive_int from metis.configuration import load_execution_config from metis.configuration import load_plugin_config from metis.plugins.c_family.codegraph import CFamilyCodeGraphProvider @@ -75,7 +76,11 @@ def __init__( if missing: raise ValueError(f"Missing required config: {', '.join(missing)}") - max_workers = cast(int, kwargs["max_workers"]) + max_workers = _required_positive_int( + kwargs, + "max_workers", + section="MetisEngine", + ) max_token_length = cast(int, kwargs["max_token_length"]) llama_query_model = cast(str, kwargs["llama_query_model"]) similarity_top_k = cast(int, kwargs["similarity_top_k"]) @@ -105,7 +110,6 @@ def __init__( kwargs.get("reachability_config") or {} ) reachability_settings = reachability_config.as_review_settings() - reachability_settings["max_workers"] = max_workers reasoning_effort = chat_model_kwargs.get("reasoning_effort") if reasoning_effort is not None: reachability_settings["reasoning_effort"] = reasoning_effort @@ -394,6 +398,7 @@ def execute_triage( return _execution_value(triage) def close(self): + self.execution.close() if self._triage_classifier is not None: self._triage_classifier.close() self.capabilities.close() diff --git a/src/metis/engine/execution/contracts.py b/src/metis/engine/execution/contracts.py index 4b3748d9..9944925f 100644 --- a/src/metis/engine/execution/contracts.py +++ b/src/metis/engine/execution/contracts.py @@ -5,6 +5,7 @@ from collections.abc import Callable from collections.abc import Mapping +from collections.abc import Sequence from dataclasses import dataclass from dataclasses import field from enum import Enum @@ -65,6 +66,20 @@ def materialize( def load(self, reference: CodeGraphReference) -> CodeGraph: ... +class NodeJobs(Protocol): + def limit(self, max_concurrency: int) -> "NodeJobs": ... + + def run[JobT, ResultT]( + self, + jobs: Sequence[JobT], + worker: Callable[[JobT], ResultT], + *, + label: str | None, + result_key: Callable[[JobT], object], + on_complete: Callable[[JobT, int, int], None] | None = None, + ) -> list[ResultT]: ... + + class EmptyNodeConfiguration(BaseModel): model_config = ConfigDict(extra="forbid", frozen=True) @@ -116,6 +131,7 @@ class NodeRuntime: chat_model_kwargs: Mapping[str, object] model_tool_max_rounds: int = 0 token_counter: Callable[[str], int] = count_tokens + jobs: NodeJobs | None = None def __post_init__(self) -> None: object.__setattr__( @@ -144,6 +160,12 @@ def __post_init__(self) -> None: MappingProxyType(dict(self.capabilities)), ) + @property + def jobs(self) -> NodeJobs: + if self.runtime.jobs is None: + raise RuntimeError("Node job scheduler is unavailable") + return self.runtime.jobs + @dataclass(frozen=True, slots=True) class NodeInvocation: diff --git a/src/metis/engine/execution/graph.py b/src/metis/engine/execution/graph.py index 3fbbc0ba..9b77ab7c 100644 --- a/src/metis/engine/execution/graph.py +++ b/src/metis/engine/execution/graph.py @@ -20,6 +20,7 @@ class ConfiguredNode(BaseModel): inputs: dict[str, InputBinding] = Field(default_factory=dict) depends_on: tuple[str, ...] = () + max_concurrency: int | None = Field(default=None, strict=True, ge=1) capabilities: tuple[str, ...] = () formats: tuple[ResultFormat, ...] | None = None filename: str | None = Field(default=None, min_length=1) diff --git a/src/metis/engine/execution/runner.py b/src/metis/engine/execution/runner.py index 0c660d6c..e3140f38 100644 --- a/src/metis/engine/execution/runner.py +++ b/src/metis/engine/execution/runner.py @@ -88,6 +88,10 @@ def run_stage( }, ) continue + max_concurrency = min( + node.definition.max_concurrency or context.runtime.max_workers, + context.runtime.max_workers, + ) with runlog.span( "node", f"{stage_name}.{node.name}", @@ -99,6 +103,7 @@ def run_stage( "capabilities": sorted(node.capabilities), "formats": node.definition.formats, "filename": node.definition.filename, + "max_concurrency": max_concurrency, }, ) as node_span: _progress( @@ -119,6 +124,11 @@ def run_stage( node_context = replace( context, capabilities=node.capabilities, + runtime=replace( + context.runtime, + max_workers=max_concurrency, + jobs=context.jobs.limit(max_concurrency), + ), ) result = node.registration.execute( NodeInvocation( diff --git a/src/metis/engine/nodes/builtins.py b/src/metis/engine/nodes/builtins.py index a190a33f..2c0cab04 100644 --- a/src/metis/engine/nodes/builtins.py +++ b/src/metis/engine/nodes/builtins.py @@ -112,7 +112,6 @@ def build_builtin_execution( triage_service = ( TriageService( - max_workers=engine_config.max_workers, triage_checkpoint_every=triage_checkpoint_every, ) if configuration.stages.triage is not None diff --git a/src/metis/engine/nodes/finding_dedup/core.py b/src/metis/engine/nodes/finding_dedup/core.py index 3edadd7e..322cb1b6 100644 --- a/src/metis/engine/nodes/finding_dedup/core.py +++ b/src/metis/engine/nodes/finding_dedup/core.py @@ -12,7 +12,7 @@ from pydantic import ConfigDict from pydantic import Field -from metis.engine.concurrency import run_jobs +from metis.engine.execution.contracts import NodeJobs from metis.engine.stages.review.models import ReviewFinding from metis.engine.stages.review.models import ReviewGroup from metis.engine.stages.review.models import StandardReviewResult @@ -48,7 +48,7 @@ def consolidate_result( | None, packet_fits: PacketFits, progress_callback: Callable[[dict[str, object]], None] | None, - max_workers: int, + jobs: NodeJobs, ) -> StandardReviewResult: original_candidates = [ ReviewCandidate(group, finding) @@ -61,7 +61,7 @@ def consolidate_result( adjudicator, packet_fits=packet_fits, progress_callback=progress_callback, - max_workers=max_workers, + jobs=jobs, ) kept = adjudicated if adjudicated is not None else candidates if len(kept) == len(original_candidates): @@ -100,7 +100,7 @@ def _apply_final_adjudication( *, packet_fits, progress_callback=None, - max_workers=1, + jobs: NodeJobs, ): if not callable(adjudicator) or not candidates: return None @@ -127,7 +127,7 @@ def _apply_final_adjudication( original_limit, progress_callback=progress_callback, phase="candidate", - max_workers=max_workers, + jobs=jobs, ): saw_valid_decision = True @@ -152,7 +152,7 @@ def _apply_final_adjudication( original_limit, progress_callback=progress_callback, phase="representative", - max_workers=max_workers, + jobs=jobs, ): saw_valid_decision = True @@ -174,7 +174,7 @@ def _run_adjudication_batches( *, progress_callback=None, phase="candidate", - max_workers=1, + jobs: NodeJobs, ): if not batches: return False @@ -182,10 +182,9 @@ def _run_adjudication_batches( total = len(batches) if progress_callback: progress_callback({"phase": phase, "completed": 0, "total": total}) - batch_results = run_jobs( + batch_results = jobs.run( list(enumerate(batches)), lambda item: (item[0], adjudicator(item[1])), - max_workers=max_workers, label="Final review deduplication", result_key=lambda item: f"{phase}:{item[0]}", on_complete=lambda _job, completed, total: ( diff --git a/src/metis/engine/nodes/finding_dedup/registration.py b/src/metis/engine/nodes/finding_dedup/registration.py index 1be24027..8a5100ab 100644 --- a/src/metis/engine/nodes/finding_dedup/registration.py +++ b/src/metis/engine/nodes/finding_dedup/registration.py @@ -110,7 +110,7 @@ def packet_fits(candidates: list[dict[str, object]]) -> bool: adjudicator=adjudicate, packet_fits=packet_fits, progress_callback=report_progress, - max_workers=invocation.context.runtime.max_workers, + jobs=invocation.context.jobs, ) if progress is not None: kept_count = sum(len(group.reviews) for group in deduped.reviews) diff --git a/src/metis/engine/nodes/reachability/incremental_review.py b/src/metis/engine/nodes/reachability/incremental_review.py index b5a4df68..233b48c4 100644 --- a/src/metis/engine/nodes/reachability/incremental_review.py +++ b/src/metis/engine/nodes/reachability/incremental_review.py @@ -26,7 +26,7 @@ from metis.engine.codegraph import CodeGraph from metis.engine.codegraph import FunctionNode -from metis.engine.concurrency import run_jobs +from metis.engine.execution.contracts import NodeJobs from metis.engine.llm_runner import JsonPromptRequest from metis.engine.llm_runner import JsonPromptRunner from metis.engine.llm_runner import rendered_prompt_token_count @@ -200,6 +200,7 @@ def review( self, graph: CodeGraph, *, + jobs: NodeJobs, options: ReachabilityReviewOptions, memory_service: MemoryService | None = None, evidence_graph: CodeGraph | None = None, @@ -232,6 +233,7 @@ def review( graph, node_ids, packets, + jobs=jobs, options=options, memory_service=memory_service, evidence_graph=discovery_graph, @@ -326,6 +328,7 @@ def _review_frontier_packets( node_ids: tuple[str, ...], initial_packets: list[_FrontierPacket], *, + jobs: NodeJobs, options: ReachabilityReviewOptions, memory_service: MemoryService | None, evidence_graph: CodeGraph | None = None, @@ -340,10 +343,9 @@ def _review_frontier_packets( while pending: unique_scheduled, representative_by_index = _unique_packet_schedule(pending) - responses = run_jobs( + responses = jobs.run( unique_scheduled, self._review_packet, - max_workers=options.max_workers, label="Reachability security review", result_key=lambda packet: packet.index, ) diff --git a/src/metis/engine/nodes/reachability/options.py b/src/metis/engine/nodes/reachability/options.py index cd02e61e..6fb55a0e 100644 --- a/src/metis/engine/nodes/reachability/options.py +++ b/src/metis/engine/nodes/reachability/options.py @@ -26,7 +26,6 @@ def as_review_settings(self) -> dict[str, Any]: @dataclass(frozen=True, slots=True) class ReachabilityReviewOptions: confirmation_model: str | None = None - max_workers: int | str | None = None max_path_length: int = DEFAULT_REACHABILITY_MAX_PATH_LENGTH progress_callback: Any = None reasoning_effort: str | None = None diff --git a/src/metis/engine/nodes/reachability/review.py b/src/metis/engine/nodes/reachability/review.py index a6d95657..f1b58b2e 100644 --- a/src/metis/engine/nodes/reachability/review.py +++ b/src/metis/engine/nodes/reachability/review.py @@ -12,6 +12,7 @@ from metis.engine.execution.contracts import CapabilityRequirement from metis.engine.execution.contracts import EmptyNodeConfiguration from metis.engine.execution.contracts import NodeInvocation +from metis.engine.execution.contracts import NodeJobs from metis.engine.execution.contracts import NodeRegistration from metis.engine.execution.contracts import NodeResult from metis.engine.nodes.reachability.options import ReachabilityReviewOptions @@ -45,6 +46,7 @@ def execute(invocation: NodeInvocation) -> NodeResult: reference = cast(CodeGraphReference, invocation.inputs["codegraph"]) review = review_service.run_review( cast(ReviewCommand, invocation.inputs["request"]), + jobs=invocation.context.jobs, codegraph=invocation.context.codegraphs.load(reference), codegraph_diagnostics=reference.diagnostics, codegraph_failed_files=reference.failed_files, @@ -96,6 +98,7 @@ def run_review( self, command: ReviewCommand, *, + jobs: NodeJobs, codegraph, codegraph_diagnostics=(), codegraph_failed_files=(), @@ -168,6 +171,7 @@ def run_review( command, supported, codegraph=codegraph, + jobs=jobs, memory_service=memory_service, progress_callback=progress_callback, diagnostics=diagnostics, @@ -185,6 +189,7 @@ def run_review( reviews.append( self._simple_llm_review.run_files( fallback, + jobs=jobs, memory_service=memory_service, index=index, progress_callback=progress_callback, @@ -203,6 +208,7 @@ def _run_reachability_review( files: tuple[str, ...], *, codegraph, + jobs: NodeJobs, memory_service: MemoryService | None, progress_callback, diagnostics: list[ReviewDiagnostic], @@ -213,6 +219,7 @@ def _run_reachability_review( reviewed, analysis = self.file_review( files[0], settings=self._settings, + jobs=jobs, progress_callback=progress_callback, diagnostic_callback=provider_diagnostics.append, codegraph=codegraph, @@ -223,6 +230,7 @@ def _run_reachability_review( codebase_groups, analysis = self.codebase_reviews( files=files, settings=self._settings, + jobs=jobs, progress_callback=progress_callback, diagnostic_callback=provider_diagnostics.append, codegraph=codegraph, @@ -283,12 +291,12 @@ def review_options( settings = dict(settings or {}) if progress_callback is not None: settings["progress_callback"] = progress_callback - settings.setdefault("max_workers", self._config.max_workers) return ReachabilityReviewOptions(**settings) def codebase_reviews( self, *, + jobs: NodeJobs, files=None, settings=None, progress_callback=None, @@ -301,6 +309,7 @@ def codebase_reviews( progress_callback=progress_callback, ) analysis = self._service.analyze_codebase( + jobs=jobs, options=options, files=files, codegraph=codegraph, @@ -319,6 +328,7 @@ def file_review( self, file_path, *, + jobs: NodeJobs, settings=None, progress_callback=None, diagnostic_callback=None, @@ -331,6 +341,7 @@ def file_review( ) analysis = self._service.analyze_file( file_path, + jobs=jobs, options=options, codegraph=codegraph, diagnostic_callback=diagnostic_callback, diff --git a/src/metis/engine/nodes/reachability/service.py b/src/metis/engine/nodes/reachability/service.py index 6607ca64..1fa4f461 100644 --- a/src/metis/engine/nodes/reachability/service.py +++ b/src/metis/engine/nodes/reachability/service.py @@ -10,6 +10,7 @@ from metis.engine.codegraph import CodeGraph from metis.engine.codegraph import CodeGraphDiagnostic +from metis.engine.execution.contracts import NodeJobs from metis.engine.concurrency import serialized_progress_callback from metis.engine.nodes.codegraph import CodeGraphService from metis.engine.threat_context_retrieval import get_threat_model_context @@ -51,6 +52,7 @@ def analyze_file( self, file_path, *, + jobs: NodeJobs, options: ReachabilityReviewOptions, codegraph: CodeGraph | None = None, diagnostic_callback: Callable[[CodeGraphDiagnostic], None] | None = None, @@ -103,6 +105,7 @@ def analyze_file( analysis_graph = _copy_graph_nodes(graph, focus.node_names) outcome = self._frontier_reviewer.review( analysis_graph, + jobs=jobs, options=options, memory_service=memory_service, evidence_graph=graph, @@ -142,6 +145,7 @@ def analyze_file( def analyze_codebase( self, *, + jobs: NodeJobs, options: ReachabilityReviewOptions, files=None, codegraph: CodeGraph | None = None, @@ -188,6 +192,7 @@ def analyze_codebase( ) outcome = self._frontier_reviewer.review( analysis_graph, + jobs=jobs, options=options, memory_service=memory_service, evidence_graph=graph, diff --git a/src/metis/engine/nodes/simple_llm_review/registration.py b/src/metis/engine/nodes/simple_llm_review/registration.py index 27a44aa3..5409af1d 100644 --- a/src/metis/engine/nodes/simple_llm_review/registration.py +++ b/src/metis/engine/nodes/simple_llm_review/registration.py @@ -31,6 +31,7 @@ def execute(invocation: NodeInvocation) -> NodeResult: ) review = service.run_review( cast(ReviewCommand, invocation.inputs["request"]), + jobs=invocation.context.jobs, memory_service=memory, index=index, progress_callback=invocation.context.callbacks.progress, diff --git a/src/metis/engine/nodes/simple_llm_review/service.py b/src/metis/engine/nodes/simple_llm_review/service.py index fca49ec4..21cd6243 100644 --- a/src/metis/engine/nodes/simple_llm_review/service.py +++ b/src/metis/engine/nodes/simple_llm_review/service.py @@ -6,15 +6,14 @@ import logging import os from collections.abc import Callable -from concurrent.futures import ThreadPoolExecutor, as_completed from dataclasses import dataclass from typing import Any from typing import TYPE_CHECKING import unidiff -from metis.usage import submit_with_current_context from metis import runlog +from metis.engine.execution.contracts import NodeJobs from metis.utils import read_file_content from metis.engine.diff_utils import process_diff_file @@ -68,6 +67,7 @@ def run_review( self, command: ReviewCommand, *, + jobs: NodeJobs, memory_service: MemoryService | None = None, index: IndexCapability | None = None, progress_callback: Callable[[dict[str, object]], None] | None = None, @@ -91,6 +91,7 @@ def run_review( ) return self.run_files( files, + jobs=jobs, memory_service=memory_service, index=index, progress_callback=progress_callback, @@ -100,6 +101,7 @@ def run_files( self, files: tuple[str, ...], *, + jobs: NodeJobs, memory_service: MemoryService | None = None, index: IndexCapability | None = None, progress_callback: Callable[[dict[str, object]], None] | None = None, @@ -107,6 +109,7 @@ def run_files( return self._run_traditional_review( files, progress_callback, + jobs=jobs, memory_service=memory_service, review_graph=self._review_graph_factory(index), ) @@ -116,11 +119,13 @@ def _run_traditional_review( files: tuple[str, ...], progress_callback, *, + jobs: NodeJobs, memory_service: MemoryService | None, review_graph: Any, ) -> ReviewRun: outcome = self.execute_standard_review_with_outcome( files, + jobs=jobs, progress_callback=progress_callback, memory_service=memory_service, review_graph=review_graph, @@ -148,6 +153,7 @@ def execute_standard_review_with_outcome( self, files, *, + jobs: NodeJobs, progress_callback=None, memory_service: MemoryService | None = None, review_graph: Any | None = None, @@ -164,6 +170,7 @@ def execute_standard_review_with_outcome( memory_service=memory_service, review_graph=review_graph, ), + jobs, ): if error is not None: logger.error(f"Error reviewing file {path}: {error}") @@ -228,20 +235,20 @@ def _review_files( self, files, review_fn, + jobs: NodeJobs, ): - with ThreadPoolExecutor(max_workers=self._config.max_workers) as executor: - future_to_path = { - submit_with_current_context( - executor, self._review_file_task, review_fn, path - ): path - for path in files - } - for future in as_completed(future_to_path): - path = future_to_path[future] - try: - yield path, future.result(), None - except Exception as exc: - yield path, None, exc + def review(path): + try: + return path, self._review_file_task(review_fn, path), None + except Exception as exc: + return path, None, exc + + return jobs.run( + files, + review, + label=None, + result_key=str, + ) def review_patch( self, diff --git a/src/metis/engine/nodes/triage/registration.py b/src/metis/engine/nodes/triage/registration.py index 8832ce3b..02808cbb 100644 --- a/src/metis/engine/nodes/triage/registration.py +++ b/src/metis/engine/nodes/triage/registration.py @@ -43,6 +43,7 @@ def execute(invocation: NodeInvocation) -> NodeResult: ) triaged = adjudicator.triage_run( cast(TriageRun, invocation.inputs["request"]), + jobs=invocation.context.jobs, classifier=partial( classifier_service.classify, navigation=navigation, diff --git a/src/metis/engine/stages/service.py b/src/metis/engine/stages/service.py index 719f42fb..b03f33a6 100644 --- a/src/metis/engine/stages/service.py +++ b/src/metis/engine/stages/service.py @@ -11,6 +11,7 @@ from typing import cast from metis import runlog +from metis.engine.concurrency import JobScheduler from metis.engine.codegraph import CodeGraphReference from metis.engine.llm_runner import JsonPromptRunner from metis.runtime_settings import TriageOptions @@ -70,6 +71,11 @@ def __init__( catalog = NodeCatalog(registrations, entry_points) self._plans: dict[StageName, StagePlan] = {} self._validate_configuration(catalog, capabilities) + self._scheduler = JobScheduler(engine_config.max_workers) + self._jobs = self._scheduler.limit(engine_config.max_workers) + + def close(self) -> None: + self._scheduler.close() def execute_initialize( self, @@ -269,6 +275,7 @@ def _context( self._engine_config.capability_settings.model_tools.max_rounds ), chat_model_kwargs=dict(self._engine_config.chat_model_kwargs), + jobs=self._jobs, ), callbacks=_node_callbacks(callbacks), triage=self._triage_adjudicator, diff --git a/src/metis/engine/stages/triage/contracts.py b/src/metis/engine/stages/triage/contracts.py index 115b5a1c..3a02da2a 100644 --- a/src/metis/engine/stages/triage/contracts.py +++ b/src/metis/engine/stages/triage/contracts.py @@ -12,6 +12,7 @@ from metis.engine.execution.contracts import CheckpointCallback from metis.engine.execution.contracts import DebugCallback +from metis.engine.execution.contracts import NodeJobs from metis.engine.execution.contracts import ProgressCallback from metis.sarif.triage import SarifFinding from metis.memory import MemoryService @@ -40,6 +41,7 @@ def triage_run( run: TriageRun, *, classifier: TriageClassifier, + jobs: NodeJobs, memory_service: MemoryService | None = None, progress_callback: ProgressCallback | None = None, debug_callback: DebugCallback | None = None, diff --git a/src/metis/engine/stages/triage/service.py b/src/metis/engine/stages/triage/service.py index 5b18350b..a075a621 100644 --- a/src/metis/engine/stages/triage/service.py +++ b/src/metis/engine/stages/triage/service.py @@ -4,11 +4,11 @@ from __future__ import annotations from collections.abc import Callable -from concurrent.futures import ThreadPoolExecutor, as_completed import logging from typing import cast from metis.engine.stages.triage.models import TriageRun +from metis.engine.execution.contracts import NodeJobs from metis.memory import MemoryService from metis.engine.threat_context_retrieval import get_threat_model_context from metis.engine.threat_context_retrieval import threat_model_scope_policy @@ -19,7 +19,6 @@ save_sarif_file, ) from metis.sarif.utils import create_fingerprint -from metis.usage import submit_with_current_context from metis import runlog @@ -33,10 +32,8 @@ class TriageService: def __init__( self, *, - max_workers: int, triage_checkpoint_every: int, ) -> None: - self.max_workers = max(1, max_workers) self.triage_checkpoint_every = triage_checkpoint_every def _invoke_callback(self, callback, *args, **kwargs) -> None: @@ -238,51 +235,54 @@ def _triage_findings_parallel( debug_callback, checkpoint_callback, classifier: TriageClassifier, + jobs: NodeJobs, memory_service: MemoryService | None, processed: int, ) -> tuple[int, set[tuple[int, int]]]: handled: set[tuple[int, int]] = set() - with ThreadPoolExecutor(max_workers=self.max_workers) as executor: - future_map = {} - for idx, finding in enumerate(findings, start=1): - self._emit_triage_progress( - progress_callback, - total, - "start", - index=idx, - finding=finding, - ) - future = submit_with_current_context( - executor, - self._triage_one_finding, + scheduled = list(enumerate(findings, start=1)) + for idx, finding in scheduled: + self._emit_triage_progress( + progress_callback, + total, + "start", + index=idx, + finding=finding, + ) + + def invoke(item): + idx, finding = item + try: + decision = self._triage_one_finding( finding, classifier=classifier, debug_callback=debug_callback, memory_service=memory_service, ) - future_map[future] = (idx, finding) + return idx, finding, decision, None + except Exception as exc: + return idx, finding, None, exc - for future in as_completed(future_map): - idx, finding = future_map[future] - try: - decision = future.result() - error = None - except Exception as exc: - decision = None - error = exc - processed, was_handled = self._handle_finding_result( - triaged_payload=triaged_payload, - finding=finding, - total=total, - idx=idx, - decision=decision, - error=error, - progress_callback=progress_callback, - checkpoint_callback=checkpoint_callback, - processed=processed, - ) - if was_handled: - handled.add((finding.run_index, finding.result_index)) + outcomes = jobs.run( + scheduled, + invoke, + label=None, + result_key=lambda item: item[0], + ) + for idx, finding, decision, error in outcomes: + processed, was_handled = self._handle_finding_result( + triaged_payload=triaged_payload, + finding=finding, + total=total, + idx=idx, + decision=decision, + error=error, + progress_callback=progress_callback, + checkpoint_callback=checkpoint_callback, + processed=processed, + ) + if was_handled: + handled.add((finding.run_index, finding.result_index)) return processed, handled def triage_run( @@ -290,6 +290,7 @@ def triage_run( run: TriageRun, *, classifier: TriageClassifier, + jobs: NodeJobs, memory_service: MemoryService | None = None, progress_callback=None, debug_callback=None, @@ -307,6 +308,7 @@ def triage_run( debug_callback=debug_callback, checkpoint_callback=checkpoint_callback, classifier=classifier, + jobs=jobs, memory_service=memory_service, processed=run.processed, ) diff --git a/src/metis/execution_nodes.py b/src/metis/execution_nodes.py index eac55565..ad39637f 100644 --- a/src/metis/execution_nodes.py +++ b/src/metis/execution_nodes.py @@ -22,6 +22,7 @@ from metis.engine.execution.contracts import NodeContext from metis.engine.execution.contracts import NodeCodeGraphs from metis.engine.execution.contracts import NodeInvocation +from metis.engine.execution.contracts import NodeJobs from metis.engine.execution.contracts import NodeRegistration from metis.engine.execution.contracts import NodeResult from metis.engine.execution.contracts import NodeRuntime @@ -53,6 +54,7 @@ "NodeContext", "NodeCodeGraphs", "NodeInvocation", + "NodeJobs", "NodeRegistration", "NodeResult", "NodeRuntime", diff --git a/tests/conftest.py b/tests/conftest.py index b077c6a8..5b084ff4 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -7,11 +7,21 @@ import pytest from metis.configuration import load_execution_config from metis.engine import MetisEngine +from metis.engine.concurrency import JobScheduler from metis.runtime_settings import ModelToolSettings from metis.runtime_settings import CapabilityRuntimeSettings from unittest.mock import Mock, MagicMock +@pytest.fixture +def node_jobs(): + scheduler = JobScheduler(2) + try: + yield scheduler.limit(2) + finally: + scheduler.close() + + @pytest.fixture def dummy_backend(): backend = Mock() @@ -110,7 +120,7 @@ def engine( ): codebase_path = tmp_path / "data" shutil.copytree(Path(__file__).parent / "data", codebase_path) - return MetisEngine( + engine = MetisEngine( codebase_path=str(codebase_path), vector_backend=dummy_backend, language_plugin="c", @@ -123,6 +133,10 @@ def engine( capability_settings=capability_settings, execution_config=execution_with_index, ) + try: + yield engine + finally: + engine.close() def pytest_addoption(parser): diff --git a/tests/test_configuration.py b/tests/test_configuration.py index 10b41164..a12cf5e0 100644 --- a/tests/test_configuration.py +++ b/tests/test_configuration.py @@ -124,6 +124,26 @@ def test_runtime_loads_codegraph_and_reachability_tuning_outside_execution( assert runtime["triage_options"] == TriageOptions(include_triaged=True) +@pytest.mark.parametrize("max_workers", (0, -1, True, "2")) +def test_runtime_rejects_invalid_max_workers(tmp_path, monkeypatch, max_workers): + config_path = _write_config( + tmp_path, + yaml.safe_dump( + { + "metis_engine": {"max_workers": max_workers}, + "llm_provider": {"name": "openai", "model": "test-model"}, + } + ), + ) + monkeypatch.setenv("OPENAI_API_KEY", "chat-key") + + with pytest.raises( + ValueError, + match="metis_engine.max_workers must be a positive integer", + ): + load_runtime_config(config_path) + + def test_runtime_passes_language_replacements_from_selected_yaml(tmp_path, monkeypatch): config_path = _write_config( tmp_path, diff --git a/tests/test_engine_core.py b/tests/test_engine_core.py index 2cd2a79f..3d333be3 100644 --- a/tests/test_engine_core.py +++ b/tests/test_engine_core.py @@ -43,6 +43,23 @@ def _execution_with_index() -> dict[str, object]: return execution +@pytest.mark.parametrize("max_workers", (0, -1, True, "2")) +def test_engine_rejects_invalid_max_workers(capability_settings, max_workers): + with pytest.raises( + ValueError, + match="MetisEngine.max_workers must be a positive integer", + ): + MetisEngine( + vector_backend=Mock(), + llm_provider=Mock(), + max_workers=max_workers, + max_token_length=2048, + llama_query_model="gpt-test", + similarity_top_k=3, + capability_settings=capability_settings, + ) + + def test_index_capability_rejects_missing_retrievers(capability_settings): bad_backend = Mock() bad_backend.init = Mock() @@ -723,5 +740,13 @@ def test_close_clears_retriever_cache_and_closes_backend(capability_settings): assert engine._state.retriever_docs is None backend.close.assert_called_once() + with pytest.raises(RuntimeError, match="cannot schedule new futures"): + engine.execution._jobs.run( + [None], + lambda _value: None, + label=None, + result_key=str, + ) + assert index.get_retrievers() == ("code-retriever", "docs-retriever") assert backend.get_retrievers.call_count == 2 diff --git a/tests/test_engine_review.py b/tests/test_engine_review.py index a229dc97..03fabff1 100644 --- a/tests/test_engine_review.py +++ b/tests/test_engine_review.py @@ -60,7 +60,10 @@ def test_simple_llm_review_processes_the_selected_scope(engine): ) ) - run = service.run_review(ReviewCommand(mode="code")) + run = service.run_review( + ReviewCommand(mode="code"), + jobs=engine.execution._jobs, + ) assert run.status is ReviewStatus.SUCCEEDED assert run.diagnostics == () @@ -78,6 +81,7 @@ def test_patch_review_uses_simple_llm_review(engine): run = service.run_review( ReviewCommand(mode="patch", target="change.patch"), + jobs=engine.execution._jobs, ) assert run.status is ReviewStatus.SUCCEEDED @@ -122,6 +126,7 @@ def test_reachability_review_falls_back_for_unsupported_files(engine, caplog): ): run = service.run_review( ReviewCommand(mode="code"), + jobs=engine.execution._jobs, codegraph=CodeGraph(), ) @@ -134,6 +139,7 @@ def test_reachability_review_falls_back_for_unsupported_files(engine, caplog): assert service.codebase_reviews.call_args.kwargs["files"] == (c_file,) fallback.run_files.assert_called_once_with( (python_file,), + jobs=engine.execution._jobs, memory_service=None, index=None, progress_callback=None, @@ -167,6 +173,7 @@ def test_reachability_does_not_run_simple_review_for_supported_files(engine): ) run = service.run_review( ReviewCommand(mode="file", target=target), + jobs=engine.execution._jobs, codegraph=CodeGraph(), ) @@ -201,6 +208,7 @@ def test_reachability_missing_graph_coverage_is_inconclusive(engine): run = service.run_review( ReviewCommand(mode="file", target=target), + jobs=engine.execution._jobs, codegraph=CodeGraph(), ) @@ -236,6 +244,7 @@ def test_reachability_operational_failure_detail_survives_product_boundary(engin run = service.run_review( ReviewCommand(mode="file", target=target), + jobs=engine.execution._jobs, codegraph=CodeGraph(), ) @@ -261,6 +270,7 @@ def test_reachability_review_rejects_symlink_outside_codebase(engine, tmp_path): run = service.run_review( ReviewCommand(mode="file", target=str(target)), + jobs=engine.execution._jobs, codegraph=CodeGraph(), ) @@ -281,7 +291,11 @@ def test_reachability_defers_fallback_to_explicit_simple_review(engine): service._repository.get_code_files = Mock(return_value=[python_file]) service.supports_file = Mock(return_value=False) - run = service.run_review(ReviewCommand(mode="code"), codegraph=CodeGraph()) + run = service.run_review( + ReviewCommand(mode="code"), + jobs=engine.execution._jobs, + codegraph=CodeGraph(), + ) assert run.status is ReviewStatus.SUCCEEDED assert run.result is not None @@ -299,7 +313,11 @@ def test_reachability_empty_scope_is_inconclusive(engine): ) service._repository.get_code_files = Mock(return_value=[]) - run = service.run_review(ReviewCommand(mode="code"), codegraph=CodeGraph()) + run = service.run_review( + ReviewCommand(mode="code"), + jobs=engine.execution._jobs, + codegraph=CodeGraph(), + ) assert run.status is ReviewStatus.INCONCLUSIVE diff --git a/tests/test_execution_graph.py b/tests/test_execution_graph.py index 56fc1677..c3a3a1eb 100644 --- a/tests/test_execution_graph.py +++ b/tests/test_execution_graph.py @@ -9,6 +9,7 @@ import pytest from pydantic import BaseModel from pydantic import ConfigDict +from pydantic import ValidationError from metis.configuration import load_execution_config from metis.engine.codegraph import CodeGraph @@ -48,6 +49,18 @@ from metis.runtime_settings import TriageOptions +class _Jobs: + def __init__(self) -> None: + self.limits: list[int] = [] + + def limit(self, max_concurrency: int): + self.limits.append(max_concurrency) + return self + + def run(self, jobs, worker, **_kwargs): + return [worker(job) for job in jobs] + + def _configuration() -> ExecutionConfiguration: return ExecutionConfiguration.model_validate( { @@ -127,7 +140,15 @@ def test_finding_dedup_combines_producers_before_result() -> None: _node_context(), stage="review", prompts=prompt_runner, - runtime=NodeRuntime("test", 2, 1000, {}, 1, lambda _text: 1), + runtime=NodeRuntime( + "test", + 2, + 1000, + {}, + 1, + lambda _text: 1, + jobs=_Jobs(), + ), ) result = finding_dedup_node.execute( @@ -167,7 +188,15 @@ def test_finding_dedup_fails_open_when_batch_exceeds_global_limit() -> None: _node_context(), stage="review", prompts=prompt_runner, - runtime=NodeRuntime("test", 1, 10, {}, 1, lambda _text: 100), + runtime=NodeRuntime( + "test", + 1, + 10, + {}, + 1, + lambda _text: 100, + jobs=_Jobs(), + ), ) result = finding_dedup_node.execute( NodeInvocation( @@ -237,11 +266,75 @@ def _node_context() -> NodeContext: capabilities={}, prompts=Mock(), codegraphs=Mock(), - runtime=NodeRuntime("test", 1, 1000, {}, 1, lambda _text: 1), + runtime=NodeRuntime( + "test", + 1, + 1000, + {}, + 1, + lambda _text: 1, + jobs=_Jobs(), + ), callbacks=NodeCallbacks(), ) +@pytest.mark.parametrize("max_concurrency", (0, -1, True, "2")) +def test_node_max_concurrency_requires_a_strict_positive_integer( + max_concurrency: object, +) -> None: + with pytest.raises(ValidationError): + StageConfiguration.model_validate( + {"nodes": {"node": {"max_concurrency": max_concurrency}}} + ) + + +@pytest.mark.parametrize( + ("configured", "expected"), + ((None, 2), (1, 1), (3, 2)), +) +def test_node_max_concurrency_caps_runtime_limit( + configured: int | None, + expected: int, +) -> None: + registration = NodeRegistration( + "node", + "initialize", + EmptyNodeConfiguration, + {}, + {"workers": int}, + lambda invocation: NodeResult( + {"workers": invocation.context.runtime.max_workers} + ), + ) + node = {} if configured is None else {"max_concurrency": configured} + stage = StageConfiguration.model_validate( + {"outputs": ["node"], "nodes": {"node": node}} + ) + context = replace( + _node_context(), + runtime=NodeRuntime( + "test", + 2, + 1000, + {}, + 1, + lambda _text: 1, + jobs=_Jobs(), + ), + ) + + result = run_stage( + "initialize", + _compile((registration,), stage), + context, + {}, + ) + + assert result.outputs == {"node": expected} + assert context.jobs.limits == [expected] + + def _value_node(name: str, value: str) -> NodeRegistration: return NodeRegistration( name, diff --git a/tests/test_execution_node_api.py b/tests/test_execution_node_api.py index 21603841..9da64a54 100644 --- a/tests/test_execution_node_api.py +++ b/tests/test_execution_node_api.py @@ -38,6 +38,7 @@ def test_execution_node_public_api_is_explicit() -> None: "NodeCodeGraphs", "NodeContext", "NodeInvocation", + "NodeJobs", "NodeRegistration", "NodeResult", "NodeRuntime", diff --git a/tests/test_navigation_capability.py b/tests/test_navigation_capability.py index 6a838935..cc0a9698 100644 --- a/tests/test_navigation_capability.py +++ b/tests/test_navigation_capability.py @@ -1,8 +1,12 @@ # SPDX-FileCopyrightText: Copyright 2026 Arm Limited and/or its affiliates # SPDX-License-Identifier: Apache-2.0 +from concurrent.futures import ThreadPoolExecutor import pytest import subprocess +from threading import Barrier +from threading import BrokenBarrierError +from threading import Lock from metis.engine.capabilities.navigation import NavigationCapability @@ -114,3 +118,35 @@ def test_shell_grep_forces_filename_prefix_for_single_file(tmp_path): out.splitlines()[0].endswith("/a.c:2:beta") or out.splitlines()[0] == "a.c:2:beta" ) + + +def test_shell_navigation_serializes_subprocesses(tmp_path, monkeypatch): + runner = NavigationCapability( + codebase_path=str(tmp_path), timeout_seconds=8, max_chars=16000 + ) + runner._has_grep = True + rendezvous = Barrier(2) + state_lock = Lock() + active = 0 + peak = 0 + + def run(argv, **_kwargs): + nonlocal active, peak + with state_lock: + active += 1 + peak = max(peak, active) + try: + rendezvous.wait(timeout=0.2) + except BrokenBarrierError: + pass + with state_lock: + active -= 1 + return subprocess.CompletedProcess(argv, 0, stdout="match", stderr="") + + monkeypatch.setattr(subprocess, "run", run) + + with ThreadPoolExecutor(max_workers=2) as executor: + results = tuple(executor.map(lambda _index: runner.grep("x", "."), range(2))) + + assert results == ("match", "match") + assert peak == 1 diff --git a/tests/test_reachability_incremental_review.py b/tests/test_reachability_incremental_review.py index 024b1902..83587e4e 100644 --- a/tests/test_reachability_incremental_review.py +++ b/tests/test_reachability_incremental_review.py @@ -79,6 +79,7 @@ def _reviewer(tmp_path: Path) -> IncrementalGraphReviewer: def test_review_runs_one_discovery_pass_with_deterministic_contracts( tmp_path: Path, + node_jobs, ) -> None: (tmp_path / "graph.c").write_text( "void root(void) { sink(); }\nvoid sink(void) {}\n", @@ -125,8 +126,8 @@ def invoke(request: Any) -> object: reviewer._runner.invoke = Mock(side_effect=invoke) outcome = reviewer.review( _graph(root, sink), + jobs=node_jobs, options=ReachabilityReviewOptions( - max_workers=1, domain_profiles=(" GPU ",), domain_hints=("CustomTerm",), ), @@ -228,7 +229,7 @@ def test_compact_packet_uses_tables_numbered_source_and_target_tags( packets, failures = reviewer._build_packets( _graph(root, helper, stop), (root.unique_name, helper.unique_name, stop.unique_name), - options=ReachabilityReviewOptions(max_workers=1), + options=ReachabilityReviewOptions(), memory_service=None, selected_node_ids={root.unique_name}, contracts={ @@ -285,7 +286,7 @@ def test_oversized_source_line_only_fails_its_function(tmp_path: Path) -> None: packets, failures = reviewer._build_packets( _graph(small, oversized), (small.unique_name, oversized.unique_name), - options=ReachabilityReviewOptions(max_workers=1), + options=ReachabilityReviewOptions(), memory_service=None, ) @@ -314,12 +315,15 @@ def test_source_chunk_runtime_error_propagates( reviewer._build_packets( _graph(root), (root.unique_name,), - options=ReachabilityReviewOptions(max_workers=1), + options=ReachabilityReviewOptions(), memory_service=None, ) -def test_invalid_discovery_packet_splits_and_reviews_children(tmp_path: Path) -> None: +def test_invalid_discovery_packet_splits_and_reviews_children( + tmp_path: Path, + node_jobs, +) -> None: (tmp_path / "graph.c").write_text( "void first(void) {}\nvoid second(void) {}\n", encoding="utf-8", @@ -352,7 +356,8 @@ def respond(_messages: object) -> dict[str, object | None]: ) as session: outcome = reviewer.review( _graph(_function("first", 1), _function("second", 2)), - options=ReachabilityReviewOptions(max_workers=1), + jobs=node_jobs, + options=ReachabilityReviewOptions(), ) assert calls == 3 @@ -373,6 +378,7 @@ def respond(_messages: object) -> dict[str, object | None]: def test_indexed_null_condition_reaches_deterministic_admission( tmp_path: Path, monkeypatch: MonkeyPatch, + node_jobs, ) -> None: source = ( "void *allocate(void);\nvoid root(void) { void *p = allocate(); consume(p); }\n" @@ -482,7 +488,8 @@ def invoke(request: Any) -> object: reviewer._runner.invoke = Mock(side_effect=invoke) outcome = reviewer.review( graph, - options=ReachabilityReviewOptions(max_workers=1), + jobs=node_jobs, + options=ReachabilityReviewOptions(), ) assert outcome.findings == () diff --git a/tests/test_reachability_treesitter.py b/tests/test_reachability_treesitter.py index b466ad45..28b8a75a 100644 --- a/tests/test_reachability_treesitter.py +++ b/tests/test_reachability_treesitter.py @@ -1190,6 +1190,7 @@ def test_scoped_review_limits_frontier_analysis_to_relevant_graph(): service._frontier_reviewer = frontier_reviewer service.analyze_codebase( + jobs=Mock(), files=["src/target.c"], codegraph=graph, options=ReachabilityReviewOptions(), @@ -1231,6 +1232,7 @@ def test_file_review_scans_full_focus(): service.analyze_file( "src/review.c", + jobs=Mock(), options=ReachabilityReviewOptions(), codegraph=graph, ) @@ -1265,6 +1267,7 @@ def test_file_review_reports_missing_codegraph_coverage(graph, expected_code): analysis = service.analyze_file( "src/review.c", + jobs=Mock(), options=ReachabilityReviewOptions(), codegraph=graph, diagnostic_callback=diagnostics.append, diff --git a/tests/test_runlog.py b/tests/test_runlog.py index d274a3a7..8f15956c 100644 --- a/tests/test_runlog.py +++ b/tests/test_runlog.py @@ -5,8 +5,11 @@ import json import threading +import time +from concurrent.futures import ThreadPoolExecutor from pathlib import Path from types import SimpleNamespace +from unittest.mock import Mock import pytest from langchain_core.messages import AIMessage @@ -15,7 +18,7 @@ from langchain_core.runnables import RunnableLambda from metis.engine import MetisEngine -from metis.engine.concurrency import run_jobs +from metis.engine.concurrency import JobScheduler from metis.engine.execution import ExecutionResult from metis.engine.execution import ExecutionStatus from metis.engine.execution.catalog import NodeCatalog @@ -314,7 +317,7 @@ def fail(_invocation): capabilities={}, prompts=object(), codegraphs=object(), - runtime=NodeRuntime("test", 1, 1000, {}, 1), + runtime=NodeRuntime("test", 1, 1000, {}, 1, jobs=Mock()), callbacks=NodeCallbacks(progress=progress.append), ) @@ -364,14 +367,17 @@ def fail(_invocation): def test_parallel_jobs_are_traced_with_compact_keys(tmp_path): - with runlog.open_runlog(_exact_config(tmp_path)) as session: - results = run_jobs( - [2, 3], - lambda value: value * value, - max_workers=2, - label="Reachability security review", - result_key=lambda value: value, - ) + scheduler = JobScheduler(2) + try: + with runlog.open_runlog(_exact_config(tmp_path)) as session: + results = scheduler.limit(2).run( + [2, 3], + lambda value: value * value, + label="Reachability security review", + result_key=lambda value: value, + ) + finally: + scheduler.close() assert sorted(results) == [4, 9] tasks = [ @@ -385,6 +391,89 @@ def test_parallel_jobs_are_traced_with_compact_keys(tmp_path): runlog.validate_runlog(session.ndjson_path) +def test_scheduler_shares_global_workers_across_node_limits() -> None: + scheduler = JobScheduler(2) + state_lock = threading.Lock() + active = 0 + peak = 0 + + def work(value: int) -> int: + nonlocal active, peak + with state_lock: + active += 1 + peak = max(peak, active) + time.sleep(0.02) + with state_lock: + active -= 1 + return value + + try: + jobs = scheduler.limit(2) + with ThreadPoolExecutor(max_workers=2) as callers: + first = callers.submit( + jobs.run, + [1, 2, 3], + work, + label=None, + result_key=str, + ) + second = callers.submit( + jobs.run, + [4, 5, 6], + work, + label=None, + result_key=str, + ) + assert sorted((*first.result(), *second.result())) == [1, 2, 3, 4, 5, 6] + finally: + scheduler.close() + + assert peak == 2 + + +def test_scheduler_limit_one_bounds_submissions_and_runs_off_coordinator() -> None: + scheduler = JobScheduler(2) + first_started = threading.Event() + second_started = threading.Event() + release_first = threading.Event() + coordinator: list[int] = [] + workers: list[int] = [] + + def work(value: int) -> int: + workers.append(threading.get_ident()) + if value == 1: + first_started.set() + release_first.wait(timeout=2) + else: + second_started.set() + return value + + def run() -> list[int]: + coordinator.append(threading.get_ident()) + return scheduler.limit(1).run( + [1, 2], + work, + label=None, + result_key=str, + ) + + try: + with ThreadPoolExecutor(max_workers=1) as caller: + result = caller.submit(run) + try: + assert first_started.wait(timeout=1) + second_was_blocked = not second_started.wait(timeout=0.1) + finally: + release_first.set() + assert result.result() == [1, 2] + finally: + scheduler.close() + + assert second_was_blocked + assert second_started.is_set() + assert all(worker != coordinator[0] for worker in workers) + + def test_prompt_retry_records_logical_attempts(tmp_path): class Provider: def get_chat_model(self, **_kwargs): diff --git a/tests/test_sarif_triage.py b/tests/test_sarif_triage.py index bef9bcd4..22215268 100644 --- a/tests/test_sarif_triage.py +++ b/tests/test_sarif_triage.py @@ -119,7 +119,7 @@ def triage(self, _request): raise RuntimeError("boom") return {"status": "valid", "reason": "confirmed"} - engine._triage_service.max_workers = 1 + engine.execution._jobs = engine.execution._scheduler.limit(1) dummy_workflow = _DummyWorkflow() monkeypatch.setattr( engine._triage_classifier, @@ -151,7 +151,7 @@ class _FailingWorkflow: def triage(self, _request): raise RuntimeError("boom") - engine._triage_service.max_workers = 1 + engine.execution._jobs = engine.execution._scheduler.limit(1) monkeypatch.setattr( engine._triage_classifier, "_workflow", @@ -189,7 +189,7 @@ def triage(self, _request): "_workflow", lambda _navigation, _rounds: _DummyWorkflow(), ) - engine._triage_service.max_workers = 1 + engine.execution._jobs = engine.execution._scheduler.limit(1) out = engine.execute_triage(payload)["sarif"] props = out["runs"][0]["results"][0]["properties"] @@ -241,7 +241,7 @@ def save_checkpoint(_path, checkpoint_payload): "metis.engine.stages.triage.service.save_sarif_file", save_checkpoint, ) - engine._triage_service.max_workers = 1 + engine.execution._jobs = engine.execution._scheduler.limit(1) engine._triage_service.triage_checkpoint_every = 2 result = engine.execute_triage( @@ -298,7 +298,7 @@ def triage(self, request): "_workflow", lambda _navigation, _rounds: _DummyWorkflow(), ) - engine._triage_service.max_workers = 1 + engine.execution._jobs = engine.execution._scheduler.limit(1) out = engine.execute_triage(payload)["sarif"] assert out["runs"][0]["results"][0]["properties"]["metisTriaged"] is True diff --git a/tests/test_usage.py b/tests/test_usage.py index 7b207a6f..f210b615 100644 --- a/tests/test_usage.py +++ b/tests/test_usage.py @@ -111,7 +111,10 @@ def _review_file(path, **_kwargs): service._review_file_standard = _review_file with engine.usage_command("review_code") as command: - results = service.execute_standard_review_with_outcome(files).result["reviews"] + results = service.execute_standard_review_with_outcome( + files, + jobs=engine.execution._jobs, + ).result["reviews"] record = engine.finalize_usage_command(command)