import os import gradio as gr from subprocess import Popen, PIPE import subprocess import logging import threading import time import queue logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s') logger = logging.getLogger(__name__) def check_gpu_available(): try: nvidia_smi = subprocess.run(["nvidia-smi"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True) logger.info(f"nvidia-smi 输出: {nvidia_smi.stdout}") if nvidia_smi.returncode == 0: logger.info("通过nvidia-smi命令检测到GPU") return True if os.environ.get("SPACE_RUNTIME_ARCH", "") == "gpu": logger.info("通过环境变量检测到GPU") return True import torch has_gpu = torch.cuda.is_available() logger.info(f"通过PyTorch检测到GPU，设备数量: {torch.cuda.device_count()}" if has_gpu else "PyTorch未检测到GPU") return has_gpu except Exception as e: logger.error(f"GPU检测失败: {str(e)}") return False MODEL_NAME = os.environ.get("MODEL_NAME", "zhangchenxu/TinyV-1.5B") API_PORT = int(os.environ.get("API_PORT", 8000)) GRADIO_PORT = int(os.environ.get("GRADIO_PORT", 7860)) API_KEY = os.environ.get("API_KEY", "token-abc123") # 默认API密钥设置为token-abc123 USE_TRANSFORMERS_IMPL = os.environ.get("USE_TRANSFORMERS_IMPL", "true").lower() == "true" ENFORCE_EAGER = os.environ.get("ENFORCE_EAGER", "true").lower() == "true" HAS_GPU = check_gpu_available() FORCE_GPU = os.environ.get("FORCE_GPU", "false").lower() == "true" if FORCE_GPU: HAS_GPU = True vllm_process = None log_queue = queue.Queue(maxsize=1000) # 用于存储日志的队列 log_thread = None stop_log_thread = False def log_reader_thread(process): """异步读取进程日志的线程""" global stop_log_thread while not stop_log_thread: # 读取stderr if process.stderr: line = process.stderr.readline() if line: log_queue.put(f"[ERROR] {line.strip()}") continue # 读取stdout if process.stdout: line = process.stdout.readline() if line: log_queue.put(line.strip()) continue # 检查进程是否还在运行 if process.poll() is not None: log_queue.put(f"进程已退出，返回码: {process.poll()}") break # 短暂休眠以减少CPU使用 time.sleep(0.1) def start_vllm_server(): global vllm_process, log_thread, stop_log_thread if vllm_process is not None: return "vLLM 服务已经在运行" os.environ["VLLM_LOGGING_LEVEL"] = "DEBUG" cmd = [ "vllm", "serve", MODEL_NAME, "--host", "0.0.0.0", "--port", str(API_PORT), "--dtype", "auto", "--trust-remote-code", "--disable-async-output-proc", "--api-key", API_KEY, # 添加API密钥 ] if ENFORCE_EAGER: cmd.append("--enforce-eager") if USE_TRANSFORMERS_IMPL: cmd.extend(["--model-impl", "transformers"]) if HAS_GPU: logger.info("使用GPU模式启动vLLM") cmd.extend(["--device", "cuda", "--max-model-len", "2048", "--gpu-memory-utilization", "0.9"]) else: logger.info("使用CPU模式启动vLLM") cmd.extend(["--device", "cpu", "--max-model-len", "1024"]) logger.info(f"启动命令: {' '.join(cmd)}") try: vllm_process = Popen(cmd, stdout=PIPE, stderr=PIPE, text=True, bufsize=1) # 启动日志读取线程 stop_log_thread = False log_thread = threading.Thread(target=log_reader_thread, args=(vllm_process,)) log_thread.daemon = True log_thread.start() return "vLLM 服务器已启动！请等待模型加载完成..." except Exception as e: logger.error(f"启动vLLM服务器时出错: {str(e)}") return f"启动vLLM服务器时出错: {str(e)}" def stop_vllm_server(): global vllm_process, stop_log_thread if vllm_process is None: return "vLLM 服务未运行" # 停止日志线程 stop_log_thread = True if log_thread and log_thread.is_alive(): log_thread.join(timeout=2) # 终止进程 vllm_process.terminate() try: vllm_process.wait(timeout=5) except subprocess.TimeoutExpired: vllm_process.kill() vllm_process = None return "vLLM 服务已停止" def check_server_status(): if vllm_process is None: return "未运行" return_code = vllm_process.poll() return "运行中" if return_code is None else f"已停止 (返回码: {return_code})" def get_server_logs(): """获取日志，不会阻塞UI""" if vllm_process is None: return "服务未运行，无日志可显示" # 从队列中获取日志 logs = [] try: # 最多获取200行日志，避免过多 for _ in range(200): if log_queue.empty(): break logs.append(log_queue.get_nowait()) log_queue.task_done() except queue.Empty: pass if logs: return "\n".join(logs) else: # 检查进程状态 if vllm_process.poll() is not None: return f"服务已停止，返回码: {vllm_process.poll()}" return "服务正在运行，暂无新日志" def serve_test_ui(): with gr.Blocks(title="vLLM OpenAI兼容API服务") as demo: with gr.Row(): with gr.Column(): gr.Markdown("# vLLM OpenAI兼容API服务控制面板") # 系统信息 gpu_info = "已检测到" if HAS_GPU else "未检测到" system_info = f""" ## 系统信息 - GPU: {gpu_info} - 运行环境: {'Hugging Face Space' if 'SPACE_ID' in os.environ else '本地环境'} - 当前加载模型: `{MODEL_NAME}` - API密钥: `{API_KEY}` """ gr.Markdown(system_info) with gr.Row(): start_btn = gr.Button("启动服务", variant="primary") stop_btn = gr.Button("停止服务", variant="stop") status_text = gr.Textbox(label="服务状态", value="未运行", interactive=False) refresh_btn = gr.Button("刷新状态") logs_text = gr.Textbox(label="服务日志", interactive=False, lines=15) logs_refresh_btn = gr.Button("刷新日志") # 高级选项 with gr.Accordion("高级选项", open=False): model_input = gr.Textbox(label="模型名称", value=MODEL_NAME, placeholder="输入模型名称，如 zhangchenxu/TinyV-1.5B") with gr.Row(): force_gpu = gr.Checkbox(label="强制使用GPU模式", value=FORCE_GPU, info="如果自动检测失败但您确定有GPU，请选中此项") use_transformers = gr.Checkbox(label="使用Transformers实现", value=USE_TRANSFORMERS_IMPL, info="使用Transformers实现而不是vLLM原生实现，可能更稳定但性能略低") enforce_eager = gr.Checkbox(label="强制Eager模式", value=ENFORCE_EAGER, info="强制使用PyTorch的Eager模式，避免CUDA图形相关问题") apply_btn = gr.Button("应用设置", variant="primary") # API使用说明 gr.Markdown("## API 使用说明") api_info = gr.Markdown(f""" ### 接口地址 在Hugging Face Space上，接口地址为您Space的URL加上`/v1`路径： ``` https://你的HF_SPACE_URL/v1 ``` ### 支持的API 1. **Chat Completions API** (`/v1/chat/completions`) - 用于聊天生成 - 与OpenAI的Chat API兼容 2. **Completions API** (`/v1/completions`) - 用于文本生成 - 与OpenAI的Completions API兼容 ### Python示例代码 ```python from openai import OpenAI # 创建客户端 client = OpenAI( base_url="https://你的HF_SPACE_URL/v1", api_key="{API_KEY}", ) # 聊天完成示例 chat_completion = client.chat.completions.create( model="{MODEL_NAME}", # 模型名称可以是任意值 messages=[ {{"role": "user", "content": "Hello!"}} ] ) print(chat_completion.choices[0].message.content) # 文本生成示例 completion = client.completions.create( model="{MODEL_NAME}", prompt="Once upon a time", max_tokens=50 ) print(completion.choices[0].text) ``` ### curl示例 ```bash # 聊天完成 curl https://你的HF_SPACE_URL/v1/chat/completions \\ -H "Content-Type: application/json" \\ -H "Authorization: Bearer {API_KEY}" \\ -d '{{"model": "{MODEL_NAME}", "messages": [{{"role": "user", "content": "Hello!"}}]}}' # 文本生成 curl https://你的HF_SPACE_URL/v1/completions \\ -H "Content-Type: application/json" \\ -H "Authorization: Bearer {API_KEY}" \\ -d '{{"model": "{MODEL_NAME}", "prompt": "Once upon a time", "max_tokens": 50}}' ``` """) # 自动刷新日志 gr.Markdown("## 日志自动刷新") auto_refresh = gr.Checkbox(label="启用日志自动刷新", value=False) # 设置事件处理 start_btn.click(start_vllm_server, inputs=[], outputs=status_text) stop_btn.click(stop_vllm_server, inputs=[], outputs=status_text) refresh_btn.click(check_server_status, inputs=[], outputs=status_text) logs_refresh_btn.click(get_server_logs, inputs=[], outputs=logs_text) # 高级选项的事件处理 def apply_settings(model_name, force_gpu_mode, use_transformers_impl, enforce_eager_mode): global MODEL_NAME, FORCE_GPU, HAS_GPU, USE_TRANSFORMERS_IMPL, ENFORCE_EAGER changed = [] if model_name.strip() and model_name != MODEL_NAME: MODEL_NAME = model_name.strip() changed.append(f"模型已更改为: {MODEL_NAME}") if force_gpu_mode != FORCE_GPU: FORCE_GPU = force_gpu_mode if FORCE_GPU: HAS_GPU = True changed.append("已强制启用GPU模式") else: HAS_GPU = check_gpu_available() changed.append(f"已恢复自动检测，GPU状态: {'已检测到' if HAS_GPU else '未检测到'}") if use_transformers_impl != USE_TRANSFORMERS_IMPL: USE_TRANSFORMERS_IMPL = use_transformers_impl changed.append(f"Transformers实现: {'已启用' if USE_TRANSFORMERS_IMPL else '已禁用'}") if enforce_eager_mode != ENFORCE_EAGER: ENFORCE_EAGER = enforce_eager_mode changed.append(f"Eager模式: {'已启用' if ENFORCE_EAGER else '已禁用'}") if changed: return "\n".join(changed) + "\n\n设置已应用。如果服务正在运行，需要重启服务以使更改生效。" else: return "没有设置被更改" def auto_refresh_logs(auto): if auto: # 每秒自动刷新日志 return gr.update(every=1) return gr.update(every=0) apply_btn.click( apply_settings, inputs=[model_input, force_gpu, use_transformers, enforce_eager], outputs=status_text ) auto_refresh.change( auto_refresh_logs, inputs=[auto_refresh], outputs=[logs_text] ) # 页面加载时自动启动状态检查 demo.load(lambda: f"系统就绪。GPU状态: {'已检测到' if HAS_GPU else '未检测到'}", inputs=[], outputs=status_text) return demo if __name__ == "__main__": demo = serve_test_ui() demo.queue().launch(server_name="0.0.0.0", server_port=GRADIO_PORT, share=True)