batch.py

Go to the documentation of this file.

"""!
@package batch
 
Defines a set of classes and a command-line interface for submitting batch jobs.
 
Supported systems include serial execution locally, a multiprocessing pool,
Slurm, LSF, and Auger.
"""
 
import os
import argparse
import subprocess
import sys
import logging
import signal
import multiprocessing
import psutil
from pathlib import Path
import socket
 
import xml.etree.ElementTree as ET
from xml.dom import minidom
from xml.sax.saxutils import unescape
from distutils.spawn import find_executable
 
from abc import ABC, abstractmethod
 
from hpsmc.job import Job, JobStore, JobScriptDatabase
 
logger = logging.getLogger("hpsmc.batch")
 
RUN_SCRIPT = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'job.py')
 
 
class Batch(ABC):
    """!
    Generic batch processing interface.
    """
 
    def __init__(self):
 
        parser = argparse.ArgumentParser(self.__class__.__name__,
                                         epilog='Available scripts: %s' % ', '.join(JobScriptDatabase().get_script_names()))
 
        parser.add_argument("-c", "--config-file", nargs='?', help="Config file", action='append')
        parser.add_argument("-l", "--log-dir", nargs='?', help="Log file output dir", required=False, default=str(Path(os.getcwd(), 'logs')))
        parser.add_argument("-d", "--run-dir", nargs='?', help="Base run dir for the jobs (must be an absolute path)", default=None)
        parser.add_argument("-D", "--debug", action='store_true', help="Enable debug settings", required=False)
        parser.add_argument("-o", "--check-output", action='store_true', required=False, help="Do not submit jobs where output files already exist")
        parser.add_argument("-s", "--job-steps", type=int, default=None, required=False)
        parser.add_argument("-r", "--job-range", nargs='?', help="Submit jobs numbers within range (e.g. '1:100')", required=False)
        parser.add_argument("script", nargs='?', help="Name of job script")
        parser.add_argument("jobstore", nargs='?', help="Job store in JSON format")
        parser.add_argument("jobids", nargs="*", type=int, help="List of individual job IDs to submit (optional)")
 
        self.parser = parser
 
    def parse_args(self, args):
        """! Parse command line arguments and perform setup."""
 
        cl = self.parser.parse_args(args)
 
        logger.debug(str(cl))
 
        if cl.script is None:
            raise Exception('The script is a required argument.')
        self.script_name = cl.script  # Name of script
        script_db = JobScriptDatabase()
        if not script_db.exists(self.script_name):
            raise Exception('The script name is not valid: %s' % self.script_name)
        self.script = script_db.get_script_path(self.script_name)  # Path to script
        if not os.path.isfile(self.script):
            raise Exception('The job script does not exist: %s' % self.script)
 
        if cl.jobstore is None:
            raise Exception('The job store file is a required argument.')
        if not os.path.isfile(cl.jobstore):
            raise Exception('The job store does not exist: %s' % cl.jobstore)
        self.jobstore = JobStore(cl.jobstore)
 
        self.debug = cl.debug
 
        # Set log dir which is for copying back log files generated by the batch system
        self.log_dir = os.path.abspath(cl.log_dir)
        logger.info('log dir: {}'.format(self.log_dir))
        if not os.path.exists(self.log_dir):
            os.makedirs(self.log_dir)
            logger.info('Created log dir: {}'.format(self.log_dir))
 
        # Set run dir which is a root directory under which job directories will be created
        self.run_dir = cl.run_dir
        if self.run_dir is not None:
            logger.info('run dir: {}'.format(self.run_dir))
            if not os.path.isabs(self.run_dir):
                # Require that the run dir is supplied as an abs path
                raise Exception("The run dir for batch processing must be an abs path.")
 
        self.check_output = cl.check_output
 
        if cl.jobids:
            self.job_ids = list(map(int, cl.jobids))
        else:
            self.job_ids = []
 
        self.job_steps = cl.job_steps
 
        if cl.job_range:
            toks = cl.job_range.split(':')
            if len(toks) != 2:
                raise ValueError('Bad format for job range: ' + cl.job_range)
            self.start_job_num = int(toks[0])
            self.end_job_num = int(toks[1])
            if self.start_job_num > self.end_job_num:
                raise ValueError("The start job number must be >= the end job num when using a range.")
            if self.start_job_num < 0 or self.end_job_num < 0:
                raise ValueError("The job range numbers must be > 0.")
        else:
            self.start_job_num = None
            self.end_job_num = None
 
        if cl.config_file:
            self.config_files = list(map(os.path.abspath, cl.config_file))
        else:
            self.config_files = []
 
        return cl
 
    @abstractmethod
    def submit_job(self, job_id):
        """!
        Submit a single batch job and return the batch ID.
 
        This is abstract as each batch system will do this differently.
 
        Some batch systems don't implement this but sub-classes should override this and make it a
        no-op so that they can be instantiated.
        """
        pass
 
    def submit(self):
        """!
        This is the generic batch submission function which gets a list of jobs to run based on command line
        arguments and submits them individually. It calls the abstract submit_job() method and prints the batch
        system ID that was returned, if any.
        """
        job_ids = self._get_filtered_job_ids()
        logger.info('Submitting jobs: %s' % str(job_ids))
        for job_id in job_ids:
            if not self.jobstore.has_job_id(job_id):
                raise Exception('Job ID was not found in job store: %s' % job_id)
            job_data = self.jobstore.get_job(job_id)
            batch_id = self.submit_job(job_id)
            logger.info(f"Submitted job {job_id} with batch ID {str(batch_id)}")
 
    def default_rundir(self, job_id=None):
        if job_id is None:
            raise Exception('Missing valid job ID')
        return str(Path(os.getcwd(), 'scratch', str(job_id)))
 
    def build_cmd(self, job_id):
        """!
        This is the basic implementation of building a command to run the job from a batch system.
        """
        cmd = [sys.executable, RUN_SCRIPT, 'run']
        logfile = self._logfile(job_id)
        cmd.extend(['-o', f"{logfile}.out",
                    '-e', f"{logfile}.err"])
        if self.run_dir:
            # Set the job's base run dir explicitly from user argument,
            # appending the job number as a subdirectory.
            job_dir = str(Path(self.run_dir, str(job_id)))
        else:
            # Set the job directory to the default.
            job_dir = self.default_rundir(job_id)
 
        logger.debug(f'job dir: {job_dir}')
        cmd.extend(['-d', job_dir])
 
        if len(self.config_files):
            for cfg in self.config_files:
                cmd.extend(['-c', cfg])
        if self.job_steps:
            cmd.extend(['--job-steps', str(self.job_steps)])
        cmd.extend(['-i', str(job_id)])
        cmd.append(self.script)
        cmd.append(os.path.abspath(self.jobstore.path))
        logger.debug("Job command: %s" % " ".join(cmd))
        return cmd
 
    def _logfile(self, job_id):
        """!
        Get the base name of a log file for the job.
        """
        return os.path.abspath(os.path.join(self.log_dir, 'job.%s' % str(job_id)))
 
    @staticmethod
    def _outputs_exist(job):
        """!
        Check if all output files exist for the given job. This is not the job ID but the full JSON job data.
 
        Return False when first missing output is found.
        """
        for src, dest in job["output_files"].items():
            if not os.path.isfile(os.path.join(job["output_dir"], dest)):
                logger.debug('Job output does not exist: %s -> %s' % (src, dest))
                return False
        return True
 
    def _get_filtered_job_ids(self):
        """!
        Get a list of job IDs to submit based on parsed command line options and whether output files are being checked.
        """
        submit_ids = self.jobstore.get_job_ids()
        logger.debug('Initial pre-filtered job IDs: {}'.format(str(submit_ids)))
        if self.start_job_num:
            submit_ids = [job_id for job_id in submit_ids
                          if int(job_id) >= self.start_job_num and int(job_id) <= self.end_job_num]
        elif len(self.job_ids):
            submit_ids = self.job_ids
        logger.debug('job IDs after range check: {}'.format(str(submit_ids)))
        if self.check_output:
            submit_ids = self._job_ids_missing_output(submit_ids)
            logger.info('job IDs after output file check: {}'.format(str(submit_ids)))
        return submit_ids
 
    def _job_ids_missing_output(self, job_ids):
        """! Get a list of IDs for jobs that are missing output files."""
        return [job_id for job_id in job_ids if not self._outputs_exist(self.jobstore.get_job(job_id))]
 
 
class BatchSystem(Batch, ABC):
    """!
    Represents a batch processing system that requires submission like Slurm or Auger.
 
    This subclasses Batch because it adds a number of different parameters which do not apply to all the
    batch system types (namely Pool and Local).
    """
 
    def __init__(self):
 
        super().__init__()
 
        self.parser.add_argument("-q", "--queue", nargs='?',
                                 help="Job queue or partition",
                                 required=False)
        self.parser.add_argument("-W", "--job-length", type=int, help="Max job length in hours", required=False, default=4)
        self.parser.add_argument("-m", "--memory", type=int, help="Max job memory allocation in MB", default=2000)
        self.parser.add_argument("-f", "--diskspace", type=int, help="Disk space needed for job in GB", default=20)
        self.parser.add_argument("-e", "--email", nargs='?', help="Email address for job notifications", required=False)
        
        self.parser.add_argument("-O", "--os", nargs='?', help="Operating system of batch nodes (Auger and LSF)")
 
        # Set site based on FQDN
        self.site = BatchSystem._site()
 
    def parse_args(self, args):
        """! Parse command line arguments and perform setup."""
 
        cl = super().parse_args(args)
 
        self.email = cl.email
        self.queue = cl.queue
        self.os = cl.os
        self.memory = cl.memory
        self.diskspace = cl.diskspace
        self.job_length = cl.job_length
 
        return cl
 
    @staticmethod
    def _site():
        fqdn = socket.getfqdn()
        site = None
        if 'slac.stanford.edu' in fqdn:
            site = 'slac'
        elif 'jlab.org' in fqdn:
            site = 'jlab'
        return site
 
 
class LSF(BatchSystem):
    """! Submit LSF batch jobs."""
 
    def __init__(self):
        super().__init__()
 
    def parse_args(self, args):
        super().parse_args(args)
        os.environ['LSB_JOB_REPORT_MAIL'] = 'Y' if self.email else 'N'
 
    def build_cmd(self, job_id):
 
        log_file = os.path.abspath(os.path.join(self.log_dir, 'job.%s.log' % str(job_id)))
 
        queue = self.queue
        if queue is None:
            queue = 'long'
 
        if self.os is not None:
            lsf_os = self.os
        else:
            lsf_os = 'centos7'
 
        cmd = ['bsub',
               '-W', str(self.job_lengthjob_length) + ':0',
               '-q', queue,
               '-R', lsf_os,
               '-o', log_file,
               '-e', log_file]
 
        if self.email:
            cmd.extend(['-u', self.email])
 
        cmd.extend(super().build_cmd(self, job_id))
 
        return cmd
 
    def submit_job(self, job_id):
        cmd = self.build_cmdbuild_cmd(job_id)
        logger.info('Submitting job %s to LSF with command: %s' % (job_id, ' '.join(cmd)))
        proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        out, err = proc.communicate()
        if err is not None and len(err):
            logger.warning(err)
        tokens = out.decode().split(" ")
        if tokens[0] != 'Job':
            raise Exception('Unexpected output from bsub command: %s' % out)
        batch_id = int(tokens[1].replace('<', '').replace('>', ''))
        return batch_id
 
 
class Slurm(BatchSystem):
    """! Submit Slurm batch jobs."""
 
    def __init__(self):
 
        super().__init__()
 
        self.parser.add_argument("-S", "--sh-dir", nargs='?', help="Directory to hold generated shell scripts for Slurm", default=str(Path(os.getcwd(), 'sh')))
        self.parser.add_argument("-E", "--env", nargs='?', help="Full path to env setup script", required=False, default=None)
        self.parser.add_argument("-A", "--account", nargs='?', help="Account name for s3df slurm jobs.", required=False, default=None)
 
    def parse_args(self, args):
 
        cl = super().parse_args(args)
 
        # Set Slurm env script
        self.env = cl.env
        self.account = cl.account
 
        # Set Slurm scripts dir
        self.sh_dir = os.path.abspath(cl.sh_dir)
        logger.info('Slurm sh dir: {}'.format(self.sh_dir))
        if not os.path.exists(self.sh_dir):
            os.makedirs(self.sh_dir)
            logger.info('Created Slurm sh dir: {}'.format(self.sh_dir))
 
    def default_rundir(self, job_id=None):
        """!
        Override the basic implementation for getting the default run directory.
        """
        if self.sitesite == 'slac':
            run_dir = '$LSCRATCH'
        elif self.sitesite == 'jlab':
            run_dir = '/scratch/slurm/$SLURM_JOBID'
        else:
            run_dir = os.getcwd() + "/scratch/$SLURM_JOBID"
        return run_dir
 
    def _default_queue(self):
        queue = self.queue
        if queue is None:
            if self.sitesite == 'slac':
                queue = 'shared'
            elif self.sitesite == 'jlab':
                queue = 'ifarm'
            else:
                raise Exception('No queue name was provided.')
        return queue
 
    def _sbatch(self, job_id):
        log_file = self._logfile(job_id)
        sbatch_cmd = ['sbatch',
                      '--time=%s' % (str(self.job_lengthjob_length) + ':00:00'),
                      '--mem=%sM' % self.memorymemory,
                      '--job-name=%i_%s' % (job_id, self.script_namescript_name),
                      '--output=%s.out' % log_file,
                      '--error=%s.err' % log_file]
        if self.queue:
            sbatch_cmd.extend([f'--partition={self.queue}'])
        if self.account:
            sbatch_cmd.extend([f'--account={self.account}'])
        if self.email:
            sbatch_cmd.extend([f'--mail-user={self.email}',
                               f'--mail-type=ALL'])
        return sbatch_cmd
 
    def _sh_filename(self, job_id):
        return self.sh_dir + '/job.%i.sh' % job_id
 
    def build_cmd(self, job_id):
        """!
        Wrap submission of Slurm jobs using a generated script.
        """
 
        # Get the sbatch command
        cmd = self._sbatch(job_id)
 
        # Get name of shell script to generate
        sh_filename = self._sh_filename(job_id)
 
        # Build the basic job command for execution
        job_cmd = super().build_cmd(job_id)
        if self.run_dir is None:
            # The superclass will have already set this if the user provided an
            # explicit run dir. Here we set a default scratch directory if none
            # was given.
            job_cmd.extend(['-d', self.default_rundirdefault_rundir()])
 
        # Write the job submission script out
        self._write_job_script(sh_filename, job_cmd)
 
        # Append job run script to Slurm command
        cmd.append(sh_filename)
 
        return cmd
 
    def _write_job_script(self, sh_filename, job_cmd):
        """!
        Write the shell script for Slurm job submission using the 'sbatch' command.
        """
 
        script_lines = ['#!/bin/bash',
                        '']
        if self.env:
            script_lines.append(f'source {self.env}')
        script_lines.extend(['echo Start time: `date`',
                             'echo PWD=`pwd`',
                             'echo ---- Start Environment ----',
                             'env | sort',
                             'echo ---- End Environment ----',
                             'time ' + ' '.join(job_cmd),
                             'echo End time: `date`'])
 
        logger.debug("Slurm submission script:\n" + str(script_lines))
 
        with open(sh_filename, 'w') as sh_file:
            for script_line in script_lines:
                sh_file.write(script_line + '\n')
 
        logger.debug('Wrote Slurm submission script to: '.format(str(Path(self.sh_dir, sh_filename))))
 
    def submit_job(self, job_id):
        cmd = self.build_cmdbuild_cmd(job_id)
        logger.info('Submitting job %s to Slurm with command: %s' % (job_id, ' '.join(cmd)))
        proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        out, err = proc.communicate()
        if err is not None and len(err):
            logger.warning(err)
        tokens = out.decode().split(" ")
        if tokens[0] != 'Submitted':
            raise Exception('Unexpected output from sbatch command: %s' % out)
        batch_id = int(tokens[3].replace('<', '').replace('>', ''))
        return batch_id
 
 
class Auger(BatchSystem):
    """!
    Submit Auger batch jobs.
 
    Auger itself is actually deprecated and unavailable but its submission XML format is supported by
    the Swif class (see below).
    """
 
    def __init__(self):
 
        super().__init__()
 
        self.setup_script = find_executable('hps-mc-env.csh')
 
        if not self.setup_script:
            raise Exception("Failed to find 'hps-mc-env.csh' in environment.")
 
    def submit_job(self, job_id):
        """!
        Make this a no-op. Auger is a bit of a special case in terms of how batch submission works with a
        generated XML file including all job IDs, so we do not implement single job submission.
        """
        pass
 
    def submit(self):
        """!
        Batch submission method for Auger.
 
        This differs from some of the other systems in that it doesn't loop over individual
        job IDs. Instead a single XML file is submitted for all the jobs at once.
        """
        xml_filename = self._create_job_xml()  # write request to XML file
        auger_ids = self._jsub(xml_filename)  # execute jsub to submit jobs
        logger.info("Submitted Auger jobs: %s" % str(auger_ids))
 
    def _create_job_xml(self):
        job_ids = self._get_filtered_job_ids()
        logger.info('Submitting jobs: %s' % str(job_ids))
        req = self._create_req(self.script_name)  # create request XML header
        for job_id in job_ids:
            if not self.jobstore.has_job_id(job_id):
                raise Exception('Job ID was not found in job store: %s' % job_id)
            job_params = self.jobstore.get_job(job_id)
            if self.check_output and Batch._outputs_exist(job_params):
                logger.warning("Skipping Auger submission for job "
                               "because outputs already exist: %d" % job_id)
            else:
                self._add_job(req, job_params)  # add job to request
        return self._write_req(req)    # write request to file
 
    def _jsub(self, xml_filename):
        cmd = ['jsub', '-xml', xml_filename]
        proc = subprocess.Popen(cmd, stdout=subprocess.PIPE)
        out, err = proc.communicate()
        auger_ids = self._get_auger_ids(out)
        return auger_ids
 
    def _get_auger_ids(self, out):
        auger_ids = []
        for line in out.splitlines():
            if line.strip().startswith(b'<jsub>'):
                j = ET.fromstring(line)
                for req in j.getchildren():
                    for child in req.getchildren():
                        if child.tag == 'jobIndex':
                            auger_id = int(child.text)
                            auger_ids.append(auger_id)
                        elif child.tag == 'error':
                            # Submission failed so raise an exception with the error msg
                            raise Exception(child.text)
                break
        return auger_ids
 
    def _write_req(self, req, filename='temp.xml'):
        pretty = unescape(minidom.parseString(ET.tostring(req)).toprettyxml(indent="  "))
        with open(filename, 'w') as f:
            f.write(pretty)
            return f.name
 
    def _create_req(self, req_name):
        req = ET.Element("Request")
        name_elem = ET.SubElement(req, "Name")
        name_elem.set("name", req_name)
        prj = ET.SubElement(req, "Project")
        prj.set("name", "hps")
        trk = ET.SubElement(req, "Track")
        if self.debug:
            # Queue arg is not used when debug flag is active.
            trk.set("name", "debug")
        else:
            # Queue name is used to set job track.
            queue = 'simulation'
            if self.queue is not None:
                queue = self.queue
            trk.set("name", queue)
        if self.emailemail:
            email = ET.SubElement(req, "Email")
            email.set("email", self.emailemail)
            email.set("request", "true")
            email.set("job", "true")
        mem = ET.SubElement(req, "Memory")
        mem.set("space", str(self.memorymemory))
        mem.set("unit", "MB")
        disk = ET.SubElement(req, "DiskSpace")
        disk.set("space", str(self.diskspacediskspace))
        disk.set("unit", "GB")
        limit = ET.SubElement(req, "TimeLimit")
        limit.set("time", str(self.job_lengthjob_length))
        limit.set("unit", "hours")
        os_elem = ET.SubElement(req, "OS")
        if self.os is not None:
            auger_os = self.os
        else:
            auger_os = 'el9'
        os_elem.set("name", auger_os)
        return req
 
    def build_cmd(self, job_id):
        cmd = [sys.executable, RUN_SCRIPT, 'run']
        if len(self.config_filesconfig_files):
            for cfg in self.config_filesconfig_files:
                cmd.extend(['-c', cfg])
        if self.job_stepsjob_steps is not None:
            cmd.extend(['--job-steps', str(self.job_stepsjob_steps)])
        cmd.extend(['-i', str(job_id)])
        cmd.append(self.scriptscript)
        cmd.append(os.path.abspath(self.jobstore.path))
        logger.debug("Job command: %s" % " ".join(cmd))
        return cmd
 
    def _create_job(self, params):
        """! Needed for resolving ptag output sources."""
        j = Job()
        j.script = self.scriptscript
        j._load_params(params)
        j._load_script()
        return j
 
    def _add_job(self, req, job_params):
        job = ET.SubElement(req, "Job")
        job_id = job_params['job_id']
        year = ''  # /todo change to number
        if 'year' in job_params.keys():
            year = job_params['year']
 
        if 'input_files' in list(job_params.keys()):
            inputfiles = job_params['input_files']
            for src, dest in inputfiles.items():
                if not src.startswith('http'):
                    input_elem = ET.SubElement(job, "Input")
                    input_elem.set("dest", dest)
                    if src.startswith("/mss"):
                        src_file = "mss:%s" % src
                    else:
                        src_file = src
                    input_elem.set("src", src_file)
                else:
                    logger.warning("http input file will not be included in XML job descr: {}".format(src))
        outputfiles = job_params["output_files"]
        outputdir = job_params["output_dir"]
        # outputdir = os.path.realpath(outputdir)
        j = self._create_job(job_params)
        for src, dest in outputfiles.items():
            output_elem = ET.SubElement(job, "Output")
            res_src = j.resolve_output_src(src)
            output_elem.set("src", res_src)
            dest_file = os.path.abspath(os.path.join(outputdir, dest))
            if dest_file.startswith("/mss"):
                dest_file = "mss:%s" % dest_file
            logger.debug('Auger dest file: {} -> {}'.format(src, dest))
            output_elem.set("dest", dest_file)
 
        job_name = ET.SubElement(job, "Name")
        job_name.set("name", '%ihps%i' % (year, job_id))
 
        job_err = ET.SubElement(job, "Stderr")
        stdout_file = os.path.abspath(os.path.join(self.log_dir, "job.%d.out" % job_id))
        stderr_file = os.path.abspath(os.path.join(self.log_dir, "job.%d.err" % job_id))
        job_err.set("dest", stderr_file)
        job_out = ET.SubElement(job, "Stdout")
        job_out.set("dest", stdout_file)
 
        cmd = ET.SubElement(job, "Command")
        cmd_lines = []
        cmd_lines.append("<![CDATA[")
 
        cmd_lines.append('pwd;\n')
        cmd_lines.append('env | sort;\n')
        cmd_lines.append('ls -lart;\n')
        cmd_lines.append("source %s;\n" % os.path.realpath(self.setup_script))
        cmd_lines.append("source %s/bin/jlab-env.csh;\n" % os.getenv('HPSMC_DIR'))
 
        job_cmd = self.build_cmdbuild_cmd(job_id)
 
        # Write log file locally so it can be copied back with Output element
        # log_file = 'job.%d.log' % job_id
        # job_cmd.extend(['-l', '$PWD/%s' % log_file])
        # log_out_elem = ET.SubElement(job, "Output")
        # log_out_elem.set('src', log_file)
        # log_out_elem.set('dest', os.path.join(self.log_dir, log_file))
 
        cmd_lines.extend(job_cmd)
        cmd_lines.append(';\n')
 
        cmd_lines.append('ls -lart; \n')
 
        cmd_lines.append("]]>")
 
        # logger.debug(cmd_lines)
 
        cmd.text = ' '.join(cmd_lines)
 
 
class Swif(Auger):
    """!
    Submit jobs to the 'swif2' workflow system at JLAB.
 
    The legacy 'swif2 add-jsub' command (which consumed an Auger XML request file) has been deprecated by
    JLAB and is being removed. This class now builds the workflow with 'swif2 create' and adds each job with
    'swif2 add-job', mapping the parameters that the parent Auger class used to encode in XML directly onto
    add-job command-line flags, then starts the workflow with 'swif2 run'.
 
    Note that 'swif2 create' fails if a workflow of the same name already exists. Existing workflows generated
    by this class should be fully cancelled and removed before resubmitting under the same name.
    """
 
    
    DEFAULT_PROJECT = 'hallb-pro'
 
    
    DEFAULT_SITE = 'jlab/enp'
 
    
    DEFAULT_OS = 'el9'
 
    
    VALID_PARTITIONS = frozenset(['production', 'ifarm', 'priority', 'jupyter', 'gpu'])
 
    
    LEGACY_TRACK_PARTITIONS = {
        'simulation': 'production',
        'analysis': 'production',
        'debug': 'priority',
        'one_pass': 'production',
    }
 
    def __init__(self):
 
        super().__init__()
 
        self.parser.add_argument("-w", "--workflow", nargs='?', help="Name of swif2 workflow", required=False)
        self.parser.add_argument("-P", "--project", nargs='?', help="Project/allocation (swif2 add-job -account)",
                                 required=False, default=Swif.DEFAULT_PROJECT)
        self.parser.add_argument("--site", nargs='?', help="swif2 site name (swif2 create -site-name)",
                                 required=False, default=Swif.DEFAULT_SITE)
        self.parser.add_argument("--max-concurrent", type=int,
                                 help="Max concurrent dispatched jobs (swif2 create -max-concurrent)",
                                 required=False, default=None)
        self.parser.add_argument("--recreate", action='store_true',
                                 help="If the workflow already exists, cancel and recreate it (destructive)")
 
    def parse_args(self, args):
        cl = super().parse_args(args)
        if cl.workflow:
            self.workflow = cl.workflow
        else:
            self.workflow = self.script_name
        self.project = cl.project
        self.swif_site = cl.site
        self.max_concurrent = cl.max_concurrent
        self.recreate = cl.recreate
        logger.debug(f'swif workflow name set to: {self.workflow}')
        return cl
 
    def _partition(self):
        """!
        Resolve the swif2 (slurm) partition. The old Auger "track" (debug mode -> 'debug', else the queue, else
        the 'production' default) is validated against the farm's real partitions: legacy track names are
        remapped to a valid partition with a warning, and anything else unknown is a hard error.
        @return  a valid slurm partition name
        """
        if self.debug:
            track = 'debug'
        elif self.queue is not None:
            track = self.queue
        else:
            track = 'production'
 
        if track in Swif.VALID_PARTITIONS:
            return track
        if track in Swif.LEGACY_TRACK_PARTITIONS:
            partition = Swif.LEGACY_TRACK_PARTITIONS[track]
            logger.warning("'%s' is a legacy Auger track, not a slurm partition; using partition '%s'. "
                           "Pass -q/--queue with one of (%s) to set it explicitly."
                           % (track, partition, ', '.join(sorted(Swif.VALID_PARTITIONS))))
            return partition
        raise Exception("Invalid slurm partition '%s'. Valid partitions are: %s."
                        % (track, ', '.join(sorted(Swif.VALID_PARTITIONS))))
 
    def submit(self):
 
        logger.info("Submitting swif workflow: {}".format(self.workflow))
 
        # Create the workflow. add-job (unlike the old add-jsub) requires the workflow to already exist.
        self._create_workflow()
 
        # Add each job to the workflow with its own 'swif2 add-job' invocation.
        job_ids = self._get_filtered_job_ids()
        logger.info('Submitting jobs: %s' % str(job_ids))
        for job_id in job_ids:
            if not self.jobstore.has_job_id(job_id):
                raise Exception('Job ID was not found in job store: %s' % job_id)
            job_params = self.jobstore.get_job(job_id)
            if self.check_output and Batch._outputs_exist(job_params):
                logger.warning("Skipping swif submission for job "
                               "because outputs already exist: %d" % job_id)
                continue
            self._run_swif2(self._add_job_cmd(job_params))
 
        # Start releasing jobs to the batch system.
        self._run_swif2(['run', self.workflow])
 
    def _run_swif2(self, args, check=True, quiet=False):
        """!
        Run a single 'swif2' subcommand, echoing its (non-empty) output.
        @param args  list of arguments following the 'swif2' executable
        @param check  if True, raise on a non-zero exit; if False, return the result for the caller to inspect
        @param quiet  if True, do not echo the command output (the caller handles it based on the outcome)
        @return  a (returncode, output) tuple
        """
        cmd = ['swif2'] + args
        logger.debug('swif2 command: %s' % ' '.join(cmd))
        proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
        out = proc.communicate()[0]
        text = out.decode()
        if not quiet:
            Swif._echo(text)
        if check and proc.returncode:
            raise Exception("swif2 command failed (exit %d): %s" % (proc.returncode, ' '.join(cmd)))
        return proc.returncode, text
 
    @staticmethod
    def _echo(text):
        """! Print the non-empty lines of some swif2 output."""
        printed = "".join([s for s in text.strip().splitlines(True) if s.strip()])
        if printed:
            print(printed)
 
    def _create_cmd(self):
        # The Auger request set Project/Track/OS on the request as a whole. In swif2 the project/allocation
        # and track/OS are applied per-job on add-job (see _add_job_cmd); the site is set here on create.
        cmd = ['create', '-workflow', self.workflow]
        if self.swif_site:
            cmd += ['-site-name', self.swif_site]
        if self.max_concurrent is not None:
            cmd += ['-max-concurrent', str(self.max_concurrent)]
        return cmd
 
    def _create_workflow(self):
        """!
        Ensure the swif2 workflow exists (add-job requires it to already exist).
 
        If a workflow of the same name already exists this is not an error: the create is skipped and jobs are
        added to the existing workflow. Pass --recreate to cancel and recreate it instead. Any other create
        failure is fatal. The create output is captured and only echoed on success or a genuine failure, so the
        benign 'already exists' case does not surface swif's scary error text.
        """
        returncode, text = self._run_swif2(self._create_cmd(), check=False, quiet=True)
        if returncode == 0:
            Swif._echo(text)
            return
        if 'already exists' in text:
            if self.recreate:
                logger.warning("Workflow '%s' already exists; cancelling and recreating it (--recreate)."
                               % self.workflow)
                self._run_swif2(['cancel', '-workflow', self.workflow], check=False)
                self._run_swif2(self._create_cmd())  # recreate; fatal if the name still cannot be reused
                return
            print("Workflow '%s' already exists; adding jobs to the existing workflow "
                  "(pass --recreate to cancel and recreate it instead)." % self.workflow)
            return
        raise Exception("swif2 create failed for workflow '%s':\n%s" % (self.workflow, text))
 
    def _add_job_cmd(self, job_params):
        """!
        Build the 'swif2 add-job' argument list for a single job.
 
        This reproduces, as add-job flags, the per-job information the parent Auger class used to write into the
        request XML: resource limits, inputs, outputs, stdout/stderr, the job name, and the command to run.
        @param job_params  raw job data dict from the job store
        @return list of arguments following the 'swif2' executable
        """
        job_id = job_params['job_id']
        year = job_params.get('year', '')
 
        cmd = ['add-job', '-workflow', self.workflow]
 
        # Unique per-workflow job name (mirrors the old Auger <Name> element).
        cmd += ['-name', '%shps%s' % (str(year), str(job_id))]
 
        # Resource limits (Auger Memory/DiskSpace/TimeLimit + one core).
        cmd += ['-cores', '1']
        cmd += ['-ram', '%dMB' % self.memorymemory]
        cmd += ['-disk', '%dGB' % self.diskspacediskspace]
        cmd += ['-time', '%dhours' % self.job_lengthjob_length]
 
        # Project/allocation -> account (add-job passes this to sbatch as '-A'). Auger encoded this as the
        # request <Project>; the working JLAB value is "hallb-pro" (see DEFAULT_PROJECT).
        cmd += ['-account', self.project]
 
        # Auger "Track" -> swif2 '-partition'; request OS -> '-constraint' (both confirmed against live jobs:
        # partition=production, constraint=el9 by default).
        cmd += ['-partition', self._partition()]
        cmd += ['-constraint', self.os if self.os is not None else Swif.DEFAULT_OS]
 
        # Redirect stdout/stderr to the log dir (Auger <Stdout>/<Stderr>).
        stdout_file = os.path.abspath(os.path.join(self.log_dir, "job.%d.out" % job_id))
        stderr_file = os.path.abspath(os.path.join(self.log_dir, "job.%d.err" % job_id))
        cmd += ['-stdout', stdout_file]
        cmd += ['-stderr', stderr_file]
 
        # Inputs: '-input <local-file> <remote-uri>' (Auger <Input dest=... src=...>).
        for dest, remote in self._job_inputs(job_params):
            cmd += ['-input', dest, remote]
 
        # Outputs: '-output <local-file> <remote-uri>' (Auger <Output src=... dest=...>).
        for local, remote in self._job_outputs(job_params):
            cmd += ['-output', local, remote]
 
        # The command to run: an executable tcsh wrapper script, passed as a single positional argument.
        # An inline command string does not survive: swif truncates it at the first ';' (only 'pwd' ran, empty
        # logs otherwise), and '-shell /bin/tcsh <string>' ran nothing at all. A script path has no shell
        # metacharacters for swif to truncate, and a bare single-token command is executed correctly.
        cmd += [self._write_job_script(job_id)]
 
        return cmd
 
    def _job_inputs(self, job_params):
        """!
        Yield (local_dest, remote_uri) pairs for a job's inputs, applying swif2 URI schemes.
        http inputs are skipped (they are not staged by swif), matching the previous Auger behavior.
        """
        pairs = []
        for src, dest in job_params.get('input_files', {}).items():
            if src.startswith('http'):
                logger.warning("http input file will not be staged by swif2: {}".format(src))
                continue
            pairs.append((dest, self._swif_uri(src)))
        return pairs
 
    def _job_outputs(self, job_params):
        """!
        Yield (local_src, remote_uri) pairs for a job's outputs, resolving ptag sources and applying schemes.
        """
        outputdir = job_params["output_dir"]
        j = self._create_job(job_params)  # needed to resolve ptag output sources
        pairs = []
        for src, dest in job_params["output_files"].items():
            local_src = j.resolve_output_src(src)
            dest_file = os.path.abspath(os.path.join(outputdir, dest))
            pairs.append((local_src, self._swif_uri(dest_file)))
        return pairs
 
    @staticmethod
    def _swif_uri(path):
        """!
        Convert a filesystem path to a swif2 remote URI. Tape paths get the 'mss:' scheme; ordinary filesystem
        paths (e.g. /lustre, /work) are left bare, matching both the old Auger behavior and the remote URIs
        observed on live JLAB workflows (swif canonicalizes tape URIs to 'mss:jlab:...' on its own). Paths that
        already carry a scheme are returned unchanged.
        """
        if path.startswith('mss:') or path.startswith('file:'):
            return path
        if path.startswith('/mss'):
            return 'mss:%s' % path
        return path
 
    def _write_job_script(self, job_id):
        """!
        Write an executable tcsh wrapper script that sets up the (c-shell) environment and runs the job, and
        return its absolute path.
 
        The script contents reproduce the old Auger <Command> CDATA block. It is written to the log dir (on
        /farm_out, which is shared with the compute nodes) and passed to add-job as the job command; passing a
        script path avoids the shell metacharacters that swif truncates out of an inline command string.
        @param job_id  job ID
        @return  absolute path to the generated script
        """
        script_path = os.path.abspath(os.path.join(self.log_dir, 'swif_job.%d.csh' % job_id))
        job_cmd = self.build_cmdbuild_cmd(job_id)
        lines = ['#!/bin/tcsh',
                 'pwd',
                 'env | sort',
                 'ls -lart',
                 'source %s' % os.path.realpath(self.setup_scriptsetup_script),
                 'source %s/bin/jlab-env.csh' % os.getenv('HPSMC_DIR'),
                 ' '.join(job_cmd),
                 'ls -lart']
        with open(script_path, 'w') as f:
            f.write('\n'.join(lines) + '\n')
        os.chmod(script_path, 0o755)
        logger.debug('Wrote swif2 job script: %s' % script_path)
        return script_path
 
 
class Local(Batch):
    """!
    Run local batch jobs sequentially.
    """
 
    def __init__(self):
        super().__init__()
 
    def submit_job(self, job_id):
        """! Run a single job locally."""
        cmd = self.build_cmd(job_id)
        if self.submit:
            logger.info(f"Executing local job: {job_id}")
            proc = subprocess.Popen(cmd, shell=False)
            proc.communicate()
            if proc.returncode:
                logger.error(f"Local execution of {job_id} returned error code: {proc.returncode}")
 
 
# Queue used to keep track of processes created by batch pool.
mp_queue = multiprocessing.Queue()
 
 
def run_job_pool(cmd):
    """! Run the command in a new process whose PID is added to a global MP queue."""
    try:
        sys.stdout.flush()
        proc = subprocess.Popen(cmd, preexec_fn=os.setsid)
        mp_queue.put(proc.pid)
        proc.wait()
        returncode = proc.returncode
    except subprocess.CalledProcessError as e:
        logger.error(str(e))
        sys.stdout.flush()
        pass
    return returncode
 
 
def is_running(proc):
    """!
    Check if a system process looks like it is still running.
    """
    return proc.status() in [psutil.STATUS_RUNNING,
                             psutil.STATUS_SLEEPING,
                             psutil.STATUS_DISK_SLEEP,
                             psutil.STATUS_IDLE]
 
 
class KillProcessQueue():
    """!
    Kill processes in the multiprocessing queue if the jobs are canceled.
    """
 
    def __init__(self, mp_queue):
        self.mp_queue = mp_queue
 
    def __enter__(self):
        return self
 
    def __exit__(self, type, val, tb):
        """! Kill processes on exit."""
        while True:
            pid = mp_queue.get()
            try:
                parent = psutil.Process(pid)
                for child in parent.children(recursive=True):
                    if is_running(child):
                        print('Killing running process: %d' % child.pid)
                        child.kill()
                if is_running(parent):
                    parent.kill()
            except Exception as e:
                # This probably just means it already finished.
                pass
 
            if mp_queue.empty():
                break
 
 
class Pool(Batch):
    """!
    Run a set of jobs in a local multiprocessing pool using Python's multiprocessing module.
 
    The number of processes to spawn can be provided using the '-p' argument.
    """
 
    # Max wait in seconds when getting results
    max_wait = 999999
 
    def __init__(self):
        super().__init__()
        self.parser.add_argument("-p", "--pool-size", type=int,
                                 help="Job pool size (only applicable when running pool)", required=False,
                                 default=multiprocessing.cpu_count())
 
    def submit_job(self, job_id):
        """!
        Make this a no-op as we do not implement single job submission for the processing pool.
        """
        pass
 
    def parse_args(self, args):
        cl = super().parse_args(args)
        self.pool_size = int(cl.pool_size)
        return cl
 
    def submit(self):
        """! Submit jobs to a local processing pool.
 
        This method will not return until all jobs are finished or execution
        is interrupted.
        """
 
        cmds = []
        for job_id in self._get_filtered_job_ids():
            cmd = self.build_cmd(job_id)
            cmds.append(cmd)
 
        # logger.debug('Running job commands in pool ...')
        # logger.debug('\n'.join([' '.join(cmd) for cmd in cmds]))
 
        if not len(cmds):
            raise Exception('No job IDs found to submit')
 
        # Run jobs in an MP pool and cleanup child processes on exit
        with KillProcessQueue(mp_queue):
            original_sigint_handler = signal.signal(signal.SIGINT, signal.SIG_IGN)
            pool = multiprocessing.Pool(self.pool_size)
            signal.signal(signal.SIGINT, original_sigint_handler)
            try:
                logger.info("Running %d jobs in pool ..." % len(cmds))
                res = pool.map_async(run_job_pool, cmds)
                # timeout must be properly set, otherwise tasks will crash
                logger.info("Pool results: " + str(res.get(Pool.max_wait)))
                logger.info("Normal termination")
                pool.close()
                pool.join()
            except KeyboardInterrupt:
                logger.fatal("Caught KeyboardInterrupt, terminating workers")
                pool.terminate()
            except Exception as e:
                logger.fatal("Caught Exception '%s', terminating workers" % (str(e)))
                pool.terminate()
            except BaseException:  # catch *all* exceptions
                e = sys.exc_info()[0]
                logger.fatal("Caught non-Python Exception '%s'" % (e))
                pool.terminate()
 
 
if __name__ == '__main__':
    system_dict = {
        "lsf": LSF,
        "slurm": Slurm,
        "auger": Auger,
        "local": Local,
        "pool": Pool,
        "swif": Swif
    }
    if len(sys.argv) > 1:
        system = sys.argv[1].lower()
        if system not in list(system_dict.keys()):
            raise Exception(f"The batch system {system} is not valid.")
        batch = system_dict[system]()
        args = sys.argv[2:]
        batch.parse_args(args)
        batch.submit()
    else:
        print("Usage: batch.py [system] [args]")
        print("    Available systems: {}".format(', '.join(list(system_dict.keys()))))