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

Miscellaneous minor tune ups, shouldn't effect functioning #100

Merged
merged 5 commits into from
Aug 10, 2018
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
91 changes: 50 additions & 41 deletions batchspawner/batchspawner.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
"""
import pwd
import os
import re

import xml.etree.ElementTree as ET

Expand Down Expand Up @@ -46,7 +47,7 @@ def format_template(template, *args, **kwargs):
"""
if isinstance(template, Template):
return template.render(*args, **kwargs)
elif '{{' in template or '{%' in template:
elif '{{' in template or '{%' in template:
return Template(template).render(*args, **kwargs)
return template.format(*args, **kwargs)

Expand Down Expand Up @@ -78,52 +79,52 @@ class BatchSpawnerBase(Spawner):
# override default server ip since batch jobs normally running remotely
ip = Unicode("0.0.0.0", help="Address for singleuser server to listen at").tag(config=True)

exec_prefix = Unicode('sudo -E -u {username}', \
exec_prefix = Unicode('sudo -E -u {username}',
help="Standard executon prefix (e.g. the default sudo -E -u {username})"
).tag(config=True)

# all these req_foo traits will be available as substvars for templated strings
req_queue = Unicode('', \
req_queue = Unicode('',
help="Queue name to submit job to resource manager"
).tag(config=True)

req_host = Unicode('', \
req_host = Unicode('',
help="Host name of batch server to submit job to resource manager"
).tag(config=True)

req_memory = Unicode('', \
req_memory = Unicode('',
help="Memory to request from resource manager"
).tag(config=True)

req_nprocs = Unicode('', \
req_nprocs = Unicode('',
help="Number of processors to request from resource manager"
).tag(config=True)

req_ngpus = Unicode('', \
req_ngpus = Unicode('',
help="Number of GPUs to request from resource manager"
).tag(config=True)

req_runtime = Unicode('', \
req_runtime = Unicode('',
help="Length of time for submitted job to run"
).tag(config=True)

req_partition = Unicode('', \
req_partition = Unicode('',
help="Partition name to submit job to resource manager"
).tag(config=True)

req_account = Unicode('', \
req_account = Unicode('',
help="Account name string to pass to the resource manager"
).tag(config=True)

req_options = Unicode('', \
req_options = Unicode('',
help="Other options to include into job submission script"
).tag(config=True)

req_prologue = Unicode('', \
req_prologue = Unicode('',
help="Script to run before single user server starts."
).tag(config=True)

req_epilogue = Unicode('', \
req_epilogue = Unicode('',
help="Script to run after single user server ends."
).tag(config=True)

Expand All @@ -143,7 +144,7 @@ def _req_homedir_default(self):
def _req_keepvars_default(self):
return ','.join(self.get_env().keys())

batch_script = Unicode('', \
batch_script = Unicode('',
help="Template for job submission script. Traits on this class named like req_xyz "
"will be substituted in the template for {xyz} using string.Formatter. "
"Must include {cmd} which will be replaced with the jupyterhub-singleuser command line."
Expand All @@ -166,7 +167,7 @@ def get_req_subvars(self):
subvars[t[4:]] = getattr(self, t)
return subvars

batch_submit_cmd = Unicode('', \
batch_submit_cmd = Unicode('',
help="Command to run to submit batch scripts. Formatted using req_xyz traits as {xyz}."
).tag(config=True)

Expand Down Expand Up @@ -233,7 +234,7 @@ def submit_batch_script(self):
return self.job_id

# Override if your batch system needs something more elaborate to read the job status
batch_query_cmd = Unicode('', \
batch_query_cmd = Unicode('',
help="Command to run to read job status. Formatted using req_xyz traits as {xyz} "
"and self.job_id as {job_id}."
).tag(config=True)
Expand Down Expand Up @@ -324,7 +325,7 @@ def poll(self):
self.clear_state()
return 1

startup_poll_interval = Float(0.5, \
startup_poll_interval = Float(0.5,
help="Polling interval (seconds) to check job state during startup"
).tag(config=True)

Expand All @@ -334,8 +335,9 @@ def start(self):
if self.user and self.user.server and self.user.server.port:
self.port = self.user.server.port
self.db.commit()
elif (jupyterhub.version_info < (0,7) and not self.user.server.port) or \
(jupyterhub.version_info >= (0,7) and not self.port):
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.db.commit()
job = yield self.submit_batch_script()
Expand All @@ -356,8 +358,8 @@ def start(self):
else:
self.log.warn('Job ' + self.job_id + ' neither pending nor running.\n' +
self.job_status)
raise RuntimeError('The Jupyter batch job has disappeared '
' while pending in the queue or died immediately '
raise RuntimeError('The Jupyter batch job has disappeared'
' while pending in the queue or died immediately'
' after starting.')
yield gen.sleep(self.startup_poll_interval)

Expand Down Expand Up @@ -394,7 +396,6 @@ def stop(self, now=False):
self.job_id, self.current_ip, self.port)
)

import re

class BatchSpawnerRegexStates(BatchSpawnerBase):
"""Subclass of BatchSpawnerBase that uses config-supplied regular expressions
Expand Down Expand Up @@ -431,13 +432,15 @@ def state_ispending(self):
assert self.state_pending_re, "Misconfigured: define state_running_re"
if self.job_status and re.search(self.state_pending_re, self.job_status):
return True
else: return False
else:
return False

def state_isrunning(self):
assert self.state_running_re, "Misconfigured: define state_running_re"
if self.job_status and re.search(self.state_running_re, self.job_status):
return True
else: return False
else:
return False

def state_gethost(self):
assert self.state_exechost_re, "Misconfigured: define state_exechost_re"
Expand All @@ -450,6 +453,7 @@ def state_gethost(self):
else:
return match.expand(self.state_exechost_exp)


class TorqueSpawner(BatchSpawnerRegexStates):
batch_script = Unicode("""#!/bin/sh
#PBS -q {queue}@{host}
Expand All @@ -473,6 +477,7 @@ class TorqueSpawner(BatchSpawnerRegexStates):
state_running_re = Unicode(r'<job_state>R</job_state>').tag(config=True)
state_exechost_re = Unicode(r'<exec_host>((?:[\w_-]+\.?)+)/\d+').tag(config=True)


class MoabSpawner(TorqueSpawner):
# outputs job id string
batch_submit_cmd = Unicode('msub').tag(config=True)
Expand All @@ -483,6 +488,7 @@ class MoabSpawner(TorqueSpawner):
state_running_re = Unicode(r'State="Running"').tag(config=True)
state_exechost_re = Unicode(r'AllocNodeList="([^\r\n\t\f :"]*)').tag(config=True)


class UserEnvMixin:
"""Mixin class that computes values for USER, SHELL and HOME in the environment passed to
the job submission subprocess in case the batch system needs these for the batch script."""
Expand All @@ -504,23 +510,8 @@ def get_env(self):
env = self.user_env(env)
return env

class SlurmSpawner(UserEnvMixin,BatchSpawnerRegexStates):
"""A Spawner that just uses Popen to start local processes."""

# all these req_foo traits will be available as substvars for templated strings
req_cluster = Unicode('', \
help="Cluster name to submit job to resource manager"
).tag(config=True)

req_qos = Unicode('', \
help="QoS name to submit job to resource manager"
).tag(config=True)

req_srun = Unicode('srun',
help="Job step wrapper, default 'srun'. Set to '' you do not want "
"to run in job step (affects environment handling)"
).tag(config=True)

class SlurmSpawner(UserEnvMixin,BatchSpawnerRegexStates):
batch_script = Unicode("""#!/bin/bash
#SBATCH --output={{homedir}}/jupyterhub_slurmspawner_%j.log
#SBATCH --job-name=spawner-jupyterhub
Expand All @@ -540,6 +531,21 @@ class SlurmSpawner(UserEnvMixin,BatchSpawnerRegexStates):
echo "jupyterhub-singleuser ended gracefully"
{{epilogue}}
""").tag(config=True)

# all these req_foo traits will be available as substvars for templated strings
req_cluster = Unicode('',
help="Cluster name to submit job to resource manager"
).tag(config=True)

req_qos = Unicode('',
help="QoS name to submit job to resource manager"
).tag(config=True)

req_srun = Unicode('srun',
help="Job step wrapper, default 'srun'. Set to '' you do not want "
"to run in job step (affects environment handling)"
).tag(config=True)

# outputs line like "Submitted batch job 209"
batch_submit_cmd = Unicode('sbatch --parsable').tag(config=True)
# outputs status and exec node like "RUNNING hostname"
Expand All @@ -561,6 +567,7 @@ def parse_job_id(self, output):
raise e
return id


class MultiSlurmSpawner(SlurmSpawner):
'''When slurm has been compiled with --enable-multiple-slurmd, the
administrator sets the name of the slurmd instance via the slurmd -N
Expand All @@ -573,6 +580,7 @@ def state_gethost(self):
host = SlurmSpawner.state_gethost(self)
return self.daemon_resolver.get(host, host)


class GridengineSpawner(BatchSpawnerBase):
batch_script = Unicode("""#!/bin/bash
#$ -j yes
Expand Down Expand Up @@ -620,6 +628,7 @@ def state_gethost(self):
self.log.error("Spawner unable to match host addr in job {0} with status {1}".format(self.job_id, self.job_status))
return


class CondorSpawner(UserEnvMixin,BatchSpawnerRegexStates):
batch_script = Unicode("""
Executable = /bin/sh
Expand Down Expand Up @@ -657,6 +666,7 @@ def parse_job_id(self, output):
def cmd_formatted_for_batch(self):
return super(CondorSpawner,self).cmd_formatted_for_batch().replace('"','""').replace("'","''")


class LsfSpawner(BatchSpawnerBase):
'''A Spawner that uses IBM's Platform Load Sharing Facility (LSF) to launch notebooks.'''

Expand Down Expand Up @@ -701,7 +711,6 @@ def state_isrunning(self):
if self.job_status:
return self.job_status.split(' ')[0].upper() == 'RUN'


def state_gethost(self):
if self.job_status:
return self.job_status.split(' ')[1].strip()
Expand Down