ENH: Improving the resource monitor -- infer PID from process name

mriqc/instrumentation/__main__.py

@@ -22,11 +22,23 @@
 #
 
 if __name__ == "__main__":
-    import sys
-    from mriqc.instrumentation.resources import ResourceRecorder
+    from mriqc.instrumentation.resources import ResourceRecorder, FindProcess
     from . import __name__ as module
+    import argparse
 
     # `python -m <module>` typically displays the command as __main__.py
-    if "__main__.py" in sys.argv[0]:
-        sys.argv[0] = "%s -m %s" % (sys.executable, module)
-    ResourceRecorder(pid=int(sys.argv[1])).start()
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-n", "--name", type=str, required=True)
+    parser.add_argument("-p", "--logfile_path", type=str, default=".")
+    args = parser.parse_args()
+
+    if "__main__.py" in argparse._sys.argv[0]:  # sys.argv[0]:
+        argparse._sys.argv[0] = "%s -m %s" % (argparse._sys.executable, module)
+
+    if args.name.isnumeric():
+        pid = args.name
+    else:
+        pid = FindProcess(args.name)
+
+    ResourceRecorder(pid, log_file=args.logfile_path + str(pid) + ".tsv").run()

mriqc/instrumentation/resources.py

@@ -29,6 +29,7 @@
 import signal
 import psutil
 
+
 _MB = 1024.0**2
 SAMPLE_ATTRS = (
     "pid",
@@ -42,6 +43,33 @@
 )
 
 
+def FindProcess(process_name):
+    """
+     Find a process by its name and returns its PID. Child processes are excluded
+     Parameters
+     ----------
+     process_name : :obj:`str`
+         The name of the process that must be found.
+    Return
+     ----------
+     PID of the process if found, False if the process is not found
+    """
+
+    for proc in psutil.process_iter():
+        try:
+            if process_name == proc.name():
+
+                parent = proc.parent()
+
+                if parent.name() != process_name:
+                    return proc.pid
+
+        except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
+            pass
+    print("Process ", process_name, " not found")
+    return False
+
+
 def sample(
     pid=None,
     recursive=True,
@@ -50,7 +78,6 @@ def sample(
 ):
     """
     Probe process tree and snapshot current resource utilization.
-
     Parameters
     ----------
     pid : :obj:`int` or :obj:`None`
@@ -61,7 +88,6 @@ def sample(
     attrs : :obj:`iterable` of :obj:`str`
         A list of :obj:`psutil.Process` attribute names that will be retrieved when
         sampling.
-
     """
     proc_list = [psutil.Process(pid)]
     if proc_list and recursive:
@@ -127,7 +153,7 @@ def __init__(
     ):
         Process.__init__(self, name="nipype_resmon", daemon=True, **process_kwargs)
 
-        self._pid = pid
+        self._pid = int(pid)
         """The process to be sampled."""
         self._logfile = str(
             Path(log_file if log_file is not None else f".prof-{pid}.tsv").absolute()
@@ -145,8 +171,8 @@ def __init__(
 
     def run(self, *args, **kwargs):
         """Core monitoring function, called by start()"""
-
         # Open file now, because it cannot be pickled.
+
         Path(self._logfile).parent.mkdir(parents=True, exist_ok=True)
         _logfile = Path(self._logfile).open("w")
 
@@ -191,6 +217,7 @@ def run(self, *args, **kwargs):
     def stop(self, *args):
         # Tear-down process
         self._done.set()
+        print("stop")
         with Path(self._logfile).open("a") as f:
             f.write(
                 f"# MRIQC Resource recorder finished "