Spaces:
Runtime error
Runtime error
File size: 7,482 Bytes
102f4d2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 | """On-demand vLLM process manager.
Starts vLLM when first needed, shuts it down after idle.
Set VLLM_ON_DEMAND=false to use an externally managed vLLM instead.
Set VLLM_IDLE_TIMEOUT=300 (seconds) to control the idle shutdown window.
"""
import os
import subprocess
import threading
import time
import requests
from loguru import logger
VLLM_MODEL = os.getenv("VLLM_MODEL", "Qwen/Qwen2.5-VL-7B-Instruct")
VLLM_PORT = int(os.getenv("VLLM_PORT", "8000"))
IDLE_TIMEOUT = int(os.getenv("VLLM_IDLE_TIMEOUT", "300")) # 5 min default
ON_DEMAND = os.getenv("VLLM_ON_DEMAND", "true").lower() == "true"
DOCKER_CONTAINER = os.getenv("VLLM_DOCKER_CONTAINER", "rocm") # container that has vllm installed
class _VLLMManager:
def __init__(self):
self._proc: subprocess.Popen | None = None
self._lock = threading.Lock()
self._last_used = 0.0
threading.Thread(target=self._watchdog, daemon=True, name="vllm-watchdog").start()
# ── Public ────────────────────────────────────────────────────────────
def is_running(self) -> bool:
if not ON_DEMAND or DOCKER_CONTAINER:
# Docker mode or external vLLM: rely solely on health endpoint
return self._check_health()
with self._lock:
if self._proc is None or self._proc.poll() is not None:
return False
return self._check_health()
def ensure_running(self, progress_cb=None) -> None:
"""Start vLLM if not running. Blocks until healthy (max 3 min)."""
if not ON_DEMAND:
return
with self._lock:
if self._check_health():
self._last_used = time.time()
return
self._start(progress_cb)
def stop(self) -> None:
if not ON_DEMAND:
return
with self._lock:
self._stop_locked()
def touch(self) -> None:
"""Reset idle timer — call after each successful vLLM API call."""
self._last_used = time.time()
def status(self) -> dict:
running = self.is_running()
idle = round(time.time() - self._last_used, 1) if self._last_used else None
return {
"running": running,
"on_demand": ON_DEMAND,
"idle_seconds": idle,
"idle_timeout": IDLE_TIMEOUT,
"model": VLLM_MODEL,
}
# ── Internal ──────────────────────────────────────────────────────────
def _health_url(self) -> str:
return f"http://localhost:{VLLM_PORT}/health"
def _check_health(self) -> bool:
try:
return requests.get(self._health_url(), timeout=2).status_code == 200
except Exception:
return False
def _start(self, progress_cb=None) -> None:
logger.info("vLLM: starting on demand…")
if progress_cb:
progress_cb("Starting AI model (Qwen2.5-VL)… ~2 min first time")
# Try Docker container first (vLLM may only be installed inside a container)
if DOCKER_CONTAINER:
self._start_via_docker(progress_cb)
else:
self._start_via_subprocess(progress_cb)
def _start_via_docker(self, progress_cb=None) -> None:
"""Start vLLM inside an existing Docker container via docker exec."""
cmd = (
f"vllm serve {VLLM_MODEL} "
f"--host 0.0.0.0 --port {VLLM_PORT} "
f"--gpu-memory-utilization 0.85 --max-model-len 4096 "
f"> /tmp/vllm_server.log 2>&1"
)
subprocess.Popen(
["docker", "exec", "-d", DOCKER_CONTAINER, "bash", "-c", cmd],
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
)
self._proc = None # process lives inside container, tracked by health check
deadline = time.time() + 200
tick = 0
while time.time() < deadline:
time.sleep(5)
tick += 1
if self._check_health():
self._last_used = time.time()
logger.info(f"vLLM (docker) ready after {tick * 5}s")
return
if progress_cb and tick % 6 == 0:
progress_cb(f"AI model loading… {tick * 5}s")
raise RuntimeError("vLLM did not start within 200s")
def _start_via_subprocess(self, progress_cb=None) -> None:
"""Start vLLM as a direct subprocess (vllm must be in current Python env)."""
import sys
self._proc = subprocess.Popen(
[
sys.executable, "-m", "vllm.entrypoints.openai.api_server",
"--model", VLLM_MODEL,
"--device", "rocm",
"--port", str(VLLM_PORT),
"--gpu-memory-utilization", "0.85",
"--max-model-len", "4096",
],
stdout=subprocess.DEVNULL,
stderr=subprocess.PIPE,
)
deadline = time.time() + 200
tick = 0
while time.time() < deadline:
time.sleep(5)
tick += 1
if self._proc.poll() is not None:
err = self._proc.stderr.read().decode()[-600:]
raise RuntimeError(f"vLLM exited during startup: {err}")
if self._check_health():
self._last_used = time.time()
logger.info(f"vLLM ready after {tick * 5}s")
return
if progress_cb and tick % 6 == 0:
progress_cb(f"AI model loading… {tick * 5}s")
raise RuntimeError("vLLM did not start within 200s")
def _stop_locked(self) -> None:
if DOCKER_CONTAINER:
subprocess.run(
["docker", "exec", DOCKER_CONTAINER, "pkill", "-f", "vllm"],
capture_output=True,
)
self._proc = None
elif self._proc and self._proc.poll() is None:
self._proc.terminate()
try:
self._proc.wait(timeout=10)
except subprocess.TimeoutExpired:
self._proc.kill()
self._proc = None
logger.info("vLLM stopped")
def _watchdog(self) -> None:
while True:
time.sleep(60)
if not ON_DEMAND or IDLE_TIMEOUT <= 0:
continue
with self._lock:
if (self._proc
and self._proc.poll() is None
and self._last_used > 0
and time.time() - self._last_used > IDLE_TIMEOUT):
logger.info(
f"vLLM idle {IDLE_TIMEOUT}s → shutting down to save GPU credits"
)
self._stop_locked()
_manager = _VLLMManager()
# ── Module-level helpers ──────────────────────────────────────────────────────
def ensure_vllm_running(progress_cb=None) -> None:
_manager.ensure_running(progress_cb)
def vllm_touch() -> None:
_manager.touch()
def vllm_stop() -> None:
_manager.stop()
def vllm_is_running() -> bool:
return _manager.is_running()
def vllm_status() -> dict:
return _manager.status()
|