#job -*- coding: utf-8 -*-
"""
Monitor the queue for torque or slurm.
Provides a class to monitor the torque, slurm, or local queues with identical
syntax.
At its simplest, you can use it like::
q = queue.Queue()
q.jobs
q.running
q.queued
q.complete
All of the above commands return a dictionary of:
{job_no: Queue.QueueJob}
Queue.QueueJob classes include information on job state, owner, queue, nodes,
threads, exitcode, etc.
Queue also defines a wait() method that takes a list of job numbers, job.Job()
objects, or JobQueue.Job objects and blocks until those jobs to complete
The default cluster environment is also defined in this file as MODE, it can be
set directly or with the get_cluster_environment() function definied here.
"""
import re
import sys
import pwd # Used to get usernames for queue
import socket # Used to get the hostname
import getpass # Used to get usernames for queue
from datetime import datetime as _dt
from time import time, sleep
from subprocess import check_output, CalledProcessError
# For parsing torque queues
import xml.etree.ElementTree as ET
###############################################################################
# Our functions #
###############################################################################
from . import run
from . import logme
from . import conf
from . import ClusterError
#########################
# Which system to use #
#########################
from . import ALLOWED_MODES
#########################################################
# The multiprocessing pool, only used in 'local' mode #
#########################################################
from . import local
# Funtions to import if requested
__all__ = ['Queue', 'wait', 'check_queue', 'get_cluster_environment']
# We only need the queue defaults
_defaults = conf.get_option('queue')
# This is set in the get_cluster_environment() function.
MODE = ''
# Define torque-to-slurm mappings
TORQUE_SLURM_STATES = {
'C': 'completed',
'E': 'completing',
'H': 'held', # Not a SLURM state
'Q': 'pending',
'R': 'running',
'T': 'suspended',
'W': 'running',
'S': 'suspended',
}
# Define job states
GOOD_STATES = ['complete', 'completed', 'special_exit']
ACTIVE_STATES = ['configuring', 'completing', 'pending',
'running']
BAD_STATES = ['boot_fail', 'cancelled', 'failed',
'node_fail', 'timeout', 'disappeared']
UNCERTAIN_STATES = ['hold', 'preempted', 'stopped',
'suspended']
ALL_STATES = GOOD_STATES + ACTIVE_STATES + BAD_STATES + UNCERTAIN_STATES
DONE_STATES = GOOD_STATES + BAD_STATES
###############################################################################
# The Queue Class #
###############################################################################
[docs]class Queue(object):
"""A wrapper for torque, slurm, or local queues.
Attributes:
jobs (dict) A dictionary of all jobs in this queue
{jobid: Queue.QueueJob}
max_jobs (int): The maximum number of jobs allowed in the queue
job_states (int): A list of the different states of jobs in this
queue
active_job_count (int): A count of all jobs that are either queued or
running in the current queue
can_submit (bool): True if total active jobs is less than max_jobs
"""
def __init__(self, user=None, partition=None, qtype=None,):
"""Can filter by user, queue type or partition on initialization.
Args:
user (str): Optional usernameto filter the queue with.
If user='self' or 'current', the current user will
be used.
partition (str): Optional partition to filter the queue with.
qtype (str): 'torque', 'slurm', or 'local', defaults to auto-detect.
"""
# Get user ID as an int UID
if user:
if user == 'self' or user == 'current':
self.user = getpass.getuser()
"""The username if defined."""
self.uid = pwd.getpwnam(self.user).pw_uid
elif user == 'ALL':
self.user = None
else:
if isinstance(user, int) \
or (isinstance(user, str) and user.isdigit()):
self.uid = int(user)
else:
self.uid = pwd.getpwnam(str(user)).pw_uid
else:
self.uid = None
self.user = pwd.getpwuid(self.uid).pw_name if self.uid else None
self.partition = partition
"""The partition if defined."""
# Set queue length
self.max_jobs = conf.get_option('queue', 'max_jobs')
"""The maximum number of jobs that can run in this queue."""
# Support python2, which hates reciprocal import
from .job import Job
from .local import Job as QJob
self._Job = Job
self._JobQueue = QJob
# Get sleep time and update time
self.queue_update_time = conf.get_option('queue', 'queue_update', 2)
self.sleep_len = conf.get_option('queue', 'sleep_len', 2)
# Set type
if qtype:
check_queue(qtype)
else:
check_queue()
self.qtype = qtype if qtype else MODE
# Allow tracking of updates to prevent too many updates
self._updating = False
# Will contain a dict of QueueJob objects indexed by ID
self.jobs = {}
"""All jobs currently in this queue."""
self._update()
####################
# Public Methods #
####################
def check_dependencies(self, dependencies):
"""Check if dependencies are running.
Args:
dependencies (list): List of job IDs
Returns:
str: 'active' if dependencies are running or queued, 'good' if
completed, 'bad' if failed, cancelled, or suspended, 'absent'
otherwise.
"""
for dep in run.listify(dependencies):
dep = int(dep)
if dep not in self.jobs:
return 'absent'
state = self.jobs[dep].state
if state in ACTIVE_STATES:
return 'active'
elif state in GOOD_STATES:
return 'good'
elif state in BAD_STATES or state in UNCERTAIN_STATES:
return 'bad'
[docs] def wait(self, jobs):
"""Block until all jobs in jobs are complete.
Update time is dependant upon the queue_update parameter in
your ~/.fyrd file.
In addition, wait() will not return until between 1 and 3
seconds after a job has completed, irrespective of queue_update
time. This allows time for any copy operations to complete after
the job exits.
Args:
jobs: A job or list of jobs to check. Can be one of: Job or
multiprocessing.pool.ApplyResult objects, job ID (int/str),
or a object or a list/tuple of multiple Jobs or job IDs.
Returns:
True on success False or None on failure.
"""
self.update()
logme.log('Queue waiting.', 'debug')
# Sanitize arguments
if not isinstance(jobs, (list, tuple)):
jobs = [jobs]
for job in jobs:
if not isinstance(job, (str, int, self.QueueJob, self._Job,
self._JobQueue)):
raise ClusterError('job must be int, string, or Job, ' +
'is {}'.format(type(job)))
# Wait for 0.1 second before checking, as jobs take a while to be
# queued sometimes
sleep(0.1)
for job in jobs:
logme.log('Checking {}'.format(job), 'debug')
qtype = job.qtype if isinstance(job, self._Job) else self.qtype
if isinstance(job, (self._Job, self._JobQueue, self.QueueJob)):
job = job.id
if qtype == 'local':
logme.log('Job is in local queue', 'debug')
try:
job = int(job)
except TypeError:
raise TypeError('Job must be a Job object or job #.')
if not local.JQUEUE \
or not local.JQUEUE.runner.is_alive():
raise ClusterError('Cannot wait on job ' + str(job) +
'JobQueue does not exist')
local.JQUEUE.wait(job)
else:
logme.log('Job is in remote queue', 'debug')
if isinstance(job, self._Job):
job = job.id
job = int(job)
not_found = 0
lgd = False
while True:
self._update()
# Allow 12 seconds to elapse before job is found in queue,
# if it is not in the queue by then, assume completion.
if job not in self.jobs:
if lgd:
logme.log('Attempt #{}/12'.format(not_found),
'debug')
else:
logme.log('{} not in queue, waiting up to 12s '
.format(job) +
'for it to appear', 'info')
lgd = True
sleep(1)
not_found += 1
if not_found == 12:
logme.log(
'{} not in queue, tried 12 times over 12s'
.format(job) + '. Job likely completed, ' +
'assuming completion, stats will be ' +
'unavailable.','warn'
)
return 'disappeared'
continue
## Actually look for job in running/queued queues
lgd = False
lgd2 = False
start = _dt.now()
res_time = conf.get_option('queue', 'res_time')
count = 0
# Get job state
job_state = self.jobs[job].state
# Check the state
if job_state in GOOD_STATES:
logme.log('Queue wait for {} complete'
.format(job), 'debug')
sleep(0.1)
break
elif job_state in ACTIVE_STATES:
if lgd:
logme.log('{} not complete yet, waiting'
.format(job), 'debug')
lgd = True
else:
logme.log('{} still not complete, waiting'
.format(job), 'verbose')
sleep(self.sleep_len)
elif job_state in BAD_STATES:
logme.log('Job {} failed with state {}'
.format(job, job_state), 'error')
return False
elif job_state in UNCERTAIN_STATES:
if not lgd2:
logme.log('Job {} in state {}, waiting {} '
.format(job, job_state, res_time) +
'seconds for resolution', 'warn')
lgd2 = True
if (_dt.now() - start).seconds > res_time:
logme.log('Job {} still in state {}, aborting'
.format(job, job_state), 'error')
return False
sleep(self.sleep_len)
else:
if count == 5:
logme.log('Job {} in unknown state {} '
.format(job, job_state) +
'cannot continue', 'critical')
raise QueueError('Unknown job state {}'
.format(job_state))
logme.log('Job {} in unknown state {} '
.format(job, job_state) +
'trying to resolve', 'debug')
count += 1
sleep(self.sleep_len)
# Sleep an extra half second to allow post-run scripts to run
sleep(0.5)
return True
[docs] def wait_to_submit(self, max_jobs=None):
"""Block until fewer running/queued jobs in queue than max_jobs.
Args:
max_jobs (int): Override self.max_jobs
"""
count = 50
written = False
if max_jobs:
self.max_jobs = int(max_jobs)
while True:
if self.can_submit:
return
if not written:
logme.log(('The queue is full, there are {} jobs running and '
'{} jobs queued. Will wait to submit, retrying '
'every {} seconds.')
.format(len(self.running), len(self.queued),
self.sleep_len),
'info')
written = True
if count == 0:
logme.log('Still waiting to submit.', 'info')
count = 50
count -= 1
sleep(self.sleep_len)
[docs] def update(self):
"""Refresh the list of jobs from the server, limit queries."""
if int(time()) - self.last_update > self.queue_update_time:
self._update()
else:
logme.log('Skipping update as last update too recent', 'debug')
return self
[docs] def get_jobs(self, key):
"""Return a dict of jobs where state matches key."""
retjobs = {}
for jobid, job in self.jobs.items():
if job.state == key.lower():
retjobs[jobid] = job
return retjobs
def get_user_jobs(self, users):
"""Filter jobs by user.
Args:
users (list): A list of users/owners
Returns:
dict: A filtered job dictionary of {job_id: QueueJob} for all jobs
owned by the queried users.
"""
try:
if isinstance(users, (str, int)):
users = [users]
else:
users = list(users)
except TypeError:
users = [users]
return {k: v for k, v in self.jobs.items() if v.owner in users}
@property
def users(self):
"""Return a list of users with jobs running."""
return [job.owner for job in self.jobs.values()]
@property
def job_states(self):
"""Return a list of job states for all jobs in the queue."""
return [job.state for job in self.jobs]
@property
def finished(self):
"""Return a list of jobs that are neither queued nor running."""
return {i: j for i, j in self.jobs.items() \
if j.state not in ACTIVE_STATES}
@property
def bad(self):
"""Return a list of jobs that have bad or uncertain states."""
return {i: j for i, j in self.jobs.items() \
if j.state in BAD_STATES or j.state in UNCERTAIN_STATES}
@property
def active_job_count(self):
"""Return a count of all queued or running jobs."""
return sum([
len(j) for j in [
self.get_jobs(i) for i in ACTIVE_STATES
]
])
@property
def can_submit(self):
"""Return True if R/Q jobs are less than max_jobs.
If max_jobs is None, default from config is used.
"""
self.update()
return self.active_job_count < self.max_jobs
######################
# Internal Functions #
######################
def _update(self):
"""Refresh the list of jobs from the server.
This is the core queue interaction function of this class.
"""
if self._updating:
return
logme.log('Queue updating', 'debug')
# Set the update time I don't care about microseconds
self.last_update = int(time())
jobs = [] # list of jobs created this session
# Mode specific initialization
if self.qtype == 'local':
if not local.JQUEUE or not local.JQUEUE.runner.is_alive():
local.JQUEUE = local.JobQueue(cores=local.THREADS)
for job_id, job_info in local.JQUEUE:
if job_id in self.jobs:
job = self.jobs[job_id]
else:
job = self.QueueJob()
job.id = job_id
job.name = job_info.function.__name__
job.owner = self.user
self.nodes = socket.gethostname()
if job_info.state == 'Not Submitted':
job.state = 'pending'
elif job_info.state == 'waiting' \
or job_info.state == 'submitted':
job.state = 'pending'
elif job_info.state == 'started' \
or job_info.state == 'running':
job.state = 'running'
elif job_info.state == 'done':
job.state = 'completed'
job.exitcode = int(job_info.exitcode)
elif job_info.state == 'queued':
job.state = 'pending'
else:
raise Exception('Unrecognized state')
# Assign the job to self.
self.jobs[job_id] = job
jobs.append(job_id)
else:
for [job_id, job_name, job_user, job_partition,
job_state, job_nodelist, job_nodecount,
job_cpus, job_exitcode] in queue_parser(self.qtype,
self.user,
self.partition):
if job_id not in self.jobs:
job = self.QueueJob()
else:
job = self.jobs[job_id]
jobs.append(job_id)
job.id = job_id
job.name = job_name
job.owner = job_user
job.queue = job_partition
job.state = job_state.lower()
job.nodes = job_nodelist
# Threads is number of nodes * jobs per node
if job_nodecount and job_cpus:
job.threads = int(job_nodecount) * int(job_cpus)
else:
job.threads = None
if job.state == 'completed' or job.state == 'failed':
job.exitcode = job_exitcode
# Assign the job to self.
self.jobs[job_id] = job
jobs.append(job_id)
# We assume that if a job just disappeared it completed
if self.jobs:
for qjob in self.jobs.values():
if qjob.id not in jobs:
qjob.state = 'completed'
qjob.disappeared = True
def __getattr__(self, key):
"""Make running and queued attributes dynamic."""
key = key.lower()
if key in TORQUE_SLURM_STATES:
key = TORQUE_SLURM_STATES[key]
if key == 'complete':
key = 'completed'
elif key == 'queued':
key = 'pending'
if key in ALL_STATES:
return self.get_jobs(key)
def __getitem__(self, key):
"""Allow direct accessing of jobs by job id."""
if isinstance(key, self._Job):
key = key.jobid
key = int(key)
try:
return self.jobs[key]
except KeyError:
return None
def __iter__(self):
"""Allow us to be iterable."""
for jb in self.jobs.values():
yield jb
def __len__(self):
"""Length is the total job count."""
return len(self.jobs)
def __repr__(self):
"""For debugging."""
self._updating = True
if self.user:
outstr = 'Queue<jobs:{};completed:{};queued:{};user={}>'.format(
len(self), len(self.complete), len(self.queued), self.user)
else:
outstr = 'Queue<jobs:{};completed:{};queued:{};user=ALL>'.format(
len(self), len(self.complete), len(self.queued))
self._updating = False
return outstr
def __str__(self):
"""A list of keys."""
return str(self.jobs.keys())
##############################################
# A simple class to hold jobs in the queue #
##############################################
class QueueJob(object):
"""A very simple class to store info about jobs in the queue.
Only used for torque and slurm queues.
Attributes:
id (int): Job ID
name (str): Job name
owner (str): User who owns the job
threads (int): Number of cores used by the job
queue (str): The queue/partition the job is running in
state (str): Current state of the job, normalized to slurm
states
nodes (list): List of nodes job is running on
exitcode (int): Exit code of completed job
disappeared (bool): Job cannot be found in the queue anymore
"""
id = None
name = None
owner = None
threads = None
queue = None
state = None
nodes = None
exitcode = None
disappeared = False
def __init__(self):
"""No initialization needed all attributes are set elsewhere."""
pass
def __repr__(self):
"""Show all info."""
outstr = ("Queue.QueueJob<{id}:{state}({name},owner:{owner}," +
"queue:{queue},nodes:{nodes},threads:{threads}," +
"exitcode:{code})").format(
id=self.id, name=self.name, owner=self.owner,
queue=self.queue, nodes=self.nodes,
code=self.exitcode, threads=self.threads,
state=self.state)
if self.disappeared:
outstr += 'DISAPPEARED>'
else:
outstr += '>'
return outstr
def __str__(self):
"""Print job ID."""
return str(self.id)
################
# Exceptions #
################
[docs]class QueueError(Exception):
"""Simple Exception wrapper."""
pass
###############################################################################
# Non-Class Functions #
###############################################################################
###################
# Queue Parsers #
###################
[docs]def queue_parser(qtype=None, user=None, partition=None):
"""Call either torque or slurm qtype parsers depending on qtype.
Args:
qtype: Either 'torque' or 'slurm', defaults to current MODE
user: optional user name to pass to queue to filter queue with
Yields:
tuple: job_id, name, userid, partition, state, nodelist, numnodes,
ntpernode, exit_code
"""
if not qtype:
qtype = get_cluster_environment()
if qtype == 'torque':
return torque_queue_parser(user, partition)
elif qtype == 'slurm':
return slurm_queue_parser(user, partition)
else:
raise ClusterError("Invalid qtype type {}, must be 'torque' or 'slurm'"
.format(qtype))
[docs]def torque_queue_parser(user=None, partition=None):
"""Iterator for torque queues.
Use the `qstat -x` command to get an XML queue for compatibility.
Args:
user: optional user name to pass to qstat to filter queue with
partiton: optional partition to filter the queue with
Yields:
tuple: job_id, name, userid, partition, state, nodelist, numnodes,
ntpernode, exit_code
numcpus is currently always 1 as most torque queues treat every core as a
node.
"""
# I am not using run.cmd because I want to catch XML errors also
try_count = 0
qargs = ['qstat', '-x']
while True:
try:
xmlqueue = ET.fromstring(check_output(qargs))
except CalledProcessError:
sleep(1)
if try_count == 5:
raise
else:
try_count += 1
except ET.ParseError:
# ElementTree throws error when string is empty
sleep(1)
if try_count == 1:
xmlqueue = None
break
else:
try_count += 1
else:
break
# Create QueueJob objects for all entries that match user
if xmlqueue is not None:
for xmljob in xmlqueue:
job_id = xmljob.find('Job_Id').text.split('.')[0]
if '[' in job_id:
job_id, array_id = job_id.split('[')
array_id = array_id.strip('[]')
if array_id:
array_id = int(array_id)
else:
array_id = 0
else:
array_id = None
try:
job_id = int(job_id)
except ValueError:
# Allow string job IDs
pass
job_owner = xmljob.find('Job_Owner').text.split('@')[0]
if user and job_owner != user:
continue
job_name = xmljob.find('Job_Name').text
job_queue = xmljob.find('queue').text
job_state = xmljob.find('job_state').text
job_state = TORQUE_SLURM_STATES[job_state]
logme.log('Job {} state: {}'.format(job_id, job_state),
'debug')
ndsx = xmljob.find('exec_host')
if ndsx:
nds = ndsx.text.split('+')
else:
nds = []
nodes = []
for node in nds:
if '-' in node:
nm, num = node.split('/')
for i in range(*[int(i) for i in num.split('-')]):
nodes.append(nm + '/' + str(i).zfill(2))
else:
nodes.append(node)
# I assume that every 'node' is a core, as that is the
# default for torque, but it isn't always true
job_threads = len(nodes)
exitcode = xmljob.find('exit_status')
if hasattr(exitcode, 'text'):
exitcode = int(exitcode.text)
if partition and job_queue != partition:
continue
yield (job_id, job_name, job_owner, job_queue, job_state,
nodes, job_threads, 1, exitcode)
[docs]def slurm_queue_parser(user=None, partition=None):
"""Iterator for slurm queues.
Use the `squeue -O` command to get standard data across implementation,
supplement this data with the results of `sacct`. sacct returns data only
for the current user but retains a much longer job history. Only jobs not
returned by squeue are added with sacct, and they are added to *the end* of
the returned queue, i.e. *out of order with respect to the actual queue*.
Args:
user: optional user name to filter queue with
partition: optional partition to filter queue with
Yields:
tuple: job_id, name, userid, partition, state, nodelist, numnodes,
ntpernode, exit_code
"""
nodequery = re.compile(r'([^\[,]+)(\[[^\[]+\])?')
qargs = ['squeue', '-h', '-O',
'jobid:400,arraytaskid:400,name:400,userid:400,partition:400,' +
'state:400,nodelist:400,numnodes:400,numcpus:400,exit_code:400']
# Parse queue info by length
squeue = [
tuple(
[k[i:i+200].rstrip() for i in range(0, 4000, 400)]
) for k in run.cmd(qargs)[1].split('\n')
]
# SLURM sometimes clears the queue extremely fast, so we use sacct
# to get old jobs by the current user
qargs = ['sacct', '-p',
'--format=jobid,jobname,user,partition,state,' +
'nodelist,reqnodes,ncpus,exitcode']
try:
sacct = [tuple(i.strip(' |').split('|')) for i in
run.cmd(qargs)[1].split('\n')]
sacct = sacct[1:]
# This command isn't super stable and we don't care that much, so I will
# just let it die no matter what
except Exception as e:
if logme.MIN_LEVEL == 'debug':
raise e
else:
sacct = []
if sacct:
if len(sacct[0]) != 9:
logme.log('sacct parsing failed unexpectedly as there are not ' +
'9 columns, aborting.', 'critical')
raise ValueError('sacct output does not have 9 columns. Has:' +
'{}: {}'.format(len(sacct[0]), sacct[0]))
jobids = [i[0] for i in squeue]
for sinfo in sacct:
# Skip job steps, only index whole jobs
if '.' in sinfo[0]:
logme.log('Skipping {} '.format(sinfo[0]) +
"in sacct processing as it is a job part.",
'verbose')
continue
# These are the values I expect
try:
[sid, sname, suser, spartition, sstate,
snodelist, snodes, scpus, scode] = sinfo
if '_' in sid:
sid, sarr = sid.split('_')
sif = '{}_{}'.format(sid, sarr)
else:
sarr = None
sif = '{}'.format(sid)
except ValueError as err:
logme.log('sacct parsing failed with error {} '.format(err) +
'due to an incorrect number of entries.\n' +
'Contents of sinfo:\n{}\n'.format(sinfo) +
'Expected 9 values\n:' +
'[sid, sname, suser, spartition, sstate, ' +
'snodelist, snodes, scpus, scode]',
'critical')
raise
# Skip jobs that were already in squeue
if sid in jobids:
logme.log('{} still in squeue output'.format(sid),
'verbose')
continue
scode = int(scode.split(':')[-1])
squeue.append((sid, sarr, sname, suser, spartition, sstate,
snodelist, snodes, scpus, scode))
else:
logme.log('No job info in sacct', 'debug')
# Sanitize data
for sinfo in squeue:
if len(sinfo) == 10:
[sid, sarr, sname, suser, spartition, sstate, sndlst,
snodes, scpus, scode] = sinfo
else:
sys.stderr.write('{}'.format(repr(sinfo)))
raise ClusterError('Queue parsing error, expected 10 items '
'in output of squeue and sacct, got {}\n'
.format(len(sinfo)))
if partition and spartition != partition:
continue
if not isinstance(sid, int):
sid = int(sid) if sid else None
if isinstance(sarr, str) and sarr.isdigit():
sarr = int(sarr)
else:
sarr = None
if not isinstance(snodes, int):
snodes = int(snodes) if snodes else None
if not isinstance(scpus, int):
scpus = int(scpus) if snodes else None
if not isinstance(scode, int):
scode = int(scode) if scode else None
# Convert user from ID to name
if suser.isdigit():
suser = pwd.getpwuid(int(suser)).pw_name
if user and suser != user:
continue
# Attempt to parse nodelist
snodelist = []
if sndlst:
if nodequery.search(sndlst):
nsplit = nodequery.findall(sndlst)
for nrg in nsplit:
node, rge = nrg
if not rge:
snodelist.append(node)
else:
for reg in rge.strip('[]').split(','):
# Node range
if '-' in reg:
start, end = [int(i) for i in reg.split('-')]
for i in range(start, end):
snodelist.append('{}{}'.format(node, i))
else:
snodelist.append('{}{}'.format(node, reg))
else:
snodelist = sndlst.split(',')
yield (sid, sname, suser, spartition, sstate, snodelist,
snodes, scpus, scode)
###########################################################
# Set the global cluster type: slurm, torque, or local #
###########################################################
[docs]def get_cluster_environment():
"""Detect the local cluster environment and set MODE globally.
Uses which to search for sbatch first, then qsub. If neither is found,
MODE is set to local.
Returns:
tuple: MODE variable ('torque', 'slurm', or 'local')
"""
global MODE
conf_queue = conf.get_option('queue', 'queue_type', 'auto')
if conf_queue not in ['torque', 'slurm', 'local', 'auto']:
logme.log('queue_type in the config file is {}, '.format(conf_queue) +
'but it should be one of torque, slurm, local, or auto. ' +
'Resetting it to auto', 'warn')
conf.set_option('queue', 'queue_type', 'auto')
conf_queue = 'auto'
if conf_queue == 'auto':
sbatch_cmnd = conf.get_option('queue', 'sbatch')
qsub_cmnd = conf.get_option('queue', 'qsub')
sbatch_cmnd = sbatch_cmnd if sbatch_cmnd else 'sbatch'
qsub_cmnd = qsub_cmnd if qsub_cmnd else 'qsub'
if run.which(sbatch_cmnd):
MODE = 'slurm'
elif run.which(qsub_cmnd):
MODE = 'torque'
else:
MODE = 'local'
else:
MODE = conf_queue
if MODE == 'slurm' or MODE == 'torque':
logme.log('{} detected, using for cluster submissions'.format(MODE),
'debug')
else:
logme.log('No cluster environment detected, using multiprocessing',
'debug')
return MODE
##############################
# Check if queue is usable #
##############################
[docs]def check_queue(qtype=None):
"""Raise exception if MODE is incorrect."""
if 'MODE' not in globals():
global MODE
MODE = get_cluster_environment()
if not MODE:
MODE = get_cluster_environment()
if qtype:
if qtype not in ALLOWED_MODES:
raise ClusterError('qtype value {} is not recognized, '
.format(qtype) +
'should be: local, torque, or slurm')
else:
if MODE not in ALLOWED_MODES:
MODE = qtype
return True
elif MODE not in ALLOWED_MODES:
raise ClusterError('MODE value {} is not recognized, '.format(MODE) +
'should be: local, torque, or slurm')
######################################################################
# Expose the Queue waiting method without requiring a Queue object #
######################################################################
def wait(jobs):
"""Wait for jobs to finish.
Args:
jobs: A single job or list of jobs to wait for. With torque or slurm,
these should be job IDs, with local mode, these are
multiprocessing job objects (returned by submit())
"""
# Support python2, which hates reciprocal import for 80's reasons
from .job import Job
from .local import JobQueue
check_queue() # Make sure the MODE is usable
# Sanitize argument
if not isinstance(jobs, (list, tuple)):
jobs = [jobs]
for job in jobs:
if not isinstance(job, (str, int, Job, JobQueue)):
raise ClusterError('job must be int, string, or Job, ' +
'is {}'.format(type(job)))
if MODE == 'local':
for job in jobs:
try:
job = int(job)
except TypeError:
raise TypeError('Job must be a Job object or job #.')
if not local.JQUEUE or not local.JQUEUE.runner.is_alive():
raise ClusterError('Cannot wait on job ' + str(job) +
'JobQueue does not exist')
local.JQUEUE.wait(job)
elif MODE == 'torque':
# Wait for 1 seconds before checking, as jobs take a while to be queued
# sometimes
sleep(1)
s = re.compile(r' +') # For splitting qstat output
# Jobs must be strings for comparison operations
jobs = [str(j) for j in jobs]
q = run.cmd('qstat -a', tries=8)[1].rstrip().split('\n')
# Check header
if not re.split(r' {2,100}', q[3])[9] == 'S':
raise ClusterError('Unrecognized torque qstat format')
# Build a list of completed jobs
complete = []
for j in q[5:]:
i = s.split(j)
if i[9] == 'C':
complete.append(i[0].split('.')[0])
# Build a list of all jobs
alljobs = [s.split(j)[0].split('.')[0] for j in q[5:]]
# Trim down job list
jobs = [j for j in jobs if j in alljobs]
jobs = [j for j in jobs if j not in complete]
if len(jobs) == 0:
return
elif MODE == 'slurm':
# Wait for 2 seconds before checking, as jobs take a while to be queued
# sometimes
sleep(2)
# Jobs must be strings for comparison operations
jobs = [str(j) for j in jobs]
while True:
# Slurm allows us to get a custom output for faster parsing
q = check_output(
['squeue', '-h', '-o', "'%A,%t'"]).decode().rstrip().split(',')
# Build a list of jobs
complete = [i[0] for i in q if i[1] == 'CD']
failed = [i[0] for i in q if i[1] == 'F']
allj = [i[0] for i in q]
# Trim down job list, ignore failures
jobs = [i for i in jobs if i not in allj]
jobs = [i for i in jobs if i not in complete]
jobs = [i for i in jobs if i not in failed]
if len(jobs) == 0:
return
sleep(2)