Coverage for /builds/ericyuan00000/ase/ase/calculators/socketio.py: 90.91%
396 statements
« prev ^ index » next coverage.py v7.5.3, created at 2025-06-18 01:20 +0000
1# fmt: off
3import os
4import socket
5from contextlib import 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 )
250 return Popen(argv, cwd=cwd, env=os.environ)
251 else:
252 # Old FileIOCalculator:
253 self.calc.write_input(atoms, properties=properties,
254 system_changes=all_changes)
256 if isinstance(profile, StandardProfile):
257 return profile.execute_nonblocking(self.calc)
259 if profile is None:
260 cmd = self.calc.command.replace('PREFIX', self.calc.prefix)
261 cmd = cmd.format(port=port, unixsocket=unixsocket)
262 elif isinstance(profile, OldShellProfile):
263 cmd = profile.command.replace("PREFIX", self.calc.prefix)
264 else:
265 raise TypeError(
266 f"Profile type {type(profile)} not supported for socketio")
268 return Popen(cmd, shell=True, cwd=cwd)
271class SocketServer(IOContext):
272 default_port = 31415
274 def __init__(self, # launch_client=None,
275 port=None, unixsocket=None, timeout=None,
276 log=None):
277 """Create server and listen for connections.
279 Parameters:
281 client_command: Shell command to launch client process, or None
282 The process will be launched immediately, if given.
283 Else the user is expected to launch a client whose connection
284 the server will then accept at any time.
285 One calculate() is called, the server will block to wait
286 for the client.
287 port: integer or None
288 Port on which to listen for INET connections. Defaults
289 to 31415 if neither this nor unixsocket is specified.
290 unixsocket: string or None
291 Filename for unix socket.
292 timeout: float or None
293 timeout in seconds, or unlimited by default.
294 This parameter is passed to the Python socket object; see
295 documentation therof
296 log: file object or None
297 useful debug messages are written to this."""
299 if unixsocket is None and port is None:
300 port = self.default_port
301 elif unixsocket is not None and port is not None:
302 raise ValueError('Specify only one of unixsocket and port')
304 self.port = port
305 self.unixsocket = unixsocket
306 self.timeout = timeout
307 self._closed = False
309 if unixsocket is not None:
310 actualsocket = actualunixsocketname(unixsocket)
311 conn_name = f'UNIX-socket {actualsocket}'
312 socket_context = bind_unixsocket(actualsocket)
313 else:
314 conn_name = f'INET port {port}'
315 socket_context = bind_inetsocket(port)
317 self.serversocket = self.closelater(socket_context)
319 if log:
320 print(f'Accepting clients on {conn_name}', file=log)
322 self.serversocket.settimeout(timeout)
324 self.serversocket.listen(1)
326 self.log = log
328 self.proc = None
330 self.protocol = None
331 self.clientsocket = None
332 self.address = None
334 # if launch_client is not None:
335 # self.proc = launch_client(port=port, unixsocket=unixsocket)
337 def _accept(self):
338 """Wait for client and establish connection."""
339 # It should perhaps be possible for process to be launched by user
340 log = self.log
341 if log:
342 print('Awaiting client', file=self.log)
344 # If we launched the subprocess, the process may crash.
345 # We want to detect this, using loop with timeouts, and
346 # raise an error rather than blocking forever.
347 if self.proc is not None:
348 self.serversocket.settimeout(1.0)
350 while True:
351 try:
352 self.clientsocket, self.address = self.serversocket.accept()
353 self.closelater(self.clientsocket)
354 except socket.timeout:
355 if self.proc is not None:
356 status = self.proc.poll()
357 if status is not None:
358 raise OSError('Subprocess terminated unexpectedly'
359 ' with status {}'.format(status))
360 else:
361 break
363 self.serversocket.settimeout(self.timeout)
364 self.clientsocket.settimeout(self.timeout)
366 if log:
367 # For unix sockets, address is b''.
368 source = ('client' if self.address == b'' else self.address)
369 print(f'Accepted connection from {source}', file=log)
371 self.protocol = IPIProtocol(self.clientsocket, txt=log)
373 def close(self):
374 if self._closed:
375 return
377 super().close()
379 if self.log:
380 print('Close socket server', file=self.log)
381 self._closed = True
383 # Proper way to close sockets?
384 # And indeed i-pi connections...
385 # if self.protocol is not None:
386 # self.protocol.end() # Send end-of-communication string
387 self.protocol = None
388 if self.proc is not None:
389 exitcode = self.proc.wait()
390 if exitcode != 0:
391 import warnings
393 # Quantum Espresso seems to always exit with status 128,
394 # even if successful.
395 # Should investigate at some point
396 warnings.warn('Subprocess exited with status {}'
397 .format(exitcode))
398 # self.log('IPI server closed')
400 def calculate(self, atoms):
401 """Send geometry to client and return calculated things as dict.
403 This will block until client has established connection, then
404 wait for the client to finish the calculation."""
405 assert not self._closed
407 # If we have not established connection yet, we must block
408 # until the client catches up:
409 if self.protocol is None:
410 self._accept()
411 return self.protocol.calculate(atoms.positions, atoms.cell)
414class SocketClient:
415 def __init__(self, host='localhost', port=None,
416 unixsocket=None, timeout=None, log=None, comm=world):
417 """Create client and connect to server.
419 Parameters:
421 host: string
422 Hostname of server. Defaults to localhost
423 port: integer or None
424 Port to which to connect. By default 31415.
425 unixsocket: string or None
426 If specified, use corresponding UNIX socket.
427 See documentation of unixsocket for SocketIOCalculator.
428 timeout: float or None
429 See documentation of timeout for SocketIOCalculator.
430 log: file object or None
431 Log events to this file
432 comm: communicator or None
433 MPI communicator object. Defaults to ase.parallel.world.
434 When ASE runs in parallel, only the process with world.rank == 0
435 will communicate over the socket. The received information
436 will then be broadcast on the communicator. The SocketClient
437 must be created on all ranks of world, and will see the same
438 Atoms objects."""
439 # Only rank0 actually does the socket work.
440 # The other ranks only need to follow.
441 #
442 # Note: We actually refrain from assigning all the
443 # socket-related things except on master
444 self.comm = comm
446 if self.comm.rank == 0:
447 if unixsocket is not None:
448 sock = socket.socket(socket.AF_UNIX)
449 actualsocket = actualunixsocketname(unixsocket)
450 sock.connect(actualsocket)
451 else:
452 if port is None:
453 port = SocketServer.default_port
454 sock = socket.socket(socket.AF_INET)
455 sock.connect((host, port))
456 sock.settimeout(timeout)
457 self.host = host
458 self.port = port
459 self.unixsocket = unixsocket
461 self.protocol = IPIProtocol(sock, txt=log)
462 self.log = self.protocol.log
463 self.closed = False
465 self.bead_index = 0
466 self.bead_initbytes = b''
467 self.state = 'READY'
469 def close(self):
470 if not self.closed:
471 self.log('Close SocketClient')
472 self.closed = True
473 self.protocol.socket.close()
475 def calculate(self, atoms, use_stress):
476 # We should also broadcast the bead index, once we support doing
477 # multiple beads.
478 self.comm.broadcast(atoms.positions, 0)
479 self.comm.broadcast(np.ascontiguousarray(atoms.cell), 0)
481 energy = atoms.get_potential_energy()
482 forces = atoms.get_forces()
483 if use_stress:
484 stress = atoms.get_stress(voigt=False)
485 virial = -atoms.get_volume() * stress
486 else:
487 virial = np.zeros((3, 3))
488 return energy, forces, virial
490 def irun(self, atoms, use_stress=None):
491 if use_stress is None:
492 use_stress = any(atoms.pbc)
494 my_irun = self.irun_rank0 if self.comm.rank == 0 else self.irun_rankN
495 return my_irun(atoms, use_stress)
497 def irun_rankN(self, atoms, use_stress=True):
498 stop_criterion = np.zeros(1, bool)
499 while True:
500 self.comm.broadcast(stop_criterion, 0)
501 if stop_criterion[0]:
502 return
504 self.calculate(atoms, use_stress)
505 yield
507 def irun_rank0(self, atoms, use_stress=True):
508 # For every step we either calculate or quit. We need to
509 # tell other MPI processes (if this is MPI-parallel) whether they
510 # should calculate or quit.
511 try:
512 while True:
513 try:
514 msg = self.protocol.recvmsg()
515 except SocketClosed:
516 # Server closed the connection, but we want to
517 # exit gracefully anyway
518 msg = 'EXIT'
520 if msg == 'EXIT':
521 # Send stop signal to clients:
522 self.comm.broadcast(np.ones(1, bool), 0)
523 # (When otherwise exiting, things crashed and we should
524 # let MPI_ABORT take care of the mess instead of trying
525 # to synchronize the exit)
526 return
527 elif msg == 'STATUS':
528 self.protocol.sendmsg(self.state)
529 elif msg == 'POSDATA':
530 assert self.state == 'READY'
531 cell, _icell, positions = self.protocol.recvposdata()
532 atoms.cell[:] = cell
533 atoms.positions[:] = positions
535 # User may wish to do something with the atoms object now.
536 # Should we provide option to yield here?
537 #
538 # (In that case we should MPI-synchronize *before*
539 # whereas now we do it after.)
541 # Send signal for other ranks to proceed with calculation:
542 self.comm.broadcast(np.zeros(1, bool), 0)
543 energy, forces, virial = self.calculate(atoms, use_stress)
545 self.state = 'HAVEDATA'
546 yield
547 elif msg == 'GETFORCE':
548 assert self.state == 'HAVEDATA', self.state
549 self.protocol.sendforce(energy, forces, virial)
550 self.state = 'NEEDINIT'
551 elif msg == 'INIT':
552 assert self.state == 'NEEDINIT'
553 bead_index, initbytes = self.protocol.recvinit()
554 self.bead_index = bead_index
555 self.bead_initbytes = initbytes
556 self.state = 'READY'
557 else:
558 raise KeyError('Bad message', msg)
559 finally:
560 self.close()
562 def run(self, atoms, use_stress=False):
563 for _ in self.irun(atoms, use_stress=use_stress):
564 pass
567class SocketIOCalculator(Calculator, IOContext):
568 implemented_properties = ['energy', 'free_energy', 'forces', 'stress']
569 supported_changes = {'positions', 'cell'}
571 def __init__(self, calc=None, port=None,
572 unixsocket=None, timeout=None, log=None, *,
573 launch_client=None, comm=world):
574 """Initialize socket I/O calculator.
576 This calculator launches a server which passes atomic
577 coordinates and unit cells to an external code via a socket,
578 and receives energy, forces, and stress in return.
580 ASE integrates this with the Quantum Espresso, FHI-aims and
581 Siesta calculators. This works with any external code that
582 supports running as a client over the i-PI protocol.
584 Parameters:
586 calc: calculator or None
588 If calc is not None, a client process will be launched
589 using calc.command, and the input file will be generated
590 using ``calc.write_input()``. Otherwise only the server will
591 run, and it is up to the user to launch a compliant client
592 process.
594 port: integer
596 port number for socket. Should normally be between 1025
597 and 65535. Typical ports for are 31415 (default) or 3141.
599 unixsocket: str or None
601 if not None, ignore host and port, creating instead a
602 unix socket using this name prefixed with ``/tmp/ipi_``.
603 The socket is deleted when the calculator is closed.
605 timeout: float >= 0 or None
607 timeout for connection, by default infinite. See
608 documentation of Python sockets. For longer jobs it is
609 recommended to set a timeout in case of undetected
610 client-side failure.
612 log: file object or None (default)
614 logfile for communication over socket. For debugging or
615 the curious.
617 In order to correctly close the sockets, it is
618 recommended to use this class within a with-block:
620 >>> from ase.calculators.socketio import SocketIOCalculator
622 >>> with SocketIOCalculator(...) as calc: # doctest:+SKIP
623 ... atoms.calc = calc
624 ... atoms.get_forces()
625 ... atoms.rattle()
626 ... atoms.get_forces()
628 It is also possible to call calc.close() after
629 use. This is best done in a finally-block."""
631 Calculator.__init__(self)
633 if calc is not None:
634 if launch_client is not None:
635 raise ValueError('Cannot pass both calc and launch_client')
636 launch_client = FileIOSocketClientLauncher(calc)
637 self.launch_client = launch_client
638 self.timeout = timeout
639 self.server = None
641 self.log = self.openfile(file=log, comm=comm)
643 # We only hold these so we can pass them on to the server.
644 # They may both be None as stored here.
645 self._port = port
646 self._unixsocket = unixsocket
648 # If there is a calculator, we will launch in calculate() because
649 # we are responsible for executing the external process, too, and
650 # should do so before blocking. Without a calculator we want to
651 # block immediately:
652 if self.launch_client is None:
653 self.server = self.launch_server()
655 def todict(self):
656 d = {'type': 'calculator',
657 'name': 'socket-driver'}
658 # if self.calc is not None:
659 # d['calc'] = self.calc.todict()
660 return d
662 def launch_server(self):
663 return self.closelater(SocketServer(
664 # launch_client=launch_client,
665 port=self._port,
666 unixsocket=self._unixsocket,
667 timeout=self.timeout, log=self.log,
668 ))
670 def calculate(self, atoms=None, properties=['energy'],
671 system_changes=all_changes):
672 bad = [change for change in system_changes
673 if change not in self.supported_changes]
675 # First time calculate() is called, system_changes will be
676 # all_changes. After that, only positions and cell may change.
677 if self.atoms is not None and any(bad):
678 raise PropertyNotImplementedError(
679 'Cannot change {} through IPI protocol. '
680 'Please create new socket calculator.'
681 .format(bad if len(bad) > 1 else bad[0]))
683 self.atoms = atoms.copy()
685 if self.server is None:
686 self.server = self.launch_server()
687 proc = self.launch_client(atoms, properties,
688 port=self._port,
689 unixsocket=self._unixsocket)
690 self.server.proc = proc # XXX nasty hack
692 results = self.server.calculate(atoms)
693 results['free_energy'] = results['energy']
694 virial = results.pop('virial')
695 if self.atoms.cell.rank == 3 and any(self.atoms.pbc):
696 vol = atoms.get_volume()
697 results['stress'] = -full_3x3_to_voigt_6_stress(virial) / vol
698 self.results.update(results)
700 def close(self):
701 self.server = None
702 super().close()
705class PySocketIOClient:
706 def __init__(self, calculator_factory):
707 self._calculator_factory = calculator_factory
709 def __call__(self, atoms, properties=None, port=None, unixsocket=None):
710 import pickle
711 import sys
713 # We pickle everything first, so we won't need to bother with the
714 # process as long as it succeeds.
715 transferbytes = pickle.dumps([
716 dict(unixsocket=unixsocket, port=port),
717 atoms.copy(),
718 self._calculator_factory,
719 ])
721 proc = Popen([sys.executable, '-m', 'ase.calculators.socketio'],
722 stdin=PIPE)
724 proc.stdin.write(transferbytes)
725 proc.stdin.close()
726 return proc
728 @staticmethod
729 def main():
730 import pickle
731 import sys
733 socketinfo, atoms, get_calculator = pickle.load(sys.stdin.buffer)
734 atoms.calc = get_calculator()
735 client = SocketClient(host='localhost',
736 unixsocket=socketinfo.get('unixsocket'),
737 port=socketinfo.get('port'))
738 # XXX In principle we could avoid calculating stress until
739 # someone requests the stress, could we not?
740 # Which would make use_stress boolean unnecessary.
741 client.run(atoms, use_stress=True)
744if __name__ == '__main__':
745 PySocketIOClient.main()