Coverage for /builds/ase/ase/ase/calculators/socketio.py: 89.78%
401 statements
« prev ^ index » next coverage.py v7.5.3, created at 2025-08-02 00:12 +0000
1# fmt: off
3import os
4import socket
5from contextlib import ExitStack, contextmanager
6from subprocess import PIPE, Popen
8import numpy as np
10import ase.units as units
11from ase.calculators.calculator import (
12 Calculator,
13 OldShellProfile,
14 PropertyNotImplementedError,
15 StandardProfile,
16 all_changes,
17)
18from ase.calculators.genericfileio import GenericFileIOCalculator
19from ase.parallel import world
20from ase.stress import full_3x3_to_voigt_6_stress
21from ase.utils import IOContext
24def actualunixsocketname(name):
25 return f'/tmp/ipi_{name}'
28class SocketClosed(OSError):
29 pass
32class IPIProtocol:
33 """Communication using IPI protocol."""
35 def __init__(self, socket, txt=None):
36 self.socket = socket
38 if txt is None:
39 def log(*args):
40 pass
41 else:
42 def log(*args):
43 print('Driver:', *args, file=txt)
44 txt.flush()
45 self.log = log
47 def sendmsg(self, msg):
48 self.log(' sendmsg', repr(msg))
49 # assert msg in self.statements, msg
50 msg = msg.encode('ascii').ljust(12)
51 self.socket.sendall(msg)
53 def _recvall(self, nbytes):
54 """Repeatedly read chunks until we have nbytes.
56 Normally we get all bytes in one read, but that is not guaranteed."""
57 remaining = nbytes
58 chunks = []
59 while remaining > 0:
60 chunk = self.socket.recv(remaining)
61 if len(chunk) == 0:
62 # (If socket is still open, recv returns at least one byte)
63 raise SocketClosed
64 chunks.append(chunk)
65 remaining -= len(chunk)
66 msg = b''.join(chunks)
67 assert len(msg) == nbytes and remaining == 0
68 return msg
70 def recvmsg(self):
71 msg = self._recvall(12)
72 if not msg:
73 raise SocketClosed
75 assert len(msg) == 12, msg
76 msg = msg.rstrip().decode('ascii')
77 # assert msg in self.responses, msg
78 self.log(' recvmsg', repr(msg))
79 return msg
81 def send(self, a, dtype):
82 buf = np.asarray(a, dtype).tobytes()
83 # self.log(' send {}'.format(np.array(a).ravel().tolist()))
84 self.log(f' send {len(buf)} bytes of {dtype}')
85 self.socket.sendall(buf)
87 def recv(self, shape, dtype):
88 a = np.empty(shape, dtype)
89 nbytes = np.dtype(dtype).itemsize * np.prod(shape)
90 buf = self._recvall(nbytes)
91 assert len(buf) == nbytes, (len(buf), nbytes)
92 self.log(f' recv {len(buf)} bytes of {dtype}')
93 # print(np.frombuffer(buf, dtype=dtype))
94 a.flat[:] = np.frombuffer(buf, dtype=dtype)
95 # self.log(' recv {}'.format(a.ravel().tolist()))
96 assert np.isfinite(a).all()
97 return a
99 def sendposdata(self, cell, icell, positions):
100 assert cell.size == 9
101 assert icell.size == 9
102 assert positions.size % 3 == 0
104 self.log(' sendposdata')
105 self.sendmsg('POSDATA')
106 self.send(cell.T / units.Bohr, np.float64)
107 self.send(icell.T * units.Bohr, np.float64)
108 self.send(len(positions), np.int32)
109 self.send(positions / units.Bohr, np.float64)
111 def recvposdata(self):
112 cell = self.recv((3, 3), np.float64).T.copy()
113 icell = self.recv((3, 3), np.float64).T.copy()
114 natoms = self.recv(1, np.int32)[0]
115 positions = self.recv((natoms, 3), np.float64)
116 return cell * units.Bohr, icell / units.Bohr, positions * units.Bohr
118 def sendrecv_force(self):
119 self.log(' sendrecv_force')
120 self.sendmsg('GETFORCE')
121 msg = self.recvmsg()
122 assert msg == 'FORCEREADY', msg
123 e = self.recv(1, np.float64)[0]
124 natoms = self.recv(1, np.int32)[0]
125 assert natoms >= 0
126 forces = self.recv((int(natoms), 3), np.float64)
127 virial = self.recv((3, 3), np.float64).T.copy()
128 nmorebytes = self.recv(1, np.int32)[0]
129 morebytes = self.recv(nmorebytes, np.byte)
130 return (e * units.Ha, (units.Ha / units.Bohr) * forces,
131 units.Ha * virial, morebytes)
133 def sendforce(self, energy, forces, virial,
134 morebytes=np.zeros(1, dtype=np.byte)):
135 assert np.array([energy]).size == 1
136 assert forces.shape[1] == 3
137 assert virial.shape == (3, 3)
139 self.log(' sendforce')
140 self.sendmsg('FORCEREADY') # mind the units
141 self.send(np.array([energy / units.Ha]), np.float64)
142 natoms = len(forces)
143 self.send(np.array([natoms]), np.int32)
144 self.send(units.Bohr / units.Ha * forces, np.float64)
145 self.send(1.0 / units.Ha * virial.T, np.float64)
146 # We prefer to always send at least one byte due to trouble with
147 # empty messages. Reading a closed socket yields 0 bytes
148 # and thus can be confused with a 0-length bytestring.
149 self.send(np.array([len(morebytes)]), np.int32)
150 self.send(morebytes, np.byte)
152 def status(self):
153 self.log(' status')
154 self.sendmsg('STATUS')
155 msg = self.recvmsg()
156 return msg
158 def end(self):
159 self.log(' end')
160 self.sendmsg('EXIT')
162 def recvinit(self):
163 self.log(' recvinit')
164 bead_index = self.recv(1, np.int32)
165 nbytes = self.recv(1, np.int32)
166 initbytes = self.recv(nbytes, np.byte)
167 return bead_index, initbytes
169 def sendinit(self):
170 # XXX Not sure what this function is supposed to send.
171 # It 'works' with QE, but for now we try not to call it.
172 self.log(' sendinit')
173 self.sendmsg('INIT')
174 self.send(0, np.int32) # 'bead index' always zero for now
175 # We send one byte, which is zero, since things may not work
176 # with 0 bytes. Apparently implementations ignore the
177 # initialization string anyway.
178 self.send(1, np.int32)
179 self.send(np.zeros(1), np.byte) # initialization string
181 def calculate(self, positions, cell):
182 self.log('calculate')
183 msg = self.status()
184 # We don't know how NEEDINIT is supposed to work, but some codes
185 # seem to be okay if we skip it and send the positions instead.
186 if msg == 'NEEDINIT':
187 self.sendinit()
188 msg = self.status()
189 assert msg == 'READY', msg
190 icell = np.linalg.pinv(cell).transpose()
191 self.sendposdata(cell, icell, positions)
192 msg = self.status()
193 assert msg == 'HAVEDATA', msg
194 e, forces, virial, morebytes = self.sendrecv_force()
195 r = dict(energy=e,
196 forces=forces,
197 virial=virial,
198 morebytes=morebytes)
199 return r
202@contextmanager
203def bind_unixsocket(socketfile):
204 assert socketfile.startswith('/tmp/ipi_'), socketfile
205 serversocket = socket.socket(socket.AF_UNIX)
206 try:
207 serversocket.bind(socketfile)
208 except OSError as err:
209 raise OSError(f'{err}: {socketfile!r}')
211 try:
212 with serversocket:
213 yield serversocket
214 finally:
215 os.unlink(socketfile)
218@contextmanager
219def bind_inetsocket(port):
220 serversocket = socket.socket(socket.AF_INET)
221 serversocket.setsockopt(socket.SOL_SOCKET,
222 socket.SO_REUSEADDR, 1)
223 serversocket.bind(('', port))
224 with serversocket:
225 yield serversocket
228class FileIOSocketClientLauncher:
229 def __init__(self, calc):
230 self.calc = calc
232 def __call__(self, atoms, properties=None, port=None, unixsocket=None):
233 assert self.calc is not None
234 cwd = self.calc.directory
236 profile = getattr(self.calc, 'profile', None)
237 if isinstance(self.calc, GenericFileIOCalculator):
238 # New GenericFileIOCalculator:
239 template = getattr(self.calc, 'template')
241 self.calc.write_inputfiles(atoms, properties)
242 if unixsocket is not None:
243 argv = template.socketio_argv(
244 profile, unixsocket=unixsocket, port=None
245 )
246 else:
247 argv = template.socketio_argv(
248 profile, unixsocket=None, port=port
249 )
251 if hasattr(self.calc.template, "outputname"):
252 with ExitStack() as stack:
253 output_path = self.calc.template.outputname
254 fd_out = stack.enter_context(open(output_path, "w"))
255 return Popen(argv, cwd=cwd, env=os.environ, stdout=fd_out)
257 return Popen(argv, cwd=cwd, env=os.environ)
258 else:
259 # Old FileIOCalculator:
260 self.calc.write_input(atoms, properties=properties,
261 system_changes=all_changes)
263 if isinstance(profile, StandardProfile):
264 return profile.execute_nonblocking(self.calc)
266 if profile is None:
267 cmd = self.calc.command.replace('PREFIX', self.calc.prefix)
268 cmd = cmd.format(port=port, unixsocket=unixsocket)
269 elif isinstance(profile, OldShellProfile):
270 cmd = profile.command.replace("PREFIX", self.calc.prefix)
271 else:
272 raise TypeError(
273 f"Profile type {type(profile)} not supported for socketio")
275 return Popen(cmd, shell=True, cwd=cwd)
278class SocketServer(IOContext):
279 default_port = 31415
281 def __init__(self, # launch_client=None,
282 port=None, unixsocket=None, timeout=None,
283 log=None):
284 """Create server and listen for connections.
286 Parameters:
288 client_command: Shell command to launch client process, or None
289 The process will be launched immediately, if given.
290 Else the user is expected to launch a client whose connection
291 the server will then accept at any time.
292 One calculate() is called, the server will block to wait
293 for the client.
294 port: integer or None
295 Port on which to listen for INET connections. Defaults
296 to 31415 if neither this nor unixsocket is specified.
297 unixsocket: string or None
298 Filename for unix socket.
299 timeout: float or None
300 timeout in seconds, or unlimited by default.
301 This parameter is passed to the Python socket object; see
302 documentation therof
303 log: file object or None
304 useful debug messages are written to this."""
306 if unixsocket is None and port is None:
307 port = self.default_port
308 elif unixsocket is not None and port is not None:
309 raise ValueError('Specify only one of unixsocket and port')
311 self.port = port
312 self.unixsocket = unixsocket
313 self.timeout = timeout
314 self._closed = False
316 if unixsocket is not None:
317 actualsocket = actualunixsocketname(unixsocket)
318 conn_name = f'UNIX-socket {actualsocket}'
319 socket_context = bind_unixsocket(actualsocket)
320 else:
321 conn_name = f'INET port {port}'
322 socket_context = bind_inetsocket(port)
324 self.serversocket = self.closelater(socket_context)
326 if log:
327 print(f'Accepting clients on {conn_name}', file=log)
329 self.serversocket.settimeout(timeout)
331 self.serversocket.listen(1)
333 self.log = log
335 self.proc = None
337 self.protocol = None
338 self.clientsocket = None
339 self.address = None
341 # if launch_client is not None:
342 # self.proc = launch_client(port=port, unixsocket=unixsocket)
344 def _accept(self):
345 """Wait for client and establish connection."""
346 # It should perhaps be possible for process to be launched by user
347 log = self.log
348 if log:
349 print('Awaiting client', file=self.log)
351 # If we launched the subprocess, the process may crash.
352 # We want to detect this, using loop with timeouts, and
353 # raise an error rather than blocking forever.
354 if self.proc is not None:
355 self.serversocket.settimeout(1.0)
357 while True:
358 try:
359 self.clientsocket, self.address = self.serversocket.accept()
360 self.closelater(self.clientsocket)
361 except socket.timeout:
362 if self.proc is not None:
363 status = self.proc.poll()
364 if status is not None:
365 raise OSError('Subprocess terminated unexpectedly'
366 ' with status {}'.format(status))
367 else:
368 break
370 self.serversocket.settimeout(self.timeout)
371 self.clientsocket.settimeout(self.timeout)
373 if log:
374 # For unix sockets, address is b''.
375 source = ('client' if self.address == b'' else self.address)
376 print(f'Accepted connection from {source}', file=log)
378 self.protocol = IPIProtocol(self.clientsocket, txt=log)
380 def close(self):
381 if self._closed:
382 return
384 super().close()
386 if self.log:
387 print('Close socket server', file=self.log)
388 self._closed = True
390 # Proper way to close sockets?
391 # And indeed i-pi connections...
392 # if self.protocol is not None:
393 # self.protocol.end() # Send end-of-communication string
394 self.protocol = None
395 if self.proc is not None:
396 exitcode = self.proc.wait()
397 if exitcode != 0:
398 import warnings
400 # Quantum Espresso seems to always exit with status 128,
401 # even if successful.
402 # Should investigate at some point
403 warnings.warn('Subprocess exited with status {}'
404 .format(exitcode))
405 # self.log('IPI server closed')
407 def calculate(self, atoms):
408 """Send geometry to client and return calculated things as dict.
410 This will block until client has established connection, then
411 wait for the client to finish the calculation."""
412 assert not self._closed
414 # If we have not established connection yet, we must block
415 # until the client catches up:
416 if self.protocol is None:
417 self._accept()
418 return self.protocol.calculate(atoms.positions, atoms.cell)
421class SocketClient:
422 def __init__(self, host='localhost', port=None,
423 unixsocket=None, timeout=None, log=None, comm=world):
424 """Create client and connect to server.
426 Parameters:
428 host: string
429 Hostname of server. Defaults to localhost
430 port: integer or None
431 Port to which to connect. By default 31415.
432 unixsocket: string or None
433 If specified, use corresponding UNIX socket.
434 See documentation of unixsocket for SocketIOCalculator.
435 timeout: float or None
436 See documentation of timeout for SocketIOCalculator.
437 log: file object or None
438 Log events to this file
439 comm: communicator or None
440 MPI communicator object. Defaults to ase.parallel.world.
441 When ASE runs in parallel, only the process with world.rank == 0
442 will communicate over the socket. The received information
443 will then be broadcast on the communicator. The SocketClient
444 must be created on all ranks of world, and will see the same
445 Atoms objects."""
446 # Only rank0 actually does the socket work.
447 # The other ranks only need to follow.
448 #
449 # Note: We actually refrain from assigning all the
450 # socket-related things except on master
451 self.comm = comm
453 if self.comm.rank == 0:
454 if unixsocket is not None:
455 sock = socket.socket(socket.AF_UNIX)
456 actualsocket = actualunixsocketname(unixsocket)
457 sock.connect(actualsocket)
458 else:
459 if port is None:
460 port = SocketServer.default_port
461 sock = socket.socket(socket.AF_INET)
462 sock.connect((host, port))
463 sock.settimeout(timeout)
464 self.host = host
465 self.port = port
466 self.unixsocket = unixsocket
468 self.protocol = IPIProtocol(sock, txt=log)
469 self.log = self.protocol.log
470 self.closed = False
472 self.bead_index = 0
473 self.bead_initbytes = b''
474 self.state = 'READY'
476 def close(self):
477 if not self.closed:
478 self.log('Close SocketClient')
479 self.closed = True
480 self.protocol.socket.close()
482 def calculate(self, atoms, use_stress):
483 # We should also broadcast the bead index, once we support doing
484 # multiple beads.
485 self.comm.broadcast(atoms.positions, 0)
486 self.comm.broadcast(np.ascontiguousarray(atoms.cell), 0)
488 energy = atoms.get_potential_energy()
489 forces = atoms.get_forces()
490 if use_stress:
491 stress = atoms.get_stress(voigt=False)
492 virial = -atoms.get_volume() * stress
493 else:
494 virial = np.zeros((3, 3))
495 return energy, forces, virial
497 def irun(self, atoms, use_stress=None):
498 if use_stress is None:
499 use_stress = any(atoms.pbc)
501 my_irun = self.irun_rank0 if self.comm.rank == 0 else self.irun_rankN
502 return my_irun(atoms, use_stress)
504 def irun_rankN(self, atoms, use_stress=True):
505 stop_criterion = np.zeros(1, bool)
506 while True:
507 self.comm.broadcast(stop_criterion, 0)
508 if stop_criterion[0]:
509 return
511 self.calculate(atoms, use_stress)
512 yield
514 def irun_rank0(self, atoms, use_stress=True):
515 # For every step we either calculate or quit. We need to
516 # tell other MPI processes (if this is MPI-parallel) whether they
517 # should calculate or quit.
518 try:
519 while True:
520 try:
521 msg = self.protocol.recvmsg()
522 except SocketClosed:
523 # Server closed the connection, but we want to
524 # exit gracefully anyway
525 msg = 'EXIT'
527 if msg == 'EXIT':
528 # Send stop signal to clients:
529 self.comm.broadcast(np.ones(1, bool), 0)
530 # (When otherwise exiting, things crashed and we should
531 # let MPI_ABORT take care of the mess instead of trying
532 # to synchronize the exit)
533 return
534 elif msg == 'STATUS':
535 self.protocol.sendmsg(self.state)
536 elif msg == 'POSDATA':
537 assert self.state == 'READY'
538 cell, _icell, positions = self.protocol.recvposdata()
539 atoms.cell[:] = cell
540 atoms.positions[:] = positions
542 # User may wish to do something with the atoms object now.
543 # Should we provide option to yield here?
544 #
545 # (In that case we should MPI-synchronize *before*
546 # whereas now we do it after.)
548 # Send signal for other ranks to proceed with calculation:
549 self.comm.broadcast(np.zeros(1, bool), 0)
550 energy, forces, virial = self.calculate(atoms, use_stress)
552 self.state = 'HAVEDATA'
553 yield
554 elif msg == 'GETFORCE':
555 assert self.state == 'HAVEDATA', self.state
556 self.protocol.sendforce(energy, forces, virial)
557 self.state = 'NEEDINIT'
558 elif msg == 'INIT':
559 assert self.state == 'NEEDINIT'
560 bead_index, initbytes = self.protocol.recvinit()
561 self.bead_index = bead_index
562 self.bead_initbytes = initbytes
563 self.state = 'READY'
564 else:
565 raise KeyError('Bad message', msg)
566 finally:
567 self.close()
569 def run(self, atoms, use_stress=False):
570 for _ in self.irun(atoms, use_stress=use_stress):
571 pass
574class SocketIOCalculator(Calculator, IOContext):
575 implemented_properties = ['energy', 'free_energy', 'forces', 'stress']
576 supported_changes = {'positions', 'cell'}
578 def __init__(self, calc=None, port=None,
579 unixsocket=None, timeout=None, log=None, *,
580 launch_client=None, comm=world):
581 """Initialize socket I/O calculator.
583 This calculator launches a server which passes atomic
584 coordinates and unit cells to an external code via a socket,
585 and receives energy, forces, and stress in return.
587 ASE integrates this with the Quantum Espresso, FHI-aims and
588 Siesta calculators. This works with any external code that
589 supports running as a client over the i-PI protocol.
591 Parameters:
593 calc: calculator or None
595 If calc is not None, a client process will be launched
596 using calc.command, and the input file will be generated
597 using ``calc.write_input()``. Otherwise only the server will
598 run, and it is up to the user to launch a compliant client
599 process.
601 port: integer
603 port number for socket. Should normally be between 1025
604 and 65535. Typical ports for are 31415 (default) or 3141.
606 unixsocket: str or None
608 if not None, ignore host and port, creating instead a
609 unix socket using this name prefixed with ``/tmp/ipi_``.
610 The socket is deleted when the calculator is closed.
612 timeout: float >= 0 or None
614 timeout for connection, by default infinite. See
615 documentation of Python sockets. For longer jobs it is
616 recommended to set a timeout in case of undetected
617 client-side failure.
619 log: file object or None (default)
621 logfile for communication over socket. For debugging or
622 the curious.
624 In order to correctly close the sockets, it is
625 recommended to use this class within a with-block:
627 >>> from ase.calculators.socketio import SocketIOCalculator
629 >>> with SocketIOCalculator(...) as calc: # doctest:+SKIP
630 ... atoms.calc = calc
631 ... atoms.get_forces()
632 ... atoms.rattle()
633 ... atoms.get_forces()
635 It is also possible to call calc.close() after
636 use. This is best done in a finally-block."""
638 Calculator.__init__(self)
640 if calc is not None:
641 if launch_client is not None:
642 raise ValueError('Cannot pass both calc and launch_client')
643 launch_client = FileIOSocketClientLauncher(calc)
644 self.launch_client = launch_client
645 self.timeout = timeout
646 self.server = None
648 self.log = self.openfile(file=log, comm=comm)
650 # We only hold these so we can pass them on to the server.
651 # They may both be None as stored here.
652 self._port = port
653 self._unixsocket = unixsocket
655 # If there is a calculator, we will launch in calculate() because
656 # we are responsible for executing the external process, too, and
657 # should do so before blocking. Without a calculator we want to
658 # block immediately:
659 if self.launch_client is None:
660 self.server = self.launch_server()
662 def todict(self):
663 d = {'type': 'calculator',
664 'name': 'socket-driver'}
665 # if self.calc is not None:
666 # d['calc'] = self.calc.todict()
667 return d
669 def launch_server(self):
670 return self.closelater(SocketServer(
671 # launch_client=launch_client,
672 port=self._port,
673 unixsocket=self._unixsocket,
674 timeout=self.timeout, log=self.log,
675 ))
677 def calculate(self, atoms=None, properties=['energy'],
678 system_changes=all_changes):
679 bad = [change for change in system_changes
680 if change not in self.supported_changes]
682 # First time calculate() is called, system_changes will be
683 # all_changes. After that, only positions and cell may change.
684 if self.atoms is not None and any(bad):
685 raise PropertyNotImplementedError(
686 'Cannot change {} through IPI protocol. '
687 'Please create new socket calculator.'
688 .format(bad if len(bad) > 1 else bad[0]))
690 self.atoms = atoms.copy()
692 if self.server is None:
693 self.server = self.launch_server()
694 proc = self.launch_client(atoms, properties,
695 port=self._port,
696 unixsocket=self._unixsocket)
697 self.server.proc = proc # XXX nasty hack
699 results = self.server.calculate(atoms)
700 results['free_energy'] = results['energy']
701 virial = results.pop('virial')
702 if self.atoms.cell.rank == 3 and any(self.atoms.pbc):
703 vol = atoms.get_volume()
704 results['stress'] = -full_3x3_to_voigt_6_stress(virial) / vol
705 self.results.update(results)
707 def close(self):
708 self.server = None
709 super().close()
712class PySocketIOClient:
713 def __init__(self, calculator_factory):
714 self._calculator_factory = calculator_factory
716 def __call__(self, atoms, properties=None, port=None, unixsocket=None):
717 import pickle
718 import sys
720 # We pickle everything first, so we won't need to bother with the
721 # process as long as it succeeds.
722 transferbytes = pickle.dumps([
723 dict(unixsocket=unixsocket, port=port),
724 atoms.copy(),
725 self._calculator_factory,
726 ])
728 proc = Popen([sys.executable, '-m', 'ase.calculators.socketio'],
729 stdin=PIPE)
731 proc.stdin.write(transferbytes)
732 proc.stdin.close()
733 return proc
735 @staticmethod
736 def main():
737 import pickle
738 import sys
740 socketinfo, atoms, get_calculator = pickle.load(sys.stdin.buffer)
741 atoms.calc = get_calculator()
742 client = SocketClient(host='localhost',
743 unixsocket=socketinfo.get('unixsocket'),
744 port=socketinfo.get('port'))
745 # XXX In principle we could avoid calculating stress until
746 # someone requests the stress, could we not?
747 # Which would make use_stress boolean unnecessary.
748 client.run(atoms, use_stress=True)
751if __name__ == '__main__':
752 PySocketIOClient.main()