BatchSpawnerBase: Add background_tasks, connect_to_job feature.

This adds the possibility to start a "connect_to_job"
background task on the hub on job start,
which establishes connectivity to the actual single user server.
An example for this can be "condor_ssh_to_job" for HTCondor batch systems.

Additionally, the background tasks are monitored:
- for successful startup. The background task is given
  some time to successfully establish connectivity.
- in poll() during job runtime
and if they fail, the job is terminated.

Author: olifre

batchspawner/batchspawner.py

@@ -172,6 +172,12 @@ def _req_keepvars_default(self):
              "specification."
         ).tag(config=True)
 
+    connect_to_job_cmd = Unicode('',
+        help="Command to connect to running batch job and forward the port "
+             "of the running notebook to the Hub. If empty, direct connectivity is assumed. "
+             "Uses self.job_id as {job_id} and the self.port as {port}."
+        ).tag(config=True)
+
     # Raw output of job submission command unless overridden
     job_id = Unicode()
 
@@ -200,6 +206,18 @@ def cmd_formatted_for_batch(self):
         """The command which is substituted inside of the batch script"""
         return ' '.join([self.batchspawner_singleuser_cmd] + self.cmd + self.get_args())
 
+    async def connect_to_job(self):
+        """This command ensures the port of the singleuser server is reachable from the
+        Batchspawner machine. By default, it does nothing, i.e. direct connectivity
+        is assumed.
+        """
+        subvars = self.get_req_subvars()
+        subvars['job_id'] = self.job_id
+        subvars['port'] = self.port
+        cmd = ' '.join((format_template(self.exec_prefix, **subvars),
+                        format_template(self.connect_to_job_cmd, **subvars)))
+        await self.run_background_command(cmd)
+
     async def run_command(self, cmd, input=None, env=None):
         proc = await asyncio.create_subprocess_shell(cmd, env=env,
                                                     stdin=asyncio.subprocess.PIPE,
@@ -243,6 +261,46 @@ async def run_command(self, cmd, input=None, env=None):
         out = out.decode().strip()
         return out
 
+    # List of running background processes, e.g. used by connect_to_job.
+    background_processes = []
+
+    async def _async_wait_process(self, sleep_time):
+        """Asynchronously sleeping process for delayed checks"""
+        await asyncio.sleep(sleep_time)
+
+    async def run_background_command(self, cmd, startup_check_delay=1, input=None, env=None):
+        """Runs the given background command, adds it to background_processes,
+        and checks if the command is still running after startup_check_delay."""
+        background_process = self.run_command(cmd, input, env)
+        success_check_delay = self._async_wait_process(startup_check_delay)
+
+        # Start up both the success check process and the actual process.
+        done, pending = await asyncio.wait([background_process, success_check_delay], return_when=asyncio.FIRST_COMPLETED)
+
+        # If the success check process is the one which exited first, all is good, else fail.
+        if list(done)[0]._coro == success_check_delay:
+            background_task = list(pending)[0]
+            self.background_processes.append(background_task)
+            return background_task
+        else:
+            self.log.error("Background command exited early: %s" % cmd)
+            gather_pending = asyncio.gather(*pending)
+            gather_pending.cancel()
+            try:
+                self.log.debug("Cancelling pending success check task...")
+                await gather_pending
+            except asyncio.CancelledError:
+                self.log.debug("Cancel was successful.")
+                pass
+
+            # Retrieve exception from "done" process.
+            try:
+                gather_done = asyncio.gather(*done)
+                await gather_done
+            except:
+                self.log.debug("Retrieving exception from failed background task...")
+                raise RuntimeError('{} failed!'.format(cmd))
+
     async def _get_batch_script(self, **subvars):
         """Format batch script from vars"""
         # Could be overridden by subclasses, but mainly useful for testing
@@ -270,6 +328,27 @@ async def submit_batch_script(self):
             self.job_id = ''
         return self.job_id
 
+    def background_tasks_ok(self):
+        # Check background processes.
+        if self.background_processes:
+            self.log.debug('Checking background processes...')
+            for background_process in self.background_processes:
+                if background_process.done():
+                    self.log.debug('Found a background process in state "done"...')
+                    try:
+                        background_exception = background_process.exception()
+                    except asyncio.CancelledError:
+                        self.log.error('Background process was cancelled!')
+                    if background_exception:
+                        self.log.error('Background process exited with an exception:')
+                        self.log.error(background_exception)
+                    self.log.error('At least one background process exited!')
+                    return False
+                else:
+                    self.log.debug('Found a not-yet-done background process...')
+            self.log.debug('All background processes still running.')
+        return True
+
     # Override if your batch system needs something more elaborate to query the job status
     batch_query_cmd = Unicode('',
         help="Command to run to query job status. Formatted using req_xyz traits as {xyz} "
@@ -314,6 +393,29 @@ async def cancel_batch_job(self):
         cmd = ' '.join((format_template(self.exec_prefix, **subvars),
                         format_template(self.batch_cancel_cmd, **subvars)))
         self.log.info('Cancelling job ' + self.job_id + ': ' + cmd)
+
+        if self.background_processes:
+            self.log.debug('Job being cancelled, cancelling background processes...')
+            for background_process in self.background_processes:
+                if not background_process.cancelled():
+                    try:
+                        background_process.cancel()
+                    except:
+                        self.log.error('Encountered an exception cancelling background process...')
+                    self.log.debug('Cancelled background process, waiting for it to finish...')
+                    try:
+                        await asyncio.wait([background_process])
+                    except asyncio.CancelledError:
+                        self.log.error('Successfully cancelled background process.')
+                        pass
+                    except:
+                        self.log.error('Background process exited with another exception!')
+                        raise
+                else:
+                    self.log.debug('Background process already cancelled...')
+            self.background_processes.clear()
+            self.log.debug('All background processes cancelled.')
+
         await self.run_command(cmd)
 
     def load_state(self, state):
@@ -361,6 +463,13 @@ async def poll(self):
         """Poll the process"""
         status = await self.query_job_status()
         if status in (JobStatus.PENDING, JobStatus.RUNNING, JobStatus.UNKNOWN):
+            if not self.background_tasks_ok():
+                self.log.debug('Going to stop job, since background tasks have failed!')
+                await self.stop(now=True)
+                status = await self.query_job_status()
+                if status not in (JobStatus.PENDING, JobStatus.RUNNING, JobStatus.UNKNOWN):
+                    self.clear_state()
+                    return 1
             return None
         else:
             self.clear_state()
@@ -420,6 +529,9 @@ async def start(self):
                         self.job_id, self.ip, self.port)
             )
 
+        if self.connect_to_job_cmd:
+            await self.connect_to_job()
+
         return self.ip, self.port
 
     async def stop(self, now=False):