feat: warn when SLURM ranks double count a node

codecarbon/core/slurm.py

@@ -0,0 +1,30 @@
+import os
+
+from codecarbon.external.logger import logger
+
+
+def warn_on_multi_rank_double_counting(tracking_mode: str) -> None:
+    """
+    Warn when a machine-mode tracker runs on a non-zero SLURM local rank.
+
+    In ``machine`` mode every tracker reports the whole node's power, so if
+    more than one rank per node starts a tracker, the job's footprint is
+    silently multiplied. Local rank 0 (the recommended tracker) and the batch
+    step (``SLURM_LOCALID`` unset) never warn.
+    """
+    if tracking_mode != "machine":
+        return
+    local_id = os.environ.get("SLURM_LOCALID", "")
+    if not local_id.isdigit() or int(local_id) == 0:
+        return
+    # Step-scoped and always set by srun; written "4", or "4(x2)".
+    tasks_per_node = os.environ.get(
+        "SLURM_STEP_TASKS_PER_NODE", os.environ.get("SLURM_TASKS_PER_NODE", "?")
+    )
+    logger.warning(
+        f"tracking_mode is 'machine' on SLURM local rank {local_id} "
+        f"(tasks per node: {tasks_per_node}). If more than one rank per node "
+        "starts a tracker, each one measures the whole node and the job's total "
+        "is multiplied. Start the tracker only on SLURM_LOCALID == 0, or use "
+        "tracking_mode='process'."
+    )

codecarbon/emissions_tracker.py

@@ -23,6 +23,7 @@
 
 from codecarbon._version import __version__
 from codecarbon.core.config import get_hierarchical_config, normalize_gpu_ids
+from codecarbon.core.slurm import warn_on_multi_rank_double_counting
 from codecarbon.core.units import Energy, Power, Time, Water
 from codecarbon.core.util import count_cpus, count_physical_cpus, suppress
 from codecarbon.external.hardware import CPU, GPU, AppleSiliconChip
@@ -620,6 +621,7 @@ def __init__(
         assert self._tracking_mode in ["machine", "process"]
         set_logger_level(self._log_level)
         set_logger_format(self._logger_preamble)
+        warn_on_multi_rank_double_counting(self._tracking_mode)
         self._initialize_runtime_state()
         self._initialize_scheduler_state()
         self._initialize_emissions_context()

tests/test_slurm.py

@@ -0,0 +1,40 @@
+import unittest
+from unittest import mock
+
+from codecarbon.core.slurm import warn_on_multi_rank_double_counting
+from codecarbon.emissions_tracker import OfflineEmissionsTracker
+
+
+class TestMultiRankWarning(unittest.TestCase):
+    def _warnings(self, tracking_mode="machine", **env):
+        with mock.patch.dict("os.environ", env, clear=True):
+            with mock.patch("codecarbon.core.slurm.logger") as mocked_logger:
+                warn_on_multi_rank_double_counting(tracking_mode)
+        return mocked_logger.warning.call_count
+
+    def test_non_zero_local_rank_in_machine_mode_warns(self):
+        self.assertEqual(
+            1, self._warnings(SLURM_LOCALID="1", SLURM_STEP_TASKS_PER_NODE="4")
+        )
+        # Heterogeneous allocations are written "4(x2)"; falls back to the job var.
+        self.assertEqual(
+            1, self._warnings(SLURM_LOCALID="3", SLURM_TASKS_PER_NODE="4(x2)")
+        )
+
+    def test_no_warning_without_double_counting(self):
+        # Local rank 0 is the recommended tracker.
+        self.assertEqual(
+            0, self._warnings(SLURM_LOCALID="0", SLURM_STEP_TASKS_PER_NODE="4")
+        )
+        # Batch step: per-node counts are exported but SLURM_LOCALID is not.
+        self.assertEqual(
+            0, self._warnings(SLURM_NTASKS_PER_NODE="4", SLURM_TASKS_PER_NODE="4")
+        )
+        self.assertEqual(0, self._warnings())
+        self.assertEqual(0, self._warnings(tracking_mode="process", SLURM_LOCALID="1"))
+
+    def test_tracker_init_warns_on_non_zero_local_rank(self):
+        with mock.patch.dict("os.environ", {"SLURM_LOCALID": "2"}):
+            with mock.patch("codecarbon.core.slurm.logger") as mocked_logger:
+                OfflineEmissionsTracker(country_iso_code="FRA", tracking_mode="machine")
+        self.assertEqual(1, mocked_logger.warning.call_count)