diff --git a/codecarbon/external/ram.py b/codecarbon/external/ram.py index b417e1c5b..d99663594 100644 --- a/codecarbon/external/ram.py +++ b/codecarbon/external/ram.py @@ -1,4 +1,5 @@ import math +import os import re import subprocess from dataclasses import dataclass @@ -34,7 +35,7 @@ class RAM(BaseHardware): def __init__( self, - pid: int = psutil.Process().pid, + pid: Optional[int] = None, children: bool = True, tracking_mode: str = "machine", force_ram_power: Optional[int] = None, @@ -46,7 +47,7 @@ def __init__( Args: pid (int, optional): Process id (with respect to which we'll look for - children). Defaults to psutil.Process().pid. + children). Defaults to the current process id. children (int, optional): Look for children of the process when computing total RAM used. Defaults to True. tracking_mode (str, optional): Whether to track "machine" or "process" RAM. @@ -55,7 +56,7 @@ def __init__( this value is used instead of estimating RAM power. Defaults to None. """ - self._pid = pid + self._pid = os.getpid() if pid is None else pid self._children = children self._tracking_mode = tracking_mode self._force_ram_power = force_ram_power diff --git a/tests/test_ram.py b/tests/test_ram.py index 6b553dd3c..6a7cb8c95 100644 --- a/tests/test_ram.py +++ b/tests/test_ram.py @@ -1,9 +1,11 @@ +import os import subprocess import unittest from textwrap import dedent from unittest import mock import numpy as np +import pytest from codecarbon.external.ram import RAM, RAM_SLOT_POWER_X86 @@ -437,3 +439,23 @@ def test_force_ram_power(self): ram_power = ram.total_power() # Verify the calculation method was not called mock_calc.assert_not_called() + + @pytest.mark.skipif(not hasattr(os, "fork"), reason="requires os.fork") + def test_default_pid_is_resolved_in_forked_child(self): + read_fd, write_fd = os.pipe() + pid = os.fork() + if pid == 0: + # Child: the module is already imported, so a default argument + # evaluated at import time would still hold the parent's pid. + try: + os.close(read_fd) + ram = RAM(tracking_mode="process") + os.write(write_fd, str(ram._pid).encode()) + os.close(write_fd) + finally: + os._exit(0) + os.close(write_fd) + with os.fdopen(read_fd) as f: + child_ram_pid = int(f.read()) + os.waitpid(pid, 0) + self.assertEqual(child_ram_pid, pid)