diff --git a/batchspawner/__init__.py b/batchspawner/__init__.py index bba3352d..e6b50a69 100644 --- a/batchspawner/__init__.py +++ b/batchspawner/__init__.py @@ -1 +1,2 @@ from .batchspawner import * +from . import api \ No newline at end of file diff --git a/batchspawner/api.py b/batchspawner/api.py new file mode 100644 index 00000000..f78693a4 --- /dev/null +++ b/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)) diff --git a/batchspawner/batchspawner.py b/batchspawner/batchspawner.py index 166d7328..1a8a77b5 100644 --- a/batchspawner/batchspawner.py +++ b/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): diff --git a/batchspawner/singleuser.py b/batchspawner/singleuser.py new file mode 100644 index 00000000..6d493262 --- /dev/null +++ b/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() \ No newline at end of file diff --git a/batchspawner/tests/test_spawners.py b/batchspawner/tests/test_spawners.py index c08d6ba8..672ce395 100644 --- a/batchspawner/tests/test_spawners.py +++ b/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) diff --git a/scripts/batchspawner-singleuser b/scripts/batchspawner-singleuser new file mode 100644 index 00000000..1b91fe47 --- /dev/null +++ b/scripts/batchspawner-singleuser @@ -0,0 +1,6 @@ +#!/usr/bin/env python3 + +from batchspawner.singleuser import main + +if __name__ == '__main__': + main() \ No newline at end of file diff --git a/setup.py b/setup.py index d84d92e1..2037998e 100755 --- a/setup.py +++ b/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.""",