From 77c6faba59759f172e9dca13c757694d9cbca4e4 Mon Sep 17 00:00:00 2001 From: David Berenstein Date: Sun, 16 Aug 2026 10:30:43 +0200 Subject: [PATCH 1/2] feat: warn when SLURM ranks double count a node With several ranks per node and tracking_mode="machine", every rank measures the whole node, so the job's reported footprint is silently multiplied by the rank count. Read SLURM_NTASKS_PER_NODE at tracker init and say so, since a docs warning does not catch anyone. Split out of #1366, where it was unrelated scope. Co-Authored-By: Claude Opus 5 (1M context) --- codecarbon/core/slurm.py | 26 ++++++++++++++++++++++++++ codecarbon/emissions_tracker.py | 2 ++ tests/test_slurm.py | 24 ++++++++++++++++++++++++ 3 files changed, 52 insertions(+) create mode 100644 codecarbon/core/slurm.py create mode 100644 tests/test_slurm.py diff --git a/codecarbon/core/slurm.py b/codecarbon/core/slurm.py new file mode 100644 index 000000000..9c45c9100 --- /dev/null +++ b/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'." + ) diff --git a/codecarbon/emissions_tracker.py b/codecarbon/emissions_tracker.py index 531c628b1..007d1e598 100644 --- a/codecarbon/emissions_tracker.py +++ b/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() diff --git a/tests/test_slurm.py b/tests/test_slurm.py new file mode 100644 index 000000000..0b72fc4f5 --- /dev/null +++ b/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") + ) From 1a45b547558dd3f140897091d716419454edf0f6 Mon Sep 17 00:00:00 2001 From: David Berenstein Date: Wed, 23 Sep 2026 16:37:39 +0900 Subject: [PATCH 2/2] fix(slurm): warn only on non-zero SLURM_LOCALID, read step tasks-per-node Co-Authored-By: Claude Opus 5.5 (1M context) --- codecarbon/core/slurm.py | 34 +++++++++++++++++++--------------- tests/test_slurm.py | 30 +++++++++++++++++++++++------- 2 files changed, 42 insertions(+), 22 deletions(-) diff --git a/codecarbon/core/slurm.py b/codecarbon/core/slurm.py index 9c45c9100..47d73d401 100644 --- a/codecarbon/core/slurm.py +++ b/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'." + ) diff --git a/tests/test_slurm.py b/tests/test_slurm.py index 0b72fc4f5..ea23163fb 100644 --- a/tests/test_slurm.py +++ b/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)