Skip to content

Commit

Permalink
BatchSpawnerBase: Extend connect_to_job for port forwarding approaches.
Browse files Browse the repository at this point in the history
This allows to use {rport} inside the connect_to_job_cmd.
If this is done, {port} is set to a random_port chosen locally
on the Hub, and {rport} is set to the original remote port.
This is useful e.g. for SSH port forwarding.

If this is used, the {rport} of the notebook is saved into
a new class variable in case it will be needed again.
  • Loading branch information
olifre committed Mar 9, 2021
1 parent 8f46940 commit 6c5e850
Showing 1 changed file with 18 additions and 3 deletions.
21 changes: 18 additions & 3 deletions batchspawner/batchspawner.py
Original file line number Diff line number Diff line change
Expand Up @@ -176,8 +176,16 @@ def _req_keepvars_default(self):
help="Command to connect to running batch job and forward the port "
"of the running notebook to the Hub. If empty, direct connectivity is assumed. "
"Uses self.job_id as {job_id} and the self.port as {port}."
"If {rport} is used in this string, it is set to self.port, "
"and a new random self.port is chosen locally and used as {port}."
"This is useful e.g. for SSH port forwarding."
).tag(config=True)

rport = Integer(0,
help="Remote port of notebook, will be set if it differs from the self.port."
"This is set by connect_to_job() if needed."
)

# Raw output of job submission command unless overridden
job_id = Unicode()

Expand Down Expand Up @@ -208,12 +216,19 @@ def cmd_formatted_for_batch(self):

async def connect_to_job(self):
"""This command ensures the port of the singleuser server is reachable from the
Batchspawner machine. By default, it does nothing, i.e. direct connectivity
is assumed.
Batchspawner machine. Only called if connect_to_job_cmd is set.
If the template string connect_to_job_cmd contains {rport},
a new random self.port is chosen locally (useful e.g. for SSH port forwarding).
"""
subvars = self.get_req_subvars()
subvars['job_id'] = self.job_id
subvars['port'] = self.port
if '{rport}' in self.connect_to_job_cmd:
self.rport = self.port
self.port = random_port()
subvars['rport'] = self.rport
subvars['port'] = self.port
else:
subvars['port'] = self.port
cmd = ' '.join((format_template(self.exec_prefix, **subvars),
format_template(self.connect_to_job_cmd, **subvars)))
await self.run_background_command(cmd)
Expand Down

0 comments on commit 6c5e850

Please sign in to comment.