climate-ref 0.5.0__py3-none-any.whl → 0.5.1__py3-none-any.whl

climate_ref/executor/local.py

@@ -1,42 +1,128 @@
+import concurrent.futures
+import time
+from concurrent.futures import Future, ProcessPoolExecutor
 from typing import Any
 
+from attrs import define
 from loguru import logger
+from tqdm import tqdm
 
 from climate_ref.config import Config
 from climate_ref.database import Database
-from climate_ref.executor import handle_execution_result
 from climate_ref.models import Execution
-from climate_ref_core.diagnostics import Diagnostic, ExecutionDefinition, ExecutionResult
-from climate_ref_core.logging import redirect_logs
-from climate_ref_core.providers import DiagnosticProvider
+from climate_ref_core.diagnostics import ExecutionDefinition, ExecutionResult
+from climate_ref_core.exceptions import ExecutionError
+from climate_ref_core.executor import execute_locally
+from climate_ref_core.logging import add_log_handler
+
+from .result_handling import handle_execution_result
+
+
+def process_result(
+    config: Config, database: Database, result: ExecutionResult, execution: Execution | None
+) -> None:
+    """
+    Process the result of a diagnostic execution
+
+    Parameters
+    ----------
+    config
+        The configuration object
+    database
+        The database object
+    result
+        The result of the diagnostic execution.
+
+        This could have either been a success or failure.
+    execution
+        A database model representing the execution of the diagnostic.
+    """
+    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.exception(f"Error running {result.definition.execution_slug()}. {info_msg}")
+
+    if execution:
+        handle_execution_result(config, database, execution, result)
+
+
+@define
+class ExecutionFuture:
+    """
+    A container to hold the future and execution definition
+    """
+
+    future: Future[ExecutionResult]
+    definition: ExecutionDefinition
+    execution_id: int | None = None
+
+
+def _process_initialiser() -> None:
+    # Setup the logging for the process
+    # This replaces the loguru default handler
+    try:
+        add_log_handler()
+    except Exception as e:
+        # Don't raise an exception here as that would kill the process pool
+        # We want to log the error and continue
+        logger.error(f"Failed to add log handler: {e}")
+
+
+def _process_run(definition: ExecutionDefinition, log_level: str) -> ExecutionResult:
+    # This is a catch-all for any exceptions that occur in the process
+    try:
+        return execute_locally(definition=definition, log_level=log_level)
+    except Exception:  # pragma: no cover
+        # This isn't expected but if it happens we want to log the error before the process exits
+        logger.exception("Error running diagnostic")
+        # This will kill the process pool
+        raise
 
 
 class LocalExecutor:
     """
-    Run a diagnostic locally, in-process.
+    Run a diagnostic locally using a process pool.
 
-    This is mainly useful for debugging and testing.
-    The production executor will run the diagnostic in a separate process or container,
-    the exact manner of which is yet to be determined.
+    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"
 
     def __init__(
-        self, *, database: Database | None = None, config: Config | None = None, **kwargs: Any
+        self,
+        *,
+        database: Database | None = None,
+        config: Config | None = None,
+        n: int | None = None,
+        pool: concurrent.futures.Executor | None = None,
+        **kwargs: Any,
     ) -> None:
         if config is None:
             config = Config.default()
         if database is None:
             database = Database.from_config(config, run_migrations=False)
+        self.n = n
 
         self.database = database
         self.config = config
 
+        if pool is not None:
+            self.pool = pool
+        else:
+            self.pool = ProcessPoolExecutor(max_workers=n, initializer=_process_initialiser)
+        self._results: list[ExecutionFuture] = []
+
     def run(
         self,
-        provider: DiagnosticProvider,
-        diagnostic: Diagnostic,
         definition: ExecutionDefinition,
         execution: Execution | None = None,
     ) -> None:
@@ -45,45 +131,92 @@ class LocalExecutor:
 
         Parameters
         ----------
-        provider
-            The provider of the diagnostic
-        diagnostic
-            Diagnostic to run
         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.
         """
-        definition.output_directory.mkdir(parents=True, exist_ok=True)
-
-        try:
-            with redirect_logs(definition, self.config.log_level):
-                result = diagnostic.run(definition=definition)
-        except Exception:
-            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.exception(f"Error running diagnostic {diagnostic.slug}. {info_msg}")
-            result = ExecutionResult.build_from_failure(definition)
-
-        if execution:
-            handle_execution_result(self.config, self.database, execution, result)
+        # Submit the execution to the process pool
+        # and track the future so we can wait for it to complete
+        future = self.pool.submit(
+            _process_run,
+            definition=definition,
+            log_level=self.config.log_level,
+        )
+        self._results.append(
+            ExecutionFuture(
+                future=future,
+                definition=definition,
+                execution_id=execution.id if execution else None,
+            )
+        )
 
     def join(self, timeout: float) -> None:
         """
         Wait for all diagnostics to finish
 
-        This returns immediately because the local executor runs diagnostics synchronously.
+        This will block until all diagnostics have completed or the timeout is reached.
+        If the timeout is reached, the method will return and raise an exception.
 
         Parameters
         ----------
         timeout
-            Timeout in seconds (Not used)
+            Timeout in seconds
+
+        Raises
+        ------
+        TimeoutError
+            If the timeout is reached
         """
-        return
+        start_time = time.time()
+        refresh_time = 0.5  # Time to wait between checking for completed tasks in seconds
+
+        results = self._results
+        t = tqdm(total=len(results), desc="Waiting for executions to complete", unit="execution")
+
+        try:
+            while results:
+                # Iterate over a copy of the list and remove finished tasks
+                for result in results[:]:
+                    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
+                            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"
+                        )
+
+                        # 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, result.future.result(), execution)
+                        logger.debug(f"Execution completed: {result}")
+                        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
+
+                if elapsed_time > timeout:
+                    raise TimeoutError("Not all tasks completed within the specified timeout")
+
+                # Wait for a short time before checking for completed executions
+                time.sleep(refresh_time)
+        finally:
+            t.close()

climate_ref/executor/result_handling.py

@@ -0,0 +1,231 @@
+"""
+Execute diagnostics in different environments
+
+We support running diagnostics in different environments, such as locally,
+in a separate process, or in a container.
+These environments are represented by `climate_ref.executor.Executor` classes.
+
+The simplest executor is the `LocalExecutor`, which runs the diagnostic in the same process.
+This is useful for local testing and debugging.
+"""
+
+import pathlib
+import shutil
+from typing import TYPE_CHECKING
+
+from loguru import logger
+from sqlalchemy import insert
+
+from climate_ref.database import Database
+from climate_ref.models import ScalarMetricValue
+from climate_ref.models.execution import Execution, ExecutionOutput, ResultOutputType
+from climate_ref_core.diagnostics import ExecutionResult, ensure_relative_path
+from climate_ref_core.exceptions import ResultValidationError
+from climate_ref_core.logging import EXECUTION_LOG_FILENAME
+from climate_ref_core.pycmec.controlled_vocabulary import CV
+from climate_ref_core.pycmec.metric import CMECMetric
+from climate_ref_core.pycmec.output import CMECOutput, OutputDict
+
+if TYPE_CHECKING:
+    from climate_ref.config import Config
+
+
+def _copy_file_to_results(
+    scratch_directory: pathlib.Path,
+    results_directory: pathlib.Path,
+    fragment: pathlib.Path | str,
+    filename: pathlib.Path | str,
+) -> None:
+    """
+    Copy a file from the scratch directory to the executions directory
+
+    Parameters
+    ----------
+    scratch_directory
+        The directory where the file is currently located
+    results_directory
+        The directory where the file should be copied to
+    fragment
+        The fragment of the executions directory where the file should be copied
+    filename
+        The name of the file to be copied
+    """
+    assert results_directory != scratch_directory
+    input_directory = scratch_directory / fragment
+    output_directory = results_directory / fragment
+
+    filename = ensure_relative_path(filename, input_directory)
+
+    if not (input_directory / filename).exists():
+        raise FileNotFoundError(f"Could not find {filename} in {input_directory}")
+
+    output_filename = output_directory / filename
+    output_filename.parent.mkdir(parents=True, exist_ok=True)
+
+    shutil.copy(input_directory / filename, output_filename)
+
+
+def handle_execution_result(
+    config: "Config",
+    database: Database,
+    execution: Execution,
+    result: "ExecutionResult",
+) -> None:
+    """
+    Handle the result of a diagnostic execution
+
+    This will update the diagnostic execution result with the output of the diagnostic execution.
+    The output will be copied from the scratch directory to the executions directory.
+
+    Parameters
+    ----------
+    config
+        The configuration to use
+    database
+        The active database session to use
+    execution
+        The diagnostic execution result DB object to update
+    result
+        The result of the diagnostic execution, either successful or failed
+    """
+    # Always copy log data
+    _copy_file_to_results(
+        config.paths.scratch,
+        config.paths.results,
+        execution.output_fragment,
+        EXECUTION_LOG_FILENAME,
+    )
+
+    if result.successful and result.metric_bundle_filename is not None:
+        logger.info(f"{execution} successful")
+
+        _copy_file_to_results(
+            config.paths.scratch,
+            config.paths.results,
+            execution.output_fragment,
+            result.metric_bundle_filename,
+        )
+        execution.mark_successful(result.as_relative_path(result.metric_bundle_filename))
+
+        if result.output_bundle_filename:
+            _copy_file_to_results(
+                config.paths.scratch,
+                config.paths.results,
+                execution.output_fragment,
+                result.output_bundle_filename,
+            )
+            _handle_output_bundle(
+                config,
+                database,
+                execution,
+                result.to_output_path(result.output_bundle_filename),
+            )
+
+        cmec_metric_bundle = CMECMetric.load_from_json(result.to_output_path(result.metric_bundle_filename))
+
+        # Check that the diagnostic values conform with the controlled vocabulary
+        try:
+            cv = CV.load_from_file(config.paths.dimensions_cv)
+            cv.validate_metrics(cmec_metric_bundle)
+        except (ResultValidationError, AssertionError):
+            logger.exception("Diagnostic values do not conform with the controlled vocabulary")
+            # TODO: Mark the diagnostic execution result as failed once the CV has stabilised
+            # execution.mark_failed()
+
+        # Perform a bulk insert of scalar values
+        # TODO: The section below will likely fail until we have agreed on a controlled vocabulary
+        # The current implementation will swallow the exception, but display a log message
+        try:
+            # Perform this in a nested transaction to (hopefully) gracefully rollback if something
+            # goes wrong
+            with database.session.begin_nested():
+                database.session.execute(
+                    insert(ScalarMetricValue),
+                    [
+                        {
+                            "execution_id": execution.id,
+                            "value": result.value,
+                            "attributes": result.attributes,
+                            **result.dimensions,
+                        }
+                        for result in cmec_metric_bundle.iter_results()
+                    ],
+                )
+        except Exception:
+            # TODO: Remove once we have settled on a controlled vocabulary
+            logger.exception("Something went wrong when ingesting diagnostic values")
+
+        # TODO Ingest the series values
+
+        # TODO: This should check if the result is the most recent for the execution,
+        # if so then update the dirty fields
+        # i.e. if there are outstanding executions don't make as clean
+        execution.execution_group.dirty = False
+    else:
+        logger.error(f"{execution} failed")
+        execution.mark_failed()
+
+
+def _handle_output_bundle(
+    config: "Config",
+    database: Database,
+    execution: Execution,
+    cmec_output_bundle_filename: pathlib.Path,
+) -> None:
+    # Extract the registered outputs
+    # Copy the content to the output directory
+    # Track in the db
+    cmec_output_bundle = CMECOutput.load_from_json(cmec_output_bundle_filename)
+    _handle_outputs(
+        cmec_output_bundle.plots,
+        output_type=ResultOutputType.Plot,
+        config=config,
+        database=database,
+        execution=execution,
+    )
+    _handle_outputs(
+        cmec_output_bundle.data,
+        output_type=ResultOutputType.Data,
+        config=config,
+        database=database,
+        execution=execution,
+    )
+    _handle_outputs(
+        cmec_output_bundle.html,
+        output_type=ResultOutputType.HTML,
+        config=config,
+        database=database,
+        execution=execution,
+    )
+
+
+def _handle_outputs(
+    outputs: dict[str, OutputDict] | None,
+    output_type: ResultOutputType,
+    config: "Config",
+    database: Database,
+    execution: Execution,
+) -> None:
+    outputs = outputs or {}
+
+    for key, output_info in outputs.items():
+        filename = ensure_relative_path(
+            output_info.filename, config.paths.scratch / execution.output_fragment
+        )
+
+        _copy_file_to_results(
+            config.paths.scratch,
+            config.paths.results,
+            execution.output_fragment,
+            filename,
+        )
+        database.session.add(
+            ExecutionOutput(
+                execution_id=execution.id,
+                output_type=output_type,
+                filename=str(filename),
+                description=output_info.description,
+                short_name=key,
+                long_name=output_info.long_name,
+            )
+        )

climate_ref/executor/synchronous.py

@@ -0,0 +1,62 @@
+from typing import Any
+
+from climate_ref.config import Config
+from climate_ref.database import Database
+from climate_ref.executor.local import process_result
+from climate_ref.models import Execution
+from climate_ref_core.diagnostics import ExecutionDefinition
+from climate_ref_core.executor import execute_locally
+
+
+class SynchronousExecutor:
+    """
+    Run a diagnostic synchronously, in-process.
+
+    This is mainly useful for debugging and testing.
+    [climate_ref.executor.LocalExecutor][] is a more general purpose executor.
+    """
+
+    name = "synchronous"
+
+    def __init__(
+        self, *, database: Database | None = None, config: Config | None = None, **kwargs: Any
+    ) -> None:
+        if config is None:
+            config = Config.default()
+        if database is None:
+            database = Database.from_config(config, run_migrations=False)
+
+        self.database = database
+        self.config = config
+
+    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.
+        """
+        result = execute_locally(definition, log_level=self.config.log_level)
+        process_result(self.config, self.database, result, execution)
+
+    def join(self, timeout: float) -> None:
+        """
+        Wait for all diagnostics to finish
+
+        This returns immediately because the executor runs diagnostics synchronously.
+
+        Parameters
+        ----------
+        timeout
+            Timeout in seconds (Not used)
+        """
+        pass

climate_ref/migrations/env.py

@@ -1,3 +1,4 @@
+import alembic_postgresql_enum  # noqa
 from alembic import context, op
 from loguru import logger
 from sqlalchemy import Connection, inspect

climate_ref/migrations/versions/2025-05-02T1418_341a4aa2551e_regenerate.py

@@ -235,38 +235,17 @@ def upgrade() -> None:
         sa.Column("attributes", sa.JSON(), nullable=False),
         sa.Column("created_at", sa.DateTime(), server_default=sa.text("(CURRENT_TIMESTAMP)"), nullable=False),
         sa.Column("updated_at", sa.DateTime(), server_default=sa.text("(CURRENT_TIMESTAMP)"), nullable=False),
-        sa.Column("model", sa.Text(), nullable=True),
-        sa.Column("source_id", sa.Text(), nullable=True),
-        sa.Column("variant_label", sa.Text(), nullable=True),
-        sa.Column("metric", sa.Text(), nullable=True),
-        sa.Column("region", sa.Text(), nullable=True),
-        sa.Column("statistic", sa.Text(), nullable=True),
         sa.ForeignKeyConstraint(
             ["execution_id"], ["execution.id"], name=op.f("fk_metric_value_execution_id_execution")
         ),
         sa.PrimaryKeyConstraint("id", name=op.f("pk_metric_value")),
     )
-    with op.batch_alter_table("metric_value", schema=None) as batch_op:
-        batch_op.create_index(batch_op.f("ix_metric_value_metric"), ["metric"], unique=False)
-        batch_op.create_index(batch_op.f("ix_metric_value_model"), ["model"], unique=False)
-        batch_op.create_index(batch_op.f("ix_metric_value_region"), ["region"], unique=False)
-        batch_op.create_index(batch_op.f("ix_metric_value_source_id"), ["source_id"], unique=False)
-        batch_op.create_index(batch_op.f("ix_metric_value_statistic"), ["statistic"], unique=False)
-        batch_op.create_index(batch_op.f("ix_metric_value_variant_label"), ["variant_label"], unique=False)
 
     # ### end Alembic commands ###
 
 
 def downgrade() -> None:
     # ### commands auto generated by Alembic - please adjust! ###
-    with op.batch_alter_table("metric_value", schema=None) as batch_op:
-        batch_op.drop_index(batch_op.f("ix_metric_value_variant_label"))
-        batch_op.drop_index(batch_op.f("ix_metric_value_statistic"))
-        batch_op.drop_index(batch_op.f("ix_metric_value_source_id"))
-        batch_op.drop_index(batch_op.f("ix_metric_value_region"))
-        batch_op.drop_index(batch_op.f("ix_metric_value_model"))
-        batch_op.drop_index(batch_op.f("ix_metric_value_metric"))
-
     op.drop_table("metric_value")
     with op.batch_alter_table("execution_output", schema=None) as batch_op:
         batch_op.drop_index(batch_op.f("ix_execution_output_output_type"))

climate_ref/migrations/versions/2025-05-09T2032_03dbb4998e49_series_metric_value.py

@@ -0,0 +1,57 @@
+"""series-metric-value
+
+Revision ID: 03dbb4998e49
+Revises: 341a4aa2551e
+Create Date: 2025-05-09 20:32:08.664426
+
+"""
+
+from collections.abc import Sequence
+from typing import Union
+
+import sqlalchemy as sa
+from alembic import op
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision: str = "03dbb4998e49"
+down_revision: Union[str, None] = "341a4aa2551e"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("metric_value", schema=None) as batch_op:
+        batch_op.add_column(sa.Column("values", sa.JSON(), nullable=True))
+        batch_op.add_column(sa.Column("index", sa.JSON(), nullable=True))
+        batch_op.add_column(sa.Column("index_name", sa.String(), nullable=True))
+        batch_op.alter_column("value", existing_type=sa.FLOAT(), nullable=True)
+
+    if sa.inspect(op.get_bind()).dialect.name == "postgresql":
+        sa.Enum("SCALAR", "SERIES", name="metricvaluetype").create(op.get_bind())
+        op.add_column(
+            "metric_value",
+            sa.Column(
+                "type",
+                postgresql.ENUM("SCALAR", "SERIES", name="metricvaluetype", create_type=False),
+                nullable=False,
+            ),
+        )
+    else:
+        with op.batch_alter_table("metric_value", schema=None) as batch_op:
+            batch_op.add_column(
+                sa.Column("type", sa.Enum("SCALAR", "SERIES", name="metricvaluetype"), nullable=False)
+            )
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("metric_value", schema=None) as batch_op:
+        batch_op.alter_column("value", existing_type=sa.FLOAT(), nullable=False)
+        batch_op.drop_column("index_name")
+        batch_op.drop_column("index")
+        batch_op.drop_column("values")
+        batch_op.drop_column("type")
+
+    # ### end Alembic commands ###

climate_ref/models/__init__.py

@@ -14,7 +14,7 @@ from climate_ref.models.execution import (
     ExecutionGroup,
     ExecutionOutput,
 )
-from climate_ref.models.metric_value import MetricValue
+from climate_ref.models.metric_value import MetricValue, ScalarMetricValue, SeriesMetricValue
 from climate_ref.models.provider import Provider
 
 Table = TypeVar("Table", bound=Base)
@@ -29,5 +29,7 @@ __all__ = [
     "ExecutionOutput",
     "MetricValue",
     "Provider",
+    "ScalarMetricValue",
+    "SeriesMetricValue",
     "Table",
 ]

climate_ref/models/base.py

@@ -12,6 +12,8 @@ class Base(DeclarativeBase):
 
     type_annotation_map = {  # noqa: RUF012
         dict[str, Any]: JSON,
+        list[float | int]: JSON,
+        list[float | int | str]: JSON,
     }
     metadata = MetaData(
         # Enforce a common naming convention for constraints