% ThirdPartySchedulerUtils - A collection of common utilities for third party schedulers/clusters.

% Copyright 2017-2019 The MathWorks, Inc.

classdef (Hidden, Sealed) ThirdPartySchedulerUtils
    
    methods (Hidden, Static)
        
        function scriptOnCluster = copyWrapperScript(wrapper, clusterOsType, dlStruct, jobLocation)
            % Deal with the copying the wrapper script into the job subdirectory.
            
            import parallel.internal.apishared.FilenameUtils
            
            if exist(wrapper, 'file') ~= 2
                error(message('parallel:cluster:JobWrapperMissing', wrapper));
            end
            
            [~, shortPart, ext] = fileparts(wrapper);
            
            [clientWrapper, scriptOnCluster] = FilenameUtils.jobSpecificFilename(...
                clusterOsType, dlStruct, jobLocation, [shortPart, ext], false);
            
            [success, msg] = copyfile(wrapper, clientWrapper);
            if ~success
                error(message('parallel:cluster:JobWrapperCopyFailed', wrapper, clientWrapper, msg));
            end
            
            % The script needs to be executable on the workers. Windows
            % workers don't care about the execute bit, and Windows clients
            % default to copying files with "+x" set for UNIX workers. Thus
            % we only need to make the script executable if the client is
            % UNIX.
            if isunix
                [success, msg] = fileattrib(clientWrapper, '+x');
                if ~success
                    error(message('parallel:cluster:JobWrapperChmodFailed', clientWrapper, msg));
                end
            end
        end
        
        function string = makeTaskIdString( taskIds )
            % makeTaskIdString - given a list of Task IDs, return the string needed
            % by LSF and Slurm Schedulers during job submission.
            validateattributes( taskIds, {'double'}, {} );
            
            numTasks = numel(taskIds);
            
            % Easy if only one ID
            if numTasks == 1
                string = sprintf('%d', taskIds);
                return;
            end
            
            % Easy if IDs consecutive
            if numTasks == taskIds(end)-taskIds(1)+1
                string = sprintf('%d-%d', taskIds(1), taskIds(end));
                return;
            end
            
            % Only get here if some tasks have been deleted/cancelled
            
            % Ensure taskIds is a row before we continue
            taskIds = taskIds(:).';
            
            % Check that the IDs were originally in a monotonic order
            [IDs, index] = sort(taskIds);
            if ~isequal(index, 1:numel(index))
                % Should never get here
                error(message('parallel:cluster:ThirdPartySchedulerInvalidTaskState'));
            end
            
            % The format of the string is 1-N1,N2,N3-N4,N5,N6 ...
            
            % Deduce the difference between subsequent ID's
            deltaIDs = diff(IDs);
            % Define a cell array to hold the output - starts of ranges will be the
            % first in a series of diffs of 1.
            out = cell(numTasks, 1);
            % Ensure we start with a range and not in a range
            lastDiff = 2;
            % Loop over each diff
            for ii = 1:numel(deltaIDs)
                % Get this ID and the difference to the next
                thisID = IDs(ii);
                thisDiff = deltaIDs(ii);
                if thisDiff > 1
                    % If the diff to the next is more than 1 then we are either at the
                    % end of a range or on a singleton number, so simply print the
                    % number followed by a comma
                    out{ii} = sprintf('%d,', thisID);
                elseif thisDiff == 1
                    % If the diff is one is this the start of a range or are we in a
                    % range. If the last wasn't 1 then this is the start of a range
                    if lastDiff > 1
                        out{ii} = sprintf('%d-', thisID);
                    end
                else
                    % Should never get here
                    error(message('parallel:cluster:ThirdPartySchedulerDuplicateTaskIds'));
                end
                % Remember the last difference
                lastDiff = thisDiff;
            end
            % Finally always add the last number, irrespective of what has happened
            out{numTasks} = sprintf('%d', IDs(numTasks));
            % Now remove all the empty cells
            out(cellfun('isempty', out)) = [];
            % And concatenate the whole lot together.
            string = strcat(out{:});
        end
        
        function [splitTasks, splitTaskIDs] = splitTaskArray(tasks, maxSubarraySize, noSingles)
            % Helper function for built-in clusters to split an array of
            % tasks and taskIDs into subarrays that are each smaller than
            % maxSubarraySize.
            
            % Specify noSingles as true if the scheduler does not allow job
            % arrays with just one element
            if nargin < 3
                noSingles = false;
            end
            
            % Ensure input is column vector
            tasks = tasks(:);
            
            % Determine number of full subarrays to use
            numTasks         = numel(tasks);
            numFullSubarrays = floor(numTasks./maxSubarraySize);
            splitSizes       = maxSubarraySize.*ones(1, numFullSubarrays);
            
            % If maxSubarraySize doesn't divide nicely into numTasks, we'll
            % have one shorter array at the end
            lastSubarraySize = rem(numTasks, maxSubarraySize);
            if lastSubarraySize > 0
                splitSizes = [splitSizes, lastSubarraySize];
            end
            
            % If the short array at the end consists of just one element,
            % move a task from the second-to-last array to the end
            if noSingles && splitSizes(end) == 1
                splitSizes(end)   = splitSizes(end)   + 1;
                splitSizes(end-1) = splitSizes(end-1) - 1;
            end
            
            % Split the array
            splitTasks = mat2cell(tasks, splitSizes);
            
            % Only bother returning the IDs if we're asked to, as this many
            % save a call to disk
            if nargout > 1
                taskIDs = cell2mat(get(tasks, {'Id'}));
                splitTaskIDs = mat2cell(taskIDs, splitSizes);
            end
        end
        
        function varargout = querySchedulerConfig(configCmd, names, defaults, parseFcn)
            % Calls the configCmd, and searches the output for names
            % If any aren't found, return the value in defaults
            
            varargout = defaults;
            
            % Provide default parsing functions
            if nargin < 4
                parseFcn = repmat({@str2double}, size(names));
            end
            
            % Retrieve scheduler config data
            [failed, out] = parallel.internal.apishared.ThirdPartySchedulerUtils.system(configCmd);
            if failed
                return;
            end
            lines = splitlines(out);
            % For 'scontrol show config' on a Slurm cluster, this will be
            % 100-200 lines, looking something like:
            % ...
            % MailProg                = /usr/bin/mail
            % MaxArraySize            = 1001
            % MaxJobCount             = 10000
            % MaxJobId                = 2147418112
            % MaxMemPerNode           = UNLIMITED
            % MaxStepCount            = 40000
            % ...
            % For a cluster running PBSPro, 'qstat -Bf' will give around
            % 30 lines looking something like:
            % Server: bat554201pbsglnxa64
            %     server_state = Active
            %     server_host = bat554201pbsglnxa64.mathworks.com
            %     scheduling = True
            % ...
            
            % Attempt to find each of the names
            for ii = 1:numel(names)
                line = lines(contains(lines, names{ii}));
                % Expecting only one matching line
                if numel(line)~=1
                    continue;
                end
                % Expecting format 'something = number' or 'something: number'
                splitline = strsplit(line{1}, {'=',':'});
                if length(splitline) < 2
                    continue;
                end
                % Apply parsing function (default is str2double)
                num = parseFcn{ii}(strip(splitline{2}));
                if isnan(num)
                    continue;
                end
                varargout{ii} = num;
            end
        end
        
        function resourceArgument = generateResourceArgument(template, numWorkers, numThreads)
            % Generate the resource argument for job submission based on an
            % input template.
            %
            % All occurrences of the '^N^' in the template will be replaced
            % by numWorkers. All occurrences of ^T^ will be replaced by
            % numThreads.
            %
            % For communicating jobs, numWorkers should be the number of
            % workers in the job.  For independent jobs/tasks, numWorkers
            % should always be 1.
            
            % Trim the template of any leading or trailing white space.
            template = strtrim(template);

            if isnumeric(numWorkers)
                numWorkers = num2str(numWorkers);
            end
            if ~isempty(template)
                resourceArgument = replace(template, {'^N^', '^T^'}, {numWorkers, num2str(numThreads)});
            else
                error(message('parallel:cluster:ThirdPartySchedulerNoResourceTemplate'));
            end
        end
        
        function [status, result] = system(cmd)
            % Many third party scheduler utility commands on unix return
            % exit codes > 127, which causes MATLAB's system() command to
            % report this as a signal. This is not the case, so to work
            % around this, simply wrap all system calls on UNIX within a
            % wrapper script which sanitizes the exit codes.
            
            persistent wrapper
            
            if isunix
                if isempty(wrapper)
                    wrapper = fullfile( toolboxdir('parallel'), ...
                        'bin', 'util', 'shellWrapper.sh' );
                end
                cmd = [wrapper, ' ', cmd];
            end
            
            [status, result] = system(cmd);
        end
        
        function onPath = isExeOnPath(baseName)
            % Checks to see if a given program is on the executable PATH.
            % On Windows, this will append extensions ".exe", ".bat" and
            % ".cmd" and check those too.
            
            exePath = textscan(getenv('PATH'), '%s', 'delimiter', pathsep);
            exePath = exePath{1};
            onPath = any(cellfun(@(x)(iExeHere(fullfile(x, baseName))), exePath));
        end
    end
end

function isHere = iExeHere(nameNoExt)
% iExeHere - is there an executable file at this location? Handles PC
% executable extensions

if ispc
    % On PC, handle the case where the input argument already specifies
    % .exe or whatever, and also append known likely executable extensions.
    trail = {'', '.exe', '.bat', '.cmd', '.com'};
else
    trail = {''};
end

isHere = any(cellfun(@(x)(exist([nameNoExt, x], 'file')), trail) == 2);
end
