Coverage for ase / calculators / socketio.py: 89.72%
399 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-03-30 08:22 +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
287 ----------
289 client_command: Shell command to launch client process, or None
290 The process will be launched immediately, if given.
291 Else the user is expected to launch a client whose connection
292 the server will then accept at any time.
293 One calculate() is called, the server will block to wait
294 for the client.
295 port: integer or None
296 Port on which to listen for INET connections. Defaults
297 to 31415 if neither this nor unixsocket is specified.
298 unixsocket: string or None
299 Filename for unix socket.
300 timeout: float or None
301 timeout in seconds, or unlimited by default.
302 This parameter is passed to the Python socket object; see
303 documentation therof
304 log: file object or None
305 useful debug messages are written to this."""
307 if unixsocket is None and port is None:
308 port = self.default_port
309 elif unixsocket is not None and port is not None:
310 raise ValueError('Specify only one of unixsocket and port')
312 self.port = port
313 self.unixsocket = unixsocket
314 self.timeout = timeout
315 self._closed = False
317 if unixsocket is not None:
318 actualsocket = actualunixsocketname(unixsocket)
319 conn_name = f'UNIX-socket {actualsocket}'
320 socket_context = bind_unixsocket(actualsocket)
321 else:
322 conn_name = f'INET port {port}'
323 socket_context = bind_inetsocket(port)
325 self.serversocket = self.closelater(socket_context)
327 if log:
328 print(f'Accepting clients on {conn_name}', file=log)
330 self.serversocket.settimeout(timeout)
332 self.serversocket.listen(1)
334 self.log = log
336 self.proc = None
338 self.protocol = None
339 self.clientsocket = None
340 self.address = None
342 # if launch_client is not None:
343 # self.proc = launch_client(port=port, unixsocket=unixsocket)
345 def _accept(self):
346 """Wait for client and establish connection."""
347 # It should perhaps be possible for process to be launched by user
348 log = self.log
349 if log:
350 print('Awaiting client', file=self.log)
352 # If we launched the subprocess, the process may crash.
353 # We want to detect this, using loop with timeouts, and
354 # raise an error rather than blocking forever.
355 if self.proc is not None:
356 self.serversocket.settimeout(1.0)
358 while True:
359 try:
360 self.clientsocket, self.address = self.serversocket.accept()
361 self.closelater(self.clientsocket)
362 except socket.timeout:
363 if self.proc is not None:
364 status = self.proc.poll()
365 if status is not None:
366 raise OSError('Subprocess terminated unexpectedly'
367 ' with status {}'.format(status))
368 else:
369 break
371 self.serversocket.settimeout(self.timeout)
372 self.clientsocket.settimeout(self.timeout)
374 if log:
375 # For unix sockets, address is b''.
376 source = ('client' if self.address == b'' else self.address)
377 print(f'Accepted connection from {source}', file=log)
379 self.protocol = IPIProtocol(self.clientsocket, txt=log)
381 def close(self):
382 if self._closed:
383 return
385 super().close()
387 if self.log:
388 print('Close socket server', file=self.log)
389 self._closed = True
391 # Proper way to close sockets?
392 # And indeed i-pi connections...
393 # if self.protocol is not None:
394 # self.protocol.end() # Send end-of-communication string
395 self.protocol = None
396 if self.proc is not None:
397 exitcode = self.proc.wait()
398 if exitcode != 0:
399 import warnings
401 # Quantum Espresso seems to always exit with status 128,
402 # even if successful.
403 # Should investigate at some point
404 warnings.warn('Subprocess exited with status {}'
405 .format(exitcode))
406 # self.log('IPI server closed')
408 def calculate(self, atoms):
409 """Send geometry to client and return calculated things as dict.
411 This will block until client has established connection, then
412 wait for the client to finish the calculation."""
413 assert not self._closed
415 # If we have not established connection yet, we must block
416 # until the client catches up:
417 if self.protocol is None:
418 self._accept()
419 return self.protocol.calculate(atoms.positions, atoms.cell)
422class SocketClient:
423 def __init__(self, host='localhost', port=None,
424 unixsocket=None, timeout=None, log=None, comm=world):
425 """Create client and connect to server.
427 Parameters
428 ----------
430 host: string
431 Hostname of server. Defaults to localhost
432 port: integer or None
433 Port to which to connect. By default 31415.
434 unixsocket: string or None
435 If specified, use corresponding UNIX socket.
436 See documentation of unixsocket for SocketIOCalculator.
437 timeout: float or None
438 See documentation of timeout for SocketIOCalculator.
439 log: file object or None
440 Log events to this file
441 comm: communicator or None
442 MPI communicator object. Defaults to ase.parallel.world.
443 When ASE runs in parallel, only the process with world.rank == 0
444 will communicate over the socket. The received information
445 will then be broadcast on the communicator. The SocketClient
446 must be created on all ranks of world, and will see the same
447 Atoms objects."""
448 # Only rank0 actually does the socket work.
449 # The other ranks only need to follow.
450 #
451 # Note: We actually refrain from assigning all the
452 # socket-related things except on master
453 self.comm = comm
455 if self.comm.rank == 0:
456 if unixsocket is not None:
457 sock = socket.socket(socket.AF_UNIX)
458 actualsocket = actualunixsocketname(unixsocket)
459 sock.connect(actualsocket)
460 else:
461 if port is None:
462 port = SocketServer.default_port
463 sock = socket.socket(socket.AF_INET)
464 sock.connect((host, port))
465 sock.settimeout(timeout)
466 self.host = host
467 self.port = port
468 self.unixsocket = unixsocket
470 self.protocol = IPIProtocol(sock, txt=log)
471 self.log = self.protocol.log
472 self.closed = False
474 self.bead_index = 0
475 self.bead_initbytes = b''
476 self.state = 'READY'
478 def close(self):
479 if not self.closed:
480 self.log('Close SocketClient')
481 self.closed = True
482 self.protocol.socket.close()
484 def calculate(self, atoms, use_stress):
485 # We should also broadcast the bead index, once we support doing
486 # multiple beads.
487 self.comm.broadcast(atoms.positions, 0)
488 self.comm.broadcast(np.ascontiguousarray(atoms.cell), 0)
490 energy = atoms.get_potential_energy()
491 forces = atoms.get_forces()
492 if use_stress:
493 stress = atoms.get_stress(voigt=False)
494 virial = -atoms.get_volume() * stress
495 else:
496 virial = np.zeros((3, 3))
497 return energy, forces, virial
499 def irun(self, atoms, use_stress=None):
500 if use_stress is None:
501 use_stress = any(atoms.pbc)
503 my_irun = self.irun_rank0 if self.comm.rank == 0 else self.irun_rankN
504 return my_irun(atoms, use_stress)
506 def irun_rankN(self, atoms, use_stress=True):
507 stop_criterion = np.zeros(1, bool)
508 while True:
509 self.comm.broadcast(stop_criterion, 0)
510 if stop_criterion[0]:
511 return
513 self.calculate(atoms, use_stress)
514 yield
516 def irun_rank0(self, atoms, use_stress=True):
517 # For every step we either calculate or quit. We need to
518 # tell other MPI processes (if this is MPI-parallel) whether they
519 # should calculate or quit.
520 try:
521 while True:
522 try:
523 msg = self.protocol.recvmsg()
524 except SocketClosed:
525 # Server closed the connection, but we want to
526 # exit gracefully anyway
527 msg = 'EXIT'
529 if msg == 'EXIT':
530 # Send stop signal to clients:
531 self.comm.broadcast(np.ones(1, bool), 0)
532 # (When otherwise exiting, things crashed and we should
533 # let MPI_ABORT take care of the mess instead of trying
534 # to synchronize the exit)
535 return
536 elif msg == 'STATUS':
537 self.protocol.sendmsg(self.state)
538 elif msg == 'POSDATA':
539 assert self.state == 'READY'
540 cell, _icell, positions = self.protocol.recvposdata()
541 atoms.cell[:] = cell
542 atoms.positions[:] = positions
544 # User may wish to do something with the atoms object now.
545 # Should we provide option to yield here?
546 #
547 # (In that case we should MPI-synchronize *before*
548 # whereas now we do it after.)
550 # Send signal for other ranks to proceed with calculation:
551 self.comm.broadcast(np.zeros(1, bool), 0)
552 energy, forces, virial = self.calculate(atoms, use_stress)
554 self.state = 'HAVEDATA'
555 yield
556 elif msg == 'GETFORCE':
557 assert self.state == 'HAVEDATA', self.state
558 self.protocol.sendforce(energy, forces, virial)
559 self.state = 'NEEDINIT'
560 elif msg == 'INIT':
561 assert self.state == 'NEEDINIT'
562 bead_index, initbytes = self.protocol.recvinit()
563 self.bead_index = bead_index
564 self.bead_initbytes = initbytes
565 self.state = 'READY'
566 else:
567 raise KeyError('Bad message', msg)
568 finally:
569 self.close()
571 def run(self, atoms, use_stress=False):
572 for _ in self.irun(atoms, use_stress=use_stress):
573 pass
576class SocketIOCalculator(Calculator, IOContext):
577 implemented_properties = ['energy', 'free_energy', 'forces', 'stress']
578 supported_changes = {'positions', 'cell'}
580 def __init__(self, calc=None, port=None,
581 unixsocket=None, timeout=None, log=None, *,
582 launch_client=None, comm=world):
583 """Initialize socket I/O calculator.
585 This calculator launches a server which passes atomic
586 coordinates and unit cells to an external code via a socket,
587 and receives energy, forces, and stress in return.
589 ASE integrates this with the Quantum Espresso, FHI-aims and
590 Siesta calculators. This works with any external code that
591 supports running as a client over the i-PI protocol.
593 Parameters
594 ----------
596 calc: calculator or None
598 If calc is not None, a client process will be launched
599 using calc.command, and the input file will be generated
600 using ``calc.write_input()``. Otherwise only the server will
601 run, and it is up to the user to launch a compliant client
602 process.
604 port: integer
606 port number for socket. Should normally be between 1025
607 and 65535. Typical ports for are 31415 (default) or 3141.
609 unixsocket: str or None
611 if not None, ignore host and port, creating instead a
612 unix socket using this name prefixed with ``/tmp/ipi_``.
613 The socket is deleted when the calculator is closed.
615 timeout: float >= 0 or None
617 timeout for connection, by default infinite. See
618 documentation of Python sockets. For longer jobs it is
619 recommended to set a timeout in case of undetected
620 client-side failure.
622 log: file object or None (default)
624 logfile for communication over socket. For debugging or
625 the curious.
627 In order to correctly close the sockets, it is
628 recommended to use this class within a with-block:
630 >>> from ase.calculators.socketio import SocketIOCalculator
632 >>> with SocketIOCalculator(...) as calc: # doctest:+SKIP
633 ... atoms.calc = calc
634 ... atoms.get_forces()
635 ... atoms.rattle()
636 ... atoms.get_forces()
638 It is also possible to call calc.close() after
639 use. This is best done in a finally-block."""
641 Calculator.__init__(self)
643 if calc is not None:
644 if launch_client is not None:
645 raise ValueError('Cannot pass both calc and launch_client')
646 launch_client = FileIOSocketClientLauncher(calc)
647 self.launch_client = launch_client
648 self.timeout = timeout
649 self.server = None
651 self.log = self.openfile(file=log, comm=comm)
653 # We only hold these so we can pass them on to the server.
654 # They may both be None as stored here.
655 self._port = port
656 self._unixsocket = unixsocket
658 # If there is a calculator, we will launch in calculate() because
659 # we are responsible for executing the external process, too, and
660 # should do so before blocking. Without a calculator we want to
661 # block immediately:
662 if self.launch_client is None:
663 self.server = self.launch_server()
665 def todict(self):
666 d = {'type': 'calculator',
667 'name': 'socket-driver'}
668 # if self.calc is not None:
669 # d['calc'] = self.calc.todict()
670 return d
672 def launch_server(self):
673 return self.closelater(SocketServer(
674 # launch_client=launch_client,
675 port=self._port,
676 unixsocket=self._unixsocket,
677 timeout=self.timeout, log=self.log,
678 ))
680 def calculate(self, atoms=None, properties=['energy'],
681 system_changes=all_changes):
682 bad = [change for change in system_changes
683 if change not in self.supported_changes]
685 # First time calculate() is called, system_changes will be
686 # all_changes. After that, only positions and cell may change.
687 if self.atoms is not None and any(bad):
688 raise PropertyNotImplementedError(
689 'Cannot change {} through IPI protocol. '
690 'Please create new socket calculator.'
691 .format(bad if len(bad) > 1 else bad[0]))
693 self.atoms = atoms.copy()
695 if self.server is None:
696 self.server = self.launch_server()
697 proc = self.launch_client(atoms, properties,
698 port=self._port,
699 unixsocket=self._unixsocket)
700 self.server.proc = proc # XXX nasty hack
702 results = self.server.calculate(atoms)
703 results['free_energy'] = results['energy']
704 virial = results.pop('virial')
705 if self.atoms.cell.rank == 3 and any(self.atoms.pbc):
706 vol = atoms.get_volume()
707 results['stress'] = -full_3x3_to_voigt_6_stress(virial) / vol
708 self.results.update(results)
710 def close(self):
711 self.server = None
712 super().close()
715class PySocketIOClient:
716 def __init__(self, calculator_factory):
717 self._calculator_factory = calculator_factory
719 def __call__(self, atoms, properties=None, port=None, unixsocket=None):
720 import pickle
721 import sys
723 # We pickle everything first, so we won't need to bother with the
724 # process as long as it succeeds.
725 transferbytes = pickle.dumps([
726 dict(unixsocket=unixsocket, port=port),
727 atoms.copy(),
728 self._calculator_factory,
729 ])
731 proc = Popen([sys.executable, '-m', 'ase.calculators.socketio'],
732 stdin=PIPE)
734 proc.stdin.write(transferbytes)
735 proc.stdin.close()
736 return proc
738 @staticmethod
739 def main():
740 import pickle
741 import sys
743 socketinfo, atoms, get_calculator = pickle.load(sys.stdin.buffer)
744 atoms.calc = get_calculator()
745 client = SocketClient(host='localhost',
746 unixsocket=socketinfo.get('unixsocket'),
747 port=socketinfo.get('port'))
748 # XXX In principle we could avoid calculating stress until
749 # someone requests the stress, could we not?
750 # Which would make use_stress boolean unnecessary.
751 client.run(atoms, use_stress=True)
754if __name__ == '__main__':
755 PySocketIOClient.main()