"""Process management for vLLM/SGLang.""" from __future__ import annotations import asyncioimport jsonimport osimport subprocessfrom pathlib import Pathfrom typing import List, Optional, Tuple import psutil from .config import settingsfrom .models import Backend, ProcessInfo, Recipe def _extract_flag(cmdline: List[str], flag: str) -> Optional[str]: """Extract value of a CLI flag.""" for i, arg in enumerate(cmdline): if arg == flag and i + 1 < len(cmdline): return cmdline[i + 1] return None def _is_inference_process(cmdline: List[str]) -> Optional[str]: """Check if cmdline is vLLM or SGLang, return backend name.""" if not cmdline: return None joined = " ".join(cmdline) if "vllm.entrypoints.openai.api_server" in joined: return "vllm" if len(cmdline) >= 2 and cmdline[0].endswith("vllm") and cmdline[1] == "serve": return "vllm" if "sglang.launch_server" in joined: return "sglang" return None def find_inference_process(port: int) -> Optional[ProcessInfo]: """Find running inference process on given port.""" for proc in psutil.process_iter(["pid", "cmdline"]): try: cmdline = proc.info.get("cmdline") or [] backend = _is_inference_process(cmdline) if not backend: continue p = _extract_flag(cmdline, "--port") if p is None or int(p) != port: continue # Extract model path model_path = _extract_flag(cmdline, "--model") or _extract_flag(cmdline, "--model-path") if not model_path and len(cmdline) >= 3 and cmdline[1] == "serve": model_path = cmdline[2] if not cmdline[2].startswith("-") else None return ProcessInfo( pid=proc.info["pid"], backend=backend, model_path=model_path, port=port, served_model_name=_extract_flag(cmdline, "--served-model-name"), ) except (psutil.NoSuchProcess, psutil.AccessDenied, ValueError): continue return None async def kill_process(pid: int, force: bool = False) -> bool: """Kill process and its children.""" try: proc = psutil.Process(pid) except psutil.NoSuchProcess: return True # Kill children first for child in proc.children(recursive=True): try: child.kill() except psutil.NoSuchProcess: pass # Terminate main process try: proc.terminate() proc.wait(timeout=10) except psutil.TimeoutExpired: if force: proc.kill() except psutil.NoSuchProcess: pass await asyncio.sleep(1) return True def build_vllm_command(recipe: Recipe) -> List[str]: """Build vLLM launch command.""" if recipe.python_path: cmd = [recipe.python_path, "-m", "vllm.entrypoints.openai.api_server"] else: cmd = ["vllm", "serve"] cmd.extend([recipe.model_path, "--host", recipe.host, "--port", str(recipe.port)]) if recipe.served_model_name: cmd.extend(["--served-model-name", recipe.served_model_name]) if recipe.tensor_parallel_size > 1: cmd.extend(["--tensor-parallel-size", str(recipe.tensor_parallel_size)]) if recipe.pipeline_parallel_size > 1: cmd.extend(["--pipeline-parallel-size", str(recipe.pipeline_parallel_size)]) cmd.extend(["--max-model-len", str(recipe.max_model_len)]) cmd.extend(["--gpu-memory-utilization", str(recipe.gpu_memory_utilization)]) cmd.extend(["--max-num-seqs", str(recipe.max_num_seqs)]) if recipe.kv_cache_dtype != "auto": cmd.extend(["--kv-cache-dtype", recipe.kv_cache_dtype]) if recipe.trust_remote_code: cmd.append("--trust-remote-code") if recipe.tool_call_parser: cmd.extend(["--tool-call-parser", recipe.tool_call_parser, "--enable-auto-tool-choice"]) if recipe.quantization: cmd.extend(["--quantization", recipe.quantization]) if recipe.dtype: cmd.extend(["--dtype", recipe.dtype]) # Extra args for key, value in recipe.extra_args.items(): flag = f"--{key.replace('_', '-')}" if value is True: cmd.append(flag) elif value not in (False, None): if isinstance(value, (dict, list)): cmd.extend([flag, json.dumps(value)]) else: cmd.extend([flag, str(value)]) return cmd def build_sglang_command(recipe: Recipe) -> List[str]: """Build SGLang launch command.""" python = recipe.python_path or settings.sglang_python or "python" cmd = [python, "-m", "sglang.launch_server"] cmd.extend(["--model-path", recipe.model_path]) cmd.extend(["--host", recipe.host, "--port", str(recipe.port)]) if recipe.served_model_name: cmd.extend(["--served-model-name", recipe.served_model_name]) if recipe.tensor_parallel_size > 1: cmd.extend(["--tp", str(recipe.tensor_parallel_size)]) cmd.extend(["--context-length", str(recipe.max_model_len)]) cmd.extend(["--mem-fraction-static", str(recipe.gpu_memory_utilization)]) if recipe.trust_remote_code: cmd.append("--trust-remote-code") if recipe.quantization: cmd.extend(["--quantization", recipe.quantization]) for key, value in recipe.extra_args.items(): flag = f"--{key.replace('_', '-')}" if value is True: cmd.append(flag) elif value not in (False, None): cmd.extend([flag, str(value)]) return cmd async def launch_model(recipe: Recipe) -> Tuple[bool, Optional[int], str]: """Launch inference server with recipe config.""" recipe.port = settings.inference_port # Override with configured port if recipe.backend == Backend.SGLANG: cmd = build_sglang_command(recipe) else: cmd = build_vllm_command(recipe) log_file = Path(f"/tmp/vllm_{recipe.id}.log") env = os.environ.copy() try: with open(log_file, "w") as log: proc = subprocess.Popen( cmd, stdout=log, stderr=subprocess.STDOUT, env=env, start_new_session=True, ) await asyncio.sleep(3) if proc.poll() is not None: tail = log_file.read_text()[-500:] if log_file.exists() else "" return False, None, f"Process exited early: {tail}" return True, proc.pid, str(log_file) except Exception as e: return False, None, str(e) async def evict_model(force: bool = False) -> Optional[int]: """Stop current running model.""" current = find_inference_process(settings.inference_port) if not current: return None await kill_process(current.pid, force=force) return current.pid async def switch_model(recipe: Recipe, force: bool = False) -> Tuple[bool, Optional[int], str]: """Switch to a new model (evict current + launch new).""" await evict_model(force=force) await asyncio.sleep(2) return await launch_model(recipe)
"""Process management for vLLM/SGLang."""
from __future__ import annotations
import asyncio
import json
import os
import subprocess
from pathlib import Path
from typing import List, Optional, Tuple
import psutil
from .config import settings
from .models import Backend, ProcessInfo, Recipe
def _extract_flag(cmdline: List[str], flag: str) -> Optional[str]:
"""Extract value of a CLI flag."""
for i, arg in enumerate(cmdline):
if arg == flag and i + 1 < len(cmdline):
return cmdline[i + 1]
return None
def _is_inference_process(cmdline: List[str]) -> Optional[str]:
"""Check if cmdline is vLLM or SGLang, return backend name."""
if not cmdline:
return None
joined = " ".join(cmdline)
if "vllm.entrypoints.openai.api_server" in joined:
return "vllm"
if len(cmdline) >= 2 and cmdline[0].endswith("vllm") and cmdline[1] == "serve":
return "vllm"
if "sglang.launch_server" in joined:
return "sglang"
return None
def find_inference_process(port: int) -> Optional[ProcessInfo]:
"""Find running inference process on given port."""
for proc in psutil.process_iter(["pid", "cmdline"]):
try:
cmdline = proc.info.get("cmdline") or []
backend = _is_inference_process(cmdline)
if not backend:
continue
p = _extract_flag(cmdline, "--port")
if p is None or int(p) != port:
continue
# Extract model path
model_path = _extract_flag(cmdline, "--model") or _extract_flag(cmdline, "--model-path")
if not model_path and len(cmdline) >= 3 and cmdline[1] == "serve":
model_path = cmdline[2] if not cmdline[2].startswith("-") else None
return ProcessInfo(
pid=proc.info["pid"],
backend=backend,
model_path=model_path,
port=port,
served_model_name=_extract_flag(cmdline, "--served-model-name"),
)
except (psutil.NoSuchProcess, psutil.AccessDenied, ValueError):
continue
return None
async def kill_process(pid: int, force: bool = False) -> bool:
"""Kill process and its children."""
try:
proc = psutil.Process(pid)
except psutil.NoSuchProcess:
return True
# Kill children first
for child in proc.children(recursive=True):
try:
child.kill()
except psutil.NoSuchProcess:
pass
# Terminate main process
try:
proc.terminate()
proc.wait(timeout=10)
except psutil.TimeoutExpired:
if force:
proc.kill()
except psutil.NoSuchProcess:
pass
await asyncio.sleep(1)
return True
def build_vllm_command(recipe: Recipe) -> List[str]:
"""Build vLLM launch command."""
if recipe.python_path:
cmd = [recipe.python_path, "-m", "vllm.entrypoints.openai.api_server"]
else:
cmd = ["vllm", "serve"]
cmd.extend([recipe.model_path, "--host", recipe.host, "--port", str(recipe.port)])
if recipe.served_model_name:
cmd.extend(["--served-model-name", recipe.served_model_name])
if recipe.tensor_parallel_size > 1:
cmd.extend(["--tensor-parallel-size", str(recipe.tensor_parallel_size)])
if recipe.pipeline_parallel_size > 1:
cmd.extend(["--pipeline-parallel-size", str(recipe.pipeline_parallel_size)])
cmd.extend(["--max-model-len", str(recipe.max_model_len)])
cmd.extend(["--gpu-memory-utilization", str(recipe.gpu_memory_utilization)])
cmd.extend(["--max-num-seqs", str(recipe.max_num_seqs)])
if recipe.kv_cache_dtype != "auto":
cmd.extend(["--kv-cache-dtype", recipe.kv_cache_dtype])
if recipe.trust_remote_code:
cmd.append("--trust-remote-code")
if recipe.tool_call_parser:
cmd.extend(["--tool-call-parser", recipe.tool_call_parser, "--enable-auto-tool-choice"])
if recipe.quantization:
cmd.extend(["--quantization", recipe.quantization])
if recipe.dtype:
cmd.extend(["--dtype", recipe.dtype])
# Extra args
for key, value in recipe.extra_args.items():
flag = f"--{key.replace('_', '-')}"
if value is True:
cmd.append(flag)
elif value not in (False, None):
if isinstance(value, (dict, list)):
cmd.extend([flag, json.dumps(value)])
else:
cmd.extend([flag, str(value)])
return cmd
def build_sglang_command(recipe: Recipe) -> List[str]:
"""Build SGLang launch command."""
python = recipe.python_path or settings.sglang_python or "python"
cmd = [python, "-m", "sglang.launch_server"]
cmd.extend(["--model-path", recipe.model_path])
cmd.extend(["--host", recipe.host, "--port", str(recipe.port)])
if recipe.served_model_name:
cmd.extend(["--served-model-name", recipe.served_model_name])
if recipe.tensor_parallel_size > 1:
cmd.extend(["--tp", str(recipe.tensor_parallel_size)])
cmd.extend(["--context-length", str(recipe.max_model_len)])
cmd.extend(["--mem-fraction-static", str(recipe.gpu_memory_utilization)])
if recipe.trust_remote_code:
cmd.append("--trust-remote-code")
if recipe.quantization:
cmd.extend(["--quantization", recipe.quantization])
for key, value in recipe.extra_args.items():
flag = f"--{key.replace('_', '-')}"
if value is True:
cmd.append(flag)
elif value not in (False, None):
cmd.extend([flag, str(value)])
return cmd
async def launch_model(recipe: Recipe) -> Tuple[bool, Optional[int], str]:
"""Launch inference server with recipe config."""
recipe.port = settings.inference_port # Override with configured port
if recipe.backend == Backend.SGLANG:
cmd = build_sglang_command(recipe)
else:
cmd = build_vllm_command(recipe)
log_file = Path(f"/tmp/vllm_{recipe.id}.log")
env = os.environ.copy()
try:
with open(log_file, "w") as log:
proc = subprocess.Popen(
cmd,
stdout=log,
stderr=subprocess.STDOUT,
env=env,
start_new_session=True,
)
await asyncio.sleep(3)
if proc.poll() is not None:
tail = log_file.read_text()[-500:] if log_file.exists() else ""
return False, None, f"Process exited early: {tail}"
return True, proc.pid, str(log_file)
except Exception as e:
return False, None, str(e)
async def evict_model(force: bool = False) -> Optional[int]:
"""Stop current running model."""
current = find_inference_process(settings.inference_port)
if not current:
return None
await kill_process(current.pid, force=force)
return current.pid
async def switch_model(recipe: Recipe, force: bool = False) -> Tuple[bool, Optional[int], str]:
"""Switch to a new model (evict current + launch new)."""
await evict_model(force=force)
await asyncio.sleep(2)
return await launch_model(recipe)