feat: warn when SLURM ranks double count a node

codecarbon/core/slurm.py

@@ -0,0 +1,26 @@
+import os
+import re
+
+from codecarbon.external.logger import logger
+
+
+def warn_on_multi_rank_double_counting(tracking_mode: str) -> None:
+    """
+    Warn when several ranks share a node and each one measures the whole node.
+
+    In ``machine`` mode every rank reports the node's power, so the job's
+    reported footprint is silently multiplied by the number of ranks per node.
+    """
+    if tracking_mode != "machine":
+        return
+    # SLURM writes this as "4", or as "4(x2)" for a heterogeneous allocation.
+    ntasks_per_node = os.environ.get("SLURM_NTASKS_PER_NODE", "")
+    match = re.match(r"\d+", ntasks_per_node)
+    if match and int(match.group()) > 1:
+        logger.warning(
+            f"SLURM_NTASKS_PER_NODE is {ntasks_per_node} and tracking_mode is "
+            "'machine': every rank measures the whole node, so the job's total "
+            "will be multiplied by the number of ranks per node. Start the "
+            "tracker on one rank per node (SLURM_LOCALID == 0), or use "
+            "tracking_mode='process'."
+        )

codecarbon/emissions_tracker.py

@@ -21,6 +21,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
@@ -600,6 +601,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,24 @@
+import unittest
+from unittest import mock
+
+from codecarbon.core.slurm import warn_on_multi_rank_double_counting
+
+
+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_several_ranks_per_node_in_machine_mode_warns(self):
+        self.assertEqual(1, self._warnings(SLURM_NTASKS_PER_NODE="4"))
+        # Heterogeneous allocations are written "4(x2)".
+        self.assertEqual(1, self._warnings(SLURM_NTASKS_PER_NODE="4(x2)"))
+
+    def test_no_warning_without_double_counting(self):
+        self.assertEqual(0, self._warnings(SLURM_NTASKS_PER_NODE="1"))
+        self.assertEqual(0, self._warnings())
+        self.assertEqual(
+            0, self._warnings(tracking_mode="process", SLURM_NTASKS_PER_NODE="4")
+        )