feat: add scheduler job metadata to output

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'."
+        )

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

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:

docs/how-to/slurm.md

@@ -160,6 +160,57 @@ tail -f logs/<job_id>.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

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.

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,,,,,,,

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)