claudson/ai/lib/python3.12/site-packages/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("Sever 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 sever")

    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)