Feat: added utility function to retrieve CPU index from PID

tests/device/cpu/test_socket_detection.py

@@ -0,0 +1,52 @@
+from __future__ import annotations
+
+import builtins
+import os
+import pytest
+from unittest.mock import patch
+
+from zeus.device.cpu import get_current_cpu_index
+
+class MockFile:
+  def __init__(self, file_path):
+    self.file_path = file_path
+
+  def read(self, *args, **kwargs):
+    with open(self.file_path, "r") as file:
+      return file.read()
+
+  def __enter__(self):
+    return self
+
+  def __exit__(self, exc_type, exc_value, traceback):
+    pass
+
+
+@pytest.fixture
+def mock_linux_files():
+  files_folder = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../files"))
+  mocked_stat_file_path = os.path.join(files_folder, "stat")
+  mocked_phys_package_path = os.path.join(files_folder, "physical_package_id")
+  mocked_stat_file = MockFile(mocked_stat_file_path)
+  mocked_phys_package_file = MockFile(mocked_phys_package_path)
+
+  real_open = builtins.open
+
+  def mock_file_open(filepath, *args, **kwargs):
+    if filepath == "/proc/515/stat":
+      return mocked_stat_file
+    elif filepath == "/sys/devices/system/cpu/cpu24/topology/physical_package_id":
+      return mocked_phys_package_file
+    else:
+      return real_open(filepath, *args, **kwargs)
+
+  patch_open = patch("builtins.open", side_effect=mock_file_open)
+
+  patch_open.start()
+
+  yield
+
+  patch_open.stop()
+
+def test_get_current_cpu_index(mock_linux_files):
+  assert(get_current_cpu_index(515) == 0)

tests/files/physical_package_id

@@ -0,0 +1 @@
+0

tests/files/stat

@@ -0,0 +1 @@
+6511 (cat) R 6501 6511 6501 34816 6511 4194304 95 0 0 0 0 0 0 0 20 0 1 0 7155 6053888 255 18446744073709551615 94288847015936 94288847031350 140727519907328 0 0 0 0 0 0 0 0 0 17 24 0 0 0 0 0 94288847043296 94288847044712 94288866566144 140727519913901 140727519913921 140727519913921 140727519916011 0

zeus/device/cpu/__init__.py

@@ -6,11 +6,34 @@
 
 from __future__ import annotations
 
+from typing import Literal
+import os
+
 from zeus.device.cpu.common import CPUs, ZeusCPUInitError
 from zeus.device.cpu.rapl import rapl_is_available, RAPLCPUs
 
 _cpus: CPUs | None = None
 
+def get_current_cpu_index(pid: int | Literal["current"] = "current") -> int:
+    """Retrieves the specific CPU index (socket) where the given PID is running.
+
+    If no PID is given or pid is "current", the CPU index returned is of the CPU running the current process.
+
+    Note: Linux schedulers can preempt and reschedule processes to different CPUs. To prevent this from happening
+    during monitoring, use `taskset` to pin processes to specific CPUs.
+    """
+    if pid == "current":
+        pid = os.getpid()
+
+    cpu_core = None
+    with open(f"/proc/{pid}/stat", "r") as stat_file:
+        cpu_core = int(stat_file.read().split()[38])
+
+    cpu_index = None
+    with open(f"/sys/devices/system/cpu/cpu{cpu_core}/topology/physical_package_id", "r") as phys_package_file:
+        cpu_index = int(phys_package_file.read().strip())
+
+    return cpu_index 
 
 def get_cpus() -> CPUs:
     """Initialize and return a singleton CPU monitoring object for INTEL CPUs.