Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Singleuserapp and 1.0 #143

Merged
merged 11 commits into from
Jun 18, 2019
16 changes: 11 additions & 5 deletions batchspawner/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -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))
32 changes: 15 additions & 17 deletions batchspawner/batchspawner.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down Expand Up @@ -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_') ]
Expand All @@ -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):
Expand Down Expand Up @@ -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

Expand All @@ -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):
Expand All @@ -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)
)


Expand Down
32 changes: 18 additions & 14 deletions batchspawner/singleuser.py
Original file line number Diff line number Diff line change
@@ -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()
14 changes: 7 additions & 7 deletions batchspawner/tests/test_spawners.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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)
Expand Down Expand Up @@ -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),
Expand Down Expand Up @@ -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),
Expand Down Expand Up @@ -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),
Expand Down Expand Up @@ -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),
Expand Down Expand Up @@ -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 = [
Expand Down