Merge pull request #58 from cmd-ntrf/remote_port

Select singleserver port number on remote host

Author: mbmilligan

batchspawner/__init__.py

@@ -1 +1,2 @@
 from .batchspawner import *
+from . import api

batchspawner/api.py

@@ -0,0 +1,16 @@
+import json
+from tornado import web
+from jupyterhub.apihandlers import  APIHandler, default_handlers
+
+class BatchSpawnerAPIHandler(APIHandler):
+    @web.authenticated
+    def post(self):
+        """POST set user's spawner port number"""
+        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"}))
+        self.set_status(201)
+
+default_handlers.append((r"/api/batchspawner", BatchSpawnerAPIHandler))

batchspawner/batchspawner.py

@@ -29,6 +29,7 @@
 from tornado.iostream import StreamClosedError
 
 from jupyterhub.spawner import Spawner
+from jupyterhub.traitlets import Command
 from traitlets import (
     Integer, Unicode, Float, Dict, default
 )
@@ -73,6 +74,9 @@ 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)
 
@@ -164,6 +168,9 @@ def _req_keepvars_default(self):
     # 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_') ]
@@ -342,14 +349,9 @@ def poll(self):
     @gen.coroutine
     def start(self):
         """Start the process"""
-        if self.user and self.user.server and self.user.server.port:
-            self.port = self.user.server.port
-            self.db.commit()
-        elif (jupyterhub.version_info < (0,7) and not self.user.server.port) or (
-             jupyterhub.version_info >= (0,7) and not self.port
-        ):
-            self.port = random_port()
-            self.db.commit()
+        if jupyterhub.version_info >= (0,8) and self.server:
+            self.server.port = self.port
+
         job = yield self.submit_batch_script()
 
         # We are called with a timeout, and if the timeout expires this function will
@@ -374,16 +376,19 @@ def start(self):
             yield gen.sleep(self.startup_poll_interval)
 
         self.current_ip = self.state_gethost()
+        while self.current_port == 0:
+            yield gen.sleep(self.startup_poll_interval)
+
         if jupyterhub.version_info < (0,7):
             # store on user for pre-jupyterhub-0.7:
-            self.user.server.port = self.port
+            self.user.server.port = self.current_port
             self.user.server.ip = self.current_ip
         self.db.commit()
         self.log.info("Notebook server job {0} started at {1}:{2}".format(
-                        self.job_id, self.current_ip, self.port)
+                        self.job_id, self.current_ip, self.current_port)
             )
 
-        return self.current_ip, self.port
+        return self.current_ip, self.current_port
 
     @gen.coroutine
     def stop(self, now=False):

batchspawner/singleuser.py

@@ -0,0 +1,21 @@
+from jupyterhub.singleuser import SingleUserNotebookApp
+from jupyterhub.utils import random_port, url_path_join
+from traitlets import default
+
+class BatchSingleUserNotebookApp(SingleUserNotebookApp):
+    @default('port')
+    def _port(self):
+        return random_port()
+
+    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()
+
+def main(argv=None):
+    return BatchSingleUserNotebookApp.launch_instance(argv)
+
+if __name__ == "__main__":
+    main()

batchspawner/tests/test_spawners.py

@@ -10,13 +10,14 @@
 from tornado import gen
 
 try:
-    from jupyterhub.objects import Hub
+    from jupyterhub.objects import Hub, Server
     from jupyterhub.user import User
 except:
     pass
 
 testhost = "userhost123"
 testjob  = "12345"
+testport = 54321
 
 class BatchDummy(BatchSpawnerRegexStates):
     exec_prefix = ''
@@ -59,8 +60,11 @@ def new_spawner(db, spawner_class=BatchDummy, **kwargs):
     else:
         hub = Hub()
         user = User(user, {})
+        server = Server()
+        kwargs.setdefault('server', server)
     kwargs.setdefault('hub', hub)
     kwargs.setdefault('user', user)
+    kwargs.setdefault('current_port', testport)
     kwargs.setdefault('INTERRUPT_TIMEOUT', 1)
     kwargs.setdefault('TERM_TIMEOUT', 1)
     kwargs.setdefault('KILL_TIMEOUT', 1)

scripts/batchspawner-singleuser

@@ -0,0 +1,6 @@
+#!/usr/bin/env python3
+
+from batchspawner.singleuser import main
+
+if __name__ == '__main__':
+    main()

setup.py

@@ -14,6 +14,7 @@
 import sys
 
 from setuptools import setup
+from glob import glob
 
 pjoin = os.path.join
 here = os.path.abspath(os.path.dirname(__file__))
@@ -28,6 +29,7 @@
 
 setup_args = dict(
     name                = 'batchspawner',
+    scripts             = glob(pjoin('scripts', '*')),
     packages            = ['batchspawner'],
     version             = version_ns['__version__'],
     description         = """Batchspawner: A spawner for Jupyterhub to spawn notebooks using batch resource managers.""",