diff --git a/codecarbon/core/schedulers.py b/codecarbon/core/schedulers.py new file mode 100644 index 000000000..96297c437 --- /dev/null +++ b/codecarbon/core/schedulers.py @@ -0,0 +1,74 @@ +""" +Detection of HPC batch scheduler job metadata. + +Schedulers export the identity of the running job into the environment of every +job step, so no scheduler library is needed: reading ``os.environ`` is enough. + +SLURM is detected automatically. Any other scheduler is supported through the +generic ``CODECARBON_JOB_*`` environment contract, which also takes precedence +over the auto-detected values, so a site can map its own scheduler in a couple +of lines of shell. +""" + +import os +import re +from typing import Dict + +from codecarbon.external.logger import logger + +# Field name on EmissionsData -> SLURM environment variable holding it. +SLURM_ENV_VARS = { + "job_id": "SLURM_JOB_ID", + "job_name": "SLURM_JOB_NAME", + "job_user": "SLURM_JOB_USER", + "job_account": "SLURM_JOB_ACCOUNT", + "job_partition": "SLURM_JOB_PARTITION", + "node_name": "SLURMD_NODENAME", +} + +JOB_METADATA_FIELDS = ("scheduler",) + tuple(SLURM_ENV_VARS) + + +def detect_job_metadata() -> Dict[str, str]: + """ + Collect the scheduler job metadata of the current process. + + :return: a dict with one entry per field of ``JOB_METADATA_FIELDS``, with + empty strings for everything that could not be detected. Outside of a + batch job every value is empty. + """ + metadata = {field: "" for field in JOB_METADATA_FIELDS} + + if os.environ.get("SLURM_JOB_ID"): + metadata["scheduler"] = "slurm" + for field, env_var in SLURM_ENV_VARS.items(): + metadata[field] = os.environ.get(env_var, "") + + # Explicit configuration always wins, whatever the scheduler. + for field in JOB_METADATA_FIELDS: + override = os.environ.get(f"CODECARBON_{field.upper()}") + if override: + metadata[field] = override + + return metadata + + +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. + match = re.match(r"\d+", os.environ.get("SLURM_NTASKS_PER_NODE", "")) + if match and int(match.group()) > 1: + logger.warning( + f"SLURM_NTASKS_PER_NODE is {os.environ['SLURM_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..733c5f716 100644 --- a/codecarbon/emissions_tracker.py +++ b/codecarbon/emissions_tracker.py @@ -21,6 +21,10 @@ from codecarbon._version import __version__ from codecarbon.core.config import get_hierarchical_config, normalize_gpu_ids +from codecarbon.core.schedulers import ( + detect_job_metadata, + 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 @@ -598,6 +602,7 @@ def __init__( ) assert self._tracking_mode in ["machine", "process"] + warn_on_multi_rank_double_counting(self._tracking_mode) set_logger_level(self._log_level) set_logger_format(self._logger_preamble) self._initialize_runtime_state() @@ -1098,6 +1103,7 @@ def _prepare_emissions_data(self) -> EmissionsData: tracking_mode=self._conf.get("tracking_mode"), pue=self._pue, wue=self._wue, + **detect_job_metadata(), ) logger.debug(total_emissions) return total_emissions diff --git a/codecarbon/output_methods/emissions_data.py b/codecarbon/output_methods/emissions_data.py index 17544aa51..2ca69fa47 100644 --- a/codecarbon/output_methods/emissions_data.py +++ b/codecarbon/output_methods/emissions_data.py @@ -47,6 +47,14 @@ class EmissionsData: on_cloud: str = "N" pue: float = 1 wue: float = 0 + # Batch scheduler job metadata, empty outside of an HPC job. + scheduler: str = "" + job_id: str = "" + job_name: str = "" + job_user: str = "" + job_account: str = "" + job_partition: str = "" + node_name: str = "" @property def values(self) -> OrderedDict: diff --git a/docs/how-to/slurm.md b/docs/how-to/slurm.md index 233097b8b..c3c364950 100644 --- a/docs/how-to/slurm.md +++ b/docs/how-to/slurm.md @@ -160,6 +160,57 @@ tail -f logs/.out sinfo ``` +## Job metadata in the output + +When CodeCarbon runs inside a SLURM job step it reads the job's identity from the +environment SLURM already provides and stores it on every emissions record. There is +nothing to enable and no code to change. + +| Output field | SLURM environment variable | +|--------------|----------------------------| +| `scheduler` | set to `slurm` when `SLURM_JOB_ID` is present | +| `job_id` | `SLURM_JOB_ID` | +| `job_name` | `SLURM_JOB_NAME` | +| `job_user` | `SLURM_JOB_USER` | +| `job_account` | `SLURM_JOB_ACCOUNT` | +| `job_partition` | `SLURM_JOB_PARTITION` | +| `node_name` | `SLURMD_NODENAME` | + +Outside of a job all of these are empty, so nothing changes for non-HPC users. + +This makes `emissions.csv` directly joinable against SLURM accounting: + +```bash +sacct -j 1234567 --format=JobID,JobName,Account,Partition,Elapsed,AllocTRES --parsable2 +``` + +!!! tip "You no longer need `CODECARBON_PROJECT_NAME=$SLURM_JOB_ID`" + Overloading the project name with the job ID used to be the only way to tell runs + apart. Keep `project_name` for your project and use `job_id` for the job. + +### Other schedulers + +Any field can be set, or overridden, with an environment variable named after it. This +is how PBS, LSF or OAR sites get the same columns without CodeCarbon needing to know +about their scheduler — map their variables onto ours in your job script: + +```bash +export CODECARBON_SCHEDULER=pbs +export CODECARBON_JOB_ID=$PBS_JOBID +export CODECARBON_JOB_NAME=$PBS_JOBNAME +export CODECARBON_NODE_NAME=$(hostname) +``` + +These variables take precedence over the auto-detected SLURM values, so they also work +for correcting a field on a site whose SLURM configuration is unusual. + +!!! warning "One tracker per node" + Power is a property of the node, not of a rank. If you launch CodeCarbon on every + rank of a multi-node job in `machine` tracking mode, each one measures the whole + node and your total is multiplied by the number of ranks. Start the tracker on one + rank per node (for example when `SLURM_LOCALID` is `0`), or use `process` tracking + mode. + ## Troubleshooting ### Error: AMD GPU detected but amdsmi is not properly configured diff --git a/docs/reference/output.md b/docs/reference/output.md index 720e0a883..b3099d1e8 100644 --- a/docs/reference/output.md +++ b/docs/reference/output.md @@ -64,6 +64,23 @@ The package has an in-built logger that logs data into a CSV file named `emissio | gpu_utilization_percent | Average GPU utilization during tracking period (%) | | ram_utilization_percent | Average RAM utilization during tracking period (%) | | ram_used_gb | Average RAM used during tracking period (GB) | +| scheduler | Batch scheduler that started the job, e.g. `slurm`. Empty outside of an HPC job | +| job_id | Scheduler job ID, e.g. the value of `SLURM_JOB_ID` | +| job_name | Scheduler job name | +| job_user | User the job runs as | +| job_account | Account/project the job is charged to | +| job_partition | Partition/queue the job runs in | +| node_name | Name of the compute node, as the scheduler knows it | + +The last seven fields are filled in automatically, see +[Using CodeCarbon on SLURM](../how-to/slurm.md#job-metadata-in-the-output). + +!!! warning "Existing `emissions.csv` files are rotated once" + + These seven columns change the CSV header. On the first run after upgrading, + CodeCarbon backs up an existing `emissions.csv` next to it and starts a new + file with the new header. Nothing is lost, but a pipeline reading a fixed + path will see a file with only the new rows in it. !!! note Developers can enhance the Output interface by implementing a custom class that extends `BaseOutput` at `codecarbon/output.py`. For example, to log into a database. diff --git a/tests/test_data/emissions_valid_headers.csv b/tests/test_data/emissions_valid_headers.csv index b7493c902..736e61ac3 100644 --- a/tests/test_data/emissions_valid_headers.csv +++ b/tests/test_data/emissions_valid_headers.csv @@ -1,2 +1,2 @@ -timestamp,project_name,run_id,experiment_id,duration,emissions,emissions_rate,cpu_power,gpu_power,ram_power,cpu_energy,gpu_energy,ram_energy,energy_consumed,water_consumed,country_name,country_iso_code,region,cloud_provider,cloud_region,os,python_version,codecarbon_version,cpu_count,cpu_model,gpu_count,gpu_model,longitude,latitude,ram_total_size,tracking_mode,cpu_utilization_percent,gpu_utilization_percent,ram_utilization_percent,ram_used_gb,on_cloud,pue,wue -2021-09-23T15:04:51,codecarbon,0a578547-1d6b-4e2f-be0c-7ad10f2f7c97,test,161.20380687713623,0.0004490989249167,0.0027859076880178,0.269999999999999,0.0,12.884901888000002,0.0,0,0.00057442898176,0.00057442898176,0.1,Morocco,MAR,casablanca-settat,,,macOS-10.15.7-x86_64-i386-64bit,3.8.0,2.1.3,12,Intel(R) Core(TM) i7-8850H CPU @ 2.60GHz,,,-7.9084,33.5932,,machine,0.0,0.0,0.0,0.0,N,1.0,0.0 +timestamp,project_name,run_id,experiment_id,duration,emissions,emissions_rate,cpu_power,gpu_power,ram_power,cpu_energy,gpu_energy,ram_energy,energy_consumed,water_consumed,country_name,country_iso_code,region,cloud_provider,cloud_region,os,python_version,codecarbon_version,cpu_count,cpu_model,gpu_count,gpu_model,longitude,latitude,ram_total_size,tracking_mode,cpu_utilization_percent,gpu_utilization_percent,ram_utilization_percent,ram_used_gb,on_cloud,pue,wue,scheduler,job_id,job_name,job_user,job_account,job_partition,node_name +2021-09-23T15:04:51,codecarbon,0a578547-1d6b-4e2f-be0c-7ad10f2f7c97,test,161.20380687713623,0.0004490989249167,0.0027859076880178,0.269999999999999,0.0,12.884901888000002,0.0,0,0.00057442898176,0.00057442898176,0.1,Morocco,MAR,casablanca-settat,,,macOS-10.15.7-x86_64-i386-64bit,3.8.0,2.1.3,12,Intel(R) Core(TM) i7-8850H CPU @ 2.60GHz,,,-7.9084,33.5932,,machine,0.0,0.0,0.0,0.0,N,1.0,0.0,,,,,,, diff --git a/tests/test_schedulers.py b/tests/test_schedulers.py new file mode 100644 index 000000000..f93e953e6 --- /dev/null +++ b/tests/test_schedulers.py @@ -0,0 +1,107 @@ +import unittest +from unittest import mock + +from codecarbon.core.schedulers import ( + JOB_METADATA_FIELDS, + detect_job_metadata, + warn_on_multi_rank_double_counting, +) + +SLURM_ENV = { + "SLURM_JOB_ID": "1234567", + "SLURM_JOB_NAME": "train", + "SLURM_JOB_USER": "researcher", + "SLURM_JOB_ACCOUNT": "proj42", + "SLURM_JOB_PARTITION": "gpu", + "SLURMD_NODENAME": "nid001", +} + + +class TestSchedulers(unittest.TestCase): + @mock.patch.dict("os.environ", SLURM_ENV, clear=True) + def test_detect_slurm_env_parses_variables(self): + self.assertEqual( + detect_job_metadata(), + { + "scheduler": "slurm", + "job_id": "1234567", + "job_name": "train", + "job_user": "researcher", + "job_account": "proj42", + "job_partition": "gpu", + "node_name": "nid001", + }, + ) + + @mock.patch.dict("os.environ", {}, clear=True) + def test_no_scheduler_env_is_inert(self): + metadata = detect_job_metadata() + self.assertEqual(set(metadata), set(JOB_METADATA_FIELDS)) + self.assertEqual(set(metadata.values()), {""}) + + @mock.patch.dict("os.environ", {"SLURM_JOB_ID": "42"}, clear=True) + def test_partial_slurm_env_leaves_others_empty(self): + metadata = detect_job_metadata() + self.assertEqual(metadata["scheduler"], "slurm") + self.assertEqual(metadata["job_id"], "42") + self.assertEqual(metadata["job_name"], "") + + @mock.patch.dict( + "os.environ", + {"CODECARBON_SCHEDULER": "pbs", "CODECARBON_JOB_ID": "99.pbsserver"}, + clear=True, + ) + def test_generic_env_contract_without_slurm(self): + metadata = detect_job_metadata() + self.assertEqual(metadata["scheduler"], "pbs") + self.assertEqual(metadata["job_id"], "99.pbsserver") + + @mock.patch.dict( + "os.environ", dict(SLURM_ENV, CODECARBON_JOB_ACCOUNT="billed-to"), clear=True + ) + def test_generic_env_contract_overrides_slurm(self): + metadata = detect_job_metadata() + self.assertEqual(metadata["job_account"], "billed-to") + self.assertEqual(metadata["job_id"], "1234567") + + +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.schedulers.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") + ) + + +class TestSchedulerMetadataOnEmissionsData(unittest.TestCase): + @mock.patch.dict("os.environ", SLURM_ENV, clear=True) + def test_job_fields_reach_the_emissions_data(self): + from codecarbon.emissions_tracker import OfflineEmissionsTracker + + tracker = OfflineEmissionsTracker( + country_iso_code="FRA", output_methods=[], allow_multiple_runs=True + ) + tracker.start() + try: + data = tracker._prepare_emissions_data() + finally: + tracker.stop() + + self.assertEqual(data.scheduler, "slurm") + self.assertEqual(data.job_id, "1234567") + self.assertEqual(data.job_account, "proj42") + self.assertEqual(data.node_name, "nid001") + # The new fields must be part of the CSV columns. + self.assertIn("job_partition", data.values)