Add API handler to post remote server port number

Author: cmd-ntrf

batchspawner/__init__.py

@@ -1 +1,3 @@
 from .batchspawner import *
+from . import singleuser
+from . import api

batchspawner/api.py

@@ -0,0 +1,17 @@
+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()
+        if user.spawner.port == 0:
+            port = data.get('port', 0)
+            user.spawner.port = int(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)
 
@@ -342,14 +346,7 @@ 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()
+        self.port = self.server.port = 0
         job = yield self.submit_batch_script()
 
         # We are called with a timeout, and if the timeout expires this function will
@@ -374,6 +371,9 @@ def start(self):
             yield gen.sleep(self.startup_poll_interval)
 
         self.current_ip = self.state_gethost()
+        while self.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

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()

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.""",