Add list mappings

Author: HCookie

src/anemoi/inference/clusters/distributed.py

@@ -7,7 +7,6 @@
 # nor does it submit to any jurisdiction.
 #
 
-import os
 
 from anemoi.inference.clusters import cluster_registry
 from anemoi.inference.clusters.mapping import EnvMapping
@@ -32,4 +31,4 @@ def __init__(self) -> None:
 
     @classmethod
     def used(cls) -> bool:
-        return DISTRIBUTED_MAPPING.global_rank in os.environ and DISTRIBUTED_MAPPING.local_rank in os.environ
+        return bool(DISTRIBUTED_MAPPING.get_env("world_size")) and bool(DISTRIBUTED_MAPPING.get_env("global_rank"))

src/anemoi/inference/clusters/mapping.py

@@ -10,6 +10,7 @@
 import dataclasses
 import logging
 import os
+from typing import Any
 
 from anemoi.inference.clusters import cluster_registry
 from anemoi.inference.clusters.client import ComputeClientFactory
@@ -19,16 +20,34 @@
 
 @dataclasses.dataclass
 class EnvMapping:
-    local_rank: str
-    global_rank: str
-    world_size: str
+    """Dataclass to hold environment variable mappings for cluster configuration.
 
-    master_addr: str
-    master_port: str
+    Elements can be either strings or lists of strings.
+    If a list is provided, the first found environment variable will be used.
+    """
+
+    local_rank: str | list[str]
+    global_rank: str | list[str]
+    world_size: str | list[str]
+
+    master_addr: str | list[str]
+    master_port: str | list[str]
 
     backend: str | None = None
     init_method: str = "env://"
 
+    def get_env(self, key: str, default: Any = None):
+        """Get the environment variable value for the given key."""
+        mapped_value = getattr(self, key)
+        if mapped_value is None:
+            return default
+
+        for env_var in (mapped_value if isinstance(mapped_value, list) else [mapped_value]):
+            value = os.environ.get(env_var)
+            if value is not None:
+                return value
+        return default
+
 
 @cluster_registry.register("custom")
 class MappingCluster(ComputeClientFactory):
@@ -88,27 +107,27 @@ def backend(self) -> str:
     @property
     def world_size(self) -> int:
         """Return the total number of processes in the cluster."""
-        return int(os.environ.get(self._mapping.world_size, 1))
+        return int(self._mapping.get_env("world_size", 1))
 
     @property
     def global_rank(self) -> int:
         """Return the rank of the current process."""
-        return int(os.environ.get(self._mapping.global_rank, 0))
+        return int(self._mapping.get_env("global_rank", 0))
 
     @property
     def local_rank(self) -> int:
         """Return the rank of the current process."""
-        return int(os.environ.get(self._mapping.local_rank, self.global_rank))
+        return int(self._mapping.get_env("local_rank", self.global_rank))
 
     @property
     def master_addr(self) -> str:
         """Return the master address."""
-        return os.environ.get(self._mapping.master_addr, "")
+        return self._mapping.get_env("master_addr", "")
 
     @property
     def master_port(self) -> int:
         """Return the master port."""
-        return int(os.environ.get(self._mapping.master_port, 0))
+        return int(self._mapping.get_env("master_port", 0))
 
     @classmethod
     def used(cls) -> bool:

src/anemoi/inference/clusters/mpi.py

@@ -8,7 +8,6 @@
 #
 
 import logging
-import os
 
 from anemoi.inference.clusters import cluster_registry
 from anemoi.inference.clusters.mapping import EnvMapping
@@ -18,9 +17,9 @@
 LOG = logging.getLogger(__name__)
 
 MPI_MAPPING = EnvMapping(
-    local_rank="OMPI_COMM_WORLD_LOCAL_RANK",
-    global_rank="OMPI_COMM_WORLD_RANK",
-    world_size="OMPI_COMM_WORLD_SIZE",
+    local_rank=["OMPI_COMM_WORLD_LOCAL_RANK", "PMI_RANK"],
+    global_rank=["OMPI_COMM_WORLD_RANK", "PMI_RANK"],
+    world_size=["OMPI_COMM_WORLD_SIZE", "PMI_SIZE"],
     master_addr="MASTER_ADDR",
     master_port="MASTER_PORT",
     init_method="tcp://{master_addr}:{master_port}",
@@ -44,7 +43,7 @@ def __init__(self, use_mpi_backend: bool = False, **kwargs) -> None:
 
     @classmethod
     def used(cls) -> bool:
-        return MPI_MAPPING.world_size in os.environ or "PMI_SIZE" in os.environ
+        return bool(MPI_MAPPING.get_env("world_size"))
 
     @property
     def backend(self) -> str:

src/anemoi/inference/clusters/slurm.py

@@ -42,7 +42,7 @@ def __init__(self) -> None:
     def used(cls) -> bool:
         # from pytorch lightning
         # https://github.com/Lightning-AI/pytorch-lightning/blob/a944e7744e57a5a2c13f3c73b9735edf2f71e329/src/lightning/fabric/plugins/environments/slurm.py
-        return SLURM_MAPPING.world_size in os.environ and os.environ.get("SLURM_JOB_NAME") not in (
+        return bool(SLURM_MAPPING.get_env("world_size")) and os.environ.get("SLURM_JOB_NAME") not in (
             "bash",
             "interactive",
         )

tests/unit/test_clusters.py

@@ -320,7 +320,7 @@ def test_distributed_cluster_used_detection(self):
             assert not DistributedCluster.used()
 
         # In distributed environment (torchrun)
-        with patch.dict(os.environ, {"RANK": "3", "LOCAL_RANK": "1"}):
+        with patch.dict(os.environ, {"RANK": "3", "WORLD_SIZE": "4"}):
             assert DistributedCluster.used()
 
     def test_distributed_cluster_initialization(self):
@@ -422,6 +422,35 @@ def test_mapping_cluster_not_used(self):
         """Test that MappingCluster.used() returns False."""
         assert not MappingCluster.used()
 
+    def test_mapping_cluster_list(self):
+        """Test MappingCluster with list mapping, such that the first found env var is used."""
+        mapping = EnvMapping(
+            local_rank=["MY_LOCAL_RANK_A", "MY_LOCAL_RANK_B"],
+            global_rank=["MY_GLOBAL_RANK_A", "MY_GLOBAL_RANK_B"],
+            world_size=["MY_WORLD_SIZE_A", "MY_WORLD_SIZE_B"],
+            master_addr=["MY_MASTER_ADDR_A", "MY_MASTER_ADDR_B"],
+            master_port=["MY_MASTER_PORT_A", "MY_MASTER_PORT_B"],
+            init_method="tcp://{master_addr}:{master_port}",
+        )
+
+        with patch.dict(
+            os.environ,
+            {
+                "MY_WORLD_SIZE_B": "16",
+                "MY_GLOBAL_RANK_A": "5",
+                "MY_LOCAL_RANK_B": "2",
+                "MY_MASTER_ADDR_A": "192.168.1.1",
+                "MY_MASTER_PORT_B": "40000",
+            },
+        ):
+            cluster = MappingCluster(mapping=mapping)
+
+            assert cluster.world_size == 16
+            assert cluster.global_rank == 5
+            assert cluster.local_rank == 2
+            assert cluster.master_addr == "192.168.1.1"
+            assert cluster.master_port == 40000
+
 
 class TestClusterRegistry:
     """Tests for cluster registry and creation."""