Merge pull request #392 from Climate-REF/update-conda-for-pmp

Specify environment variables when running PMP diagnostic

Author: bouweandela

changelog/392.feature.md

@@ -0,0 +1 @@
+Made it possible to run PMP diagnostics without having conda installed.

packages/climate-ref-core/src/climate_ref_core/providers.py

@@ -232,6 +232,27 @@ def _get_micromamba_url() -> str:
 class CondaDiagnosticProvider(CommandLineDiagnosticProvider):
     """
     A provider for diagnostics that can be run from the command line in a conda environment.
+
+    Parameters
+    ----------
+    name
+        The name of the provider.
+    version
+        The version of the provider.
+    slug
+        A slugified version of the name.
+    repo
+        URL of the git repository to install a development version of the package from.
+    tag_or_commit
+        Tag or commit to install from the `repo` repository.
+
+    Attributes
+    ----------
+    env_vars
+        Environment variables to set when running commands in the conda environment.
+    url
+        URL to install a development version of the package from.
+
     """
 
     def __init__(
@@ -246,6 +267,7 @@ def __init__(
         self._conda_exe: Path | None = None
         self._prefix: Path | None = None
         self.url = f"git+{repo}@{tag_or_commit}" if repo and tag_or_commit else None
+        self.env_vars: dict[str, str] = {}
 
     @property
     def prefix(self) -> Path:
@@ -404,6 +426,8 @@ def run(self, cmd: Iterable[str]) -> None:
             *cmd,
         ]
         logger.info(f"Running '{' '.join(cmd)}'")
+        env_vars = os.environ.copy()
+        env_vars.update(self.env_vars)
         try:
             # This captures the log output until the execution is complete
             # We could poll using `subprocess.Popen` if we want something more responsive
@@ -413,6 +437,7 @@ def run(self, cmd: Iterable[str]) -> None:
                 stdout=subprocess.PIPE,
                 stderr=subprocess.STDOUT,
                 text=True,
+                env=env_vars,
             )
             logger.info("Command output: \n" + res.stdout)
             logger.info("Command execution successful")

packages/climate-ref-core/tests/unit/test_providers.py

@@ -294,6 +294,12 @@ def test_run(self, mocker, tmp_path, provider, env_exists):
             ):
                 provider.run(["mock-command"])
         else:
+            mocker.patch.object(
+                climate_ref_core.providers.os.environ,
+                "copy",
+                return_value={"existing_var": "existing_value"},
+            )
+            provider.env_vars = {"test_var": "test_value"}
             provider.run(["mock-command"])
 
             run.assert_called_with(
@@ -308,4 +314,5 @@ def test_run(self, mocker, tmp_path, provider, env_exists):
                 stdout=subprocess.PIPE,
                 stderr=subprocess.STDOUT,
                 text=True,
+                env={"existing_var": "existing_value", "test_var": "test_value"},
             )

packages/climate-ref-pmp/src/climate_ref_pmp/__init__.py

@@ -2,17 +2,45 @@
 Rapid evaluating CMIP data
 """
 
+from __future__ import annotations
+
 import importlib.metadata
+import os
+from typing import TYPE_CHECKING
+
+from loguru import logger
 
 from climate_ref_core.dataset_registry import DATASET_URL, dataset_registry_manager
 from climate_ref_core.providers import CondaDiagnosticProvider
 from climate_ref_pmp.diagnostics import ENSO, AnnualCycle, ExtratropicalModesOfVariability
 
+if TYPE_CHECKING:
+    from climate_ref.config import Config
+
 __version__ = importlib.metadata.version("climate-ref-pmp")
 
+
 # Create the PMP diagnostics provider
 # PMP uses a conda environment to run the diagnostics
-provider = CondaDiagnosticProvider("PMP", __version__)
+class PMPDiagnosticProvider(CondaDiagnosticProvider):
+    """
+    Provider for PMP diagnostics.
+    """
+
+    def configure(self, config: Config) -> None:
+        """Configure the provider."""
+        super().configure(config)
+        self.env_vars["PCMDI_CONDA_EXE"] = str(self.get_conda_exe())
+        # This is a workaround for a fatal error in internal_Finalize of MPICH
+        # when running in a conda environment on MacOS.
+        # It is not clear if this is a bug in MPICH or a problem with the conda environment.
+        if "FI_PROVIDER" not in os.environ:  # pragma: no branch
+            logger.debug("Setting env variable 'FI_PROVIDER=tcp'")
+            self.env_vars["FI_PROVIDER"] = "tcp"
+
+
+provider = PMPDiagnosticProvider("PMP", __version__)
+
 
 # Annual cycle diagnostics and metrics
 provider.register(AnnualCycle())

packages/climate-ref-pmp/src/climate_ref_pmp/pmp_driver.py

@@ -159,13 +159,6 @@ def build_pmp_command(
     # Note this uses the driver script from the REF env *not* the PMP conda env
     _parameter_file = _get_resource("climate_ref_pmp.params", parameter_file, use_resources=True)
 
-    # This is a workaround for a fatal error in internal_Finalize of MPICH
-    # when running in a conda environment on MacOS.
-    # It is not clear if this is a bug in MPICH or a problem with the conda environment.
-    if "FI_PROVIDER" not in os.environ:  # pragma: no branch
-        logger.debug("Setting env variable 'FI_PROVIDER=tcp'")
-        os.environ["FI_PROVIDER"] = "tcp"
-
     # Run the driver script inside the PMP conda environment
     cmd = [
         driver_file,