fix(slurm): warn only on non-zero SLURM_LOCALID, read step tasks-per-node

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

Author: davidberenstein1957

codecarbon/core/slurm.py

@@ -1,26 +1,30 @@
 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.
+    Warn when a machine-mode tracker runs on a non-zero SLURM local rank.
 
-    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.
+    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
-    # 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'."
-        )
+    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'."
+    )

tests/test_slurm.py

@@ -2,6 +2,7 @@
 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):
@@ -11,14 +12,29 @@ def _warnings(self, tracking_mode="machine", **env):
                 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_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):
-        self.assertEqual(0, self._warnings(SLURM_NTASKS_PER_NODE="1"))
-        self.assertEqual(0, self._warnings())
+        # Local rank 0 is the recommended tracker.
         self.assertEqual(
-            0, self._warnings(tracking_mode="process", SLURM_NTASKS_PER_NODE="4")
+            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)