% LsfUtils - a collection of utilities for LSF schedulers/clusters.

% Copyright 2011-2019 The MathWorks, Inc.

classdef ( Hidden, Sealed ) LsfUtils

    methods ( Hidden, Static )

        function aHandler = messageHandler(jobID, taskID, initialMessageNumber)
            %messageHandler Return a new LSF message handler
            if nargin < 3
                initialMessageNumber = 0;
            end

            if isempty( taskID )
                % Parallel execution
                initialMessageNumber = labindex-1;
                messageIncrement = numlabs;
                jobString = jobID;
            else
                messageIncrement = 1;
                jobString = sprintf('%s[%s]', jobID, taskID);
            end
            messageNumber = initialMessageNumber;
            aHandler = @nMessageHandler;
            if ispc
                quote = '"';
            else
                quote = '''';
            end

            function nMessageHandler(stringToPost)
                import parallel.internal.apishared.ThirdPartySchedulerUtils

                % Display the message so that it ends up in the debug log
                disp(stringToPost);
                % Strip characters that would cause problems with the
                % command line - any quotes or newlines
                stringToPost = regexprep(stringToPost, '[''"\n\r]', '');
                ThirdPartySchedulerUtils.system( ...
                    sprintf('bpost -i %d -d %s%s%s "%s"', messageNumber, quote, stringToPost, quote, jobString));
                messageNumber = messageNumber + messageIncrement;
            end
        end
        
        function environmentSetup()
            % Some environment variables will cause our LSF integration to
            % fail. These must be unset rather than set to empty before we
            % can proceed with job submission
            pctunsetenv('BSUB_QUIET');
        end

        function [masterName, clusterName] = getMasterAndClusterNames()
            % Better check that there actually is an lsf scheduler
            % available - call the LSF lsid command to get information
            % about my cluster
            
            import parallel.internal.apishared.ThirdPartySchedulerUtils
            
            [FAILED, out] = ThirdPartySchedulerUtils.system('lsid');
            % Possible that lsid isn't a file I can execute
            if FAILED
                lsidIsOnPath = ThirdPartySchedulerUtils.isExeOnPath('lsid');
                if lsidIsOnPath
                    error(message('parallel:cluster:LsfErrorRunningLsid', out));
                else
                    error(message('parallel:cluster:LsfUnableToFindService'));
                end
            end

            allEndIndex = regexp(out, '\n') - 1;
            % Lets try and parse the output from out to get the relevant information
            clusterStr = 'My cluster name is ';
            masterStr  = 'My master name is ';

            startClusterIndex = regexp(out, clusterStr) + numel(clusterStr);
            endClusterIndex   = min(allEndIndex(allEndIndex > startClusterIndex));

            startMasterIndex  = regexp(out, masterStr) + numel(masterStr);
            endMasterIndex    = min(allEndIndex(allEndIndex > startMasterIndex));

            clusterName = out(startClusterIndex:endClusterIndex);
            masterName  = out(startMasterIndex:endMasterIndex);
        end

        function OK = isAcceptableBkillError( bkillString )
            OK = false;
            % Job already finished is OK because we might be destroying or
            % cancelling a job that has finished
            OK = OK || ~isempty(regexp(bkillString, ': Job has already finished', 'once'));
            % After a time LSF forgets about jobs - we don't. Thus LSF
            % might be out of date with us.
            OK = OK || ~isempty(regexp(bkillString, ': No matching job found', 'once'));
            % LSF seems to incorrectly indicate an error when this is returned
            OK = OK || ~isempty(regexp(bkillString, ': Operation is in progress', 'once'));
        end

        function [FAILED, numPending, numRunning, numFailed, numFinished, failedLsfIndex] = getStateInformation( lsfIDs )
            import parallel.internal.apishared.ThirdPartySchedulerUtils
            
            numPending = 0;
            numRunning = 0;
            numFailed = 0;
            numFinished = 0;
            failedLsfIndex = [];
            
            % Pad IDs with double quotes and join them up with spaces.
            % We expect to have multiple IDs for large job arrays that
            % needed to be submitted with multiple calls to bsub.
            lsfIDsStr = strjoin( strcat( '"', lsfIDs, '"' ) );

            % Ask LSF about this job
            [FAILED, out] = ThirdPartySchedulerUtils.system(sprintf('bjobs %s', lsfIDsStr));
            if FAILED
                warning(message('parallel:cluster:LsfBjobsFailed', out));
                return
            end
            
            % How many PEND, PSUSP, USUSP, SSUSP, WAIT
            numPending = numel(regexp(out, 'PEND|PSUSP|USUSP|SSUSP|WAIT'));
            % How many RUN strings - UNKWN started running and then comms was lost
            % with the sbatchd process.
            numRunning = numel(regexp(out, 'RUN|UNKWN'));
            % How many EXIT, ZOMBI strings
            numFailed = numel(regexp(out, 'EXIT|ZOMBI'));
            % How many DONE
            numFinished = numel(regexp(out, 'DONE'));

            if numFailed > 0
                % Get the LSF task ID of the failed tasks so we can set
                % those tasks to failed.
                failedLsfIndex = iParseBjobsOutputForExit(out, numFailed);
            end
        end

        function lsfID = extractLsfJobId( out )
            % Extract the LSF job ID as a numeric value.
            jobNumberStr = regexp(out, 'Job <(\d+)>', 'tokens', 'once');
            if isempty( jobNumberStr )
                error(message('parallel:cluster:LsfFailedToParseJobId', out));
            end
            lsfID = jobNumberStr{1};
        end

        function [filesToCopy, remoteTaskLoc] = ...
                calcFilesToCopy( storageLocation, jobLocation )
            % TODO:later reconcile with JobFilesInfo
            localJobLoc   = [storageLocation filesep jobLocation];
            remoteJobLoc  = ['%J.%I.mps/' jobLocation];
            localTaskLoc  = [localJobLoc filesep 'Task%I'];
            remoteTaskLoc = [remoteJobLoc '/Task%I'];
            % Construct the required file transfer commands. The remote
            % paths are all relative to the job's current working directory
            % when it is executed because we do not know what directories
            % exist on the execution host.  The Job and Task zip files are
            % copied to LSB_JOBID.LSB_JOBINDEX.Job.zip and
            % LSB_JOBID.LSB_JOBINDEX.Task.zip at the start of the job. The
            % worker unzips them in into a directory called
            % LSB_JOBID.LSB_JOBINDEX.mps (see
            % parallel.internal.decode.lsfSingleZippedTask). At the end of
            % job, all the output files in LSB_JOBID.LSB_JOBINDEX.mps are
            % copied back to the client.
            filesToCopy = {...
                ['-f "' localJobLoc '.zip > %J.%I.Job.zip"'] ...
                ['-f "' localJobLoc filesep 'Task%I.zip > %J.%I.Task.zip"'] ...
                ['-f "' localJobLoc  '.common.mat  < ' remoteJobLoc  '.common.mat"'] ...
                ['-f "' localJobLoc  '.state.mat   < ' remoteJobLoc  '.state.mat"'] ...
                ['-f "' localTaskLoc '.common.mat  < ' remoteTaskLoc '.common.mat"'] ...
                ['-f "' localTaskLoc '.out.mat     < ' remoteTaskLoc '.out.mat"'] ...
                ['-f "' localTaskLoc '.state.mat   < ' remoteTaskLoc '.state.mat"'] ...
                ['-f "' localTaskLoc '.log         < ' remoteTaskLoc '.log"'] ...
                ['-f "' localTaskLoc '.diary.txt   < ' remoteTaskLoc '.diary.txt"'] ...
                };
            filesToCopy = sprintf('%s ', filesToCopy{:});
        end

        function lsfID = callBsub( cmdLine )

            import parallel.internal.apishared.LsfUtils
            import parallel.internal.apishared.ThirdPartySchedulerUtils

            storedEnv = distcomp.pClearEnvironmentBeforeSubmission();
            cleanup = onCleanup( @() distcomp.pRestoreEnvironmentAfterSubmission( storedEnv ) );

            [FAILED, out] = ThirdPartySchedulerUtils.system(cmdLine);
            if FAILED
                error(message('parallel:cluster:LsfBsubInvocationFailed', out));
            end
            lsfID = LsfUtils.extractLsfJobId( out );
        end
        function callBresume( lsfID, optTaskIdString )
            % Now set the job actually running by moving it from the PSUSP
            % state to the PEND state
            
            import parallel.internal.apishared.ThirdPartySchedulerUtils
            if nargin == 1
                arg = lsfID;
            else
                arg = sprintf( '%s[%s]', lsfID, optTaskIdString );
            end

            [FAILED, out] = ThirdPartySchedulerUtils.system(sprintf('bresume "%s"', arg));
            if FAILED
                error(message('parallel:cluster:LsfBresumeInvocationFailed', out));
            end

        end
    end
end

function iFailedLsfIndex = iParseBjobsOutputForExit(out, numFailed)
% We are expecting out to look something like this
%
% JOBID   USER    STAT  QUEUE      FROM_HOST   EXEC_HOST   JOB_NAME   SUBMIT_TIME
% 7005    jlmarti EXIT  normal     uk-martinj-    -        Job2[1]    Jul 26 11:16
% 7005    jlmarti EXIT  normal     uk-martinj-    -        Job2[2]    Jul 26 11:16
% 7005    jlmarti EXIT  normal     uk-martinj-    -        Job2[3]    Jul 26 11:16
% 7005    jlmarti EXIT  normal     uk-martinj-    -        Job2[4]    Jul 26 11:16

% Convert the output to a cell array of lines
lout = regexp(out, '.*?\n', 'match');
% Remove the header line
lout(1) = [];
% Loop over the lines looking for EXIT or ZOMBI
iFailedLsfIndex = zeros(1, numFailed);
iFailedCount = 1;
for i = 1:numel(lout)
    thisLine = lout{i};
    if ~isempty(regexp(thisLine, 'EXIT|ZOMBI', 'once'))
        % Split into whitespace delimited cells
        thisLineCell = regexp(thisLine, '[^\s]+', 'match');
        jobIndexCellStr = {};
        % Need to check that thisLineCell{7} has [] with a number inside
        if numel(thisLineCell) >= 7
            jobIndexCellStr = regexp(thisLineCell{7}, '\[([0-9]+)\]', 'tokens', 'once');
        end
        % Did our first attempt to get the jobIndexCellStr work?
        if isempty(jobIndexCellStr)
            jobIndexCellStr = regexp(thisLine, '\[([0-9]+)\]', 'tokens', 'once');
        end
        try
            iFailedLsfIndex(iFailedCount) = str2double(jobIndexCellStr{1});
            iFailedCount = iFailedCount + 1;
        catch err %#ok<NASGU>
            % Remove this one from the possible list since we failed to get
            % an LSF job array index
            iFailedLsfIndex(iFailedCount) = [];
        end
    end
end
end
