diff --git a/docs/advanced/input_files/input-main.md b/docs/advanced/input_files/input-main.md index ccfa7048a8e..4b987618fd6 100644 --- a/docs/advanced/input_files/input-main.md +++ b/docs/advanced/input_files/input-main.md @@ -9,6 +9,7 @@ - [suffix](#suffix) - [ntype](#ntype) - [calculation](#calculation) + - [socket\_driver](#socket_driver) - [esolver\_type](#esolver_type) - [symmetry](#symmetry) - [symmetry\_prec](#symmetry_prec) @@ -592,6 +593,18 @@ - test_neighbour: obtain information of neighboring atoms (for LCAO basis only), please specify a positive search_radius manually - **Default**: scf +### socket_driver + +- **Type**: Boolean +- **Description**: If set to True, ABACUS keeps the calculation type as scf and receives atomic positions from an external driver through the i-PI socket protocol. + + > Note: Use calculation = scf with socket_driver = True. ABACUS connects to the external i-PI server selected by ABACUS_SOCKET_ADDRESS. If ABACUS_SOCKET_ADDRESS is unset, ABACUS uses localhost:31415. The value can use one of two forms: + + - host:port, for example localhost:31415 or 127.0.0.1:31415, opens a TCP connection to that host and port. Use this when the i-PI server listens on a TCP port. + - path:UNIX, for example /tmp/ipi_abacus_si:UNIX, opens a Unix-domain socket at the given filesystem path. The :UNIX suffix tells ABACUS that the preceding value is a local socket path rather than a TCP host name. This form only works on the same machine. + When using the ASE AbacusSocketIO interface, this environment variable is set automatically from the port or unixsocket calculator argument. +- **Default**: False + ### esolver_type - **Type**: String @@ -840,7 +853,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 and socket_driver calculations, 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/advanced/interface/ase.md b/docs/advanced/interface/ase.md index e9b7b062809..b9735279d1c 100644 --- a/docs/advanced/interface/ase.md +++ b/docs/advanced/interface/ase.md @@ -103,6 +103,77 @@ In the new implementation, we limit the range of functionalties supported to mai Please read the examples in `interfaces/ASE_interface/examples/` for more details. +### Socket I/O with ASE + +For socket-driven ASE workflows, use the `AbacusSocketIO` calculator. ASE runs the i-PI socket server, while ABACUS keeps `calculation=scf` and is launched with `socket_driver=1` as the client. The protocol is simple: ASE sends atomic positions and cell data to ABACUS; ABACUS evaluates one SCF step for that structure and returns energy, forces, and virial. See the [ASE socket I/O documentation](https://ase-lib.org/ase/calculators/socketio/socketio.html) and the i-PI reference paper, [Ceriotti et al., Comput. Phys. Commun. 185, 1019-1026 (2014)](https://doi.org/10.1016/j.cpc.2013.10.027), for the protocol background. + +Build ABACUS as usual before using this interface. PW-only builds work with `basis_type=pw`; LCAO socket calculations require an LCAO-enabled executable. No extra socket library is required. + +With CMake, choose the executable according to the basis: + +```bash +cmake -S . -B build-pw -DENABLE_MPI=ON -DENABLE_LCAO=OFF +cmake --build build-pw --target abacus_pw_para -j + +cmake -S . -B build-lcao -DENABLE_MPI=ON -DENABLE_LCAO=ON +cmake --build build-lcao --target abacus_basic_para -j +``` + +With the ABACUS toolchain workflow, build the normal ABACUS executable with LCAO support when `basis_type=lcao` is needed, then pass that executable to `AbacusProfile(command=...)`. The command can include an MPI launcher, for example `mpirun -np 4 /path/to/abacus`; ABACUS rank 0 opens the socket connection and broadcasts the i-PI data to the other ranks internally. On managed clusters, keep scheduler-specific launch options outside the calculator when possible and test the exact launcher command on a compute node. + +For PW calculations on CUDA/ROCm with multiple MPI ranks, use a k-point layout compatible with ABACUS' GPU parallelization. In practice, make sure each k-point pool contains one MPI rank; for example, a 4-rank PW GPU socket calculation should use at least four k-points so the default GPU `kpar` adjustment can assign one rank per pool. A one-k-point PW GPU job with several MPI ranks can fail in the PW GPU transform path; reduce the rank count or use a denser k-point mesh such as a smaller `kspacing`. + +The ASE interface can be installed from this repository with: + +```bash +cd interfaces/ASE_interface +pip install . +``` + +A minimal socket calculator setup is: + +```python +from ase.optimize import BFGS +from abacuslite import AbacusProfile, AbacusSocketIO + +aprof = AbacusProfile( + command="mpirun -np 4 /path/to/abacus", + pseudo_dir="/path/to/pseudopotentials", + orbital_dir="/path/to/orbitals", + omp_num_threads=1, +) + +abacus = AbacusSocketIO( + profile=aprof, + directory="socketio", + unixsocket="abacus_si", + pseudopotentials={"Si": "Si_ONCV_PBE-1.0.upf"}, + basissets={"Si": "Si_gga_8au_100Ry_2s2p1d.orb"}, + inp={"calculation": "scf", "basis_type": "lcao", "kspacing": 0.1}, +) + +with abacus as calc: + atoms.calc = calc + BFGS(atoms).run(fmax=0.05) +``` + +`AbacusSocketIO` sets `socket_driver=1` and `cal_force=1` automatically. It also selects the socket endpoint and passes it to ABACUS through `ABACUS_SOCKET_ADDRESS`, so users normally do not set this environment variable by hand when using abacuslite. + +There are two endpoint styles: + +- `unixsocket="abacus_si"` uses a local Unix-domain socket. ASE creates and listens on `/tmp/ipi_abacus_si`; abacuslite launches ABACUS with `ABACUS_SOCKET_ADDRESS=/tmp/ipi_abacus_si:UNIX`. The `:UNIX` suffix is part of ABACUS' address syntax and means that `/tmp/ipi_abacus_si` is a filesystem socket path, not a TCP host. This is usually the best choice when ASE and ABACUS run on the same node because it avoids TCP port conflicts. +- `port=31415` uses a TCP socket. abacuslite launches ABACUS with `ABACUS_SOCKET_ADDRESS=localhost:31415`, meaning host `localhost` and TCP port `31415`. Use this style when the socket server should listen on a TCP port. If ABACUS is launched manually instead of through `AbacusSocketIO`, set `ABACUS_SOCKET_ADDRESS` yourself to the same `host:port` or `path:UNIX` endpoint. + +Calling `atoms.get_potential_energy()` is supported, but the ABACUS client still computes forces because the i-PI `GETFORCE` exchange returns energy, forces, and virial as one response. + +A socket calculator owns one ABACUS process initialized from one fixed `INPUT`/`STRU` setup. Reuse the same `AbacusSocketIO` instance only for position updates under the same electronic-structure settings and the same cell. Do not change `kpts`, `kspacing`, `nspin`, `basis_type`, `basissets`, pseudopotentials, species, atom count, cell, or other core `INPUT`/`STRU` parameters through an existing socket calculator; create a new `AbacusSocketIO` instance and a new ABACUS client process for those changes. `AbacusSocketIO` rejects cell changes before sending them to ABACUS, and the ABACUS socket driver also checks incoming POSDATA cells against the initial `STRU` cell and exits if they differ. + +In socket mode, ABACUS keeps one client process alive. All SCF evaluations produced by the same `AbacusSocketIO` instance are appended to the same `OUT.ABACUS/running_scf.log`, because the ABACUS calculation type remains `scf`. The authoritative per-step energy and force results are returned through the i-PI socket to ASE. Use ASE trajectory and optimizer log files, such as `BFGS(atoms, trajectory="opt.traj", logfile="opt.log")`, when each optimizer or MD step should be saved separately. Treat `running_scf.log` mainly as the ABACUS diagnostic log for the socket client, not as one independent FileIO result per structure. + +The i-PI protocol does not transmit element symbols. `AbacusSocketIO` therefore sorts the internal socket atoms with the same first-occurrence species grouping used when writing `STRU`, and maps returned forces back to the original ASE `Atoms` order. This avoids silent force/atom mismatches when structures are read from CIF, extxyz, POSCAR, or other formats whose atom order is not already grouped for ABACUS. Users should not manually reorder atoms for socket I/O; pass the physical ASE `Atoms` object directly to the calculator. + +A complete fixed-cell validation and benchmark example is available in `interfaces/ASE_interface/examples/socketio.py`. + ## SPAP Analysis [SPAP](https://github.com/chuanxun/StructurePrototypeAnalysisPackage) (Structure Prototype Analysis Package) is written by Dr. Chuanxun Su to analyze symmetry and compare similarity of large amount of atomic structures. The coordination characterization function (CCF) is used to diff --git a/docs/parameters.yaml b/docs/parameters.yaml index 1e540b12cba..c5a775f471c 100644 --- a/docs/parameters.yaml +++ b/docs/parameters.yaml @@ -38,6 +38,19 @@ parameters: default_value: scf unit: "" availability: "" + - name: socket_driver + category: System variables + type: Boolean + description: | + If set to True, ABACUS keeps the calculation type as scf and receives atomic positions from an external driver through the i-PI socket protocol. + + [NOTE] Use calculation = scf with socket_driver = True. ABACUS connects to the external i-PI server selected by ABACUS_SOCKET_ADDRESS. If ABACUS_SOCKET_ADDRESS is unset, ABACUS uses localhost:31415. The value can use one of two forms: + * host:port, for example localhost:31415 or 127.0.0.1:31415, opens a TCP connection to that host and port. Use this when the i-PI server listens on a TCP port. + * path:UNIX, for example /tmp/ipi_abacus_si:UNIX, opens a Unix-domain socket at the given filesystem path. The :UNIX suffix tells ABACUS that the preceding value is a local socket path rather than a TCP host name. This form only works on the same machine. + When using the ASE AbacusSocketIO interface, this environment variable is set automatically from the port or unixsocket calculator argument. + default_value: "False" + unit: "" + availability: "" - name: esolver_type category: System variables type: String @@ -329,7 +342,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 and socket_driver calculations, 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/interfaces/ASE_interface/README.md b/interfaces/ASE_interface/README.md index 824e69900de..a76e68e177e 100644 --- a/interfaces/ASE_interface/README.md +++ b/interfaces/ASE_interface/README.md @@ -32,6 +32,7 @@ Please refer to the example scripts in the `examples` folder. Recommended learni 7. **constraintmd.py** - Constrained molecular dynamics simulation 8. **metadynamics.py** - Metadynamics simulation 9. **neb.py** - Nudged Elastic Band (NEB) calculation +10. **socketio.py** - ASE optimization with the AbacusSocketIO calculator running ABACUS as an i-PI socket client More usage examples will be provided in future versions. diff --git a/interfaces/ASE_interface/abacuslite/core.py b/interfaces/ASE_interface/abacuslite/core.py index e285db03385..ca3ead454a7 100644 --- a/interfaces/ASE_interface/abacuslite/core.py +++ b/interfaces/ASE_interface/abacuslite/core.py @@ -45,6 +45,7 @@ GenericFileIOCalculator, read_stdout ) +from ase.calculators.socketio import SocketIOCalculator from ase.atoms import Atoms from ase.dft.kpoints import BandPath from ase.io import read @@ -124,17 +125,35 @@ def __init__(self, @staticmethod def parse_version(stdout) -> str: - # up to the ABACUS version v3.9.0.17, the run of command - # `abacus --version` would returns the information organized - # in the following way: - # ABACUS version v3.9.0.17 - return re.match(r'ABACUS version (\S+)', stdout).group(1) + # MPI launchers may add informational lines before ABACUS output. + match = re.search(r'ABACUS version (\S+)', stdout or '') + if match is None: + raise RuntimeError( + 'Could not parse ABACUS version from command output. ' + 'Expected a line like "ABACUS version vX.Y.Z".' + ) + return match.group(1) def get_calculator_command(self, inputfile) -> List[str]: # because ABACUS run in the folder where there are INPUT files, so the # additional inputfile argument is not used. return [] + def socketio_argv_inet(self, port: Optional[int] = None) -> List[str]: + port = 31415 if port is None else port + return [ + 'env', + f'ABACUS_SOCKET_ADDRESS=localhost:{port}', + *self._split_command, + ] + + def socketio_argv_unix(self, socket: str) -> List[str]: + return [ + 'env', + f'ABACUS_SOCKET_ADDRESS=/tmp/ipi_{socket}:UNIX', + *self._split_command, + ] + def version(self) -> str: '''get the abacus version information''' cmd_ = [*self._split_command, '--version'] @@ -443,6 +462,17 @@ def __init__(self, directory=directory, ) + def write_input(self, atoms, properties=None, system_changes=None): + if properties is None: + properties = self.template.implemented_properties + self.template.write_input( + profile=self.profile, + directory=Path(self.directory), + atoms=atoms, + parameters=self.parameters, + properties=properties, + ) + @classmethod def restart(cls, profile=None, directory='.', **kwargs): '''instantiate one ABACUS calculator from an existing job directory, @@ -558,11 +588,208 @@ def band_structure(self, efermi=None): from ase.spectrum.band_structure import get_band_structure return get_band_structure(calc=self, reference=efermi) +class AbacusSocketIO(SocketIOCalculator): + """ASE socket I/O calculator that launches ABACUS as an i-PI client. + + A socket calculator owns one ABACUS process with one fixed INPUT/STRU + setup. The i-PI protocol can update positions, but electronic-structure + parameters such as k-points, spin, basis, pseudopotentials, and species + require a new calculator instance. Energy-only ASE calls are accepted, but + ABACUS still computes forces because i-PI GETFORCE returns energy, forces, + and virial together. + """ + + def __init__(self, + profile=None, + directory='.', + port=None, + unixsocket=None, + timeout=None, + log=None, + **kwargs): + inp = self._socket_inp(kwargs.pop('inp', {})) + self.abacus = Abacus( + profile=profile, + directory=directory, + inp=inp, + **kwargs, + ) + self._reference_cell = None + super().__init__( + port=port, + unixsocket=unixsocket, + timeout=timeout, + log=log, + launch_client=self._launch_client, + ) + + def calculate(self, atoms=None, properties=['energy'], system_changes=None): + from ase.calculators.calculator import ( + PropertyNotImplementedError, + all_changes, + ) + from ase.stress import full_3x3_to_voigt_6_stress + + if system_changes is None: + system_changes = all_changes + if atoms is None: + atoms = self.atoms + if atoms is None: + raise ValueError('AbacusSocketIO.calculate requires atoms') + + bad = [change for change in system_changes + if change not in self.supported_changes] + if self.atoms is not None and any(bad): + raise PropertyNotImplementedError( + 'Cannot change {} through IPI protocol. ' + 'Please create new socket calculator.' + .format(bad if len(bad) > 1 else bad[0])) + + self._check_fixed_cell(atoms) + order = self._socket_sort_indices(atoms) + socket_atoms = atoms[order] + self.atoms = atoms.copy() + + if self.server is None: + self.server = self.launch_server() + proc = self.launch_client(socket_atoms, properties, + port=self._port, + unixsocket=self._unixsocket) + self.server.proc = proc + + results = self.server.calculate(socket_atoms) + results['free_energy'] = results['energy'] + virial = results.pop('virial') + if self.atoms.cell.rank == 3 and any(self.atoms.pbc): + vol = atoms.get_volume() + results['stress'] = -full_3x3_to_voigt_6_stress(virial) / vol + if 'forces' in results: + results['forces'] = self._forces_to_input_order( + results['forces'], order) + self.results.update(results) + + def _check_fixed_cell(self, atoms): + from ase.calculators.calculator import PropertyNotImplementedError + + cell = atoms.cell.array.copy() + if self._reference_cell is None: + self._reference_cell = cell + return + max_delta = np.max(np.abs(cell - self._reference_cell)) + if max_delta > 1.0e-10: + raise PropertyNotImplementedError( + 'AbacusSocketIO is fixed-cell only; create a new socket ' + 'calculator for a changed cell, or use the normal Abacus ' + 'FileIO calculator for variable-cell workflows.' + ) + + def set(self, **kwargs): + if kwargs: + raise ValueError( + 'AbacusSocketIO input parameters are fixed after construction; ' + 'create a new AbacusSocketIO calculator to change k-points, ' + 'spin, basis, pseudopotentials, species, or other INPUT/STRU ' + 'settings.' + ) + return super().set(**kwargs) + + def _launch_client(self, atoms, properties=None, port=None, unixsocket=None): + from subprocess import Popen + + if properties is None: + properties = self.abacus.template.implemented_properties + + directory = Path(self.abacus.directory) + directory.mkdir(exist_ok=True, parents=True) + + if hasattr(self.abacus, 'write_inputfiles'): + self.abacus.write_inputfiles(atoms, properties) + else: + self.abacus.write_input(atoms, properties=properties) + + if unixsocket is not None: + argv = self.abacus.profile.socketio_argv_unix(socket=unixsocket) + else: + argv = self.abacus.profile.socketio_argv_inet(port=port) + + stdout = open(directory / self.abacus.template.outputname, 'w') + stderr = open(directory / self.abacus.template.errorname, 'w') + try: + return Popen(argv, cwd=directory, env=os.environ, + stdout=stdout, stderr=stderr) + finally: + stdout.close() + stderr.close() + + @staticmethod + def _socket_inp(inp): + inp = dict(inp) + calculation = inp.get('calculation', 'scf') + if calculation != 'scf': + raise ValueError('ABACUS socket I/O requires calculation="scf"') + inp.update({ + 'calculation': 'scf', + 'socket_driver': 1, + 'cal_force': 1, + }) + return inp + + @staticmethod + def _socket_sort_indices(atoms): + return species_group_indices(atoms.get_chemical_symbols()) + + @staticmethod + def _forces_to_input_order(forces, order): + reordered = np.empty_like(forces) + for sorted_index, original_index in enumerate(order): + reordered[original_index] = forces[sorted_index] + return reordered + + class TestAbacusCalculator(unittest.TestCase): here = Path(__file__).parent pporb = here.parent.parent.parent / 'tests' / 'PP_ORB' + def test_socketio_species_order_mapping(self): + atoms = Atoms(symbols=['Si', 'O', 'C', 'Si', 'O', 'C']) + order = AbacusSocketIO._socket_sort_indices(atoms) + self.assertEqual(order, [0, 3, 1, 4, 2, 5]) + + socket_forces = np.arange(18).reshape(6, 3) + input_forces = AbacusSocketIO._forces_to_input_order( + socket_forces, order) + + expected = np.empty_like(socket_forces) + for sorted_index, original_index in enumerate(order): + expected[original_index] = socket_forces[sorted_index] + np.testing.assert_array_equal(input_forces, expected) + + def test_socketio_rejects_parameter_changes(self): + calc = object.__new__(AbacusSocketIO) + with self.assertRaisesRegex(ValueError, 'fixed after construction'): + calc.set(kpts={'mode': 'mp-sampling', 'nk': [2, 2, 2]}) + + def test_socketio_rejects_cell_changes(self): + from ase.calculators.calculator import PropertyNotImplementedError + + calc = object.__new__(AbacusSocketIO) + calc.atoms = Atoms('Si', cell=[5.0, 5.0, 5.0], pbc=True) + calc._reference_cell = calc.atoms.cell.array.copy() + + changed = calc.atoms.copy() + changed.cell[0, 0] = 5.1 + with self.assertRaisesRegex(PropertyNotImplementedError, 'fixed-cell'): + calc._check_fixed_cell(changed) + + def test_parse_version_allows_launcher_noise(self): + stdout = 'launcher info\nABACUS version v3.11.0-beta6\n' + self.assertEqual(AbacusProfile.parse_version(stdout), 'v3.11.0-beta6') + + def test_parse_version_rejects_missing_version(self): + with self.assertRaisesRegex(RuntimeError, 'ABACUS version'): + AbacusProfile.parse_version('launcher failed before abacus started') + def test_calculator_results(self): from ase.build.bulk import bulk silicon = bulk('Si', crystalstructure='diamond', a=5.43) diff --git a/interfaces/ASE_interface/examples/socketio.py b/interfaces/ASE_interface/examples/socketio.py new file mode 100644 index 00000000000..e3087453ffc --- /dev/null +++ b/interfaces/ASE_interface/examples/socketio.py @@ -0,0 +1,157 @@ +""" +This example validates and benchmarks ABACUS socket I/O from ASE. + +ASE runs as the i-PI socket server and ABACUS runs as the socket client. +ABACUS keeps calculation=scf and enables socket_driver internally. + +The script checks two PR-review relevant points: +1. socket SCF gives the same energy and forces as a normal non-socket SCF; +2. repeated socket calculations avoid relaunching ABACUS and are faster than + the normal FileIO calculator for a sequence of SCF force evaluations. + +The i-PI protocol does not carry element symbols. AbacusSocketIO handles +the required STRU/socket atom-order alignment internally and returns forces +in the original ASE Atoms order. +""" +import os +import shutil +import time +from pathlib import Path + +import numpy as np +from ase import Atoms +from abacuslite import Abacus, AbacusProfile, AbacusSocketIO + +here = Path(__file__).parent +pporb = here.parent.parent.parent / 'tests' / 'PP_ORB' + +aprof = AbacusProfile( + command=os.environ.get('ABACUS_COMMAND', 'mpirun -np 4 abacus'), + pseudo_dir=pporb, + orbital_dir=pporb, + omp_num_threads=1, +) + +common_kwargs = { + 'pseudopotentials': {'Si': 'Si_ONCV_PBE-1.0.upf'}, + 'basissets': {'Si': 'Si_gga_8au_100Ry_2s2p1d.orb'}, + 'inp': { + 'calculation': 'scf', + 'nspin': 1, + 'basis_type': 'lcao', + 'ks_solver': 'scalapack_gvx', + 'ecutwfc': 30, + 'symmetry': 0, + 'kspacing': 0.5, + 'scf_thr': 1e-8, + 'scf_nmax': 40, + 'chg_extrap': 'atomic', + 'cal_force': 1, + }, +} + +base_atoms = Atoms( + 'Si2', + positions=[[0.0, 0.0, 0.0], [1.25, 1.25, 1.25]], + cell=[5.43, 5.43, 5.43], + pbc=True, +) + + +def clean(directory): + shutil.rmtree(directory, ignore_errors=True) + + +def run_fileio(atoms, directory): + clean(directory) + calc = Abacus(profile=aprof, directory=str(directory), **common_kwargs) + atoms = atoms.copy() + atoms.calc = calc + forces = atoms.get_forces() + energy = atoms.get_potential_energy() + return energy, forces + + +def run_socketio(atoms, directory, socket_name): + clean(directory) + calc = AbacusSocketIO( + profile=aprof, + directory=str(directory), + unixsocket=socket_name, + timeout=120, + **common_kwargs, + ) + atoms = atoms.copy() + with calc: + atoms.calc = calc + energy = atoms.get_potential_energy() + forces = atoms.get_forces() + return energy, forces + + +def displaced_structures(): + structures = [] + for scale in (0.00, 0.03, -0.02, 0.05): + atoms = base_atoms.copy() + atoms.positions[1] += scale + structures.append(atoms) + return structures + + +fileio_dir = here / 'socketio_fileio_scf' +socket_dir = here / 'socketio_socket_scf' +bench_fileio_dir = here / 'socketio_bench_fileio' +bench_socket_dir = here / 'socketio_bench_socket' + +try: + reference_energy, reference_forces = run_fileio(base_atoms, fileio_dir) + socket_energy, socket_forces = run_socketio( + base_atoms, socket_dir, 'abacus_si_check') + + energy_diff = abs(socket_energy - reference_energy) + force_diff = np.max(np.abs(socket_forces - reference_forces)) + print(f'FileIO SCF energy: {reference_energy:.12f} eV') + print(f'Socket SCF energy: {socket_energy:.12f} eV') + print(f'|dE|: {energy_diff:.3e} eV') + print(f'max |dF|: {force_diff:.3e} eV/Angstrom') + assert energy_diff < 1e-4 + assert force_diff < 1e-5 + + structures = displaced_structures() + + clean(bench_fileio_dir) + fileio_calc = Abacus( + profile=aprof, + directory=str(bench_fileio_dir), + **common_kwargs, + ) + t0 = time.perf_counter() + for atoms in structures: + atoms = atoms.copy() + atoms.calc = fileio_calc + atoms.get_forces() + fileio_seconds = time.perf_counter() - t0 + + clean(bench_socket_dir) + socket_calc = AbacusSocketIO( + profile=aprof, + directory=str(bench_socket_dir), + unixsocket='abacus_si_bench', + timeout=120, + **common_kwargs, + ) + t0 = time.perf_counter() + with socket_calc as calc: + for atoms in structures: + atoms = atoms.copy() + atoms.calc = calc + atoms.get_forces() + socket_seconds = time.perf_counter() - t0 + + speedup = fileio_seconds / socket_seconds + print(f'FileIO repeated SCF force time: {fileio_seconds:.2f} s') + print(f'Socket repeated SCF force time: {socket_seconds:.2f} s') + print(f'Socket speedup vs FileIO: {speedup:.2f}x') +finally: + for directory in (fileio_dir, socket_dir, bench_fileio_dir, bench_socket_dir): + clean(directory) diff --git a/source/source_io/module_parameter/input_parameter.h b/source/source_io/module_parameter/input_parameter.h index 3045d663361..c7cb990da21 100644 --- a/source/source_io/module_parameter/input_parameter.h +++ b/source/source_io/module_parameter/input_parameter.h @@ -19,6 +19,7 @@ struct Input_para std::string calculation = "scf"; ///< "scf" : self consistent calculation. ///< "nscf" : non-self consistent calculation. ///< "relax" : cell relaxations + bool socket_driver = false; ///< run ABACUS as an i-PI socket client std::string esolver_type = "ksdft"; ///< the energy solver: ksdft, sdft, ofdft, tddft, lj, dp /* symmetry level: -1, no symmetry at all; 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..10444a1b4bd 100644 --- a/source/source_io/module_parameter/read_input_item_system.cpp +++ b/source/source_io/module_parameter/read_input_item_system.cpp @@ -136,6 +136,27 @@ void ReadInput::item_system() sync_string(input.calculation); this->add_item(item); } + { + Input_Item item("socket_driver"); + item.annotation = "run as a socket client for external drivers using the i-PI protocol"; + item.category = "System variables"; + item.type = "Boolean"; + item.description = R"(If set to True, ABACUS keeps the calculation type as scf and receives atomic positions from an external driver through the i-PI socket protocol. + +[NOTE] Use calculation = scf with socket_driver = True. ABACUS connects to the external i-PI server selected by ABACUS_SOCKET_ADDRESS. If ABACUS_SOCKET_ADDRESS is unset, ABACUS uses localhost:31415. The value can use one of two forms: +* host:port, for example localhost:31415 or 127.0.0.1:31415, opens a TCP connection to that host and port. Use this when the i-PI server listens on a TCP port. +* path:UNIX, for example /tmp/ipi_abacus_si:UNIX, opens a Unix-domain socket at the given filesystem path. The :UNIX suffix tells ABACUS that the preceding value is a local socket path rather than a TCP host name. This form only works on the same machine. +When using the ASE AbacusSocketIO interface, this environment variable is set automatically from the port or unixsocket calculator argument.)"; + item.default_value = "False"; + read_sync_bool(input.socket_driver); + item.check_value = [](const Input_Item& item, const Parameter& para) { + if (para.input.socket_driver && para.input.calculation != "scf") + { + ModuleBase::WARNING_QUIT("ReadInput", "socket_driver is only supported with calculation = scf."); + } + }; + this->add_item(item); + } { Input_Item item("esolver_type"); item.annotation = "the energy solver: ksdft, sdft, ofdft, tdofdft, tddft, lj, dp, ks-lr, lr"; @@ -268,7 +289,7 @@ void ReadInput::item_system() item.reset_value = [](const Input_Item& item, Parameter& para) { std::vector use_force = {"cell-relax", "relax", "md"}; 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()) + if (para.input.socket_driver || std::find(use_force.begin(), use_force.end(), para.input.calculation) != use_force.end()) { if (!para.input.cal_force) { @@ -873,7 +894,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 and socket_driver calculations, 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 +908,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.socket_driver)) { 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..0a53fa13d85 100644 --- a/source/source_io/test_serial/read_input_item_test.cpp +++ b/source/source_io/test_serial/read_input_item_test.cpp @@ -64,6 +64,27 @@ TEST_F(InputTest, Item_test) EXPECT_EXIT(it->second.check_value(it->second, param), ::testing::ExitedWithCode(1), ""); output = testing::internal::GetCapturedStdout(); EXPECT_THAT(output, testing::HasSubstr("NOTICE")); + + param.input.calculation = "socket"; + testing::internal::CaptureStdout(); + EXPECT_EXIT(it->second.check_value(it->second, param), ::testing::ExitedWithCode(1), ""); + output = testing::internal::GetCapturedStdout(); + EXPECT_THAT(output, testing::HasSubstr("NOTICE")); + } + + { // socket_driver + auto it = find_label("socket_driver", readinput.input_lists); + param.input.socket_driver = true; + param.input.calculation = "nscf"; + testing::internal::CaptureStdout(); + EXPECT_EXIT(it->second.check_value(it->second, param), ::testing::ExitedWithCode(1), ""); + output = testing::internal::GetCapturedStdout(); + EXPECT_THAT(output, testing::HasSubstr("NOTICE")); + + param.input.socket_driver = true; + param.input.calculation = "scf"; + EXPECT_NO_THROW(it->second.check_value(it->second, param)); + param.input.socket_driver = false; } { // esolver_type @@ -246,6 +267,7 @@ TEST_F(InputTest, Item_test) auto it = find_label("cal_force", readinput.input_lists); param.input.calculation = "cell-relax"; param.input.cal_force = false; + param.input.socket_driver = false; it->second.reset_value(it->second, param); EXPECT_EQ(param.input.cal_force, true); @@ -253,6 +275,13 @@ TEST_F(InputTest, Item_test) param.input.cal_force = true; it->second.reset_value(it->second, param); EXPECT_EQ(param.input.cal_force, false); + + param.input.calculation = "scf"; + param.input.socket_driver = true; + param.input.cal_force = false; + it->second.reset_value(it->second, param); + EXPECT_EQ(param.input.cal_force, true); + param.input.socket_driver = false; } { // ecutrho auto it = find_label("ecutrho", readinput.input_lists); @@ -362,8 +391,15 @@ 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 = "scf"; + param.input.socket_driver = true; + it->second.reset_value(it->second, param); + EXPECT_EQ(param.input.chg_extrap, "first-order"); + param.input.chg_extrap = "default"; param.input.calculation = "none"; + param.input.socket_driver = false; it->second.reset_value(it->second, param); EXPECT_EQ(param.input.chg_extrap, "atomic"); diff --git a/source/source_main/driver_run.cpp b/source/source_main/driver_run.cpp index 9dc2935c64d..74f5e521192 100644 --- a/source/source_main/driver_run.cpp +++ b/source/source_main/driver_run.cpp @@ -74,9 +74,8 @@ void Driver::driver_run() 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); diff --git a/source/source_relax/CMakeLists.txt b/source/source_relax/CMakeLists.txt index 9e7ef96b0e2..77079eb6a1f 100644 --- a/source/source_relax/CMakeLists.txt +++ b/source/source_relax/CMakeLists.txt @@ -2,6 +2,8 @@ add_library( relax OBJECT relax_data.cpp + socket_ipi.cpp + socket_driver.cpp cg_base.cpp relax_driver.cpp relax_sync.cpp diff --git a/source/source_relax/relax_driver.cpp b/source/source_relax/relax_driver.cpp index 26128bf0408..922a230d882 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.socket_driver) + { + 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 00000000000..42fafda5e63 --- /dev/null +++ b/source/source_relax/socket_driver.cpp @@ -0,0 +1,533 @@ +#include "socket_driver.h" + +#include "source_relax/socket_ipi.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 socket_address() +{ + const char* env = std::getenv("ABACUS_SOCKET_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_driver 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 = socket_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 00000000000..86c180fc42e --- /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/socket_ipi.cpp b/source/source_relax/socket_ipi.cpp new file mode 100644 index 00000000000..2bb09385b0d --- /dev/null +++ b/source/source_relax/socket_ipi.cpp @@ -0,0 +1,270 @@ +#include "source_relax/socket_ipi.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; + } + if (errno == ECONNRESET && done == 0) + { + throw IpiSocketClosed("i-PI socket peer reset before next header"); + } + 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/socket_ipi.h b/source/source_relax/socket_ipi.h new file mode 100644 index 00000000000..eac15eb41ac --- /dev/null +++ b/source/source_relax/socket_ipi.h @@ -0,0 +1,47 @@ +#ifndef ABACUS_SOCKET_IPI_H +#define ABACUS_SOCKET_IPI_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/test/CMakeLists.txt b/source/source_relax/test/CMakeLists.txt index 2262fefb419..51dc9a82d3b 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_socket_ipi_test + SOURCES socket_ipi_test.cpp ../socket_ipi.cpp +) + AddTest( TARGET MODULE_RELAX_relax_new_line_search LIBS parameter diff --git a/source/source_relax/test/socket_ipi_test.cpp b/source/source_relax/test/socket_ipi_test.cpp new file mode 100644 index 00000000000..930ab1fe629 --- /dev/null +++ b/source/source_relax/test/socket_ipi_test.cpp @@ -0,0 +1,256 @@ +#include "../socket_ipi.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_socket_ipi_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); +}