Close database connections on every entrypoint

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.

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)

packages/climate-ref-celery/tests/unit/test_worker_tasks.py

@@ -1,11 +1,37 @@
-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
 from climate_ref.models import Execution, ExecutionGroup
 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)

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

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)

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)."
         )

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

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():

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