[Entrypoints][Refactor] Refactor llumnix entrypoints to be more modular

configs/base.yml

@@ -1,6 +1,6 @@
 SERVER:
   HOST: '127.0.0.1'
-  PORT: 37000
+  PORT: 1234
   QUEUE_TYPE: "rayqueue"
 
 RAY:

examlpes/offline_inference.py

@@ -27,7 +27,7 @@
 ray_cluster_port=6379
 
 # Note: launch_ray_cluster will stop current ray cluster first, then init a new one.
-launch_ray_cluster(ray_cluster_port=ray_cluster_port)
+launch_ray_cluster(port=ray_cluster_port)
 connect_to_ray_cluster(port=ray_cluster_port)
 
 # Set manager args and engine args.

llumnix/__init__.py

@@ -15,8 +15,10 @@
 from vllm import *
 
 from llumnix.server_info import ServerInfo
-from llumnix.entrypoints.llumnix_utils import (launch_ray_cluster, connect_to_ray_cluster,
-                                               init_manager, init_llumlets)
+from llumnix.entrypoints.utils import (launch_ray_cluster,
+                                       connect_to_ray_cluster,
+                                       init_manager,
+                                       init_llumlets)
 from llumnix.arg_utils import EngineManagerArgs
 from llumnix.llm_engine_manager import LLMEngineManager
 from llumnix.llumlet.llumlet import Llumlet

llumnix/arg_utils.py

@@ -22,6 +22,77 @@
 from llumnix.config import LlumnixConfig, get_llumnix_config
 from llumnix.config.default import _C
 
+
+class LlumnixArgumentParser(argparse.ArgumentParser):
+    def __init__(self, *args, **kwargs):
+        self.cur_namespace = "llumnix"
+        super().__init__(*args, **kwargs)
+
+    def set_namespace(self, namespace: str):
+        self.cur_namespace = namespace
+
+    def add_argument(self, *args, **kwargs):
+        if self.cur_namespace == 'llumnix' and "--help" not in args:
+            assert 'default' not in kwargs or kwargs['default'] is None, \
+                f"Do not set the default value for '{args[0]}' in CLI, or set default value to None. " \
+                f"The default value will be retrieved from config/default.py in get_llumnix_config."
+            if kwargs.get('action') == 'store_true':
+                kwargs['default'] = None
+        super().add_argument(*args, **kwargs)
+
+
+# All the default values of llumnix arguments are setted in default.py. So all the arguments here are setted to None.
+
+@dataclass
+class LlumnixEntrypointsArgs:
+    launch_ray_cluster: bool = None
+    ray_cluster_port: int = None
+    queue_type: str = None
+    request_output_queue_port: int = None
+    disable_log_requests_server: bool = None
+    log_request_timestamps: bool = None
+    config_file: bool = None
+
+    def __post_init__(self):
+        for attr in dataclasses.fields(self):
+            if getattr(self, attr.name) is None:
+                setattr(self, attr.name, getattr(_C.SERVER, attr.name.upper()))
+
+    @classmethod
+    def from_llumnix_config(cls, cfg: LlumnixConfig = get_llumnix_config()) -> 'LlumnixEntrypointsArgs':
+        # Get the list of attributes of this dataclass.
+        attrs = [attr.name for attr in dataclasses.fields(cls)]
+        # Set the attributes from the parsed arguments.
+        # The defalut values of attributes are defined in default.py.
+        llumnix_entrypoints_args = cls(**{attr: getattr(cfg.SERVER, attr.upper()) for attr in attrs})
+        return llumnix_entrypoints_args
+
+    @staticmethod
+    def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
+        parser.add_argument('--launch-ray-cluster',
+                            action='store_true',
+                            help='if launch ray cluster in api server')
+        parser.add_argument("--ray-cluster-port",
+                            type=int,
+                            help='ray cluster port')
+        parser.add_argument("--queue-type",
+                            type=str,
+                            choices=['rayqueue', 'zmq'],
+                            help='queue type for request output queue')
+        parser.add_argument("--request-output-queue-port",
+                            type=int,
+                            help='port for zmq')
+        parser.add_argument('--disable-log-requests-server',
+                            action='store_true',
+                            help='disable logging requests in server')
+        parser.add_argument("--log-request-timestamps",
+                            action='store_true',
+                            help='if log request timestamps')
+        parser.add_argument("--config-file",
+                            type=str,
+                            help="path to config file")
+        return parser
+
 @dataclass
 class EngineManagerArgs:
     disable_init_instance_by_manager: bool = None
@@ -106,6 +177,7 @@ def from_llumnix_config(cls, cfg: LlumnixConfig = get_llumnix_config()) -> 'Engi
         # Get the list of attributes of this dataclass.
         attrs = [attr.name for attr in dataclasses.fields(cls)]
         # Set the attributes from the parsed arguments.
+        # The defalut values of attributes are defined in default.py.
         engine_manager_args = cls(**{attr: getattr(cfg.MANAGER, attr.upper()) for attr in attrs})
         return engine_manager_args
 

llumnix/backends/vllm/llm_engine.py

@@ -39,7 +39,7 @@
 from llumnix.server_info import ServerInfo
 from llumnix.internal_config import MigrationConfig
 from llumnix.queue.queue_client_base import QueueClientBase
-from llumnix.queue.utils import get_output_queue_client, QueueType
+from llumnix.queue.utils import init_output_queue_client, QueueType
 
 logger = init_logger(__name__)
 
@@ -48,7 +48,7 @@ class AsyncPutQueueActor:
     def __init__(self, instance_id, output_queue_type: QueueType):
         self.instance_id = instance_id
         self.output_queue_type = output_queue_type
-        self.request_output_queue_client: QueueClientBase = get_output_queue_client(output_queue_type)
+        self.request_output_queue_client: QueueClientBase = init_output_queue_client(output_queue_type)
         self.engine_actor_handle = None
 
     async def put_nowait_to_servers(self,
@@ -225,7 +225,7 @@ def step(self) -> None:
             tot_blocks = set(tot_blocks)
             instance_info.num_blocks_last_running_request = len(tot_blocks)
         if request_outputs:
-            self.put_queue_args_queue.put((request_outputs, server_infos))
+            self.put_queue_args_queue.put_nowait((request_outputs, server_infos))
         self.instance_info = instance_info
         for request_output in request_outputs:
             if hasattr(request_output, 'request_timestamps'):

llumnix/backends/vllm/sequence.py

@@ -25,6 +25,9 @@ def __init__(self, request_id, server_info, expected_steps: int, *args, **kwargs
     def prompt_len(self) -> int:
         return self.get_seqs()[0].get_prompt_len()
 
+    def is_finished(self) -> bool:
+        return self.get_seqs()[0].is_finished()
+
     @property
     def request_len(self) -> int:
         return self.get_seqs()[0].get_len()

llumnix/config/default.py

@@ -21,34 +21,35 @@
 _C = LC()
 
 # -----------------------------------------------------------------------------
-# API SERVER CONFIGURATION
+# SERVER CONFIGURATION
 # -----------------------------------------------------------------------------
 _C.SERVER = LC()
 # Hostname for the server
 _C.SERVER.HOST = "localhost"
 # Port number for the server
 _C.SERVER.PORT = 8000
-# Queue type for request output queue
-_C.SERVER.QUEUE_TYPE = "rayqueue"
-# Port number for the request output queue
-_C.SERVER.REQUEST_OUTPUT_QUEUE_PORT = 1234
 # Path to SSL key file for secure connections
 _C.SERVER.SSL_KEYFILE = None
 # Path to SSL certificate file for secure connections
 _C.SERVER.SSL_CERTFILE = None
+# Queue type for request output queue
+_C.SERVER.QUEUE_TYPE = "rayqueue"
+# Port number for the request output queue
+_C.SERVER.REQUEST_OUTPUT_QUEUE_PORT = 1234
 # Disable logging requests in server
 _C.SERVER.DISABLE_LOG_REQUESTS_SERVER = False
 # Enable logging request timestamp
 _C.SERVER.LOG_REQUEST_TIMESTAMPS = False
+# Config file of Llumnix arguments
+_C.SERVER.CONFIG_FILE = None
 
 # -----------------------------------------------------------------------------
 # RAY CONFIGURATION
 # -----------------------------------------------------------------------------
-_C.RAY = LC()
-# Port number for the Ray cluster
-_C.RAY.RAY_CLUSTER_PORT = 6379
 # If True, launch Ray cluster in API server
-_C.RAY.LAUNCH_RAY_CLUSTER = False
+_C.SERVER.LAUNCH_RAY_CLUSTER = False
+# Port number for the Ray cluster
+_C.SERVER.RAY_CLUSTER_PORT = 6379
 
 # -----------------------------------------------------------------------------
 # MANAGER CONFIGURATION