Update the new test cases

Author: ApostaC

bootstrapper.py

@@ -83,12 +83,18 @@ class LocalLMCacheServerBootstrapper(Bootstrapper):
     """
     Bootstraps a local lmcache server
     """
-    def __init__(self, config: Bootstrapper, log_dir = "/tmp"):
+    def __init__(self, config: BootstrapConfig, log_dir = "/tmp"):
         super().__init__(config)
         server_config = self.parse_lmcache_server_config(self.config.lmcache_config.config_path)
         self.handle = None
         self.started = False
 
+        match config.lmcache_config.remote_device:
+            case None:
+                self.remote_device = "cpu"
+            case path:
+                self.remote_device = path
+
         if server_config is None:
             self.is_needed = False
         else:
@@ -126,7 +132,7 @@ def start(self):
         if not self.is_needed:
             return
 
-        cmd = f"python3 -um lmcache_server.server {self.host} {self.port}"
+        cmd = f"python3 -um lmcache_server.server {self.host} {self.port} {self.remote_device}"
         self.handle = run_command(cmd, self.stdout_log, self.stderr_log, detach=True)
         self.started = True
 
@@ -136,9 +142,10 @@ def wait_until_ready(self, timeout = 60) -> bool:
             return True
 
         if not self.is_healthy():
+            logger.error("LMCacheServer is dead!")
             return False
 
-        self._monitor_file_output([self.stdout_log, self.stderr_log], "Server started at", timeout=timeout)
+        return self._monitor_file_output([self.stdout_log, self.stderr_log], "Server started at", timeout=timeout)
 
         return True
 
@@ -185,6 +192,17 @@ def get_or_create(
 
         return cls._instances[instance_id]
 
+    @classmethod
+    def close_servers(cls):
+        """
+        Close and remove all the active lmcache servers
+        """
+        for instance_id, instance in cls._instances.items():
+            instance.close()
+        cls._instances = {}
+        cls._engine_types = {}
+
+
 
 class LocalVllmBootstrapper(Bootstrapper):
     """
@@ -200,11 +218,12 @@ def __init__(self, config: BootstrapConfig, log_dir = "/tmp"):
         self.lmcache_server_handler = LMCacheServerManager.get_or_create(config)
 
     def get_cmdline(self) -> str:
-        return f"python3 -m vllm.entrypoints.openai.api_server {self.config.vllm_config.cmdargs()} {self.config.vllm_optional_config.cmdargs()} {self.config.lmcache_config.cmdargs()}"
+        extra_args = "--trust-remote-code"
+        return f"python3 -m vllm.entrypoints.openai.api_server {self.config.vllm_config.cmdargs()} {self.config.vllm_optional_config.cmdargs()} {self.config.lmcache_config.cmdargs()} {extra_args}"
 
     def start(self):
         self.lmcache_server_handler.start()
-        self.lmcache_server_handler.wait_until_ready(timeout = 5)
+        self.lmcache_server_handler.wait_until_ready(timeout = 10)
 
         self.handle = run_command(
                 self.command, 
@@ -214,6 +233,7 @@ def start(self):
     def wait_until_ready(self, timeout = 60) -> bool:
         # Try reading the log file to see if the server is ready
         if not self.is_healthy():
+            logger.error(f"VLLM or lmcache server is dead!")
             return False
 
         if not os.path.exists(self.stdout_log):
@@ -224,6 +244,7 @@ def wait_until_ready(self, timeout = 60) -> bool:
 
     def is_healthy(self) -> bool:
         if not self.lmcache_server_handler.is_healthy():
+            logger.warn(f"LMCache Server is dead during vLLM's check!")
             return False 
 
         if self.handle is not None:

configs.py

@@ -51,6 +51,8 @@ class LMCacheConfig(Config):
     # Path to the lmcache configuration
     config_path: str
 
+    remote_device: Optional[str] = None
+
     def cmdargs(self) -> str:
         return f"--lmcache-config-file {self.config_path}" if self.config_path is not None else ""
 

configs/lmcache_remote_cachegen_pipeline.yaml

@@ -1,6 +1,6 @@
 chunk_size: 256
 local_device: null
-remote_url: "lm://localhost:65431"
+remote_url: "lm://localhost:65430"
 remote_serde: "cachegen"
 
 # Whether retrieve() is pipelined or not

configs/lmcache_remote_safetensor_pipeline.yaml

@@ -1,6 +1,6 @@
 chunk_size: 256
 local_device: null
-remote_url: "lm://localhost:65431"
+remote_url: "lm://localhost:65430"
 remote_serde: "safetensor"
 
 # Whether retrieve() is pipelined or not

driver.py

@@ -7,7 +7,7 @@
 
 from configs import BootstrapConfig, WorkloadConfig, Usecase
 from test_cases import TestCase
-from bootstrapper import CreateBootstrapper, Bootstrapper
+from bootstrapper import CreateBootstrapper, Bootstrapper, LMCacheServerManager
 from workload import CreateWorkloadGenerator, Request
 from utils import read_gpu_memory
 
@@ -106,27 +106,32 @@ def execute_openai_request(request: Request, model: str, client: openai.Client)
     #time.sleep(t)
     #return t, t
 
-    chat_completion = client.chat.completions.create(
-            messages = messages,
-            model = model,
-            temperature = 0,
-            stream = True,
-        )
-
-
-    start_time = time.perf_counter()
-    first_token_time = None
-    ntokens = 0
-    for chunk in chat_completion:
-        chunk_message = chunk.choices[0].delta.content
-        if chunk_message is not None:
-            if first_token_time is None:
-                first_token_time = time.perf_counter()
-            ntokens += 1
-    end_time = time.perf_counter()
-
-    ttft = first_token_time - start_time
-    throughput = ntokens / (end_time - first_token_time)
+    
+    try:
+        chat_completion = client.chat.completions.create(
+                messages = messages,
+                model = model,
+                temperature = 0,
+                stream = True,
+            )
+
+        start_time = time.perf_counter()
+        first_token_time = None
+        ntokens = 0
+        for chunk in chat_completion:
+            chunk_message = chunk.choices[0].delta.content
+            if chunk_message is not None:
+                if first_token_time is None:
+                    first_token_time = time.perf_counter()
+                ntokens += 1
+        end_time = time.perf_counter()
+
+        ttft = first_token_time - start_time
+        throughput = ntokens / (end_time - first_token_time)
+    except Exception as e:
+        logger.error(f"OpenAI request failed: {e}")
+        return -1, -1
+
     return ttft, throughput
 
 
@@ -151,6 +156,9 @@ def cleanup(bootstrappers: List[Bootstrapper]):
         logger.info("Cleanning up the engine processes")
         for bootstrapper in bootstrappers:
             bootstrapper.close()
+        LMCacheServerManager.close_servers()
+
+    logger.info(f"Running experiment: {workload_config.desc()} {usecase}")
 
     # Create the workloads
     workload_generators = [CreateWorkloadGenerator(workload_config, usecase) for _ in engine_configs]
@@ -165,7 +173,7 @@ def cleanup(bootstrappers: List[Bootstrapper]):
     try:
         # Wait for the engines to be ready
         for bootstrapper in bootstrappers:
-            ready = bootstrapper.wait_until_ready(timeout = 120)
+            ready = bootstrapper.wait_until_ready(timeout = 180)
             if not ready:
                 logger.error(f"Engine {bootstrapper} is not ready")
                 cleanup(bootstrappers)
@@ -180,7 +188,7 @@ def cleanup(bootstrappers: List[Bootstrapper]):
         executor.schedule_requests(workloads, clients, models)
         results = executor.execute_all()
 
-        print(results)
+        #print(results)
 
         # Read GPU memory utilization
         gpu_usage = read_gpu_memory()

log.py

@@ -1,12 +1,53 @@
 import logging
 from logging import Logger
 
-logging.basicConfig(
-        format='\033[33m%(levelname)s: \033[0m%(message)s [%(asctime)s.%(msecs)03d]', 
-        level=logging.INFO,
-    )
+#logging.basicConfig(
+#        format='\033[33m%(levelname)s: \033[0m%(message)s [%(asctime)s.%(msecs)03d]', 
+#        level=logging.INFO,
+#    )
+
+def build_format(color):
+    reset = "\x1b[0m"
+    underline = "\x1b[3m"
+    return f"{color}[%(asctime)s] %(levelname)s:{reset} %(message)s {underline}(%(filename)s:%(lineno)d:%(name)s){reset}"
+
+class CustomFormatter(logging.Formatter):
+
+    grey = "\x1b[1m"
+    green = "\x1b[32;20m"
+    yellow = "\x1b[33;20m"
+    red = "\x1b[31;20m"
+    bold_red = "\x1b[31;1m"
+    reset = "\x1b[0m"
+
+    FORMATS = {
+        logging.DEBUG: build_format(grey),
+        logging.INFO: build_format(green),
+        logging.WARNING: build_format(yellow),
+        logging.ERROR: build_format(red),
+        logging.CRITICAL: build_format(bold_red),
+    }
+
+    def format(self, record):
+        log_fmt = self.FORMATS.get(record.levelno)
+        formatter = logging.Formatter(log_fmt)
+        return formatter.format(record)
 
 def init_logger(name: str) -> Logger:
     logger = logging.getLogger(name)
+
+    ch = logging.StreamHandler()
+    ch.setLevel(logging.DEBUG)
+    ch.setFormatter(CustomFormatter())
+    logger.addHandler(ch)
     logger.setLevel(logging.DEBUG)
+
     return logger
+
+if __name__ == "__main__":
+    logger = init_logger(__name__)
+    logger.debug("Debug message")
+    logger.info("Info message")
+    logger.warning("Warning message")
+    logger.error("Error message")
+    logger.critical("Critical message")

main.py

@@ -26,7 +26,7 @@ def CreateSingleLocalBootstrapConfig(
         envs = {"CUDA_VISIBLE_DEVICES": str(gpu_id)}
     )
 
-def CreateDummyExperiment(num_requests, context_length, gap_between_requests = 10):
+def CreateDummyExperiment(num_requests, context_length, gap_between_requests = 8):
     """
     Create some requests for DUMMY usecase
     The query length will be 16
@@ -86,7 +86,7 @@ def test_lmcache_local_cpu() -> pd.DataFrame:
 
     # Experiments: 8K, 16K, 24K shared context, each experiments has 5 queries
     lengths = [8192, 16384, 24576]
-    experiments = [CreateDummyExperiment(5, length) for length in lengths]
+    experiments = [CreateDummyExperiment(10, length) for length in lengths]
 
     test_case = TestCase(
             experiments = experiments,
@@ -103,7 +103,7 @@ def test_lmcache_local_disk() -> pd.DataFrame:
 
     # Experiments: 8K, 16K, 24K shared context, each experiments has 5 queries
     lengths = [8192, 16384, 24576]
-    experiments = [CreateDummyExperiment(5, length) for length in lengths]
+    experiments = [CreateDummyExperiment(10, length) for length in lengths]
 
     test_case = TestCase(
             experiments = experiments,
@@ -120,7 +120,7 @@ def test_lmcache_remote_cachegen() -> pd.DataFrame:
 
     # Experiments: 8K, 16K, 24K shared context, each experiments has 5 queries
     lengths = [8192, 16384, 24576]
-    experiments = [CreateDummyExperiment(5, length) for length in lengths]
+    experiments = [CreateDummyExperiment(10, length) for length in lengths]
 
     test_case = TestCase(
             experiments = experiments,
@@ -137,7 +137,26 @@ def test_lmcache_remote_safetensor() -> pd.DataFrame:
 
     # Experiments: 8K, 16K, 24K shared context, each experiments has 5 queries
     lengths = [8192, 16384, 24576]
-    experiments = [CreateDummyExperiment(5, length) for length in lengths]
+    experiments = [CreateDummyExperiment(10, length) for length in lengths]
+
+    test_case = TestCase(
+            experiments = experiments,
+            engines = [config1, config2])
+
+    # Run test case
+    final_result = run_test_case(test_case)
+    return final_result
+
+def test_lmcache_remote_disk() -> pd.DataFrame:
+    # Start two servers: with lmcache and without lmcache
+    config1 = CreateSingleLocalBootstrapConfig(8000, 0, "mistralai/Mistral-7B-Instruct-v0.2", "configs/lmcache_remote_cachegen.yaml")
+    config2 = CreateSingleLocalBootstrapConfig(8001, 1, "mistralai/Mistral-7B-Instruct-v0.2", None)
+
+    config1.lmcache_config.remote_device = "/local/lmcache-tests/lmcache-server"
+
+    # Experiments: 8K, 16K, 24K shared context, each experiments has 5 queries
+    lengths = [8192, 16384, 24576]
+    experiments = [CreateDummyExperiment(10, length) for length in lengths]
 
     test_case = TestCase(
             experiments = experiments,
@@ -147,11 +166,46 @@ def test_lmcache_remote_safetensor() -> pd.DataFrame:
     final_result = run_test_case(test_case)
     return final_result
 
+def test_lmcache_chatglm() -> pd.DataFrame:
+    # Start two servers: with lmcache and without lmcache
+    config1 = CreateSingleLocalBootstrapConfig(8000, 0, "THUDM/glm-4-9b-chat", "configs/lmcache_remote_cachegen.yaml")
+    config2 = CreateSingleLocalBootstrapConfig(8001, 1, "THUDM/glm-4-9b-chat", None)
+
+    config1.vllm_config.tensor_parallel_size = 2
+    config1.vllm_config.gpu_memory_utilization = 0.8
+    config1.envs = {}
+    config1.vllm_optional_config["trust_remote_code"] = ""
+
+    config2.vllm_config.tensor_parallel_size = 2
+    config2.vllm_config.gpu_memory_utilization = 0.8
+    config2.envs = {}
+    config2.vllm_optional_config["trust_remote_code"] = ""
+
+    # Experiments: 8K, 16K, 24K shared context, each experiments has 5 queries
+    lengths = [8192, 16384, 24576]
+    experiments = [CreateDummyExperiment(10, length) for length in lengths]
+
+    test_case1 = TestCase(
+            experiments = experiments,
+            engines = [config1])
+
+    test_case2 = TestCase(
+            experiments = experiments,
+            engines = [config2])
+
+    # Run test case
+    final_result1 = run_test_case(test_case1)
+    final_result2 = run_test_case(test_case2)
+    final_result1["engine_id"] = 0
+    final_result2["engine_id"] = 1
+    return pd.concat([final_result1, final_result2])
 
 if __name__ == "__main__":
     print("Start running test cases")
     #wrapped_runner(test_lmcache_local_gpu, "outputs/test_lmcache_local_gpu.csv")
     #wrapped_runner(test_lmcache_local_cpu, "outputs/test_lmcache_local_cpu.csv")
-    wrapped_runner(test_lmcache_local_disk, "outputs/test_lmcache_local_disk.csv")
-    wrapped_runner(test_lmcache_remote_cachegen, "outputs/test_lmcache_remote_cachegen.csv")
-    wrapped_runner(test_lmcache_remote_cachegen, "outputs/test_lmcache_remote_safetensor.csv")
+    #wrapped_runner(test_lmcache_local_disk, "outputs/test_lmcache_local_disk.csv")
+    #wrapped_runner(test_lmcache_remote_safetensor, "outputs/test_lmcache_remote_safetensor.csv")
+    #wrapped_runner(test_lmcache_remote_cachegen, "outputs/test_lmcache_remote_cachegen.csv")
+    #wrapped_runner(test_lmcache_remote_disk, "outputs/test_lmcache_remote_disk.csv")
+    wrapped_runner(test_lmcache_chatglm, "outputs/test_lmcache_chatglm.csv")