mpi_cluster.hpp

Go to the documentation of this file.

#ifdef STAN_MPI
 
#ifndef STAN_MATH_PRIM_FUNCTOR_MPI_CLUSTER_HPP
#define STAN_MATH_PRIM_FUNCTOR_MPI_CLUSTER_HPP
 
#include <stan/math/prim/functor/mpi_command.hpp>
 
#include <boost/mpi/allocator.hpp>
#include <boost/mpi/collectives.hpp>
#include <boost/mpi/communicator.hpp>
#include <boost/mpi/datatype.hpp>
#include <boost/mpi/environment.hpp>
#include <boost/mpi/nonblocking.hpp>
#include <boost/mpi/operations.hpp>
 
#include <boost/serialization/access.hpp>
#include <boost/serialization/base_object.hpp>
#include <boost/serialization/export.hpp>
#include <boost/serialization/shared_ptr.hpp>
 
#include <mutex>
#include <vector>
#include <memory>
 
namespace stan {
namespace math {
 
class mpi_stop_listen : public std::exception {
  virtual const char* what() const throw() {
    return "Stopping MPI listening mode.";
  }
};
 
class mpi_is_in_use : public std::exception {
  virtual const char* what() const throw() { return "MPI resource is in use."; }
};
 
struct mpi_stop_worker : public mpi_command {
  friend class boost::serialization::access;
  template <class Archive>
  void serialize(Archive& ar, const unsigned int version) {
    ar& BOOST_SERIALIZATION_BASE_OBJECT_NVP(mpi_command);
  }
  void run() const {
    boost::mpi::communicator world;
    throw mpi_stop_listen();
  }
};
 
inline std::vector<int> mpi_map_chunks(std::size_t num_jobs,
                                       std::size_t chunk_size = 1) {
  boost::mpi::communicator world;
  const std::size_t world_size = world.size();
 
  std::vector<int> chunks(world_size, num_jobs / world_size);
 
  const std::size_t delta_r = chunks[0] == 0 ? 0 : 1;
 
  for (std::size_t r = 0; r != num_jobs % world_size; ++r)
    ++chunks[r + delta_r];
 
  for (std::size_t i = 0; i != world_size; ++i)
    chunks[i] *= chunk_size;
 
  return chunks;
}
 
template <typename T>
std::unique_lock<std::mutex> mpi_broadcast_command();
 
struct mpi_cluster {
  boost::mpi::environment env;
  boost::mpi::communicator world_;
  std::size_t const rank_ = world_.rank();
 
  mpi_cluster() {}
 
  ~mpi_cluster() {
    // the destructor will ensure that the childs are being
    // shutdown
    stop_listen();
  }
 
  void listen() {
    listening_status() = true;
    if (rank_ == 0) {
      return;
    }
 
    try {
      // lock on the workers the cluster as MPI commands must be
      // initiated from the root and any attempt to do this on the
      // workers must fail
      std::unique_lock<std::mutex> worker_lock(in_use());
      while (1) {
        std::shared_ptr<mpi_command> work;
 
        boost::mpi::broadcast(world_, work, 0);
 
        work->run();
      }
    } catch (const mpi_stop_listen& e) {
    }
  }
 
  void stop_listen() {
    if (rank_ == 0 && listening_status()) {
      mpi_broadcast_command<mpi_stop_worker>();
    }
    listening_status() = false;
  }
 
  static bool& listening_status() {
    static bool listening_status = false;
    return listening_status;
  }
 
  static std::mutex& in_use() {
    static std::mutex in_use_mutex;
    return in_use_mutex;
  }
};
 
inline std::unique_lock<std::mutex> mpi_broadcast_command(
    std::shared_ptr<mpi_command>& command) {
  boost::mpi::communicator world;
 
  if (world.rank() != 0)
    throw std::runtime_error("only root may broadcast commands.");
 
  if (!mpi_cluster::listening_status())
    throw std::runtime_error("cluster is not listening to commands.");
 
  std::unique_lock<std::mutex> cluster_lock(mpi_cluster::in_use(),
                                            std::try_to_lock);
 
  if (!cluster_lock.owns_lock())
    throw mpi_is_in_use();
 
  boost::mpi::broadcast(world, command, 0);
 
  return cluster_lock;
}
 
template <typename T>
std::unique_lock<std::mutex> mpi_broadcast_command() {
  std::shared_ptr<mpi_command> command(new T);
 
  return mpi_broadcast_command(command);
}
 
}  // namespace math
}  // namespace stan
 
#endif
 
#endif