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 96ed00c91..ba6c28652 100644 --- a/codecarbon/emissions_tracker.py +++ b/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() 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") + )