diff --git a/ravenframework/CustomModes/ClusterMode.py b/ravenframework/CustomModes/ClusterMode.py new file mode 100644 index 0000000000..9e46d10658 --- /dev/null +++ b/ravenframework/CustomModes/ClusterMode.py @@ -0,0 +1,100 @@ +# Copyright 2017 Battelle Energy Alliance, LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" + Shared base class for cluster (node-file based) SimulationModes such as the + Slurm and PBS modes. This module intentionally defines NO modeName / + modeClassName, so the CustomModes discovery mechanism skips it. +""" + +from ravenframework import Simulation +from ravenframework.CustomModes import ClusterUtils + + +class ClusterSimulationMode(Simulation.SimulationMode): + """ + Base class for SimulationModes that distribute runs over the nodes of a + scheduler allocation described by a node file (one line per processor). + Subclasses perform the scheduler-specific node discovery and remote + submission; the batch sizing, node-file splitting and precommand assembly + live here (previously triplicated across the Slurm/PBS/MPI modes). + """ + + def _modifyInfoForCluster(self, runInfoDict, nodeFileName, mpiParams=None, + createPrecommand=True, splitNodeFiles=True, + clusterName="the cluster"): + """ + Shared implementation of modifyInfo for node-file based cluster modes. + @ In, runInfoDict, dict, the original runInfo + @ In, nodeFileName, str or None, path of the node file describing the + allocation (None: not inside an allocation, run on the local machine) + @ In, mpiParams, list(str), optional, extra parameters for mpiexec + @ In, createPrecommand, bool, optional, whether to prepend the mpiexec + precommand (False leaves the existing precommand untouched) + @ In, splitNodeFiles, bool, optional, whether to write the per-batch + node_%INDEX% files when batchSize > 1 (False when the launcher, e.g. + srun, assigns resources itself) + @ Out, newRunInfo, dict, of modified values + """ + newRunInfo = {} + newRunInfo['batchSize'] = runInfoDict['batchSize'] + numMPI = runInfoDict['NumMPI'] + if nodeFileName is not None: + self.raiseADebug('Setting up remote nodes based on "{}"'.format(nodeFileName)) + lines = ClusterUtils.readNodeFile(nodeFileName) + if len(lines) == 0: + self.raiseAnError(IOError, 'Node file "{}" is empty! Cannot determine ' + 'the nodes available on {}.'.format(nodeFileName, clusterName)) + #XXX This is an undocumented way to pass information back + # (JobHandler strips these lines, so keep the newline-terminated form) + newRunInfo['Nodes'] = [line + "\n" for line in lines] + oldBatchsize = runInfoDict['batchSize'] + newBatchsize, changed = ClusterUtils.computeBatchSize(len(lines), numMPI, oldBatchsize) + if changed: + newRunInfo['batchSize'] = newBatchsize + self.raiseAWarning("changing batchsize from "+str(oldBatchsize)+" to " + +str(newBatchsize)+" to fit on "+str(len(lines))+" processors") + newBatchsize = newRunInfo['batchSize'] + self.raiseADebug('Batch size is "{}"'.format(newBatchsize)) + if newBatchsize > 1 and splitNodeFiles: + #need to split node lines so that numMPI processors are available per run, + #then give each index a separate file + ClusterUtils.writeNodeSubFiles(lines, newBatchsize, numMPI, runInfoDict['WorkingDir']) + nodeCommand = runInfoDict["NodeParameter"]+" %BASE_WORKING_DIR%/node_%INDEX% " + elif splitNodeFiles: + #If only one batch just use the original node file + nodeCommand = runInfoDict["NodeParameter"]+" "+nodeFileName + else: + #the launcher (e.g. srun) assigns the resources itself + nodeCommand = " " + else: + #Not inside an allocation and no node file supplied in the input. + #TODO, we don't have a way to know which machines it can run on + # in this case, so just distribute it over the local machine: + nodeCommand = " " + + # Create the mpiexec pre command + # Note, with defaults the precommand is "mpiexec -f nodeFile -n numMPI" + if createPrecommand: + newRunInfo['precommand'] = ClusterUtils.buildMPIPrecommand( + runInfoDict["MPIExec"], mpiParams or [], nodeCommand, numMPI, + runInfoDict['precommand']) + else: + newRunInfo['precommand'] = runInfoDict['precommand'] + if runInfoDict['NumThreads'] > 1: + newRunInfo['threadParameter'] = runInfoDict['threadParameter'] + #add number of threads to the post command. + newRunInfo['postcommand'] = " {} {}".format(newRunInfo['threadParameter'], runInfoDict['postcommand']) + self.raiseAMessage("precommand: "+newRunInfo['precommand']+", postcommand: " + +newRunInfo.get('postcommand', runInfoDict['postcommand'])) + return newRunInfo diff --git a/ravenframework/CustomModes/ClusterUtils.py b/ravenframework/CustomModes/ClusterUtils.py new file mode 100644 index 0000000000..f64c2ae41f --- /dev/null +++ b/ravenframework/CustomModes/ClusterUtils.py @@ -0,0 +1,228 @@ +# Copyright 2017 Battelle Energy Alliance, LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" + Shared, framework-independent utilities for the cluster SimulationModes + (Slurm, PBS, MPI legacy). + + NOTE: this module intentionally has NO imports from the ravenframework + package so that it can be unit tested without a built RAVEN installation. +""" + +import os +import re +import math +import string +import subprocess + + +def parseSlurmTasksPerNode(spec): + """ + Parses the SLURM_TASKS_PER_NODE environment variable format into a flat + list of per-node task counts. + The format is a comma separated list of entries, where each entry is + either "N" (N tasks on one node) or "N(xM)" (N tasks on each of M nodes). + For example "36(x2),20" -> [36, 36, 20]. + @ In, spec, str, the SLURM_TASKS_PER_NODE-style specification + @ Out, counts, list(int), one entry per node with the task count + """ + counts = [] + if spec is None: + return counts + for entry in spec.split(","): + entry = entry.strip() + if not entry: + continue + match = re.match(r"^(\d+)(\(x(\d+)\))?$", entry) + if match is None: + raise ValueError(f'Unparsable SLURM_TASKS_PER_NODE entry "{entry}" in "{spec}"') + tasks = int(match.group(1)) + repeat = int(match.group(3)) if match.group(3) is not None else 1 + counts.extend([tasks] * repeat) + return counts + + +def slurmNodeListFromScontrol(nodeList=None, tasksPerNode=None, runner=subprocess.run): + """ + Expands a Slurm node list (e.g. "node[01-03],gpu01") into one line per + task using "scontrol show hostnames" and the SLURM_TASKS_PER_NODE + specification. This is the fallback node-discovery mechanism used when + "srun hostname" is unavailable or fails. + @ In, nodeList, str, optional, node list (defaults to $SLURM_JOB_NODELIST) + @ In, tasksPerNode, str, optional, tasks-per-node spec (defaults to + $SLURM_TASKS_PER_NODE, then $SLURM_CPUS_ON_NODE, then 1 per node) + @ In, runner, callable, optional, subprocess.run-compatible callable + (injectable for testing) + @ Out, lines, list(str), one hostname entry per task (no trailing newline) + """ + if nodeList is None: + nodeList = os.environ.get("SLURM_JOB_NODELIST") + if nodeList is None: + return [] + result = runner(["scontrol", "show", "hostnames", nodeList], + capture_output=True, text=True, timeout=60) + if result.returncode != 0: + raise RuntimeError(f'"scontrol show hostnames {nodeList}" failed with code ' + f'{result.returncode}: {result.stderr}') + hosts = [line.strip() for line in result.stdout.splitlines() if line.strip()] + if tasksPerNode is None: + tasksPerNode = os.environ.get("SLURM_TASKS_PER_NODE") + if tasksPerNode is not None: + counts = parseSlurmTasksPerNode(tasksPerNode) + else: + cpusOnNode = os.environ.get("SLURM_CPUS_ON_NODE") + perNode = int(cpusOnNode) if cpusOnNode is not None else 1 + counts = [perNode] * len(hosts) + if len(counts) < len(hosts): + # be forgiving: pad with the last known count + counts = counts + [counts[-1] if counts else 1] * (len(hosts) - len(counts)) + lines = [] + for host, count in zip(hosts, counts): + lines.extend([host] * count) + return lines + + +def readNodeFile(path): + """ + Reads a node file (one node name per line, one line per available + task/processor) and returns the stripped, non-empty lines. + @ In, path, str, the path of the node file + @ Out, lines, list(str), the node names (one entry per processor) + """ + with open(path, "r") as nodeFileObject: + return [line.strip() for line in nodeFileObject if line.strip()] + + +def computeBatchSize(numProcessors, numMPI, requestedBatchSize): + """ + Computes the batch size that fits in the given number of processors. + The batch size is the number of processors divided by numMPI (processors + per run); floor/max keep the numbers reasonable. + @ In, numProcessors, int, number of available processors (node file lines) + @ In, numMPI, int, number of MPI processes per run + @ In, requestedBatchSize, int, the batch size requested in the input + @ Out, (batchSize, changed), (int, bool), the usable batch size and whether + it had to be reduced from the requested one + """ + maxBatchsize = max(int(math.floor(numProcessors / numMPI)), 1) + if maxBatchsize < requestedBatchSize: + return maxBatchsize, True + return requestedBatchSize, False + + +def writeNodeSubFiles(lines, batchSize, numMPI, workingDir, prefix="node_"): + """ + Splits the node list so that numMPI processors are available per batch + slot, writing one node file per slot (node_0, node_1, ...): the files + referenced by the "%BASE_WORKING_DIR%/node_%INDEX%" placeholder command. + @ In, lines, list(str), node names, one entry per processor + @ In, batchSize, int, number of batch slots + @ In, numMPI, int, number of MPI processes per run + @ In, workingDir, str, directory in which to write the files + @ In, prefix, str, optional, file name prefix + @ Out, written, list(str), the paths written + """ + written = [] + for i in range(batchSize): + subFileName = os.path.join(workingDir, f"{prefix}{i}") + with open(subFileName, "w") as subNodeFile: + for line in lines[i*numMPI : (i+1)*numMPI]: + subNodeFile.write(line.rstrip("\n") + "\n") + written.append(subFileName) + return written + + +def buildMPIPrecommand(mpiExec, mpiParams, nodeCommand, numMPI, existingPrecommand): + """ + Creates the mpiexec precommand. With defaults the precommand is + "mpiexec -f nodeFile -n numMPI ". + @ In, mpiExec, str, the mpi executor (e.g. "mpiexec") + @ In, mpiParams, list(str), extra parameters for mpi + @ In, nodeCommand, str, the node-selection portion (e.g. "-f nodefile"), + or " " when running on the local machine only + @ In, numMPI, int, number of MPI processes per run + @ In, existingPrecommand, str, the pre-existing precommand to append + @ Out, precommand, str, the assembled precommand + """ + mpiParamsStr = (" ".join(mpiParams) + " ") if mpiParams else "" + return mpiExec + " " + mpiParamsStr + nodeCommand + " -n " + str(numMPI) + " " + existingPrecommand + + +def sanitizeJobName(jobName, maxLength=None): + """ + Validates and (optionally) shortens a scheduler job name. Only + alphanumeric characters, "_" and "-" are allowed. When maxLength is given + and exceeded, the name is shortened keeping the head and the last 4 + characters (e.g. maxLength=15: first 10 + "-" + last 4). + @ In, jobName, str, the requested job name + @ In, maxLength, int, optional, maximum allowed length + @ Out, jobName, str, the validated (possibly shortened) job name + """ + validChars = set(string.ascii_letters) | set(string.digits) | set('-_') + if any(char not in validChars for char in jobName): + raise ValueError('JobName can only contain alphanumeric, "_" and "-" ' + 'characters! Received: ' + jobName) + if maxLength is not None and len(jobName) > maxLength: + jobName = jobName[:maxLength-5] + '-' + jobName[-4:] + return jobName + + +def buildSrunPrecommand(numMPI, mpiParams, existingPrecommand): + """ + Creates a Slurm-native "srun" precommand instead of the mpiexec+nodefile + one. Slurm tracks per-step resource assignment itself, so no node files + are needed: "--exact" gives each step exactly the requested tasks and + "--overlap" allows the steps of a batch to share the allocation. Works + with all major MPI stacks via PMI/PMIx. + @ In, numMPI, int, number of MPI processes (tasks) per run + @ In, mpiParams, list(str), extra parameters passed to srun + @ In, existingPrecommand, str, the pre-existing precommand to append + @ Out, precommand, str, the assembled srun precommand + """ + mpiParamsStr = (" ".join(mpiParams) + " ") if mpiParams else "" + return "srun --overlap --exact -n " + str(numMPI) + " " + mpiParamsStr + existingPrecommand + + +def assembleDaskJobqueueKwargs(config, runInfoDict): + """ + Assembles the constructor keyword arguments for a dask_jobqueue cluster + (SLURMCluster / PBSCluster) from the RunInfo configuration, + deriving sensible defaults from the RAVEN run info. Pure function so it + can be unit tested without dask_jobqueue installed. + @ In, config, dict, with 'scheduler' ("slurm" or "pbs") and 'options' + (dict of dask_jobqueue options from the XML attributes; 'memory' is + required by dask_jobqueue, 'cores' defaults to numProcByRun, 'jobs' + controls how many scheduler jobs to scale to and defaults to batchSize) + @ In, runInfoDict, dict, the RAVEN run info dictionary + @ Out, (clusterClassName, kwargs, jobs), (str, dict, int), the + dask_jobqueue class name, its constructor kwargs, and the number of + scheduler jobs to scale the cluster to + """ + clusterClasses = {'slurm': 'SLURMCluster', 'pbs': 'PBSCluster'} + scheduler = str(config.get('scheduler', '')).strip().lower() + if scheduler not in clusterClasses: + raise ValueError(f' scheduler must be one of {sorted(clusterClasses)}, ' + f'got "{scheduler}"') + options = dict(config.get('options') or {}) + memory = options.pop('memory', None) + if memory is None: + raise ValueError(' requires a "memory" attribute (per scheduler job), ' + 'e.g. memory="4GB", because dask_jobqueue requires it') + cores = int(options.pop('cores', max(1, int(runInfoDict.get('numProcByRun', 1))))) + jobs = int(options.pop('jobs', max(1, int(runInfoDict.get('batchSize', 1))))) + kwargs = dict(cores=cores, memory=memory) + if 'walltime' not in options and runInfoDict.get('expectedTime'): + kwargs['walltime'] = runInfoDict['expectedTime'] + # any remaining attributes (queue, account, interface, ...) pass through verbatim + kwargs.update(options) + return clusterClasses[scheduler], kwargs, jobs diff --git a/ravenframework/CustomModes/MPILegacySimulationMode.py b/ravenframework/CustomModes/MPILegacySimulationMode.py index 56322753e2..bc1a527e4d 100644 --- a/ravenframework/CustomModes/MPILegacySimulationMode.py +++ b/ravenframework/CustomModes/MPILegacySimulationMode.py @@ -229,7 +229,11 @@ def XMLread(self, xmlNode): for child in xmlNode: child_tag = child.tag.lower() if child.tag == "nodefileenv": - self.__nodefile = os.environ[child.text.strip()] + envName = child.text.strip() + if envName not in os.environ: + self.raiseAnError(IOError, f' environment variable "{envName}" ' + 'is not defined in the current environment!') + self.__nodefile = os.environ[envName] elif child.tag == "nodefile": self.__nodefile = child.text.strip() elif child_tag == "runqsub": diff --git a/ravenframework/CustomModes/PBSSimulationMode.py b/ravenframework/CustomModes/PBSSimulationMode.py index e90f2becd6..f1b486b001 100644 --- a/ravenframework/CustomModes/PBSSimulationMode.py +++ b/ravenframework/CustomModes/PBSSimulationMode.py @@ -22,12 +22,14 @@ import math import string from ravenframework import Simulation +from ravenframework.CustomModes import ClusterUtils +from ravenframework.CustomModes.ClusterMode import ClusterSimulationMode #For the mode information modeName = ["mpi","pbs"] modeClassName = "PBSSimulationMode" -class PBSSimulationMode(Simulation.SimulationMode): +class PBSSimulationMode(ClusterSimulationMode): """ PBSSimulationMode is a specialized class of SimulationMode. It is aimed to distribute the runs using the MPI protocol on PBS @@ -57,72 +59,24 @@ def modifyInfo(self, runInfoDict): @ In, runInfoDict, dict, the original runInfo @ Out, newRunInfo, dict, of modified values """ - newRunInfo = {} - newRunInfo['batchSize'] = runInfoDict['batchSize'] + nodeFileName = None if self.__nodefile or self.__inPbs: if not self.__nodefile: #Figure out number of nodes and use for batchsize - nodefile = os.environ["PBS_NODEFILE"] + nodeFileName = os.environ["PBS_NODEFILE"] else: - nodefile = self.__nodefile - self.raiseADebug('Setting up remote nodes based on "{}"'.format(nodefile)) - lines = open(nodefile,"r").readlines() - #XXX This is an undocumented way to pass information back - newRunInfo['Nodes'] = list(lines) - numMPI = runInfoDict['NumMPI'] - oldBatchsize = runInfoDict['batchSize'] - #the batchsize is just the number of nodes of which there is one - # per line in the nodefile divided by the numMPI (which is per run) - # and the floor and int and max make sure that the numbers are reasonable - maxBatchsize = max(int(math.floor(len(lines) / numMPI)), 1) - - if maxBatchsize < oldBatchsize: - newRunInfo['batchSize'] = maxBatchsize - self.raiseAWarning("changing batchsize from "+str(oldBatchsize)+" to "+str(maxBatchsize)+" to fit on "+str(len(lines))+" processors") - newBatchsize = newRunInfo['batchSize'] - self.raiseADebug('Batch size is "{}"'.format(newBatchsize)) - if newBatchsize > 1: - #need to split node lines so that numMPI nodes are available per run - workingDir = runInfoDict['WorkingDir'] - for i in range(newBatchsize): - nodeFile = open(os.path.join(workingDir, f"node_{i}"), "w") - for line in lines[i*numMPI : (i+1) * numMPI]: - nodeFile.write(line) - nodeFile.close() - #then give each index a separate file. - nodeCommand = runInfoDict["NodeParameter"]+" %BASE_WORKING_DIR%/node_%INDEX% " - else: - #If only one batch just use original node file - nodeCommand = runInfoDict["NodeParameter"]+" "+nodefile - - else: - #Not in PBS, so can't look at PBS_NODEFILE and none supplied in input - newBatchsize = newRunInfo['batchSize'] - numMPI = runInfoDict['NumMPI'] - #TODO, we don't have a way to know which machines it can run on - # when not in PBS so just distribute it over the local machine: - nodeCommand = " " + nodeFileName = self.__nodefile #Disable MPI processor affinity, which causes multiple processes # to be forced to the same thread. os.environ["MV2_ENABLE_AFFINITY"] = "0" - if len(self.__mpiparams) > 0: - mpiParams = " ".join(self.__mpiparams)+" " - else: - mpiParams = "" - # Create the mpiexec pre command - # Note, with defaults the precommand is "mpiexec -f nodeFile -n numMPI" - if self.__createPrecommand: - newRunInfo['precommand'] = runInfoDict["MPIExec"]+" "+mpiParams+nodeCommand+" -n "+str(numMPI)+" "+runInfoDict['precommand'] - else: - newRunInfo['precommand'] = runInfoDict['precommand'] - if runInfoDict['NumThreads'] > 1: - newRunInfo['threadParameter'] = runInfoDict['threadParameter'] - #add number of threads to the post command. - newRunInfo['postcommand'] =" {} {}".format(newRunInfo['threadParameter'],runInfoDict['postcommand']) - self.raiseAMessage("precommand: "+newRunInfo['precommand']+", postcommand: "+newRunInfo.get('postcommand',runInfoDict['postcommand'])) - return newRunInfo + #the batch sizing, node-file splitting and precommand assembly are shared + # with the other cluster modes (see ClusterMode.ClusterSimulationMode) + return self._modifyInfoForCluster(runInfoDict, nodeFileName, + mpiParams=self.__mpiparams, + createPrecommand=self.__createPrecommand, + clusterName="this PBS allocation") def __createAndRunQSUB(self, runInfoDict): """ @@ -148,14 +102,14 @@ def __createAndRunQSUB(self, runInfoDict): ncpus = runInfoDict['NumThreads'] # job title jobName = runInfoDict['JobName'] if 'JobName' in runInfoDict.keys() else 'raven_qsub' - ## fix up job title - validChars = set(string.ascii_letters).union(set(string.digits)).union(set('-_')) - if any(char not in validChars for char in jobName): - raise IOError('JobName can only contain alphanumeric and "_", "-" characters! Received'+jobName) - #check jobName for length - if len(jobName) > 15: - jobName = jobName[:10]+'-'+jobName[-4:] - print('JobName is limited to 15 characters; truncating to '+jobName) + ## fix up job title (shared validator; PBS limits names to 15 characters) + try: + shortJobName = ClusterUtils.sanitizeJobName(jobName, maxLength=15) + except ValueError as err: + self.raiseAnError(IOError, str(err)) + if shortJobName != jobName: + self.raiseAMessage('JobName is limited to 15 characters; truncating to '+shortJobName) + jobName = shortJobName # Generate the qsub command needed to run input ## raven_framework location raven = os.path.abspath(os.path.join(frameworkDir,'..','raven_framework')) @@ -201,7 +155,11 @@ def XMLread(self, xmlNode): """ for child in xmlNode: if child.tag == "nodefileenv": - self.__nodefile = os.environ[child.text.strip()] + envName = child.text.strip() + if envName not in os.environ: + self.raiseAnError(IOError, f' environment variable "{envName}" ' + 'is not defined in the current environment!') + self.__nodefile = os.environ[envName] elif child.tag == "nodefile": self.__nodefile = child.text.strip() elif child.tag == "memory": diff --git a/ravenframework/CustomModes/SlurmSimulationMode.py b/ravenframework/CustomModes/SlurmSimulationMode.py index ac26cbb5b4..09f66b4adb 100644 --- a/ravenframework/CustomModes/SlurmSimulationMode.py +++ b/ravenframework/CustomModes/SlurmSimulationMode.py @@ -18,14 +18,17 @@ import os import math import string +import subprocess from ravenframework import Simulation from ravenframework.utils import InputData, InputTypes +from ravenframework.CustomModes import ClusterUtils +from ravenframework.CustomModes.ClusterMode import ClusterSimulationMode #For the mode information modeName = "slurm" modeClassName = "SlurmSimulationMode" -class SlurmSimulationMode(Simulation.SimulationMode): +class SlurmSimulationMode(ClusterSimulationMode): """ SlurmSimulationMode is a specialized class of SimulationMode. It is aimed to distribute the runs on a Slurm cluster @@ -46,6 +49,8 @@ def __init__(self, *args): self.__partition = None #If not none, use this for partition= self.__mpiparams = [] #Paramaters to give to mpi self.__createPrecommand = True #If true, do create precommand. + self.__runSbatch = False #If true, submit this run via sbatch when outside Slurm. + self.__useSrun = False #If true, launch runs with native srun instead of mpiexec+nodefiles. self.printTag = 'SLURM SIMULATION MODE' def modifyInfo(self, runInfoDict): @@ -55,68 +60,68 @@ def modifyInfo(self, runInfoDict): @ In, runInfoDict, dict, the original runInfo @ Out, newRunInfo, dict, of modified values """ - newRunInfo = {} - newRunInfo['batchSize'] = runInfoDict['batchSize'] workingDir = runInfoDict['WorkingDir'] + nodeFileName = None if self.__nodeFile or self.__inSlurm: if not self.__nodeFile: self.__nodeFile = os.path.join(workingDir,"slurmNodeFile_"+str(os.getpid())) - #generate nodeFile - os.system("srun --overlap -- hostname > "+self.__nodeFile) - self.raiseADebug('Setting up remote nodes based on "{}"'.format(self.__nodeFile)) - lines = open(self.__nodeFile,"r").readlines() - #XXX This is an undocumented way to pass information back - newRunInfo['Nodes'] = list(lines) - numMPI = runInfoDict['NumMPI'] - oldBatchsize = runInfoDict['batchSize'] - #the batchsize is just the number of nodes of which there is one - # per line in the nodeFile divided by the numMPI (which is per run) - # and the floor and int and max make sure that the numbers are reasonable - maxBatchsize = max(int(math.floor(len(lines) / numMPI)), 1) + #generate nodeFile (checked srun, with scontrol-based fallback) + self.__generateNodeFile(self.__nodeFile) + nodeFileName = self.__nodeFile + if self.__useSrun: + #srun-native launch: Slurm assigns per-step resources itself, so no + # per-batch node files are needed; keep node discovery for the Nodes + # bookkeeping (used e.g. by the internal-parallel cluster bring-up) + newRunInfo = self._modifyInfoForCluster(runInfoDict, nodeFileName, + mpiParams=None, + createPrecommand=False, + splitNodeFiles=False, + clusterName="this Slurm allocation") + if self.__createPrecommand: + newRunInfo['precommand'] = ClusterUtils.buildSrunPrecommand( + runInfoDict['NumMPI'], self.__mpiparams, runInfoDict['precommand']) + self.raiseAMessage("srun precommand: "+newRunInfo['precommand']) + return newRunInfo + #the batch sizing, node-file splitting and precommand assembly are shared + # with the other cluster modes (see ClusterMode.ClusterSimulationMode) + return self._modifyInfoForCluster(runInfoDict, nodeFileName, + mpiParams=self.__mpiparams, + createPrecommand=self.__createPrecommand, + clusterName="this Slurm allocation") - if maxBatchsize < oldBatchsize: - newRunInfo['batchSize'] = maxBatchsize - self.raiseAWarning("changing batchsize from "+str(oldBatchsize)+" to "+str(maxBatchsize)+" to fit on "+str(len(lines))+" processors") - newBatchsize = newRunInfo['batchSize'] - self.raiseADebug('Batch size is "{}"'.format(newBatchsize)) - if newBatchsize > 1: - #need to split node lines so that numMPI nodes are available per run - workingDir = runInfoDict['WorkingDir'] - for i in range(newBatchsize): - subNodeFile = open(os.path.join(workingDir, f"node_{i}"), "w") - for line in lines[i*numMPI : (i+1) * numMPI]: - subNodeFile.write(line) - subNodeFile.close() - #then give each index a separate file. - nodeCommand = runInfoDict["NodeParameter"]+" %BASE_WORKING_DIR%/node_%INDEX% " + def __generateNodeFile(self, nodeFileName): + """ + Generates the node file (one line per available task/processor) for the + current Slurm allocation. Tries "srun hostname" first and falls back to + expanding $SLURM_JOB_NODELIST via "scontrol show hostnames". + @ In, nodeFileName, str, the path of the node file to write + @ Out, None + """ + lines = None + try: + result = subprocess.run(["srun", "--overlap", "--", "hostname"], + capture_output=True, text=True, timeout=300) + if result.returncode == 0 and result.stdout.strip(): + lines = [line for line in result.stdout.splitlines() if line.strip()] else: - #If only one batch just use original node file - nodeCommand = runInfoDict["NodeParameter"]+" "+self.__nodeFile - - else: - #Not in PBS, so can't look at PBS_NODEFILE and none supplied in input - newBatchsize = newRunInfo['batchSize'] - numMPI = runInfoDict['NumMPI'] - #TODO, we don't have a way to know which machines it can run on - # when not in PBS so just distribute it over the local machine: - nodeCommand = " " - - if len(self.__mpiparams) > 0: - mpiParams = " ".join(self.__mpiparams)+" " - else: - mpiParams = "" - # Create the mpiexec pre command - # Note, with defaults the precommand is "mpiexec -f nodeFile -n numMPI" - if self.__createPrecommand: - newRunInfo['precommand'] = runInfoDict["MPIExec"]+" "+mpiParams+nodeCommand+" -n "+str(numMPI)+" "+runInfoDict['precommand'] - else: - newRunInfo['precommand'] = runInfoDict['precommand'] - if runInfoDict['NumThreads'] > 1: - newRunInfo['threadParameter'] = runInfoDict['threadParameter'] - #add number of threads to the post command. - newRunInfo['postcommand'] =" {} {}".format(newRunInfo['threadParameter'],runInfoDict['postcommand']) - self.raiseAMessage("precommand: "+newRunInfo['precommand']+", postcommand: "+newRunInfo.get('postcommand',runInfoDict['postcommand'])) - return newRunInfo + self.raiseAWarning('"srun --overlap -- hostname" failed (return code ' + f'{result.returncode}): {result.stderr.strip()}') + except (OSError, subprocess.TimeoutExpired, subprocess.SubprocessError) as exc: + self.raiseAWarning(f'Unable to run "srun --overlap -- hostname": {exc}') + if lines is None: + # fall back to scontrol-based expansion of the allocation node list + self.raiseADebug('Falling back to "scontrol show hostnames" for node discovery') + try: + lines = ClusterUtils.slurmNodeListFromScontrol() + except (OSError, RuntimeError, ValueError, subprocess.SubprocessError) as exc: + self.raiseAnError(RuntimeError, 'Could not determine the nodes of this Slurm ' + f'allocation with either srun or scontrol: {exc}') + if not lines: + self.raiseAnError(RuntimeError, 'Slurm node discovery returned no nodes! ' + 'Check that RAVEN is running inside a valid allocation.') + with open(nodeFileName, "w") as nodeFileObject: + for line in lines: + nodeFileObject.write(line.strip() + "\n") def __createAndRunSbatch(self, runInfoDict): """ @@ -144,10 +149,11 @@ def __createAndRunSbatch(self, runInfoDict): ncpus = runInfoDict['NumThreads'] # job title jobName = runInfoDict['JobName'] if 'JobName' in runInfoDict.keys() else 'raven_qsub' - ## fix up job title - validChars = set(string.ascii_letters).union(set(string.digits)).union(set('_')) - if any(char not in validChars for char in jobName): - raise IOError('JobName can only contain alphanumeric and "_" characters! Received'+jobName) + ## fix up job title (shared validator; alphanumeric, "_" and "-" allowed) + try: + jobName = ClusterUtils.sanitizeJobName(jobName) + except ValueError as err: + self.raiseAnError(IOError, str(err)) #--job-name= # Generate the sbatch command needed to run input ## raven_framework location @@ -172,11 +178,10 @@ def __createAndRunSbatch(self, runInfoDict): remoteRunCommand["cwd"] = runInfoDict['InputDir'] ## command to run in that directory remoteRunCommand["args"] = command - print("remoteRunCommand",remoteRunCommand) - print("COMMAND", command_env["COMMAND"]) - print("RAVEN_FRAMEWORK_DIR", command_env["RAVEN_FRAMEWORK_DIR"]) + self.raiseAMessage("remoteRunCommand: "+str(remoteRunCommand)) + self.raiseADebug("COMMAND: "+command_env["COMMAND"]) + self.raiseADebug("RAVEN_FRAMEWORK_DIR: "+command_env["RAVEN_FRAMEWORK_DIR"]) remoteRunCommand["env"] = command_env - ## print out for debugging return remoteRunCommand def remoteRunCommand(self, runInfoDict): @@ -203,11 +208,14 @@ class cls. """ inputSpecification = InputData.parameterInputFactory("mode", ordered=False, contentType=InputTypes.StringType) inputSpecification.addSub(InputData.parameterInputFactory("runSbatch")) + inputSpecification.addSub(InputData.parameterInputFactory("nodefile", contentType=InputTypes.StringType)) + inputSpecification.addSub(InputData.parameterInputFactory("nodefileenv", contentType=InputTypes.StringType)) inputSpecification.addSub(InputData.parameterInputFactory("memory", contentType=InputTypes.StringType)) inputSpecification.addSub(InputData.parameterInputFactory("coresneeded", contentType=InputTypes.IntegerType)) inputSpecification.addSub(InputData.parameterInputFactory("partition", contentType=InputTypes.StringType)) inputSpecification.addSub(InputData.parameterInputFactory("MPIParam", contentType=InputTypes.StringType)) inputSpecification.addSub(InputData.parameterInputFactory("noprecommand")) + inputSpecification.addSub(InputData.parameterInputFactory("useSrun")) return inputSpecification def handleInput(self, paramInput): @@ -217,17 +225,29 @@ def handleInput(self, paramInput): @ Out, None """ for child in paramInput.subparts: - if child.getName() == "nodefile": + childName = child.getName().lower() + if childName == "nodefile": self.__nodeFile = child.value.strip() - elif child.getName() == "memory": + elif childName == "nodefileenv": + envName = child.value.strip() + if envName not in os.environ: + self.raiseAnError(IOError, f' environment variable "{envName}" ' + 'is not defined in the current environment!') + self.__nodeFile = os.environ[envName] + elif childName == "memory": self.__memNeeded = child.value.strip() - elif child.getName() == "coresneeded": + elif childName == "coresneeded": self.__coresNeeded = child.value - elif child.getName() == "partition": + elif childName == "partition": self.__partition = child.value.strip() - elif child.getName() == "runSbatch": + elif childName == "runsbatch": self.__runSbatch = True - elif child.getName() == "MPIParam": + elif childName == "mpiparam": self.__mpiparams.append(child.value.strip()) - elif child.getName() == "noPrecommand": + elif childName == "noprecommand": self.__createPrecommand = False + elif childName == "usesrun": + self.__useSrun = True + else: + self.raiseAWarning(f'Unrecognized option "{child.getName()}" ignored ' + 'by the Slurm simulation mode.') diff --git a/ravenframework/CustomModes/tests/test_cluster_utils.py b/ravenframework/CustomModes/tests/test_cluster_utils.py new file mode 100644 index 0000000000..a814d60634 --- /dev/null +++ b/ravenframework/CustomModes/tests/test_cluster_utils.py @@ -0,0 +1,224 @@ +# Copyright 2017 Battelle Energy Alliance, LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" + Unit tests for CustomModes/ClusterUtils.py (node discovery, batch sizing, + node-file splitting, precommand construction, job-name sanitation). + + These tests deliberately load ClusterUtils directly from its file path so + they can run WITHOUT a built RAVEN installation (no Crow, no framework + imports): + + python3 ravenframework/CustomModes/tests/test_cluster_utils.py +""" + +import importlib.util +import os +import shutil +import tempfile +import unittest + +HERE = os.path.dirname(os.path.abspath(__file__)) +MODULE_PATH = os.path.abspath(os.path.join(HERE, os.pardir, "ClusterUtils.py")) +spec = importlib.util.spec_from_file_location("ClusterUtils", MODULE_PATH) +ClusterUtils = importlib.util.module_from_spec(spec) +spec.loader.exec_module(ClusterUtils) + + +class FakeCompletedProcess: + """ Minimal stand-in for subprocess.CompletedProcess """ + def __init__(self, returncode=0, stdout="", stderr=""): + self.returncode = returncode + self.stdout = stdout + self.stderr = stderr + + +class TestParseSlurmTasksPerNode(unittest.TestCase): + """ Tests for parseSlurmTasksPerNode """ + + def testSimple(self): + self.assertEqual(ClusterUtils.parseSlurmTasksPerNode("4"), [4]) + + def testRepeated(self): + self.assertEqual(ClusterUtils.parseSlurmTasksPerNode("36(x2),20"), [36, 36, 20]) + + def testMixed(self): + self.assertEqual(ClusterUtils.parseSlurmTasksPerNode("2,3(x3),1"), [2, 3, 3, 3, 1]) + + def testNone(self): + self.assertEqual(ClusterUtils.parseSlurmTasksPerNode(None), []) + + def testInvalid(self): + with self.assertRaises(ValueError): + ClusterUtils.parseSlurmTasksPerNode("bogus(x)") + + +class TestSlurmNodeListFromScontrol(unittest.TestCase): + """ Tests for slurmNodeListFromScontrol with an injected fake scontrol """ + + def testExpansion(self): + def fakeRunner(cmd, capture_output, text, timeout): + self.assertEqual(cmd, ["scontrol", "show", "hostnames", "node[01-02]"]) + return FakeCompletedProcess(stdout="node01\nnode02\n") + lines = ClusterUtils.slurmNodeListFromScontrol(nodeList="node[01-02]", + tasksPerNode="2(x2)", + runner=fakeRunner) + self.assertEqual(lines, ["node01", "node01", "node02", "node02"]) + + def testFailureRaises(self): + def fakeRunner(cmd, capture_output, text, timeout): + return FakeCompletedProcess(returncode=1, stderr="boom") + with self.assertRaises(RuntimeError): + ClusterUtils.slurmNodeListFromScontrol(nodeList="nodeXX", + tasksPerNode="1", + runner=fakeRunner) + + def testNoNodeList(self): + # no node list available at all -> empty result, no crash + old = os.environ.pop("SLURM_JOB_NODELIST", None) + try: + self.assertEqual(ClusterUtils.slurmNodeListFromScontrol(nodeList=None, + tasksPerNode="1", + runner=None), []) + finally: + if old is not None: + os.environ["SLURM_JOB_NODELIST"] = old + + +class TestNodeFiles(unittest.TestCase): + """ Tests for readNodeFile / writeNodeSubFiles """ + + def setUp(self): + self.workDir = tempfile.mkdtemp() + + def tearDown(self): + shutil.rmtree(self.workDir, ignore_errors=True) + + def testReadNodeFileSkipsBlanks(self): + path = os.path.join(self.workDir, "nodes") + with open(path, "w") as f: + f.write("a\n\nb\n \nc\n") + self.assertEqual(ClusterUtils.readNodeFile(path), ["a", "b", "c"]) + + def testWriteNodeSubFiles(self): + # 6 processors, numMPI=2, batchSize=3 -> node_0..node_2 with 2 lines each + lines = ["n1", "n1", "n2", "n2", "n3", "n3"] + written = ClusterUtils.writeNodeSubFiles(lines, 3, 2, self.workDir) + self.assertEqual([os.path.basename(p) for p in written], + ["node_0", "node_1", "node_2"]) + self.assertEqual(ClusterUtils.readNodeFile(written[0]), ["n1", "n1"]) + self.assertEqual(ClusterUtils.readNodeFile(written[1]), ["n2", "n2"]) + self.assertEqual(ClusterUtils.readNodeFile(written[2]), ["n3", "n3"]) + + +class TestComputeBatchSize(unittest.TestCase): + """ Tests for computeBatchSize """ + + def testFits(self): + self.assertEqual(ClusterUtils.computeBatchSize(8, 2, 4), (4, False)) + + def testClamped(self): + self.assertEqual(ClusterUtils.computeBatchSize(8, 2, 10), (4, True)) + + def testAtLeastOne(self): + self.assertEqual(ClusterUtils.computeBatchSize(1, 4, 2), (1, True)) + + def testExactFit(self): + self.assertEqual(ClusterUtils.computeBatchSize(4, 4, 1), (1, False)) + + +class TestBuildMPIPrecommand(unittest.TestCase): + """ Tests for buildMPIPrecommand (matches the legacy string construction) """ + + def testDefault(self): + pre = ClusterUtils.buildMPIPrecommand("mpiexec", [], "-f /wd/nodes", 4, "") + self.assertEqual(pre, "mpiexec -f /wd/nodes -n 4 ") + + def testWithParamsAndExisting(self): + pre = ClusterUtils.buildMPIPrecommand("mpiexec", ["--bind-to core"], + "-f %BASE_WORKING_DIR%/node_%INDEX% ", + 2, "oldpre") + self.assertEqual(pre, "mpiexec --bind-to core -f %BASE_WORKING_DIR%/node_%INDEX% -n 2 oldpre") + + def testLocalMachine(self): + # not in a cluster: nodeCommand is a single space (legacy behavior) + pre = ClusterUtils.buildMPIPrecommand("mpiexec", [], " ", 2, "") + self.assertEqual(pre, "mpiexec -n 2 ") + + +class TestBuildSrunPrecommand(unittest.TestCase): + """ Tests for buildSrunPrecommand (Slurm-native launch, quick-ref #9) """ + + def testDefault(self): + pre = ClusterUtils.buildSrunPrecommand(4, [], "") + self.assertEqual(pre, "srun --overlap --exact -n 4 ") + + def testWithParamsAndExisting(self): + pre = ClusterUtils.buildSrunPrecommand(2, ["--mpi=pmix"], "oldpre") + self.assertEqual(pre, "srun --overlap --exact -n 2 --mpi=pmix oldpre") + + +class TestAssembleDaskJobqueueKwargs(unittest.TestCase): + """ Tests for assembleDaskJobqueueKwargs (quick-ref #10) """ + + def testSlurmDefaults(self): + runInfo = {'numProcByRun': 4, 'batchSize': 3, 'expectedTime': '2:00:00'} + config = {'scheduler': 'slurm', 'options': {'memory': '4GB'}} + name, kwargs, jobs = ClusterUtils.assembleDaskJobqueueKwargs(config, runInfo) + self.assertEqual(name, 'SLURMCluster') + self.assertEqual(kwargs, {'cores': 4, 'memory': '4GB', 'walltime': '2:00:00'}) + self.assertEqual(jobs, 3) + + def testPbsOverridesAndPassthrough(self): + runInfo = {'numProcByRun': 4, 'batchSize': 3, 'expectedTime': '2:00:00'} + config = {'scheduler': 'pbs', 'options': {'memory': '8GB', 'cores': '16', + 'jobs': '2', 'walltime': '0:30:00', + 'queue': 'short', 'account': 'proj1'}} + name, kwargs, jobs = ClusterUtils.assembleDaskJobqueueKwargs(config, runInfo) + self.assertEqual(name, 'PBSCluster') + self.assertEqual(jobs, 2) + self.assertEqual(kwargs, {'cores': 16, 'memory': '8GB', 'walltime': '0:30:00', + 'queue': 'short', 'account': 'proj1'}) + + def testMissingMemoryRaises(self): + with self.assertRaises(ValueError): + ClusterUtils.assembleDaskJobqueueKwargs({'scheduler': 'slurm', 'options': {}}, {}) + + def testUnknownSchedulerRaises(self): + with self.assertRaises(ValueError): + ClusterUtils.assembleDaskJobqueueKwargs({'scheduler': 'lsf', + 'options': {'memory': '1GB'}}, {}) + + +class TestSanitizeJobName(unittest.TestCase): + """ Tests for sanitizeJobName """ + + def testValid(self): + self.assertEqual(ClusterUtils.sanitizeJobName("my_job-1"), "my_job-1") + + def testInvalidRaises(self): + with self.assertRaises(ValueError): + ClusterUtils.sanitizeJobName("bad name!") + + def testTruncation(self): + # PBS behavior: 15-char limit -> first 10 + '-' + last 4 + name = "abcdefghijklmnopqrst" + self.assertEqual(ClusterUtils.sanitizeJobName(name, maxLength=15), + "abcdefghij-qrst") + + def testNoTruncationWhenShort(self): + self.assertEqual(ClusterUtils.sanitizeJobName("short", maxLength=15), "short") + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/ravenframework/JobHandler.py b/ravenframework/JobHandler.py index a78d7df56c..d3f21fcfc6 100644 --- a/ravenframework/JobHandler.py +++ b/ravenframework/JobHandler.py @@ -201,6 +201,8 @@ def __init__(self): self.remoteServers = None self.daskSchedulerFile = None self._daskScheduler = None + self._daskJobqueueCluster = None # dask_jobqueue cluster object (if is used) + self._headDaskWorker = None # Popen of the dask worker started on the head node (if any) def __getstate__(self): """ @@ -216,6 +218,8 @@ def __getstate__(self): #This will be reinitialized from a schedulerFile. if self._parallelLib == ParallelLibEnum.dask and '_server' in state: state.pop('_server') + # the dask_jobqueue cluster object is not picklable + state.pop('_daskJobqueueCluster', None) return state def __setstate__(self, d): @@ -225,6 +229,7 @@ def __setstate__(self, d): @ Out, None """ self.__dict__.update(d) + self.__dict__.setdefault('_daskJobqueueCluster', None) self.__queueLock = threading.RLock() # Reinitialize the per-job event registry (lost during pickling). # New events will be created when jobs are submitted via reAddJob(). @@ -318,6 +323,16 @@ def __checkAndRemoveFinished(self, running): # FIXME: The running.command was always internal now, so I removed it. # We should probably find a way to give more pertinent information. self.raiseAMessage(f" Process Failed {running.identifier}:{running} internal returnCode {returnCode}") + # surface the failure details (e.g. the remote/threaded traceback), if + # the runner recorded any, both in the log and in the failed-job metadata + failureInfo = getattr(running, 'getFailureInfo', lambda: None)() + if failureInfo: + self.raiseAMessage(f' Failure details for job "{running.identifier}":\n{failureInfo}') + if isinstance(metadataToKeep, dict): + metadataToKeep = dict(metadataToKeep) + metadataToKeep['failureInfo'] = failureInfo + elif metadataToKeep is None: + metadataToKeep = {'failureInfo': failureInfo} self.__failedJobs[running.identifier]=(returnCode,copy.deepcopy(metadataToKeep)) def __initializeDistributed(self): @@ -347,6 +362,11 @@ def __initializeDistributed(self): # is ray instanciated outside? self.rayInstanciatedOutside = 'headNode' in self.runInfoDict self.daskInstanciatedOutside = 'schedulerFile' in self.runInfoDict + # dask-jobqueue managed cluster (workers submitted as scheduler jobs)? + if self._parallelLib == ParallelLibEnum.dask and self.runInfoDict.get('daskJobqueue'): + self.__initializeDaskJobqueue(self.runInfoDict['daskJobqueue']) + self.__isDistributedInitialized = True + return if len(self.runInfoDict['Nodes']) > 0 or self.rayInstanciatedOutside or self.daskInstanciatedOutside: availableNodes = [nodeId.strip() for nodeId in self.runInfoDict['Nodes']] uniqueN = list(set(availableNodes)) @@ -424,12 +444,14 @@ def __initializeDistributed(self): else: self.raiseAWarning("parallellib creation not handled") if self._parallelLib == ParallelLibEnum.ray: - self.raiseADebug("Head node IP address: ", self._server.address_info['node_ip_address']) - self.raiseADebug("Redis address : ", self._server.address_info['redis_address']) - self.raiseADebug("Object store address: ", self._server.address_info['object_store_address']) - self.raiseADebug("Raylet socket name : ", self._server.address_info['raylet_socket_name']) - self.raiseADebug("Session directory : ", self._server.address_info['session_dir']) - self.raiseADebug("GCS Address : ", self._server.address_info['gcs_address']) + # use .get: some keys (e.g. redis_address) were removed in Ray 2.x + addressInfo = getattr(self._server, 'address_info', {}) or {} + self.raiseADebug("Head node IP address: ", addressInfo.get('node_ip_address', 'N/A')) + self.raiseADebug("Redis address : ", addressInfo.get('redis_address', 'N/A')) + self.raiseADebug("Object store address: ", addressInfo.get('object_store_address', 'N/A')) + self.raiseADebug("Raylet socket name : ", addressInfo.get('raylet_socket_name', 'N/A')) + self.raiseADebug("Session directory : ", addressInfo.get('session_dir', 'N/A')) + self.raiseADebug("GCS Address : ", addressInfo.get('gcs_address', 'N/A')) if servers: self.raiseADebug("# of remote servers : ", str(len(servers))) self.raiseADebug("Remote servers : ", " , ".join(servers)) @@ -442,6 +464,38 @@ def __initializeDistributed(self): # ray or dask is initialized self.__isDistributedInitialized = True + def __initializeDaskJobqueue(self, config): + """ + Initializes a dask-jobqueue managed cluster: the Dask scheduler runs + locally and the workers are submitted AS scheduler jobs (sbatch/qsub) by + dask_jobqueue, so no inter-node ssh and no hand-rolled bring-up scripts + are needed. Configured via the RunInfo element. + @ In, config, dict, with 'scheduler' ("slurm"/"pbs") and 'options' + (dask_jobqueue constructor options; see ClusterUtils.assembleDaskJobqueueKwargs) + @ Out, None + """ + from ravenframework.CustomModes import ClusterUtils + try: + import dask_jobqueue + except ImportError: + self.raiseAnError(RuntimeError, 'The option requires the ' + '"dask_jobqueue" package (e.g. pip install dask-jobqueue), ' + 'which could not be imported!') + try: + clusterClassName, kwargs, jobs = ClusterUtils.assembleDaskJobqueueKwargs(config, self.runInfoDict) + except ValueError as err: + self.raiseAnError(IOError, str(err)) + clusterClass = getattr(dask_jobqueue, clusterClassName) + self.raiseAMessage(f'Starting dask_jobqueue.{clusterClassName} with options {kwargs}, ' + f'scaling to {jobs} scheduler job(s)') + self._daskJobqueueCluster = clusterClass(**kwargs) + self._daskJobqueueCluster.scale(jobs=jobs) + self._server = dask.distributed.Client(self._daskJobqueueCluster) + # RAVEN owns this cluster (teardown happens in __shutdownParallel) + self.daskInstanciatedOutside = False + self.raiseADebug('dask-jobqueue dashboard: ' + +str(getattr(self._daskJobqueueCluster, 'dashboard_link', 'N/A'))) + def __getLocalAndRemoteMachineNames(self): """ Method to get the qualified host and remote nodes' names @@ -492,10 +546,39 @@ def __shutdownParallel(self): rayTerminate.wait() if rayTerminate.returncode != 0: self.raiseAWarning("RAY FAILED TO TERMINATE ON NODE: "+nodeAddress) - elif self._parallelLib == ParallelLibEnum.dask and self._server is not None and not self.rayInstanciatedOutside: - self._server.close() - if self._daskScheduler is not None: - self._daskScheduler.terminate() + elif self._parallelLib == ParallelLibEnum.dask and self._server is not None: + if self._daskJobqueueCluster is not None: + # dask-jobqueue managed cluster: closing the cluster cancels the + # worker scheduler jobs (sbatch/qsub) and stops the scheduler + try: + self._server.close() + finally: + try: + self._daskJobqueueCluster.close() + except Exception as exc: + self.raiseAWarning("dask-jobqueue cluster close raised: "+repr(exc)) + self._daskJobqueueCluster = None + elif not self.daskInstanciatedOutside: + # We own this cluster: shut down the scheduler and ALL (local and + # remote) workers, not just the client connection. Client.shutdown() + # asks the scheduler to retire every worker before closing, which + # prevents zombie "dask worker" processes on remote nodes. + try: + self._server.shutdown() + except Exception as exc: + self.raiseAWarning("Dask cluster shutdown raised: "+repr(exc)) + try: + self._server.close() + except Exception: + pass + if self._headDaskWorker is not None and self._headDaskWorker.poll() is None: + self._headDaskWorker.terminate() + if self._daskScheduler is not None and self._daskScheduler.poll() is None: + self._daskScheduler.terminate() + else: + # Externally managed cluster: only disconnect our client; leave the + # scheduler and workers to their owner. + self._server.close() def __runHeadNode(self, nProcs, port=None): """ @@ -514,10 +597,9 @@ def __runHeadNode(self, nProcs, port=None): command.append("--num-cpus="+str(nProcs)) if port is not None: command.append("--port="+str(port)) - outFile = open("ray_head.ip", 'w') - rayStart = utils.pickleSafeSubprocessPopen(command,shell=False,stdout=outFile, stderr=outFile, env=localEnv) - rayStart.wait() - outFile.close() + with open("ray_head.ip", 'w') as outFile: + rayStart = utils.pickleSafeSubprocessPopen(command,shell=False,stdout=outFile, stderr=outFile, env=localEnv) + rayStart.wait() if rayStart.returncode != 0: self.raiseAnError(RuntimeError, f"RAY failed to start on the --head node! Return code is {rayStart.returncode}") else: @@ -559,13 +641,15 @@ def __runHeadNode(self, nProcs, port=None): if succeeded: #do equivelent of dask worker start in start_dask.sh: # dask worker --nworkers $NUM_CPUS --scheduler-file $SCHEDULER_FILE >> $OUTFILE - outFile = open(os.path.join(self.runInfoDict['WorkingDir'], - "server_debug_"+self.__getLocalHost()),'w') command = ["dask","worker","--scheduler-file",self.daskSchedulerFile] if nProcs is not None: command.extend(("--nworkers",str(nProcs))) - headDaskWorker = utils.pickleSafeSubprocessPopen(command,shell=False, - stdout=outFile, stderr=outFile, env=localEnv) + # the subprocess duplicates the file descriptor, so the Python-side + # handle can (and should) be closed right after spawning + with open(os.path.join(self.runInfoDict['WorkingDir'], + "server_debug_"+self.__getLocalHost()),'w') as outFile: + self._headDaskWorker = utils.pickleSafeSubprocessPopen(command,shell=False, + stdout=outFile, stderr=outFile, env=localEnv) return address def __getRayInfoFromStart(self, rayLog): @@ -1213,7 +1297,7 @@ def fillJobQueue(self): if infoKey in self.runInfoDict: kwargs[infoKey] = self.runInfoDict[infoKey] kwargs['INDEX'] = str(i) - kwargs['INDEX1'] = str(i+i) + kwargs['INDEX1'] = str(i+1) kwargs['CURRENT_ID'] = str(self.__nextId) kwargs['CURRENT_ID1'] = str(self.__nextId+1) kwargs['SCRIPT_DIR'] = self.runInfoDict['ScriptDir'] diff --git a/ravenframework/RemoteNodeScripts/start_dask.sh b/ravenframework/RemoteNodeScripts/start_dask.sh index eb53c6d576..8e97e50772 100755 --- a/ravenframework/RemoteNodeScripts/start_dask.sh +++ b/ravenframework/RemoteNodeScripts/start_dask.sh @@ -10,13 +10,17 @@ WORKING_DIR=$6 PYTHONPATH=$7 export PYTHONPATH -echo starting >> $OUTFILE +echo starting >> "$OUTFILE" -cd $WORKING_DIR -source $REMOTE_BASH >> $OUTFILE 2>&1 +cd "$WORKING_DIR" || { echo "cannot cd to $WORKING_DIR" >> "$OUTFILE"; exit 1; } +# the remote bash profile is optional +if [ -n "$REMOTE_BASH" ] && [ -f "$REMOTE_BASH" ] + then + source "$REMOTE_BASH" >> "$OUTFILE" 2>&1 +fi -which dask >> $OUTFILE 2>&1 -hostname >> $OUTFILE -echo PYTHONPATH $PYTHONPATH >> $OUTFILE +which dask >> "$OUTFILE" 2>&1 +hostname >> "$OUTFILE" +echo PYTHONPATH "$PYTHONPATH" >> "$OUTFILE" -dask worker --nworkers $NUM_CPUS --scheduler-file $SCHEDULER_FILE >> $OUTFILE 2>&1 & +dask worker --nworkers "$NUM_CPUS" --scheduler-file "$SCHEDULER_FILE" >> "$OUTFILE" 2>&1 & diff --git a/ravenframework/RemoteNodeScripts/start_ray.sh b/ravenframework/RemoteNodeScripts/start_ray.sh index 74ade05aad..3329251ff2 100755 --- a/ravenframework/RemoteNodeScripts/start_ray.sh +++ b/ravenframework/RemoteNodeScripts/start_ray.sh @@ -21,24 +21,31 @@ # start ray # OUTPUT FILE FOR LOGGING OUTFILE=$1 -# BASH_PROFILE +# HEAD NODE ADDRESS HEAD_ADDRESS=$2 NUM_CPUS=$3 -RAVEN_FRAMEWORK_DIR=$4 +NUM_GPUS=$4 +RAVEN_FRAMEWORK_DIR=$5 -echo starting >> $OUTFILE +echo starting >> "$OUTFILE" -if [ $# -eq 5 ] +if [ $# -ge 6 ] && [ -n "$6" ] then - REMOTE_BASH=$5 - source $REMOTE_BASH >> $OUTFILE 2>&1 + REMOTE_BASH=$6 + source "$REMOTE_BASH" >> "$OUTFILE" 2>&1 fi -which ray >> $OUTFILE 2>&1 -hostname >> $OUTFILE +which ray >> "$OUTFILE" 2>&1 +hostname >> "$OUTFILE" -echo loaded >> $OUTFILE -command -v ray >> $OUTFILE 2>&1 +echo loaded >> "$OUTFILE" +command -v ray >> "$OUTFILE" 2>&1 mkdir -p /tmp/ray -echo ray start --verbose --address=$HEAD_ADDRESS --num-cpus $NUM_CPUS >> $OUTFILE 2>&1 -ray start --verbose --address=$HEAD_ADDRESS --num-cpus $NUM_CPUS >> $OUTFILE 2>&1 +# only pass --num-gpus when a non-negative count was requested +GPU_ARGS=() +if [ -n "$NUM_GPUS" ] && [ "$NUM_GPUS" -ge 0 ] 2>/dev/null + then + GPU_ARGS=(--num-gpus "$NUM_GPUS") +fi +echo ray start --verbose --address="$HEAD_ADDRESS" --num-cpus "$NUM_CPUS" "${GPU_ARGS[@]}" >> "$OUTFILE" 2>&1 +ray start --verbose --address="$HEAD_ADDRESS" --num-cpus "$NUM_CPUS" "${GPU_ARGS[@]}" >> "$OUTFILE" 2>&1 diff --git a/ravenframework/RemoteNodeScripts/start_remote_servers.sh b/ravenframework/RemoteNodeScripts/start_remote_servers.sh index 3397ce153e..cbc8edd813 100755 --- a/ravenframework/RemoteNodeScripts/start_remote_servers.sh +++ b/ravenframework/RemoteNodeScripts/start_remote_servers.sh @@ -92,7 +92,7 @@ do case "$1" in --help) display_usage - return + exit 0 ;; --remote-node-address) shift @@ -134,25 +134,25 @@ echo $REMOTE_ADDRESS if [[ "$REMOTE_ADDRESS" == "" ]]; then echo ... ERROR: --remote-node-address argument must be inputted ! - exit + exit 1 fi if [[ "$HEAD_ADDRESS" == "" ]]; then echo ... ERROR: --address argument must be inputted ! - exit + exit 1 fi if [[ "$PYTHONPATH" == "" ]]; then echo ... ERROR: --python-path argument must be inputted ! - exit + exit 1 fi if [[ "$WORKINGDIR" == "" ]]; then echo ... ERROR: --working-dir argument must be inputted ! - exit + exit 1 fi echo RAVEN_FRAMEWORK_DIR $RAVEN_FRAMEWORK_DIR @@ -160,10 +160,13 @@ echo RAVEN_FRAMEWORK_DIR $RAVEN_FRAMEWORK_DIR # ssh in the remote node and run the ray servers CWD=`pwd` OUTPUT=$CWD/server_debug_$REMOTE_ADDRESS +START_OUTPUT=${OUTPUT}_start.log +# NUM_GPUS is forwarded to start_ray.sh (it is ignored there when negative); +# REMOTE_BASH is passed as an optional trailing argument if [[ "$REMOTE_BASH" == "" ]]; then - ssh $REMOTE_ADDRESS $ECE_SCRIPT_DIR/server_start.py ${WORKINGDIR} ${OUTPUT} ${PYTHONPATH} "${ECE_SCRIPT_DIR}/start_ray.sh $OUTPUT $HEAD_ADDRESS $NUM_CPUS $RAVEN_FRAMEWORK_DIR" 2>&1 | tee $START_OUTPUT + ssh "$REMOTE_ADDRESS" "$ECE_SCRIPT_DIR/server_start.py" "${WORKINGDIR}" "${OUTPUT}" "${PYTHONPATH}" "${ECE_SCRIPT_DIR}/start_ray.sh $OUTPUT $HEAD_ADDRESS $NUM_CPUS $NUM_GPUS $RAVEN_FRAMEWORK_DIR" 2>&1 | tee "$START_OUTPUT" else - ssh $REMOTE_ADDRESS $ECE_SCRIPT_DIR/server_start.py ${WORKINGDIR} ${OUTPUT} ${PYTHONPATH} "${ECE_SCRIPT_DIR}/start_ray.sh $OUTPUT $HEAD_ADDRESS $NUM_CPUS $RAVEN_FRAMEWORK_DIR $REMOTE_BASH" 2>&1 | tee $START_OUTPUT + ssh "$REMOTE_ADDRESS" "$ECE_SCRIPT_DIR/server_start.py" "${WORKINGDIR}" "${OUTPUT}" "${PYTHONPATH}" "${ECE_SCRIPT_DIR}/start_ray.sh $OUTPUT $HEAD_ADDRESS $NUM_CPUS $NUM_GPUS $RAVEN_FRAMEWORK_DIR $REMOTE_BASH" 2>&1 | tee "$START_OUTPUT" fi diff --git a/ravenframework/Runners/DaskRunner.py b/ravenframework/Runners/DaskRunner.py index 765d9efe28..f37918f6e3 100644 --- a/ravenframework/Runners/DaskRunner.py +++ b/ravenframework/Runners/DaskRunner.py @@ -18,6 +18,7 @@ """ #External Modules------------------------------------------------------------------------------------ import sys +import traceback import gc import copy import threading @@ -136,15 +137,17 @@ def _collectRunnerResponse(self): with self.__funcLock: if not self.hasBeenAdded: if self.__func is not None: - #if the function threw an exception, result will rethrow it. + #if the function threw an exception, result() will rethrow it here. try: self.runReturn = self.__func.result() + self.runSucceeded = True except Exception as ae: self.runReturn = None self.hasBeenAdded = True self.returnCode = -1 - self.raiseAWarning(self.__class__.__name__ + " job "+self.identifier+" failed with error:"+ str(ae) +" !",'ExceptedErrorInCollect') - raise ae + self.runSucceeded = False + self.failureInfo = traceback.format_exc() + self.raiseAWarning(self.__class__.__name__ + " job "+self.identifier+" failed with error:"+ str(ae) +" !\n"+self.failureInfo,'ExceptedErrorInCollect') else: self.runReturn = None self.hasBeenAdded = True @@ -163,11 +166,9 @@ def start(self): return except Exception as ae: - #Uncomment if you need the traceback self.exceptionTrace = sys.exc_info() - #exc_type, exc_value, exc_traceback = sys.exc_info() - #import traceback - #traceback.print_exception(exc_type, exc_value, exc_traceback) + self.failureInfo = traceback.format_exc() + self.runSucceeded = False self.raiseAWarning(self.__class__.__name__ + " job "+self.identifier+" failed with error:"+ str(ae) +" !",'ExceptedErrorInStart') self.returnCode = -1 @@ -178,6 +179,14 @@ def kill(self): @ Out, None """ with self.__funcLock: + if self.__func is not None: + # actually cancel the task on the cluster; simply dropping the Future + # leaves the task running and consuming cluster resources + try: + self.__func.cancel() + except Exception as exc: + self.raiseAWarning('Unable to cancel dask future for job "' + +self.identifier+'": '+repr(exc)) del self.__func self.__func = None self.returnCode = -1 diff --git a/ravenframework/Runners/InternalRunner.py b/ravenframework/Runners/InternalRunner.py index e0e51072a6..781089ce07 100644 --- a/ravenframework/Runners/InternalRunner.py +++ b/ravenframework/Runners/InternalRunner.py @@ -48,6 +48,12 @@ def __init__(self, args, functionToRun, **kwargs): self.hasBeenAdded = False self.returnCode = 0 self.exceptionTrace = None # sys.exc_info() if an error occurred while running + self.runSucceeded = None # tri-state: None (unknown/legacy), True (function + # completed without exception), False (function failed). + # Allows functions legitimately returning None to be + # distinguished from failed evaluations. + self.failureInfo = None # human-readable failure description (e.g. formatted + # traceback of the remote/threaded exception) ## These things cannot be deep copied self.skipOnCopy = ['functionToRun','thread','__queueLock', '_InternalRunner__queueLock'] @@ -94,9 +100,25 @@ def getEvaluation(self): """ if self.isDone(): self._collectRunnerResponse() + if self.runSucceeded is False: + self.returnCode = -1 + return Error() if self.runReturn is None: + # a None return means either no outcome was recorded, or the wrapped + # function itself returned None as its failure signal without raising + # (e.g. Models.Code.evaluateSample on a non-zero process return code); + # in both cases, treat the run as failed self.returnCode = -1 return Error() return self.runReturn else: return Error() + + def getFailureInfo(self): + """ + Returns a human-readable description of the failure (e.g. the formatted + traceback raised by the evaluated function), if any. + @ In, None + @ Out, failureInfo, str or None, the failure description + """ + return self.failureInfo diff --git a/ravenframework/Runners/RayRunner.py b/ravenframework/Runners/RayRunner.py index 452c7969f2..fb8beea973 100644 --- a/ravenframework/Runners/RayRunner.py +++ b/ravenframework/Runners/RayRunner.py @@ -18,6 +18,7 @@ """ #External Modules------------------------------------------------------------------------------------ import sys +import traceback import gc import copy import threading @@ -102,7 +103,11 @@ def isDone(self): runReturn = ray.get(self.__func, timeout=waitTimeOut) self.runReturn = runReturn self.hasBeenAdded = True + self.runSucceeded = True if self.runReturn is None: + # the function returned None as its own failure signal without + # raising (e.g. Models.Code.evaluateSample on a non-zero process + # return code); treat as failed self.returnCode = -1 return True except ray.exceptions.GetTimeoutError: @@ -113,6 +118,10 @@ def isDone(self): # I assume it means the task has unfixably died, # and so is done, and set return code to failed. self.raiseAWarning("RayTaskError: "+str(rte)) + self.failureInfo = str(rte) + self.runSucceeded = False + self.hasBeenAdded = True + self.runReturn = None self.returnCode = -1 return True #Alternative that was tried: @@ -129,7 +138,17 @@ def _collectRunnerResponse(self): with self.__funcLock: if not self.hasBeenAdded: if self.__func is not None: - self.runReturn = ray.get(self.__func) + try: + self.runReturn = ray.get(self.__func) + self.runSucceeded = True + if self.runReturn is None: + self.returnCode = -1 + except ray.exceptions.RayTaskError as rte: + self.raiseAWarning("RayTaskError: "+str(rte)) + self.failureInfo = str(rte) + self.runSucceeded = False + self.runReturn = None + self.returnCode = -1 else: self.runReturn = None self.hasBeenAdded = True @@ -148,11 +167,9 @@ def start(self): return except Exception as ae: - #Uncomment if you need the traceback self.exceptionTrace = sys.exc_info() - #exc_type, exc_value, exc_traceback = sys.exc_info() - #import traceback - #traceback.print_exception(exc_type, exc_value, exc_traceback) + self.failureInfo = traceback.format_exc() + self.runSucceeded = False self.raiseAWarning(self.__class__.__name__ + " job "+self.identifier+" failed with error:"+ str(ae) +" !",'ExceptedError') self.returnCode = -1 @@ -163,6 +180,14 @@ def kill(self): @ Out, None """ with self.__funcLock: + if self.__func is not None: + # actually cancel the remote task; simply dropping the ObjectRef + # leaves the task running and consuming cluster resources + try: + ray.cancel(self.__func, force=True, recursive=True) + except Exception as exc: + self.raiseAWarning('Unable to cancel remote ray task for job "' + +self.identifier+'": '+repr(exc)) del self.__func self.__func = None self.returnCode = -1 diff --git a/ravenframework/Runners/SharedMemoryRunner.py b/ravenframework/Runners/SharedMemoryRunner.py index cb8e3e354f..d7c54488bc 100644 --- a/ravenframework/Runners/SharedMemoryRunner.py +++ b/ravenframework/Runners/SharedMemoryRunner.py @@ -25,6 +25,7 @@ import ctypes import inspect import threading +import traceback #External Modules End-------------------------------------------------------------------------------- @@ -82,9 +83,13 @@ def getReturnCode(self): """ if not self.hasBeenAdded: self._collectRunnerResponse() - ## Is this necessary and sufficient for all failed runs? - if len(self.subque) == 0 and self.runReturn is None: - self.runReturn = None + if self.runSucceeded is False: + self.returnCode = -1 + elif self.runReturn is None: + ## Either the wrapper never recorded an outcome (e.g. the thread was + ## killed before completing), or the wrapped function returned None as + ## its own failure signal without raising (e.g. Models.Code.evaluateSample + ## on a non-zero process return code). Treat as failed either way. self.returnCode = -1 return self.returnCode @@ -111,8 +116,30 @@ def start(self): @ In, None @ Out, None """ + def _runFunction(q, *arg): + """ + Thread target: runs the function, capturing exceptions and recording + the outcome explicitly (so that a legitimate None return value is not + mistaken for a failure, and the traceback is preserved). + @ In, q, collections.deque, queue collecting the result + @ In, arg, tuple, arguments for the function + @ Out, None + """ + try: + result = self.functionToRun(*arg) + except Exception: + self.exceptionTrace = sys.exc_info() + self.failureInfo = traceback.format_exc() + self.runSucceeded = False + self.returnCode = -1 + self.raiseAWarning(self.__class__.__name__ + " job "+self.identifier + +" failed with error:\n"+self.failureInfo, 'ExceptedError') + return + q.append(result) + self.runSucceeded = True + try: - self.thread = InterruptibleThread(target = lambda q, *arg : q.append(self.functionToRun(*arg)), + self.thread = InterruptibleThread(target = _runFunction, name = self.identifier, args=(self.subque,) + tuple(self.args)) @@ -122,6 +149,8 @@ def start(self): self.started = True except Exception as ae: self.exceptionTrace = sys.exc_info() + self.failureInfo = traceback.format_exc() + self.runSucceeded = False self.raiseAWarning(self.__class__.__name__ + " job "+self.identifier+" failed with error:"+ str(ae) +" !",'ExceptedError') self.returnCode = -1 @@ -133,9 +162,25 @@ def kill(self): """ if self.thread is not None: self.raiseADebug('Terminating job thread "{}" and RAVEN identifier "{}"'.format(self.thread.ident, self.identifier)) - while self.thread is not None and self.thread.is_alive(): - time.sleep(0.1) + ## NOTE: raising an asynchronous exception in a thread is inherently + ## unreliable: it is silently ignored while the thread is blocked inside + ## C extension code or system calls (exactly where external models spend + ## most of their time). The previous unbounded loop could therefore spin + ## forever. Bound the attempts and, if the thread will not die, warn and + ## move on: the thread is a daemon, so it cannot keep the process alive. + killTimeout = 10.0 # seconds + waited = 0.0 + while self.thread is not None and self.thread.is_alive() and waited < killTimeout: self.thread.kill() + time.sleep(0.1) + waited += 0.1 + if self.thread is not None and self.thread.is_alive(): + self.raiseAWarning('Job thread "{}" (RAVEN identifier "{}") did not terminate ' + 'within {} s; it is likely blocked in native code. ' + 'Abandoning it as a daemon thread.'.format( + self.thread.ident, self.identifier, killTimeout)) + self.runSucceeded = False + self.returnCode = -1 self.trackTime('runner_killed') ## The following code is extracted from stack overflow with some minor cosmetic diff --git a/ravenframework/Simulation.py b/ravenframework/Simulation.py index 4d8ebdea5c..6b7f56fdd0 100644 --- a/ravenframework/Simulation.py +++ b/ravenframework/Simulation.py @@ -739,6 +739,17 @@ def __readRunInfo(self, xmlNode, runInfoSkip, xmlFilename): self.runInfoDict['remoteNodes'] = [el.strip() for el in element.text.strip().split(',')] elif element.tag == 'schedulerFile': self.runInfoDict['schedulerFile'] = element.text.strip() + elif element.tag == 'daskJobqueue': + # dask-jobqueue managed cluster: workers are submitted as scheduler + # jobs (SLURMCluster/PBSCluster). The element text selects the + # scheduler ("slurm" or "pbs"); the XML attributes are passed to the + # dask_jobqueue cluster constructor (memory is required; cores, jobs, + # queue, account, walltime, interface, ... are optional). + # Example: slurm + self.runInfoDict['daskJobqueue'] = { + 'scheduler': element.text.strip().lower() if element.text is not None else '', + 'options': dict(element.attrib), + } elif element.tag == 'PYTHONPATH': self.runInfoDict['UPDATE_PYTHONPATH'] = element.text.strip() elif element.tag == 'delSucLogFiles' :