| # Copyright lowRISC contributors. |
| # Licensed under the Apache License, Version 2.0, see LICENSE for details. |
| # SPDX-License-Identifier: Apache-2.0 |
| # ---------------------------------- |
| # SGE.py |
| # _JobData Class |
| # ---------------------------------- |
| import logging |
| import pwd |
| import os |
| import re |
| import subprocess |
| import time |
| |
| # |
| from qsubopts import qsubOptions |
| sgeJob_t = qsubOptions() |
| |
| |
| class _JobData: |
| """ |
| Internal helper class to manage job data from qstat |
| """ |
| def __init__(self, qstat_job_line): |
| # The format of the line goes like |
| # job-ID prior name user state submit/start at queue slots ja-task-ID |
| |
| tokens = qstat_job_line.split() |
| |
| assert len(tokens) >= 8, 'Not a valid qstat line: ' + qstat_job_line |
| # Job line must have at least 8 tokens |
| try: |
| self.id = int(tokens[0]) |
| except ValueError: |
| error_msg = "Could not convert data to an integer." |
| raise ValueError(error_msg) |
| self.priority = float(tokens[1]) |
| self.name = tokens[2] |
| self.user = tokens[3] |
| self.state = tokens[4] |
| self.time = ' '.join(tokens[5:7]) |
| |
| if '@' in qstat_job_line: |
| # Has queue assigned, e.g. core2.q@q2.caspian.mit.edu |
| self.queue = tokens[7] |
| self.slots = tokens[8] |
| if len(tokens) == 9: |
| self.ja_task_id = None |
| elif len(tokens) == 10: |
| self.ja_task_id = tokens[-1] |
| else: |
| raise ValueError(f"Could not parse line: {qstat_job_line}") |
| else: |
| # No queue assigned |
| self.slots = tokens[7] |
| if len(tokens) == 8: |
| self.ja_task_id = None |
| elif len(tokens) == 9: |
| self.ja_task_id = tokens[-1] |
| else: |
| raise ValueError(f"Could not parse line: {qstat_job_line}") |
| |
| # Convert array indices ja_task_id into python list format |
| ja_task_id = [] |
| if self.ja_task_id is not None: |
| for blob in self.ja_task_id.split(','): |
| # Parse data of form '193-349:1' or '1934' |
| x = blob.split(':') |
| if len(x) == 1: |
| ja_task_id += x |
| else: |
| subjobs, step = x |
| begin, end = subjobs.split('-') |
| ja_task_id += range(int(begin), int(end) + 1, int(step)) |
| |
| self._ja_tasklist = ja_task_id |
| |
| def __repr__(self): |
| repr = ['{'] |
| for key, value in self.__dict__.items(): |
| if key[0] != '_': |
| repr.append(key + '=' + str(value)) |
| repr.append('}') |
| return '\n'.join(repr) |
| |
| |
| class JobList: |
| """ |
| Internal helper class to manage job lists |
| """ |
| |
| def __init__(self, qstat_output=None): |
| self._joblist = [] |
| for line in qstat_output.split('\n')[2:-1]: |
| self._joblist.append(_JobData(line)) |
| |
| def __iter__(self): |
| for job in self._joblist: |
| yield job |
| |
| def __repr__(self): |
| return '\n'.join([str(job) for job in self._joblist]) |
| |
| |
| class SGE: |
| """External system call handler for Sun Grid Engine environment.""" |
| |
| def __init__(self, q=None, path='', ): |
| |
| logger = logging.getLogger('SGE.__init__') |
| |
| if q is None: |
| # No queue specified. By default, submit to all available queues. |
| self.cmd_qconf = os.path.join(path, 'qconf') |
| |
| try: |
| qliststr = _exec(self.cmd_qconf + ' -sql') |
| except IOError: |
| error_msg = 'Error querying queue configuration' |
| logger.error(error_msg) |
| raise IOError(error_msg) |
| |
| self.q = qliststr.replace('\n', ',')[:-1] |
| |
| logger.info("""Sun Grid Engine handler initialized |
| Queues detected: %s""", self.q) |
| |
| else: |
| self.q = q |
| |
| self.cmd_qsub = os.path.join(path, 'qsub') |
| self.cmd_qstat = os.path.join(path, 'qstat') |
| |
| def wait(self, jobid, interval=10, name=None, pbar=None, |
| pbar_mode=None): |
| """Waits for job running on SGE Grid Engine environment to finish. |
| |
| If you are just waiting for one job, this becomes a dumb substitute for |
| the -sync y option which can be specified to qsub. |
| |
| Inputs: |
| |
| jobid |
| interval - Polling interval of SGE queue, in seconds. (Default: 10) |
| """ |
| |
| logger = logging.getLogger('SGE.wait') |
| |
| dowait = True |
| while dowait: |
| p = subprocess.Popen(self.cmd_qstat, shell=True, |
| stdout=subprocess.PIPE) |
| pout, _ = p.communicate() |
| |
| if pbar is not None: |
| logger.error('Progress bar handling not implemented') |
| |
| dowait = False |
| for line in pout.split('\n'): |
| t = line.split() |
| if len(t) >= 5 and t[0] == str(jobid): |
| # Find a line with useful info |
| if re.search(t[4], 'qwrt'): |
| # Job must be queued, running or being transferred |
| dowait = True |
| break |
| if re.search(t[4], 'acuE'): |
| # Job or host in error state |
| logger.warning('Job %d in error state', str(jobid)) |
| |
| if dowait: |
| time.sleep(interval) |
| if name is None: |
| logger.info("Time %s: waiting for jobid %s to finish", |
| time.ctime(), str(jobid)) |
| else: |
| logger.info("Time %s: waiting for job '%s' (jobid %s) to \ |
| finish", time.ctime(), name, str(jobid)) |
| |
| def submit(self, job, array=False, useenvironment=True, usecwd=True, |
| name=None, stdin=None, stdout=None, stderr=None, |
| joinstdouterr=True, nproc=1, wait=True, lammpi=True): |
| """ |
| Submits a job to SGE |
| Returns jobid as a number |
| """ |
| |
| logger = logging.getLogger('SGE.submit') |
| |
| logger.info("Submitting job: " + str(job) + " stdout: %s \ |
| Stderr: %s", stdout, stderr) |
| |
| # Parameters to qsub specified as the header of the job specified on |
| # STDIN |
| lamstring = lammpi and f" -pe lammpi {nproc}" or "" |
| qsuboptslist = ['-cwd -V ', lamstring] |
| |
| if name is not None: |
| qsuboptslist.append('-N ' + name) |
| if stdin is not None: |
| qsuboptslist.append('-i ' + stdin) |
| if stdout is not None: |
| qsuboptslist.append('-o ' + stdout) |
| if stderr is not None: |
| qsuboptslist.append('-e ' + stderr) |
| if joinstdouterr: |
| qsuboptslist.append('-j') |
| if wait: |
| qsuboptslist.append('-sync y') |
| if usecwd: |
| qsuboptslist.append('-cwd') |
| if useenvironment: |
| qsuboptslist.append('-V') |
| if array is not False: |
| try: |
| n = int(array[0]) |
| except IndexError: |
| n = int(array) |
| raise IndexError("List is empty!") |
| except ValueError: |
| error_msg = "array[0] being an out of bounds access." |
| logger.error(error_msg) |
| raise ValueError(error_msg) |
| try: |
| m = int(array[1]) |
| except ValueError: |
| m = None |
| raise ValueError("Could not convert data to an integer.") |
| except IndexError: |
| m = None |
| raise IndexError("array[1] being an out of bounds access.") |
| try: |
| s = int(array[2]) |
| except IndexError: |
| s = None |
| raise IndexError("array[2] being an out of bounds access.") |
| except ValueError: |
| s = None |
| raise ValueError("Could not convert data to an integer.") |
| if m == s is None: |
| qsuboptslist.append('-t %d' % n) |
| elif s is None: |
| qsuboptslist.append('-t %d-%d' % (n, m)) |
| else: |
| qsuboptslist.append('-t %d-%d:%d' % (n, m, s)) |
| |
| qsubopts = ' '.join(qsuboptslist) |
| |
| pout = _exec(self.cmd_qsub, stdin=qsubopts + '\n' + job, |
| print_command=False) |
| |
| try: |
| # Next to last line should be |
| # "Your job 1389 (name) has been submitted" |
| # parse for job id |
| jobid = int(pout.split('\n')[-2].split()[2]) |
| return jobid |
| # except (ValueErrorValueError, IndexError, AttributeError) (e): |
| except (ValueError, IndexError, AttributeError): |
| error_msg = """Error submitting SGE job: |
| %s |
| %s |
| |
| Output was: |
| %s""" % (qsubopts, job, pout) |
| logger.error(error_msg) |
| raise IOError(error_msg) |
| |
| def getuserjobs(self, user=pwd.getpwuid(os.getuid())[0]): |
| """Returns a list of SGE jobids run by a specific user |
| |
| Inputs |
| user - SGE user to poll (Default = '', i.e. current user) |
| qstat - path to qstat binary (Default = 'qstat') |
| """ |
| |
| p = subprocess.Popen(self.cmd_qstat + " -u " + user, shell=True, |
| stdout=subprocess.PIPE) |
| qstat_output, _ = p.communicate() |
| joblist = JobList(qstat_output) |
| return [job for job in joblist if job.user == user] |
| |
| def run_job(self, command, name='default', logfnm='default.log', |
| wait=True): |
| """Run job on SGE with piped logging.""" |
| jobid = self.submit(command, name=name, stdout=logfnm, |
| stderr=logfnm, wait=wait) |
| return jobid |
| |
| def get_queue_instance_status(self): |
| """ |
| Get loads for each queue instance |
| """ |
| output = _exec(' '.join([self.cmd_qstat, '-f'])) |
| |
| data = [] |
| for line in output.split('\n')[1:]: |
| t = line.split() |
| if len(t) != 5: |
| continue |
| |
| nodename = t[0].split('@')[1].split('.')[0] |
| maxslots = int(t[2].split('/')[2]) |
| load = float(t[3]) |
| |
| data.append({'name': nodename, 'maxslots': maxslots, 'load': load}) |
| |
| return data |
| |
| |
| def _exec(command, print_to_screen=False, logfnm=None, stdin='', |
| print_command=False): |
| """ |
| Runs command line using subprocess, optionally returning stdout |
| """ |
| def _call_cmd(command, stdin=''): |
| p = subprocess.Popen(command, shell=True, stdin=subprocess.PIPE, |
| stdout=subprocess.PIPE, |
| stderr=subprocess.STDOUT) |
| output, _ = p.communicate(stdin) |
| return output |
| |
| logger = logging.getLogger('_exec') |
| |
| if print_command: |
| logger.info("Executing process: \x1b[1;92m%-50s\x1b[0m Logfile: %s", |
| command, logfnm) |
| |
| output = "" |
| if logfnm is not None: |
| try: |
| with open(logfnm, 'a') as f: |
| if print_command: |
| print(f, "Executing process: %s" % command) |
| output = _call_cmd(command, stdin) |
| f.write(output) |
| except IOError: |
| error_msg = 'Error: File: ' + str(logfnm) + ' does not appear to exist.' |
| logger.error(error_msg) |
| raise IOError(error_msg) |
| else: |
| output = _call_cmd(command, stdin) |
| |
| logger.info('Output of command is:\n%s', output) |
| |
| if print_to_screen: |
| print(output) |
| |
| return output |
