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.

Author: jianrui geng

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

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();

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")
     {

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

source/source_main/ipi_socket.cpp source/source_relax/ipi_socket.cppsource/source_main/ipi_socket.cpp renamed to source/source_relax/ipi_socket.cpp

@@ -1,4 +1,4 @@
-#include "ipi_socket.h"
+#include "source_relax/ipi_socket.h"
 
 #include <arpa/inet.h>
 #include <cerrno>
@@ -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<std::size_t>(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)
         {

source/source_main/ipi_socket.h source/source_relax/ipi_socket.hsource/source_main/ipi_socket.h renamed to source/source_relax/ipi_socket.h

@@ -2,9 +2,16 @@
 #define ABACUS_IPI_SOCKET_H
 
 #include <cstddef>
+#include <stdexcept>
 #include <string>
 #include <vector>
 
+class IpiSocketClosed : public std::runtime_error
+{
+  public:
+    explicit IpiSocketClosed(const std::string& message);
+};
+
 class IpiSocket
 {
   public:

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)