Added checks for python deps and added method using builtin std library functions to get system memory

Author: pintohutch

nipype/interfaces/tests/test_runtime_profiler.py

@@ -10,6 +10,15 @@
 import unittest
 from nipype.interfaces.base import traits, CommandLine, CommandLineInputSpec
 
+try:
+    import psutil
+    import memory_profiler
+    run_profiler = True
+    skip_profile_msg = 'Run profiler tests'
+except ImportError as exc:
+    skip_profile_msg = 'Missing python packages for runtime profiling, skipping...\n'\
+                       'Error: %s' % exc
+    run_profiler = False
 
 # UseResources inputspec
 class UseResourcesInputSpec(CommandLineInputSpec):
@@ -151,6 +160,7 @@ def _run_workflow(self):
         return finish_str
 
     # Test resources were used as expected
+    @unittest.skipIf(run_profiler == False, skip_profile_msg)
     def test_wf_logfile(self):
         '''
         Test runtime profiler correctly records workflow RAM/CPUs consumption

nipype/pipeline/plugins/multiproc.py

@@ -9,13 +9,13 @@
 # Import packages
 from multiprocessing import Process, Pool, cpu_count, pool
 from traceback import format_exception
+import os
 import sys
 
 import numpy as np
 from copy import deepcopy
 from ..engine import MapNode
 from ...utils.misc import str2bool
-import psutil
 from ... import logging
 import semaphore_singleton
 from .base import (DistributedPluginBase, report_crash)
@@ -78,6 +78,34 @@ def release_lock(args):
     semaphore_singleton.semaphore.release()
 
 
+# Get total system RAM
+def get_system_total_memory_gb():
+    """Function to get the total RAM of the running system in GB
+    """
+
+    # Import packages
+    import os
+    import sys
+
+    # Get memory
+    if 'linux' in sys.platform:
+        with open('/proc/meminfo', 'r') as f_in:
+            meminfo_lines = f_in.readlines()
+            mem_total_line = [line for line in meminfo_lines \
+                              if 'MemTotal' in line][0]
+            mem_total = float(mem_total_line.split()[1])
+            memory_gb = mem_total/(1024.0**2)
+    elif 'darwin' in sys.platform:
+        mem_str = os.popen('sysctl hw.memsize').read().strip().split(' ')[-1]
+        memory_gb = float(mem_str)/(1024.0**3)
+    else:
+        err_msg = 'System platform: %s is not supported'
+        raise Exception(err_msg)
+
+    # Return memory
+    return memory_gb
+
+
 class MultiProcPlugin(DistributedPluginBase):
     """Execute workflow with multiprocessing, not sending more jobs at once
     than the system can support.
@@ -102,22 +130,24 @@ class MultiProcPlugin(DistributedPluginBase):
     """
 
     def __init__(self, plugin_args=None):
+        # Init variables and instance attributes
         super(MultiProcPlugin, self).__init__(plugin_args=plugin_args)
         self._taskresult = {}
         self._taskid = 0
         non_daemon = True
         self.plugin_args = plugin_args
         self.processors = cpu_count()
-        memory = psutil.virtual_memory()
-        self.memory = float(memory.total) / (1024.0**3)
+        self.memory_gb = get_system_total_memory_gb()
+
+        # Check plugin args
         if self.plugin_args:
             if 'non_daemon' in self.plugin_args:
                 non_daemon = plugin_args['non_daemon']
             if 'n_procs' in self.plugin_args:
                 self.processors = self.plugin_args['n_procs']
             if 'memory' in self.plugin_args:
                 self.memory = self.plugin_args['memory']
-
+        # Instantiate different thread pools for non-daemon processes
         if non_daemon:
             # run the execution using the non-daemon pool subclass
             self.pool = NonDaemonPool(processes=self.processors)
@@ -172,40 +202,39 @@ def _send_procs_to_workers(self, updatehash=False, graph=None):
         jobids = np.flatnonzero((self.proc_pending == True) & \
                                 (self.depidx.sum(axis=0) == 0).__array__())
 
-        #check available system resources by summing all threads and memory used
-        busy_memory = 0
+        # Check available system resources by summing all threads and memory used
+        busy_memory_gb = 0
         busy_processors = 0
         for jobid in jobids:
-            busy_memory += self.procs[jobid]._interface.estimated_memory_gb
+            busy_memory_gb += self.procs[jobid]._interface.estimated_memory_gb
             busy_processors += self.procs[jobid]._interface.num_threads
 
-        free_memory = self.memory - busy_memory
+        free_memory_gb = self.memory_gb - busy_memory_gb
         free_processors = self.processors - busy_processors
 
 
-        #check all jobs without dependency not run
+        # Check all jobs without dependency not run
         jobids = np.flatnonzero((self.proc_done == False) & \
                                 (self.depidx.sum(axis=0) == 0).__array__())
 
 
-        #sort jobs ready to run first by memory and then by number of threads
-        #The most resource consuming jobs run first
+        # Sort jobs ready to run first by memory and then by number of threads
+        # The most resource consuming jobs run first
         jobids = sorted(jobids,
                         key=lambda item: (self.procs[item]._interface.estimated_memory_gb,
                                           self.procs[item]._interface.num_threads))
 
         logger.debug('Free memory (GB): %d, Free processors: %d',
-                     free_memory, free_processors)
-
+                     free_memory_gb, free_processors)
 
-        #while have enough memory and processors for first job
-        #submit first job on the list
+        # While have enough memory and processors for first job
+        # Submit first job on the list
         for jobid in jobids:
             logger.debug('Next Job: %d, memory (GB): %d, threads: %d' \
                          % (jobid, self.procs[jobid]._interface.estimated_memory_gb,
                             self.procs[jobid]._interface.num_threads))
 
-            if self.procs[jobid]._interface.estimated_memory_gb <= free_memory and \
+            if self.procs[jobid]._interface.estimated_memory_gb <= free_memory_gb and \
                self.procs[jobid]._interface.num_threads <= free_processors:
                 logger.info('Executing: %s ID: %d' %(self.procs[jobid]._id, jobid))
                 executing_now.append(self.procs[jobid])
@@ -226,7 +255,7 @@ def _send_procs_to_workers(self, updatehash=False, graph=None):
                 self.proc_done[jobid] = True
                 self.proc_pending[jobid] = True
 
-                free_memory -= self.procs[jobid]._interface.estimated_memory_gb
+                free_memory_gb -= self.procs[jobid]._interface.estimated_memory_gb
                 free_processors -= self.procs[jobid]._interface.num_threads
 
                 # Send job to task manager and add to pending tasks

nipype/pipeline/plugins/tests/test_multiproc.py

@@ -3,13 +3,13 @@
 from tempfile import mkdtemp
 from shutil import rmtree
 from multiprocessing import cpu_count
-import psutil
 
 import nipype.interfaces.base as nib
 from nipype.utils import draw_gantt_chart
 from nipype.testing import assert_equal
 import nipype.pipeline.engine as pe
 from nipype.pipeline.plugins.callback_log import log_nodes_cb
+from nipype.pipeline.plugins.multiproc import get_system_total_memory_gb
 
 class InputSpec(nib.TraitedSpec):
     input1 = nib.traits.Int(desc='a random int')
@@ -222,7 +222,7 @@ def test_do_not_use_more_threads_then_specified():
 
     yield assert_equal, result, True, "using more threads than specified"
 
-    max_memory = psutil.virtual_memory().total / (1024*1024)
+    max_memory = get_system_total_memory_gb()
     result = True
     for m in memory:
         if m > max_memory: