import logging import math import os import re import time from typing import Any, Dict, Optional import typeguard from parsl.jobs.states import JobState, JobStatus from parsl.launchers import SingleNodeLauncher from parsl.launchers.base import Launcher from parsl.providers.cluster_provider import ClusterProvider from parsl.providers.errors import SubmitException from parsl.providers.slurm.template import template_string from parsl.utils import RepresentationMixin, wtime_to_minutes logger = logging.getLogger(__name__) # From https://slurm.schedmd.com/sacct.html#SECTION_JOB-STATE-CODES sacct_translate_table = { 'PENDING': JobState.PENDING, 'RUNNING': JobState.RUNNING, 'CANCELLED': JobState.CANCELLED, 'COMPLETED': JobState.COMPLETED, 'FAILED': JobState.FAILED, 'NODE_FAIL': JobState.FAILED, 'BOOT_FAIL': JobState.FAILED, 'DEADLINE': JobState.TIMEOUT, 'TIMEOUT': JobState.TIMEOUT, 'REVOKED': JobState.FAILED, 'OUT_OF_MEMORY': JobState.FAILED, 'SUSPENDED': JobState.HELD, 'PREEMPTED': JobState.TIMEOUT, 'REQUEUED': JobState.PENDING } squeue_translate_table = { 'PD': JobState.PENDING, 'R': JobState.RUNNING, 'CA': JobState.CANCELLED, 'CF': JobState.PENDING, # (configuring), 'CG': JobState.RUNNING, # (completing), 'CD': JobState.COMPLETED, 'F': JobState.FAILED, # (failed), 'TO': JobState.TIMEOUT, # (timeout), 'NF': JobState.FAILED, # (node failure), 'RV': JobState.FAILED, # (revoked) and 'SE': JobState.FAILED # (special exit state) } class SlurmProvider(ClusterProvider, RepresentationMixin): """Slurm Execution Provider This provider uses sbatch to submit, sacct for status and scancel to cancel jobs. The sbatch script to be used is created from a template file in this same module. Parameters ---------- partition : str Slurm partition to request blocks from. If unspecified or ``None``, no partition slurm directive will be specified. account : str Slurm account to which to charge resources used by the job. If unspecified or ``None``, the job will use the user's default account. qos : str Slurm queue to place job in. If unspecified or ``None``, no queue slurm directive will be specified. constraint : str Slurm job constraint, often used to choose cpu or gpu type. If unspecified or ``None``, no constraint slurm directive will be added. clusters : str Slurm cluster name, or comma seperated cluster list, used to choose between different clusters in a federated Slurm instance. If unspecified or ``None``, no slurm directive for clusters will be added. nodes_per_block : int Nodes to provision per block. cores_per_node : int Specify the number of cores to provision per node. If set to None, executors will assume all cores on the node are available for computation. Default is None. mem_per_node : int Specify the real memory to provision per node in GB. If set to None, no explicit request to the scheduler will be made. Default is None. init_blocks : int Number of blocks to provision at the start of the run. Default is 1. min_blocks : int Minimum number of blocks to maintain. max_blocks : int Maximum number of blocks to maintain. parallelism : float Ratio of provisioned task slots to active tasks. A parallelism value of 1 represents aggressive scaling where as many resources as possible are used; parallelism close to 0 represents the opposite situation in which as few resources as possible (i.e., min_blocks) are used. walltime : str Walltime requested per block in HH:MM:SS. scheduler_options : str String to prepend to the #SBATCH blocks in the submit script to the scheduler. regex_job_id : str The regular expression used to extract the job ID from the ``sbatch`` standard output. The default is ``r"Submitted batch job (?P\\S*)"``, where ``id`` is the regular expression symbolic group for the job ID. worker_init : str Command to be run before starting a worker, such as 'module load Anaconda; source activate env'. exclusive : bool (Default = True) Requests nodes which are not shared with other running jobs. launcher : Launcher Launcher for this provider. Possible launchers include :class:`~parsl.launchers.SingleNodeLauncher` (the default), :class:`~parsl.launchers.SrunLauncher`, or :class:`~parsl.launchers.AprunLauncher` """ @typeguard.typechecked def __init__(self, partition: Optional[str] = None, account: Optional[str] = None, qos: Optional[str] = None, constraint: Optional[str] = None, clusters: Optional[str] = None, nodes_per_block: int = 1, cores_per_node: Optional[int] = None, mem_per_node: Optional[int] = None, init_blocks: int = 1, min_blocks: int = 0, max_blocks: int = 1, parallelism: float = 1, walltime: str = "00:10:00", scheduler_options: str = '', regex_job_id: str = r"Submitted batch job (?P\S*)", worker_init: str = '', cmd_timeout: int = 10, exclusive: bool = True, launcher: Launcher = SingleNodeLauncher()): label = 'slurm' super().__init__(label, nodes_per_block, init_blocks, min_blocks, max_blocks, parallelism, walltime, cmd_timeout=cmd_timeout, launcher=launcher) self.partition = partition self.cores_per_node = cores_per_node self.mem_per_node = mem_per_node self.exclusive = exclusive self.account = account self.qos = qos self.constraint = constraint self.clusters = clusters self.scheduler_options = scheduler_options + '\n' if exclusive: self.scheduler_options += "#SBATCH --exclusive\n" if partition: self.scheduler_options += "#SBATCH --partition={}\n".format(partition) if account: self.scheduler_options += "#SBATCH --account={}\n".format(account) if qos: self.scheduler_options += "#SBATCH --qos={}\n".format(qos) if constraint: self.scheduler_options += "#SBATCH --constraint={}\n".format(constraint) if clusters: self.scheduler_options += "#SBATCH --clusters={}\n".format(clusters) self.regex_job_id = regex_job_id self.worker_init = worker_init + '\n' # Check if sacct works and if not fall back to squeue cmd = "sacct -X" logger.debug("Executing %s", cmd) retcode, stdout, stderr = self.execute_wait(cmd) # If sacct fails it should return retcode=1 stderr="Slurm accounting storage is disabled" logger.debug(f"sacct returned retcode={retcode} stderr={stderr}") if retcode == 0: logger.debug("using sacct to get job status") _cmd = "sacct" # Add clusters option to sacct if provided if self.clusters: _cmd += f" --clusters={self.clusters}" # Using state%20 to get enough characters to not truncate output # of the state. Without output can look like " CANCELLED+" self._cmd = _cmd + " -X --noheader --format=jobid,state%20 --job '{0}'" self._translate_table = sacct_translate_table else: logger.debug(f"sacct failed with retcode={retcode}") logger.debug("falling back to using squeue to get job status") _cmd = "squeue" # Add clusters option to squeue if provided if self.clusters: _cmd += f" --clusters={self.clusters}" self._cmd = _cmd + " --noheader --format='%i %t' --job '{0}'" self._translate_table = squeue_translate_table def _status(self): '''Returns the status list for a list of job_ids Args: self Returns: [status...] : Status list of all jobs ''' job_id_list = ','.join( [jid for jid, job in self.resources.items() if not job['status'].terminal] ) if not job_id_list: logger.debug('No active jobs, skipping status update') return cmd = self._cmd.format(job_id_list) logger.debug("Executing %s", cmd) retcode, stdout, stderr = self.execute_wait(cmd) logger.debug("sacct/squeue returned %s %s", stdout, stderr) # Execute_wait failed. Do no update if retcode != 0: logger.warning("sacct/squeue failed with non-zero exit code {}".format(retcode)) return jobs_missing = set(self.resources.keys()) for line in stdout.split('\n'): if not line: # Blank line continue # Sacct includes extra information in some outputs # For example " CANCELLED by " # This splits and ignores anything past the first two unpacked values job_id, slurm_state, *ignore = line.split() if slurm_state not in self._translate_table: logger.warning(f"Slurm status {slurm_state} is not recognized") status = self._translate_table.get(slurm_state, JobState.UNKNOWN) logger.debug("Updating job {} with slurm status {} to parsl state {!s}".format(job_id, slurm_state, status)) self.resources[job_id]['status'] = JobStatus(status, stdout_path=self.resources[job_id]['job_stdout_path'], stderr_path=self.resources[job_id]['job_stderr_path']) jobs_missing.remove(job_id) # sacct can get job info after jobs have completed so this path shouldn't be hit # squeue does not report on jobs that are not running. So we are filling in the # blanks for missing jobs, we might lose some information about why the jobs failed. for missing_job in jobs_missing: logger.debug("Updating missing job {} to completed status".format(missing_job)) self.resources[missing_job]['status'] = JobStatus( JobState.COMPLETED, stdout_path=self.resources[missing_job]['job_stdout_path'], stderr_path=self.resources[missing_job]['job_stderr_path']) def submit(self, command: str, tasks_per_node: int, job_name="parsl.slurm") -> str: """Submit the command as a slurm job. Parameters ---------- command : str Command to be made on the remote side. tasks_per_node : int Command invocations to be launched per node job_name : str Name for the job Returns ------- job id : str A string identifier for the job """ scheduler_options = self.scheduler_options worker_init = self.worker_init if self.mem_per_node is not None: scheduler_options += '#SBATCH --mem={}g\n'.format(self.mem_per_node) worker_init += 'export PARSL_MEMORY_GB={}\n'.format(self.mem_per_node) if self.cores_per_node is not None: cpus_per_task = math.floor(self.cores_per_node / tasks_per_node) scheduler_options += '#SBATCH --cpus-per-task={}'.format(cpus_per_task) worker_init += 'export PARSL_CORES={}\n'.format(cpus_per_task) job_name = "{0}.{1}".format(job_name, time.time()) assert self.script_dir, "Expected script_dir to be set" script_path = os.path.join(self.script_dir, job_name) script_path = os.path.abspath(script_path) job_stdout_path = script_path + ".stdout" job_stderr_path = script_path + ".stderr" logger.debug("Requesting one block with {} nodes".format(self.nodes_per_block)) job_config: Dict[str, Any] = {} job_config["submit_script_dir"] = self.script_dir job_config["nodes"] = self.nodes_per_block job_config["tasks_per_node"] = tasks_per_node job_config["walltime"] = wtime_to_minutes(self.walltime) job_config["scheduler_options"] = scheduler_options job_config["worker_init"] = worker_init job_config["user_script"] = command job_config["job_stdout_path"] = job_stdout_path job_config["job_stderr_path"] = job_stderr_path # Wrap the command job_config["user_script"] = self.launcher(command, tasks_per_node, self.nodes_per_block) logger.debug("Writing submit script") self._write_submit_script(template_string, script_path, job_name, job_config) retcode, stdout, stderr = self.execute_wait("sbatch {0}".format(script_path)) if retcode == 0: for line in stdout.split('\n'): match = re.match(self.regex_job_id, line) if match: job_id = match.group("id") self.resources[job_id] = {'job_id': job_id, 'status': JobStatus(JobState.PENDING), 'job_stdout_path': job_stdout_path, 'job_stderr_path': job_stderr_path, } return job_id else: logger.error("Could not read job ID from submit command standard output.") logger.error("Retcode:%s STDOUT:%s STDERR:%s", retcode, stdout.strip(), stderr.strip()) raise SubmitException( job_name, "Could not read job ID from submit command standard output", stdout=stdout, stderr=stderr, retcode=retcode ) else: logger.error("Submit command failed") logger.error("Retcode:%s STDOUT:%s STDERR:%s", retcode, stdout.strip(), stderr.strip()) raise SubmitException( job_name, "Could not read job ID from submit command standard output", stdout=stdout, stderr=stderr, retcode=retcode ) def cancel(self, job_ids): ''' Cancels the jobs specified by a list of job ids Args: job_ids : [ ...] Returns : [True/False...] : If the cancel operation fails the entire list will be False. ''' job_id_list = ' '.join(job_ids) # Make the command to cancel jobs _cmd = "scancel" if self.clusters: _cmd += f" --clusters={self.clusters}" _cmd += " {0}" retcode, stdout, stderr = self.execute_wait(_cmd.format(job_id_list)) rets = None if retcode == 0: for jid in job_ids: self.resources[jid]['status'] = JobStatus(JobState.CANCELLED) # Setting state to cancelled rets = [True for i in job_ids] else: rets = [False for i in job_ids] return rets @property def status_polling_interval(self): return 60