diff --git a/changelog/879.fix.md b/changelog/879.fix.md new file mode 100644 index 000000000..defb90f94 --- /dev/null +++ b/changelog/879.fix.md @@ -0,0 +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. 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..c549fb6cc 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,11 @@ def join(self, timeout: float) -> None: # 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 + self._fail_outstanding(results, t) + self.pool.shutdown(wait=False, cancel_futures=True) + raise finally: t.close() @@ -304,7 +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()} 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/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)." ) 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..6519a2c34 100644 --- a/packages/climate-ref/src/climate_ref/migrations/env.py +++ b/packages/climate-ref/src/climate_ref/migrations/env.py @@ -116,26 +116,32 @@ def run_migrations_online() -> None: """ connectable = config.attributes.get("connection", None) + # A database opened here is ours to close. + 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..42ab811dd 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 + + @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. + """ + 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=ExceptionCls("ingest broke")) + failure_spy = mocker.patch("climate_ref.executor.result_handling.process_result") + + with pytest.raises(ExceptionCls, 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