final work getting xserver wroking on login and slurm nodes

Signed-off-by: vsoch <vsoch@users.noreply.github.com>

CHANGELOG.md

@@ -14,7 +14,7 @@ and **Merged pull requests**. Critical items to know are:
 The versions coincide with releases on pip. Only major versions will be released as tags on Github.
 
 ## [0.0.x](https://github.com/tunel-apps/tunel/tree/main) (0.0.x)
- - support for xserver (primarily singularity to start) (0.0.18)
+ - support for xserver (singularity and slurm) (0.0.18)
  - bugfix in slurm stop and adding username to proxy (0.0.17)
  - bugfix template generation (0.0.16)
  - support for docker launcher (and intending to test podman) (0.0.15)

tunel/apps/xserver/neurodesk/app.yaml

@@ -1,13 +1,16 @@
 launcher: singularity
 launchers_supported:
  - singularity
+ - slurm
 script: app.sh
 description: Launch neurodesk (or an associated application)
 needs:
   xserver: true
 examples: |
-  # Run app on login node with singularity, custom tag dev
+  # Run app on login node with singularity
   tunel run-app waffles neurodocker
+  # Run app on an interactive node (also with singularity)
+  tunel run-app waffles neurodocker --launcher slurm
   # Force a new pull
   tunel run-app waffles neurodocker --pull
 args:

tunel/launcher/slurm.py

@@ -107,13 +107,22 @@ def run_app(self, app):
         os.remove(tmpfile)
 
         # Assemble the command
-        command = [
-            "sbatch",
-            "--job-name=%s" % app.job_name,
-            "--output=%s" % render["log_output"],
-            "--error=%s" % render["log_error"],
-            remote_script,
-        ]
+        if app.has_xserver:
+            command = [
+                "srun",
+                "--job-name=%s" % app.job_name,
+                "--pty",
+                "/bin/bash",
+                remote_script,
+            ]
+        else:
+            command = [
+                "sbatch",
+                "--job-name=%s" % app.job_name,
+                "--output=%s" % render["log_output"],
+                "--error=%s" % render["log_error"],
+                remote_script,
+            ]
 
         # Launch with command
         if not self.previous_job_exists(app.job_name):
@@ -239,13 +248,11 @@ def run(self, cmd, job_name=None, logs_prefix=None, app=None, socket=None):
             self.ssh.shell("srun --pty bash", interactive=True)
 
         else:
-            if "sbatch" not in cmd:
-                cmd = ["sbatch"] + cmd
-            res = self.ssh.execute(" ".join(cmd))
+            res = self.ssh.execute(" ".join(cmd), xserver=app.has_xserver)
             self.ssh.print_output(res)
 
             # A successful submission should:
-            if res["return_code"] == 0:
+            if res["return_code"] == 0 and not app.has_xserver:
 
                 # 1. Show the user how to quickly get logs (if logs_prefix provided)
                 if logs_prefix:
@@ -256,15 +263,13 @@ def run(self, cmd, job_name=None, logs_prefix=None, app=None, socket=None):
 
                 # An xserver launches the app directly
                 time.sleep(10)
-                if not app.has_xserver:
-                    self.print_session_instructions(job_name)
+                self.print_session_instructions(job_name)
 
                 # Create another process to check logs?
                 if logs_prefix:
                     self.print_updated_logs(logs_prefix, app, socket=socket)
-                self.ssh.tunnel(
-                    machine, socket=socket, app=app, xserver=app.has_xserver
-                )
+
+                self.ssh.tunnel(machine, socket=socket, app=app)
 
     def print_session_instructions(self, job_name):
         """

tunel/ssh/commands.py

@@ -110,49 +110,18 @@ def _tunnel_isolated_port(self, machine):
     self._tunnel_wait()
 
 
-def _tunnel_xserver(self, machine):
-    """
-    Create a tunnel with an xserver
-    """
-    # TODO do we need to close up connections on login node?
-    connection = "%s:localhost:%s" % (self.local_port, self.remote_port)
-    cmd = [
-        "-f",
-        "-X",
-        "-o",
-        "ForwardX11=yes",
-        "-L",
-        connection,
-        self.server,
-        "ssh",
-        "-L",
-        connection,
-        "-X",
-        "-o",
-        "ForwardX11=yes",
-        "-N",
-        machine,
-    ]
-    self.execute(cmd)
-    self._tunnel_wait()
-
-
 # Tunnels to login node
 
 
-def _tunnel_login(self, xserver=False):
+def _tunnel_login(self):
     """
     Create a simple tunnel to the login node (assumes not isolated nodes)
     """
     socket_file = self._get_socket_path()
     cmd = ["-K", "-f", "-M"]
-    if xserver:
-        cmd += ["-X", "-o", "ForwardX11=yes"]
 
     # Add the socket file
     cmd += ["-S", socket_file, "-L"]
-    if xserver:
-        cmd += ["-X", "-o", "ForwardX11=yes"]
     cmd += ["%s:%s:%s" % (self.local_port, self.server, self.remote_port), "-N"]
     self.execute(cmd)
     self._tunnel_wait(socket_file)

tunel/ssh/ssh.py

@@ -63,12 +63,11 @@ def execute(self, cmd, login_shell=False, quiet=False, stream=False, xserver=Fal
             cmd = shlex.split(cmd)
         if stream:
             command = ["ssh", "-t", self.server]
+        elif xserver:
+            command = ["ssh", "-XY", "-o", "ForwardX11=yes", self.server]
         else:
             command = ["ssh", self.server]
 
-        # Do we want an xserver?
-        if xserver:
-            command.append("-X")
         cmd = command + cmd
         return tunel.utils.run_command(cmd, quiet=quiet, stream=stream)
 
@@ -133,26 +132,21 @@ def tunnel(
         remote_port=None,
         socket=None,
         app=None,
-        xserver=False,
     ):
         """
         Given a remote and local port, open a tunnel. If an isolated node ssh is
         done, the name of the machine is required too.
         """
         # If no machine, we have to do a login node
         if not machine:
-            return self._tunnel_login(xserver)
+            return self._tunnel_login()
 
         # The app requires a socket
         if app.needs.get("socket", False):
             if not socket:
                 logger.exit("A socket path is required.")
             return self._tunnel_isolated_socket(machine, socket=socket)
 
-        # If we are using an xserver
-        if xserver:
-            return self._tunnel_xserver(machine)
-
         port = port or self.local_port
         remote_port = remote_port or self.remote_port
         logger.info(
@@ -201,7 +195,6 @@ def shell(self, cmd=None, interactive=False):
 Tunnel._tunnel_isolated_sockets = commands._tunnel_isolated_sockets
 Tunnel._tunnel_isolated_proxyjump_sockets = commands._tunnel_isolated_proxyjump_sockets
 Tunnel._tunnel_isolated_port = commands._tunnel_isolated_port
-Tunnel._tunnel_xserver = commands._tunnel_xserver
 Tunnel._tunnel_login = commands._tunnel_login
 Tunnel._tunnel_login_node_port = commands._tunnel_login_node_port
 Tunnel._tunnel_login_node_socket = commands._tunnel_login_node_socket