diff --git a/README.md b/README.md index b126f827..c5c65675 100644 --- a/README.md +++ b/README.md @@ -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) diff --git a/batchspawner/batchspawner.py b/batchspawner/batchspawner.py index 6d71af69..81890691 100644 --- a/batchspawner/batchspawner.py +++ b/batchspawner/batchspawner.py @@ -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 @@ -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_') ] @@ -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() diff --git a/batchspawner/tests/test_spawners.py b/batchspawner/tests/test_spawners.py index 01873486..64f6fbf7 100644 --- a/batchspawner/tests/test_spawners.py +++ b/batchspawner/tests/test_spawners.py @@ -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): @@ -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 diff --git a/batchspawner/utils.py b/batchspawner/utils.py new file mode 100644 index 00000000..a59a2202 --- /dev/null +++ b/batchspawner/utils.py @@ -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)