#include "module.h"

#include <ctranslate2/devices.h>

#include "utils.h"

namespace ctranslate2 {
  namespace python {

    void register_mpi(py::module& m) {
      py::class_<ScopedMPISetter>(
        m, "MpiInfo",
        R"pbdoc(
            An object to manage the MPI communication between processes.
            It provides information about MPI connexion.
        )pbdoc")

        .def_static("getNRanks", &ScopedMPISetter::getNRanks,
                             "Get the number of gpus running for the current model.")

        .def_static("getCurRank", &ScopedMPISetter::getCurRank,
                             "Get the current rank of process.")

        .def_static("getLocalRank", &ScopedMPISetter::getLocalRank,
                             "Get the current GPU id used by process.")
        ;
    }

  }
}