Skip to content

climate_ref.executor.local #

ExecutionFuture #

A container linking a submitted future to its execution metadata.

Source code in packages/climate-ref/src/climate_ref/executor/result_handling.py
@define
class ExecutionFuture:
    """A container linking a submitted future to its execution metadata."""

    future: Future[ExecutionResult]
    """
    The future representing the asynchronous execution of this task.

    This future has a base-class of ``concurrent.futures.Future``,
    but the concrete class of the instance will depend on the executor.
    """

    definition: ExecutionDefinition
    """
    The execution definition associated with this future.
    """

    execution_id: int | None = None
    """
    The ID of the execution associated with this future, or ``None`` if not yet assigned.
    """

    submitted_at: float = 0.0
    """
    Wall-clock time (``time.time()``) at which this future was submitted to the executor.
    """

    started_at: float | None = None
    """
    Wall-clock time (``time.time()``) at which a worker was first observed running this future,
    or ``None`` while it is still queued.

    The per-task timeout is budgeted against this rather than ``submitted_at``
    so that time spent waiting in the pool queue is not counted as execution time.
    """

    def __repr__(self) -> str:
        return f"<ExecutionFuture {self.definition.execution_slug()} execution_id={self.execution_id}>"

    @property
    def queue_seconds(self) -> float | None:
        """
        Time spent waiting between submission and the start of execution, in seconds.

        ``None`` when the start of execution could not be established.
        """
        if self.started_at is None:
            return None
        return max(0.0, self.started_at - self.submitted_at)

    def infer_started_at(self, result: ExecutionResult, observed_at: float) -> None:
        """
        Place ``started_at`` using the worker's own wall time, if a poll never observed the start.

        A task submitted and finished between two polls is never seen running,
        which would leave the queue latency unrecorded for every short execution.
        The worker measured how long it ran,
        so subtracting that from the moment completion was observed puts the start
        within one poll interval of the truth.

        Parameters
        ----------
        result
            The finished result, whose ``resource_usage`` carries the measured wall time.
        observed_at
            Wall-clock time (``time.time()``) at which the future was seen to be done.
        """
        if self.started_at is not None:
            return
        usage = result.resource_usage
        if usage is None:
            return
        self.started_at = max(self.submitted_at, observed_at - usage.wall_seconds)

definition instance-attribute #

The execution definition associated with this future.

execution_id = None class-attribute instance-attribute #

The ID of the execution associated with this future, or None if not yet assigned.

future instance-attribute #

The future representing the asynchronous execution of this task.

This future has a base-class of concurrent.futures.Future, but the concrete class of the instance will depend on the executor.

queue_seconds property #

Time spent waiting between submission and the start of execution, in seconds.

None when the start of execution could not be established.

started_at = None class-attribute instance-attribute #

Wall-clock time (time.time()) at which a worker was first observed running this future, or None while it is still queued.

The per-task timeout is budgeted against this rather than submitted_at so that time spent waiting in the pool queue is not counted as execution time.

submitted_at = 0.0 class-attribute instance-attribute #

Wall-clock time (time.time()) at which this future was submitted to the executor.

infer_started_at(result, observed_at) #

Place started_at using the worker's own wall time, if a poll never observed the start.

A task submitted and finished between two polls is never seen running, which would leave the queue latency unrecorded for every short execution. The worker measured how long it ran, so subtracting that from the moment completion was observed puts the start within one poll interval of the truth.

Parameters:

Name Type Description Default
result ExecutionResult

The finished result, whose resource_usage carries the measured wall time.

required
observed_at float

Wall-clock time (time.time()) at which the future was seen to be done.

required
Source code in packages/climate-ref/src/climate_ref/executor/result_handling.py
def infer_started_at(self, result: ExecutionResult, observed_at: float) -> None:
    """
    Place ``started_at`` using the worker's own wall time, if a poll never observed the start.

    A task submitted and finished between two polls is never seen running,
    which would leave the queue latency unrecorded for every short execution.
    The worker measured how long it ran,
    so subtracting that from the moment completion was observed puts the start
    within one poll interval of the truth.

    Parameters
    ----------
    result
        The finished result, whose ``resource_usage`` carries the measured wall time.
    observed_at
        Wall-clock time (``time.time()``) at which the future was seen to be done.
    """
    if self.started_at is not None:
        return
    usage = result.resource_usage
    if usage is None:
        return
    self.started_at = max(self.submitted_at, observed_at - usage.wall_seconds)

LocalExecutor #

Run a diagnostic locally using a process pool.

This performs the diagnostic executions in parallel using different processes. The maximum number of processes is determined by the n parameter and default to the number of CPUs.

This executor is the default executor and is used when no other executor is specified.

Source code in packages/climate-ref/src/climate_ref/executor/local.py
class LocalExecutor:
    """
    Run a diagnostic locally using a process pool.

    This performs the diagnostic executions in parallel using different processes.
    The maximum number of processes is determined by the `n` parameter and default to the number of CPUs.

    This executor is the default executor and is used when no other executor is specified.
    """

    name = "local"

    # Results are copied to the results directory and ingested only during ``join``,
    # so skipping ``join`` (queue-and-exit) would orphan completed work in scratch.
    collects_results_on_join = True

    # Diagnostics run on the solving host, so a health check there reflects where they run.
    runs_on_solve_host = True

    def __init__(
        self,
        *,
        database: Database | None = None,
        config: Config | None = None,
        n: int | None = None,
        pool: concurrent.futures.Executor | None = None,
        task_timeout: float = 6 * 60 * 60,
        **kwargs: Any,
    ) -> None:
        if config is None:
            config = Config.default()
        if database is None:
            database = Database.from_config(config, run_migrations=False)
            # A database we opened is ours to close, so it does not outlive this executor.
            weakref.finalize(self, database.close)
        self.n = n

        self.database = database
        self.config = config

        # Per-task wall-clock budget (default 6 hours, matching the Celery task_time_limit).
        # This budgets *execution* time measured from when a worker starts the task (instead of submission),
        # Diagnostics that hang past this while running are considered lost.
        # Set to ``0`` to disable.
        self.task_timeout = task_timeout

        if pool is not None:
            self.pool = pool
            if config.executor.measure_resources and getattr(pool, "_max_tasks_per_child", None) != 1:
                logger.warning(
                    "The supplied pool does not recycle its worker after every task, "
                    "so a recorded peak memory can include the footprint of an earlier execution "
                    "in the same worker"
                )
        else:
            self.pool = ProcessPoolExecutor(
                max_workers=n,
                initializer=_process_initialiser,
                # Explicitly set the context to "spawn" to avoid issues with hanging on MacOS
                mp_context=multiprocessing.get_context("spawn"),
                # A fresh process per execution for the right rusage peak-memory
                max_tasks_per_child=1,
            )

        # Every worker shares this process's cgroup,
        # so only a single-worker pool can attribute a cgroup reading to one execution.
        self._exclusive = _pool_is_serial(self.pool)
        self._results: list[ExecutionFuture] = []

    def run(
        self,
        definition: ExecutionDefinition,
        execution: Execution | None = None,
    ) -> None:
        """
        Run a diagnostic in process

        Parameters
        ----------
        definition
            A description of the information needed for this execution of the diagnostic
        execution
            A database model representing the execution of the diagnostic.
            If provided, the result will be updated in the database when completed.
        """
        submitted_at = time.time()
        future = self.pool.submit(
            _process_run,
            definition=definition,
            log_level=self.config.log_level,
            measure=self.config.executor.measure_resources,
            exclusive=self._exclusive,
        )
        self._results.append(
            ExecutionFuture(
                future=future,
                definition=definition,
                execution_id=execution.id if execution else None,
                submitted_at=submitted_at,
            )
        )

    def join(self, timeout: float) -> None:
        """
        Wait for all diagnostics to finish

        This will block until all diagnostics have completed or the timeout is reached.

        Each individual execution is also bounded by ``self.task_timeout``,
        measured from when a worker starts running it,
        so that a hung diagnostic cannot block the pool indefinitely.
        Tasks still waiting in the pool queue are never culled by this budget,
        only the overall ``timeout`` applies to them.

        Outstanding executions are always marked as failed-retryable before this method returns or raises,
        so the next solve can pick them up rather than seeing them stuck with ``successful=None``.

        Parameters
        ----------
        timeout
            Overall wall-clock timeout in seconds for the whole join.
            A non-positive value (``<= 0``) waits indefinitely.

        Raises
        ------
        TimeoutError
            If a positive overall timeout is reached
        """
        start_time = time.time()
        refresh_time = 0.5  # Time to wait between checking for completed tasks in seconds

        results = self._results
        t = LoggedProgress(
            total=len(results),
            desc="Waiting for executions to complete",
            unit="execution",
        )

        try:
            while results:
                now = time.time()

                # Iterate over a copy of the list and remove finished tasks
                for result in results[:]:
                    # Record when a worker first picks the task up.
                    # ``submitted_at`` only marks enqueue time
                    if result.started_at is None and result.future.running():
                        result.started_at = now

                    if result.future.done():
                        try:
                            execution_result = result.future.result(timeout=0)
                        except Exception as e:
                            # Something went wrong when attempting to run the execution
                            # This is likely a failure in the execution itself not the diagnostic
                            fail_future(self.database, self.config, result, detail=f"execution raised {e!r}")
                            results.remove(result)
                            raise ExecutionError(
                                f"Failed to execute {result.definition.execution_slug()!r}"
                            ) from e

                        assert execution_result is not None, "Execution result should not be None"
                        assert isinstance(execution_result, ExecutionResult), (
                            "Execution result should be of type ExecutionResult"
                        )

                        result.infer_started_at(execution_result, now)

                        # Process the result in the main process
                        # The results should be committed after each execution
                        with self.database.session.begin():
                            execution = (
                                self.database.session.get(Execution, result.execution_id)
                                if result.execution_id
                                else None
                            )
                            process_result(
                                self.config,
                                self.database,
                                execution_result,
                                execution,
                                queue_seconds=result.queue_seconds,
                            )
                        logger.debug(f"Execution completed: {result}")
                        t.update(n=1)
                        results.remove(result)
                        continue

                    # Per-task timeout: a runaway *running* diagnostic cannot block the pool forever.
                    # Cancel its future and mark the row failed-retryable.
                    # Tasks still queued (``started_at is None``) are left alone here,
                    # but may be reaped by the overall timeout.
                    if (
                        self.task_timeout > 0
                        and result.started_at is not None
                        and now - result.started_at > self.task_timeout
                    ):
                        logger.error(
                            f"Execution {result.definition.execution_slug()!r} exceeded per-task "
                            f"timeout of {self.task_timeout}s; marking failed-retryable"
                        )
                        result.future.cancel()
                        fail_future(
                            self.database,
                            self.config,
                            result,
                            detail=f"exceeded per-task timeout of {self.task_timeout}s",
                        )
                        t.update(n=1)
                        results.remove(result)

                # Break early to avoid waiting for one more sleep cycle
                if len(results) == 0:
                    break

                elapsed_time = time.time() - start_time

                # ``timeout <= 0`` means wait indefinitely (no overall deadline).
                if timeout > 0 and elapsed_time > timeout:
                    fail_outstanding(
                        self.database,
                        self.config,
                        results,
                        t,
                        detail=f"overall timeout of {timeout}s reached",
                    )
                    self.pool.shutdown(wait=False, cancel_futures=True)
                    raise TimeoutError("Not all tasks completed within the specified timeout")

                # Wait for a short time before checking for completed executions
                time.sleep(refresh_time)
        except BaseException:
            # Handle Ctrl-C and SystemExit exceptions to mark in-flight executions as retryable
            fail_outstanding(
                self.database,
                self.config,
                results,
                t,
                detail="execution was not collected (interrupted or crashed)",
            )
            self.pool.shutdown(wait=False, cancel_futures=True)
            raise
        finally:
            t.close()

        logger.info("All executions finished")

join(timeout) #

Wait for all diagnostics to finish

This will block until all diagnostics have completed or the timeout is reached.

Each individual execution is also bounded by self.task_timeout, measured from when a worker starts running it, so that a hung diagnostic cannot block the pool indefinitely. Tasks still waiting in the pool queue are never culled by this budget, only the overall timeout applies to them.

Outstanding executions are always marked as failed-retryable before this method returns or raises, so the next solve can pick them up rather than seeing them stuck with successful=None.

Parameters:

Name Type Description Default
timeout float

Overall wall-clock timeout in seconds for the whole join. A non-positive value (<= 0) waits indefinitely.

required

Raises:

Type Description
TimeoutError

If a positive overall timeout is reached

Source code in packages/climate-ref/src/climate_ref/executor/local.py
def join(self, timeout: float) -> None:
    """
    Wait for all diagnostics to finish

    This will block until all diagnostics have completed or the timeout is reached.

    Each individual execution is also bounded by ``self.task_timeout``,
    measured from when a worker starts running it,
    so that a hung diagnostic cannot block the pool indefinitely.
    Tasks still waiting in the pool queue are never culled by this budget,
    only the overall ``timeout`` applies to them.

    Outstanding executions are always marked as failed-retryable before this method returns or raises,
    so the next solve can pick them up rather than seeing them stuck with ``successful=None``.

    Parameters
    ----------
    timeout
        Overall wall-clock timeout in seconds for the whole join.
        A non-positive value (``<= 0``) waits indefinitely.

    Raises
    ------
    TimeoutError
        If a positive overall timeout is reached
    """
    start_time = time.time()
    refresh_time = 0.5  # Time to wait between checking for completed tasks in seconds

    results = self._results
    t = LoggedProgress(
        total=len(results),
        desc="Waiting for executions to complete",
        unit="execution",
    )

    try:
        while results:
            now = time.time()

            # Iterate over a copy of the list and remove finished tasks
            for result in results[:]:
                # Record when a worker first picks the task up.
                # ``submitted_at`` only marks enqueue time
                if result.started_at is None and result.future.running():
                    result.started_at = now

                if result.future.done():
                    try:
                        execution_result = result.future.result(timeout=0)
                    except Exception as e:
                        # Something went wrong when attempting to run the execution
                        # This is likely a failure in the execution itself not the diagnostic
                        fail_future(self.database, self.config, result, detail=f"execution raised {e!r}")
                        results.remove(result)
                        raise ExecutionError(
                            f"Failed to execute {result.definition.execution_slug()!r}"
                        ) from e

                    assert execution_result is not None, "Execution result should not be None"
                    assert isinstance(execution_result, ExecutionResult), (
                        "Execution result should be of type ExecutionResult"
                    )

                    result.infer_started_at(execution_result, now)

                    # Process the result in the main process
                    # The results should be committed after each execution
                    with self.database.session.begin():
                        execution = (
                            self.database.session.get(Execution, result.execution_id)
                            if result.execution_id
                            else None
                        )
                        process_result(
                            self.config,
                            self.database,
                            execution_result,
                            execution,
                            queue_seconds=result.queue_seconds,
                        )
                    logger.debug(f"Execution completed: {result}")
                    t.update(n=1)
                    results.remove(result)
                    continue

                # Per-task timeout: a runaway *running* diagnostic cannot block the pool forever.
                # Cancel its future and mark the row failed-retryable.
                # Tasks still queued (``started_at is None``) are left alone here,
                # but may be reaped by the overall timeout.
                if (
                    self.task_timeout > 0
                    and result.started_at is not None
                    and now - result.started_at > self.task_timeout
                ):
                    logger.error(
                        f"Execution {result.definition.execution_slug()!r} exceeded per-task "
                        f"timeout of {self.task_timeout}s; marking failed-retryable"
                    )
                    result.future.cancel()
                    fail_future(
                        self.database,
                        self.config,
                        result,
                        detail=f"exceeded per-task timeout of {self.task_timeout}s",
                    )
                    t.update(n=1)
                    results.remove(result)

            # Break early to avoid waiting for one more sleep cycle
            if len(results) == 0:
                break

            elapsed_time = time.time() - start_time

            # ``timeout <= 0`` means wait indefinitely (no overall deadline).
            if timeout > 0 and elapsed_time > timeout:
                fail_outstanding(
                    self.database,
                    self.config,
                    results,
                    t,
                    detail=f"overall timeout of {timeout}s reached",
                )
                self.pool.shutdown(wait=False, cancel_futures=True)
                raise TimeoutError("Not all tasks completed within the specified timeout")

            # Wait for a short time before checking for completed executions
            time.sleep(refresh_time)
    except BaseException:
        # Handle Ctrl-C and SystemExit exceptions to mark in-flight executions as retryable
        fail_outstanding(
            self.database,
            self.config,
            results,
            t,
            detail="execution was not collected (interrupted or crashed)",
        )
        self.pool.shutdown(wait=False, cancel_futures=True)
        raise
    finally:
        t.close()

    logger.info("All executions finished")

run(definition, execution=None) #

Run a diagnostic in process

Parameters:

Name Type Description Default
definition ExecutionDefinition

A description of the information needed for this execution of the diagnostic

required
execution Execution | None

A database model representing the execution of the diagnostic. If provided, the result will be updated in the database when completed.

None
Source code in packages/climate-ref/src/climate_ref/executor/local.py
def run(
    self,
    definition: ExecutionDefinition,
    execution: Execution | None = None,
) -> None:
    """
    Run a diagnostic in process

    Parameters
    ----------
    definition
        A description of the information needed for this execution of the diagnostic
    execution
        A database model representing the execution of the diagnostic.
        If provided, the result will be updated in the database when completed.
    """
    submitted_at = time.time()
    future = self.pool.submit(
        _process_run,
        definition=definition,
        log_level=self.config.log_level,
        measure=self.config.executor.measure_resources,
        exclusive=self._exclusive,
    )
    self._results.append(
        ExecutionFuture(
            future=future,
            definition=definition,
            execution_id=execution.id if execution else None,
            submitted_at=submitted_at,
        )
    )

process_result(config, database, result, execution, *, queue_seconds=None, detail=None) #

Process the result of a diagnostic execution, persisting outcome to the DB.

queue_seconds is supplied by the submitting executor, because only the parent knows when the task was submitted. detail says why a failed execution failed and is written to its results log.

Source code in packages/climate-ref/src/climate_ref/executor/result_handling.py
def process_result(  # noqa: PLR0913
    config: "Config",
    database: Database,
    result: ExecutionResult,
    execution: Execution | None,
    *,
    queue_seconds: float | None = None,
    detail: str | None = None,
) -> None:
    """Process the result of a diagnostic execution, persisting outcome to the DB.

    ``queue_seconds`` is supplied by the submitting executor,
    because only the parent knows when the task was submitted.
    ``detail`` says why a failed execution failed and is written to its results log.
    """
    if not result.successful:
        if execution is not None:  # pragma: no branch
            info_msg = (
                f"\nAdditional information about this execution can be viewed using: "
                f"ref executions inspect {execution.execution_group_id}"
            )
        else:
            info_msg = ""
        logger.error(f"Error running {result.definition.execution_slug()}. {info_msg}")

    if execution:
        handle_execution_result(
            config, database, execution, result, queue_seconds=queue_seconds, detail=detail
        )