inference/server/plain_server.py

import argparse
import random
import socket
import time
import traceback

import json
import logging
import torch
from transformers import AutoModelForCausalLM, AutoProcessor, TextIteratorStreamer
from threading import Thread
from multiprocessing import Process, Queue

EOS_FLAG = "<EOS>"
SEPARATOR = "<SEP>"


def get_logger(name):
    logger = logging.getLogger(name)
    logger.setLevel(logging.INFO)
    handler = logging.StreamHandler()
    formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(name)s - %(message)s")
    handler.setFormatter(formatter)
    logger.addHandler(handler)
    return logger


class Streamer(object):

    def __init__(self, timeout=None):
        self.timeout = timeout
        self.queue = Queue(maxsize=1024)
        self.stop_signal = EOS_FLAG

    def put(self, value):
        self.queue.put(value)

    def __iter__(self):
        return self

    def __next__(self):
        try:
            value = self.queue.get(timeout=self.timeout)
        except:
            raise StopIteration()

        if value == self.stop_signal:
            raise StopIteration()
        else:
            return value


class PenguinVLQwen3PlainClient(object):

    def __init__(self, host="localhost", port=16666):
        self.host = host
        self.port = port

        self.input_buffer = Queue(maxsize=1024)
        self.streamers = dict()
        self.logger = get_logger("penguinvl_qwen3.client")

        client_thread = Thread(target=self._client_worker)
        client_thread.deamon = True
        client_thread.start()        

    def _receive_worker(self, server_socket):
        try:
            while True:
                data = server_socket.recv(8192)
                if not data:
                    self.logger.info(f"Connection has been terminated.")
                    for streamer in self.streamers.values():
                        streamer.put(streamer.stop_signal)
                    break

                for sub_data in data.decode("utf-8").split(SEPARATOR):
                    if len(sub_data) == 0:
                        continue

                    try:
                        sub_data = json.loads(sub_data)
                    except:
                        self.logger.info(f"Failed to parse data: {sub_data}")
                        continue

                    self.logger.info(f"Received: {sub_data['data']}")
                    self.streamers[sub_data["id"]].put(sub_data["data"])

                    if sub_data["data"] == EOS_FLAG:
                        self.streamers.pop(sub_data["id"])

        except ConnectionResetError:
            self.logger.info(f"Connection has been terminated.")

    def _send_worker(self, server_socket):
        while True:
            request_id, conversation = self.input_buffer.get()
            data = json.dumps({"id": request_id, "data": conversation}) + SEPARATOR
            server_socket.sendall(data.encode("utf-8"))
            self.logger.info(f"Sent: {data}")

    def _client_worker(self):
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as server_socket:
            while True:
                try:
                    server_socket.connect((self.host, self.port))
                    break
                except ConnectionRefusedError:
                    self.logger.info("Waiting for the server to start...")
                    time.sleep(1)
                    continue

            self.logger.info("Connected to server.")
            receive_thread = Thread(target=self._receive_worker, args=(server_socket,))
            receive_thread.daemon = True
            receive_thread.start()

            send_thread = Thread(target=self._send_worker, args=(server_socket,))
            send_thread.daemon = True
            send_thread.start()

            receive_thread.join()

    def submit(self, conversation):
        request_id = random.randint(0, 4294967295)
        streamer = Streamer()
        self.streamers[request_id] = streamer
        self.input_buffer.put((request_id, conversation))
        return streamer


class PenguinVLQwen3PlainServer(object):

    def __init__(
        self,
        model_path,
        torch_dtype=torch.bfloat16,
        attn_implementation="flash_attention_2",
        num_processes=1,
        buffer_size=2,
        host="localhost",
        port=16666,
    ):
        self.model_path = model_path
        self.torch_dtype = torch_dtype
        self.attn_implementation = attn_implementation
        self.num_processes = num_processes
        self.buffer_size = buffer_size

        self.host = host
        self.port = port

    def _model_worker(self, input_buffer, output_buffer, device_map, rank):
        logger = get_logger(f"penguinvl_qwen3.server.worker_{rank}")
        logger.info(f"Loading model from {self.model_path}...")

        model = AutoModelForCausalLM.from_pretrained(
            self.model_path,
            trust_remote_code=True,
            torch_dtype=self.torch_dtype,
            attn_implementation=self.attn_implementation,
            device_map=device_map,
        )
        processor = AutoProcessor.from_pretrained(self.model_path, trust_remote_code=True)
        logger.info(f"Successfully loaded model.")

        while True:
            logger.info("Waiting for input...")
            request_id, data = input_buffer.get()
            try:
                inputs = processor(
                    conversation=data["conversation"],
                    add_system_prompt=True,
                    add_generation_prompt=True,
                    return_tensors="pt"
                )
                inputs = {k: v.to(f"cuda:{rank}") if isinstance(v, torch.Tensor) else v for k, v in inputs.items()}
                if "pixel_values" in inputs:
                    inputs["pixel_values"] = inputs["pixel_values"].to(torch.bfloat16)

                streamer = TextIteratorStreamer(processor.tokenizer, skip_prompt=True, skip_special_tokens=True)
                generation_kwargs = {
                    **inputs,
                    **data["generation_config"],
                    "streamer": streamer,
                }

                thread = Thread(target=model.generate, kwargs=generation_kwargs)
                thread.deamon = True
                thread.start()

                for token in streamer:
                    output_buffer.put((request_id, token))
                output_buffer.put((request_id, EOS_FLAG))

            except:
                logger.error(f"An error occurred: {traceback.format_exc()}")
                output_buffer.put((request_id, "Server error! Please check the server logs and retry."))
                output_buffer.put((request_id, EOS_FLAG))

    def _receive_worker(self, logger, input_buffer, client_socket, client_address):
        try:
            while True:
                data = client_socket.recv(8192)
                if not data:
                    logger.info(f"Connection from {client_address} has been terminated.")
                    break

                for sub_data in data.decode("utf-8").split(SEPARATOR):
                    if len(sub_data) == 0:
                        continue

                    try:
                        sub_data = json.loads(sub_data)
                    except:
                        logger.info(f"Failed to parse data: {sub_data}")
                        continue

                    logger.info(f"Received from {client_address}: {sub_data}")
                    input_buffer.put((sub_data["id"], sub_data["data"]))

        except ConnectionResetError:
            logger.info(f"Connection from {client_address} has been terminated.")

    def _send_worker(self, logger, output_buffer, client_socket, client_address):
        try:
            while True:
                request_id, token = output_buffer.get()
                data = json.dumps({"id": request_id, "data": token}) + SEPARATOR
                client_socket.sendall(data.encode("utf-8"))

        except ConnectionResetError:
            logger.info(f"Connection from {client_address} has been terminated.")

    def launch(self):
        logger = get_logger(f"penguinvl_qwen3.server.controller")

        input_buffer = Queue(maxsize=self.num_processes * self.buffer_size)
        output_buffer = Queue(maxsize=self.num_processes * 1024)

        for i in range(self.num_processes):
            device_map = {"": f"cuda:{i}"}
            process = Process(target=self._model_worker, args=(input_buffer, output_buffer, device_map, i))
            process.start()

        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as server_socket:
            server_socket.bind((self.host, self.port))
            server_socket.listen(1)
            logger.info("Waiting for connection...")

            while True:
                client_socket, client_address = server_socket.accept()
                logger.info(f"Connected to {client_address}.")

                receive_thread = Thread(target=self._receive_worker, args=(logger, input_buffer, client_socket, client_address))
                receive_thread.deamon = True
                receive_thread.start()

                send_thread = Thread(target=self._send_worker, args=(logger, output_buffer, client_socket, client_address))
                send_thread.deamon = True
                send_thread.start()


if __name__ == "__main__":
    torch.multiprocessing.set_start_method("spawn")

    parser = argparse.ArgumentParser()
    parser.add_argument("--model-path", "--model_path", type=str, required=True)
    parser.add_argument("--nproc", type=int, default=8)
    parser.add_argument("--port", type=int, default=16666)
    args = parser.parse_args()

    server = PenguinVLQwen3PlainServer(
        model_path=args.model_path,
        num_processes=args.nproc,
        port=args.port,
    )
    server.launch()