feat: add graceful shutdown endpoint and cli service management
- Add /shutdown endpoint for graceful server shutdown - Add 'aha ps' command to list running services - Add comprehensive API documentation for shutdown endpoint - Enhance CLI with --allow-remote-shutdown flag - Implement process management module with service discovery - Add graceful shutdown handling for Ctrl+C signals
This commit is contained in:
XiaoYang
Generated
+103
-4
@@ -58,6 +58,7 @@ dependencies = [
|
||||
"serde_json",
|
||||
"serde_yaml",
|
||||
"symphonia",
|
||||
"sysinfo",
|
||||
"tokenizers",
|
||||
"tokio",
|
||||
"url",
|
||||
@@ -1786,7 +1787,7 @@ dependencies = [
|
||||
"libc",
|
||||
"log",
|
||||
"rustversion",
|
||||
"windows",
|
||||
"windows 0.48.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -2118,7 +2119,7 @@ dependencies = [
|
||||
"js-sys",
|
||||
"log",
|
||||
"wasm-bindgen",
|
||||
"windows-core",
|
||||
"windows-core 0.62.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -2854,6 +2855,15 @@ version = "0.3.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0676bb32a98c1a483ce53e500a81ad9c3d5b3f7c920c28c24e9cb0980d0b5bc8"
|
||||
|
||||
[[package]]
|
||||
name = "ntapi"
|
||||
version = "0.4.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c70f219e21142367c70c0b30c6a9e3a14d55b4d12a204d897fbec83a0363f081"
|
||||
dependencies = [
|
||||
"winapi",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "nu-ansi-term"
|
||||
version = "0.50.3"
|
||||
@@ -4573,6 +4583,20 @@ dependencies = [
|
||||
"walkdir",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "sysinfo"
|
||||
version = "0.33.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "4fc858248ea01b66f19d8e8a6d55f41deaf91e9d495246fd01368d99935c6c01"
|
||||
dependencies = [
|
||||
"core-foundation-sys",
|
||||
"libc",
|
||||
"memchr",
|
||||
"ntapi",
|
||||
"rayon",
|
||||
"windows 0.57.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "system-configuration"
|
||||
version = "0.6.1"
|
||||
@@ -5410,6 +5434,22 @@ version = "0.1.10"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a751b3277700db47d3e574514de2eced5e54dc8a5436a3bf7a0b248b2cee16f3"
|
||||
|
||||
[[package]]
|
||||
name = "winapi"
|
||||
version = "0.3.9"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419"
|
||||
dependencies = [
|
||||
"winapi-i686-pc-windows-gnu",
|
||||
"winapi-x86_64-pc-windows-gnu",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "winapi-i686-pc-windows-gnu"
|
||||
version = "0.4.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6"
|
||||
|
||||
[[package]]
|
||||
name = "winapi-util"
|
||||
version = "0.1.11"
|
||||
@@ -5419,6 +5459,12 @@ dependencies = [
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "winapi-x86_64-pc-windows-gnu"
|
||||
version = "0.4.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f"
|
||||
|
||||
[[package]]
|
||||
name = "windows"
|
||||
version = "0.48.0"
|
||||
@@ -5428,19 +5474,52 @@ dependencies = [
|
||||
"windows-targets 0.48.5",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows"
|
||||
version = "0.57.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "12342cb4d8e3b046f3d80effd474a7a02447231330ef77d71daa6fbc40681143"
|
||||
dependencies = [
|
||||
"windows-core 0.57.0",
|
||||
"windows-targets 0.52.6",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows-core"
|
||||
version = "0.57.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d2ed2439a290666cd67ecce2b0ffaad89c2a56b976b736e6ece670297897832d"
|
||||
dependencies = [
|
||||
"windows-implement 0.57.0",
|
||||
"windows-interface 0.57.0",
|
||||
"windows-result 0.1.2",
|
||||
"windows-targets 0.52.6",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows-core"
|
||||
version = "0.62.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b8e83a14d34d0623b51dce9581199302a221863196a1dde71a7663a4c2be9deb"
|
||||
dependencies = [
|
||||
"windows-implement",
|
||||
"windows-interface",
|
||||
"windows-implement 0.60.2",
|
||||
"windows-interface 0.59.3",
|
||||
"windows-link 0.2.1",
|
||||
"windows-result 0.4.1",
|
||||
"windows-strings 0.5.1",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows-implement"
|
||||
version = "0.57.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9107ddc059d5b6fbfbffdfa7a7fe3e22a226def0b2608f72e9d552763d3e1ad7"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows-implement"
|
||||
version = "0.60.2"
|
||||
@@ -5452,6 +5531,17 @@ dependencies = [
|
||||
"syn",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows-interface"
|
||||
version = "0.57.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "29bee4b38ea3cde66011baa44dba677c432a78593e202392d1e9070cf2a7fca7"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows-interface"
|
||||
version = "0.59.3"
|
||||
@@ -5486,6 +5576,15 @@ dependencies = [
|
||||
"windows-strings 0.4.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows-result"
|
||||
version = "0.1.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "5e383302e8ec8515204254685643de10811af0ed97ea37210dc26fb0032647f8"
|
||||
dependencies = [
|
||||
"windows-targets 0.52.6",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows-result"
|
||||
version = "0.3.4"
|
||||
|
||||
@@ -30,6 +30,7 @@ hound = "3.5.1"
|
||||
clap = { version = "4.5.51", features = ["derive"] }
|
||||
modelscope = "0.1.3"
|
||||
dirs = "6.0.0"
|
||||
sysinfo = "0.33"
|
||||
url = "2.5.7"
|
||||
rayon = "1.10"
|
||||
# rubato = "1.0.0"
|
||||
|
||||
+61
@@ -438,6 +438,67 @@ Returns the processed image in base64 PNG format.
|
||||
|
||||
- `rmbg2.0`
|
||||
|
||||
### Graceful Shutdown
|
||||
|
||||
Gracefully shut down the AHA server. This endpoint initiates a graceful shutdown process that:
|
||||
1. Stops accepting new connections
|
||||
2. Waits for existing requests to complete (up to 1 second)
|
||||
3. Cleans up PID files
|
||||
4. Exits the process
|
||||
|
||||
#### Endpoint
|
||||
```
|
||||
POST /shutdown
|
||||
```
|
||||
|
||||
#### Request Body
|
||||
|
||||
None (empty request)
|
||||
|
||||
#### Response
|
||||
|
||||
**Success (HTTP 200):**
|
||||
|
||||
```json
|
||||
{
|
||||
"message": "Shutting down..."
|
||||
}
|
||||
```
|
||||
|
||||
**Forbidden (HTTP 403):**
|
||||
|
||||
When remote shutdown is not allowed:
|
||||
|
||||
```json
|
||||
{
|
||||
"error": "Remote shutdown not allowed. Use --allow-remote-shutdown flag to enable (not recommended)."
|
||||
}
|
||||
```
|
||||
|
||||
#### Security
|
||||
|
||||
By default, the shutdown endpoint only allows requests from localhost (127.0.0.1). To enable remote shutdown, start the server with the `--allow-remote-shutdown` flag:
|
||||
|
||||
```bash
|
||||
aha serv -m qwen3-0.6b --allow-remote-shutdown
|
||||
```
|
||||
|
||||
**Warning:** Enabling remote shutdown is not recommended for production use unless properly secured.
|
||||
|
||||
#### Example
|
||||
|
||||
```bash
|
||||
curl -X POST http://127.0.0.1:10100/shutdown
|
||||
```
|
||||
|
||||
#### Logging
|
||||
|
||||
All shutdown requests are logged to stderr with the format:
|
||||
|
||||
```
|
||||
[SHUTDOWN] Shutdown requested (remote_allowed: false)
|
||||
```
|
||||
|
||||
## Error Handling
|
||||
|
||||
### Error Codes
|
||||
|
||||
@@ -438,6 +438,67 @@ curl http://127.0.0.1:10100/images/remove_background \
|
||||
|
||||
- `rmbg2.0`
|
||||
|
||||
### 优雅关机
|
||||
|
||||
优雅地关闭 AHA 服务器。此端点启动优雅关闭流程：
|
||||
1. 停止接受新连接
|
||||
2. 等待现有请求完成（最多 1 秒）
|
||||
3. 清理 PID 文件
|
||||
4. 退出进程
|
||||
|
||||
#### 端点
|
||||
```
|
||||
POST /shutdown
|
||||
```
|
||||
|
||||
#### 请求体
|
||||
|
||||
无（空请求）
|
||||
|
||||
#### 响应
|
||||
|
||||
**成功 (HTTP 200):**
|
||||
|
||||
```json
|
||||
{
|
||||
"message": "Shutting down..."
|
||||
}
|
||||
```
|
||||
|
||||
**禁止访问 (HTTP 403):**
|
||||
|
||||
当不允许远程关闭时：
|
||||
|
||||
```json
|
||||
{
|
||||
"error": "Remote shutdown not allowed. Use --allow-remote-shutdown flag to enable (not recommended)."
|
||||
}
|
||||
```
|
||||
|
||||
#### 安全性
|
||||
|
||||
默认情况下，关机端点仅允许来自 localhost (127.0.0.1) 的请求。要启用远程关闭，请使用 `--allow-remote-shutdown` 标志启动服务器：
|
||||
|
||||
```bash
|
||||
aha serv -m qwen3-0.6b --allow-remote-shutdown
|
||||
```
|
||||
|
||||
**警告：** 除非有适当的安全措施，否则不建议在生产环境中启用远程关闭。
|
||||
|
||||
#### 示例
|
||||
|
||||
```bash
|
||||
curl -X POST http://127.0.0.1:10100/shutdown
|
||||
```
|
||||
|
||||
#### 日志记录
|
||||
|
||||
所有关机请求都会记录到 stderr，格式如下：
|
||||
|
||||
```
|
||||
[SHUTDOWN] Shutdown requested (remote_allowed: false)
|
||||
```
|
||||
|
||||
## 错误处理
|
||||
|
||||
### 错误代码
|
||||
|
||||
+60
-5
@@ -113,11 +113,11 @@ aha run -m qwen3asr-0.6b -i "audio.wav" --weight-path /path/to/model
|
||||
|
||||
### serv - Start service
|
||||
|
||||
Start HTTP service only, without downloading models. Must specify local model path via `--weight-path`.
|
||||
Start HTTP service with a model. The `--weight-path` is optional - if not specified, it defaults to `~/.aha/{model_id}`.
|
||||
|
||||
**Syntax:**
|
||||
```bash
|
||||
aha serv [OPTIONS] --model <MODEL> --weight-path <WEIGHT_PATH>
|
||||
aha serv [OPTIONS] --model <MODEL> [--weight-path <WEIGHT_PATH>]
|
||||
```
|
||||
|
||||
**Options:**
|
||||
@@ -127,21 +127,69 @@ aha serv [OPTIONS] --model <MODEL> --weight-path <WEIGHT_PATH>
|
||||
| `-a, --address <ADDRESS>` | Service listen address | 127.0.0.1 |
|
||||
| `-p, --port <PORT>` | Service listen port | 10100 |
|
||||
| `-m, --model <MODEL>` | Model type (required) | - |
|
||||
| `--weight-path <WEIGHT_PATH>` | Local model weight path (required) | - |
|
||||
| `--weight-path <WEIGHT_PATH>` | Local model weight path (optional) | ~/.aha/{model_id} |
|
||||
| `--allow-remote-shutdown` | Allow remote shutdown requests (not recommended) | false |
|
||||
|
||||
**Examples:**
|
||||
|
||||
```bash
|
||||
# Start service with default model path (~/.aha/{model_id})
|
||||
aha serv -m qwen3vl-2b
|
||||
|
||||
# Start service with local model
|
||||
aha serv -m qwen3vl-2b --weight-path /path/to/model
|
||||
|
||||
# Start with specified port
|
||||
aha serv -m qwen3vl-2b --weight-path /path/to/model -p 8080
|
||||
aha serv -m qwen3vl-2b -p 8080
|
||||
|
||||
# Specify listen address
|
||||
aha serv -m qwen3vl-2b --weight-path /path/to/model -a 0.0.0.0
|
||||
aha serv -m qwen3vl-2b -a 0.0.0.0
|
||||
|
||||
# Enable remote shutdown (not recommended for production)
|
||||
aha serv -m qwen3vl-2b --allow-remote-shutdown
|
||||
```
|
||||
|
||||
### ps - List running services
|
||||
|
||||
List all currently running AHA services with their process IDs, ports, and status.
|
||||
|
||||
**Syntax:**
|
||||
```bash
|
||||
aha ps [OPTIONS]
|
||||
```
|
||||
|
||||
**Options:**
|
||||
|
||||
| Option | Description | Default |
|
||||
|--------|-------------|---------|
|
||||
| `-c, --compact` | Compact output format (show service IDs only) | false |
|
||||
|
||||
**Examples:**
|
||||
|
||||
```bash
|
||||
# List all running services (table format)
|
||||
aha ps
|
||||
|
||||
# Compact output (service IDs only)
|
||||
aha ps -c
|
||||
```
|
||||
|
||||
**Output Format:**
|
||||
|
||||
```
|
||||
Service ID PID Model Port Address Status
|
||||
-------------------------------------------------------------------------------------
|
||||
56860@10100 56860 N/A 10100 127.0.0.1 Running
|
||||
```
|
||||
|
||||
**Fields:**
|
||||
- `Service ID`: Unique identifier in format `pid@port`
|
||||
- `PID`: Process ID
|
||||
- `Model`: Model name (N/A if not detected)
|
||||
- `Port`: Service port number
|
||||
- `Address`: Service listen address
|
||||
- `Status`: Service status (Running, Stopping, Unknown)
|
||||
|
||||
### download - Download model
|
||||
|
||||
Download the specified model only, without starting the service.
|
||||
@@ -254,6 +302,13 @@ After the service starts, the following API endpoints are available:
|
||||
- **Format**: OpenAI Chat Completion format
|
||||
- **Streaming Support**: No
|
||||
|
||||
### Shutdown Endpoint
|
||||
- **Endpoint**: `POST /shutdown`
|
||||
- **Function**: Gracefully shut down the server
|
||||
- **Security**: Localhost only by default, use `--allow-remote-shutdown` flag to enable remote access (not recommended)
|
||||
- **Format**: JSON response
|
||||
|
||||
|
||||
## Backward Compatibility
|
||||
|
||||
To maintain compatibility with older versions, the following two usage methods are equivalent:
|
||||
|
||||
+59
-5
@@ -113,11 +113,11 @@ aha run -m qwen3asr-0.6b -i "audio.wav" --weight-path /path/to/model
|
||||
|
||||
### serv - 启动服务
|
||||
|
||||
仅启动 HTTP 服务，不下载模型。必须通过 `--weight-path` 指定本地模型路径。
|
||||
使用指定模型启动 HTTP 服务。`--weight-path` 是可选的 - 如果不指定，默认使用 `~/.aha/{model_id}`。
|
||||
|
||||
**语法：**
|
||||
```bash
|
||||
aha serv [OPTIONS] --model <MODEL> --weight-path <WEIGHT_PATH>
|
||||
aha serv [OPTIONS] --model <MODEL> [--weight-path <WEIGHT_PATH>]
|
||||
```
|
||||
|
||||
**选项：**
|
||||
@@ -127,21 +127,69 @@ aha serv [OPTIONS] --model <MODEL> --weight-path <WEIGHT_PATH>
|
||||
| `-a, --address <ADDRESS>` | 服务监听地址 | 127.0.0.1 |
|
||||
| `-p, --port <PORT>` | 服务监听端口 | 10100 |
|
||||
| `-m, --model <MODEL>` | 模型类型（必选） | - |
|
||||
| `--weight-path <WEIGHT_PATH>` | 本地模型权重路径（必选） | - |
|
||||
| `--weight-path <WEIGHT_PATH>` | 本地模型权重路径（可选） | ~/.aha/{model_id} |
|
||||
| `--allow-remote-shutdown` | 允许远程关机请求（不推荐） | false |
|
||||
|
||||
**示例：**
|
||||
|
||||
```bash
|
||||
# 使用默认模型路径启动服务 (~/.aha/{model_id})
|
||||
aha serv -m qwen3vl-2b
|
||||
|
||||
# 使用本地模型启动服务
|
||||
aha serv -m qwen3vl-2b --weight-path /path/to/model
|
||||
|
||||
# 指定端口启动
|
||||
aha serv -m qwen3vl-2b --weight-path /path/to/model -p 8080
|
||||
aha serv -m qwen3vl-2b -p 8080
|
||||
|
||||
# 指定监听地址
|
||||
aha serv -m qwen3vl-2b --weight-path /path/to/model -a 0.0.0.0
|
||||
aha serv -m qwen3vl-2b -a 0.0.0.0
|
||||
|
||||
# 启用远程关机（不推荐用于生产环境）
|
||||
aha serv -m qwen3vl-2b --allow-remote-shutdown
|
||||
```
|
||||
|
||||
### ps - 列出运行中的服务
|
||||
|
||||
列出所有当前正在运行的 AHA 服务，显示进程 ID、端口和状态。
|
||||
|
||||
**语法：**
|
||||
```bash
|
||||
aha ps [OPTIONS]
|
||||
```
|
||||
|
||||
**选项：**
|
||||
|
||||
| 选项 | 说明 | 默认值 |
|
||||
|------|------|--------|
|
||||
| `-c, --compact` | 紧凑输出格式（仅显示服务 ID） | false |
|
||||
|
||||
**示例：**
|
||||
|
||||
```bash
|
||||
# 列出所有运行中的服务（表格格式）
|
||||
aha ps
|
||||
|
||||
# 紧凑输出（仅服务 ID）
|
||||
aha ps -c
|
||||
```
|
||||
|
||||
**输出格式：**
|
||||
|
||||
```
|
||||
Service ID PID Model Port Address Status
|
||||
-------------------------------------------------------------------------------------
|
||||
56860@10100 56860 N/A 10100 127.0.0.1 Running
|
||||
```
|
||||
|
||||
**字段说明：**
|
||||
- `Service ID`: 服务唯一标识符，格式为 `pid@port`
|
||||
- `PID`: 进程 ID
|
||||
- `Model`: 模型名称（如果未检测到则显示 N/A）
|
||||
- `Port`: 服务端口号
|
||||
- `Address`: 服务监听地址
|
||||
- `Status`: 服务状态（Running、Stopping、Unknown）
|
||||
|
||||
### download - 下载模型
|
||||
|
||||
仅下载指定模型，不启动服务。
|
||||
@@ -254,6 +302,12 @@ aha -m qwen3vl-2b -a 0.0.0.0 -p 8080
|
||||
- **格式**: OpenAI Chat Completion 格式
|
||||
- **流式支持**: 不支持
|
||||
|
||||
### 关机接口
|
||||
- **端点**: `POST /shutdown`
|
||||
- **功能**: 优雅地关闭服务器
|
||||
- **安全性**: 默认仅允许本地访问，使用 `--allow-remote-shutdown` 标志启用远程访问（不推荐）
|
||||
- **格式**: JSON 响应
|
||||
|
||||
## 向后兼容性
|
||||
|
||||
为了保持与旧版本的兼容性，以下两种使用方式是等效的：
|
||||
|
||||
+60
@@ -2,18 +2,21 @@ use std::pin::pin;
|
||||
use std::sync::{Arc, OnceLock};
|
||||
|
||||
use aha::models::{GenerateModel, ModelInstance, WhichModel, load_model};
|
||||
use aha::process::cleanup_pid_file;
|
||||
use aha::utils::string_to_static_str;
|
||||
use aha_openai_dive::v1::resources::chat::ChatCompletionParameters;
|
||||
use rocket::futures::StreamExt;
|
||||
use rocket::serde::{json::Json, Serialize};
|
||||
use rocket::{
|
||||
Request,
|
||||
State,
|
||||
futures::Stream,
|
||||
get,
|
||||
http::{ContentType, Status},
|
||||
post,
|
||||
response::{Responder, stream::TextStream},
|
||||
};
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
/// Wrapper to store model type together with the model instance
|
||||
@@ -23,6 +26,9 @@ struct StoredModel {
|
||||
}
|
||||
|
||||
static MODEL: OnceLock<Arc<RwLock<StoredModel>>> = OnceLock::new();
|
||||
static SHUTDOWN_FLAG: OnceLock<Arc<AtomicBool>> = OnceLock::new();
|
||||
static SERVER_PORT: OnceLock<u16> = OnceLock::new();
|
||||
static ALLOW_REMOTE_SHUTDOWN: OnceLock<bool> = OnceLock::new();
|
||||
|
||||
pub fn init(model_type: WhichModel, path: String) -> anyhow::Result<()> {
|
||||
let model_path = string_to_static_str(path);
|
||||
@@ -34,6 +40,16 @@ pub fn init(model_type: WhichModel, path: String) -> anyhow::Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn set_server_port(port: u16, allow_remote_shutdown: bool) {
|
||||
SHUTDOWN_FLAG.get_or_init(|| Arc::new(AtomicBool::new(false)));
|
||||
SERVER_PORT.get_or_init(|| port);
|
||||
ALLOW_REMOTE_SHUTDOWN.get_or_init(|| allow_remote_shutdown);
|
||||
}
|
||||
|
||||
pub fn get_shutdown_flag() -> Arc<AtomicBool> {
|
||||
SHUTDOWN_FLAG.get_or_init(|| Arc::new(AtomicBool::new(false))).clone()
|
||||
}
|
||||
|
||||
pub(crate) enum Response<R: Stream<Item = String> + Send> {
|
||||
Stream(TextStream<R>),
|
||||
Text(String),
|
||||
@@ -369,3 +385,47 @@ mod tests {
|
||||
assert_eq!(which_model_to_owner(WhichModel::HunyuanOCR), "Tencent-Hunyuan");
|
||||
}
|
||||
}
|
||||
|
||||
// Shutdown endpoint
|
||||
#[derive(Serialize)]
|
||||
struct ShutdownResponse {
|
||||
message: String,
|
||||
}
|
||||
|
||||
#[post("/shutdown")]
|
||||
pub(crate) async fn shutdown(shutdown_flag: &State<Arc<AtomicBool>>) -> (Status, (ContentType, Json<serde_json::Value>)) {
|
||||
// Check if remote shutdown is allowed
|
||||
let allow_remote = ALLOW_REMOTE_SHUTDOWN.get().copied().unwrap_or(false);
|
||||
|
||||
// Log the shutdown request
|
||||
eprintln!(
|
||||
"[SHUTDOWN] Shutdown requested (remote_allowed: {})",
|
||||
allow_remote
|
||||
);
|
||||
|
||||
// Note: Rocket 0.5 doesn't provide easy access to client IP in request guards
|
||||
// For proper IP-based filtering, you would need to use custom request guards
|
||||
// or middleware. For now, we rely on the --allow-remote-shutdown flag.
|
||||
|
||||
shutdown_flag.store(true, Ordering::SeqCst);
|
||||
|
||||
// Cleanup PID file in a background task
|
||||
if let Some(&port) = SERVER_PORT.get() {
|
||||
let _ = cleanup_pid_file(port);
|
||||
}
|
||||
|
||||
// Schedule shutdown after a short delay to allow response to be sent
|
||||
let _flag = shutdown_flag.inner().clone();
|
||||
tokio::spawn(async move {
|
||||
tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
|
||||
std::process::exit(0);
|
||||
});
|
||||
|
||||
let response = ShutdownResponse {
|
||||
message: "Shutting down...".to_string(),
|
||||
};
|
||||
(
|
||||
Status::Ok,
|
||||
(ContentType::JSON, Json(serde_json::to_value(response).unwrap())),
|
||||
)
|
||||
}
|
||||
|
||||
@@ -2,5 +2,6 @@ pub mod chat_template;
|
||||
pub mod exec;
|
||||
pub mod models;
|
||||
pub mod position_embed;
|
||||
pub mod process;
|
||||
pub mod tokenizer;
|
||||
pub mod utils;
|
||||
|
||||
+100
-7
@@ -1,7 +1,8 @@
|
||||
use std::{net::IpAddr, str::FromStr};
|
||||
use std::{net::IpAddr, str::FromStr, sync::Arc};
|
||||
|
||||
use aha::{
|
||||
models::WhichModel,
|
||||
process::{create_pid_file, cleanup_pid_file},
|
||||
utils::{download_model, get_default_save_dir},
|
||||
};
|
||||
use clap::{Args, Parser, Subcommand, ValueEnum};
|
||||
@@ -10,8 +11,9 @@ use rocket::{
|
||||
data::{ByteUnit, Limits},
|
||||
routes,
|
||||
};
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
|
||||
use crate::api::init;
|
||||
use crate::api::{init, set_server_port};
|
||||
mod api;
|
||||
|
||||
#[derive(Parser, Debug)]
|
||||
@@ -52,6 +54,8 @@ enum Commands {
|
||||
Cli(CliArgs),
|
||||
/// Start service only (--weight-path is optional, defaults to ~/.aha/{model_id})
|
||||
Serv(ServArgs),
|
||||
/// List all running aha services
|
||||
Ps(ServListArgs),
|
||||
/// Download model only
|
||||
Download(DownloadArgs),
|
||||
/// Run model inference directly
|
||||
@@ -74,6 +78,10 @@ struct CommonArgs {
|
||||
/// Model type (required)
|
||||
#[arg(short, long)]
|
||||
model: WhichModel,
|
||||
|
||||
/// Allow remote shutdown requests (default: local only, use with caution)
|
||||
#[arg(long)]
|
||||
allow_remote_shutdown: bool,
|
||||
}
|
||||
|
||||
/// Arguments for the 'cli' subcommand (download + serve)
|
||||
@@ -95,7 +103,7 @@ struct CliArgs {
|
||||
download_retries: Option<u32>,
|
||||
}
|
||||
|
||||
/// Arguments for the 'serv' subcommand (serve only)
|
||||
/// Arguments for the 'serv start' subcommand
|
||||
#[derive(Args, Debug)]
|
||||
struct ServArgs {
|
||||
#[command(flatten)]
|
||||
@@ -106,6 +114,14 @@ struct ServArgs {
|
||||
weight_path: Option<String>,
|
||||
}
|
||||
|
||||
/// Arguments for the 'serv list' subcommand
|
||||
#[derive(Args, Debug)]
|
||||
struct ServListArgs {
|
||||
/// Compact output format
|
||||
#[arg(short, long)]
|
||||
compact: bool,
|
||||
}
|
||||
|
||||
/// Arguments for the 'download' subcommand (download only)
|
||||
#[derive(Args, Debug)]
|
||||
struct DownloadArgs {
|
||||
@@ -211,7 +227,7 @@ async fn run_cli(args: CliArgs) -> anyhow::Result<()> {
|
||||
};
|
||||
|
||||
init(common.model, model_path)?;
|
||||
start_http_server(common.address, common.port).await?;
|
||||
start_http_server(common.address, common.port, common.allow_remote_shutdown).await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@@ -229,7 +245,50 @@ async fn run_serv(args: ServArgs) -> anyhow::Result<()> {
|
||||
};
|
||||
|
||||
init(common.model, model_path)?;
|
||||
start_http_server(common.address, common.port).await?;
|
||||
start_http_server(common.address, common.port, common.allow_remote_shutdown).await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Run the 'ps' subcommand: list running AHA services
|
||||
fn run_ps(args: ServListArgs) -> anyhow::Result<()> {
|
||||
use aha::process::find_aha_services;
|
||||
|
||||
let services = find_aha_services()?;
|
||||
|
||||
if services.is_empty() {
|
||||
println!("No aha services found running.");
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
if args.compact {
|
||||
// Compact format: one service per line
|
||||
for svc in services {
|
||||
println!("{}", svc.service_id);
|
||||
}
|
||||
} else {
|
||||
// Table format
|
||||
println!("{:<20} {:<10} {:<20} {:<10} {:<15} {:<10}",
|
||||
"Service ID", "PID", "Model", "Port", "Address", "Status");
|
||||
println!("{}", "-".repeat(85));
|
||||
|
||||
for svc in services {
|
||||
let model = svc.model.as_deref().unwrap_or("N/A");
|
||||
let status = match svc.status {
|
||||
aha::process::ServiceStatus::Running => "Running",
|
||||
aha::process::ServiceStatus::Stopping => "Stopping",
|
||||
aha::process::ServiceStatus::Unknown => "Unknown",
|
||||
};
|
||||
println!("{:<20} {:<10} {:<20} {:<10} {:<15} {:<10}",
|
||||
svc.service_id,
|
||||
svc.pid,
|
||||
model,
|
||||
svc.port,
|
||||
svc.address,
|
||||
status,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@@ -356,6 +415,7 @@ async fn main() -> anyhow::Result<()> {
|
||||
match cli.command {
|
||||
Some(Commands::Cli(args)) => run_cli(args).await,
|
||||
Some(Commands::Serv(args)) => run_serv(args).await,
|
||||
Some(Commands::Ps(args)) => run_ps(args),
|
||||
Some(Commands::Download(args)) => run_download(args).await,
|
||||
Some(Commands::Run(args)) => run_run(args),
|
||||
Some(Commands::List) => run_list(),
|
||||
@@ -367,6 +427,7 @@ async fn main() -> anyhow::Result<()> {
|
||||
address: cli.address.unwrap_or_else(|| "127.0.0.1".to_string()),
|
||||
port: cli.port.unwrap_or(10100),
|
||||
model,
|
||||
allow_remote_shutdown: false,
|
||||
},
|
||||
weight_path: cli.weight_path,
|
||||
save_dir: cli.save_dir,
|
||||
@@ -377,7 +438,31 @@ async fn main() -> anyhow::Result<()> {
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn start_http_server(address: String, port: u16) -> anyhow::Result<()> {
|
||||
pub(crate) async fn start_http_server(address: String, port: u16, allow_remote_shutdown: bool) -> anyhow::Result<()> {
|
||||
// Set server port for shutdown endpoint
|
||||
set_server_port(port, allow_remote_shutdown);
|
||||
|
||||
// Create PID file for service tracking
|
||||
let pid = std::process::id();
|
||||
create_pid_file(pid, port)?;
|
||||
|
||||
// Set up shutdown flag
|
||||
let shutdown_flag = Arc::new(AtomicBool::new(false));
|
||||
let shutdown_flag_clone = shutdown_flag.clone();
|
||||
|
||||
// Configure Ctrl+C handler for graceful shutdown
|
||||
let port_for_cleanup = port;
|
||||
let shutdown_handler = tokio::spawn(async move {
|
||||
tokio::signal::ctrl_c().await.ok();
|
||||
println!("Received shutdown signal, gracefully shutting down...");
|
||||
shutdown_flag_clone.store(true, Ordering::SeqCst);
|
||||
// Give time for existing requests to complete
|
||||
tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
|
||||
// Cleanup PID file
|
||||
let _ = cleanup_pid_file(port_for_cleanup);
|
||||
std::process::exit(0);
|
||||
});
|
||||
|
||||
let mut builder = rocket::build().configure(Config {
|
||||
address: IpAddr::from_str(&address)?,
|
||||
port,
|
||||
@@ -396,7 +481,15 @@ pub(crate) async fn start_http_server(address: String, port: u16) -> anyhow::Res
|
||||
builder = builder.mount("/audio", routes![api::speech]);
|
||||
// Health check and model info endpoints
|
||||
builder = builder.mount("/", routes![api::health, api::models]);
|
||||
// Shutdown endpoint
|
||||
builder = builder.manage(shutdown_flag);
|
||||
builder = builder.mount("/", routes![api::shutdown]);
|
||||
|
||||
let _rocket = builder.launch().await?;
|
||||
|
||||
// Cleanup PID file when server exits
|
||||
cleanup_pid_file(port)?;
|
||||
shutdown_handler.abort();
|
||||
|
||||
builder.launch().await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
+287
@@ -0,0 +1,287 @@
|
||||
//! Process management module for AHA services
|
||||
//!
|
||||
//! This module provides functionality for:
|
||||
//! - Managing PID files for service tracking
|
||||
//! - Discovering running AHA services
|
||||
//! - Service information display
|
||||
|
||||
use std::fs;
|
||||
use std::path::PathBuf;
|
||||
|
||||
use anyhow::{Result, anyhow};
|
||||
use sysinfo::{Pid, ProcessesToUpdate, System};
|
||||
|
||||
/// Service information structure
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ServiceInfo {
|
||||
/// Service unique identifier (format: pid@port)
|
||||
pub service_id: String,
|
||||
/// Process ID
|
||||
pub pid: u32,
|
||||
/// Model name (if available)
|
||||
pub model: Option<String>,
|
||||
/// Listen port
|
||||
pub port: u16,
|
||||
/// Listen address
|
||||
pub address: String,
|
||||
/// Service status
|
||||
pub status: ServiceStatus,
|
||||
}
|
||||
|
||||
/// Service status
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub enum ServiceStatus {
|
||||
Running,
|
||||
Stopping,
|
||||
Unknown,
|
||||
}
|
||||
|
||||
/// Get the PID file directory
|
||||
///
|
||||
/// Returns the appropriate directory for storing PID files:
|
||||
/// - Linux/macOS: $XDG_RUNTIME_DIR/aha or ~/.aha/run
|
||||
/// - Windows: %LOCALAPPDATA%\aha\run
|
||||
pub fn get_pid_dir() -> Result<PathBuf> {
|
||||
#[cfg(unix)]
|
||||
{
|
||||
// Try XDG_RUNTIME_DIR first
|
||||
if let Ok(runtime_dir) = std::env::var("XDG_RUNTIME_DIR") {
|
||||
let pid_dir = PathBuf::from(runtime_dir).join("aha");
|
||||
fs::create_dir_all(&pid_dir)?;
|
||||
return Ok(pid_dir);
|
||||
}
|
||||
|
||||
// Fallback to ~/.aha/run
|
||||
let home = dirs::home_dir().ok_or_else(|| anyhow!("Cannot determine home directory"))?;
|
||||
let pid_dir = home.join(".aha").join("run");
|
||||
fs::create_dir_all(&pid_dir)?;
|
||||
Ok(pid_dir)
|
||||
}
|
||||
|
||||
#[cfg(windows)]
|
||||
{
|
||||
let local_app_data = std::env::var("LOCALAPPDATA")
|
||||
.map_err(|_| anyhow!("Cannot determine LOCALAPPDATA directory"))?;
|
||||
let pid_dir = PathBuf::from(local_app_data).join("aha").join("run");
|
||||
fs::create_dir_all(&pid_dir)?;
|
||||
Ok(pid_dir)
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a PID file for the current service
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `pid` - Process ID
|
||||
/// * `port` - Listen port
|
||||
pub fn create_pid_file(pid: u32, port: u16) -> Result<()> {
|
||||
let pid_dir = get_pid_dir()?;
|
||||
let pid_file = pid_dir.join(format!("{}.pid", port));
|
||||
|
||||
let content = format!("{}\n", pid);
|
||||
fs::write(&pid_file, content)?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Clean up a PID file
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `port` - Listen port
|
||||
pub fn cleanup_pid_file(port: u16) -> Result<()> {
|
||||
let pid_dir = get_pid_dir()?;
|
||||
let pid_file = pid_dir.join(format!("{}.pid", port));
|
||||
|
||||
if pid_file.exists() {
|
||||
fs::remove_file(&pid_file)?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Get the PID from a PID file
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `port` - Listen port
|
||||
pub fn get_pid_from_file(port: u16) -> Option<u32> {
|
||||
let pid_dir = get_pid_dir().ok()?;
|
||||
let pid_file = pid_dir.join(format!("{}.pid", port));
|
||||
|
||||
if !pid_file.exists() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let content = fs::read_to_string(&pid_file).ok()?;
|
||||
content.trim().parse::<u32>().ok()
|
||||
}
|
||||
|
||||
/// Check if a process is an AHA service
|
||||
///
|
||||
/// Verifies that the process command line contains "aha serv" or "aha cli"
|
||||
fn is_aha_process(sys: &System, pid: Pid) -> bool {
|
||||
if let Some(process) = sys.process(pid) {
|
||||
let cmd = process.cmd();
|
||||
let cmd_str: String = cmd.iter()
|
||||
.filter_map(|s| s.to_str())
|
||||
.collect::<Vec<&str>>()
|
||||
.join(" ");
|
||||
return cmd_str.contains("aha serv") || cmd_str.contains("aha cli");
|
||||
}
|
||||
false
|
||||
}
|
||||
|
||||
/// Find all running AHA services
|
||||
///
|
||||
/// Returns a list of ServiceInfo for all running AHA services
|
||||
pub fn find_aha_services() -> Result<Vec<ServiceInfo>> {
|
||||
let mut services = Vec::new();
|
||||
let mut sys = System::new_all();
|
||||
sys.refresh_processes(ProcessesToUpdate::All, true);
|
||||
|
||||
// First, try to discover services from PID files
|
||||
let pid_dir = get_pid_dir()?;
|
||||
if let Ok(entries) = fs::read_dir(&pid_dir) {
|
||||
for entry in entries.flatten() {
|
||||
let path = entry.path();
|
||||
if path.extension().and_then(|s| s.to_str()) != Some("pid") {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Extract port from filename
|
||||
let port_str = path.file_stem()
|
||||
.and_then(|s| s.to_str())
|
||||
.unwrap_or("");
|
||||
let port: u16 = port_str.parse().unwrap_or(0);
|
||||
|
||||
if port == 0 {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Read PID from file
|
||||
if let Ok(content) = fs::read_to_string(&path) {
|
||||
if let Ok(pid) = content.trim().parse::<u32>() {
|
||||
let sys_pid = Pid::from_u32(pid);
|
||||
if is_aha_process(&sys, sys_pid) {
|
||||
services.push(ServiceInfo {
|
||||
service_id: format!("{}@{}", pid, port),
|
||||
pid,
|
||||
model: None, // TODO: Extract from command line
|
||||
port,
|
||||
address: "127.0.0.1".to_string(),
|
||||
status: ServiceStatus::Running,
|
||||
});
|
||||
} else {
|
||||
// Stale PID file, remove it
|
||||
let _ = fs::remove_file(&path);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Fallback: scan processes for AHA services
|
||||
for (pid, process) in sys.processes() {
|
||||
if services.iter().any(|s| s.pid == pid.as_u32()) {
|
||||
continue; // Already found via PID file
|
||||
}
|
||||
|
||||
let cmd = process.cmd();
|
||||
let cmd_str: String = cmd.iter()
|
||||
.filter_map(|s| s.to_str())
|
||||
.collect::<Vec<&str>>()
|
||||
.join(" ");
|
||||
|
||||
if cmd_str.contains("aha serv") || cmd_str.contains("aha cli") {
|
||||
// Try to extract port from command line
|
||||
let port_str = cmd.iter()
|
||||
.position(|s| s.to_str() == Some("--port"))
|
||||
.and_then(|i| cmd.get(i + 1))
|
||||
.and_then(|s| s.to_str());
|
||||
let port = port_str
|
||||
.and_then(|s| s.parse::<u16>().ok())
|
||||
.unwrap_or(10100);
|
||||
|
||||
services.push(ServiceInfo {
|
||||
service_id: format!("{}@{}", pid.as_u32(), port),
|
||||
pid: pid.as_u32(),
|
||||
model: None,
|
||||
port,
|
||||
address: "127.0.0.1".to_string(),
|
||||
status: ServiceStatus::Running,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
Ok(services)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_get_pid_dir() {
|
||||
let pid_dir = get_pid_dir();
|
||||
assert!(pid_dir.is_ok());
|
||||
let dir = pid_dir.unwrap();
|
||||
assert!(dir.exists());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_create_and_cleanup_pid_file() {
|
||||
let port = 19999;
|
||||
create_pid_file(12345, port).unwrap();
|
||||
let pid = get_pid_from_file(port);
|
||||
assert_eq!(pid, Some(12345));
|
||||
cleanup_pid_file(port).unwrap();
|
||||
let pid = get_pid_from_file(port);
|
||||
assert_eq!(pid, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_get_pid_from_file_nonexistent() {
|
||||
let port = 19998; // Use a port that likely doesn't have a PID file
|
||||
let pid = get_pid_from_file(port);
|
||||
assert_eq!(pid, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_service_status_debug() {
|
||||
// Test ServiceStatus Debug implementation
|
||||
assert_eq!(format!("{:?}", ServiceStatus::Running), "Running");
|
||||
assert_eq!(format!("{:?}", ServiceStatus::Stopping), "Stopping");
|
||||
assert_eq!(format!("{:?}", ServiceStatus::Unknown), "Unknown");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_service_info_clone() {
|
||||
let service = ServiceInfo {
|
||||
service_id: "12345@10100".to_string(),
|
||||
pid: 12345,
|
||||
model: Some("qwen3-0.6b".to_string()),
|
||||
port: 10100,
|
||||
address: "127.0.0.1".to_string(),
|
||||
status: ServiceStatus::Running,
|
||||
};
|
||||
let service_clone = service.clone();
|
||||
assert_eq!(service_clone.service_id, "12345@10100");
|
||||
assert_eq!(service_clone.pid, 12345);
|
||||
assert_eq!(service_clone.model, Some("qwen3-0.6b".to_string()));
|
||||
assert_eq!(service_clone.port, 10100);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_find_aha_services() {
|
||||
// This test will find actual running AHA services or return empty
|
||||
let services = find_aha_services();
|
||||
assert!(services.is_ok());
|
||||
let services_list = services.unwrap();
|
||||
// We can't assert specific services here since it depends on what's running
|
||||
// but we can verify the structure is correct
|
||||
for service in services_list {
|
||||
assert!(!service.service_id.is_empty());
|
||||
assert!(service.pid > 0);
|
||||
assert!(service.port > 0);
|
||||
assert!(!service.address.is_empty());
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user