climate-ref 0.6.6__py3-none-any.whl → 0.8.0__py3-none-any.whl

climate_ref/executor/hpc.py

@@ -19,8 +19,9 @@ except ImportError:  # pragma: no cover
     )
 
 import os
+import re
 import time
-from typing import Any
+from typing import Annotated, Any, Literal
 
 import parsl
 from loguru import logger
@@ -29,6 +30,7 @@ from parsl.config import Config as ParslConfig
 from parsl.executors import HighThroughputExecutor
 from parsl.launchers import SimpleLauncher, SrunLauncher
 from parsl.providers import SlurmProvider
+from pydantic import BaseModel, Field, StrictBool, field_validator, model_validator
 from tqdm import tqdm
 
 from climate_ref.config import Config
@@ -43,6 +45,72 @@ from .local import ExecutionFuture, process_result
 from .pbs_scheduler import SmartPBSProvider
 
 
+class SlurmConfig(BaseModel):
+    """Slurm Configurations"""
+
+    scheduler: Literal["slurm"]
+    account: str
+    username: str
+    partition: str | None = None
+    log_dir: str = "runinfo"
+    qos: str | None = None
+    req_nodes: Annotated[int, Field(strict=True, ge=1, le=1000)] = 1
+    cores_per_worker: Annotated[int, Field(strict=True, ge=1, le=1000)] = 1
+    mem_per_worker: Annotated[float, Field(strict=True, gt=0, lt=1000.0)] | None = None
+    max_workers_per_node: Annotated[int, Field(strict=True, ge=1, le=1000)] = 16
+    validation: StrictBool = False
+    walltime: str = "00:30:00"
+    scheduler_options: str = ""
+    retries: Annotated[int, Field(strict=True, ge=1, le=3)] = 2
+    max_blocks: Annotated[int, Field(strict=True, ge=1)] = 1  # one block mean one job?
+    worker_init: str = ""
+    overrides: str = ""
+    cmd_timeout: Annotated[int, Field(strict=True, ge=0)] = 120
+    cpu_affinity: str = "none"
+
+    @model_validator(mode="before")
+    def _check_parition_qos(cls, data: Any) -> Any:
+        if not ("partition" in data or "qos" in data):
+            raise ValueError("partition or qos is needed")
+        return data
+
+    @field_validator("scheduler_options")
+    def _validate_sbatch_syntax(cls, v: str | None) -> Any:
+        if not v:
+            return v
+
+        sbatch_pattern = re.compile(
+            r"^\s*#SBATCH\s+"  # Start with #SBATCH
+            r"(?:-\w+\s+[^\s]+"  # Option-value pairs
+            r"(?:\s+-\w+\s+[^\s]+)*)"  # Additional options
+            r"\s*$",
+            re.IGNORECASE | re.MULTILINE,
+        )
+
+        invalid_lines = [
+            line
+            for line in v.split("\n")
+            if not (line.strip().upper().startswith("#SBATCH") and sbatch_pattern.match(line.strip()))
+        ]
+
+        if invalid_lines:
+            error_msg = (
+                "Invalid SBATCH directives:\n"
+                + "\n".join(invalid_lines)
+                + "\n"
+                + "Expected format: '#SBATCH -option value [-option value ...]'"
+            )
+            raise ValueError(error_msg)
+        return v
+
+    @field_validator("walltime")
+    def _validate_walltime(cls, v: str) -> str:
+        pattern = r"^(\d+-)?\d{1,5}:[0-5][0-9]:[0-5][0-9]$"
+        if not re.match(pattern, v):
+            raise ValueError("Walltime must be in `D-HH:MM:SS/HH:MM:SS` format")
+        return v
+
+
 @python_app
 def _process_run(definition: ExecutionDefinition, log_level: str) -> ExecutionResult:
     """Run the function on computer nodes"""
@@ -112,13 +180,18 @@ class HPCExecutor:
         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(":"))
+        if self.scheduler == "slurm":
+            self.slurm_config = SlurmConfig.model_validate(executor_config)
+            hours, minutes, seconds = map(int, self.slurm_config.walltime.split(":"))
+
+            if self.slurm_config.validation and HAS_REAL_SLURM:
+                self._validate_slurm_params()
+        else:
+            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] = []
@@ -131,45 +204,52 @@ class HPCExecutor:
             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")
+        if self.slurm_config.account and not slurm_checker.get_account_info(self.slurm_config.account):
+            raise ValueError(f"Account: {self.slurm_config.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 self.slurm_config.partition:
+            if not slurm_checker.get_partition_info(self.slurm_config.partition):
+                raise ValueError(f"Partition: {self.slurm_config.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}")
+            if not slurm_checker.can_account_use_partition(
+                self.slurm_config.account, self.slurm_config.partition
+            ):
+                raise ValueError(
+                    f"Account: {self.slurm_config.account}"
+                    f" cannot access partiton: {self.slurm_config.partition}"
+                )
 
-            partition_limits = slurm_checker.get_partition_limits(self.partition)
-            node_info = slurm_checker.get_node_from_partition(self.partition)
+            partition_limits = slurm_checker.get_partition_limits(self.slurm_config.partition)
+            node_info = slurm_checker.get_node_from_partition(self.slurm_config.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 self.slurm_config.qos:
+            if not slurm_checker.get_qos_info(self.slurm_config.qos):
+                raise ValueError(f"QOS: {self.slurm_config.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}")
+            if not slurm_checker.can_account_use_qos(self.slurm_config.account, self.slurm_config.qos):
+                raise ValueError(
+                    f"Account: {self.slurm_config.account} cannot access qos: {self.slurm_config.qos}"
+                )
 
-            qos_limits = slurm_checker.get_qos_limits(self.qos)
+            qos_limits = slurm_checker.get_qos_limits(self.slurm_config.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:
+        if max_cores_per_node and self.slurm_config.cores_per_worker:
+            if self.slurm_config.cores_per_worker > max_cores_per_node:
                 raise ValueError(
-                    f"cores_per_work:{self.cores_per_worker}"
+                    f"cores_per_work:{self.slurm_config.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:
+        if max_mem_per_node and self.slurm_config.mem_per_worker:
+            if self.slurm_config.mem_per_worker > max_mem_per_node:
                 raise ValueError(
-                    f"mem_per_work:{self.mem_per_worker}"
+                    f"mem_per_work:{self.slurm_config.mem_per_worker}"
                     f"larger than the maximum mem in a node {max_mem_per_node}"
                 )
 
@@ -182,8 +262,8 @@ class HPCExecutor:
 
         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}"
+                f"Walltime: {self.slurm_config.walltime} exceed the maximum time "
+                f"{max_walltime_minutes} allowed by {self.slurm_config.partition} and {self.slurm_config.qos}"
             )
 
     def _initialize_parsl(self) -> None:
@@ -192,19 +272,34 @@ class HPCExecutor:
         provider: SlurmProvider | SmartPBSProvider
         if self.scheduler == "slurm":
             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"),
+                account=self.slurm_config.account,
+                partition=self.slurm_config.partition,
+                qos=self.slurm_config.qos,
+                nodes_per_block=self.slurm_config.req_nodes,
+                max_blocks=self.slurm_config.max_blocks,
+                scheduler_options=self.slurm_config.scheduler_options,
+                worker_init=self.slurm_config.worker_init,
                 launcher=SrunLauncher(
                     debug=True,
-                    overrides=executor_config.get("overrides", ""),
+                    overrides=self.slurm_config.overrides,
                 ),
-                walltime=self.walltime,
-                cmd_timeout=int(executor_config.get("cmd_timeout", 120)),
+                walltime=self.slurm_config.walltime,
+                cmd_timeout=self.slurm_config.cmd_timeout,
+            )
+
+            executor = HighThroughputExecutor(
+                label="ref_hpc_executor",
+                cores_per_worker=self.slurm_config.cores_per_worker,
+                mem_per_worker=self.slurm_config.mem_per_worker,
+                max_workers_per_node=self.slurm_config.max_workers_per_node,
+                cpu_affinity=self.slurm_config.cpu_affinity,
+                provider=provider,
+            )
+
+            hpc_config = ParslConfig(
+                run_dir=self.slurm_config.log_dir,
+                executors=[executor],
+                retries=self.slurm_config.retries,
             )
 
         elif self.scheduler == "pbs":
@@ -227,23 +322,24 @@ class HPCExecutor:
                 walltime=self.walltime,
                 cmd_timeout=int(executor_config.get("cmd_timeout", 120)),
             )
-        else:
-            raise ValueError(f"Unsupported scheduler: {self.scheduler}")
 
-        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,
-        )
+            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)),
-        )
+            hpc_config = ParslConfig(
+                run_dir=self.log_dir,
+                executors=[executor],
+                retries=int(executor_config.get("retries", 2)),
+            )
+
+        else:
+            raise ValueError(f"Unsupported scheduler: {self.scheduler}")
 
         parsl.load(hpc_config)
 

climate_ref/executor/local.py

@@ -88,8 +88,7 @@ def _process_run(definition: ExecutionDefinition, log_level: str) -> ExecutionRe
     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
+        return ExecutionResult.build_from_failure(definition)
 
 
 class LocalExecutor:

climate_ref/executor/result_handling.py

@@ -154,6 +154,8 @@ def _process_execution_series(
                 "execution_id": execution.id,
                 "values": series_result.values,
                 "attributes": series_result.attributes,
+                "index": series_result.index,
+                "index_name": series_result.index_name,
                 **series_result.dimensions,
             }
             for series_result in series_values
@@ -195,12 +197,19 @@ def handle_execution_result(
         The result of the diagnostic execution, either successful or failed
     """
     # Always copy log data to the results directory
-    _copy_file_to_results(
-        config.paths.scratch,
-        config.paths.results,
-        execution.output_fragment,
-        EXECUTION_LOG_FILENAME,
-    )
+    try:
+        _copy_file_to_results(
+            config.paths.scratch,
+            config.paths.results,
+            execution.output_fragment,
+            EXECUTION_LOG_FILENAME,
+        )
+    except FileNotFoundError:
+        logger.error(
+            f"Could not find log file {EXECUTION_LOG_FILENAME} in scratch directory: {config.paths.scratch}"
+        )
+        execution.mark_failed()
+        return
 
     if not result.successful or result.metric_bundle_filename is None:
         logger.error(f"{execution} failed")
@@ -304,12 +313,13 @@ def _handle_outputs(
             filename,
         )
         database.session.add(
-            ExecutionOutput(
+            ExecutionOutput.build(
                 execution_id=execution.id,
                 output_type=output_type,
                 filename=str(filename),
                 description=output_info.description,
                 short_name=key,
                 long_name=output_info.long_name,
+                dimensions=output_info.dimensions or {},
             )
         )

climate_ref/migrations/env.py

@@ -4,7 +4,10 @@ from sqlalchemy import Connection, inspect
 
 from climate_ref.config import Config
 from climate_ref.database import Database
-from climate_ref.models import Base, MetricValue
+from climate_ref.models import Base
+from climate_ref.models.execution import ExecutionOutput
+from climate_ref.models.metric_value import MetricValue
+from climate_ref.models.mixins import DimensionMixin
 from climate_ref_core.logging import capture_logging
 from climate_ref_core.pycmec.controlled_vocabulary import CV
 
@@ -33,7 +36,7 @@ target_metadata = Base.metadata
 # Custom migration functions that are run on every migration
 
 
-def _add_metric_value_columns(connection: Connection) -> None:
+def _add_dimension_columns(connection: Connection, table: str, Cls: type[DimensionMixin]) -> None:
     """
     Add any missing columns in the current CV to the database
 
@@ -44,27 +47,25 @@ def _add_metric_value_columns(connection: Connection) -> None:
     connection
         Open connection to the database
     """
-    metric_value_table = "metric_value"
-
     inspector = inspect(connection)
 
     # Check if table already exists
     # Skip if it doesn't
     tables = inspector.get_table_names()
-    if metric_value_table not in tables:
-        logger.warning(f"No table named {metric_value_table!r} found")
+    if table not in tables:
+        logger.warning(f"No table named {table!r} found")
         return
 
     # Extract the current columns in the DB
-    existing_columns = [c["name"] for c in inspector.get_columns(metric_value_table)]
+    existing_columns = [c["name"] for c in inspector.get_columns(table)]
 
     cv_file = ref_config.paths.dimensions_cv
     cv = CV.load_from_file(cv_file)
 
     for dimension in cv.dimensions:
         if dimension.name not in existing_columns:
-            logger.info(f"Adding missing metric value dimension: {dimension.name!r}")
-            op.add_column(metric_value_table, MetricValue.build_dimension_column(dimension))
+            logger.info(f"Adding missing value dimension: {dimension.name!r}")
+            op.add_column(table, Cls.build_dimension_column(dimension))
 
 
 def include_object(object_, name: str, type_, reflected, compare_to) -> bool:
@@ -134,7 +135,8 @@ def run_migrations_online() -> None:
             # Set up the Operations context
             # This is needed to alter the tables
             with op.Operations.context(context.get_context()):  # type: ignore
-                _add_metric_value_columns(connection)
+                _add_dimension_columns(connection, "metric_value", MetricValue)
+                _add_dimension_columns(connection, "execution_output", ExecutionOutput)
 
 
 if context.is_offline_mode():

climate_ref/migrations/versions/2025-09-10T1358_2f6e36738e06_use_version_as_version_facet_for_.py

@@ -0,0 +1,35 @@
+"""use 'version' as version facet for obs4MIPs
+
+Revision ID: 2f6e36738e06
+Revises: 8d28e5e0f9c3
+Create Date: 2025-09-10 13:58:40.660076
+
+"""
+
+from collections.abc import Sequence
+from typing import Union
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "2f6e36738e06"
+down_revision: Union[str, None] = "8d28e5e0f9c3"
+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("obs4mips_dataset", schema=None) as batch_op:
+        batch_op.add_column(sa.Column("version", sa.String(), nullable=False))
+
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("obs4mips_dataset", schema=None) as batch_op:
+        batch_op.drop_column("version")
+
+    # ### end Alembic commands ###

climate_ref/migrations/versions/2025-09-22T2359_20cd136a5b04_add_pmp_version.py

@@ -0,0 +1,35 @@
+"""add pmp version
+
+Revision ID: 20cd136a5b04
+Revises: 2f6e36738e06
+Create Date: 2025-09-22 23:59:42.724007
+
+"""
+
+from collections.abc import Sequence
+from typing import Union
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "20cd136a5b04"
+down_revision: Union[str, None] = "2f6e36738e06"
+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("pmp_climatology_dataset", schema=None) as batch_op:
+        batch_op.add_column(sa.Column("version", sa.String(), nullable=False))
+
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("pmp_climatology_dataset", schema=None) as batch_op:
+        batch_op.drop_column("version")
+
+    # ### end Alembic commands ###

climate_ref/models/__init__.py

@@ -4,9 +4,7 @@ Declaration of the models used by the REF.
 These models are used to represent the data that is stored in the database.
 """
 
-from typing import TypeVar
-
-from climate_ref.models.base import Base
+from climate_ref.models.base import Base, Table
 from climate_ref.models.dataset import Dataset
 from climate_ref.models.diagnostic import Diagnostic
 from climate_ref.models.execution import (
@@ -17,9 +15,6 @@ from climate_ref.models.execution import (
 from climate_ref.models.metric_value import MetricValue, ScalarMetricValue, SeriesMetricValue
 from climate_ref.models.provider import Provider
 
-Table = TypeVar("Table", bound=Base)
-
-
 __all__ = [
     "Base",
     "Dataset",

climate_ref/models/base.py

@@ -1,8 +1,7 @@
-import datetime
-from typing import Any
+from typing import Any, TypeVar
 
-from sqlalchemy import JSON, MetaData, func
-from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
+from sqlalchemy import JSON, MetaData
+from sqlalchemy.orm import DeclarativeBase
 
 
 class Base(DeclarativeBase):
@@ -28,19 +27,4 @@ class Base(DeclarativeBase):
     )
 
 
-class CreatedUpdatedMixin:
-    """
-    Mixin for models that have a created_at and updated_at fields
-    """
-
-    created_at: Mapped[datetime.datetime] = mapped_column(server_default=func.now())
-    """
-    When the dataset was added to the database
-    """
-
-    updated_at: Mapped[datetime.datetime] = mapped_column(
-        server_default=func.now(), onupdate=func.now(), index=True
-    )
-    """
-    When the dataset was updated.
-    """
+Table = TypeVar("Table", bound=Base)

climate_ref/models/dataset.py

@@ -172,6 +172,7 @@ class Obs4MIPsDataset(Dataset):
     units: Mapped[str] = mapped_column()
     variable_id: Mapped[str] = mapped_column()
     variant_label: Mapped[str] = mapped_column()
+    version: Mapped[str] = mapped_column()
     vertical_levels: Mapped[int] = mapped_column()
     source_version_number: Mapped[str] = mapped_column()
 
@@ -206,6 +207,7 @@ class PMPClimatologyDataset(Dataset):
     units: Mapped[str] = mapped_column()
     variable_id: Mapped[str] = mapped_column()
     variant_label: Mapped[str] = mapped_column()
+    version: Mapped[str] = mapped_column()
     vertical_levels: Mapped[int] = mapped_column()
     source_version_number: Mapped[str] = mapped_column()
 

climate_ref/models/diagnostic.py

@@ -3,7 +3,8 @@ from typing import TYPE_CHECKING
 from sqlalchemy import ForeignKey, UniqueConstraint
 from sqlalchemy.orm import Mapped, mapped_column, relationship
 
-from climate_ref.models.base import Base, CreatedUpdatedMixin
+from climate_ref.models.base import Base
+from climate_ref.models.mixins import CreatedUpdatedMixin
 
 if TYPE_CHECKING:
     from climate_ref.models.execution import ExecutionGroup