From a4ba054aa088cb0eed3ffd836583eff445e2c586 Mon Sep 17 00:00:00 2001
From: emilyhunt <emily.hunt.physics@gmail.com>
Date: Wed, 29 Nov 2023 16:39:05 +0100
Subject: [PATCH] Started adding a SmartProcess

---
 pyproject.toml                         |  4 +-
 src/smartmultiprocessing/_sphandler.py |  3 +-
 src/smartmultiprocessing/process.py    | 86 ++++++++++++++++++++++++++
 3 files changed, 90 insertions(+), 3 deletions(-)
 create mode 100644 src/smartmultiprocessing/process.py

diff --git a/pyproject.toml b/pyproject.toml
index da844ec..b678da8 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -26,8 +26,8 @@ classifiers = [  # Optional
   "Programming Language :: Python :: 3.12",
   "Programming Language :: Python :: 3 :: Only",
 ]
-dependencies = [ # Optional
-  "peppercorn"
+dependencies = [
+  "psutil>=5.0,<6.0",
 ]
 
 
diff --git a/src/smartmultiprocessing/_sphandler.py b/src/smartmultiprocessing/_sphandler.py
index 89bb139..1345931 100644
--- a/src/smartmultiprocessing/_sphandler.py
+++ b/src/smartmultiprocessing/_sphandler.py
@@ -1,4 +1,4 @@
-"""Definition of the main class of the module."""
+"""Definition of the main class of the old module."""
 
 import numpy as np
 import pandas as pd
@@ -34,6 +34,7 @@ class SubprocessHandler:
             Todo: update this fstring as I code more
 
         """
+        warnings.warn("SubprocessHandler will be deprecated in a future version.", DeprecationWarning)
         self.config = config
         self.config['logging_dir'].mkdir(exist_ok=True, parents=True)
         self._main_logfile = Logfile(self.config['logging_dir'] / f"main {timestamp(trailing_space=False)}.log")
diff --git a/src/smartmultiprocessing/process.py b/src/smartmultiprocessing/process.py
new file mode 100644
index 0000000..d8a2030
--- /dev/null
+++ b/src/smartmultiprocessing/process.py
@@ -0,0 +1,86 @@
+"""Definition of the SmartProcess class, which extends multiprocessing.Process to include memory monitoring."""
+import multiprocessing
+import psutil
+import time
+from typing import Optional, Iterable, Any, Mapping, Union
+
+
+class SmartProcess:
+    def __init__(
+        self,
+        target=None,
+        name: Optional[str] = None,
+        args: Iterable[Any] = (),
+        kwargs: Mapping[str, Any] = {},
+        group: None = None,
+        daemon: Optional[bool] = None,
+    ) -> None:
+        self.process = multiprocessing.Process(
+            group=group,
+            target=target,
+            name=name,
+            args=args,
+            kwargs=kwargs,
+            daemon=daemon,
+        )
+        self.psutil_process = None
+
+    def _get_children(self):
+        return self.psutil_process.children(recursive=True)
+
+    def memory_usage(self, children: Union[Iterable, bool] = False) -> int:
+        # Proportional set size of parent thread. Best guess at the memory usage of this
+        # thread, as shared with others.
+        memory_usage = self.psutil_process.memory_full_info().pss
+        if not children:
+            return memory_usage
+
+        # Optionally also get the memory usage of all child processes of this process.
+        if children == True:
+            children = self._get_children()
+        for child_process in children:
+            try:
+                memory_usage += child_process.memory_info().rss
+            except psutil.NoSuchProcess:
+                pass
+        return memory_usage
+
+    def cpu_usage(
+        self,
+        interval: Optional[Union[int, float]] = None,
+        children: Union[Iterable, bool] = False,
+    ) -> int:
+        if not children:
+            return self.psutil_process.cpu_percent(interval=interval)
+        cpu_percent = self.psutil_process.cpu_percent(interval=None)
+
+        for child_process in children:
+            try:
+                cpu_percent += child_process.cpu_percent(interval=None)
+            except psutil.NoSuchProcess:
+                pass
+
+        if interval is None:
+            return cpu_percent
+
+    def run(self):
+        self.process.run()
+
+    def start(self):
+        self.process.start()
+        self.psutil_process = psutil.Process(self.process.pid)
+
+    def join(self, timeout: Optional[float] = None):
+        self.process.join(timeout)
+
+    def is_alive(self) -> bool:
+        return self.process.is_alive()
+
+    def terminate(self):
+        self.process.terminate()
+
+    def kill(self):
+        self.process.kill()
+
+    def close(self):
+        self.process.close()
-- 
2.54.0