HeterogeneousCore/MPICore/plugins/MPISource.cc - master - githubmirror/cmssw

HeterogeneousCore/MPICore/plugins/MPISource.cc
400 строк15 KB
Andrea Bocci
Redesign the MPI event and stream synchronisation
27 апр 2026, 15:05
27 апр 2026, 15:057eadb80
// C++ headers
#include <memory>
#include <stdexcept>
#include <string>
#include <vector>

// MPI headers
#include <mpi.h>

// CMSSW headers
#include "DataFormats/Provenance/interface/BranchListIndex.h"
#include "DataFormats/Provenance/interface/EventAuxiliary.h"
#include "DataFormats/Provenance/interface/EventSelectionID.h"
#include "DataFormats/Provenance/interface/EventToProcessBlockIndexes.h"
#include "DataFormats/Provenance/interface/LuminosityBlockAuxiliary.h"
#include "DataFormats/Provenance/interface/ProcessHistory.h"
#include "DataFormats/Provenance/interface/ProcessHistoryRegistry.h"
#include "DataFormats/Provenance/interface/RunAuxiliary.h"
#include "FWCore/Framework/interface/Event.h"
#include "FWCore/Framework/interface/EventPrincipal.h"
#include "FWCore/Framework/interface/InputSource.h"
#include "FWCore/Framework/interface/InputSourceDescription.h"
#include "FWCore/Framework/interface/InputSourceMacros.h"
#include "FWCore/Framework/interface/ProductProvenanceRetriever.h"
#include "FWCore/Framework/interface/TriggerNamesService.h"
#include "FWCore/MessageLogger/interface/ErrorObj.h"
#include "FWCore/MessageLogger/interface/MessageLogger.h"
#include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
#include "FWCore/ParameterSet/interface/EmptyGroupDescription.h"
#include "FWCore/ParameterSet/interface/ParameterSet.h"
#include "FWCore/ParameterSet/interface/ParameterSetDescription.h"
#include "FWCore/ParameterSet/interface/ParameterSetDescriptionFiller.h"
#include "FWCore/ServiceRegistry/interface/Service.h"
#include "FWCore/Sources/interface/ProducerSourceBase.h"
#include "FWCore/Utilities/interface/EDMException.h"
#include "FWCore/Utilities/interface/StreamID.h"
#include "HeterogeneousCore/MPICore/interface/MPIChannel.h"
#include "HeterogeneousCore/MPICore/interface/MPIToken.h"
#include "HeterogeneousCore/MPICore/interface/conversion.h"
#include "HeterogeneousCore/MPICore/interface/messages.h"
#include "HeterogeneousCore/MPIServices/interface/MPIService.h"

class MPISource : public edm::ProducerSourceBase {
public:
  explicit MPISource(edm::ParameterSet const& config, edm::InputSourceDescription const& desc);
  ~MPISource() override;
  using InputSource::processHistoryRegistryForUpdate;
  using InputSource::productRegistryUpdate;

  static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);

private:
  bool setRunAndEventInfo(edm::EventID& id, edm::TimeValue_t& time, edm::EventAuxiliary::ExperimentType&) override;
  void produce(edm::Event&) override;

  enum Mode { kInvalid = 0, kCommWorld, kIntercommunicator };
  static constexpr const char* ModeDescription[] = {"Invalid", "CommWorld", "Intercommunicator"};
  Mode parseMode(std::string const& label) {
    if (label == ModeDescription[kCommWorld])
      return kCommWorld;
    else if (label == ModeDescription[kIntercommunicator])
      return kIntercommunicator;
    else
      return kInvalid;
  }

  char port_[MPI_MAX_PORT_NAME];
  MPI_Comm comm_ = MPI_COMM_NULL;
  MPIChannel controller_;
  std::vector<std::unique_ptr<MPIChannel>> channels_;
  edm::EDPutTokenT<MPIToken> token_;
  Mode mode_;

  edm::ProcessHistory history_;

  // temporary value used to pass information from setRunAndEventInfo() to produce()
  MPIChannel* channel_ = nullptr;
};

MPISource::MPISource(edm::ParameterSet const& config, edm::InputSourceDescription const& desc)
    :  // note that almost all configuration parameters passed to IDGeneratorSourceBase via ProducerSourceBase will
       // effectively be ignored, because this ConfigurableSource will explicitly set the run, lumi, and event
       // numbers, the timestamp, and the event type
      edm::ProducerSourceBase(config, desc, false),
      token_(produces<MPIToken>()),
      mode_(parseMode(config.getUntrackedParameter<std::string>("mode")))  //
{
  // Make sure that MPI is initialised.
  MPIService::required();

  // Make sure the EDM MPI types are available.
  EDM_MPI_build_types();

  if (mode_ == kCommWorld) {
    // All processes are in MPI_COMM_WORLD.
    edm::LogInfo("MPI") << "MPISource in " << ModeDescription[mode_] << " mode.";

    // Check how many processes are there in MPI_COMM_WORLD
    int size;
    MPI_Comm_size(MPI_COMM_WORLD, &size);

    // Check the rank of this process.
    int rank;
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);

    edm::LogInfo("MPI") << "MPIController Comm World size: " << size;

    // All processes exchange the hashes of their names.
    // One follower process has to make one communication channel with the controller process
    // If controller process is not unique, error is thrown
    auto controller_name = config.getParameter<std::string>("controllerProcessName");
    if (controller_name.empty()) {
      throw edm::Exception(edm::errors::Configuration)
          << "ERROR: Controller process name cannot be empty. Aborting MPISource...";
    }

    edm::Service<edm::service::TriggerNamesService> tns;
    std::string const& this_process_name = tns->getProcessName();
    if (controller_name == this_process_name) {
      throw edm::Exception(edm::errors::Configuration)
          << "ERROR: controller and follower processes cannot have the same name. Aborting MPISource...";
    }

    edm::Service<MPIService> mpiservice;
    std::vector<int> controller_indices = mpiservice->getRanksByProcessName(controller_name);
    int remote = -1;
    if (controller_indices.empty()) {
      throw edm::Exception(edm::errors::Configuration)
          << "ERROR: No controller process with name " << controller_name << " found. Aborting...";
    } else if (controller_indices.size() == 1) {
      remote = controller_indices[0];
    } else {
      throw edm::Exception(edm::errors::Configuration)
          << "ERROR: Multiple controller processes with name " << controller_name
          << " were found. Currently, only one controller process is supported. Aborting...";
    }

    // Create a new communicator that spans only this process and the one with the given remote rank.
    int ranks[2] = {remote, rank};
    MPI_Group world_group, comm_group;
    MPI_Comm_group(MPI_COMM_WORLD, &world_group);
    MPI_Group_incl(world_group, 2, ranks, &comm_group);
    MPI_Comm_create_group(MPI_COMM_WORLD, comm_group, 0, &comm_);
    MPI_Group_free(&world_group);
    MPI_Group_free(&comm_group);
    edm::LogInfo("MPI") << "The MPIController process and MPISource have ranks " << remote << ", " << rank
                        << " in MPI_COMM_WORLD, mapped to ranks 0, 1 in their private communicator.";
    // The remote process always has rank 0 in the new communicator.
    remote = 0;
    controller_ = MPIChannel(comm_, remote);
  } else if (mode_ == kIntercommunicator) {
    // Use an intercommunicator to let two groups of processes communicate with each other.
    // The current implementation supports only two processes: one controller and one source.
    edm::LogInfo("MPI") << "MPISource in " << ModeDescription[mode_] << " mode.";

    // Check how many processes are there in MPI_COMM_WORLD
    int size;
    MPI_Comm_size(MPI_COMM_WORLD, &size);
    if (size != 1) {
      throw edm::Exception(edm::errors::Configuration)
          << "The current implementation supports only two processes: one controller and one source.";
    }

    // Open a server-side port.
    MPI_Open_port(MPI_INFO_NULL, port_);

    // Publish the port under the name indicated by the parameter "server".
    std::string name = config.getUntrackedParameter<std::string>("name", "server");
    MPI_Info port_info;
    MPI_Info_create(&port_info);
    MPI_Info_set(port_info, "ompi_global_scope", "true");
    MPI_Info_set(port_info, "ompi_unique", "true");
    MPI_Publish_name(name.c_str(), port_info, port_);

    // Create an intercommunicator and accept a client connection.
    edm::LogInfo("MPI") << "Waiting for a connection to the MPI server at port " << port_;

    MPI_Comm_accept(port_, MPI_INFO_NULL, 0, MPI_COMM_SELF, &comm_);
    edm::LogInfo("MPI") << "Connection accepted.";
    controller_ = MPIChannel(comm_, 0);
  } else {
    // Invalid mode.
    throw edm::Exception(edm::errors::Configuration)
        << "Invalid mode \"" << config.getUntrackedParameter<std::string>("mode") << "\"";
  }

  // Wait for a client to connect.
  MPI_Status status;
  EDM_MPI_Empty_t buffer;
  MPI_Recv(&buffer, 1, EDM_MPI_Empty, MPI_ANY_SOURCE, EDM_MPI_Connect, comm_, &status);
  edm::LogInfo("MPI") << "connected from " << status.MPI_SOURCE;
}

MPISource::~MPISource() {
  // Disconnect the communicators.
  for (auto& channel : channels_) {
    if (channel) {
      channel->reset();
    }
  }
  controller_.reset();

  if (mode_ == kIntercommunicator) {
    // Close the intercommunicator.
    MPI_Comm_disconnect(&comm_);

    // Unpublish and close the port.
    MPI_Info port_info;
    MPI_Info_create(&port_info);
    MPI_Info_set(port_info, "ompi_global_scope", "true");
    MPI_Info_set(port_info, "ompi_unique", "true");
    MPI_Unpublish_name("server", port_info, port_);
    MPI_Close_port(port_);
  }
}

//MPISource::ItemTypeInfo MPISource::getNextItemType() {
bool MPISource::setRunAndEventInfo(edm::EventID& event,
                                   edm::TimeValue_t& time,
                                   edm::EventAuxiliary::ExperimentType& type) {
  while (true) {
    MPI_Status status;
    MPI_Message message;
    MPI_Mprobe(MPI_ANY_SOURCE, MPI_ANY_TAG, comm_, &message, &status);
    switch (status.MPI_TAG) {
      // Connect message
      case EDM_MPI_Connect: {
        // receive the message header
        EDM_MPI_Empty_t buffer;
        MPI_Mrecv(&buffer, 1, EDM_MPI_Empty, &message, &status);

        // the Connect message is unexpected here (see above)
        throw cms::Exception("InvalidValue")
            << "The MPISource has received an EDM_MPI_Connect message after the initial connection";
        return false;
      }

      // Disconnect message
      case EDM_MPI_Disconnect: {
        // receive the message header
        EDM_MPI_Empty_t buffer;
        MPI_Mrecv(&buffer, 1, EDM_MPI_Empty, &message, &status);

        // signal the end of the input data
        return false;
      }

      // BeginStream message
      case EDM_MPI_BeginStream: {
        // receive the message header
        EDM_MPI_Empty_t buffer;
        MPI_Mrecv(&buffer, 1, EDM_MPI_Empty, &message, &status);

        // receive the next message
        break;
      }

      // EndStream message
      case EDM_MPI_EndStream: {
        // receive the message header
        EDM_MPI_Empty_t buffer;
        MPI_Mrecv(&buffer, 1, EDM_MPI_Empty, &message, &status);

        // receive the next message
        break;
      }

      // BeginRun message
      case EDM_MPI_BeginRun: {
        // receive the RunAuxiliary
        EDM_MPI_RunAuxiliary_t buffer;
        MPI_Mrecv(&buffer, 1, EDM_MPI_RunAuxiliary, &message, &status);
        // TODO this is currently not used
        edm::RunAuxiliary runAuxiliary;
        edmFromBuffer(buffer, runAuxiliary);

        // receive the ProcessHistory
        history_.clear();
        controller_.receiveProduct(0, history_);
        history_.initializeTransients();
        /*
        if (processHistoryRegistryForUpdate().registerProcessHistory(history_)) {
          edm::LogInfo("MPI") << "new ProcessHistory registered: " << history_;
        }
        */

        // receive the next message
        break;
      }

      // EndRun message
      case EDM_MPI_EndRun: {
        // receive the RunAuxiliary message
        EDM_MPI_RunAuxiliary_t buffer;
        MPI_Mrecv(&buffer, 1, EDM_MPI_RunAuxiliary, &message, &status);

        // receive the next message
        break;
      }

      // BeginLuminosityBlock message
      case EDM_MPI_BeginLuminosityBlock: {
        // receive the LuminosityBlockAuxiliary
        EDM_MPI_LuminosityBlockAuxiliary_t buffer;
        MPI_Mrecv(&buffer, 1, EDM_MPI_LuminosityBlockAuxiliary, &message, &status);
        // TODO this is currently not used
        edm::LuminosityBlockAuxiliary luminosityBlockAuxiliary;
        edmFromBuffer(buffer, luminosityBlockAuxiliary);

        // receive the next message
        break;
      }

      // EndLuminosityBlock message
      case EDM_MPI_EndLuminosityBlock: {
        // receive the LuminosityBlockAuxiliary
        EDM_MPI_LuminosityBlockAuxiliary_t buffer;
        MPI_Mrecv(&buffer, 1, EDM_MPI_LuminosityBlockAuxiliary, &message, &status);

        // receive the next message
        break;
      }

      // ProcessEvent message
      case EDM_MPI_ProcessEvent: {
        // receive the EventAuxiliary
        edm::EventAuxiliary aux;
        unsigned int slot;
        status = controller_.receiveEvent(aux, slot, message);

        // use the same communicator that the MPIController will use for this event
        if (slot >= channels_.size()) {
          channels_.resize(slot + 1);
        }
        if (not channels_[slot]) {
          channels_[slot] = controller_.duplicate(slot);
        }

        // store the channel to use it in produce()
        channel_ = channels_[slot].get();

        // extract the rank of the other process (currently unused)
        int source = status.MPI_SOURCE;
        (void)source;

        // fill the event details
        event = aux.id();
        time = aux.time().value();
        type = aux.experimentType();

        // signal a new event
        return true;
      }

      // unexpected message
      default: {
        throw cms::Exception("InvalidValue")
            << "The MPISource has received an unknown message with tag " << status.MPI_TAG;
        return false;
      }
    }
  }
}

void MPISource::produce(edm::Event& event) {
  // Wait for the barrier to be cleared by the MPI software in the local process.
  channel_->wait();

  // The destructor of the last copy of the token will call channel_->sync().
  // The channel is ready to receive a new event after the call is made by both local and remote processes.
  event.emplace(token_, *channel_);
  channel_ = nullptr;
}

void MPISource::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
  descriptions.setComment(
      "This module connects to an \"MPIController\" in a separate CMSSW job, receives all Run, LuminosityBlock and "
      "Event transitions from the remote process and reproduces them in the local one.");

  edm::ParameterSetDescription desc;
  edm::ProducerSourceBase::fillDescription(desc);
  desc.ifValue(
          edm::ParameterDescription<std::string>("mode", "CommWorld", false),
          ModeDescription[kCommWorld] >>
                  edm::ParameterDescription<std::string>(
                      "controllerProcessName",
                      "",
                      true,
                      edm::Comment("Process name of the controller process corresponding to this MPISource.\n"
                                   "Only one process with this name is expected.\n")) or
              ModeDescription[kIntercommunicator] >> edm::ParameterDescription<std::string>("name", "server", false))
      ->setComment(
          "Valid modes are CommWorld (use MPI_COMM_WORLD) and Intercommunicator (use an MPI name server to setup an "
          "intercommunicator).");

  descriptions.add("source", desc);
}

#include "FWCore/Framework/interface/InputSourceMacros.h"
DEFINE_FWK_INPUT_SOURCE(MPISource);