text: #   Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  from __future__ import print_function
  from . import core
  from . import framework
  from . import executor
  from . import compiler
  import sys
  
  __all__ = ['ParallelExecutor']
  
  ExecutionStrategy = core.ParallelExecutor.ExecutionStrategy
  BuildStrategy = core.ParallelExecutor.BuildStrategy
  
  
  class ParallelExecutor(object):
      """
      The ParallelExecutor is an upgraded version of :code:`fluid.Executor` that supports multi-node model
      training and testing based on the data-parallel mode. In data-parallel mode,
      ParallelExecutor will broadcast the parameters from Node0 to other nodes during
      construction and copy the input Program to other nodes from Node0 to make sure
      that the initial state on each node is the same. Each node runs the model independently
      and the parameters' gradient is aggregated between those nodes during backward
      computation, and then each node independently updates its parameters. If you use
      the GPU to run the model, i.e. use_cuda=True, the node refers to the GPU,
      ParallelExecutor will automatically get the GPU resources available on the
      current machine, users can also set the available GPU resources in the environment
      variable, for example: want to use GPU0, GPU1, export CUDA_VISIBLEDEVICES=0,1;
      If the operation is performed on the CPU, i.e. use_cuda=False, the node refers to the CPU.
      **Note: At this time, the user needs to manually add CPU_NUM to the environment variable
      and set the number of CPU devices. For example, export CPU_NUM=4, if the environment
      variable is not set, the executor will add the variable to the environment variable
      and set it to 1.**
  
  
      Args:
          use_cuda (bool): Whether to use CUDA or not.
          loss_name (str): This parameter is the name of the loss variable of the
              model. **Note: If it is data-parallel model training, you must set loss_name,
              otherwise, the results may be wrong**. The default is None.
          main_program (Program): This parameter represents the Program to be executed.
              If this parameter is not provided, that parameter is None, the program will
              be set to :code:`fluid.default_main_program()`. The default is None.
          share_vars_from(ParallelExecutor): If share_vars_from is set, the current
              ParallelExecutor will share the parameters with the ParallelExecutor
              specified by share_vars_from. This parameter needs to be set when model testing
              is required during model training, and the data parallel mode is used for
              training and testing. Since ParallelExecutor will only distribute parameter
              variables to other devices when it is first executed, the ParallelExecutor
              specified by share_vars_from must be run before the current ParallelExecutor.
              The default is None.
          exec_strategy(ExecutionStrategy): exec_strategy specifies the options that can
              be changed when running the current model, such as the thread pool size.
              For more information about exec_strategy, please refer to :code:`fluid.ExecutionStrategy`.
              The default is None.
          build_strategy(BuildStrategy): By configuring build_strategy, we can
              optimize the computational graph, such as operators' fusion in the
              computational graph and memory optimization during the execution
              of the computational graph. For more information about build_strategy,
              please refer to :code:`fluid.BuildStrategy`.  The default is None.
          num_trainers(int): This parameter needs to be set in GPU distributed training.
              If the parameter value is greater than 1, NCCL will be initialized by multi-level
              nodes. Each node should have the same number of GPUs. The default is 1.
          trainer_id(int): This parameter needs to be set when performing GPU distributed
              training. This parameter must be used with the num_trainers parameter.
              Trainer_id indicates the "rank" of the current node. The trainer_id starts
              counting from 0. The default is 0.
          scope(Scope): Specifies the scope in which the program is executed.
              The default is fluid.global_scope().
  
      Returns:
          ParallelExecutor: The initialized ParallelExecutor object.
  
      Raises:
          TypeError: If share_vars_from is provided, but not ParallelExecutor object.
  
      NOTES:
  
          1. If you only use ParallelExecutor to do multi-card test, you don't need to set loss_name
             and share_vars_from.
  
          2. If you need to train and test the model with ParallelExecutor, the share_vars_from
             must be set when building the ParallelExecutor corresponding to the model test.
             Otherwise, the parameters used in the model test and the model training are inconsistent.
  
      Examples:
          .. code-block:: python
  
            import paddle.fluid as fluid
            import numpy
            import os
  
            use_cuda = True
            place = fluid.CUDAPlace(0) if use_cuda else fluid.CPUPlace()
  
            # NOTE: If you use CPU to run the program, you need
            # to specify the CPU_NUM, otherwise, fluid will use
            # all the number of the logic core as the CPU_NUM,
            # in that case, the batch size of the input should be
            # greater than CPU_NUM, if not, the process will be
            # failed by an exception.
            if not use_cuda:
                os.environ['CPU_NUM'] = str(2)
  
            exe = fluid.Executor(place)
  
            train_program = fluid.Program()
            startup_program = fluid.Program()
            with fluid.program_guard(train_program, startup_program):
                data = fluid.data(name='X', shape=[None, 1], dtype='float32')
                hidden = fluid.layers.fc(input=data, size=10)
                loss = fluid.layers.mean(hidden)
                test_program = fluid.default_main_program().clone(for_test=True)
                fluid.optimizer.SGD(learning_rate=0.01).minimize(loss)
  
            startup_program.random_seed=1
            exe.run(startup_program)
  
            train_exe = fluid.ParallelExecutor(use_cuda=use_cuda,
                                               main_program=train_program,
                                               loss_name=loss.name)
            test_exe = fluid.ParallelExecutor(use_cuda=use_cuda,
                                              main_program=test_program,
                                              share_vars_from=train_exe)
  
            x = numpy.random.random(size=(10, 1)).astype('float32')
            loss_data, = train_exe.run(feed={"X": x},
                                       fetch_list=[loss.name])
  
            loss_data, = test_exe.run(feed={"X": x},
                                      fetch_list=[loss.name])
  
      """
  
      def __init__(self,
                   use_cuda,
                   loss_name=None,
                   main_program=None,
                   share_vars_from=None,
                   exec_strategy=None,
                   build_strategy=None,
                   num_trainers=1,
                   trainer_id=0,
                   scope=None):
          if build_strategy is None:
              build_strategy = BuildStrategy()
  
          # TODO(paddle-dev): trainer_id and num_trainers should be removed from parameter list.
          if num_trainers != 1 and build_strategy.num_trainers != num_trainers:
              sys.stderr.write(
                  'The value of build_strategy.num_trainers[%d] is overwritten '
                  'by the passed num_trainers[%d].\n' %
                  (build_strategy.num_trainers, num_trainers))
              build_strategy.num_trainers = num_trainers
          if trainer_id != 0 and build_strategy.trainer_id != trainer_id:
              sys.stderr.write(
                  'The value of build_strategy.trainer_id[%d] is overwritten '
                  'by the passed trainer_id[%d].\n' %
                  (build_strategy.trainer_id, trainer_id))
              build_strategy.trainer_id = trainer_id
  
          self._places = framework.cuda_places(
          ) if use_cuda else framework.cpu_places()
          self._scope = scope if scope is not None else executor.global_scope()
  
          if main_program is not None and main_program._enable_dgc:
              assert build_strategy.num_trainers > 1, "dgc is not useful when num_trainers <= 1"
              assert build_strategy.reduce_strategy == BuildStrategy.ReduceStrategy.AllReduce, "dgc \
                  only used for allreduce"
  
              assert build_strategy.num_trainers * len(
                  self._places) > 1, "dgc is not useful for single card training"
              assert use_cuda, "dgc only used under cuda"
  
          main_program = main_program if main_program is not None \
              else framework.default_main_program()
  
          self._compiled_program = compiler.CompiledProgram(main_program)
          if share_vars_from:
              assert isinstance(
                  share_vars_from, ParallelExecutor
              ), "The share_vars_from should be ParallelExecutor."
  
          self._compiled_program.with_data_parallel(
              loss_name=loss_name,
              build_strategy=build_strategy,
              exec_strategy=exec_strategy,
              share_vars_from=share_vars_from._compiled_program
              if share_vars_from else None)
  
          self._place = core.CUDAPlace(0) if use_cuda else core.CPUPlace()
          self._exe = executor.Executor(self._place)
  
      def run(self, fetch_list, feed=None, feed_dict=None, return_numpy=True):
          """
          This interface is used to run the current model. It should be noted
          that the executor will execute all the operators in the Program,
          and will not prune some operators in the Program according to the
          fetch_list.
  
          Args:
              fetch_list(list): This parameter represents the variables that need to be returned
                  after the model runs. The default is None.
              feed(list|dict): This parameter represents the input variables of the model.
                  If it is single card training, the feed is dict type, and if it is multi-card
                  training, the parameter feed can be dict or list type variable. If the
                  parameter type is dict, the data in the feed will be split and sent to
                  multiple devices (CPU/GPU), that is to say, the input data will be evenly
                  sent to different devices, so you should make sure the number of samples of
                  the current mini-batch must be greater than the number of places;
                  if the parameter type is list, those data are copied directly to each device,
                  so the length of this list should be equal to the number of places.
                  The default is None.
              feed_dict: Alias for feed parameter, for backward compatibility.
                  This parameter has been deprecated. Default None.
              return_numpy(bool): This parameter indicates whether convert the fetched variables
                  (the variable specified in the fetch list) to numpy.ndarray. if it is False,
                  the type of the return value is a list of :code:`LoDTensor`. The default is True.
  
          Returns:
              List: The fetched result list.
  
          Raises:
              ValueError: If the feed is a list, but its length is not equal the
                  length of active places, or its element's is not dict.
  
          NOTES:
              1. If the feed parameter is dict type, the input data will be evenly distributed
                 to different cards. For example, using two GPUs to run the model, the input
                 sample number is 3, that is, [0, 1, 2], the sample number on GPU0 is 1,
                 that is, [0], and the sample number on GPU1 is 2, that is, [1, 2].
                 If the number of samples is less than the number of devices, the program will
                 throw an exception, so when running the model, you should make sure that the
                 number of samples of the last batch of the data set should be greater than the
                 number of CPU cores or GPU cards, if it is less than, it is recommended that
                 the batch be discarded.
              2. If the number of CPU cores or GPU cards available is greater than 1, the fetch
                 results are spliced together in dimension 0 for the same variable values
                 (variables in fetch_list) on different devices.
  
  
          Examples:
              .. code-block:: python
  
                import paddle.fluid as fluid
                import numpy
                import os
  
                use_cuda = True
                place = fluid.CUDAPlace(0) if use_cuda else fluid.CPUPlace()
  
                # NOTE: If you use CPU to run the program, you need
                # to specify the CPU_NUM, otherwise, fluid will use
                # all the number of the logic core as the CPU_NUM,
                # in that case, the batch size of the input should be
                # greater than CPU_NUM, if not, the process will be
                # failed by an exception.
                if not use_cuda:
                    os.environ['CPU_NUM'] = str(2)
  
                exe = fluid.Executor(place)
  
                train_program = fluid.Program()
                startup_program = fluid.Program()
                with fluid.program_guard(train_program, startup_program):
                    data = fluid.data(name='X', shape=[None, 1], dtype='float32')
                    hidden = fluid.layers.fc(input=data, size=10)
                    loss = fluid.layers.mean(hidden)
                    fluid.optimizer.SGD(learning_rate=0.01).minimize(loss)
  
                exe.run(startup_program)
  
                train_exe = fluid.ParallelExecutor(use_cuda=use_cuda,
                                                   main_program=train_program,
                                                   loss_name=loss.name)
  
                # If the feed is a dict:
                # the image will be splitted into devices. If there is two devices
                # each device will process an image with shape (5, 1)
                x = numpy.random.random(size=(10, 1)).astype('float32')
                loss_data, = train_exe.run(feed={"X": x},
                                           fetch_list=[loss.name])
  
                # If the feed is a list:
                # each device will process each element in the list.
                # the 1st device will process an image with shape (10, 1)
                # the 2nd device will process an image with shape (9, 1)
                #
                # you can use exe.device_count to get the device number.
                x2 = numpy.random.random(size=(9, 1)).astype('float32')
                loss_data, = train_exe.run(feed=[{"X": x}, {"X": x2}],
                                           fetch_list=[loss.name])
  
          """
          return self._exe.run(program=self._compiled_program,
                               scope=self._scope,
                               feed=feed,
                               fetch_list=fetch_list,
                               return_numpy=return_numpy)
  
      @property
      def device_count(self):
          return len(self._places)
  
      def drop_local_exe_scopes(self):
          """
          Drop the local execution scopes immediately. In order to avoid frequently
          application and release of temporary variables, the strategy adopted by
          ParallelExecutor is to drop the local execution scopes after several iterations.
          ParallelExecutor provides the num_iteration_per_drop_scope option in
          :code:`fluid.ExecutionStrategy`, which indicates how many iterations are intervened to
          drop the local execution scopes. If the num_iteration_per_drop_scope value
          is 100, but you want to drop the local execution scopes after 50 iterations,
          you can call the interface manually.
  
          Returns:
              None
  
          Examples:
              .. code-block:: python
  
                import paddle.fluid as fluid
                import numpy
                import os
  
                use_cuda = True
                # NOTE: If you use CPU to run the program, you need
                # to specify the CPU_NUM, otherwise, fluid will use
                # all the number of the logic core as the CPU_NUM,
                # in that case, the batch size of the input should be
                # greater than CPU_NUM, if not, the process will be
                # failed by an exception.
                if not use_cuda:
                    os.environ['CPU_NUM'] = str(2)
  
                train_program = fluid.Program()
                startup_program = fluid.Program()
                with fluid.program_guard(train_program, startup_program):
                    data = fluid.data(name='X', shape=[None, 1], dtype='float32')
                    hidden = fluid.layers.fc(input=data, size=10)
                    loss = fluid.layers.mean(hidden)
  
                place = fluid.CUDAPlace(0) if use_cuda else fluid.CPUPlace()
                exe = fluid.Executor(place)
                exe.run(startup_program)
  
                parallel_exe = fluid.ParallelExecutor(use_cuda=use_cuda,
                                                   main_program=train_program,
                                                   loss_name=loss.name)
  
                x = numpy.random.random(size=(10, 1)).astype('float32')
                loss_data, = parallel_exe.run(feed={"X": x},
                                           fetch_list=[loss.name])
  
                parallel_exe.drop_local_exe_scopes()
          """
          assert isinstance(
              self._compiled_program._executor,
              core.ParallelExecutor), "The Executor should be ParallelExecutor."
          self._compiled_program._executor.drop_local_exe_scopes()
  
      # This API is used to check whether DropLocalExeScopes can work.
      def _need_create_local_exe_scopes(self):
          assert isinstance(
              self._compiled_program._executor,
              core.ParallelExecutor), "The Executor should be ParallelExecutor."
          return self._compiled_program._executor._need_create_local_exe_scopes()
  