climate-ref 0.6.0__tar.gz → 0.6.1__tar.gz

{climate_ref-0.6.0 → climate_ref-0.6.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: climate-ref
-Version: 0.6.0
+Version: 0.6.1
 Summary: Application which runs the CMIP Rapid Evaluation Framework
 Author-email: Jared Lewis <jared.lewis@climate-resource.com>, Mika Pflueger <mika.pflueger@climate-resource.com>, Bouwe Andela <b.andela@esciencecenter.nl>, Jiwoo Lee <lee1043@llnl.gov>, Min Xu <xum1@ornl.gov>, Nathan Collier <collierno@ornl.gov>, Dora Hegedus <dora.hegedus@stfc.ac.uk>
 License-Expression: Apache-2.0
@@ -25,6 +25,7 @@ Requires-Dist: climate-ref-core
 Requires-Dist: ecgtools>=2024.7.31
 Requires-Dist: environs>=11.0.0
 Requires-Dist: loguru>=0.7.2
+Requires-Dist: parsl>=2025.5.19
 Requires-Dist: platformdirs>=4.3.6
 Requires-Dist: sqlalchemy>=2.0.36
 Requires-Dist: tomlkit>=0.13.2

{climate_ref-0.6.0 → climate_ref-0.6.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "climate-ref"
-version = "0.6.0"
+version = "0.6.1"
 description = "Application which runs the CMIP Rapid Evaluation Framework"
 readme = "README.md"
 authors = [
@@ -39,7 +39,8 @@ dependencies = [
     "loguru>=0.7.2",
     "ecgtools>=2024.7.31",
     "platformdirs>=4.3.6",
-    "tqdm>=4.67.1"
+    "tqdm>=4.67.1",
+    "parsl>=2025.5.19"
 ]
 
 [project.optional-dependencies]

{climate_ref-0.6.0 → climate_ref-0.6.1}/src/climate_ref/executor/__init__.py

@@ -9,8 +9,9 @@ The simplest executor is the `LocalExecutor`, which runs the diagnostic in the s
 This is useful for local testing and debugging.
 """
 
+from .hpc import HPCExecutor
 from .local import LocalExecutor
 from .result_handling import handle_execution_result
 from .synchronous import SynchronousExecutor
 
-__all__ = ["LocalExecutor", "SynchronousExecutor", "handle_execution_result"]
+__all__ = ["HPCExecutor", "LocalExecutor", "SynchronousExecutor", "handle_execution_result"]

climate_ref-0.6.1/src/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-0.6.0 → climate_ref-0.6.1}/src/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
@@ -124,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(

climate_ref-0.6.1/src/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-0.6.1/tests/unit/executor/test_hpc_executor.py

@@ -0,0 +1,85 @@
+import re
+from unittest.mock import MagicMock, patch
+
+import parsl
+import pytest
+from parsl.dataflow import futures
+
+from climate_ref.executor.hpc import HPCExecutor, execute_locally
+from climate_ref.executor.local import ExecutionFuture
+from climate_ref_core.diagnostics import ExecutionResult
+from climate_ref_core.exceptions import DiagnosticError, ExecutionError
+from climate_ref_core.executor import Executor
+
+
+def test_execute_locally_failed(definition_factory, mock_diagnostic):
+    mock_diagnostic.run = lambda definition: 1 / 0
+
+    # execution raises an exception
+    with pytest.raises(DiagnosticError):
+        result = execute_locally(
+            definition_factory(diagnostic=mock_diagnostic),
+            log_level="DEBUG",
+            raise_error=True,
+        )
+
+        assert result is None
+
+
+class TestHPCExecutor:
+    def test_is_executor(self, tmp_path):
+        executor = HPCExecutor(log_dir=tmp_path / "parsl_runinfo")
+
+        assert executor.name == "hpc"
+        assert isinstance(executor, Executor)
+        parsl.dfk().cleanup()
+
+    def test_run_metric(self, metric_definition, provider, mock_diagnostic, mocker, caplog, tmp_path):
+        with patch.object(HPCExecutor, "run", autospec=True) as mock_run:
+            # Configure the mock to behave similarly to the original
+            mock_run.side_effect = lambda self, definition, execution=None: (
+                self.parsl_results.append(
+                    ExecutionFuture(
+                        future=MagicMock(),  # Mock the future object
+                        definition=definition,
+                        execution_id=execution.id if execution else None,
+                    )
+                )
+            )
+
+            executor = HPCExecutor(log_dir=tmp_path / "parsl_runinfo")
+
+            # shall have the SerializationError, but not raised
+            executor.run(metric_definition, None)
+            assert len(executor.parsl_results) == 1
+            assert executor.parsl_results[0].definition == metric_definition
+            assert executor.parsl_results[0].execution_id is None
+
+        parsl.dfk().cleanup()
+
+    def test_join(self, metric_definition, tmp_path):
+        executor = HPCExecutor(log_dir=tmp_path / "parsl_runinfo")
+        future = futures.AppFuture(1)
+        executor.parsl_results = [ExecutionFuture(future, definition=metric_definition, execution_id=None)]
+
+        future.set_result(
+            ExecutionResult(
+                definition=metric_definition,
+                successful=False,
+                output_bundle_filename=None,
+                metric_bundle_filename=None,
+            )
+        )
+        executor.join(0.1)
+
+        assert len(executor.parsl_results) == 0
+
+    def test_join_exception(self, metric_definition, tmp_path):
+        executor = HPCExecutor(log_dir=tmp_path / "parsl_runinfo")
+        future = futures.AppFuture(1)
+        executor.parsl_results = [ExecutionFuture(future, definition=metric_definition, execution_id=None)]
+
+        future.set_exception(ValueError("Some thing bad went wrong"))
+
+        with pytest.raises(ExecutionError, match=re.escape("Failed to execute 'mock_provider/mock/key'")):
+            executor.join(0.1)

climate_ref-0.6.1/tests/unit/test_slurm.py

@@ -0,0 +1,359 @@
+from unittest.mock import MagicMock
+
+import pytest
+
+from climate_ref.slurm import SlurmChecker
+
+
+@pytest.fixture
+def create_mock_association():
+    """Factory fixture for creating mock objects"""
+
+    def _create_mock_association(**kwargs):
+        """
+        Creates a mock PySlurm association object with configurable attributes.
+
+        Args:
+            **kwargs: Key-value pairs to set as attributes on the mock association
+
+        Returns:
+            MagicMock: Configured mock association object
+        """
+        # Default values for association attributes
+        defaults = {
+            "account": "test_account",
+            "cluster": "test_cluster",
+            "comment": None,
+            "group_jobs": "UNLIMITED",
+            "group_jobs_accrue": "UNLIMITED",
+            "group_submit_jobs": "UNLIMITED",
+            "group_tres": None,
+            "group_tres_mins": None,
+            "group_tres_run_mins": None,
+            "group_wall_time": "UNLIMITED",
+            "is_default": False,
+            "max_jobs_accrue": 4,
+            "max_submit_jobs": "UNLIMITED",
+            "max_tres_mins_per_job": None,
+            "max_tres_per_job": None,
+            "max_tres_per_node": None,
+            "max_tres_run_mins_per_user": None,
+            "max_wall_time_per_job": "UNLIMITED",
+            "max_jobs": 10,
+            "max_wall_pj": "1-00:00:00",
+            "min_priority_threshold": "UNLIMITED",
+            "partition": "test_partition",
+            "shares": 100,
+            "qos": ["normal"],
+            "priority": 1000,
+            "id": 1,
+            "lft": 0,
+            "parent_acct": None,
+            "parent_account_id": 4,
+            "rgt": 0,
+            "user": "test_user",
+        }
+
+        # Update defaults with any provided kwargs
+        defaults.update(kwargs)
+
+        # Create the mock association
+        # mock_assoc = MagicMock(spec=['account', 'partition'])
+        mock_assoc = MagicMock()
+
+        # Set attributes on the mock
+        for key, value in defaults.items():
+            setattr(mock_assoc, key, value)
+
+        return mock_assoc
+
+    return _create_mock_association
+
+
+@pytest.fixture
+def create_mock_partition():
+    """Factory fixture for creating mock objects"""
+
+    def _create_mock_partition(**kwargs):
+        """
+        Creates a mock PySlurm partition object with configurable attributes.
+
+        Args:
+            **kwargs: Key-value pairs to set as attributes on the mock partition
+
+        Returns:
+            MagicMock: Configured mock partition object
+        """
+
+        # Default values for partition attributes
+        defaults = {
+            "allow_root_jobs": False,
+            "allowed_accounts": ["ALL"],
+            "allowed_groups": ["ALL"],
+            "allowed_qos": ["ALL"],
+            "allowed_submit_nodes": ["ALL"],
+            "alternate": None,
+            "cluster": "test_cluster",
+            "cpu_binding": None,
+            "default_cpus_per_gpu": None,
+            "default_memory_per_cpu": None,
+            "default_memory_per_gpu": None,
+            "default_memory_per_node": "UNLIMITED",
+            "default_time": 720,
+            "denied_accounts": ["ALL"],
+            "denied_qos": ["ALL"],
+            "is_default": False,
+            "is_hidden": False,
+            "is_root_only": False,
+            "is_user_exclusive": False,
+            "least_loaded_nodes_scheduling": False,
+            "max_cpus_per_node": "UNLIMITED",
+            "max_cpus_per_socket": "UNLIMITED",
+            "max_memory_per_cpu": None,
+            "max_memory_per_node": "UNLIMITED",
+            "max_nodes": "UNLIMITED",
+            "max_time": 7200,
+            "min_nodes": 0,
+            "name": "batch",
+            "nodes": "baseline[3-140]",
+            "nodesets": [],
+            "over_time_limit": None,
+            "oversubscribe": "EXCLUSIVE",
+            "preempt_mode": "OFF",
+            "preemption_grace_time": None,
+            "priority_job_factor": 1,
+            "priority_tier": 1,
+            "qos": None,
+            "requires_reservation": False,
+            "select_type_parameters": [],
+            "state": "UP",
+            "total_cpus": 17664,
+            "total_nodes": 138,
+        }
+
+        # Update defaults with any provided kwargs
+        defaults.update(kwargs)
+
+        # Create the mock partition
+        mock_part = MagicMock()
+
+        # Set attributes on the mock
+        for key, value in defaults.items():
+            setattr(mock_part, key, value)
+
+        def mock_to_dict():
+            return defaults
+
+        mock_part.to_dict.side_effect = mock_to_dict
+
+        return mock_part
+
+    return _create_mock_partition
+
+
+@pytest.fixture
+def create_mock_qos():
+    """Factory fixture for creating mock objects"""
+
+    def _create_mock_qos(**kwargs):
+        """
+        Creates a mock PySlurm QoS object with configurable attributes.
+
+        Args:
+            **kwargs: Key-value pairs to set as attributes on the mock QoS
+
+        Returns:
+            MagicMock: Configured mock QoS object
+        """
+        # Default values for QoS attributes
+        defaults = {
+            "name": "normal",
+            "description": "Default Quality of Service",
+            "flags": 0,
+            "grace_time": 0,
+            "grp_jobs": 4294967295,
+            "grp_submit_jobs": 4294967295,
+            "grp_tres": None,
+            "grp_tres_mins": None,
+            "grp_tres_run_mins": None,
+            "grp_wall": 4294967295,
+            "max_jobs_pu": 4294967295,
+            "max_submit_jobs_pu": 4294967295,
+            "max_tres_mins_pj": None,
+            "max_tres_pj": None,
+            "max_tres_pn": None,
+            "max_tres_pu": None,
+            "max_tres_run_mins_pu": None,
+            "max_wall_pj": 4294967295,
+            "min_tres_pj": None,
+            "preempt_mode": "OFF",
+            "priority": 0,
+            "usage_factor": 1.0,
+            "usage_thres": 4294967295.0,
+        }
+
+        # Update defaults with any provided kwargs
+        defaults.update(kwargs)
+
+        # Create the mock QoS
+        mock_qos = MagicMock()
+
+        # Set attributes on the mock
+        for key, value in defaults.items():
+            setattr(mock_qos, key, value)
+
+        return mock_qos
+
+    return _create_mock_qos
+
+
+@pytest.fixture
+def create_mock_node():
+    """Factory fixture for creating mock objects"""
+
+    def _create_mock_node(**kwargs):
+        """
+        Creates a mock PySlurm node object with configurable attributes.
+
+        Args:
+            **kwargs: Key-value pairs to set as attributes on the mock node
+
+        Returns:
+            MagicMock: Configured mock node object
+        """
+        # Default values for node attributes
+        defaults = {
+            "active_features": [],
+            "address": "test_address",
+            "allocated_cpus": 0,
+            "allocated_gres": {},
+            "allocated_memory": 0,
+            "allocated_tres": {},
+            "architecture": "x86_64",
+            "available_features": [],
+            "avg_watts": 0,
+            "bcast_address": None,
+            "boards": 1,
+            "boot_time": 1747765065,
+            "cap_watts": 0,
+            "cluster": "baseline",
+            "comment": None,
+            "configured_gres": {},
+            "cores_per_socket": 64,
+            "cores_reserved_for_system": None,
+            "cpu_binding": None,
+            "cpu_load": 0.18,
+            "current_watts": 0,
+            "effective_cpus": 128,
+            "external_sensors": {"joules_total": None, "current_watts": None, "temperature": None},
+            "extra": None,
+            "free_memory": 231147,
+            "hostname": "baseline101",
+            "idle_cpus": 128,
+            "idle_memory": 256000,
+            "last_busy_time": 1747789216,
+            "mcs_label": None,
+            "memory_reserved_for_system": None,
+            "name": "baseline101",
+            "next_state": None,
+            "operating_system": "Linux 4.18.0-553.46.1.el8_10.x86_64 #1 SMP Sat Mar 15 01:37:33 EDT 2025",
+            "owner": None,
+            "partitions": ["batch", "batch_low_memory"],
+            "real_memory": 256000,
+            "reason": None,
+            "reason_time": None,
+            "reason_user": None,
+            "slurm_version": "23.02.7",
+            "slurmd_port": 6818,
+            "slurmd_start_time": 1747765241,
+            "sockets": 2,
+            "state": "IDLE",
+            "temporary_disk": None,
+            "threads_per_core": 1,
+            "total_cpus": 128,
+            "weight": 1,
+        }
+
+        # Update defaults with any provided kwargs
+        defaults.update(kwargs)
+
+        # Create the mock node
+        mock_node = MagicMock()
+
+        # Set attributes on the mock
+        for key, value in defaults.items():
+            setattr(mock_node, key, value)
+
+        return mock_node
+
+    return _create_mock_node
+
+
+def test_slurm_checker(
+    create_mock_association, create_mock_partition, create_mock_qos, create_mock_node, mocker
+):
+    # Setup mock data
+    associations = {
+        1: create_mock_association(account="climate_ref1", partition="cpu"),
+        2: create_mock_association(account="climate_ref2", partition="gpu"),
+    }
+
+    partitions = {
+        "normal": create_mock_partition(name="nomral"),
+        "cpu": create_mock_partition(name="cpu"),
+        "batch": create_mock_partition(name="batch"),
+        "gpu": create_mock_partition(name="gpu"),
+    }
+
+    qoss = {
+        "normal": create_mock_qos(name="normal"),
+    }
+    nodes = {
+        "node0001": create_mock_node(name="node0001"),
+    }
+
+    mock_pyslurm = MagicMock()
+
+    # Set up the nested structure
+    mock_pyslurm.db = MagicMock()
+    mock_pyslurm.db.Associations = MagicMock()
+    mock_pyslurm.db.Associations.load.return_value = associations
+
+    mock_pyslurm.Partitions = MagicMock()
+    mock_pyslurm.Partitions.load.return_value = partitions
+
+    mock_pyslurm.Nodes = MagicMock()
+    mock_pyslurm.Nodes.load.return_value = nodes
+
+    # Mock QoS
+    mock_qos_instance = MagicMock()
+    mock_qos_instance.get.return_value = qoss
+    mock_pyslurm.qos.return_value = mock_qos_instance
+
+    # Patch the module before import
+    mocker.patch.dict("sys.modules", {"pyslurm": mock_pyslurm})
+
+    # Create and test the checker
+    checker = SlurmChecker(intest=True)
+
+    assert checker.can_account_use_partition("climate_ref1", "cpu") is True
+    assert checker.can_account_use_partition("climate_ref2", "gpu") is True
+    assert checker.can_account_use_partition("climate", "nonexistent") is False
+
+    assert checker.can_account_use_qos("climate_ref1", "normal") is True
+    assert checker.can_account_use_qos("climate_ref3", "normal") is False
+
+    assert checker.get_partition_limits("cpu") == {
+        "max_time_minutes": 7200,
+        "default_time_minutes": 720,
+        "max_nodes": 1,
+        "total_nodes": 138,
+        "total_cpus": 17664,
+    }
+    assert checker.check_account_partition_access_with_limits("climate_ref2", "gpu") == {
+        "account_exists": True,
+        "partition_exists": True,
+        "has_access": True,
+        "time_limits": 7200,
+        "error": "none",
+    }