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

climate_ref/executor/__init__.py

@@ -9,266 +9,8 @@ The simplest executor is the `LocalExecutor`, which runs the diagnostic in the s
 This is useful for local testing and debugging.
 """
 
-import importlib
-import pathlib
-import shutil
-from typing import TYPE_CHECKING
+from .local import LocalExecutor
+from .result_handling import handle_execution_result
+from .synchronous import SynchronousExecutor
 
-from loguru import logger
-from sqlalchemy import insert
-
-from climate_ref.database import Database
-from climate_ref.models.execution import Execution, ExecutionOutput, ResultOutputType
-from climate_ref.models.metric_value import MetricValue
-from climate_ref_core.diagnostics import ExecutionResult, ensure_relative_path
-from climate_ref_core.exceptions import InvalidExecutorException, ResultValidationError
-from climate_ref_core.executor import EXECUTION_LOG_FILENAME, Executor
-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 import_executor_cls(fqn: str) -> type[Executor]:
-    """
-    Import an executor using a fully qualified module path
-
-    Parameters
-    ----------
-    fqn
-        Full package and attribute name of the executor to import
-
-        For example: `climate_ref_example.executor` will use the `executor` attribute from the
-        `climate_ref_example` package.
-
-    Raises
-    ------
-    climate_ref_core.exceptions.InvalidExecutorException
-        If the executor cannot be imported
-
-        If the executor isn't a valid `DiagnosticProvider`.
-
-    Returns
-    -------
-    :
-        Executor instance
-    """
-    module, attribute_name = fqn.rsplit(".", 1)
-
-    try:
-        imp = importlib.import_module(module)
-        executor: type[Executor] = getattr(imp, attribute_name)
-
-        # We can't really check if the executor is a subclass of Executor here
-        # Protocols can't be used with issubclass if they have non-method members
-        # We have to check this at class instantiation time
-
-        return executor
-    except ModuleNotFoundError:
-        logger.error(f"Package '{fqn}' not found")
-        raise InvalidExecutorException(fqn, f"Module '{module}' not found")
-    except AttributeError:
-        logger.error(f"Provider '{fqn}' not found")
-        raise InvalidExecutorException(fqn, f"Executor '{attribute_name}' not found in {module}")
-
-
-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  # noqa
-    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 a diagnostic bundle
-        # 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(MetricValue),
-                    [
-                        {
-                            "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: 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:
-    if outputs is None:
-        return
-
-    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,
-            )
-        )
+__all__ = ["LocalExecutor", "SynchronousExecutor", "handle_execution_result"]

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