diff --git a/Dockerfile b/Dockerfile index b6c3cb7..2fb9b1b 100644 --- a/Dockerfile +++ b/Dockerfile @@ -9,6 +9,6 @@ COPY app /workspace/app ENV PYTHONUNBUFFERED=1 -EXPOSE 8000 8001 +EXPOSE 8001 -CMD ["python", "-m", "app.start_api"] +CMD ["python", "-m", "app.start_openai"] diff --git a/README.md b/README.md index 389f051..f164fa6 100644 --- a/README.md +++ b/README.md @@ -6,7 +6,6 @@ - 提供可容器化部署的模型推理 API - 使用 vLLM + ROCm 在 AMD GPU 上执行推理 -- 提供健康检查、鉴权和参数化配置能力 - 暴露 `8001` OpenAI 标准协议接口,兼容 OpenClaw 调用 ## 目录结构 @@ -15,10 +14,7 @@ . ├── app │ ├── config.py -│ ├── engine.py -│ ├── main.py │ ├── model_catalog.py -│ ├── start_api.py │ ├── start_openai.py │ └── schemas.py ├── .dockerignore @@ -28,37 +24,12 @@ └── requirements.txt ``` -## 接口说明 - -### 1) 健康检查 - -- 方法:`GET /health` -- 返回:服务状态与当前模型名 - -### 2) 文本生成 - -- 方法:`POST /v1/generate` -- Header:可选 `x-api-key`(配置 `API_KEY` 后必填) -- 请求体: - -```json -{ - "prompt": "请介绍一下ROCm和vLLM的关系", - "max_tokens": 256, - "temperature": 0.7, - "top_p": 0.95, - "repetition_penalty": 1.0, - "stop": [""] -} -``` - ## 配置项 项目只读取一个配置文件:`config.json`。 -- `services.api.host` / `services.api.port`:自定义推理服务监听地址与端口(默认 `0.0.0.0:8000`) - `services.openai.host` / `services.openai.port`:OpenAI 协议服务监听地址与端口(默认 `0.0.0.0:8001`) -- `api_key`:接口访问密钥(同时用于 8000 与 8001) +- `api_key`:OpenAI 接口访问密钥 - `tensor_parallel_size`:张量并行数,双卡建议 `2` - `dtype`:推理精度,默认 `bfloat16` - `model_root`:本地模型根目录,建议 `/opt/model` @@ -99,26 +70,12 @@ docker compose up -d --build ``` -3. 验证自定义推理服务: - -```bash -curl http://localhost:/health -``` - -4. 验证 OpenAI 协议服务: +3. 验证 OpenAI 协议服务: ```bash curl http://localhost:/v1/models ``` -## 推理请求示例 - -```bash -curl -X POST "http://localhost:/v1/generate" \ - -H "Content-Type: application/json" \ - -d "{\"prompt\":\"请用三句话介绍大模型推理优化\",\"max_tokens\":128,\"temperature\":0.7}" -``` - ## OpenAI 协议示例(8001) ```bash diff --git a/app/engine.py b/app/engine.py index b18ca09..0379882 100644 --- a/app/engine.py +++ b/app/engine.py @@ -1,6 +1,4 @@ -from threading import Lock - -from vllm import LLM, SamplingParams +import httpx from app.config import Settings from app.schemas import GenerateRequest, GenerateResponse @@ -9,36 +7,35 @@ from app.schemas import GenerateRequest, GenerateResponse class InferenceEngine: def __init__(self, settings: Settings) -> None: self.settings = settings - self.lock = Lock() - self.model = LLM( - model=settings.model_name, - tensor_parallel_size=settings.tensor_parallel_size, - max_model_len=settings.max_model_len, - gpu_memory_utilization=settings.gpu_memory_utilization, - max_num_seqs=settings.max_num_seqs, - dtype=settings.dtype, - enforce_eager=settings.enforce_eager, - trust_remote_code=settings.trust_remote_code, - revision=settings.revision, - ) + self.client = httpx.Client(timeout=300.0) def close(self) -> None: - return None + self.client.close() def generate(self, req: GenerateRequest) -> GenerateResponse: - sampling_params = SamplingParams( - temperature=req.temperature, - top_p=req.top_p, - max_tokens=req.max_tokens, - repetition_penalty=req.repetition_penalty, - stop=req.stop, + headers = {"Content-Type": "application/json"} + if self.settings.api_key: + headers["Authorization"] = f"Bearer {self.settings.api_key}" + payload = { + "model": self.settings.served_model_name or self.settings.model_name, + "messages": [{"role": "user", "content": req.prompt}], + "max_tokens": req.max_tokens, + "temperature": req.temperature, + "top_p": req.top_p, + } + if req.stop: + payload["stop"] = req.stop + response = self.client.post( + f"{self.settings.vllm_openai_internal_url}/chat/completions", + headers=headers, + json=payload, ) - with self.lock: - outputs = self.model.generate([req.prompt], sampling_params, use_tqdm=False) - output = outputs[0] - completion = output.outputs[0].text - usage_prompt = len(output.prompt_token_ids) - usage_completion = len(output.outputs[0].token_ids) + response.raise_for_status() + body = response.json() + completion = body["choices"][0]["message"]["content"] + usage = body.get("usage", {}) + usage_prompt = int(usage.get("prompt_tokens", 0)) + usage_completion = int(usage.get("completion_tokens", 0)) return GenerateResponse( text=completion, prompt=req.prompt, diff --git a/docker-compose.yml b/docker-compose.yml index 4dbaf37..141a252 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -1,30 +1,5 @@ services: - inference: - build: - context: . - dockerfile: Dockerfile - image: rocm-vllm-inference:latest - container_name: rocm-vllm-inference - entrypoint: ["python"] - command: ["-m", "app.start_api"] - ports: - - "8002:8002" - volumes: - - /opt/model:/opt/model:ro - - ./config.json:/workspace/config.json:ro - devices: - - /dev/kfd - - /dev/dri - group_add: - - video - ipc: host - shm_size: 16g - cap_add: - - SYS_PTRACE - security_opt: - - seccomp=unconfined - restart: unless-stopped - openai_api: + vllm: build: context: . dockerfile: Dockerfile diff --git a/requirements.txt b/requirements.txt index 3ae98d3..3512bce 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,4 @@ fastapi==0.116.1 uvicorn==0.35.0 pydantic==2.11.7 +httpx==0.28.1