diff --git a/colcon_parallel_executor/executor/parallel.py b/colcon_parallel_executor/executor/parallel.py index 5c4cc87..4c378a0 100644 --- a/colcon_parallel_executor/executor/parallel.py +++ b/colcon_parallel_executor/executor/parallel.py @@ -19,6 +19,14 @@ from colcon_core.subprocess import new_event_loop from colcon_core.subprocess import SIGINT_RESULT from colcon_parallel_executor.event.executor import ParallelStatus +from colcon_parallel_executor.resource_guard import \ + add_resource_guard_arguments +from colcon_parallel_executor.resource_guard import \ + initialize_resource_guard_extensions +from colcon_parallel_executor.resource_guard import \ + run_guarded_job +from colcon_parallel_executor.resource_guard.execution_policy import \ + ExecutionPolicyGuard logger = colcon_logger.getChild(__name__) @@ -62,6 +70,8 @@ def add_arguments(self, *, parser): # noqa: D102 "or '0' for no limit " '(default: {max_workers_default})'.format_map(locals())) + add_resource_guard_arguments(parser=parser) + def execute(self, args, jobs, *, on_error=OnError.interrupt): # noqa: D102 # avoid debug message from asyncio when colcon uses debug log level asyncio_logger = logging.getLogger('asyncio') @@ -111,6 +121,11 @@ def execute(self, args, jobs, *, on_error=OnError.interrupt): # noqa: D102 return result async def _execute(self, args, jobs, *, on_error): + guards = await initialize_resource_guard_extensions(args) + + policy_guard = ExecutionPolicyGuard(on_error) + guards.append(policy_guard) + # count the number of dependent jobs for each job # in order to process jobs with more dependent jobs first recursive_dependent_counts = {} @@ -142,11 +157,6 @@ async def _execute(self, args, jobs, *, on_error): # take "ready" jobs take_jobs = [] for package_name, job, _ in ready_jobs: - # don't schedule more jobs then workers - # to prevent starting further jobs when a job fails - if args.parallel_workers: - if len(futures) + len(take_jobs) >= args.parallel_workers: - break take_jobs.append((package_name, job)) del jobs[package_name] @@ -154,7 +164,7 @@ async def _execute(self, args, jobs, *, on_error): for package_name, job in take_jobs: assert iscoroutinefunction(job.__call__), \ 'Job is not a coroutine' - future = asyncio.ensure_future(job()) + future = asyncio.ensure_future(run_guarded_job(job, guards)) futures[future] = job # wait for futures diff --git a/colcon_parallel_executor/resource_guard/__init__.py b/colcon_parallel_executor/resource_guard/__init__.py new file mode 100644 index 0000000..575bea0 --- /dev/null +++ b/colcon_parallel_executor/resource_guard/__init__.py @@ -0,0 +1,150 @@ +# Copyright 2026 Open Source Robotics Foundation, Inc. +# Licensed under the Apache License, Version 2.0 + +import asyncio +from contextlib import AsyncExitStack +import traceback + +from colcon_core.logging import colcon_logger +from colcon_core.plugin_system import instantiate_extensions +from colcon_core.plugin_system import order_extensions_grouped_by_priority + +logger = colcon_logger.getChild(__name__) + + +class JobIncompleteError(Exception): + """ + Raised when a job does not finish successfully. + + If `result` is None, the job was cleanly skipped before execution. + If `result` is not None, the job failed with a non-zero exit code. + """ + + def __init__(self, result=None): # noqa: D107 + super().__init__() + self.result = result + + +class ResourceGuardExtensionPoint: + """ + The interface for parallel execution resource guard extensions. + + A resource guard provider provides asynchronous context managers that + act as gatekeepers to throttle job execution. + """ + + """The version of the resource guard extension interface.""" + EXTENSION_POINT_VERSION = '1.0' + + """The priority of resource guard extensions.""" + PRIORITY = 100 + + def __init__(self): # noqa: D107 + super().__init__() + + def add_arguments(self, *, parser): + """ + Add command line arguments specific to the resource guard. + + :param parser: The argument parser + """ + pass + + async def initialize(self, args): + """ + Initialize the resource guard provider before starting job execution. + + :param args: The parsed command line arguments + """ + pass + + def get_guard(self, job): + """ + Get the asynchronous context manager (guard) for the specified job. + + If this guard does not apply to the job, return None. + + :param job: The job object + :returns: An object implementing __aenter__ and __aexit__, or None + """ + return None + + +def get_resource_guard_extensions(*, group_name=None): + """ + Get the available resource guard extensions. + + :rtype: OrderedDict + """ + if group_name is None: + group_name = 'colcon_parallel_executor.resource_guard' + + try: + extensions = instantiate_extensions(group_name) + except Exception as e: # noqa: B902, F841 + exc = traceback.format_exc() + logger.error( + 'Exception in resource guard discovery: {e}\n{exc}' + .format_map(locals())) + return {} + + for name, extension in extensions.items(): + extension.RESOURCE_GUARD_NAME = name + return order_extensions_grouped_by_priority(extensions) + + +def add_resource_guard_arguments(parser, *, extensions=None): + """Add command line arguments for the resource guard extensions.""" + if extensions is None: + extensions = get_resource_guard_extensions() + for priority in sorted(extensions.keys()): + for name, guard in extensions[priority].items(): + try: + retval = guard.add_arguments(parser=parser) + assert retval is None, 'add_arguments() should return None' + except Exception as e: # noqa: B902 + logger.error( + 'Failed to add arguments for resource guard: %s', e) + + +async def initialize_resource_guard_extensions(args, *, extensions=None): + """ + Initialize the resource guard extensions. + + :param args: The parsed command line arguments + """ + if extensions is None: + extensions = get_resource_guard_extensions() + guards = [] + for priority in sorted(extensions.keys()): + for name, guard in extensions[priority].items(): + try: + retval = await guard.initialize(args) + assert retval is None, 'initialize() should return None' + guards.append(guard) + except Exception as e: # noqa: B902 + logger.error( + 'Failed to initialize resource guard: %s', e) + return guards + + +async def run_guarded_job(job, guards): + """ + Acquire resource guards and execute the job inside their context. + + :param job: The job coroutine function + :param guards: List of resource guard extension providers + """ + try: + async with AsyncExitStack() as stack: + for provider in guards: + guard = provider.get_guard(job) + if guard is not None: + await stack.enter_async_context(guard) + await asyncio.sleep(0) + result = await job() + if result: + raise JobIncompleteError(result) + return result + except JobIncompleteError as e: + return e.result diff --git a/colcon_parallel_executor/resource_guard/execution_policy.py b/colcon_parallel_executor/resource_guard/execution_policy.py new file mode 100644 index 0000000..8fd1f11 --- /dev/null +++ b/colcon_parallel_executor/resource_guard/execution_policy.py @@ -0,0 +1,36 @@ +# Copyright 2026 Open Source Robotics Foundation, Inc. +# Licensed under the Apache License, Version 2.0 + +from colcon_core.executor import OnError +from colcon_parallel_executor.resource_guard import JobIncompleteError +from colcon_parallel_executor.resource_guard import \ + ResourceGuardExtensionPoint + + +class ExecutionPolicyGuard(ResourceGuardExtensionPoint): + """ + An implicit resource guard that cleanly skips pending jobs. + + This is activated when a failure has occurred during execution. + """ + + def __init__(self, on_error): # noqa: D107 + super().__init__() + self.on_error = on_error + self.skip_all = False + + def get_guard(self, job): # noqa: D102 + return self + + async def __aenter__(self): # noqa: D105 + if self.skip_all: + raise JobIncompleteError() + + async def __aexit__(self, exc_type, exc_val, exc_tb): # noqa: D105 + if ( + exc_type is not None and + self.on_error in (OnError.interrupt, OnError.skip_pending) and + (exc_type is not JobIncompleteError or exc_val.result is not None) + ): + self.skip_all = True + return False diff --git a/colcon_parallel_executor/resource_guard/worker_limiter.py b/colcon_parallel_executor/resource_guard/worker_limiter.py new file mode 100644 index 0000000..b142680 --- /dev/null +++ b/colcon_parallel_executor/resource_guard/worker_limiter.py @@ -0,0 +1,25 @@ +# Copyright 2026 Open Source Robotics Foundation, Inc. +# Licensed under the Apache License, Version 2.0 + +import asyncio + +from colcon_core.plugin_system import satisfies_version +from colcon_parallel_executor.resource_guard import ResourceGuardExtensionPoint + + +class WorkerLimiterGuard(ResourceGuardExtensionPoint): + """Limits concurrent jobs using a semaphore.""" + + PRIORITY = 100 + + def __init__(self): # noqa: D107 + super().__init__() + satisfies_version( + ResourceGuardExtensionPoint.EXTENSION_POINT_VERSION, '^1.0') + + async def initialize(self, args): # noqa: D102 + workers = getattr(args, 'parallel_workers', 0) + self._semaphore = asyncio.Semaphore(workers) if workers > 0 else None + + def get_guard(self, job): # noqa: D102 + return self._semaphore diff --git a/setup.cfg b/setup.cfg index e13d5ae..fb23a5e 100644 --- a/setup.cfg +++ b/setup.cfg @@ -65,6 +65,10 @@ colcon_core.event_handler = parallel_status = colcon_parallel_executor.event_handler.parallel_status:ParallelStatusEventHandler colcon_core.executor = parallel = colcon_parallel_executor.executor.parallel:ParallelExecutorExtension +colcon_core.extension_point = + colcon_parallel_executor.resource_guard = colcon_parallel_executor.resource_guard:ResourceGuardExtensionPoint +colcon_parallel_executor.resource_guard = + default_worker_limiter = colcon_parallel_executor.resource_guard.worker_limiter:WorkerLimiterGuard [flake8] import-order-style = google diff --git a/test/run_until_complete.py b/test/run_until_complete.py new file mode 100644 index 0000000..377f8e3 --- /dev/null +++ b/test/run_until_complete.py @@ -0,0 +1,16 @@ +# Copyright 2016-2018 Dirk Thomas +# Licensed under the Apache License, Version 2.0 + +import asyncio + +from colcon_core.subprocess import new_event_loop + + +def run_until_complete(coroutine): + loop = new_event_loop() + asyncio.set_event_loop(loop) + try: + return loop.run_until_complete(coroutine) + finally: + loop.close() + assert loop.is_closed() diff --git a/test/spell_check.words b/test/spell_check.words index e1048c6..1b4aafb 100644 --- a/test/spell_check.words +++ b/test/spell_check.words @@ -1,6 +1,9 @@ +aenter +aexit apache argparse asyncio +autouse capsys colcon contextlib @@ -20,6 +23,8 @@ pydocstyle pytest readouterr returncode +retval +rtype scspell setuptools sigint diff --git a/test/test_executor_parallel.py b/test/test_executor_parallel.py index 4ba1d82..31e5edf 100644 --- a/test/test_executor_parallel.py +++ b/test/test_executor_parallel.py @@ -19,11 +19,25 @@ from colcon_parallel_executor.executor.parallel import counting_number from colcon_parallel_executor.executor.parallel \ import ParallelExecutorExtension +from colcon_parallel_executor.resource_guard.worker_limiter import \ + WorkerLimiterGuard import pytest ran_jobs = [] +@pytest.fixture(autouse=True, scope='module') +def default_execution_policies(): + """Mock extension discovery to isolate executor tests.""" + guard = WorkerLimiterGuard() + with patch( + 'colcon_parallel_executor.resource_guard.' + 'instantiate_extensions', + return_value={'default_worker_limiter': guard} + ): + yield + + class Job1(Job): def __init__(self, identifier='job1'): diff --git a/test/test_resource_guard.py b/test/test_resource_guard.py new file mode 100644 index 0000000..d251d3b --- /dev/null +++ b/test/test_resource_guard.py @@ -0,0 +1,198 @@ +# Copyright 2026 Open Source Robotics Foundation, Inc. +# Licensed under the Apache License, Version 2.0 + +import argparse +import asyncio +from collections import OrderedDict +from types import SimpleNamespace +from unittest.mock import patch + +from colcon_core.executor import Job +from colcon_parallel_executor.executor.parallel \ + import ParallelExecutorExtension +from colcon_parallel_executor.resource_guard import \ + add_resource_guard_arguments +from colcon_parallel_executor.resource_guard import \ + initialize_resource_guard_extensions +from colcon_parallel_executor.resource_guard import \ + ResourceGuardExtensionPoint +from colcon_parallel_executor.resource_guard.worker_limiter import \ + WorkerLimiterGuard +from run_until_complete import run_until_complete + +ran_jobs = [] + + +class Job1(Job): + + def __init__(self, identifier='job1', dependencies=None): + super().__init__( + identifier=identifier, + dependencies=dependencies or set(), + task=None, + task_context=None + ) + + async def __call__(self, *args, **kwargs): + ran_jobs.append(self.identifier) + + +def test_worker_limiter_guard(): + """Verify WorkerLimiterGuard correctly limits concurrency.""" + guard_provider = WorkerLimiterGuard() + args = SimpleNamespace(parallel_workers=2) + run_until_complete(guard_provider.initialize(args)) + + # Returns an asyncio.Semaphore instance + sem = guard_provider.get_guard(Job1('job1')) + assert isinstance(sem, asyncio.Semaphore) + assert sem._value == 2 + + # If parallel_workers is 0, no semaphore (None) is returned + args_zero = SimpleNamespace(parallel_workers=0) + run_until_complete(guard_provider.initialize(args_zero)) + sem = guard_provider.get_guard(Job1('job1')) + assert sem is None + + +def test_worker_limiter_guard_edge_cases(): + """Verify WorkerLimiterGuard argument parsing and empty edge cases.""" + guard_provider = WorkerLimiterGuard() + # add_arguments should return None + parser = argparse.ArgumentParser() + assert guard_provider.add_arguments(parser=parser) is None + + # Missing parallel_workers attribute should fall back to 0 (no limit) + run_until_complete(guard_provider.initialize(SimpleNamespace())) + sem = guard_provider.get_guard(Job1('job1')) + assert sem is None + + +def test_parallel_no_extensions(): + """Verify execution is unrestricted when no extensions are enabled.""" + extension = ParallelExecutorExtension() + args = SimpleNamespace(parallel_workers=2) + jobs = OrderedDict() + jobs['one'] = Job1('job1') + jobs['two'] = Job1('job2') + + with patch( + 'colcon_parallel_executor.resource_guard.' + 'instantiate_extensions', + return_value={} + ): + rc = extension.execute(args, jobs) + assert rc == 0 + assert set(ran_jobs) == {'job1', 'job2'} + ran_jobs.clear() + + +def test_custom_resource_guard(): + """Verify a custom capacity limiter can withhold jobs via async context.""" + class CustomThrottler(ResourceGuardExtensionPoint): + """Custom limiter that blocks the second job until first completes.""" + + PRIORITY = 1500 + + async def initialize(self, args): + self.allow_second = False + self.first_job_future = asyncio.Future() + + def get_guard(self, job): + return CustomGuardContext(self, job.identifier) + + class CustomGuardContext: + + def __init__(self, provider, package_name): + self.provider = provider + self.package_name = package_name + + async def __aenter__(self): + if self.package_name == 'two': + # Block second job until the first job finishes and + # resolves the future + await self.provider.first_job_future + return self + + async def __aexit__(self, exc_type, exc_val, exc_tb): + if self.package_name == 'one': + # First job finished, wake up the second job + self.provider.allow_second = True + self.provider.first_job_future.set_result(True) + + extension = ParallelExecutorExtension() + args = SimpleNamespace(parallel_workers=2) + jobs = OrderedDict() + jobs['one'] = Job1('job1') + jobs['two'] = Job1('job2') + + throttler = CustomThrottler() + + with patch( + 'colcon_parallel_executor.resource_guard.' + 'instantiate_extensions', + return_value={'custom_throttler': throttler} + ): + rc = extension.execute(args, jobs) + + assert rc == 0 + assert set(ran_jobs) == {'job1', 'job2'} + ran_jobs.clear() + + +def test_initialize_resource_guard_extensions(): + """Verify initialize_resource_guard_extensions initializes correctly.""" + class MockGuard(ResourceGuardExtensionPoint): + + def __init__(self): + super().__init__() + self.called = False + + async def initialize(self, args): + self.called = True + + mock_guard = MockGuard() + extensions = OrderedDict([ + (100, {'mock_guard': mock_guard}) + ]) + + args = SimpleNamespace() + guards = run_until_complete( + initialize_resource_guard_extensions(args, extensions=extensions)) + assert guards == [mock_guard] + assert mock_guard.called + + +def test_initialize_resource_guard_extensions_exception(): + """Verify initialize_resource_guard_extensions handles exceptions.""" + class MockGuard(ResourceGuardExtensionPoint): + + async def initialize(self, args): + raise RuntimeError('init error') + + mock_guard = MockGuard() + extensions = OrderedDict([ + (100, {'mock_guard': mock_guard}) + ]) + + args = SimpleNamespace() + guards = run_until_complete( + initialize_resource_guard_extensions(args, extensions=extensions)) + assert guards == [] + + +def test_add_resource_guard_arguments_exception(): + """Verify add_resource_guard_arguments handles exceptions gracefully.""" + class MockGuard(ResourceGuardExtensionPoint): + + def add_arguments(self, *, parser): + raise RuntimeError('parser error') + + mock_guard = MockGuard() + extensions = OrderedDict([ + (100, {'mock_guard': mock_guard}) + ]) + + parser = object() + # This should handle the exception internally and not raise + add_resource_guard_arguments(parser, extensions=extensions)