From a506ba5ce0c5c0e1b110c3bcc9e2f99672fdb0e3 Mon Sep 17 00:00:00 2001 From: Jared Lewis Date: Fri, 14 Aug 2026 14:10:02 +1000 Subject: [PATCH 1/3] fix: close database connections on every entrypoint Celery workers opened a database per task and never closed it, so a long-running worker churned a connection for every result it handled. Workers now open one database per process and reuse it, keyed on the pid so a forked child never uses connections it inherited. Also closes the databases the executors open for themselves, and the one the standalone alembic environment opens. An execution still in flight when `LocalExecutor.join` raises is now marked failed-retryable, which is what the docstring already promised. Only the timeout path did this before, so an ingestion failure left the remaining rows stuck with `successful=None` until the stale-execution reaper caught them hours later. --- changelog/879.fix.md | 4 ++ .../src/climate_ref_celery/worker_tasks.py | 62 ++++++++++++++++++- .../tests/unit/test_worker_tasks.py | 28 ++++++++- .../src/climate_ref/executor/hpc.py | 6 +- .../src/climate_ref/executor/local.py | 13 +++- .../src/climate_ref/executor/synchronous.py | 3 + .../src/climate_ref/migrations/env.py | 39 +++++++----- .../unit/executor/test_local_executor.py | 35 +++++++++++ 8 files changed, 169 insertions(+), 21 deletions(-) create mode 100644 changelog/879.fix.md diff --git a/changelog/879.fix.md b/changelog/879.fix.md new file mode 100644 index 000000000..005397bd3 --- /dev/null +++ b/changelog/879.fix.md @@ -0,0 +1,4 @@ +Celery workers now reuse a single database connection for the lifetime of a worker process +instead of opening a new one for every result they handle. +An execution that is still in flight when `ref solve` fails is now recorded as failed-retryable, +so the next solve picks it up rather than waiting for the stale-execution reaper. diff --git a/packages/climate-ref-celery/src/climate_ref_celery/worker_tasks.py b/packages/climate-ref-celery/src/climate_ref_celery/worker_tasks.py index c45ce2f39..381a9d47f 100644 --- a/packages/climate-ref-celery/src/climate_ref_celery/worker_tasks.py +++ b/packages/climate-ref-celery/src/climate_ref_celery/worker_tasks.py @@ -2,7 +2,10 @@ Celery worker tasks for handling diagnostic execution executions. """ +import os + from celery import current_app +from celery.signals import worker_process_shutdown from loguru import logger from climate_ref.config import Config @@ -12,6 +15,61 @@ from climate_ref_core.diagnostics import ExecutionResult +class _WorkerDatabase: + """ + The database a worker process uses, opened once and reused across tasks + + A worker handles many tasks, so opening an engine per task churns a connection each time. + The record of which process opened it means a forked child opens its own + rather than using connections it inherited, which are not safe to share across processes. + """ + + def __init__(self) -> None: + self._opened: tuple[int, str, Database] | None = None + + def get(self, config: Config) -> Database: + """ + Open this process's database, or return the one already open for this URL + + Parameters + ---------- + config + REF configuration describing where the database lives. + + Returns + ------- + : + The database for this process. + """ + key = (os.getpid(), config.db.database_url) + if self._opened is not None and self._opened[:2] == key: + return self._opened[2] + + self.close() + self._opened = (*key, Database.from_config(config, run_migrations=False)) + return self._opened[2] + + def close(self) -> None: + """ + Release the connections held by this process + + A database opened by another process is dropped rather than closed, + because its connections belong to the process that opened them. + """ + if self._opened is not None and self._opened[0] == os.getpid(): + self._opened[2].close() + self._opened = None + + +_worker_database = _WorkerDatabase() + + +@worker_process_shutdown.connect +def _close_worker_database(**kwargs: object) -> None: # pragma: no cover + """Release this process's database connections as the worker process goes away.""" + _worker_database.close() + + @current_app.task(max_retries=0) def handle_result(result: ExecutionResult, execution_id: int) -> None: """ @@ -29,7 +87,7 @@ def handle_result(result: ExecutionResult, execution_id: int) -> None: logger.info(f"Handling result for execution {execution_id} + {result}") config = Config.default() - db = Database.from_config(config, run_migrations=False) + db = _worker_database.get(config) with db.session.begin(): execution = db.session.get(Execution, execution_id) @@ -69,7 +127,7 @@ def handle_failure(task_id: str, execution_id: int) -> None: ) config = Config.default() - db = Database.from_config(config, run_migrations=False) + db = _worker_database.get(config) with db.session.begin(): execution = db.session.get(Execution, execution_id) diff --git a/packages/climate-ref-celery/tests/unit/test_worker_tasks.py b/packages/climate-ref-celery/tests/unit/test_worker_tasks.py index b80902c27..0bf86ae62 100644 --- a/packages/climate-ref-celery/tests/unit/test_worker_tasks.py +++ b/packages/climate-ref-celery/tests/unit/test_worker_tasks.py @@ -1,4 +1,6 @@ -from climate_ref_celery.worker_tasks import handle_failure, handle_result +import pytest +from climate_ref_celery import worker_tasks +from climate_ref_celery.worker_tasks import _worker_database, handle_failure, handle_result from climate_ref_example import provider as example_provider from climate_ref.database import Database @@ -6,6 +8,30 @@ from climate_ref.provider_registry import _register_provider +@pytest.fixture(autouse=True) +def _clear_worker_database(): + worker_tasks._worker_database.close() + yield + worker_tasks._worker_database.close() + + +def test_worker_database_is_reused_across_tasks(config): + first = _worker_database.get(config) + second = _worker_database.get(config) + + assert second is first + + +def test_worker_database_reopens_for_a_different_url(config, tmp_path): + first = _worker_database.get(config) + + config.db.database_url = f"sqlite:///{tmp_path}/other.db" + second = _worker_database.get(config) + + assert second is not first + assert second.url == config.db.database_url + + def test_worker_task(mocker, config): mock_handle_result = mocker.patch("climate_ref_celery.worker_tasks.handle_execution_result") db = Database.from_config(config, run_migrations=True) diff --git a/packages/climate-ref/src/climate_ref/executor/hpc.py b/packages/climate-ref/src/climate_ref/executor/hpc.py index 48b0bffd5..22ed7800f 100644 --- a/packages/climate-ref/src/climate_ref/executor/hpc.py +++ b/packages/climate-ref/src/climate_ref/executor/hpc.py @@ -23,6 +23,7 @@ import re import resource import time +import weakref from collections.abc import Callable from typing import Annotated, Any, Literal, TypeVar, cast @@ -215,7 +216,10 @@ def __init__( **executor_config: str | float | int, ) -> None: config = config or Config.default() - database = database or Database.from_config(config, run_migrations=False) + 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.config = config self.database = database diff --git a/packages/climate-ref/src/climate_ref/executor/local.py b/packages/climate-ref/src/climate_ref/executor/local.py index a8108f43c..1c656ae8e 100644 --- a/packages/climate-ref/src/climate_ref/executor/local.py +++ b/packages/climate-ref/src/climate_ref/executor/local.py @@ -1,6 +1,7 @@ import concurrent.futures import multiprocessing import time +import weakref from concurrent.futures import ProcessPoolExecutor from typing import Any @@ -102,6 +103,8 @@ def __init__( 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 @@ -287,6 +290,13 @@ def join(self, timeout: float) -> None: # Wait for a short time before checking for completed executions time.sleep(refresh_time) + except BaseException: + # Whatever went wrong, the executions still in flight are not going to be collected, + # so record them as retryable rather than leaving them stuck with ``successful=None``. + # ``_fail_outstanding`` empties the list, so the timeout path does not do this twice. + self._fail_outstanding(results, t) + self.pool.shutdown(wait=False, cancel_futures=True) + raise finally: t.close() @@ -304,7 +314,8 @@ def _mark_failed(self, result: ExecutionFuture, *, retryable: bool) -> None: def _fail_outstanding(self, results: list[ExecutionFuture], progress: Any) -> None: for outstanding in list(results): logger.warning( - f"Execution {outstanding.definition.execution_slug()} did not complete within the timeout" + f"Execution {outstanding.definition.execution_slug()} was not collected; " + "marking it failed-retryable" ) self._mark_failed(outstanding, retryable=True) progress.update(n=1) diff --git a/packages/climate-ref/src/climate_ref/executor/synchronous.py b/packages/climate-ref/src/climate_ref/executor/synchronous.py index b7851d5cb..ed78d3af0 100644 --- a/packages/climate-ref/src/climate_ref/executor/synchronous.py +++ b/packages/climate-ref/src/climate_ref/executor/synchronous.py @@ -1,3 +1,4 @@ +import weakref from typing import Any from loguru import logger @@ -30,6 +31,8 @@ def __init__( 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.database = database self.config = config diff --git a/packages/climate-ref/src/climate_ref/migrations/env.py b/packages/climate-ref/src/climate_ref/migrations/env.py index e3bc1982c..35bd72f86 100644 --- a/packages/climate-ref/src/climate_ref/migrations/env.py +++ b/packages/climate-ref/src/climate_ref/migrations/env.py @@ -116,26 +116,33 @@ def run_migrations_online() -> None: """ connectable = config.attributes.get("connection", None) + # A database opened here is ours to close. One passed in belongs to the caller, + # which keeps using it after the migration returns. + db = None if connectable is None: db = Database.from_config(ref_config, run_migrations=False) connectable = db._engine - with connectable.connect() as connection: - context.configure( - connection=connection, - target_metadata=target_metadata, - render_as_batch=True, - include_object=include_object, - ) - - with context.begin_transaction(): - context.run_migrations() - - # Set up the Operations context - # This is needed to alter the tables - with op.Operations.context(context.get_context()): # type: ignore - _add_dimension_columns(connection, "metric_value", MetricValue) - _add_dimension_columns(connection, "execution_output", ExecutionOutput) + try: + with connectable.connect() as connection: + context.configure( + connection=connection, + target_metadata=target_metadata, + render_as_batch=True, + include_object=include_object, + ) + + with context.begin_transaction(): + context.run_migrations() + + # Set up the Operations context + # This is needed to alter the tables + with op.Operations.context(context.get_context()): # type: ignore + _add_dimension_columns(connection, "metric_value", MetricValue) + _add_dimension_columns(connection, "execution_output", ExecutionOutput) + finally: + if db is not None: + db.close() if context.is_offline_mode(): diff --git a/packages/climate-ref/tests/unit/executor/test_local_executor.py b/packages/climate-ref/tests/unit/executor/test_local_executor.py index d2cef189b..19bf7db3b 100644 --- a/packages/climate-ref/tests/unit/executor/test_local_executor.py +++ b/packages/climate-ref/tests/unit/executor/test_local_executor.py @@ -203,3 +203,38 @@ def test_join_exception(self, metric_definition, mocker): forwarded = process_spy.call_args.args[2] assert isinstance(forwarded, ExecutionResult) assert forwarded.retryable is True + + def test_join_marks_outstanding_when_ingestion_fails(self, metric_definition, mocker): + """ + An execution still in flight when ``join`` raises is recorded as retryable. + + Otherwise it keeps ``successful=None`` and only the stale-execution reaper recovers it, + hours later. + """ + executor = LocalExecutor(n=1) + completed = Future() + completed.set_result( + ExecutionResult( + definition=metric_definition, + successful=True, + output_bundle_filename=None, + metric_bundle_filename=None, + ) + ) + executor._results = [ + ExecutionFuture(completed, definition=metric_definition, execution_id=None), + ExecutionFuture(Future(), definition=metric_definition, execution_id=None), + ] + + mocker.patch("climate_ref.executor.local.process_result", side_effect=ValueError("ingest broke")) + failure_spy = mocker.patch("climate_ref.executor.result_handling.process_result") + + with pytest.raises(ValueError, match="ingest broke"): + executor.join(0) + + # Both are recorded: the completed one raised before it could be ingested. + assert len(executor._results) == 0 + assert failure_spy.call_count == 2 + for call in failure_spy.call_args_list: + assert call.args[2].successful is False + assert call.args[2].retryable is True From fdf6eaeaf2193d565a20416cc0cd3300fead717f Mon Sep 17 00:00:00 2001 From: Jared Lewis Date: Fri, 14 Aug 2026 14:29:22 +1000 Subject: [PATCH 2/3] chore: add tests for interrupts --- changelog/879.fix.md | 3 +-- packages/climate-ref/src/climate_ref/executor/local.py | 8 +++----- packages/climate-ref/src/climate_ref/migrations/env.py | 3 +-- .../tests/unit/executor/test_local_executor.py | 10 +++++----- 4 files changed, 10 insertions(+), 14 deletions(-) diff --git a/changelog/879.fix.md b/changelog/879.fix.md index 005397bd3..defb90f94 100644 --- a/changelog/879.fix.md +++ b/changelog/879.fix.md @@ -1,4 +1,3 @@ Celery workers now reuse a single database connection for the lifetime of a worker process instead of opening a new one for every result they handle. -An execution that is still in flight when `ref solve` fails is now recorded as failed-retryable, -so the next solve picks it up rather than waiting for the stale-execution reaper. +An execution that is still in flight when `ref solve` fails is now recorded as failed-retryable. diff --git a/packages/climate-ref/src/climate_ref/executor/local.py b/packages/climate-ref/src/climate_ref/executor/local.py index 1c656ae8e..c5b255f68 100644 --- a/packages/climate-ref/src/climate_ref/executor/local.py +++ b/packages/climate-ref/src/climate_ref/executor/local.py @@ -291,9 +291,7 @@ def join(self, timeout: float) -> None: # Wait for a short time before checking for completed executions time.sleep(refresh_time) except BaseException: - # Whatever went wrong, the executions still in flight are not going to be collected, - # so record them as retryable rather than leaving them stuck with ``successful=None``. - # ``_fail_outstanding`` empties the list, so the timeout path does not do this twice. + # Handle Ctrl-C and SystemExit exceptions to mark in-flight executions as retryable self._fail_outstanding(results, t) self.pool.shutdown(wait=False, cancel_futures=True) raise @@ -314,8 +312,8 @@ def _mark_failed(self, result: ExecutionFuture, *, retryable: bool) -> None: def _fail_outstanding(self, results: list[ExecutionFuture], progress: Any) -> None: for outstanding in list(results): logger.warning( - f"Execution {outstanding.definition.execution_slug()} was not collected; " - "marking it failed-retryable" + f"Execution {outstanding.definition.execution_slug()} was not collected." + "Marking it failed-retryable" ) self._mark_failed(outstanding, retryable=True) progress.update(n=1) diff --git a/packages/climate-ref/src/climate_ref/migrations/env.py b/packages/climate-ref/src/climate_ref/migrations/env.py index 35bd72f86..6519a2c34 100644 --- a/packages/climate-ref/src/climate_ref/migrations/env.py +++ b/packages/climate-ref/src/climate_ref/migrations/env.py @@ -116,8 +116,7 @@ def run_migrations_online() -> None: """ connectable = config.attributes.get("connection", None) - # A database opened here is ours to close. One passed in belongs to the caller, - # which keeps using it after the migration returns. + # A database opened here is ours to close. db = None if connectable is None: db = Database.from_config(ref_config, run_migrations=False) diff --git a/packages/climate-ref/tests/unit/executor/test_local_executor.py b/packages/climate-ref/tests/unit/executor/test_local_executor.py index 19bf7db3b..42ab811dd 100644 --- a/packages/climate-ref/tests/unit/executor/test_local_executor.py +++ b/packages/climate-ref/tests/unit/executor/test_local_executor.py @@ -204,12 +204,12 @@ def test_join_exception(self, metric_definition, mocker): assert isinstance(forwarded, ExecutionResult) assert forwarded.retryable is True - def test_join_marks_outstanding_when_ingestion_fails(self, metric_definition, mocker): + @pytest.mark.parametrize("ExceptionCls", [ValueError, SystemExit, KeyboardInterrupt]) + def test_join_marks_outstanding_when_ingestion_fails(self, ExceptionCls, metric_definition, mocker): """ An execution still in flight when ``join`` raises is recorded as retryable. - Otherwise it keeps ``successful=None`` and only the stale-execution reaper recovers it, - hours later. + Otherwise it keeps ``successful=None`` and only the stale-execution reaper recovers it hours later. """ executor = LocalExecutor(n=1) completed = Future() @@ -226,10 +226,10 @@ def test_join_marks_outstanding_when_ingestion_fails(self, metric_definition, mo ExecutionFuture(Future(), definition=metric_definition, execution_id=None), ] - mocker.patch("climate_ref.executor.local.process_result", side_effect=ValueError("ingest broke")) + mocker.patch("climate_ref.executor.local.process_result", side_effect=ExceptionCls("ingest broke")) failure_spy = mocker.patch("climate_ref.executor.result_handling.process_result") - with pytest.raises(ValueError, match="ingest broke"): + with pytest.raises(ExceptionCls, match="ingest broke"): executor.join(0) # Both are recorded: the completed one raised before it could be ingested. From 121f830f9bed24623925fe9000e4a8b2bbcc4355 Mon Sep 17 00:00:00 2001 From: Jared Lewis Date: Fri, 14 Aug 2026 14:51:44 +1000 Subject: [PATCH 3/3] fix: quieten the log for an execution that never started An execution abandoned before a worker picked it up has no output directory, so the missing log file is expected rather than a system error worth shouting about. Reports it at debug in that case, and keeps the error for an execution that did start. Also adds the missing space between the two halves of the uncollected-execution warning. --- packages/climate-ref/src/climate_ref/executor/local.py | 2 +- .../climate-ref/src/climate_ref/executor/result_handling.py | 6 +++++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/packages/climate-ref/src/climate_ref/executor/local.py b/packages/climate-ref/src/climate_ref/executor/local.py index c5b255f68..c549fb6cc 100644 --- a/packages/climate-ref/src/climate_ref/executor/local.py +++ b/packages/climate-ref/src/climate_ref/executor/local.py @@ -312,7 +312,7 @@ def _mark_failed(self, result: ExecutionFuture, *, retryable: bool) -> None: def _fail_outstanding(self, results: list[ExecutionFuture], progress: Any) -> None: for outstanding in list(results): logger.warning( - f"Execution {outstanding.definition.execution_slug()} was not collected." + f"Execution {outstanding.definition.execution_slug()} was not collected. " "Marking it failed-retryable" ) self._mark_failed(outstanding, retryable=True) diff --git a/packages/climate-ref/src/climate_ref/executor/result_handling.py b/packages/climate-ref/src/climate_ref/executor/result_handling.py index df9bf5e0f..b92e60485 100644 --- a/packages/climate-ref/src/climate_ref/executor/result_handling.py +++ b/packages/climate-ref/src/climate_ref/executor/result_handling.py @@ -495,7 +495,11 @@ def handle_execution_result( # noqa: PLR0913 EXECUTION_LOG_FILENAME, ) except FileNotFoundError: - logger.error( + # An execution abandoned before a worker picked it up has no output directory, + # so a missing log is expected rather than a sign that something went wrong. + started = (config.paths.scratch / execution.output_fragment).exists() + report = logger.error if started else logger.debug + report( f"Could not find log file {EXECUTION_LOG_FILENAME} in scratch directory: {config.paths.scratch}. " f"This is likely a system error (will be retried on next solve)." )