% SlurmUtils - A collection of utilities for SlurmUtils schedulers/clusters.

% Copyright 2017-2020 The MathWorks, Inc.

classdef (Hidden, Sealed) SlurmUtils
    
    methods (Hidden, Static)
        
        function clusterName = getClusterName()
            import parallel.internal.apishared.ThirdPartySchedulerUtils
            
            % Call sinfo to try to get information out about the cluster
            [failed, out] = system('sinfo --clusters all --noheader --format ""');
            
            if failed
                % On some Slurm configurations the command
                % 'sinfo --clusters all' will fail, but a plain 'sinfo'
                % will succeed.
                [failed, out] = system('sinfo');
                
                if failed
                    % Is sinfo on the path?
                    sinfoIsOnPath = ThirdPartySchedulerUtils.isExeOnPath('sinfo');
                    if sinfoIsOnPath
                        % Error out with the output from sinfo.
                        error(message('parallel:cluster:SlurmCommandLineUtilityFailed', 'sinfo', out));
                    end
                    % Otherwise error saying that sinfo is not on the path.
                    error(message('parallel:cluster:SlurmUnableToFindService'));
                end
                
                % If we get here the plain sinfo command succeeded but we
                % cannot determine the Cluster's name so just return an empty
                % char.
                clusterName = '';
                return
            end
            
            % Parse the output from out to try to find the cluster name
            newLineIndices = regexp(out, '\n') - 1;
            clusterSearchStr = 'CLUSTER: ';
            
            startClusterIndex = regexp(out, clusterSearchStr) + numel(clusterSearchStr);
            if numel(startClusterIndex) ~= 1
                % Either sinfo has returned information on more than one
                % cluster, or has not returned the cluster name. Either way
                % we cannot determine the cluster name so return an empty
                % string.
                clusterName = '';
            else
                endClusterIndex = min(newLineIndices(newLineIndices > startClusterIndex));
                clusterName = out(startClusterIndex:endClusterIndex);
            end
        end
        
        function tf = canUseJobArrays()
            % Job arrays are only supported in Slurm version 2.6.0 and
            % later
            [failed, out] = system('sinfo --version');
            if failed
                warning(message('parallel:cluster:SlurmCommandLineUtilityFailed', 'sinfo', out));
                tf = false;
                return
            end
            
            % Expect output to be of the form 'slurm 2.3.4'
            versionNumber = strsplit(out(7:end), '.');
            if str2double(versionNumber{1}) == 1
                tf = false;
                return
            end
            if str2double(versionNumber{1}) == 2
                if str2double(versionNumber{2}) < 6
                    tf = false;
                    return
                end
            end
            tf = true;
        end
        
        function [useArrays, maxJobArraySize, maxJobCount] = getSlurmConfig()
            % Checks the Slurm version and configuration to determine
            % whether job arrays can be used, how large the arrays can be
            % and the maximum number of jobs that the queue can hold
            
            % Check this version supports arrays
            useArrays = parallel.internal.apishared.SlurmUtils.canUseJobArrays();
            
            % Look up the max job array size and queue capacity
            % If we can't find max array size, assume the default
            % If we can't find max job count, just assume we won't hit it
            [maxJobArraySize, maxJobCount] = parallel.internal.apishared.ThirdPartySchedulerUtils.querySchedulerConfig( ...
                'scontrol show config', ...
                {'MaxArraySize', 'MaxJobCount'}, ...
                {1001, inf} );
            
            % scontrol outputs exclusive upper bound, we want inclusive
            maxJobArraySize = maxJobArraySize - 1;
            
            % Only use job arrays if we can get sufficient use out of them.
            % Due to the solution for g1934770, the MaxArraySize needs to
            % be at least 10 for job arrays to be useful as we loop the
            % Slurm array index between 0 and the largest power of 10 less
            % than MaxArraySize.
            % See parallel.cluster.Slurm/hSubmitIndependentJob
            useArrays = useArrays && (maxJobArraySize > 9);
        end
        
        function slurmID = callSbatch(cmdLine)
            
            import parallel.internal.apishared.SlurmUtils
            
            storedEnv = distcomp.pClearEnvironmentBeforeSubmission();
            cleanup = onCleanup(@()distcomp.pRestoreEnvironmentAfterSubmission(storedEnv));
            
            [FAILED, out] = system(cmdLine);
            if FAILED
                error(message('parallel:cluster:SlurmCommandLineUtilityFailed', 'sbatch', out));
            end
            slurmID = SlurmUtils.extractSlurmJobId(out);
        end
        
        function slurmID = extractSlurmJobId(sbatchCommandOutput)
            % Output from sbatch expected to be in the following format:
            %   'Submitted batch job 12345'
            %
            % sbatch could also attach a warning to the output, such as:
            %
            %   'sbatch: Warning: can't run 1 processes on 3 nodes, setting nnodes to 1
            %    Submitted batch job 12346'
            
            % Trim sbatch command output for use in debug message
            trimmedCommandOutput = strtrim(sbatchCommandOutput);
            
            % Ignore anything before or after 'Submitted batch job ###', and extract the numeric value.
            searchPattern = '.*Submitted batch job (\d+).*';
            
            % When we match searchPattern, matchedTokens is a single entry cell array containing the jobID.
            % Otherwise we failed to match searchPattern, so matchedTokens is an empty cell array.
            matchedTokens = regexp(sbatchCommandOutput, searchPattern, 'tokens', 'once');
            
            if isempty(matchedTokens)
                % Callers check for error in extracting Job ID using isempty() on return value.
                dctSchedulerMessage(0, '%s: Failed to extract Job ID from sbatch output: \n\t%s', mfilename, trimmedCommandOutput);
                error(message('parallel:cluster:SlurmFailedToParseJobId', sbatchCommandOutput));
            else
                slurmID = matchedTokens{1};
                dctSchedulerMessage(6, '%s: Job ID %s was extracted from sbatch output: \n\t%s', mfilename, slurmID, trimmedCommandOutput);
            end
        end
        
        function callScontrolRelease(slurmIDs, optTaskIdString)
            % Removes the hold on the job so it can run.
            
            import parallel.internal.apishared.SlurmUtils
            
            if nargin == 2 && ~isempty(optTaskIdString)
                arg = sprintf('%s_[%s]', SlurmUtils.slurmIDsToString(slurmIDs), optTaskIdString);
            else
                arg = sprintf('%s', SlurmUtils.slurmIDsToString(slurmIDs));
            end
            
            [FAILED, out] = system(sprintf('scontrol release "%s"', arg));
            if FAILED
                error(message('parallel:cluster:SlurmCommandLineUtilityFailed', 'scontrol', out));
            end
        end
        
        function [failed, out] = callSqueue(slurmIDsStr, noHeader)
            import parallel.internal.apishared.SlurmUtils
            
            cmd = sprintf('squeue -j %s --states=all --Format=jobarrayid,state --noheader --array', slurmIDsStr);
            if noHeader
                cmd = [cmd, ' --noheader'];
            end
            [failed, out] = system(cmd);
            
            if failed
                % Do not warn about squeue's failure if it was because it
                % could not find the job.
                if ~contains(out, 'slurm_load_jobs error: Invalid job id specified')
                    warning(message('parallel:cluster:SlurmCommandLineUtilityFailed', 'squeue', out));
                end
            end
        end
        
        function [failed, out] = callSacct(slurmIDsStr, noHeader)
            import parallel.internal.apishared.SlurmUtils
            
            % --allocations shows cumulative statistics for each job, not
            % the intermediate steps.
            % --format=jobid%-50 prints 50 characters of jobid left
            % justified.
            cmd = sprintf('sacct --allocations --format=jobid%%-50,state --jobs %s', slurmIDsStr);
            if noHeader
                cmd = [cmd, ' --noheader'];
            end
            [failed, out] = system(cmd);
            
            if failed
                warning(message('parallel:cluster:SlurmCommandLineUtilityFailed', 'sacct', out));
            end
        end
        
        function [failed, out] = getSlurmJobStatusOutput(slurmIDsStr, expectedNumOutputs, noHeader)
            import parallel.internal.apishared.SlurmUtils
            
            % Neither 'squeue' nor 'sacct' on their own are a reliable way
            % of finding the state of a job: There is a a short period of
            % time after a job has been submitted that it will not show up
            % in sacct, whilst squeue does not return information on old
            % jobs that have completed.  So we shall try 'squeue' first and
            % if that does not return the expected number of outputs, then
            % we shall try 'sacct'.
            
            [failed, out] = SlurmUtils.callSqueue(slurmIDsStr, noHeader);
            
            numOutputs = numel(splitlines(strtrim(out)));
            
            if ~noHeader
                % One of the lines will be the header, so take away 1 from
                % the total.
                numOutputs = numOutputs - 1;
            end
            
            if failed || numOutputs < expectedNumOutputs
                % Try sacct instead.
                [failed, out] = SlurmUtils.callSacct(slurmIDsStr, noHeader);
            end
        end
        
        function [failedCommand, numPending, numRunning, numFinished, numFailed, failedSlurmJobIds] = getStateInformation(slurmIDs, expectedNumOutputs)
            import parallel.internal.apishared.SlurmUtils
            
            noHeader = true;
            slurmIDsStr = SlurmUtils.slurmIDsToString(slurmIDs);
            [failedCommand, out] = SlurmUtils.getSlurmJobStatusOutput(slurmIDsStr, expectedNumOutputs, noHeader);
            
            numPending  = numel(regexp(out, 'PENDING|SPECIAL_EXIT'));
            numRunning  = numel(regexp(out, 'CONFIGURING|COMPLETING|RUNNING|STOPPED|SUSPENDED|RESIZING'));
            numFinished = numel(regexp(out, 'COMPLETED'));
            numFailed   = numel(regexp(out, 'BOOT_FAIL|CANCELLED|DEADLINE|FAILED|NODE_FAIL|REVOKED|PREEMPTED|TIMEOUT'));
            
            if numFailed > 0
                % Get the Slurm Job IDs of the failed jobs
                failedSlurmJobIds = iGetFailedSlurmJobIds(out, numFailed);
            else
                failedSlurmJobIds = {};
            end
        end
        
        function slurmIDstring = slurmIDsToString(slurmIDs)
            % Remove any empty cells.  Cells can be empty if we're not
            % using job arrays and tasks were deleted from an independent
            % job before submission.
            slurmIDs = slurmIDs(~cellfun('isempty',slurmIDs));
            slurmIDstring = strjoin(slurmIDs, ',');
        end
        
        function OK = isAcceptableScancelError(scancelString)
            % If the job has already finished or squeue has forgotten about
            % the job, scancel may (depending on the version of Slurm)
            % return a string of the form: "scancel: error: Kill job error
            % on job id 515: Invalid job id specified".  This is okay
            % because one way or another the job is no longer running.
            OK = ~isempty(regexp(scancelString, ': Invalid job id specified', 'once'));
        end
        
        function taskIDs = getTaskIDsFromSlurmJobArrayIDs(slurmJobIds)
            % Extracts the Slurm taskIDs from a cell array of strings of
            % the following forms: '123_4', '123_[4,5]', '123_[4,6-8,10]',
            % where 123 is the Slurm jobID of the array, and the numbers
            % after the underscore (and in the square brackets) are the
            % taskIDs.
            
            taskIDStrings = regexp(slurmJobIds, '\d+\_(\d+)|\d+\_\[([\d,-]+)\]', 'tokens', 'once');
            
            % Flatten cell array structure
            taskIDStrings = cat(2, taskIDStrings{:});
            
            if isempty(taskIDStrings)
                % We didn't find any matches
                taskIDs = [];
                return;
            end
            % Split strings by ','
            taskIDStrings = cellfun(@(x) strsplit(x, ','), taskIDStrings, 'UniformOutput', false);
            taskIDStrings = cat(2, taskIDStrings{:});
            
            % Find out which strings contain a '-'
            containsHyphen = cell2mat(cellfun(@(x) contains(x, '-'), taskIDStrings, 'UniformOutput', false));
            
            taskIDStringsWithHyphen = taskIDStrings(containsHyphen);
            taskIDStringsWithoutHyphen = taskIDStrings(~containsHyphen);
            
            % Convert the strings that don't contain a '-' to numeric
            taskIDs = cell2mat(cellfun(@str2double, taskIDStringsWithoutHyphen, 'UniformOutput', false));
            
            if ~isempty(taskIDStringsWithHyphen)
                lowerAndUpperBounds = regexp(taskIDStringsWithHyphen, '[0-9]+', 'match');
                for ii = 1:numel(lowerAndUpperBounds)
                    lowerBound = str2double(lowerAndUpperBounds{ii}{1});
                    upperBound = str2double(lowerAndUpperBounds{ii}{2});
                    taskIDs = [taskIDs, lowerBound:upperBound]; %#ok<AGROW>
                end
            end
        end
    end
end

function failedSlurmJobIds = iGetFailedSlurmJobIds(output, numFailed)
% If we are using Slurm job arrays, we are expecting output to look
% something like:
% '1233_2         COMPLETED
%  1233_3         FAILED
%  1233_1         COMPLETED
% '
%
% For a short period of time after the last tasks in a job array have been
% cancelled, sacct's output may look like this:
% '1708_[4,6-8]   CANCELLED+
%  1708_1         COMPLETED
%  1708_2         COMPLETED
% '
%
% Otherwise we are expecting output to look like:
% '1233           COMPLETED
%  1234           FAILED
%  1235           COMPLETED
% '

slurmIDSearchString = '(\d+\_\d+)|(\d+\_\[[\d,-]+\])|(\d+)';

% Remove trailing whitespace
output = strtrim(output);

% Split into lines
lout = splitlines(output);

% Loop over the lines looking for BOOT_FAIL, CANCELLED, FAILED, TIMEOUT,
% PREEMPTED, NODE_FAIL, OUT_OF
failedSlurmJobIds = cell(1, numFailed);
failedCount = 1;
for i = 1:numel(lout)
    thisLine = lout{i};
    if ~isempty(regexp(thisLine, 'BOOT_FAIL|CANCELLED|FAILED|TIMEOUT|PREEMPTED|NODE_FAIL|OUT_OF', 'once'))
        % Split into whitespace delimited cells
        thisLineCell = regexp(thisLine, '[^\s]+', 'match');
        jobIndexCellStr = {};
        
        % Need to check that thisLineCell{1} is of the form e.g. 1233_2,
        % where 1233 is the Slurm job number and 2 is the task number.
        if numel(thisLineCell) >= 1
            jobIndexCellStr = regexp(thisLineCell{1}, slurmIDSearchString, 'tokens', 'once');
        end
        
        % Did our first attempt to get the jobIndexCellStr work?
        if isempty(jobIndexCellStr)
            jobIndexCellStr = regexp(thisLine, slurmIDSearchString, 'tokens', 'once');
        end
        try
            failedSlurmJobIds{failedCount} = jobIndexCellStr{1};
            failedCount = failedCount + 1;
        catch err %#ok<NASGU>
            % Remove this one from the possible list since we failed to get
            % a Slurm job array index
            failedSlurmJobIds{failedCount} = [];
        end
    end
end
end
