Multiple pid tracking

.flake8

@@ -5,4 +5,4 @@
 select = C,E,F,W,B
 ignore = E203, E501, W503, B008
 copyright-check = True
-max-complexity = 15
+max-complexity = 16

codecarbon/core/cpu.py

@@ -817,6 +817,7 @@ def get_cpu_details(self, duration: Time) -> Dict:
 
             for rapl_file in self._rapl_files:
                 logger.debug(rapl_file)
+                # Delta, if total external/cpu will break
                 cpu_details[rapl_file.name] = rapl_file.energy_delta.kWh
                 # We fake the name used by Power Gadget when using RAPL
                 if "Energy" in rapl_file.name:

codecarbon/core/measure.py

@@ -1,6 +1,7 @@
 from time import perf_counter
 
-from codecarbon.external.hardware import CPU, GPU, RAM, AppleSiliconChip
+from codecarbon.external.cpu import CPU
+from codecarbon.external.hardware import GPU, RAM, AppleSiliconChip
 from codecarbon.external.logger import logger
 
 

codecarbon/core/resource_tracker.py

@@ -4,7 +4,8 @@
 from codecarbon.core import cpu, gpu, powermetrics
 from codecarbon.core.config import parse_gpu_ids
 from codecarbon.core.util import detect_cpu_model, is_linux_os, is_mac_os, is_windows_os
-from codecarbon.external.hardware import CPU, GPU, MODE_CPU_LOAD, AppleSiliconChip
+from codecarbon.external.cpu import CPU, MODE_CPU_LOAD
+from codecarbon.external.hardware import GPU, AppleSiliconChip
 from codecarbon.external.logger import logger
 from codecarbon.external.ram import RAM
 
@@ -28,6 +29,7 @@ def set_RAM_tracking(self):
             self.ram_tracker = "RAM power estimation model"
         ram = RAM(
             tracking_mode=self.tracker._tracking_mode,
+            tracking_pids=self.tracker._tracking_pids,
             force_ram_power=self.tracker._force_ram_power,
         )
         self.tracker._conf["ram_total_size"] = ram.machine_memory_GB
@@ -46,6 +48,7 @@ def _setup_cpu_load_mode(self, tdp, max_power):
             model,
             max_power,
             tracking_mode=self.tracker._tracking_mode,
+            tracking_pids=self.tracker._tracking_pids,
         )
         self.cpu_tracker = MODE_CPU_LOAD
         self.tracker._conf["cpu_model"] = hardware_cpu.get_model()
@@ -56,7 +59,12 @@ def _setup_power_gadget(self):
         """Set up CPU tracking using Intel Power Gadget."""
         logger.info("Tracking Intel CPU via Power Gadget")
         self.cpu_tracker = "Power Gadget"
-        hardware_cpu = CPU.from_utils(self.tracker._output_dir, "intel_power_gadget")
+        hardware_cpu = CPU.from_utils(
+            self.tracker._output_dir,
+            "intel_power_gadget",
+            tracking_mode=self.tracker._tracking_mode,
+            tracking_pids=self.tracker._tracking_pids,
+        )
         self.tracker._hardware.append(hardware_cpu)
         self.tracker._conf["cpu_model"] = hardware_cpu.get_model()
         return True
@@ -70,6 +78,8 @@ def _setup_rapl(self):
             mode="intel_rapl",
             rapl_include_dram=self.tracker._rapl_include_dram,
             rapl_prefer_psys=self.tracker._rapl_prefer_psys,
+            tracking_mode=self.tracker._tracking_mode,
+            tracking_pids=self.tracker._tracking_pids,
         )
         self.tracker._hardware.append(hardware_cpu)
         self.tracker._conf["cpu_model"] = hardware_cpu.get_model()
@@ -81,6 +91,11 @@ def _setup_powermetrics(self):
         self.gpu_tracker = "PowerMetrics"
         self.cpu_tracker = "PowerMetrics"
 
+        if self.tracker._tracking_mode != "machine":
+            logger.warning(
+                "PowerMetrics only supports 'machine' tracking mode. Overriding tracking mode to 'machine'."
+            )
+
         hardware_cpu = AppleSiliconChip.from_utils(
             self.tracker._output_dir, chip_part="CPU"
         )
@@ -141,14 +156,20 @@ def _setup_fallback_tracking(self, tdp, max_power):
                     model,
                     max_power,
                     tracking_mode=self.tracker._tracking_mode,
+                    tracking_pids=self.tracker._tracking_pids,
                 )
                 self.cpu_tracker = MODE_CPU_LOAD
             else:
                 logger.warning(
                     "No CPU tracking mode found. Falling back on CPU constant mode."
                 )
                 hardware_cpu = CPU.from_utils(
-                    self.tracker._output_dir, "constant", model, max_power
+                    self.tracker._output_dir,
+                    "constant",
+                    model,
+                    max_power,
+                    tracking_mode=self.tracker._tracking_mode,
+                    tracking_pids=self.tracker._tracking_pids,
                 )
                 self.cpu_tracker = "global constant"
             self.tracker._hardware.append(hardware_cpu)
@@ -163,14 +184,20 @@ def _setup_fallback_tracking(self, tdp, max_power):
                     model,
                     max_power,
                     tracking_mode=self.tracker._tracking_mode,
+                    tracking_pids=self.tracker._tracking_pids,
                 )
                 self.cpu_tracker = MODE_CPU_LOAD
             else:
                 logger.warning(
                     "Failed to match CPU TDP constant. Falling back on a global constant."
                 )
                 self.cpu_tracker = "global constant"
-                hardware_cpu = CPU.from_utils(self.tracker._output_dir, "constant")
+                hardware_cpu = CPU.from_utils(
+                    self.tracker._output_dir,
+                    "constant",
+                    tracking_mode=self.tracker._tracking_mode,
+                    tracking_pids=self.tracker._tracking_pids,
+                )
             self.tracker._hardware.append(hardware_cpu)
 
     def set_CPU_tracking(self):

codecarbon/emissions_tracker.py

@@ -22,8 +22,9 @@
 from codecarbon.core.resource_tracker import ResourceTracker
 from codecarbon.core.units import Energy, Power, Time, Water
 from codecarbon.core.util import count_cpus, count_physical_cpus, suppress
+from codecarbon.external.cpu import CPU
 from codecarbon.external.geography import CloudMetadata, GeoMetadata
-from codecarbon.external.hardware import CPU, GPU, AppleSiliconChip
+from codecarbon.external.hardware import GPU, AppleSiliconChip
 from codecarbon.external.logger import logger, set_logger_format, set_logger_level
 from codecarbon.external.ram import RAM
 from codecarbon.external.scheduler import PeriodicScheduler
@@ -179,6 +180,7 @@ def __init__(
             str
         ] = _sentinel,  # Deprecated, use electricitymaps_api_token
         tracking_mode: Optional[str] = _sentinel,
+        tracking_pids: Optional[List[int]] = _sentinel,
         log_level: Optional[Union[int, str]] = _sentinel,
         on_csv_write: Optional[str] = _sentinel,
         logger_preamble: Optional[str] = _sentinel,
@@ -234,6 +236,8 @@ def __init__(
                               power consumption due to the entire machine or to try and
                               isolate the tracked processe's in isolation.
                               Defaults to "machine".
+        :param tracking_pids: PID of the process to be tracked when using "process" mode.
+                             Defaults to None, which means the current process.
         :param log_level: Global codecarbon log level. Accepts one of:
                             {"debug", "info", "warning", "error", "critical"}.
                           Defaults to "info".
@@ -317,6 +321,30 @@ def __init__(
         self._set_from_conf(prometheus_url, "prometheus_url", "localhost:9091")
         self._set_from_conf(output_handlers, "output_handlers", [])
         self._set_from_conf(tracking_mode, "tracking_mode", "machine")
+        self._set_from_conf(
+            tracking_pids, "tracking_pids", [psutil.Process().pid], List[int]
+        )
+        if self._tracking_pids is None or len(self._tracking_pids) == 0:
+            self._tracking_pids = [psutil.Process().pid]
+            tracking_pids = self._tracking_pids
+
+        # Check if tracking pids are child of each other
+        pid_check = set()
+        for pid in self._tracking_pids:
+            try:
+                process = psutil.Process(pid)
+                pids_to_track = {pid} | {
+                    child.pid for child in process.children(recursive=True)
+                }
+            except psutil.NoSuchProcess:
+                continue
+
+            duplicates = pids_to_track & pid_check
+            for dup_pid in duplicates:
+                logger.warning(f"Process with pid {dup_pid} is already being tracked.")
+
+            pid_check.update(pids_to_track)
+
         self._set_from_conf(on_csv_write, "on_csv_write", "append")
         self._set_from_conf(logger_preamble, "logger_preamble", "")
         self._set_from_conf(force_cpu_power, "force_cpu_power", None, float)
@@ -391,6 +419,12 @@ def __init__(
         else:
             logger.info(f"  GPU model: {self._conf.get('gpu_model')}")
 
+        if self._tracking_mode == "process":
+            logger.info("  Tracking mode: process")
+            logger.info("  Tracked PIDs: " + str(self._tracking_pids))
+        else:
+            logger.info("  Tracking mode: machine")
+
         # Run `self._measure_power_and_energy` every `measure_power_secs` seconds in a
         # background thread
         self._scheduler = PeriodicScheduler(
@@ -1233,6 +1267,7 @@ def track_emissions(
         str
     ] = _sentinel,  # Deprecated, use electricitymaps_api_token
     tracking_mode: Optional[str] = _sentinel,
+    tracking_pids: Optional[List[int]] = _sentinel,
     log_level: Optional[Union[int, str]] = _sentinel,
     on_csv_write: Optional[str] = _sentinel,
     logger_preamble: Optional[str] = _sentinel,
@@ -1293,6 +1328,8 @@ def track_emissions(
                           power consumption due to the entire machine or to try and
                           isolate the tracked processe's in isolation.
                           Defaults to "machine".
+    :param tracking_pids: PID of the process to be tracked when using "process" mode.
+                          Defaults to None, which means the current process.
     :param log_level: Global codecarbon log level. Accepts one of:
                       {"debug", "info", "warning", "error", "critical"}.
                       Defaults to "info".
@@ -1367,6 +1404,7 @@ def wrapped_fn(*args, **kwargs):
                     gpu_ids=gpu_ids,
                     electricitymaps_api_token=_electricitymaps_token,
                     tracking_mode=tracking_mode,
+                    tracking_pids=tracking_pids,
                     log_level=log_level,
                     on_csv_write=on_csv_write,
                     logger_preamble=logger_preamble,
@@ -1406,6 +1444,7 @@ def wrapped_fn(*args, **kwargs):
                     experiment_name=experiment_name,
                     electricitymaps_api_token=_electricitymaps_token,
                     tracking_mode=tracking_mode,
+                    tracking_pids=tracking_pids,
                     log_level=log_level,
                     on_csv_write=on_csv_write,
                     logger_preamble=logger_preamble,