diff --git a/codecarbon/core/slurm.py b/codecarbon/core/slurm.py new file mode 100644 index 000000000..47d73d401 --- /dev/null +++ b/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'." + ) 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..ea23163fb --- /dev/null +++ b/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)