Added memory monitoring for non-dockerised tools (#1006)

Author: mb1069

cwltool/job.py

@@ -2,28 +2,29 @@
 
 import datetime
 import functools
+import itertools
 import logging
 import os
 import re
 import shutil
 import stat
 import sys
-import time
 import tempfile
+import time
 import uuid
 from abc import ABCMeta, abstractmethod
 from io import IOBase, open  # pylint: disable=redefined-builtin
 from threading import Timer
 from typing import (IO, Any, AnyStr, Callable, Dict, Iterable, List, Tuple,
                     MutableMapping, MutableSequence, Optional, Union, cast)
 
-import shellescape
 import psutil
+import shellescape
 from prov.model import PROV
+from schema_salad.sourceline import SourceLine
 from six import PY2, with_metaclass
 from typing_extensions import (TYPE_CHECKING,  # pylint: disable=unused-import
                                Text)
-from schema_salad.sourceline import SourceLine
 
 from .builder import Builder, HasReqsHints  # pylint: disable=unused-import
 from .context import RuntimeContext  # pylint: disable=unused-import
@@ -231,11 +232,11 @@ def _setup(self, runtimeContext):  # type: (RuntimeContext) -> None
                                 for p in self.generatemapper.files()}, indent=4))
 
     def _execute(self,
-                 runtime,               # type: List[Text]
-                 env,                   # type: MutableMapping[Text, Text]
-                 runtimeContext,        # type: RuntimeContext
-                 monitor_function=None  # type: Optional[Callable]
-                 ):  # type: (...) -> None
+                 runtime,                # type: List[Text]
+                 env,                    # type: MutableMapping[Text, Text]
+                 runtimeContext,         # type: RuntimeContext
+                 monitor_function=None,  # type: Optional[Callable]
+                 ):                      # type: (...) -> None
 
         scr, _ = self.get_requirement("ShellCommandRequirement")
 
@@ -383,6 +384,30 @@ def _execute(self,
             _logger.debug(u"[job %s] Removing temporary directory %s", self.name, self.tmpdir)
             shutil.rmtree(self.tmpdir, True)
 
+    def process_monitor(self, sproc):
+        monitor = psutil.Process(sproc.pid)
+        memory_usage = [None]  # Value must be list rather than integer to utilise pass-by-reference in python
+
+        def get_tree_mem_usage(memory_usage):
+            children = monitor.children()
+            rss = monitor.memory_info().rss
+            while len(children):
+                rss += sum([process.memory_info().rss for process in children])
+                children = list(itertools.chain(*[process.children() for process in children]))
+            if memory_usage[0] is None or rss > memory_usage[0]:
+                memory_usage[0] = rss
+
+        mem_tm = Timer(interval=1, function=get_tree_mem_usage, args=(memory_usage,))
+        mem_tm.daemon = True
+        mem_tm.start()
+        sproc.wait()
+        mem_tm.cancel()
+        if memory_usage[0] is not None:
+            _logger.info(u"[job %s] Max memory used: %iMiB", self.name,
+                         round(memory_usage[0] / (2 ** 20)))
+        else:
+            _logger.info(u"Could not collect memory usage, job ended before monitoring began.")
+
 
 class CommandLineJob(JobBase):
     def run(self,
@@ -419,7 +444,9 @@ def run(self,
                 self.generatemapper, self.outdir, self.builder.outdir,
                 inplace_update=self.inplace_update)
 
-        self._execute([], env, runtimeContext)
+        monitor_function = functools.partial(self.process_monitor)
+
+        self._execute([], env, runtimeContext, monitor_function)
 
 
 CONTROL_CODE_RE = r'\x1b\[[0-9;]*[a-zA-Z]'
@@ -621,6 +648,8 @@ def run(self, runtimeContext):  # type: (RuntimeContext) -> None
             monitor_function = functools.partial(
                 self.docker_monitor, cidfile, runtimeContext.tmpdir_prefix,
                 not bool(runtimeContext.cidfile_dir))
+        elif runtimeContext.user_space_docker_cmd:
+            monitor_function = functools.partial(self.process_monitor)
         self._execute(runtime, env, runtimeContext, monitor_function)
 
     @staticmethod

tests/test_udocker.py

@@ -3,32 +3,65 @@
 import os
 import subprocess
 from .util import get_data, get_main_output
+import tempfile
+import shutil
+from psutil.tests import TRAVIS
 
 LINUX = sys.platform in ('linux', 'linux2')
 
 
 @pytest.mark.skipif(not LINUX, reason="LINUX only")
-def test_udocker_usage_should_not_write_cid_file(tmpdir):
-    cwd = tmpdir.chdir()
-    install_cmds = [
-        "curl https://raw.githubusercontent.com/indigo-dc/udocker/master/udocker.py -o ./udocker",
-        "chmod u+rx ./udocker",
-        "./udocker install"]
-    os.environ['UDOCKER_DIR'] = os.path.join(str(tmpdir), ".udocker")
-
-    assert sum([subprocess.call(cmd.split()) for cmd in install_cmds]) == 0
-
-    docker_path = os.path.join(str(tmpdir), 'udocker')
-
-    test_file = "tests/wf/wc-tool.cwl"
-    job_file = "tests/wf/wc-job.json"
-    error_code, stdout, stderr = get_main_output(
-        ["--debug", "--default-container", "debian", "--user-space-docker-cmd=" + docker_path,
-         get_data(test_file), get_data(job_file)])
-    cwd.chdir()
-    cidfiles_count = sum(1 for _ in tmpdir.visit(fil="*.cid"))
-
-    tmpdir.remove(ignore_errors=True)
-
-    assert "completed success" in stderr
-    assert cidfiles_count == 0
+class TestUdocker:
+    udocker_path = None
+
+    @classmethod
+    def setup_class(cls):
+        install_cmds = [
+            "curl https://raw.githubusercontent.com/indigo-dc/udocker/master/udocker.py -o ./udocker",
+            "chmod u+rx ./udocker",
+            "./udocker install"]
+
+        test_cwd = os.getcwd()
+
+        cls.docker_install_dir = tempfile.mkdtemp()
+        os.chdir(cls.docker_install_dir)
+
+        os.environ['UDOCKER_DIR'] = os.path.join(cls.docker_install_dir, ".udocker")
+
+        assert sum([subprocess.call(cmd.split()) for cmd in install_cmds]) == 0
+
+        cls.udocker_path = os.path.join(cls.docker_install_dir, 'udocker')
+        os.chdir(test_cwd)
+        print('Udocker install dir: ' + cls.docker_install_dir)
+
+    @classmethod
+    def teardown_class(cls):
+        shutil.rmtree(cls.docker_install_dir)
+
+    def test_udocker_usage_should_not_write_cid_file(self, tmpdir):
+        cwd = tmpdir.chdir()
+
+        test_file = "tests/wf/wc-tool.cwl"
+        job_file = "tests/wf/wc-job.json"
+        error_code, stdout, stderr = get_main_output(
+            ["--debug", "--default-container", "debian", "--user-space-docker-cmd=" + self.udocker_path,
+             get_data(test_file), get_data(job_file)])
+        cwd.chdir()
+        cidfiles_count = sum(1 for _ in tmpdir.visit(fil="*.cid"))
+
+        tmpdir.remove(ignore_errors=True)
+
+        assert "completed success" in stderr
+        assert cidfiles_count == 0
+
+    @pytest.mark.skipif(TRAVIS, reason='Not reliable on single threaded test on travis.')
+    def test_udocker_should_display_memory_usage(self, tmpdir):
+        cwd = tmpdir.chdir()
+        error_code, stdout, stderr = get_main_output(
+            ["--default-container=debian",  "--user-space-docker-cmd=" + self.udocker_path,
+             get_data("tests/wf/timelimit.cwl"), "--sleep_time", "10"])
+        cwd.chdir()
+        tmpdir.remove(ignore_errors=True)
+
+        assert "completed success" in stderr
+        assert "Max memory" in stderr