Skip to content

Commit

Permalink
Split off random port selection into a separate callable.
Browse files Browse the repository at this point in the history
- Allows selection of random ports within a certain range only as a
  first use case.
- Closes: jupyterhub#102 via a different strategy.
- Currently uses the traitlet Any, but should be Callable - but this
  is not in upstream traitlets, it is special to kubespawner.  So
  leave Any for now and change later.
  • Loading branch information
rkdarst committed Aug 7, 2018
1 parent c91dc0b commit 3afc87e
Show file tree
Hide file tree
Showing 4 changed files with 53 additions and 3 deletions.
1 change: 1 addition & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -163,6 +163,7 @@ clusters, as well as an option to run a local notebook directly on the jupyterhu
* Pass the environment dictionary to the queue and cancel commands as well. This is mostly user environment, but may be useful to these commands as well in some cases. #108, #111
* SlurmSpawner: add the `req_reservation` option. #
* Improve debugging on failed submission by raising errors including error messages from the commands. #106
* Allow selecting random ports via different logic, for example select a random port only within a certain range via: `c.BatchSpawner.random_port = batchspawner.utils.random_port_range(low, high)`.
* Many other non-user or developer visible changes. #107 #106 #100

### v0.8.1 (bugfix release)
Expand Down
9 changes: 7 additions & 2 deletions batchspawner/batchspawner.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@

from jupyterhub.spawner import Spawner
from traitlets import (
Integer, Unicode, Float, Dict, default
Integer, Unicode, Float, Dict, Any, default
)

from jupyterhub.utils import random_port
Expand Down Expand Up @@ -164,6 +164,11 @@ def _req_keepvars_default(self):
# Will get the address of the server as reported by job manager
current_ip = Unicode()

# Random port function
random_port = Any(random_port,
help="Function to call to request a random port for singleuser spawer."
).tag(config=True)

# 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 Down Expand Up @@ -348,7 +353,7 @@ def start(self):
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.port = self.random_port()
self.db.commit()
job = yield self.submit_batch_script()

Expand Down
31 changes: 30 additions & 1 deletion batchspawner/tests/test_spawners.py
Original file line number Diff line number Diff line change
Expand Up @@ -271,7 +271,7 @@ def run_command(self, cmd, input=None, env=None):
# batch_poll_cmd
status = io_loop.run_sync(spawner.poll, timeout=5)
assert status == 1

return spawner


def test_torque(db, io_loop):
Expand Down Expand Up @@ -475,3 +475,32 @@ def test_keepvars(db, io_loop):
run_typical_slurm_spawner(db, io_loop,
spawner_kwargs=spawner_kwargs,
batch_script_re_list=batch_script_re_list)


def test_random_port(db, io_loop):
# Test single fixed port (probably don't ever use this!)
spawner_kwargs = {
'random_port': lambda: 12345,
}
spawner = run_typical_slurm_spawner(db, io_loop,
spawner_kwargs=spawner_kwargs)
assert spawner.port == 12345

# Random port range
import batchspawner.utils as utils
spawner_kwargs = {
'random_port': utils.random_port_range(12300, 12301),
}
spawner = run_typical_slurm_spawner(db, io_loop,
spawner_kwargs=spawner_kwargs)
assert 12300 <= spawner.port and spawner.port <= 12301

# Arbitrary function
def port_555():
return 555
spawner_kwargs = {
'random_port': port_555,
}
spawner = run_typical_slurm_spawner(db, io_loop,
spawner_kwargs=spawner_kwargs)
assert spawner.port == 555
15 changes: 15 additions & 0 deletions batchspawner/utils.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,15 @@
# Miscellaneous utilities

import random

def random_port_range(low, high):
"""Factory function to select a random port number in the range [low,high].
Usage: c.BatchSpawner.random_port = random_port_range(low, high)
"""
# TODO: This does not prevent port number conflicts like
# jupyterhub/utils.random_port tries to do. But we actually
# can't, because we run on a different host, so
# jupyterhub.utils.random_port actually doesn't work for us. #58
# tries to do this better.
return lambda: random.randint(low, high)

0 comments on commit 3afc87e

Please sign in to comment.