From 5225369481289bef488061a65c6eed395304ee79 Mon Sep 17 00:00:00 2001 From: jianrui geng Date: Tue, 7 Jul 2026 23:26:33 +0800 Subject: [PATCH 1/4] Add fixed-cell i-PI socket driver --- docs/advanced/input_files/input-main.md | 8 +- docs/parameters.yaml | 8 +- source/CMakeLists.txt | 4 +- .../read_input_item_system.cpp | 15 +- .../test_serial/read_input_item_test.cpp | 5 + source/source_main/driver.h | 1 + source/source_main/driver_ipi.cpp | 507 ++++++++++++++++++ source/source_main/driver_run.cpp | 9 +- source/source_main/ipi_socket.cpp | 241 +++++++++ source/source_main/ipi_socket.h | 40 ++ 10 files changed, 829 insertions(+), 9 deletions(-) create mode 100644 source/source_main/driver_ipi.cpp create mode 100644 source/source_main/ipi_socket.cpp create mode 100644 source/source_main/ipi_socket.h diff --git a/docs/advanced/input_files/input-main.md b/docs/advanced/input_files/input-main.md index 334144f44c5..c6c1b732991 100644 --- a/docs/advanced/input_files/input-main.md +++ b/docs/advanced/input_files/input-main.md @@ -584,6 +584,7 @@ - relax: perform structure relaxation calculations, the relax_nmax parameter depicts the maximal number of ionic iterations - cell-relax: perform cell relaxation calculations - md: perform molecular dynamics simulations + - socket: run as a socket client for external drivers using the i-PI protocol - get_pchg: obtain partial (band-decomposed) charge densities (for LCAO basis only). See out_pchg for more information - get_wf: obtain real space wave functions (for LCAO basis only). See out_wfc_norm and out_wfc_re_im for more information - get_s: obtain the overlap matrix formed by localized orbitals (for LCAO basis with multiple k points). the file name is SR.csr with file format being the same as that generated by out_mat_hs2 @@ -840,7 +841,12 @@ ### chg_extrap - **Type**: String -- **Description**: Charge extrapolation method for MD and relaxation calculations. +- **Description**: Charge extrapolation method for MD, relaxation, and socket-driven calculations. + + When set to default, ABACUS chooses second-order for md, first-order for + relax/cell-relax/socket, and atomic for other calculations. Socket-driven + molecular dynamics can explicitly set second-order if the external driver + updates structures smoothly enough for second-order extrapolation. - **Default**: default ### nb2d diff --git a/docs/parameters.yaml b/docs/parameters.yaml index 00ef9b9dcf7..d183783ad81 100644 --- a/docs/parameters.yaml +++ b/docs/parameters.yaml @@ -28,6 +28,7 @@ parameters: * relax: perform structure relaxation calculations, the relax_nmax parameter depicts the maximal number of ionic iterations * cell-relax: perform cell relaxation calculations * md: perform molecular dynamics simulations + * socket: run as a socket client for external drivers using the i-PI protocol * get_pchg: obtain partial (band-decomposed) charge densities (for LCAO basis only). See out_pchg for more information * get_wf: obtain real space wave functions (for LCAO basis only). See out_wfc_norm and out_wfc_re_im for more information * get_s: obtain the overlap matrix formed by localized orbitals (for LCAO basis with multiple k points). the file name is SR.csr with file format being the same as that generated by out_mat_hs2 @@ -328,7 +329,12 @@ parameters: category: System variables type: String description: | - Charge extrapolation method for MD and relaxation calculations. + Charge extrapolation method for MD, relaxation, and socket-driven calculations. + + When set to default, ABACUS chooses second-order for md, first-order for + relax/cell-relax/socket, and atomic for other calculations. Socket-driven + molecular dynamics can explicitly set second-order if the external driver + updates structures smoothly enough for second-order extrapolation. default_value: default unit: "" availability: "" diff --git a/source/CMakeLists.txt b/source/CMakeLists.txt index 0337707b61a..a5b78bd3b50 100644 --- a/source/CMakeLists.txt +++ b/source/CMakeLists.txt @@ -477,7 +477,9 @@ add_library( driver OBJECT source_main/driver.cpp - source_main/driver_run.cpp) + source_main/driver_run.cpp + source_main/driver_ipi.cpp + source_main/ipi_socket.cpp) list(APPEND device_srcs source_pw/module_pwdft/kernels/nonlocal_op.cpp diff --git a/source/source_io/module_parameter/read_input_item_system.cpp b/source/source_io/module_parameter/read_input_item_system.cpp index 3b0ae011ec4..6ed572b21f9 100644 --- a/source/source_io/module_parameter/read_input_item_system.cpp +++ b/source/source_io/module_parameter/read_input_item_system.cpp @@ -75,7 +75,7 @@ void ReadInput::item_system() } { Input_Item item("calculation"); - item.annotation = "scf; relax; md; cell-relax; nscf; get_s; get_wf; get_pchg; gen_bessel; gen_opt_abfs; test_memory; test_neighbour"; + item.annotation = "scf; relax; md; socket; cell-relax; nscf; get_s; get_wf; get_pchg; gen_bessel; gen_opt_abfs; test_memory; test_neighbour"; item.category = "System variables"; item.type = "String"; item.description = R"(Specify the type of calculation. @@ -85,6 +85,7 @@ void ReadInput::item_system() * relax: perform structure relaxation calculations, the relax_nmax parameter depicts the maximal number of ionic iterations * cell-relax: perform cell relaxation calculations * md: perform molecular dynamics simulations +* socket: run as a socket client for external drivers using the i-PI protocol * get_pchg: obtain partial (band-decomposed) charge densities (for LCAO basis only). See out_pchg for more information * get_wf: obtain real space wave functions (for LCAO basis only). See out_wfc_norm and out_wfc_re_im for more information * get_s: obtain the overlap matrix formed by localized orbitals (for LCAO basis with multiple k points). the file name is SR.csr with file format being the same as that generated by out_mat_hs2 @@ -102,6 +103,7 @@ void ReadInput::item_system() std::vector callist = {"scf", "relax", "md", + "socket", "cell-relax", "nscf", "get_s", @@ -266,7 +268,7 @@ void ReadInput::item_system() item.description = "If set to True, calculate the force at the end of the electronic iteration."; item.default_value = "False"; item.reset_value = [](const Input_Item& item, Parameter& para) { - std::vector use_force = {"cell-relax", "relax", "md"}; + std::vector use_force = {"cell-relax", "relax", "md", "socket"}; std::vector not_use_force = {"get_wf", "get_pchg", "get_s"}; if (std::find(use_force.begin(), use_force.end(), para.input.calculation) != use_force.end()) { @@ -873,7 +875,12 @@ Available options are: item.annotation = "atomic; first-order; second-order; dm:coefficients of SIA"; item.category = "System variables"; item.type = "String"; - item.description = "Charge extrapolation method for MD and relaxation calculations."; + item.description = R"(Charge extrapolation method for MD, relaxation, and socket-driven calculations. + +When set to default, ABACUS chooses second-order for md, first-order for +relax/cell-relax/socket, and atomic for other calculations. Socket-driven +molecular dynamics can explicitly set second-order if the external driver +updates structures smoothly enough for second-order extrapolation.)"; item.default_value = "default"; read_sync_string(input.chg_extrap); item.reset_value = [](const Input_Item& item, Parameter& para) { @@ -882,7 +889,7 @@ Available options are: para.input.chg_extrap = "second-order"; } else if (para.input.chg_extrap == "default" - && (para.input.calculation == "relax" || para.input.calculation == "cell-relax")) + && (para.input.calculation == "relax" || para.input.calculation == "cell-relax" || para.input.calculation == "socket")) { para.input.chg_extrap = "first-order"; } diff --git a/source/source_io/test_serial/read_input_item_test.cpp b/source/source_io/test_serial/read_input_item_test.cpp index 61673a418c7..eb043fe1b03 100644 --- a/source/source_io/test_serial/read_input_item_test.cpp +++ b/source/source_io/test_serial/read_input_item_test.cpp @@ -362,6 +362,11 @@ TEST_F(InputTest, Item_test) it->second.reset_value(it->second, param); EXPECT_EQ(param.input.chg_extrap, "first-order"); + param.input.chg_extrap = "default"; + param.input.calculation = "socket"; + it->second.reset_value(it->second, param); + EXPECT_EQ(param.input.chg_extrap, "first-order"); + param.input.chg_extrap = "default"; param.input.calculation = "none"; it->second.reset_value(it->second, param); diff --git a/source/source_main/driver.h b/source/source_main/driver.h index e0204d079ca..5a51701e26c 100644 --- a/source/source_main/driver.h +++ b/source/source_main/driver.h @@ -36,6 +36,7 @@ class Driver // the actual calculations void driver_run(); + void driver_ipi_run(); // Init harewares according to Input parameters void init_hardware(); diff --git a/source/source_main/driver_ipi.cpp b/source/source_main/driver_ipi.cpp new file mode 100644 index 00000000000..7be2105a9b3 --- /dev/null +++ b/source/source_main/driver_ipi.cpp @@ -0,0 +1,507 @@ +#include "source_main/driver.h" + +#include "source_main/ipi_socket.h" +#include "source_base/global_function.h" +#include "source_base/mathzone.h" +#include "source_cell/check_atomic_stru.h" +#include "source_cell/update_cell.h" +#include "source_esolver/esolver.h" +#include "source_base/global_variable.h" +#include "source_io/module_json/para_json.h" +#include "source_io/module_output/print_info.h" +#include "source_io/module_parameter/parameter.h" + +#ifdef __MPI +#include +#endif + +#include +#include +#include +#include +#include +#include +#include + +namespace +{ +constexpr double RY_TO_HARTREE = 0.5; +constexpr int IPI_RANK_ROOT = 0; + +bool is_root() +{ + return GlobalV::MY_RANK == IPI_RANK_ROOT; +} + +void bcast_int(int& value) +{ +#ifdef __MPI + MPI_Bcast(&value, 1, MPI_INT, IPI_RANK_ROOT, MPI_COMM_WORLD); +#else + (void)value; +#endif +} + +void bcast_double_vector(std::vector& values) +{ +#ifdef __MPI + MPI_Bcast(values.data(), static_cast(values.size()), MPI_DOUBLE, IPI_RANK_ROOT, MPI_COMM_WORLD); +#else + (void)values; +#endif +} + +std::string bcast_string(std::string value) +{ + int nbytes = static_cast(value.size()); + bcast_int(nbytes); + if (nbytes < 0) + { + throw std::runtime_error("negative string length in i-PI broadcast"); + } + if (!is_root()) + { + value.assign(static_cast(nbytes), '\0'); + } +#ifdef __MPI + if (nbytes > 0) + { + MPI_Bcast(&value[0], nbytes, MPI_CHAR, IPI_RANK_ROOT, MPI_COMM_WORLD); + } +#endif + return value; +} + +void throw_if_root_io_failed(int root_failed, const std::string& root_message) +{ + bcast_int(root_failed); + const std::string message = bcast_string(root_message); + if (root_failed != 0) + { + throw std::runtime_error(message.empty() ? "i-PI socket I/O failed" : message); + } +} + +std::string bcast_header(const std::string& root_header) +{ + char buffer[13] = {' ', ' ', ' ', ' ', ' ', ' ', ' ', ' ', ' ', ' ', ' ', ' ', '\0'}; + if (is_root()) + { + const std::size_t n = root_header.size() > 12 ? 12 : root_header.size(); + for (std::size_t i = 0; i < n; ++i) + { + buffer[i] = root_header[i]; + } + } +#ifdef __MPI + MPI_Bcast(buffer, 12, MPI_CHAR, IPI_RANK_ROOT, MPI_COMM_WORLD); +#endif + std::string out(buffer, 12); + while (!out.empty() && out.back() == ' ') + { + out.pop_back(); + } + return out; +} + +std::string ipi_address() +{ + const char* env = std::getenv("ABACUS_IPI_ADDRESS"); + if (env == nullptr || std::string(env).empty()) + { + return "localhost:31415"; + } + return std::string(env); +} + +std::vector ipi_cell_bohr_from_unitcell(const UnitCell& ucell) +{ + const double lat0 = ucell.lat0; + // ASE/i-PI sends POSDATA cell as cell.T in C order. ABACUS stores + // lattice vectors as rows in latvec, so use the transposed order here. + return { + ucell.latvec.e11 * lat0, ucell.latvec.e21 * lat0, ucell.latvec.e31 * lat0, + ucell.latvec.e12 * lat0, ucell.latvec.e22 * lat0, ucell.latvec.e32 * lat0, + ucell.latvec.e13 * lat0, ucell.latvec.e23 * lat0, ucell.latvec.e33 * lat0, + }; +} + +double max_abs_delta(const std::vector& a, const std::vector& b) +{ + if (a.size() != b.size()) + { + return 1.0e99; + } + double out = 0.0; + for (std::size_t i = 0; i < a.size(); ++i) + { + out = std::max(out, std::abs(a[i] - b[i])); + } + return out; +} + +void set_positions_from_ipi_bohr(UnitCell& ucell, const std::vector& positions_bohr) +{ + if (positions_bohr.size() != static_cast(3 * ucell.nat)) + { + throw std::runtime_error("POSDATA atom count does not match STRU."); + } + int iat = 0; + for (int it = 0; it < ucell.ntype; ++it) + { + Atom* atom = &ucell.atoms[it]; + for (int ia = 0; ia < atom->na; ++ia) + { + const double tau_x = positions_bohr[3 * iat + 0] / ucell.lat0; + const double tau_y = positions_bohr[3 * iat + 1] / ucell.lat0; + const double tau_z = positions_bohr[3 * iat + 2] / ucell.lat0; + + double dx = 0.0; + double dy = 0.0; + double dz = 0.0; + ModuleBase::Mathzone::Cartesian_to_Direct(tau_x, + tau_y, + tau_z, + ucell.latvec.e11, + ucell.latvec.e12, + ucell.latvec.e13, + ucell.latvec.e21, + ucell.latvec.e22, + ucell.latvec.e23, + ucell.latvec.e31, + ucell.latvec.e32, + ucell.latvec.e33, + dx, + dy, + dz); + + atom->dis[ia].x = dx - atom->taud[ia].x; + atom->dis[ia].y = dy - atom->taud[ia].y; + atom->dis[ia].z = dz - atom->taud[ia].z; + atom->taud[ia].x = dx; + atom->taud[ia].y = dy; + atom->taud[ia].z = dz; + atom->tau[ia].x = tau_x; + atom->tau[ia].y = tau_y; + atom->tau[ia].z = tau_z; + ++iat; + } + } + unitcell::periodic_boundary_adjustment(ucell.atoms, ucell.latvec, ucell.ntype); + ucell.ionic_position_updated = true; + ucell.cell_parameter_updated = false; +} + +std::vector flatten_forces_hartree_per_bohr(const ModuleBase::matrix& force) +{ + std::vector out(static_cast(force.nr * force.nc)); + for (int iat = 0; iat < force.nr; ++iat) + { + for (int idir = 0; idir < force.nc; ++idir) + { + out[static_cast(3 * iat + idir)] = force(iat, idir) * RY_TO_HARTREE; + } + } + return out; +} + +class CalculationModeGuard +{ + public: + explicit CalculationModeGuard(const std::string& inner_calculation) + : outer_calculation_(PARAM.inp.calculation) + { + const_cast(PARAM.inp.calculation) = inner_calculation; + } + + ~CalculationModeGuard() + { + const_cast(PARAM.inp.calculation) = outer_calculation_; + } + + CalculationModeGuard(const CalculationModeGuard&) = delete; + CalculationModeGuard& operator=(const CalculationModeGuard&) = delete; + + private: + std::string outer_calculation_; +}; +} // namespace + +void Driver::driver_ipi_run() +{ + ModuleBase::TITLE("Driver", "driver_ipi_run"); + + // "socket" is an outer driver mode. The KS/LCAO ESolver internals use the + // standard SCF code path for each POSDATA request from the i-PI protocol. + CalculationModeGuard calculation_guard("scf"); + + UnitCell ucell; + ucell.setup(PARAM.inp.latname, PARAM.inp.ntype, PARAM.inp.lmaxmax, PARAM.inp.init_vel, PARAM.inp.fixed_axes); + ucell.setup_cell(PARAM.globalv.global_in_stru, GlobalV::ofs_running); + unitcell::check_atomic_stru(ucell, PARAM.inp.min_dist_coef); + + IpiSocket socket; + std::unique_ptr p_esolver; + bool hardware_initialized = false; + bool esolver_ready = false; + bool runner_completed = false; + std::string pending_error; + + try + { + this->init_hardware(); + hardware_initialized = true; + + p_esolver.reset(ModuleESolver::init_esolver(PARAM.inp, ucell)); + p_esolver->before_all_runners(ucell, PARAM.inp); + esolver_ready = true; + +#ifdef __RAPIDJSON + Json::gen_stru_wrapper(&ucell); +#endif + + int io_failed = 0; + std::string io_message; + if (is_root()) + { + try + { + const std::string address = ipi_address(); + GlobalV::ofs_running << " ABACUS socket driver connecting to i-PI endpoint " << address << std::endl; + socket.connect(address); + } + catch (const std::exception& exc) + { + io_failed = 1; + io_message = exc.what(); + } + } + throw_if_root_io_failed(io_failed, io_message); + + bool isinit = false; + bool hasdata = false; + int istep = 0; + int nat_return = ucell.nat; + double energy_hartree = 0.0; + std::vector forces_hartree_bohr(static_cast(3 * ucell.nat), 0.0); + std::vector virial_hartree(9, 0.0); + + const std::vector reference_cell = ipi_cell_bohr_from_unitcell(ucell); + + while (true) + { + std::string header; + io_failed = 0; + io_message.clear(); + if (is_root()) + { + try + { + header = socket.read_header(); + } + catch (const std::exception& exc) + { + io_failed = 1; + io_message = exc.what(); + } + } + throw_if_root_io_failed(io_failed, io_message); + header = bcast_header(header); + + if (header == "STATUS") + { + io_failed = 0; + io_message.clear(); + if (is_root()) + { + try + { + if (hasdata) + { + socket.write_header("HAVEDATA"); + } + else if (isinit) + { + socket.write_header("READY"); + } + else + { + socket.write_header("NEEDINIT"); + } + } + catch (const std::exception& exc) + { + io_failed = 1; + io_message = exc.what(); + } + } + throw_if_root_io_failed(io_failed, io_message); + } + else if (header == "INIT") + { + int rid = 0; + int nbytes = 0; + std::string params; + io_failed = 0; + io_message.clear(); + if (is_root()) + { + try + { + rid = socket.read_int(); + nbytes = socket.read_int(); + if (nbytes < 0) + { + throw std::runtime_error("negative INIT payload length from i-PI socket"); + } + if (nbytes > 0) + { + params = socket.read_string(static_cast(nbytes)); + } + } + catch (const std::exception& exc) + { + io_failed = 1; + io_message = exc.what(); + } + } + throw_if_root_io_failed(io_failed, io_message); + bcast_int(rid); + bcast_int(nbytes); + if (nbytes > 0 && is_root()) + { + GlobalV::ofs_running << " ABACUS socket INIT params " << params << std::endl; + } + isinit = true; + if (is_root()) + { + GlobalV::ofs_running << " ABACUS socket INIT replica " << rid << std::endl; + } + } + else if (header == "POSDATA") + { + std::vector cell(9, 0.0); + std::vector inv_cell(9, 0.0); + int nat_socket = 0; + std::vector positions; + io_failed = 0; + io_message.clear(); + if (is_root()) + { + try + { + cell = socket.read_doubles(9); + inv_cell = socket.read_doubles(9); + nat_socket = socket.read_int(); + if (nat_socket < 0) + { + throw std::runtime_error("negative POSDATA atom count from i-PI socket"); + } + positions = socket.read_doubles(static_cast(3 * nat_socket)); + } + catch (const std::exception& exc) + { + io_failed = 1; + io_message = exc.what(); + } + } + throw_if_root_io_failed(io_failed, io_message); + bcast_double_vector(cell); + bcast_double_vector(inv_cell); + bcast_int(nat_socket); + if (!is_root()) + { + positions.assign(static_cast(3 * nat_socket), 0.0); + } + bcast_double_vector(positions); + + if (nat_socket != ucell.nat) + { + throw std::runtime_error("POSDATA atom count does not match STRU."); + } + const double max_cell_delta_bohr = max_abs_delta(cell, reference_cell); + if (max_cell_delta_bohr > 1.0e-6) + { + throw std::runtime_error("variable-cell socket updates are not supported yet."); + } + + set_positions_from_ipi_bohr(ucell, positions); + p_esolver->runner(ucell, istep); + runner_completed = true; + energy_hartree = p_esolver->cal_energy() * RY_TO_HARTREE; + ModuleBase::matrix force; + if (PARAM.inp.cal_force) + { + p_esolver->cal_force(ucell, force); + forces_hartree_bohr = flatten_forces_hartree_per_bohr(force); + } + ++istep; + hasdata = true; + } + else if (header == "GETFORCE") + { + io_failed = 0; + io_message.clear(); + if (is_root()) + { + try + { + socket.write_header("FORCEREADY"); + socket.write_double(energy_hartree); + socket.write_int(nat_return); + socket.write_doubles(forces_hartree_bohr); + socket.write_doubles(virial_hartree); + socket.write_int(0); + } + catch (const std::exception& exc) + { + io_failed = 1; + io_message = exc.what(); + } + } + throw_if_root_io_failed(io_failed, io_message); + isinit = false; + hasdata = false; + } + else + { + if (is_root()) + { + GlobalV::ofs_running << " ABACUS socket driver exiting on header " << header << std::endl; + } + break; + } + } + } + catch (const std::exception& exc) + { + pending_error = exc.what(); + if (is_root()) + { + GlobalV::ofs_running << " ABACUS socket driver ended with error: " << pending_error << std::endl; + } + } + + if (is_root()) + { + socket.close(); + } + if (esolver_ready && runner_completed && p_esolver) + { + p_esolver->after_all_runners(ucell); + } + p_esolver.reset(); + if (hardware_initialized) + { + this->finalize_hardware(); + } + +#ifdef __RAPIDJSON + Json::create_Json(&ucell, PARAM); +#endif + + if (!pending_error.empty()) + { + ModuleBase::WARNING_QUIT("ABACUS socket", pending_error); + } +} diff --git a/source/source_main/driver_run.cpp b/source/source_main/driver_run.cpp index 4911b133de0..9285f9b2330 100644 --- a/source/source_main/driver_run.cpp +++ b/source/source_main/driver_run.cpp @@ -38,6 +38,13 @@ void Driver::driver_run() { ModuleBase::TITLE("Driver", "driver_run"); + const std::string cal = PARAM.inp.calculation; + if (cal == "socket") + { + this->driver_ipi_run(); + return; + } + //! 1: setup cell and atom information // this warning should not be here, mohan 2024-05-22 #ifndef __LCAO @@ -71,8 +78,6 @@ void Driver::driver_run() Json::gen_stru_wrapper(&ucell); #endif - const std::string cal = PARAM.inp.calculation; - //! 4: different types of calculations if (cal == "md") { diff --git a/source/source_main/ipi_socket.cpp b/source/source_main/ipi_socket.cpp new file mode 100644 index 00000000000..eca053566ff --- /dev/null +++ b/source/source_main/ipi_socket.cpp @@ -0,0 +1,241 @@ +#include "ipi_socket.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace +{ +constexpr std::size_t IPI_HEADER_LEN = 12; + +std::string errno_message(const std::string& prefix) +{ + return prefix + ": " + std::strerror(errno); +} + +std::string trim_header(const char* data) +{ + std::string value(data, IPI_HEADER_LEN); + while (!value.empty() && value.back() == ' ') + { + value.pop_back(); + } + return value; +} + +std::string padded_header(const std::string& header) +{ + if (header.size() > IPI_HEADER_LEN) + { + throw std::runtime_error("i-PI header is longer than 12 bytes: " + header); + } + std::string out = header; + out.resize(IPI_HEADER_LEN, ' '); + return out; +} +} // namespace + +IpiSocket::~IpiSocket() +{ + this->close(); +} + +void IpiSocket::connect(const std::string& address) +{ + this->close(); + const std::size_t colon = address.rfind(':'); + if (colon == std::string::npos) + { + throw std::runtime_error("i-PI address must be host:port or path:UNIX, got " + address); + } + const std::string host = address.substr(0, colon); + const std::string service = address.substr(colon + 1); + + if (service == "UNIX") + { + fd_ = ::socket(AF_UNIX, SOCK_STREAM, 0); + if (fd_ < 0) + { + throw std::runtime_error(errno_message("failed to create UNIX socket")); + } + sockaddr_un addr; + std::memset(&addr, 0, sizeof(addr)); + addr.sun_family = AF_UNIX; + if (host.size() >= sizeof(addr.sun_path)) + { + this->close(); + throw std::runtime_error("UNIX socket path too long: " + host); + } + std::strncpy(addr.sun_path, host.c_str(), sizeof(addr.sun_path) - 1); + if (::connect(fd_, reinterpret_cast(&addr), sizeof(addr)) != 0) + { + const std::string msg = errno_message("failed to connect UNIX i-PI socket " + host); + this->close(); + throw std::runtime_error(msg); + } + return; + } + + addrinfo hints; + std::memset(&hints, 0, sizeof(hints)); + hints.ai_family = AF_UNSPEC; + hints.ai_socktype = SOCK_STREAM; + + addrinfo* result = nullptr; + const int gai = ::getaddrinfo(host.c_str(), service.c_str(), &hints, &result); + if (gai != 0) + { + throw std::runtime_error("failed to resolve i-PI socket " + address + ": " + ::gai_strerror(gai)); + } + + std::string last_error; + for (addrinfo* rp = result; rp != nullptr; rp = rp->ai_next) + { + fd_ = ::socket(rp->ai_family, rp->ai_socktype, rp->ai_protocol); + if (fd_ < 0) + { + last_error = errno_message("failed to create INET socket"); + continue; + } + if (::connect(fd_, rp->ai_addr, rp->ai_addrlen) == 0) + { + ::freeaddrinfo(result); + return; + } + last_error = errno_message("failed to connect INET i-PI socket " + address); + this->close(); + } + ::freeaddrinfo(result); + throw std::runtime_error(last_error.empty() ? "failed to connect i-PI socket " + address : last_error); +} + +void IpiSocket::close() +{ + if (fd_ >= 0) + { + ::close(fd_); + fd_ = -1; + } +} + +std::string IpiSocket::read_header() +{ + char header[IPI_HEADER_LEN]; + this->read_exact(header, sizeof(header)); + return trim_header(header); +} + +void IpiSocket::write_header(const std::string& header) +{ + const std::string padded = padded_header(header); + this->write_exact(padded.data(), padded.size()); +} + +int IpiSocket::read_int() +{ + int value = 0; + this->read_exact(&value, sizeof(value)); + return value; +} + +void IpiSocket::write_int(int value) +{ + this->write_exact(&value, sizeof(value)); +} + +double IpiSocket::read_double() +{ + double value = 0.0; + this->read_exact(&value, sizeof(value)); + return value; +} + +void IpiSocket::write_double(double value) +{ + this->write_exact(&value, sizeof(value)); +} + +std::vector IpiSocket::read_doubles(std::size_t n) +{ + std::vector values(n); + if (!values.empty()) + { + this->read_exact(values.data(), values.size() * sizeof(double)); + } + return values; +} + +void IpiSocket::write_doubles(const std::vector& values) +{ + if (!values.empty()) + { + this->write_exact(values.data(), values.size() * sizeof(double)); + } +} + +std::string IpiSocket::read_string(std::size_t nbytes) +{ + std::string value(nbytes, '\0'); + if (nbytes > 0) + { + this->read_exact(&value[0], nbytes); + } + return value; +} + +void IpiSocket::read_exact(void* data, std::size_t nbytes) +{ + char* cursor = static_cast(data); + std::size_t done = 0; + while (done < nbytes) + { + const ssize_t nread = ::recv(fd_, cursor + done, nbytes - done, 0); + if (nread == 0) + { + throw std::runtime_error("i-PI socket closed while reading"); + } + if (nread < 0) + { + if (errno == EINTR) + { + continue; + } + throw std::runtime_error(errno_message("i-PI socket read failed")); + } + done += static_cast(nread); + } +} + +void IpiSocket::write_exact(const void* data, std::size_t nbytes) +{ + const char* cursor = static_cast(data); + std::size_t done = 0; + while (done < nbytes) + { +#ifdef MSG_NOSIGNAL + const int flags = MSG_NOSIGNAL; +#else + const int flags = 0; +#endif + const ssize_t nwritten = ::send(fd_, cursor + done, nbytes - done, flags); + if (nwritten == 0) + { + throw std::runtime_error("i-PI socket closed while writing"); + } + if (nwritten < 0) + { + if (errno == EINTR) + { + continue; + } + throw std::runtime_error(errno_message("i-PI socket write failed")); + } + done += static_cast(nwritten); + } +} diff --git a/source/source_main/ipi_socket.h b/source/source_main/ipi_socket.h new file mode 100644 index 00000000000..3e1713e868d --- /dev/null +++ b/source/source_main/ipi_socket.h @@ -0,0 +1,40 @@ +#ifndef ABACUS_IPI_SOCKET_H +#define ABACUS_IPI_SOCKET_H + +#include +#include +#include + +class IpiSocket +{ + public: + IpiSocket() = default; + ~IpiSocket(); + + IpiSocket(const IpiSocket&) = delete; + IpiSocket& operator=(const IpiSocket&) = delete; + + void connect(const std::string& address); + void close(); + + std::string read_header(); + void write_header(const std::string& header); + + int read_int(); + void write_int(int value); + + double read_double(); + void write_double(double value); + + std::vector read_doubles(std::size_t n); + void write_doubles(const std::vector& values); + std::string read_string(std::size_t nbytes); + + private: + int fd_ = -1; + + void read_exact(void* data, std::size_t nbytes); + void write_exact(const void* data, std::size_t nbytes); +}; + +#endif From e0f70ca031422c4821d5304050c0035f1be85ec7 Mon Sep 17 00:00:00 2001 From: jianrui geng Date: Thu, 9 Jul 2026 09:13:53 +0800 Subject: [PATCH 2/4] Refactor i-PI socket driver into source_relax Route calculation=socket through Relax_Driver, keep ESolver lifecycle in driver_run, and add source_relax socket transport tests. --- source/CMakeLists.txt | 4 +- source/source_main/driver.h | 1 - source/source_main/driver_run.cpp | 22 +- source/source_relax/CMakeLists.txt | 2 + .../ipi_socket.cpp | 31 ++- .../ipi_socket.h | 7 + source/source_relax/relax_driver.cpp | 9 + .../socket_driver.cpp} | 243 ++++++----------- source/source_relax/socket_driver.h | 22 ++ source/source_relax/test/CMakeLists.txt | 6 + source/source_relax/test/ipi_socket_test.cpp | 256 ++++++++++++++++++ 11 files changed, 431 insertions(+), 172 deletions(-) rename source/{source_main => source_relax}/ipi_socket.cpp (87%) rename source/{source_main => source_relax}/ipi_socket.h (85%) rename source/{source_main/driver_ipi.cpp => source_relax/socket_driver.cpp} (64%) create mode 100644 source/source_relax/socket_driver.h create mode 100644 source/source_relax/test/ipi_socket_test.cpp diff --git a/source/CMakeLists.txt b/source/CMakeLists.txt index a5b78bd3b50..0337707b61a 100644 --- a/source/CMakeLists.txt +++ b/source/CMakeLists.txt @@ -477,9 +477,7 @@ add_library( driver OBJECT source_main/driver.cpp - source_main/driver_run.cpp - source_main/driver_ipi.cpp - source_main/ipi_socket.cpp) + source_main/driver_run.cpp) list(APPEND device_srcs source_pw/module_pwdft/kernels/nonlocal_op.cpp diff --git a/source/source_main/driver.h b/source/source_main/driver.h index 5a51701e26c..e0204d079ca 100644 --- a/source/source_main/driver.h +++ b/source/source_main/driver.h @@ -36,7 +36,6 @@ class Driver // the actual calculations void driver_run(); - void driver_ipi_run(); // Init harewares according to Input parameters void init_hardware(); diff --git a/source/source_main/driver_run.cpp b/source/source_main/driver_run.cpp index 9285f9b2330..3e7d669b6c2 100644 --- a/source/source_main/driver_run.cpp +++ b/source/source_main/driver_run.cpp @@ -39,11 +39,7 @@ void Driver::driver_run() ModuleBase::TITLE("Driver", "driver_run"); const std::string cal = PARAM.inp.calculation; - if (cal == "socket") - { - this->driver_ipi_run(); - return; - } + const bool socket_mode = (cal == "socket"); //! 1: setup cell and atom information // this warning should not be here, mohan 2024-05-22 @@ -66,12 +62,20 @@ void Driver::driver_run() unitcell::check_atomic_stru(ucell, PARAM.inp.min_dist_coef); //! 2: initialize the ESolver (depends on a set-up ucell after `setup_cell`) + Input_para esolver_inp = PARAM.inp; + Input_para driver_inp = PARAM.inp; + if (socket_mode) + { + esolver_inp.calculation = "scf"; + driver_inp.calculation = cal; + } + this->init_hardware(); - ModuleESolver::ESolver* p_esolver = ModuleESolver::init_esolver(PARAM.inp, ucell); + ModuleESolver::ESolver* p_esolver = ModuleESolver::init_esolver(esolver_inp, ucell); //! 3: initialize Esolver and fill json-structure - p_esolver->before_all_runners(ucell, PARAM.inp); + p_esolver->before_all_runners(ucell, esolver_inp); // this Json part should be moved to before_all_runners, mohan 2024-05-12 #ifdef __RAPIDJSON @@ -83,10 +87,10 @@ void Driver::driver_run() { Run_MD::md_line(ucell, p_esolver, PARAM); } - else if (cal == "scf" || cal == "relax" || cal == "cell-relax" || cal == "nscf") + else if (cal == "scf" || cal == "relax" || cal == "cell-relax" || cal == "nscf" || cal == "socket") { Relax_Driver rl_driver; - rl_driver.relax_driver(p_esolver, ucell, PARAM.inp, GlobalV::ofs_running); + rl_driver.relax_driver(p_esolver, ucell, driver_inp, GlobalV::ofs_running); } else if (cal == "get_s") { diff --git a/source/source_relax/CMakeLists.txt b/source/source_relax/CMakeLists.txt index 9e7ef96b0e2..84f69c06c80 100644 --- a/source/source_relax/CMakeLists.txt +++ b/source/source_relax/CMakeLists.txt @@ -2,6 +2,8 @@ add_library( relax OBJECT relax_data.cpp + ipi_socket.cpp + socket_driver.cpp cg_base.cpp relax_driver.cpp relax_sync.cpp diff --git a/source/source_main/ipi_socket.cpp b/source/source_relax/ipi_socket.cpp similarity index 87% rename from source/source_main/ipi_socket.cpp rename to source/source_relax/ipi_socket.cpp index eca053566ff..304a0181813 100644 --- a/source/source_main/ipi_socket.cpp +++ b/source/source_relax/ipi_socket.cpp @@ -1,4 +1,4 @@ -#include "ipi_socket.h" +#include "source_relax/ipi_socket.h" #include #include @@ -41,6 +41,10 @@ std::string padded_header(const std::string& header) } } // namespace +IpiSocketClosed::IpiSocketClosed(const std::string& message) : std::runtime_error(message) +{ +} + IpiSocket::~IpiSocket() { this->close(); @@ -127,7 +131,28 @@ void IpiSocket::close() std::string IpiSocket::read_header() { char header[IPI_HEADER_LEN]; - this->read_exact(header, sizeof(header)); + std::size_t done = 0; + while (done < sizeof(header)) + { + const ssize_t nread = ::recv(fd_, header + done, sizeof(header) - done, 0); + if (nread == 0) + { + if (done == 0) + { + throw IpiSocketClosed("i-PI socket closed before next header"); + } + throw std::runtime_error("i-PI socket closed while reading header"); + } + if (nread < 0) + { + if (errno == EINTR) + { + continue; + } + throw std::runtime_error(errno_message("i-PI socket header read failed")); + } + done += static_cast(nread); + } return trim_header(header); } @@ -198,7 +223,7 @@ void IpiSocket::read_exact(void* data, std::size_t nbytes) const ssize_t nread = ::recv(fd_, cursor + done, nbytes - done, 0); if (nread == 0) { - throw std::runtime_error("i-PI socket closed while reading"); + throw IpiSocketClosed("i-PI socket closed while reading"); } if (nread < 0) { diff --git a/source/source_main/ipi_socket.h b/source/source_relax/ipi_socket.h similarity index 85% rename from source/source_main/ipi_socket.h rename to source/source_relax/ipi_socket.h index 3e1713e868d..28ce2cb81ff 100644 --- a/source/source_main/ipi_socket.h +++ b/source/source_relax/ipi_socket.h @@ -2,9 +2,16 @@ #define ABACUS_IPI_SOCKET_H #include +#include #include #include +class IpiSocketClosed : public std::runtime_error +{ + public: + explicit IpiSocketClosed(const std::string& message); +}; + class IpiSocket { public: diff --git a/source/source_relax/relax_driver.cpp b/source/source_relax/relax_driver.cpp index 26128bf0408..0cada1bfc3e 100644 --- a/source/source_relax/relax_driver.cpp +++ b/source/source_relax/relax_driver.cpp @@ -1,4 +1,5 @@ #include "relax_driver.h" +#include "socket_driver.h" #include "source_base/global_file.h" #include "source_io/module_output/cif_io.h" #include "source_io/module_json/output_info.h" @@ -17,6 +18,14 @@ void Relax_Driver::relax_driver( ModuleBase::TITLE("Relax_Driver", "relax_driver"); ModuleBase::timer::start("Relax_Driver", "relax_driver"); + if (inp.calculation == "socket") + { + Socket_Driver socket_driver; + socket_driver.socket_driver(p_esolver, ucell, inp, ofs_running); + ModuleBase::timer::end("Relax_Driver", "relax_driver"); + return; + } + this->init_relax(ucell.nat, inp); // steps[0]: istep (main iteration step) diff --git a/source/source_main/driver_ipi.cpp b/source/source_relax/socket_driver.cpp similarity index 64% rename from source/source_main/driver_ipi.cpp rename to source/source_relax/socket_driver.cpp index 7be2105a9b3..02928fec9d3 100644 --- a/source/source_main/driver_ipi.cpp +++ b/source/source_relax/socket_driver.cpp @@ -1,25 +1,17 @@ -#include "source_main/driver.h" +#include "socket_driver.h" -#include "source_main/ipi_socket.h" +#include "source_relax/ipi_socket.h" #include "source_base/global_function.h" +#include "source_base/global_variable.h" #include "source_base/mathzone.h" -#include "source_cell/check_atomic_stru.h" +#include "source_base/parallel_common.h" +#include "source_base/timer.h" #include "source_cell/update_cell.h" -#include "source_esolver/esolver.h" -#include "source_base/global_variable.h" -#include "source_io/module_json/para_json.h" -#include "source_io/module_output/print_info.h" -#include "source_io/module_parameter/parameter.h" - -#ifdef __MPI -#include -#endif #include #include #include -#include -#include +#include #include #include @@ -33,75 +25,42 @@ bool is_root() return GlobalV::MY_RANK == IPI_RANK_ROOT; } -void bcast_int(int& value) -{ -#ifdef __MPI - MPI_Bcast(&value, 1, MPI_INT, IPI_RANK_ROOT, MPI_COMM_WORLD); -#else - (void)value; -#endif -} - void bcast_double_vector(std::vector& values) { -#ifdef __MPI - MPI_Bcast(values.data(), static_cast(values.size()), MPI_DOUBLE, IPI_RANK_ROOT, MPI_COMM_WORLD); -#else - (void)values; -#endif + if (!values.empty()) + { + Parallel_Common::bcast_double(values.data(), static_cast(values.size())); + } } -std::string bcast_string(std::string value) +void bcast_socket_string(std::string& value) { - int nbytes = static_cast(value.size()); - bcast_int(nbytes); - if (nbytes < 0) - { - throw std::runtime_error("negative string length in i-PI broadcast"); - } + int size = static_cast(value.size()); + Parallel_Common::bcast_int(size); if (!is_root()) { - value.assign(static_cast(nbytes), '\0'); + value.resize(static_cast(size)); } -#ifdef __MPI - if (nbytes > 0) + if (size > 0) { - MPI_Bcast(&value[0], nbytes, MPI_CHAR, IPI_RANK_ROOT, MPI_COMM_WORLD); + Parallel_Common::bcast_char(&value[0], size); } -#endif - return value; } -void throw_if_root_io_failed(int root_failed, const std::string& root_message) +void quit_if_root_io_failed(int root_failed, std::string root_message) { - bcast_int(root_failed); - const std::string message = bcast_string(root_message); + Parallel_Common::bcast_int(root_failed); + bcast_socket_string(root_message); if (root_failed != 0) { - throw std::runtime_error(message.empty() ? "i-PI socket I/O failed" : message); + ModuleBase::WARNING_QUIT("ABACUS socket", root_message.empty() ? "i-PI socket I/O failed" : root_message); } } -std::string bcast_header(const std::string& root_header) +std::string bcast_header(std::string header) { - char buffer[13] = {' ', ' ', ' ', ' ', ' ', ' ', ' ', ' ', ' ', ' ', ' ', ' ', '\0'}; - if (is_root()) - { - const std::size_t n = root_header.size() > 12 ? 12 : root_header.size(); - for (std::size_t i = 0; i < n; ++i) - { - buffer[i] = root_header[i]; - } - } -#ifdef __MPI - MPI_Bcast(buffer, 12, MPI_CHAR, IPI_RANK_ROOT, MPI_COMM_WORLD); -#endif - std::string out(buffer, 12); - while (!out.empty() && out.back() == ' ') - { - out.pop_back(); - } - return out; + bcast_socket_string(header); + return header; } std::string ipi_address() @@ -144,8 +103,9 @@ void set_positions_from_ipi_bohr(UnitCell& ucell, const std::vector& pos { if (positions_bohr.size() != static_cast(3 * ucell.nat)) { - throw std::runtime_error("POSDATA atom count does not match STRU."); + ModuleBase::WARNING_QUIT("ABACUS socket", "POSDATA atom count does not match STRU."); } + int iat = 0; for (int it = 0; it < ucell.ntype; ++it) { @@ -204,62 +164,29 @@ std::vector flatten_forces_hartree_per_bohr(const ModuleBase::matrix& fo } return out; } +} // namespace -class CalculationModeGuard +void Socket_Driver::socket_driver(ModuleESolver::ESolver* p_esolver, + UnitCell& ucell, + const Input_para& inp, + std::ofstream& ofs_running) { - public: - explicit CalculationModeGuard(const std::string& inner_calculation) - : outer_calculation_(PARAM.inp.calculation) + ModuleBase::TITLE("Socket_Driver", "socket_driver"); + ModuleBase::timer::start("Socket_Driver", "socket_driver"); + + if (p_esolver == nullptr) { - const_cast(PARAM.inp.calculation) = inner_calculation; + ModuleBase::WARNING_QUIT("ABACUS socket", "socket driver requires a valid ESolver."); } - - ~CalculationModeGuard() + if (!inp.cal_force) { - const_cast(PARAM.inp.calculation) = outer_calculation_; + ModuleBase::WARNING_QUIT("ABACUS socket", "socket calculation requires cal_force=1 for i-PI GETFORCE."); } - CalculationModeGuard(const CalculationModeGuard&) = delete; - CalculationModeGuard& operator=(const CalculationModeGuard&) = delete; - - private: - std::string outer_calculation_; -}; -} // namespace - -void Driver::driver_ipi_run() -{ - ModuleBase::TITLE("Driver", "driver_ipi_run"); - - // "socket" is an outer driver mode. The KS/LCAO ESolver internals use the - // standard SCF code path for each POSDATA request from the i-PI protocol. - CalculationModeGuard calculation_guard("scf"); - - UnitCell ucell; - ucell.setup(PARAM.inp.latname, PARAM.inp.ntype, PARAM.inp.lmaxmax, PARAM.inp.init_vel, PARAM.inp.fixed_axes); - ucell.setup_cell(PARAM.globalv.global_in_stru, GlobalV::ofs_running); - unitcell::check_atomic_stru(ucell, PARAM.inp.min_dist_coef); - IpiSocket socket; - std::unique_ptr p_esolver; - bool hardware_initialized = false; - bool esolver_ready = false; - bool runner_completed = false; - std::string pending_error; try { - this->init_hardware(); - hardware_initialized = true; - - p_esolver.reset(ModuleESolver::init_esolver(PARAM.inp, ucell)); - p_esolver->before_all_runners(ucell, PARAM.inp); - esolver_ready = true; - -#ifdef __RAPIDJSON - Json::gen_stru_wrapper(&ucell); -#endif - int io_failed = 0; std::string io_message; if (is_root()) @@ -267,7 +194,7 @@ void Driver::driver_ipi_run() try { const std::string address = ipi_address(); - GlobalV::ofs_running << " ABACUS socket driver connecting to i-PI endpoint " << address << std::endl; + ofs_running << " ABACUS socket driver connecting to i-PI endpoint " << address << std::endl; socket.connect(address); } catch (const std::exception& exc) @@ -276,12 +203,12 @@ void Driver::driver_ipi_run() io_message = exc.what(); } } - throw_if_root_io_failed(io_failed, io_message); + quit_if_root_io_failed(io_failed, io_message); bool isinit = false; bool hasdata = false; int istep = 0; - int nat_return = ucell.nat; + const int nat_return = ucell.nat; double energy_hartree = 0.0; std::vector forces_hartree_bohr(static_cast(3 * ucell.nat), 0.0); std::vector virial_hartree(9, 0.0); @@ -299,16 +226,28 @@ void Driver::driver_ipi_run() { header = socket.read_header(); } + catch (const IpiSocketClosed&) + { + header.clear(); + } catch (const std::exception& exc) { io_failed = 1; io_message = exc.what(); } } - throw_if_root_io_failed(io_failed, io_message); + quit_if_root_io_failed(io_failed, io_message); header = bcast_header(header); - if (header == "STATUS") + if (header.empty()) + { + if (is_root()) + { + ofs_running << " ABACUS socket driver exiting after peer closed connection" << std::endl; + } + break; + } + else if (header == "STATUS") { io_failed = 0; io_message.clear(); @@ -335,7 +274,7 @@ void Driver::driver_ipi_run() io_message = exc.what(); } } - throw_if_root_io_failed(io_failed, io_message); + quit_if_root_io_failed(io_failed, io_message); } else if (header == "INIT") { @@ -352,9 +291,10 @@ void Driver::driver_ipi_run() nbytes = socket.read_int(); if (nbytes < 0) { - throw std::runtime_error("negative INIT payload length from i-PI socket"); + io_failed = 1; + io_message = "negative INIT payload length from i-PI socket"; } - if (nbytes > 0) + else if (nbytes > 0) { params = socket.read_string(static_cast(nbytes)); } @@ -365,17 +305,17 @@ void Driver::driver_ipi_run() io_message = exc.what(); } } - throw_if_root_io_failed(io_failed, io_message); - bcast_int(rid); - bcast_int(nbytes); + quit_if_root_io_failed(io_failed, io_message); + Parallel_Common::bcast_int(rid); + Parallel_Common::bcast_int(nbytes); if (nbytes > 0 && is_root()) { - GlobalV::ofs_running << " ABACUS socket INIT params " << params << std::endl; + ofs_running << " ABACUS socket INIT params bytes " << nbytes << std::endl; } isinit = true; if (is_root()) { - GlobalV::ofs_running << " ABACUS socket INIT replica " << rid << std::endl; + ofs_running << " ABACUS socket INIT replica " << rid << std::endl; } } else if (header == "POSDATA") @@ -395,9 +335,13 @@ void Driver::driver_ipi_run() nat_socket = socket.read_int(); if (nat_socket < 0) { - throw std::runtime_error("negative POSDATA atom count from i-PI socket"); + io_failed = 1; + io_message = "negative POSDATA atom count from i-PI socket"; + } + else + { + positions = socket.read_doubles(static_cast(3 * nat_socket)); } - positions = socket.read_doubles(static_cast(3 * nat_socket)); } catch (const std::exception& exc) { @@ -405,10 +349,10 @@ void Driver::driver_ipi_run() io_message = exc.what(); } } - throw_if_root_io_failed(io_failed, io_message); + quit_if_root_io_failed(io_failed, io_message); bcast_double_vector(cell); bcast_double_vector(inv_cell); - bcast_int(nat_socket); + Parallel_Common::bcast_int(nat_socket); if (!is_root()) { positions.assign(static_cast(3 * nat_socket), 0.0); @@ -417,20 +361,27 @@ void Driver::driver_ipi_run() if (nat_socket != ucell.nat) { - throw std::runtime_error("POSDATA atom count does not match STRU."); + ModuleBase::WARNING_QUIT("ABACUS socket", "POSDATA atom count does not match STRU."); } const double max_cell_delta_bohr = max_abs_delta(cell, reference_cell); if (max_cell_delta_bohr > 1.0e-6) { - throw std::runtime_error("variable-cell socket updates are not supported yet."); + ModuleBase::WARNING_QUIT("ABACUS socket", "variable-cell socket updates are not supported yet."); } set_positions_from_ipi_bohr(ucell, positions); p_esolver->runner(ucell, istep); - runner_completed = true; - energy_hartree = p_esolver->cal_energy() * RY_TO_HARTREE; + const double energy_ry = p_esolver->cal_energy(); + energy_hartree = energy_ry * RY_TO_HARTREE; + if (is_root()) + { + ofs_running << " ABACUS socket return energy " + << energy_ry << " Ry, " + << energy_ry * ModuleBase::Ry_to_eV << " eV, " + << energy_hartree << " Ha" << std::endl; + } ModuleBase::matrix force; - if (PARAM.inp.cal_force) + if (inp.cal_force) { p_esolver->cal_force(ucell, force); forces_hartree_bohr = flatten_forces_hartree_per_bohr(force); @@ -459,7 +410,7 @@ void Driver::driver_ipi_run() io_message = exc.what(); } } - throw_if_root_io_failed(io_failed, io_message); + quit_if_root_io_failed(io_failed, io_message); isinit = false; hasdata = false; } @@ -467,7 +418,7 @@ void Driver::driver_ipi_run() { if (is_root()) { - GlobalV::ofs_running << " ABACUS socket driver exiting on header " << header << std::endl; + ofs_running << " ABACUS socket driver exiting on header " << header << std::endl; } break; } @@ -475,33 +426,13 @@ void Driver::driver_ipi_run() } catch (const std::exception& exc) { - pending_error = exc.what(); - if (is_root()) - { - GlobalV::ofs_running << " ABACUS socket driver ended with error: " << pending_error << std::endl; - } + ModuleBase::WARNING_QUIT("ABACUS socket", exc.what()); } if (is_root()) { socket.close(); } - if (esolver_ready && runner_completed && p_esolver) - { - p_esolver->after_all_runners(ucell); - } - p_esolver.reset(); - if (hardware_initialized) - { - this->finalize_hardware(); - } -#ifdef __RAPIDJSON - Json::create_Json(&ucell, PARAM); -#endif - - if (!pending_error.empty()) - { - ModuleBase::WARNING_QUIT("ABACUS socket", pending_error); - } + ModuleBase::timer::end("Socket_Driver", "socket_driver"); } diff --git a/source/source_relax/socket_driver.h b/source/source_relax/socket_driver.h new file mode 100644 index 00000000000..51b75289cb3 --- /dev/null +++ b/source/source_relax/socket_driver.h @@ -0,0 +1,22 @@ +#ifndef ABACUS_SOURCE_RELAX_SOCKET_DRIVER_H +#define ABACUS_SOURCE_RELAX_SOCKET_DRIVER_H + +#include "source_cell/unitcell.h" +#include "source_esolver/esolver.h" +#include "source_io/module_parameter/input_parameter.h" + +#include + +class Socket_Driver +{ + public: + Socket_Driver() = default; + ~Socket_Driver() = default; + + void socket_driver(ModuleESolver::ESolver* p_esolver, + UnitCell& ucell, + const Input_para& inp, + std::ofstream& ofs_running); +}; + +#endif diff --git a/source/source_relax/test/CMakeLists.txt b/source/source_relax/test/CMakeLists.txt index 2262fefb419..d1c867dc690 100644 --- a/source/source_relax/test/CMakeLists.txt +++ b/source/source_relax/test/CMakeLists.txt @@ -6,6 +6,12 @@ abacus_disable_feature_definitions(__ROCM) install(DIRECTORY support DESTINATION ${CMAKE_CURRENT_BINARY_DIR}) + +AddTest( + TARGET MODULE_RELAX_ipi_socket_test + SOURCES ipi_socket_test.cpp ../ipi_socket.cpp +) + AddTest( TARGET MODULE_RELAX_relax_new_line_search LIBS parameter diff --git a/source/source_relax/test/ipi_socket_test.cpp b/source/source_relax/test/ipi_socket_test.cpp new file mode 100644 index 00000000000..21db9be2b90 --- /dev/null +++ b/source/source_relax/test/ipi_socket_test.cpp @@ -0,0 +1,256 @@ +#include "../ipi_socket.h" + +#include "gtest/gtest.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace +{ +constexpr std::size_t IPI_HEADER_LEN = 12; + +std::string errno_message(const std::string& prefix) +{ + return prefix + ": " + std::strerror(errno); +} + +void send_all(int fd, const void* data, std::size_t nbytes) +{ + const char* cursor = static_cast(data); + std::size_t done = 0; + while (done < nbytes) + { + const ssize_t sent = ::send(fd, cursor + done, nbytes - done, 0); + if (sent < 0) + { + if (errno == EINTR) + { + continue; + } + throw std::runtime_error(errno_message("send failed")); + } + if (sent == 0) + { + throw std::runtime_error("send returned zero"); + } + done += static_cast(sent); + } +} + +void recv_all(int fd, void* data, std::size_t nbytes) +{ + char* cursor = static_cast(data); + std::size_t done = 0; + while (done < nbytes) + { + const ssize_t received = ::recv(fd, cursor + done, nbytes - done, 0); + if (received < 0) + { + if (errno == EINTR) + { + continue; + } + throw std::runtime_error(errno_message("recv failed")); + } + if (received == 0) + { + throw std::runtime_error("socket closed while receiving test data"); + } + done += static_cast(received); + } +} + +std::string padded_header(const std::string& header) +{ + std::string padded = header; + padded.resize(IPI_HEADER_LEN, ' '); + return padded; +} + +class UnixSocketServer +{ + public: + UnixSocketServer() + { + char dir_template[] = "/tmp/abacus_ipi_socket_test_XXXXXX"; + char* made_dir = ::mkdtemp(dir_template); + if (made_dir == nullptr) + { + throw std::runtime_error(errno_message("mkdtemp failed")); + } + dir_ = made_dir; + path_ = dir_ + "/ipi.sock"; + + listen_fd_ = ::socket(AF_UNIX, SOCK_STREAM, 0); + if (listen_fd_ < 0) + { + throw std::runtime_error(errno_message("socket failed")); + } + + sockaddr_un addr; + std::memset(&addr, 0, sizeof(addr)); + addr.sun_family = AF_UNIX; + std::strncpy(addr.sun_path, path_.c_str(), sizeof(addr.sun_path) - 1); + if (::bind(listen_fd_, reinterpret_cast(&addr), sizeof(addr)) != 0) + { + throw std::runtime_error(errno_message("bind failed")); + } + if (::listen(listen_fd_, 1) != 0) + { + throw std::runtime_error(errno_message("listen failed")); + } + } + + ~UnixSocketServer() + { + if (listen_fd_ >= 0) + { + ::close(listen_fd_); + } + if (!path_.empty()) + { + ::unlink(path_.c_str()); + } + if (!dir_.empty()) + { + ::rmdir(dir_.c_str()); + } + } + + UnixSocketServer(const UnixSocketServer&) = delete; + UnixSocketServer& operator=(const UnixSocketServer&) = delete; + + std::string address() const + { + return path_ + ":UNIX"; + } + + int accept_once() + { + const int fd = ::accept(listen_fd_, nullptr, nullptr); + if (fd < 0) + { + throw std::runtime_error(errno_message("accept failed")); + } + return fd; + } + + private: + int listen_fd_ = -1; + std::string dir_; + std::string path_; +}; + +void rethrow_thread_error(const std::exception_ptr& thread_error) +{ + if (thread_error) + { + std::rethrow_exception(thread_error); + } +} +} // namespace + +TEST(IpiSocketTest, WriteHeaderPadsToTwelveBytes) +{ + UnixSocketServer server; + std::string received; + std::exception_ptr thread_error; + std::thread peer([&]() { + try + { + const int fd = server.accept_once(); + char buffer[IPI_HEADER_LEN]; + recv_all(fd, buffer, sizeof(buffer)); + received.assign(buffer, sizeof(buffer)); + ::close(fd); + } + catch (...) + { + thread_error = std::current_exception(); + } + }); + + IpiSocket socket; + socket.connect(server.address()); + socket.write_header("READY"); + socket.close(); + + peer.join(); + rethrow_thread_error(thread_error); + EXPECT_EQ(padded_header("READY"), received); +} + +TEST(IpiSocketTest, CleanPeerCloseBeforeNextHeaderThrowsDedicatedSignal) +{ + UnixSocketServer server; + std::exception_ptr thread_error; + std::thread peer([&]() { + try + { + const int fd = server.accept_once(); + const std::string header = padded_header("STATUS"); + send_all(fd, header.data(), header.size()); + ::close(fd); + } + catch (...) + { + thread_error = std::current_exception(); + } + }); + + IpiSocket socket; + socket.connect(server.address()); + EXPECT_EQ("STATUS", socket.read_header()); + EXPECT_THROW(socket.read_header(), IpiSocketClosed); + socket.close(); + + peer.join(); + rethrow_thread_error(thread_error); +} + +TEST(IpiSocketTest, PartialHeaderCloseStaysRuntimeError) +{ + UnixSocketServer server; + std::exception_ptr thread_error; + std::thread peer([&]() { + try + { + const int fd = server.accept_once(); + const std::string partial = "STAT"; + send_all(fd, partial.data(), partial.size()); + ::close(fd); + } + catch (...) + { + thread_error = std::current_exception(); + } + }); + + IpiSocket socket; + socket.connect(server.address()); + try + { + static_cast(socket.read_header()); + FAIL() << "partial header EOF should throw"; + } + catch (const IpiSocketClosed&) + { + FAIL() << "partial header EOF must not be treated as clean peer close"; + } + catch (const std::runtime_error& exc) + { + EXPECT_NE(std::string::npos, std::string(exc.what()).find("closed while reading header")); + } + socket.close(); + + peer.join(); + rethrow_thread_error(thread_error); +} From fe2c2a3cc61995b3383aaf7f2a6abf5ab1882907 Mon Sep 17 00:00:00 2001 From: jianrui geng Date: Thu, 9 Jul 2026 11:12:06 +0800 Subject: [PATCH 3/4] Fix socket driver CI regressions --- source/source_main/driver_run.cpp | 15 ++-- source/source_relax/socket_driver.cpp | 99 +++++++++++++++++++++++++-- 2 files changed, 101 insertions(+), 13 deletions(-) diff --git a/source/source_main/driver_run.cpp b/source/source_main/driver_run.cpp index 3e7d669b6c2..408b507b902 100644 --- a/source/source_main/driver_run.cpp +++ b/source/source_main/driver_run.cpp @@ -38,9 +38,6 @@ void Driver::driver_run() { ModuleBase::TITLE("Driver", "driver_run"); - const std::string cal = PARAM.inp.calculation; - const bool socket_mode = (cal == "socket"); - //! 1: setup cell and atom information // this warning should not be here, mohan 2024-05-22 #ifndef __LCAO @@ -62,13 +59,17 @@ void Driver::driver_run() unitcell::check_atomic_stru(ucell, PARAM.inp.min_dist_coef); //! 2: initialize the ESolver (depends on a set-up ucell after `setup_cell`) - Input_para esolver_inp = PARAM.inp; - Input_para driver_inp = PARAM.inp; + const std::string cal = PARAM.inp.calculation; + const bool socket_mode = (cal == "socket"); + Input_para socket_esolver_inp = PARAM.inp; + Input_para socket_driver_inp = PARAM.inp; if (socket_mode) { - esolver_inp.calculation = "scf"; - driver_inp.calculation = cal; + socket_esolver_inp.calculation = "scf"; + socket_driver_inp.calculation = cal; } + const Input_para& esolver_inp = socket_mode ? socket_esolver_inp : PARAM.inp; + const Input_para& driver_inp = socket_mode ? socket_driver_inp : PARAM.inp; this->init_hardware(); diff --git a/source/source_relax/socket_driver.cpp b/source/source_relax/socket_driver.cpp index 02928fec9d3..bb0a6757f50 100644 --- a/source/source_relax/socket_driver.cpp +++ b/source/source_relax/socket_driver.cpp @@ -27,29 +27,52 @@ bool is_root() void bcast_double_vector(std::vector& values) { +#ifdef __MPI if (!values.empty()) { Parallel_Common::bcast_double(values.data(), static_cast(values.size())); } +#else + (void)values; +#endif +} + +void bcast_socket_int(int& value) +{ +#ifdef __MPI + Parallel_Common::bcast_int(value); +#else + (void)value; +#endif +} + +void bcast_socket_chars(char* value, const int size) +{ +#ifdef __MPI + Parallel_Common::bcast_char(value, size); +#else + (void)value; + (void)size; +#endif } void bcast_socket_string(std::string& value) { int size = static_cast(value.size()); - Parallel_Common::bcast_int(size); + bcast_socket_int(size); if (!is_root()) { value.resize(static_cast(size)); } if (size > 0) { - Parallel_Common::bcast_char(&value[0], size); + bcast_socket_chars(&value[0], size); } } void quit_if_root_io_failed(int root_failed, std::string root_message) { - Parallel_Common::bcast_int(root_failed); + bcast_socket_int(root_failed); bcast_socket_string(root_message); if (root_failed != 0) { @@ -85,6 +108,58 @@ std::vector ipi_cell_bohr_from_unitcell(const UnitCell& ucell) }; } +double max_wrapped_direct_delta_from_unitcell(const UnitCell& ucell, const std::vector& positions_bohr) +{ + if (positions_bohr.size() != static_cast(3 * ucell.nat)) + { + return 1.0e99; + } + + double out = 0.0; + int iat = 0; + for (int it = 0; it < ucell.ntype; ++it) + { + const Atom* atom = &ucell.atoms[it]; + for (int ia = 0; ia < atom->na; ++ia) + { + const double tau_x = positions_bohr[3 * iat + 0] / ucell.lat0; + const double tau_y = positions_bohr[3 * iat + 1] / ucell.lat0; + const double tau_z = positions_bohr[3 * iat + 2] / ucell.lat0; + + double dx = 0.0; + double dy = 0.0; + double dz = 0.0; + ModuleBase::Mathzone::Cartesian_to_Direct(tau_x, + tau_y, + tau_z, + ucell.latvec.e11, + ucell.latvec.e12, + ucell.latvec.e13, + ucell.latvec.e21, + ucell.latvec.e22, + ucell.latvec.e23, + ucell.latvec.e31, + ucell.latvec.e32, + ucell.latvec.e33, + dx, + dy, + dz); + + double ddx = dx - atom->taud[ia].x; + double ddy = dy - atom->taud[ia].y; + double ddz = dz - atom->taud[ia].z; + ddx -= std::round(ddx); + ddy -= std::round(ddy); + ddz -= std::round(ddz); + out = std::max(out, std::abs(ddx)); + out = std::max(out, std::abs(ddy)); + out = std::max(out, std::abs(ddz)); + ++iat; + } + } + return out; +} + double max_abs_delta(const std::vector& a, const std::vector& b) { if (a.size() != b.size()) @@ -214,6 +289,7 @@ void Socket_Driver::socket_driver(ModuleESolver::ESolver* p_esolver, std::vector virial_hartree(9, 0.0); const std::vector reference_cell = ipi_cell_bohr_from_unitcell(ucell); + bool checked_initial_positions = false; while (true) { @@ -306,8 +382,8 @@ void Socket_Driver::socket_driver(ModuleESolver::ESolver* p_esolver, } } quit_if_root_io_failed(io_failed, io_message); - Parallel_Common::bcast_int(rid); - Parallel_Common::bcast_int(nbytes); + bcast_socket_int(rid); + bcast_socket_int(nbytes); if (nbytes > 0 && is_root()) { ofs_running << " ABACUS socket INIT params bytes " << nbytes << std::endl; @@ -352,7 +428,7 @@ void Socket_Driver::socket_driver(ModuleESolver::ESolver* p_esolver, quit_if_root_io_failed(io_failed, io_message); bcast_double_vector(cell); bcast_double_vector(inv_cell); - Parallel_Common::bcast_int(nat_socket); + bcast_socket_int(nat_socket); if (!is_root()) { positions.assign(static_cast(3 * nat_socket), 0.0); @@ -368,6 +444,17 @@ void Socket_Driver::socket_driver(ModuleESolver::ESolver* p_esolver, { ModuleBase::WARNING_QUIT("ABACUS socket", "variable-cell socket updates are not supported yet."); } + if (!checked_initial_positions) + { + checked_initial_positions = true; + if (max_wrapped_direct_delta_from_unitcell(ucell, positions) > 1.0e-5 && is_root()) + { + ModuleBase::WARNING( + "ABACUS socket", + "first POSDATA positions are not PBC-equivalent to STRU atom order; " + "i-PI POSDATA carries no species, so the client atoms should use the same atom order as STRU."); + } + } set_positions_from_ipi_bohr(ucell, positions); p_esolver->runner(ucell, istep); From 36353412e892d8d3a955d58f057285e886d1659d Mon Sep 17 00:00:00 2001 From: jianrui geng Date: Thu, 9 Jul 2026 11:23:16 +0800 Subject: [PATCH 4/4] Reduce socket driver global dependencies --- source/source_main/driver_run.cpp | 9 +++------ source/source_relax/socket_driver.cpp | 12 ++++++++++-- source/source_relax/socket_driver.h | 12 ++++++++---- 3 files changed, 21 insertions(+), 12 deletions(-) diff --git a/source/source_main/driver_run.cpp b/source/source_main/driver_run.cpp index 408b507b902..e565871d604 100644 --- a/source/source_main/driver_run.cpp +++ b/source/source_main/driver_run.cpp @@ -59,17 +59,13 @@ void Driver::driver_run() unitcell::check_atomic_stru(ucell, PARAM.inp.min_dist_coef); //! 2: initialize the ESolver (depends on a set-up ucell after `setup_cell`) - const std::string cal = PARAM.inp.calculation; - const bool socket_mode = (cal == "socket"); Input_para socket_esolver_inp = PARAM.inp; - Input_para socket_driver_inp = PARAM.inp; + const bool socket_mode = (socket_esolver_inp.calculation == "socket"); if (socket_mode) { socket_esolver_inp.calculation = "scf"; - socket_driver_inp.calculation = cal; } const Input_para& esolver_inp = socket_mode ? socket_esolver_inp : PARAM.inp; - const Input_para& driver_inp = socket_mode ? socket_driver_inp : PARAM.inp; this->init_hardware(); @@ -84,6 +80,7 @@ void Driver::driver_run() #endif //! 4: different types of calculations + const std::string cal = PARAM.inp.calculation; if (cal == "md") { Run_MD::md_line(ucell, p_esolver, PARAM); @@ -91,7 +88,7 @@ void Driver::driver_run() else if (cal == "scf" || cal == "relax" || cal == "cell-relax" || cal == "nscf" || cal == "socket") { Relax_Driver rl_driver; - rl_driver.relax_driver(p_esolver, ucell, driver_inp, GlobalV::ofs_running); + rl_driver.relax_driver(p_esolver, ucell, PARAM.inp, GlobalV::ofs_running); } else if (cal == "get_s") { diff --git a/source/source_relax/socket_driver.cpp b/source/source_relax/socket_driver.cpp index bb0a6757f50..ecd4d1be27e 100644 --- a/source/source_relax/socket_driver.cpp +++ b/source/source_relax/socket_driver.cpp @@ -2,11 +2,13 @@ #include "source_relax/ipi_socket.h" #include "source_base/global_function.h" -#include "source_base/global_variable.h" #include "source_base/mathzone.h" #include "source_base/parallel_common.h" #include "source_base/timer.h" +#include "source_cell/unitcell.h" #include "source_cell/update_cell.h" +#include "source_esolver/esolver.h" +#include "source_io/module_parameter/input_parameter.h" #include #include @@ -22,7 +24,13 @@ constexpr int IPI_RANK_ROOT = 0; bool is_root() { - return GlobalV::MY_RANK == IPI_RANK_ROOT; +#ifdef __MPI + int rank = IPI_RANK_ROOT; + MPI_Comm_rank(MPI_COMM_WORLD, &rank); + return rank == IPI_RANK_ROOT; +#else + return true; +#endif } void bcast_double_vector(std::vector& values) diff --git a/source/source_relax/socket_driver.h b/source/source_relax/socket_driver.h index 51b75289cb3..86c180fc42e 100644 --- a/source/source_relax/socket_driver.h +++ b/source/source_relax/socket_driver.h @@ -1,11 +1,15 @@ #ifndef ABACUS_SOURCE_RELAX_SOCKET_DRIVER_H #define ABACUS_SOURCE_RELAX_SOCKET_DRIVER_H -#include "source_cell/unitcell.h" -#include "source_esolver/esolver.h" -#include "source_io/module_parameter/input_parameter.h" +#include -#include +class UnitCell; +struct Input_para; + +namespace ModuleESolver +{ +class ESolver; +} class Socket_Driver {