climate-ref 0.5.5__py3-none-any.whl → 0.6.1__py3-none-any.whl

climate_ref/executor/hpc.py

@@ -0,0 +1,308 @@
+"""
+HPC-based Executor to use job schedulers.
+
+If you want to
+- run REF under the HPC workflows
+- run REF in multiple nodes
+
+"""
+
+import os
+import time
+from typing import Any, cast
+
+import parsl
+from loguru import logger
+from parsl import python_app
+from parsl.config import Config as ParslConfig
+from parsl.executors import HighThroughputExecutor
+from parsl.launchers import SrunLauncher
+from parsl.providers import SlurmProvider
+from tqdm import tqdm
+
+from climate_ref.config import Config
+from climate_ref.database import Database
+from climate_ref.models import Execution
+from climate_ref.slurm import HAS_REAL_SLURM, SlurmChecker
+from climate_ref_core.diagnostics import ExecutionDefinition, ExecutionResult
+from climate_ref_core.exceptions import DiagnosticError, ExecutionError
+from climate_ref_core.executor import execute_locally
+
+from .local import ExecutionFuture, process_result
+
+
+@python_app
+def _process_run(definition: ExecutionDefinition, log_level: str) -> ExecutionResult:
+    """Run the function on computer nodes"""
+    # This is a catch-all for any exceptions that occur in the process and need to raise for
+    # parsl retries to work
+    try:
+        return execute_locally(definition=definition, log_level=log_level, raise_error=True)
+    except DiagnosticError as e:  # pragma: no cover
+        # any diagnostic error will be caught here
+        logger.exception("Error running diagnostic")
+        return cast(ExecutionResult, e.result)
+
+
+def _to_float(x: Any) -> float | None:
+    if x is None:
+        return None
+    if isinstance(x, int | float):
+        return float(x)
+    try:
+        return float(x)
+    except (ValueError, TypeError):
+        return None
+
+
+def _to_int(x: Any) -> int | None:
+    if x is None:
+        return None
+    if isinstance(x, int):
+        return x
+    try:
+        return int(float(x))  # Handles both "123" and "123.0"
+    except (ValueError, TypeError):
+        return None
+
+
+class HPCExecutor:
+    """
+    Run diagnostics by submitting a job script
+
+    """
+
+    name = "hpc"
+
+    def __init__(
+        self,
+        *,
+        database: Database | None = None,
+        config: Config | None = None,
+        **executor_config: str | float | int,
+    ) -> None:
+        config = config or Config.default()
+        database = database or Database.from_config(config, run_migrations=False)
+
+        self.config = config
+        self.database = database
+
+        self.scheduler = executor_config.get("scheduler", "slurm")
+        self.account = str(executor_config.get("account", os.environ.get("USER")))
+        self.username = executor_config.get("username", os.environ.get("USER"))
+        self.partition = str(executor_config.get("partition")) if executor_config.get("partition") else None
+        self.qos = str(executor_config.get("qos")) if executor_config.get("qos") else None
+        self.req_nodes = int(executor_config.get("req_nodes", 1))
+        self.walltime = str(executor_config.get("walltime", "00:10:00"))
+        self.log_dir = str(executor_config.get("log_dir", "runinfo"))
+
+        self.cores_per_worker = _to_int(executor_config.get("cores_per_worker"))
+        self.mem_per_worker = _to_float(executor_config.get("mem_per_worker"))
+
+        hours, minutes, seconds = map(int, self.walltime.split(":"))
+        total_minutes = hours * 60 + minutes + seconds / 60
+        self.total_minutes = total_minutes
+
+        if executor_config.get("validation") and HAS_REAL_SLURM:
+            self._validate_slurm_params()
+
+        self._initialize_parsl()
+
+        self.parsl_results: list[ExecutionFuture] = []
+
+    def _validate_slurm_params(self) -> None:
+        """Validate the Slurm configuration using SlurmChecker.
+
+        Raises
+        ------
+            ValueError: If account, partition or QOS are invalid or inaccessible.
+        """
+        slurm_checker = SlurmChecker()
+        if self.account and not slurm_checker.get_account_info(self.account):
+            raise ValueError(f"Account: {self.account} not valid")
+
+        partition_limits = None
+        node_info = None
+
+        if self.partition:
+            if not slurm_checker.get_partition_info(self.partition):
+                raise ValueError(f"Partition: {self.partition} not valid")
+
+            if not slurm_checker.can_account_use_partition(self.account, self.partition):
+                raise ValueError(f"Account: {self.account} cannot access partiton: {self.partition}")
+
+            partition_limits = slurm_checker.get_partition_limits(self.partition)
+            node_info = slurm_checker.get_node_from_partition(self.partition)
+
+        qos_limits = None
+        if self.qos:
+            if not slurm_checker.get_qos_info(self.qos):
+                raise ValueError(f"QOS: {self.qos} not valid")
+
+            if not slurm_checker.can_account_use_qos(self.account, self.qos):
+                raise ValueError(f"Account: {self.account} cannot access qos: {self.qos}")
+
+            qos_limits = slurm_checker.get_qos_limits(self.qos)
+
+        max_cores_per_node = int(node_info["cpus"]) if node_info else None
+        if max_cores_per_node and self.cores_per_worker:
+            if self.cores_per_worker > max_cores_per_node:
+                raise ValueError(
+                    f"cores_per_work:{self.cores_per_worker}"
+                    f"larger than the maximum in a node {max_cores_per_node}"
+                )
+
+        max_mem_per_node = float(node_info["real_memory"]) if node_info else None
+        if max_mem_per_node and self.mem_per_worker:
+            if self.mem_per_worker > max_mem_per_node:
+                raise ValueError(
+                    f"mem_per_work:{self.mem_per_worker}"
+                    f"larger than the maximum mem in a node {max_mem_per_node}"
+                )
+
+        max_walltime_partition = (
+            partition_limits["max_time_minutes"] if partition_limits else self.total_minutes
+        )
+        max_walltime_qos = qos_limits["max_time_minutes"] if qos_limits else self.total_minutes
+
+        max_walltime_minutes = min(float(max_walltime_partition), float(max_walltime_qos))
+
+        if self.total_minutes > float(max_walltime_minutes):
+            raise ValueError(
+                f"Walltime: {self.walltime} exceed the maximum time "
+                f"{max_walltime_minutes} allowed by {self.partition} and {self.qos}"
+            )
+
+    def _initialize_parsl(self) -> None:
+        executor_config = self.config.executor.config
+
+        provider = SlurmProvider(
+            account=self.account,
+            partition=self.partition,
+            qos=self.qos,
+            nodes_per_block=self.req_nodes,
+            max_blocks=int(executor_config.get("max_blocks", 1)),
+            scheduler_options=executor_config.get("scheduler_options", "#SBATCH -C cpu"),
+            worker_init=executor_config.get("worker_init", "source .venv/bin/activate"),
+            launcher=SrunLauncher(
+                debug=True,
+                overrides=executor_config.get("overrides", ""),
+            ),
+            walltime=self.walltime,
+            cmd_timeout=int(executor_config.get("cmd_timeout", 120)),
+        )
+        executor = HighThroughputExecutor(
+            label="ref_hpc_executor",
+            cores_per_worker=self.cores_per_worker if self.cores_per_worker else 1,
+            mem_per_worker=self.mem_per_worker,
+            max_workers_per_node=_to_int(executor_config.get("max_workers_per_node", 16)),
+            cpu_affinity=str(executor_config.get("cpu_affinity")),
+            provider=provider,
+        )
+
+        hpc_config = ParslConfig(
+            run_dir=self.log_dir, executors=[executor], retries=int(executor_config.get("retries", 2))
+        )
+        parsl.load(hpc_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.
+        """
+        # Submit the execution to the process pool
+        # and track the future so we can wait for it to complete
+        future = _process_run(
+            definition=definition,
+            log_level=self.config.log_level,
+        )
+
+        self.parsl_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 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 (won't used in HPCExecutor)
+
+        Raises
+        ------
+        TimeoutError
+            If the timeout is reached
+        """
+        start_time = time.time()
+        refresh_time = 0.5
+
+        results = self.parsl_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 > self.total_minutes * 60:
+                    logger.debug(f"Time elasped {elapsed_time} for joining the results")
+
+                # Wait for a short time before checking for completed executions
+                time.sleep(refresh_time)
+        finally:
+            t.close()
+            if parsl.dfk():
+                parsl.dfk().cleanup()

climate_ref/executor/local.py

@@ -1,4 +1,5 @@
 import concurrent.futures
+import multiprocessing
 import time
 from concurrent.futures import Future, ProcessPoolExecutor
 from typing import Any
@@ -13,7 +14,7 @@ from climate_ref.models import Execution
 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 climate_ref_core.logging import initialise_logging
 
 from .result_handling import handle_execution_result
 
@@ -63,11 +64,17 @@ class ExecutionFuture:
     execution_id: int | None = None
 
 
-def _process_initialiser() -> None:
+def _process_initialiser() -> None:  # pragma: no cover
     # Setup the logging for the process
     # This replaces the loguru default handler
     try:
-        add_log_handler()
+        logger.remove()
+        config = Config.default()
+        initialise_logging(
+            level=config.log_level,
+            format=config.log_format,
+            log_directory=config.paths.log,
+        )
     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
@@ -118,7 +125,12 @@ class LocalExecutor:
         if pool is not None:
             self.pool = pool
         else:
-            self.pool = ProcessPoolExecutor(max_workers=n, initializer=_process_initialiser)
+            self.pool = ProcessPoolExecutor(
+                max_workers=n,
+                initializer=_process_initialiser,
+                # Explicitly set the context to "spawn" to avoid issues with hanging on MacOS
+                mp_context=multiprocessing.get_context("spawn"),
+            )
         self._results: list[ExecutionFuture] = []
 
     def run(
@@ -214,9 +226,17 @@ class LocalExecutor:
                 elapsed_time = time.time() - start_time
 
                 if elapsed_time > timeout:
+                    for result in results:
+                        logger.warning(
+                            f"Execution {result.definition.execution_slug()} "
+                            f"did not complete within the timeout"
+                        )
+                    self.pool.shutdown(wait=False, cancel_futures=True)
                     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()
+
+        logger.info("All executions completed successfully")

climate_ref/executor/result_handling.py

@@ -129,7 +129,6 @@ def handle_execution_result(
             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

climate_ref/slurm.py

@@ -0,0 +1,192 @@
+import importlib.util
+from typing import Any
+
+HAS_REAL_SLURM = importlib.util.find_spec("pyslurm") is not None
+
+
+class SlurmChecker:
+    """Check and get slurm settings."""
+
+    def __init__(self, intest: bool = False) -> None:
+        if HAS_REAL_SLURM:
+            import pyslurm  # type: ignore
+
+            self.slurm_association: dict[int, Any] | None = pyslurm.db.Associations.load()
+            self.slurm_partition: dict[str, Any] | None = pyslurm.Partitions.load()
+            self.slurm_qos: dict[str, Any] | None = pyslurm.qos().get()
+            self.slurm_node: dict[str, Any] | None = pyslurm.Nodes.load()
+        elif intest:
+            import pyslurm
+
+            self.slurm_association = pyslurm.db.Associations.load()  # dict [num -> Association]
+            self.slurm_partition = pyslurm.Partitions.load()  # collection
+            self.slurm_qos = pyslurm.qos().get()  # dict
+            self.slurm_node = pyslurm.Nodes.load()  # dict
+        else:
+            print("Warning: pyslurm not found. Skipping HPCExecutor config validations")
+            self.slurm_association = None
+            self.slurm_partition = None
+            self.slurm_qos = None
+            self.slurm_node = None
+
+    def get_partition_info(self, partition_name: str) -> Any:
+        """Check if a partition exists in the Slurm configuration."""
+        return self.slurm_partition.get(partition_name) if self.slurm_partition else None
+
+    def get_qos_info(self, qos_name: str) -> Any:
+        """Check if a qos exists in the Slurm configuration."""
+        return self.slurm_qos.get(qos_name) if self.slurm_qos else None
+
+    def get_account_info(self, account_name: str) -> list[Any]:
+        """Get all associations for an account"""
+        if self.slurm_association:
+            return [a for a in self.slurm_association.values() if a.account == account_name]
+        else:
+            return [None]
+
+    def can_account_use_partition(self, account_name: str, partition_name: str) -> bool:
+        """
+        Check if an account has access to a specific partition.
+
+        Returns
+        -------
+            bool: True if accessible, False if not accessible or error occurred
+        """
+        account_info = self.get_account_info(account_name)
+        if not account_info:
+            return False
+
+        partition_info = self.get_partition_info(partition_name)
+
+        if not partition_info:
+            return False
+
+        allowed_partitions = account_info[0].partition
+        if allowed_partitions is None:
+            return True
+        else:
+            return partition_name in allowed_partitions
+
+    def can_account_use_qos(self, account_name: str, qos_name: str) -> bool:
+        """
+        Check if an account has access to a specific qos.
+
+        Returns
+        -------
+            bool: True if accessible, False if not accessible or error occurred
+        """
+        account_info = self.get_account_info(account_name)
+
+        if not account_info:
+            return False
+
+        qos_info = self.get_qos_info(qos_name)
+        if not qos_info:
+            return False
+
+        sample_acc = account_info[0]
+        for acc in account_info:
+            if acc.user == "minxu":
+                sample_acc = acc
+                break
+
+        allowed_qoss = sample_acc.qos
+        if allowed_qoss is None:
+            return True
+        else:
+            return qos_name in allowed_qoss
+
+    def get_partition_limits(self, partition_name: str) -> dict[str, str | int] | None:
+        """
+        Get time limits for a specific partition.
+
+        Returns
+        -------
+            Dict with 'max_time' and 'default_time' (strings or UNLIMITED)
+            or None if partition doesn't exist or error occurred
+        """
+        partition_info = self.get_partition_info(partition_name)
+        if not partition_info:
+            return None
+
+        return {
+            "max_time_minutes": partition_info.to_dict().get("max_time", 0),  # in minutes
+            "default_time_minutes": partition_info.to_dict().get("default_time", 30),  # in minutes
+            "max_nodes": partition_info.to_dict().get("max_node", 1),
+            "total_nodes": partition_info.to_dict().get("total_nodes", 0),
+            "total_cpus": partition_info.to_dict().get("total_cpus", 0),
+        }
+
+    def get_node_from_partition(self, partition_name: str) -> dict[str, str | int] | None:
+        """
+        Get the node information for a specific partition.
+
+        Returns
+        -------
+            Dicts
+        """
+        partition_info = self.get_partition_info(partition_name)
+        if not partition_info:
+            return None
+
+        sample_node = None
+
+        if self.slurm_node:
+            for node in self.slurm_node.values():
+                if partition_name in node.partitions and "cpu" in node.available_features:
+                    sample_node = node
+                    break
+
+        return {
+            "cpus": int(sample_node.total_cpus) if sample_node is not None else 1,
+            "cores_per_socket": int(sample_node.cores_per_socket) if sample_node is not None else 1,
+            "sockets": int(sample_node.sockets) if sample_node is not None else 1,
+            "threads_per_core": int(sample_node.threads_per_core) if sample_node is not None else 1,
+            "real_memory": int(sample_node.real_memory) if sample_node is not None else 215,
+            "node_names": sample_node.name if sample_node is not None else "unknown",
+        }
+
+    def get_qos_limits(self, qos_name: str) -> dict[str, str | int]:
+        """
+        Get time limits for a specific qos.
+
+        Returns
+        -------
+            Dict with 'max_time' and 'default_time' (strings or UNLIMITED)
+            or None if partition doesn't exist or error occurred
+        """
+        qos_info = self.get_qos_info(qos_name)
+
+        return {
+            "max_time_minutes": qos_info.get("max_wall_pj", 1.0e6),
+            "max_jobs_pu": qos_info.get("max_jobs_pu", 1.0e6),
+            "max_submit_jobs_pu": qos_info.get("max_submit_jobs_pu", 1.0e6),
+            "max_tres_pj": qos_info.get("max_tres_pj").split("=")[0],
+            "default_time_minutes": 120,
+        }
+
+    def check_account_partition_access_with_limits(
+        self, account_name: str, partition_name: str
+    ) -> dict[str, Any]:
+        """
+        Comprehensive check of account access and partition limits.
+
+        Returns dictionary with all relevant information.
+        """
+        result = {
+            "account_exists": True if self.get_account_info(account_name) else False,
+            "partition_exists": True if self.get_partition_info(partition_name) else False,
+            "has_access": False,
+            "time_limits": None,
+            "error": "none",
+        }
+
+        try:
+            if result["account_exists"] and result["partition_exists"]:
+                result["has_access"] = self.can_account_use_partition(account_name, partition_name)
+                if result["has_access"]:
+                    result["time_limits"] = self.get_partition_info(partition_name).to_dict().get("max_time")
+        except Exception as e:
+            result["error"] = str(e)
+
+        return result

climate_ref/solver.py

@@ -245,6 +245,57 @@ def _solve_from_data_requirements(
         )
 
 
+@define
+class SolveFilterOptions:
+    """
+    Options to filter the diagnostics that are solved
+    """
+
+    diagnostic: list[str] | None = None
+    """
+    Check if the diagnostic slug contains any of the provided values
+    """
+    provider: list[str] | None = None
+    """
+    Check if the provider slug contains any of the provided values
+    """
+
+
+def matches_filter(diagnostic: Diagnostic, filters: SolveFilterOptions | None) -> bool:
+    """
+    Check if a diagnostic matches the provided filters
+
+    Each filter is optional and a diagnostic will match if it satisfies all the provided filters.
+    i.e. the filters are ANDed together.
+
+    Parameters
+    ----------
+    diagnostic
+        Diagnostic to check against the filters
+    filters
+        Collection of filters to apply to the diagnostic
+
+        If no filters are provided, the diagnostic is considered to match
+
+    Returns
+    -------
+        True if the diagnostic matches the filters, False otherwise
+    """
+    if filters is None:
+        return True
+
+    diagnostic_slug = diagnostic.slug
+    provider_slug = diagnostic.provider.slug
+
+    if filters.provider and not any([f.lower() in provider_slug for f in filters.provider]):
+        return False
+
+    if filters.diagnostic and not any([f.lower() in diagnostic_slug for f in filters.diagnostic]):
+        return False
+
+    return True
+
+
 @define
 class ExecutionSolver:
     """
@@ -278,7 +329,9 @@ class ExecutionSolver:
             },
         )
 
-    def solve(self) -> typing.Generator[DiagnosticExecution, None, None]:
+    def solve(
+        self, filters: SolveFilterOptions | None = None
+    ) -> typing.Generator[DiagnosticExecution, None, None]:
         """
         Solve which executions need to be calculated for a dataset
 
@@ -293,17 +346,23 @@ class ExecutionSolver:
         """
         for provider in self.provider_registry.providers:
             for diagnostic in provider.diagnostics():
+                # Filter the diagnostic based on the provided filters
+                if not matches_filter(diagnostic, filters):
+                    logger.debug(f"Skipping {diagnostic.full_slug()} due to filter")
+                    continue
                 yield from solve_executions(self.data_catalog, diagnostic, provider)
 
 
 def solve_required_executions(  # noqa: PLR0913
     db: Database,
     dry_run: bool = False,
+    execute: bool = True,
     solver: ExecutionSolver | None = None,
     config: Config | None = None,
     timeout: int = 60,
     one_per_provider: bool = False,
     one_per_diagnostic: bool = False,
+    filters: SolveFilterOptions | None = None,
 ) -> None:
     """
     Solve for executions that require recalculation
@@ -328,7 +387,7 @@ def solve_required_executions(  # noqa: PLR0913
     diagnostic_count = {}
     provider_count = {}
 
-    for potential_execution in solver.solve():
+    for potential_execution in solver.solve(filters):
         # The diagnostic output is first written to the scratch directory
         definition = potential_execution.build_execution_definition(output_root=config.paths.scratch)
 
@@ -371,6 +430,7 @@ def solve_required_executions(  # noqa: PLR0913
                 logger.info(f"Created new execution group: {potential_execution.execution_slug()!r}")
                 db.session.flush()
 
+            # TODO: Move this logic to the solver
             # Check if we should run given the one_per_provider or one_per_diagnostic flags
             one_of_check_failed = (
                 one_per_provider and provider_count.get(diagnostic.provider.slug, 0) > 0
@@ -403,10 +463,11 @@ def solve_required_executions(  # noqa: PLR0913
                 # Add links to the datasets used in the execution
                 execution.register_datasets(db, definition.datasets)
 
-                executor.run(
-                    definition=definition,
-                    execution=execution,
-                )
+                if execute:
+                    executor.run(
+                        definition=definition,
+                        execution=execution,
+                    )
 
                 provider_count[diagnostic.provider.slug] += 1
                 diagnostic_count[diagnostic.full_slug()] += 1