Merge pull request #143 from rkdarst/singleuserapp_and_1.0

Singleuserapp and 1.0

batchspawner/api.py

@@ -5,12 +5,18 @@
 class BatchSpawnerAPIHandler(APIHandler):
     @web.authenticated
     def post(self):
-        """POST set user's spawner port number"""
-        user = self.get_current_user()
+        """POST set user spawner data"""
+        if hasattr(self, 'current_user'):
+            # Jupyterhub compatability, (september 2018, d79a99323ef1d)
+            user = self.current_user
+        else:
+            # Previous jupyterhub, 0.9.4 and before.
+            user = self.get_current_user()
         data = self.get_json_body()
-        port = int(data.get('port', 0))
-        user.spawner.current_port = port
-        self.finish(json.dumps({"message": "BatchSpawner port configured"}))
+        for key, value in data.items():
+            if hasattr(user.spawner, key):
+                setattr(user.spawner, key, value)
+        self.finish(json.dumps({"message": "BatchSpawner data configured"}))
         self.set_status(201)
 
 default_handlers.append((r"/api/batchspawner", BatchSpawnerAPIHandler))

batchspawner/batchspawner.py

@@ -74,9 +74,6 @@ class BatchSpawnerBase(Spawner):
         state_gethost
     """
 
-    # override default since will need to set the listening port using the api
-    cmd = Command(['batchspawner-singleuser'], allow_none=True).tag(config=True)
-
     # override default since batch systems typically need longer
     start_timeout = Integer(300).tag(config=True)
 
@@ -165,12 +162,6 @@ def _req_keepvars_default(self):
     # Will get the raw output of the job status command unless overridden
     job_status = Unicode()
 
-    # Will get the address of the server as reported by job manager
-    current_ip = Unicode()
-
-    # Will get the port of the server as reported by singleserver
-    current_port = Integer()
-
     # Prepare substitution variables for templates using req_xyz traits
     def get_req_subvars(self):
         reqlist = [ t for t in self.trait_names() if t.startswith('req_') ]
@@ -190,7 +181,7 @@ def parse_job_id(self, output):
         return output
 
     def cmd_formatted_for_batch(self):
-        return ' '.join(self.cmd + self.get_args())
+        return ' '.join(['batchspawner-singleuser'] + self.cmd + self.get_args())
 
     @gen.coroutine
     def run_command(self, cmd, input=None, env=None):
@@ -349,6 +340,9 @@ def poll(self):
     @gen.coroutine
     def start(self):
         """Start the process"""
+        self.ip = self.traits()['ip'].default_value
+        self.port = self.traits()['port'].default_value
+
         if jupyterhub.version_info >= (0,8) and self.server:
             self.server.port = self.port
 
@@ -375,20 +369,24 @@ def start(self):
                            ' after starting.')
             yield gen.sleep(self.startup_poll_interval)
 
-        self.current_ip = self.state_gethost()
-        while self.current_port == 0:
+        self.ip = self.state_gethost()
+        while self.port == 0:
             yield gen.sleep(self.startup_poll_interval)
+            # Test framework: For testing, mock_port is set because we
+            # don't actually run the single-user server yet.
+            if hasattr(self, 'mock_port'):
+                self.port = self.mock_port
 
         if jupyterhub.version_info < (0,7):
             # store on user for pre-jupyterhub-0.7:
-            self.user.server.port = self.current_port
-            self.user.server.ip = self.current_ip
+            self.user.server.port = self.port
+            self.user.server.ip = self.ip
         self.db.commit()
         self.log.info("Notebook server job {0} started at {1}:{2}".format(
-                        self.job_id, self.current_ip, self.current_port)
+                        self.job_id, self.ip, self.port)
             )
 
-        return self.current_ip, self.current_port
+        return self.ip, self.port
 
     @gen.coroutine
     def stop(self, now=False):
@@ -408,7 +406,7 @@ def stop(self, now=False):
             yield gen.sleep(1.0)
         if self.job_id:
             self.log.warn("Notebook server job {0} at {1}:{2} possibly failed to terminate".format(
-                             self.job_id, self.current_ip, self.port)
+                             self.job_id, self.ip, self.port)
                 )
 
 

batchspawner/singleuser.py

@@ -1,21 +1,25 @@
-from jupyterhub.singleuser import SingleUserNotebookApp
-from jupyterhub.utils import random_port, url_path_join
-from traitlets import default
+import os
+import sys
 
-class BatchSingleUserNotebookApp(SingleUserNotebookApp):
-    @default('port')
-    def _port(self):
-        return random_port()
+from runpy import run_path
+from shutil import which
 
-    def start(self):
-        # Send Notebook app's port number to remote Spawner
-        self.hub_auth._api_request(method='POST',
-                                   url=url_path_join(self.hub_api_url, 'batchspawner'),
-                                   json={'port' : self.port})
-        super().start()
+from jupyterhub.utils import random_port, url_path_join
+from jupyterhub.services.auth import HubAuth
 
 def main(argv=None):
-    return BatchSingleUserNotebookApp.launch_instance(argv)
+    port = random_port()
+    hub_auth = HubAuth()
+    hub_auth.client_ca = os.environ.get('JUPYTERHUB_SSL_CLIENT_CA', '')
+    hub_auth.certfile = os.environ.get('JUPYTERHUB_SSL_CERTFILE', '')
+    hub_auth.keyfile = os.environ.get('JUPYTERHUB_SSL_KEYFILE', '')
+    hub_auth._api_request(method='POST',
+                          url=url_path_join(hub_auth.api_url, 'batchspawner'),
+                          json={'port' : port})
+
+    cmd_path = which(sys.argv[1])
+    sys.argv = sys.argv[1:] + ['--port={}'.format(port)]
+    run_path(cmd_path, run_name="__main__")
 
 if __name__ == "__main__":
     main()

batchspawner/tests/test_spawners.py

@@ -64,7 +64,7 @@ def new_spawner(db, spawner_class=BatchDummy, **kwargs):
         kwargs.setdefault('server', server)
     kwargs.setdefault('hub', hub)
     kwargs.setdefault('user', user)
-    kwargs.setdefault('current_port', testport)
+    kwargs.setdefault('mock_port', testport)
     kwargs.setdefault('INTERRUPT_TIMEOUT', 1)
     kwargs.setdefault('TERM_TIMEOUT', 1)
     kwargs.setdefault('KILL_TIMEOUT', 1)
@@ -84,7 +84,7 @@ def check_ip(spawner, value):
     if version_info < (0,7):
         assert spawner.user.server.ip == value
     else:
-        assert spawner.current_ip == value
+        assert spawner.ip == value
 
 def test_spawner_start_stop_poll(db, io_loop):
     spawner = new_spawner(db=db)
@@ -287,7 +287,7 @@ def test_torque(db, io_loop):
         'req_epilogue': 'EPILOGUE',
         }
     batch_script_re_list = [
-        re.compile(r'^PROLOGUE.*^singleuser_command.*^EPILOGUE', re.S|re.M),
+        re.compile(r'^PROLOGUE.*^batchspawner-singleuser singleuser_command.*^EPILOGUE', re.S|re.M),
         re.compile(r'mem=5678'),
         re.compile(r'ppn=5'),
         re.compile(r'^#PBS some_option_asdf', re.M),
@@ -315,7 +315,7 @@ def test_moab(db, io_loop):
         'req_epilogue': 'EPILOGUE',
         }
     batch_script_re_list = [
-        re.compile(r'^PROLOGUE.*^singleuser_command.*^EPILOGUE', re.S|re.M),
+        re.compile(r'^PROLOGUE.*^batchspawner-singleuser singleuser_command.*^EPILOGUE', re.S|re.M),
         re.compile(r'mem=5678'),
         re.compile(r'ppn=5'),
         re.compile(r'^#PBS some_option_asdf', re.M),
@@ -376,7 +376,7 @@ def test_slurm(db, io_loop):
         'req_reservation': 'RES123',
         }
     batch_script_re_list = [
-        re.compile(r'PROLOGUE.*srun singleuser_command.*EPILOGUE', re.S),
+        re.compile(r'PROLOGUE.*srun batchspawner-singleuser singleuser_command.*EPILOGUE', re.S),
         re.compile(r'^#SBATCH \s+ --cpus-per-task=5', re.X|re.M),
         re.compile(r'^#SBATCH \s+ --time=3-05:10:10', re.X|re.M),
         re.compile(r'^#SBATCH \s+ some_option_asdf', re.X|re.M),
@@ -441,7 +441,7 @@ def test_condor(db, io_loop):
         'req_options': 'some_option_asdf',
         }
     batch_script_re_list = [
-        re.compile(r'exec singleuser_command'),
+        re.compile(r'exec batchspawner-singleuser singleuser_command'),
         re.compile(r'RequestCpus = 5'),
         re.compile(r'RequestMemory = 5678'),
         re.compile(r'^some_option_asdf', re.M),
@@ -470,7 +470,7 @@ def test_lfs(db, io_loop):
         'req_epilogue': 'EPILOGUE',
         }
     batch_script_re_list = [
-        re.compile(r'^PROLOGUE.*^singleuser_command.*^EPILOGUE', re.S|re.M),
+        re.compile(r'^PROLOGUE.*^batchspawner-singleuser singleuser_command.*^EPILOGUE', re.S|re.M),
         re.compile(r'#BSUB\s+-q\s+some_queue', re.M),
         ]
     script = [