%ClusterParforEngine PARFOR implementation that submits independent tasks to a cluster.

% Copyright 2018-2020 The MathWorks, Inc.

classdef ClusterParforEngine < parallel.internal.parfor.Engine

    properties (GetAccess = private, SetAccess = immutable)
        Cluster
        AutoAttachFiles
        AutoAddClientPath
        AttachedFiles
        AdditionalPaths
        InitData
        ParforFcn
    end

    properties (Access = private)
        Intervals
        Job = parallel.Job.empty()
        UnfinishedTasks = parallel.Task.empty();
        UnconsumedFinishedTasks = parallel.Task.empty();
        AddedFiles = false
    end

    properties (GetAccess = private, Constant)
        ForbiddenTypes = { 'parallel.pool.Constant', ...
                           'parallel.pool.DataQueue', ...
                           'parallel.pool.PollableDataQueue', ...
                           'Composite', ...
                           'distributed' };
    end

    methods (Access = private)

        function errorIfInvalidSerialization(~, serializedClassesList)
            invalidSerializations = intersect(parallel.internal.cluster.ClusterParforEngine.ForbiddenTypes, ...
                                              serializedClassesList);
            if ~isempty(invalidSerializations)
                error(message('parallel:lang:parfor:ClusterParforUnsupportedClass', ...
                              strjoin(invalidSerializations)));
            end
        end

        function submitPendingIntervals(obj, moreFilesToAttach)
        % Submit all pending intervals as a new job.

            [listener, getListFcn] = ...
                parallel.internal.general.SerializationNotifier.createAccumulatingListener(); %#ok<ASGLU>

            attachedFiles = [obj.AttachedFiles, moreFilesToAttach];

            % If AutoAttachFiles is true, we want to try and attach everything we can find
            % relating to the ParforFcn function.
            if obj.AutoAttachFiles
                loopFunction = obj.ParforFcn;
                loopFunctionInfo = functions(loopFunction);
                loopFunctionFile = loopFunctionInfo.file;
                if ~isempty(loopFunctionFile)
                    attachedFiles{end+1} = loopFunctionFile;
                end
            end

            tag = getString(message('parallel:lang:parfor:ParforJobDescription'));
            jobData = {pwd, obj.InitData}; % Set as JobData the working directory of the client, and the initData.

            % Each task takes as input arguments the 'tag' of the interval, and the interval data.
            vecArgs = {keys(obj.Intervals), values(obj.Intervals)};

            obj.Job = obj.Cluster.hBatchVec(@iEvaluateInterval, 3, vecArgs, ...
                'AttachedFiles', attachedFiles, ...
                'AutoAttachFiles', obj.AutoAttachFiles, ...
                'AutoAddClientPath', obj.AutoAddClientPath, ...
                'AdditionalPaths', obj.AdditionalPaths, ...
                'ApiTag', 'Created_by_ClusterParforEngine', ...
                'Tag', tag, ...
                'JobData', jobData, ...
                'CaptureDiary', false, ...
                'SlicedIndices', [1 2]);
            obj.UnfinishedTasks = obj.Job.Tasks;
            
            obj.UnconsumedFinishedTasks = parallel.Task.empty();

            % Throw an error if we've observed serialization of any problematic classes.
            obj.errorIfInvalidSerialization(getListFcn());

            dctSchedulerMessage(1, 'Submitted job with ID: %d with %d tasks', ...
                                obj.Job.ID, numel(obj.UnfinishedTasks));
        end

        function [shouldRetry, err, filesToAttach] = shouldRetryAfterFailure(obj, err)
        %shouldRetryAfterFailure True if we should attach some files and try again.

            filesToAttach = {};
            % Try again if we haven't already added files, and AutoAttachFiles is true.
            shouldRetry = ~obj.AddedFiles && obj.AutoAttachFiles;
            if shouldRetry
                % Try and work out some files to attach for next time. This might transform
                % the exception if it is a missing source exception.
                [err, filesToAttach] = ...
                    parallel.internal.parfor.maybeTransformMissingSourceException(...
                        err, obj.ParforFcn);

                % Only retry if there are actually files to attach.
                shouldRetry = ~isempty(filesToAttach);
            end
        end

        function tasks = awaitTasks(obj, maxNumTasks)
        % awaitTasks - wait for up to the given maximum number of tasks to
        % be complete and return them
        
            while isempty(obj.UnconsumedFinishedTasks)
                taskStates = obj.UnfinishedTasks.hGetStateEnum();
                if iscell(taskStates)
                    taskStates = [taskStates{:}]';
                end
                newlyFinishedTaskIndices = arrayfun(@(taskState) taskState >= parallel.internal.types.States.Finished, taskStates);
                obj.UnconsumedFinishedTasks = obj.UnfinishedTasks(newlyFinishedTaskIndices);
                obj.UnfinishedTasks(newlyFinishedTaskIndices) = [];
                if ~any(newlyFinishedTaskIndices)
                    pause(1);
                end
            end
            
            numTasksToReturn = min(maxNumTasks, numel(obj.UnconsumedFinishedTasks));
            tasks = obj.UnconsumedFinishedTasks(1:numTasksToReturn);
            obj.UnconsumedFinishedTasks(1:numTasksToReturn) = [];
        end
    end

    methods
        function obj = ClusterParforEngine(partitionMethod, partitionSize, ...
                                           cluster, autoAttachFiles, autoAddClientPath, ...
                                           attachedFiles, additionalPaths, ...
                                           initData, parforFcn)
            obj@parallel.internal.parfor.Engine(partitionMethod, partitionSize);

            obj.Cluster = cluster;

            obj.AutoAttachFiles = autoAttachFiles;
            obj.AutoAddClientPath = autoAddClientPath;
            obj.AttachedFiles = attachedFiles;
            obj.AdditionalPaths = additionalPaths;

            obj.InitData = initData;
            obj.ParforFcn  = parforFcn;

            obj.Intervals = containers.Map('KeyType', 'int32', ... % 'tag'
                                           'ValueType', 'any');
        end

        function ok = addInterval(obj, tag, varargin)
        %addInterval stores interval data for later submission to the cluster.
            ok = true;
            obj.Intervals(tag) = varargin;
        end

        function allIntervalsAdded(obj)
        % Submit a single job to execute everything.
            obj.submitPendingIntervals({});
        end

        function [tags, results] = getCompleteIntervals(obj, num)
            tags = nan(num, 1);
            results = cell(num, 2);

            % Repeatedly wait for up to the given number of tasks to finish
            % and retrieve their completed intervals
            idx = 1;
            while idx <= num
                numRemainingTasksToAwait = num - idx + 1;
                tasks = awaitTasks(obj, numRemainingTasksToAwait);
                taskProps = tasks.hGetProperty({'Error', 'OutputArguments', 'InputArguments'});
                
                for ii = 1:numel(tasks)
                    task = tasks(ii);
                    taskError = taskProps{ii, 1};
                    taskOutputArguments = taskProps{ii, 2};
                    taskInputArguments = taskProps{ii, 3};
                
                    if isempty(taskError)
                        if isempty(taskOutputArguments)
                            % Maybe get here if the worker crashed.
                            disp(task.Diary);
                            error(message('parallel:lang:parfor:TaskExecutionFailed', task.ID));
                        end
                        % Print the captured output text
                        taskText = taskOutputArguments{3};
                        fprintf(1, '%s', taskText);

                        % caughtError is any exception confined by iEvaluateInterval.
                        caughtError = taskOutputArguments{1};

                        if ~isempty(caughtError)
                            % shouldRetryAfterFailure tells us if we should retry - and if we should,
                            % it tells us the files we should attach. It might transform the 
                            % exception too.
                            [shouldRetry, caughtError, possibleSources] = obj.shouldRetryAfterFailure(caughtError);
                            if ~shouldRetry
                                throwAsCaller(caughtError);
                            else
                                dctSchedulerMessage(0, 'Attaching additional files and retrying.');
                                dctSchedulerMessage(1, 'Additional files are:\n%s', ...
                                                    strjoin(possibleSources, pathsep));
                                delete(obj.Job);
                                obj.submitPendingIntervals(possibleSources);

                                % Make sure we don't make a second attempt to attach files.
                                obj.AddedFiles = true;

                                % Didn't consume a result, but we need to go back to waiting for a task to
                                % complete, so break.
                                break
                            end
                        end
                        results(idx, :) = taskOutputArguments{2};
                        tags(idx) = taskInputArguments{1};
                        remove(obj.Intervals, tags(idx));
                        idx = 1 + idx;
                    else
                        % This doesn't normally happen - but we can get here if e.g. an MJS task gets
                        % cancelled because it exhausts re-run attempts.
                        dctSchedulerMessage(0, 'Unexpected task execution failure.', taskError);
                        E = MException(message('parallel:lang:parfor:TaskExecutionFailed', task.ID));
                        throw(addCause(E, taskError));
                    end
                end
            end
        end

        function N = getNumWorkers(obj)
        % getNumWorkers - the number of workers used for determining subrange partitioning
        %    Here we must use the maximum number of workers supported by the cluster,
        %    rather than the current number of workers. This is to support
        %    resizable clusters where NumWorkers does not reflect how many
        %    workers can be made available to run tasks.
            N = obj.Cluster.hGetMaxNumWorkers();
        end

        function [init, remain] = getDispatchSizes(~, numIntervals, W)
            init = iClamp(Inf, 1, numIntervals);
            remain = 2 * W;
        end

        function complete(obj, errorWasThrown)
            if isempty(obj.Job) || ~isvalid(obj.Job)
                return
            end
            if errorWasThrown
                [~, undoc] = pctconfig;
                preserve = undoc.preservejobs;
            else
                preserve = false;
            end
            try
                if preserve
                    dctSchedulerMessage(1, 'Preserving job with ID: %d', ...
                                        obj.Job.ID);
                    % Simply cancel the job, but leave it around.
                    cancel(obj.Job);
                else
                    delete(obj.Job);
                end
            catch E
                % Simply log any problems during cancel/delete of the job.
                dctSchedulerMessage(0, sprintf('An error occurred cancelling/deleting job with ID: %d', obj.Job.ID), E);
            end
        end
        function delete(obj)
            if ~isempty(obj.Job) && isvalid(obj.Job)
                % Normally, we would expect parallel_function to have called "complete", and
                % so we would not enter this block. However, in the case of
                % CTRL-C, that might not happen. So, here we ensure that the
                % cleanup happens, and we treat this as though an error
                % occurred.
                dctSchedulerMessage(1, 'In ClusterParforEngine destructor with valid job.');
                errorWasThrown = true;
                obj.complete(errorWasThrown);
            end
        end
    end
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function [ERROR, out, outText] = iEvaluateInterval(~, intervalData) %#ok<INUSD> evalc
    job = getCurrentJob();
    [clientWD, loopInitData] = deal(job.JobData{:});
    if exist(clientWD, 'dir') == 7
        try
            cd(clientWD);
        catch E
            % Failed to change directory - log and continue.
            dctSchedulerMessage(0, sprintf('Failed to change directory to %s:', clientWD), E);
        end
    else
        % No access to client WD - just log
        dctSchedulerMessage(1, 'Worker cannot see client''s WD: %s', clientWD);
    end
    try
        [channelMaker, parfor_C] = deal(loopInitData{:});
        channel = channelMaker(parfor_C); %#ok<NASGU> evalc
        outText = evalc('out = feval(channel, intervalData{:});');
        ERROR = MException.empty();
    catch ERROR
        outText = '';
        out = [];
    end
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function val = iClamp(inVal, minBound, maxBound)
    val = min(max(inVal, minBound), maxBound);
end
