diff --git a/docs/advanced/input_files/input-main.md b/docs/advanced/input_files/input-main.md index 334144f44c..c6c1b73299 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 00ef9b9dcf..d183783ad8 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/source_io/module_parameter/read_input_item_system.cpp b/source/source_io/module_parameter/read_input_item_system.cpp index 3b0ae011ec..6ed572b21f 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 61673a418c..eb043fe1b0 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_run.cpp b/source/source_main/driver_run.cpp index 4911b133de..e565871d60 100644 --- a/source/source_main/driver_run.cpp +++ b/source/source_main/driver_run.cpp @@ -59,26 +59,33 @@ 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 socket_esolver_inp = PARAM.inp; + const bool socket_mode = (socket_esolver_inp.calculation == "socket"); + if (socket_mode) + { + socket_esolver_inp.calculation = "scf"; + } + const Input_para& esolver_inp = socket_mode ? socket_esolver_inp : PARAM.inp; + 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 Json::gen_stru_wrapper(&ucell); #endif - const std::string cal = PARAM.inp.calculation; - //! 4: different types of calculations + const std::string cal = PARAM.inp.calculation; if (cal == "md") { 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); diff --git a/source/source_relax/CMakeLists.txt b/source/source_relax/CMakeLists.txt index 9e7ef96b0e..84f69c06c8 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_relax/ipi_socket.cpp b/source/source_relax/ipi_socket.cpp new file mode 100644 index 0000000000..304a018181 --- /dev/null +++ b/source/source_relax/ipi_socket.cpp @@ -0,0 +1,266 @@ +#include "source_relax/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 + +IpiSocketClosed::IpiSocketClosed(const std::string& message) : std::runtime_error(message) +{ +} + +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]; + 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); +} + +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 IpiSocketClosed("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_relax/ipi_socket.h b/source/source_relax/ipi_socket.h new file mode 100644 index 0000000000..28ce2cb81f --- /dev/null +++ b/source/source_relax/ipi_socket.h @@ -0,0 +1,47 @@ +#ifndef ABACUS_IPI_SOCKET_H +#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: + 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 diff --git a/source/source_relax/relax_driver.cpp b/source/source_relax/relax_driver.cpp index 26128bf040..0cada1bfc3 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_relax/socket_driver.cpp b/source/source_relax/socket_driver.cpp new file mode 100644 index 0000000000..ecd4d1be27 --- /dev/null +++ b/source/source_relax/socket_driver.cpp @@ -0,0 +1,533 @@ +#include "socket_driver.h" + +#include "source_relax/ipi_socket.h" +#include "source_base/global_function.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 +#include +#include +#include +#include + +namespace +{ +constexpr double RY_TO_HARTREE = 0.5; +constexpr int IPI_RANK_ROOT = 0; + +bool is_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) +{ +#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()); + bcast_socket_int(size); + if (!is_root()) + { + value.resize(static_cast(size)); + } + if (size > 0) + { + bcast_socket_chars(&value[0], size); + } +} + +void quit_if_root_io_failed(int root_failed, std::string root_message) +{ + bcast_socket_int(root_failed); + bcast_socket_string(root_message); + if (root_failed != 0) + { + ModuleBase::WARNING_QUIT("ABACUS socket", root_message.empty() ? "i-PI socket I/O failed" : root_message); + } +} + +std::string bcast_header(std::string header) +{ + bcast_socket_string(header); + return header; +} + +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_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()) + { + 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)) + { + ModuleBase::WARNING_QUIT("ABACUS socket", "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; +} +} // namespace + +void Socket_Driver::socket_driver(ModuleESolver::ESolver* p_esolver, + UnitCell& ucell, + const Input_para& inp, + std::ofstream& ofs_running) +{ + ModuleBase::TITLE("Socket_Driver", "socket_driver"); + ModuleBase::timer::start("Socket_Driver", "socket_driver"); + + if (p_esolver == nullptr) + { + ModuleBase::WARNING_QUIT("ABACUS socket", "socket driver requires a valid ESolver."); + } + if (!inp.cal_force) + { + ModuleBase::WARNING_QUIT("ABACUS socket", "socket calculation requires cal_force=1 for i-PI GETFORCE."); + } + + IpiSocket socket; + + try + { + int io_failed = 0; + std::string io_message; + if (is_root()) + { + try + { + const std::string address = ipi_address(); + 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(); + } + } + quit_if_root_io_failed(io_failed, io_message); + + bool isinit = false; + bool hasdata = false; + int istep = 0; + 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); + + const std::vector reference_cell = ipi_cell_bohr_from_unitcell(ucell); + bool checked_initial_positions = false; + + while (true) + { + std::string header; + io_failed = 0; + io_message.clear(); + if (is_root()) + { + try + { + header = socket.read_header(); + } + catch (const IpiSocketClosed&) + { + header.clear(); + } + catch (const std::exception& exc) + { + io_failed = 1; + io_message = exc.what(); + } + } + quit_if_root_io_failed(io_failed, io_message); + header = bcast_header(header); + + 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(); + 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(); + } + } + quit_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) + { + io_failed = 1; + io_message = "negative INIT payload length from i-PI socket"; + } + else if (nbytes > 0) + { + params = socket.read_string(static_cast(nbytes)); + } + } + catch (const std::exception& exc) + { + io_failed = 1; + io_message = exc.what(); + } + } + quit_if_root_io_failed(io_failed, io_message); + bcast_socket_int(rid); + bcast_socket_int(nbytes); + if (nbytes > 0 && is_root()) + { + ofs_running << " ABACUS socket INIT params bytes " << nbytes << std::endl; + } + isinit = true; + if (is_root()) + { + 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) + { + io_failed = 1; + io_message = "negative POSDATA atom count from i-PI socket"; + } + else + { + positions = socket.read_doubles(static_cast(3 * nat_socket)); + } + } + catch (const std::exception& exc) + { + io_failed = 1; + io_message = exc.what(); + } + } + quit_if_root_io_failed(io_failed, io_message); + bcast_double_vector(cell); + bcast_double_vector(inv_cell); + bcast_socket_int(nat_socket); + if (!is_root()) + { + positions.assign(static_cast(3 * nat_socket), 0.0); + } + bcast_double_vector(positions); + + if (nat_socket != ucell.nat) + { + 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) + { + 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); + 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 (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(); + } + } + quit_if_root_io_failed(io_failed, io_message); + isinit = false; + hasdata = false; + } + else + { + if (is_root()) + { + ofs_running << " ABACUS socket driver exiting on header " << header << std::endl; + } + break; + } + } + } + catch (const std::exception& exc) + { + ModuleBase::WARNING_QUIT("ABACUS socket", exc.what()); + } + + if (is_root()) + { + socket.close(); + } + + 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 0000000000..86c180fc42 --- /dev/null +++ b/source/source_relax/socket_driver.h @@ -0,0 +1,26 @@ +#ifndef ABACUS_SOURCE_RELAX_SOCKET_DRIVER_H +#define ABACUS_SOURCE_RELAX_SOCKET_DRIVER_H + +#include + +class UnitCell; +struct Input_para; + +namespace ModuleESolver +{ +class ESolver; +} + +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 2262fefb41..d1c867dc69 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 0000000000..21db9be2b9 --- /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); +}