Source code for ase.calculators.socketio

import os
import socket
from contextlib import contextmanager
from subprocess import PIPE, Popen

import numpy as np

import ase.units as units
from ase.calculators.calculator import (
    Calculator,
    OldShellProfile,
    PropertyNotImplementedError,
    StandardProfile,
    all_changes,
)
from ase.calculators.genericfileio import GenericFileIOCalculator
from ase.parallel import world
from ase.stress import full_3x3_to_voigt_6_stress
from ase.utils import IOContext


def actualunixsocketname(name):
    return f'/tmp/ipi_{name}'


class SocketClosed(OSError):
    pass


class IPIProtocol:
    """Communication using IPI protocol."""

    def __init__(self, socket, txt=None):
        self.socket = socket

        if txt is None:
            def log(*args):
                pass
        else:
            def log(*args):
                print('Driver:', *args, file=txt)
                txt.flush()
        self.log = log

    def sendmsg(self, msg):
        self.log('  sendmsg', repr(msg))
        # assert msg in self.statements, msg
        msg = msg.encode('ascii').ljust(12)
        self.socket.sendall(msg)

    def _recvall(self, nbytes):
        """Repeatedly read chunks until we have nbytes.

        Normally we get all bytes in one read, but that is not guaranteed."""
        remaining = nbytes
        chunks = []
        while remaining > 0:
            chunk = self.socket.recv(remaining)
            if len(chunk) == 0:
                # (If socket is still open, recv returns at least one byte)
                raise SocketClosed
            chunks.append(chunk)
            remaining -= len(chunk)
        msg = b''.join(chunks)
        assert len(msg) == nbytes and remaining == 0
        return msg

    def recvmsg(self):
        msg = self._recvall(12)
        if not msg:
            raise SocketClosed

        assert len(msg) == 12, msg
        msg = msg.rstrip().decode('ascii')
        # assert msg in self.responses, msg
        self.log('  recvmsg', repr(msg))
        return msg

    def send(self, a, dtype):
        buf = np.asarray(a, dtype).tobytes()
        # self.log('  send {}'.format(np.array(a).ravel().tolist()))
        self.log(f'  send {len(buf)} bytes of {dtype}')
        self.socket.sendall(buf)

    def recv(self, shape, dtype):
        a = np.empty(shape, dtype)
        nbytes = np.dtype(dtype).itemsize * np.prod(shape)
        buf = self._recvall(nbytes)
        assert len(buf) == nbytes, (len(buf), nbytes)
        self.log(f'  recv {len(buf)} bytes of {dtype}')
        # print(np.frombuffer(buf, dtype=dtype))
        a.flat[:] = np.frombuffer(buf, dtype=dtype)
        # self.log('  recv {}'.format(a.ravel().tolist()))
        assert np.isfinite(a).all()
        return a

    def sendposdata(self, cell, icell, positions):
        assert cell.size == 9
        assert icell.size == 9
        assert positions.size % 3 == 0

        self.log(' sendposdata')
        self.sendmsg('POSDATA')
        self.send(cell.T / units.Bohr, np.float64)
        self.send(icell.T * units.Bohr, np.float64)
        self.send(len(positions), np.int32)
        self.send(positions / units.Bohr, np.float64)

    def recvposdata(self):
        cell = self.recv((3, 3), np.float64).T.copy()
        icell = self.recv((3, 3), np.float64).T.copy()
        natoms = self.recv(1, np.int32)[0]
        positions = self.recv((natoms, 3), np.float64)
        return cell * units.Bohr, icell / units.Bohr, positions * units.Bohr

    def sendrecv_force(self):
        self.log(' sendrecv_force')
        self.sendmsg('GETFORCE')
        msg = self.recvmsg()
        assert msg == 'FORCEREADY', msg
        e = self.recv(1, np.float64)[0]
        natoms = self.recv(1, np.int32)[0]
        assert natoms >= 0
        forces = self.recv((int(natoms), 3), np.float64)
        virial = self.recv((3, 3), np.float64).T.copy()
        nmorebytes = self.recv(1, np.int32)[0]
        morebytes = self.recv(nmorebytes, np.byte)
        return (e * units.Ha, (units.Ha / units.Bohr) * forces,
                units.Ha * virial, morebytes)

    def sendforce(self, energy, forces, virial,
                  morebytes=np.zeros(1, dtype=np.byte)):
        assert np.array([energy]).size == 1
        assert forces.shape[1] == 3
        assert virial.shape == (3, 3)

        self.log(' sendforce')
        self.sendmsg('FORCEREADY')  # mind the units
        self.send(np.array([energy / units.Ha]), np.float64)
        natoms = len(forces)
        self.send(np.array([natoms]), np.int32)
        self.send(units.Bohr / units.Ha * forces, np.float64)
        self.send(1.0 / units.Ha * virial.T, np.float64)
        # We prefer to always send at least one byte due to trouble with
        # empty messages.  Reading a closed socket yields 0 bytes
        # and thus can be confused with a 0-length bytestring.
        self.send(np.array([len(morebytes)]), np.int32)
        self.send(morebytes, np.byte)

    def status(self):
        self.log(' status')
        self.sendmsg('STATUS')
        msg = self.recvmsg()
        return msg

    def end(self):
        self.log(' end')
        self.sendmsg('EXIT')

    def recvinit(self):
        self.log(' recvinit')
        bead_index = self.recv(1, np.int32)
        nbytes = self.recv(1, np.int32)
        initbytes = self.recv(nbytes, np.byte)
        return bead_index, initbytes

    def sendinit(self):
        # XXX Not sure what this function is supposed to send.
        # It 'works' with QE, but for now we try not to call it.
        self.log(' sendinit')
        self.sendmsg('INIT')
        self.send(0, np.int32)  # 'bead index' always zero for now
        # We send one byte, which is zero, since things may not work
        # with 0 bytes.  Apparently implementations ignore the
        # initialization string anyway.
        self.send(1, np.int32)
        self.send(np.zeros(1), np.byte)  # initialization string

    def calculate(self, positions, cell):
        self.log('calculate')
        msg = self.status()
        # We don't know how NEEDINIT is supposed to work, but some codes
        # seem to be okay if we skip it and send the positions instead.
        if msg == 'NEEDINIT':
            self.sendinit()
            msg = self.status()
        assert msg == 'READY', msg
        icell = np.linalg.pinv(cell).transpose()
        self.sendposdata(cell, icell, positions)
        msg = self.status()
        assert msg == 'HAVEDATA', msg
        e, forces, virial, morebytes = self.sendrecv_force()
        r = dict(energy=e,
                 forces=forces,
                 virial=virial,
                 morebytes=morebytes)
        return r


@contextmanager
def bind_unixsocket(socketfile):
    assert socketfile.startswith('/tmp/ipi_'), socketfile
    serversocket = socket.socket(socket.AF_UNIX)
    try:
        serversocket.bind(socketfile)
    except OSError as err:
        raise OSError(f'{err}: {socketfile!r}')

    try:
        with serversocket:
            yield serversocket
    finally:
        os.unlink(socketfile)


@contextmanager
def bind_inetsocket(port):
    serversocket = socket.socket(socket.AF_INET)
    serversocket.setsockopt(socket.SOL_SOCKET,
                            socket.SO_REUSEADDR, 1)
    serversocket.bind(('', port))
    with serversocket:
        yield serversocket


class FileIOSocketClientLauncher:
    def __init__(self, calc):
        self.calc = calc

    def __call__(self, atoms, properties=None, port=None, unixsocket=None):
        assert self.calc is not None
        cwd = self.calc.directory

        profile = getattr(self.calc, 'profile', None)
        if isinstance(self.calc, GenericFileIOCalculator):
            # New GenericFileIOCalculator:
            template = getattr(self.calc, 'template')

            self.calc.write_inputfiles(atoms, properties)
            if unixsocket is not None:
                argv = template.socketio_argv(
                    profile, unixsocket=unixsocket, port=None
                )
            else:
                argv = template.socketio_argv(
                    profile, unixsocket=None, port=port
                )
            return Popen(argv, cwd=cwd, env=os.environ)
        else:
            # Old FileIOCalculator:
            self.calc.write_input(atoms, properties=properties,
                                  system_changes=all_changes)

            if isinstance(profile, StandardProfile):
                return profile.execute_nonblocking(self.calc)

            if profile is None:
                cmd = self.calc.command.replace('PREFIX', self.calc.prefix)
                cmd = cmd.format(port=port, unixsocket=unixsocket)
            elif isinstance(profile, OldShellProfile):
                cmd = profile.command.replace("PREFIX", self.calc.prefix)
            else:
                raise TypeError(
                    f"Profile type {type(profile)} not supported for socketio")

            return Popen(cmd, shell=True, cwd=cwd)


[docs] class SocketServer(IOContext): default_port = 31415 def __init__(self, # launch_client=None, port=None, unixsocket=None, timeout=None, log=None): """Create server and listen for connections. Parameters: client_command: Shell command to launch client process, or None The process will be launched immediately, if given. Else the user is expected to launch a client whose connection the server will then accept at any time. One calculate() is called, the server will block to wait for the client. port: integer or None Port on which to listen for INET connections. Defaults to 31415 if neither this nor unixsocket is specified. unixsocket: string or None Filename for unix socket. timeout: float or None timeout in seconds, or unlimited by default. This parameter is passed to the Python socket object; see documentation therof log: file object or None useful debug messages are written to this.""" if unixsocket is None and port is None: port = self.default_port elif unixsocket is not None and port is not None: raise ValueError('Specify only one of unixsocket and port') self.port = port self.unixsocket = unixsocket self.timeout = timeout self._closed = False if unixsocket is not None: actualsocket = actualunixsocketname(unixsocket) conn_name = f'UNIX-socket {actualsocket}' socket_context = bind_unixsocket(actualsocket) else: conn_name = f'INET port {port}' socket_context = bind_inetsocket(port) self.serversocket = self.closelater(socket_context) if log: print(f'Accepting clients on {conn_name}', file=log) self.serversocket.settimeout(timeout) self.serversocket.listen(1) self.log = log self.proc = None self.protocol = None self.clientsocket = None self.address = None # if launch_client is not None: # self.proc = launch_client(port=port, unixsocket=unixsocket) def _accept(self): """Wait for client and establish connection.""" # It should perhaps be possible for process to be launched by user log = self.log if log: print('Awaiting client', file=self.log) # If we launched the subprocess, the process may crash. # We want to detect this, using loop with timeouts, and # raise an error rather than blocking forever. if self.proc is not None: self.serversocket.settimeout(1.0) while True: try: self.clientsocket, self.address = self.serversocket.accept() self.closelater(self.clientsocket) except socket.timeout: if self.proc is not None: status = self.proc.poll() if status is not None: raise OSError('Subprocess terminated unexpectedly' ' with status {}'.format(status)) else: break self.serversocket.settimeout(self.timeout) self.clientsocket.settimeout(self.timeout) if log: # For unix sockets, address is b''. source = ('client' if self.address == b'' else self.address) print(f'Accepted connection from {source}', file=log) self.protocol = IPIProtocol(self.clientsocket, txt=log) def close(self): if self._closed: return super().close() if self.log: print('Close socket server', file=self.log) self._closed = True # Proper way to close sockets? # And indeed i-pi connections... # if self.protocol is not None: # self.protocol.end() # Send end-of-communication string self.protocol = None if self.proc is not None: exitcode = self.proc.wait() if exitcode != 0: import warnings # Quantum Espresso seems to always exit with status 128, # even if successful. # Should investigate at some point warnings.warn('Subprocess exited with status {}' .format(exitcode)) # self.log('IPI server closed') def calculate(self, atoms): """Send geometry to client and return calculated things as dict. This will block until client has established connection, then wait for the client to finish the calculation.""" assert not self._closed # If we have not established connection yet, we must block # until the client catches up: if self.protocol is None: self._accept() return self.protocol.calculate(atoms.positions, atoms.cell)
[docs] class SocketClient: def __init__(self, host='localhost', port=None, unixsocket=None, timeout=None, log=None, comm=world): """Create client and connect to server. Parameters: host: string Hostname of server. Defaults to localhost port: integer or None Port to which to connect. By default 31415. unixsocket: string or None If specified, use corresponding UNIX socket. See documentation of unixsocket for SocketIOCalculator. timeout: float or None See documentation of timeout for SocketIOCalculator. log: file object or None Log events to this file comm: communicator or None MPI communicator object. Defaults to ase.parallel.world. When ASE runs in parallel, only the process with world.rank == 0 will communicate over the socket. The received information will then be broadcast on the communicator. The SocketClient must be created on all ranks of world, and will see the same Atoms objects.""" # Only rank0 actually does the socket work. # The other ranks only need to follow. # # Note: We actually refrain from assigning all the # socket-related things except on master self.comm = comm if self.comm.rank == 0: if unixsocket is not None: sock = socket.socket(socket.AF_UNIX) actualsocket = actualunixsocketname(unixsocket) sock.connect(actualsocket) else: if port is None: port = SocketServer.default_port sock = socket.socket(socket.AF_INET) sock.connect((host, port)) sock.settimeout(timeout) self.host = host self.port = port self.unixsocket = unixsocket self.protocol = IPIProtocol(sock, txt=log) self.log = self.protocol.log self.closed = False self.bead_index = 0 self.bead_initbytes = b'' self.state = 'READY' def close(self): if not self.closed: self.log('Close SocketClient') self.closed = True self.protocol.socket.close() def calculate(self, atoms, use_stress): # We should also broadcast the bead index, once we support doing # multiple beads. self.comm.broadcast(atoms.positions, 0) self.comm.broadcast(np.ascontiguousarray(atoms.cell), 0) energy = atoms.get_potential_energy() forces = atoms.get_forces() if use_stress: stress = atoms.get_stress(voigt=False) virial = -atoms.get_volume() * stress else: virial = np.zeros((3, 3)) return energy, forces, virial def irun(self, atoms, use_stress=None): if use_stress is None: use_stress = any(atoms.pbc) my_irun = self.irun_rank0 if self.comm.rank == 0 else self.irun_rankN return my_irun(atoms, use_stress) def irun_rankN(self, atoms, use_stress=True): stop_criterion = np.zeros(1, bool) while True: self.comm.broadcast(stop_criterion, 0) if stop_criterion[0]: return self.calculate(atoms, use_stress) yield def irun_rank0(self, atoms, use_stress=True): # For every step we either calculate or quit. We need to # tell other MPI processes (if this is MPI-parallel) whether they # should calculate or quit. try: while True: try: msg = self.protocol.recvmsg() except SocketClosed: # Server closed the connection, but we want to # exit gracefully anyway msg = 'EXIT' if msg == 'EXIT': # Send stop signal to clients: self.comm.broadcast(np.ones(1, bool), 0) # (When otherwise exiting, things crashed and we should # let MPI_ABORT take care of the mess instead of trying # to synchronize the exit) return elif msg == 'STATUS': self.protocol.sendmsg(self.state) elif msg == 'POSDATA': assert self.state == 'READY' cell, icell, positions = self.protocol.recvposdata() atoms.cell[:] = cell atoms.positions[:] = positions # User may wish to do something with the atoms object now. # Should we provide option to yield here? # # (In that case we should MPI-synchronize *before* # whereas now we do it after.) # Send signal for other ranks to proceed with calculation: self.comm.broadcast(np.zeros(1, bool), 0) energy, forces, virial = self.calculate(atoms, use_stress) self.state = 'HAVEDATA' yield elif msg == 'GETFORCE': assert self.state == 'HAVEDATA', self.state self.protocol.sendforce(energy, forces, virial) self.state = 'NEEDINIT' elif msg == 'INIT': assert self.state == 'NEEDINIT' bead_index, initbytes = self.protocol.recvinit() self.bead_index = bead_index self.bead_initbytes = initbytes self.state = 'READY' else: raise KeyError('Bad message', msg) finally: self.close() def run(self, atoms, use_stress=False): for _ in self.irun(atoms, use_stress=use_stress): pass
[docs] class SocketIOCalculator(Calculator, IOContext): implemented_properties = ['energy', 'free_energy', 'forces', 'stress'] supported_changes = {'positions', 'cell'} def __init__(self, calc=None, port=None, unixsocket=None, timeout=None, log=None, *, launch_client=None, comm=world): """Initialize socket I/O calculator. This calculator launches a server which passes atomic coordinates and unit cells to an external code via a socket, and receives energy, forces, and stress in return. ASE integrates this with the Quantum Espresso, FHI-aims and Siesta calculators. This works with any external code that supports running as a client over the i-PI protocol. Parameters: calc: calculator or None If calc is not None, a client process will be launched using calc.command, and the input file will be generated using ``calc.write_input()``. Otherwise only the server will run, and it is up to the user to launch a compliant client process. port: integer port number for socket. Should normally be between 1025 and 65535. Typical ports for are 31415 (default) or 3141. unixsocket: str or None if not None, ignore host and port, creating instead a unix socket using this name prefixed with ``/tmp/ipi_``. The socket is deleted when the calculator is closed. timeout: float >= 0 or None timeout for connection, by default infinite. See documentation of Python sockets. For longer jobs it is recommended to set a timeout in case of undetected client-side failure. log: file object or None (default) logfile for communication over socket. For debugging or the curious. In order to correctly close the sockets, it is recommended to use this class within a with-block: >>> from ase.calculators.socketio import SocketIOCalculator >>> with SocketIOCalculator(...) as calc: # doctest:+SKIP ... atoms.calc = calc ... atoms.get_forces() ... atoms.rattle() ... atoms.get_forces() It is also possible to call calc.close() after use. This is best done in a finally-block.""" Calculator.__init__(self) if calc is not None: if launch_client is not None: raise ValueError('Cannot pass both calc and launch_client') launch_client = FileIOSocketClientLauncher(calc) self.launch_client = launch_client self.timeout = timeout self.server = None self.log = self.openfile(file=log, comm=comm) # We only hold these so we can pass them on to the server. # They may both be None as stored here. self._port = port self._unixsocket = unixsocket # If there is a calculator, we will launch in calculate() because # we are responsible for executing the external process, too, and # should do so before blocking. Without a calculator we want to # block immediately: if self.launch_client is None: self.server = self.launch_server() def todict(self): d = {'type': 'calculator', 'name': 'socket-driver'} # if self.calc is not None: # d['calc'] = self.calc.todict() return d def launch_server(self): return self.closelater(SocketServer( # launch_client=launch_client, port=self._port, unixsocket=self._unixsocket, timeout=self.timeout, log=self.log, )) def calculate(self, atoms=None, properties=['energy'], system_changes=all_changes): bad = [change for change in system_changes if change not in self.supported_changes] # First time calculate() is called, system_changes will be # all_changes. After that, only positions and cell may change. 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.atoms = atoms.copy() if self.server is None: self.server = self.launch_server() proc = self.launch_client(atoms, properties, port=self._port, unixsocket=self._unixsocket) self.server.proc = proc # XXX nasty hack results = self.server.calculate(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 self.results.update(results) def close(self): self.server = None super().close()
class PySocketIOClient: def __init__(self, calculator_factory): self._calculator_factory = calculator_factory def __call__(self, atoms, properties=None, port=None, unixsocket=None): import pickle import sys # We pickle everything first, so we won't need to bother with the # process as long as it succeeds. transferbytes = pickle.dumps([ dict(unixsocket=unixsocket, port=port), atoms.copy(), self._calculator_factory, ]) proc = Popen([sys.executable, '-m', 'ase.calculators.socketio'], stdin=PIPE) proc.stdin.write(transferbytes) proc.stdin.close() return proc @staticmethod def main(): import pickle import sys socketinfo, atoms, get_calculator = pickle.load(sys.stdin.buffer) atoms.calc = get_calculator() client = SocketClient(host='localhost', unixsocket=socketinfo.get('unixsocket'), port=socketinfo.get('port')) # XXX In principle we could avoid calculating stress until # someone requests the stress, could we not? # Which would make use_stress boolean unnecessary. client.run(atoms, use_stress=True) if __name__ == '__main__': PySocketIOClient.main()