%CommPoolExecutor
% Parallel pool implementation of the PartitionArrayExecutor interface
% using all capabilities exposed by communication and MPI.
%

%   Copyright 2019-2020 The MathWorks, Inc.

classdef (Sealed) CommPoolExecutor < parallel.internal.bigdata.ParallelPoolExecutor
    properties (GetAccess = private, SetAccess = immutable)
        % A logical scalar that indicates whether all MATLAB Workers are
        % local to the client.
        AreWorkersLocal;
    end
    
    properties (Access = private)
        % The underlying cache manager.
        CacheManager;
    end
    
    methods
        % The main constructor.
        function obj = CommPoolExecutor(pool)
            validateattributes(pool, {'parallel.Pool'}, {'nonempty'}, 'ParallelPoolExecutor', 'pool');
            if ~isvalid(pool)
                error(message('MATLAB:class:InvalidHandle'));
            elseif ~pool.Connected
                error(message('parallel:lang:pool:InvalidPool'));
            elseif ~pool.SpmdEnabled
                error(message('parallel:mapreduce:SpmdEnabledRequired'));
            end
            obj = obj@parallel.internal.bigdata.ParallelPoolExecutor(pool);
            
            obj.AreWorkersLocal = isa(obj.Pool.Cluster, 'parallel.cluster.Local');
            
            import parallel.internal.bigdata.ParallelPoolCacheManager;
            obj.CacheManager = ParallelPoolCacheManager(pool);
            addlistener(pool, 'ObjectBeingDestroyed', @obj.poolBeingDeleted);
        end
    end
    
    % Implementation of the internal interface of ParallelPoolExecutor.
    % See parallel.internal.bigdata.ParallelPoolExecutor for more
    % information on these methods and the respective contract.
    methods (Access = protected)
        % Convert the input task graph into an array of independent tasks
        % that can be executed one by one.
        function [stageTasks, broadcastMap, additionalArguments] ...
                = buildStageTasks(obj, taskGraph, outputHandlers, progressReporter)
            import matlab.bigdata.internal.executor.ReadFailureAccumulator;
            import matlab.bigdata.internal.util.TempFolder;
            import parallel.internal.bigdata.DataQueueSession;
            import parallel.internal.bigdata.PoolOutputDispatcher;
            import parallel.internal.bigdata.ProgressReportDispatcher;
            
            outputSession = DataQueueSession();
            
            % Create the communication paths by which tasks can stream data
            % back to the client. This is only done if streaming is
            % required (e.g. for clientforeach). This is because the
            % alternative is communication via broadcast, which is more
            % efficient.
            if any([outputHandlers.IsStreamingHandler])
                outputHandlerDispatcher = PoolOutputDispatcher(outputHandlers);
                outputSession.addDispatcher(outputHandlerDispatcher);
                outputStreamFactoryFunction = @(task) outputHandlerDispatcher.createStreamFactory(task.Id, outputSession);
            else
                outputStreamFactoryFunction = [];
            end
            
            % Create the communication paths by which tasks can stream
            % progress reports back to the client.
            progressReportDispatcher = ProgressReportDispatcher(progressReporter);
            outputSession.addDispatcher(progressReportDispatcher);
            progressReportFunction = progressReportDispatcher.createProgressReportFunction(outputSession);
            
            % A map of broadcasted constants shared between client
            % and workers.
            import parallel.internal.bigdata.ParallelPoolBroadcastMap;
            broadcastMap = ParallelPoolBroadcastMap();
            
            cacheManager = obj.CacheManager;
            shuffleStrategyMap = containers.Map('KeyType', 'char', 'ValueType', 'any');
            stageTasks = matlab.bigdata.internal.executor.convertToIndependentTasks(taskGraph, ...
                'CreateShuffleStorageFunction', @(task) obj.createShuffleStorage(task, shuffleStrategyMap), ...
                'CreateStreamFactoryFunction', outputStreamFactoryFunction, ...
                'GetCacheStoreFunction', @cacheManager.getCacheStore, ...
                'ResolvePartitionStrategyFunction', @obj.resolvePartitionStrategy);
            additionalArguments = {progressReportFunction, shuffleStrategyMap, outputSession};
        end
        
        % Execute one stage task using the pool. This performs all
        % operations that require the pool itself.
        function readFailureSummary = executeStageTaskRemotely(obj, ...
                stageTask, broadcastMap, ~, progressReportFunction, shuffleStrategyMap, outputSession)
            import matlab.bigdata.internal.executor.ExecutionTask;
            import matlab.bigdata.internal.executor.ProgressReporter;
            import matlab.bigdata.internal.executor.ReadFailureAccumulator;
            import parallel.internal.bigdata.PartitionedArrayTask;
            import parallel.internal.mapreduce.ParfevalTaskRunner;
            import parallel.internal.mapreduce.ShuffleSortTask;
            
            partitionStrategy = stageTask.ExecutionPartitionStrategy;
            partitionStrategy = obj.resolvePartitionStrategy(partitionStrategy);
            dataProcessorFactory = stageTask.DataProcessorFactory;
            numExecutorPartitions = numpartitions(partitionStrategy);
            if numExecutorPartitions > 1
                % This is to avoid duplicates of the constant data being
                % made when there exists multiple parfeval calls.
                broadcastMap.setProcessor(dataProcessorFactory);
                dataProcessorFactory = @(varargin) feval(broadcastMap.getProcessor(), varargin{:});
            end
            
            cachedIndicesMap = obj.CacheManager.setupForExecution(stageTask.CacheEntryKeys);
            broadcastMap.synchronizeClientBroadcasts({stageTask.InputBroadcasts.Id});
            
            shuffleInputs = stageTask.InputShuffles;
            shuffleOutputs = stageTask.OutputShuffles;
            if obj.AreWorkersLocal
                % If all workers are local to the client, then we do not
                % need to perform any shuffle work.
                shuffleInputs = ExecutionTask.empty();
                shuffleOutputs = ExecutionTask.empty();
            end
            
            % Decide how to schedule the work among workers. This depends on
            % whether workers might already have some of the data cached,
            % as well as whether previous stage tasks generated location
            % specific output.
            if ~isempty(shuffleInputs)
                % This task has shuffle inputs. By this point, all shuffle
                % communication has been completed and all inputs for a
                % given partition exists on one specific worker. We must
                % schedule the work to match.
                lockedFunctor = iCreateStripedScheduleMap(obj.Pool.NumWorkers, numExecutorPartitions);
                freeIndices = [];
            else
                % We are free to schedule partitions to any partitions. For
                % the purposes of caching, we choose to schedule partitions
                % for which cache entries exist to the corresponding
                % worker. All other tasks are scheduled in a way that
                % allows for load balancing.
                [lockedIndices, freeIndices] = iDetermineScheduling(cachedIndicesMap, numExecutorPartitions);
                
                if all(cellfun(@isempty, lockedIndices))
                    lockedFunctor = [];
                else
                    lockedFunctor = @(id) lockedIndices{id};
                end
            end
            
            numTasks = ~isempty(lockedFunctor) + numel(freeIndices) + numel(shuffleOutputs);
            taskRunner = ParfevalTaskRunner(obj.Pool, numTasks);
            
            readFailureAccumulator = ReadFailureAccumulator(partitionStrategy.MaxNumReadFailures);
            
            outputSession.reset();
            taskRunner.OutputServer = outputSession;
            taskRunner.TaskResultCombiner = readFailureAccumulator;
            
            % Schedule the execution that depends on location.
            if ~isempty(lockedFunctor)
                taskRunner.launchTask(PartitionedArrayTask(dataProcessorFactory, ...
                    lockedFunctor, partitionStrategy, partitionStrategy.MaxNumReadFailures, ...
                    progressReportFunction, broadcastMap));
            end
            
            % Schedule any remaining execution that does not depend on location.
            % In this case, send the corresponding partition to each worker
            % instead of the partition strategy with a full copy of the
            % data.
            for partitionIndex = freeIndices(:)'
                taskRunner.launchTask(PartitionedArrayTask(dataProcessorFactory, ...
                    partitionIndex, ...
                    partitionStrategy.createPartition(partitionIndex, numExecutorPartitions), ...
                    partitionStrategy.MaxNumReadFailures, progressReportFunction, broadcastMap));
            end
            
            % Schedule the communication part of the work.
            isFinalizeRequired = true;
            for shuffleExecutionTask = shuffleOutputs(:)'
                shuffleStrategy = shuffleStrategyMap(shuffleExecutionTask.Id);
                
                if shuffleExecutionTask.OutputPartitionStrategy.IsBroadcast
                    % If broadcast output, we need to broadcast the single
                    % partition to all workers.
                    partitionIndexMap = @(~)1;
                else
                    % Otherwise we decide at this point to send each
                    % partition to a predetermined worker.
                    numWorkers = obj.Pool.NumWorkers;
                    numPartitions = obj.numPartitions(shuffleExecutionTask.OutputPartitionStrategy);
                    partitionIndexMap = iCreateStripedScheduleMap(numWorkers, numPartitions);
                end
                taskRunner.launchTask(ShuffleSortTask(partitionIndexMap, isFinalizeRequired, shuffleStrategy));
            end
            
            try
                taskRunner.wait();
            catch err
                if isa(err, 'ParallelException')
                    err = err.remotecause{1};
                elseif any(err.identifier == ["parallel:mapreduce:UndefinedFunctionOrHandleOnWorker", "parallel:mapreduce:UndefinedFunctionOnWorker"])
                    err = matlab.bigdata.BigDataException.build(err);
                end
                matlab.bigdata.internal.util.assertNotInternal(err);
                % Need to ensure all errors issued on workers are modified
                % to appear as if they were issued locally. This is to
                % avoid including parallel internals from the worker.
                err = rebaseToClient(err);
                updateAndRethrow(err);
            end
            
            % Before returning successfully, make sure that all dispatchers
            % received everything expected.
            outputSession.flush();
            outputSession.check();
            
            broadcastMap.synchronizeWorkerBroadcasts({stageTask.OutputBroadcasts.Id});
            
            readFailureSummary = readFailureAccumulator.Summary;
        end
        
        % Get a user visible string that represents this execution
        % environment.
        function name = getName(obj)
            profileName = obj.Pool.Cluster.Profile;
            if isempty(profileName)
                name = getString(message('parallel:bigdata:ParallelPoolName'));
            else
                name = getString(message('parallel:bigdata:ParallelPoolWithProfileName', profileName));
            end
        end
    end
    
    methods (Access = private)
        function poolBeingDeleted(obj, varargin)
            % Drop references to the CacheManager to avoid shutdown ordering problems in
            % deployed applications.
            obj.CacheManager = [];
        end
        
        function [writerFactory, readerFactory] = createShuffleStorage(obj, task, shuffleStrategyMap)
            % Create a shuffle point where data is stored to an intermediate
            % storage, then read in a "shuffled" order by the next task.
            strategy = task.OutputPartitionStrategy;
            strategy = obj.resolvePartitionStrategy(strategy);
            useSharedStorage = obj.AreWorkersLocal;
            [shuffleSortStrategy, writerFactory, readerFactory] ...
                = obj.createMatArrayShuffleStorage(...
                task.OutputCommunicationType, strategy, useSharedStorage);
            shuffleStrategyMap(task.Id) = shuffleSortStrategy; %#ok<NASGU>
        end
    end
end

% Create a map from Worker ID to array of partition ids that apply a
% striped pattern.
function fcn = iCreateStripedScheduleMap(numWorkers, numPartitions)
fcn = @nWorkerIdToPartitionIdMap;
    function partitionIndices = nWorkerIdToPartitionIdMap(id)
        partitionIndices = id : numWorkers : numPartitions;
    end
end

% Helper function that determines the scheduling of partitions
% 1:numPartitions based on the given array of cached partition indices
% from each worker.
function [lockedIndices, freeIndices] = iDetermineScheduling(cachedPartitionIndices, numPartitions)
lockedIndices = cachedPartitionIndices;

allIndices = [];
for ii = 1:numel(lockedIndices)
    lockedIndices{ii} = setdiff(cachedPartitionIndices{ii}, allIndices);
    allIndices = union(allIndices, lockedIndices{ii});
end
freeIndices = setdiff(1:numPartitions, allIndices);
end
