Initial vLLM agent strategy
This commit is contained in:
10
.dockerignore
Normal file
10
.dockerignore
Normal file
@@ -0,0 +1,10 @@
|
||||
__pycache__/
|
||||
*.pyc
|
||||
*.pyo
|
||||
*.pyd
|
||||
.pytest_cache/
|
||||
.env
|
||||
config.local.json
|
||||
state.db
|
||||
/data/
|
||||
.git/
|
||||
19
.gitignore
vendored
Normal file
19
.gitignore
vendored
Normal file
@@ -0,0 +1,19 @@
|
||||
# 密钥和本地配置(包含 token)
|
||||
config.local.json
|
||||
|
||||
# SQLite 数据库
|
||||
state*.db
|
||||
|
||||
# Python
|
||||
__pycache__/
|
||||
*.pyc
|
||||
*.pyo
|
||||
.venv/
|
||||
venv/
|
||||
|
||||
# 环境变量文件
|
||||
.env
|
||||
|
||||
# IDE
|
||||
.idea/
|
||||
.vscode/
|
||||
15
Dockerfile
Normal file
15
Dockerfile
Normal file
@@ -0,0 +1,15 @@
|
||||
FROM modelhubxc-4pd.tencentcloudcr.com/xc_agent_platform/python:3.11-slim
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
COPY requirements.txt ./
|
||||
RUN pip install --no-cache-dir -r requirements.txt
|
||||
|
||||
COPY app ./app
|
||||
COPY seeds ./seeds
|
||||
COPY scripts ./scripts
|
||||
|
||||
EXPOSE 8080
|
||||
|
||||
ENV PYTHONUNBUFFERED=1
|
||||
CMD ["python", "-m", "app.main"]
|
||||
221
README.md
Normal file
221
README.md
Normal file
@@ -0,0 +1,221 @@
|
||||
# vLLM Agent Strategy
|
||||
|
||||
全新版 vLLM-only 智能体策略。目标是最大化验证成功数量:优先跑悬赏模型,其次 Promote 已入库模型,再复用别人已下载模型,最后才自己下载新模型。
|
||||
|
||||
## 优先级
|
||||
|
||||
1. `bounty` / 悬赏模型:最高优先级,默认可直接提交。
|
||||
2. `promote`:已经在 ModelHub 入库/验证成功的模型。
|
||||
3. `downloaded_by_others`:别人已经下载成功的模型。
|
||||
4. `self_download`:自己下载,最多 8 个 active 下载任务。
|
||||
|
||||
任务粒度是 `model_id + gpu_type + engine`,而不是整个 model。`MAX_USER_TASKS` 默认 2000。
|
||||
|
||||
## 配置文件
|
||||
|
||||
不用每次手动设置一堆环境变量。复制示例配置:
|
||||
|
||||
```bash
|
||||
copy config.example.json config.local.json
|
||||
```
|
||||
|
||||
然后编辑 `config.local.json`,填入 token、用户信息、目标算力等配置。
|
||||
|
||||
默认会按顺序读取:
|
||||
|
||||
```text
|
||||
config.local.json
|
||||
config.json
|
||||
```
|
||||
|
||||
也可以用环境变量指定配置文件:
|
||||
|
||||
```bash
|
||||
set STRATEGY_CONFIG_FILE=D:\4paradigm\vllm_agent_strategy\config.local.json
|
||||
```
|
||||
|
||||
环境变量仍然可用,并且优先级高于配置文件,适合临时覆盖:
|
||||
|
||||
```bash
|
||||
set SUBMIT_DRY_RUN=false
|
||||
```
|
||||
|
||||
重要配置:
|
||||
|
||||
```json
|
||||
{
|
||||
"AUTH_TOKEN": "",
|
||||
"HF_TOKEN": "",
|
||||
"CONFTEST_API_TOKEN": "",
|
||||
"EMAIL": "",
|
||||
"USER_ID": "",
|
||||
"CONTRIBUTORS": "",
|
||||
"STRATEGY_ID": "",
|
||||
"SUBMIT_DRY_RUN": true,
|
||||
"TARGET_MACHINE_NAMES": "MTT S4000,Hygon K100,Kunlunxin P800"
|
||||
}
|
||||
```
|
||||
|
||||
如果平台实际字段不是 `strategyId`,改 `CONTEST_TASK_STRATEGY_FIELD` 即可。
|
||||
|
||||
## 本地 dry run
|
||||
|
||||
```bash
|
||||
python -m scripts.init_db
|
||||
python -m scripts.load_seed --file seeds/bounty_models.example.csv --origin manual_bounty
|
||||
python -m scripts.dry_run_plan --limit 50
|
||||
python -m scripts.smoke_check --no-submit
|
||||
```
|
||||
|
||||
## 启动服务
|
||||
|
||||
```bash
|
||||
python -m app.main
|
||||
```
|
||||
|
||||
健康检查:
|
||||
|
||||
```text
|
||||
GET /health -> {"status":"ok"}
|
||||
```
|
||||
|
||||
## Docker
|
||||
|
||||
```bash
|
||||
docker build -t vllm-agent-strategy:local .
|
||||
docker run --rm -p 8080:8080 --env-file .env vllm-agent-strategy:local
|
||||
```
|
||||
|
||||
## 平台部署
|
||||
|
||||
平台部署流程:
|
||||
|
||||
```text
|
||||
推代码到 dev.modelhub.org.cn/Gitea
|
||||
-> 打 tag
|
||||
-> 调 agent_platform create 创建策略
|
||||
-> 拿返回的 id 作为 STRATEGY_ID
|
||||
-> 等构建 ready
|
||||
-> deploy
|
||||
-> logs 查看运行日志
|
||||
```
|
||||
|
||||
### 1. 推代码并打 tag
|
||||
|
||||
把本目录作为策略仓库根目录推到 `dev.modelhub.org.cn`,确保仓库根目录有:
|
||||
|
||||
```text
|
||||
Dockerfile
|
||||
requirements.txt
|
||||
app/
|
||||
seeds/
|
||||
scripts/
|
||||
```
|
||||
|
||||
示例:
|
||||
|
||||
```bash
|
||||
git init
|
||||
git add .
|
||||
git commit -m "Initial vLLM agent strategy"
|
||||
git remote add origin https://dev.modelhub.org.cn/<user>/vllm_agent_strategy.git
|
||||
git tag v1.0.0
|
||||
git push origin main --tags
|
||||
```
|
||||
|
||||
如果已有仓库,按实际分支名和远端地址调整。
|
||||
|
||||
### 2. 创建策略并获取 STRATEGY_ID
|
||||
|
||||
可以使用脚本:
|
||||
|
||||
```bash
|
||||
python -m scripts.agent_platform_api create \
|
||||
--name vllm-agent-strategy \
|
||||
--repo-url https://dev.modelhub.org.cn/<user>/vllm_agent_strategy \
|
||||
--tag v1.0.0
|
||||
```
|
||||
|
||||
脚本默认从 `AGENT_PLATFORM_TOKEN` 或 `AUTH_TOKEN` 读取 token,也可以显式传:
|
||||
|
||||
```bash
|
||||
python -m scripts.agent_platform_api --token <TOKEN> create \
|
||||
--name vllm-agent-strategy \
|
||||
--repo-url https://dev.modelhub.org.cn/<user>/vllm_agent_strategy \
|
||||
--tag v1.0.0
|
||||
```
|
||||
|
||||
返回 JSON 中的:
|
||||
|
||||
```text
|
||||
id
|
||||
```
|
||||
|
||||
就是 `STRATEGY_ID`。
|
||||
|
||||
### 3. 查看构建状态
|
||||
|
||||
```bash
|
||||
python -m scripts.agent_platform_api get --strategy-id <STRATEGY_ID>
|
||||
python -m scripts.agent_platform_api sync --strategy-id <STRATEGY_ID>
|
||||
```
|
||||
|
||||
等状态变成:
|
||||
|
||||
```text
|
||||
ready
|
||||
```
|
||||
|
||||
### 4. 部署策略
|
||||
|
||||
```bash
|
||||
python -m scripts.agent_platform_api deploy --strategy-id <STRATEGY_ID>
|
||||
```
|
||||
|
||||
### 5. 查看日志
|
||||
|
||||
```bash
|
||||
python -m scripts.agent_platform_api logs --strategy-id <STRATEGY_ID> --limit 100
|
||||
```
|
||||
|
||||
### 6. 停止策略
|
||||
|
||||
```bash
|
||||
python -m scripts.agent_platform_api stop --strategy-id <STRATEGY_ID>
|
||||
```
|
||||
|
||||
### 7. 本地和平台 STRATEGY_ID 的区别
|
||||
|
||||
本地 dry-run 可以在 `config.local.json` 里填:
|
||||
|
||||
```json
|
||||
"STRATEGY_ID": "local-dry-run"
|
||||
```
|
||||
|
||||
平台部署时,平台会注入真实 `STRATEGY_ID` 环境变量。环境变量优先级高于配置文件。
|
||||
|
||||
## Agent Platform API helper
|
||||
|
||||
新增脚本:
|
||||
|
||||
```text
|
||||
scripts/agent_platform_api.py
|
||||
```
|
||||
|
||||
支持:
|
||||
|
||||
```bash
|
||||
python -m scripts.agent_platform_api me
|
||||
python -m scripts.agent_platform_api list
|
||||
python -m scripts.agent_platform_api create --name ... --repo-url ... --tag ...
|
||||
python -m scripts.agent_platform_api get --strategy-id ...
|
||||
python -m scripts.agent_platform_api sync --strategy-id ...
|
||||
python -m scripts.agent_platform_api deploy --strategy-id ...
|
||||
python -m scripts.agent_platform_api logs --strategy-id ...
|
||||
python -m scripts.agent_platform_api stop --strategy-id ...
|
||||
python -m scripts.agent_platform_api delete --strategy-id ...
|
||||
```
|
||||
|
||||
## GGUF / llamacpp
|
||||
|
||||
本版本只实现 vLLM。GGUF/llamacpp 后续可基于 `tasks.engine` 增加分支。
|
||||
573
STRATEGY_DESIGN.md
Normal file
573
STRATEGY_DESIGN.md
Normal file
@@ -0,0 +1,573 @@
|
||||
# vLLM/未来 GGUF 智能体提交策略设计
|
||||
|
||||
## 1. 背景和目标
|
||||
|
||||
当前目标不是“尽量多下载模型”,而是“尽量多产生验证成功的提交”。现在只有智能体提交的任务才会触发验证,因此本地脚本主要用于理解接口和策略,最终提交逻辑应在智能体服务中执行。
|
||||
|
||||
旧流程大致是:
|
||||
|
||||
```text
|
||||
fetchModel_db.py -> Download_db.py -> Submit_db.py / Promote_db.py
|
||||
```
|
||||
|
||||
主要问题:
|
||||
|
||||
1. `Download_db.py` 以下载为主线,下载池满时会阻塞,导致后续模型无法继续扫描。
|
||||
2. `Promote_db.py` 思路正确,但候选来源过窄,只依赖本地 `downloaded` 表。
|
||||
3. `Submit_db.py / Promote_db.py` 使用模型级 `in_db` 状态,不适合按 `model_id + gpu_type` 最大化 2000 个任务额度。
|
||||
4. 页面上的“批量模型悬赏”没有直接给模型 ID,只提示目标算力,因此不能直接爬出模型任务。
|
||||
|
||||
新版策略应该是:
|
||||
|
||||
```text
|
||||
手动设置悬赏/目标算力
|
||||
-> 爬 ModelHub 已入库模型
|
||||
-> 找出目标算力未适配的模型
|
||||
-> 按模型类型选择 vLLM 或未来 GGUF/llamacpp 路线
|
||||
-> 生成 model_id + gpu_type 任务
|
||||
-> 智能体优先提交高成功率任务
|
||||
```
|
||||
|
||||
当前实现阶段只做 vLLM;GGUF/llamacpp 先在文档里明确策略,后续再实现。
|
||||
|
||||
---
|
||||
|
||||
## 2. 总体优先级
|
||||
|
||||
推荐优先级如下:
|
||||
|
||||
```text
|
||||
P0: 手动配置的悬赏目标算力 + ModelHub 已入库 + 目标算力未适配的模型
|
||||
P1: 普通目标算力未适配的已入库模型
|
||||
P2: admin/tools/model-download 页面中别人已经下载成功的模型
|
||||
P3: 自己下载 HuggingFace 新模型
|
||||
```
|
||||
|
||||
其中 P0/P1 是 Promote 思路的升级版:模型已经在平台存在,说明至少已经被验证/入库过;如果某个目标算力还没有适配成功,就提交到该目标算力。
|
||||
|
||||
P2 只能说明模型文件已成功下载,不一定说明模型已验证成功,因此优先级低于已入库未适配模型。
|
||||
|
||||
P3 自己下载风险最高、成本最高,放最后。
|
||||
|
||||
---
|
||||
|
||||
## 3. Step 1:目标算力手动配置
|
||||
|
||||
无需自动爬取 `https://modelhub.org.cn/#/adaptation-model` 页面里的“批量模型悬赏”。原因:
|
||||
|
||||
1. 页面没有直接给模型 ID。
|
||||
2. 手动看一眼目标算力更快、更稳定。
|
||||
3. 自动爬页面还要解析 hash route / XHR,维护成本高。
|
||||
|
||||
建议通过环境变量或配置文件手动指定目标算力。
|
||||
|
||||
### 3.1 建议配置
|
||||
|
||||
```text
|
||||
TARGET_MACHINE_NAMES=MTT S4000,Hygon K100,Kunlunxin P800,Biren 166M
|
||||
TARGET_TASK_LEVEL=文本生成
|
||||
```
|
||||
|
||||
同时需要维护一张页面展示名到提交 GPU alias 的映射表。
|
||||
|
||||
示例:
|
||||
|
||||
```python
|
||||
MACHINE_TO_GPU_ALIAS = {
|
||||
"MTT S4000": "s4000",
|
||||
"Hygon K100": "k100",
|
||||
"Kunlunxin P800": "p800",
|
||||
"Biren 166M": "biren166m",
|
||||
"Ascend 910B": "910b",
|
||||
"Iluvatar BI100": "bi100",
|
||||
"Iluvatar BI150": "bi150",
|
||||
"MetaX C500": "c500",
|
||||
"Iluvatar MRV100": "mrv100",
|
||||
"Cambricon MLU370-X4": "mlu370-x4",
|
||||
"Cambricon MLU370-X8": "mlu370-x8",
|
||||
}
|
||||
```
|
||||
|
||||
这张表需要根据 `/api/computility/models/list/page/vo` 返回的 `machines[].machineName` 实际值校准。
|
||||
|
||||
---
|
||||
|
||||
## 4. Step 2:爬 ModelHub 已入库模型
|
||||
|
||||
参考 `D:\4paradigm\fetch_not_adapted_models.py`。
|
||||
|
||||
使用接口:
|
||||
|
||||
```text
|
||||
POST https://modelhub.org.cn/api/computility/models/list/page/vo
|
||||
```
|
||||
|
||||
请求示例:
|
||||
|
||||
```json
|
||||
{
|
||||
"current": 1,
|
||||
"pageSize": 100,
|
||||
"searchText": ""
|
||||
}
|
||||
```
|
||||
|
||||
返回里重点字段:
|
||||
|
||||
```text
|
||||
modelId
|
||||
machines[].machineName
|
||||
taskLevelChineseName
|
||||
taskLevelsInfo.taskLevelChineseName
|
||||
```
|
||||
|
||||
扫描逻辑:
|
||||
|
||||
1. 分页拉取平台已入库模型。
|
||||
2. 只保留文本生成模型:
|
||||
|
||||
```python
|
||||
model_task = (
|
||||
model.get("taskLevelChineseName")
|
||||
or model.get("taskLevelsInfo", {}).get("taskLevelChineseName")
|
||||
or ""
|
||||
)
|
||||
|
||||
if "文本生成" not in model_task:
|
||||
skip
|
||||
```
|
||||
|
||||
3. 读取已适配算力:
|
||||
|
||||
```python
|
||||
machines = model.get("machines", [])
|
||||
adapted_machines = [m.get("machineName") for m in machines if m.get("machineName")]
|
||||
```
|
||||
|
||||
4. 对每个手动配置的目标算力,如果目标算力不在 `adapted_machines` 中,则生成候选任务。
|
||||
|
||||
```text
|
||||
model_id + target_gpu_alias
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 5. Step 3 / Step 4:区分 vLLM 模型和 GGUF 模型
|
||||
|
||||
这里需要纠正之前的“直接过滤 GGUF”策略。
|
||||
|
||||
当前阶段只实现 vLLM,所以 GGUF 先不提交,但不能简单认为 GGUF 没用。GGUF 模型应该走 llamacpp,而且根据经验成功率更高的芯片是:
|
||||
|
||||
```text
|
||||
910b
|
||||
s4000
|
||||
```
|
||||
|
||||
因此模型类型判断应该是:
|
||||
|
||||
```text
|
||||
如果 model_id 或相关权重文件包含 gguf:
|
||||
标记为 engine=llamacpp_candidate
|
||||
当前 vLLM-only 阶段不提交
|
||||
后续 GGUF 优先提交到 910b、s4000
|
||||
否则:
|
||||
标记为 engine=vllm
|
||||
当前阶段可提交
|
||||
```
|
||||
|
||||
### 5.1 当前 vLLM-only 阶段
|
||||
|
||||
当前阶段策略:
|
||||
|
||||
```text
|
||||
非 GGUF 模型:进入 vLLM 任务队列
|
||||
GGUF 模型:写入候选表,但 status=skipped_gguf 或 engine=llamacpp_candidate,不提交
|
||||
```
|
||||
|
||||
不要把 GGUF 完全丢弃,避免后续 GGUF/llamacpp 优化时还要重新爬。
|
||||
|
||||
### 5.2 未来 GGUF/llamacpp 阶段
|
||||
|
||||
后续实现 GGUF 后,策略应为:
|
||||
|
||||
```text
|
||||
GGUF 模型优先目标 GPU:910b, s4000
|
||||
框架:llamacpp
|
||||
权重文件:需要具体 .gguf 文件名
|
||||
```
|
||||
|
||||
可以参考旧文件:
|
||||
|
||||
- `D:\4paradigm\fetchModel2_gguf.py`
|
||||
- `D:\4paradigm\run_llamacpp.py`
|
||||
- `D:\4paradigm\Utils_db.py` 中的 `choose_llamacpp_config()`
|
||||
|
||||
GGUF 任务粒度未来应是:
|
||||
|
||||
```text
|
||||
model_id + weight_file + gpu_type + engine=llamacpp
|
||||
```
|
||||
|
||||
而不是只用 `model_id + gpu_type`。
|
||||
|
||||
---
|
||||
|
||||
## 6. Step 5:别人已下载成功模型来源
|
||||
|
||||
别人已下载成功的模型可以从页面:
|
||||
|
||||
```text
|
||||
https://modelhub.org.cn/admin/tools/model-download
|
||||
```
|
||||
|
||||
页面中选择:
|
||||
|
||||
```text
|
||||
状态 = 成功
|
||||
pageSize = 50
|
||||
```
|
||||
|
||||
可以看到类似:
|
||||
|
||||
```text
|
||||
Aimiliya0011/lora_409
|
||||
```
|
||||
|
||||
这类模型说明文件已经在平台下载成功,可以作为 P2 候选。
|
||||
|
||||
### 6.1 已确认接口
|
||||
|
||||
页面背后的接口已确认:
|
||||
|
||||
```text
|
||||
GET https://modelhub.org.cn/adminApi/async/task/model-download-task
|
||||
```
|
||||
|
||||
查询参数:
|
||||
|
||||
```text
|
||||
current=1
|
||||
pageSize=50
|
||||
status=SUCCESS
|
||||
```
|
||||
|
||||
示例:
|
||||
|
||||
```text
|
||||
https://modelhub.org.cn/adminApi/async/task/model-download-task?current=1&pageSize=50&status=SUCCESS
|
||||
```
|
||||
|
||||
该接口需要认证请求头。直接在浏览器地址栏打开通常会返回 401:
|
||||
|
||||
```json
|
||||
{"code":401,"data":null,"message":"请提供有效的JWT Token或Xc-Token"}
|
||||
```
|
||||
|
||||
代码里需要像旧版 `Utils_db.py` 一样携带:
|
||||
|
||||
```text
|
||||
Authorization: Bearer <AUTH_TOKEN>
|
||||
Xc-Token: <AUTH_TOKEN>
|
||||
Content-Type: application/json
|
||||
```
|
||||
|
||||
Response 结构示例:
|
||||
|
||||
```json
|
||||
{
|
||||
"code": 0,
|
||||
"data": {
|
||||
"records": [
|
||||
{
|
||||
"id": "3464438",
|
||||
"name": "DOWNLOAD_MODEL",
|
||||
"args": {
|
||||
"hf_token": "hf_***gE",
|
||||
"storage_type": "juicefs",
|
||||
"user_email": "adminProxy@4paradigm.com",
|
||||
"user_id": 82,
|
||||
"source": "HuggingFace",
|
||||
"model_id": "bryn107/llama3-8b-indonesian-legal-bot"
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
关键字段:
|
||||
|
||||
```text
|
||||
records 路径: data.records
|
||||
model_id 路径: records[].args.model_id
|
||||
source 路径: records[].args.source
|
||||
任务 id 路径: records[].id
|
||||
```
|
||||
|
||||
注意字段名是 `args.model_id`,不是顶层 `modelId`。
|
||||
|
||||
分页停止条件:
|
||||
|
||||
```text
|
||||
优先使用 data.total 判断是否结束;
|
||||
如果没有 total,则 records 为空或 len(records) < pageSize 时结束。
|
||||
```
|
||||
|
||||
### 6.2 source 规范化
|
||||
|
||||
接口返回的来源可能是页面展示值,例如:
|
||||
|
||||
```text
|
||||
HuggingFace
|
||||
ModelScope
|
||||
```
|
||||
|
||||
新版内部建议规范化为旧脚本使用的枚举:
|
||||
|
||||
```python
|
||||
SOURCE_MAP = {
|
||||
"HuggingFace": "HUGGING_FACE",
|
||||
"ModelScope": "MODEL_SCOPE",
|
||||
}
|
||||
```
|
||||
|
||||
### 6.3 爬取后处理
|
||||
|
||||
爬到别人已下载成功模型后:
|
||||
|
||||
```text
|
||||
origin=downloaded_success_page
|
||||
priority=20
|
||||
known_downloaded_by_others=1
|
||||
```
|
||||
|
||||
然后根据模型类型:
|
||||
|
||||
```text
|
||||
非 GGUF -> engine=vllm,进入 vLLM 候选
|
||||
GGUF -> engine=llamacpp_candidate,当前阶段暂存,不提交
|
||||
```
|
||||
|
||||
注意:别人已下载成功不代表已验证成功,所以优先级低于 ModelHub 已入库未适配模型。
|
||||
|
||||
---
|
||||
|
||||
## 7. 下载策略
|
||||
|
||||
自己下载仍然保留,但只作为最后补充。
|
||||
|
||||
规则:
|
||||
|
||||
```text
|
||||
MAX_SELF_DOWNLOADS=8
|
||||
```
|
||||
|
||||
即每个人后台同时下载的模型最多 8 个。
|
||||
|
||||
旧问题是:下载池满时 `Download_db.py` 会 sleep,导致后面的模型无法继续扫描。
|
||||
|
||||
新版规则:
|
||||
|
||||
```text
|
||||
下载池满,只停止创建新的自下载任务;不停止:
|
||||
- 爬 ModelHub 已入库模型
|
||||
- 爬别人下载成功模型
|
||||
- 提交 bounty/promote 任务
|
||||
- 扫描后续候选
|
||||
```
|
||||
|
||||
遇到下载相关返回:
|
||||
|
||||
```text
|
||||
code=0 -> 自己创建下载任务,active_self_slot +1
|
||||
code=40000 -> 别人/已有任务正在下载,不占自己的 8 个槽,后续轮询
|
||||
code=60004 -> 已下载过,查状态,成功则进入提交候选
|
||||
code=60005 -> 自己下载池满,当前任务 backoff,不全局 sleep
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 8. 提交策略
|
||||
|
||||
所有正式提交都应由智能体服务执行。
|
||||
|
||||
任务粒度:
|
||||
|
||||
```text
|
||||
vLLM: model_id + gpu_type + engine=vllm
|
||||
llamacpp: model_id + weight_file + gpu_type + engine=llamacpp # 未来
|
||||
```
|
||||
|
||||
每人最多 2000 个任务,建议保守按 materialized task 数量控制,避免超量。
|
||||
|
||||
提交返回处理参考 `Utils_db.submit()`:
|
||||
|
||||
```text
|
||||
code=0 && message=ok -> submitted
|
||||
code=40000 && message 包含 正在验证中 -> conflict / running
|
||||
code=60007 -> 当前 task/gpu queue_full_backoff
|
||||
其他 -> failed 或短 backoff
|
||||
```
|
||||
|
||||
关键原则:
|
||||
|
||||
```text
|
||||
任何一个 GPU 队列满,都不能让整个策略 sleep。
|
||||
只给当前 task 或当前 gpu_type 设置 backoff,然后继续提交其他 GPU / 其他模型。
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 9. 推荐新版候选来源和状态
|
||||
|
||||
### 9.1 candidate origin
|
||||
|
||||
```text
|
||||
bounty_target_not_adapted # 手动目标算力 + 已入库模型 + 目标算力未适配
|
||||
normal_target_not_adapted # 普通目标算力未适配
|
||||
model_download_success_page # admin/tools/model-download 状态成功页面
|
||||
hf_self_download_seed # HuggingFace 新模型自下载候选
|
||||
manual_seed # 手动补充
|
||||
```
|
||||
|
||||
### 9.2 engine 类型
|
||||
|
||||
```text
|
||||
vllm
|
||||
llamacpp_candidate
|
||||
llamacpp # 未来正式实现
|
||||
```
|
||||
|
||||
### 9.3 task status
|
||||
|
||||
```text
|
||||
pending
|
||||
ready_to_submit
|
||||
waiting_download
|
||||
submitted
|
||||
conflict
|
||||
queue_full_backoff
|
||||
failed
|
||||
skipped_gguf
|
||||
skipped_unsupported_gpu
|
||||
budget_blocked
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 10. 当前代码需要如何调整
|
||||
|
||||
当前 `vllm_agent_strategy` 已经实现了 seed fallback 和 vLLM 调度骨架,但需要根据本设计做以下调整。
|
||||
|
||||
### 10.1 新增 target-not-adapted crawler
|
||||
|
||||
新增模块:
|
||||
|
||||
```text
|
||||
app/clients/modelhub.py
|
||||
list_modelhub_models(page, page_size, search_text="")
|
||||
fetch_not_adapted_models(machine_name, task_level="文本生成")
|
||||
```
|
||||
|
||||
或者单独:
|
||||
|
||||
```text
|
||||
app/crawlers/not_adapted.py
|
||||
```
|
||||
|
||||
功能:
|
||||
|
||||
```text
|
||||
读取 TARGET_MACHINE_NAMES
|
||||
分页调用 /api/computility/models/list/page/vo
|
||||
筛选文本生成
|
||||
按 machines 判断目标算力未适配
|
||||
区分 GGUF / 非 GGUF
|
||||
写入 candidates
|
||||
```
|
||||
|
||||
### 10.2 新增 machine mapping
|
||||
|
||||
新增:
|
||||
|
||||
```text
|
||||
app/domain/machines.py
|
||||
```
|
||||
|
||||
包含:
|
||||
|
||||
```python
|
||||
MACHINE_TO_GPU_ALIAS = {...}
|
||||
```
|
||||
|
||||
### 10.3 调整 GGUF 处理
|
||||
|
||||
当前 vLLM-only 阶段:
|
||||
|
||||
```text
|
||||
GGUF 不提交,但写入 candidates/tasks 或 gguf_candidates 表
|
||||
状态为 skipped_gguf / llamacpp_candidate
|
||||
```
|
||||
|
||||
后续 GGUF 实现时,只处理:
|
||||
|
||||
```text
|
||||
gpu_alias in [910b, s4000]
|
||||
engine=llamacpp
|
||||
```
|
||||
|
||||
### 10.4 新增 model-download success crawler
|
||||
|
||||
新增模块:
|
||||
|
||||
```text
|
||||
app/crawlers/download_success.py
|
||||
```
|
||||
|
||||
接口已确认:
|
||||
|
||||
```text
|
||||
GET /adminApi/async/task/model-download-task?current=<page>&pageSize=50&status=SUCCESS
|
||||
```
|
||||
|
||||
功能:
|
||||
|
||||
```text
|
||||
携带 Authorization / Xc-Token 认证头
|
||||
按 current/pageSize/status=SUCCESS 分页
|
||||
从 data.records[].args.model_id 提取 model_id
|
||||
从 data.records[].args.source 提取 source 并规范化
|
||||
区分 GGUF / 非 GGUF
|
||||
写入 P2 candidates
|
||||
```
|
||||
|
||||
推荐在 `app/clients/modelhub.py` 增加:
|
||||
|
||||
```text
|
||||
list_success_download_tasks(current, page_size)
|
||||
fetch_success_downloaded_models(page_size=50, max_pages=None)
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 11. 推荐实施顺序
|
||||
|
||||
1. 先不要自动爬 adaptation-model 页面目标算力,改为手动配置 `TARGET_MACHINE_NAMES`。
|
||||
2. 实现 `fetch_not_adapted_models.py` 的逻辑到新版智能体目录。
|
||||
3. 新增 machineName -> gpu_alias 映射。
|
||||
4. 增加 GGUF 识别:当前 vLLM-only 阶段不提交 GGUF,但保留候选。
|
||||
5. 把 P0 任务来源从 `manual_bounty seed` 改为 `bounty_target_not_adapted crawler`。
|
||||
6. 保留 manual seed 作为兜底。
|
||||
7. 基于已确认的 `/adminApi/async/task/model-download-task?status=SUCCESS` 接口,实现别人下载成功模型爬取。
|
||||
8. 最后再考虑 GGUF/llamacpp:优先 910b 和 s4000。
|
||||
|
||||
---
|
||||
|
||||
## 12. 当前最重要的结论
|
||||
|
||||
1. 不需要爬取“批量模型悬赏”的目标算力,手动配置更快更稳。
|
||||
2. 真正有价值的是用手动目标算力去爬 ModelHub 已入库模型,筛选“目标算力未适配”的模型。
|
||||
3. GGUF 不能简单丢弃;当前 vLLM 阶段跳过提交,但应保留给未来 llamacpp。
|
||||
4. GGUF/llamacpp 后续优先跑 910b 和 s4000。
|
||||
5. 别人已下载成功页面是 P2 候选来源,接口已确认为 `GET /adminApi/async/task/model-download-task?current=<page>&pageSize=50&status=SUCCESS`,模型 ID 在 `data.records[].args.model_id`。
|
||||
6. 所有正式验证提交都应由智能体服务执行,不再依赖本地 Submit/Promote 脚本。
|
||||
1
app/__init__.py
Normal file
1
app/__init__.py
Normal file
@@ -0,0 +1 @@
|
||||
"""vLLM intelligent-agent strategy package."""
|
||||
1
app/clients/__init__.py
Normal file
1
app/clients/__init__.py
Normal file
@@ -0,0 +1 @@
|
||||
"""API clients."""
|
||||
153
app/clients/modelhub.py
Normal file
153
app/clients/modelhub.py
Normal file
@@ -0,0 +1,153 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
import logging
|
||||
|
||||
import requests
|
||||
|
||||
from app.settings import Settings
|
||||
|
||||
LOG = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ApiResult:
|
||||
result: str
|
||||
code: str | None = None
|
||||
message: str | None = None
|
||||
payload: dict[str, Any] | None = None
|
||||
|
||||
|
||||
class ModelHubClient:
|
||||
def __init__(self, settings: Settings):
|
||||
self.settings = settings
|
||||
self.session = requests.Session()
|
||||
|
||||
def _headers(self) -> dict[str, str]:
|
||||
return {
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {self.settings.auth_token}",
|
||||
"Xc-Token": self.settings.auth_token,
|
||||
}
|
||||
|
||||
def is_model_absent_from_modelhub(self, model_id: str) -> bool:
|
||||
url = self.settings.modelhub_api_base + "/computility/models/list/page/vo"
|
||||
data = {"current": 1, "pageSize": 20, "searchText": model_id}
|
||||
try:
|
||||
response = self.session.post(url, headers=self._headers(), json=data, timeout=self.settings.request_timeout_seconds)
|
||||
response.raise_for_status()
|
||||
total = int(response.json().get("data", {}).get("total", 0))
|
||||
return total == 0
|
||||
except Exception as exc:
|
||||
LOG.warning("modelhub existence check failed for %s: %s", model_id, exc)
|
||||
return True
|
||||
|
||||
def list_modelhub_models(self, current: int = 1, page_size: int = 100, search_text: str = "") -> dict[str, Any]:
|
||||
url = self.settings.modelhub_api_base + "/computility/models/list/page/vo"
|
||||
data = {"current": current, "pageSize": page_size, "searchText": search_text}
|
||||
response = self.session.post(url, headers=self._headers(), json=data, timeout=self.settings.request_timeout_seconds)
|
||||
response.raise_for_status()
|
||||
body = response.json()
|
||||
if body.get("code") != 0:
|
||||
raise RuntimeError(body.get("message") or body.get("msg") or "list_modelhub_models failed")
|
||||
return body.get("data") or {}
|
||||
|
||||
def list_success_download_tasks(self, current: int = 1, page_size: int = 50) -> dict[str, Any]:
|
||||
url = self.settings.modelhub_adminapi_base + "/async/task/model-download-task"
|
||||
params = {"current": current, "pageSize": page_size, "status": "SUCCESS"}
|
||||
response = self.session.get(url, headers=self._headers(), params=params, timeout=self.settings.request_timeout_seconds)
|
||||
response.raise_for_status()
|
||||
body = response.json()
|
||||
if body.get("code") != 0:
|
||||
raise RuntimeError(body.get("message") or "list_success_download_tasks failed")
|
||||
return body.get("data") or {}
|
||||
|
||||
def sync_from_huggingface(self, model_id: str) -> ApiResult:
|
||||
url = self.settings.modelhub_adminapi_base + "/computility/models/sync-from-hugging-face"
|
||||
data = {"modelId": model_id, "forceUpdate": False, "operatorEmail": self.settings.email}
|
||||
try:
|
||||
response = self.session.post(url, headers=self._headers(), json=data, timeout=self.settings.request_timeout_seconds)
|
||||
payload = response.json()
|
||||
return ApiResult("success" if payload.get("code") == 0 else "failed", str(payload.get("code")), payload.get("message"), payload)
|
||||
except Exception as exc:
|
||||
return ApiResult("failed", None, f"sync exception: {exc}", None)
|
||||
|
||||
def create_download_task(self, model_id: str, source: str = "HUGGING_FACE", weight_file: str | None = None) -> ApiResult:
|
||||
url = self.settings.modelhub_url + "/adminApi/async/task/model-download-task"
|
||||
data: dict[str, Any] = {
|
||||
"hfToken": self.settings.hf_token,
|
||||
"modelId": model_id,
|
||||
"source": source,
|
||||
"stillDownloadAlreadySuccessDownloadedModel": False,
|
||||
}
|
||||
if weight_file:
|
||||
data["allowPatterns"] = [weight_file]
|
||||
try:
|
||||
response = self.session.post(url, headers=self._headers(), json=data, timeout=self.settings.request_timeout_seconds)
|
||||
payload = response.json()
|
||||
code = payload.get("code")
|
||||
msg = payload.get("message", "")
|
||||
if code == 0:
|
||||
result = "created"
|
||||
elif code == 40000:
|
||||
result = "running_elsewhere"
|
||||
elif code == 60004:
|
||||
result = "already_downloaded"
|
||||
elif code == 60005:
|
||||
result = "download_pool_full"
|
||||
else:
|
||||
result = "failed"
|
||||
return ApiResult(result, str(code), msg, payload)
|
||||
except Exception as exc:
|
||||
return ApiResult("failed", None, f"download exception: {exc}", None)
|
||||
|
||||
def get_download_status(self, model_id: str) -> str:
|
||||
url = self.settings.modelhub_url + "/adminApi/async/task/model-download-task"
|
||||
params = {"modelId": model_id, "userEmail": self.settings.email, "userId": self.settings.user_id}
|
||||
try:
|
||||
response = self.session.get(url, headers=self._headers(), params=params, timeout=self.settings.request_timeout_seconds)
|
||||
response.raise_for_status()
|
||||
records = response.json().get("data", {}).get("records", [])
|
||||
if records:
|
||||
return records[0].get("status") or "None"
|
||||
return "None"
|
||||
except Exception as exc:
|
||||
LOG.warning("download status check failed for %s: %s", model_id, exc)
|
||||
return "QUERY_FAILED"
|
||||
|
||||
def build_contest_payload(self, model_id: str, gpu_type: str, config: str) -> dict[str, Any]:
|
||||
payload: dict[str, Any] = {
|
||||
"contestApiToken": self.settings.contest_api_token,
|
||||
"contributors": self.settings.contributors,
|
||||
"gpuTypes": [gpu_type],
|
||||
"modelId": model_id,
|
||||
"submissionConfig": [{"gpuType": gpu_type, "config": config}],
|
||||
"taskType": "text-generation",
|
||||
}
|
||||
if self.settings.inject_strategy_id:
|
||||
payload[self.settings.contest_task_strategy_field] = self.settings.strategy_id
|
||||
return payload
|
||||
|
||||
def create_contest_task(self, model_id: str, gpu_type: str, config: str) -> ApiResult:
|
||||
payload = self.build_contest_payload(model_id, gpu_type, config)
|
||||
if self.settings.submit_dry_run:
|
||||
LOG.info("dry-run submit payload: model=%s gpu=%s strategy_field=%s", model_id, gpu_type, self.settings.contest_task_strategy_field)
|
||||
return ApiResult("success", "DRY_RUN", "dry run", {"payload": payload})
|
||||
url = self.settings.modelhub_adminapi_base + "/async/task/create-contest-task"
|
||||
try:
|
||||
response = self.session.post(url, headers=self._headers(), json=payload, timeout=self.settings.request_timeout_seconds)
|
||||
body = response.json()
|
||||
code = body.get("code")
|
||||
message = body.get("message", "")
|
||||
if code == 0 and message == "ok":
|
||||
result = "success"
|
||||
elif code == 40000 and "正在验证中" in message:
|
||||
result = "conflict"
|
||||
elif code == 60007:
|
||||
result = "queue_full"
|
||||
else:
|
||||
result = "failed"
|
||||
return ApiResult(result, str(code), message, body)
|
||||
except Exception as exc:
|
||||
return ApiResult("failed", None, f"submit exception: {exc}", None)
|
||||
1
app/crawlers/__init__.py
Normal file
1
app/crawlers/__init__.py
Normal file
@@ -0,0 +1 @@
|
||||
"""Candidate crawlers."""
|
||||
85
app/crawlers/download_success.py
Normal file
85
app/crawlers/download_success.py
Normal file
@@ -0,0 +1,85 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
|
||||
from app.clients.modelhub import ModelHubClient
|
||||
from app.domain.model_type import engine_for_model
|
||||
from app.domain.priorities import PRIORITY_OTHERS_DOWNLOADED
|
||||
from app.settings import Settings
|
||||
from app.storage.repositories import Repository
|
||||
|
||||
LOG = logging.getLogger(__name__)
|
||||
|
||||
SOURCE_MAP = {
|
||||
"HuggingFace": "HUGGING_FACE",
|
||||
"HUGGING_FACE": "HUGGING_FACE",
|
||||
"ModelScope": "MODEL_SCOPE",
|
||||
"MODEL_SCOPE": "MODEL_SCOPE",
|
||||
}
|
||||
|
||||
|
||||
def normalize_source(raw: str | None) -> str:
|
||||
if not raw:
|
||||
return "HUGGING_FACE"
|
||||
return SOURCE_MAP.get(raw, raw)
|
||||
|
||||
|
||||
class DownloadSuccessCrawler:
|
||||
def __init__(self, settings: Settings, repo: Repository, client: ModelHubClient):
|
||||
self.settings = settings
|
||||
self.repo = repo
|
||||
self.client = client
|
||||
|
||||
def refresh(self) -> int:
|
||||
if not self.settings.enable_download_success_crawler:
|
||||
return 0
|
||||
imported = 0
|
||||
seen: set[str] = set()
|
||||
page = 1
|
||||
while True:
|
||||
data = self.client.list_success_download_tasks(page, self.settings.download_success_page_size)
|
||||
records = data.get("records") or []
|
||||
total = int(data.get("total") or 0)
|
||||
if not records:
|
||||
break
|
||||
|
||||
for record in records:
|
||||
args = record.get("args") or {}
|
||||
model_id = args.get("model_id") or args.get("modelId")
|
||||
if not model_id or model_id in seen:
|
||||
continue
|
||||
seen.add(model_id)
|
||||
source = normalize_source(args.get("source"))
|
||||
engine = engine_for_model(model_id)
|
||||
if engine != "vllm":
|
||||
self.repo.upsert_candidate(
|
||||
model_id=model_id,
|
||||
source=source,
|
||||
origin="model_download_success_page_gguf",
|
||||
priority=PRIORITY_OTHERS_DOWNLOADED,
|
||||
known_downloaded_by_others=True,
|
||||
self_download_allowed=False,
|
||||
notes=f"download task {record.get('id')} SUCCESS; GGUF/llamacpp candidate",
|
||||
)
|
||||
else:
|
||||
self.repo.upsert_candidate(
|
||||
model_id=model_id,
|
||||
source=source,
|
||||
origin="model_download_success_page",
|
||||
priority=PRIORITY_OTHERS_DOWNLOADED,
|
||||
known_downloaded_by_others=True,
|
||||
self_download_allowed=False,
|
||||
notes=f"download task {record.get('id')} SUCCESS",
|
||||
)
|
||||
imported += 1
|
||||
|
||||
if self.settings.download_success_max_pages and page >= self.settings.download_success_max_pages:
|
||||
break
|
||||
if total and page * self.settings.download_success_page_size >= total:
|
||||
break
|
||||
if len(records) < self.settings.download_success_page_size:
|
||||
break
|
||||
page += 1
|
||||
|
||||
LOG.info("download-success crawler imported/updated %s candidates", imported)
|
||||
return imported
|
||||
109
app/crawlers/not_adapted.py
Normal file
109
app/crawlers/not_adapted.py
Normal file
@@ -0,0 +1,109 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
|
||||
from app.clients.modelhub import ModelHubClient
|
||||
from app.domain.machines import machine_to_gpu_alias, parse_machine_names
|
||||
from app.domain.model_type import engine_for_model
|
||||
from app.domain.priorities import PRIORITY_BOUNTY
|
||||
from app.settings import Settings
|
||||
from app.storage.repositories import Repository
|
||||
|
||||
LOG = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _model_task_name(model: dict) -> str:
|
||||
task_info = model.get("taskLevelsInfo") or {}
|
||||
return model.get("taskLevelChineseName") or task_info.get("taskLevelChineseName") or ""
|
||||
|
||||
|
||||
def _adapted_machine_names(model: dict) -> set[str]:
|
||||
machines = model.get("machines") or []
|
||||
return {m.get("machineName") for m in machines if m.get("machineName")}
|
||||
|
||||
|
||||
class NotAdaptedCrawler:
|
||||
def __init__(self, settings: Settings, repo: Repository, client: ModelHubClient):
|
||||
self.settings = settings
|
||||
self.repo = repo
|
||||
self.client = client
|
||||
|
||||
def refresh(self) -> int:
|
||||
if not self.settings.enable_not_adapted_crawler:
|
||||
return 0
|
||||
machine_names = parse_machine_names(self.settings.target_machine_names)
|
||||
if not machine_names:
|
||||
LOG.info("not-adapted crawler skipped: TARGET_MACHINE_NAMES is empty")
|
||||
return 0
|
||||
|
||||
machine_alias_pairs: list[tuple[str, str]] = []
|
||||
for machine_name in machine_names:
|
||||
alias = machine_to_gpu_alias(machine_name)
|
||||
if not alias:
|
||||
LOG.warning("unknown target machine name, skipped: %s", machine_name)
|
||||
continue
|
||||
machine_alias_pairs.append((machine_name, alias))
|
||||
if not machine_alias_pairs:
|
||||
return 0
|
||||
|
||||
imported = 0
|
||||
page = 1
|
||||
while True:
|
||||
data = self.client.list_modelhub_models(page, self.settings.modelhub_page_size, "")
|
||||
records = data.get("records") or []
|
||||
total = int(data.get("total") or 0)
|
||||
if not records:
|
||||
break
|
||||
|
||||
for model in records:
|
||||
model_id = model.get("modelId")
|
||||
if not model_id:
|
||||
continue
|
||||
task_name = _model_task_name(model)
|
||||
if self.settings.target_task_level and self.settings.target_task_level not in task_name:
|
||||
continue
|
||||
|
||||
adapted = _adapted_machine_names(model)
|
||||
engine = engine_for_model(model_id)
|
||||
for machine_name, alias in machine_alias_pairs:
|
||||
if machine_name in adapted:
|
||||
continue
|
||||
if engine != "vllm":
|
||||
self.repo.upsert_candidate(
|
||||
model_id=model_id,
|
||||
source="HUGGING_FACE",
|
||||
origin="bounty_target_not_adapted_gguf",
|
||||
priority=PRIORITY_BOUNTY,
|
||||
is_bounty=True,
|
||||
is_promote=True,
|
||||
known_downloaded_by_others=True,
|
||||
self_download_allowed=False,
|
||||
target_gpu_aliases=alias,
|
||||
notes=f"GGUF/llamacpp candidate for {machine_name}; vLLM phase skips submit",
|
||||
)
|
||||
imported += 1
|
||||
continue
|
||||
self.repo.upsert_candidate(
|
||||
model_id=model_id,
|
||||
source="HUGGING_FACE",
|
||||
origin="bounty_target_not_adapted",
|
||||
priority=PRIORITY_BOUNTY,
|
||||
is_bounty=True,
|
||||
is_promote=True,
|
||||
known_downloaded_by_others=True,
|
||||
self_download_allowed=False,
|
||||
target_gpu_aliases=alias,
|
||||
notes=f"not adapted on {machine_name}",
|
||||
)
|
||||
imported += 1
|
||||
|
||||
if self.settings.crawler_max_pages and page >= self.settings.crawler_max_pages:
|
||||
break
|
||||
if total and page * self.settings.modelhub_page_size >= total:
|
||||
break
|
||||
if len(records) < self.settings.modelhub_page_size:
|
||||
break
|
||||
page += 1
|
||||
|
||||
LOG.info("not-adapted crawler imported/updated %s candidates", imported)
|
||||
return imported
|
||||
1
app/domain/__init__.py
Normal file
1
app/domain/__init__.py
Normal file
@@ -0,0 +1 @@
|
||||
"""Domain helpers for model submission strategy."""
|
||||
41
app/domain/gpu.py
Normal file
41
app/domain/gpu.py
Normal file
@@ -0,0 +1,41 @@
|
||||
from __future__ import annotations
|
||||
|
||||
VLLM_GPU_ALIASES = {
|
||||
"910b": ["Ascend_910-b3", "Ascend_910-b4"],
|
||||
"s4000": ["Mthreads_s4000"],
|
||||
"k100": ["hygon_k100-ai"],
|
||||
"mrv100": ["Iluvatar_mrv-100"],
|
||||
"bi100": ["Iluvatar_bi-100"],
|
||||
"bi150": ["Iluvatar_bi-150"],
|
||||
"c500": ["MetaX_c-500"],
|
||||
"mlu370-x4": ["Cambricon_mlu-370-x4"],
|
||||
"mlu370-x8": ["Cambricon_mlu-370-x8"],
|
||||
"p800": ["Kunlunxin_p-800"],
|
||||
"166m": ["Biren_166m"],
|
||||
"biren166m": ["Biren_166m"], # backward-compatible alias from the old config.py
|
||||
}
|
||||
|
||||
# Sunrise s2 currently uses SGLang in the old config, so vLLM-only v1 skips it.
|
||||
UNSUPPORTED_VLLM_ALIASES = {"s2"}
|
||||
|
||||
|
||||
def parse_aliases(raw: str | None, defaults: list[str]) -> list[str]:
|
||||
if not raw:
|
||||
return defaults
|
||||
return [item.strip() for item in raw.split(",") if item.strip()]
|
||||
|
||||
|
||||
def expand_gpu_alias(alias: str) -> list[str]:
|
||||
return VLLM_GPU_ALIASES.get(alias, [])
|
||||
|
||||
|
||||
def expand_gpu_aliases(aliases: list[str]) -> list[tuple[str, str]]:
|
||||
pairs: list[tuple[str, str]] = []
|
||||
for alias in aliases:
|
||||
for gpu_type in expand_gpu_alias(alias):
|
||||
pairs.append((alias, gpu_type))
|
||||
return pairs
|
||||
|
||||
|
||||
def is_supported_alias(alias: str) -> bool:
|
||||
return alias in VLLM_GPU_ALIASES
|
||||
32
app/domain/machines.py
Normal file
32
app/domain/machines.py
Normal file
@@ -0,0 +1,32 @@
|
||||
from __future__ import annotations
|
||||
|
||||
MACHINE_TO_GPU_ALIAS = {
|
||||
"MTT S4000": "s4000",
|
||||
"Mthreads S4000": "s4000",
|
||||
"摩尔线程 S4000": "s4000",
|
||||
"Hygon K100": "k100",
|
||||
"海光 K100": "k100",
|
||||
"Kunlunxin P800": "p800",
|
||||
"Kunlunxin P-800": "p800",
|
||||
"昆仑芯 P800": "p800",
|
||||
"Biren 166M": "166m",
|
||||
"Biren 166m": "166m",
|
||||
"壁砺 166M": "166m",
|
||||
"壁砺166M": "166m",
|
||||
"Ascend 910B": "910b",
|
||||
"昇腾 910B": "910b",
|
||||
"Iluvatar BI100": "bi100",
|
||||
"Iluvatar BI150": "bi150",
|
||||
"MetaX C500": "c500",
|
||||
"Iluvatar MRV100": "mrv100",
|
||||
"Cambricon MLU370-X4": "mlu370-x4",
|
||||
"Cambricon MLU370-X8": "mlu370-x8",
|
||||
}
|
||||
|
||||
|
||||
def machine_to_gpu_alias(machine_name: str) -> str | None:
|
||||
return MACHINE_TO_GPU_ALIAS.get(machine_name.strip())
|
||||
|
||||
|
||||
def parse_machine_names(raw: str) -> list[str]:
|
||||
return [item.strip() for item in raw.split(",") if item.strip()]
|
||||
12
app/domain/model_type.py
Normal file
12
app/domain/model_type.py
Normal file
@@ -0,0 +1,12 @@
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
def is_gguf_model(model_id: str, extra_text: str = "") -> bool:
|
||||
haystack = f"{model_id} {extra_text}".lower()
|
||||
return "gguf" in haystack or ".gguf" in haystack
|
||||
|
||||
|
||||
def engine_for_model(model_id: str, extra_text: str = "") -> str:
|
||||
if is_gguf_model(model_id, extra_text):
|
||||
return "llamacpp_candidate"
|
||||
return "vllm"
|
||||
16
app/domain/priorities.py
Normal file
16
app/domain/priorities.py
Normal file
@@ -0,0 +1,16 @@
|
||||
PRIORITY_BOUNTY = 0
|
||||
PRIORITY_PROMOTE = 10
|
||||
PRIORITY_OTHERS_DOWNLOADED = 20
|
||||
PRIORITY_SELF_DOWNLOAD = 30
|
||||
PRIORITY_LOW_CONFIDENCE = 90
|
||||
|
||||
ORIGIN_TO_PRIORITY = {
|
||||
"manual_bounty": PRIORITY_BOUNTY,
|
||||
"bounty": PRIORITY_BOUNTY,
|
||||
"manual_promote": PRIORITY_PROMOTE,
|
||||
"promote": PRIORITY_PROMOTE,
|
||||
"manual_downloaded_by_others": PRIORITY_OTHERS_DOWNLOADED,
|
||||
"downloaded_by_others": PRIORITY_OTHERS_DOWNLOADED,
|
||||
"hf_seed": PRIORITY_SELF_DOWNLOAD,
|
||||
"self_download": PRIORITY_SELF_DOWNLOAD,
|
||||
}
|
||||
449
app/domain/vllm_configs.py
Normal file
449
app/domain/vllm_configs.py
Normal file
@@ -0,0 +1,449 @@
|
||||
from __future__ import annotations
|
||||
|
||||
# Templates are copied from config部分更新.py and use __MODEL_ID__ tokens instead of
|
||||
# str.format placeholders so YAML inline maps like {name: ..., value: ...} stay intact.
|
||||
|
||||
VLLM_MRV100_CONFIG = """docker_image: harbor.4pd.io/hardcore-tech/iluvatar/llm-infer-iluvatar-mr:v0
|
||||
nv_docker_image: harbor.4pd.io/dooke/vllm/vllm/vllm-openai:v0.11.0
|
||||
framework: vllm
|
||||
storage: gpfs
|
||||
modelhub_options:
|
||||
srcRelativePath: leaderboard/modelHubXC/__MODEL_ID__
|
||||
mountPoint: /model
|
||||
max_model_len: 4096
|
||||
lang: en
|
||||
temperature: 0.4
|
||||
repetition_penalty: 1.1
|
||||
top_p: 0.9
|
||||
sut_config:
|
||||
gpu_num: 1
|
||||
values:
|
||||
command: [vllm, serve, /model, --port, '20644', --served-model-name, llm, --max-model-len,
|
||||
'4096', --gpu-memory-utilization, '0.9', --enforce-eager, --trust-remote-code,
|
||||
-tp, '1']
|
||||
ref_config:
|
||||
gpu_num: 1
|
||||
values:
|
||||
command: [vllm, serve, /model, --port, '80', --served-model-name, llm, --max-model-len,
|
||||
'4096', --enforce-eager, --trust-remote-code, -tp, '1']
|
||||
"""
|
||||
|
||||
VLLM_S4000_CONFIG = """docker_image: git.modelhub.org.cn:9443/enginex-mthreads/vllm-musa-qy2-py310:v0.8.4-release
|
||||
nv_docker_image: harbor.4pd.io/dooke/vllm/vllm/vllm-openai:v0.11.0
|
||||
framework: vllm
|
||||
storage: gpfs
|
||||
modelhub_options:
|
||||
srcRelativePath: leaderboard/modelHubXC/__MODEL_ID__
|
||||
mountPoint: /model
|
||||
api: completion
|
||||
lang: en
|
||||
max_model_len: 4096
|
||||
max_tokens: 1024
|
||||
temperature: 0.7
|
||||
repetition_penalty: 1.1
|
||||
top_p: 0.9
|
||||
sut_config:
|
||||
gpu_num: 1
|
||||
values:
|
||||
command:
|
||||
- vllm
|
||||
- serve
|
||||
- /model
|
||||
- --port
|
||||
- '8000'
|
||||
- --served-model-name
|
||||
- llm
|
||||
- --max-model-len
|
||||
- '4096'
|
||||
- --dtype
|
||||
- auto
|
||||
- --gpu-memory-utilization
|
||||
- '0.95'
|
||||
- -tp
|
||||
- '1'
|
||||
- --enforce-eager
|
||||
- --trust-remote-code
|
||||
ref_config:
|
||||
gpu_num: 1
|
||||
values:
|
||||
command:
|
||||
- vllm
|
||||
- serve
|
||||
- /model
|
||||
- --port
|
||||
- '8000'
|
||||
- --served-model-name
|
||||
- llm
|
||||
- --max-model-len
|
||||
- '4096'
|
||||
- --dtype
|
||||
- auto
|
||||
- --gpu-memory-utilization
|
||||
- '0.95'
|
||||
- -tp
|
||||
- '1'
|
||||
- --enforce-eager
|
||||
- --trust-remote-code
|
||||
"""
|
||||
|
||||
VLLM_910B_CONFIG = """
|
||||
framework: vllm
|
||||
docker_image: git.modelhub.org.cn:9443/enginex-ascend/vllm-ascend:v0.11.0rc0
|
||||
nv_docker_image: harbor.4pd.io/dooke/vllm/vllm/vllm-openai:v0.11.0
|
||||
storage: gpfs
|
||||
modelhub_options:
|
||||
srcRelativePath: leaderboard/modelHubXC/__MODEL_ID__
|
||||
mountPoint: /model
|
||||
api: completion
|
||||
max_tokens: 1024
|
||||
temperature: 0.7
|
||||
repetition_penalty: 1.2
|
||||
top_p: 0.9
|
||||
lang: zh
|
||||
max_model_len: 2048
|
||||
sut_config:
|
||||
gpu_num: 1
|
||||
values:
|
||||
command:
|
||||
- vllm
|
||||
- serve
|
||||
- /model
|
||||
- --port
|
||||
- '8000'
|
||||
- --served-model-name
|
||||
- llm
|
||||
- --max-model-len
|
||||
- '2048'
|
||||
- --dtype
|
||||
- auto
|
||||
- --gpu-memory-utilization
|
||||
- '0.95'
|
||||
- -tp
|
||||
- '1'
|
||||
- --enforce-eager
|
||||
- --trust-remote-code
|
||||
ref_config:
|
||||
gpu_num: 1
|
||||
values:
|
||||
command:
|
||||
- vllm
|
||||
- serve
|
||||
- /model
|
||||
- --port
|
||||
- '8000'
|
||||
- --served-model-name
|
||||
- llm
|
||||
- --max-model-len
|
||||
- '2048'
|
||||
- --dtype
|
||||
- auto
|
||||
- --gpu-memory-utilization
|
||||
- '0.95'
|
||||
- -tp
|
||||
- '1'
|
||||
- --enforce-eager
|
||||
- --trust-remote-code
|
||||
"""
|
||||
|
||||
VLLM_BI100_CONFIG = """docker_image: git.modelhub.org.cn:9443/enginex-iluvatar-bi100/vllm:0.6.3
|
||||
nv_docker_image: harbor.4pd.io/dooke/vllm/vllm/vllm-openai:v0.11.0
|
||||
framework: vllm
|
||||
storage: gpfs
|
||||
modelhub_options:
|
||||
srcRelativePath: leaderboard/modelHubXC/__MODEL_ID__
|
||||
mountPoint: /model
|
||||
lang: en
|
||||
max_model_len: 4096
|
||||
sut_config:
|
||||
gpu_num: 1
|
||||
environment:
|
||||
MAX_MODEL_LEN: "4096"
|
||||
values:
|
||||
command: [/workspace/launch_service, --port, '80']
|
||||
ref_config:
|
||||
gpu_num: 1
|
||||
environment:
|
||||
MAX_MODEL_LEN: "4096"
|
||||
values:
|
||||
command: [vllm, serve, /model, --port, '80', --served-model-name, llm, --max-model-len,
|
||||
'4096', --enforce-eager, --trust-remote-code, -tp, '1']
|
||||
"""
|
||||
|
||||
VLLM_BI150_CONFIG = """docker_image: harbor.4pd.io/modelhubxc/sunruoxi/enginex-iluvatar-bi150-vllm-fix-tokenizer:v0.8.3
|
||||
nv_docker_image: harbor-contest.4pd.io/sunruoxi/vllm-openai-fix-tokenizer:v0.11.0
|
||||
framework: vllm_fix_tokenizer
|
||||
nv_framework: vllm_fix_tokenizer
|
||||
lang: en
|
||||
max_model_len: 4096
|
||||
modelhub_options:
|
||||
srcRelativePath: leaderboard/modelHubXC/__MODEL_ID__
|
||||
mountPoint: /model
|
||||
sut_config:
|
||||
gpu_num: 1
|
||||
values:
|
||||
command:
|
||||
- /opt/entrypoint.sh
|
||||
- /model
|
||||
- --port
|
||||
- '80'
|
||||
- --served-model-name
|
||||
- llm
|
||||
- --max-model-len
|
||||
- '4096'
|
||||
- --enforce-eager
|
||||
- --trust-remote-code
|
||||
- -tp
|
||||
- '1'
|
||||
ref_config:
|
||||
gpu_num: 1
|
||||
values:
|
||||
command:
|
||||
- /opt/entrypoint.sh
|
||||
- /model
|
||||
- --port
|
||||
- '80'
|
||||
- --served-model-name
|
||||
- llm
|
||||
- --max-model-len
|
||||
- '4096'
|
||||
- --enforce-eager
|
||||
- --trust-remote-code
|
||||
- -tp
|
||||
- '1'
|
||||
"""
|
||||
|
||||
VLLM_C500_CONFIG = """docker_image: git.modelhub.org.cn:9443/enginex-metax/vllm:0.9.1
|
||||
nv_docker_image: harbor.4pd.io/dooke/vllm/vllm/vllm-openai:v0.11.0
|
||||
framework: vllm
|
||||
storage: gpfs
|
||||
modelhub_options:
|
||||
srcRelativePath: leaderboard/modelHubXC/__MODEL_ID__
|
||||
mountPoint: /model
|
||||
api: completion
|
||||
lang: en
|
||||
max_model_len: 4096
|
||||
max_tokens: 1024
|
||||
temperature: 0.7
|
||||
repetition_penalty: 1.1
|
||||
top_p: 0.9
|
||||
sut_config:
|
||||
gpu_num: 1
|
||||
values:
|
||||
command: [/opt/conda/bin/vllm, serve, /model, --port, '20644', --served-model-name,
|
||||
llm, --max-model-len, '4096', --gpu-memory-utilization, '0.9', -tp, '1', --enforce-eager,
|
||||
--trust-remote-code]
|
||||
ref_config:
|
||||
gpu_num: 1
|
||||
values:
|
||||
command: [vllm, serve, /model, --port, '80', --served-model-name, llm, --max-model-len,
|
||||
'4096', -tp, '1', --enforce-eager, --trust-remote-code]
|
||||
"""
|
||||
|
||||
VLLM_K100_CONFIG = """docker_image: harbor.4pd.io/modelhubxc/i-peixingyu/k100-vllm-patched-v2.0:v2.0.0
|
||||
nv_docker_image: harbor-contest.4pd.io/sunruoxi/vllm-openai-fix-tokenizer:v0.11.0
|
||||
framework: vllm-patch-tokenizer
|
||||
nv_framework: vllm_fix_tokenizer
|
||||
max_model_len: 4096
|
||||
api: completion
|
||||
modelhub_options:
|
||||
srcRelativePath: leaderboard/modelHubXC/__MODEL_ID__
|
||||
mountPoint: /model
|
||||
sut_config:
|
||||
gpu_num: 1
|
||||
values:
|
||||
command:
|
||||
- /opt/entrypoint.sh
|
||||
- --port
|
||||
- '20644'
|
||||
- --served-model-name
|
||||
- llm
|
||||
- --max-model-len
|
||||
- '4096'
|
||||
- --enforce-eager
|
||||
- --trust-remote-code
|
||||
- -tp
|
||||
- '1'
|
||||
ref_config:
|
||||
gpu_num: 1
|
||||
values:
|
||||
command:
|
||||
- /opt/entrypoint.sh
|
||||
- /model
|
||||
- --port
|
||||
- '80'
|
||||
- --served-model-name
|
||||
- llm
|
||||
- --max-model-len
|
||||
- '4096'
|
||||
- --enforce-eager
|
||||
- --trust-remote-code
|
||||
- -tp
|
||||
- '1'
|
||||
"""
|
||||
|
||||
VLLM_MLU370_X4_CONFIG = """docker_image: harbor.4pd.io/hardcore-tech/cambricon-mlu370-pytorch:v25.01-torch2.5.0-torchmlu1.24.1-ubuntu22.04-py310
|
||||
nv_docker_image: harbor.4pd.io/dooke/vllm/vllm/vllm-openai:v0.11.0
|
||||
framework: vllm-mlu
|
||||
modelhub_options:
|
||||
srcRelativePath: leaderboard/modelHubXC/__MODEL_ID__
|
||||
mountPoint: /model
|
||||
sut_config:
|
||||
values:
|
||||
gpu_num: 1
|
||||
env:
|
||||
- {name: MAX_MODEL_LEN, value: 4096}
|
||||
command: [vllm, serve, /model, --port, '8000', --served-model-name, llm, --max-model-len,
|
||||
'4096', --trust-remote-code, --dtype, float16]
|
||||
ref_config:
|
||||
values:
|
||||
cpu_num: 2
|
||||
gpu_num: 1
|
||||
env:
|
||||
- {name: MAX_MODEL_LEN, value: 4096}
|
||||
command: [vllm, serve, /model, --port, '80', --served-model-name, llm, --max-model-len,
|
||||
'4096', --trust-remote-code, --dtype, float16]
|
||||
"""
|
||||
|
||||
VLLM_MLU370_X8_CONFIG = """docker_image: harbor.4pd.io/hardcore-tech/cambricon-mlu370-pytorch:v25.01-20260204
|
||||
nv_docker_image: harbor.4pd.io/dooke/vllm/vllm/vllm-openai:v0.11.0
|
||||
framework: vllm-mlu
|
||||
api: completion
|
||||
lang: en
|
||||
temperature: 0.4
|
||||
repetition_penalty: 1.1
|
||||
top_p: 0.9
|
||||
modelhub_options:
|
||||
srcRelativePath: leaderboard/modelHubXC/__MODEL_ID__
|
||||
mountPoint: /model
|
||||
sut_config:
|
||||
gpu_num: 1
|
||||
values:
|
||||
env:
|
||||
- {name: MAX_MODEL_LEN, value: 4096}
|
||||
- {name: VLLM_ALLOW_LONG_MAX_MODEL_LEN, value: '1'}
|
||||
command: [vllm, serve, /model, --port, '8000', --served-model-name, llm, --max-model-len,
|
||||
'4096', --trust-remote-code, --dtype, float16]
|
||||
ref_config:
|
||||
cpu_num: 2
|
||||
gpu_num: 1
|
||||
values:
|
||||
env:
|
||||
- {name: MAX_MODEL_LEN, value: 4096}
|
||||
- {name: VLLM_ALLOW_LONG_MAX_MODEL_LEN, value: '1'}
|
||||
command: [vllm, serve, /model, --port, '80', --served-model-name, llm, --max-model-len,
|
||||
'4096', --trust-remote-code, --dtype, float16]
|
||||
"""
|
||||
|
||||
VLLM_P800_CONFIG = """docker_image: harbor.4pd.io/modelhubxc/zheng/vllm-kunlunxin-p-800-tokenizer-patch:v1.0
|
||||
nv_docker_image: harbor-contest.4pd.io/sunruoxi/vllm-openai-fix-tokenizer:v0.11.0
|
||||
framework: vllm_tokenizer_patch
|
||||
nv_framework: vllm_fix_tokenizer
|
||||
api: completion
|
||||
modelhub_options:
|
||||
srcRelativePath: leaderboard/modelHubXC/__MODEL_ID__
|
||||
mountPoint: /model
|
||||
max_model_len: 4096
|
||||
sut_config:
|
||||
gpu_num: 1
|
||||
values:
|
||||
command:
|
||||
- vllm
|
||||
- serve
|
||||
- /model
|
||||
- --port
|
||||
- '8000'
|
||||
- --served-model-name
|
||||
- llm
|
||||
- --max-model-len
|
||||
- '4096'
|
||||
- --gpu-memory-utilization
|
||||
- '0.9'
|
||||
- --enforce-eager
|
||||
- --trust-remote-code
|
||||
- -tp
|
||||
- '1'
|
||||
ref_config:
|
||||
gpu_num: 1
|
||||
values:
|
||||
command:
|
||||
- /opt/entrypoint.sh
|
||||
- /model
|
||||
- --port
|
||||
- '80'
|
||||
- --served-model-name
|
||||
- llm
|
||||
- --max-model-len
|
||||
- '4096'
|
||||
- --enforce-eager
|
||||
- --trust-remote-code
|
||||
- -tp
|
||||
- '1'
|
||||
"""
|
||||
|
||||
VLLM_166M_CONFIG = """docker_image: harbor.4pd.io/modelhubxc/sunruoxi/enginex-llm-biren166m-fix-tokenizer:v26.01
|
||||
nv_docker_image: harbor.4pd.io/modelhubxc/enginex-nvidia/vllm:0.11.0-patch-tokenizer
|
||||
modelhub_options:
|
||||
srcRelativePath: leaderboard/modelHubXC/__MODEL_ID__
|
||||
mountPoint: /model
|
||||
framework: vllm_fix_tokenizer
|
||||
nv_framework: vllm_fix_tokenizer
|
||||
lang: en
|
||||
max_model_len: 4096
|
||||
sut_config:
|
||||
gpu_num: 1
|
||||
values:
|
||||
command:
|
||||
- /opt/entrypoint.sh
|
||||
- /model
|
||||
- --port
|
||||
- '80'
|
||||
- --served-model-name
|
||||
- llm
|
||||
- --max-model-len
|
||||
- '4096'
|
||||
- --enforce-eager
|
||||
- --trust-remote-code
|
||||
- -tp
|
||||
- '1'
|
||||
ref_config:
|
||||
gpu_num: 1
|
||||
values:
|
||||
command:
|
||||
- /opt/entrypoint.sh
|
||||
- /model
|
||||
- --port
|
||||
- '80'
|
||||
- --served-model-name
|
||||
- llm
|
||||
- --max-model-len
|
||||
- '4096'
|
||||
- --enforce-eager
|
||||
- --trust-remote-code
|
||||
- -tp
|
||||
- '1'
|
||||
"""
|
||||
|
||||
VLLM_GPU_CONFIG_DICT = {
|
||||
"Ascend_910-b3": VLLM_910B_CONFIG,
|
||||
"Ascend_910-b4": VLLM_910B_CONFIG,
|
||||
"Mthreads_s4000": VLLM_S4000_CONFIG,
|
||||
"hygon_k100-ai": VLLM_K100_CONFIG,
|
||||
"Iluvatar_mrv-100": VLLM_MRV100_CONFIG,
|
||||
"Iluvatar_bi-100": VLLM_BI100_CONFIG,
|
||||
"Iluvatar_bi-150": VLLM_BI150_CONFIG,
|
||||
"MetaX_c-500": VLLM_C500_CONFIG,
|
||||
"Cambricon_mlu-370-x4": VLLM_MLU370_X4_CONFIG,
|
||||
"Cambricon_mlu-370-x8": VLLM_MLU370_X8_CONFIG,
|
||||
"Kunlunxin_p-800": VLLM_P800_CONFIG,
|
||||
"Biren_166m": VLLM_166M_CONFIG,
|
||||
}
|
||||
|
||||
|
||||
def supported_gpu_types() -> set[str]:
|
||||
return set(VLLM_GPU_CONFIG_DICT)
|
||||
|
||||
|
||||
def gen_vllm_config(gpu_type: str, model_id: str, max_model_len: int = 1024) -> str:
|
||||
if gpu_type not in VLLM_GPU_CONFIG_DICT:
|
||||
raise ValueError(f"Unsupported vLLM gpu_type: {gpu_type}")
|
||||
return VLLM_GPU_CONFIG_DICT[gpu_type].replace("__MODEL_ID__", model_id)
|
||||
58
app/main.py
Normal file
58
app/main.py
Normal file
@@ -0,0 +1,58 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import threading
|
||||
|
||||
from flask import Flask, jsonify
|
||||
|
||||
from app.clients.modelhub import ModelHubClient
|
||||
from app.scheduler.loop import StrategyLoop
|
||||
from app.settings import load_settings
|
||||
from app.shutdown import install_signal_handlers, stop_event
|
||||
from app.storage.db import connect
|
||||
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s %(message)s")
|
||||
LOG = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def create_app() -> Flask:
|
||||
app = Flask(__name__)
|
||||
|
||||
@app.get("/health")
|
||||
def health():
|
||||
return jsonify({"status": "ok"}), 200
|
||||
|
||||
@app.get("/status")
|
||||
def status():
|
||||
return jsonify({"status": "ok", "stopping": stop_event.is_set()}), 200
|
||||
|
||||
return app
|
||||
|
||||
|
||||
def main():
|
||||
install_signal_handlers()
|
||||
settings = load_settings(require_secrets=True)
|
||||
LOG.info(
|
||||
"starting vLLM agent strategy dry_run=%s inject_strategy_id=%s strategy_field=%s db=%s",
|
||||
settings.submit_dry_run,
|
||||
settings.inject_strategy_id,
|
||||
settings.contest_task_strategy_field,
|
||||
settings.db_path,
|
||||
)
|
||||
conn = connect(settings.db_path)
|
||||
client = ModelHubClient(settings)
|
||||
loop = StrategyLoop(settings, __import__("app.storage.repositories", fromlist=["Repository"]).Repository(conn), client)
|
||||
worker = threading.Thread(target=loop.run_forever, args=(stop_event,), name="strategy-loop", daemon=True)
|
||||
worker.start()
|
||||
app = create_app()
|
||||
try:
|
||||
app.run(host=settings.host, port=settings.port, threaded=True)
|
||||
finally:
|
||||
stop_event.set()
|
||||
worker.join(timeout=25)
|
||||
conn.close()
|
||||
LOG.info("shutdown complete")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
1
app/scheduler/__init__.py
Normal file
1
app/scheduler/__init__.py
Normal file
@@ -0,0 +1 @@
|
||||
"""Scheduler components."""
|
||||
13
app/scheduler/budget.py
Normal file
13
app/scheduler/budget.py
Normal file
@@ -0,0 +1,13 @@
|
||||
from app.storage.repositories import Repository
|
||||
|
||||
|
||||
class Budget:
|
||||
def __init__(self, repo: Repository, max_tasks: int):
|
||||
self.repo = repo
|
||||
self.max_tasks = max_tasks
|
||||
|
||||
def remaining(self) -> int:
|
||||
return max(0, self.max_tasks - self.repo.count_budgeted_tasks())
|
||||
|
||||
def can_create_task(self) -> bool:
|
||||
return self.remaining() > 0
|
||||
53
app/scheduler/downloader.py
Normal file
53
app/scheduler/downloader.py
Normal file
@@ -0,0 +1,53 @@
|
||||
from app.clients.modelhub import ModelHubClient
|
||||
from app.storage.repositories import Repository, future_seconds
|
||||
|
||||
|
||||
class Downloader:
|
||||
def __init__(self, repo: Repository, client: ModelHubClient, max_self_downloads: int, poll_interval_seconds: int):
|
||||
self.repo = repo
|
||||
self.client = client
|
||||
self.max_self_downloads = max_self_downloads
|
||||
self.poll_interval_seconds = poll_interval_seconds
|
||||
|
||||
def poll_active_downloads(self) -> int:
|
||||
changed = 0
|
||||
for download in self.repo.active_downloads_due(limit=20):
|
||||
status = self.client.get_download_status(download["model_id"])
|
||||
active = status in {"WAITING", "RUNNING", "QUERY_FAILED"}
|
||||
self.repo.upsert_download(
|
||||
download["model_id"],
|
||||
download["source"],
|
||||
status,
|
||||
"self",
|
||||
active_self_slot=active,
|
||||
next_poll_at=future_seconds(self.poll_interval_seconds) if active else None,
|
||||
)
|
||||
changed += 1
|
||||
return changed
|
||||
|
||||
def maybe_start_self_downloads(self, limit: int = 5) -> int:
|
||||
started = 0
|
||||
for task in self.repo.due_tasks(("waiting_download",), limit=limit):
|
||||
if self.repo.count_active_self_downloads() >= self.max_self_downloads:
|
||||
self.repo.mark_task_status(task["id"], "waiting_download", next_attempt_at=future_seconds(300))
|
||||
continue
|
||||
sync_result = self.client.sync_from_huggingface(task["model_id"])
|
||||
if sync_result.result != "success":
|
||||
self.repo.mark_task_status(task["id"], "waiting_download", sync_result.code, sync_result.message, future_seconds(600), increment_attempt=True)
|
||||
continue
|
||||
dl = self.client.create_download_task(task["model_id"], task["source"] or "HUGGING_FACE")
|
||||
if dl.result == "created":
|
||||
self.repo.upsert_download(task["model_id"], task["source"] or "HUGGING_FACE", "WAITING", "self", True, dl.code, dl.message, future_seconds(self.poll_interval_seconds))
|
||||
self.repo.mark_task_status(task["id"], "waiting_download", dl.code, dl.message, future_seconds(self.poll_interval_seconds), increment_attempt=True)
|
||||
started += 1
|
||||
elif dl.result == "already_downloaded":
|
||||
self.repo.upsert_download(task["model_id"], task["source"] or "HUGGING_FACE", "SUCCESS", "others", False, dl.code, dl.message)
|
||||
self.repo.mark_task_status(task["id"], "ready_to_submit", dl.code, dl.message)
|
||||
elif dl.result == "running_elsewhere":
|
||||
self.repo.upsert_download(task["model_id"], task["source"] or "HUGGING_FACE", "RUNNING", "others", False, dl.code, dl.message, future_seconds(self.poll_interval_seconds))
|
||||
self.repo.mark_task_status(task["id"], "waiting_download", dl.code, dl.message, future_seconds(self.poll_interval_seconds), increment_attempt=True)
|
||||
elif dl.result == "download_pool_full":
|
||||
self.repo.mark_task_status(task["id"], "waiting_download", dl.code, dl.message, future_seconds(300), increment_attempt=True)
|
||||
else:
|
||||
self.repo.mark_task_status(task["id"], "failed", dl.code, dl.message, increment_attempt=True)
|
||||
return started
|
||||
118
app/scheduler/loop.py
Normal file
118
app/scheduler/loop.py
Normal file
@@ -0,0 +1,118 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from app.clients.modelhub import ModelHubClient
|
||||
from app.crawlers.download_success import DownloadSuccessCrawler
|
||||
from app.crawlers.not_adapted import NotAdaptedCrawler
|
||||
from app.scheduler.budget import Budget
|
||||
from app.scheduler.downloader import Downloader
|
||||
from app.scheduler.planner import Planner
|
||||
from app.scheduler.submitter import Submitter
|
||||
from app.settings import Settings
|
||||
from app.storage.repositories import Repository, future_seconds
|
||||
|
||||
LOG = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class StrategyLoop:
|
||||
def __init__(self, settings: Settings, repo: Repository, client: ModelHubClient):
|
||||
self.settings = settings
|
||||
self.repo = repo
|
||||
self.client = client
|
||||
self.budget = Budget(repo, settings.max_user_tasks)
|
||||
self.planner = Planner(repo, self.budget, settings.default_gpu_alias_list)
|
||||
self.submitter = Submitter(repo, client)
|
||||
self.downloader = Downloader(repo, client, settings.max_self_downloads, settings.download_poll_interval_seconds)
|
||||
self.not_adapted_crawler = NotAdaptedCrawler(settings, repo, client)
|
||||
self.download_success_crawler = DownloadSuccessCrawler(settings, repo, client)
|
||||
self._last_crawler_refresh = None
|
||||
|
||||
def refresh_crawlers_if_due(self) -> dict[str, int]:
|
||||
now = datetime.now(timezone.utc)
|
||||
if self._last_crawler_refresh is not None:
|
||||
elapsed = (now - self._last_crawler_refresh).total_seconds()
|
||||
if elapsed < self.settings.crawler_refresh_seconds:
|
||||
return {"not_adapted": 0, "download_success": 0}
|
||||
self._last_crawler_refresh = now
|
||||
stats = {"not_adapted": 0, "download_success": 0}
|
||||
try:
|
||||
stats["not_adapted"] = self.not_adapted_crawler.refresh()
|
||||
except Exception:
|
||||
LOG.exception("not-adapted crawler failed")
|
||||
try:
|
||||
stats["download_success"] = self.download_success_crawler.refresh()
|
||||
except Exception:
|
||||
LOG.exception("download-success crawler failed")
|
||||
return stats
|
||||
|
||||
def mark_ready_tasks(self, limit: int = 50) -> int:
|
||||
changed = 0
|
||||
for task in self.repo.due_tasks(("pending",), limit=limit):
|
||||
if task["is_bounty"]:
|
||||
if self.settings.bounty_allow_direct_submit:
|
||||
self.repo.mark_task_status(task["id"], "ready_to_submit")
|
||||
changed += 1
|
||||
else:
|
||||
status = self.client.get_download_status(task["model_id"])
|
||||
if status == "SUCCESS":
|
||||
self.repo.upsert_download(task["model_id"], task["source"], "SUCCESS", "others")
|
||||
self.repo.mark_task_status(task["id"], "ready_to_submit")
|
||||
changed += 1
|
||||
else:
|
||||
self.repo.mark_task_status(task["id"], "pending", "DOWNLOAD_STATUS", status, future_seconds(300))
|
||||
elif task["is_promote"]:
|
||||
absent = self.client.is_model_absent_from_modelhub(task["model_id"])
|
||||
if not absent:
|
||||
self.repo.mark_task_status(task["id"], "ready_to_submit")
|
||||
changed += 1
|
||||
else:
|
||||
self.repo.mark_task_status(task["id"], "pending", "NOT_IN_MODELHUB", "model not adapted yet", future_seconds(1800))
|
||||
elif task["known_downloaded_by_others"]:
|
||||
status = self.client.get_download_status(task["model_id"])
|
||||
if status == "SUCCESS":
|
||||
self.repo.upsert_download(task["model_id"], task["source"], "SUCCESS", "others")
|
||||
self.repo.mark_task_status(task["id"], "ready_to_submit")
|
||||
changed += 1
|
||||
elif task["self_download_allowed"]:
|
||||
self.repo.mark_task_status(task["id"], "waiting_download", "DOWNLOAD_STATUS", status, future_seconds(300))
|
||||
else:
|
||||
self.repo.mark_task_status(task["id"], "pending", "DOWNLOAD_STATUS", status, future_seconds(300))
|
||||
elif task["self_download_allowed"]:
|
||||
self.repo.mark_task_status(task["id"], "waiting_download")
|
||||
else:
|
||||
self.repo.mark_task_status(task["id"], "skipped", "NO_READY_PATH", "candidate is not bounty/promote/downloaded/self-download")
|
||||
return changed
|
||||
|
||||
def run_once(self) -> dict[str, int]:
|
||||
crawler_stats = self.refresh_crawlers_if_due()
|
||||
created = self.planner.materialize_tasks()
|
||||
polled = self.downloader.poll_active_downloads()
|
||||
released = self.submitter.release_queue_full_backoff()
|
||||
ready = self.mark_ready_tasks()
|
||||
submitted = self.submitter.submit_due()
|
||||
started_downloads = 0
|
||||
if submitted == 0 and ready == 0:
|
||||
started_downloads = self.downloader.maybe_start_self_downloads()
|
||||
return {
|
||||
"crawled_not_adapted": crawler_stats["not_adapted"],
|
||||
"crawled_download_success": crawler_stats["download_success"],
|
||||
"created_tasks": created,
|
||||
"polled_downloads": polled,
|
||||
"released_backoff": released,
|
||||
"ready_tasks": ready,
|
||||
"submitted_tasks": submitted,
|
||||
"started_downloads": started_downloads,
|
||||
}
|
||||
|
||||
def run_forever(self, stop_event):
|
||||
LOG.info("strategy loop started")
|
||||
while not stop_event.is_set():
|
||||
try:
|
||||
stats = self.run_once()
|
||||
LOG.info("loop stats: %s", stats)
|
||||
except Exception:
|
||||
LOG.exception("strategy loop iteration failed")
|
||||
stop_event.wait(self.settings.poll_interval_seconds)
|
||||
LOG.info("strategy loop stopped")
|
||||
24
app/scheduler/planner.py
Normal file
24
app/scheduler/planner.py
Normal file
@@ -0,0 +1,24 @@
|
||||
from app.domain.gpu import expand_gpu_aliases, parse_aliases
|
||||
from app.domain.model_type import engine_for_model
|
||||
from app.storage.repositories import Repository
|
||||
from app.scheduler.budget import Budget
|
||||
|
||||
|
||||
class Planner:
|
||||
def __init__(self, repo: Repository, budget: Budget, default_aliases: list[str]):
|
||||
self.repo = repo
|
||||
self.budget = budget
|
||||
self.default_aliases = default_aliases
|
||||
|
||||
def materialize_tasks(self) -> int:
|
||||
created = 0
|
||||
for candidate in self.repo.list_candidates():
|
||||
if engine_for_model(candidate["model_id"]) != "vllm":
|
||||
continue
|
||||
aliases = parse_aliases(candidate["target_gpu_aliases"], self.default_aliases)
|
||||
for gpu_alias, gpu_type in expand_gpu_aliases(aliases):
|
||||
if not self.budget.can_create_task():
|
||||
return created
|
||||
if self.repo.insert_task_if_absent(candidate["model_id"], gpu_alias, gpu_type, candidate["priority"]):
|
||||
created += 1
|
||||
return created
|
||||
44
app/scheduler/submitter.py
Normal file
44
app/scheduler/submitter.py
Normal file
@@ -0,0 +1,44 @@
|
||||
from app.clients.modelhub import ModelHubClient
|
||||
from app.domain.vllm_configs import gen_vllm_config
|
||||
from app.storage.repositories import Repository, future_seconds
|
||||
|
||||
|
||||
class Submitter:
|
||||
def __init__(self, repo: Repository, client: ModelHubClient):
|
||||
self.repo = repo
|
||||
self.client = client
|
||||
|
||||
def submit_due(self, limit: int = 20) -> int:
|
||||
submitted = 0
|
||||
for task in self.repo.due_tasks(("ready_to_submit",), limit=limit):
|
||||
config = gen_vllm_config(task["gpu_type"], task["model_id"], task["max_model_len"])
|
||||
result = self.client.create_contest_task(task["model_id"], task["gpu_type"], config)
|
||||
self.repo.record_submission_attempt(
|
||||
task["id"],
|
||||
task["model_id"],
|
||||
task["gpu_type"],
|
||||
self.client.settings.strategy_id,
|
||||
self.client.settings.contest_task_strategy_field if self.client.settings.inject_strategy_id else "",
|
||||
result.code,
|
||||
result.message,
|
||||
result.result,
|
||||
)
|
||||
if result.result == "success":
|
||||
self.repo.mark_task_status(task["id"], "submitted", result.code, result.message, increment_attempt=True, increment_submit=True)
|
||||
submitted += 1
|
||||
elif result.result == "conflict":
|
||||
self.repo.mark_task_status(task["id"], "conflict", result.code, result.message, increment_attempt=True)
|
||||
elif result.result == "queue_full":
|
||||
backoff = min(3600, 300 * (task["attempt_count"] + 1))
|
||||
self.repo.mark_task_status(task["id"], "queue_full_backoff", result.code, result.message, future_seconds(backoff), increment_attempt=True)
|
||||
else:
|
||||
backoff = min(1800, 120 * (task["attempt_count"] + 1))
|
||||
self.repo.mark_task_status(task["id"], "pending", result.code, result.message, future_seconds(backoff), increment_attempt=True)
|
||||
return submitted
|
||||
|
||||
def release_queue_full_backoff(self) -> int:
|
||||
released = 0
|
||||
for task in self.repo.due_tasks(("queue_full_backoff",), limit=50):
|
||||
self.repo.mark_task_status(task["id"], "ready_to_submit")
|
||||
released += 1
|
||||
return released
|
||||
165
app/settings.py
Normal file
165
app/settings.py
Normal file
@@ -0,0 +1,165 @@
|
||||
from dataclasses import dataclass
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
|
||||
def _load_config_file() -> dict[str, Any]:
|
||||
explicit = os.getenv("STRATEGY_CONFIG_FILE")
|
||||
candidates = [explicit] if explicit else ["config.local.json", "config.json"]
|
||||
for item in candidates:
|
||||
if not item:
|
||||
continue
|
||||
path = Path(item)
|
||||
if path.exists():
|
||||
with path.open("r", encoding="utf-8") as fh:
|
||||
return json.load(fh)
|
||||
return {}
|
||||
|
||||
|
||||
def _value(config: dict[str, Any], name: str, default: Any = None) -> Any:
|
||||
if name in os.environ:
|
||||
return os.environ[name]
|
||||
if name in config:
|
||||
return config[name]
|
||||
lower = name.lower()
|
||||
if lower in config:
|
||||
return config[lower]
|
||||
return default
|
||||
|
||||
|
||||
def _bool_value(config: dict[str, Any], name: str, default: bool) -> bool:
|
||||
raw = _value(config, name, default)
|
||||
if isinstance(raw, bool):
|
||||
return raw
|
||||
if raw is None:
|
||||
return default
|
||||
return str(raw).strip().lower() in {"1", "true", "yes", "y", "on"}
|
||||
|
||||
|
||||
def _int_value(config: dict[str, Any], name: str, default: int) -> int:
|
||||
raw = _value(config, name, default)
|
||||
if raw is None or raw == "":
|
||||
return default
|
||||
return int(raw)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Settings:
|
||||
auth_token: str
|
||||
hf_token: str
|
||||
contest_api_token: str
|
||||
email: str
|
||||
user_id: str
|
||||
contributors: str
|
||||
strategy_id: str
|
||||
|
||||
modelhub_url: str = "https://modelhub.org.cn"
|
||||
modelhub_api_base: str = "https://modelhub.org.cn/api"
|
||||
modelhub_adminapi_base: str = "https://modelhub.org.cn/adminApi"
|
||||
|
||||
db_path: str = "./state.db"
|
||||
host: str = "0.0.0.0"
|
||||
port: int = 8080
|
||||
poll_interval_seconds: int = 60
|
||||
download_poll_interval_seconds: int = 300
|
||||
request_timeout_seconds: int = 60
|
||||
max_self_downloads: int = 8
|
||||
max_user_tasks: int = 2000
|
||||
default_gpu_aliases: str = "910b,k100,p800,166m,bi100,bi150,c500,s4000,mrv100,mlu370-x4,mlu370-x8"
|
||||
default_max_model_len: int = 1024
|
||||
|
||||
submit_dry_run: bool = False
|
||||
inject_strategy_id: bool = True
|
||||
contest_task_strategy_field: str = "strategyId"
|
||||
bounty_allow_direct_submit: bool = True
|
||||
auto_load_seed_path: str = ""
|
||||
|
||||
enable_not_adapted_crawler: bool = True
|
||||
target_machine_names: str = ""
|
||||
target_task_level: str = "文本生成"
|
||||
modelhub_page_size: int = 100
|
||||
crawler_refresh_seconds: int = 3600
|
||||
crawler_max_pages: int = 0
|
||||
|
||||
enable_download_success_crawler: bool = True
|
||||
download_success_page_size: int = 50
|
||||
download_success_max_pages: int = 0
|
||||
|
||||
config_file_loaded: str = ""
|
||||
|
||||
@property
|
||||
def default_gpu_alias_list(self) -> list[str]:
|
||||
return [item.strip() for item in self.default_gpu_aliases.split(",") if item.strip()]
|
||||
|
||||
|
||||
def load_settings(require_secrets: bool = True) -> Settings:
|
||||
config = _load_config_file()
|
||||
explicit = os.getenv("STRATEGY_CONFIG_FILE")
|
||||
loaded_path = ""
|
||||
if explicit and Path(explicit).exists():
|
||||
loaded_path = explicit
|
||||
elif Path("config.local.json").exists():
|
||||
loaded_path = "config.local.json"
|
||||
elif Path("config.json").exists():
|
||||
loaded_path = "config.json"
|
||||
|
||||
hf_token = _value(config, "HF_TOKEN", None)
|
||||
if hf_token is None:
|
||||
hf_token = _value(config, "HFTOKEN", "")
|
||||
|
||||
settings = Settings(
|
||||
auth_token=str(_value(config, "AUTH_TOKEN", "") or ""),
|
||||
hf_token=str(hf_token or ""),
|
||||
contest_api_token=str(_value(config, "CONFTEST_API_TOKEN", "") or ""),
|
||||
email=str(_value(config, "EMAIL", "") or ""),
|
||||
user_id=str(_value(config, "USER_ID", "") or ""),
|
||||
contributors=str(_value(config, "CONTRIBUTORS", "") or ""),
|
||||
strategy_id=str(_value(config, "STRATEGY_ID", "") or ""),
|
||||
modelhub_url=str(_value(config, "MODELHUB_URL", "https://modelhub.org.cn")),
|
||||
modelhub_api_base=str(_value(config, "MODELHUB_API_BASE", "https://modelhub.org.cn/api")),
|
||||
modelhub_adminapi_base=str(_value(config, "MODELHUB_ADMINAPI_BASE", "https://modelhub.org.cn/adminApi")),
|
||||
db_path=str(_value(config, "DB_PATH", "./state.db")),
|
||||
host=str(_value(config, "HOST", "0.0.0.0")),
|
||||
port=_int_value(config, "PORT", 8080),
|
||||
poll_interval_seconds=_int_value(config, "POLL_INTERVAL_SECONDS", 60),
|
||||
download_poll_interval_seconds=_int_value(config, "DOWNLOAD_POLL_INTERVAL_SECONDS", 300),
|
||||
request_timeout_seconds=_int_value(config, "REQUEST_TIMEOUT_SECONDS", 60),
|
||||
max_self_downloads=_int_value(config, "MAX_SELF_DOWNLOADS", 8),
|
||||
max_user_tasks=_int_value(config, "MAX_USER_TASKS", 2000),
|
||||
default_gpu_aliases=str(_value(config, "DEFAULT_GPU_ALIASES", "910b,k100,p800,166m,bi100,bi150,c500,s4000,mrv100,mlu370-x4,mlu370-x8")),
|
||||
default_max_model_len=_int_value(config, "DEFAULT_MAX_MODEL_LEN", 1024),
|
||||
submit_dry_run=_bool_value(config, "SUBMIT_DRY_RUN", False),
|
||||
inject_strategy_id=_bool_value(config, "INJECT_STRATEGY_ID", True),
|
||||
contest_task_strategy_field=str(_value(config, "CONTEST_TASK_STRATEGY_FIELD", "strategyId")),
|
||||
bounty_allow_direct_submit=_bool_value(config, "BOUNTY_ALLOW_DIRECT_SUBMIT", True),
|
||||
auto_load_seed_path=str(_value(config, "AUTO_LOAD_SEED_PATH", "") or ""),
|
||||
enable_not_adapted_crawler=_bool_value(config, "ENABLE_NOT_ADAPTED_CRAWLER", True),
|
||||
target_machine_names=str(_value(config, "TARGET_MACHINE_NAMES", "") or ""),
|
||||
target_task_level=str(_value(config, "TARGET_TASK_LEVEL", "文本生成")),
|
||||
modelhub_page_size=_int_value(config, "MODELHUB_PAGE_SIZE", 100),
|
||||
crawler_refresh_seconds=_int_value(config, "CRAWLER_REFRESH_SECONDS", 3600),
|
||||
crawler_max_pages=_int_value(config, "CRAWLER_MAX_PAGES", 0),
|
||||
enable_download_success_crawler=_bool_value(config, "ENABLE_DOWNLOAD_SUCCESS_CRAWLER", True),
|
||||
download_success_page_size=_int_value(config, "DOWNLOAD_SUCCESS_PAGE_SIZE", 50),
|
||||
download_success_max_pages=_int_value(config, "DOWNLOAD_SUCCESS_MAX_PAGES", 0),
|
||||
config_file_loaded=loaded_path,
|
||||
)
|
||||
if require_secrets:
|
||||
missing = []
|
||||
required_names = [
|
||||
("AUTH_TOKEN", settings.auth_token),
|
||||
("CONFTEST_API_TOKEN", settings.contest_api_token),
|
||||
("EMAIL", settings.email),
|
||||
("USER_ID", settings.user_id),
|
||||
("CONTRIBUTORS", settings.contributors),
|
||||
]
|
||||
for name, value in required_names:
|
||||
if not value:
|
||||
missing.append(name)
|
||||
if settings.inject_strategy_id and not settings.strategy_id:
|
||||
missing.append("STRATEGY_ID")
|
||||
if missing:
|
||||
raise RuntimeError(f"Missing required settings: {', '.join(missing)}")
|
||||
return settings
|
||||
14
app/shutdown.py
Normal file
14
app/shutdown.py
Normal file
@@ -0,0 +1,14 @@
|
||||
import signal
|
||||
import threading
|
||||
|
||||
|
||||
stop_event = threading.Event()
|
||||
|
||||
|
||||
def request_shutdown(signum=None, frame=None):
|
||||
stop_event.set()
|
||||
|
||||
|
||||
def install_signal_handlers():
|
||||
signal.signal(signal.SIGTERM, request_shutdown)
|
||||
signal.signal(signal.SIGINT, request_shutdown)
|
||||
1
app/storage/__init__.py
Normal file
1
app/storage/__init__.py
Normal file
@@ -0,0 +1 @@
|
||||
"""SQLite storage helpers."""
|
||||
18
app/storage/db.py
Normal file
18
app/storage/db.py
Normal file
@@ -0,0 +1,18 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import sqlite3
|
||||
|
||||
from app.storage.schema import init_schema
|
||||
|
||||
|
||||
def connect(db_path: str) -> sqlite3.Connection:
|
||||
parent = os.path.dirname(os.path.abspath(db_path))
|
||||
if parent and not os.path.exists(parent):
|
||||
os.makedirs(parent, exist_ok=True)
|
||||
conn = sqlite3.connect(db_path, timeout=30)
|
||||
conn.row_factory = sqlite3.Row
|
||||
conn.execute("PRAGMA journal_mode=WAL")
|
||||
conn.execute("PRAGMA busy_timeout=30000")
|
||||
init_schema(conn)
|
||||
return conn
|
||||
266
app/storage/repositories.py
Normal file
266
app/storage/repositories.py
Normal file
@@ -0,0 +1,266 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime, timezone, timedelta
|
||||
import sqlite3
|
||||
|
||||
|
||||
def utcnow() -> str:
|
||||
return datetime.now(timezone.utc).isoformat()
|
||||
|
||||
|
||||
def future_seconds(seconds: int) -> str:
|
||||
return (datetime.now(timezone.utc) + timedelta(seconds=seconds)).isoformat()
|
||||
|
||||
|
||||
class Repository:
|
||||
def __init__(self, conn: sqlite3.Connection):
|
||||
self.conn = conn
|
||||
|
||||
def upsert_candidate(
|
||||
self,
|
||||
model_id: str,
|
||||
source: str,
|
||||
origin: str,
|
||||
priority: int,
|
||||
is_bounty: bool = False,
|
||||
is_promote: bool = False,
|
||||
known_downloaded_by_others: bool = False,
|
||||
self_download_allowed: bool = False,
|
||||
max_model_len: int = 1024,
|
||||
target_gpu_aliases: str | None = None,
|
||||
notes: str | None = None,
|
||||
) -> None:
|
||||
now = utcnow()
|
||||
existing = self.conn.execute(
|
||||
"SELECT target_gpu_aliases FROM candidates WHERE model_id = ?",
|
||||
(model_id,),
|
||||
).fetchone()
|
||||
if existing and target_gpu_aliases:
|
||||
old_aliases = [item.strip() for item in (existing["target_gpu_aliases"] or "").split(",") if item.strip()]
|
||||
new_aliases = [item.strip() for item in target_gpu_aliases.split(",") if item.strip()]
|
||||
merged = []
|
||||
for alias in old_aliases + new_aliases:
|
||||
if alias not in merged:
|
||||
merged.append(alias)
|
||||
target_gpu_aliases = ",".join(merged)
|
||||
self.conn.execute(
|
||||
"""
|
||||
INSERT INTO candidates (
|
||||
model_id, source, origin, priority, is_bounty, is_promote,
|
||||
known_downloaded_by_others, self_download_allowed, max_model_len,
|
||||
target_gpu_aliases, notes, created_at, updated_at
|
||||
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
ON CONFLICT(model_id) DO UPDATE SET
|
||||
source=excluded.source,
|
||||
origin=excluded.origin,
|
||||
priority=min(candidates.priority, excluded.priority),
|
||||
is_bounty=max(candidates.is_bounty, excluded.is_bounty),
|
||||
is_promote=max(candidates.is_promote, excluded.is_promote),
|
||||
known_downloaded_by_others=max(candidates.known_downloaded_by_others, excluded.known_downloaded_by_others),
|
||||
self_download_allowed=max(candidates.self_download_allowed, excluded.self_download_allowed),
|
||||
max_model_len=excluded.max_model_len,
|
||||
target_gpu_aliases=excluded.target_gpu_aliases,
|
||||
notes=excluded.notes,
|
||||
updated_at=excluded.updated_at
|
||||
""",
|
||||
(
|
||||
model_id,
|
||||
source,
|
||||
origin,
|
||||
priority,
|
||||
int(is_bounty),
|
||||
int(is_promote),
|
||||
int(known_downloaded_by_others),
|
||||
int(self_download_allowed),
|
||||
max_model_len,
|
||||
target_gpu_aliases,
|
||||
notes,
|
||||
now,
|
||||
now,
|
||||
),
|
||||
)
|
||||
self.conn.commit()
|
||||
|
||||
def list_candidates(self) -> list[sqlite3.Row]:
|
||||
cursor = self.conn.execute("SELECT * FROM candidates ORDER BY priority, created_at, model_id")
|
||||
return list(cursor.fetchall())
|
||||
|
||||
def insert_task_if_absent(self, model_id: str, gpu_alias: str, gpu_type: str, priority: int) -> bool:
|
||||
now = utcnow()
|
||||
cur = self.conn.execute(
|
||||
"""
|
||||
INSERT OR IGNORE INTO tasks (
|
||||
model_id, gpu_alias, gpu_type, engine, priority, status,
|
||||
created_at, updated_at
|
||||
) VALUES (?, ?, ?, 'vllm', ?, 'pending', ?, ?)
|
||||
""",
|
||||
(model_id, gpu_alias, gpu_type, priority, now, now),
|
||||
)
|
||||
self.conn.commit()
|
||||
return cur.rowcount > 0
|
||||
|
||||
def count_budgeted_tasks(self) -> int:
|
||||
cursor = self.conn.execute("SELECT COUNT(*) FROM tasks WHERE status != 'skipped'")
|
||||
return int(cursor.fetchone()[0])
|
||||
|
||||
def mark_task_status(
|
||||
self,
|
||||
task_id: int,
|
||||
status: str,
|
||||
error_code: str | None = None,
|
||||
error_message: str | None = None,
|
||||
next_attempt_at: str | None = None,
|
||||
increment_attempt: bool = False,
|
||||
increment_submit: bool = False,
|
||||
) -> None:
|
||||
now = utcnow()
|
||||
self.conn.execute(
|
||||
"""
|
||||
UPDATE tasks SET
|
||||
status = ?,
|
||||
last_error_code = ?,
|
||||
last_error_message = ?,
|
||||
next_attempt_at = ?,
|
||||
attempt_count = attempt_count + ?,
|
||||
submit_count = submit_count + ?,
|
||||
submitted_at = CASE WHEN ? = 'submitted' THEN ? ELSE submitted_at END,
|
||||
completed_at = CASE WHEN ? IN ('submitted', 'conflict', 'failed', 'skipped') THEN ? ELSE completed_at END,
|
||||
updated_at = ?
|
||||
WHERE id = ?
|
||||
""",
|
||||
(
|
||||
status,
|
||||
error_code,
|
||||
error_message,
|
||||
next_attempt_at,
|
||||
1 if increment_attempt else 0,
|
||||
1 if increment_submit else 0,
|
||||
status,
|
||||
now,
|
||||
status,
|
||||
now,
|
||||
now,
|
||||
task_id,
|
||||
),
|
||||
)
|
||||
self.conn.commit()
|
||||
|
||||
def due_tasks(self, statuses: tuple[str, ...], limit: int = 50) -> list[sqlite3.Row]:
|
||||
placeholders = ",".join("?" for _ in statuses)
|
||||
now = utcnow()
|
||||
cursor = self.conn.execute(
|
||||
f"""
|
||||
SELECT t.*, c.is_bounty, c.is_promote, c.known_downloaded_by_others,
|
||||
c.self_download_allowed, c.max_model_len, c.source
|
||||
FROM tasks t
|
||||
JOIN candidates c ON c.model_id = t.model_id
|
||||
WHERE t.status IN ({placeholders})
|
||||
AND (t.next_attempt_at IS NULL OR t.next_attempt_at <= ?)
|
||||
ORDER BY t.priority, t.attempt_count, t.created_at, t.model_id, t.gpu_type
|
||||
LIMIT ?
|
||||
""",
|
||||
(*statuses, now, limit),
|
||||
)
|
||||
return list(cursor.fetchall())
|
||||
|
||||
def count_active_self_downloads(self) -> int:
|
||||
cursor = self.conn.execute(
|
||||
"SELECT COUNT(*) FROM downloads WHERE active_self_slot = 1 AND status IN ('WAITING', 'RUNNING', 'unknown')"
|
||||
)
|
||||
return int(cursor.fetchone()[0])
|
||||
|
||||
def upsert_download(
|
||||
self,
|
||||
model_id: str,
|
||||
source: str,
|
||||
status: str,
|
||||
owner: str,
|
||||
active_self_slot: bool = False,
|
||||
download_code: str | None = None,
|
||||
download_message: str | None = None,
|
||||
next_poll_at: str | None = None,
|
||||
) -> None:
|
||||
now = utcnow()
|
||||
self.conn.execute(
|
||||
"""
|
||||
INSERT INTO downloads (
|
||||
model_id, source, status, owner, active_self_slot, download_code,
|
||||
download_message, last_polled_at, next_poll_at, created_at, updated_at
|
||||
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
ON CONFLICT(model_id) DO UPDATE SET
|
||||
source=excluded.source,
|
||||
status=excluded.status,
|
||||
owner=excluded.owner,
|
||||
active_self_slot=excluded.active_self_slot,
|
||||
download_code=excluded.download_code,
|
||||
download_message=excluded.download_message,
|
||||
last_polled_at=excluded.last_polled_at,
|
||||
next_poll_at=excluded.next_poll_at,
|
||||
updated_at=excluded.updated_at
|
||||
""",
|
||||
(
|
||||
model_id,
|
||||
source,
|
||||
status,
|
||||
owner,
|
||||
int(active_self_slot),
|
||||
download_code,
|
||||
download_message,
|
||||
now,
|
||||
next_poll_at,
|
||||
now,
|
||||
now,
|
||||
),
|
||||
)
|
||||
self.conn.commit()
|
||||
|
||||
def get_download(self, model_id: str):
|
||||
cursor = self.conn.execute("SELECT * FROM downloads WHERE model_id = ?", (model_id,))
|
||||
return cursor.fetchone()
|
||||
|
||||
def active_downloads_due(self, limit: int = 20) -> list[sqlite3.Row]:
|
||||
now = utcnow()
|
||||
cursor = self.conn.execute(
|
||||
"""
|
||||
SELECT * FROM downloads
|
||||
WHERE active_self_slot = 1
|
||||
AND status IN ('WAITING', 'RUNNING', 'unknown')
|
||||
AND (next_poll_at IS NULL OR next_poll_at <= ?)
|
||||
ORDER BY next_poll_at, created_at
|
||||
LIMIT ?
|
||||
""",
|
||||
(now, limit),
|
||||
)
|
||||
return list(cursor.fetchall())
|
||||
|
||||
def record_submission_attempt(
|
||||
self,
|
||||
task_id: int,
|
||||
model_id: str,
|
||||
gpu_type: str,
|
||||
strategy_id: str,
|
||||
payload_strategy_field: str,
|
||||
response_code: str | None,
|
||||
response_message: str | None,
|
||||
result: str,
|
||||
) -> None:
|
||||
self.conn.execute(
|
||||
"""
|
||||
INSERT INTO submission_attempts (
|
||||
task_id, model_id, gpu_type, strategy_id, payload_strategy_field,
|
||||
response_code, response_message, result, created_at
|
||||
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
""",
|
||||
(
|
||||
task_id,
|
||||
model_id,
|
||||
gpu_type,
|
||||
strategy_id,
|
||||
payload_strategy_field,
|
||||
response_code,
|
||||
response_message,
|
||||
result,
|
||||
utcnow(),
|
||||
),
|
||||
)
|
||||
self.conn.commit()
|
||||
87
app/storage/schema.py
Normal file
87
app/storage/schema.py
Normal file
@@ -0,0 +1,87 @@
|
||||
SCHEMA_SQL = """
|
||||
CREATE TABLE IF NOT EXISTS candidates (
|
||||
model_id TEXT PRIMARY KEY,
|
||||
source TEXT NOT NULL DEFAULT 'HUGGING_FACE',
|
||||
origin TEXT NOT NULL,
|
||||
priority INTEGER NOT NULL,
|
||||
is_bounty INTEGER NOT NULL DEFAULT 0,
|
||||
is_promote INTEGER NOT NULL DEFAULT 0,
|
||||
known_downloaded_by_others INTEGER NOT NULL DEFAULT 0,
|
||||
self_download_allowed INTEGER NOT NULL DEFAULT 0,
|
||||
max_model_len INTEGER NOT NULL DEFAULT 1024,
|
||||
target_gpu_aliases TEXT,
|
||||
notes TEXT,
|
||||
created_at TEXT NOT NULL,
|
||||
updated_at TEXT NOT NULL
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS tasks (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
model_id TEXT NOT NULL,
|
||||
gpu_alias TEXT NOT NULL,
|
||||
gpu_type TEXT NOT NULL,
|
||||
engine TEXT NOT NULL DEFAULT 'vllm',
|
||||
priority INTEGER NOT NULL,
|
||||
status TEXT NOT NULL DEFAULT 'pending',
|
||||
attempt_count INTEGER NOT NULL DEFAULT 0,
|
||||
submit_count INTEGER NOT NULL DEFAULT 0,
|
||||
last_error_code TEXT,
|
||||
last_error_message TEXT,
|
||||
next_attempt_at TEXT,
|
||||
submitted_at TEXT,
|
||||
completed_at TEXT,
|
||||
created_at TEXT NOT NULL,
|
||||
updated_at TEXT NOT NULL,
|
||||
UNIQUE(model_id, gpu_type, engine)
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS downloads (
|
||||
model_id TEXT PRIMARY KEY,
|
||||
source TEXT NOT NULL,
|
||||
status TEXT NOT NULL DEFAULT 'unknown',
|
||||
owner TEXT NOT NULL DEFAULT 'unknown',
|
||||
active_self_slot INTEGER NOT NULL DEFAULT 0,
|
||||
download_code TEXT,
|
||||
download_message TEXT,
|
||||
waiting_since TEXT,
|
||||
last_polled_at TEXT,
|
||||
next_poll_at TEXT,
|
||||
created_at TEXT NOT NULL,
|
||||
updated_at TEXT NOT NULL
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS submission_attempts (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
task_id INTEGER NOT NULL,
|
||||
model_id TEXT NOT NULL,
|
||||
gpu_type TEXT NOT NULL,
|
||||
strategy_id TEXT,
|
||||
payload_strategy_field TEXT,
|
||||
response_code TEXT,
|
||||
response_message TEXT,
|
||||
result TEXT NOT NULL,
|
||||
created_at TEXT NOT NULL
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS seed_imports (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
file_path TEXT NOT NULL,
|
||||
file_hash TEXT,
|
||||
imported_count INTEGER NOT NULL,
|
||||
created_at TEXT NOT NULL
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_tasks_sched
|
||||
ON tasks(status, priority, next_attempt_at, created_at);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_tasks_model
|
||||
ON tasks(model_id);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_downloads_poll
|
||||
ON downloads(status, next_poll_at);
|
||||
"""
|
||||
|
||||
|
||||
def init_schema(conn):
|
||||
conn.executescript(SCHEMA_SQL)
|
||||
conn.commit()
|
||||
43
config.example.json
Normal file
43
config.example.json
Normal file
@@ -0,0 +1,43 @@
|
||||
{
|
||||
"AUTH_TOKEN": "",
|
||||
"HF_TOKEN": "",
|
||||
"CONFTEST_API_TOKEN": "",
|
||||
"EMAIL": "",
|
||||
"USER_ID": "",
|
||||
"CONTRIBUTORS": "",
|
||||
"STRATEGY_ID": "",
|
||||
|
||||
"MODELHUB_URL": "https://modelhub.org.cn",
|
||||
"MODELHUB_API_BASE": "https://modelhub.org.cn/api",
|
||||
"MODELHUB_ADMINAPI_BASE": "https://modelhub.org.cn/adminApi",
|
||||
|
||||
"DB_PATH": "./state.db",
|
||||
"HOST": "0.0.0.0",
|
||||
"PORT": 8080,
|
||||
|
||||
"SUBMIT_DRY_RUN": true,
|
||||
"INJECT_STRATEGY_ID": true,
|
||||
"CONTEST_TASK_STRATEGY_FIELD": "strategyId",
|
||||
|
||||
"MAX_USER_TASKS": 2000,
|
||||
"MAX_SELF_DOWNLOADS": 8,
|
||||
"POLL_INTERVAL_SECONDS": 60,
|
||||
"DOWNLOAD_POLL_INTERVAL_SECONDS": 300,
|
||||
"CRAWLER_REFRESH_SECONDS": 3600,
|
||||
|
||||
"DEFAULT_GPU_ALIASES": "910b,k100,p800,166m,bi100,bi150,c500,s4000,mrv100,mlu370-x4,mlu370-x8",
|
||||
"DEFAULT_MAX_MODEL_LEN": 1024,
|
||||
|
||||
"ENABLE_NOT_ADAPTED_CRAWLER": true,
|
||||
"TARGET_MACHINE_NAMES": "MTT S4000,Hygon K100,Kunlunxin P800",
|
||||
"TARGET_TASK_LEVEL": "文本生成",
|
||||
"MODELHUB_PAGE_SIZE": 100,
|
||||
"CRAWLER_MAX_PAGES": 0,
|
||||
|
||||
"ENABLE_DOWNLOAD_SUCCESS_CRAWLER": true,
|
||||
"DOWNLOAD_SUCCESS_PAGE_SIZE": 50,
|
||||
"DOWNLOAD_SUCCESS_MAX_PAGES": 0,
|
||||
|
||||
"BOUNTY_ALLOW_DIRECT_SUBMIT": true,
|
||||
"AUTO_LOAD_SEED_PATH": ""
|
||||
}
|
||||
3
requirements.txt
Normal file
3
requirements.txt
Normal file
@@ -0,0 +1,3 @@
|
||||
flask>=3.0.0,<4
|
||||
requests>=2.31.0,<3
|
||||
pytest>=8.0.0,<9
|
||||
1
scripts/__init__.py
Normal file
1
scripts/__init__.py
Normal file
@@ -0,0 +1 @@
|
||||
"""Utility scripts for the vLLM agent strategy."""
|
||||
91
scripts/agent_platform_api.py
Normal file
91
scripts/agent_platform_api.py
Normal file
@@ -0,0 +1,91 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from typing import Any
|
||||
|
||||
import requests
|
||||
|
||||
BASE_URL = "https://modelhub.org.cn/agent_platform"
|
||||
|
||||
|
||||
def get_token(args) -> str:
|
||||
token = args.token or os.getenv("AGENT_PLATFORM_TOKEN") or os.getenv("AUTH_TOKEN")
|
||||
if not token:
|
||||
raise SystemExit("Missing token. Use --token, AGENT_PLATFORM_TOKEN, or AUTH_TOKEN.")
|
||||
return token
|
||||
|
||||
|
||||
def request(method: str, path: str, token: str, **kwargs) -> Any:
|
||||
url = BASE_URL.rstrip("/") + path
|
||||
headers = kwargs.pop("headers", {})
|
||||
headers["Authorization"] = f"Bearer {token}"
|
||||
if "json" in kwargs:
|
||||
headers["Content-Type"] = "application/json"
|
||||
resp = requests.request(method, url, headers=headers, timeout=60, **kwargs)
|
||||
if resp.status_code == 204:
|
||||
return {"status": 204, "message": "no content"}
|
||||
try:
|
||||
body = resp.json()
|
||||
except Exception:
|
||||
body = {"status_code": resp.status_code, "text": resp.text}
|
||||
if resp.status_code >= 400:
|
||||
raise SystemExit(json.dumps(body, ensure_ascii=False, indent=2))
|
||||
return body
|
||||
|
||||
|
||||
def print_json(data: Any) -> None:
|
||||
print(json.dumps(data, ensure_ascii=False, indent=2))
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Agent platform API helper")
|
||||
parser.add_argument("--token", default="", help="Bearer token; defaults to AGENT_PLATFORM_TOKEN or AUTH_TOKEN")
|
||||
sub = parser.add_subparsers(dest="cmd", required=True)
|
||||
|
||||
sub.add_parser("me")
|
||||
sub.add_parser("list")
|
||||
|
||||
create = sub.add_parser("create")
|
||||
create.add_argument("--name", required=True)
|
||||
create.add_argument("--repo-url", required=True)
|
||||
create.add_argument("--tag", required=True)
|
||||
|
||||
for name in ["get", "sync", "deploy", "logs", "stop", "delete"]:
|
||||
p = sub.add_parser(name)
|
||||
p.add_argument("--strategy-id", required=True)
|
||||
if name == "logs":
|
||||
p.add_argument("--limit", type=int, default=100)
|
||||
|
||||
args = parser.parse_args()
|
||||
token = get_token(args)
|
||||
|
||||
if args.cmd == "me":
|
||||
print_json(request("GET", "/auth/me", token))
|
||||
elif args.cmd == "list":
|
||||
print_json(request("GET", "/strategies", token))
|
||||
elif args.cmd == "create":
|
||||
body = {"name": args.name, "repo_url": args.repo_url, "tag": args.tag}
|
||||
data = request("POST", "/strategies", token, json=body)
|
||||
print_json(data)
|
||||
strategy_id = data.get("id")
|
||||
if strategy_id:
|
||||
print(f"\nSTRATEGY_ID={strategy_id}", file=sys.stderr)
|
||||
elif args.cmd == "get":
|
||||
print_json(request("GET", f"/strategies/{args.strategy_id}", token))
|
||||
elif args.cmd == "sync":
|
||||
print_json(request("POST", f"/strategies/{args.strategy_id}/sync", token))
|
||||
elif args.cmd == "deploy":
|
||||
print_json(request("POST", f"/strategies/{args.strategy_id}/deploy", token))
|
||||
elif args.cmd == "logs":
|
||||
print_json(request("GET", f"/strategies/{args.strategy_id}/logs", token, params={"limit": args.limit}))
|
||||
elif args.cmd == "stop":
|
||||
print_json(request("POST", f"/strategies/{args.strategy_id}/stop", token))
|
||||
elif args.cmd == "delete":
|
||||
print_json(request("DELETE", f"/strategies/{args.strategy_id}", token))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
31
scripts/dry_run_plan.py
Normal file
31
scripts/dry_run_plan.py
Normal file
@@ -0,0 +1,31 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
|
||||
from app.scheduler.budget import Budget
|
||||
from app.scheduler.planner import Planner
|
||||
from app.settings import load_settings
|
||||
from app.storage.db import connect
|
||||
from app.storage.repositories import Repository
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--limit", type=int, default=50)
|
||||
args = parser.parse_args()
|
||||
|
||||
settings = load_settings(require_secrets=False)
|
||||
conn = connect(settings.db_path)
|
||||
try:
|
||||
repo = Repository(conn)
|
||||
Planner(repo, Budget(repo, settings.max_user_tasks), settings.default_gpu_alias_list).materialize_tasks()
|
||||
rows = repo.due_tasks(("pending", "ready_to_submit", "waiting_download", "queue_full_backoff"), limit=args.limit)
|
||||
for row in rows:
|
||||
print(f"#{row['id']} priority={row['priority']} status={row['status']} {row['model_id']} -> {row['gpu_type']} ({row['gpu_alias']})")
|
||||
print(f"budget used={repo.count_budgeted_tasks()} max={settings.max_user_tasks}")
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
13
scripts/init_db.py
Normal file
13
scripts/init_db.py
Normal file
@@ -0,0 +1,13 @@
|
||||
from app.settings import load_settings
|
||||
from app.storage.db import connect
|
||||
|
||||
|
||||
def main():
|
||||
settings = load_settings(require_secrets=False)
|
||||
conn = connect(settings.db_path)
|
||||
conn.close()
|
||||
print(f"initialized db: {settings.db_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
63
scripts/load_seed.py
Normal file
63
scripts/load_seed.py
Normal file
@@ -0,0 +1,63 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import csv
|
||||
|
||||
from app.domain.priorities import ORIGIN_TO_PRIORITY, PRIORITY_LOW_CONFIDENCE
|
||||
from app.settings import load_settings
|
||||
from app.storage.db import connect
|
||||
from app.storage.repositories import Repository
|
||||
|
||||
|
||||
def flags_for_origin(origin: str) -> tuple[bool, bool, bool, bool]:
|
||||
is_bounty = origin in {"manual_bounty", "bounty"}
|
||||
is_promote = origin in {"manual_promote", "promote"}
|
||||
downloaded = is_bounty or origin in {"manual_downloaded_by_others", "downloaded_by_others"}
|
||||
self_download = origin in {"hf_seed", "self_download"}
|
||||
return is_bounty, is_promote, downloaded, self_download
|
||||
|
||||
|
||||
def load_seed(path: str, origin: str, repo: Repository) -> int:
|
||||
priority = ORIGIN_TO_PRIORITY.get(origin, PRIORITY_LOW_CONFIDENCE)
|
||||
count = 0
|
||||
with open(path, newline="", encoding="utf-8") as fh:
|
||||
reader = csv.DictReader(fh)
|
||||
for row in reader:
|
||||
model_id = (row.get("model_id") or "").strip()
|
||||
if not model_id:
|
||||
continue
|
||||
is_bounty, is_promote, downloaded, self_download = flags_for_origin(origin)
|
||||
repo.upsert_candidate(
|
||||
model_id=model_id,
|
||||
source=(row.get("source") or "HUGGING_FACE").strip(),
|
||||
origin=origin,
|
||||
priority=priority,
|
||||
is_bounty=is_bounty,
|
||||
is_promote=is_promote,
|
||||
known_downloaded_by_others=downloaded,
|
||||
self_download_allowed=self_download,
|
||||
max_model_len=int(row.get("max_model_len") or 1024),
|
||||
target_gpu_aliases=(row.get("gpu_aliases") or "").strip() or None,
|
||||
notes=row.get("notes"),
|
||||
)
|
||||
count += 1
|
||||
return count
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--file", required=True)
|
||||
parser.add_argument("--origin", required=True, choices=sorted(ORIGIN_TO_PRIORITY))
|
||||
args = parser.parse_args()
|
||||
|
||||
settings = load_settings(require_secrets=False)
|
||||
conn = connect(settings.db_path)
|
||||
try:
|
||||
count = load_seed(args.file, args.origin, Repository(conn))
|
||||
print(f"imported {count} candidates from {args.file} as {args.origin}")
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
42
scripts/smoke_check.py
Normal file
42
scripts/smoke_check.py
Normal file
@@ -0,0 +1,42 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
|
||||
from app.clients.modelhub import ModelHubClient
|
||||
from app.domain.gpu import expand_gpu_aliases
|
||||
from app.domain.vllm_configs import gen_vllm_config
|
||||
from app.settings import load_settings
|
||||
|
||||
|
||||
def redact(payload: dict) -> dict:
|
||||
clone = dict(payload)
|
||||
if "contestApiToken" in clone:
|
||||
clone["contestApiToken"] = "***"
|
||||
return clone
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--no-submit", action="store_true", help="Only print a redacted payload")
|
||||
parser.add_argument("--model-id", default="Qwen/Qwen2.5-0.5B-Instruct")
|
||||
parser.add_argument("--gpu-alias", default="k100")
|
||||
args = parser.parse_args()
|
||||
|
||||
settings = load_settings(require_secrets=not args.no_submit)
|
||||
client = ModelHubClient(settings)
|
||||
pairs = expand_gpu_aliases([args.gpu_alias])
|
||||
if not pairs:
|
||||
raise SystemExit(f"unsupported gpu alias for vLLM: {args.gpu_alias}")
|
||||
_, gpu_type = pairs[0]
|
||||
config = gen_vllm_config(gpu_type, args.model_id, settings.default_max_model_len)
|
||||
payload = client.build_contest_payload(args.model_id, gpu_type, config)
|
||||
print(json.dumps(redact(payload), ensure_ascii=False, indent=2))
|
||||
if args.no_submit:
|
||||
return
|
||||
result = client.create_contest_task(args.model_id, gpu_type, config)
|
||||
print(result)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
20
seeds/README.md
Normal file
20
seeds/README.md
Normal file
@@ -0,0 +1,20 @@
|
||||
# Seed files
|
||||
|
||||
当前代码库没有找到悬赏模型 API,因此 v1 使用 CSV seed 文件。
|
||||
|
||||
CSV 字段:
|
||||
|
||||
```csv
|
||||
model_id,source,gpu_aliases,max_model_len,notes
|
||||
```
|
||||
|
||||
- `source`: 一般是 `HUGGING_FACE`。
|
||||
- `gpu_aliases`: 例如 `910b,k100,p800,biren166m`。
|
||||
- `max_model_len`: 默认 1024。
|
||||
|
||||
导入时用 `--origin` 控制优先级:
|
||||
|
||||
- `manual_bounty`: 悬赏模型,最高优先级。
|
||||
- `manual_promote`: 已入库/已验证成功模型。
|
||||
- `manual_downloaded_by_others`: 别人已下载模型。
|
||||
- `hf_seed`: 低优先级自下载候选。
|
||||
2
seeds/bounty_models.example.csv
Normal file
2
seeds/bounty_models.example.csv
Normal file
@@ -0,0 +1,2 @@
|
||||
model_id,source,gpu_aliases,max_model_len,notes
|
||||
Qwen/Qwen2.5-0.5B-Instruct,HUGGING_FACE,"910b,k100,p800",1024,"manual bounty seed example"
|
||||
|
2
seeds/downloaded_by_others.example.csv
Normal file
2
seeds/downloaded_by_others.example.csv
Normal file
@@ -0,0 +1,2 @@
|
||||
model_id,source,gpu_aliases,max_model_len,notes
|
||||
Qwen/Qwen2.5-0.5B-Instruct,HUGGING_FACE,"bi100,c500",1024,"manual downloaded-by-others seed example"
|
||||
|
2
seeds/promote_models.example.csv
Normal file
2
seeds/promote_models.example.csv
Normal file
@@ -0,0 +1,2 @@
|
||||
model_id,source,gpu_aliases,max_model_len,notes
|
||||
Qwen/Qwen2.5-0.5B-Instruct,HUGGING_FACE,"k100,biren166m",1024,"manual promote seed example"
|
||||
|
0
tests/__init__.py
Normal file
0
tests/__init__.py
Normal file
12
tests/test_download_slot_policy.py
Normal file
12
tests/test_download_slot_policy.py
Normal file
@@ -0,0 +1,12 @@
|
||||
from app.storage.db import connect
|
||||
from app.storage.repositories import Repository
|
||||
|
||||
|
||||
def test_active_self_download_count(tmp_path):
|
||||
conn = connect(str(tmp_path / "state.db"))
|
||||
repo = Repository(conn)
|
||||
for i in range(8):
|
||||
repo.upsert_download(f"owner/model-{i}", "HUGGING_FACE", "RUNNING", "self", True)
|
||||
repo.upsert_download("owner/other", "HUGGING_FACE", "RUNNING", "others", False)
|
||||
assert repo.count_active_self_downloads() == 8
|
||||
conn.close()
|
||||
5
tests/test_priority_order.py
Normal file
5
tests/test_priority_order.py
Normal file
@@ -0,0 +1,5 @@
|
||||
from app.domain.priorities import PRIORITY_BOUNTY, PRIORITY_PROMOTE, PRIORITY_OTHERS_DOWNLOADED, PRIORITY_SELF_DOWNLOAD
|
||||
|
||||
|
||||
def test_priority_order():
|
||||
assert PRIORITY_BOUNTY < PRIORITY_PROMOTE < PRIORITY_OTHERS_DOWNLOADED < PRIORITY_SELF_DOWNLOAD
|
||||
21
tests/test_strategy_id_payload.py
Normal file
21
tests/test_strategy_id_payload.py
Normal file
@@ -0,0 +1,21 @@
|
||||
from app.clients.modelhub import ModelHubClient
|
||||
from app.settings import Settings
|
||||
|
||||
|
||||
def test_strategy_id_payload_field_configurable():
|
||||
settings = Settings(
|
||||
auth_token="a", hf_token="h", contest_api_token="c", email="e", user_id="1",
|
||||
contributors="u", strategy_id="sid", contest_task_strategy_field="strategy_id"
|
||||
)
|
||||
payload = ModelHubClient(settings).build_contest_payload("owner/model", "hygon_k100-ai", "cfg")
|
||||
assert payload["strategy_id"] == "sid"
|
||||
assert payload["taskType"] == "text-generation"
|
||||
|
||||
|
||||
def test_strategy_id_payload_can_disable_injection():
|
||||
settings = Settings(
|
||||
auth_token="a", hf_token="h", contest_api_token="c", email="e", user_id="1",
|
||||
contributors="u", strategy_id="sid", inject_strategy_id=False
|
||||
)
|
||||
payload = ModelHubClient(settings).build_contest_payload("owner/model", "hygon_k100-ai", "cfg")
|
||||
assert "strategyId" not in payload
|
||||
38
tests/test_submit_result_mapping.py
Normal file
38
tests/test_submit_result_mapping.py
Normal file
@@ -0,0 +1,38 @@
|
||||
from app.clients.modelhub import ModelHubClient
|
||||
from app.settings import Settings
|
||||
|
||||
|
||||
class FakeResponse:
|
||||
def __init__(self, body):
|
||||
self.body = body
|
||||
def json(self):
|
||||
return self.body
|
||||
|
||||
|
||||
class FakeSession:
|
||||
def __init__(self, body):
|
||||
self.body = body
|
||||
def post(self, *args, **kwargs):
|
||||
return FakeResponse(self.body)
|
||||
|
||||
|
||||
def make_client(body):
|
||||
settings = Settings(
|
||||
auth_token="a", hf_token="h", contest_api_token="c", email="e", user_id="1",
|
||||
contributors="u", strategy_id="s", submit_dry_run=False
|
||||
)
|
||||
client = ModelHubClient(settings)
|
||||
client.session = FakeSession(body)
|
||||
return client
|
||||
|
||||
|
||||
def test_submit_success_mapping():
|
||||
assert make_client({"code": 0, "message": "ok"}).create_contest_task("m", "hygon_k100-ai", "cfg").result == "success"
|
||||
|
||||
|
||||
def test_submit_queue_full_mapping():
|
||||
assert make_client({"code": 60007, "message": "full"}).create_contest_task("m", "hygon_k100-ai", "cfg").result == "queue_full"
|
||||
|
||||
|
||||
def test_submit_conflict_mapping():
|
||||
assert make_client({"code": 40000, "message": "正在验证中"}).create_contest_task("m", "hygon_k100-ai", "cfg").result == "conflict"
|
||||
16
tests/test_task_granularity.py
Normal file
16
tests/test_task_granularity.py
Normal file
@@ -0,0 +1,16 @@
|
||||
from app.scheduler.budget import Budget
|
||||
from app.scheduler.planner import Planner
|
||||
from app.storage.db import connect
|
||||
from app.storage.repositories import Repository
|
||||
|
||||
|
||||
def test_task_granularity_per_model_gpu(tmp_path):
|
||||
conn = connect(str(tmp_path / "state.db"))
|
||||
repo = Repository(conn)
|
||||
repo.upsert_candidate("owner/model", "HUGGING_FACE", "manual_bounty", 0, True, target_gpu_aliases="k100,p800")
|
||||
created = Planner(repo, Budget(repo, 2000), ["k100"]).materialize_tasks()
|
||||
assert created == 2
|
||||
rows = repo.due_tasks(("pending",), limit=10)
|
||||
assert {row["gpu_type"] for row in rows} == {"hygon_k100-ai", "Kunlunxin_p-800"}
|
||||
assert Planner(repo, Budget(repo, 2000), ["k100"]).materialize_tasks() == 0
|
||||
conn.close()
|
||||
Reference in New Issue
Block a user