faiss/contrib/rpc.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

"""
Simplistic RPC implementation.
Exposes all functions of a Server object.

This code is for demonstration purposes only, and does not include certain
security protections. It is not meant to be run on an untrusted network or
in a production environment.
"""

import importlib
import os
import pickle
import sys
import _thread
import traceback
import socket
import logging

LOG = logging.getLogger(__name__)

# default
PORT = 12032

safe_modules = {
    "numpy",
    "numpy.core.multiarray",
}


class RestrictedUnpickler(pickle.Unpickler):

    def find_class(self, module, name):
        # Only allow safe modules.
        if module in safe_modules:
            return getattr(importlib.import_module(module), name)
        # Forbid everything else.
        raise pickle.UnpicklingError(
            "global '%s.%s' is forbidden" % (module, name)
        )


class FileSock:
    "wraps a socket so that it is usable by pickle/cPickle"

    def __init__(self, sock):
        self.sock = sock
        self.nr = 0

    def write(self, buf):
        # print("sending %d bytes"%len(buf))
        # self.sock.sendall(buf)
        # print("...done")
        bs = 512 * 1024
        ns = 0
        while ns < len(buf):
            sent = self.sock.send(buf[ns : ns + bs])
            ns += sent

    def read(self, bs=512 * 1024):
        # if self.nr==10000: pdb.set_trace()
        self.nr += 1
        # print("read bs=%d"%bs)
        b = []
        nb = 0
        while len(b) < bs:
            # print('   loop')
            rb = self.sock.recv(bs - nb)
            if not rb:
                break
            b.append(rb)
            nb += len(rb)
        return b"".join(b)

    def readline(self):
        # print("readline!")
        """may be optimized..."""
        s = bytes()
        while True:
            c = self.read(1)
            s += c
        if len(c) == 0 or chr(c[0]) == "\n":
            return s


class ClientExit(Exception):
    pass


class ServerException(Exception):
    pass


class Server:
    """
    server protocol. Methods from classes that subclass Server can be called
    transparently from a client
    """

    def __init__(self, s, logf=sys.stderr, log_prefix=""):
        self.logf = logf
        self.log_prefix = log_prefix

        # connection

        self.conn = s
        self.fs = FileSock(s)

    def log(self, s):
        self.logf.write("Server log %s: %s\n" % (self.log_prefix, s))

    def one_function(self):
        """
        Executes a single function with associated I/O.
        Protocol:
        - the arguments and results are serialized with the pickle protocol
        - client sends : (fname,args)
            fname = method name to call
            args = tuple of arguments
        - server sends result: (rid,st,ret)
            rid = request id
            st = None, or exception if there was during execution
            ret = return value or None if st!=None
        """

        try:
            (fname, args) = RestrictedUnpickler(self.fs).load()
        except EOFError:
            raise ClientExit("read args")
        self.log("executing method %s" % (fname))
        st = None
        ret = None
        try:
            f = getattr(self, fname)
        except AttributeError:
            st = AttributeError("unknown method " + fname)
            self.log("unknown method")

        try:
            ret = f(*args)
        except Exception as e:
            # due to a bug (in mod_python?), ServerException cannot be
            # unpickled, so send the string and make the exception on the
            # client side

            # st=ServerException(
            #  "".join(traceback.format_tb(sys.exc_info()[2]))+
            #  str(e))
            st = "".join(traceback.format_tb(sys.exc_info()[2])) + str(e)
            self.log("exception in method")
            traceback.print_exc(50, self.logf)
            self.logf.flush()

        LOG.info("return")
        try:
            pickle.dump((st, ret), self.fs, protocol=4)
        except EOFError:
            raise ClientExit("function return")

    def exec_loop(self):
        """main execution loop. Loops and handles exit states"""

        self.log("in exec_loop")
        try:
            while True:
                self.one_function()
        except ClientExit as e:
            self.log("ClientExit %s" % e)
        except socket.error as e:
            self.log("socket error %s" % e)
            traceback.print_exc(50, self.logf)
        except EOFError:
            self.log("EOF during communication")
            traceback.print_exc(50, self.logf)
        except BaseException:
            # unexpected
            traceback.print_exc(50, sys.stderr)
            sys.exit(1)

        LOG.info("exit server")

    def exec_loop_cleanup(self):
        pass

    ###################################################################
    # spying stuff

    def get_ps_stats(self):
        ret = ""
        f = os.popen(
            "echo ============ `hostname` uptime:; uptime;"
            + "echo ============ self:; "
            + "ps -p %d -o pid,vsize,rss,%%cpu,nlwp,psr; " % os.getpid()
            + "echo ============ run queue:;"
            + "ps ar -o user,pid,%cpu,%mem,ni,nlwp,psr,vsz,rss,cputime,command"
        )
        for l in f:
            ret += l
        return ret


class Client:
    """
    Methods of the server object can be called transparently. Exceptions are
    re-raised.
    """

    def __init__(self, HOST, port=PORT, v6=False):
        socktype = socket.AF_INET6 if v6 else socket.AF_INET

        sock = socket.socket(socktype, socket.SOCK_STREAM)
        LOG.info("connecting to %s:%d, socket type: %s", HOST, port, socktype)
        sock.connect((HOST, port))
        self.sock = sock
        self.fs = FileSock(sock)

    def generic_fun(self, fname, args):
        # int "gen fun",fname
        pickle.dump((fname, args), self.fs, protocol=4)
        return self.get_result()

    def get_result(self):
        (st, ret) = RestrictedUnpickler(self.fs).load()
        if st != None:
            raise ServerException(st)
        else:
            return ret

    def __getattr__(self, name):
        return lambda *x: self.generic_fun(name, x)


def run_server(new_handler, port=PORT, report_to_file=None, v6=False):

    HOST = ""  # Symbolic name meaning the local host
    socktype = socket.AF_INET6 if v6 else socket.AF_INET
    s = socket.socket(socktype, socket.SOCK_STREAM)
    s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)

    LOG.info("bind %s:%d", HOST, port)
    s.bind((HOST, port))
    s.listen(5)

    LOG.info("accepting connections")
    if report_to_file is not None:
        LOG.info("storing host+port in %s", report_to_file)
        open(report_to_file, "w").write("%s:%d " % (socket.gethostname(), port))

    while True:
        try:
            conn, addr = s.accept()
        except socket.error as e:
            if e[1] == "Interrupted system call":
                continue
            raise

        LOG.info("Connected to %s", addr)

        ibs = new_handler(conn)

        tid = _thread.start_new_thread(ibs.exec_loop, ())

        LOG.debug("Thread ID: %d", tid)