Python zmq.utils.jsonapi.loads() Examples

def server_status(self):
        """
            Get the current status of the server connected to this client

        :return: a dictionary contains the current status of the server connected to this client
        :rtype: dict[str, str]

        """
        try:
            self.receiver.setsockopt(zmq.RCVTIMEO, self.timeout)
            self._send(b'SHOW_CONFIG')
            return jsonapi.loads(self._recv().content[1])
        except zmq.error.Again as _e:
            t_e = TimeoutError(
                'no response from the server (with "timeout"=%d ms), '
                'is the server on-line? is network broken? are "port" and "port_out" correct?' % self.timeout)
            if _py2:
                raise t_e
            else:
                raise t_e from _e
        finally:
            self.receiver.setsockopt(zmq.RCVTIMEO, -1) 

def recv_json(self, flags=0, **kwargs):
        """Receive a Python object as a message using json to serialize.

        Keyword arguments are passed on to json.loads
        
        Parameters
        ----------
        flags : int
            Any valid flags for :func:`Socket.recv`.

        Returns
        -------
        obj : Python object
            The Python object that arrives as a message.

        Raises
        ------
        ZMQError
            for any of the reasons :func:`~Socket.recv` might fail
        """
        msg = self.recv(flags)
        return self._deserialize(msg, lambda buf: jsonapi.loads(buf, **kwargs)) 

def recv_pyobj(self, flags=0):
        """Receive a Python object as a message using pickle to serialize.

        Parameters
        ----------
        flags : int
            Any valid flags for :func:`Socket.recv`.

        Returns
        -------
        obj : Python object
            The Python object that arrives as a message.

        Raises
        ------
        ZMQError
            for any of the reasons :func:`~Socket.recv` might fail
        """
        msg = self.recv(flags)
        return self._deserialize(msg, pickle.loads) 

def server_status(self):
        """
            Get the current status of the server connected to this client

        :return: a dictionary contains the current status of the server connected to this client
        :rtype: dict[str, str]

        """
        try:
            self.receiver.setsockopt(zmq.RCVTIMEO, self.timeout)
            self._send(b'SHOW_CONFIG')
            return jsonapi.loads(self._recv().content[1])
        except zmq.error.Again as _e:
            t_e = TimeoutError(
                'no response from the server (with "timeout"=%d ms), '
                'is the server on-line? is network broken? are "port" and "port_out" correct?' % self.timeout)
            if _py2:
                raise t_e
            else:
                raise t_e from _e
        finally:
            self.receiver.setsockopt(zmq.RCVTIMEO, -1) 

def recv_json(self, flags=0, **kwargs):
        """Receive a Python object as a message using json to serialize.

        Keyword arguments are passed on to json.loads
        
        Parameters
        ----------
        flags : int
            Any valid flags for :func:`Socket.recv`.

        Returns
        -------
        obj : Python object
            The Python object that arrives as a message.

        Raises
        ------
        ZMQError
            for any of the reasons :func:`~Socket.recv` might fail
        """
        msg = self.recv(flags)
        return self._deserialize(msg, lambda buf: jsonapi.loads(buf, **kwargs)) 

def recv_pyobj(self, flags=0):
        """Receive a Python object as a message using pickle to serialize.

        Parameters
        ----------
        flags : int
            Any valid flags for :func:`Socket.recv`.

        Returns
        -------
        obj : Python object
            The Python object that arrives as a message.

        Raises
        ------
        ZMQError
            for any of the reasons :func:`~Socket.recv` might fail
        """
        msg = self.recv(flags)
        return self._deserialize(msg, pickle.loads) 

def recv_json(self, flags=0):
        """receive a Python object as a message using json to serialize

        Parameters
        ----------
        flags : int
            Any valid recv flag.

        Returns
        -------
        obj : Python object
            The Python object that arrives as a message.
        """
        if jsonapi.jsonmod is None:
            raise ImportError('jsonlib{1,2}, json or simplejson library is required.')
        else:
            msg = self.recv(flags)
            return jsonapi.loads(msg) 

def recv_json(self, flags=0, **kwargs):
        """Receive a Python object as a message using json to serialize.

        Keyword arguments are passed on to json.loads
        
        Parameters
        ----------
        flags : int
            Any valid flags for :func:`Socket.recv`.

        Returns
        -------
        obj : Python object
            The Python object that arrives as a message.

        Raises
        ------
        ZMQError
            for any of the reasons :func:`~Socket.recv` might fail
        """
        from zmq.utils import jsonapi
        msg = self.recv(flags)
        return self._deserialize(msg, lambda buf: jsonapi.loads(buf, **kwargs)) 

def recv_pyobj(self, flags=0):
        """Receive a Python object as a message using pickle to serialize.

        Parameters
        ----------
        flags : int
            Any valid flags for :func:`Socket.recv`.

        Returns
        -------
        obj : Python object
            The Python object that arrives as a message.

        Raises
        ------
        ZMQError
            for any of the reasons :func:`~Socket.recv` might fail
        """
        msg = self.recv(flags)
        return self._deserialize(msg, pickle.loads) 

def _recv_ndarray(self, wait_for_req_id=None):
        request_id, response = self._recv(wait_for_req_id)
        arr_info, arr_val = jsonapi.loads(response[1]), response[2]
        X = np.frombuffer(_buffer(arr_val), dtype=str(arr_info['dtype']))
        return Response(request_id, self.formatter(X.reshape(arr_info['shape'])), arr_info.get('tokens', '')) 

def _recv_ndarray(self):
        request_id, response = self._recv()
        arr_info, arr_val = jsonapi.loads(response[1]), response[2]
        X = np.frombuffer(_buffer(arr_val), dtype=str(arr_info['dtype']))
        return Response(request_id, self.formatter(X.reshape(arr_info['shape']))) 

def server_status(self):
        """
            Get the current status of the server connected to this client

        :return: a dictionary contains the current status of the server connected to this client
        :rtype: dict[str, str]

        """
        req_id = self._send(b'SHOW_CONFIG')
        return jsonapi.loads(self._recv(req_id).content[1]) 

def _recv_scores(self, wait_for_req_id=None):
        request_id, response = self._recv(wait_for_req_id)
        scores = jsonapi.loads(response[1])
        return Response(request_id, scores) 

def input_fn_builder(self, worker):
        def gen():
            while not self.exit_flag.is_set():
                client_id, msg = worker.recv_multipart()
                msg = jsonapi.loads(msg)
                self.logger.info('new job %s, size: %d' % (client_id, len(msg)))
                if BertClient.is_valid_input(msg):
                    tmp_f = list(convert_lst_to_features(msg, self.max_seq_len, self.tokenizer))
                    yield {
                        'client_id': client_id,
                        'input_ids': [f.input_ids for f in tmp_f],
                        'input_mask': [f.input_mask for f in tmp_f],
                        'input_type_ids': [f.input_type_ids for f in tmp_f]
                    }
                else:
                    self.logger.error('unsupported type of job %s! sending back None' % client_id)

        def input_fn():
            return (tf.data.Dataset.from_generator(
                gen,
                output_types={'input_ids': tf.int32,
                              'input_mask': tf.int32,
                              'input_type_ids': tf.int32,
                              'client_id': tf.string},
                output_shapes={
                    'client_id': (),
                    'input_ids': (None, self.max_seq_len),
                    'input_mask': (None, self.max_seq_len),
                    'input_type_ids': (None, self.max_seq_len)}))

        return input_fn 

def _recv_ndarray(self):
        request_id, response = self._recv()
        arr_info, arr_val = jsonapi.loads(response[1]), response[2]
        X = np.frombuffer(_buffer(arr_val), dtype=str(arr_info['dtype']))
        return Response(request_id, self.formatter(X.reshape(arr_info['shape']))) 

def recv_ndarray(self):
        response = self.recv()
        arr_info, arr_val = jsonapi.loads(response[1]), response[2]
        X = np.frombuffer(_buffer(arr_val), dtype=arr_info['dtype'])
        return self.formatter(X.reshape(arr_info['shape'])) 

def on_message(self, msg):
        msg = jsonapi.loads(msg)
        self.session.send(self.zmq_stream, msg) 

def recv_pyobj(self, flags=0):
        """receive a Python object as a message using pickle to serialize

        Parameters
        ----------
        flags : int
            Any valid recv flag.

        Returns
        -------
        obj : Python object
            The Python object that arrives as a message.
        """
        s = self.recv(flags)
        return pickle.loads(s) 

def server_status(self):
        """
            Get the current status of the server connected to this client

        :return: a dictionary contains the current status of the server connected to this client
        :rtype: dict[str, str]

        """
        req_id = self._send(b'SHOW_STATUS')
        return jsonapi.loads(self._recv(req_id).content[1]) 

def server_config(self):
        """
            Get the current configuration of the server connected to this client

        :return: a dictionary contains the current configuration of the server connected to this client
        :rtype: dict[str, str]

        """
        req_id = self._send(b'SHOW_CONFIG')
        return jsonapi.loads(self._recv(req_id).content[1]) 

def _handle_pipe(self):
        """
        Handle a message from front-end API.
        """
        terminate = False

        # Get the whole message off the pipe in one go
        msg = self.pipe.recv_multipart()

        if msg is None:
            terminate = True
            return terminate

        command = msg[0]
        self.log.debug("auth received API command %r", command)

        if command == b'ALLOW':
            addresses = [u(m, self.encoding) for m in msg[1:]]
            try:
                self.authenticator.allow(*addresses)
            except Exception as e:
                self.log.exception("Failed to allow %s", addresses)

        elif command == b'DENY':
            addresses = [u(m, self.encoding) for m in msg[1:]]
            try:
                self.authenticator.deny(*addresses)
            except Exception as e:
                self.log.exception("Failed to deny %s", addresses)

        elif command == b'PLAIN':
            domain = u(msg[1], self.encoding)
            json_passwords = msg[2]
            self.authenticator.configure_plain(domain, jsonapi.loads(json_passwords))

        elif command == b'CURVE':
            # For now we don't do anything with domains
            domain = u(msg[1], self.encoding)

            # If location is CURVE_ALLOW_ANY, allow all clients. Otherwise
            # treat location as a directory that holds the certificates.
            location = u(msg[2], self.encoding)
            self.authenticator.configure_curve(domain, location)

        elif command == b'TERMINATE':
            terminate = True

        else:
            self.log.error("Invalid auth command from API: %r", command)

        return terminate 

def _handle_pipe(self):
        """
        Handle a message from front-end API.
        """
        terminate = False

        # Get the whole message off the pipe in one go
        msg = self.pipe.recv_multipart()

        if msg is None:
            terminate = True
            return terminate

        command = msg[0]
        self.log.debug("auth received API command %r", command)

        if command == b'ALLOW':
            addresses = [u(m, self.encoding) for m in msg[1:]]
            try:
                self.authenticator.allow(*addresses)
            except Exception as e:
                self.log.exception("Failed to allow %s", addresses)

        elif command == b'DENY':
            addresses = [u(m, self.encoding) for m in msg[1:]]
            try:
                self.authenticator.deny(*addresses)
            except Exception as e:
                self.log.exception("Failed to deny %s", addresses)

        elif command == b'PLAIN':
            domain = u(msg[1], self.encoding)
            json_passwords = msg[2]
            self.authenticator.configure_plain(domain, jsonapi.loads(json_passwords))

        elif command == b'CURVE':
            # For now we don't do anything with domains
            domain = u(msg[1], self.encoding)

            # If location is CURVE_ALLOW_ANY, allow all clients. Otherwise
            # treat location as a directory that holds the certificates.
            location = u(msg[2], self.encoding)
            self.authenticator.configure_curve(domain, location)

        elif command == b'TERMINATE':
            terminate = True

        else:
            self.log.error("Invalid auth command from API: %r", command)

        return terminate 

def input_fn_builder(self, socks, tf, sink):
        from .bert.extract_features import convert_lst_to_features
        from .bert.tokenization import FullTokenizer

        def gen():
            # Windows does not support logger in MP environment, thus get a new logger
            # inside the process for better compatibility
            logger = set_logger(colored('WORKER-%d' % self.worker_id, 'yellow'), self.verbose)
            tokenizer = FullTokenizer(vocab_file=os.path.join(self.model_dir, 'vocab.txt'), do_lower_case=self.do_lower_case)

            poller = zmq.Poller()
            for sock in socks:
                poller.register(sock, zmq.POLLIN)

            logger.info('ready and listening!')
            self.is_ready.set()

            while not self.exit_flag.is_set():
                events = dict(poller.poll())
                for sock_idx, sock in enumerate(socks):
                    if sock in events:
                        client_id, raw_msg = sock.recv_multipart()
                        msg = jsonapi.loads(raw_msg)
                        logger.info('new job\tsocket: %d\tsize: %d\tclient: %s' % (sock_idx, len(msg), client_id))
                        # check if msg is a list of list, if yes consider the input is already tokenized
                        is_tokenized = all(isinstance(el, list) for el in msg)
                        tmp_f = list(convert_lst_to_features(msg, self.max_seq_len,
                                                             self.bert_config.max_position_embeddings,
                                                             tokenizer, logger,
                                                             is_tokenized, self.mask_cls_sep, self.no_special_token))
                        if self.show_tokens_to_client:
                            sink.send_multipart([client_id, jsonapi.dumps([f.tokens for f in tmp_f]),
                                                 b'', ServerCmd.data_token])
                        yield {
                            'client_id': client_id,
                            'input_ids': [f.input_ids for f in tmp_f],
                            'input_mask': [f.input_mask for f in tmp_f],
                            'input_type_ids': [f.input_type_ids for f in tmp_f]
                        }

        def input_fn():
            return (tf.data.Dataset.from_generator(
                gen,
                output_types={'input_ids': tf.int32,
                              'input_mask': tf.int32,
                              'input_type_ids': tf.int32,
                              'client_id': tf.string},
                output_shapes={
                    'client_id': (),
                    'input_ids': (None, None),
                    'input_mask': (None, None),
                    'input_type_ids': (None, None)}).prefetch(self.prefetch_size))

        return input_fn 

def scorer(self, socks, sink):
        logger = set_logger(colored('WORKER-%d' % self.worker_id, 'yellow'), self.verbose)

        poller = zmq.Poller()
        for sock in socks:
            poller.register(sock, zmq.POLLIN)

        logger.info('ready and listening!')
        
#        summary = ['this is a cat.']
#        references_text = [['this is a cat.'], ['this is a lovely cat.']]
#        logger.info(rouge_n(summary, references_text, 1, alpha=0))
        
#        system = ['This is test summary']
#        references = ['This is ref summary two','this is test summary']        
#                                
#        score = sent_mover_score(references, system * len(references))
#        logger.info(np.mean(score))

        while not self.exit_flag.is_set():
            events = dict(poller.poll())
            for sock_idx, sock in enumerate(socks):
                if sock in events:
                    client_id, raw_msg = sock.recv_multipart()
                    msg = jsonapi.loads(raw_msg)
                    logger.info('new job\tsocket: %d\tsize: %d\tclient: %s' % (sock_idx, len(msg), client_id))
                    # check if msg is a list of list, if yes consider the input is already tokenized

                    for index,_ in enumerate(msg):
                        system = msg[index][0]
                        references = msg[index][1]
                        eval_metric = msg[index][2]
                        if eval_metric[:-1] == 'rouge_':
                            if eval_metric[-1] == '1' or eval_metric[-1] == 'n':
                                score = rouge_n(system, [[r] for r in references], 1, alpha=0)
                            if eval_metric[-1] == '2':
                                score = rouge_n(system, [[r] for r in references], 2, alpha=0)                         
                        if eval_metric[:-1] == 'wmd_':
                            if eval_metric[-1] == '1':
                                score = np.mean(word_mover_score(1, references, system * len(references)))
                            if eval_metric[-1] == '2':
                                score = np.mean(word_mover_score(2, references, system * len(references)))
                        if eval_metric == 'smd':
                            score = np.mean(sent_mover_score(references, system * len(references)))
                            
                        yield {
                            'client_id': client_id,
                            'score': score
                        } 

def _handle_pipe(self):
        """
        Handle a message from front-end API.
        """
        terminate = False

        # Get the whole message off the pipe in one go
        msg = self.pipe.recv_multipart()

        if msg is None:
            terminate = True
            return terminate

        command = msg[0]
        self.log.debug("auth received API command %r", command)

        if command == b'ALLOW':
            addresses = [u(m, self.encoding) for m in msg[1:]]
            try:
                self.authenticator.allow(*addresses)
            except Exception as e:
                self.log.exception("Failed to allow %s", addresses)

        elif command == b'DENY':
            addresses = [u(m, self.encoding) for m in msg[1:]]
            try:
                self.authenticator.deny(*addresses)
            except Exception as e:
                self.log.exception("Failed to deny %s", addresses)

        elif command == b'PLAIN':
            domain = u(msg[1], self.encoding)
            json_passwords = msg[2]
            self.authenticator.configure_plain(domain, jsonapi.loads(json_passwords))

        elif command == b'CURVE':
            # For now we don't do anything with domains
            domain = u(msg[1], self.encoding)

            # If location is CURVE_ALLOW_ANY, allow all clients. Otherwise
            # treat location as a directory that holds the certificates.
            location = u(msg[2], self.encoding)
            self.authenticator.configure_curve(domain, location)

        elif command == b'TERMINATE':
            terminate = True

        else:
            self.log.error("Invalid auth command from API: %r", command)

        return terminate 

def run(self):
        available_gpus = range(self.num_worker)
        run_on_gpu = True
        num_req = 0
        try:
            import GPUtil
            available_gpus = GPUtil.getAvailable(limit=self.num_worker)
            if len(available_gpus) < self.num_worker:
                self.logger.warn('only %d GPU(s) is available, but ask for %d' % (len(available_gpus), self.num_worker))
        except FileNotFoundError:
            self.logger.warn('nvidia-smi is missing, often means no gpu found on this machine. '
                             'will run service on cpu instead')
            run_on_gpu = False

        # start the backend processes
        for i in available_gpus:
            process = BertWorker(i, self.args, self.addr_backend, self.addr_sink)
            self.processes.append(process)
            process.start()

        try:
            while True:
                client, msg = self.frontend.recv_multipart()
                if msg == ServerCommand.show_config:
                    self.sink.send_multipart([client, msg,
                                              jsonapi.dumps({**{'client': client.decode('ascii'),
                                                                'num_subprocess': len(self.processes),
                                                                'frontend -> backend': self.addr_backend,
                                                                'backend -> sink': self.addr_sink,
                                                                'frontend <-> sink': self.addr_front2sink,
                                                                'server_current_time': str(datetime.now()),
                                                                'run_on_gpu': run_on_gpu,
                                                                'num_request': num_req},
                                                             **self.args_dict})])
                    continue

                num_req += 1
                client = client + b'#' + str(uuid.uuid4()).encode('ascii')
                seqs = jsonapi.loads(msg)
                num_seqs = len(seqs)
                # tell sink to collect a new job
                self.sink.send_multipart([client, ServerCommand.new_job, b'%d' % num_seqs])

                if num_seqs > self.max_batch_size:
                    # divide the large batch into small batches
                    s_idx = 0
                    while s_idx < num_seqs:
                        tmp = seqs[s_idx: (s_idx + self.max_batch_size)]
                        if tmp:
                            # get the worker with minimum workload
                            client_partial_id = client + b'@%d' % s_idx
                            self.backend.send_multipart([client_partial_id, jsonapi.dumps(tmp)])
                        s_idx += len(tmp)
                else:
                    self.backend.send_multipart([client, msg])
        except zmq.error.ContextTerminated:
            self.logger.error('context is closed!') 

def _handle_pipe(self):
        """
        Handle a message from front-end API.
        """
        terminate = False

        # Get the whole message off the pipe in one go
        msg = self.pipe.recv_multipart()

        if msg is None:
            terminate = True
            return terminate

        command = msg[0]
        self.log.debug("auth received API command", command=command)

        if command == b'ALLOW':
            addresses = [u(m, self.encoding) for m in msg[1:]]
            try:
                self.authenticator.allow(*addresses)
            except Exception as e:
                self.log.exception("Failed to allow", addresses=addresses, e=e)

        elif command == b'DENY':
            addresses = [u(m, self.encoding) for m in msg[1:]]
            try:
                self.authenticator.deny(*addresses)
            except Exception as e:
                self.log.exception("Failed to deny", addresses=addresses, e=e)

        elif command == b'PLAIN':
            domain = u(msg[1], self.encoding)
            json_passwords = msg[2]
            self.authenticator.configure_plain(domain, jsonapi.loads(json_passwords))

        elif command == b'CURVE':
            # For now we don't do anything with domains
            domain = u(msg[1], self.encoding)

            # If location is CURVE_ALLOW_ANY, allow all clients. Otherwise
            # treat location as a directory that holds the certificates.
            location = u(msg[2], self.encoding)
            self.authenticator.configure_curve(domain, location)

        elif command == b'TERMINATE':
            terminate = True

        else:
            self.log.error("Invalid auth command from API", command=command)

        return terminate