Merge pull request #1156 from oesteban/fix/RedirectX

[FIX] Improve redirections of X

Author: satra

nipype/algorithms/mesh.py

@@ -78,6 +78,7 @@ class ComputeMeshWarp(BaseInterface):
 
     input_spec = ComputeMeshWarpInputSpec
     output_spec = ComputeMeshWarpOutputSpec
+    _redirect_x = True
 
     def _triangle_area(self, A, B, C):
         A = np.array(A)
@@ -221,6 +222,7 @@ class MeshWarpMaths(BaseInterface):
 
     input_spec = MeshWarpMathsInputSpec
     output_spec = MeshWarpMathsOutputSpec
+    _redirect_x = True
 
     def _run_interface(self, runtime):
         try:

nipype/interfaces/base.py

@@ -53,25 +53,34 @@ def __init__(self, value):
     def __str__(self):
         return repr(self.value)
 
-def _lock_files():
-    tmpdir = '/tmp'
-    pattern = '.X*-lock'
-    names = fnmatch.filter(os.listdir(tmpdir), pattern)
-    ls = [os.path.join(tmpdir, child) for child in names]
-    ls = [p for p in ls if os.path.isfile(p)]
-    return ls
-
-def _search_for_free_display():
-    ls = [int(x.split('X')[1].split('-')[0]) for x in _lock_files()]
-    min_display_num = 1000
-    if len(ls):
-        display_num = max(min_display_num, max(ls) + 1)
-    else:
-        display_num = min_display_num
-    random.seed()
-    display_num += random.randint(0, 100)
-    return display_num
 
+def _unlock_display(ndisplay):
+    lockf = os.path.join('/tmp', '.X%d-lock' % ndisplay)
+    try:
+        os.remove(lockf)
+    except:
+        return False
+
+    return True
+
+def _exists_in_path(cmd, environ):
+    '''
+    Based on a code snippet from
+     http://orip.org/2009/08/python-checking-if-executable-exists-in.html
+    '''
+
+    if 'PATH' in environ:
+        input_environ = environ.get("PATH")
+    else:
+        input_environ = os.environ.get("PATH", "")
+    extensions = os.environ.get("PATHEXT", "").split(os.pathsep)
+    for directory in input_environ.split(os.pathsep):
+        base = os.path.join(directory, cmd)
+        options = [base] + [(base + ext) for ext in extensions]
+        for filename in options:
+            if os.path.exists(filename):
+                return True, filename
+    return False, None
 
 def load_template(name):
     """Load a template from the script_templates directory
@@ -946,6 +955,36 @@ def _check_version_requirements(self, trait_object, raise_exception=True):
                                      version, max_ver))
         return unavailable_traits
 
+    def _run_wrapper(self, runtime):
+        sysdisplay = os.getenv('DISPLAY')
+        if self._redirect_x:
+            try:
+                from xvfbwrapper import Xvfb
+            except ImportError:
+                iflogger.error('Xvfb wrapper could not be imported')
+                raise
+
+            vdisp = Xvfb(nolisten='tcp')
+            vdisp.start()
+            vdisp_num = vdisp.vdisplay_num
+
+            iflogger.info('Redirecting X to :%d' % vdisp_num)
+            runtime.environ['DISPLAY'] = ':%d' % vdisp_num
+
+        runtime = self._run_interface(runtime)
+
+        if self._redirect_x:
+            if sysdisplay is None:
+                os.unsetenv('DISPLAY')
+            else:
+                os.environ['DISPLAY'] = sysdisplay
+
+            iflogger.info('Freeing X :%d' % vdisp_num)
+            vdisp.stop()
+            _unlock_display(vdisp_num)
+
+        return runtime
+
     def _run_interface(self, runtime):
         """ Core function that executes interface
         """
@@ -982,34 +1021,7 @@ def run(self, **inputs):
                         hostname=getfqdn(),
                         version=self.version)
         try:
-            if self._redirect_x:
-                exist_val, _ = self._exists_in_path('Xvfb',
-                                                    runtime.environ)
-                if not exist_val:
-                    raise IOError("Xvfb could not be found on host %s" %
-                                  (runtime.hostname))
-                else:
-                    vdisplay_num = _search_for_free_display()
-                    xvfb_cmd = ['Xvfb', ':%d' % vdisplay_num]
-                    xvfb_proc = subprocess.Popen(xvfb_cmd,
-                                                 stdout=open(os.devnull),
-                                                 stderr=open(os.devnull))
-                    wait_step = 0.2
-                    wait_time = 0
-                    while xvfb_proc.poll() is not None:
-                        if wait_time > config.get('execution', 'xvfb_max_wait'):
-                            raise Exception('Error: Xvfb did not start')
-                        time.sleep(wait_step)  # give Xvfb time to start
-                        wait_time += wait_step
-
-                    runtime.environ['DISPLAY'] = ':%s' % vdisplay_num
-
-            runtime = self._run_interface(runtime)
-
-            if self._redirect_x:
-                xvfb_proc.kill()
-                xvfb_proc.wait()
-
+            runtime = self._run_wrapper(runtime)
             outputs = self.aggregate_outputs(runtime)
             runtime.endTime = dt.isoformat(dt.utcnow())
             timediff = parseutc(runtime.endTime) - parseutc(runtime.startTime)
@@ -1169,32 +1181,39 @@ def _read(self, drain):
         self._lastidx = len(self._rows)
 
 
-def run_command(runtime, output=None, timeout=0.01):
+def run_command(runtime, output=None, timeout=0.01, redirect_x=False):
     """Run a command, read stdout and stderr, prefix with timestamp.
 
     The returned runtime contains a merged stdout+stderr log with timestamps
     """
     PIPE = subprocess.PIPE
 
+    cmdline = runtime.cmdline
+    if redirect_x:
+        exist_xvfb, _ = _exists_in_path('xvfb-run', runtime.environ)
+        if not exist_xvfb:
+            raise RuntimeError('Xvfb was not found, X redirection aborted')
+        cmdline = 'xvfb-run -a ' + cmdline
+
     if output == 'file':
         errfile = os.path.join(runtime.cwd, 'stderr.nipype')
         outfile = os.path.join(runtime.cwd, 'stdout.nipype')
         stderr = open(errfile, 'wt')
         stdout = open(outfile, 'wt')
 
-        proc = subprocess.Popen(runtime.cmdline,
+        proc = subprocess.Popen(cmdline,
                                 stdout=stdout,
                                 stderr=stderr,
                                 shell=True,
                                 cwd=runtime.cwd,
                                 env=runtime.environ)
     else:
-        proc = subprocess.Popen(runtime.cmdline,
-                                 stdout=PIPE,
-                                 stderr=PIPE,
-                                 shell=True,
-                                 cwd=runtime.cwd,
-                                 env=runtime.environ)
+        proc = subprocess.Popen(cmdline,
+                                stdout=PIPE,
+                                stderr=PIPE,
+                                shell=True,
+                                cwd=runtime.cwd,
+                                env=runtime.environ)
     result = {}
     errfile = os.path.join(runtime.cwd, 'stderr.nipype')
     outfile = os.path.join(runtime.cwd, 'stdout.nipype')
@@ -1412,7 +1431,7 @@ def _get_environ(self):
 
     def version_from_command(self, flag='-v'):
         cmdname = self.cmd.split()[0]
-        if self._exists_in_path(cmdname):
+        if _exists_in_path(cmdname):
             env = deepcopy(os.environ.data)
             out_environ = self._get_environ()
             env.update(out_environ)
@@ -1425,6 +1444,10 @@ def version_from_command(self, flag='-v'):
             o, e = proc.communicate()
             return o
 
+    def _run_wrapper(self, runtime):
+        runtime = self._run_interface(runtime)
+        return runtime
+
     def _run_interface(self, runtime, correct_return_codes=[0]):
         """Execute command via subprocess
 
@@ -1444,40 +1467,22 @@ def _run_interface(self, runtime, correct_return_codes=[0]):
         out_environ = self._get_environ()
         runtime.environ.update(out_environ)
         executable_name = self.cmd.split()[0]
-        exist_val, cmd_path = self._exists_in_path(executable_name,
-                                                   runtime.environ)
+        exist_val, cmd_path = _exists_in_path(executable_name,
+                                              runtime.environ)
         if not exist_val:
             raise IOError("%s could not be found on host %s" %
                           (self.cmd.split()[0], runtime.hostname))
         setattr(runtime, 'command_path', cmd_path)
         setattr(runtime, 'dependencies', get_dependencies(executable_name,
                                                           runtime.environ))
-        runtime = run_command(runtime, output=self.inputs.terminal_output)
+        runtime = run_command(runtime, output=self.inputs.terminal_output,
+                              redirect_x=self._redirect_x)
         if runtime.returncode is None or \
                         runtime.returncode not in correct_return_codes:
             self.raise_exception(runtime)
 
         return runtime
 
-    def _exists_in_path(self, cmd, environ):
-        '''
-        Based on a code snippet from
-         http://orip.org/2009/08/python-checking-if-executable-exists-in.html
-        '''
-
-        if 'PATH' in environ:
-            input_environ = environ.get("PATH")
-        else:
-            input_environ = os.environ.get("PATH", "")
-        extensions = os.environ.get("PATHEXT", "").split(os.pathsep)
-        for directory in input_environ.split(os.pathsep):
-            base = os.path.join(directory, cmd)
-            options = [base] + [(base + ext) for ext in extensions]
-            for filename in options:
-                if os.path.exists(filename):
-                    return True, filename
-        return False, None
-
     def _format_arg(self, name, trait_spec, value):
         """A helper function for _parse_inputs