add explicit Init and Finalize methods and export them to python

zml · Mar 19, 2024 · 23508eb · 23508eb
1 parent b74bbb9
commit 23508eb
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 5 deletions.
diff --git a/xla/pjrt/cpu/mpi_collectives.cc b/xla/pjrt/cpu/mpi_collectives.cc
@@ -236,15 +236,15 @@ absl::Status MpiCollectivesCommunicator::ReduceScatter(
       input_buffer, output_buffer, recvcounts.data(), type, op, comm_));
 }
 
-MpiCollectives::MpiCollectives() {
+void MpiCollectives::Init() {
   int provided;
   MPI_Init_thread(NULL, NULL, MPI_THREAD_FUNNELED, &provided);
   MPI_Comm_rank(MPI_COMM_WORLD, &mpi_world_rank_);
   MPI_Comm_size(MPI_COMM_WORLD, &mpi_world_size_);
   VLOG(1) << "MPI rank=" << mpi_world_rank_ << " size=" << mpi_world_size_;
 }
 
-MpiCollectives::~MpiCollectives() {
+void MpiCollectives::Finalize() {
   contexts_.clear();
   MPI_Finalize();
 }

diff --git a/xla/pjrt/cpu/mpi_collectives.h b/xla/pjrt/cpu/mpi_collectives.h
@@ -71,8 +71,9 @@ class MpiCollectivesCommunicator : public CollectivesCommunicator {
 
 class MpiCollectives : public CollectivesInterface {
  public:
-  MpiCollectives();
-  ~MpiCollectives() override;
+
+  void Init();
+  void Finalize();
 
   absl::StatusOr<std::shared_ptr<CollectivesCommunicator>> GetCommunicator(
       absl::Span<GlobalDeviceId const> global_devices, int rank) override;

diff --git a/xla/python/xla.cc b/xla/python/xla.cc
@@ -273,9 +273,15 @@ NB_MODULE(xla_extension, m_nb) {
       nb::arg("distributed_client"), nb::arg("hostname").none() = std::nullopt,
       nb::arg("interface").none() = std::nullopt);
 
+#ifndef _WIN32
+  nb::class_<cpu::MpiCollectives> mpi_collectives(m_nb, "MpiCollectives", cpu_collectives);
+  mpi_collectives.def("Init", &cpu::MpiCollectives::Init);
+  mpi_collectives.def("Finalize", &cpu::MpiCollectives::Finalize);
+#endif // _WIN32
+
   m_nb.def(
       "make_mpi_collectives",
-      []() -> std::shared_ptr<xla::cpu::CollectivesInterface> {
+      []() -> std::shared_ptr<cpu::MpiCollectives> {
 #ifndef _WIN32
         return std::make_shared<cpu::MpiCollectives>();
 #else  // _WIN32