speedy-utils 1.0.5__py3-none-any.whl → 1.0.11__py3-none-any.whl

llm_utils/scripts/vllm_serve.py

@@ -9,19 +9,17 @@ Serve a base model:
 svllm serve --model MODEL_NAME --gpus GPU_GROUPS
 
 Add a LoRA to a served model:
-svllm add-lora --lora LORA_NAME LORA_PATH --host_port host:port (if add then the port must be specify)
+svllm add-lora --lora LORA_NAME LORA_PATH --host_port host:port
+(if add then the port must be specify)
 """
 
-from glob import glob
 import os
 import subprocess
-import time
-from typing import List, Literal, Optional
-from fastcore.script import call_parse
-from loguru import logger
+from typing import List, Optional
 import argparse
 import requests
 import openai
+from loguru import logger
 
 from speedy_utils.common.utils_io import load_by_ext
 
@@ -32,63 +30,22 @@ HF_HOME: str = os.environ.get("HF_HOME", os.path.expanduser("~/.cache/huggingfac
 logger.info(f"LORA_DIR: {LORA_DIR}")
 
 
-def model_list(host_port, api_key="abc"):
+def model_list(host_port: str, api_key: str = "abc") -> None:
+    """List models from the vLLM server."""
     client = openai.OpenAI(base_url=f"http://{host_port}/v1", api_key=api_key)
     models = client.models.list()
     for model in models:
         print(f"Model ID: {model.id}")
 
 
-def kill_existing_vllm(vllm_binary: Optional[str] = None) -> None:
-    """Kill selected vLLM processes using fzf."""
-    if not vllm_binary:
-        vllm_binary = get_vllm()
-
-    # List running vLLM processes
-    result = subprocess.run(
-        f"ps aux | grep {vllm_binary} | grep -v grep",
-        shell=True,
-        capture_output=True,
-        text=True,
-    )
-    processes = result.stdout.strip().split("\n")
-
-    if not processes or processes == [""]:
-        print("No running vLLM processes found.")
-        return
-
-    # Use fzf to select processes to kill
-    fzf = subprocess.Popen(
-        ["fzf", "--multi"],
-        stdin=subprocess.PIPE,
-        stdout=subprocess.PIPE,
-        text=True,
-    )
-    selected, _ = fzf.communicate("\n".join(processes))
-
-    if not selected:
-        print("No processes selected.")
-        return
-
-    # Extract PIDs and kill selected processes
-    pids = [line.split()[1] for line in selected.strip().split("\n")]
-    for pid in pids:
-        subprocess.run(
-            f"kill -9 {pid}",
-            shell=True,
-            stdout=subprocess.DEVNULL,
-            stderr=subprocess.DEVNULL,
-        )
-    print(f"Killed processes: {', '.join(pids)}")
-
-
 def add_lora(
     lora_name_or_path: str,
     host_port: str,
     url: str = "http://HOST:PORT/v1/load_lora_adapter",
     served_model_name: Optional[str] = None,
-    lora_module: Optional[str] = None,  # Added parameter
+    lora_module: Optional[str] = None,
 ) -> dict:
+    """Add a LoRA adapter to a running vLLM server."""
     url = url.replace("HOST:PORT", host_port)
     headers = {"Content-Type": "application/json"}
 
@@ -96,15 +53,12 @@ def add_lora(
         "lora_name": served_model_name,
         "lora_path": os.path.abspath(lora_name_or_path),
     }
-    if lora_module:  # Include lora_module if provided
+    if lora_module:
         data["lora_module"] = lora_module
     logger.info(f"{data=}, {headers}, {url=}")
-    # logger.warning(f"Failed to unload LoRA adapter: {str(e)}")
     try:
-        response = requests.post(url, headers=headers, json=data)
+        response = requests.post(url, headers=headers, json=data, timeout=10)
         response.raise_for_status()
-
-        # Handle potential non-JSON responses
         try:
             return response.json()
         except ValueError:
@@ -116,113 +70,100 @@ def add_lora(
                     else "Request completed with empty response"
                 ),
             }
-
     except requests.exceptions.RequestException as e:
         logger.error(f"Request failed: {str(e)}")
         return {"error": f"Request failed: {str(e)}"}
 
 
-def unload_lora(lora_name, host_port):
+def unload_lora(lora_name: str, host_port: str) -> Optional[dict]:
+    """Unload a LoRA adapter from a running vLLM server."""
     try:
         url = f"http://{host_port}/v1/unload_lora_adapter"
         logger.info(f"{url=}")
         headers = {"Content-Type": "application/json"}
         data = {"lora_name": lora_name}
         logger.info(f"Unloading LoRA adapter: {data=}")
-        response = requests.post(url, headers=headers, json=data)
+        response = requests.post(url, headers=headers, json=data, timeout=10)
         response.raise_for_status()
         logger.success(f"Unloaded LoRA adapter: {lora_name}")
     except requests.exceptions.RequestException as e:
         return {"error": f"Request failed: {str(e)}"}
 
 
-def serve(
-    model: str,
-    gpu_groups: str,
-    served_model_name: Optional[str] = None,
-    port_start: int = 8155,
-    gpu_memory_utilization: float = 0.93,
-    dtype: str = "bfloat16",
-    max_model_len: int = 8192,
-    enable_lora: bool = False,
-    is_bnb: bool = False,
-    eager: bool = False,
-    lora_modules: Optional[List[str]] = None,  # Updated type
-) -> None:
-    """Main function to start or kill vLLM containers."""
-
+def serve(args) -> None:
     """Start vLLM containers with dynamic args."""
     print("Starting vLLM containers...,")
-    gpu_groups_arr: List[str] = gpu_groups.split(",")
-    VLLM_BINARY: str = get_vllm()
-    if enable_lora:
-        VLLM_BINARY = "VLLM_ALLOW_RUNTIME_LORA_UPDATING=True " + VLLM_BINARY
-
-    # Auto-detect quantization based on model name if not explicitly set
-    if not is_bnb and model and ("bnb" in model.lower() or "4bit" in model.lower()):
-        is_bnb = True
-        print(f"Auto-detected quantization for model: {model}")
-
-    # Set environment variables for LoRA if needed
-    if enable_lora:
+    gpu_groups_arr: List[str] = args.gpu_groups.split(",")
+    vllm_binary: str = get_vllm()
+    if args.enable_lora:
+        vllm_binary = "VLLM_ALLOW_RUNTIME_LORA_UPDATING=True " + vllm_binary
+
+    if (
+        not args.bnb
+        and args.model
+        and ("bnb" in args.model.lower() or "4bit" in args.model.lower())
+    ):
+        args.bnb = True
+        print(f"Auto-detected quantization for model: {args.model}")
+
+    if args.enable_lora:
         os.environ["VLLM_ALLOW_RUNTIME_LORA_UPDATING"] = "True"
         print("Enabled runtime LoRA updating")
 
     for i, gpu_group in enumerate(gpu_groups_arr):
-        port = port_start + i
+        port = int(args.host_port.split(":")[-1]) + i
         gpu_group = ",".join([str(x) for x in gpu_group])
         tensor_parallel = len(gpu_group.split(","))
 
         cmd = [
             f"CUDA_VISIBLE_DEVICES={gpu_group}",
-            VLLM_BINARY,
+            vllm_binary,
             "serve",
-            model,
+            args.model,
             "--port",
             str(port),
             "--tensor-parallel",
             str(tensor_parallel),
             "--gpu-memory-utilization",
-            str(gpu_memory_utilization),
+            str(args.gpu_memory_utilization),
             "--dtype",
-            dtype,
+            args.dtype,
             "--max-model-len",
-            str(max_model_len),
+            str(args.max_model_len),
             "--enable-prefix-caching",
             "--disable-log-requests",
             "--uvicorn-log-level critical",
         ]
         if HF_HOME:
-            # insert
             cmd.insert(0, f"HF_HOME={HF_HOME}")
-        if eager:
+        if args.eager:
             cmd.append("--enforce-eager")
 
-        if served_model_name:
-            cmd.extend(["--served-model-name", served_model_name])
+        if args.served_model_name:
+            cmd.extend(["--served-model-name", args.served_model_name])
 
-        if is_bnb:
+        if args.bnb:
             cmd.extend(
                 ["--quantization", "bitsandbytes", "--load-format", "bitsandbytes"]
             )
 
-        if enable_lora:
+        if args.enable_lora:
             cmd.extend(["--fully-sharded-loras", "--enable-lora"])
 
-        if lora_modules:
-            # for lora_module in lora_modules:
-            # len must be even and we will join tuple with `=`
-            assert len(lora_modules) % 2 == 0, "lora_modules must be even"
-            # lora_modulle = [f'{name}={module}' for name, module in zip(lora_module[::2], lora_module[1::2])]
-            # import ipdb;ipdb.set_trace()
+        if args.lora_modules:
+            assert len(args.lora_modules) % 2 == 0, "lora_modules must be even"
             s = ""
-            for i in range(0, len(lora_modules), 2):
-                name = lora_modules[i]
-                module = lora_modules[i + 1]
+            for i in range(0, len(args.lora_modules), 2):
+                name = args.lora_modules[i]
+                module = args.lora_modules[i + 1]
                 s += f"{name}={module} "
-
             cmd.extend(["--lora-modules", s])
-        # add kwargs
+
+        if hasattr(args, "enable_reasoning") and args.enable_reasoning:
+            cmd.extend(["--enable-reasoning", "--reasoning-parser", "deepseek_r1"])
+            # Add VLLM_USE_V1=0 to the environment for reasoning mode
+            cmd.insert(0, "VLLM_USE_V1=0")
+
         final_cmd = " ".join(cmd)
         log_file = f"/tmp/vllm_{port}.txt"
         final_cmd_with_log = f'"{final_cmd} 2>&1 | tee {log_file}"'
@@ -235,14 +176,15 @@ def serve(
         os.system(run_in_tmux)
 
 
-def get_vllm():
-    VLLM_BINARY = subprocess.check_output("which vllm", shell=True, text=True).strip()
-    VLLM_BINARY = os.getenv("VLLM_BINARY", VLLM_BINARY)
-    logger.info(f"vLLM binary: {VLLM_BINARY}")
+def get_vllm() -> str:
+    """Get the vLLM binary path."""
+    vllm_binary = subprocess.check_output("which vllm", shell=True, text=True).strip()
+    vllm_binary = os.getenv("VLLM_BINARY", vllm_binary)
+    logger.info(f"vLLM binary: {vllm_binary}")
     assert os.path.exists(
-        VLLM_BINARY
-    ), f"vLLM binary not found at {VLLM_BINARY}, please set VLLM_BINARY env variable"
-    return VLLM_BINARY
+        vllm_binary
+    ), f"vLLM binary not found at {vllm_binary}, please set VLLM_BINARY env variable"
+    return vllm_binary
 
 
 def get_args():
@@ -330,6 +272,9 @@ def get_args():
         type=str,
         help="List of LoRA modules in the format lora_name lora_module",
     )
+    parser.add_argument(
+        "--enable-reasoning", action="store_true", help="Enable reasoning"
+    )
     return parser.parse_args()
 
 
@@ -371,23 +316,8 @@ def main():
                 logger.info(f"Model name from LoRA config: {model_name}")
                 args.model = model_name
         # port_start from hostport
-        port_start = int(args.host_port.split(":")[-1])
-        serve(
-            args.model,
-            args.gpu_groups,
-            args.served_model_name,
-            port_start,
-            args.gpu_memory_utilization,
-            args.dtype,
-            args.max_model_len,
-            args.enable_lora,
-            args.bnb,
-            args.eager,
-            args.lora_modules,
-        )
+        serve(args)
 
-    elif args.mode == "kill":
-        kill_existing_vllm(args.vllm_binary)
     elif args.mode == "add_lora":
         if args.lora:
             lora_name, lora_path = args.lora

{speedy_utils-1.0.5.dist-info → speedy_utils-1.0.11.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: speedy-utils
-Version: 1.0.5
+Version: 1.0.11
 Summary: Fast and easy-to-use package for data science
 Author: AnhVTH
 Author-email: anhvth.226@gmail.com
@@ -54,9 +54,6 @@ Description-Content-Type: text/markdown
   - [Data Manipulation](#data-manipulation)
   - [Utility Functions](#utility-functions)
 - [Testing](#testing)
-- [Deployment](#deployment)
-- [Contributing](#contributing)
-- [License](#license)
 
 ## Features
 
@@ -84,6 +81,18 @@ cd speedy-utils
 pip install .
 ```
 
+## Updating from previous versions
+
+To update from previous versions or switch to v1.x, first uninstall any old
+packages, then install the latest version:
+
+```bash
+pip uninstall speedy_llm_utils speedy_utils
+pip install -e ./  # for local development
+# or
+pip install speedy_utils -U  # for PyPI upgrade
+```
+
 ## Usage
 
 Below are examples demonstrating how to utilize various features of **Speedy Utils**.
@@ -262,7 +271,6 @@ Ensure all dependencies are installed before running tests:
 pip install -r requirements.txt
 ```
 
-
 Run the script to parse and display the arguments:
 
 ```bash

{speedy_utils-1.0.5.dist-info → speedy_utils-1.0.11.dist-info}/RECORD

@@ -1,11 +1,14 @@
-llm_utils/__init__.py,sha256=2g0XXQLj9WsGmWE4UQ9YaaYGfF5ZUyvc1hfR2OQIVSo,679
-llm_utils/chat_format.py,sha256=ZY2HYv3FPL2xiMxbbO-huIwT5LZrcJm_if_us-2eSZ4,15094
+llm_utils/__init__.py,sha256=ibEVUPkL11M4htL-3uXkSyyUZiIO-TZD6IzWVmi8QYw,697
+llm_utils/chat_format/__init__.py,sha256=8dBIUqFJvkgQYedxBtcyxt-4tt8JxAKVap2JlTXmgaM,737
+llm_utils/chat_format/display.py,sha256=a3zWzo47SUf4i-uic-dwf-vxtu6gZWLbnJrszjjZjQ8,9801
+llm_utils/chat_format/transform.py,sha256=328V18FOgRQzljAl9Mh8NF4Tl-N3cZZIPmAwHQspXCY,5461
+llm_utils/chat_format/utils.py,sha256=xTxN4HrLHcRO2PfCTR43nH1M5zCa7v0kTTdzAcGkZg0,1229
 llm_utils/group_messages.py,sha256=wyiZzs7O8yK2lyIakV2x-1CrrWVT12sjnP1vVnmPet4,3606
-llm_utils/lm/__init__.py,sha256=a4N_hh0JuB2FEEzxNiE8GwTIbd8enm6mj_lw4TBxreI,76
-llm_utils/lm/base_lm.py,sha256=ZbuLagAPZsA4Oa2eIkoDnsldqMf4pKtz442LtWgMrMk,10704
+llm_utils/lm/__init__.py,sha256=vXFILZLBmmpg39cy5XniQPSMzoFQCE3wdfz39EtqDKU,71
+llm_utils/lm/lm.py,sha256=4bEo4nnyCi_ybTOYfzrJz9AwpxJNkzRFAUPq7KpBklw,16695
 llm_utils/lm/utils.py,sha256=-fDNueiXKQI6RDoNHJYNyORomf2XlCf2doJZ3GEV2Io,4762
-llm_utils/scripts/vllm_load_balancer.py,sha256=uSjGd_jOmI9W9eVOhiOXbeUnZkQq9xG4bCVzhmpupcA,16096
-llm_utils/scripts/vllm_serve.py,sha256=uFS5kNXZ7kZ9rQms63LnliGEVV3rATT6dEppGTgoR0s,13910
+llm_utils/scripts/vllm_load_balancer.py,sha256=MgMnnoKWJQc-l2fspUSkyA9wxL1RkXd7wdBLJNQBlr4,17384
+llm_utils/scripts/vllm_serve.py,sha256=LlrkwfWLxdMDhfOJ-eL1VJnA4AY1Beh_cI8U6l9Xl-A,11975
 speedy_utils/__init__.py,sha256=I2bSfDIE9yRF77tnHW0vqfExDA2m1gUx4AH8C9XmGtg,1707
 speedy_utils/all.py,sha256=A9jiKGjo950eg1pscS9x38OWAjKGyusoAN5mrfweY4E,3090
 speedy_utils/common/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -21,7 +24,7 @@ speedy_utils/multi_worker/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJW
 speedy_utils/multi_worker/process.py,sha256=XwQlffxzRFnCVeKjDNBZDwFfUQHiJiuFA12MRGJVru8,6708
 speedy_utils/multi_worker/thread.py,sha256=9pXjvgjD0s0Hp0cZ6I3M0ndp1OlYZ1yvqbs_bcun_Kw,12775
 speedy_utils/scripts/mpython.py,sha256=ZzkBWI5Xw3vPoMx8xQt2x4mOFRjtwWqfvAJ5_ngyWgw,3816
-speedy_utils-1.0.5.dist-info/METADATA,sha256=QAwtVoX05-q-aPwtJN7fg6AvFRF9MF47ycEwsq-7z-0,7165
-speedy_utils-1.0.5.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-speedy_utils-1.0.5.dist-info/entry_points.txt,sha256=fsv8_lMg62BeswoUHrqfj2u6q2l4YcDCw7AgQFg6GRw,61
-speedy_utils-1.0.5.dist-info/RECORD,,
+speedy_utils-1.0.11.dist-info/METADATA,sha256=F48tr0hmL3k-r9O2tPbUdfbBU5JHnwxVGB547eQXElU,7392
+speedy_utils-1.0.11.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+speedy_utils-1.0.11.dist-info/entry_points.txt,sha256=rP43satgw1uHcKUAlmVxS-MTAQImL-03-WwLIB5a300,165
+speedy_utils-1.0.11.dist-info/RECORD,,

speedy_utils-1.0.11.dist-info/entry_points.txt

@@ -0,0 +1,5 @@
+[console_scripts]
+mpython=speedy_utils.scripts.mpython:main
+svllm=llm_utils.scripts.vllm_serve:main
+svllm-lb=llm_utils.scripts.vllm_load_balancer:run_load_balancer
+