From a4121c2396c84974bcb1e905c080451b3c6ce0d9 Mon Sep 17 00:00:00 2001 From: Richard Darst Date: Wed, 8 Aug 2018 01:33:38 +0300 Subject: [PATCH] Split off random port selection into a separate callable. - Allows selection of random ports within a certain range only as a first use case. - Closes: #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. --- README.md | 1 + batchspawner/batchspawner.py | 9 +++++++-- batchspawner/tests/test_spawners.py | 31 ++++++++++++++++++++++++++++- batchspawner/utils.py | 15 ++++++++++++++ 4 files changed, 53 insertions(+), 3 deletions(-) create mode 100644 batchspawner/utils.py diff --git a/README.md b/README.md index 3feb88f4..391b22d9 100644 --- a/README.md +++ b/README.md @@ -178,6 +178,7 @@ Changed Fixed * 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 * In Travis CI, blacklist jsonschema=3.0.0a1 because it breaks tests diff --git a/batchspawner/batchspawner.py b/batchspawner/batchspawner.py index 9f34854c..8a784f2e 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)