Compare commits
20 Commits
aff97f8643
...
main
| Author | SHA1 | Date | |
|---|---|---|---|
| 4a5afb7ed7 | |||
| ac9c80ed8c | |||
| 938f84f1ba | |||
| 686739fea2 | |||
| 63963eba29 | |||
| a6c66995e3 | |||
| 70f52a907b | |||
| 054ed3612d | |||
| 5d0d849c28 | |||
| e3dcfe526f | |||
| 6d3fefc7a6 | |||
| 7b89d60b3b | |||
| c06f61ed03 | |||
| f02059126e | |||
| ce8c2128b4 | |||
| 4dacac24f0 | |||
| 45b6a5f05e | |||
| 64558c8cea | |||
| 2a1168c43c | |||
| c405adc4fc |
5
.gitignore
vendored
@@ -1,3 +1,8 @@
|
||||
verl/
|
||||
skypilot-ssh-test/
|
||||
ray_in_docker/
|
||||
__pycache__/
|
||||
.venv/
|
||||
.pytest_cache/
|
||||
.coverage
|
||||
htmlcov/
|
||||
165
README.md
Normal file
@@ -0,0 +1,165 @@
|
||||
# Argus Cluster
|
||||
|
||||
本仓库主要代码位于 `src/mvp/py/`,单元测试位于 `src/mvp/py/tests/`(`pytest.ini` 已配置 `testpaths`)。
|
||||
|
||||
## 使用 uv 创建/激活虚拟环境并安装依赖
|
||||
|
||||
前置条件:已安装 `uv`(Astral 的 Python 包/虚拟环境工具)。
|
||||
|
||||
1) 在仓库根目录创建虚拟环境:
|
||||
|
||||
```bash
|
||||
uv venv .venv
|
||||
```
|
||||
|
||||
2) 激活虚拟环境:
|
||||
|
||||
```bash
|
||||
source .venv/bin/activate
|
||||
```
|
||||
|
||||
3) 在虚拟环境中安装依赖(运行与测试依赖):
|
||||
|
||||
```bash
|
||||
uv pip install -r src/mvp/py/requirements.txt -r src/mvp/py/requirements-dev.txt
|
||||
```
|
||||
|
||||
## 运行单元测试
|
||||
|
||||
在仓库根目录执行:
|
||||
|
||||
```bash
|
||||
pytest
|
||||
```
|
||||
|
||||
如需显式使用 Python 模块方式:
|
||||
|
||||
```bash
|
||||
python -m pytest
|
||||
```
|
||||
|
||||
## 远端开发(h1):同步代码、构建镜像、初始化共享目录、拉起 Ray 集群
|
||||
|
||||
`src/mvp/docker-compose.yaml` 以 `src/mvp/` 为工作目录,并挂载 `../../shared` 与 `../../verl`,因此推荐远端目录结构如下:
|
||||
|
||||
- `/home2/argus/infra/mvp/src/mvp/`:本仓库的 `src/mvp/` 内容
|
||||
- `/home2/argus/infra/mvp/shared/`:共享目录(模拟/对齐 NFS)
|
||||
- `/home2/argus/infra/mvp/verl/`:训练代码仓库(脚本会检查该目录存在)
|
||||
|
||||
1) 同步代码到远端(只同步 `src/mvp/`,确保相对路径挂载正确):
|
||||
|
||||
```bash
|
||||
ssh argus@h1 "mkdir -p /home2/argus/infra/mvp/src/mvp"
|
||||
rsync -av --delete src/mvp/ argus@h1:/home2/argus/infra/mvp/src/mvp/
|
||||
```
|
||||
|
||||
2) 在远端准备 `verl` 仓库(若已存在可跳过):
|
||||
|
||||
```bash
|
||||
ssh argus@h1 "mkdir -p /home2/argus/infra/mvp && test -d /home2/argus/infra/mvp/verl || echo '缺少 /home2/argus/infra/mvp/verl(请先 git clone)'"
|
||||
|
||||
# 下载verl
|
||||
ssh argus@h1 "cd /home2/argus/infra/mvp && git clone https://github.com/volcengine/verl.git"
|
||||
```
|
||||
|
||||
3) 登录远端并初始化共享目录环境(会创建 `../../shared/...` 等目录):
|
||||
|
||||
```bash
|
||||
ssh argus@h1
|
||||
cd /home2/argus/infra/mvp/src/mvp
|
||||
./scripts/00_prereq_check.sh
|
||||
```
|
||||
|
||||
4) 构建 Ray 节点镜像并拉起集群(首次或镜像不存在时建议强制构建):
|
||||
|
||||
```bash
|
||||
cd /home2/argus/infra/mvp/src/mvp
|
||||
BUILD=1 ./scripts/01_up.sh
|
||||
```
|
||||
|
||||
5) 验证集群状态:
|
||||
|
||||
```bash
|
||||
./scripts/50_status.sh
|
||||
curl -sS http://127.0.0.1:8265/api/version | head
|
||||
docker ps --format 'table {{.Names}}\t{{.Status}}\t{{.Ports}}' | grep -E 'argus-ray-|sftpgo|wandb' || true
|
||||
```
|
||||
|
||||
6) 关闭集群:
|
||||
|
||||
```bash
|
||||
./scripts/02_down.sh
|
||||
```
|
||||
|
||||
## 远端(h1):端口映射、创建 W&B user用户
|
||||
使用以下命令将h1的端口映射到本机
|
||||
```
|
||||
ssh -p 12022 -L 8265:127.0.0.1:8265 -L 8080:127.0.0.1:8080 -L 8081:127.0.0.1:8081 -L 2022:127.0.0.1:2022 -L 8090:127.0.0.1:8090 -L 8000:127.0.0.1:8000 -o ProxyJump=ssh@jump.nasp.fit:36022 nasp@192.168.20.121
|
||||
```
|
||||
其中:
|
||||
- 8265: ray dashboard
|
||||
- 8080: webui & API 服务端口
|
||||
- 8081: sftpgo web client
|
||||
- 2022: sftpgo sftp协议端口
|
||||
- 8090: weight & bias 网站端口
|
||||
- 8000: model serving openai服务端口
|
||||
|
||||
|
||||
## 远端(h1):预先准备数据/模型,启动与关闭 API Server
|
||||
|
||||
下面命令均在远端执行:
|
||||
|
||||
```bash
|
||||
ssh argus@h1
|
||||
cd /home2/argus/infra/mvp/src/mvp
|
||||
```
|
||||
|
||||
1) 预先下载数据集与模型(写入共享目录,幂等可重复执行):
|
||||
|
||||
```bash
|
||||
# 可选:指定要缓存的模型(默认 Qwen/Qwen2.5-0.5B-Instruct)
|
||||
MODEL_ID="Qwen/Qwen2.5-0.5B-Instruct" ./scripts/30_prepare_data_and_model.sh
|
||||
```
|
||||
|
||||
2) 安装 API 依赖(在 head 容器内 best-effort 安装 `fastapi/uvicorn/yaml` 等):
|
||||
|
||||
```bash
|
||||
./scripts/12_install_api_deps.sh
|
||||
```
|
||||
|
||||
3) 启动 API Server(监听 `:8080`,需要设置鉴权 token):
|
||||
|
||||
```bash
|
||||
export MVP_INTERNAL_TOKEN="your-dev-token"
|
||||
./scripts/60_start_api.sh
|
||||
curl -sS http://127.0.0.1:8080/docs | head
|
||||
```
|
||||
|
||||
查看 API 进程状态(基于 pid 文件):
|
||||
|
||||
```bash
|
||||
./scripts/62_status_api.sh
|
||||
```
|
||||
|
||||
4) 关闭 API Server:
|
||||
|
||||
```bash
|
||||
./scripts/61_stop_api.sh
|
||||
```
|
||||
|
||||
## 远端(h1):W&B Local(wandb)License 激活常见问题(端口转发)
|
||||
|
||||
如果你通过 SSH 端口转发访问 W&B UI,但在 “Add license” 页面点击 “Update” 时提示:
|
||||
`Unable to reach the backend api`,通常是因为 W&B 前端使用了“绝对 URL”去请求后端 API,
|
||||
而该 URL 的 host/port 与你当前端口转发的地址不一致(例如 UI 打开在 `http://localhost:8090`,
|
||||
但前端去请求 `http://localhost:8080`)。
|
||||
|
||||
建议二选一:
|
||||
|
||||
- 方式 A(推荐,最简单):把本地 `8080` 转发到远端 `8090`,然后用 `http://localhost:8080` 打开 UI:
|
||||
- `ssh -L 8080:127.0.0.1:8090 argus@h1`
|
||||
- 方式 B:在远端启动容器前设置 `WANDB_HOST`(与浏览器地址一致),并重建 `wandb` 容器:
|
||||
- `export WANDB_HOST=http://localhost:8090`
|
||||
- `docker compose -f /home2/argus/infra/mvp/src/mvp/docker-compose.yaml up -d --force-recreate wandb`
|
||||
|
||||
也可以通过设置 `WANDB_LICENSE=...` 在容器启动时注入 license,绕过 UI 更新步骤。
|
||||
2133
doc/arch.excalidraw
|
Before Width: | Height: | Size: 5.2 MiB After Width: | Height: | Size: 5.2 MiB |
|
Before Width: | Height: | Size: 5.2 MiB After Width: | Height: | Size: 5.2 MiB |
|
Before Width: | Height: | Size: 5.6 MiB After Width: | Height: | Size: 5.6 MiB |
|
Before Width: | Height: | Size: 5.6 MiB After Width: | Height: | Size: 5.6 MiB |
|
Before Width: | Height: | Size: 5.1 MiB After Width: | Height: | Size: 5.1 MiB |
|
Before Width: | Height: | Size: 5.8 MiB After Width: | Height: | Size: 5.8 MiB |
|
Before Width: | Height: | Size: 5.6 MiB After Width: | Height: | Size: 5.6 MiB |
279
doc/hl_design.md
@@ -1,279 +0,0 @@
|
||||
# AI Infra 训练平台建设方案
|
||||
|
||||
## 1. 愿景与目标
|
||||
|
||||
### 1.1 愿景
|
||||
|
||||
构建一套**端到端的智能化AI训练平台**,将分散的训练框架、资源调度、监控运维、数据管理能力整合为统一的标准化流水线,让大模型训练算法团队**专注于模型创新而非基础设施运维**,同时与现有运维智能体深度协同,实现训练任务的智能化运维闭环。
|
||||
|
||||
### 1.2 核心目标
|
||||
|
||||
| 目标维度 | 描述 |
|
||||
|---------|------|
|
||||
| **效率提升** | 训练任务从准备到启动时间缩短 70%,故障恢复时间缩短 50% |
|
||||
| **标准化** | 建立统一的训练流程规范,消除"人人一套环境"的混乱局面 |
|
||||
| **可观测性** | 全链路监控覆盖,训练状态、资源利用、异常事件一目了然 |
|
||||
| **智能运维** | 与运维智能体对接,实现断训自动分析、故障智能诊断 |
|
||||
|
||||
|
||||
---
|
||||
|
||||
## 2. 整体架构概览
|
||||
|
||||
```
|
||||
┌─────────────────────────────────────────────────────────────────────┐
|
||||
│ 用户交互层 │
|
||||
│ ┌─────────────┐ ┌─────────────┐ ┌─────────────┐ │
|
||||
│ │ Web 前端 │ │ CLI 工具 │ │ API 接口 │ │
|
||||
│ │ (任务提交) │ │ (高级用户) │ │ (自动化集成) │ │
|
||||
│ └─────────────┘ └─────────────┘ └─────────────┘ │
|
||||
└─────────────────────────────────────────────────────────────────────┘
|
||||
│
|
||||
┌─────────────────────────────────────────────────────────────────────┐
|
||||
│ 平台服务层 │
|
||||
│ ┌───────────────┐ ┌───────────────┐ ┌───────────────┐ │
|
||||
│ │ 任务调度 │ │ 数据管理 │ │ 模型管理 │ │
|
||||
│ │ (SkyPilot) │ │(schema/dataset)│ │ (版本/产物) │ │
|
||||
│ └───────────────┘ └───────────────┘ └───────────────┘ │
|
||||
│ ┌───────────────┐ ┌───────────────┐ ┌───────────────┐ │
|
||||
│ │ 镜像管理 │ │ 指标追踪 │ │ 日志中心 │ │
|
||||
│ │(Local Registry)│ │ (W&B) │ │ (集中采集) │ │
|
||||
│ └───────────────┘ └───────────────┘ └───────────────┘ │
|
||||
└─────────────────────────────────────────────────────────────────────┘
|
||||
│
|
||||
┌─────────────────────────────────────────────────────────────────────┐
|
||||
│ 智能运维层 │
|
||||
│ ┌─────────────────────────────────────────────────────────────┐ │
|
||||
│ │ 运维智能体对接 │ │
|
||||
│ │ • 断训自动分析 • 故障根因定位 • 资源利用率优化建议 │ │
|
||||
│ └─────────────────────────────────────────────────────────────┘ │
|
||||
└─────────────────────────────────────────────────────────────────────┘
|
||||
│
|
||||
┌─────────────────────────────────────────────────────────────────────┐
|
||||
│ 基础设施层 │
|
||||
│ ┌───────────────┐ ┌───────────────┐ ┌───────────────┐ │
|
||||
│ │ Kubernetes │ │ GPU 集群 │ │ 分布式存储 │ │
|
||||
│ │ (容器编排) │ │ H20/A6000/H100│ │ (JuiceFS) │ │
|
||||
│ └───────────────┘ └───────────────┘ └───────────────┘ │
|
||||
└─────────────────────────────────────────────────────────────────────┘
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 3. 用户故事
|
||||
|
||||
### 3.1 算法工程师视角
|
||||
|
||||
> **作为**一名算法工程师,
|
||||
> **我希望**通过简单的界面配置就能提交一个多节点 RLHF 训练任务,
|
||||
> **以便于**我可以专注于模型和数据本身,而不是花大量时间在环境配置和资源协调上。
|
||||
|
||||
**验收标准:**
|
||||
- [ ] 在 Web 界面上选择数据集、模型、训练配置
|
||||
- [ ] 一键提交后,系统自动完成资源分配、镜像拉取、任务启动
|
||||
- [ ] 实时查看训练进度曲线和关键指标
|
||||
|
||||
---
|
||||
|
||||
> **作为**一名算法工程师,
|
||||
> **我希望**训练中断时能快速定位问题原因,
|
||||
> **以便于**减少排查时间,尽快恢复训练。
|
||||
|
||||
**验收标准:**
|
||||
- [ ] 系统自动检测训练中断事件
|
||||
- [ ] 智能体自动分析中断原因(OOM、网络故障、硬件异常等)
|
||||
- [ ] 提供可操作的恢复建议
|
||||
|
||||
---
|
||||
|
||||
> **作为**一名算法工程师,
|
||||
> **我希望**启动一个 Notebook 环境调试代码.
|
||||
> **以便于**小规模试跑训练,测试训练数据集、调整模型参数
|
||||
|
||||
**验收标准:**
|
||||
- [ ] Notebook容器启动速度
|
||||
- [ ] 开发容器内置依赖包完善度,按照新包
|
||||
|
||||
---
|
||||
|
||||
### 3.2 团队负责人视角
|
||||
|
||||
> **作为**团队负责人,
|
||||
> **我希望**能够看到所有训练任务的整体资源利用情况,
|
||||
> **以便于**合理规划算力资源,识别资源浪费。
|
||||
|
||||
**验收标准:**
|
||||
- [ ] 仪表盘展示各集群 GPU 利用率趋势
|
||||
- [ ] 任务队列可视化,等待/运行/完成状态一目了然
|
||||
- [ ] 资源使用报表按项目/用户统计
|
||||
|
||||
---
|
||||
|
||||
### 3.3 运维工程师视角
|
||||
|
||||
> **作为**运维工程师,
|
||||
> **我希望**训练任务的监控数据能自动接入现有运维系统,
|
||||
> **以便于**统一管理,减少割裂的监控工具。
|
||||
|
||||
**验收标准:**
|
||||
- [ ] 训练任务指标自动推送到运维智能体
|
||||
- [ ] 异常告警自动触发智能体分析流程
|
||||
- [ ] 与现有运维系统数据互通
|
||||
|
||||
---
|
||||
|
||||
## 4. 里程碑规划
|
||||
|
||||
### 里程碑总览
|
||||
|
||||
```
|
||||
M1 M2 M3 M4
|
||||
│ │ │ │
|
||||
────●──────────────────●──────────────────●──────────────────●────────▶
|
||||
│ │ │ │
|
||||
基础设施就绪 训练流水线上线 监控运维闭环 智能化运维
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### M1: 基础设施就绪
|
||||
|
||||
**目标:** 完成底层平台搭建,具备运行训练任务的基础能力
|
||||
|
||||
| 交付物 | 说明 |
|
||||
|-------|------|
|
||||
| K8S 集群 | H20 集群上部署 Kubernetes,支持 GPU 调度 |
|
||||
| 本地 Registry | 内网镜像仓库,解决镜像拉取问题 |
|
||||
| JuiceFS/MinIO | 分布式存储,数据集和模型 checkpoint 持久化 |
|
||||
| 基础镜像 | veRL 训练镜像,预置常用依赖 |
|
||||
|
||||
**关键验证点:**
|
||||
- 能够手动在 K8S 上启动单节点训练任务
|
||||
- 数据从 JuiceFS 正常读写
|
||||
- 镜像从本地 Registry 正常拉取
|
||||
- 引入 Volcano 或 Kueue,并配置 Gang Scheduling 策略,实现All-or-Nothing 的资源分配
|
||||
- 确认JuiceFS 的本地 SSD 缓存策略,在其中一台机器部署MiniIO单节点,另外两台机器上部署JuiceFS client
|
||||
- 网络通信支持 RoCE/InfiniBand
|
||||
- Notebook 交互式开发环境
|
||||
|
||||
---
|
||||
|
||||
### M2: 训练流水线上线
|
||||
|
||||
**目标:** 用户可通过前端提交和管理训练任务
|
||||
|
||||
| 交付物 | 说明 |
|
||||
|-------|------|
|
||||
| SkyPilot 集成 | 任务调度与资源编排 |
|
||||
| W&B 本地服务 | 训练指标追踪与可视化 |
|
||||
| 任务管理前端 | 数据上传、任务提交、进度查看、日志查看 |
|
||||
| 数据管理模块 | 支持从 HuggingFace 链接或 FTP 导入数据集 |
|
||||
|
||||
**关键验证点:**
|
||||
- 端到端完成一次多节点 SFT 训练
|
||||
- 通过前端提交任务,查看 W&B 训练曲线
|
||||
- 训练日志完整保存并可查询
|
||||
- 多租户、项目制配额管理功能
|
||||
|
||||
---
|
||||
|
||||
### M3: 监控运维闭环
|
||||
|
||||
**目标:** 实现任务全生命周期监控,与运维智能体初步对接
|
||||
|
||||
| 交付物 | 说明 |
|
||||
|-------|------|
|
||||
| 资源监控 | GPU 利用率、显存、网络带宽实时采集 |
|
||||
| 日志采集 | 训练日志集中存储,支持检索 |
|
||||
| 智能体对接 | 断训事件自动推送,触发智能体分析 |
|
||||
| 告警机制 | 异常状态(OOM、任务卡死等)自动告警 |
|
||||
|
||||
**关键验证点:**
|
||||
- 训练任务异常时,5 分钟内收到告警
|
||||
- 断训事件自动生成分析报告
|
||||
- Grafana 仪表盘展示集群整体健康状态
|
||||
- sidecar方式部署 DCGM Exporter 来获取细粒度指标,自动采集到Prometheus
|
||||
- 断训时“保留现场”机制,供人工/智能体排查介入
|
||||
|
||||
---
|
||||
|
||||
### M4: 智能化运维
|
||||
|
||||
**目标:** 深度整合运维智能体,实现智能调度与自愈
|
||||
|
||||
| 交付物 | 说明 |
|
||||
|-------|------|
|
||||
| 故障自愈 | 常见故障自动处理(如重新调度到健康节点) |
|
||||
| 智能调度 | 基于历史数据优化任务资源分配 |
|
||||
| 根因分析 | 复杂故障场景的深度分析能力 |
|
||||
| 容量预测 | 基于任务趋势预测算力需求 |
|
||||
|
||||
---
|
||||
|
||||
## 5. 约束与风险
|
||||
|
||||
### 5.1 已知约束
|
||||
|
||||
| 约束项 | 影响 | 应对策略 |
|
||||
|-------|------|---------|
|
||||
| **内网环境** | 无法直接访问HF/Dockerhub/Github资源(模型、数据集) | 本地 Registry + 数据导入工具 |
|
||||
| **算力平台限制** | 现有平台调度能力有限 | 引入 SkyPilot 作为上层调度 |
|
||||
| **数据持久化** | 需要可靠的分布式存储 | JuiceFS + MinIO 方案 |
|
||||
|
||||
### 5.2 潜在风险
|
||||
|
||||
| 风险 | 可能性 | 影响 | 缓解措施 |
|
||||
|-----|-------|------|---------|
|
||||
| K8S 与现有系统集成复杂 | 中 | 高 | 先在 H20 集群小范围验证 |
|
||||
| 智能体接口适配工作量大 | 中 | 中 | 早期明确接口规范,持续对齐 |
|
||||
| 用户习惯迁移阻力 | 低 | 中 | 渐进式推广,保留手动模式 |
|
||||
|
||||
---
|
||||
|
||||
## 6. 资源与依赖
|
||||
|
||||
### 6.1 硬件资源
|
||||
|
||||
| 集群 | 配置 | 用途 |
|
||||
|-----|------|------|
|
||||
| H20 集群 | 2 节点 × 8 卡 = 16 卡 | 主力训练集群,首期部署目标 |
|
||||
| A6000 集群 | 2 节点 × 4 卡 = 8 卡 | 开发测试、小规模实验 |
|
||||
| H100 集群 | 多节点 | 目前仅提供容器方式,不确定能否提供KubeConfig接入,大规模训练 |
|
||||
|
||||
### 6.2 外部依赖
|
||||
|
||||
| 依赖项 | 状态 | 负责方 |
|
||||
|-------|------|-------|
|
||||
| yd运维智能体接口 | 已有基础 | |
|
||||
| argus运维系统 | 已有 | 运维团队 |
|
||||
|
||||
---
|
||||
|
||||
## 7. 成功标准
|
||||
|
||||
### 阶段一完成标准(M1 + M2)
|
||||
|
||||
- [ ] 算法工程师可通过 Web 界面完成 SFT/RLHF 训练全流程
|
||||
- [ ] 任务提交到开始训练时间 < 10 分钟
|
||||
- [ ] 训练指标实时可视化,延迟 < 1 分钟
|
||||
- [ ] 至少完成 3 个实际项目的验证使用
|
||||
|
||||
### 阶段二完成标准(M3 + M4)
|
||||
|
||||
- [ ] 断训事件 100% 自动检测并推送智能体
|
||||
- [ ] 常见故障(OOM、节点失联)自动生成分析报告
|
||||
- [ ] GPU 整体利用率提升 20%(通过更好的调度)
|
||||
- [ ] 平均故障恢复时间(MTTR)缩短 50%
|
||||
|
||||
---
|
||||
|
||||
## 附录:关键技术选型
|
||||
|
||||
| 领域 | 选型 | 选型理由 |
|
||||
|-----|------|---------|
|
||||
| 容器编排 | Kubernetes | 业界标准,生态成熟 |
|
||||
| 任务调度 | SkyPilot | 专为 ML 场景设计,支持多集群 |
|
||||
| 分布式存储 | JuiceFS + MinIO | 兼容 POSIX,适合训练场景 |
|
||||
| 实验追踪 | W&B (自部署) | 功能完善,团队已有使用经验 |
|
||||
| 镜像仓库 | Harbor / Registry | 内网环境必需 |
|
||||
| 训练框架 | veRL / Megatron | 支持 RLHF,与现有工作对齐 |
|
||||
@@ -1,18 +0,0 @@
|
||||
|
||||
目标:设计一套AI任务调度和管理系统
|
||||
环境:在K8S集群(商业化封装的算力平台)上运行若干个GPU容器和CPU容器
|
||||
约束:CPU容器可以暴露端口供desktop访问,CPU容器可以通过ssh访问GPU容器,容器的生命周期在单独的算力平台控制台上管理;所有容器
|
||||
有共享的NFS文件系统
|
||||
目标:在CPU容器上部署一个web server 向desktop提供服务,后端运行一个类似skypilot的程序,能够注册和管理GPU容器实例进来。
|
||||
- 用户在desktop上通过网站提交执行模型训练的任务,任务管理程序会将任务加入调度队列,等资源池有空闲资源的时候,就会调度任务到对应的GPU节点上执行任务。任务有单节点任务,和多节点任务两类,预先的节点也提前配置好了,多节点任务的容器之间是专门
|
||||
走的IB网络。
|
||||
任务:完成技术选型,主要是普通的skypilot不支持现成的ssh节点管理。你调研是否有节点满足需求的任务调度管理方案。
|
||||
补充:运行的任务一般是 verl / skyRL 强化学习训练框架,以及llamafactory 框架等。GPU节点容器假设是已经预先安装好了对应的环境依赖,任务管理器只需要调度和查看任务执行情况,还有监控任务是否健康,有没有任务中断。
|
||||
未来规划:
|
||||
- 统一数据集管理,包括模型库、数据集库
|
||||
- 统一日志采集&规范
|
||||
- 统一容器GPU/CPU指标采集
|
||||
- 统一的 weight & bias 指标服务器
|
||||
结构图:见附件
|
||||
|
||||
|
||||
@@ -1,315 +0,0 @@
|
||||
# 关于MinIO + JuiceFS 在H20部署
|
||||
|
||||
这是一个架构设计的关键决策点。在 **H20 节点运行 K3s** 的前提下,我建议采取 **“存储组件独立部署(Outside),客户端插件化接入(Inside)”** 的混合策略。
|
||||
|
||||
简单来说:**MinIO 和 Redis 跑在 K3s 外面(用 Docker),JuiceFS Client 通过 CSI 跑在 K3s 里面。**
|
||||
|
||||
以下是详细的架构拆解和理由:
|
||||
|
||||
### 1\. 架构图解
|
||||
|
||||
```mermaid
|
||||
graph TD
|
||||
subgraph "物理层 / Host OS (H20 Node)"
|
||||
NVMe[NVMe SSD (物理盘)]
|
||||
end
|
||||
|
||||
subgraph "K3s 集群 (容器层)"
|
||||
Pod[训练任务 Pod]
|
||||
CSI[JuiceFS CSI Driver (DaemonSet)]
|
||||
|
||||
Pod -- "PVC 挂载" --> CSI
|
||||
end
|
||||
|
||||
subgraph "独立 Docker 容器 (非 K3s)"
|
||||
Redis[Redis 容器]
|
||||
MinIO[MinIO 容器]
|
||||
end
|
||||
|
||||
%% 关键数据流
|
||||
CSI -- "1. 读写缓存" --> NVMe
|
||||
CSI -- "2. 网络 IO" --> MinIO
|
||||
CSI -- "3. 元数据 IO" --> Redis
|
||||
|
||||
%% 避免环路
|
||||
MinIO -. "数据持久化" .- NVMe
|
||||
```
|
||||
|
||||
-----
|
||||
|
||||
### 2\. 组件部署位置详解
|
||||
|
||||
#### A. MinIO & Redis:建议 **独立于 K3s 部署 (Outside)**
|
||||
|
||||
**方式**:直接在宿主机(Host OS)上使用 `docker run` 或 `docker-compose` 启动,或者使用 `systemd` 管理。
|
||||
|
||||
**理由**:
|
||||
|
||||
1. **爆炸半径隔离(最重要)**:AI 实验初期,K3s 集群可能会因为 CNI 网络插件配置错误、Device Plugin 崩溃、或者各种误操作导致集群不可用(Crash)。如果存储后端(MinIO)跑在 K3s 里,K3s 一挂,你的数据就“锁”在里面了,连日志都拷不出来。
|
||||
2. **资源竞争保护**:K8s 的调度器虽然有 QoS,但在只有 2 个节点的情况下,H20 的高负载很容易导致控制面不稳。把 MinIO 独立出来,保证存储服务的 CPU/内存资源不受 K3s 内部 Pod 挤压。
|
||||
3. **网络直通**:宿主机直接部署 MinIO,使用的是 `host` 网络,少了一层 K8s Service/Ingress 转发,对于大带宽吞吐稍微有一点点性能优势。
|
||||
|
||||
#### B. JuiceFS Client:建议 **集成在 K3s 内部 (Inside)**
|
||||
|
||||
**方式**:使用 **JuiceFS CSI Driver**。
|
||||
|
||||
**理由**:
|
||||
|
||||
1. **符合“平台化”愿景**:你的目标是做一个 "AI Infra 平台"。对于上层用户(算法工程师),他们不应该关心宿主机挂载了什么目录。他们提交的 YAML 只需要写:
|
||||
```yaml
|
||||
volumes:
|
||||
- name: data
|
||||
persistentVolumeClaim:
|
||||
claimName: pvc-juicefs
|
||||
```
|
||||
这是标准的云原生体验。
|
||||
2. **权限与隔离**:CSI Driver 可以更好地处理容器内的挂载点权限问题。如果手动在宿主机 `mount` 然后用 `hostPath` 映射进容器,经常会遇到 `Permission Denied` 或者 `root` 用户归属问题,非常头疼。
|
||||
3. **自动缓存管理**:CSI Driver 可以通过配置 `StorageClass`,自动帮你处理缓存路径的挂载和清理。
|
||||
|
||||
-----
|
||||
|
||||
### 3\. 实操配置指南 (关键点)
|
||||
|
||||
既然你只有两台 H20 并且要上 K3s,以下是落地的关键配置:
|
||||
|
||||
#### 第一步:在宿主机启动 MinIO 和 Redis (后端)
|
||||
|
||||
不要用 Helm,直接写一个 `docker-compose.yml` 放在 `/opt/ai-infra/` 下运行:
|
||||
|
||||
```yaml
|
||||
version: '3'
|
||||
services:
|
||||
minio:
|
||||
image: minio/minio
|
||||
network_mode: "host" # 直接用宿主机网络,性能最好
|
||||
volumes:
|
||||
- /mnt/data1/minio_data:/data # 挂载到大容量HDD或SATA SSD
|
||||
command: server /data --console-address ":9001"
|
||||
|
||||
redis:
|
||||
image: redis:7
|
||||
network_mode: "host"
|
||||
volumes:
|
||||
- /mnt/data1/redis_data:/data
|
||||
command: redis-server --appendonly yes
|
||||
```
|
||||
|
||||
#### 第二步:在 K3s 中安装 JuiceFS CSI Driver (前端)
|
||||
|
||||
1. **安装 CSI Driver**:
|
||||
|
||||
```bash
|
||||
helm repo add juicefs https://juicedata.github.io/charts/
|
||||
helm install juicefs-csi-driver juicefs/juicefs-csi-driver -n kube-system
|
||||
```
|
||||
|
||||
2. **创建 Secret (连接信息)**:
|
||||
|
||||
```yaml
|
||||
apiVersion: v1
|
||||
kind: Secret
|
||||
metadata:
|
||||
name: juicefs-sc-secret
|
||||
namespace: default
|
||||
type: Opaque
|
||||
stringData:
|
||||
name: "myjfs"
|
||||
metaurl: "redis://<宿主机IP>:6379/1" # 注意这里填宿主机内网IP
|
||||
storage: "minio"
|
||||
bucket: "http://<宿主机IP>:9000/mybucket"
|
||||
access-key: "minioadmin"
|
||||
secret-key: "minioadmin"
|
||||
```
|
||||
|
||||
3. **创建 StorageClass (核心中的核心)**:
|
||||
这里要配置**本地缓存**,让 JuiceFS 能够利用 H20 的 NVMe。
|
||||
|
||||
```yaml
|
||||
apiVersion: storage.k8s.io/v1
|
||||
kind: StorageClass
|
||||
metadata:
|
||||
name: juicefs-sc
|
||||
provisioner: csi.juicedata.com
|
||||
parameters:
|
||||
csi.storage.k8s.io/node-publish-secret-name: juicefs-sc-secret
|
||||
csi.storage.k8s.io/node-publish-secret-namespace: default
|
||||
csi.storage.k8s.io/provisioner-secret-name: juicefs-sc-secret
|
||||
csi.storage.k8s.io/provisioner-secret-namespace: default
|
||||
|
||||
# === 关键性能参数 ===
|
||||
cache-dir: /mnt/nvme0/juicefs-cache # H20 宿主机上的 NVMe 路径
|
||||
cache-size: "2048000" # 2TB 缓存上限
|
||||
buffer-size: "512" # 写缓冲
|
||||
prefetch: "2" # 预读并发度
|
||||
```
|
||||
|
||||
### 4\. 备选方案:如果 CSI 搞不定怎么办?
|
||||
|
||||
如果 M1 阶段你发现 CSI Driver 调试太麻烦(有时候 K3s 的路径和标准 K8s 不一样,导致 CSI 挂载失败),可以退回到 **“Host Mount + HostPath”** 模式:
|
||||
|
||||
1. **Outside**: 在宿主机直接运行 `juicefs mount ... /mnt/ai-data`。
|
||||
2. **Inside**: 在 Pod YAML 里直接映射:
|
||||
```yaml
|
||||
volumes:
|
||||
- name: data
|
||||
hostPath:
|
||||
path: /mnt/ai-data
|
||||
type: Directory
|
||||
```
|
||||
|
||||
**结论:**
|
||||
优先尝试 **CSI Driver** 方案,这才是正规军的做法,能为你的“平台化”愿景打下基础。MinIO/Redis 请务必保持独立,不要让 K3s 的波动影响到底层数据安全。
|
||||
|
||||
|
||||
----------------
|
||||
|
||||
# 关于Minio扩容
|
||||
这是一个非常关键的运维问题。
|
||||
|
||||
首先要泼一盆冷水(也是为了避坑):**MinIO 不支持从“单机模式(Standalone)”无缝平滑扩容到“分布式模式(Distributed)”**。
|
||||
|
||||
* **单机模式**下,MinIO 是直接把文件存成普通文件。
|
||||
* **分布式模式**下,MinIO 会启用**纠删码(Erasure Coding)**,把文件切片打散存储。
|
||||
|
||||
底层数据结构完全不同,因此你**不能**简单地把新节点加入旧集群。你必须通过\*\*“建立新集群 -\> 数据迁移 -\> 切换流量”\*\*的方式来完成。
|
||||
|
||||
鉴于你使用了 JuiceFS,这个迁移过程其实比纯 MinIO 迁移要灵活得多。以下是针对你 H20/K3s 环境的最佳实战方案:
|
||||
|
||||
-----
|
||||
|
||||
### 方案概览:蓝绿部署迁移法
|
||||
|
||||
我们不修改旧的 MinIO,而是搭建一个新的 3 节点集群,同步数据,然后改一下 K8s Secret 指向新集群。
|
||||
|
||||
#### 前置条件
|
||||
|
||||
假设你现有的环境是:
|
||||
|
||||
* **Old MinIO**: 跑在 Node A (192.168.1.10)
|
||||
* **New Target**: 准备在 Node A, Node B, Node C 上跑 3 节点分布式 MinIO。
|
||||
|
||||
-----
|
||||
|
||||
### 第一步:搭建全新的 3 节点 MinIO 集群
|
||||
|
||||
由于旧的 MinIO 还在通过 `host` 网络运行(占用 9000 端口),新集群如果不关掉旧的,需要用**不同端口**(比如 9100)或者部署在不同机器上。
|
||||
|
||||
假设你新增了机器,或者错开端口。以下是 3 节点分布式 MinIO 的 `docker-compose.yml` 示例(需要在三台机器上都运行):
|
||||
|
||||
```yaml
|
||||
version: '3'
|
||||
services:
|
||||
minio-distributed:
|
||||
image: minio/minio
|
||||
network_mode: "host"
|
||||
hostname: "node1" # 另外两台改为 node2, node3
|
||||
# 关键:分布式启动命令,必须列出所有节点
|
||||
command: server http://192.168.1.10:9100/data http://192.168.1.11:9100/data http://192.168.1.12:9100/data --console-address ":9101"
|
||||
volumes:
|
||||
- /mnt/data_new:/data # 挂载新的数据盘(或者旧盘的新目录)
|
||||
environment:
|
||||
MINIO_ROOT_USER: "admin"
|
||||
MINIO_ROOT_PASSWORD: "strongpassword"
|
||||
```
|
||||
|
||||
*注意:3 节点 MinIO 允许挂掉 1 台机器而不丢失数据。*
|
||||
|
||||
-----
|
||||
|
||||
### 第二步:数据迁移 (两种路径)
|
||||
|
||||
鉴于你用的是 JuiceFS,这里有**两条路**可选:
|
||||
|
||||
#### 路径 A:底座迁移(推荐,速度快,原汁原味)
|
||||
|
||||
直接搬运 MinIO 里的对象块(Block)。因为 JuiceFS 把数据切成了固定的 Block 存在 MinIO 里,我们只需要把这些 Block 从旧 MinIO 搬到新 MinIO,**不需要经过 JuiceFS 客户端**。
|
||||
|
||||
1. **安装 `mc` (MinIO Client)** 命令行工具。
|
||||
2. **配置别名**:
|
||||
```bash
|
||||
mc alias set oldm http://192.168.1.10:9000 minioadmin minioadmin
|
||||
mc alias set newm http://192.168.1.10:9100 admin strongpassword
|
||||
```
|
||||
3. **全量镜像 (Mirror)**:
|
||||
```bash
|
||||
# 创建新桶
|
||||
mc mb newm/mybucket
|
||||
|
||||
# 开始同步数据 (将旧桶数据镜像到新桶)
|
||||
# --watch 参数可以持续监听增量数据,适合不停机迁移
|
||||
mc mirror --watch oldm/mybucket newm/mybucket
|
||||
```
|
||||
|
||||
#### 路径 B:JuiceFS 层面迁移(适合要换云厂商/存储类型)
|
||||
|
||||
如果你想顺便整理数据碎片,或者从 MinIO 迁移到 阿里云 OSS,可以用这个。但在你的场景下,路径 A 更快。
|
||||
|
||||
* 命令:`juicefs sync minio://... minio://...` (不推荐,因为需要解密再加密,消耗 CPU)。
|
||||
|
||||
-----
|
||||
|
||||
### 第三步:停机切换 (Cutover)
|
||||
|
||||
为了保证 100% 数据一致性,建议申请一个短时间的维护窗口(10-20分钟)。
|
||||
|
||||
1. **停止训练任务**:Scale down 所有的 Training Job。
|
||||
2. **停止旧 MinIO 写入**:
|
||||
* 确保 `mc mirror` 已经追平了数据(没有 pending)。
|
||||
* 你可以把旧 MinIO 设为只读,或者直接停止旧容器。
|
||||
3. **最后一次 Sync**:
|
||||
* 运行 `mc mirror --overwrite` 确保最后一点差异也同步了。
|
||||
4. **修改 K8s Secret**:
|
||||
* 这是 JuiceFS 的魔力所在。你**不需要**重新格式化 JuiceFS,也**不需要**修改元数据引擎(Redis)。你只需要告诉 JuiceFS:“底下的对象存储换地方了”。
|
||||
* 编辑 K8s 里的 Secret:
|
||||
<!-- end list -->
|
||||
```bash
|
||||
kubectl edit secret juicefs-sc-secret
|
||||
```
|
||||
* 修改 `bucket` 和 `access-key/secret-key`:
|
||||
<!-- end list -->
|
||||
```yaml
|
||||
stringData:
|
||||
# Redis 配置保持不变!元数据还在那里
|
||||
metaurl: "redis://192.168.1.10:6379/1"
|
||||
|
||||
# 修改 bucket 地址为新集群地址
|
||||
# 注意:MinIO 分布式集群通常前面会加一个 Nginx 做负载均衡
|
||||
# 或者你可以填任意一个节点的 IP (MinIO 会自动转发)
|
||||
bucket: "http://192.168.1.10:9100/mybucket"
|
||||
access-key: "admin"
|
||||
secret-key: "strongpassword"
|
||||
```
|
||||
5. **重启 CSI Driver**:
|
||||
* 删除 JuiceFS CSI 的 Pod,让它们重建并加载新的 Secret。
|
||||
* 删除使用该 PVC 的 Pod。
|
||||
|
||||
-----
|
||||
|
||||
### 第四步:验证与回收
|
||||
|
||||
1. **验证**:
|
||||
* 启动一个测试 Pod,挂载 PVC。
|
||||
* `ls` 一下目录,确认文件都在。
|
||||
* `cat` 一个文件,确认能读(JuiceFS 会去新 MinIO 拿数据块)。
|
||||
2. **回收**:
|
||||
* 确认一切正常运行几天后。
|
||||
* 下线旧的单节点 MinIO 容器。
|
||||
* 释放旧的磁盘空间。
|
||||
|
||||
### 专家建议 (Pro Tips)
|
||||
|
||||
1. **负载均衡 (Load Balancer)**:
|
||||
|
||||
* 到了 3 节点分布式阶段,建议在 3 台 MinIO 前面架设一个 **Nginx** 或者 **HAProxy**。
|
||||
* 这样 JuiceFS Client 连接的是 `http://nginx:9000`,Nginx 轮询转发给后端 3 台 MinIO。
|
||||
* **好处**:如果某台 MinIO 挂了,Nginx 自动剔除,训练任务完全无感。如果直连某台 IP,那台挂了训练就断了。
|
||||
|
||||
2. **元数据备份**:
|
||||
|
||||
* 在做这一切之前,**务必备份 Redis 的 dump.rdb**。
|
||||
* 只要 Redis 在,JuiceFS 的文件结构就在。对象存储里的数据块丢了还能通过文件名找回来(虽然内容坏了),但 Redis 丢了,对象存储里那一堆 `chunk-xxx` 的文件就是一堆毫无意义的二进制垃圾,神仙难救。
|
||||
|
||||
3. **拓扑限制**:
|
||||
|
||||
* MinIO 扩容通常是“倍增”或者“对等扩容”。比如 4 节点扩容,通常是再加 4 节点(变成 2 个 Server Pool)。
|
||||
* 所以,规划 3 节点时,最好磁盘大小一致,网络环境一致。
|
||||
11966
doc/sw_arch.excalidraw
Normal file
3
pytest.ini
Normal file
@@ -0,0 +1,3 @@
|
||||
[pytest]
|
||||
testpaths = src/mvp/py/tests
|
||||
addopts = --maxfail=1 --cov=argus --cov=server --cov-report=term-missing --cov-fail-under=90
|
||||
@@ -1,120 +0,0 @@
|
||||
# MVP 计划(V1)
|
||||
|
||||
本文档目标:把当前“口述/实验记录”整理成**可复现、可验收**的 MVP 计划,并明确下一步最小闭环。
|
||||
|
||||
## 1. 背景与目标
|
||||
|
||||
我们要验证的最小闭环是:
|
||||
|
||||
1) 在“Head(CPU 容器)+ Worker(GPU 容器)”的 Ray Cluster 上,能够跑通一次 `verl` 的 PPO 训练。
|
||||
2) 训练所需的 **数据集 / 模型缓存 / 训练产物(checkpoint)/ 日志** 不落在容器临时文件系统里,而是落在**共享存储(NFS)**,容器重启后可继续使用。
|
||||
3) 所有步骤能写成一套**清晰命令/脚本**,新人可照着复现。
|
||||
|
||||
## 2. 环境与假设
|
||||
|
||||
- 机器:H20 机器(具体规格由算力平台提供)
|
||||
- 访问方式:通过 `ssh h1a` 远程登录(进入算力平台/宿主访问入口)
|
||||
- 容器:算力平台可申请 CPU 容器(对外暴露端口)与若干 GPU 容器(可 SSH 互通)
|
||||
- 共享存储:所有容器可挂载同一套 NFS(在 `specs/hl_design_v2.md` 中假设为 `/mnt/shared`)
|
||||
|
||||
## 3. 已验证现状(现有实验)
|
||||
|
||||
目录 `ray_in_docker/` 已经做过一次可运行的实验(偏“本地/示例级别”):
|
||||
|
||||
- 用 `docker-compose` 起了 2 个 `verl` 镜像容器:
|
||||
- `verl-head`:作为 Ray Head(Dashboard 端口 `8265`)
|
||||
- `verl-worker`:作为 Ray Worker
|
||||
- 在容器中执行:
|
||||
- 下载 GSM8K 数据集(`examples/data_preprocess/gsm8k.py`)
|
||||
- 拉取 HuggingFace 模型(示例:`Qwen/Qwen2.5-0.5B-Instruct`)
|
||||
- `ray start --head` + `ray start --address=...`
|
||||
- 通过 `ray job submit ... python -m verl.trainer.main_ppo ...` 提交 PPO 训练任务(见 `ray_in_docker/ray_example/ppo_train.sh`)
|
||||
|
||||
结论:**训练脚本可以跑通**。
|
||||
|
||||
## 4. 当前主要问题(从实验到平台化 MVP 的差距)
|
||||
|
||||
1) **数据 / 模型 / 输出落在容器内**:容器重启/替换后不可复用;也不利于多人共享与审计。
|
||||
2) **缓存路径不规范**:HuggingFace cache、Ray 临时目录、Hydra 输出目录等可能分散在容器默认路径。
|
||||
3) **可复现不足**:缺少明确的目录规范、统一的启动/提交流程、验收口径。
|
||||
4) Ray 节点**打标签/亲和性调度**的方法未固化:需要明确是否统一用 `ray start --resources`,以及命名规范如何设计。
|
||||
|
||||
## 5. MVP V1:最小闭环定义
|
||||
|
||||
以 `specs/hl_design_v2.md` 的方向为准,但 V1 **只做最小可运行原型**,暂不做完整 Web/调度系统。
|
||||
|
||||
### 5.1 目录规范(统一落到 NFS)
|
||||
|
||||
约定所有容器统一挂载 NFS 到 `/mnt/shared`,并在其中固定目录结构:
|
||||
|
||||
- `/mnt/shared/code/`:代码(可选:按版本/分支隔离)
|
||||
- `/mnt/shared/datasets/`:数据集(如 `gsm8k/`)
|
||||
- `/mnt/shared/hf/`:HuggingFace 缓存(设置 `HF_HOME=/mnt/shared/hf`)
|
||||
- `/mnt/shared/ray/`:Ray 运行期临时目录(可选:设置 `RAY_TMPDIR=/mnt/shared/ray/tmp`)
|
||||
- `/mnt/shared/outputs/`:训练输出根目录(Hydra/日志/ckpt 统一落这里)
|
||||
- `/mnt/shared/outputs/logs/<job_id>/`
|
||||
- `/mnt/shared/outputs/checkpoints/<job_id>/`
|
||||
|
||||
### 5.2 最小集群形态
|
||||
|
||||
- 1 个 Head(CPU 容器)
|
||||
- 跑 `ray start --head --dashboard-host=0.0.0.0`
|
||||
- 暴露 `8265` 给 Desktop/用户查看 Job 状态
|
||||
- 1~2 个 Worker(GPU 容器)
|
||||
- 跑 `ray start --address=<head_ip>:6379` 加入集群
|
||||
- (可选)通过 `--resources='{\"gpu_pool_a\": 1}'` 给节点打标签
|
||||
|
||||
### 5.3 最小训练任务
|
||||
|
||||
- 目标任务:跑通一次 `verl.trainer.main_ppo`(以 GSM8K 为例)
|
||||
- 要求:
|
||||
- `data.train_files` / `data.val_files` 指向 `/mnt/shared/datasets/...`
|
||||
- HuggingFace 模型下载缓存落到 `/mnt/shared/hf`
|
||||
- 训练输出(Hydra outputs、checkpoint、stdout/stderr)落到 `/mnt/shared/outputs/...`
|
||||
|
||||
建议在提交命令里显式覆盖 Hydra 输出目录(示例,具体目录名按需调整):
|
||||
|
||||
- `hydra.run.dir=/mnt/shared/outputs/logs/${JOB_TAG}`
|
||||
|
||||
## 6. 实施步骤(Checklist)
|
||||
|
||||
### 6.1 一次性准备
|
||||
|
||||
- [ ] 确认所有容器已挂载 NFS 到同一路径:`/mnt/shared`
|
||||
- [ ] 在 `/mnt/shared/` 下创建目录:`datasets/ hf/ outputs/ ray/`
|
||||
- [ ] 在所有容器中设置/注入环境变量(推荐写入统一脚本):
|
||||
- `HF_HOME=/mnt/shared/hf`
|
||||
- `HF_ENDPOINT=https://hf-mirror.com`(如需)
|
||||
- `RAY_TMPDIR=/mnt/shared/ray/tmp`(可选)
|
||||
|
||||
### 6.2 启动集群(Head + Worker)
|
||||
|
||||
- [ ] 在 Head 容器启动 Ray Head(并记录 `head_ip:6379`)
|
||||
- [ ] 在每个 Worker 容器执行 `ray start --address=...` 加入集群
|
||||
- [ ] 在 Head 上通过 `ray status` / Dashboard 验证节点已注册
|
||||
|
||||
### 6.3 准备数据与模型
|
||||
|
||||
- [ ] 数据集下载到:`/mnt/shared/datasets/gsm8k/`
|
||||
- [ ] 模型缓存落到:`/mnt/shared/hf/`(拉取一次即可多任务复用)
|
||||
|
||||
### 6.4 提交训练任务
|
||||
|
||||
- [ ] 用 `ray job submit --address=http://<head>:8265 ...` 提交 PPO 训练
|
||||
- [ ] 训练日志与 checkpoint 在 `/mnt/shared/outputs/` 可见
|
||||
|
||||
## 7. 验收标准(V1)
|
||||
|
||||
- [ ] Ray Head/Worker 能稳定加入同一集群(Dashboard 可见)
|
||||
- [ ] PPO 训练任务可提交并跑通(至少完成若干 step/epoch)
|
||||
- [ ] 数据集、HF 缓存、训练输出均在 `/mnt/shared/` 下可复用(容器重启后仍在)
|
||||
- [ ] 有一份“从零到跑通”的命令清单(或脚本)可复现
|
||||
|
||||
## 8. 未决问题(记录待补齐)
|
||||
|
||||
- [ ] Ray 节点标签/亲和性调度:是否统一用 `ray start --resources`,以及命名规范如何设计
|
||||
- [ ] RL workload 的 Driver 放置策略:先按 `verl` 默认即可,后续再按 `specs/hl_design_v2.md` 收敛到“Driver-on-Head / Placement Group”等模式
|
||||
|
||||
## 9. 下一步(进入 V2)
|
||||
|
||||
当 V1 达到“可复现 + 产物可落盘”的验收标准后,下一阶段工作见:`specs/mvp_plan_v2.md`。
|
||||
@@ -1,111 +0,0 @@
|
||||
# MVP V1 远程实验行动文档(待确认后执行)
|
||||
|
||||
## 1. 任务复述(我理解的需求)
|
||||
|
||||
你希望我在远程机器 `argus@h1` 上,进入目录 `/home2/argus/infra/mvp`,把 MVP V1 的“原本流程”**手动完整跑一遍并验证**。要求:
|
||||
|
||||
1) 在宿主机上编写脚本,脚本通过 `docker exec` 在容器内执行命令,负责协调启动顺序(先 head、后 worker)。
|
||||
2) 集群拓扑改为:
|
||||
- 1 个 Ray Head:**没有 GPU**,并且 Head 的 Ray 资源 `CPU=0`(防止 Ray 把训练任务调度到 head)。
|
||||
- 2 个 Ray Worker:各自 **4 GPU**(总 8 GPU)。
|
||||
3) PPO 训练需要“轻量化”,把 `total_epochs` 改为 `1`。
|
||||
4) 先在本地仓库 `src/mvp/v1/` 写好脚本与 compose 文件;再拷贝到远程目录执行与验证。
|
||||
5) 在你确认这份行动文档没问题之前,我**不执行**远程操作。
|
||||
|
||||
## 2. 本地已准备的文件(在本仓库内)
|
||||
|
||||
- `src/mvp/v1/docker-compose.yaml`:3 容器(head + 2 worker),head 不使用 nvidia runtime;worker0/1 各限制 4 GPU。
|
||||
- `src/mvp/v1/scripts/`:宿主机脚本(内部全部用 `docker exec`)
|
||||
- `01_up.sh`:起容器
|
||||
- `20_start_head.sh`:启动 Ray head(`--num-cpus=0 --num-gpus=0`)
|
||||
- `21_start_workers.sh`:启动 Ray worker 加入集群
|
||||
- `30_prepare_data_and_model.sh`:准备 GSM8K 数据与预下载模型
|
||||
- `40_submit_ppo_epoch1.sh`:提交 PPO(`trainer.total_epochs=1`,并设置 `nnodes=2, n_gpus_per_node=4`)
|
||||
- `run_all.sh`:按顺序一键执行
|
||||
|
||||
## 3. 远程环境前置条件(需要你确认/保证)
|
||||
|
||||
在 `argus@h1` 上:
|
||||
|
||||
- Docker 可用,且有 `docker compose` 插件(Compose v2)。
|
||||
- NVIDIA runtime 可用(worker 容器需要 `runtime: nvidia`),宿主机有至少 8 张 GPU。
|
||||
- 不强制要求提前准备 `./verl`:脚本会在宿主机侧检查 `${PWD}/verl`,如果不存在会自动执行:
|
||||
- `git clone https://github.com/volcengine/verl.git`
|
||||
|
||||
此外本实验默认写入持久化目录:`/home2/argus/infra/mvp/shared`(会自动创建)。
|
||||
|
||||
## 4. 拷贝到远程(我执行前会再次征求你确认)
|
||||
|
||||
从本地(本机)同步到远程:
|
||||
|
||||
1) 同步脚本与 compose:
|
||||
- `rsync -av ./src/mvp/v1/ argus@h1:/home2/argus/infra/mvp/src/mvp/v1/`
|
||||
- `rsync -av ./specs/mvp/v1_action.md argus@h1:/home2/argus/infra/mvp/specs/mvp/v1_action.md`
|
||||
2) `verl/` 默认不需要同步(远程会 clone)。如果你更希望固定版本/避免网络波动,也可以手动同步:
|
||||
- `rsync -av --delete ./verl/ argus@h1:/home2/argus/infra/mvp/verl/`
|
||||
|
||||
## 5. 远程执行步骤(在宿主机上)
|
||||
|
||||
在远程机器执行:
|
||||
|
||||
1) 进入目录:
|
||||
- `cd /home2/argus/infra/mvp`
|
||||
2) 确保脚本可执行(首次同步后需要做一次):
|
||||
- `chmod +x ./src/mvp/v1/scripts/*.sh`
|
||||
3) 启动容器:
|
||||
- `./src/mvp/v1/scripts/01_up.sh`
|
||||
4) 安装 editable 版 `verl`(保证 `python -m verl...` 可用):
|
||||
- `./src/mvp/v1/scripts/10_install_verl_editable.sh`
|
||||
5) 启动 Ray Head(禁止调度到 head):
|
||||
- `./src/mvp/v1/scripts/20_start_head.sh`
|
||||
6) 启动两个 Ray Worker 加入集群:
|
||||
- `./src/mvp/v1/scripts/21_start_workers.sh`
|
||||
7) 准备数据 + 预下载模型(落到 `./shared`):
|
||||
- `./src/mvp/v1/scripts/30_prepare_data_and_model.sh`
|
||||
8) 提交 PPO(`total_epochs=1`,必须用 `ray job submit` 在 head 提交;通过 `--entrypoint-resources` 强制 driver 调度到 worker):
|
||||
- `./src/mvp/v1/scripts/40_submit_ppo_epoch1.sh`
|
||||
9) 观察状态:
|
||||
- `./src/mvp/v1/scripts/50_status.sh`
|
||||
- 打开 Ray Dashboard:`http://<h1宿主机IP>:8265`
|
||||
|
||||
也可以一键跑:
|
||||
- `./src/mvp/v1/scripts/run_all.sh`
|
||||
|
||||
## 6. 验收与验证点(执行时我会逐项检查)
|
||||
|
||||
1) Head 节点无 GPU:在 head 容器内 `nvidia-smi` 应不可用或无设备(worker 内可见 4 张)。
|
||||
2) Head 的 Ray 逻辑资源为 `CPU=0, GPU=0`:head 不应承载训练任务调度资源(通过 `ray start --num-cpus=0 --num-gpus=0`)。
|
||||
3) 集群节点数量正确:`ray status` 中应看到 1 head + 2 worker。
|
||||
4) PPO driver 不在 head:`ray job list` 里该 `submission_id` 的 `driver_info.node_ip_address` 应该是 worker 的 IP(`172.19.0.3/172.19.0.4`),不能是 head(`172.19.0.2`)。
|
||||
5) PPO 训练只跑 1 个 epoch:提交参数包含 `trainer.total_epochs=1`。
|
||||
6) checkpoint 落盘:`/mnt/shared/jobs/<job_id>/checkpoints/` 有产物(脚本通过 `trainer.default_local_dir` 强制指向该目录;不设置 `trainer.default_hdfs_dir`)。
|
||||
7) 数据与缓存落盘:`/home2/argus/infra/mvp/shared/` 下出现 datasets/hf/jobs 等目录。
|
||||
|
||||
补充(磁盘保护):
|
||||
- checkpoint 不要每步保存(会非常占空间);当前脚本默认 `trainer.save_freq=10`(每 10 step 保存一次)。
|
||||
|
||||
## 10. 目录命名约定(submission id)
|
||||
|
||||
- 脚本默认会显式指定 `ray job submit --submission-id=$SUBMISSION_ID`,并使用同一个值作为输出目录名:
|
||||
- 输出目录:`/mnt/shared/jobs/$SUBMISSION_ID/`
|
||||
- 你可以在提交时自定义 ID(推荐这样便于检索):
|
||||
- `SUBMISSION_ID=my_run_20251219_001 ./src/mvp/v1/scripts/40_submit_ppo_epoch1.sh`
|
||||
|
||||
## 7. 风险点与兜底
|
||||
|
||||
- 如果 `runtime: nvidia` 在该环境不生效:需要改成 compose 的 `gpus:` 写法(我会按远程 docker 版本调整)。
|
||||
- 如果 Ray Jobs 的 driver 必须在 head 启动(Ray 机制如此):这不影响“训练任务不调度到 head”,但 head 仍会有一个 job driver 进程。
|
||||
- 如果 `verl` 在镜像内已安装但版本不匹配:脚本会优先 `pip install -e /workspace/verl` 以保证行为一致。
|
||||
|
||||
## 8. 你需要确认的 3 个问题(你已确认,我按此执行)
|
||||
|
||||
1) `verl/`:脚本会在远程自动 `git clone https://github.com/volcengine/verl.git`(如你希望固定版本,可改成同步或 checkout tag/commit)。
|
||||
2) GPU:`0-7` 可用(worker0 用 `0-3`,worker1 用 `4-7`)。
|
||||
3) PPO:用满 8 GPU(`nnodes=2, n_gpus_per_node=4`)。
|
||||
|
||||
## 9. 你新增的关键要求(我已纳入脚本)
|
||||
|
||||
- 数据与模型必须落在 `/mnt/shared`(由宿主机 `./shared` bind mount 提供),并且具备**幂等**:
|
||||
- 数据:如果 `train.parquet/test.parquet` 已存在则跳过下载。
|
||||
- 模型:优先检测本地 cache(`HF_HOME=/mnt/shared/hf`);存在则跳过,否则才下载。
|
||||
- 提交 job 时显式注入 `HF_HOME/HUGGINGFACE_HUB_CACHE/TRANSFORMERS_CACHE`,确保训练使用持久化缓存与数据路径。
|
||||
32
src/mvp/README.md
Normal file
@@ -0,0 +1,32 @@
|
||||
# MVP(模块化结构)
|
||||
|
||||
本目录是 MVP 的“按功能模块”组织后的最终入口(替代原 `v1.1/`、`v2.0/` 版本目录)。
|
||||
|
||||
目录:
|
||||
- `py/`:Python 代码(`argus.*` 包)
|
||||
- `argus.ray`:Ray Jobs SDK submitter core(提交/查询/停止/日志)
|
||||
- `argus.service`:API + scheduler + SQLite(服务层队列/重试/状态聚合)
|
||||
- `argus.cli`:CLI 工具(用于脚本化提交/调试)
|
||||
- `configs/`:配置(`ray:` 与 `service:` 两段,纯 YAML 风格)
|
||||
- `taskspecs/`:TaskSpec(训练任务语义参数)
|
||||
- `scripts/`:宿主机脚本(docker compose/exec 编排)
|
||||
- `docker-compose.yaml`:dev 环境容器定义(head+workers)
|
||||
|
||||
快速开始:
|
||||
- CLI 提交流程:`scripts/run_all_cli.sh`
|
||||
- API 提交流程:`scripts/run_all_api.sh`
|
||||
- v2.5(Stateless worker + user 隔离 jobs)E2E:`scripts/run_all_v25_api.sh`
|
||||
- v3.0(WebUI + SFTPGo + user datasets/models + jobs retention)E2E:`scripts/run_all_v30_api.sh`
|
||||
|
||||
v3.0 访问入口(dev/h1):
|
||||
- WebUI:`http://127.0.0.1:8080/ui`
|
||||
- Ray Dashboard:`http://127.0.0.1:8265`
|
||||
- SFTPGo:
|
||||
- SFTP:`127.0.0.1:2022`
|
||||
- Admin API/UI:`http://127.0.0.1:8081`(容器内 8080,host 映射到 8081 避免与 API server 冲突)
|
||||
|
||||
v3.8(Ray Serve LLM / vLLM 模型服务):
|
||||
- 推理端口:`8000`(Ray Serve HTTP)
|
||||
- OpenAI-compatible endpoint:`http://127.0.0.1:8000/v1`
|
||||
- 注意:v3.8 推理接口**不做鉴权**
|
||||
- E2E 脚本:`scripts/run_all_v38_serving.sh`
|
||||
83
src/mvp/configs/dev.yaml
Normal file
@@ -0,0 +1,83 @@
|
||||
ray:
|
||||
# Ray Job server address (head 容器内视角)
|
||||
address: "http://127.0.0.1:8265"
|
||||
|
||||
# 共享根路径(容器内统一 /private,对齐生产)
|
||||
shared_root: "/private"
|
||||
|
||||
# 强制 driver 落 worker(head 不跑训练)
|
||||
entrypoint_num_cpus: 1
|
||||
entrypoint_resources:
|
||||
worker_node: 1
|
||||
|
||||
# 所有 job 通用 runtime env
|
||||
runtime_env:
|
||||
env_vars:
|
||||
HF_ENDPOINT: "https://hf-mirror.com"
|
||||
PYTHONUNBUFFERED: "1"
|
||||
# v3.7: forbid HuggingFace Hub network access from Ray jobs (use cached snapshots).
|
||||
HF_HUB_OFFLINE: "1"
|
||||
# Unify cache dirs so `from_pretrained("org/name")` resolves from the same on-disk cache in offline mode.
|
||||
HF_HOME: "/private/hf"
|
||||
HUGGINGFACE_HUB_CACHE: "/private/hf/hub"
|
||||
TRANSFORMERS_CACHE: "/private/hf/hub"
|
||||
|
||||
# 用户自定义代码目录(可被 PYTHONPATH 注入)
|
||||
user_code_path: "/private/user/code"
|
||||
|
||||
service:
|
||||
api:
|
||||
host: "0.0.0.0"
|
||||
port: 8080
|
||||
auth:
|
||||
token_env: "MVP_INTERNAL_TOKEN"
|
||||
sqlite:
|
||||
db_path: "/private/common/db/mvp.sqlite3"
|
||||
scheduler:
|
||||
tick_s: 5
|
||||
retry_interval_s: 60
|
||||
max_running_tasks: 1
|
||||
|
||||
tracking:
|
||||
wandb:
|
||||
# v3.6: enable Weights & Biases integration (best-effort).
|
||||
# - If WANDB_API_KEY is missing, the platform degrades to console logging and records a warning.
|
||||
enabled: true
|
||||
# For dev compose, use docker bridge gateway + published port for stability.
|
||||
base_url: "http://172.22.0.1:8090"
|
||||
api_key_env: "WANDB_API_KEY"
|
||||
project_suffix: "_project"
|
||||
|
||||
# v3.0: user data management (filesystem + SFTPGo)
|
||||
data:
|
||||
# All user writable data is placed under this root:
|
||||
# /private/users/<user_id>/{datasets,models,code,jobs,trash/jobs}
|
||||
user_root: "/private/users"
|
||||
|
||||
# SFTPGo is optional in dev; when enabled, admin endpoints will call SFTPGo admin API.
|
||||
# Admin password is provided by env var `data.sftpgo.admin_password_env`.
|
||||
sftpgo:
|
||||
enabled: true
|
||||
host: "" # shown to users via GET /api/v2/me
|
||||
sftp_port: 2022
|
||||
# NOTE: the Ray head container image sometimes fails to resolve docker-internal DNS names
|
||||
# (e.g. sftpgo/argus-sftpgo). Use the docker bridge gateway + published port for stability.
|
||||
# - host port 8081 -> sftpgo container 8080
|
||||
# - 172.22.0.1 is the gateway of `mvp_argus-ray-net` in the dev compose
|
||||
admin_api_base: "http://172.22.0.1:8081/api/v2" # head 容器内访问 SFTPGo admin API
|
||||
admin_user: "admin"
|
||||
admin_password_env: "SFTPGO_ADMIN_PASSWORD"
|
||||
|
||||
# Jobs retention policy (v3.0 janitor): move to trash after 3d, purge after 7d.
|
||||
retention:
|
||||
jobs_trash_after_days: 3
|
||||
jobs_purge_after_days: 7
|
||||
janitor_interval_s: 3600
|
||||
|
||||
# v3.8: model serving via Ray Serve LLM (vLLM backend)
|
||||
serving:
|
||||
serve:
|
||||
http_port: 8000
|
||||
proxy_location: HeadOnly
|
||||
llm:
|
||||
accelerator_type: H20
|
||||
182
src/mvp/docker-compose.yaml
Normal file
@@ -0,0 +1,182 @@
|
||||
services:
|
||||
ray_head:
|
||||
image: argus/argus-ray-node:vllm011.latest
|
||||
build:
|
||||
context: .
|
||||
dockerfile: images/argus-ray-node/Dockerfile
|
||||
args:
|
||||
BASE_IMAGE: verlai/verl:vllm011.latest
|
||||
container_name: argus-ray-head
|
||||
ports:
|
||||
- "8265:8265"
|
||||
- "8080:8080"
|
||||
- "8000:8000"
|
||||
volumes:
|
||||
# NOTE: this compose file is intended for the dev env layout like:
|
||||
# /home2/argus/infra/mvp/{shared,verl,src/mvp}
|
||||
# so from `src/mvp/` the shared/verl dirs are `../../shared` and `../../verl`.
|
||||
- ../../verl:/workspace/verl
|
||||
- ../../shared:/private
|
||||
- .:/workspace/mvp
|
||||
shm_size: "10g"
|
||||
ulimits:
|
||||
nofile:
|
||||
soft: 65536
|
||||
hard: 65536
|
||||
cap_add:
|
||||
- SYS_ADMIN
|
||||
- SYS_PTRACE
|
||||
networks:
|
||||
- argus-ray-net
|
||||
environment:
|
||||
ARGUS_ROLE: "head"
|
||||
ARGUS_SHARED_ROOT: "/private"
|
||||
ARGUS_CLUSTER_NAME: "argus-ray"
|
||||
ARGUS_LOG_DIR: "/private/common/logs"
|
||||
# Make SFTPGo admin password available to API server started via `docker exec`.
|
||||
# Keep it consistent with the sftpgo container default below.
|
||||
SFTPGO_ADMIN_PASSWORD: "${SFTPGO_ADMIN_PASSWORD:-my-dev-sftpgo-admin}"
|
||||
HF_HOME: "/private/hf"
|
||||
HUGGINGFACE_HUB_CACHE: "/private/hf/hub"
|
||||
TRANSFORMERS_CACHE: "/private/hf/transformers"
|
||||
HF_ENDPOINT: "https://hf-mirror.com"
|
||||
PYTHONUNBUFFERED: "1"
|
||||
|
||||
# v3.0: Data management service (SFTPGo).
|
||||
# - SFTP: 2022
|
||||
# - Admin API/UI: 8080 (mapped to host 8081 to avoid collision with MVP API server on 8080)
|
||||
#
|
||||
# NOTE: This is for dev / h1 E2E. In prod you may use an internal mirror image/tag and different ports.
|
||||
sftpgo:
|
||||
image: drakkan/sftpgo:latest
|
||||
container_name: argus-sftpgo
|
||||
user: "0:0"
|
||||
ports:
|
||||
- "2022:2022"
|
||||
- "8081:8080"
|
||||
volumes:
|
||||
- ../../shared:/private
|
||||
- ../../shared/common/sftpgo:/var/lib/sftpgo
|
||||
networks:
|
||||
argus-ray-net:
|
||||
aliases:
|
||||
- sftpgo
|
||||
- argus-sftpgo
|
||||
environment:
|
||||
# Create a default admin on first start (used by API server to manage users).
|
||||
# Override on host as needed:
|
||||
# export SFTPGO_ADMIN_PASSWORD=...
|
||||
SFTPGO_DATA_PROVIDER__CREATE_DEFAULT_ADMIN: "true"
|
||||
# Persist the sqlite DB under the mounted metadata dir; otherwise it defaults to a relative path.
|
||||
SFTPGO_DATA_PROVIDER__NAME: "/var/lib/sftpgo/sftpgo.db"
|
||||
SFTPGO_DEFAULT_ADMIN_USERNAME: "admin"
|
||||
SFTPGO_DEFAULT_ADMIN_PASSWORD: "${SFTPGO_ADMIN_PASSWORD:-my-dev-sftpgo-admin}"
|
||||
# Explicitly pin default ports via env-var schema (double-underscore for nesting).
|
||||
SFTPGO_HTTPD__BINDINGS__0__PORT: "8080"
|
||||
SFTPGO_SFTPD__BINDINGS__0__PORT: "2022"
|
||||
|
||||
# v3.6: Weights & Biases local server (shared single account; per-user projects).
|
||||
#
|
||||
# - UI: 8080 in container, mapped to host 8090
|
||||
# - Data: persist under /private/common/wandb (NFS/shared)
|
||||
#
|
||||
# NOTE:
|
||||
# - The license/initial setup is done via the W&B UI.
|
||||
# - Ray jobs use WANDB_BASE_URL from config (recommended: docker bridge gateway + host published port).
|
||||
wandb:
|
||||
image: wandb/local:latest
|
||||
container_name: argus-wandb
|
||||
ports:
|
||||
- "8090:8080"
|
||||
volumes:
|
||||
- ../../shared/common/wandb:/vol
|
||||
networks:
|
||||
argus-ray-net:
|
||||
aliases:
|
||||
- wandb
|
||||
- argus-wandb
|
||||
environment:
|
||||
# W&B Local uses this as its externally-reachable base URL for the web app/backend.
|
||||
# If you access via SSH port-forwarding, set WANDB_HOST to the exact URL you open in the browser
|
||||
# (e.g. http://localhost:8090 or http://localhost:8080).
|
||||
HOST: "${WANDB_HOST:-http://localhost:8090}"
|
||||
# Optional: provide license at container start to avoid UI activation flakiness.
|
||||
LICENSE: "${WANDB_LICENSE:-}"
|
||||
|
||||
ray_worker_0:
|
||||
image: argus/argus-ray-node:vllm011.latest
|
||||
build:
|
||||
context: .
|
||||
dockerfile: images/argus-ray-node/Dockerfile
|
||||
args:
|
||||
BASE_IMAGE: verlai/verl:vllm011.latest
|
||||
container_name: argus-ray-worker-0
|
||||
volumes:
|
||||
- ../../verl:/workspace/verl
|
||||
- ../../shared:/private
|
||||
- .:/workspace/mvp
|
||||
shm_size: "10g"
|
||||
ulimits:
|
||||
nofile:
|
||||
soft: 65536
|
||||
hard: 65536
|
||||
cap_add:
|
||||
- SYS_ADMIN
|
||||
- SYS_PTRACE
|
||||
networks:
|
||||
- argus-ray-net
|
||||
runtime: nvidia
|
||||
environment:
|
||||
ARGUS_ROLE: "worker"
|
||||
ARGUS_SHARED_ROOT: "/private"
|
||||
ARGUS_CLUSTER_NAME: "argus-ray"
|
||||
ARGUS_LOG_DIR: "/private/common/logs"
|
||||
ARGUS_WORKER_RESOURCES_KV: "worker_node=100"
|
||||
NVIDIA_VISIBLE_DEVICES: "0,1,2,3"
|
||||
NVIDIA_DRIVER_CAPABILITIES: "all"
|
||||
HF_HOME: "/private/hf"
|
||||
HUGGINGFACE_HUB_CACHE: "/private/hf/hub"
|
||||
TRANSFORMERS_CACHE: "/private/hf/transformers"
|
||||
HF_ENDPOINT: "https://hf-mirror.com"
|
||||
PYTHONUNBUFFERED: "1"
|
||||
|
||||
ray_worker_1:
|
||||
image: argus/argus-ray-node:vllm011.latest
|
||||
build:
|
||||
context: .
|
||||
dockerfile: images/argus-ray-node/Dockerfile
|
||||
args:
|
||||
BASE_IMAGE: verlai/verl:vllm011.latest
|
||||
container_name: argus-ray-worker-1
|
||||
volumes:
|
||||
- ../../verl:/workspace/verl
|
||||
- ../../shared:/private
|
||||
- .:/workspace/mvp
|
||||
shm_size: "10g"
|
||||
ulimits:
|
||||
nofile:
|
||||
soft: 65536
|
||||
hard: 65536
|
||||
cap_add:
|
||||
- SYS_ADMIN
|
||||
- SYS_PTRACE
|
||||
networks:
|
||||
- argus-ray-net
|
||||
runtime: nvidia
|
||||
environment:
|
||||
ARGUS_ROLE: "worker"
|
||||
ARGUS_SHARED_ROOT: "/private"
|
||||
ARGUS_CLUSTER_NAME: "argus-ray"
|
||||
ARGUS_LOG_DIR: "/private/common/logs"
|
||||
ARGUS_WORKER_RESOURCES_KV: "worker_node=100"
|
||||
NVIDIA_VISIBLE_DEVICES: "4,5,6,7"
|
||||
NVIDIA_DRIVER_CAPABILITIES: "all"
|
||||
HF_HOME: "/private/hf"
|
||||
HUGGINGFACE_HUB_CACHE: "/private/hf/hub"
|
||||
TRANSFORMERS_CACHE: "/private/hf/transformers"
|
||||
HF_ENDPOINT: "https://hf-mirror.com"
|
||||
PYTHONUNBUFFERED: "1"
|
||||
|
||||
networks:
|
||||
argus-ray-net:
|
||||
driver: bridge
|
||||
35
src/mvp/images/argus-ray-node/Dockerfile
Normal file
@@ -0,0 +1,35 @@
|
||||
ARG BASE_IMAGE=verlai/verl:vllm011.latest
|
||||
FROM ${BASE_IMAGE}
|
||||
|
||||
SHELL ["/bin/bash", "-lc"]
|
||||
|
||||
# Install supervisord (prefer pip to avoid relying on distro package manager).
|
||||
RUN python3 -m pip install --no-cache-dir supervisor
|
||||
|
||||
# v3.8: Ray Serve LLM deps (keep Ray version pinned to what's already in the base image).
|
||||
# NOTE: base image already includes Ray; we only add extras.
|
||||
RUN RAY_VER="$(python3 -c 'import ray; print(ray.__version__)')" && \
|
||||
python3 -m pip install --no-cache-dir "ray[serve,llm]==${RAY_VER}"
|
||||
# Ray Serve LLM's import chain currently pulls in ray.rllib which requires extra deps.
|
||||
# Install them explicitly to make `from ray.serve.llm import ...` work reliably.
|
||||
RUN python3 -m pip install --no-cache-dir gymnasium dm-tree && \
|
||||
python3 -c "from ray.serve.llm import LLMConfig, build_openai_app; print('ray_serve_llm_ok')"
|
||||
|
||||
RUN mkdir -p /opt/argus/py/argus/ray
|
||||
|
||||
# Minimal embedded code for stateless pool (API code is intentionally excluded).
|
||||
COPY py/argus/__init__.py /opt/argus/py/argus/__init__.py
|
||||
COPY py/argus/ray/__init__.py /opt/argus/py/argus/ray/__init__.py
|
||||
COPY py/argus/ray/discovery.py /opt/argus/py/argus/ray/discovery.py
|
||||
COPY py/argus/ray/head_publisher.py /opt/argus/py/argus/ray/head_publisher.py
|
||||
COPY py/argus/ray/worker_watchdog.py /opt/argus/py/argus/ray/worker_watchdog.py
|
||||
|
||||
COPY images/argus-ray-node/entrypoint.sh /usr/local/bin/argus-entrypoint.sh
|
||||
COPY images/argus-ray-node/argus-head-ray.sh /usr/local/bin/argus-head-ray.sh
|
||||
COPY images/argus-ray-node/argus-head-publisher.sh /usr/local/bin/argus-head-publisher.sh
|
||||
COPY images/argus-ray-node/argus-worker-watchdog.sh /usr/local/bin/argus-worker-watchdog.sh
|
||||
RUN chmod +x /usr/local/bin/argus-entrypoint.sh /usr/local/bin/argus-head-ray.sh /usr/local/bin/argus-head-publisher.sh /usr/local/bin/argus-worker-watchdog.sh
|
||||
|
||||
ENV PYTHONUNBUFFERED=1
|
||||
|
||||
ENTRYPOINT ["/usr/local/bin/argus-entrypoint.sh"]
|
||||
20
src/mvp/images/argus-ray-node/argus-head-publisher.sh
Normal file
@@ -0,0 +1,20 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
cluster_name="${ARGUS_CLUSTER_NAME:-argus-ray}"
|
||||
head_ip_file="${ARGUS_HEAD_IP_FILE:?missing ARGUS_HEAD_IP_FILE}"
|
||||
node_ip="${ARGUS_NODE_IP:?missing ARGUS_NODE_IP}"
|
||||
ray_port="${ARGUS_RAY_PORT:-6379}"
|
||||
dashboard_port="${ARGUS_DASHBOARD_PORT:-8265}"
|
||||
ttl_s="${ARGUS_TTL_S:-60}"
|
||||
refresh_s="${ARGUS_REFRESH_S:-10}"
|
||||
|
||||
exec python3 -m argus.ray.head_publisher \
|
||||
--cluster-name "${cluster_name}" \
|
||||
--head-ip-file "${head_ip_file}" \
|
||||
--head-ip "${node_ip}" \
|
||||
--gcs-port "${ray_port}" \
|
||||
--dashboard-port "${dashboard_port}" \
|
||||
--ttl-s "${ttl_s}" \
|
||||
--refresh-s "${refresh_s}"
|
||||
|
||||
23
src/mvp/images/argus-ray-node/argus-head-ray.sh
Normal file
@@ -0,0 +1,23 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
node_ip="${ARGUS_NODE_IP:?missing ARGUS_NODE_IP}"
|
||||
ray_port="${ARGUS_RAY_PORT:-6379}"
|
||||
dashboard_port="${ARGUS_DASHBOARD_PORT:-8265}"
|
||||
|
||||
export RAY_DISABLE_USAGE_STATS=1
|
||||
|
||||
ray stop --force || true
|
||||
|
||||
# Head should not run training workloads.
|
||||
exec ray start \
|
||||
--head \
|
||||
--node-ip-address="${node_ip}" \
|
||||
--port="${ray_port}" \
|
||||
--dashboard-host=0.0.0.0 \
|
||||
--dashboard-port="${dashboard_port}" \
|
||||
--num-cpus="${ARGUS_HEAD_NUM_CPUS:-1}" \
|
||||
--num-gpus=0 \
|
||||
--disable-usage-stats \
|
||||
--block \
|
||||
${ARGUS_RAY_EXTRA_ARGS:-}
|
||||
24
src/mvp/images/argus-ray-node/argus-worker-watchdog.sh
Normal file
@@ -0,0 +1,24 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
head_ip_file="${ARGUS_HEAD_IP_FILE:?missing ARGUS_HEAD_IP_FILE}"
|
||||
node_ip="${ARGUS_NODE_IP:?missing ARGUS_NODE_IP}"
|
||||
poll_s="${ARGUS_POLL_S:-5}"
|
||||
|
||||
resources_kv="${ARGUS_WORKER_RESOURCES_KV:-worker_node=100}"
|
||||
|
||||
args=()
|
||||
IFS=',' read -r -a parts <<<"${resources_kv}"
|
||||
for kv in "${parts[@]}"; do
|
||||
kv="$(echo "${kv}" | xargs || true)"
|
||||
[[ -z "${kv}" ]] && continue
|
||||
args+=( "--resources-kv" "${kv}" )
|
||||
done
|
||||
|
||||
exec python3 -m argus.ray.worker_watchdog \
|
||||
--head-ip-file "${head_ip_file}" \
|
||||
--node-ip "${node_ip}" \
|
||||
"${args[@]}" \
|
||||
--poll-s "${poll_s}" \
|
||||
--ray-start-block
|
||||
|
||||
93
src/mvp/images/argus-ray-node/entrypoint.sh
Normal file
@@ -0,0 +1,93 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
role="${ARGUS_ROLE:-}"
|
||||
if [[ -z "${role}" ]]; then
|
||||
echo "ERROR: ARGUS_ROLE must be set to head|worker" >&2
|
||||
exit 1
|
||||
fi
|
||||
if [[ "${role}" != "head" && "${role}" != "worker" ]]; then
|
||||
echo "ERROR: invalid ARGUS_ROLE=${role} (expected head|worker)" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
shared_root="${ARGUS_SHARED_ROOT:-/private}"
|
||||
cluster_name="${ARGUS_CLUSTER_NAME:-argus-ray}"
|
||||
head_ip_file="${ARGUS_HEAD_IP_FILE:-${shared_root}/ray/discovery/${cluster_name}/head.json}"
|
||||
|
||||
ray_port="${ARGUS_RAY_PORT:-6379}"
|
||||
dashboard_port="${ARGUS_DASHBOARD_PORT:-8265}"
|
||||
|
||||
ttl_s="${ARGUS_TTL_S:-60}"
|
||||
refresh_s="${ARGUS_REFRESH_S:-10}"
|
||||
poll_s="${ARGUS_POLL_S:-5}"
|
||||
|
||||
log_dir="${ARGUS_LOG_DIR:-${shared_root}/common/logs}"
|
||||
mkdir -p "${log_dir}"
|
||||
|
||||
node_ip="${ARGUS_NODE_IP:-}"
|
||||
if [[ -z "${node_ip}" ]]; then
|
||||
# Prefer the first IP returned by `hostname -i`.
|
||||
node_ip="$(hostname -i | awk '{print $1}')"
|
||||
fi
|
||||
if [[ -z "${node_ip}" ]]; then
|
||||
echo "ERROR: failed to determine container IP; set ARGUS_NODE_IP" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
export PYTHONPATH="/opt/argus/py:${PYTHONPATH:-}"
|
||||
export ARGUS_SHARED_ROOT="${shared_root}"
|
||||
export ARGUS_CLUSTER_NAME="${cluster_name}"
|
||||
export ARGUS_HEAD_IP_FILE="${head_ip_file}"
|
||||
export ARGUS_NODE_IP="${node_ip}"
|
||||
export ARGUS_RAY_PORT="${ray_port}"
|
||||
export ARGUS_DASHBOARD_PORT="${dashboard_port}"
|
||||
export ARGUS_TTL_S="${ttl_s}"
|
||||
export ARGUS_REFRESH_S="${refresh_s}"
|
||||
export ARGUS_POLL_S="${poll_s}"
|
||||
export ARGUS_LOG_DIR="${log_dir}"
|
||||
|
||||
supervisor_conf="/tmp/supervisord.conf"
|
||||
cat >"${supervisor_conf}" <<'CONF'
|
||||
[supervisord]
|
||||
nodaemon=true
|
||||
pidfile=/tmp/supervisord.pid
|
||||
logfile=/dev/null
|
||||
childlogdir=/tmp
|
||||
user=root
|
||||
|
||||
[supervisorctl]
|
||||
serverurl=unix:///tmp/supervisor.sock
|
||||
|
||||
[unix_http_server]
|
||||
file=/tmp/supervisor.sock
|
||||
chmod=0700
|
||||
CONF
|
||||
|
||||
append_program() {
|
||||
local name="$1"
|
||||
local cmd="$2"
|
||||
local stdout="$3"
|
||||
local stderr="$4"
|
||||
cat >>"${supervisor_conf}" <<CONF
|
||||
|
||||
[program:${name}]
|
||||
command=${cmd}
|
||||
autorestart=true
|
||||
startsecs=1
|
||||
startretries=999
|
||||
stdout_logfile=${stdout}
|
||||
stdout_logfile_maxbytes=0
|
||||
stderr_logfile=${stderr}
|
||||
stderr_logfile_maxbytes=0
|
||||
CONF
|
||||
}
|
||||
|
||||
if [[ "${role}" == "head" ]]; then
|
||||
append_program "argus_head_ray" "/usr/local/bin/argus-head-ray.sh" "/dev/fd/1" "/dev/fd/2"
|
||||
append_program "argus_head_publisher" "/usr/local/bin/argus-head-publisher.sh" "${log_dir}/argus_head_publisher.log" "${log_dir}/argus_head_publisher.err.log"
|
||||
else
|
||||
append_program "argus_worker_watchdog" "/usr/local/bin/argus-worker-watchdog.sh" "${log_dir}/argus_worker_watchdog.log" "${log_dir}/argus_worker_watchdog.err.log"
|
||||
fi
|
||||
|
||||
exec supervisord -n -c "${supervisor_conf}"
|
||||
2
src/mvp/py/argus/__init__.py
Normal file
@@ -0,0 +1,2 @@
|
||||
__all__ = []
|
||||
|
||||
2
src/mvp/py/argus/cli/__init__.py
Normal file
@@ -0,0 +1,2 @@
|
||||
__all__ = []
|
||||
|
||||
72
src/mvp/py/argus/cli/run.py
Normal file
@@ -0,0 +1,72 @@
|
||||
#!/usr/bin/env python3
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
|
||||
|
||||
def _ensure_import_path() -> None:
|
||||
# Allow running as a file path (not just `python3 -m argus.cli.run`).
|
||||
# Ensure `/workspace/mvp/py` is on sys.path so `import argus.*` works.
|
||||
here = os.path.dirname(os.path.abspath(__file__))
|
||||
py_root = os.path.dirname(os.path.dirname(here)) # .../py
|
||||
if py_root not in sys.path:
|
||||
sys.path.insert(0, py_root)
|
||||
|
||||
|
||||
def main() -> int:
|
||||
_ensure_import_path()
|
||||
|
||||
from argus.ray.models import JobSpec, RayConfig
|
||||
from argus.ray.ray_job_tool import RayJobTool
|
||||
from argus.ray.yaml_io import load_yaml
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--config", required=True, help="MVP config yaml (ray/service)")
|
||||
parser.add_argument("--taskspec", help="Training TaskSpec yaml (required for submit)")
|
||||
parser.add_argument("--jobspec", dest="taskspec", help=argparse.SUPPRESS) # backward compatible alias
|
||||
parser.add_argument("--action", required=True, choices=["submit", "status", "stop", "logs", "list"])
|
||||
parser.add_argument("--submission-id", help="For status/stop/logs")
|
||||
parser.add_argument("--no-wait", action="store_true", help="Submit and return immediately")
|
||||
parser.add_argument("--tail", type=int, default=0, help="Tail N lines for logs")
|
||||
args = parser.parse_args()
|
||||
|
||||
cfg = RayConfig.from_dict(load_yaml(args.config))
|
||||
tool = RayJobTool(cfg)
|
||||
|
||||
if args.action == "submit":
|
||||
if not args.taskspec:
|
||||
raise SystemExit("--taskspec is required for submit")
|
||||
spec = JobSpec.from_dict(load_yaml(args.taskspec))
|
||||
submitted = tool.submit(spec, no_wait=args.no_wait)
|
||||
print(submitted)
|
||||
return 0
|
||||
|
||||
if args.action in ("status", "stop", "logs"):
|
||||
sid = args.submission_id or ""
|
||||
if not sid:
|
||||
raise SystemExit("--submission-id is required for status/stop/logs")
|
||||
if args.action == "status":
|
||||
print(tool.status(sid))
|
||||
return 0
|
||||
if args.action == "stop":
|
||||
print(tool.stop(sid))
|
||||
return 0
|
||||
logs = tool.logs(sid)
|
||||
if args.tail and args.tail > 0:
|
||||
lines = logs.splitlines()
|
||||
logs = "\n".join(lines[-args.tail :]) + ("\n" if lines else "")
|
||||
print(logs, end="")
|
||||
return 0
|
||||
|
||||
if args.action == "list":
|
||||
print(json.dumps(tool.list(), indent=2))
|
||||
return 0
|
||||
|
||||
raise SystemExit(f"unknown action: {args.action}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
2
src/mvp/py/argus/core/__init__.py
Normal file
@@ -0,0 +1,2 @@
|
||||
__all__ = []
|
||||
|
||||
44
src/mvp/py/argus/core/ids.py
Normal file
@@ -0,0 +1,44 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import secrets
|
||||
from datetime import datetime
|
||||
import re
|
||||
|
||||
|
||||
_USER_SAFE_RE = re.compile(r"[^a-z0-9_-]+")
|
||||
|
||||
|
||||
def _normalize_user_id(user_id: str) -> str:
|
||||
s = (user_id or "").strip().lower()
|
||||
s = _USER_SAFE_RE.sub("-", s)
|
||||
s = re.sub(r"-{2,}", "-", s).strip("-")
|
||||
return s[:24] if s else ""
|
||||
|
||||
|
||||
def new_task_id(workload: str, *, user_id: str | None = None) -> str:
|
||||
ts = datetime.now().strftime("%Y%m%d-%H%M%S")
|
||||
suffix = secrets.token_hex(2)
|
||||
u = _normalize_user_id(user_id or "")
|
||||
if u:
|
||||
return f"mvp2-{u}-{workload}-{ts}-{suffix}"
|
||||
return f"mvp2-{workload}-{ts}-{suffix}"
|
||||
|
||||
|
||||
def attempt_submission_id(task_id: str, attempt_no: int) -> str:
|
||||
return f"{task_id}--a{attempt_no:02d}"
|
||||
|
||||
|
||||
def new_model_key(*, user_id: str) -> str:
|
||||
"""
|
||||
Internal identifier for a serving model record.
|
||||
|
||||
Note:
|
||||
- model_id is the OpenAI-facing name (user_id + timestamp prefix + suffix).
|
||||
- model_key is used for stable DB identity and API resource path.
|
||||
"""
|
||||
ts = datetime.now().strftime("%Y%m%d-%H%M%S")
|
||||
suffix = secrets.token_hex(2)
|
||||
u = _normalize_user_id(user_id)
|
||||
if not u:
|
||||
raise ValueError("user_id is required")
|
||||
return f"mvp2-{u}-serve-{ts}-{suffix}"
|
||||
1
src/mvp/py/argus/ray/__init__.py
Normal file
@@ -0,0 +1 @@
|
||||
|
||||
108
src/mvp/py/argus/ray/builders.py
Normal file
@@ -0,0 +1,108 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
from .models import JobSpec
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BuiltCommand:
|
||||
argv: list[str]
|
||||
|
||||
|
||||
def build_training_argv(spec: JobSpec, submission_id: str, job_dir: str) -> BuiltCommand:
|
||||
"""
|
||||
Returns argv for the actual training process (Hydra overrides preserved).
|
||||
This argv is executed by a lightweight Python driver entrypoint.
|
||||
"""
|
||||
if spec.workload in ("ppo", "grpo"):
|
||||
algo_overrides: list[str] = []
|
||||
if spec.workload == "grpo":
|
||||
algo_overrides.append("algorithm.adv_estimator=grpo")
|
||||
|
||||
test_freq = spec.test_freq if spec.test_freq is not None else -1
|
||||
maybe_total_steps: list[str] = []
|
||||
if spec.total_training_steps is not None:
|
||||
maybe_total_steps.append(f"trainer.total_training_steps={spec.total_training_steps}")
|
||||
|
||||
argv = [
|
||||
"python3",
|
||||
"-m",
|
||||
"verl.trainer.main_ppo",
|
||||
f"data.train_files={spec.train_file}",
|
||||
f"data.val_files={spec.val_file}",
|
||||
"data.train_batch_size=256",
|
||||
"data.max_prompt_length=512",
|
||||
"data.max_response_length=512",
|
||||
f"actor_rollout_ref.model.path={spec.model_id}",
|
||||
"actor_rollout_ref.actor.optim.lr=1e-6",
|
||||
"actor_rollout_ref.actor.ppo_mini_batch_size=64",
|
||||
"actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=4",
|
||||
"actor_rollout_ref.rollout.name=vllm",
|
||||
"actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=8",
|
||||
"actor_rollout_ref.rollout.tensor_model_parallel_size=1",
|
||||
"actor_rollout_ref.rollout.gpu_memory_utilization=0.4",
|
||||
"actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4",
|
||||
"critic.optim.lr=1e-5",
|
||||
f"critic.model.path={spec.model_id}",
|
||||
"critic.ppo_micro_batch_size_per_gpu=4",
|
||||
"algorithm.kl_ctrl.kl_coef=0.001",
|
||||
*algo_overrides,
|
||||
"trainer.logger=console",
|
||||
"trainer.val_before_train=False",
|
||||
f"trainer.n_gpus_per_node={spec.n_gpus_per_node}",
|
||||
f"trainer.nnodes={spec.nnodes}",
|
||||
f"trainer.save_freq={spec.save_freq}",
|
||||
f"trainer.test_freq={test_freq}",
|
||||
f"trainer.total_epochs={spec.total_epochs}",
|
||||
*maybe_total_steps,
|
||||
"trainer.resume_mode=disable",
|
||||
f"trainer.default_local_dir={job_dir}/checkpoints",
|
||||
"+ray_kwargs.ray_init.address=auto",
|
||||
f"hydra.run.dir={job_dir}/logs/hydra",
|
||||
]
|
||||
return BuiltCommand(argv=argv)
|
||||
|
||||
if spec.workload == "sft":
|
||||
trainer_device = spec.trainer_device or "cpu"
|
||||
maybe_total_steps: list[str] = []
|
||||
if spec.total_training_steps is not None:
|
||||
maybe_total_steps.append(f"trainer.total_training_steps={spec.total_training_steps}")
|
||||
|
||||
argv = [
|
||||
"python3",
|
||||
"-m",
|
||||
"verl.trainer.sft_trainer_ray",
|
||||
f"model.path={spec.model_id}",
|
||||
f"data.train_files={spec.train_file}",
|
||||
f"data.val_files={spec.val_file}",
|
||||
"data.train_batch_size=64",
|
||||
"data.micro_batch_size_per_gpu=1",
|
||||
"data.max_token_len_per_gpu=2048",
|
||||
"data.max_length=1024",
|
||||
"trainer.logger=console",
|
||||
"trainer.project_name=mvp11-sft",
|
||||
f"trainer.experiment_name={submission_id}",
|
||||
f"trainer.total_epochs={spec.total_epochs}",
|
||||
*maybe_total_steps,
|
||||
f"trainer.save_freq={spec.save_freq}",
|
||||
"trainer.test_freq=-1",
|
||||
"trainer.resume_mode=disable",
|
||||
f"trainer.device={trainer_device}",
|
||||
f"trainer.default_local_dir={job_dir}/checkpoints",
|
||||
f"trainer.nnodes={spec.nnodes}",
|
||||
f"trainer.n_gpus_per_node={spec.n_gpus_per_node}",
|
||||
f"hydra.run.dir={job_dir}/logs/hydra",
|
||||
]
|
||||
return BuiltCommand(argv=argv)
|
||||
|
||||
raise ValueError(f"unsupported workload: {spec.workload}")
|
||||
|
||||
|
||||
def build_advanced_argv(command: str) -> BuiltCommand:
|
||||
"""
|
||||
Execute an arbitrary user-provided command under bash for maximum shell compatibility
|
||||
(multiline, env vars, line continuations, etc.).
|
||||
"""
|
||||
|
||||
return BuiltCommand(argv=["bash", "-lc", command])
|
||||
115
src/mvp/py/argus/ray/discovery.py
Normal file
@@ -0,0 +1,115 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
|
||||
def _utc_now() -> datetime:
|
||||
return datetime.now(timezone.utc)
|
||||
|
||||
|
||||
def _parse_utc(ts: str) -> datetime:
|
||||
# Accept ISO8601 with trailing Z
|
||||
s = ts.strip()
|
||||
if s.endswith("Z"):
|
||||
s = s[:-1] + "+00:00"
|
||||
return datetime.fromisoformat(s).astimezone(timezone.utc)
|
||||
|
||||
|
||||
def _iso_z(dt: datetime) -> str:
|
||||
return dt.astimezone(timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z")
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class HeadRecord:
|
||||
cluster_name: str
|
||||
head_ip: str
|
||||
gcs_port: int
|
||||
dashboard_port: int
|
||||
job_server_url: str
|
||||
updated_at: str
|
||||
expires_at: str
|
||||
|
||||
@staticmethod
|
||||
def from_dict(d: dict[str, Any]) -> "HeadRecord":
|
||||
return HeadRecord(
|
||||
cluster_name=str(d["cluster_name"]),
|
||||
head_ip=str(d["head_ip"]),
|
||||
gcs_port=int(d.get("gcs_port", 6379)),
|
||||
dashboard_port=int(d.get("dashboard_port", 8265)),
|
||||
job_server_url=str(d.get("job_server_url") or f"http://{d['head_ip']}:{int(d.get('dashboard_port', 8265))}"),
|
||||
updated_at=str(d["updated_at"]),
|
||||
expires_at=str(d["expires_at"]),
|
||||
)
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return {
|
||||
"cluster_name": self.cluster_name,
|
||||
"head_ip": self.head_ip,
|
||||
"gcs_port": self.gcs_port,
|
||||
"dashboard_port": self.dashboard_port,
|
||||
"job_server_url": self.job_server_url,
|
||||
"updated_at": self.updated_at,
|
||||
"expires_at": self.expires_at,
|
||||
}
|
||||
|
||||
def head_addr(self) -> str:
|
||||
return f"{self.head_ip}:{self.gcs_port}"
|
||||
|
||||
|
||||
def build_head_record(
|
||||
*,
|
||||
cluster_name: str,
|
||||
head_ip: str,
|
||||
gcs_port: int = 6379,
|
||||
dashboard_port: int = 8265,
|
||||
ttl_s: int = 60,
|
||||
now: datetime | None = None,
|
||||
) -> HeadRecord:
|
||||
now_dt = now or _utc_now()
|
||||
expires = now_dt + timedelta(seconds=int(ttl_s))
|
||||
updated_at = _iso_z(now_dt)
|
||||
expires_at = _iso_z(expires)
|
||||
return HeadRecord(
|
||||
cluster_name=cluster_name,
|
||||
head_ip=head_ip,
|
||||
gcs_port=int(gcs_port),
|
||||
dashboard_port=int(dashboard_port),
|
||||
job_server_url=f"http://{head_ip}:{int(dashboard_port)}",
|
||||
updated_at=updated_at,
|
||||
expires_at=expires_at,
|
||||
)
|
||||
|
||||
|
||||
def load_head_record(path: str, *, now: datetime | None = None) -> HeadRecord | None:
|
||||
p = Path(path)
|
||||
if not p.exists():
|
||||
return None
|
||||
try:
|
||||
obj = json.loads(p.read_text(encoding="utf-8"))
|
||||
except Exception:
|
||||
return None
|
||||
if not isinstance(obj, dict):
|
||||
return None
|
||||
try:
|
||||
rec = HeadRecord.from_dict(obj)
|
||||
expires = _parse_utc(rec.expires_at)
|
||||
except Exception:
|
||||
return None
|
||||
now_dt = now or _utc_now()
|
||||
if expires <= now_dt:
|
||||
return None
|
||||
return rec
|
||||
|
||||
|
||||
def write_head_record_atomic(path: str, rec: HeadRecord) -> None:
|
||||
p = Path(path)
|
||||
os.makedirs(p.parent, exist_ok=True)
|
||||
tmp = p.with_suffix(p.suffix + ".tmp")
|
||||
tmp.write_text(json.dumps(rec.to_dict(), indent=2, sort_keys=True) + "\n", encoding="utf-8")
|
||||
tmp.replace(p)
|
||||
|
||||
65
src/mvp/py/argus/ray/driver_entrypoint.py
Normal file
@@ -0,0 +1,65 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import shlex
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def _preflight() -> None:
|
||||
print("MVP_PRECHECK_PYTHON:", sys.executable, flush=True)
|
||||
print("MVP_PRECHECK_PYTHONPATH:", os.environ.get("PYTHONPATH"), flush=True)
|
||||
print("MVP_PRECHECK_MVP_CODE_PATH:", os.environ.get("MVP_CODE_PATH"), flush=True)
|
||||
try:
|
||||
import verl # type: ignore
|
||||
|
||||
print("MVP_PRECHECK_VERL_FILE:", getattr(verl, "__file__", None), flush=True)
|
||||
except Exception as e:
|
||||
print("MVP_PRECHECK_VERL_IMPORT_ERROR:", repr(e), flush=True)
|
||||
|
||||
try:
|
||||
import mvp_marker # type: ignore
|
||||
|
||||
print("MVP_PRECHECK_MARKER:", getattr(mvp_marker, "MARKER", None), flush=True)
|
||||
except Exception as e:
|
||||
print("MVP_PRECHECK_MARKER_MISSING:", repr(e), flush=True)
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--job-dir", required=True)
|
||||
parser.add_argument("cmd", nargs=argparse.REMAINDER)
|
||||
args = parser.parse_args()
|
||||
|
||||
job_dir = Path(args.job_dir)
|
||||
job_dir.mkdir(parents=True, exist_ok=True)
|
||||
os.environ.setdefault("MVP_JOB_DIR", str(job_dir))
|
||||
os.chdir(job_dir)
|
||||
|
||||
_preflight()
|
||||
|
||||
if not args.cmd:
|
||||
print("no command provided", file=sys.stderr)
|
||||
return 2
|
||||
|
||||
# argparse includes the leading "--" if the caller uses it; strip it.
|
||||
cmd = list(args.cmd)
|
||||
if cmd and cmd[0] == "--":
|
||||
cmd = cmd[1:]
|
||||
if not cmd:
|
||||
print("no command provided", file=sys.stderr)
|
||||
return 2
|
||||
|
||||
# Execute training command as a subprocess so that logs are captured by Ray job logs.
|
||||
cmd_str = " ".join(shlex.quote(x) for x in cmd)
|
||||
print("MVP_DRIVER_EXEC:", cmd_str, flush=True)
|
||||
|
||||
proc = subprocess.run(cmd, check=False)
|
||||
print("MVP_DRIVER_EXIT_CODE:", proc.returncode, flush=True)
|
||||
return proc.returncode
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
65
src/mvp/py/argus/ray/head_publisher.py
Normal file
@@ -0,0 +1,65 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import time
|
||||
|
||||
from .discovery import build_head_record, write_head_record_atomic
|
||||
|
||||
|
||||
def publish_once(
|
||||
*,
|
||||
cluster_name: str,
|
||||
head_ip_file: str,
|
||||
head_ip: str,
|
||||
gcs_port: int,
|
||||
dashboard_port: int,
|
||||
ttl_s: int,
|
||||
) -> None:
|
||||
rec = build_head_record(
|
||||
cluster_name=cluster_name,
|
||||
head_ip=head_ip,
|
||||
gcs_port=gcs_port,
|
||||
dashboard_port=dashboard_port,
|
||||
ttl_s=ttl_s,
|
||||
)
|
||||
write_head_record_atomic(head_ip_file, rec)
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
ap = argparse.ArgumentParser(description="Publish Ray head address to shared storage (head.json).")
|
||||
ap.add_argument("--cluster-name", required=True)
|
||||
ap.add_argument("--head-ip-file", required=True)
|
||||
ap.add_argument("--head-ip", required=True)
|
||||
ap.add_argument("--gcs-port", type=int, default=6379)
|
||||
ap.add_argument("--dashboard-port", type=int, default=8265)
|
||||
ap.add_argument("--ttl-s", type=int, default=60)
|
||||
ap.add_argument("--refresh-s", type=int, default=10)
|
||||
ap.add_argument("--once", action="store_true", help="Write once then exit (for testing/debug).")
|
||||
args = ap.parse_args(argv)
|
||||
|
||||
if args.once:
|
||||
publish_once(
|
||||
cluster_name=args.cluster_name,
|
||||
head_ip_file=args.head_ip_file,
|
||||
head_ip=args.head_ip,
|
||||
gcs_port=args.gcs_port,
|
||||
dashboard_port=args.dashboard_port,
|
||||
ttl_s=args.ttl_s,
|
||||
)
|
||||
return 0
|
||||
|
||||
refresh_s = max(1, int(args.refresh_s))
|
||||
while True:
|
||||
publish_once(
|
||||
cluster_name=args.cluster_name,
|
||||
head_ip_file=args.head_ip_file,
|
||||
head_ip=args.head_ip,
|
||||
gcs_port=args.gcs_port,
|
||||
dashboard_port=args.dashboard_port,
|
||||
ttl_s=args.ttl_s,
|
||||
)
|
||||
time.sleep(refresh_s)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
183
src/mvp/py/argus/ray/models.py
Normal file
@@ -0,0 +1,183 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
|
||||
def _require(d: dict[str, Any], key: str) -> Any:
|
||||
if key not in d or d[key] in (None, ""):
|
||||
raise ValueError(f"missing required field: {key}")
|
||||
return d[key]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class RayConfig:
|
||||
address: str
|
||||
shared_root: str
|
||||
entrypoint_num_cpus: float
|
||||
entrypoint_resources: dict[str, float]
|
||||
runtime_env_env_vars: dict[str, str]
|
||||
user_code_path: str
|
||||
verl_code_path: str
|
||||
|
||||
@staticmethod
|
||||
def from_dict(d: dict[str, Any]) -> "RayConfig":
|
||||
# New format: root.ray
|
||||
# Backward compatible: root (flat)
|
||||
if isinstance(d.get("ray"), dict):
|
||||
d = d["ray"] # type: ignore[assignment]
|
||||
|
||||
runtime_env = d.get("runtime_env") or {}
|
||||
env_vars = (runtime_env.get("env_vars") or {}) if isinstance(runtime_env, dict) else {}
|
||||
if not isinstance(env_vars, dict):
|
||||
raise ValueError("runtime_env.env_vars must be a mapping")
|
||||
|
||||
entrypoint_resources = d.get("entrypoint_resources") or {}
|
||||
if not isinstance(entrypoint_resources, dict):
|
||||
raise ValueError("entrypoint_resources must be a mapping")
|
||||
|
||||
return RayConfig(
|
||||
address=str(_require(d, "address")),
|
||||
shared_root=str(_require(d, "shared_root")),
|
||||
entrypoint_num_cpus=float(d.get("entrypoint_num_cpus", 1)),
|
||||
entrypoint_resources={str(k): float(v) for k, v in entrypoint_resources.items()},
|
||||
runtime_env_env_vars={str(k): str(v) for k, v in env_vars.items()},
|
||||
user_code_path=str(d.get("user_code_path", f"{_require(d, 'shared_root')}/user/code")),
|
||||
verl_code_path=str(d.get("verl_code_path", f"{_require(d, 'shared_root')}/common/code/verl/verl_repo")),
|
||||
)
|
||||
|
||||
def to_public_dict(self) -> dict[str, Any]:
|
||||
return {
|
||||
"address": self.address,
|
||||
"shared_root": self.shared_root,
|
||||
"entrypoint_num_cpus": self.entrypoint_num_cpus,
|
||||
"entrypoint_resources": self.entrypoint_resources,
|
||||
"runtime_env": {"env_vars": self.runtime_env_env_vars},
|
||||
"user_code_path": self.user_code_path,
|
||||
"verl_code_path": self.verl_code_path,
|
||||
}
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class JobSpec:
|
||||
workload: str # ppo|grpo|sft
|
||||
submission_id: str | None
|
||||
code_path: str
|
||||
model_id: str
|
||||
|
||||
train_file: str
|
||||
val_file: str
|
||||
|
||||
nnodes: int
|
||||
n_gpus_per_node: int
|
||||
|
||||
total_epochs: int
|
||||
total_training_steps: int | None
|
||||
|
||||
save_freq: int
|
||||
test_freq: int | None
|
||||
|
||||
trainer_device: str | None # only for sft (driver-side device)
|
||||
|
||||
@staticmethod
|
||||
def from_dict(d: dict[str, Any]) -> "JobSpec":
|
||||
workload = str(_require(d, "workload"))
|
||||
if workload not in ("ppo", "grpo", "sft"):
|
||||
raise ValueError(f"unsupported workload: {workload}")
|
||||
|
||||
test_freq = d.get("test_freq", None)
|
||||
if test_freq in ("", "null"):
|
||||
test_freq = None
|
||||
|
||||
total_training_steps = d.get("total_training_steps", None)
|
||||
if total_training_steps in ("", "null"):
|
||||
total_training_steps = None
|
||||
|
||||
return JobSpec(
|
||||
workload=workload,
|
||||
submission_id=(str(d["submission_id"]) if d.get("submission_id") else None),
|
||||
code_path=str(_require(d, "code_path")),
|
||||
model_id=str(_require(d, "model_id")),
|
||||
train_file=str(_require(d, "train_file")),
|
||||
val_file=str(_require(d, "val_file")),
|
||||
nnodes=int(d.get("nnodes", 2)),
|
||||
n_gpus_per_node=int(d.get("n_gpus_per_node", 4)),
|
||||
total_epochs=int(d.get("total_epochs", 1)),
|
||||
# If omitted or explicitly null: do not override trainer.total_training_steps
|
||||
# (let VERL derive it from epochs and dataset length).
|
||||
total_training_steps=(int(total_training_steps) if total_training_steps is not None else None),
|
||||
save_freq=int(d.get("save_freq", 10)),
|
||||
test_freq=(int(test_freq) if test_freq is not None else None),
|
||||
trainer_device=(str(d.get("trainer_device")) if d.get("trainer_device") else None),
|
||||
)
|
||||
|
||||
def to_public_dict(self) -> dict[str, Any]:
|
||||
out: dict[str, Any] = {
|
||||
"workload": self.workload,
|
||||
"submission_id": self.submission_id or "",
|
||||
"code_path": self.code_path,
|
||||
"model_id": self.model_id,
|
||||
"train_file": self.train_file,
|
||||
"val_file": self.val_file,
|
||||
"nnodes": self.nnodes,
|
||||
"n_gpus_per_node": self.n_gpus_per_node,
|
||||
"total_epochs": self.total_epochs,
|
||||
"total_training_steps": self.total_training_steps,
|
||||
"save_freq": self.save_freq,
|
||||
"test_freq": self.test_freq,
|
||||
}
|
||||
if self.workload == "sft":
|
||||
out["trainer_device"] = self.trainer_device or "cpu"
|
||||
return out
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class AdvancedTaskSpec:
|
||||
kind: str # advanced
|
||||
workload: str # always "advanced" for classification/ID prefix
|
||||
submission_id: str | None
|
||||
nnodes: int
|
||||
n_gpus_per_node: int
|
||||
command: str
|
||||
|
||||
@staticmethod
|
||||
def from_dict(d: dict[str, Any]) -> "AdvancedTaskSpec":
|
||||
kind = str(_require(d, "kind"))
|
||||
if kind != "advanced":
|
||||
raise ValueError(f"unsupported taskspec kind: {kind}")
|
||||
|
||||
command = str(_require(d, "command"))
|
||||
|
||||
return AdvancedTaskSpec(
|
||||
kind=kind,
|
||||
workload="advanced",
|
||||
submission_id=(str(d["submission_id"]) if d.get("submission_id") else None),
|
||||
nnodes=int(_require(d, "nnodes")),
|
||||
n_gpus_per_node=int(_require(d, "n_gpus_per_node")),
|
||||
command=command,
|
||||
)
|
||||
|
||||
def to_public_dict(self) -> dict[str, Any]:
|
||||
return {
|
||||
"kind": self.kind,
|
||||
"workload": self.workload,
|
||||
"submission_id": self.submission_id or "",
|
||||
"nnodes": self.nnodes,
|
||||
"n_gpus_per_node": self.n_gpus_per_node,
|
||||
"command": self.command,
|
||||
}
|
||||
|
||||
|
||||
def parse_taskspec(d: dict[str, Any]) -> JobSpec | AdvancedTaskSpec:
|
||||
"""
|
||||
v3.x compatibility rule:
|
||||
- If no "kind": treat as legacy/basic JobSpec.
|
||||
- If kind == "advanced": parse AdvancedTaskSpec.
|
||||
"""
|
||||
|
||||
kind = d.get("kind")
|
||||
if kind in (None, ""):
|
||||
return JobSpec.from_dict(d)
|
||||
if kind == "advanced":
|
||||
return AdvancedTaskSpec.from_dict(d)
|
||||
raise ValueError(f"unsupported taskspec kind: {kind}")
|
||||
289
src/mvp/py/argus/ray/ray_job_tool.py
Normal file
@@ -0,0 +1,289 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import shlex
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import ray
|
||||
from ray.job_submission import JobSubmissionClient
|
||||
|
||||
from .builders import BuiltCommand, build_advanced_argv, build_training_argv
|
||||
from .models import AdvancedTaskSpec, JobSpec, RayConfig
|
||||
from .yaml_io import dump_yaml
|
||||
|
||||
|
||||
def _ts() -> str:
|
||||
return datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
|
||||
|
||||
def _mkdir(p: Path) -> None:
|
||||
p.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
|
||||
def _write_text(p: Path, content: str) -> None:
|
||||
_mkdir(p.parent)
|
||||
p.write_text(content, encoding="utf-8")
|
||||
|
||||
|
||||
def _write_json(p: Path, obj: Any) -> None:
|
||||
_write_text(p, json.dumps(obj, indent=2, ensure_ascii=False) + "\n")
|
||||
|
||||
|
||||
def _safe_basename(path: str) -> str:
|
||||
return path.rstrip("/").split("/")[-1]
|
||||
|
||||
|
||||
class RayJobTool:
|
||||
def __init__(self, cfg: RayConfig):
|
||||
self.cfg = cfg
|
||||
self.client = JobSubmissionClient(cfg.address)
|
||||
|
||||
def _job_dir(self, submission_id: str) -> str:
|
||||
return f"{self.cfg.shared_root}/jobs/{submission_id}"
|
||||
|
||||
def _runtime_env(self, spec: JobSpec, *, extra_env_vars: dict[str, str] | None = None) -> dict[str, Any]:
|
||||
return self._runtime_env_for(code_path=spec.code_path, workload=spec.workload, extra_env_vars=extra_env_vars)
|
||||
|
||||
def _runtime_env_for(
|
||||
self,
|
||||
*,
|
||||
code_path: str,
|
||||
workload: str,
|
||||
extra_pythonpath: list[str] | None = None,
|
||||
extra_env_vars: dict[str, str] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
env_vars = dict(self.cfg.runtime_env_env_vars)
|
||||
|
||||
# Default HF cache
|
||||
env_vars.setdefault("HF_HOME", f"{self.cfg.shared_root}/hf")
|
||||
env_vars.setdefault("HUGGINGFACE_HUB_CACHE", f"{self.cfg.shared_root}/hf/hub")
|
||||
env_vars.setdefault("TRANSFORMERS_CACHE", f"{self.cfg.shared_root}/hf/transformers")
|
||||
env_vars.setdefault("PYTHONUNBUFFERED", "1")
|
||||
|
||||
# Tool code path must be importable on workers (compose mounts `/workspace/mvp` into all containers).
|
||||
# Place it before verl code to avoid interfering with verl import priority.
|
||||
tool_code_path = os.environ.get("MVP_TOOL_CODE_PATH", "/workspace/mvp/py")
|
||||
|
||||
existing = env_vars.get("PYTHONPATH", "")
|
||||
base_parts = [tool_code_path, code_path, self.cfg.user_code_path]
|
||||
if extra_pythonpath:
|
||||
base_parts.extend(extra_pythonpath)
|
||||
prefix = ":".join(base_parts)
|
||||
env_vars["PYTHONPATH"] = f"{prefix}:{existing}" if existing else prefix
|
||||
|
||||
# For debugging / log visibility
|
||||
env_vars["MVP_CODE_PATH"] = code_path
|
||||
|
||||
# SFT: ensure ray.init() connects to the cluster
|
||||
if workload == "sft":
|
||||
env_vars.setdefault("RAY_ADDRESS", "auto")
|
||||
|
||||
# Per-job extra env vars (e.g. W&B). Do not override existing keys.
|
||||
if extra_env_vars:
|
||||
for k, v in dict(extra_env_vars).items():
|
||||
if k not in env_vars:
|
||||
env_vars[k] = str(v)
|
||||
|
||||
return {"env_vars": env_vars}
|
||||
|
||||
def submit(
|
||||
self,
|
||||
spec: JobSpec,
|
||||
no_wait: bool,
|
||||
job_dir: str | None = None,
|
||||
*,
|
||||
extra_env_vars: dict[str, str] | None = None,
|
||||
extra_overrides: list[str] | None = None,
|
||||
) -> str:
|
||||
submission_id = spec.submission_id or f"mvp11_{spec.workload}_{_ts()}_{os.getpid()}"
|
||||
job_dir = job_dir or self._job_dir(submission_id)
|
||||
|
||||
built = build_training_argv(spec, submission_id=submission_id, job_dir=job_dir)
|
||||
if extra_overrides:
|
||||
built = BuiltCommand(argv=[*built.argv, *list(extra_overrides)])
|
||||
entrypoint_argv = [
|
||||
"python3",
|
||||
"-m",
|
||||
"argus.ray.driver_entrypoint",
|
||||
"--job-dir",
|
||||
job_dir,
|
||||
*built.argv,
|
||||
]
|
||||
entrypoint = " ".join(shlex.quote(x) for x in entrypoint_argv)
|
||||
|
||||
runtime_env = self._runtime_env(spec, extra_env_vars=extra_env_vars)
|
||||
|
||||
# Prepare job artifacts directory
|
||||
job_root = Path(job_dir)
|
||||
_mkdir(job_root / "config")
|
||||
_mkdir(job_root / "logs")
|
||||
_mkdir(job_root / "debug")
|
||||
_mkdir(job_root / "checkpoints")
|
||||
|
||||
_write_text(job_root / "config" / "ray_config.yaml", dump_yaml(self.cfg.to_public_dict()))
|
||||
_write_text(job_root / "config" / "jobspec.yaml", dump_yaml(spec.to_public_dict()))
|
||||
_write_json(job_root / "config" / "submit_payload.json", {
|
||||
"submission_id": submission_id,
|
||||
"address": self.cfg.address,
|
||||
"entrypoint": entrypoint,
|
||||
"entrypoint_num_cpus": self.cfg.entrypoint_num_cpus,
|
||||
"entrypoint_resources": self.cfg.entrypoint_resources,
|
||||
"runtime_env": runtime_env,
|
||||
})
|
||||
|
||||
# Pre-submit debug snapshot (ray cluster resources via ray.init)
|
||||
try:
|
||||
ray.init(address="auto", ignore_reinit_error=True, log_to_driver=False)
|
||||
_write_json(job_root / "debug" / "ray_cluster_resources_pre.json", ray.cluster_resources())
|
||||
_write_json(job_root / "debug" / "ray_available_resources_pre.json", ray.available_resources())
|
||||
except Exception as e:
|
||||
_write_text(job_root / "debug" / "ray_resources_pre.error.txt", repr(e) + "\n")
|
||||
|
||||
try:
|
||||
submitted = self.client.submit_job(
|
||||
entrypoint=entrypoint,
|
||||
submission_id=submission_id,
|
||||
runtime_env=runtime_env,
|
||||
entrypoint_num_cpus=self.cfg.entrypoint_num_cpus,
|
||||
entrypoint_resources=self.cfg.entrypoint_resources,
|
||||
)
|
||||
except Exception as e:
|
||||
_write_text(job_root / "logs" / "submit.error.txt", repr(e) + "\n")
|
||||
raise
|
||||
|
||||
_write_text(job_root / "config" / "ray_submission_id.txt", submitted + "\n")
|
||||
|
||||
# Post-submit debug snapshot via SDK
|
||||
try:
|
||||
jobs = self.client.list_jobs()
|
||||
_write_text(
|
||||
job_root / "debug" / "ray_job_list_post.json",
|
||||
json.dumps([_job_details_to_dict(j) for j in jobs], indent=2) + "\n",
|
||||
)
|
||||
except Exception as e:
|
||||
_write_text(job_root / "debug" / "ray_job_list_post.error.txt", repr(e) + "\n")
|
||||
|
||||
if not no_wait:
|
||||
# caller can separately wait; keep submit non-blocking by default in scripts
|
||||
pass
|
||||
|
||||
return submitted
|
||||
|
||||
def submit_advanced(
|
||||
self,
|
||||
spec: AdvancedTaskSpec,
|
||||
no_wait: bool,
|
||||
job_dir: str | None = None,
|
||||
user_id: str | None = None,
|
||||
*,
|
||||
extra_env_vars: dict[str, str] | None = None,
|
||||
) -> str:
|
||||
submission_id = spec.submission_id or f"mvp11_{spec.workload}_{_ts()}_{os.getpid()}"
|
||||
job_dir = job_dir or self._job_dir(submission_id)
|
||||
|
||||
built = build_advanced_argv(spec.command)
|
||||
entrypoint_argv = [
|
||||
"python3",
|
||||
"-m",
|
||||
"argus.ray.driver_entrypoint",
|
||||
"--job-dir",
|
||||
job_dir,
|
||||
*built.argv,
|
||||
]
|
||||
entrypoint = " ".join(shlex.quote(x) for x in entrypoint_argv)
|
||||
|
||||
extra_pythonpath: list[str] = []
|
||||
if user_id:
|
||||
extra_pythonpath.append(f"{self.cfg.shared_root}/users/{user_id}/code")
|
||||
|
||||
runtime_env = self._runtime_env_for(
|
||||
code_path=self.cfg.verl_code_path,
|
||||
workload=spec.workload,
|
||||
extra_pythonpath=extra_pythonpath,
|
||||
extra_env_vars=extra_env_vars,
|
||||
)
|
||||
|
||||
# Prepare job artifacts directory
|
||||
job_root = Path(job_dir)
|
||||
_mkdir(job_root / "config")
|
||||
_mkdir(job_root / "logs")
|
||||
_mkdir(job_root / "debug")
|
||||
_mkdir(job_root / "checkpoints")
|
||||
|
||||
_write_text(job_root / "config" / "ray_config.yaml", dump_yaml(self.cfg.to_public_dict()))
|
||||
_write_text(job_root / "config" / "taskspec.yaml", dump_yaml(spec.to_public_dict()))
|
||||
_write_json(
|
||||
job_root / "config" / "submit_payload.json",
|
||||
{
|
||||
"submission_id": submission_id,
|
||||
"address": self.cfg.address,
|
||||
"entrypoint": entrypoint,
|
||||
"entrypoint_num_cpus": self.cfg.entrypoint_num_cpus,
|
||||
"entrypoint_resources": self.cfg.entrypoint_resources,
|
||||
"runtime_env": runtime_env,
|
||||
},
|
||||
)
|
||||
|
||||
# Pre-submit debug snapshot (ray cluster resources via ray.init)
|
||||
try:
|
||||
ray.init(address="auto", ignore_reinit_error=True, log_to_driver=False)
|
||||
_write_json(job_root / "debug" / "ray_cluster_resources_pre.json", ray.cluster_resources())
|
||||
_write_json(job_root / "debug" / "ray_available_resources_pre.json", ray.available_resources())
|
||||
except Exception as e:
|
||||
_write_text(job_root / "debug" / "ray_resources_pre.error.txt", repr(e) + "\n")
|
||||
|
||||
try:
|
||||
submitted = self.client.submit_job(
|
||||
entrypoint=entrypoint,
|
||||
submission_id=submission_id,
|
||||
runtime_env=runtime_env,
|
||||
entrypoint_num_cpus=self.cfg.entrypoint_num_cpus,
|
||||
entrypoint_resources=self.cfg.entrypoint_resources,
|
||||
)
|
||||
except Exception as e:
|
||||
_write_text(job_root / "logs" / "submit.error.txt", repr(e) + "\n")
|
||||
raise
|
||||
|
||||
_write_text(job_root / "config" / "ray_submission_id.txt", submitted + "\n")
|
||||
|
||||
# Post-submit debug snapshot via SDK
|
||||
try:
|
||||
jobs = self.client.list_jobs()
|
||||
_write_text(
|
||||
job_root / "debug" / "ray_job_list_post.json",
|
||||
json.dumps([_job_details_to_dict(j) for j in jobs], indent=2) + "\n",
|
||||
)
|
||||
except Exception as e:
|
||||
_write_text(job_root / "debug" / "ray_job_list_post.error.txt", repr(e) + "\n")
|
||||
|
||||
if not no_wait:
|
||||
pass
|
||||
|
||||
return submitted
|
||||
|
||||
def status(self, submission_id: str) -> str:
|
||||
return str(self.client.get_job_status(submission_id))
|
||||
|
||||
def stop(self, submission_id: str) -> bool:
|
||||
return bool(self.client.stop_job(submission_id))
|
||||
|
||||
def logs(self, submission_id: str) -> str:
|
||||
return self.client.get_job_logs(submission_id)
|
||||
|
||||
def list(self) -> list[dict[str, Any]]:
|
||||
return [_job_details_to_dict(j) for j in self.client.list_jobs()]
|
||||
|
||||
|
||||
def _job_details_to_dict(obj: Any) -> dict[str, Any]:
|
||||
# Ray uses pydantic models internally, but depending on bundled pydantic version
|
||||
# we might get `.model_dump()` (v2) or `.dict()` (v1).
|
||||
if hasattr(obj, "model_dump"):
|
||||
return obj.model_dump() # type: ignore[no-any-return]
|
||||
if hasattr(obj, "dict"):
|
||||
return obj.dict() # type: ignore[no-any-return]
|
||||
if hasattr(obj, "__dict__"):
|
||||
return dict(obj.__dict__)
|
||||
return {"repr": repr(obj)}
|
||||
166
src/mvp/py/argus/ray/worker_watchdog.py
Normal file
@@ -0,0 +1,166 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from typing import Callable
|
||||
|
||||
from .discovery import HeadRecord, load_head_record
|
||||
|
||||
|
||||
Runner = Callable[[list[str]], int]
|
||||
|
||||
|
||||
def _default_runner(argv: list[str]) -> int:
|
||||
proc = subprocess.run(argv, check=False)
|
||||
return int(proc.returncode)
|
||||
|
||||
|
||||
def _ray_stop_cmd() -> list[str]:
|
||||
return ["ray", "stop", "--force"]
|
||||
|
||||
|
||||
def _ray_start_cmd(*, head_addr: str, node_ip: str | None, resources_json: str, block: bool) -> list[str]:
|
||||
argv = ["ray", "start", f"--address={head_addr}", f"--resources={resources_json}", "--disable-usage-stats"]
|
||||
if node_ip:
|
||||
argv.append(f"--node-ip-address={node_ip}")
|
||||
if block:
|
||||
argv.append("--block")
|
||||
return argv
|
||||
|
||||
|
||||
@dataclass
|
||||
class Watchdog:
|
||||
head_ip_file: str
|
||||
node_ip: str | None
|
||||
resources_json: str
|
||||
poll_s: int
|
||||
runner: Runner
|
||||
ray_start_block: bool
|
||||
|
||||
_current_head_addr: str | None = None
|
||||
_ray_proc: subprocess.Popen[bytes] | None = None
|
||||
|
||||
def _kill_ray_proc(self) -> None:
|
||||
p = self._ray_proc
|
||||
self._ray_proc = None
|
||||
if not p:
|
||||
return
|
||||
try:
|
||||
p.terminate()
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
p.wait(timeout=10)
|
||||
except Exception:
|
||||
try:
|
||||
p.kill()
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
p.wait(timeout=5)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def _restart_ray(self, desired: str) -> None:
|
||||
# Best-effort stop then start.
|
||||
self._kill_ray_proc()
|
||||
self.runner(_ray_stop_cmd())
|
||||
if self.ray_start_block:
|
||||
argv = _ray_start_cmd(
|
||||
head_addr=desired,
|
||||
node_ip=self.node_ip,
|
||||
resources_json=self.resources_json,
|
||||
block=True,
|
||||
)
|
||||
# Use Popen so watchdog can notice exits and reconnect if needed.
|
||||
self._ray_proc = subprocess.Popen(argv, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
|
||||
else:
|
||||
self.runner(
|
||||
_ray_start_cmd(
|
||||
head_addr=desired,
|
||||
node_ip=self.node_ip,
|
||||
resources_json=self.resources_json,
|
||||
block=False,
|
||||
)
|
||||
)
|
||||
self._current_head_addr = desired
|
||||
|
||||
def tick_once(self) -> HeadRecord | None:
|
||||
rec = load_head_record(self.head_ip_file)
|
||||
if rec is None:
|
||||
return None
|
||||
desired = rec.head_addr()
|
||||
if self._current_head_addr != desired:
|
||||
self._restart_ray(desired)
|
||||
return rec
|
||||
|
||||
# If we're in block mode and the ray process died, restart it against the same head.
|
||||
if self.ray_start_block and self._ray_proc is not None:
|
||||
if self._ray_proc.poll() is not None:
|
||||
self._restart_ray(desired)
|
||||
return rec
|
||||
|
||||
def run_forever(self) -> None:
|
||||
while True:
|
||||
try:
|
||||
self.tick_once()
|
||||
except Exception:
|
||||
# keep watchdog alive
|
||||
pass
|
||||
time.sleep(max(1, int(self.poll_s)))
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
ap = argparse.ArgumentParser(description="Stateless Ray worker watchdog (auto-connect + reconnect).")
|
||||
ap.add_argument("--head-ip-file", required=True)
|
||||
ap.add_argument("--node-ip", default=None)
|
||||
ap.add_argument(
|
||||
"--resources-kv",
|
||||
action="append",
|
||||
default=[],
|
||||
help='Repeatable k=v resources, e.g. --resources-kv worker_node=100',
|
||||
)
|
||||
ap.add_argument("--poll-s", type=int, default=5)
|
||||
ap.add_argument("--ray-start-block", action="store_true", help="Run `ray start ... --block` and self-heal on exit.")
|
||||
ap.add_argument("--once", action="store_true", help="Run one tick then exit (for testing/debug).")
|
||||
args = ap.parse_args(argv)
|
||||
|
||||
resources: dict[str, float] = {"worker_node": 100.0}
|
||||
if args.resources_kv:
|
||||
resources = {}
|
||||
for item in args.resources_kv:
|
||||
if "=" not in item:
|
||||
raise SystemExit(f"invalid --resources-kv (expected k=v): {item!r}")
|
||||
k, v = item.split("=", 1)
|
||||
k = k.strip()
|
||||
v = v.strip()
|
||||
if not k:
|
||||
raise SystemExit(f"invalid --resources-kv key: {item!r}")
|
||||
resources[k] = float(v)
|
||||
|
||||
node_ip = args.node_ip
|
||||
if node_ip in (None, ""):
|
||||
node_ip = os.environ.get("ARGUS_NODE_IP") or None
|
||||
|
||||
wd = Watchdog(
|
||||
head_ip_file=args.head_ip_file,
|
||||
node_ip=node_ip,
|
||||
resources_json=json.dumps(resources),
|
||||
poll_s=int(args.poll_s),
|
||||
runner=_default_runner,
|
||||
ray_start_block=bool(args.ray_start_block),
|
||||
)
|
||||
if args.once:
|
||||
wd.tick_once()
|
||||
return 0
|
||||
|
||||
wd.run_forever()
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
21
src/mvp/py/argus/ray/yaml_io.py
Normal file
@@ -0,0 +1,21 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import yaml
|
||||
|
||||
|
||||
def load_yaml(path: str) -> dict[str, Any]:
|
||||
p = Path(path)
|
||||
data = yaml.safe_load(p.read_text(encoding="utf-8"))
|
||||
if data is None:
|
||||
return {}
|
||||
if not isinstance(data, dict):
|
||||
raise ValueError(f"yaml root must be a mapping: {path}")
|
||||
return data
|
||||
|
||||
|
||||
def dump_yaml(data: dict[str, Any]) -> str:
|
||||
return yaml.safe_dump(data, sort_keys=False, allow_unicode=True)
|
||||
|
||||
2
src/mvp/py/argus/service/__init__.py
Normal file
@@ -0,0 +1,2 @@
|
||||
__all__ = []
|
||||
|
||||
76
src/mvp/py/argus/service/advanced_command.py
Normal file
@@ -0,0 +1,76 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
|
||||
|
||||
_RE_PRIVATE_USER = re.compile(r"/private/users/([^/\s\"']+)")
|
||||
_RE_TRAIN_FILES = re.compile(r"data\.train_files=([^\s\\]+)")
|
||||
_RE_VAL_FILES = re.compile(r"data\.val_files=([^\s\\]+)")
|
||||
_RE_REWARD_PATH = re.compile(r"custom_reward_function\.path=([^\s\\]+)")
|
||||
|
||||
|
||||
def _strip_quotes(s: str) -> str:
|
||||
s = s.strip()
|
||||
if len(s) >= 2 and s[0] == s[-1] and s[0] in ("'", '"'):
|
||||
return s[1:-1]
|
||||
return s
|
||||
|
||||
|
||||
def expand_advanced_command(*, user_id: str, command: str) -> str:
|
||||
"""
|
||||
Expand user-facing macros into real container paths.
|
||||
|
||||
Important: `$HOME/common/...` is a SFTPGo virtual folder; the training process
|
||||
cannot see it. We map those paths to real shared mounts.
|
||||
"""
|
||||
|
||||
home = f"/private/users/{user_id}"
|
||||
expanded = command
|
||||
expanded = expanded.replace("$HOME/common/datasets", "/private/datasets")
|
||||
expanded = expanded.replace("$HOME/common/hf", "/private/hf")
|
||||
expanded = expanded.replace("$HOME", home)
|
||||
return expanded
|
||||
|
||||
|
||||
def validate_advanced_command(*, user_id: str, expanded_command: str) -> None:
|
||||
"""
|
||||
Best-effort validation for advanced commands.
|
||||
|
||||
Strong constraints (raise ValueError on failure):
|
||||
- Must include a python3 VERL trainer entrypoint
|
||||
- Must not reference other users' /private/users/<other>/... paths
|
||||
- If train/val files are present, must be under allowed roots
|
||||
- If reward path is present, must be under /private/users/<me>/code/...
|
||||
"""
|
||||
|
||||
# v3.5: advanced commands must still be scoped to VERL entrypoints (not arbitrary shell).
|
||||
# Allow non-trainer utilities like `verl.model_merger`.
|
||||
if "python3" not in expanded_command or "-m verl." not in expanded_command:
|
||||
raise ValueError("command must include 'python3' and '-m verl.'")
|
||||
|
||||
for other in _RE_PRIVATE_USER.findall(expanded_command):
|
||||
if other != user_id:
|
||||
raise ValueError(f"command references another user's path: /private/users/{other}/...")
|
||||
|
||||
allowed_dataset_prefixes = (
|
||||
f"/private/users/{user_id}/datasets/",
|
||||
"/private/datasets/",
|
||||
# Backward compatible with v3.0 paths used by existing templates/users.
|
||||
"/private/common/datasets/",
|
||||
)
|
||||
|
||||
for m in _RE_TRAIN_FILES.finditer(expanded_command):
|
||||
path = _strip_quotes(m.group(1))
|
||||
if not path.startswith(allowed_dataset_prefixes):
|
||||
raise ValueError(f"train_file must be under allowed roots, got: {path}")
|
||||
|
||||
for m in _RE_VAL_FILES.finditer(expanded_command):
|
||||
path = _strip_quotes(m.group(1))
|
||||
if not path.startswith(allowed_dataset_prefixes):
|
||||
raise ValueError(f"val_file must be under allowed roots, got: {path}")
|
||||
|
||||
for m in _RE_REWARD_PATH.finditer(expanded_command):
|
||||
path = _strip_quotes(m.group(1))
|
||||
allow = f"/private/users/{user_id}/code/"
|
||||
if not path.startswith(allow):
|
||||
raise ValueError(f"reward path must be under {allow}, got: {path}")
|
||||
786
src/mvp/py/argus/service/app.py
Normal file
@@ -0,0 +1,786 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import secrets
|
||||
import threading
|
||||
from typing import Any
|
||||
import json
|
||||
from dataclasses import asdict
|
||||
|
||||
import yaml
|
||||
from fastapi import FastAPI, HTTPException, Request, Response
|
||||
|
||||
from argus.core.ids import new_model_key, new_task_id
|
||||
from argus.ray.models import AdvancedTaskSpec, JobSpec, RayConfig, parse_taskspec
|
||||
|
||||
from .advanced_command import expand_advanced_command, validate_advanced_command
|
||||
from .config import V2Config
|
||||
from .db import Db
|
||||
from .janitor import JobsJanitor
|
||||
from .scheduler import Scheduler
|
||||
from .serving_spec import ServingSpec, parse_serving_spec, resolve_serving_spec
|
||||
from .sftpgo import SFTPGoAdminClient, SFTPGoError
|
||||
from .ui import register_ui_routes
|
||||
|
||||
|
||||
def _utc_now_iso() -> str:
|
||||
from datetime import datetime
|
||||
|
||||
return datetime.utcnow().replace(microsecond=0).isoformat() + "Z"
|
||||
|
||||
|
||||
def _load_yaml_file(path: str) -> dict[str, Any]:
|
||||
with open(path, "r", encoding="utf-8") as f:
|
||||
obj = yaml.safe_load(f) or {}
|
||||
if not isinstance(obj, dict):
|
||||
raise ValueError("config yaml must be a mapping")
|
||||
return obj
|
||||
|
||||
|
||||
def create_app(config_path: str) -> FastAPI:
|
||||
root = _load_yaml_file(config_path)
|
||||
ray_cfg = RayConfig.from_dict(root)
|
||||
v2_cfg = V2Config.from_root_dict(root)
|
||||
|
||||
db = Db(v2_cfg.sqlite.db_path)
|
||||
db.init()
|
||||
|
||||
scheduler = Scheduler(db=db, ray_cfg=ray_cfg, v2_cfg=v2_cfg)
|
||||
janitor = JobsJanitor(
|
||||
db=db,
|
||||
user_root=v2_cfg.data.user_root,
|
||||
trash_after_days=v2_cfg.data.retention.jobs_trash_after_days,
|
||||
purge_after_days=v2_cfg.data.retention.jobs_purge_after_days,
|
||||
interval_s=v2_cfg.data.retention.janitor_interval_s,
|
||||
)
|
||||
stop_flag = threading.Event()
|
||||
tool = scheduler.tool
|
||||
|
||||
app = FastAPI(title="mvp-v2", version="2.0")
|
||||
|
||||
def _user_home(user_id: str) -> str:
|
||||
base = v2_cfg.data.user_root.rstrip("/")
|
||||
return f"{base}/{user_id}"
|
||||
|
||||
def _ensure_user_dirs(user_id: str) -> None:
|
||||
home = _user_home(user_id)
|
||||
# Note: create a physical /common directory under each user's home so that
|
||||
# SFTPGo WebClient shows a "common" entry at the root. The actual shared
|
||||
# content is exposed via SFTPGo virtual folders under /common/{datasets,models}.
|
||||
for rel in ("datasets", "models", "code", "jobs", "trash/jobs", "common"):
|
||||
os.makedirs(f"{home}/{rel}", exist_ok=True)
|
||||
|
||||
def _sftpgo_enabled() -> bool:
|
||||
return bool(v2_cfg.data.sftpgo.enabled)
|
||||
|
||||
def _sftpgo_client() -> SFTPGoAdminClient:
|
||||
cfg = v2_cfg.data.sftpgo
|
||||
if not cfg.admin_api_base:
|
||||
raise HTTPException(status_code=500, detail="sftpgo enabled but data.sftpgo.admin_api_base is empty")
|
||||
pw = os.environ.get(cfg.admin_password_env, "")
|
||||
if not pw:
|
||||
raise HTTPException(status_code=500, detail=f"missing env: {cfg.admin_password_env}")
|
||||
shared_root = ray_cfg.shared_root.rstrip("/")
|
||||
return SFTPGoAdminClient(
|
||||
admin_api_base=str(cfg.admin_api_base),
|
||||
admin_user=str(cfg.admin_user),
|
||||
admin_password=pw,
|
||||
common_root=f"{shared_root}/common",
|
||||
)
|
||||
|
||||
def _serving_enabled() -> bool:
|
||||
return bool(v2_cfg.serving.enabled)
|
||||
|
||||
def _openai_base_url_for_model(req: Request, *, model_key: str) -> str:
|
||||
# Prefer forwarded headers if present; otherwise fall back to Host.
|
||||
host = req.headers.get("x-forwarded-host") or req.headers.get("host") or req.url.hostname or "127.0.0.1"
|
||||
hostname = host
|
||||
if hostname.startswith("[") and "]" in hostname:
|
||||
hostname = hostname.split("]")[0] + "]"
|
||||
else:
|
||||
hostname = hostname.split(":")[0]
|
||||
scheme = req.headers.get("x-forwarded-proto") or req.url.scheme or "http"
|
||||
port = int(v2_cfg.serving.serve.http_port)
|
||||
return f"{scheme}://{hostname}:{port}/serve/{model_key}/v1"
|
||||
|
||||
def _dump_yaml(obj: Any) -> str:
|
||||
return yaml.safe_dump(obj, sort_keys=False)
|
||||
|
||||
def _serving_spec_to_dict(spec: ServingSpec) -> dict[str, Any]:
|
||||
return {
|
||||
"model_id": spec.model_id,
|
||||
"model_source": spec.model_source,
|
||||
"num_replicas": int(spec.num_replicas),
|
||||
"gpus_per_replica": int(spec.gpus_per_replica),
|
||||
"engine_kwargs": spec.engine_kwargs,
|
||||
}
|
||||
|
||||
def _serve_model_public(row: dict[str, Any], *, req: Request) -> dict[str, Any]:
|
||||
num_replicas = int(row.get("num_replicas") or 0)
|
||||
gpus_per_replica = int(row.get("gpus_per_replica") or 0)
|
||||
total_gpus = num_replicas * gpus_per_replica
|
||||
model_id = str(row.get("model_id") or "")
|
||||
model_key = str(row.get("model_key") or "")
|
||||
return {
|
||||
"model_key": model_key,
|
||||
"user_id": str(row.get("user_id") or ""),
|
||||
"model_id": model_id,
|
||||
"model_id_suffix": str(row.get("model_id_suffix") or ""),
|
||||
"model_id_prefix": str(row.get("model_id_prefix") or ""),
|
||||
"model_source": str(row.get("model_source") or ""),
|
||||
"num_replicas": num_replicas,
|
||||
"gpus_per_replica": gpus_per_replica,
|
||||
"total_gpus": total_gpus,
|
||||
"state": str(row.get("state") or ""),
|
||||
"error_summary": row.get("error_summary"),
|
||||
"created_at": str(row.get("created_at") or ""),
|
||||
"updated_at": str(row.get("updated_at") or ""),
|
||||
"deleted_at": row.get("deleted_at"),
|
||||
"endpoint": {
|
||||
"openai_base_url": _openai_base_url_for_model(req, model_key=model_key),
|
||||
"model": model_id,
|
||||
},
|
||||
}
|
||||
|
||||
def _auth(req: Request) -> dict[str, Any]:
|
||||
token_env = v2_cfg.auth.token_env
|
||||
admin_token = os.environ.get(token_env, "")
|
||||
if not admin_token:
|
||||
raise HTTPException(status_code=500, detail=f"missing token env: {token_env}")
|
||||
|
||||
auth = req.headers.get("authorization") or ""
|
||||
if not auth.startswith("Bearer "):
|
||||
raise HTTPException(status_code=401, detail="missing bearer token")
|
||||
got = auth.removeprefix("Bearer ").strip()
|
||||
|
||||
if got == admin_token:
|
||||
return {"user_id": "admin", "is_admin": True}
|
||||
|
||||
info = db.resolve_token(got)
|
||||
if not info:
|
||||
raise HTTPException(status_code=401, detail="invalid token")
|
||||
if info.get("state") != "ACTIVE":
|
||||
raise HTTPException(status_code=403, detail="user disabled")
|
||||
return {"user_id": info["user_id"], "is_admin": False}
|
||||
|
||||
def _require_admin(req: Request) -> dict[str, Any]:
|
||||
subject = _auth(req)
|
||||
if not subject.get("is_admin"):
|
||||
raise HTTPException(status_code=403, detail="admin required")
|
||||
return subject
|
||||
|
||||
@app.on_event("startup")
|
||||
def _startup() -> None:
|
||||
t = threading.Thread(target=scheduler.run_forever, args=(stop_flag,), daemon=True)
|
||||
t.start()
|
||||
if int(janitor.interval_s) > 0:
|
||||
tj = threading.Thread(target=janitor.run_forever, args=(stop_flag,), daemon=True)
|
||||
tj.start()
|
||||
|
||||
@app.on_event("shutdown")
|
||||
def _shutdown() -> None:
|
||||
stop_flag.set()
|
||||
|
||||
@app.post("/api/v2/users")
|
||||
async def create_user(req: Request) -> dict[str, Any]:
|
||||
_require_admin(req)
|
||||
obj = await req.json()
|
||||
if not isinstance(obj, dict):
|
||||
raise HTTPException(status_code=400, detail="body must be a JSON object")
|
||||
user_id = str(obj.get("user_id") or "").strip()
|
||||
if not user_id:
|
||||
raise HTTPException(status_code=400, detail="missing user_id")
|
||||
display_name = obj.get("display_name")
|
||||
try:
|
||||
row = db.create_user(user_id=user_id, display_name=str(display_name) if display_name is not None else None)
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=409, detail=f"user create failed: {e!r}")
|
||||
|
||||
# v3.0: create user dir structure (datasets/models/code/jobs/trash).
|
||||
try:
|
||||
_ensure_user_dirs(user_id)
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=f"failed to create user dirs: {e!r}")
|
||||
|
||||
out: dict[str, Any] = {"user_id": row.get("user_id", user_id), "state": row.get("state", "ACTIVE")}
|
||||
|
||||
# v3.0: scheme A (password) SFTPGo integration.
|
||||
if _sftpgo_enabled():
|
||||
pw = secrets.token_urlsafe(12)
|
||||
try:
|
||||
_sftpgo_client().create_user(username=user_id, password=pw, home_dir=_user_home(user_id))
|
||||
except SFTPGoError as e:
|
||||
# Make create user idempotent for retries:
|
||||
# If the user already exists in SFTPGo, reset password and enable it.
|
||||
if "http error: 409" in str(e):
|
||||
try:
|
||||
_sftpgo_client().reset_password(username=user_id, new_password=pw, home_dir=_user_home(user_id))
|
||||
_sftpgo_client().enable_user(username=user_id, home_dir=_user_home(user_id))
|
||||
except Exception as e2:
|
||||
raise HTTPException(status_code=502, detail=f"sftpgo upsert user failed: {e2!r}")
|
||||
else:
|
||||
raise HTTPException(status_code=502, detail=f"sftpgo create user failed: {e!r}")
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=502, detail=f"sftpgo create user failed: {e!r}")
|
||||
out["sftp"] = {
|
||||
"username": user_id,
|
||||
"password": pw, # one-time return to admin; do not persist plaintext
|
||||
"host": v2_cfg.data.sftpgo.host,
|
||||
"port": int(v2_cfg.data.sftpgo.sftp_port),
|
||||
}
|
||||
return out
|
||||
|
||||
@app.post("/api/v2/users/{user_id}/tokens")
|
||||
async def issue_token(user_id: str, req: Request) -> dict[str, Any]:
|
||||
_require_admin(req)
|
||||
u = db.get_user(user_id)
|
||||
if not u:
|
||||
raise HTTPException(status_code=404, detail="user not found")
|
||||
token = db.issue_token(user_id=user_id)
|
||||
return {"user_id": user_id, "token": token}
|
||||
|
||||
@app.get("/api/v2/users")
|
||||
async def list_users(req: Request, limit: int = 200) -> dict[str, Any]:
|
||||
_require_admin(req)
|
||||
lim = max(1, min(int(limit), 1000))
|
||||
return {"users": db.list_users(limit=lim), "limit": lim}
|
||||
|
||||
@app.post("/api/v2/users/{user_id}:disable")
|
||||
async def disable_user(user_id: str, req: Request) -> dict[str, Any]:
|
||||
_require_admin(req)
|
||||
u = db.get_user(user_id)
|
||||
if not u:
|
||||
raise HTTPException(status_code=404, detail="user not found")
|
||||
db.disable_user(user_id=user_id)
|
||||
if _sftpgo_enabled():
|
||||
try:
|
||||
_sftpgo_client().disable_user(username=user_id, home_dir=_user_home(user_id))
|
||||
except Exception:
|
||||
# Best-effort: DB state is source of truth for API auth; SFTPGo sync can lag.
|
||||
pass
|
||||
return {"user_id": user_id, "state": "DISABLED"}
|
||||
|
||||
@app.post("/api/v2/users/{user_id}/sftp:reset_password")
|
||||
async def reset_sftp_password(user_id: str, req: Request) -> dict[str, Any]:
|
||||
_require_admin(req)
|
||||
u = db.get_user(user_id)
|
||||
if not u:
|
||||
raise HTTPException(status_code=404, detail="user not found")
|
||||
if not _sftpgo_enabled():
|
||||
raise HTTPException(status_code=400, detail="sftpgo is not enabled")
|
||||
pw = secrets.token_urlsafe(12)
|
||||
try:
|
||||
_sftpgo_client().reset_password(username=user_id, new_password=pw, home_dir=_user_home(user_id))
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=502, detail=f"sftpgo reset password failed: {e!r}")
|
||||
return {"user_id": user_id, "password": pw}
|
||||
|
||||
@app.post("/api/v2/me/sftp:reset_password")
|
||||
async def reset_my_sftp_password(req: Request) -> dict[str, Any]:
|
||||
"""
|
||||
v3.0 WebUI: allow a user to rotate their own SFTPGo password and copy it.
|
||||
Note: SFTPGo does not allow reading the current password, so this endpoint returns a new one-time password.
|
||||
"""
|
||||
subject = _auth(req)
|
||||
user_id = str(subject["user_id"])
|
||||
u = db.get_user(user_id)
|
||||
if not u:
|
||||
raise HTTPException(status_code=404, detail="user not found")
|
||||
if not _sftpgo_enabled():
|
||||
raise HTTPException(status_code=400, detail="sftpgo is not enabled")
|
||||
pw = secrets.token_urlsafe(12)
|
||||
try:
|
||||
_sftpgo_client().reset_password(username=user_id, new_password=pw, home_dir=_user_home(user_id))
|
||||
_sftpgo_client().enable_user(username=user_id, home_dir=_user_home(user_id))
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=502, detail=f"sftpgo reset password failed: {e!r}")
|
||||
return {"user_id": user_id, "password": pw}
|
||||
|
||||
@app.get("/api/v2/me")
|
||||
async def me(req: Request) -> dict[str, Any]:
|
||||
subject = _auth(req)
|
||||
user_id = str(subject["user_id"])
|
||||
try:
|
||||
_ensure_user_dirs(user_id)
|
||||
except Exception:
|
||||
# Best-effort: user may still be able to use API even if FS init fails.
|
||||
pass
|
||||
# Best-effort: reconcile SFTPGo user to include /common read-only mounts.
|
||||
# This makes the virtual folders visible without requiring a password reset.
|
||||
if _sftpgo_enabled() and not subject.get("is_admin"):
|
||||
try:
|
||||
_sftpgo_client().enable_user(username=user_id, home_dir=_user_home(user_id))
|
||||
except Exception:
|
||||
pass
|
||||
home = _user_home(user_id)
|
||||
out: dict[str, Any] = {
|
||||
"user_id": user_id,
|
||||
"is_admin": bool(subject.get("is_admin")),
|
||||
"paths": {
|
||||
"home": home,
|
||||
"datasets": f"{home}/datasets",
|
||||
"models": f"{home}/models",
|
||||
"code": f"{home}/code",
|
||||
"jobs": f"{home}/jobs",
|
||||
"trash_jobs": f"{home}/trash/jobs",
|
||||
},
|
||||
"retention": {
|
||||
"jobs_trash_after_days": int(v2_cfg.data.retention.jobs_trash_after_days),
|
||||
"jobs_purge_after_days": int(v2_cfg.data.retention.jobs_purge_after_days),
|
||||
},
|
||||
}
|
||||
if v2_cfg.tracking.wandb.enabled:
|
||||
suffix = v2_cfg.tracking.wandb.project_suffix or "_project"
|
||||
out["wandb"] = {
|
||||
"enabled": True,
|
||||
# Training uses base_url; UI can also open the host port directly.
|
||||
"base_url": v2_cfg.tracking.wandb.base_url,
|
||||
"project_name": f"{user_id}{suffix}",
|
||||
}
|
||||
if _sftpgo_enabled():
|
||||
out["sftp"] = {
|
||||
"host": v2_cfg.data.sftpgo.host,
|
||||
"port": int(v2_cfg.data.sftpgo.sftp_port),
|
||||
"username": user_id,
|
||||
}
|
||||
return out
|
||||
|
||||
@app.get("/api/v2/tasks")
|
||||
async def list_tasks(req: Request, limit: int = 200, offset: int = 0, states: str | None = None) -> dict[str, Any]:
|
||||
subject = _auth(req)
|
||||
lim = max(1, min(int(limit), 1000))
|
||||
off = max(0, int(offset))
|
||||
state_list: list[str] | None = None
|
||||
if states:
|
||||
raw = [s.strip() for s in str(states).split(",") if s.strip()]
|
||||
# Keep it strict to avoid typos silently returning empty results.
|
||||
allowed = {
|
||||
"QUEUED",
|
||||
"PENDING_RESOURCES",
|
||||
"SUBMITTING",
|
||||
"SUBMITTED",
|
||||
"RUNNING",
|
||||
"SUCCEEDED",
|
||||
"FAILED",
|
||||
"CANCELED",
|
||||
}
|
||||
for s in raw:
|
||||
if s not in allowed:
|
||||
raise HTTPException(status_code=400, detail=f"invalid state: {s}")
|
||||
state_list = raw or None
|
||||
if subject.get("is_admin"):
|
||||
tasks = db.list_tasks(user_id=None, states=state_list, limit=lim, offset=off)
|
||||
else:
|
||||
tasks = db.list_tasks(user_id=str(subject["user_id"]), states=state_list, limit=lim, offset=off)
|
||||
return {"tasks": tasks, "limit": lim, "offset": off, "has_more": bool(len(tasks) == lim)}
|
||||
|
||||
@app.post("/api/v2/tasks")
|
||||
async def submit_task(req: Request) -> dict[str, Any]:
|
||||
subject = _auth(req)
|
||||
body = (await req.body()).decode("utf-8")
|
||||
try:
|
||||
obj = yaml.safe_load(body) or {}
|
||||
except Exception as e:
|
||||
# yaml.YAMLError and similar should be a user-facing 400, not a 500.
|
||||
raise HTTPException(status_code=400, detail=f"invalid YAML: {e!r}")
|
||||
if not isinstance(obj, dict):
|
||||
raise HTTPException(status_code=400, detail="jobspec must be a YAML mapping")
|
||||
|
||||
try:
|
||||
spec = parse_taskspec(obj)
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=400, detail=f"invalid jobspec: {e!r}")
|
||||
|
||||
root = ray_cfg.shared_root.rstrip("/")
|
||||
user_id = str(subject["user_id"]).strip()
|
||||
task_id = new_task_id(str(spec.workload), user_id=user_id)
|
||||
|
||||
# v3.5: Advanced TaskSpec (custom command)
|
||||
if isinstance(spec, AdvancedTaskSpec):
|
||||
expanded = expand_advanced_command(user_id=user_id, command=spec.command)
|
||||
try:
|
||||
validate_advanced_command(user_id=user_id, expanded_command=expanded)
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=400, detail=f"invalid command: {e!r}")
|
||||
|
||||
stored = spec.to_public_dict()
|
||||
stored["command"] = expanded
|
||||
stored_yaml = yaml.safe_dump(stored, sort_keys=False)
|
||||
|
||||
db.create_task_v25(
|
||||
task_id=task_id,
|
||||
user_id=user_id,
|
||||
workload=spec.workload,
|
||||
jobspec_yaml=stored_yaml,
|
||||
nnodes=spec.nnodes,
|
||||
n_gpus_per_node=spec.n_gpus_per_node,
|
||||
)
|
||||
return {"task_id": task_id, "state": "QUEUED"}
|
||||
|
||||
# v3.0 path policy (Basic JobSpec):
|
||||
# - code_path: only allow /private/common/...
|
||||
# - train/val: allow /private/common/datasets/... OR /private/users/<me>/datasets/...
|
||||
# - model_id: if it looks like a local path (/private/...), allow only models dirs:
|
||||
# /private/common/models/... OR /private/users/<me>/models/...
|
||||
common_prefix = f"{root}/common/"
|
||||
user_prefix = f"{root}/users/{user_id}/"
|
||||
|
||||
common_datasets_prefix = f"{common_prefix}datasets/"
|
||||
user_datasets_prefix = f"{user_prefix}datasets/"
|
||||
common_models_prefix = f"{common_prefix}models/"
|
||||
user_models_prefix = f"{user_prefix}models/"
|
||||
|
||||
if not str(spec.code_path).startswith(common_prefix):
|
||||
raise HTTPException(status_code=400, detail=f"code_path must start with {common_prefix}")
|
||||
|
||||
for k, v in (("train_file", spec.train_file), ("val_file", spec.val_file)):
|
||||
if v is None:
|
||||
continue
|
||||
sv = str(v)
|
||||
if not (sv.startswith(common_datasets_prefix) or sv.startswith(user_datasets_prefix)):
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=f"{k} must start with {common_datasets_prefix} or {user_datasets_prefix}",
|
||||
)
|
||||
|
||||
model_id = str(spec.model_id)
|
||||
if model_id.startswith(f"{root}/"):
|
||||
if not (model_id.startswith(common_models_prefix) or model_id.startswith(user_models_prefix)):
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=f"model_id local path must start with {common_models_prefix} or {user_models_prefix}",
|
||||
)
|
||||
|
||||
try:
|
||||
db.create_task_v25(
|
||||
task_id=task_id,
|
||||
user_id=user_id,
|
||||
workload=spec.workload,
|
||||
jobspec_yaml=body,
|
||||
nnodes=spec.nnodes,
|
||||
n_gpus_per_node=spec.n_gpus_per_node,
|
||||
)
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=f"db create task failed: {e!r}")
|
||||
return {"task_id": task_id, "state": "QUEUED"}
|
||||
|
||||
@app.get("/api/v2/tasks/{task_id}")
|
||||
async def get_task(task_id: str, req: Request) -> dict[str, Any]:
|
||||
subject = _auth(req)
|
||||
row = db.get_task(task_id)
|
||||
if not row:
|
||||
raise HTTPException(status_code=404, detail="task not found")
|
||||
if not subject.get("is_admin"):
|
||||
if str(row.get("user_id") or "") != str(subject["user_id"]):
|
||||
raise HTTPException(status_code=404, detail="task not found")
|
||||
attempts = db.list_attempts(task_id)
|
||||
latest_attempt = attempts[-1] if attempts else None
|
||||
desired = {
|
||||
"nnodes": int(row["nnodes"]),
|
||||
"n_gpus_per_node": int(row["n_gpus_per_node"]),
|
||||
"total_gpus": int(row["nnodes"]) * int(row["n_gpus_per_node"]),
|
||||
}
|
||||
out: dict[str, Any] = {
|
||||
"task_id": row["task_id"],
|
||||
"workload": row["workload"],
|
||||
"state": row["state"],
|
||||
"created_at": row.get("created_at"),
|
||||
"updated_at": row.get("updated_at"),
|
||||
"desired_resources": desired,
|
||||
"error_summary": row.get("error_summary"),
|
||||
}
|
||||
if latest_attempt:
|
||||
latest_out: dict[str, Any] = {
|
||||
"attempt_no": latest_attempt["attempt_no"],
|
||||
"ray_submission_id": latest_attempt["ray_submission_id"],
|
||||
"ray_status": latest_attempt.get("ray_status"),
|
||||
"start_time": latest_attempt.get("start_time"),
|
||||
"end_time": latest_attempt.get("end_time"),
|
||||
"failure_kind": latest_attempt.get("failure_kind"),
|
||||
"message": latest_attempt.get("message"),
|
||||
}
|
||||
if v2_cfg.tracking.wandb.enabled:
|
||||
suffix = v2_cfg.tracking.wandb.project_suffix or "_project"
|
||||
# For v3.6: shared W&B account, user-level projects.
|
||||
# Use attempt's ray_submission_id as run name for stable mapping.
|
||||
latest_out["wandb"] = {
|
||||
"base_url": v2_cfg.tracking.wandb.base_url,
|
||||
"project_name": f"{str(row.get('user_id') or subject.get('user_id'))}{suffix}",
|
||||
"run_name": str(latest_attempt.get("ray_submission_id") or ""),
|
||||
}
|
||||
out["latest_attempt"] = latest_out
|
||||
return out
|
||||
|
||||
@app.get("/api/v2/tasks/{task_id}/spec")
|
||||
async def get_task_spec(task_id: str, req: Request) -> Response:
|
||||
"""
|
||||
Returns the TaskSpec YAML for this task.
|
||||
|
||||
Note: We prefer returning a "resolved" TaskSpec (includes default values) so the
|
||||
UI can show a complete spec even if the user submitted a minimal YAML.
|
||||
"""
|
||||
subject = _auth(req)
|
||||
row = db.get_task(task_id)
|
||||
if not row:
|
||||
raise HTTPException(status_code=404, detail="task not found")
|
||||
if not subject.get("is_admin"):
|
||||
if str(row.get("user_id") or "") != str(subject["user_id"]):
|
||||
raise HTTPException(status_code=404, detail="task not found")
|
||||
raw = str(row.get("jobspec_yaml") or "")
|
||||
try:
|
||||
obj = yaml.safe_load(raw) or {}
|
||||
if not isinstance(obj, dict):
|
||||
raise ValueError("jobspec must be a YAML mapping")
|
||||
spec = parse_taskspec(obj)
|
||||
resolved = spec.to_public_dict()
|
||||
attempts = db.list_attempts(task_id)
|
||||
if attempts:
|
||||
resolved["submission_id"] = str(attempts[-1].get("ray_submission_id") or resolved.get("submission_id") or "")
|
||||
text = yaml.safe_dump(resolved, sort_keys=False)
|
||||
except Exception:
|
||||
# Backward compatibility / best-effort: if old rows contain invalid YAML,
|
||||
# fall back to the original content.
|
||||
text = raw
|
||||
return Response(content=text, media_type="text/plain")
|
||||
|
||||
@app.get("/api/v2/tasks/{task_id}/attempts")
|
||||
async def get_attempts(task_id: str, req: Request) -> dict[str, Any]:
|
||||
subject = _auth(req)
|
||||
row = db.get_task(task_id)
|
||||
if not row:
|
||||
raise HTTPException(status_code=404, detail="task not found")
|
||||
if not subject.get("is_admin"):
|
||||
if str(row.get("user_id") or "") != str(subject["user_id"]):
|
||||
raise HTTPException(status_code=404, detail="task not found")
|
||||
return {"task_id": task_id, "attempts": db.list_attempts(task_id)}
|
||||
|
||||
@app.post("/api/v2/tasks/{task_id}:cancel")
|
||||
async def cancel(task_id: str, req: Request) -> dict[str, Any]:
|
||||
subject = _auth(req)
|
||||
row = db.get_task(task_id)
|
||||
if not row:
|
||||
raise HTTPException(status_code=404, detail="task not found")
|
||||
if not subject.get("is_admin"):
|
||||
if str(row.get("user_id") or "") != str(subject["user_id"]):
|
||||
raise HTTPException(status_code=404, detail="task not found")
|
||||
|
||||
state = str(row["state"])
|
||||
if state in ("SUCCEEDED", "FAILED", "CANCELED"):
|
||||
raise HTTPException(status_code=409, detail=f"task already terminal: {state}")
|
||||
|
||||
attempts = db.list_attempts(task_id)
|
||||
if attempts:
|
||||
ray_sid = str(attempts[-1]["ray_submission_id"])
|
||||
try:
|
||||
tool.stop(ray_sid)
|
||||
except Exception:
|
||||
pass
|
||||
# Mark attempt as canceled on the service side so that API doesn't keep reporting RUNNING.
|
||||
# Ray stop is async; we deliberately reflect the user's intent here.
|
||||
db.update_attempt(
|
||||
task_id=task_id,
|
||||
attempt_no=int(attempts[-1]["attempt_no"]),
|
||||
ray_status="STOPPED",
|
||||
failure_kind="CANCELED",
|
||||
message="Canceled by user via API (Ray stop requested).",
|
||||
end_time=_utc_now_iso(),
|
||||
)
|
||||
|
||||
db.set_task_state(task_id=task_id, state="CANCELED", event_type="CANCELED")
|
||||
return {"task_id": task_id, "state": "CANCELED"}
|
||||
|
||||
@app.get("/api/v2/tasks/{task_id}/logs")
|
||||
async def logs(task_id: str, req: Request, tail: int = 2000, attempt: str = "latest") -> Response:
|
||||
subject = _auth(req)
|
||||
row = db.get_task(task_id)
|
||||
if not row:
|
||||
raise HTTPException(status_code=404, detail="task not found")
|
||||
if not subject.get("is_admin"):
|
||||
if str(row.get("user_id") or "") != str(subject["user_id"]):
|
||||
raise HTTPException(status_code=404, detail="task not found")
|
||||
attempts = db.list_attempts(task_id)
|
||||
if not attempts:
|
||||
raise HTTPException(status_code=404, detail="no attempts yet")
|
||||
a = attempts[-1] if attempt == "latest" else None
|
||||
if a is None:
|
||||
raise HTTPException(status_code=400, detail="only attempt=latest supported in v2.0")
|
||||
ray_sid = str(a["ray_submission_id"])
|
||||
text = tool.logs(ray_sid)
|
||||
if tail and tail > 0:
|
||||
lines = text.splitlines()
|
||||
text = "\n".join(lines[-tail:]) + ("\n" if lines else "")
|
||||
return Response(content=text, media_type="text/plain")
|
||||
|
||||
@app.get("/api/v2/queue")
|
||||
async def queue(req: Request) -> dict[str, Any]:
|
||||
subject = _auth(req)
|
||||
if subject.get("is_admin"):
|
||||
return db.list_queue()
|
||||
return db.list_queue(user_id=str(subject["user_id"]))
|
||||
|
||||
# v3.8: Model serving (Ray Serve LLM) management APIs.
|
||||
@app.post("/api/v2/serve/models")
|
||||
async def create_serve_model(req: Request) -> dict[str, Any]:
|
||||
subject = _auth(req)
|
||||
if not _serving_enabled():
|
||||
raise HTTPException(status_code=400, detail="serving is not enabled")
|
||||
|
||||
body = (await req.body()).decode("utf-8")
|
||||
try:
|
||||
obj = yaml.safe_load(body) or {}
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=400, detail=f"invalid YAML: {e!r}")
|
||||
if not isinstance(obj, dict):
|
||||
raise HTTPException(status_code=400, detail="serving spec must be a YAML mapping")
|
||||
|
||||
user_id = str(subject["user_id"]).strip()
|
||||
try:
|
||||
spec = parse_serving_spec(obj)
|
||||
resolved = resolve_serving_spec(spec=spec, user_id=user_id)
|
||||
except PermissionError as e:
|
||||
raise HTTPException(status_code=403, detail=str(e))
|
||||
except ValueError as e:
|
||||
msg = str(e)
|
||||
code = 422 if ("num_replicas" in msg or "gpus_per_replica" in msg) else 400
|
||||
raise HTTPException(status_code=code, detail=f"invalid serving spec: {e!r}")
|
||||
|
||||
model_key = new_model_key(user_id=user_id)
|
||||
try:
|
||||
engine_kwargs_json = json.dumps(resolved.engine_kwargs, sort_keys=True) if resolved.engine_kwargs is not None else None
|
||||
except TypeError as e:
|
||||
raise HTTPException(status_code=400, detail=f"engine_kwargs must be JSON-serializable: {e!r}")
|
||||
|
||||
spec_yaml = _dump_yaml(_serving_spec_to_dict(spec))
|
||||
resolved_spec_yaml = _dump_yaml(asdict(resolved))
|
||||
|
||||
db.create_serve_model(
|
||||
model_key=model_key,
|
||||
user_id=user_id,
|
||||
model_id_suffix=resolved.model_id_suffix,
|
||||
model_id_prefix=resolved.model_id_prefix,
|
||||
model_id=resolved.model_id,
|
||||
model_source=resolved.model_source,
|
||||
num_replicas=resolved.num_replicas,
|
||||
gpus_per_replica=resolved.gpus_per_replica,
|
||||
engine_kwargs_json=engine_kwargs_json,
|
||||
spec_yaml=spec_yaml,
|
||||
resolved_spec_yaml=resolved_spec_yaml,
|
||||
)
|
||||
return {"model_key": model_key, "state": "QUEUED"}
|
||||
|
||||
@app.get("/api/v2/serve/models")
|
||||
async def list_serve_models(req: Request, limit: int = 200, offset: int = 0, include_deleted: int = 0) -> dict[str, Any]:
|
||||
subject = _auth(req)
|
||||
if not _serving_enabled():
|
||||
raise HTTPException(status_code=400, detail="serving is not enabled")
|
||||
|
||||
lim = max(1, min(int(limit), 1000))
|
||||
off = max(0, int(offset))
|
||||
inc = bool(int(include_deleted))
|
||||
user_id = str(subject["user_id"])
|
||||
|
||||
items = db.list_serve_models(user_id=user_id, include_deleted=inc, limit=lim, offset=off)
|
||||
out = [_serve_model_public(i, req=req) for i in items]
|
||||
return {
|
||||
"items": out,
|
||||
# Per-model route_prefix; see each item.endpoint.openai_base_url.
|
||||
"openai_base_url": None,
|
||||
"limit": lim,
|
||||
"offset": off,
|
||||
"has_more": bool(len(items) == lim),
|
||||
}
|
||||
|
||||
@app.get("/api/v2/serve/models/{model_key}")
|
||||
async def get_serve_model(model_key: str, req: Request) -> dict[str, Any]:
|
||||
subject = _auth(req)
|
||||
if not _serving_enabled():
|
||||
raise HTTPException(status_code=400, detail="serving is not enabled")
|
||||
|
||||
row = db.get_serve_model(model_key)
|
||||
if not row:
|
||||
raise HTTPException(status_code=404, detail="model not found")
|
||||
if not subject.get("is_admin"):
|
||||
if str(row.get("user_id") or "") != str(subject["user_id"]):
|
||||
raise HTTPException(status_code=404, detail="model not found")
|
||||
|
||||
events = db.list_serve_events(model_key, limit=200, offset=0)
|
||||
ev_out = [
|
||||
{
|
||||
"id": int(e.get("id") or 0),
|
||||
"model_key": str(e.get("model_key") or ""),
|
||||
"created_at": str(e.get("ts") or ""),
|
||||
"event_type": str(e.get("event_type") or ""),
|
||||
"payload_json": e.get("payload_json"),
|
||||
}
|
||||
for e in events
|
||||
]
|
||||
return {
|
||||
"model": _serve_model_public(row, req=req),
|
||||
"resolved_spec_yaml": str(row.get("resolved_spec_yaml") or ""),
|
||||
"events": ev_out,
|
||||
"serve_status": None,
|
||||
}
|
||||
|
||||
@app.patch("/api/v2/serve/models/{model_key}")
|
||||
async def patch_serve_model(model_key: str, req: Request) -> dict[str, Any]:
|
||||
subject = _auth(req)
|
||||
if not _serving_enabled():
|
||||
raise HTTPException(status_code=400, detail="serving is not enabled")
|
||||
|
||||
row = db.get_serve_model(model_key)
|
||||
if not row:
|
||||
raise HTTPException(status_code=404, detail="model not found")
|
||||
if not subject.get("is_admin"):
|
||||
if str(row.get("user_id") or "") != str(subject["user_id"]):
|
||||
raise HTTPException(status_code=404, detail="model not found")
|
||||
|
||||
obj = await req.json()
|
||||
if not isinstance(obj, dict):
|
||||
raise HTTPException(status_code=400, detail="body must be a JSON object")
|
||||
if "num_replicas" not in obj:
|
||||
raise HTTPException(status_code=400, detail="missing num_replicas")
|
||||
num_replicas = obj.get("num_replicas")
|
||||
if not isinstance(num_replicas, int) or int(num_replicas) < 1:
|
||||
raise HTTPException(status_code=422, detail="num_replicas must be an integer >= 1")
|
||||
|
||||
db.update_serve_model_num_replicas(model_key=model_key, num_replicas=int(num_replicas))
|
||||
return {"model_key": model_key, "state": "QUEUED"}
|
||||
|
||||
@app.delete("/api/v2/serve/models/{model_key}")
|
||||
async def delete_serve_model(model_key: str, req: Request) -> dict[str, Any]:
|
||||
subject = _auth(req)
|
||||
if not _serving_enabled():
|
||||
raise HTTPException(status_code=400, detail="serving is not enabled")
|
||||
|
||||
row = db.get_serve_model(model_key)
|
||||
if not row:
|
||||
raise HTTPException(status_code=404, detail="model not found")
|
||||
if not subject.get("is_admin"):
|
||||
if str(row.get("user_id") or "") != str(subject["user_id"]):
|
||||
raise HTTPException(status_code=404, detail="model not found")
|
||||
|
||||
db.set_serve_model_state(model_key=model_key, state="DELETING", event_type="SERVE_DELETE_REQUESTED")
|
||||
return {"model_key": model_key, "state": "DELETING"}
|
||||
|
||||
@app.get("/api/v2/serve/status")
|
||||
async def serve_status(req: Request) -> dict[str, Any]:
|
||||
_require_admin(req)
|
||||
if not _serving_enabled():
|
||||
raise HTTPException(status_code=400, detail="serving is not enabled")
|
||||
return {
|
||||
"enabled": True,
|
||||
# Per-model route_prefix; see /api/v2/serve/models items[].endpoint.openai_base_url.
|
||||
"openai_base_url": None,
|
||||
"http_port": int(v2_cfg.serving.serve.http_port),
|
||||
"proxy_location": str(v2_cfg.serving.serve.proxy_location),
|
||||
"accelerator_type": str(v2_cfg.serving.llm.accelerator_type),
|
||||
}
|
||||
|
||||
# v3.0: minimal WebUI (no server-side session; token stored in browser localStorage).
|
||||
register_ui_routes(app)
|
||||
|
||||
return app
|
||||
199
src/mvp/py/argus/service/config.py
Normal file
@@ -0,0 +1,199 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class V2ApiConfig:
|
||||
host: str = "0.0.0.0"
|
||||
port: int = 8080
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class V2AuthConfig:
|
||||
token_env: str = "MVP_INTERNAL_TOKEN"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class V2SqliteConfig:
|
||||
db_path: str
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class V2SchedulerConfig:
|
||||
tick_s: int = 5
|
||||
retry_interval_s: int = 60
|
||||
max_running_tasks: int = 1
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class V2RetentionConfig:
|
||||
jobs_trash_after_days: int = 3
|
||||
jobs_purge_after_days: int = 7
|
||||
janitor_interval_s: int = 3600
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class V2WandbConfig:
|
||||
enabled: bool = False
|
||||
base_url: str = ""
|
||||
api_key_env: str = "WANDB_API_KEY"
|
||||
project_suffix: str = "_project"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class V2TrackingConfig:
|
||||
wandb: V2WandbConfig
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class V2SFTPGoConfig:
|
||||
enabled: bool = False
|
||||
host: str = ""
|
||||
sftp_port: int = 2022
|
||||
admin_api_base: str = ""
|
||||
admin_user: str = "admin"
|
||||
admin_password_env: str = "SFTPGO_ADMIN_PASSWORD"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class V2ServingServeConfig:
|
||||
http_port: int = 8000
|
||||
proxy_location: str = "HeadOnly"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class V2ServingLLMConfig:
|
||||
accelerator_type: str = ""
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class V2ServingConfig:
|
||||
enabled: bool = False
|
||||
serve: V2ServingServeConfig = V2ServingServeConfig()
|
||||
llm: V2ServingLLMConfig = V2ServingLLMConfig()
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class V2DataConfig:
|
||||
user_root: str
|
||||
sftpgo: V2SFTPGoConfig
|
||||
retention: V2RetentionConfig
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class V2Config:
|
||||
api: V2ApiConfig
|
||||
auth: V2AuthConfig
|
||||
sqlite: V2SqliteConfig
|
||||
scheduler: V2SchedulerConfig
|
||||
tracking: V2TrackingConfig
|
||||
data: V2DataConfig
|
||||
serving: V2ServingConfig
|
||||
|
||||
@staticmethod
|
||||
def from_root_dict(root: dict[str, Any]) -> "V2Config":
|
||||
# New format: root.service
|
||||
# Backward compatible: root.v2
|
||||
service = root.get("service")
|
||||
if service is None:
|
||||
service = root.get("v2") or {}
|
||||
if not isinstance(service, dict):
|
||||
raise ValueError("config.service must be a mapping")
|
||||
|
||||
api = service.get("api") or {}
|
||||
auth = service.get("auth") or {}
|
||||
sqlite = service.get("sqlite") or {}
|
||||
scheduler = service.get("scheduler") or {}
|
||||
|
||||
if not isinstance(api, dict) or not isinstance(auth, dict) or not isinstance(sqlite, dict) or not isinstance(scheduler, dict):
|
||||
raise ValueError("config.service.{api,auth,sqlite,scheduler} must be mappings")
|
||||
|
||||
ray = root.get("ray")
|
||||
if isinstance(ray, dict) and ray.get("shared_root"):
|
||||
shared_root = str(ray.get("shared_root"))
|
||||
else:
|
||||
shared_root = str(root.get("shared_root") or "/private")
|
||||
|
||||
tracking = root.get("tracking") or {}
|
||||
if not isinstance(tracking, dict):
|
||||
raise ValueError("config.tracking must be a mapping")
|
||||
wandb = tracking.get("wandb") or {}
|
||||
if not isinstance(wandb, dict):
|
||||
raise ValueError("config.tracking.wandb must be a mapping")
|
||||
|
||||
data = root.get("data") or {}
|
||||
if not isinstance(data, dict):
|
||||
raise ValueError("config.data must be a mapping")
|
||||
sftpgo = data.get("sftpgo") or {}
|
||||
retention = data.get("retention") or {}
|
||||
if not isinstance(sftpgo, dict) or not isinstance(retention, dict):
|
||||
raise ValueError("config.data.{sftpgo,retention} must be mappings")
|
||||
|
||||
serving = root.get("serving") or {}
|
||||
if not isinstance(serving, dict):
|
||||
raise ValueError("config.serving must be a mapping")
|
||||
serving_enabled = bool(serving.get("enabled")) if "enabled" in serving else bool(serving)
|
||||
serving_serve = serving.get("serve") or {}
|
||||
serving_llm = serving.get("llm") or {}
|
||||
if not isinstance(serving_serve, dict) or not isinstance(serving_llm, dict):
|
||||
raise ValueError("config.serving.{serve,llm} must be mappings")
|
||||
|
||||
default_db_path = f"{shared_root}/common/db/mvp.sqlite3"
|
||||
db_path = str(sqlite.get("db_path") or default_db_path)
|
||||
|
||||
user_root = str(data.get("user_root") or f"{shared_root}/users")
|
||||
|
||||
wandb_enabled = bool(wandb.get("enabled") or False)
|
||||
wandb_base_url = str(wandb.get("base_url") or "")
|
||||
if wandb_enabled and not wandb_base_url:
|
||||
raise ValueError("tracking.wandb.base_url is required when tracking.wandb.enabled=true")
|
||||
|
||||
return V2Config(
|
||||
api=V2ApiConfig(
|
||||
host=str(api.get("host") or "0.0.0.0"),
|
||||
port=int(api.get("port") or 8080),
|
||||
),
|
||||
auth=V2AuthConfig(token_env=str(auth.get("token_env") or "MVP_INTERNAL_TOKEN")),
|
||||
sqlite=V2SqliteConfig(db_path=db_path),
|
||||
scheduler=V2SchedulerConfig(
|
||||
tick_s=int(scheduler.get("tick_s") or 5),
|
||||
retry_interval_s=int(scheduler.get("retry_interval_s") or 60),
|
||||
max_running_tasks=int(scheduler.get("max_running_tasks") or 1),
|
||||
),
|
||||
tracking=V2TrackingConfig(
|
||||
wandb=V2WandbConfig(
|
||||
enabled=wandb_enabled,
|
||||
base_url=wandb_base_url,
|
||||
api_key_env=str(wandb.get("api_key_env") or "WANDB_API_KEY"),
|
||||
project_suffix=str(wandb.get("project_suffix") or "_project"),
|
||||
)
|
||||
),
|
||||
data=V2DataConfig(
|
||||
user_root=user_root,
|
||||
sftpgo=V2SFTPGoConfig(
|
||||
enabled=bool(sftpgo.get("enabled") or False),
|
||||
host=str(sftpgo.get("host") or ""),
|
||||
sftp_port=int(sftpgo.get("sftp_port") or 2022),
|
||||
admin_api_base=str(sftpgo.get("admin_api_base") or ""),
|
||||
admin_user=str(sftpgo.get("admin_user") or "admin"),
|
||||
admin_password_env=str(sftpgo.get("admin_password_env") or "SFTPGO_ADMIN_PASSWORD"),
|
||||
),
|
||||
retention=V2RetentionConfig(
|
||||
jobs_trash_after_days=int(retention.get("jobs_trash_after_days") or 3),
|
||||
jobs_purge_after_days=int(retention.get("jobs_purge_after_days") or 7),
|
||||
janitor_interval_s=int(retention.get("janitor_interval_s") or 3600),
|
||||
),
|
||||
),
|
||||
serving=V2ServingConfig(
|
||||
enabled=serving_enabled,
|
||||
serve=V2ServingServeConfig(
|
||||
http_port=int(serving_serve.get("http_port") or 8000),
|
||||
proxy_location=str(serving_serve.get("proxy_location") or "HeadOnly"),
|
||||
),
|
||||
llm=V2ServingLLMConfig(
|
||||
accelerator_type=str(serving_llm.get("accelerator_type") or ""),
|
||||
),
|
||||
),
|
||||
)
|
||||
768
src/mvp/py/argus/service/db.py
Normal file
@@ -0,0 +1,768 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import sqlite3
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
import hashlib
|
||||
import secrets
|
||||
from typing import Any, Iterator
|
||||
|
||||
|
||||
def _utc_now_iso() -> str:
|
||||
# Keep it simple; wall-clock ordering only.
|
||||
import datetime as _dt
|
||||
|
||||
return _dt.datetime.utcnow().replace(microsecond=0).isoformat() + "Z"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Db:
|
||||
db_path: str
|
||||
|
||||
def _hash_token(self, token: str) -> str:
|
||||
# Internal tokens only; store hash to avoid plaintext at rest.
|
||||
return hashlib.sha256(token.encode("utf-8")).hexdigest()
|
||||
|
||||
def _connect(self) -> sqlite3.Connection:
|
||||
os.makedirs(os.path.dirname(self.db_path), exist_ok=True)
|
||||
conn = sqlite3.connect(self.db_path, timeout=30, isolation_level=None)
|
||||
conn.row_factory = sqlite3.Row
|
||||
conn.execute("PRAGMA journal_mode=WAL;")
|
||||
conn.execute("PRAGMA foreign_keys=ON;")
|
||||
return conn
|
||||
|
||||
def init(self) -> None:
|
||||
with self._connect() as conn:
|
||||
conn.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS users (
|
||||
user_id TEXT PRIMARY KEY,
|
||||
display_name TEXT,
|
||||
state TEXT NOT NULL,
|
||||
created_at TEXT NOT NULL,
|
||||
updated_at TEXT NOT NULL
|
||||
)
|
||||
"""
|
||||
)
|
||||
conn.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS api_tokens (
|
||||
token_hash TEXT PRIMARY KEY,
|
||||
user_id TEXT NOT NULL,
|
||||
token_plain TEXT NOT NULL,
|
||||
created_at TEXT NOT NULL,
|
||||
last_used_at TEXT,
|
||||
FOREIGN KEY (user_id) REFERENCES users(user_id) ON DELETE CASCADE
|
||||
)
|
||||
"""
|
||||
)
|
||||
conn.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS tasks (
|
||||
task_id TEXT PRIMARY KEY,
|
||||
workload TEXT NOT NULL,
|
||||
state TEXT NOT NULL,
|
||||
jobspec_yaml TEXT NOT NULL,
|
||||
nnodes INTEGER NOT NULL,
|
||||
n_gpus_per_node INTEGER NOT NULL,
|
||||
created_at TEXT NOT NULL,
|
||||
updated_at TEXT NOT NULL,
|
||||
next_run_at TEXT,
|
||||
error_summary TEXT,
|
||||
latest_attempt_no INTEGER NOT NULL DEFAULT 0
|
||||
)
|
||||
"""
|
||||
)
|
||||
# Best-effort: add user_id column for forward compatibility (v2.5+).
|
||||
# v2.0 tables may exist without it; SQLite supports ADD COLUMN.
|
||||
try:
|
||||
conn.execute("ALTER TABLE tasks ADD COLUMN user_id TEXT")
|
||||
except sqlite3.OperationalError:
|
||||
pass
|
||||
# Best-effort: add missing columns (forward compatibility).
|
||||
try:
|
||||
conn.execute("ALTER TABLE users ADD COLUMN updated_at TEXT")
|
||||
except sqlite3.OperationalError:
|
||||
pass
|
||||
try:
|
||||
conn.execute("ALTER TABLE api_tokens ADD COLUMN token_plain TEXT")
|
||||
except sqlite3.OperationalError:
|
||||
pass
|
||||
conn.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS attempts (
|
||||
task_id TEXT NOT NULL,
|
||||
attempt_no INTEGER NOT NULL,
|
||||
ray_submission_id TEXT NOT NULL UNIQUE,
|
||||
ray_status TEXT,
|
||||
failure_kind TEXT,
|
||||
message TEXT,
|
||||
start_time TEXT NOT NULL,
|
||||
end_time TEXT,
|
||||
PRIMARY KEY (task_id, attempt_no),
|
||||
FOREIGN KEY (task_id) REFERENCES tasks(task_id) ON DELETE CASCADE
|
||||
)
|
||||
"""
|
||||
)
|
||||
conn.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS events (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
task_id TEXT,
|
||||
ts TEXT NOT NULL,
|
||||
event_type TEXT NOT NULL,
|
||||
payload_json TEXT,
|
||||
FOREIGN KEY (task_id) REFERENCES tasks(task_id) ON DELETE CASCADE
|
||||
)
|
||||
"""
|
||||
)
|
||||
conn.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS serve_models (
|
||||
model_key TEXT PRIMARY KEY,
|
||||
user_id TEXT NOT NULL,
|
||||
model_id_suffix TEXT NOT NULL,
|
||||
model_id_prefix TEXT NOT NULL,
|
||||
model_id TEXT NOT NULL,
|
||||
model_source TEXT NOT NULL,
|
||||
num_replicas INTEGER NOT NULL,
|
||||
gpus_per_replica INTEGER NOT NULL,
|
||||
engine_kwargs_json TEXT,
|
||||
state TEXT NOT NULL,
|
||||
spec_yaml TEXT NOT NULL,
|
||||
resolved_spec_yaml TEXT NOT NULL,
|
||||
error_summary TEXT,
|
||||
created_at TEXT NOT NULL,
|
||||
updated_at TEXT NOT NULL,
|
||||
deleted_at TEXT
|
||||
)
|
||||
"""
|
||||
)
|
||||
conn.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS serve_events (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
model_key TEXT NOT NULL,
|
||||
ts TEXT NOT NULL,
|
||||
event_type TEXT NOT NULL,
|
||||
payload_json TEXT,
|
||||
FOREIGN KEY (model_key) REFERENCES serve_models(model_key) ON DELETE CASCADE
|
||||
)
|
||||
"""
|
||||
)
|
||||
conn.execute("CREATE INDEX IF NOT EXISTS idx_serve_models_user ON serve_models(user_id)")
|
||||
conn.execute("CREATE INDEX IF NOT EXISTS idx_serve_models_state ON serve_models(state)")
|
||||
conn.execute("CREATE INDEX IF NOT EXISTS idx_serve_events_model ON serve_events(model_key)")
|
||||
|
||||
@contextmanager
|
||||
def tx(self) -> Iterator[sqlite3.Connection]:
|
||||
conn = self._connect()
|
||||
try:
|
||||
conn.execute("BEGIN IMMEDIATE;")
|
||||
yield conn
|
||||
conn.execute("COMMIT;")
|
||||
except Exception:
|
||||
conn.execute("ROLLBACK;")
|
||||
raise
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
def create_task(self, *, task_id: str, workload: str, jobspec_yaml: str, nnodes: int, n_gpus_per_node: int) -> dict[str, Any]:
|
||||
now = _utc_now_iso()
|
||||
with self.tx() as conn:
|
||||
conn.execute(
|
||||
"""
|
||||
INSERT INTO tasks (task_id, workload, state, jobspec_yaml, nnodes, n_gpus_per_node, created_at, updated_at, user_id)
|
||||
VALUES (?, ?, 'QUEUED', ?, ?, ?, ?, ?, ?)
|
||||
""",
|
||||
(task_id, workload, jobspec_yaml, nnodes, n_gpus_per_node, now, now, None),
|
||||
)
|
||||
conn.execute(
|
||||
"INSERT INTO events (task_id, ts, event_type, payload_json) VALUES (?, ?, 'TASK_CREATED', ?)",
|
||||
(task_id, now, None),
|
||||
)
|
||||
row = conn.execute("SELECT * FROM tasks WHERE task_id = ?", (task_id,)).fetchone()
|
||||
return dict(row) if row else {}
|
||||
|
||||
def create_task_v25(
|
||||
self,
|
||||
*,
|
||||
task_id: str,
|
||||
user_id: str,
|
||||
workload: str,
|
||||
jobspec_yaml: str,
|
||||
nnodes: int,
|
||||
n_gpus_per_node: int,
|
||||
) -> dict[str, Any]:
|
||||
now = _utc_now_iso()
|
||||
with self.tx() as conn:
|
||||
conn.execute(
|
||||
"""
|
||||
INSERT INTO tasks (task_id, user_id, workload, state, jobspec_yaml, nnodes, n_gpus_per_node, created_at, updated_at)
|
||||
VALUES (?, ?, ?, 'QUEUED', ?, ?, ?, ?, ?)
|
||||
""",
|
||||
(task_id, user_id, workload, jobspec_yaml, nnodes, n_gpus_per_node, now, now),
|
||||
)
|
||||
conn.execute(
|
||||
"INSERT INTO events (task_id, ts, event_type, payload_json) VALUES (?, ?, 'TASK_CREATED', ?)",
|
||||
(task_id, now, None),
|
||||
)
|
||||
row = conn.execute("SELECT * FROM tasks WHERE task_id = ?", (task_id,)).fetchone()
|
||||
return dict(row) if row else {}
|
||||
|
||||
def create_user(self, *, user_id: str, display_name: str | None = None) -> dict[str, Any]:
|
||||
now = _utc_now_iso()
|
||||
with self.tx() as conn:
|
||||
conn.execute(
|
||||
"""
|
||||
INSERT INTO users (user_id, display_name, state, created_at, updated_at)
|
||||
VALUES (?, ?, 'ACTIVE', ?, ?)
|
||||
""",
|
||||
(user_id, display_name, now, now),
|
||||
)
|
||||
conn.execute(
|
||||
"INSERT INTO events (task_id, ts, event_type, payload_json) VALUES (NULL, ?, 'USER_CREATED', ?)",
|
||||
(now, user_id),
|
||||
)
|
||||
row = conn.execute("SELECT * FROM users WHERE user_id = ?", (user_id,)).fetchone()
|
||||
return dict(row) if row else {}
|
||||
|
||||
def disable_user(self, *, user_id: str) -> None:
|
||||
now = _utc_now_iso()
|
||||
with self.tx() as conn:
|
||||
conn.execute("UPDATE users SET state = 'DISABLED', updated_at = ? WHERE user_id = ?", (now, user_id))
|
||||
conn.execute(
|
||||
"INSERT INTO events (task_id, ts, event_type, payload_json) VALUES (NULL, ?, 'USER_DISABLED', ?)",
|
||||
(now, user_id),
|
||||
)
|
||||
|
||||
def get_user(self, user_id: str) -> dict[str, Any] | None:
|
||||
with self._connect() as conn:
|
||||
row = conn.execute("SELECT * FROM users WHERE user_id = ?", (user_id,)).fetchone()
|
||||
return dict(row) if row else None
|
||||
|
||||
def issue_token(self, *, user_id: str) -> str:
|
||||
# Returns plaintext token once.
|
||||
# Note: For v3.0 WebUI admin convenience we persist the plaintext token so admins
|
||||
# can re-copy it later. This is acceptable only for internal/dev deployments.
|
||||
now = _utc_now_iso()
|
||||
token = f"mvp_u_{user_id}_{secrets.token_urlsafe(18)}"
|
||||
token_hash = self._hash_token(token)
|
||||
with self.tx() as conn:
|
||||
conn.execute(
|
||||
"INSERT INTO api_tokens (token_hash, user_id, token_plain, created_at) VALUES (?, ?, ?, ?)",
|
||||
(token_hash, user_id, token, now),
|
||||
)
|
||||
conn.execute("UPDATE users SET updated_at = ? WHERE user_id = ?", (now, user_id))
|
||||
conn.execute(
|
||||
"INSERT INTO events (task_id, ts, event_type, payload_json) VALUES (NULL, ?, 'TOKEN_ISSUED', ?)",
|
||||
(now, user_id),
|
||||
)
|
||||
return token
|
||||
|
||||
def resolve_token(self, token: str) -> dict[str, Any] | None:
|
||||
token_hash = self._hash_token(token)
|
||||
with self.tx() as conn:
|
||||
row = conn.execute(
|
||||
"""
|
||||
SELECT t.user_id, u.state
|
||||
FROM api_tokens t
|
||||
JOIN users u ON u.user_id = t.user_id
|
||||
WHERE t.token_hash = ?
|
||||
""",
|
||||
(token_hash,),
|
||||
).fetchone()
|
||||
if not row:
|
||||
return None
|
||||
now = _utc_now_iso()
|
||||
conn.execute("UPDATE api_tokens SET last_used_at = ? WHERE token_hash = ?", (now, token_hash))
|
||||
conn.execute("UPDATE users SET updated_at = ? WHERE user_id = ?", (now, str(row["user_id"])))
|
||||
return {"user_id": row["user_id"], "state": row["state"]}
|
||||
|
||||
def list_users(self, limit: int = 200) -> list[dict[str, Any]]:
|
||||
with self._connect() as conn:
|
||||
rows = conn.execute(
|
||||
"""
|
||||
SELECT
|
||||
u.user_id,
|
||||
u.display_name,
|
||||
u.state,
|
||||
u.created_at,
|
||||
u.updated_at,
|
||||
(
|
||||
SELECT t.token_plain
|
||||
FROM api_tokens t
|
||||
WHERE t.user_id = u.user_id
|
||||
ORDER BY t.created_at DESC
|
||||
LIMIT 1
|
||||
) AS token,
|
||||
(
|
||||
SELECT t.created_at
|
||||
FROM api_tokens t
|
||||
WHERE t.user_id = u.user_id
|
||||
ORDER BY t.created_at DESC
|
||||
LIMIT 1
|
||||
) AS token_created_at,
|
||||
(
|
||||
SELECT t.last_used_at
|
||||
FROM api_tokens t
|
||||
WHERE t.user_id = u.user_id
|
||||
ORDER BY t.created_at DESC
|
||||
LIMIT 1
|
||||
) AS token_last_used_at
|
||||
FROM users u
|
||||
ORDER BY u.created_at DESC
|
||||
LIMIT ?
|
||||
""",
|
||||
(int(limit),),
|
||||
).fetchall()
|
||||
return [dict(r) for r in rows]
|
||||
|
||||
def get_task(self, task_id: str) -> dict[str, Any] | None:
|
||||
with self._connect() as conn:
|
||||
row = conn.execute("SELECT * FROM tasks WHERE task_id = ?", (task_id,)).fetchone()
|
||||
return dict(row) if row else None
|
||||
|
||||
def list_attempts(self, task_id: str) -> list[dict[str, Any]]:
|
||||
with self._connect() as conn:
|
||||
rows = conn.execute(
|
||||
"SELECT * FROM attempts WHERE task_id = ? ORDER BY attempt_no ASC", (task_id,)
|
||||
).fetchall()
|
||||
return [dict(r) for r in rows]
|
||||
|
||||
def list_queue(self, *, user_id: str | None = None) -> dict[str, list[dict[str, Any]]]:
|
||||
with self._connect() as conn:
|
||||
params: list[Any] = []
|
||||
user_filter_sql = ""
|
||||
if user_id is not None:
|
||||
user_filter_sql = " AND user_id = ?"
|
||||
params = [user_id]
|
||||
|
||||
pending_sql = (
|
||||
"SELECT task_id, workload, state, nnodes, n_gpus_per_node, next_run_at, created_at, updated_at "
|
||||
"FROM tasks "
|
||||
"WHERE state IN ('QUEUED','PENDING_RESOURCES')"
|
||||
f"{user_filter_sql} "
|
||||
"ORDER BY created_at ASC "
|
||||
"LIMIT 200"
|
||||
)
|
||||
running_sql = (
|
||||
"SELECT task_id, workload, state, nnodes, n_gpus_per_node, latest_attempt_no, created_at, updated_at "
|
||||
"FROM tasks "
|
||||
"WHERE state IN ('SUBMITTING','SUBMITTED','RUNNING')"
|
||||
f"{user_filter_sql} "
|
||||
"ORDER BY updated_at ASC "
|
||||
"LIMIT 200"
|
||||
)
|
||||
|
||||
pending = conn.execute(pending_sql, tuple(params)).fetchall()
|
||||
running = conn.execute(running_sql, tuple(params)).fetchall()
|
||||
return {"pending": [dict(r) for r in pending], "running": [dict(r) for r in running]}
|
||||
|
||||
def list_tasks(
|
||||
self,
|
||||
*,
|
||||
user_id: str | None = None,
|
||||
states: list[str] | None = None,
|
||||
limit: int = 200,
|
||||
offset: int = 0,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""
|
||||
Returns recent tasks (including terminal tasks).
|
||||
"""
|
||||
with self._connect() as conn:
|
||||
params: list[Any] = []
|
||||
where_clauses: list[str] = []
|
||||
if user_id is not None:
|
||||
where_clauses.append("user_id = ?")
|
||||
params.append(user_id)
|
||||
if states:
|
||||
placeholders = ",".join(["?"] * len(states))
|
||||
where_clauses.append(f"state IN ({placeholders})")
|
||||
params.extend(states)
|
||||
where_sql = f" WHERE {' AND '.join(where_clauses)}" if where_clauses else ""
|
||||
sql = (
|
||||
"SELECT task_id, user_id, workload, state, nnodes, n_gpus_per_node, latest_attempt_no, created_at, updated_at, error_summary "
|
||||
"FROM tasks"
|
||||
f"{where_sql} "
|
||||
"ORDER BY created_at DESC "
|
||||
"LIMIT ? OFFSET ?"
|
||||
)
|
||||
params.append(int(limit))
|
||||
params.append(max(0, int(offset)))
|
||||
rows = conn.execute(sql, tuple(params)).fetchall()
|
||||
return [dict(r) for r in rows]
|
||||
|
||||
def count_running(self) -> int:
|
||||
with self._connect() as conn:
|
||||
row = conn.execute(
|
||||
"SELECT COUNT(1) AS n FROM tasks WHERE state IN ('SUBMITTING','SUBMITTED','RUNNING')"
|
||||
).fetchone()
|
||||
return int(row["n"]) if row else 0
|
||||
|
||||
def list_active_tasks(self, limit: int = 50) -> list[dict[str, Any]]:
|
||||
with self._connect() as conn:
|
||||
rows = conn.execute(
|
||||
"SELECT * FROM tasks WHERE state IN ('SUBMITTING','SUBMITTED','RUNNING') ORDER BY updated_at ASC LIMIT ?",
|
||||
(int(limit),),
|
||||
).fetchall()
|
||||
return [dict(r) for r in rows]
|
||||
|
||||
def pick_next_runnable_task(self) -> dict[str, Any] | None:
|
||||
now = _utc_now_iso()
|
||||
with self._connect() as conn:
|
||||
row = conn.execute(
|
||||
"""
|
||||
SELECT *
|
||||
FROM tasks
|
||||
WHERE state IN ('QUEUED','PENDING_RESOURCES')
|
||||
AND (next_run_at IS NULL OR next_run_at <= ?)
|
||||
ORDER BY created_at ASC
|
||||
LIMIT 1
|
||||
""",
|
||||
(now,),
|
||||
).fetchone()
|
||||
return dict(row) if row else None
|
||||
|
||||
def set_task_state(
|
||||
self,
|
||||
*,
|
||||
task_id: str,
|
||||
state: str,
|
||||
error_summary: str | None = None,
|
||||
next_run_at: str | None = None,
|
||||
latest_attempt_no: int | None = None,
|
||||
event_type: str = "STATE_UPDATE",
|
||||
payload_json: str | None = None,
|
||||
) -> None:
|
||||
now = _utc_now_iso()
|
||||
with self.tx() as conn:
|
||||
sets = ["state = ?", "updated_at = ?"]
|
||||
params: list[Any] = [state, now]
|
||||
if error_summary is not None:
|
||||
sets.append("error_summary = ?")
|
||||
params.append(error_summary)
|
||||
if next_run_at is not None:
|
||||
sets.append("next_run_at = ?")
|
||||
params.append(next_run_at)
|
||||
if latest_attempt_no is not None:
|
||||
sets.append("latest_attempt_no = ?")
|
||||
params.append(int(latest_attempt_no))
|
||||
params.append(task_id)
|
||||
conn.execute(f"UPDATE tasks SET {', '.join(sets)} WHERE task_id = ?", tuple(params))
|
||||
conn.execute(
|
||||
"INSERT INTO events (task_id, ts, event_type, payload_json) VALUES (?, ?, ?, ?)",
|
||||
(task_id, now, event_type, payload_json),
|
||||
)
|
||||
|
||||
def create_attempt(self, *, task_id: str, attempt_no: int, ray_submission_id: str) -> None:
|
||||
now = _utc_now_iso()
|
||||
with self.tx() as conn:
|
||||
conn.execute(
|
||||
"""
|
||||
INSERT INTO attempts (task_id, attempt_no, ray_submission_id, ray_status, start_time)
|
||||
VALUES (?, ?, ?, ?, ?)
|
||||
""",
|
||||
(task_id, attempt_no, ray_submission_id, "SUBMITTING", now),
|
||||
)
|
||||
conn.execute(
|
||||
"INSERT INTO events (task_id, ts, event_type, payload_json) VALUES (?, ?, 'ATTEMPT_CREATED', ?)",
|
||||
(task_id, now, None),
|
||||
)
|
||||
|
||||
def update_attempt(
|
||||
self,
|
||||
*,
|
||||
task_id: str,
|
||||
attempt_no: int,
|
||||
ray_status: str | None = None,
|
||||
failure_kind: str | None = None,
|
||||
message: str | None = None,
|
||||
end_time: str | None = None,
|
||||
) -> None:
|
||||
now = _utc_now_iso()
|
||||
with self.tx() as conn:
|
||||
sets = []
|
||||
params: list[Any] = []
|
||||
if ray_status is not None:
|
||||
sets.append("ray_status = ?")
|
||||
params.append(ray_status)
|
||||
if failure_kind is not None:
|
||||
sets.append("failure_kind = ?")
|
||||
params.append(failure_kind)
|
||||
if message is not None:
|
||||
sets.append("message = ?")
|
||||
params.append(message)
|
||||
if end_time is not None:
|
||||
sets.append("end_time = ?")
|
||||
params.append(end_time)
|
||||
if not sets:
|
||||
return
|
||||
params.extend([task_id, attempt_no])
|
||||
conn.execute(
|
||||
f"UPDATE attempts SET {', '.join(sets)} WHERE task_id = ? AND attempt_no = ?",
|
||||
tuple(params),
|
||||
)
|
||||
conn.execute(
|
||||
"INSERT INTO events (task_id, ts, event_type, payload_json) VALUES (?, ?, 'ATTEMPT_UPDATE', ?)",
|
||||
(task_id, now, None),
|
||||
)
|
||||
|
||||
def list_ended_attempts_before(self, *, end_time_le: str, limit: int = 2000) -> list[dict[str, Any]]:
|
||||
"""
|
||||
Used by the jobs janitor:
|
||||
- Only considers tasks with a non-null user_id (v2.5+).
|
||||
- Returns attempts that have end_time <= end_time_le.
|
||||
"""
|
||||
with self._connect() as conn:
|
||||
rows = conn.execute(
|
||||
"""
|
||||
SELECT t.task_id, t.user_id, a.ray_submission_id, a.end_time
|
||||
FROM attempts a
|
||||
JOIN tasks t ON t.task_id = a.task_id
|
||||
WHERE t.user_id IS NOT NULL
|
||||
AND a.end_time IS NOT NULL
|
||||
AND a.end_time <= ?
|
||||
ORDER BY a.end_time ASC
|
||||
LIMIT ?
|
||||
""",
|
||||
(str(end_time_le), int(limit)),
|
||||
).fetchall()
|
||||
return [dict(r) for r in rows]
|
||||
|
||||
def create_serve_model(
|
||||
self,
|
||||
*,
|
||||
model_key: str,
|
||||
user_id: str,
|
||||
model_id_suffix: str,
|
||||
model_id_prefix: str,
|
||||
model_id: str,
|
||||
model_source: str,
|
||||
num_replicas: int,
|
||||
gpus_per_replica: int,
|
||||
spec_yaml: str,
|
||||
resolved_spec_yaml: str,
|
||||
engine_kwargs_json: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
now = _utc_now_iso()
|
||||
with self.tx() as conn:
|
||||
conn.execute(
|
||||
"""
|
||||
INSERT INTO serve_models (
|
||||
model_key,
|
||||
user_id,
|
||||
model_id_suffix,
|
||||
model_id_prefix,
|
||||
model_id,
|
||||
model_source,
|
||||
num_replicas,
|
||||
gpus_per_replica,
|
||||
engine_kwargs_json,
|
||||
state,
|
||||
spec_yaml,
|
||||
resolved_spec_yaml,
|
||||
created_at,
|
||||
updated_at
|
||||
)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, 'QUEUED', ?, ?, ?, ?)
|
||||
""",
|
||||
(
|
||||
model_key,
|
||||
user_id,
|
||||
model_id_suffix,
|
||||
model_id_prefix,
|
||||
model_id,
|
||||
model_source,
|
||||
int(num_replicas),
|
||||
int(gpus_per_replica),
|
||||
engine_kwargs_json,
|
||||
spec_yaml,
|
||||
resolved_spec_yaml,
|
||||
now,
|
||||
now,
|
||||
),
|
||||
)
|
||||
conn.execute(
|
||||
"INSERT INTO serve_events (model_key, ts, event_type, payload_json) VALUES (?, ?, 'SERVE_MODEL_CREATED', ?)",
|
||||
(model_key, now, None),
|
||||
)
|
||||
row = conn.execute("SELECT * FROM serve_models WHERE model_key = ?", (model_key,)).fetchone()
|
||||
return dict(row) if row else {}
|
||||
|
||||
def list_serve_models(
|
||||
self,
|
||||
*,
|
||||
user_id: str,
|
||||
include_deleted: bool = False,
|
||||
limit: int = 200,
|
||||
offset: int = 0,
|
||||
) -> list[dict[str, Any]]:
|
||||
with self._connect() as conn:
|
||||
where_sql = "WHERE user_id = ?"
|
||||
params: list[Any] = [user_id]
|
||||
if not include_deleted:
|
||||
where_sql += " AND deleted_at IS NULL"
|
||||
params.append(int(limit))
|
||||
params.append(max(0, int(offset)))
|
||||
rows = conn.execute(
|
||||
f"""
|
||||
SELECT
|
||||
model_key,
|
||||
user_id,
|
||||
model_id_suffix,
|
||||
model_id_prefix,
|
||||
model_id,
|
||||
model_source,
|
||||
num_replicas,
|
||||
gpus_per_replica,
|
||||
engine_kwargs_json,
|
||||
state,
|
||||
error_summary,
|
||||
created_at,
|
||||
updated_at,
|
||||
deleted_at
|
||||
FROM serve_models
|
||||
{where_sql}
|
||||
ORDER BY created_at DESC, model_key DESC
|
||||
LIMIT ? OFFSET ?
|
||||
""",
|
||||
tuple(params),
|
||||
).fetchall()
|
||||
return [dict(r) for r in rows]
|
||||
|
||||
def list_all_serve_models(
|
||||
self,
|
||||
*,
|
||||
include_deleted: bool = False,
|
||||
limit: int = 2000,
|
||||
offset: int = 0,
|
||||
) -> list[dict[str, Any]]:
|
||||
with self._connect() as conn:
|
||||
where_sql = ""
|
||||
if not include_deleted:
|
||||
where_sql = "WHERE deleted_at IS NULL"
|
||||
rows = conn.execute(
|
||||
f"""
|
||||
SELECT
|
||||
model_key,
|
||||
user_id,
|
||||
model_id_suffix,
|
||||
model_id_prefix,
|
||||
model_id,
|
||||
model_source,
|
||||
num_replicas,
|
||||
gpus_per_replica,
|
||||
engine_kwargs_json,
|
||||
state,
|
||||
error_summary,
|
||||
spec_yaml,
|
||||
resolved_spec_yaml,
|
||||
created_at,
|
||||
updated_at,
|
||||
deleted_at
|
||||
FROM serve_models
|
||||
{where_sql}
|
||||
ORDER BY created_at ASC, model_key ASC
|
||||
LIMIT ? OFFSET ?
|
||||
""",
|
||||
(int(limit), max(0, int(offset))),
|
||||
).fetchall()
|
||||
return [dict(r) for r in rows]
|
||||
|
||||
def get_serve_model(self, model_key: str) -> dict[str, Any] | None:
|
||||
with self._connect() as conn:
|
||||
row = conn.execute("SELECT * FROM serve_models WHERE model_key = ?", (model_key,)).fetchone()
|
||||
return dict(row) if row else None
|
||||
|
||||
def list_serve_events(self, model_key: str, limit: int = 200, offset: int = 0) -> list[dict[str, Any]]:
|
||||
with self._connect() as conn:
|
||||
rows = conn.execute(
|
||||
"""
|
||||
SELECT id, model_key, ts, event_type, payload_json
|
||||
FROM serve_events
|
||||
WHERE model_key = ?
|
||||
ORDER BY id DESC
|
||||
LIMIT ? OFFSET ?
|
||||
""",
|
||||
(model_key, int(limit), max(0, int(offset))),
|
||||
).fetchall()
|
||||
return [dict(r) for r in rows]
|
||||
|
||||
def append_serve_event(self, *, model_key: str, event_type: str, payload_json: str | None = None) -> None:
|
||||
now = _utc_now_iso()
|
||||
with self.tx() as conn:
|
||||
conn.execute(
|
||||
"INSERT INTO serve_events (model_key, ts, event_type, payload_json) VALUES (?, ?, ?, ?)",
|
||||
(model_key, now, event_type, payload_json),
|
||||
)
|
||||
|
||||
def set_serve_model_state(
|
||||
self,
|
||||
*,
|
||||
model_key: str,
|
||||
state: str,
|
||||
error_summary: str | None = None,
|
||||
event_type: str = "SERVE_STATE_UPDATE",
|
||||
payload_json: str | None = None,
|
||||
) -> None:
|
||||
now = _utc_now_iso()
|
||||
with self.tx() as conn:
|
||||
sets = ["state = ?", "updated_at = ?"]
|
||||
params: list[Any] = [state, now]
|
||||
if error_summary is not None:
|
||||
sets.append("error_summary = ?")
|
||||
params.append(error_summary)
|
||||
if state == "DELETED":
|
||||
sets.append("deleted_at = ?")
|
||||
params.append(now)
|
||||
params.append(model_key)
|
||||
conn.execute(f"UPDATE serve_models SET {', '.join(sets)} WHERE model_key = ?", tuple(params))
|
||||
conn.execute(
|
||||
"INSERT INTO serve_events (model_key, ts, event_type, payload_json) VALUES (?, ?, ?, ?)",
|
||||
(model_key, now, event_type, payload_json),
|
||||
)
|
||||
|
||||
def update_serve_model_num_replicas(self, *, model_key: str, num_replicas: int) -> None:
|
||||
if not isinstance(num_replicas, int) or num_replicas < 1:
|
||||
raise ValueError("num_replicas must be an integer >= 1")
|
||||
now = _utc_now_iso()
|
||||
with self.tx() as conn:
|
||||
conn.execute(
|
||||
"""
|
||||
UPDATE serve_models
|
||||
SET num_replicas = ?, state = 'QUEUED', error_summary = NULL, updated_at = ?
|
||||
WHERE model_key = ?
|
||||
""",
|
||||
(int(num_replicas), now, model_key),
|
||||
)
|
||||
conn.execute(
|
||||
"INSERT INTO serve_events (model_key, ts, event_type, payload_json) VALUES (?, ?, 'SERVE_PATCH_NUM_REPLICAS', ?)",
|
||||
(model_key, now, str(num_replicas)),
|
||||
)
|
||||
|
||||
def pick_next_runnable_serve_change(self) -> dict[str, Any] | None:
|
||||
"""
|
||||
Returns the next serve model that needs reconciliation.
|
||||
|
||||
Minimal state machine for now:
|
||||
- QUEUED: needs (re)apply
|
||||
- DELETING: needs removal
|
||||
"""
|
||||
with self._connect() as conn:
|
||||
row = conn.execute(
|
||||
"""
|
||||
SELECT *
|
||||
FROM serve_models
|
||||
WHERE deleted_at IS NULL
|
||||
AND state IN ('QUEUED','DELETING')
|
||||
ORDER BY updated_at ASC
|
||||
LIMIT 1
|
||||
"""
|
||||
).fetchone()
|
||||
return dict(row) if row else None
|
||||
|
||||
# Backward compatible naming (v3.8 docs originally used "upsert").
|
||||
def upsert_serve_model(self, **kwargs: Any) -> dict[str, Any]:
|
||||
return self.create_serve_model(**kwargs)
|
||||
105
src/mvp/py/argus/service/janitor.py
Normal file
@@ -0,0 +1,105 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import shutil
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
from .db import Db
|
||||
|
||||
|
||||
def _parse_iso_z(ts: str) -> datetime:
|
||||
# Stored as "YYYY-MM-DDTHH:MM:SSZ" (naive UTC). Parse into aware UTC.
|
||||
return datetime.fromisoformat(ts.replace("Z", "+00:00")).astimezone(timezone.utc)
|
||||
|
||||
|
||||
def _iso_z(dt: datetime) -> str:
|
||||
return dt.astimezone(timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z")
|
||||
|
||||
|
||||
@dataclass
|
||||
class JobsJanitor:
|
||||
db: Db
|
||||
user_root: str
|
||||
trash_after_days: int = 3
|
||||
purge_after_days: int = 7
|
||||
interval_s: int = 3600
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if int(self.trash_after_days) < 0 or int(self.purge_after_days) < 0:
|
||||
raise ValueError("retention days must be non-negative")
|
||||
if int(self.purge_after_days) and int(self.purge_after_days) < int(self.trash_after_days):
|
||||
raise ValueError("purge_after_days must be >= trash_after_days")
|
||||
|
||||
def _job_dir(self, *, user_id: str, ray_submission_id: str) -> str:
|
||||
base = self.user_root.rstrip("/")
|
||||
return f"{base}/{user_id}/jobs/{ray_submission_id}"
|
||||
|
||||
def _trash_dir(self, *, user_id: str, ray_submission_id: str) -> str:
|
||||
base = self.user_root.rstrip("/")
|
||||
return f"{base}/{user_id}/trash/jobs/{ray_submission_id}"
|
||||
|
||||
def tick_once(self, *, now: datetime | None = None, limit: int = 2000) -> None:
|
||||
if int(self.trash_after_days) <= 0 and int(self.purge_after_days) <= 0:
|
||||
return
|
||||
|
||||
now_dt = now or datetime.now(timezone.utc)
|
||||
move_cutoff = now_dt - timedelta(days=int(self.trash_after_days))
|
||||
move_cutoff_iso = _iso_z(move_cutoff)
|
||||
|
||||
rows = self.db.list_ended_attempts_before(end_time_le=move_cutoff_iso, limit=int(limit))
|
||||
for r in rows:
|
||||
user_id = str(r.get("user_id") or "").strip()
|
||||
sid = str(r.get("ray_submission_id") or "").strip()
|
||||
end_time = str(r.get("end_time") or "").strip()
|
||||
if not user_id or not sid or not end_time:
|
||||
continue
|
||||
|
||||
try:
|
||||
ended_at = _parse_iso_z(end_time)
|
||||
except Exception:
|
||||
continue
|
||||
|
||||
age_days = (now_dt - ended_at).total_seconds() / 86400.0
|
||||
src = self._job_dir(user_id=user_id, ray_submission_id=sid)
|
||||
dst = self._trash_dir(user_id=user_id, ray_submission_id=sid)
|
||||
dst_parent = os.path.dirname(dst)
|
||||
|
||||
# Between trash and purge: ensure in trash.
|
||||
if age_days >= float(self.trash_after_days) and age_days < float(self.purge_after_days or 10**9):
|
||||
if os.path.exists(src) and not os.path.exists(dst):
|
||||
os.makedirs(dst_parent, exist_ok=True)
|
||||
try:
|
||||
shutil.move(src, dst)
|
||||
except Exception:
|
||||
pass
|
||||
continue
|
||||
|
||||
# Purge: move to trash (if still in jobs) then delete from trash.
|
||||
if int(self.purge_after_days) > 0 and age_days >= float(self.purge_after_days):
|
||||
if os.path.exists(src) and not os.path.exists(dst):
|
||||
os.makedirs(dst_parent, exist_ok=True)
|
||||
try:
|
||||
shutil.move(src, dst)
|
||||
except Exception:
|
||||
pass
|
||||
if os.path.exists(dst):
|
||||
try:
|
||||
shutil.rmtree(dst)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def run_forever(self, stop_flag: object) -> None:
|
||||
try:
|
||||
is_set = stop_flag.is_set # type: ignore[attr-defined]
|
||||
except Exception:
|
||||
raise ValueError("stop_flag must be a threading.Event-like object with is_set()")
|
||||
|
||||
while not is_set():
|
||||
try:
|
||||
self.tick_once()
|
||||
except Exception:
|
||||
pass
|
||||
time.sleep(max(1, int(self.interval_s)))
|
||||
|
||||
37
src/mvp/py/argus/service/ray_resources.py
Normal file
@@ -0,0 +1,37 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
import ray
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ClusterAvailable:
|
||||
total_available_gpus: float
|
||||
total_available_npus: float
|
||||
|
||||
|
||||
def get_cluster_available() -> ClusterAvailable:
|
||||
# Align with verl's fail-fast check which uses ray._private.state.available_resources_per_node().
|
||||
# This is a best-effort internal API and may change with Ray versions.
|
||||
try:
|
||||
import ray._private.state # type: ignore
|
||||
|
||||
per_node = ray._private.state.available_resources_per_node()
|
||||
except Exception:
|
||||
# If we cannot fetch per-node resources, conservatively return 0.
|
||||
return ClusterAvailable(total_available_gpus=0.0, total_available_npus=0.0)
|
||||
|
||||
total_gpu = 0.0
|
||||
total_npu = 0.0
|
||||
for _, info in per_node.items():
|
||||
if not isinstance(info, dict):
|
||||
continue
|
||||
total_gpu += float(info.get("GPU", 0) or 0)
|
||||
total_npu += float(info.get("NPU", 0) or 0)
|
||||
return ClusterAvailable(total_available_gpus=total_gpu, total_available_npus=total_npu)
|
||||
|
||||
|
||||
def ensure_ray_connected() -> None:
|
||||
ray.init(address="auto", ignore_reinit_error=True, log_to_driver=False)
|
||||
|
||||
309
src/mvp/py/argus/service/scheduler.py
Normal file
@@ -0,0 +1,309 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timedelta
|
||||
from typing import Any
|
||||
|
||||
import yaml
|
||||
|
||||
from argus.core.ids import attempt_submission_id
|
||||
from argus.ray.models import AdvancedTaskSpec, JobSpec, RayConfig, parse_taskspec
|
||||
from argus.ray.ray_job_tool import RayJobTool
|
||||
|
||||
from .config import V2Config
|
||||
from .db import Db
|
||||
from .ray_resources import ensure_ray_connected, get_cluster_available
|
||||
from .serve_client import RayServeClient
|
||||
from .serving_reconciler import ServingReconciler
|
||||
|
||||
|
||||
_INSUFFICIENT_RE = re.compile(r"Total available GPUs\\s+\\d+\\s+is less than total desired GPUs\\s+\\d+")
|
||||
|
||||
|
||||
def _utc_now_iso() -> str:
|
||||
return datetime.utcnow().replace(microsecond=0).isoformat() + "Z"
|
||||
|
||||
|
||||
def _utc_after_s(seconds: int) -> str:
|
||||
return (datetime.utcnow() + timedelta(seconds=seconds)).replace(microsecond=0).isoformat() + "Z"
|
||||
|
||||
|
||||
@dataclass
|
||||
class Scheduler:
|
||||
db: Db
|
||||
ray_cfg: RayConfig
|
||||
v2_cfg: V2Config
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
self.tool = RayJobTool(self.ray_cfg)
|
||||
self._serving: ServingReconciler | None = None
|
||||
if bool(self.v2_cfg.serving.enabled):
|
||||
self._serving = ServingReconciler(
|
||||
db=self.db,
|
||||
v2_cfg=self.v2_cfg,
|
||||
ray_runtime_env_env_vars=self.ray_cfg.runtime_env_env_vars,
|
||||
serve_client=RayServeClient(
|
||||
http_port=int(self.v2_cfg.serving.serve.http_port),
|
||||
proxy_location=str(self.v2_cfg.serving.serve.proxy_location),
|
||||
ray_init_address="auto",
|
||||
),
|
||||
)
|
||||
|
||||
def _job_dir_for_task(self, *, user_id: str | None, ray_submission_id: str) -> str:
|
||||
root = self.ray_cfg.shared_root.rstrip("/")
|
||||
if user_id:
|
||||
return f"{root}/users/{user_id}/jobs/{ray_submission_id}"
|
||||
return f"{root}/jobs/{ray_submission_id}"
|
||||
|
||||
def _resources_sufficient(self, *, nnodes: int, n_gpus_per_node: int) -> bool:
|
||||
avail = get_cluster_available()
|
||||
required = float(nnodes * n_gpus_per_node)
|
||||
return avail.total_available_gpus >= required
|
||||
|
||||
def _wandb_env_vars(self, *, ray_submission_id: str, job_dir: str) -> tuple[dict[str, str], str | None]:
|
||||
"""
|
||||
Returns (env_vars, warn_message).
|
||||
|
||||
env_vars is always non-empty so that AdvancedTaskSpec examples can safely
|
||||
reference injected env vars without requiring the user to edit their command:
|
||||
- MVP_TRAINER_LOGGER: "console" or "[console,wandb]"
|
||||
- MVP_WANDB_PROJECT: "<user_id>_project"
|
||||
- MVP_WANDB_RUN: "<ray_submission_id>"
|
||||
|
||||
When W&B is effectively enabled (config enabled + API key present),
|
||||
WANDB_* env vars are included as well.
|
||||
"""
|
||||
env, warn, _enabled = self._tracking_env_vars(user_id=None, ray_submission_id=ray_submission_id, job_dir=job_dir)
|
||||
return env, warn
|
||||
|
||||
def _tracking_env_vars(
|
||||
self, *, user_id: str | None, ray_submission_id: str, job_dir: str
|
||||
) -> tuple[dict[str, str], str | None, bool]:
|
||||
"""
|
||||
Returns (env_vars, warn_message, wandb_enabled_effective).
|
||||
|
||||
- env_vars always includes MVP_* vars so advanced commands can use them.
|
||||
- wandb_enabled_effective is True only when tracking.wandb.enabled AND api key exists.
|
||||
"""
|
||||
cfg = self.v2_cfg.tracking.wandb
|
||||
project_suffix = cfg.project_suffix or "_project"
|
||||
project_name = f"{user_id}{project_suffix}" if user_id else f"mvp{project_suffix}"
|
||||
|
||||
api_key = os.environ.get(cfg.api_key_env, "").strip()
|
||||
wandb_enabled_effective = bool(cfg.enabled and api_key)
|
||||
warn: str | None = None
|
||||
if cfg.enabled and not api_key:
|
||||
warn = f"WARN_WANDB_DISABLED_MISSING_ENV:{cfg.api_key_env}"
|
||||
|
||||
env: dict[str, str] = {
|
||||
"MVP_TRAINER_LOGGER": ("[console,wandb]" if wandb_enabled_effective else "console"),
|
||||
"MVP_WANDB_PROJECT": project_name,
|
||||
"MVP_WANDB_RUN": ray_submission_id,
|
||||
}
|
||||
|
||||
if wandb_enabled_effective:
|
||||
env.update(
|
||||
{
|
||||
"WANDB_BASE_URL": cfg.base_url,
|
||||
"WANDB_API_KEY": api_key,
|
||||
"WANDB_MODE": "online",
|
||||
"WANDB_DIR": f"{job_dir.rstrip('/')}/wandb",
|
||||
}
|
||||
)
|
||||
|
||||
return env, warn, wandb_enabled_effective
|
||||
|
||||
def _wandb_training_overrides(self, *, user_id: str | None, ray_submission_id: str) -> list[str] | None:
|
||||
"""
|
||||
Returns hydra overrides to enable W&B logging for built-in training workloads.
|
||||
|
||||
Only call this when W&B is effectively enabled for the job (i.e. API key exists).
|
||||
"""
|
||||
cfg = self.v2_cfg.tracking.wandb
|
||||
project_suffix = cfg.project_suffix or "_project"
|
||||
project_name = f"{user_id}{project_suffix}" if user_id else f"mvp{project_suffix}"
|
||||
return [
|
||||
"trainer.logger=[console,wandb]",
|
||||
f"trainer.project_name={project_name}",
|
||||
f"trainer.experiment_name={ray_submission_id}",
|
||||
]
|
||||
|
||||
def _parse_taskspec(self, jobspec_yaml: str) -> JobSpec | AdvancedTaskSpec:
|
||||
obj = yaml.safe_load(jobspec_yaml) or {}
|
||||
if not isinstance(obj, dict):
|
||||
raise ValueError("jobspec must be a YAML mapping")
|
||||
return parse_taskspec(obj)
|
||||
|
||||
def _submit_one(self, task_row: dict[str, Any]) -> None:
|
||||
task_id = str(task_row["task_id"])
|
||||
jobspec_yaml = str(task_row["jobspec_yaml"])
|
||||
user_id = task_row.get("user_id")
|
||||
user_id_s = str(user_id) if user_id not in (None, "") else None
|
||||
|
||||
spec = self._parse_taskspec(jobspec_yaml)
|
||||
attempt_no = int(task_row.get("latest_attempt_no", 0)) + 1
|
||||
ray_sid = attempt_submission_id(task_id, attempt_no)
|
||||
job_dir = self._job_dir_for_task(user_id=user_id_s, ray_submission_id=ray_sid)
|
||||
|
||||
# Record attempt first so that we can surface it even if submit crashes.
|
||||
self.db.create_attempt(task_id=task_id, attempt_no=attempt_no, ray_submission_id=ray_sid)
|
||||
self.db.set_task_state(task_id=task_id, state="SUBMITTING", latest_attempt_no=attempt_no)
|
||||
|
||||
# Override submission_id in taskspec (v1.1 compatible)
|
||||
tracking_env, wandb_warn, wandb_enabled = self._tracking_env_vars(
|
||||
user_id=user_id_s, ray_submission_id=ray_sid, job_dir=job_dir
|
||||
)
|
||||
if isinstance(spec, JobSpec):
|
||||
d = spec.to_public_dict()
|
||||
d["submission_id"] = ray_sid
|
||||
spec2 = JobSpec.from_dict(d)
|
||||
wandb_overrides = self._wandb_training_overrides(user_id=user_id_s, ray_submission_id=ray_sid) if wandb_enabled else None
|
||||
submit = lambda: self.tool.submit(
|
||||
spec2,
|
||||
no_wait=True,
|
||||
job_dir=job_dir,
|
||||
extra_env_vars=tracking_env,
|
||||
extra_overrides=wandb_overrides,
|
||||
)
|
||||
else:
|
||||
d = spec.to_public_dict()
|
||||
d["submission_id"] = ray_sid
|
||||
spec2 = AdvancedTaskSpec.from_dict(d)
|
||||
submit = lambda: self.tool.submit_advanced(
|
||||
spec2, no_wait=True, job_dir=job_dir, user_id=user_id_s, extra_env_vars=tracking_env
|
||||
)
|
||||
|
||||
try:
|
||||
submitted = submit()
|
||||
# submitted should equal ray_sid; keep as source of truth.
|
||||
self.db.update_attempt(
|
||||
task_id=task_id,
|
||||
attempt_no=attempt_no,
|
||||
ray_status="SUBMITTED",
|
||||
message=wandb_warn,
|
||||
)
|
||||
self.db.set_task_state(task_id=task_id, state="SUBMITTED")
|
||||
if submitted != ray_sid:
|
||||
self.db.set_task_state(task_id=task_id, state="SUBMITTED", event_type="WARN_SUBMISSION_ID_MISMATCH")
|
||||
except Exception as e:
|
||||
msg = repr(e)
|
||||
self.db.update_attempt(task_id=task_id, attempt_no=attempt_no, ray_status="FAILED", failure_kind="UNKNOWN", message=msg, end_time=_utc_now_iso())
|
||||
self.db.set_task_state(task_id=task_id, state="FAILED", error_summary=msg)
|
||||
|
||||
def _sync_one_running(self, task_row: dict[str, Any]) -> None:
|
||||
task_id = str(task_row["task_id"])
|
||||
latest_attempt_no = int(task_row.get("latest_attempt_no", 0))
|
||||
if latest_attempt_no <= 0:
|
||||
return
|
||||
|
||||
# Look up ray_submission_id
|
||||
attempts = self.db.list_attempts(task_id)
|
||||
if not attempts:
|
||||
return
|
||||
ray_sid = str(attempts[-1]["ray_submission_id"])
|
||||
|
||||
try:
|
||||
st = self.tool.status(ray_sid)
|
||||
except Exception as e:
|
||||
# Keep current state; transient failures should not flap tasks.
|
||||
self.db.set_task_state(task_id=task_id, state=str(task_row["state"]), event_type="RAY_STATUS_ERROR", payload_json=repr(e))
|
||||
return
|
||||
|
||||
st_s = str(st)
|
||||
if st_s in ("PENDING", "RUNNING"):
|
||||
self.db.update_attempt(task_id=task_id, attempt_no=latest_attempt_no, ray_status=st_s)
|
||||
self.db.set_task_state(task_id=task_id, state=("RUNNING" if st_s == "RUNNING" else "SUBMITTED"))
|
||||
return
|
||||
|
||||
if st_s in ("SUCCEEDED",):
|
||||
self.db.update_attempt(task_id=task_id, attempt_no=latest_attempt_no, ray_status=st_s, end_time=_utc_now_iso())
|
||||
self.db.set_task_state(task_id=task_id, state="SUCCEEDED")
|
||||
return
|
||||
|
||||
if st_s in ("FAILED", "STOPPED"):
|
||||
logs = ""
|
||||
try:
|
||||
logs = self.tool.logs(ray_sid)
|
||||
except Exception:
|
||||
logs = ""
|
||||
failure_kind = "UNKNOWN"
|
||||
msg = ""
|
||||
if _INSUFFICIENT_RE.search(logs):
|
||||
failure_kind = "INSUFFICIENT_RESOURCES"
|
||||
msg = "Insufficient resources (verl fail-fast): " + (_INSUFFICIENT_RE.search(logs).group(0))
|
||||
self.db.update_attempt(
|
||||
task_id=task_id,
|
||||
attempt_no=latest_attempt_no,
|
||||
ray_status=st_s,
|
||||
failure_kind=failure_kind,
|
||||
message=msg,
|
||||
end_time=_utc_now_iso(),
|
||||
)
|
||||
self.db.set_task_state(
|
||||
task_id=task_id,
|
||||
state="PENDING_RESOURCES",
|
||||
error_summary=msg,
|
||||
next_run_at=_utc_after_s(self.v2_cfg.scheduler.retry_interval_s),
|
||||
event_type="RETRY_SCHEDULED",
|
||||
)
|
||||
return
|
||||
|
||||
msg = f"Ray job {st_s}"
|
||||
self.db.update_attempt(
|
||||
task_id=task_id,
|
||||
attempt_no=latest_attempt_no,
|
||||
ray_status=st_s,
|
||||
failure_kind=failure_kind,
|
||||
message=msg,
|
||||
end_time=_utc_now_iso(),
|
||||
)
|
||||
self.db.set_task_state(task_id=task_id, state="FAILED", error_summary=msg)
|
||||
|
||||
def tick(self) -> None:
|
||||
ensure_ray_connected()
|
||||
|
||||
# v3.8: reconcile serve_models (best-effort).
|
||||
if self._serving is not None:
|
||||
try:
|
||||
self._serving.tick()
|
||||
except Exception:
|
||||
# Keep scheduler alive even if serving tick fails.
|
||||
pass
|
||||
|
||||
# Sync active tasks
|
||||
for row in self.db.list_active_tasks(limit=50):
|
||||
self._sync_one_running(row)
|
||||
|
||||
# Submit new tasks if capacity allows
|
||||
if self.db.count_running() >= self.v2_cfg.scheduler.max_running_tasks:
|
||||
return
|
||||
|
||||
row = self.db.pick_next_runnable_task()
|
||||
if not row:
|
||||
return
|
||||
|
||||
nnodes = int(row["nnodes"])
|
||||
n_gpus_per_node = int(row["n_gpus_per_node"])
|
||||
|
||||
if not self._resources_sufficient(nnodes=nnodes, n_gpus_per_node=n_gpus_per_node):
|
||||
self.db.set_task_state(
|
||||
task_id=str(row["task_id"]),
|
||||
state="PENDING_RESOURCES",
|
||||
next_run_at=_utc_after_s(self.v2_cfg.scheduler.retry_interval_s),
|
||||
event_type="PENDING_RESOURCES",
|
||||
)
|
||||
return
|
||||
|
||||
self._submit_one(row)
|
||||
|
||||
def run_forever(self, stop_flag: Any) -> None:
|
||||
while not stop_flag.is_set():
|
||||
try:
|
||||
self.tick()
|
||||
except Exception:
|
||||
# Best-effort: don't crash the scheduler loop
|
||||
pass
|
||||
time.sleep(max(1, int(self.v2_cfg.scheduler.tick_s)))
|
||||
50
src/mvp/py/argus/service/serve_client.py
Normal file
@@ -0,0 +1,50 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class RayServeClient:
|
||||
"""
|
||||
Minimal Ray Serve client wrapper.
|
||||
|
||||
This is intentionally tiny and uses runtime imports so that:
|
||||
- unit tests can stub `ray` modules without needing real Ray installed
|
||||
- production can run with the real Ray Serve stack (v3.8+)
|
||||
"""
|
||||
|
||||
http_port: int = 8000
|
||||
proxy_location: str = "HeadOnly"
|
||||
ray_init_address: str = "auto"
|
||||
|
||||
def ensure_started(self) -> None:
|
||||
import ray # runtime import
|
||||
|
||||
# Scheduler already calls ray.init(); make this idempotent.
|
||||
ray.init(address=self.ray_init_address, ignore_reinit_error=True, log_to_driver=False) # type: ignore[call-arg]
|
||||
|
||||
# Import serve lazily to allow tests to stub it.
|
||||
from ray import serve # type: ignore
|
||||
|
||||
serve.start(proxy_location=self.proxy_location, http_options={"host": "0.0.0.0", "port": int(self.http_port)})
|
||||
|
||||
def apply_app(self, *, app: Any, app_name: str, route_prefix: str = "/") -> Any:
|
||||
from ray import serve # type: ignore
|
||||
|
||||
# If Ray Serve LLM isn't available, callers may pass a plain dict placeholder.
|
||||
# Running that through serve.run() results in a confusing TypeError; fail fast.
|
||||
if isinstance(app, dict):
|
||||
raise ValueError("invalid serve app object (Ray Serve LLM not available or build_openai_app failed)")
|
||||
|
||||
return serve.run(app, name=app_name, route_prefix=route_prefix)
|
||||
|
||||
def get_status(self) -> Any:
|
||||
from ray import serve # type: ignore
|
||||
|
||||
return serve.status()
|
||||
|
||||
def delete_app(self, *, app_name: str) -> Any:
|
||||
from ray import serve # type: ignore
|
||||
|
||||
return serve.delete(app_name)
|
||||
63
src/mvp/py/argus/service/serve_llm_config.py
Normal file
@@ -0,0 +1,63 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from .serving_spec import ResolvedServingSpec
|
||||
|
||||
|
||||
def _ensure_hf_env_defaults(env: dict[str, str]) -> dict[str, str]:
|
||||
out = dict(env or {})
|
||||
# Prefer existing values if present, but always force offline mode in the platform.
|
||||
out.setdefault("HF_HOME", "/private/hf")
|
||||
out.setdefault("HUGGINGFACE_HUB_CACHE", "/private/hf/hub")
|
||||
out.setdefault("TRANSFORMERS_CACHE", "/private/hf/transformers")
|
||||
out["HF_HUB_OFFLINE"] = "1"
|
||||
return out
|
||||
|
||||
|
||||
def build_llm_config_dict(
|
||||
resolved: ResolvedServingSpec,
|
||||
*,
|
||||
accelerator_type: str,
|
||||
runtime_env_env_vars: dict[str, str] | None,
|
||||
cpu_per_gpu: float = 1.0,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Pure builder: maps a platform ResolvedServingSpec to a Ray Serve LLM-like config.
|
||||
|
||||
We return a plain dict here to keep this layer unit-testable without depending on
|
||||
a specific Ray Serve LLM version. The reconciler (later milestone) can choose to
|
||||
instantiate `ray.serve.llm.LLMConfig` using this dict.
|
||||
"""
|
||||
if not accelerator_type:
|
||||
raise ValueError("accelerator_type is required")
|
||||
if resolved.num_replicas < 1:
|
||||
raise ValueError("num_replicas must be >= 1")
|
||||
if resolved.gpus_per_replica < 1:
|
||||
raise ValueError("gpus_per_replica must be >= 1")
|
||||
if cpu_per_gpu <= 0:
|
||||
raise ValueError("cpu_per_gpu must be > 0")
|
||||
|
||||
engine_kwargs: dict[str, Any] = dict(resolved.engine_kwargs or {})
|
||||
# Enforce tensor parallel mapping; user-provided value must not contradict requested GPUs.
|
||||
engine_kwargs["tensor_parallel_size"] = int(resolved.gpus_per_replica)
|
||||
|
||||
# Ray Serve LLM (Ray 2.49.x) exposes `resources_per_bundle` instead of the older
|
||||
# `placement_group_config`. Use a single bundle that reserves the full GPU set
|
||||
# required by tensor-parallel execution.
|
||||
resources_per_bundle = {
|
||||
"GPU": float(resolved.gpus_per_replica),
|
||||
"CPU": float(cpu_per_gpu) * float(resolved.gpus_per_replica),
|
||||
}
|
||||
|
||||
env_vars = _ensure_hf_env_defaults(dict(runtime_env_env_vars or {}))
|
||||
|
||||
return {
|
||||
# Ray Serve LLM expects `model_loading_config` with model_id/model_source.
|
||||
"model_loading_config": {"model_id": resolved.model_id, "model_source": resolved.model_source},
|
||||
"accelerator_type": accelerator_type,
|
||||
"deployment_config": {"num_replicas": int(resolved.num_replicas)},
|
||||
"engine_kwargs": engine_kwargs,
|
||||
"resources_per_bundle": resources_per_bundle,
|
||||
"runtime_env": {"env_vars": env_vars},
|
||||
}
|
||||
150
src/mvp/py/argus/service/serving_reconciler.py
Normal file
@@ -0,0 +1,150 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import traceback
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Protocol
|
||||
|
||||
from argus.service.ray_resources import ClusterAvailable, get_cluster_available
|
||||
|
||||
from .config import V2Config
|
||||
from .db import Db
|
||||
from .serve_llm_config import build_llm_config_dict
|
||||
from .serving_spec import ResolvedServingSpec
|
||||
|
||||
|
||||
class ServeClient(Protocol):
|
||||
def ensure_started(self) -> None: ...
|
||||
|
||||
def apply_app(self, *, app: Any, app_name: str, route_prefix: str = "/") -> Any: ...
|
||||
|
||||
def delete_app(self, *, app_name: str) -> Any: ...
|
||||
|
||||
def get_status(self) -> Any: ...
|
||||
|
||||
|
||||
def _parse_engine_kwargs(row: dict[str, Any]) -> dict[str, Any] | None:
|
||||
raw = row.get("engine_kwargs_json")
|
||||
if raw in (None, ""):
|
||||
return None
|
||||
try:
|
||||
obj = json.loads(str(raw))
|
||||
return obj if isinstance(obj, dict) else None
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def _row_to_resolved_spec(row: dict[str, Any]) -> ResolvedServingSpec:
|
||||
return ResolvedServingSpec(
|
||||
user_id=str(row["user_id"]),
|
||||
model_id_suffix=str(row["model_id_suffix"]),
|
||||
model_id_prefix=str(row["model_id_prefix"]),
|
||||
model_id=str(row["model_id"]),
|
||||
model_source=str(row["model_source"]),
|
||||
num_replicas=int(row["num_replicas"]),
|
||||
gpus_per_replica=int(row["gpus_per_replica"]),
|
||||
engine_kwargs=_parse_engine_kwargs(row),
|
||||
)
|
||||
|
||||
|
||||
def _needed_model_gpus(row: dict[str, Any]) -> int:
|
||||
return int(row.get("num_replicas") or 0) * int(row.get("gpus_per_replica") or 0)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ServingReconciler:
|
||||
"""
|
||||
v3.8 (per-model apps): reconcile serve_models (SQLite) into per-model Ray Serve apps.
|
||||
|
||||
This reconciler is intentionally conservative:
|
||||
- Only acts on models in states QUEUED/DELETING.
|
||||
- Performs a minimal GPU precheck using ray available GPU totals (per-model).
|
||||
- Writes events and state transitions for explainability.
|
||||
"""
|
||||
|
||||
db: Db
|
||||
v2_cfg: V2Config
|
||||
ray_runtime_env_env_vars: dict[str, str]
|
||||
serve_client: ServeClient
|
||||
cpu_per_gpu: float = 1.0
|
||||
get_available_fn: Any = get_cluster_available
|
||||
|
||||
def tick(self) -> None:
|
||||
# Pick the next desired change.
|
||||
change = self.db.pick_next_runnable_serve_change()
|
||||
if not change:
|
||||
return
|
||||
|
||||
model_key = str(change["model_key"])
|
||||
state = str(change.get("state") or "")
|
||||
|
||||
# Ensure Ray (and Serve) can be started before doing anything else.
|
||||
try:
|
||||
self.serve_client.ensure_started()
|
||||
except Exception as e:
|
||||
self.db.append_serve_event(model_key=model_key, event_type="SERVE_START_ERROR", payload_json=repr(e))
|
||||
return
|
||||
|
||||
# Deleting is per-model: remove the model's app by name. This avoids touching other models.
|
||||
if state == "DELETING":
|
||||
try:
|
||||
self.db.append_serve_event(model_key=model_key, event_type="SERVE_DELETE_REQUESTED", payload_json=None)
|
||||
self.serve_client.delete_app(app_name=model_key)
|
||||
except Exception as e:
|
||||
err = f"{type(e).__name__}: {e}"
|
||||
tb = traceback.format_exc(limit=10)
|
||||
self.db.set_serve_model_state(
|
||||
model_key=model_key, state="FAILED", error_summary=err, event_type="SERVE_DELETE_FAILED", payload_json=tb
|
||||
)
|
||||
return
|
||||
self.db.set_serve_model_state(model_key=model_key, state="DELETED", event_type="SERVE_DELETE_APPLIED")
|
||||
return
|
||||
|
||||
# Precheck resources: per-model apply needs enough GPUs for this model only.
|
||||
needed = _needed_model_gpus(change)
|
||||
avail: ClusterAvailable = self.get_available_fn()
|
||||
if float(avail.total_available_gpus) < float(needed):
|
||||
msg = f"Insufficient GPUs: need {needed}, available {avail.total_available_gpus}"
|
||||
self.db.append_serve_event(model_key=model_key, event_type="SERVE_PENDING_RESOURCES", payload_json=msg)
|
||||
return
|
||||
|
||||
# Build the single-model LLM config (dict form for unit tests).
|
||||
accelerator_type = str(self.v2_cfg.serving.llm.accelerator_type or "")
|
||||
resolved = _row_to_resolved_spec(change)
|
||||
llm_cfg_dicts = [
|
||||
build_llm_config_dict(
|
||||
resolved,
|
||||
accelerator_type=accelerator_type,
|
||||
runtime_env_env_vars=self.ray_runtime_env_env_vars,
|
||||
cpu_per_gpu=self.cpu_per_gpu,
|
||||
)
|
||||
]
|
||||
|
||||
# Build a Ray Serve OpenAI-compatible app if Ray Serve LLM is available.
|
||||
# Fall back to a plain dict so unit tests can run without real Ray Serve.
|
||||
app_obj: Any
|
||||
try:
|
||||
from ray.serve.llm import LLMConfig, build_openai_app # type: ignore
|
||||
|
||||
llm_cfgs = [LLMConfig(**d) for d in llm_cfg_dicts]
|
||||
app_obj = build_openai_app({"llm_configs": llm_cfgs})
|
||||
except Exception as e:
|
||||
self.db.append_serve_event(model_key=model_key, event_type="SERVE_LLM_IMPORT_ERROR", payload_json=repr(e))
|
||||
app_obj = {"llm_configs": llm_cfg_dicts}
|
||||
|
||||
try:
|
||||
self.db.append_serve_event(model_key=model_key, event_type="SERVE_APPLY_REQUESTED", payload_json="1")
|
||||
self.serve_client.apply_app(app=app_obj, app_name=model_key, route_prefix=f"/serve/{model_key}")
|
||||
except Exception as e:
|
||||
err = f"{type(e).__name__}: {e}"
|
||||
tb = traceback.format_exc(limit=10)
|
||||
self.db.set_serve_model_state(model_key=model_key, state="FAILED", error_summary=err, event_type="SERVE_APPLY_FAILED", payload_json=tb)
|
||||
return
|
||||
|
||||
# Mark as deploying; best-effort status probe can promote to RUNNING.
|
||||
self.db.set_serve_model_state(model_key=model_key, state="DEPLOYING", event_type="SERVE_DEPLOYING")
|
||||
try:
|
||||
_ = self.serve_client.get_status()
|
||||
self.db.set_serve_model_state(model_key=model_key, state="RUNNING", event_type="SERVE_RUNNING")
|
||||
except Exception as e:
|
||||
self.db.append_serve_event(model_key=model_key, event_type="SERVE_STATUS_ERROR", payload_json=repr(e))
|
||||
144
src/mvp/py/argus/service/serving_spec.py
Normal file
@@ -0,0 +1,144 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
|
||||
|
||||
_MODEL_ID_SUFFIX_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,63}$")
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ServingSpec:
|
||||
model_id: str
|
||||
model_source: str
|
||||
num_replicas: int = 1
|
||||
gpus_per_replica: int = 1
|
||||
engine_kwargs: dict[str, Any] | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ResolvedServingSpec:
|
||||
user_id: str
|
||||
model_id_suffix: str
|
||||
model_id_prefix: str
|
||||
model_id: str
|
||||
model_source: str
|
||||
num_replicas: int
|
||||
gpus_per_replica: int
|
||||
engine_kwargs: dict[str, Any] | None
|
||||
|
||||
|
||||
def validate_model_id_suffix(suffix: str) -> None:
|
||||
if not isinstance(suffix, str):
|
||||
raise ValueError("model_id must be a string")
|
||||
s = suffix.strip()
|
||||
if s != suffix:
|
||||
raise ValueError("model_id must not contain leading/trailing whitespace")
|
||||
if not s:
|
||||
raise ValueError("model_id is required")
|
||||
if not _MODEL_ID_SUFFIX_RE.match(s):
|
||||
raise ValueError("model_id must match regex: ^[A-Za-z0-9][A-Za-z0-9._-]{0,63}$")
|
||||
if ".." in s:
|
||||
raise ValueError("model_id must not contain '..'")
|
||||
|
||||
|
||||
def make_model_id_prefix(*, user_id: str, now_utc: datetime | None = None) -> str:
|
||||
if not user_id or not isinstance(user_id, str):
|
||||
raise ValueError("user_id is required")
|
||||
if "/" in user_id:
|
||||
raise ValueError("user_id must not contain '/'")
|
||||
|
||||
dt = now_utc or datetime.now(timezone.utc)
|
||||
if dt.tzinfo is None:
|
||||
dt = dt.replace(tzinfo=timezone.utc)
|
||||
dt = dt.astimezone(timezone.utc)
|
||||
stamp = dt.strftime("%Y%m%d%H%M")
|
||||
return f"{user_id}-{stamp}"
|
||||
|
||||
|
||||
def expand_home_macros(*, user_id: str, text: str) -> str:
|
||||
if not isinstance(text, str):
|
||||
raise ValueError("model_source must be a string")
|
||||
if not text:
|
||||
raise ValueError("model_source is required")
|
||||
|
||||
out = text
|
||||
out = out.replace("$HOME/common/hf", "/private/hf")
|
||||
out = out.replace("$HOME/common/datasets", "/private/datasets")
|
||||
out = out.replace("$HOME", f"/private/users/{user_id}")
|
||||
return out
|
||||
|
||||
|
||||
def validate_model_source_path(*, user_id: str, model_source: str) -> None:
|
||||
if not isinstance(model_source, str):
|
||||
raise ValueError("model_source must be a string")
|
||||
if not model_source.startswith("/"):
|
||||
raise ValueError("model_source must be an absolute path")
|
||||
if not model_source.startswith("/private/"):
|
||||
raise ValueError("model_source must be under /private")
|
||||
if "\x00" in model_source:
|
||||
raise ValueError("model_source contains null byte")
|
||||
parts = [p for p in model_source.split("/") if p]
|
||||
if any(p == ".." for p in parts):
|
||||
raise ValueError("model_source must not contain '..'")
|
||||
|
||||
allowed_user_prefix = f"/private/users/{user_id}/"
|
||||
allowed = model_source.startswith("/private/hf/") or model_source.startswith(allowed_user_prefix)
|
||||
if not allowed:
|
||||
raise PermissionError("model_source is not allowed (must be under /private/hf or your /private/users/<user_id>)")
|
||||
|
||||
|
||||
def parse_serving_spec(obj: Any) -> ServingSpec:
|
||||
if not isinstance(obj, dict):
|
||||
raise ValueError("serving spec must be a mapping")
|
||||
|
||||
model_id = obj.get("model_id")
|
||||
model_source = obj.get("model_source")
|
||||
num_replicas = obj.get("num_replicas", 1)
|
||||
gpus_per_replica = obj.get("gpus_per_replica", 1)
|
||||
engine_kwargs = obj.get("engine_kwargs", None)
|
||||
|
||||
if not isinstance(model_id, str):
|
||||
raise ValueError("missing required field: model_id")
|
||||
validate_model_id_suffix(model_id)
|
||||
|
||||
if not isinstance(model_source, str) or not model_source:
|
||||
raise ValueError("missing required field: model_source")
|
||||
|
||||
if not isinstance(num_replicas, int) or num_replicas < 1:
|
||||
raise ValueError("num_replicas must be an integer >= 1")
|
||||
if not isinstance(gpus_per_replica, int) or gpus_per_replica < 1:
|
||||
raise ValueError("gpus_per_replica must be an integer >= 1")
|
||||
|
||||
if engine_kwargs is not None and not isinstance(engine_kwargs, dict):
|
||||
raise ValueError("engine_kwargs must be a mapping when provided")
|
||||
|
||||
return ServingSpec(
|
||||
model_id=model_id,
|
||||
model_source=model_source,
|
||||
num_replicas=num_replicas,
|
||||
gpus_per_replica=gpus_per_replica,
|
||||
engine_kwargs=engine_kwargs,
|
||||
)
|
||||
|
||||
|
||||
def resolve_serving_spec(*, spec: ServingSpec, user_id: str, now_utc: datetime | None = None) -> ResolvedServingSpec:
|
||||
validate_model_id_suffix(spec.model_id)
|
||||
prefix = make_model_id_prefix(user_id=user_id, now_utc=now_utc)
|
||||
full_model_id = f"{prefix}-{spec.model_id}"
|
||||
|
||||
resolved_source = expand_home_macros(user_id=user_id, text=spec.model_source)
|
||||
validate_model_source_path(user_id=user_id, model_source=resolved_source)
|
||||
|
||||
return ResolvedServingSpec(
|
||||
user_id=user_id,
|
||||
model_id_suffix=spec.model_id,
|
||||
model_id_prefix=prefix,
|
||||
model_id=full_model_id,
|
||||
model_source=resolved_source,
|
||||
num_replicas=spec.num_replicas,
|
||||
gpus_per_replica=spec.gpus_per_replica,
|
||||
engine_kwargs=spec.engine_kwargs,
|
||||
)
|
||||
228
src/mvp/py/argus/service/sftpgo.py
Normal file
@@ -0,0 +1,228 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import json
|
||||
from dataclasses import dataclass
|
||||
from urllib.error import HTTPError, URLError
|
||||
from urllib.request import Request, urlopen
|
||||
|
||||
|
||||
class SFTPGoError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SFTPGoAdminClient:
|
||||
"""
|
||||
Minimal SFTPGo admin API client (v3.0).
|
||||
|
||||
SFTPGo OpenAPI documents the admin token flow:
|
||||
- GET /api/v2/token with HTTP BasicAuth -> returns {"access_token": "..."}.
|
||||
- Use Authorization: Bearer <token> for admin endpoints like POST /api/v2/users.
|
||||
|
||||
See upstream OpenAPI for details:
|
||||
https://raw.githubusercontent.com/drakkan/sftpgo/main/openapi/openapi.yaml
|
||||
"""
|
||||
|
||||
admin_api_base: str
|
||||
admin_user: str
|
||||
admin_password: str
|
||||
common_root: str = "/private/common"
|
||||
|
||||
def _url(self, path: str) -> str:
|
||||
base = self.admin_api_base.rstrip("/")
|
||||
p = path if path.startswith("/") else f"/{path}"
|
||||
return f"{base}{p}"
|
||||
|
||||
def _basic_auth_header(self) -> str:
|
||||
raw = f"{self.admin_user}:{self.admin_password}".encode("utf-8")
|
||||
return "Basic " + base64.b64encode(raw).decode("ascii")
|
||||
|
||||
def _get_json(self, url: str, headers: dict[str, str]) -> dict:
|
||||
req = Request(url, headers=headers, method="GET")
|
||||
try:
|
||||
with urlopen(req, timeout=10) as resp:
|
||||
return json.loads(resp.read().decode("utf-8"))
|
||||
except HTTPError as e:
|
||||
raise SFTPGoError(f"sftpgo http error: {e.code} {e.reason}") from e
|
||||
except URLError as e:
|
||||
raise SFTPGoError(f"sftpgo connection error: {e!r}") from e
|
||||
|
||||
def _post_json(self, url: str, payload: dict, headers: dict[str, str]) -> None:
|
||||
data = json.dumps(payload).encode("utf-8")
|
||||
req = Request(url, data=data, headers=headers, method="POST")
|
||||
try:
|
||||
with urlopen(req, timeout=10) as resp:
|
||||
_ = resp.read()
|
||||
except HTTPError as e:
|
||||
raise SFTPGoError(f"sftpgo http error: {e.code} {e.reason}") from e
|
||||
except URLError as e:
|
||||
raise SFTPGoError(f"sftpgo connection error: {e!r}") from e
|
||||
|
||||
def _put_json(self, url: str, payload: dict, headers: dict[str, str]) -> None:
|
||||
data = json.dumps(payload).encode("utf-8")
|
||||
req = Request(url, data=data, headers=headers, method="PUT")
|
||||
try:
|
||||
with urlopen(req, timeout=10) as resp:
|
||||
_ = resp.read()
|
||||
except HTTPError as e:
|
||||
raise SFTPGoError(f"sftpgo http error: {e.code} {e.reason}") from e
|
||||
except URLError as e:
|
||||
raise SFTPGoError(f"sftpgo connection error: {e!r}") from e
|
||||
|
||||
def _admin_token(self) -> str:
|
||||
url = self._url("/token")
|
||||
obj = self._get_json(url, headers={"Authorization": self._basic_auth_header()})
|
||||
tok = str(obj.get("access_token") or "").strip()
|
||||
if not tok:
|
||||
raise SFTPGoError("sftpgo token response missing access_token")
|
||||
return tok
|
||||
|
||||
def _auth_headers(self, tok: str) -> dict[str, str]:
|
||||
return {"Authorization": f"Bearer {tok}", "Content-Type": "application/json"}
|
||||
|
||||
def _ensure_folder(self, *, tok: str, name: str, mapped_path: str) -> None:
|
||||
url = self._url("/folders")
|
||||
try:
|
||||
self._post_json(url, {"name": name, "mapped_path": mapped_path}, headers=self._auth_headers(tok))
|
||||
except SFTPGoError as e:
|
||||
# Idempotent + self-healing:
|
||||
# If it already exists, update mapped_path to the desired value.
|
||||
if "409" in str(e):
|
||||
self._put_json(
|
||||
self._url(f"/folders/{name}"),
|
||||
{"name": name, "mapped_path": mapped_path},
|
||||
headers=self._auth_headers(tok),
|
||||
)
|
||||
return
|
||||
raise
|
||||
|
||||
def _ensure_common_folders(self, *, tok: str) -> None:
|
||||
# Important: do NOT map datasets to /private/common/datasets because that path is
|
||||
# a symlink farm (e.g. gsm8k -> /private/datasets/gsm8k) and SFTPGo WebClient
|
||||
# will treat symlink traversal as escaping the virtual folder root, resulting
|
||||
# in "permission denied". Map directly to the real datasets root.
|
||||
common = self.common_root.rstrip("/")
|
||||
if common.endswith("/common"):
|
||||
shared_root = common[: -len("/common")] or "/private"
|
||||
else:
|
||||
shared_root = "/private"
|
||||
|
||||
self._ensure_folder(tok=tok, name="common_datasets", mapped_path=f"{shared_root}/datasets")
|
||||
# Expose HF cache read-only so users can inspect downloaded models/datasets.
|
||||
self._ensure_folder(tok=tok, name="common_hf", mapped_path=f"{shared_root}/hf")
|
||||
|
||||
def _get_user(self, *, tok: str, username: str) -> dict:
|
||||
url = self._url(f"/users/{username}")
|
||||
return self._get_json(url, headers={"Authorization": f"Bearer {tok}"})
|
||||
|
||||
def _put_user(self, *, tok: str, username: str, payload: dict) -> None:
|
||||
url = self._url(f"/users/{username}")
|
||||
self._put_json(url, payload, headers=self._auth_headers(tok))
|
||||
|
||||
def _apply_common_readonly(self, user_payload: dict) -> dict:
|
||||
# Path-based permissions: make /common/* read-only while keeping home writeable.
|
||||
perms = dict(user_payload.get("permissions") or {"/": ["*"]})
|
||||
# Ensure /common is visible as a directory and can be traversed.
|
||||
perms["/common"] = ["list"]
|
||||
# SFTPGo permissions are path-scoped. In practice, granting on the directory itself may
|
||||
# not always cover nested paths depending on client/API behavior, so include common
|
||||
# wildcard forms to ensure users can read files under these virtual folders.
|
||||
perms["/common/datasets"] = ["list"]
|
||||
perms["/common/datasets/*"] = ["list", "download"]
|
||||
perms["/common/hf"] = ["list"]
|
||||
perms["/common/hf/*"] = ["list", "download"]
|
||||
user_payload["permissions"] = perms
|
||||
|
||||
desired_vf = [
|
||||
{"name": "common_datasets", "virtual_path": "/common/datasets"},
|
||||
{"name": "common_hf", "virtual_path": "/common/hf"},
|
||||
]
|
||||
existing = user_payload.get("virtual_folders") or []
|
||||
if not isinstance(existing, list):
|
||||
existing = []
|
||||
seen = {(vf.get("name"), vf.get("virtual_path")) for vf in existing if isinstance(vf, dict)}
|
||||
merged = list(existing)
|
||||
for vf in desired_vf:
|
||||
key = (vf["name"], vf["virtual_path"])
|
||||
if key not in seen:
|
||||
merged.append(vf)
|
||||
user_payload["virtual_folders"] = merged
|
||||
return user_payload
|
||||
|
||||
def create_user(self, *, username: str, password: str, home_dir: str) -> None:
|
||||
tok = self._admin_token()
|
||||
self._ensure_common_folders(tok=tok)
|
||||
url = self._url("/users")
|
||||
payload: dict = {
|
||||
"status": 1,
|
||||
"username": username,
|
||||
"password": password,
|
||||
"home_dir": home_dir,
|
||||
"permissions": {
|
||||
"/": ["*"],
|
||||
"/common": ["list"],
|
||||
"/common/datasets": ["list"],
|
||||
"/common/datasets/*": ["list", "download"],
|
||||
"/common/hf": ["list"],
|
||||
"/common/hf/*": ["list", "download"],
|
||||
},
|
||||
"virtual_folders": [
|
||||
{"name": "common_datasets", "virtual_path": "/common/datasets"},
|
||||
{"name": "common_hf", "virtual_path": "/common/hf"},
|
||||
],
|
||||
}
|
||||
self._post_json(
|
||||
url,
|
||||
payload,
|
||||
headers=self._auth_headers(tok),
|
||||
)
|
||||
|
||||
def disable_user(self, *, username: str, home_dir: str) -> None:
|
||||
tok = self._admin_token()
|
||||
self._ensure_common_folders(tok=tok)
|
||||
cur = self._get_user(tok=tok, username=username)
|
||||
payload: dict = {
|
||||
"username": username,
|
||||
"status": 0,
|
||||
"home_dir": home_dir,
|
||||
"uid": cur.get("uid", 0),
|
||||
"gid": cur.get("gid", 0),
|
||||
"permissions": cur.get("permissions") or {"/": ["*"]},
|
||||
"virtual_folders": cur.get("virtual_folders") or [],
|
||||
}
|
||||
self._apply_common_readonly(payload)
|
||||
self._put_user(tok=tok, username=username, payload=payload)
|
||||
|
||||
def enable_user(self, *, username: str, home_dir: str) -> None:
|
||||
tok = self._admin_token()
|
||||
self._ensure_common_folders(tok=tok)
|
||||
cur = self._get_user(tok=tok, username=username)
|
||||
payload: dict = {
|
||||
"username": username,
|
||||
"status": 1,
|
||||
"home_dir": home_dir,
|
||||
"uid": cur.get("uid", 0),
|
||||
"gid": cur.get("gid", 0),
|
||||
"permissions": cur.get("permissions") or {"/": ["*"]},
|
||||
"virtual_folders": cur.get("virtual_folders") or [],
|
||||
}
|
||||
self._apply_common_readonly(payload)
|
||||
self._put_user(tok=tok, username=username, payload=payload)
|
||||
|
||||
def reset_password(self, *, username: str, new_password: str, home_dir: str) -> None:
|
||||
tok = self._admin_token()
|
||||
self._ensure_common_folders(tok=tok)
|
||||
cur = self._get_user(tok=tok, username=username)
|
||||
payload: dict = {
|
||||
"username": username,
|
||||
"status": 1,
|
||||
"home_dir": home_dir,
|
||||
"uid": cur.get("uid", 0),
|
||||
"gid": cur.get("gid", 0),
|
||||
"permissions": cur.get("permissions") or {"/": ["*"]},
|
||||
"virtual_folders": cur.get("virtual_folders") or [],
|
||||
"password": new_password,
|
||||
}
|
||||
self._apply_common_readonly(payload)
|
||||
self._put_user(tok=tok, username=username, payload=payload)
|
||||
11
src/mvp/py/argus/service/ui.py
Normal file
@@ -0,0 +1,11 @@
|
||||
from __future__ import annotations
|
||||
|
||||
# Backward-compatible import path.
|
||||
#
|
||||
# The UI implementation lives under `argus.ui` (split into multiple modules for maintainability),
|
||||
# while `argus.service.app` continues to import `argus.service.ui.register_ui_routes`.
|
||||
|
||||
from argus.ui import register_ui_routes
|
||||
|
||||
__all__ = ["register_ui_routes"]
|
||||
|
||||
6
src/mvp/py/argus/ui/__init__.py
Normal file
@@ -0,0 +1,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from .routes import register_ui_routes
|
||||
|
||||
__all__ = ["register_ui_routes"]
|
||||
|
||||
2
src/mvp/py/argus/ui/assets/__init__.py
Normal file
@@ -0,0 +1,2 @@
|
||||
from __future__ import annotations
|
||||
|
||||
33
src/mvp/py/argus/ui/assets/base_css.py
Normal file
@@ -0,0 +1,33 @@
|
||||
from __future__ import annotations
|
||||
|
||||
BASE_CSS = """
|
||||
:root { --bg:#0b1020; --panel:#111a33; --muted:#95a3c6; --fg:#e8eeff; --accent:#7aa2ff; --danger:#ff6b6b; --ok:#3ddc97; }
|
||||
* { box-sizing: border-box; }
|
||||
body { margin:0; font-family: ui-sans-serif, system-ui, -apple-system, Segoe UI, Roboto, Helvetica, Arial; background:var(--bg); color:var(--fg); }
|
||||
a { color:var(--accent); text-decoration:none; }
|
||||
.layout { display:flex; min-height:100vh; }
|
||||
.nav { width: 240px; padding:16px; background: linear-gradient(180deg, #0e1630, #0b1020); border-right: 1px solid rgba(255,255,255,0.06); }
|
||||
.brand { font-weight: 700; letter-spacing: .2px; margin-bottom: 12px; }
|
||||
.nav a { display:block; padding:10px 10px; border-radius:10px; color: var(--fg); opacity: .9; }
|
||||
.nav a.active { background: rgba(122,162,255,0.14); border: 1px solid rgba(122,162,255,0.22); }
|
||||
.nav a:hover { background: rgba(255,255,255,0.06); }
|
||||
.main { flex:1; padding: 20px 24px; }
|
||||
.card { background: rgba(255,255,255,0.04); border: 1px solid rgba(255,255,255,0.08); border-radius: 14px; padding: 16px; }
|
||||
.row { display:flex; gap: 12px; align-items:center; flex-wrap: wrap; }
|
||||
.muted { color: var(--muted); }
|
||||
.btn { border: 1px solid rgba(255,255,255,0.16); background: rgba(255,255,255,0.06); color: var(--fg); padding: 10px 12px; border-radius: 10px; cursor: pointer; }
|
||||
.btn:hover { background: rgba(255,255,255,0.10); }
|
||||
.btn.active { background: rgba(122,162,255,0.14); border-color: rgba(122,162,255,0.22); }
|
||||
.btn.danger { border-color: rgba(255,107,107,0.35); background: rgba(255,107,107,0.10); }
|
||||
.pill { display:inline-block; padding: 2px 10px; border-radius: 999px; border: 1px solid rgba(255,255,255,0.16); font-size: 12px; }
|
||||
.pill.ok { border-color: rgba(61,220,151,0.35); background: rgba(61,220,151,0.12); }
|
||||
.pill.bad { border-color: rgba(255,107,107,0.35); background: rgba(255,107,107,0.12); }
|
||||
textarea, input { width: 100%; color: var(--fg); background: rgba(255,255,255,0.06); border: 1px solid rgba(255,255,255,0.12); border-radius: 12px; padding: 10px 12px; outline: none; }
|
||||
select { color: var(--fg); background: rgba(255,255,255,0.06); border: 1px solid rgba(255,255,255,0.12); border-radius: 12px; padding: 10px 12px; outline: none; }
|
||||
select option { color: var(--fg); background: #0e1630; }
|
||||
button:disabled { opacity: .45; cursor: not-allowed; }
|
||||
pre { white-space: pre-wrap; word-break: break-word; }
|
||||
table { width:100%; border-collapse: collapse; }
|
||||
th, td { padding: 10px 8px; border-bottom: 1px solid rgba(255,255,255,0.08); text-align:left; }
|
||||
""".strip()
|
||||
|
||||
70
src/mvp/py/argus/ui/assets/base_js.py
Normal file
@@ -0,0 +1,70 @@
|
||||
from __future__ import annotations
|
||||
|
||||
BASE_JS = """
|
||||
function mvpTokenGet() {
|
||||
return (localStorage.getItem("mvp_token") || "").trim();
|
||||
}
|
||||
function mvpTokenSet(v) {
|
||||
localStorage.setItem("mvp_token", (v || "").trim());
|
||||
}
|
||||
function mvpSftpPasswordGet() {
|
||||
return (localStorage.getItem("mvp_sftp_password") || "").trim();
|
||||
}
|
||||
function mvpSftpPasswordSet(v) {
|
||||
localStorage.setItem("mvp_sftp_password", (v || "").trim());
|
||||
}
|
||||
async function apiFetch(path, opts) {
|
||||
opts = opts || {};
|
||||
opts.headers = opts.headers || {};
|
||||
const tok = mvpTokenGet();
|
||||
if (tok) opts.headers["Authorization"] = "Bearer " + tok;
|
||||
return fetch(path, opts);
|
||||
}
|
||||
async function apiJson(path, opts) {
|
||||
const resp = await apiFetch(path, opts);
|
||||
const text = await resp.text();
|
||||
if (!resp.ok) {
|
||||
const err = new Error("HTTP " + resp.status);
|
||||
err.status = resp.status;
|
||||
err.body = text;
|
||||
throw err;
|
||||
}
|
||||
return JSON.parse(text);
|
||||
}
|
||||
function fmtJson(obj) {
|
||||
try { return JSON.stringify(obj, null, 2); } catch (e) { return String(obj); }
|
||||
}
|
||||
function curOriginWithPort(port) {
|
||||
const proto = window.location.protocol;
|
||||
const host = window.location.hostname;
|
||||
return proto + "//" + host + ":" + port;
|
||||
}
|
||||
async function copyText(v) {
|
||||
if (!v) return false;
|
||||
try {
|
||||
await navigator.clipboard.writeText(v);
|
||||
return true;
|
||||
} catch (e) {
|
||||
// Fallback for non-secure contexts (http) or older browsers.
|
||||
try {
|
||||
const ta = document.createElement("textarea");
|
||||
ta.value = v;
|
||||
ta.style.position = "fixed";
|
||||
ta.style.opacity = "0";
|
||||
document.body.appendChild(ta);
|
||||
ta.focus();
|
||||
ta.select();
|
||||
const ok = document.execCommand("copy");
|
||||
document.body.removeChild(ta);
|
||||
return ok;
|
||||
} catch (e2) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
}
|
||||
document.addEventListener("DOMContentLoaded", () => {
|
||||
const el = document.getElementById("nav-ray-dashboard");
|
||||
if (el) el.href = curOriginWithPort(8265);
|
||||
});
|
||||
""".strip()
|
||||
|
||||
2
src/mvp/py/argus/ui/layout/__init__.py
Normal file
@@ -0,0 +1,2 @@
|
||||
from __future__ import annotations
|
||||
|
||||
24
src/mvp/py/argus/ui/layout/nav.py
Normal file
@@ -0,0 +1,24 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import html
|
||||
|
||||
|
||||
def nav(active: str) -> str:
|
||||
links = [
|
||||
("login", "/ui/login", "Login"),
|
||||
("tasks", "/ui/tasks", "Tasks"),
|
||||
("serving", "/ui/serving", "Serving"),
|
||||
("new", "/ui/tasks/new", "New Task"),
|
||||
("data", "/ui/data", "Data"),
|
||||
("admin", "/ui/admin", "Admin"),
|
||||
("ray", "#", "Ray Dashboard"),
|
||||
]
|
||||
items: list[str] = []
|
||||
for key, href, label in links:
|
||||
cls = "active" if key == active else ""
|
||||
extra = ""
|
||||
if key == "ray":
|
||||
extra = ' id="nav-ray-dashboard" target="_blank" rel="noopener"'
|
||||
items.append(f'<a class="{cls}" href="{href}"{extra}>{html.escape(label)}</a>')
|
||||
return "\n".join(items)
|
||||
|
||||
36
src/mvp/py/argus/ui/layout/page.py
Normal file
@@ -0,0 +1,36 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import html
|
||||
|
||||
from ..assets.base_css import BASE_CSS
|
||||
from ..assets.base_js import BASE_JS
|
||||
from .nav import nav
|
||||
|
||||
|
||||
def page(title: str, active: str, body: str, script: str = "") -> str:
|
||||
return f"""<!doctype html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="utf-8" />
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1" />
|
||||
<title>{html.escape(title)}</title>
|
||||
<style>{BASE_CSS}</style>
|
||||
</head>
|
||||
<body>
|
||||
<div class="layout">
|
||||
<nav class="nav">
|
||||
<div class="brand">Argus MVP</div>
|
||||
{nav(active)}
|
||||
<div style="margin-top:12px" class="muted">
|
||||
Token stored in browser localStorage as <code>mvp_token</code>.
|
||||
</div>
|
||||
</nav>
|
||||
<main class="main">
|
||||
{body}
|
||||
</main>
|
||||
</div>
|
||||
<script>{BASE_JS}</script>
|
||||
<script>{script}</script>
|
||||
</body>
|
||||
</html>"""
|
||||
|
||||
2
src/mvp/py/argus/ui/pages/__init__.py
Normal file
@@ -0,0 +1,2 @@
|
||||
from __future__ import annotations
|
||||
|
||||
124
src/mvp/py/argus/ui/pages/admin.py
Normal file
@@ -0,0 +1,124 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from fastapi import FastAPI
|
||||
from fastapi.responses import HTMLResponse
|
||||
|
||||
from ..layout.page import page
|
||||
|
||||
|
||||
def register(app: FastAPI) -> None:
|
||||
@app.get("/ui/admin")
|
||||
async def ui_admin() -> HTMLResponse:
|
||||
body = """
|
||||
<h1>Admin</h1>
|
||||
<div class="card">
|
||||
<div class="muted">
|
||||
This page requires the <code>admin</code> token (set it in <a href="/ui/login">Login</a>).
|
||||
</div>
|
||||
<div style="height:14px"></div>
|
||||
|
||||
<h3>Create user</h3>
|
||||
<div class="row">
|
||||
<input id="new-user-id" placeholder="user_id (e.g. alice)" style="max-width:320px" />
|
||||
<input id="new-display-name" placeholder="display_name (optional)" style="max-width:320px" />
|
||||
<button class="btn" id="create-user">Create</button>
|
||||
</div>
|
||||
<div style="height:10px"></div>
|
||||
<pre id="create-msg" class="muted"></pre>
|
||||
|
||||
<div style="height:14px"></div>
|
||||
<div class="row">
|
||||
<button class="btn" id="refresh">Refresh</button>
|
||||
</div>
|
||||
<div style="height:10px"></div>
|
||||
<div id="out" class="muted">Loading...</div>
|
||||
</div>
|
||||
""".strip()
|
||||
script = """
|
||||
const out = document.getElementById("out");
|
||||
const createMsg = document.getElementById("create-msg");
|
||||
const userIdEl = document.getElementById("new-user-id");
|
||||
const displayNameEl = document.getElementById("new-display-name");
|
||||
|
||||
function esc(s) {
|
||||
s = String(s || "");
|
||||
return s.replaceAll("&","&").replaceAll("<","<").replaceAll(">",">");
|
||||
}
|
||||
|
||||
async function refresh() {
|
||||
out.textContent = "Loading...";
|
||||
try {
|
||||
const obj = await apiJson("/api/v2/users?limit=200");
|
||||
const users = (obj.users || []);
|
||||
function row(u) {
|
||||
const uid = u.user_id;
|
||||
const tok = u.token || "";
|
||||
const tokShort = tok ? (tok.length > 18 ? (tok.slice(0, 18) + "…") : tok) : "";
|
||||
const created = u.created_at || "";
|
||||
const updated = u.updated_at || "";
|
||||
const tCreated = u.token_created_at || "";
|
||||
const tUsed = u.token_last_used_at || "";
|
||||
return `<tr>
|
||||
<td><code>${esc(uid)}</code></td>
|
||||
<td class="muted">${esc(created)}</td>
|
||||
<td class="muted">${esc(updated)}</td>
|
||||
<td>
|
||||
<div class="row" style="gap:8px">
|
||||
<code title="${esc(tok)}">${esc(tokShort)}</code>
|
||||
<button class="btn" data-copy="${esc(tok)}">Copy</button>
|
||||
<button class="btn" data-issue="${esc(uid)}">Issue token</button>
|
||||
</div>
|
||||
<div class="muted" style="margin-top:6px">token_created_at: ${esc(tCreated)}; last_used_at: ${esc(tUsed)}</div>
|
||||
</td>
|
||||
</tr>`;
|
||||
}
|
||||
out.innerHTML = `
|
||||
<table>
|
||||
<thead><tr><th>User</th><th>Created</th><th>Updated</th><th>Token</th></tr></thead>
|
||||
<tbody>${users.map(row).join("") || "<tr><td colspan=4 class=muted>(none)</td></tr>"}</tbody>
|
||||
</table>
|
||||
`;
|
||||
for (const btn of out.querySelectorAll("button[data-copy]")) {
|
||||
btn.onclick = async () => { await copyText(btn.getAttribute("data-copy") || ""); };
|
||||
}
|
||||
for (const btn of out.querySelectorAll("button[data-issue]")) {
|
||||
btn.onclick = async () => {
|
||||
const uid = btn.getAttribute("data-issue");
|
||||
if (!uid) return;
|
||||
try {
|
||||
const r = await apiJson("/api/v2/users/" + encodeURIComponent(uid) + "/tokens", { method: "POST" });
|
||||
createMsg.textContent = "Issued token:\\n" + fmtJson(r);
|
||||
await refresh();
|
||||
} catch (e) {
|
||||
createMsg.textContent = "Error issuing token: " + (e.status || "") + "\\n" + (e.body || String(e));
|
||||
}
|
||||
};
|
||||
}
|
||||
} catch (e) {
|
||||
out.textContent = "Error: " + (e.status || "") + "\\n" + (e.body || String(e));
|
||||
}
|
||||
}
|
||||
|
||||
document.getElementById("refresh").onclick = refresh;
|
||||
document.getElementById("create-user").onclick = async () => {
|
||||
createMsg.textContent = "Creating...";
|
||||
const user_id = (userIdEl.value || "").trim();
|
||||
const display_name = (displayNameEl.value || "").trim();
|
||||
if (!user_id) { createMsg.textContent = "user_id is required"; return; }
|
||||
const payload = { user_id: user_id };
|
||||
if (display_name) payload.display_name = display_name;
|
||||
try {
|
||||
const r = await apiJson("/api/v2/users", { method: "POST", headers: {"Content-Type":"application/json"}, body: JSON.stringify(payload) });
|
||||
createMsg.textContent = "Created:\\n" + fmtJson(r);
|
||||
userIdEl.value = "";
|
||||
displayNameEl.value = "";
|
||||
await refresh();
|
||||
} catch (e) {
|
||||
createMsg.textContent = "Error: " + (e.status || "") + "\\n" + (e.body || String(e));
|
||||
}
|
||||
};
|
||||
|
||||
refresh();
|
||||
""".strip()
|
||||
return HTMLResponse(content=page("Admin", "admin", body, script))
|
||||
|
||||
93
src/mvp/py/argus/ui/pages/data.py
Normal file
@@ -0,0 +1,93 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from fastapi import FastAPI
|
||||
from fastapi.responses import HTMLResponse
|
||||
|
||||
from ..layout.page import page
|
||||
|
||||
|
||||
def register(app: FastAPI) -> None:
|
||||
@app.get("/ui/data")
|
||||
async def ui_data() -> HTMLResponse:
|
||||
body = """
|
||||
<h1>Data</h1>
|
||||
<div class="card">
|
||||
<div class="muted">User files live under your home directory. Keep long-term artifacts in <code>models/</code> or <code>datasets/</code>.</div>
|
||||
<div style="height:14px"></div>
|
||||
|
||||
<div class="row">
|
||||
<div style="flex:1; min-width:260px">
|
||||
<div class="muted">Username</div>
|
||||
<div style="height:6px"></div>
|
||||
<div class="row" style="gap:8px">
|
||||
<input id="u" readonly />
|
||||
<button class="btn" id="copy-u">Copy</button>
|
||||
</div>
|
||||
</div>
|
||||
<div style="flex:1; min-width:260px">
|
||||
<div class="muted">SFTPGo password</div>
|
||||
<div style="height:6px"></div>
|
||||
<div class="row" style="gap:8px">
|
||||
<input id="p" placeholder="Click Reset to generate..." />
|
||||
<button class="btn" id="copy-p">Copy</button>
|
||||
<button class="btn" id="reset-p">Reset</button>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div style="height:12px"></div>
|
||||
<div class="row">
|
||||
<a class="btn" id="sftp-web" target="_blank" rel="noopener" href="#">Open SFTPGo Web Client (:8081)</a>
|
||||
</div>
|
||||
|
||||
<div style="height:12px"></div>
|
||||
<div class="muted">
|
||||
You can also use an SFTP client (e.g. FileZilla) with the same username/password.
|
||||
Host: <code id="sftp-host"></code>, Port: <code id="sftp-port"></code>.
|
||||
</div>
|
||||
|
||||
<div style="height:14px"></div>
|
||||
<div id="msg" class="muted"></div>
|
||||
</div>
|
||||
""".strip()
|
||||
script = """
|
||||
const msg = document.getElementById("msg");
|
||||
document.getElementById("nav-ray-dashboard").href = curOriginWithPort(8265);
|
||||
const u = document.getElementById("u");
|
||||
const p = document.getElementById("p");
|
||||
const sftpWeb = document.getElementById("sftp-web");
|
||||
const sftpHost = document.getElementById("sftp-host");
|
||||
const sftpPort = document.getElementById("sftp-port");
|
||||
document.getElementById("copy-u").onclick = async () => { await copyText(u.value || ""); };
|
||||
document.getElementById("copy-p").onclick = async () => { await copyText(p.value || ""); };
|
||||
|
||||
async function refresh() {
|
||||
const resp = await apiFetch("/api/v2/me");
|
||||
const text = await resp.text();
|
||||
if (!resp.ok) { msg.textContent = "Error: " + resp.status + "\\n" + text; return; }
|
||||
const obj = JSON.parse(text);
|
||||
u.value = (obj.user_id || "");
|
||||
const cached = mvpSftpPasswordGet();
|
||||
if (cached) p.value = cached;
|
||||
const host = curOriginWithPort(8081);
|
||||
sftpWeb.href = host + "/web/client";
|
||||
sftpHost.textContent = (obj.sftp && obj.sftp.host) ? obj.sftp.host : window.location.hostname;
|
||||
sftpPort.textContent = (obj.sftp && obj.sftp.port) ? String(obj.sftp.port) : "2022";
|
||||
msg.textContent = "";
|
||||
}
|
||||
document.getElementById("reset-p").onclick = async () => {
|
||||
p.value = "";
|
||||
mvpSftpPasswordSet("");
|
||||
msg.textContent = "Resetting...";
|
||||
const resp = await apiFetch("/api/v2/me/sftp:reset_password", { method: "POST" });
|
||||
const text = await resp.text();
|
||||
if (!resp.ok) { msg.textContent = "Error: " + resp.status + "\\n" + text; return; }
|
||||
const obj = JSON.parse(text);
|
||||
p.value = obj.password || "";
|
||||
mvpSftpPasswordSet(p.value);
|
||||
msg.textContent = "SFTPGo password rotated.";
|
||||
};
|
||||
refresh();
|
||||
""".strip()
|
||||
return HTMLResponse(content=page("Data", "data", body, script))
|
||||
|
||||
88
src/mvp/py/argus/ui/pages/login.py
Normal file
@@ -0,0 +1,88 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from fastapi import FastAPI
|
||||
from fastapi.responses import HTMLResponse
|
||||
|
||||
from ..layout.page import page
|
||||
|
||||
|
||||
def register(app: FastAPI) -> None:
|
||||
@app.get("/ui/login")
|
||||
async def ui_login() -> HTMLResponse:
|
||||
body = """
|
||||
<h1>Login</h1>
|
||||
<div class="card">
|
||||
<div class="muted">Paste your API token (without the <code>Bearer</code> prefix).</div>
|
||||
<div style="height:10px"></div>
|
||||
<input id="tok" placeholder="token..." />
|
||||
<div style="height:10px"></div>
|
||||
<div class="row">
|
||||
<button class="btn" id="save">Save</button>
|
||||
<button class="btn" id="clear">Clear</button>
|
||||
<a href="/ui/tasks" class="btn" style="display:inline-block">Go to Tasks</a>
|
||||
</div>
|
||||
<div style="height:10px"></div>
|
||||
<div id="msg" class="muted"></div>
|
||||
</div>
|
||||
<div style="height:12px"></div>
|
||||
<div class="card">
|
||||
<h3 style="margin-top:0">User Info</h3>
|
||||
<div class="muted">Shown after login via <code>/api/v2/me</code>.</div>
|
||||
<div style="height:10px"></div>
|
||||
<div class="row">
|
||||
<button class="btn" id="me-refresh">Refresh</button>
|
||||
</div>
|
||||
<div style="height:10px"></div>
|
||||
<pre id="me" class="muted">(not loaded)</pre>
|
||||
</div>
|
||||
<div style="height:12px"></div>
|
||||
<div class="card">
|
||||
<h3 style="margin-top:0">W&B</h3>
|
||||
<div class="muted">Weights & Biases local server (v3.6). Metrics are written by training jobs; this UI is for viewing.</div>
|
||||
<div style="height:10px"></div>
|
||||
<div class="row">
|
||||
<a class="btn" id="wandb-open" target="_blank" rel="noopener" href="#">Open W&B (:8090)</a>
|
||||
<button class="btn" id="wandb-copy-project">Copy project</button>
|
||||
</div>
|
||||
<div style="height:10px"></div>
|
||||
<div class="muted">project: <code id="wandb-project">(unknown)</code></div>
|
||||
<div class="muted">base_url (job runtime): <code id="wandb-base-url">(unknown)</code></div>
|
||||
</div>
|
||||
""".strip()
|
||||
script = """
|
||||
const tokEl = document.getElementById("tok");
|
||||
const msg = document.getElementById("msg");
|
||||
const me = document.getElementById("me");
|
||||
const wandbOpen = document.getElementById("wandb-open");
|
||||
const wandbProject = document.getElementById("wandb-project");
|
||||
const wandbBaseUrl = document.getElementById("wandb-base-url");
|
||||
document.getElementById("wandb-copy-project").onclick = async () => { await copyText(wandbProject.textContent || ""); };
|
||||
wandbOpen.href = curOriginWithPort(8090);
|
||||
tokEl.value = mvpTokenGet();
|
||||
|
||||
async function refreshMe() {
|
||||
me.textContent = "Loading...";
|
||||
try {
|
||||
const obj = await apiJson("/api/v2/me");
|
||||
me.textContent = fmtJson(obj);
|
||||
if (obj.wandb && obj.wandb.enabled) {
|
||||
wandbProject.textContent = obj.wandb.project_name || "(unknown)";
|
||||
wandbBaseUrl.textContent = obj.wandb.base_url || "(unknown)";
|
||||
} else {
|
||||
wandbProject.textContent = "(disabled)";
|
||||
wandbBaseUrl.textContent = "(disabled)";
|
||||
}
|
||||
} catch (e) {
|
||||
me.textContent = "Error: " + (e.status || "") + "\\n" + (e.body || String(e));
|
||||
wandbProject.textContent = "(error)";
|
||||
wandbBaseUrl.textContent = "(error)";
|
||||
}
|
||||
}
|
||||
|
||||
document.getElementById("me-refresh").onclick = refreshMe;
|
||||
document.getElementById("save").onclick = () => { mvpTokenSet(tokEl.value); msg.textContent = "Saved."; refreshMe(); };
|
||||
document.getElementById("clear").onclick = () => { mvpTokenSet(""); tokEl.value = ""; msg.textContent = "Cleared."; me.textContent = "(not loaded)"; };
|
||||
if (mvpTokenGet()) refreshMe();
|
||||
""".strip()
|
||||
return HTMLResponse(content=page("Login", "login", body, script))
|
||||
|
||||
261
src/mvp/py/argus/ui/pages/serving.py
Normal file
@@ -0,0 +1,261 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import html
|
||||
import json
|
||||
|
||||
from fastapi import FastAPI
|
||||
from fastapi.responses import HTMLResponse
|
||||
|
||||
from ..layout.page import page
|
||||
|
||||
|
||||
def register(app: FastAPI) -> None:
|
||||
@app.get("/ui/serving")
|
||||
async def ui_serving() -> HTMLResponse:
|
||||
body = """
|
||||
<h1>Serving</h1>
|
||||
<div class="card">
|
||||
<div class="row">
|
||||
<button class="btn" id="refresh">Refresh</button>
|
||||
<a class="btn" href="/ui/serving/new" style="display:inline-block">New Model</a>
|
||||
</div>
|
||||
<div style="height:10px"></div>
|
||||
<div id="out" class="muted">Loading...</div>
|
||||
</div>
|
||||
""".strip()
|
||||
script = """
|
||||
document.getElementById("nav-ray-dashboard").href = curOriginWithPort(8265);
|
||||
const out = document.getElementById("out");
|
||||
|
||||
function pill(state) {
|
||||
const s = String(state || "");
|
||||
if (s === "RUNNING") return `<span class="pill ok">${s}</span>`;
|
||||
if (s === "FAILED") return `<span class="pill bad">${s}</span>`;
|
||||
return `<span class="pill">${s}</span>`;
|
||||
}
|
||||
|
||||
async function refresh() {
|
||||
out.textContent = "Loading...";
|
||||
try {
|
||||
const lim = 50;
|
||||
const off = Number(localStorage.getItem("mvp_serving_offset") || "0") || 0;
|
||||
const resp = await apiJson("/api/v2/serve/models?limit=" + lim + "&offset=" + off + "&include_deleted=0");
|
||||
const items = resp.items || [];
|
||||
const hasMore = !!resp.has_more;
|
||||
const pageNo = Math.floor(off / lim) + 1;
|
||||
const prevDisabled = off <= 0;
|
||||
const nextDisabled = !hasMore;
|
||||
|
||||
const baseHint = curOriginWithPort(8000) + "/serve/<model_key>/v1";
|
||||
function row(m) {
|
||||
const endpoint = (m.endpoint && m.endpoint.openai_base_url) || (curOriginWithPort(8000) + "/serve/" + encodeURIComponent(m.model_key) + "/v1");
|
||||
return `<tr>
|
||||
<td><a href="/ui/serving/${m.model_key}">${m.model_key}</a></td>
|
||||
<td><code>${m.model_id}</code></td>
|
||||
<td>${pill(m.state)}</td>
|
||||
<td>${m.num_replicas} × ${m.gpus_per_replica} GPU</td>
|
||||
<td><code>${endpoint}</code></td>
|
||||
<td>${m.updated_at || ""}</td>
|
||||
</tr>`;
|
||||
}
|
||||
const rows = items.map(row).join("");
|
||||
|
||||
out.innerHTML = `
|
||||
<div class="row" style="justify-content: space-between; margin-bottom: 8px;">
|
||||
<div class="muted">Per-model OpenAI base: <code>${baseHint}</code></div>
|
||||
<div class="row">
|
||||
<span class="muted">Page ${pageNo}</span>
|
||||
<button class="btn" id="prev" ${prevDisabled ? "disabled" : ""}>Prev</button>
|
||||
<button class="btn" id="next" ${nextDisabled ? "disabled" : ""}>Next</button>
|
||||
</div>
|
||||
</div>
|
||||
<table>
|
||||
<thead><tr><th>Model Key</th><th>Model ID</th><th>State</th><th>Resources</th><th>Endpoint</th><th>Updated</th></tr></thead>
|
||||
<tbody>${rows || "<tr><td colspan=6 class=muted>(none)</td></tr>"}</tbody>
|
||||
</table>
|
||||
`;
|
||||
|
||||
const prevBtn = document.getElementById("prev");
|
||||
const nextBtn = document.getElementById("next");
|
||||
if (prevBtn) prevBtn.onclick = () => { localStorage.setItem("mvp_serving_offset", String(Math.max(0, off - lim))); refresh(); };
|
||||
if (nextBtn) nextBtn.onclick = () => { localStorage.setItem("mvp_serving_offset", String(off + lim)); refresh(); };
|
||||
} catch (e) {
|
||||
let text = "Error: " + (e.status || "") + "\\n" + (e.body || String(e));
|
||||
if (e.body && String(e.body).includes("serving is not enabled")) {
|
||||
text = "Serving is not enabled in server config.\\nAsk admin to enable `serving:` in dev.yaml.";
|
||||
}
|
||||
out.textContent = text;
|
||||
}
|
||||
}
|
||||
|
||||
document.getElementById("refresh").onclick = refresh;
|
||||
refresh();
|
||||
""".strip()
|
||||
return HTMLResponse(content=page("Serving", "serving", body, script))
|
||||
|
||||
@app.get("/ui/serving/new")
|
||||
async def ui_serving_new() -> HTMLResponse:
|
||||
example = """# ServingSpec (YAML)
|
||||
# 说明:
|
||||
# - model_id: 这里是 suffix(平台会自动加前缀:<user_id>-<YYYYMMDDHHMM>-<suffix>)
|
||||
# - model_source: 本地模型路径(支持 $HOME 宏;推荐使用 $HOME/common/hf 指向共享 HF cache)
|
||||
#
|
||||
# 常用路径:
|
||||
# - $HOME/common/hf -> /private/hf
|
||||
# - $HOME -> /private/users/<user_id>
|
||||
#
|
||||
model_id: qwen-0.5b
|
||||
model_source: $HOME/common/hf/hub/models--Qwen--Qwen2.5-0.5B-Instruct/snapshots/<SNAPSHOT_HASH>
|
||||
num_replicas: 1
|
||||
gpus_per_replica: 1
|
||||
|
||||
# engine_kwargs: # 可选:透传给 vLLM
|
||||
# gpu_memory_utilization: 0.4
|
||||
""".strip()
|
||||
body = f"""
|
||||
<h1>New Model</h1>
|
||||
<div class="card">
|
||||
<div class="muted">Paste ServingSpec YAML and submit to <code>/api/v2/serve/models</code>.</div>
|
||||
<div style="height:10px"></div>
|
||||
<textarea id="yaml" rows="14">{html.escape(example)}</textarea>
|
||||
<div style="height:10px"></div>
|
||||
<div class="row">
|
||||
<button class="btn" id="submit">Submit</button>
|
||||
<a class="btn" href="/ui/serving" style="display:inline-block">Back</a>
|
||||
</div>
|
||||
<div style="height:10px"></div>
|
||||
<pre id="out" class="muted"></pre>
|
||||
</div>
|
||||
""".strip()
|
||||
script = """
|
||||
document.getElementById("nav-ray-dashboard").href = curOriginWithPort(8265);
|
||||
const out = document.getElementById("out");
|
||||
document.getElementById("submit").onclick = async () => {
|
||||
out.textContent = "Submitting...";
|
||||
const yaml = document.getElementById("yaml").value || "";
|
||||
try {
|
||||
const resp = await apiJson("/api/v2/serve/models", { method: "POST", headers: { "Content-Type": "application/yaml" }, body: yaml });
|
||||
out.textContent = "Created: " + resp.model_key + "\\nState: " + resp.state;
|
||||
if (resp.model_key) window.location.href = "/ui/serving/" + encodeURIComponent(resp.model_key);
|
||||
} catch (e) {
|
||||
out.textContent = "Error: " + (e.status || "") + "\\n" + (e.body || String(e));
|
||||
}
|
||||
};
|
||||
""".strip()
|
||||
return HTMLResponse(content=page("New Model", "serving", body, script))
|
||||
|
||||
@app.get("/ui/serving/{model_key}")
|
||||
async def ui_serving_detail(model_key: str) -> HTMLResponse:
|
||||
body = f"""
|
||||
<h1>Model</h1>
|
||||
<div class="card">
|
||||
<div class="row" style="justify-content: space-between;">
|
||||
<div class="muted">model_key: <code>{html.escape(model_key)}</code></div>
|
||||
<div class="row">
|
||||
<a class="btn" href="/ui/serving" style="display:inline-block">Back</a>
|
||||
<a class="btn" id="openai-models" target="_blank" rel="noopener" href="#">OpenAI /v1/models</a>
|
||||
</div>
|
||||
</div>
|
||||
<div style="height:10px"></div>
|
||||
<div class="row">
|
||||
<label class="muted" style="min-width:120px">Scale replicas</label>
|
||||
<input id="replicas" type="number" min="1" step="1" value="1" style="max-width: 180px" />
|
||||
<button class="btn" id="scale">Apply</button>
|
||||
<button class="btn danger" id="delete">Delete</button>
|
||||
</div>
|
||||
<div style="height:10px"></div>
|
||||
<div id="meta" class="muted">Loading...</div>
|
||||
<div style="height:12px"></div>
|
||||
<h3 style="margin-top:0">Resolved Spec (YAML)</h3>
|
||||
<pre id="spec" class="muted">(loading)</pre>
|
||||
<div style="height:12px"></div>
|
||||
<h3 style="margin-top:0">Events</h3>
|
||||
<div id="events" class="muted">(loading)</div>
|
||||
<div style="height:12px"></div>
|
||||
<h3 style="margin-top:0">OpenAI Example</h3>
|
||||
<pre id="example" class="muted">(loading)</pre>
|
||||
</div>
|
||||
""".strip()
|
||||
script = f"""
|
||||
document.getElementById("nav-ray-dashboard").href = curOriginWithPort(8265);
|
||||
const modelKey = {json.dumps(model_key)};
|
||||
document.getElementById("openai-models").href = curOriginWithPort(8000) + "/serve/" + encodeURIComponent(modelKey) + "/v1/models";
|
||||
const meta = document.getElementById("meta");
|
||||
const spec = document.getElementById("spec");
|
||||
const eventsEl = document.getElementById("events");
|
||||
const example = document.getElementById("example");
|
||||
const replicas = document.getElementById("replicas");
|
||||
|
||||
function pill(state) {{
|
||||
const s = String(state || "");
|
||||
if (s === "RUNNING") return `<span class="pill ok">${{s}}</span>`;
|
||||
if (s === "FAILED") return `<span class="pill bad">${{s}}</span>`;
|
||||
return `<span class="pill">${{s}}</span>`;
|
||||
}}
|
||||
|
||||
function renderEvents(events) {{
|
||||
if (!events || !events.length) return "<div class=muted>(none)</div>";
|
||||
const rows = events.map(e => {{
|
||||
const payload = (e.payload_json || "");
|
||||
const short = String(payload).length > 240 ? String(payload).slice(0, 240) + "..." : String(payload);
|
||||
return `<tr><td>${{e.created_at || ""}}</td><td><code>${{e.event_type}}</code></td><td><pre class=muted style=\\"margin:0\\">${{short}}</pre></td></tr>`;
|
||||
}}).join("");
|
||||
return `<table><thead><tr><th>Time</th><th>Type</th><th>Payload</th></tr></thead><tbody>${{rows}}</tbody></table>`;
|
||||
}}
|
||||
|
||||
async function refresh() {{
|
||||
meta.textContent = "Loading...";
|
||||
spec.textContent = "(loading)";
|
||||
eventsEl.textContent = "(loading)";
|
||||
example.textContent = "(loading)";
|
||||
try {{
|
||||
const obj = await apiJson("/api/v2/serve/models/" + encodeURIComponent(modelKey));
|
||||
const m = obj.model || {{}};
|
||||
replicas.value = String(m.num_replicas || 1);
|
||||
const endpoint = (m.endpoint && m.endpoint.openai_base_url) || (curOriginWithPort(8000) + "/serve/" + encodeURIComponent(modelKey) + "/v1");
|
||||
meta.innerHTML = `
|
||||
<div class=row>
|
||||
<div>state: ${{pill(m.state)}}</div>
|
||||
<div class=muted>model_id: <code>${{m.model_id || ""}}</code></div>
|
||||
<div class=muted>source: <code>${{m.model_source || ""}}</code></div>
|
||||
</div>
|
||||
<div class=muted>endpoint: <code>${{endpoint}}</code></div>
|
||||
`;
|
||||
spec.textContent = obj.resolved_spec_yaml || "";
|
||||
eventsEl.innerHTML = renderEvents(obj.events || []);
|
||||
const base = endpoint;
|
||||
const mid = m.model_id || "";
|
||||
example.textContent = `curl -sS -H 'Content-Type: application/json' \\\\\\n -X POST ${{base}}/chat/completions \\\\\\n --data-binary '{{\\"model\\":\\"${{mid}}\\",\\"messages\\":[{{\\"role\\":\\"user\\",\\"content\\":\\"hello\\"}}],\\"max_tokens\\":16,\\"stream\\":false}}' | python3 -m json.tool`;
|
||||
}} catch (e) {{
|
||||
meta.textContent = "Error: " + (e.status || "") + "\\n" + (e.body || String(e));
|
||||
spec.textContent = "";
|
||||
eventsEl.textContent = "";
|
||||
example.textContent = "";
|
||||
}}
|
||||
}}
|
||||
|
||||
document.getElementById("scale").onclick = async () => {{
|
||||
const n = Number(replicas.value || "1");
|
||||
if (!Number.isFinite(n) || n < 1) return;
|
||||
try {{
|
||||
await apiJson("/api/v2/serve/models/" + encodeURIComponent(modelKey), {{ method: "PATCH", headers: {{ "Content-Type": "application/json" }}, body: JSON.stringify({{ num_replicas: n }}) }});
|
||||
await refresh();
|
||||
}} catch (e) {{
|
||||
meta.textContent = "Error: " + (e.status || "") + "\\n" + (e.body || String(e));
|
||||
}}
|
||||
}};
|
||||
|
||||
document.getElementById("delete").onclick = async () => {{
|
||||
if (!confirm("Delete this model?")) return;
|
||||
try {{
|
||||
await apiJson("/api/v2/serve/models/" + encodeURIComponent(modelKey), {{ method: "DELETE" }});
|
||||
await refresh();
|
||||
}} catch (e) {{
|
||||
meta.textContent = "Error: " + (e.status || "") + "\\n" + (e.body || String(e));
|
||||
}}
|
||||
}};
|
||||
|
||||
refresh();
|
||||
""".strip()
|
||||
return HTMLResponse(content=page("Model", "serving", body, script))
|
||||
|
||||
88
src/mvp/py/argus/ui/pages/task_detail.py
Normal file
@@ -0,0 +1,88 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import html
|
||||
|
||||
from fastapi import FastAPI
|
||||
from fastapi.responses import HTMLResponse
|
||||
|
||||
from ..layout.page import page
|
||||
|
||||
|
||||
def register(app: FastAPI) -> None:
|
||||
@app.get("/ui/tasks/{task_id}")
|
||||
async def ui_task_detail(task_id: str) -> HTMLResponse:
|
||||
safe_id = html.escape(task_id)
|
||||
body = f"""
|
||||
<h1>Task: <code>{safe_id}</code></h1>
|
||||
<div class="card">
|
||||
<div class="row">
|
||||
<a class="btn" href="/ui/tasks/{safe_id}/logs" style="display:inline-block">Logs</a>
|
||||
<button class="btn" id="refresh">Refresh</button>
|
||||
<button class="btn danger" id="cancel">Cancel</button>
|
||||
<a class="btn" href="/ui/tasks" style="display:inline-block">Back</a>
|
||||
</div>
|
||||
<div style="height:10px"></div>
|
||||
<pre id="out" class="muted">Loading...</pre>
|
||||
</div>
|
||||
<div style="height:12px"></div>
|
||||
<div class="card">
|
||||
<h3 style="margin-top:0">W&B</h3>
|
||||
<div class="muted">This is a best-effort hint. Run name maps to <code>ray_submission_id</code> (attempt).</div>
|
||||
<div style="height:10px"></div>
|
||||
<div class="row">
|
||||
<a class="btn" id="wandb-open" target="_blank" rel="noopener" href="#">Open W&B (:8090)</a>
|
||||
<button class="btn" id="wandb-copy-run">Copy run</button>
|
||||
</div>
|
||||
<div style="height:10px"></div>
|
||||
<div class="muted">project: <code id="wandb-project">(unknown)</code></div>
|
||||
<div class="muted">run: <code id="wandb-run">(unknown)</code></div>
|
||||
</div>
|
||||
<div style="height:12px"></div>
|
||||
<div class="card">
|
||||
<h3 style="margin-top:0">TaskSpec (YAML)</h3>
|
||||
<div class="muted">Resolved TaskSpec (includes default values; submission_id reflects latest attempt when available).</div>
|
||||
<div style="height:10px"></div>
|
||||
<pre id="spec" class="muted">Loading...</pre>
|
||||
</div>
|
||||
""".strip()
|
||||
script = f"""
|
||||
document.getElementById("nav-ray-dashboard").href = curOriginWithPort(8265);
|
||||
document.getElementById("wandb-open").href = curOriginWithPort(8090);
|
||||
document.getElementById("wandb-copy-run").onclick = async () => {{ await copyText((document.getElementById("wandb-run").textContent || \"\").trim()); }};
|
||||
const out = document.getElementById("out");
|
||||
const spec = document.getElementById("spec");
|
||||
async function refresh() {{
|
||||
out.textContent = "Loading...";
|
||||
spec.textContent = "Loading...";
|
||||
const resp = await apiFetch("/api/v2/tasks/{task_id}");
|
||||
const text = await resp.text();
|
||||
if (!resp.ok) {{ out.textContent = "Error: " + resp.status + "\\n" + text; return; }}
|
||||
const obj = JSON.parse(text);
|
||||
out.textContent = fmtJson(obj);
|
||||
const p = document.getElementById("wandb-project");
|
||||
const r = document.getElementById("wandb-run");
|
||||
const w = (obj.latest_attempt && obj.latest_attempt.wandb) ? obj.latest_attempt.wandb : null;
|
||||
if (w) {{
|
||||
p.textContent = w.project_name || "(unknown)";
|
||||
r.textContent = w.run_name || "(unknown)";
|
||||
}} else {{
|
||||
p.textContent = "(not available)";
|
||||
r.textContent = "(not available)";
|
||||
}}
|
||||
|
||||
const resp2 = await apiFetch("/api/v2/tasks/{task_id}/spec");
|
||||
const text2 = await resp2.text();
|
||||
spec.textContent = resp2.ok ? text2 : ("Error: " + resp2.status + "\\n" + text2);
|
||||
}}
|
||||
document.getElementById("refresh").onclick = refresh;
|
||||
document.getElementById("cancel").onclick = async () => {{
|
||||
if (!confirm("Cancel this task?")) return;
|
||||
const resp = await apiFetch("/api/v2/tasks/{task_id}:cancel", {{ method: "POST" }});
|
||||
const text = await resp.text();
|
||||
out.textContent = (resp.ok ? "Canceled.\\n" : "Error: " + resp.status + "\\n") + text;
|
||||
setTimeout(refresh, 800);
|
||||
}};
|
||||
refresh();
|
||||
""".strip()
|
||||
return HTMLResponse(content=page(f"Task {task_id}", "tasks", body, script))
|
||||
|
||||
48
src/mvp/py/argus/ui/pages/task_logs.py
Normal file
@@ -0,0 +1,48 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import html
|
||||
|
||||
from fastapi import FastAPI
|
||||
from fastapi.responses import HTMLResponse
|
||||
|
||||
from ..layout.page import page
|
||||
|
||||
|
||||
def register(app: FastAPI) -> None:
|
||||
@app.get("/ui/tasks/{task_id}/logs")
|
||||
async def ui_task_logs(task_id: str) -> HTMLResponse:
|
||||
safe_id = html.escape(task_id)
|
||||
body = f"""
|
||||
<h1>Logs: <code>{safe_id}</code></h1>
|
||||
<div class="card">
|
||||
<div class="row">
|
||||
<button class="btn" id="refresh">Refresh</button>
|
||||
<label class="muted">Auto refresh <input type="checkbox" id="auto" /></label>
|
||||
<a class="btn" href="/ui/tasks/{safe_id}" style="display:inline-block">Back</a>
|
||||
</div>
|
||||
<div style="height:10px"></div>
|
||||
<pre id="out" class="muted">Loading...</pre>
|
||||
</div>
|
||||
""".strip()
|
||||
script = f"""
|
||||
document.getElementById("nav-ray-dashboard").href = curOriginWithPort(8265);
|
||||
const out = document.getElementById("out");
|
||||
let timer = null;
|
||||
async function refresh() {{
|
||||
const resp = await apiFetch("/api/v2/tasks/{task_id}/logs?tail=4000");
|
||||
const text = await resp.text();
|
||||
out.textContent = resp.ok ? text : ("Error: " + resp.status + "\\n" + text);
|
||||
}}
|
||||
document.getElementById("refresh").onclick = refresh;
|
||||
document.getElementById("auto").onchange = (e) => {{
|
||||
if (e.target.checked) {{
|
||||
timer = setInterval(refresh, 2000);
|
||||
}} else {{
|
||||
if (timer) clearInterval(timer);
|
||||
timer = null;
|
||||
}}
|
||||
}};
|
||||
refresh();
|
||||
""".strip()
|
||||
return HTMLResponse(content=page(f"Logs {task_id}", "tasks", body, script))
|
||||
|
||||
538
src/mvp/py/argus/ui/pages/task_new.py
Normal file
@@ -0,0 +1,538 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import html
|
||||
|
||||
from fastapi import FastAPI
|
||||
from fastapi.responses import HTMLResponse
|
||||
|
||||
from ..layout.page import page
|
||||
|
||||
|
||||
def register(app: FastAPI) -> None:
|
||||
@app.get("/ui/tasks/new")
|
||||
async def ui_new_task() -> HTMLResponse:
|
||||
ppo = """# PPO TaskSpec (YAML)
|
||||
workload: ppo # 任务类型(必填):ppo|grpo|sft
|
||||
code_path: /private/common/code/verl/verl_repo # 代码路径(必填):v3.0 固定使用 common 下的 verl 快照(不支持用户自定义代码)
|
||||
model_id: Qwen/Qwen2.5-0.5B-Instruct # 基础模型(必填):HuggingFace 模型 ID 或 /private/... 本地模型路径
|
||||
train_file: /private/common/datasets/gsm8k/train.parquet # 训练数据(必填):parquet 文件路径(支持 /private/common/datasets 或 /private/users/<user>/datasets)
|
||||
val_file: /private/common/datasets/gsm8k/test.parquet # 验证数据(必填):parquet 文件路径(VERL 侧会用来构建 val dataset,不能为 null)
|
||||
|
||||
# nnodes: 2 # 训练节点数(可选,默认:2)
|
||||
# n_gpus_per_node: 4 # 每节点 GPU 数(可选,默认:4)
|
||||
|
||||
# total_epochs: 1 # 总训练 epoch(可选,默认:1)
|
||||
# total_training_steps: null # 总训练 step(可选,默认:null;不传则让 VERL 按 epochs 和数据长度自动推导)
|
||||
# save_freq: 10 # checkpoint 保存频率(step)(可选,默认:10)
|
||||
# test_freq: null # 验证频率(step)(可选,默认:null;训练端会当成 -1=不验证)
|
||||
|
||||
# submission_id: "" # Ray submission_id(可选,默认空;通常由服务自动生成,无需填写)
|
||||
""".strip()
|
||||
grpo = """# GRPO TaskSpec (YAML)
|
||||
workload: grpo # 任务类型(必填):ppo|grpo|sft(grpo 会自动启用对应的算法配置)
|
||||
code_path: /private/common/code/verl/verl_repo # 代码路径(必填):v3.0 固定使用 common 下的 verl 快照(不支持用户自定义代码)
|
||||
model_id: Qwen/Qwen2.5-0.5B-Instruct # 基础模型(必填):HuggingFace 模型 ID 或 /private/... 本地模型路径
|
||||
train_file: /private/common/datasets/gsm8k/train.parquet # 训练数据(必填):parquet 文件路径(支持 /private/common/datasets 或 /private/users/<user>/datasets)
|
||||
val_file: /private/common/datasets/gsm8k/test.parquet # 验证数据(必填):parquet 文件路径(VERL 侧会用来构建 val dataset,不能为 null)
|
||||
|
||||
# nnodes: 2 # 训练节点数(可选,默认:2)
|
||||
# n_gpus_per_node: 4 # 每节点 GPU 数(可选,默认:4)
|
||||
|
||||
# total_epochs: 1 # 总训练 epoch(可选,默认:1)
|
||||
# total_training_steps: null # 总训练 step(可选,默认:null;不传则让 VERL 按 epochs 和数据长度自动推导)
|
||||
# save_freq: 10 # checkpoint 保存频率(step)(可选,默认:10)
|
||||
# test_freq: null # 验证频率(step)(可选,默认:null;训练端会当成 -1=不验证)
|
||||
|
||||
# submission_id: "" # Ray submission_id(可选,默认空;通常由服务自动生成,无需填写)
|
||||
""".strip()
|
||||
sft = """# SFT TaskSpec (YAML)
|
||||
workload: sft # 任务类型(必填):ppo|grpo|sft
|
||||
code_path: /private/common/code/verl/verl_repo # 代码路径(必填):v3.0 固定使用 common 下的 verl 快照(不支持用户自定义代码)
|
||||
model_id: Qwen/Qwen2.5-0.5B-Instruct # 基础模型(必填):HuggingFace 模型 ID 或 /private/... 本地模型路径
|
||||
train_file: /private/common/datasets/gsm8k_sft/train.parquet # 训练数据(必填):parquet 文件路径(支持 /private/common/datasets 或 /private/users/<user>/datasets)
|
||||
val_file: /private/common/datasets/gsm8k_sft/test.parquet # 验证数据(必填):parquet 文件路径(VERL 侧会用来构建 val dataset,不能为 null)
|
||||
|
||||
# nnodes: 2 # 训练节点数(可选,默认:2;单机可设 1)
|
||||
# n_gpus_per_node: 4 # 每节点 GPU 数(可选,默认:4;单卡可设 1)
|
||||
|
||||
# total_epochs: 1 # 总训练 epoch(可选,默认:1)
|
||||
# total_training_steps: null # 总训练 step(可选,默认:null;不传则让 VERL 按 epochs 和数据长度自动推导)
|
||||
# save_freq: 10 # checkpoint 保存频率(step)(可选,默认:10)
|
||||
# test_freq: null # 验证频率(step)(可选,默认:null;训练端会当成 -1=不验证)
|
||||
|
||||
# trainer_device: cpu # 仅 SFT 生效:driver 侧 device(可选,默认:cpu)
|
||||
# submission_id: "" # Ray submission_id(可选,默认空;通常由服务自动生成,无需填写)
|
||||
""".strip()
|
||||
adv = """# Advanced TaskSpec (YAML) - v3.6
|
||||
kind: advanced # 任务类型(必填):advanced(自定义 command)
|
||||
# 说明:平台统一按 "advanced" 做任务分类与 task_id 命名(不按 ppo/grpo/sft 细分)。
|
||||
#
|
||||
# 自定义训练命令:平台会做 $HOME 宏替换:
|
||||
# - $HOME -> /private/users/<user>
|
||||
# - $HOME/common/datasets -> /private/datasets(共享只读数据)
|
||||
# - $HOME/common/hf -> /private/hf(共享只读 HF cache)
|
||||
#
|
||||
# W&B(v3.6):
|
||||
# - 平台会在 runtime_env 注入 WANDB_BASE_URL/WANDB_API_KEY/WANDB_DIR
|
||||
# - 以及注入以下 env,供 Advanced command 使用(无需用户手动修改 project):
|
||||
# - MVP_TRAINER_LOGGER: "console" 或 "[console,wandb]"
|
||||
# - MVP_WANDB_PROJECT: "<user_id>_project"
|
||||
# - MVP_WANDB_RUN: "<ray_submission_id>"
|
||||
#
|
||||
nnodes: 2 # 训练节点数(必填):用于平台队列调度与资源预检查
|
||||
n_gpus_per_node: 4 # 每节点 GPU 数(必填):用于平台队列调度与资源预检查
|
||||
|
||||
command: |
|
||||
PYTHONUNBUFFERED=1 \
|
||||
python3 -m verl.trainer.main_ppo \
|
||||
data.train_files=$HOME/common/datasets/gsm8k/train.parquet \
|
||||
data.val_files=$HOME/common/datasets/gsm8k/test.parquet \
|
||||
data.train_batch_size=256 \
|
||||
data.max_prompt_length=512 \
|
||||
data.max_response_length=512 \
|
||||
actor_rollout_ref.model.path=Qwen/Qwen2.5-0.5B-Instruct \
|
||||
actor_rollout_ref.actor.optim.lr=1e-6 \
|
||||
actor_rollout_ref.actor.ppo_mini_batch_size=64 \
|
||||
actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=4 \
|
||||
actor_rollout_ref.rollout.name=vllm \
|
||||
actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=8 \
|
||||
actor_rollout_ref.rollout.tensor_model_parallel_size=1 \
|
||||
actor_rollout_ref.rollout.gpu_memory_utilization=0.4 \
|
||||
actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \
|
||||
critic.optim.lr=1e-5 \
|
||||
critic.model.path=Qwen/Qwen2.5-0.5B-Instruct \
|
||||
critic.ppo_micro_batch_size_per_gpu=4 \
|
||||
algorithm.kl_ctrl.kl_coef=0.001 \
|
||||
trainer.logger=${MVP_TRAINER_LOGGER} \
|
||||
trainer.project_name=${MVP_WANDB_PROJECT} \
|
||||
trainer.experiment_name=${MVP_WANDB_RUN} \
|
||||
trainer.val_before_train=False \
|
||||
trainer.nnodes=2 \
|
||||
trainer.n_gpus_per_node=4 \
|
||||
trainer.total_epochs=1 \
|
||||
trainer.total_training_steps=10 \
|
||||
trainer.save_freq=10 \
|
||||
trainer.test_freq=-1 \
|
||||
trainer.resume_mode=disable \
|
||||
trainer.default_local_dir=checkpoints \
|
||||
+ray_kwargs.ray_init.address=auto \
|
||||
hydra.run.dir=logs/hydra
|
||||
|
||||
# 可选:自定义 reward(方式 A:直接写在 command 里)
|
||||
# command 里增加如下 overrides:
|
||||
# custom_reward_function.path=$HOME/code/reward.py
|
||||
# custom_reward_function.name=compute_score
|
||||
""".strip()
|
||||
merge = """# Model Merge (YAML) - v3.5 (Advanced command)
|
||||
# 用途:将 VERL 训练产生的 FSDP 分片 checkpoint 合并为 HuggingFace 格式目录。
|
||||
#
|
||||
# 你需要把 <RAY_SUBMISSION_ID> 和 <GLOBAL_STEP> 替换成真实值:
|
||||
# - submission id:在 Tasks 详情页里看到的 `ray_submission_id`(如 mvp2-...--a01)
|
||||
# - global_step:对应 checkpoints 下的 global_step_xxx 目录(如 global_step_10)
|
||||
#
|
||||
# 注意:这里不需要 GPU,建议把 n_gpus_per_node 设为 0,这样不受训练任务占用 GPU 的影响。
|
||||
kind: advanced
|
||||
nnodes: 1
|
||||
n_gpus_per_node: 0
|
||||
|
||||
command: |
|
||||
python3 -m verl.model_merger merge \
|
||||
--backend fsdp \
|
||||
--local_dir $HOME/jobs/<RAY_SUBMISSION_ID>/checkpoints/<GLOBAL_STEP>/actor \
|
||||
--target_dir $HOME/jobs/<RAY_SUBMISSION_ID>/checkpoints/<GLOBAL_STEP>/actor/huggingface
|
||||
""".strip()
|
||||
evaluation = """# Evaluation (YAML) - v3.6 (Advanced command)
|
||||
# 用途:对“已生成的结果 parquet”做离线评估(基于 custom reward function),示例使用 VERL 内置 main_eval。
|
||||
#
|
||||
# 说明:
|
||||
# - 你需要准备一个包含生成 responses 的 parquet,并替换 <EVAL_PARQUET>(示例放在某个 job 目录下)。
|
||||
# - main_eval 会在 Ray 上并发计算,因此这里也用 +ray_kwargs 连接已有 Ray 集群。
|
||||
#
|
||||
kind: advanced
|
||||
nnodes: 1
|
||||
n_gpus_per_node: 0
|
||||
|
||||
command: |
|
||||
PYTHONUNBUFFERED=1 \
|
||||
python3 -m verl.trainer.main_eval \
|
||||
data.path=$HOME/jobs/<RAY_SUBMISSION_ID>/outputs/<EVAL_PARQUET>.parquet \
|
||||
+ray_kwargs.ray_init.address=auto
|
||||
|
||||
# 可选:指定 parquet schema(按你的 parquet 列名调整)
|
||||
# data.response_key=responses
|
||||
# data.data_source_key=data_source
|
||||
# data.reward_model_key=reward_model
|
||||
#
|
||||
# 可选:自定义 reward(方式 A:用户自行提供 reward.py)
|
||||
# custom_reward_function.path=$HOME/code/reward.py
|
||||
# custom_reward_function.name=compute_score
|
||||
""".strip()
|
||||
body = f"""
|
||||
<h1>New Task</h1>
|
||||
<div class="card">
|
||||
<div class="muted">
|
||||
Paste TaskSpec YAML and submit to API server.
|
||||
Basic tasks require <code>code_path</code>; Advanced tasks use <code>kind: advanced</code> with a custom <code>command</code>.
|
||||
</div>
|
||||
<div style="height:10px"></div>
|
||||
<div class="row">
|
||||
<button class="btn" id="mode-yaml">YAML 模式</button>
|
||||
<button class="btn" id="mode-form">表单模式</button>
|
||||
</div>
|
||||
<div style="height:10px"></div>
|
||||
|
||||
<div id="panel-yaml">
|
||||
<div class="row">
|
||||
<button class="btn" id="tpl-ppo">PPO example</button>
|
||||
<button class="btn" id="tpl-grpo">GRPO example</button>
|
||||
<button class="btn" id="tpl-sft">SFT example</button>
|
||||
<button class="btn" id="tpl-adv">Advanced example</button>
|
||||
<button class="btn" id="tpl-eval">Evaluation example</button>
|
||||
<button class="btn" id="tpl-merge">Model merge example</button>
|
||||
</div>
|
||||
<div style="height:10px"></div>
|
||||
<textarea id="yaml" rows="16">{html.escape(ppo)}</textarea>
|
||||
</div>
|
||||
|
||||
<div id="panel-form" style="display:none">
|
||||
<div class="row">
|
||||
<div style="min-width:240px">
|
||||
<div class="muted">模板</div>
|
||||
<div style="height:6px"></div>
|
||||
<select id="f-tpl" style="width:240px; padding:10px 12px; border-radius: 12px; background: rgba(255,255,255,0.06); color: var(--fg); border: 1px solid rgba(255,255,255,0.12);">
|
||||
<option value="ppo">PPO</option>
|
||||
<option value="grpo">GRPO</option>
|
||||
<option value="sft">SFT</option>
|
||||
<option value="advanced">Advanced</option>
|
||||
<option value="merge">Model Merge</option>
|
||||
</select>
|
||||
</div>
|
||||
<div class="muted" style="flex:1; min-width:260px">
|
||||
表单模式会生成 TaskSpec YAML 并用同一个 API 提交;可随时切回 YAML 模式查看/编辑生成结果。
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div style="height:14px"></div>
|
||||
|
||||
<div id="f-basic">
|
||||
<div class="row">
|
||||
<div style="flex:1; min-width:260px">
|
||||
<div class="muted">code_path</div>
|
||||
<div style="height:6px"></div>
|
||||
<input id="f-code-path" />
|
||||
</div>
|
||||
<div style="flex:1; min-width:260px">
|
||||
<div class="muted">model_id</div>
|
||||
<div style="height:6px"></div>
|
||||
<input id="f-model-id" />
|
||||
</div>
|
||||
</div>
|
||||
<div style="height:10px"></div>
|
||||
<div class="row">
|
||||
<div style="flex:1; min-width:260px">
|
||||
<div class="muted">train_file</div>
|
||||
<div style="height:6px"></div>
|
||||
<input id="f-train-file" />
|
||||
</div>
|
||||
<div style="flex:1; min-width:260px">
|
||||
<div class="muted">val_file</div>
|
||||
<div style="height:6px"></div>
|
||||
<input id="f-val-file" />
|
||||
</div>
|
||||
</div>
|
||||
<div style="height:10px"></div>
|
||||
<div class="row">
|
||||
<div style="min-width:220px">
|
||||
<div class="muted">nnodes</div>
|
||||
<div style="height:6px"></div>
|
||||
<div class="row" style="gap:8px">
|
||||
<button class="btn" id="f-nnodes-dec">-</button>
|
||||
<input id="f-nnodes" type="number" min="1" step="1" style="width:120px" />
|
||||
<button class="btn" id="f-nnodes-inc">+</button>
|
||||
</div>
|
||||
</div>
|
||||
<div style="min-width:240px">
|
||||
<div class="muted">n_gpus_per_node</div>
|
||||
<div style="height:6px"></div>
|
||||
<div class="row" style="gap:8px">
|
||||
<button class="btn" id="f-gpn-dec">-</button>
|
||||
<input id="f-gpn" type="number" min="0" step="1" style="width:120px" />
|
||||
<button class="btn" id="f-gpn-inc">+</button>
|
||||
</div>
|
||||
</div>
|
||||
<div style="min-width:220px">
|
||||
<div class="muted">total_epochs</div>
|
||||
<div style="height:6px"></div>
|
||||
<div class="row" style="gap:8px">
|
||||
<button class="btn" id="f-epochs-dec">-</button>
|
||||
<input id="f-epochs" type="number" min="1" step="1" style="width:120px" />
|
||||
<button class="btn" id="f-epochs-inc">+</button>
|
||||
</div>
|
||||
</div>
|
||||
<div style="min-width:260px">
|
||||
<div class="muted">total_training_steps(可选)</div>
|
||||
<div style="height:6px"></div>
|
||||
<input id="f-steps" type="number" min="1" step="1" placeholder="留空表示不传(让 VERL 自动推导)" />
|
||||
</div>
|
||||
</div>
|
||||
<div style="height:10px"></div>
|
||||
<div class="row">
|
||||
<div style="min-width:220px">
|
||||
<div class="muted">save_freq</div>
|
||||
<div style="height:6px"></div>
|
||||
<input id="f-save-freq" type="number" min="1" step="1" style="width:160px" />
|
||||
</div>
|
||||
<div style="min-width:260px">
|
||||
<div class="muted">test_freq(可选)</div>
|
||||
<div style="height:6px"></div>
|
||||
<input id="f-test-freq" type="number" min="1" step="1" placeholder="留空表示不传(训练端会当成 -1=不验证)" />
|
||||
</div>
|
||||
<div id="f-sft-device-wrap" style="min-width:260px; display:none">
|
||||
<div class="muted">trainer_device(仅 SFT)</div>
|
||||
<div style="height:6px"></div>
|
||||
<select id="f-trainer-device" style="width:200px; padding:10px 12px; border-radius: 12px; background: rgba(255,255,255,0.06); color: var(--fg); border: 1px solid rgba(255,255,255,0.12);">
|
||||
<option value="cpu">cpu</option>
|
||||
<option value="cuda">cuda</option>
|
||||
</select>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div id="f-adv" style="display:none">
|
||||
<div class="row">
|
||||
<div style="min-width:220px">
|
||||
<div class="muted">nnodes</div>
|
||||
<div style="height:6px"></div>
|
||||
<div class="row" style="gap:8px">
|
||||
<button class="btn" id="f-adv-nnodes-dec">-</button>
|
||||
<input id="f-adv-nnodes" type="number" min="1" step="1" style="width:120px" />
|
||||
<button class="btn" id="f-adv-nnodes-inc">+</button>
|
||||
</div>
|
||||
</div>
|
||||
<div style="min-width:240px">
|
||||
<div class="muted">n_gpus_per_node</div>
|
||||
<div style="height:6px"></div>
|
||||
<div class="row" style="gap:8px">
|
||||
<button class="btn" id="f-adv-gpn-dec">-</button>
|
||||
<input id="f-adv-gpn" type="number" min="0" step="1" style="width:120px" />
|
||||
<button class="btn" id="f-adv-gpn-inc">+</button>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
<div style="height:10px"></div>
|
||||
<div class="muted">command</div>
|
||||
<div style="height:6px"></div>
|
||||
<textarea id="f-command" rows="10"></textarea>
|
||||
<div style="height:10px"></div>
|
||||
<div class="muted">
|
||||
Tips:
|
||||
Advanced 模式只允许执行 <code>python3 -m verl.*</code>。
|
||||
如果你需要自定义 reward,可把 <code>reward.py</code> 上传到 <code>$HOME/code/</code> 并在 command 里引用。
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div style="height:10px"></div>
|
||||
<div class="row">
|
||||
<button class="btn" id="submit">Submit</button>
|
||||
<a class="btn" href="/ui/tasks" style="display:inline-block">Back</a>
|
||||
</div>
|
||||
<div style="height:10px"></div>
|
||||
<pre id="msg" class="muted"></pre>
|
||||
</div>
|
||||
""".strip()
|
||||
|
||||
tpl_ppo = ppo.replace("\n", "\\n").replace('"', '\\"')
|
||||
tpl_grpo = grpo.replace("\n", "\\n").replace('"', '\\"')
|
||||
tpl_sft = sft.replace("\n", "\\n").replace('"', '\\"')
|
||||
tpl_adv = adv.replace("\n", "\\n").replace('"', '\\"')
|
||||
tpl_eval = evaluation.replace("\n", "\\n").replace('"', '\\"')
|
||||
tpl_merge = merge.replace("\n", "\\n").replace('"', '\\"')
|
||||
|
||||
script = (
|
||||
"""
|
||||
document.getElementById("nav-ray-dashboard").href = curOriginWithPort(8265);
|
||||
const yamlEl = document.getElementById("yaml");
|
||||
const msg = document.getElementById("msg");
|
||||
|
||||
const TPL_PPO = "__TPL_PPO__";
|
||||
const TPL_GRPO = "__TPL_GRPO__";
|
||||
const TPL_SFT = "__TPL_SFT__";
|
||||
const TPL_ADV = "__TPL_ADV__";
|
||||
const TPL_EVAL = "__TPL_EVAL__";
|
||||
const TPL_MERGE = "__TPL_MERGE__";
|
||||
|
||||
function showMode(mode) {
|
||||
const a = document.getElementById("mode-yaml");
|
||||
const b = document.getElementById("mode-form");
|
||||
const py = document.getElementById("panel-yaml");
|
||||
const pf = document.getElementById("panel-form");
|
||||
if (mode === "yaml") {
|
||||
a.classList.add("active"); b.classList.remove("active");
|
||||
py.style.display = ""; pf.style.display = "none";
|
||||
} else {
|
||||
b.classList.add("active"); a.classList.remove("active");
|
||||
pf.style.display = ""; py.style.display = "none";
|
||||
}
|
||||
}
|
||||
|
||||
function yamlQuote(v) {
|
||||
v = String(v || "");
|
||||
if (!v) return '""';
|
||||
if (/^[A-Za-z0-9_./:-]+$/.test(v) && !v.includes("#")) return v;
|
||||
return '"' + v.replaceAll("\\\\","\\\\\\\\").replaceAll('"','\\\\"') + '"';
|
||||
}
|
||||
|
||||
function updateFormVisibility() {
|
||||
const tpl = document.getElementById("f-tpl").value;
|
||||
const basic = document.getElementById("f-basic");
|
||||
const adv = document.getElementById("f-adv");
|
||||
const sftDev = document.getElementById("f-sft-device-wrap");
|
||||
if (tpl === "advanced" || tpl === "merge") {
|
||||
basic.style.display = "none";
|
||||
adv.style.display = "";
|
||||
sftDev.style.display = "none";
|
||||
} else {
|
||||
basic.style.display = "";
|
||||
adv.style.display = "none";
|
||||
sftDev.style.display = (tpl === "sft") ? "" : "none";
|
||||
}
|
||||
}
|
||||
|
||||
function updatePreview() {
|
||||
const tpl = document.getElementById("f-tpl").value;
|
||||
if (tpl === "advanced" || tpl === "merge") {
|
||||
const nnodes = Number(document.getElementById("f-adv-nnodes").value || "1") || 1;
|
||||
const gpn = Number(document.getElementById("f-adv-gpn").value || "0") || 0;
|
||||
const cmd = (document.getElementById("f-command").value || "").trimEnd();
|
||||
const out = `kind: advanced\\n` +
|
||||
`nnodes: ${nnodes}\\n` +
|
||||
`n_gpus_per_node: ${gpn}\\n\\n` +
|
||||
`command: |\\n` +
|
||||
cmd.split("\\n").map(line => " " + line).join("\\n") + "\\n";
|
||||
return out;
|
||||
}
|
||||
|
||||
const workload = tpl;
|
||||
const codePath = document.getElementById("f-code-path").value || "";
|
||||
const modelId = document.getElementById("f-model-id").value || "";
|
||||
const trainFile = document.getElementById("f-train-file").value || "";
|
||||
const valFile = document.getElementById("f-val-file").value || "";
|
||||
const nnodes = Number(document.getElementById("f-nnodes").value || "2") || 2;
|
||||
const gpn = Number(document.getElementById("f-gpn").value || "4") || 4;
|
||||
const epochs = Number(document.getElementById("f-epochs").value || "1") || 1;
|
||||
const steps = (document.getElementById("f-steps").value || "").trim();
|
||||
const saveFreq = Number(document.getElementById("f-save-freq").value || "10") || 10;
|
||||
const testFreq = (document.getElementById("f-test-freq").value || "").trim();
|
||||
const trainerDevice = document.getElementById("f-trainer-device").value || "cpu";
|
||||
|
||||
let y = "";
|
||||
y += `workload: ${workload}\\n`;
|
||||
y += `code_path: ${yamlQuote(codePath)}\\n`;
|
||||
y += `model_id: ${yamlQuote(modelId)}\\n`;
|
||||
y += `train_file: ${yamlQuote(trainFile)}\\n`;
|
||||
y += `val_file: ${yamlQuote(valFile)}\\n`;
|
||||
y += `nnodes: ${nnodes}\\n`;
|
||||
y += `n_gpus_per_node: ${gpn}\\n`;
|
||||
y += `total_epochs: ${epochs}\\n`;
|
||||
if (steps) y += `total_training_steps: ${Number(steps)}\\n`;
|
||||
y += `save_freq: ${saveFreq}\\n`;
|
||||
if (testFreq) y += `test_freq: ${Number(testFreq)}\\n`;
|
||||
if (workload === "sft") y += `trainer_device: ${trainerDevice}\\n`;
|
||||
return y;
|
||||
}
|
||||
|
||||
function bindStepper(inputId, decId, incId, min) {
|
||||
const inp = document.getElementById(inputId);
|
||||
document.getElementById(decId).onclick = () => { inp.value = String(Math.max(min, (Number(inp.value || "0") || 0) - 1)); inp.dispatchEvent(new Event("input")); };
|
||||
document.getElementById(incId).onclick = () => { inp.value = String((Number(inp.value || "0") || 0) + 1); inp.dispatchEvent(new Event("input")); };
|
||||
}
|
||||
|
||||
document.getElementById("tpl-ppo").onclick = () => { yamlEl.value = TPL_PPO; };
|
||||
document.getElementById("tpl-grpo").onclick = () => { yamlEl.value = TPL_GRPO; };
|
||||
document.getElementById("tpl-sft").onclick = () => { yamlEl.value = TPL_SFT; };
|
||||
document.getElementById("tpl-adv").onclick = () => { yamlEl.value = TPL_ADV; };
|
||||
document.getElementById("tpl-eval").onclick = () => { yamlEl.value = TPL_EVAL; };
|
||||
document.getElementById("tpl-merge").onclick = () => { yamlEl.value = TPL_MERGE; };
|
||||
|
||||
document.getElementById("f-code-path").value = "/private/common/code/verl/verl_repo";
|
||||
document.getElementById("f-model-id").value = "Qwen/Qwen2.5-0.5B-Instruct";
|
||||
document.getElementById("f-train-file").value = "/private/common/datasets/gsm8k/train.parquet";
|
||||
document.getElementById("f-val-file").value = "/private/common/datasets/gsm8k/test.parquet";
|
||||
document.getElementById("f-nnodes").value = "2";
|
||||
document.getElementById("f-gpn").value = "4";
|
||||
document.getElementById("f-epochs").value = "1";
|
||||
document.getElementById("f-save-freq").value = "10";
|
||||
document.getElementById("f-trainer-device").value = "cpu";
|
||||
|
||||
document.getElementById("f-adv-nnodes").value = "2";
|
||||
document.getElementById("f-adv-gpn").value = "4";
|
||||
document.getElementById("f-command").value = `PYTHONUNBUFFERED=1 \\\\\\npython3 -m verl.trainer.main_ppo \\\\\\n data.train_files=$HOME/common/datasets/gsm8k/train.parquet \\\\\\n data.val_files=$HOME/common/datasets/gsm8k/test.parquet \\\\\\n +ray_kwargs.ray_init.address=auto`;
|
||||
|
||||
bindStepper("f-nnodes", "f-nnodes-dec", "f-nnodes-inc", 1);
|
||||
bindStepper("f-gpn", "f-gpn-dec", "f-gpn-inc", 0);
|
||||
bindStepper("f-epochs", "f-epochs-dec", "f-epochs-inc", 1);
|
||||
bindStepper("f-adv-nnodes", "f-adv-nnodes-dec", "f-adv-nnodes-inc", 1);
|
||||
bindStepper("f-adv-gpn", "f-adv-gpn-dec", "f-adv-gpn-inc", 0);
|
||||
|
||||
document.getElementById("f-tpl").onchange = () => {
|
||||
const tpl = document.getElementById("f-tpl").value;
|
||||
if (tpl === "ppo") {
|
||||
document.getElementById("f-train-file").value = "/private/common/datasets/gsm8k/train.parquet";
|
||||
document.getElementById("f-val-file").value = "/private/common/datasets/gsm8k/test.parquet";
|
||||
} else if (tpl === "grpo") {
|
||||
document.getElementById("f-train-file").value = "/private/common/datasets/gsm8k/train.parquet";
|
||||
document.getElementById("f-val-file").value = "/private/common/datasets/gsm8k/test.parquet";
|
||||
} else if (tpl === "sft") {
|
||||
document.getElementById("f-train-file").value = "/private/common/datasets/gsm8k_sft/train.parquet";
|
||||
document.getElementById("f-val-file").value = "/private/common/datasets/gsm8k_sft/test.parquet";
|
||||
} else if (tpl === "advanced") {
|
||||
document.getElementById("f-adv-nnodes").value = "2";
|
||||
document.getElementById("f-adv-gpn").value = "4";
|
||||
document.getElementById("f-command").value = TPL_ADV.split("command: |")[1]?.trimStart() || document.getElementById("f-command").value;
|
||||
} else if (tpl === "merge") {
|
||||
document.getElementById("f-adv-nnodes").value = "1";
|
||||
document.getElementById("f-adv-gpn").value = "0";
|
||||
document.getElementById("f-command").value = TPL_MERGE.split("command: |")[1]?.trimStart() || document.getElementById("f-command").value;
|
||||
}
|
||||
updateFormVisibility();
|
||||
updatePreview();
|
||||
};
|
||||
|
||||
for (const id of ["f-code-path","f-model-id","f-train-file","f-val-file","f-nnodes","f-gpn","f-epochs","f-steps","f-save-freq","f-test-freq","f-trainer-device","f-adv-nnodes","f-adv-gpn","f-command"]) {
|
||||
const el = document.getElementById(id);
|
||||
if (el) el.oninput = () => updatePreview();
|
||||
}
|
||||
|
||||
document.getElementById("mode-yaml").onclick = () => {
|
||||
// When switching to YAML, sync from form if form is visible.
|
||||
const formVisible = document.getElementById("panel-form").style.display !== "none";
|
||||
if (formVisible) yamlEl.value = updatePreview();
|
||||
showMode("yaml");
|
||||
};
|
||||
document.getElementById("mode-form").onclick = () => {
|
||||
showMode("form");
|
||||
};
|
||||
|
||||
showMode("yaml");
|
||||
|
||||
document.getElementById("submit").onclick = async () => {
|
||||
msg.textContent = "Submitting...";
|
||||
const body = (document.getElementById("panel-form").style.display !== "none") ? updatePreview() : yamlEl.value;
|
||||
const resp = await apiFetch("/api/v2/tasks", { method: "POST", headers: {"Content-Type":"text/plain"}, body });
|
||||
const text = await resp.text();
|
||||
if (!resp.ok) { msg.textContent = "Error: " + resp.status + "\\n" + text; return; }
|
||||
const obj = JSON.parse(text);
|
||||
msg.textContent = "OK: " + fmtJson(obj);
|
||||
if (obj.task_id) window.location.href = "/ui/tasks/" + obj.task_id;
|
||||
};
|
||||
""".strip()
|
||||
.replace("__TPL_PPO__", tpl_ppo)
|
||||
.replace("__TPL_GRPO__", tpl_grpo)
|
||||
.replace("__TPL_SFT__", tpl_sft)
|
||||
.replace("__TPL_ADV__", tpl_adv)
|
||||
.replace("__TPL_EVAL__", tpl_eval)
|
||||
.replace("__TPL_MERGE__", tpl_merge)
|
||||
)
|
||||
return HTMLResponse(content=page("New Task", "new", body, script))
|
||||
|
||||
101
src/mvp/py/argus/ui/pages/tasks.py
Normal file
@@ -0,0 +1,101 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from fastapi import FastAPI
|
||||
from fastapi.responses import HTMLResponse
|
||||
|
||||
from ..layout.page import page
|
||||
|
||||
|
||||
def register(app: FastAPI) -> None:
|
||||
@app.get("/ui/tasks")
|
||||
async def ui_tasks() -> HTMLResponse:
|
||||
body = """
|
||||
<h1>Tasks</h1>
|
||||
<div class="card">
|
||||
<div class="row">
|
||||
<button class="btn" id="refresh">Refresh</button>
|
||||
<a class="btn" href="/ui/tasks/new" style="display:inline-block">New Task</a>
|
||||
</div>
|
||||
<div style="height:10px"></div>
|
||||
<div id="out" class="muted">Loading...</div>
|
||||
</div>
|
||||
""".strip()
|
||||
script = """
|
||||
const out = document.getElementById("out");
|
||||
async function refresh() {
|
||||
out.textContent = "Loading...";
|
||||
try {
|
||||
const q = await apiJson("/api/v2/queue");
|
||||
const completedLimit = 25;
|
||||
const completedOffset = Number(localStorage.getItem("mvp_completed_offset") || "0") || 0;
|
||||
const done = await apiJson("/api/v2/tasks?limit=" + completedLimit + "&offset=" + completedOffset + "&states=SUCCEEDED,FAILED,CANCELED");
|
||||
|
||||
function pill(state) {
|
||||
const s = String(state || "");
|
||||
if (s === "SUCCEEDED") return `<span class="pill ok">${s}</span>`;
|
||||
if (s === "FAILED") return `<span class="pill bad">${s}</span>`;
|
||||
if (s === "CANCELED") return `<span class="pill">${s}</span>`;
|
||||
if (s === "RUNNING") return `<span class="pill ok">${s}</span>`;
|
||||
if (s === "QUEUED" || s === "PENDING_RESOURCES" || s === "SUBMITTING" || s === "SUBMITTED") return `<span class="pill">${s}</span>`;
|
||||
return `<span class="pill">${s}</span>`;
|
||||
}
|
||||
function row(t) {
|
||||
const id = t.task_id;
|
||||
return `<tr>
|
||||
<td><a href="/ui/tasks/${id}">${id}</a></td>
|
||||
<td>${t.workload}</td>
|
||||
<td>${pill(t.state)}</td>
|
||||
<td>${t.nnodes} x ${t.n_gpus_per_node} GPU</td>
|
||||
<td>${t.updated_at || ""}</td>
|
||||
</tr>`;
|
||||
}
|
||||
|
||||
const running = (q.running || []).map(row).join("");
|
||||
const pending = (q.pending || []).map(row).join("");
|
||||
const doneRows = (done.tasks || []).map(row).join("");
|
||||
const pageNo = Math.floor(completedOffset / completedLimit) + 1;
|
||||
const prevDisabled = completedOffset <= 0;
|
||||
const nextDisabled = !done.has_more;
|
||||
out.innerHTML = `
|
||||
<div class="muted">Tip: configure token in <a href="/ui/login">Login</a>.</div>
|
||||
<div style="height:10px"></div>
|
||||
<h3>Running</h3>
|
||||
<table><thead><tr><th>Task</th><th>Workload</th><th>State</th><th>Resources</th><th>Updated</th></tr></thead><tbody>${running || "<tr><td colspan=5 class=muted>(none)</td></tr>"}</tbody></table>
|
||||
<div style="height:12px"></div>
|
||||
<h3>Pending</h3>
|
||||
<table><thead><tr><th>Task</th><th>Workload</th><th>State</th><th>Resources</th><th>Updated</th></tr></thead><tbody>${pending || "<tr><td colspan=5 class=muted>(none)</td></tr>"}</tbody></table>
|
||||
<div style="height:12px"></div>
|
||||
<h3>Completed</h3>
|
||||
<div class="row" style="justify-content: space-between; margin-bottom: 8px;">
|
||||
<div class="muted">Page ${pageNo}</div>
|
||||
<div class="row">
|
||||
<button class="btn" id="done-prev" ${prevDisabled ? "disabled" : ""}>Prev</button>
|
||||
<button class="btn" id="done-next" ${nextDisabled ? "disabled" : ""}>Next</button>
|
||||
</div>
|
||||
</div>
|
||||
<table><thead><tr><th>Task</th><th>Workload</th><th>State</th><th>Resources</th><th>Updated</th></tr></thead><tbody>${doneRows || "<tr><td colspan=5 class=muted>(none)</td></tr>"}</tbody></table>
|
||||
`;
|
||||
|
||||
const prevBtn = document.getElementById("done-prev");
|
||||
const nextBtn = document.getElementById("done-next");
|
||||
if (prevBtn) prevBtn.onclick = () => {
|
||||
const cur = Number(localStorage.getItem("mvp_completed_offset") || "0") || 0;
|
||||
const next = Math.max(0, cur - completedLimit);
|
||||
localStorage.setItem("mvp_completed_offset", String(next));
|
||||
refresh();
|
||||
};
|
||||
if (nextBtn) nextBtn.onclick = () => {
|
||||
const cur = Number(localStorage.getItem("mvp_completed_offset") || "0") || 0;
|
||||
const next = cur + completedLimit;
|
||||
localStorage.setItem("mvp_completed_offset", String(next));
|
||||
refresh();
|
||||
};
|
||||
} catch (e) {
|
||||
out.textContent = "Error: " + (e.status || "") + "\\n" + (e.body || String(e));
|
||||
}
|
||||
}
|
||||
document.getElementById("refresh").onclick = refresh;
|
||||
refresh();
|
||||
""".strip()
|
||||
return HTMLResponse(content=page("Tasks", "tasks", body, script))
|
||||
|
||||
22
src/mvp/py/argus/ui/routes.py
Normal file
@@ -0,0 +1,22 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from fastapi import FastAPI
|
||||
from fastapi.responses import RedirectResponse
|
||||
|
||||
from .pages import admin, data, login, serving, task_detail, task_logs, task_new, tasks
|
||||
|
||||
|
||||
def register_ui_routes(app: FastAPI) -> None:
|
||||
@app.get("/ui")
|
||||
async def ui_root() -> RedirectResponse:
|
||||
return RedirectResponse(url="/ui/tasks")
|
||||
|
||||
login.register(app)
|
||||
tasks.register(app)
|
||||
task_new.register(app)
|
||||
task_detail.register(app)
|
||||
task_logs.register(app)
|
||||
serving.register(app)
|
||||
data.register(app)
|
||||
admin.register(app)
|
||||
|
||||
4
src/mvp/py/requirements-dev.txt
Normal file
@@ -0,0 +1,4 @@
|
||||
pytest==8.4.1
|
||||
pytest-cov==6.3.0
|
||||
httpx==0.28.1
|
||||
|
||||
4
src/mvp/py/requirements.txt
Normal file
@@ -0,0 +1,4 @@
|
||||
fastapi==0.115.6
|
||||
uvicorn==0.30.6
|
||||
PyYAML==6.0.2
|
||||
|
||||
32
src/mvp/py/server.py
Normal file
@@ -0,0 +1,32 @@
|
||||
#!/usr/bin/env python3
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
|
||||
import uvicorn
|
||||
|
||||
from argus.service.app import create_app
|
||||
from argus.service.config import V2Config
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--config", required=True, help="Path to MVP config YAML (ray:/service:)")
|
||||
args = parser.parse_args()
|
||||
|
||||
# Load app and read v2.api host/port from config.
|
||||
import yaml
|
||||
|
||||
with open(args.config, "r", encoding="utf-8") as f:
|
||||
root = yaml.safe_load(f) or {}
|
||||
if not isinstance(root, dict):
|
||||
raise SystemExit("config yaml must be a mapping")
|
||||
v2 = V2Config.from_root_dict(root)
|
||||
|
||||
app = create_app(args.config)
|
||||
uvicorn.run(app, host=v2.api.host, port=v2.api.port, log_level="info")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
77
src/mvp/py/tests/conftest.py
Normal file
@@ -0,0 +1,77 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import sys
|
||||
import types
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def _ensure_mvp_py_on_path() -> None:
|
||||
repo_root = Path(__file__).resolve().parents[4]
|
||||
py_root = repo_root / "src" / "mvp" / "py"
|
||||
if str(py_root) not in sys.path:
|
||||
sys.path.insert(0, str(py_root))
|
||||
|
||||
|
||||
def _install_ray_stub() -> None:
|
||||
if "ray" in sys.modules:
|
||||
return
|
||||
|
||||
ray = types.ModuleType("ray")
|
||||
ray.__path__ = [] # type: ignore[attr-defined]
|
||||
|
||||
def _init(*args: object, **kwargs: object) -> None:
|
||||
return None
|
||||
|
||||
ray.init = _init # type: ignore[attr-defined]
|
||||
ray.cluster_resources = lambda: {} # type: ignore[attr-defined]
|
||||
ray.available_resources = lambda: {} # type: ignore[attr-defined]
|
||||
sys.modules["ray"] = ray
|
||||
|
||||
job_submission = types.ModuleType("ray.job_submission")
|
||||
job_submission.__path__ = [] # type: ignore[attr-defined]
|
||||
|
||||
class JobSubmissionClient: # minimal stub; tests can monkeypatch methods
|
||||
def __init__(self, address: str):
|
||||
self.address = address
|
||||
|
||||
def submit_job(self, **kwargs: object) -> str:
|
||||
raise NotImplementedError
|
||||
|
||||
def get_job_status(self, submission_id: str) -> str:
|
||||
raise NotImplementedError
|
||||
|
||||
def stop_job(self, submission_id: str) -> bool:
|
||||
raise NotImplementedError
|
||||
|
||||
def get_job_logs(self, submission_id: str) -> str:
|
||||
raise NotImplementedError
|
||||
|
||||
def list_jobs(self) -> list[object]:
|
||||
return []
|
||||
|
||||
job_submission.JobSubmissionClient = JobSubmissionClient # type: ignore[attr-defined]
|
||||
sys.modules["ray.job_submission"] = job_submission
|
||||
ray.job_submission = job_submission # type: ignore[attr-defined]
|
||||
|
||||
private = types.ModuleType("ray._private")
|
||||
private.__path__ = [] # type: ignore[attr-defined]
|
||||
state = types.ModuleType("ray._private.state")
|
||||
|
||||
def available_resources_per_node() -> dict[str, object]:
|
||||
return {}
|
||||
|
||||
state.available_resources_per_node = available_resources_per_node # type: ignore[attr-defined]
|
||||
sys.modules["ray._private"] = private
|
||||
sys.modules["ray._private.state"] = state
|
||||
private.state = state # type: ignore[attr-defined]
|
||||
ray._private = private # type: ignore[attr-defined]
|
||||
|
||||
|
||||
_ensure_mvp_py_on_path()
|
||||
_install_ray_stub()
|
||||
|
||||
|
||||
def pytest_configure(config: object) -> None:
|
||||
os.environ.setdefault("PYTHONUTF8", "1")
|
||||
|
||||
73
src/mvp/py/tests/test_advanced_command.py
Normal file
@@ -0,0 +1,73 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from argus.service.advanced_command import expand_advanced_command, validate_advanced_command
|
||||
|
||||
|
||||
def test_expand_maps_home_common_datasets_to_private_datasets():
|
||||
cmd = "python3 -m verl.trainer.main_ppo data.train_files=$HOME/common/datasets/gsm8k/train.parquet"
|
||||
expanded = expand_advanced_command(user_id="alice", command=cmd)
|
||||
assert "/private/datasets/gsm8k/train.parquet" in expanded
|
||||
assert "$HOME/common/datasets" not in expanded
|
||||
|
||||
|
||||
def test_expand_maps_home_common_hf_to_private_hf():
|
||||
cmd = "python3 -m verl.trainer.main_ppo HF_HOME=$HOME/common/hf"
|
||||
expanded = expand_advanced_command(user_id="alice", command=cmd)
|
||||
assert "HF_HOME=/private/hf" in expanded
|
||||
|
||||
|
||||
def test_expand_maps_home_to_private_users():
|
||||
cmd = "python3 -m verl.trainer.main_ppo data.train_files=$HOME/datasets/x.parquet"
|
||||
expanded = expand_advanced_command(user_id="alice", command=cmd)
|
||||
assert "data.train_files=/private/users/alice/datasets/x.parquet" in expanded
|
||||
|
||||
|
||||
def test_validate_requires_verl_trainer_entry():
|
||||
with pytest.raises(ValueError, match="-m verl\\."):
|
||||
validate_advanced_command(user_id="alice", expanded_command="echo hi")
|
||||
|
||||
|
||||
def test_validate_rejects_cross_user_paths():
|
||||
cmd = "python3 -m verl.trainer.main_ppo x=/private/users/bob/datasets/x"
|
||||
with pytest.raises(ValueError, match="another user's path"):
|
||||
validate_advanced_command(user_id="alice", expanded_command=cmd)
|
||||
|
||||
|
||||
def test_validate_allows_train_val_under_allowed_roots():
|
||||
cmd = (
|
||||
"python3 -m verl.trainer.main_ppo "
|
||||
"data.train_files=/private/datasets/gsm8k/train.parquet "
|
||||
"data.val_files=/private/users/alice/datasets/gsm8k/test.parquet"
|
||||
)
|
||||
validate_advanced_command(user_id="alice", expanded_command=cmd)
|
||||
|
||||
|
||||
def test_validate_rejects_train_val_outside_allowed_roots():
|
||||
cmd = (
|
||||
"python3 -m verl.trainer.main_ppo "
|
||||
"data.train_files=/etc/passwd "
|
||||
"data.val_files=/private/datasets/gsm8k/test.parquet"
|
||||
)
|
||||
with pytest.raises(ValueError, match="train_file must be under allowed roots"):
|
||||
validate_advanced_command(user_id="alice", expanded_command=cmd)
|
||||
|
||||
|
||||
def test_validate_allows_reward_path_under_user_code():
|
||||
cmd = (
|
||||
"python3 -m verl.trainer.main_ppo "
|
||||
"custom_reward_function.path=/private/users/alice/code/reward.py "
|
||||
"custom_reward_function.name=compute_score"
|
||||
)
|
||||
validate_advanced_command(user_id="alice", expanded_command=cmd)
|
||||
|
||||
|
||||
def test_validate_rejects_reward_path_outside_user_code():
|
||||
cmd = (
|
||||
"python3 -m verl.trainer.main_ppo "
|
||||
"custom_reward_function.path=/private/datasets/reward.py "
|
||||
"custom_reward_function.name=compute_score"
|
||||
)
|
||||
with pytest.raises(ValueError, match="reward path must be under"):
|
||||
validate_advanced_command(user_id="alice", expanded_command=cmd)
|
||||
452
src/mvp/py/tests/test_app.py
Normal file
@@ -0,0 +1,452 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
import yaml
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
|
||||
def _write_config(tmp_path: Path) -> Path:
|
||||
cfg = {
|
||||
"ray": {
|
||||
"address": "http://127.0.0.1:8265",
|
||||
"shared_root": "/private",
|
||||
"entrypoint_resources": {"worker_node": 1},
|
||||
"runtime_env": {"env_vars": {}},
|
||||
},
|
||||
"data": {
|
||||
# Avoid touching real /private in tests. Keep ray.shared_root as /private
|
||||
# so existing path validation tests remain unchanged.
|
||||
"user_root": str(tmp_path / "users"),
|
||||
},
|
||||
"service": {
|
||||
"api": {"host": "127.0.0.1", "port": 0},
|
||||
"auth": {"token_env": "MVP_INTERNAL_TOKEN"},
|
||||
"sqlite": {"db_path": str(tmp_path / "mvp.sqlite3")},
|
||||
"scheduler": {"tick_s": 1, "retry_interval_s": 1, "max_running_tasks": 1},
|
||||
},
|
||||
}
|
||||
p = tmp_path / "cfg.yaml"
|
||||
p.write_text(yaml.safe_dump(cfg), encoding="utf-8")
|
||||
return p
|
||||
|
||||
|
||||
def test_auth_requires_token_env(tmp_path: Path, monkeypatch):
|
||||
from argus.service import app as app_mod
|
||||
|
||||
cfg_path = _write_config(tmp_path)
|
||||
monkeypatch.delenv("MVP_INTERNAL_TOKEN", raising=False)
|
||||
|
||||
class _Scheduler:
|
||||
def __init__(self, **kwargs):
|
||||
self.tool = object()
|
||||
|
||||
def run_forever(self, stop_flag):
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(app_mod, "Scheduler", _Scheduler)
|
||||
app = app_mod.create_app(str(cfg_path))
|
||||
|
||||
with TestClient(app) as c:
|
||||
r = c.get("/api/v2/queue")
|
||||
assert r.status_code == 500
|
||||
|
||||
|
||||
def test_task_submit_get_cancel_logs_queue(tmp_path: Path, monkeypatch):
|
||||
from argus.service import app as app_mod
|
||||
|
||||
cfg_path = _write_config(tmp_path)
|
||||
monkeypatch.setenv("MVP_INTERNAL_TOKEN", "token1")
|
||||
monkeypatch.setattr(app_mod, "new_task_id", lambda workload, **kwargs: "tid1")
|
||||
|
||||
class _Tool:
|
||||
def __init__(self):
|
||||
self.stopped = []
|
||||
|
||||
def stop(self, sid: str):
|
||||
self.stopped.append(sid)
|
||||
return True
|
||||
|
||||
def logs(self, sid: str):
|
||||
return "a\nb\nc\n"
|
||||
|
||||
class _Scheduler:
|
||||
def __init__(self, **kwargs):
|
||||
self.tool = _Tool()
|
||||
|
||||
def run_forever(self, stop_flag):
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(app_mod, "Scheduler", _Scheduler)
|
||||
app = app_mod.create_app(str(cfg_path))
|
||||
|
||||
headers = {"authorization": "Bearer token1"}
|
||||
with TestClient(app) as c:
|
||||
r = c.post(
|
||||
"/api/v2/tasks",
|
||||
headers=headers,
|
||||
data="workload: ppo\ncode_path: /private/common/code/verl\nmodel_id: m\ntrain_file: /private/common/datasets/t\nval_file: /private/common/datasets/v\n",
|
||||
)
|
||||
assert r.status_code == 200
|
||||
assert r.json()["task_id"] == "tid1"
|
||||
|
||||
r2 = c.get("/api/v2/tasks/tid1", headers=headers)
|
||||
assert r2.status_code == 200
|
||||
assert r2.json()["desired_resources"]["total_gpus"] == 8
|
||||
r2s = c.get("/api/v2/tasks/tid1/spec", headers=headers)
|
||||
assert r2s.status_code == 200
|
||||
assert "workload: ppo" in r2s.text
|
||||
assert "code_path:" in r2s.text
|
||||
# Spec endpoint returns a resolved TaskSpec, i.e. includes default values.
|
||||
assert "nnodes: 2" in r2s.text
|
||||
assert "n_gpus_per_node: 4" in r2s.text
|
||||
|
||||
# Seed an attempt and ensure submission_id is reflected in the resolved spec.
|
||||
from argus.service.db import Db
|
||||
from argus.service.config import V2Config
|
||||
from argus.ray.models import RayConfig
|
||||
|
||||
root = yaml.safe_load(cfg_path.read_text(encoding="utf-8"))
|
||||
ray_cfg = RayConfig.from_dict(root)
|
||||
v2_cfg = V2Config.from_root_dict(root)
|
||||
db = Db(v2_cfg.sqlite.db_path)
|
||||
db.create_attempt(task_id="tid1", attempt_no=1, ray_submission_id="sid1")
|
||||
r2s2 = c.get("/api/v2/tasks/tid1/spec", headers=headers)
|
||||
assert r2s2.status_code == 200
|
||||
assert "submission_id: sid1" in r2s2.text
|
||||
|
||||
r3 = c.get("/api/v2/queue", headers=headers)
|
||||
assert r3.status_code == 200
|
||||
assert "pending" in r3.json()
|
||||
|
||||
r3b = c.get("/api/v2/tasks?limit=10", headers=headers)
|
||||
assert r3b.status_code == 200
|
||||
assert any(t.get("task_id") == "tid1" for t in r3b.json().get("tasks", []))
|
||||
|
||||
r3c = c.get("/api/v2/tasks?limit=10&offset=0&states=QUEUED", headers=headers)
|
||||
assert r3c.status_code == 200
|
||||
assert all(t.get("state") == "QUEUED" for t in r3c.json().get("tasks", []))
|
||||
|
||||
r3d = c.get("/api/v2/tasks?states=NOPE", headers=headers)
|
||||
assert r3d.status_code == 400
|
||||
|
||||
r4 = c.post("/api/v2/tasks/tid1:cancel", headers=headers)
|
||||
assert r4.status_code == 200
|
||||
assert r4.json()["state"] == "CANCELED"
|
||||
|
||||
# Seed an attempt then fetch logs
|
||||
root = yaml.safe_load(cfg_path.read_text(encoding="utf-8"))
|
||||
ray_cfg = RayConfig.from_dict(root)
|
||||
v2_cfg = V2Config.from_root_dict(root)
|
||||
db = Db(v2_cfg.sqlite.db_path)
|
||||
db.create_task(
|
||||
task_id="tid2",
|
||||
workload="ppo",
|
||||
jobspec_yaml="workload: ppo\ncode_path: /private/common/code/verl\nmodel_id: m\ntrain_file: /private/common/datasets/t\nval_file: /private/common/datasets/v\n",
|
||||
nnodes=2,
|
||||
n_gpus_per_node=4,
|
||||
)
|
||||
db.create_attempt(task_id="tid2", attempt_no=1, ray_submission_id="sid2")
|
||||
db.set_task_state(task_id="tid2", state="RUNNING", latest_attempt_no=1)
|
||||
|
||||
r6 = c.get("/api/v2/tasks?limit=1&offset=0&states=RUNNING", headers=headers)
|
||||
assert r6.status_code == 200
|
||||
assert any(t.get("task_id") == "tid2" for t in r6.json().get("tasks", []))
|
||||
|
||||
r7 = c.get("/api/v2/tasks?limit=1&offset=1&states=RUNNING", headers=headers)
|
||||
assert r7.status_code == 200
|
||||
assert "has_more" in r7.json()
|
||||
|
||||
r5 = c.get("/api/v2/tasks/tid2/logs?tail=1", headers=headers)
|
||||
assert r5.status_code == 200
|
||||
assert r5.text.strip() == "c"
|
||||
|
||||
|
||||
def test_me_includes_wandb_when_enabled(tmp_path: Path, monkeypatch):
|
||||
from argus.service import app as app_mod
|
||||
|
||||
cfg_path = _write_config(tmp_path)
|
||||
# Enable tracking.wandb
|
||||
cfg = yaml.safe_load(cfg_path.read_text(encoding="utf-8"))
|
||||
cfg["tracking"] = {"wandb": {"enabled": True, "base_url": "http://172.22.0.1:8090"}}
|
||||
cfg_path.write_text(yaml.safe_dump(cfg), encoding="utf-8")
|
||||
|
||||
monkeypatch.setenv("MVP_INTERNAL_TOKEN", "token1")
|
||||
|
||||
class _Scheduler:
|
||||
def __init__(self, **kwargs):
|
||||
self.tool = object()
|
||||
|
||||
def run_forever(self, stop_flag):
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(app_mod, "Scheduler", _Scheduler)
|
||||
app = app_mod.create_app(str(cfg_path))
|
||||
|
||||
with TestClient(app) as c:
|
||||
r = c.get("/api/v2/me", headers={"authorization": "Bearer token1"})
|
||||
assert r.status_code == 200
|
||||
obj = r.json()
|
||||
assert obj["user_id"] == "admin"
|
||||
assert obj.get("wandb", {}).get("enabled") is True
|
||||
assert obj["wandb"]["base_url"] == "http://172.22.0.1:8090"
|
||||
assert obj["wandb"]["project_name"] == "admin_project"
|
||||
|
||||
|
||||
def test_submit_rejects_non_mapping_jobspec(tmp_path: Path, monkeypatch):
|
||||
from argus.service import app as app_mod
|
||||
|
||||
cfg_path = _write_config(tmp_path)
|
||||
monkeypatch.setenv("MVP_INTERNAL_TOKEN", "token1")
|
||||
|
||||
class _Scheduler:
|
||||
def __init__(self, **kwargs):
|
||||
self.tool = object()
|
||||
|
||||
def run_forever(self, stop_flag):
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(app_mod, "Scheduler", _Scheduler)
|
||||
app = app_mod.create_app(str(cfg_path))
|
||||
|
||||
with TestClient(app) as c:
|
||||
r = c.post("/api/v2/tasks", headers={"authorization": "Bearer token1"}, data="- 1\n- 2\n")
|
||||
assert r.status_code == 400
|
||||
|
||||
|
||||
def test_submit_rejects_invalid_yaml_returns_400(tmp_path: Path, monkeypatch):
|
||||
from argus.service import app as app_mod
|
||||
|
||||
cfg_path = _write_config(tmp_path)
|
||||
monkeypatch.setenv("MVP_INTERNAL_TOKEN", "token1")
|
||||
|
||||
class _Scheduler:
|
||||
def __init__(self, **kwargs):
|
||||
self.tool = object()
|
||||
|
||||
def run_forever(self, stop_flag):
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(app_mod, "Scheduler", _Scheduler)
|
||||
app = app_mod.create_app(str(cfg_path))
|
||||
|
||||
with TestClient(app) as c:
|
||||
# YAML parse error: unclosed quote
|
||||
r = c.post("/api/v2/tasks", headers={"authorization": "Bearer token1"}, data="workload: \"ppo\n")
|
||||
assert r.status_code == 400
|
||||
assert "invalid YAML" in r.text
|
||||
|
||||
|
||||
def test_submit_rejects_invalid_jobspec(tmp_path: Path, monkeypatch):
|
||||
from argus.service import app as app_mod
|
||||
|
||||
cfg_path = _write_config(tmp_path)
|
||||
monkeypatch.setenv("MVP_INTERNAL_TOKEN", "token1")
|
||||
|
||||
class _Scheduler:
|
||||
def __init__(self, **kwargs):
|
||||
self.tool = object()
|
||||
|
||||
def run_forever(self, stop_flag):
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(app_mod, "Scheduler", _Scheduler)
|
||||
app = app_mod.create_app(str(cfg_path))
|
||||
|
||||
with TestClient(app) as c:
|
||||
r = c.post("/api/v2/tasks", headers={"authorization": "Bearer token1"}, data="workload: nope\n")
|
||||
assert r.status_code == 400
|
||||
|
||||
|
||||
def test_submit_advanced_ok_and_command_expanded(tmp_path: Path, monkeypatch):
|
||||
from argus.service import app as app_mod
|
||||
|
||||
cfg_path = _write_config(tmp_path)
|
||||
monkeypatch.setenv("MVP_INTERNAL_TOKEN", "admin-token")
|
||||
monkeypatch.setattr(app_mod, "new_task_id", lambda workload, **kwargs: "tid_adv")
|
||||
|
||||
class _Scheduler:
|
||||
def __init__(self, **kwargs):
|
||||
self.tool = object()
|
||||
|
||||
def run_forever(self, stop_flag):
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(app_mod, "Scheduler", _Scheduler)
|
||||
app = app_mod.create_app(str(cfg_path))
|
||||
|
||||
# seed user + token
|
||||
from argus.service.config import V2Config
|
||||
from argus.service.db import Db
|
||||
|
||||
root = yaml.safe_load(cfg_path.read_text(encoding="utf-8"))
|
||||
v2_cfg = V2Config.from_root_dict(root)
|
||||
db = Db(v2_cfg.sqlite.db_path)
|
||||
db.create_user(user_id="u1", display_name=None)
|
||||
token = db.issue_token(user_id="u1")
|
||||
|
||||
spec = """
|
||||
kind: advanced
|
||||
workload: ppo
|
||||
nnodes: 2
|
||||
n_gpus_per_node: 4
|
||||
command: |
|
||||
python3 -m verl.trainer.main_ppo \\
|
||||
data.train_files=$HOME/common/datasets/gsm8k/train.parquet \\
|
||||
data.val_files=$HOME/datasets/gsm8k/test.parquet \\
|
||||
+ray_kwargs.ray_init.address=auto
|
||||
""".lstrip()
|
||||
|
||||
with TestClient(app) as c:
|
||||
r = c.post("/api/v2/tasks", headers={"authorization": f"Bearer {token}"}, data=spec)
|
||||
assert r.status_code == 200
|
||||
assert r.json()["task_id"] == "tid_adv"
|
||||
|
||||
r2 = c.get("/api/v2/tasks/tid_adv/spec", headers={"authorization": f"Bearer {token}"})
|
||||
assert r2.status_code == 200
|
||||
assert "kind: advanced" in r2.text
|
||||
# Expanded mapping for SFTPGo "virtual common" folders:
|
||||
assert "/private/datasets/gsm8k/train.parquet" in r2.text
|
||||
assert "/private/users/u1/datasets/gsm8k/test.parquet" in r2.text
|
||||
|
||||
|
||||
def test_submit_advanced_rejects_invalid_command(tmp_path: Path, monkeypatch):
|
||||
from argus.service import app as app_mod
|
||||
|
||||
cfg_path = _write_config(tmp_path)
|
||||
monkeypatch.setenv("MVP_INTERNAL_TOKEN", "admin-token")
|
||||
|
||||
class _Scheduler:
|
||||
def __init__(self, **kwargs):
|
||||
self.tool = object()
|
||||
|
||||
def run_forever(self, stop_flag):
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(app_mod, "Scheduler", _Scheduler)
|
||||
app = app_mod.create_app(str(cfg_path))
|
||||
|
||||
# seed user + token
|
||||
from argus.service.config import V2Config
|
||||
from argus.service.db import Db
|
||||
|
||||
root = yaml.safe_load(cfg_path.read_text(encoding="utf-8"))
|
||||
v2_cfg = V2Config.from_root_dict(root)
|
||||
db = Db(v2_cfg.sqlite.db_path)
|
||||
db.create_user(user_id="u1", display_name=None)
|
||||
token = db.issue_token(user_id="u1")
|
||||
|
||||
bad = """
|
||||
kind: advanced
|
||||
workload: ppo
|
||||
nnodes: 2
|
||||
n_gpus_per_node: 4
|
||||
command: |
|
||||
echo hello
|
||||
""".lstrip()
|
||||
|
||||
with TestClient(app) as c:
|
||||
r = c.post("/api/v2/tasks", headers={"authorization": f"Bearer {token}"}, data=bad)
|
||||
assert r.status_code == 400
|
||||
assert "invalid command" in r.text
|
||||
|
||||
|
||||
def test_me_sftp_reset_password_disabled_returns_400(tmp_path: Path, monkeypatch):
|
||||
from argus.service import app as app_mod
|
||||
|
||||
cfg_path = _write_config(tmp_path)
|
||||
monkeypatch.setenv("MVP_INTERNAL_TOKEN", "token1")
|
||||
|
||||
class _Scheduler:
|
||||
def __init__(self, **kwargs):
|
||||
self.tool = object()
|
||||
|
||||
def run_forever(self, stop_flag):
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(app_mod, "Scheduler", _Scheduler)
|
||||
app = app_mod.create_app(str(cfg_path))
|
||||
|
||||
# seed user + token
|
||||
from argus.service.config import V2Config
|
||||
from argus.service.db import Db
|
||||
|
||||
root = yaml.safe_load(cfg_path.read_text(encoding="utf-8"))
|
||||
v2_cfg = V2Config.from_root_dict(root)
|
||||
db = Db(v2_cfg.sqlite.db_path)
|
||||
db.init()
|
||||
db.create_user(user_id="u1", display_name=None)
|
||||
token = db.issue_token(user_id="u1")
|
||||
|
||||
with TestClient(app) as c:
|
||||
r = c.post("/api/v2/me/sftp:reset_password", headers={"authorization": f"Bearer {token}"})
|
||||
assert r.status_code == 400
|
||||
|
||||
|
||||
def test_me_sftp_reset_password_enabled_returns_password(tmp_path: Path, monkeypatch):
|
||||
from argus.service import app as app_mod
|
||||
|
||||
cfg = yaml.safe_load(_write_config(tmp_path).read_text(encoding="utf-8"))
|
||||
cfg["data"]["sftpgo"] = {
|
||||
"enabled": True,
|
||||
"host": "127.0.0.1",
|
||||
"sftp_port": 2022,
|
||||
"admin_api_base": "http://127.0.0.1:8081",
|
||||
"admin_user": "admin",
|
||||
"admin_password_env": "SFTPGO_ADMIN_PASSWORD",
|
||||
}
|
||||
cfg_path = tmp_path / "cfg_sftp.yaml"
|
||||
cfg_path.write_text(yaml.safe_dump(cfg), encoding="utf-8")
|
||||
|
||||
monkeypatch.setenv("MVP_INTERNAL_TOKEN", "token1")
|
||||
monkeypatch.setenv("SFTPGO_ADMIN_PASSWORD", "pw1")
|
||||
|
||||
class _FakeSFTPGo:
|
||||
def __init__(self, **kwargs):
|
||||
self.reset = []
|
||||
self.enabled = []
|
||||
|
||||
def reset_password(self, username: str, new_password: str, home_dir: str):
|
||||
assert username
|
||||
assert new_password
|
||||
assert home_dir
|
||||
self.reset.append((username, home_dir))
|
||||
|
||||
def enable_user(self, username: str, home_dir: str):
|
||||
self.enabled.append((username, home_dir))
|
||||
|
||||
fake_client = _FakeSFTPGo()
|
||||
|
||||
class _FakeSFTPGoFactory:
|
||||
def __call__(self, **kwargs):
|
||||
return fake_client
|
||||
|
||||
class _Scheduler:
|
||||
def __init__(self, **kwargs):
|
||||
self.tool = object()
|
||||
|
||||
def run_forever(self, stop_flag):
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(app_mod, "Scheduler", _Scheduler)
|
||||
monkeypatch.setattr(app_mod, "SFTPGoAdminClient", _FakeSFTPGoFactory())
|
||||
app = app_mod.create_app(str(cfg_path))
|
||||
|
||||
# seed user in DB
|
||||
from argus.service.db import Db
|
||||
from argus.service.config import V2Config
|
||||
|
||||
v2_cfg = V2Config.from_root_dict(cfg)
|
||||
db = Db(v2_cfg.sqlite.db_path)
|
||||
db.init()
|
||||
db.create_user(user_id="u1", display_name=None)
|
||||
token = db.issue_token(user_id="u1")
|
||||
|
||||
with TestClient(app) as c:
|
||||
r = c.post("/api/v2/me/sftp:reset_password", headers={"authorization": f"Bearer {token}"})
|
||||
assert r.status_code == 200
|
||||
j = r.json()
|
||||
assert j["user_id"] == "u1"
|
||||
assert isinstance(j["password"], str) and len(j["password"]) >= 8
|
||||
293
src/mvp/py/tests/test_app_serving_api.py
Normal file
@@ -0,0 +1,293 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import yaml
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
|
||||
def _write_config(tmp_path: Path) -> Path:
|
||||
cfg = {
|
||||
"ray": {
|
||||
"address": "http://127.0.0.1:8265",
|
||||
"shared_root": "/private",
|
||||
"entrypoint_resources": {"worker_node": 1},
|
||||
"runtime_env": {"env_vars": {}},
|
||||
},
|
||||
"data": {
|
||||
"user_root": str(tmp_path / "users"),
|
||||
},
|
||||
"service": {
|
||||
"api": {"host": "127.0.0.1", "port": 0},
|
||||
"auth": {"token_env": "MVP_INTERNAL_TOKEN"},
|
||||
"sqlite": {"db_path": str(tmp_path / "mvp.sqlite3")},
|
||||
"scheduler": {"tick_s": 1, "retry_interval_s": 1, "max_running_tasks": 1},
|
||||
},
|
||||
"serving": {
|
||||
"serve": {"http_port": 8000, "proxy_location": "HeadOnly"},
|
||||
"llm": {"accelerator_type": "H20"},
|
||||
},
|
||||
}
|
||||
p = tmp_path / "cfg.yaml"
|
||||
p.write_text(yaml.safe_dump(cfg), encoding="utf-8")
|
||||
return p
|
||||
|
||||
|
||||
def test_serving_api_crud_flow(tmp_path: Path, monkeypatch):
|
||||
from argus.service import app as app_mod
|
||||
|
||||
cfg_path = _write_config(tmp_path)
|
||||
monkeypatch.setenv("MVP_INTERNAL_TOKEN", "admin-token")
|
||||
|
||||
class _Scheduler:
|
||||
def __init__(self, **kwargs):
|
||||
self.tool = object()
|
||||
|
||||
def run_forever(self, stop_flag):
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(app_mod, "Scheduler", _Scheduler)
|
||||
monkeypatch.setattr(app_mod, "new_model_key", lambda user_id: f"mk-{user_id}")
|
||||
|
||||
app = app_mod.create_app(str(cfg_path))
|
||||
|
||||
admin_headers = {"authorization": "Bearer admin-token"}
|
||||
with TestClient(app) as c:
|
||||
r = c.post("/api/v2/users", headers=admin_headers, json={"user_id": "alice"})
|
||||
assert r.status_code == 200
|
||||
r2 = c.post("/api/v2/users/alice/tokens", headers=admin_headers)
|
||||
assert r2.status_code == 200
|
||||
user_token = r2.json()["token"]
|
||||
|
||||
headers = {"authorization": f"Bearer {user_token}"}
|
||||
|
||||
spec_yaml = (
|
||||
"model_id: qwen-0.5b\n"
|
||||
"model_source: $HOME/common/hf/models--Qwen--Qwen2.5-0.5B-Instruct/snapshots/sha\n"
|
||||
"num_replicas: 1\n"
|
||||
"gpus_per_replica: 1\n"
|
||||
)
|
||||
r3 = c.post("/api/v2/serve/models", headers=headers, data=spec_yaml)
|
||||
assert r3.status_code == 200
|
||||
assert r3.json()["model_key"] == "mk-alice"
|
||||
assert r3.json()["state"] == "QUEUED"
|
||||
|
||||
r4 = c.get("/api/v2/serve/models?limit=10&offset=0", headers=headers)
|
||||
assert r4.status_code == 200
|
||||
obj = r4.json()
|
||||
assert obj["openai_base_url"] is None
|
||||
assert len(obj["items"]) == 1
|
||||
assert obj["items"][0]["model_key"] == "mk-alice"
|
||||
assert obj["items"][0]["endpoint"]["openai_base_url"] == "http://testserver:8000/serve/mk-alice/v1"
|
||||
|
||||
r5 = c.get("/api/v2/serve/models/mk-alice", headers=headers)
|
||||
assert r5.status_code == 200
|
||||
detail = r5.json()
|
||||
assert detail["model"]["model_key"] == "mk-alice"
|
||||
assert "model_id_prefix" in detail["model"]
|
||||
assert "resolved_spec_yaml" in detail
|
||||
assert isinstance(detail.get("events"), list)
|
||||
|
||||
r6 = c.patch("/api/v2/serve/models/mk-alice", headers=headers, json={"num_replicas": 2})
|
||||
assert r6.status_code == 200
|
||||
assert r6.json()["state"] == "QUEUED"
|
||||
|
||||
r7 = c.delete("/api/v2/serve/models/mk-alice", headers=headers)
|
||||
assert r7.status_code == 200
|
||||
assert r7.json()["state"] == "DELETING"
|
||||
|
||||
# Admin status endpoint
|
||||
r8 = c.get("/api/v2/serve/status", headers=admin_headers)
|
||||
assert r8.status_code == 200
|
||||
assert r8.json()["http_port"] == 8000
|
||||
assert r8.json()["openai_base_url"] is None
|
||||
|
||||
|
||||
def test_serving_api_rejects_path_outside_user_and_hf(tmp_path: Path, monkeypatch):
|
||||
from argus.service import app as app_mod
|
||||
|
||||
cfg_path = _write_config(tmp_path)
|
||||
monkeypatch.setenv("MVP_INTERNAL_TOKEN", "admin-token")
|
||||
|
||||
class _Scheduler:
|
||||
def __init__(self, **kwargs):
|
||||
self.tool = object()
|
||||
|
||||
def run_forever(self, stop_flag):
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(app_mod, "Scheduler", _Scheduler)
|
||||
monkeypatch.setattr(app_mod, "new_model_key", lambda user_id: f"mk-{user_id}")
|
||||
|
||||
app = app_mod.create_app(str(cfg_path))
|
||||
|
||||
admin_headers = {"authorization": "Bearer admin-token"}
|
||||
with TestClient(app) as c:
|
||||
c.post("/api/v2/users", headers=admin_headers, json={"user_id": "alice"})
|
||||
r2 = c.post("/api/v2/users/alice/tokens", headers=admin_headers)
|
||||
user_token = r2.json()["token"]
|
||||
headers = {"authorization": f"Bearer {user_token}"}
|
||||
|
||||
spec_yaml = (
|
||||
"model_id: x\n"
|
||||
"model_source: /private/users/bob/models/evil\n"
|
||||
"num_replicas: 1\n"
|
||||
"gpus_per_replica: 1\n"
|
||||
)
|
||||
r3 = c.post("/api/v2/serve/models", headers=headers, data=spec_yaml)
|
||||
assert r3.status_code == 403
|
||||
|
||||
|
||||
def test_serving_api_invalid_yaml_and_non_mapping(tmp_path: Path, monkeypatch):
|
||||
from argus.service import app as app_mod
|
||||
|
||||
cfg_path = _write_config(tmp_path)
|
||||
monkeypatch.setenv("MVP_INTERNAL_TOKEN", "admin-token")
|
||||
|
||||
class _Scheduler:
|
||||
def __init__(self, **kwargs):
|
||||
self.tool = object()
|
||||
|
||||
def run_forever(self, stop_flag):
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(app_mod, "Scheduler", _Scheduler)
|
||||
monkeypatch.setattr(app_mod, "new_model_key", lambda user_id: f"mk-{user_id}")
|
||||
app = app_mod.create_app(str(cfg_path))
|
||||
|
||||
with TestClient(app) as c:
|
||||
# Create a user token
|
||||
admin_headers = {"authorization": "Bearer admin-token"}
|
||||
c.post("/api/v2/users", headers=admin_headers, json={"user_id": "alice"})
|
||||
token = c.post("/api/v2/users/alice/tokens", headers=admin_headers).json()["token"]
|
||||
headers = {"authorization": f"Bearer {token}"}
|
||||
|
||||
r = c.post("/api/v2/serve/models", headers=headers, data=": bad\n")
|
||||
assert r.status_code == 400
|
||||
|
||||
r2 = c.post("/api/v2/serve/models", headers=headers, data="- 1\n- 2\n")
|
||||
assert r2.status_code == 400
|
||||
|
||||
|
||||
def test_serving_api_engine_kwargs_binary_rejected(tmp_path: Path, monkeypatch):
|
||||
"""
|
||||
yaml !!binary is parsed as bytes, which is not JSON-serializable.
|
||||
"""
|
||||
from argus.service import app as app_mod
|
||||
|
||||
cfg_path = _write_config(tmp_path)
|
||||
monkeypatch.setenv("MVP_INTERNAL_TOKEN", "admin-token")
|
||||
|
||||
class _Scheduler:
|
||||
def __init__(self, **kwargs):
|
||||
self.tool = object()
|
||||
|
||||
def run_forever(self, stop_flag):
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(app_mod, "Scheduler", _Scheduler)
|
||||
monkeypatch.setattr(app_mod, "new_model_key", lambda user_id: f"mk-{user_id}")
|
||||
app = app_mod.create_app(str(cfg_path))
|
||||
|
||||
admin_headers = {"authorization": "Bearer admin-token"}
|
||||
with TestClient(app) as c:
|
||||
c.post("/api/v2/users", headers=admin_headers, json={"user_id": "alice"})
|
||||
token = c.post("/api/v2/users/alice/tokens", headers=admin_headers).json()["token"]
|
||||
headers = {"authorization": f"Bearer {token}"}
|
||||
|
||||
spec_yaml = (
|
||||
"model_id: x\n"
|
||||
"model_source: $HOME/common/hf/x\n"
|
||||
"engine_kwargs:\n"
|
||||
" blob: !!binary \"AQID\"\n"
|
||||
)
|
||||
r = c.post("/api/v2/serve/models", headers=headers, data=spec_yaml)
|
||||
assert r.status_code == 400
|
||||
|
||||
|
||||
def test_serving_api_list_include_deleted_and_forwarded_base_url(tmp_path: Path, monkeypatch):
|
||||
from argus.service import app as app_mod
|
||||
from argus.service.config import V2Config
|
||||
from argus.service.db import Db
|
||||
|
||||
cfg_path = _write_config(tmp_path)
|
||||
monkeypatch.setenv("MVP_INTERNAL_TOKEN", "admin-token")
|
||||
|
||||
class _Scheduler:
|
||||
def __init__(self, **kwargs):
|
||||
self.tool = object()
|
||||
|
||||
def run_forever(self, stop_flag):
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(app_mod, "Scheduler", _Scheduler)
|
||||
keys = iter(["mk-alice-1", "mk-alice-2"])
|
||||
monkeypatch.setattr(app_mod, "new_model_key", lambda user_id: next(keys))
|
||||
|
||||
app = app_mod.create_app(str(cfg_path))
|
||||
|
||||
admin_headers = {"authorization": "Bearer admin-token"}
|
||||
with TestClient(app) as c:
|
||||
c.post("/api/v2/users", headers=admin_headers, json={"user_id": "alice"})
|
||||
token = c.post("/api/v2/users/alice/tokens", headers=admin_headers).json()["token"]
|
||||
headers = {"authorization": f"Bearer {token}"}
|
||||
|
||||
spec_yaml = "model_id: x\nmodel_source: $HOME/common/hf/x\n"
|
||||
c.post("/api/v2/serve/models", headers=headers, data=spec_yaml)
|
||||
c.post("/api/v2/serve/models", headers=headers, data=spec_yaml)
|
||||
|
||||
# Mark one model as DELETED directly in DB (sets deleted_at).
|
||||
root = yaml.safe_load(cfg_path.read_text(encoding="utf-8"))
|
||||
v2_cfg = V2Config.from_root_dict(root)
|
||||
db = Db(v2_cfg.sqlite.db_path)
|
||||
db.set_serve_model_state(model_key="mk-alice-2", state="DELETED")
|
||||
|
||||
r1 = c.get(
|
||||
"/api/v2/serve/models?limit=10&offset=0&include_deleted=0",
|
||||
headers={**headers, "x-forwarded-host": "example.com:8080", "x-forwarded-proto": "https"},
|
||||
)
|
||||
assert r1.status_code == 200
|
||||
assert r1.json()["openai_base_url"] is None
|
||||
assert {m["model_key"] for m in r1.json()["items"]} == {"mk-alice-1"}
|
||||
assert r1.json()["items"][0]["endpoint"]["openai_base_url"] == "https://example.com:8000/serve/mk-alice-1/v1"
|
||||
|
||||
# IPv6 bracketed host should be preserved.
|
||||
r1b = c.get(
|
||||
"/api/v2/serve/models?limit=10&offset=0&include_deleted=0",
|
||||
headers={**headers, "x-forwarded-host": "[::1]:8080", "x-forwarded-proto": "https"},
|
||||
)
|
||||
assert r1b.status_code == 200
|
||||
assert r1b.json()["items"][0]["endpoint"]["openai_base_url"] == "https://[::1]:8000/serve/mk-alice-1/v1"
|
||||
|
||||
r2 = c.get("/api/v2/serve/models?include_deleted=1", headers=headers)
|
||||
assert r2.status_code == 200
|
||||
assert {m["model_key"] for m in r2.json()["items"]} == {"mk-alice-1", "mk-alice-2"}
|
||||
|
||||
|
||||
def test_serving_api_patch_invalid_num_replicas(tmp_path: Path, monkeypatch):
|
||||
from argus.service import app as app_mod
|
||||
|
||||
cfg_path = _write_config(tmp_path)
|
||||
monkeypatch.setenv("MVP_INTERNAL_TOKEN", "admin-token")
|
||||
|
||||
class _Scheduler:
|
||||
def __init__(self, **kwargs):
|
||||
self.tool = object()
|
||||
|
||||
def run_forever(self, stop_flag):
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(app_mod, "Scheduler", _Scheduler)
|
||||
monkeypatch.setattr(app_mod, "new_model_key", lambda user_id: "mk-alice")
|
||||
|
||||
app = app_mod.create_app(str(cfg_path))
|
||||
|
||||
admin_headers = {"authorization": "Bearer admin-token"}
|
||||
with TestClient(app) as c:
|
||||
c.post("/api/v2/users", headers=admin_headers, json={"user_id": "alice"})
|
||||
token = c.post("/api/v2/users/alice/tokens", headers=admin_headers).json()["token"]
|
||||
headers = {"authorization": f"Bearer {token}"}
|
||||
|
||||
c.post("/api/v2/serve/models", headers=headers, data="model_id: x\nmodel_source: $HOME/common/hf/x\n")
|
||||
r = c.patch("/api/v2/serve/models/mk-alice", headers=headers, json={"num_replicas": 0})
|
||||
assert r.status_code == 422
|
||||
68
src/mvp/py/tests/test_builders.py
Normal file
@@ -0,0 +1,68 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from argus.ray.builders import build_training_argv
|
||||
from argus.ray.models import JobSpec
|
||||
|
||||
|
||||
def _base_spec(workload: str, *, total_training_steps: int | None = 10) -> JobSpec:
|
||||
return JobSpec(
|
||||
workload=workload,
|
||||
submission_id=None,
|
||||
code_path="/code",
|
||||
model_id="m",
|
||||
train_file="train.jsonl",
|
||||
val_file="val.jsonl",
|
||||
nnodes=2,
|
||||
n_gpus_per_node=4,
|
||||
total_epochs=1,
|
||||
total_training_steps=total_training_steps,
|
||||
save_freq=10,
|
||||
test_freq=None,
|
||||
trainer_device=None,
|
||||
)
|
||||
|
||||
|
||||
def test_build_training_argv_ppo_smoke():
|
||||
spec = _base_spec("ppo")
|
||||
built = build_training_argv(spec, submission_id="sid", job_dir="/job")
|
||||
assert built.argv[:3] == ["python3", "-m", "verl.trainer.main_ppo"]
|
||||
assert "data.val_files=val.jsonl" in built.argv
|
||||
# v3.7: default rollout backend switches from sglang -> vllm.
|
||||
assert "actor_rollout_ref.rollout.name=vllm" in built.argv
|
||||
assert "actor_rollout_ref.rollout.name=sglang" not in built.argv
|
||||
assert "trainer.test_freq=-1" in built.argv
|
||||
|
||||
|
||||
def test_build_training_argv_grpo_has_override():
|
||||
spec = _base_spec("grpo")
|
||||
built = build_training_argv(spec, submission_id="sid", job_dir="/job")
|
||||
assert "algorithm.adv_estimator=grpo" in built.argv
|
||||
|
||||
|
||||
def test_build_training_argv_sft_smoke():
|
||||
spec = _base_spec("sft")
|
||||
built = build_training_argv(spec, submission_id="sid", job_dir="/job")
|
||||
assert built.argv[:3] == ["python3", "-m", "verl.trainer.sft_trainer_ray"]
|
||||
assert "trainer.device=cpu" in built.argv
|
||||
assert "data.val_files=val.jsonl" in built.argv
|
||||
|
||||
|
||||
def test_build_training_argv_omits_total_training_steps_when_unset():
|
||||
spec = _base_spec("ppo", total_training_steps=None)
|
||||
built = build_training_argv(spec, submission_id="sid", job_dir="/job")
|
||||
assert not any(a.startswith("trainer.total_training_steps=") for a in built.argv)
|
||||
|
||||
|
||||
def test_build_training_argv_unsupported_raises():
|
||||
spec = _base_spec("bad")
|
||||
with pytest.raises(ValueError, match="unsupported workload"):
|
||||
build_training_argv(spec, submission_id="sid", job_dir="/job")
|
||||
|
||||
|
||||
def test_build_advanced_argv_uses_bash_lc():
|
||||
from argus.ray.builders import build_advanced_argv
|
||||
|
||||
built = build_advanced_argv("echo hi")
|
||||
assert built.argv[:2] == ["bash", "-lc"]
|
||||
77
src/mvp/py/tests/test_cli_run.py
Normal file
@@ -0,0 +1,77 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
def test_cli_submit_status_logs_list(monkeypatch, tmp_path: Path, capsys):
|
||||
from argus.ray import ray_job_tool as tool_mod
|
||||
|
||||
class _Tool:
|
||||
def __init__(self, cfg):
|
||||
self.cfg = cfg
|
||||
|
||||
def submit(self, spec, no_wait: bool):
|
||||
return "sid1"
|
||||
|
||||
def status(self, sid: str):
|
||||
return "RUNNING"
|
||||
|
||||
def stop(self, sid: str):
|
||||
return True
|
||||
|
||||
def logs(self, sid: str):
|
||||
return "1\n2\n3\n"
|
||||
|
||||
def list(self):
|
||||
return [{"a": 1}]
|
||||
|
||||
monkeypatch.setattr(tool_mod, "RayJobTool", _Tool)
|
||||
|
||||
cfg = tmp_path / "cfg.yaml"
|
||||
cfg.write_text(
|
||||
"ray:\n address: http://127.0.0.1:8265\n shared_root: /private\n entrypoint_resources: {worker_node: 1}\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
spec = tmp_path / "spec.yaml"
|
||||
spec.write_text("workload: ppo\ncode_path: /c\nmodel_id: m\ntrain_file: t\nval_file: v\n", encoding="utf-8")
|
||||
|
||||
from argus.cli.run import main
|
||||
|
||||
monkeypatch.setattr(sys, "argv", ["run.py", "--config", str(cfg), "--taskspec", str(spec), "--action", "submit"])
|
||||
assert main() == 0
|
||||
assert capsys.readouterr().out.strip() == "sid1"
|
||||
|
||||
monkeypatch.setattr(sys, "argv", ["run.py", "--config", str(cfg), "--action", "status", "--submission-id", "sid1"])
|
||||
assert main() == 0
|
||||
assert capsys.readouterr().out.strip() == "RUNNING"
|
||||
|
||||
monkeypatch.setattr(
|
||||
sys, "argv", ["run.py", "--config", str(cfg), "--action", "logs", "--submission-id", "sid1", "--tail", "1"]
|
||||
)
|
||||
assert main() == 0
|
||||
assert capsys.readouterr().out.strip() == "3"
|
||||
|
||||
monkeypatch.setattr(sys, "argv", ["run.py", "--config", str(cfg), "--action", "list"])
|
||||
assert main() == 0
|
||||
out = capsys.readouterr().out
|
||||
assert json.loads(out)[0]["a"] == 1
|
||||
|
||||
|
||||
def test_cli_requires_submission_id_for_status():
|
||||
from argus.cli.run import main
|
||||
|
||||
tmp = Path(__import__("tempfile").gettempdir()) / "mvp_test_cfg.yaml"
|
||||
tmp.write_text(
|
||||
"ray:\n address: http://127.0.0.1:8265\n shared_root: /private\n entrypoint_resources: {worker_node: 1}\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
try:
|
||||
with pytest.raises(SystemExit):
|
||||
sys.argv = ["run.py", "--config", str(tmp), "--action", "status"]
|
||||
main()
|
||||
finally:
|
||||
tmp.unlink(missing_ok=True)
|
||||
42
src/mvp/py/tests/test_db.py
Normal file
@@ -0,0 +1,42 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def test_db_lifecycle_and_basic_queries(tmp_path: Path):
|
||||
from argus.service.db import Db
|
||||
|
||||
db = Db(str(tmp_path / "mvp.sqlite3"))
|
||||
db.init()
|
||||
|
||||
t = db.create_task(
|
||||
task_id="t1",
|
||||
workload="ppo",
|
||||
jobspec_yaml="workload: ppo\ncode_path: /c\nmodel_id: m\ntrain_file: t\nval_file: v\n",
|
||||
nnodes=2,
|
||||
n_gpus_per_node=4,
|
||||
)
|
||||
assert t["task_id"] == "t1"
|
||||
assert db.get_task("t1") is not None
|
||||
|
||||
q = db.list_queue()
|
||||
assert len(q["pending"]) == 1
|
||||
|
||||
db.set_task_state(task_id="t1", state="PENDING_RESOURCES", next_run_at="2099-01-01T00:00:00Z")
|
||||
assert db.pick_next_runnable_task() is None
|
||||
|
||||
# next_run_at is sticky unless explicitly updated; a future value keeps it non-runnable.
|
||||
db.set_task_state(task_id="t1", state="QUEUED", next_run_at=None)
|
||||
assert db.pick_next_runnable_task() is None
|
||||
|
||||
# Allow it by setting next_run_at into the past.
|
||||
db.set_task_state(task_id="t1", state="QUEUED", next_run_at="2000-01-01T00:00:00Z")
|
||||
assert db.pick_next_runnable_task() is not None
|
||||
|
||||
db.create_attempt(task_id="t1", attempt_no=1, ray_submission_id="sid1")
|
||||
db.update_attempt(task_id="t1", attempt_no=1, ray_status="RUNNING")
|
||||
attempts = db.list_attempts("t1")
|
||||
assert attempts[-1]["ray_status"] == "RUNNING"
|
||||
|
||||
# No-op update is allowed
|
||||
db.update_attempt(task_id="t1", attempt_no=1)
|
||||
79
src/mvp/py/tests/test_db_serving.py
Normal file
@@ -0,0 +1,79 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def test_db_serving_model_crud_and_events(tmp_path: Path) -> None:
|
||||
from argus.service.db import Db
|
||||
|
||||
db = Db(str(tmp_path / "mvp.sqlite3"))
|
||||
db.init()
|
||||
|
||||
m1 = db.create_serve_model(
|
||||
model_key="svc-001",
|
||||
user_id="alice",
|
||||
model_id_suffix="qwen-0.5b",
|
||||
model_id_prefix="alice-202601061235",
|
||||
model_id="alice-202601061235-qwen-0.5b",
|
||||
model_source="/private/hf/hub/models--Qwen--Qwen2.5-0.5B-Instruct/snapshots/sha",
|
||||
num_replicas=1,
|
||||
gpus_per_replica=1,
|
||||
engine_kwargs_json=json.dumps({"max_model_len": 8192}),
|
||||
spec_yaml="model_id: qwen-0.5b\nmodel_source: $HOME/common/hf/...\n",
|
||||
resolved_spec_yaml="model_id: alice-202601061235-qwen-0.5b\nmodel_source: /private/hf/...\n",
|
||||
)
|
||||
assert m1["model_key"] == "svc-001"
|
||||
assert m1["state"] == "QUEUED"
|
||||
|
||||
# Same suffix may be created again; model_key is the identity.
|
||||
m2 = db.create_serve_model(
|
||||
model_key="svc-002",
|
||||
user_id="alice",
|
||||
model_id_suffix="qwen-0.5b",
|
||||
model_id_prefix="alice-202601061236",
|
||||
model_id="alice-202601061236-qwen-0.5b",
|
||||
model_source="/private/hf/hub/models--Qwen--Qwen2.5-0.5B-Instruct/snapshots/sha",
|
||||
num_replicas=1,
|
||||
gpus_per_replica=2,
|
||||
engine_kwargs_json=None,
|
||||
spec_yaml="model_id: qwen-0.5b\nmodel_source: $HOME/common/hf/...\n",
|
||||
resolved_spec_yaml="model_id: alice-202601061236-qwen-0.5b\nmodel_source: /private/hf/...\n",
|
||||
)
|
||||
assert m2["model_key"] == "svc-002"
|
||||
assert m2["model_id"] != m1["model_id"]
|
||||
|
||||
got = db.get_serve_model("svc-001")
|
||||
assert got is not None
|
||||
assert got["gpus_per_replica"] == 1
|
||||
|
||||
items = db.list_serve_models(user_id="alice")
|
||||
assert {i["model_key"] for i in items} == {"svc-001", "svc-002"}
|
||||
|
||||
# State transition writes a serve event.
|
||||
db.set_serve_model_state(model_key="svc-001", state="DEPLOYING")
|
||||
got2 = db.get_serve_model("svc-001")
|
||||
assert got2 is not None
|
||||
assert got2["state"] == "DEPLOYING"
|
||||
|
||||
events = db.list_serve_events("svc-001", limit=50)
|
||||
assert len(events) >= 2
|
||||
assert {e["event_type"] for e in events}.issuperset({"SERVE_MODEL_CREATED", "SERVE_STATE_UPDATE"})
|
||||
|
||||
# Reconciler pick: QUEUED/DELETING only.
|
||||
picked = db.pick_next_runnable_serve_change()
|
||||
assert picked is not None
|
||||
assert picked["state"] == "QUEUED"
|
||||
|
||||
db.set_serve_model_state(model_key="svc-002", state="DELETING")
|
||||
picked2 = db.pick_next_runnable_serve_change()
|
||||
assert picked2 is not None
|
||||
assert picked2["state"] in ("QUEUED", "DELETING")
|
||||
|
||||
# Deleted models are hidden unless include_deleted.
|
||||
db.set_serve_model_state(model_key="svc-002", state="DELETED")
|
||||
items2 = db.list_serve_models(user_id="alice", include_deleted=False)
|
||||
assert {i["model_key"] for i in items2} == {"svc-001"}
|
||||
items3 = db.list_serve_models(user_id="alice", include_deleted=True)
|
||||
assert {i["model_key"] for i in items3} == {"svc-001", "svc-002"}
|
||||
|
||||
50
src/mvp/py/tests/test_discovery.py
Normal file
@@ -0,0 +1,50 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from pathlib import Path
|
||||
|
||||
from argus.ray.discovery import build_head_record, load_head_record
|
||||
|
||||
|
||||
def test_load_head_record_rejects_missing(tmp_path: Path):
|
||||
assert load_head_record(str(tmp_path / "nope.json")) is None
|
||||
|
||||
|
||||
def test_load_head_record_rejects_expired(tmp_path: Path):
|
||||
p = tmp_path / "head.json"
|
||||
now = datetime.now(timezone.utc).replace(microsecond=0)
|
||||
rec = build_head_record(cluster_name="c", head_ip="1.2.3.4", ttl_s=1, now=now)
|
||||
p.write_text(json.dumps(rec.to_dict()), encoding="utf-8")
|
||||
assert load_head_record(str(p), now=now + timedelta(seconds=2)) is None
|
||||
|
||||
|
||||
def test_load_head_record_accepts_fresh(tmp_path: Path):
|
||||
p = tmp_path / "head.json"
|
||||
now = datetime.now(timezone.utc).replace(microsecond=0)
|
||||
rec = build_head_record(cluster_name="c", head_ip="1.2.3.4", ttl_s=60, now=now)
|
||||
p.write_text(json.dumps(rec.to_dict()), encoding="utf-8")
|
||||
got = load_head_record(str(p), now=now + timedelta(seconds=1))
|
||||
assert got is not None
|
||||
assert got.head_addr() == "1.2.3.4:6379"
|
||||
|
||||
|
||||
def test_head_publisher_main_once_writes_file(tmp_path: Path):
|
||||
from argus.ray import head_publisher as mod
|
||||
|
||||
p = tmp_path / "head.json"
|
||||
rc = mod.main(
|
||||
[
|
||||
"--cluster-name",
|
||||
"c",
|
||||
"--head-ip-file",
|
||||
str(p),
|
||||
"--head-ip",
|
||||
"9.9.9.9",
|
||||
"--ttl-s",
|
||||
"60",
|
||||
"--once",
|
||||
]
|
||||
)
|
||||
assert rc == 0
|
||||
assert "9.9.9.9" in p.read_text(encoding="utf-8")
|
||||
27
src/mvp/py/tests/test_driver_entrypoint.py
Normal file
@@ -0,0 +1,27 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def test_driver_entrypoint_missing_cmd_returns_2(monkeypatch, tmp_path: Path):
|
||||
from argus.ray.driver_entrypoint import main
|
||||
|
||||
monkeypatch.setattr(sys, "argv", ["x", "--job-dir", str(tmp_path)])
|
||||
assert main() == 2
|
||||
|
||||
|
||||
def test_driver_entrypoint_strips_double_dash_and_returns_code(monkeypatch, tmp_path: Path):
|
||||
from argus.ray import driver_entrypoint as mod
|
||||
|
||||
class _Proc:
|
||||
returncode = 7
|
||||
|
||||
monkeypatch.setattr(mod.subprocess, "run", lambda cmd, check: _Proc())
|
||||
got = {}
|
||||
monkeypatch.setattr(mod.os, "chdir", lambda p: got.update({"cwd": str(p)}))
|
||||
monkeypatch.setattr(sys, "argv", ["x", "--job-dir", str(tmp_path), "--", "echo", "hi"])
|
||||
|
||||
assert mod.main() == 7
|
||||
assert (tmp_path).exists()
|
||||
assert got["cwd"] == str(tmp_path)
|
||||
75
src/mvp/py/tests/test_ids.py
Normal file
@@ -0,0 +1,75 @@
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
def test_new_task_id_is_deterministic_with_patches(monkeypatch):
|
||||
import argus.core.ids as ids
|
||||
|
||||
class _FakeDatetime:
|
||||
@staticmethod
|
||||
def now():
|
||||
class _DT:
|
||||
def strftime(self, fmt: str) -> str:
|
||||
assert fmt == "%Y%m%d-%H%M%S"
|
||||
return "20250101-010203"
|
||||
|
||||
return _DT()
|
||||
|
||||
monkeypatch.setattr(ids, "datetime", _FakeDatetime)
|
||||
monkeypatch.setattr(ids.secrets, "token_hex", lambda n: "abcd")
|
||||
|
||||
assert ids.new_task_id("ppo") == "mvp2-ppo-20250101-010203-abcd"
|
||||
|
||||
|
||||
def test_new_task_id_includes_user_when_provided(monkeypatch):
|
||||
import argus.core.ids as ids
|
||||
|
||||
class _FakeDatetime:
|
||||
@staticmethod
|
||||
def now():
|
||||
class _DT:
|
||||
def strftime(self, fmt: str) -> str:
|
||||
assert fmt == "%Y%m%d-%H%M%S"
|
||||
return "20250101-010203"
|
||||
|
||||
return _DT()
|
||||
|
||||
monkeypatch.setattr(ids, "datetime", _FakeDatetime)
|
||||
monkeypatch.setattr(ids.secrets, "token_hex", lambda n: "abcd")
|
||||
|
||||
assert ids.new_task_id("ppo", user_id="Alice_01") == "mvp2-alice_01-ppo-20250101-010203-abcd"
|
||||
|
||||
|
||||
def test_attempt_submission_id_format():
|
||||
from argus.core.ids import attempt_submission_id
|
||||
|
||||
assert attempt_submission_id("t", 1) == "t--a01"
|
||||
assert attempt_submission_id("t", 12) == "t--a12"
|
||||
|
||||
|
||||
def test_new_model_key_includes_user(monkeypatch):
|
||||
import argus.core.ids as ids
|
||||
|
||||
class _FakeDatetime:
|
||||
@staticmethod
|
||||
def now():
|
||||
class _DT:
|
||||
def strftime(self, fmt: str) -> str:
|
||||
assert fmt == "%Y%m%d-%H%M%S"
|
||||
return "20250101-010203"
|
||||
|
||||
return _DT()
|
||||
|
||||
monkeypatch.setattr(ids, "datetime", _FakeDatetime)
|
||||
monkeypatch.setattr(ids.secrets, "token_hex", lambda n: "abcd")
|
||||
|
||||
assert ids.new_model_key(user_id="Alice_01") == "mvp2-alice_01-serve-20250101-010203-abcd"
|
||||
|
||||
|
||||
def test_new_model_key_requires_user_id():
|
||||
from argus.core.ids import new_model_key
|
||||
|
||||
try:
|
||||
new_model_key(user_id="")
|
||||
assert False, "expected ValueError"
|
||||
except ValueError as e:
|
||||
assert "user_id is required" in str(e)
|
||||
225
src/mvp/py/tests/test_janitor.py
Normal file
@@ -0,0 +1,225 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from pathlib import Path
|
||||
|
||||
from argus.service.db import Db
|
||||
from argus.service.janitor import JobsJanitor
|
||||
|
||||
|
||||
def _iso_z(dt: datetime) -> str:
|
||||
return dt.astimezone(timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z")
|
||||
|
||||
|
||||
def _mk_job_dir(user_root: Path, user_id: str, sid: str) -> Path:
|
||||
p = user_root / user_id / "jobs" / sid
|
||||
p.mkdir(parents=True, exist_ok=True)
|
||||
(p / "marker.txt").write_text("x", encoding="utf-8")
|
||||
return p
|
||||
|
||||
|
||||
def _mk_trash_dir(user_root: Path, user_id: str, sid: str) -> Path:
|
||||
p = user_root / user_id / "trash" / "jobs" / sid
|
||||
p.mkdir(parents=True, exist_ok=True)
|
||||
(p / "marker.txt").write_text("x", encoding="utf-8")
|
||||
return p
|
||||
|
||||
|
||||
def test_janitor_moves_jobs_to_trash_after_3_days(tmp_path: Path) -> None:
|
||||
db_path = tmp_path / "mvp.sqlite3"
|
||||
user_root = tmp_path / "users"
|
||||
db = Db(str(db_path))
|
||||
db.init()
|
||||
|
||||
task_id = "t1"
|
||||
user_id = "alice"
|
||||
sid = "sid-a01"
|
||||
db.create_task_v25(task_id=task_id, user_id=user_id, workload="sft", jobspec_yaml="{}", nnodes=1, n_gpus_per_node=1)
|
||||
db.create_attempt(task_id=task_id, attempt_no=1, ray_submission_id=sid)
|
||||
|
||||
now = datetime(2025, 1, 10, tzinfo=timezone.utc)
|
||||
ended = now - timedelta(days=4)
|
||||
db.update_attempt(task_id=task_id, attempt_no=1, end_time=_iso_z(ended), ray_status="SUCCEEDED")
|
||||
|
||||
src = _mk_job_dir(user_root, user_id, sid)
|
||||
jan = JobsJanitor(db=db, user_root=str(user_root), trash_after_days=3, purge_after_days=7, interval_s=1)
|
||||
jan.tick_once(now=now)
|
||||
|
||||
assert not src.exists()
|
||||
dst = user_root / user_id / "trash" / "jobs" / sid
|
||||
assert dst.exists()
|
||||
assert (dst / "marker.txt").read_text(encoding="utf-8") == "x"
|
||||
|
||||
|
||||
def test_janitor_purges_from_trash_after_7_days(tmp_path: Path) -> None:
|
||||
db_path = tmp_path / "mvp.sqlite3"
|
||||
user_root = tmp_path / "users"
|
||||
db = Db(str(db_path))
|
||||
db.init()
|
||||
|
||||
task_id = "t2"
|
||||
user_id = "alice"
|
||||
sid = "sid-a01"
|
||||
db.create_task_v25(task_id=task_id, user_id=user_id, workload="ppo", jobspec_yaml="{}", nnodes=1, n_gpus_per_node=1)
|
||||
db.create_attempt(task_id=task_id, attempt_no=1, ray_submission_id=sid)
|
||||
|
||||
now = datetime(2025, 1, 10, tzinfo=timezone.utc)
|
||||
ended = now - timedelta(days=8)
|
||||
db.update_attempt(task_id=task_id, attempt_no=1, end_time=_iso_z(ended), ray_status="FAILED")
|
||||
|
||||
_mk_trash_dir(user_root, user_id, sid)
|
||||
jan = JobsJanitor(db=db, user_root=str(user_root), trash_after_days=3, purge_after_days=7, interval_s=1)
|
||||
jan.tick_once(now=now)
|
||||
|
||||
dst = user_root / user_id / "trash" / "jobs" / sid
|
||||
assert not dst.exists()
|
||||
|
||||
|
||||
def test_janitor_does_not_touch_recent_jobs(tmp_path: Path) -> None:
|
||||
db_path = tmp_path / "mvp.sqlite3"
|
||||
user_root = tmp_path / "users"
|
||||
db = Db(str(db_path))
|
||||
db.init()
|
||||
|
||||
task_id = "t3"
|
||||
user_id = "alice"
|
||||
sid = "sid-a01"
|
||||
db.create_task_v25(task_id=task_id, user_id=user_id, workload="grpo", jobspec_yaml="{}", nnodes=1, n_gpus_per_node=1)
|
||||
db.create_attempt(task_id=task_id, attempt_no=1, ray_submission_id=sid)
|
||||
|
||||
now = datetime(2025, 1, 10, tzinfo=timezone.utc)
|
||||
ended = now - timedelta(days=1)
|
||||
db.update_attempt(task_id=task_id, attempt_no=1, end_time=_iso_z(ended), ray_status="FAILED")
|
||||
|
||||
src = _mk_job_dir(user_root, user_id, sid)
|
||||
jan = JobsJanitor(db=db, user_root=str(user_root), trash_after_days=3, purge_after_days=7, interval_s=1)
|
||||
jan.tick_once(now=now)
|
||||
|
||||
assert src.exists()
|
||||
assert not (user_root / user_id / "trash" / "jobs" / sid).exists()
|
||||
|
||||
|
||||
def test_janitor_skips_tasks_without_user_id(tmp_path: Path) -> None:
|
||||
db_path = tmp_path / "mvp.sqlite3"
|
||||
user_root = tmp_path / "users"
|
||||
db = Db(str(db_path))
|
||||
db.init()
|
||||
|
||||
task_id = "legacy"
|
||||
sid = "sid-legacy"
|
||||
db.create_task(task_id=task_id, workload="sft", jobspec_yaml="{}", nnodes=1, n_gpus_per_node=1)
|
||||
db.create_attempt(task_id=task_id, attempt_no=1, ray_submission_id=sid)
|
||||
|
||||
now = datetime(2025, 1, 10, tzinfo=timezone.utc)
|
||||
ended = now - timedelta(days=10)
|
||||
db.update_attempt(task_id=task_id, attempt_no=1, end_time=_iso_z(ended), ray_status="SUCCEEDED")
|
||||
|
||||
# Even if someone created a matching directory under user_root, janitor shouldn't touch it because user_id is NULL.
|
||||
src = _mk_job_dir(user_root, "alice", sid)
|
||||
jan = JobsJanitor(db=db, user_root=str(user_root), trash_after_days=3, purge_after_days=7, interval_s=1)
|
||||
jan.tick_once(now=now)
|
||||
assert src.exists()
|
||||
|
||||
|
||||
def test_janitor_validates_retention_days(tmp_path: Path) -> None:
|
||||
db = Db(str(tmp_path / "mvp.sqlite3"))
|
||||
db.init()
|
||||
try:
|
||||
JobsJanitor(db=db, user_root="/tmp", trash_after_days=-1, purge_after_days=7, interval_s=1)
|
||||
raise AssertionError("expected ValueError")
|
||||
except ValueError:
|
||||
pass
|
||||
try:
|
||||
JobsJanitor(db=db, user_root="/tmp", trash_after_days=3, purge_after_days=1, interval_s=1)
|
||||
raise AssertionError("expected ValueError")
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
|
||||
def test_janitor_noop_when_disabled(tmp_path: Path) -> None:
|
||||
db = Db(str(tmp_path / "mvp.sqlite3"))
|
||||
db.init()
|
||||
user_root = tmp_path / "users"
|
||||
jan = JobsJanitor(db=db, user_root=str(user_root), trash_after_days=0, purge_after_days=0, interval_s=1)
|
||||
jan.tick_once(now=datetime(2025, 1, 1, tzinfo=timezone.utc))
|
||||
|
||||
|
||||
def test_janitor_handles_invalid_end_time_and_missing_fields(tmp_path: Path) -> None:
|
||||
db_path = tmp_path / "mvp.sqlite3"
|
||||
user_root = tmp_path / "users"
|
||||
db = Db(str(db_path))
|
||||
db.init()
|
||||
|
||||
now = datetime(2025, 1, 10, tzinfo=timezone.utc)
|
||||
cutoff_ended = now - timedelta(days=10)
|
||||
|
||||
# Missing end_time (empty string) => should be skipped.
|
||||
db.create_task_v25(task_id="t4", user_id="alice", workload="sft", jobspec_yaml="{}", nnodes=1, n_gpus_per_node=1)
|
||||
db.create_attempt(task_id="t4", attempt_no=1, ray_submission_id="sid-empty")
|
||||
db.update_attempt(task_id="t4", attempt_no=1, end_time="", ray_status="SUCCEEDED")
|
||||
|
||||
# Invalid ISO but still lexicographically <= cutoff => should be skipped.
|
||||
db.create_task_v25(task_id="t5", user_id="alice", workload="sft", jobspec_yaml="{}", nnodes=1, n_gpus_per_node=1)
|
||||
db.create_attempt(task_id="t5", attempt_no=1, ray_submission_id="sid-bad")
|
||||
db.update_attempt(task_id="t5", attempt_no=1, end_time="2025-01-01T00:00:00ZZ", ray_status="FAILED")
|
||||
|
||||
_mk_job_dir(user_root, "alice", "sid-empty")
|
||||
_mk_job_dir(user_root, "alice", "sid-bad")
|
||||
jan = JobsJanitor(db=db, user_root=str(user_root), trash_after_days=3, purge_after_days=7, interval_s=1)
|
||||
jan.tick_once(now=cutoff_ended + timedelta(days=10))
|
||||
|
||||
assert (user_root / "alice" / "jobs" / "sid-empty").exists()
|
||||
assert (user_root / "alice" / "jobs" / "sid-bad").exists()
|
||||
|
||||
|
||||
def test_janitor_purge_moves_from_jobs_then_deletes(tmp_path: Path, monkeypatch) -> None:
|
||||
db_path = tmp_path / "mvp.sqlite3"
|
||||
user_root = tmp_path / "users"
|
||||
db = Db(str(db_path))
|
||||
db.init()
|
||||
|
||||
task_id = "t6"
|
||||
user_id = "alice"
|
||||
sid = "sid-a01"
|
||||
db.create_task_v25(task_id=task_id, user_id=user_id, workload="ppo", jobspec_yaml="{}", nnodes=1, n_gpus_per_node=1)
|
||||
db.create_attempt(task_id=task_id, attempt_no=1, ray_submission_id=sid)
|
||||
|
||||
now = datetime(2025, 1, 10, tzinfo=timezone.utc)
|
||||
ended = now - timedelta(days=9)
|
||||
db.update_attempt(task_id=task_id, attempt_no=1, end_time=_iso_z(ended), ray_status="SUCCEEDED")
|
||||
|
||||
src = _mk_job_dir(user_root, user_id, sid)
|
||||
jan = JobsJanitor(db=db, user_root=str(user_root), trash_after_days=3, purge_after_days=7, interval_s=1)
|
||||
|
||||
jan.tick_once(now=now)
|
||||
assert not src.exists()
|
||||
assert not (user_root / user_id / "trash" / "jobs" / sid).exists()
|
||||
|
||||
|
||||
def test_janitor_run_forever_requires_event_like(tmp_path: Path) -> None:
|
||||
db = Db(str(tmp_path / "mvp.sqlite3"))
|
||||
db.init()
|
||||
jan = JobsJanitor(db=db, user_root=str(tmp_path / "users"), trash_after_days=3, purge_after_days=7, interval_s=1)
|
||||
try:
|
||||
jan.run_forever(object())
|
||||
raise AssertionError("expected ValueError")
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
|
||||
def test_janitor_run_forever_survives_tick_exceptions(tmp_path: Path, monkeypatch) -> None:
|
||||
db = Db(str(tmp_path / "mvp.sqlite3"))
|
||||
db.init()
|
||||
jan = JobsJanitor(db=db, user_root=str(tmp_path / "users"), trash_after_days=3, purge_after_days=7, interval_s=1)
|
||||
|
||||
class Flag:
|
||||
def __init__(self) -> None:
|
||||
self.n = 0
|
||||
|
||||
def is_set(self) -> bool:
|
||||
self.n += 1
|
||||
return self.n >= 2
|
||||
|
||||
monkeypatch.setattr(jan, "tick_once", lambda **_: (_ for _ in ()).throw(RuntimeError("boom")))
|
||||
monkeypatch.setattr("argus.service.janitor.time.sleep", lambda *_: None)
|
||||
jan.run_forever(Flag())
|
||||
78
src/mvp/py/tests/test_llm_config_builder.py
Normal file
@@ -0,0 +1,78 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
def test_build_llm_config_dict_maps_tp_and_bundles():
|
||||
from argus.service.serve_llm_config import build_llm_config_dict
|
||||
from argus.service.serving_spec import ResolvedServingSpec
|
||||
|
||||
resolved = ResolvedServingSpec(
|
||||
user_id="alice",
|
||||
model_id_suffix="qwen-0.5b",
|
||||
model_id_prefix="alice-202601061235",
|
||||
model_id="alice-202601061235-qwen-0.5b",
|
||||
model_source="/private/hf/x",
|
||||
num_replicas=2,
|
||||
gpus_per_replica=4,
|
||||
engine_kwargs={"gpu_memory_utilization": 0.9},
|
||||
)
|
||||
|
||||
cfg = build_llm_config_dict(
|
||||
resolved,
|
||||
accelerator_type="H20",
|
||||
runtime_env_env_vars={"HF_ENDPOINT": "https://hf-mirror.com"},
|
||||
cpu_per_gpu=2.0,
|
||||
)
|
||||
assert cfg["model_loading_config"]["model_id"] == "alice-202601061235-qwen-0.5b"
|
||||
assert cfg["model_loading_config"]["model_source"] == "/private/hf/x"
|
||||
assert cfg["accelerator_type"] == "H20"
|
||||
assert cfg["deployment_config"]["num_replicas"] == 2
|
||||
|
||||
# gpus_per_replica -> tensor_parallel_size
|
||||
assert cfg["engine_kwargs"]["tensor_parallel_size"] == 4
|
||||
assert cfg["engine_kwargs"]["gpu_memory_utilization"] == 0.9
|
||||
|
||||
# resources_per_bundle reserves the full TP GPU set for each replica.
|
||||
bundle = cfg["resources_per_bundle"]
|
||||
assert bundle["GPU"] == 4.0
|
||||
assert bundle["CPU"] == 8.0
|
||||
|
||||
|
||||
def test_build_llm_config_dict_injects_hf_offline_defaults():
|
||||
from argus.service.serve_llm_config import build_llm_config_dict
|
||||
from argus.service.serving_spec import ResolvedServingSpec
|
||||
|
||||
resolved = ResolvedServingSpec(
|
||||
user_id="alice",
|
||||
model_id_suffix="x",
|
||||
model_id_prefix="alice-202601061235",
|
||||
model_id="alice-202601061235-x",
|
||||
model_source="/private/users/alice/models/x",
|
||||
num_replicas=1,
|
||||
gpus_per_replica=1,
|
||||
engine_kwargs=None,
|
||||
)
|
||||
cfg = build_llm_config_dict(resolved, accelerator_type="H20", runtime_env_env_vars={})
|
||||
env = cfg["runtime_env"]["env_vars"]
|
||||
assert env["HF_HUB_OFFLINE"] == "1"
|
||||
assert env["HF_HOME"] == "/private/hf"
|
||||
assert env["HUGGINGFACE_HUB_CACHE"].startswith("/private/hf/")
|
||||
|
||||
|
||||
def test_build_llm_config_dict_requires_accelerator_type():
|
||||
from argus.service.serve_llm_config import build_llm_config_dict
|
||||
from argus.service.serving_spec import ResolvedServingSpec
|
||||
|
||||
resolved = ResolvedServingSpec(
|
||||
user_id="alice",
|
||||
model_id_suffix="x",
|
||||
model_id_prefix="alice-202601061235",
|
||||
model_id="alice-202601061235-x",
|
||||
model_source="/private/hf/x",
|
||||
num_replicas=1,
|
||||
gpus_per_replica=1,
|
||||
engine_kwargs=None,
|
||||
)
|
||||
with pytest.raises(ValueError, match="accelerator_type is required"):
|
||||
build_llm_config_dict(resolved, accelerator_type="", runtime_env_env_vars={})
|
||||
156
src/mvp/py/tests/test_models.py
Normal file
@@ -0,0 +1,156 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
def test_require_missing_raises():
|
||||
from argus.ray.models import _require
|
||||
|
||||
with pytest.raises(ValueError, match="missing required field: x"):
|
||||
_require({}, "x")
|
||||
with pytest.raises(ValueError, match="missing required field: x"):
|
||||
_require({"x": ""}, "x")
|
||||
|
||||
|
||||
def test_ray_config_from_dict_new_format_and_defaults():
|
||||
from argus.ray.models import RayConfig
|
||||
|
||||
cfg = RayConfig.from_dict(
|
||||
{
|
||||
"ray": {
|
||||
"address": "http://127.0.0.1:8265",
|
||||
"shared_root": "/private",
|
||||
"entrypoint_resources": {"worker_node": 1},
|
||||
"runtime_env": {"env_vars": {"HF_ENDPOINT": "x"}},
|
||||
}
|
||||
}
|
||||
)
|
||||
assert cfg.address.endswith("8265")
|
||||
assert cfg.shared_root == "/private"
|
||||
assert cfg.entrypoint_num_cpus == 1.0
|
||||
assert cfg.entrypoint_resources["worker_node"] == 1.0
|
||||
assert cfg.runtime_env_env_vars["HF_ENDPOINT"] == "x"
|
||||
assert cfg.user_code_path == "/private/user/code"
|
||||
assert cfg.verl_code_path == "/private/common/code/verl/verl_repo"
|
||||
|
||||
public = cfg.to_public_dict()
|
||||
assert public["runtime_env"]["env_vars"]["HF_ENDPOINT"] == "x"
|
||||
|
||||
|
||||
def test_ray_config_from_dict_requires_mappings():
|
||||
from argus.ray.models import RayConfig
|
||||
|
||||
with pytest.raises(ValueError, match="runtime_env\\.env_vars must be a mapping"):
|
||||
RayConfig.from_dict(
|
||||
{
|
||||
"address": "x",
|
||||
"shared_root": "/p",
|
||||
"entrypoint_resources": {},
|
||||
"runtime_env": {"env_vars": ["nope"]},
|
||||
}
|
||||
)
|
||||
with pytest.raises(ValueError, match="entrypoint_resources must be a mapping"):
|
||||
RayConfig.from_dict(
|
||||
{
|
||||
"address": "x",
|
||||
"shared_root": "/p",
|
||||
"entrypoint_resources": ["nope"],
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def test_jobspec_validation_and_null_coercion():
|
||||
from argus.ray.models import JobSpec
|
||||
|
||||
spec = JobSpec.from_dict(
|
||||
{
|
||||
"workload": "ppo",
|
||||
"code_path": "/code",
|
||||
"model_id": "m",
|
||||
"train_file": "train.jsonl",
|
||||
"val_file": "val.jsonl",
|
||||
"test_freq": "",
|
||||
}
|
||||
)
|
||||
assert spec.workload == "ppo"
|
||||
assert spec.val_file == "val.jsonl"
|
||||
assert spec.test_freq is None
|
||||
assert spec.total_training_steps is None
|
||||
assert spec.nnodes == 2
|
||||
assert spec.n_gpus_per_node == 4
|
||||
|
||||
pub = spec.to_public_dict()
|
||||
assert pub["submission_id"] == ""
|
||||
assert "trainer_device" not in pub
|
||||
|
||||
|
||||
def test_jobspec_sft_adds_trainer_device_default():
|
||||
from argus.ray.models import JobSpec
|
||||
|
||||
spec = JobSpec.from_dict(
|
||||
{
|
||||
"workload": "sft",
|
||||
"code_path": "/code",
|
||||
"model_id": "m",
|
||||
"train_file": "train.jsonl",
|
||||
"val_file": "val.jsonl",
|
||||
}
|
||||
)
|
||||
pub = spec.to_public_dict()
|
||||
assert pub["trainer_device"] == "cpu"
|
||||
|
||||
|
||||
def test_jobspec_unsupported_workload():
|
||||
from argus.ray.models import JobSpec
|
||||
|
||||
with pytest.raises(ValueError, match="unsupported workload"):
|
||||
JobSpec.from_dict(
|
||||
{"workload": "nope", "code_path": "x", "model_id": "m", "train_file": "t", "val_file": "v"}
|
||||
)
|
||||
|
||||
|
||||
def test_parse_taskspec_basic_no_kind_compat():
|
||||
from argus.ray.models import JobSpec, parse_taskspec
|
||||
|
||||
got = parse_taskspec(
|
||||
{
|
||||
"workload": "ppo",
|
||||
"code_path": "/code",
|
||||
"model_id": "m",
|
||||
"train_file": "train.jsonl",
|
||||
"val_file": "val.jsonl",
|
||||
}
|
||||
)
|
||||
assert isinstance(got, JobSpec)
|
||||
assert got.workload == "ppo"
|
||||
|
||||
|
||||
def test_parse_taskspec_advanced_smoke():
|
||||
from argus.ray.models import AdvancedTaskSpec, parse_taskspec
|
||||
|
||||
got = parse_taskspec(
|
||||
{
|
||||
"kind": "advanced",
|
||||
"nnodes": 2,
|
||||
"n_gpus_per_node": 4,
|
||||
"command": "python3 -m verl.trainer.main_ppo +ray_kwargs.ray_init.address=auto",
|
||||
}
|
||||
)
|
||||
assert isinstance(got, AdvancedTaskSpec)
|
||||
assert got.kind == "advanced"
|
||||
assert got.workload == "advanced"
|
||||
assert got.nnodes == 2
|
||||
assert got.n_gpus_per_node == 4
|
||||
|
||||
|
||||
def test_parse_taskspec_advanced_requires_command_nnodes_gpus():
|
||||
from argus.ray.models import parse_taskspec
|
||||
|
||||
with pytest.raises(ValueError, match="missing required field: command"):
|
||||
parse_taskspec({"kind": "advanced", "nnodes": 1, "n_gpus_per_node": 1})
|
||||
|
||||
with pytest.raises(ValueError, match="missing required field: nnodes"):
|
||||
parse_taskspec({"kind": "advanced", "n_gpus_per_node": 1, "command": "python3 -m verl.trainer.main_ppo"})
|
||||
|
||||
with pytest.raises(ValueError, match="missing required field: n_gpus_per_node"):
|
||||
parse_taskspec({"kind": "advanced", "nnodes": 1, "command": "python3 -m verl.trainer.main_ppo"})
|
||||
218
src/mvp/py/tests/test_ray_job_tool.py
Normal file
@@ -0,0 +1,218 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from argus.ray.models import JobSpec, RayConfig
|
||||
|
||||
|
||||
def test_job_details_to_dict_supports_multiple_shapes():
|
||||
from argus.ray.ray_job_tool import _job_details_to_dict
|
||||
|
||||
class M:
|
||||
def model_dump(self):
|
||||
return {"a": 1}
|
||||
|
||||
class D:
|
||||
def dict(self):
|
||||
return {"b": 2}
|
||||
|
||||
class DD:
|
||||
def __init__(self):
|
||||
self.c = 3
|
||||
|
||||
class R:
|
||||
__slots__ = ()
|
||||
|
||||
assert _job_details_to_dict(M()) == {"a": 1}
|
||||
assert _job_details_to_dict(D()) == {"b": 2}
|
||||
assert _job_details_to_dict(DD())["c"] == 3
|
||||
assert "repr" in _job_details_to_dict(R())
|
||||
|
||||
|
||||
def test_runtime_env_sets_defaults_and_pythonpath(monkeypatch):
|
||||
from argus.ray.ray_job_tool import RayJobTool
|
||||
|
||||
cfg = RayConfig.from_dict(
|
||||
{
|
||||
"address": "http://127.0.0.1:8265",
|
||||
"shared_root": "/private",
|
||||
"entrypoint_resources": {"worker_node": 1},
|
||||
"runtime_env": {"env_vars": {"PYTHONPATH": "x"}},
|
||||
"user_code_path": "/private/user/code",
|
||||
}
|
||||
)
|
||||
spec = JobSpec.from_dict(
|
||||
{"workload": "sft", "code_path": "/c", "model_id": "m", "train_file": "t", "val_file": "v"}
|
||||
)
|
||||
monkeypatch.setenv("MVP_TOOL_CODE_PATH", "/tool")
|
||||
|
||||
tool = RayJobTool(cfg)
|
||||
env = tool._runtime_env(spec)["env_vars"]
|
||||
assert env["HF_HOME"].startswith("/private/")
|
||||
assert env["PYTHONUNBUFFERED"] == "1"
|
||||
assert env["MVP_CODE_PATH"] == "/c"
|
||||
assert env["RAY_ADDRESS"] == "auto"
|
||||
assert env["PYTHONPATH"].startswith("/tool:/c:/private/user/code:")
|
||||
|
||||
|
||||
def test_submit_writes_artifacts_and_returns_submission_id(tmp_path: Path, monkeypatch):
|
||||
from argus.ray import ray_job_tool as mod
|
||||
|
||||
class _FakeClient:
|
||||
def __init__(self, address: str):
|
||||
self.address = address
|
||||
self.last_submit_kwargs = None
|
||||
|
||||
def submit_job(self, **kwargs):
|
||||
self.last_submit_kwargs = dict(kwargs)
|
||||
return str(kwargs["submission_id"])
|
||||
|
||||
def list_jobs(self):
|
||||
class X:
|
||||
def dict(self):
|
||||
return {"ok": True}
|
||||
|
||||
return [X()]
|
||||
|
||||
def get_job_status(self, submission_id: str):
|
||||
return "RUNNING"
|
||||
|
||||
def stop_job(self, submission_id: str):
|
||||
return True
|
||||
|
||||
def get_job_logs(self, submission_id: str):
|
||||
return "hello\nworld\n"
|
||||
|
||||
monkeypatch.setattr(mod, "JobSubmissionClient", _FakeClient)
|
||||
monkeypatch.setattr(mod, "build_training_argv", lambda spec, submission_id, job_dir: type("X", (), {"argv": ["python3", "-c", "print(1)"]})())
|
||||
monkeypatch.setattr(mod.ray, "init", lambda **kwargs: (_ for _ in ()).throw(RuntimeError("no ray")))
|
||||
|
||||
cfg = RayConfig.from_dict(
|
||||
{
|
||||
"address": "http://127.0.0.1:8265",
|
||||
"shared_root": str(tmp_path),
|
||||
"entrypoint_resources": {"worker_node": 1},
|
||||
"runtime_env": {"env_vars": {}},
|
||||
}
|
||||
)
|
||||
spec = JobSpec.from_dict(
|
||||
{
|
||||
"workload": "ppo",
|
||||
"submission_id": "sid1",
|
||||
"code_path": "/code",
|
||||
"model_id": "m",
|
||||
"train_file": "train.jsonl",
|
||||
"val_file": "val.jsonl",
|
||||
}
|
||||
)
|
||||
|
||||
tool = mod.RayJobTool(cfg)
|
||||
submitted = tool.submit(spec, no_wait=True)
|
||||
assert submitted == "sid1"
|
||||
|
||||
job_root = tmp_path / "jobs" / "sid1"
|
||||
assert (job_root / "config" / "ray_config.yaml").exists()
|
||||
assert (job_root / "config" / "jobspec.yaml").exists()
|
||||
assert (job_root / "config" / "ray_submission_id.txt").read_text(encoding="utf-8").strip() == "sid1"
|
||||
|
||||
payload = json.loads((job_root / "config" / "submit_payload.json").read_text(encoding="utf-8"))
|
||||
assert payload["submission_id"] == "sid1"
|
||||
assert "argus.ray.driver_entrypoint" in payload["entrypoint"]
|
||||
|
||||
assert (job_root / "debug" / "ray_resources_pre.error.txt").exists()
|
||||
assert (job_root / "debug" / "ray_job_list_post.json").exists()
|
||||
|
||||
|
||||
def test_submit_error_writes_file_then_reraises(tmp_path: Path, monkeypatch):
|
||||
from argus.ray import ray_job_tool as mod
|
||||
|
||||
class _FakeClient:
|
||||
def __init__(self, address: str):
|
||||
self.address = address
|
||||
|
||||
def submit_job(self, **kwargs):
|
||||
raise RuntimeError("boom")
|
||||
|
||||
def list_jobs(self):
|
||||
return []
|
||||
|
||||
monkeypatch.setattr(mod, "JobSubmissionClient", _FakeClient)
|
||||
monkeypatch.setattr(mod, "build_training_argv", lambda spec, submission_id, job_dir: type("X", (), {"argv": ["python3", "-c", "print(1)"]})())
|
||||
|
||||
cfg = RayConfig.from_dict(
|
||||
{
|
||||
"address": "http://127.0.0.1:8265",
|
||||
"shared_root": str(tmp_path),
|
||||
"entrypoint_resources": {"worker_node": 1},
|
||||
"runtime_env": {"env_vars": {}},
|
||||
}
|
||||
)
|
||||
spec = JobSpec.from_dict(
|
||||
{"workload": "ppo", "submission_id": "sid2", "code_path": "/code", "model_id": "m", "train_file": "t", "val_file": "v"}
|
||||
)
|
||||
|
||||
tool = mod.RayJobTool(cfg)
|
||||
with pytest.raises(RuntimeError, match="boom"):
|
||||
tool.submit(spec, no_wait=True)
|
||||
|
||||
err = tmp_path / "jobs" / "sid2" / "logs" / "submit.error.txt"
|
||||
assert err.exists()
|
||||
|
||||
|
||||
def test_submit_advanced_writes_artifacts_and_returns_submission_id(tmp_path: Path, monkeypatch):
|
||||
from argus.ray import ray_job_tool as mod
|
||||
|
||||
class _FakeClient:
|
||||
def __init__(self, address: str):
|
||||
self.address = address
|
||||
self.last_submit_kwargs = None
|
||||
|
||||
def submit_job(self, **kwargs):
|
||||
self.last_submit_kwargs = dict(kwargs)
|
||||
return str(kwargs["submission_id"])
|
||||
|
||||
def list_jobs(self):
|
||||
class X:
|
||||
def dict(self):
|
||||
return {"ok": True}
|
||||
|
||||
return [X()]
|
||||
|
||||
monkeypatch.setattr(mod, "JobSubmissionClient", _FakeClient)
|
||||
monkeypatch.setattr(mod.ray, "init", lambda **kwargs: (_ for _ in ()).throw(RuntimeError("no ray")))
|
||||
|
||||
cfg = RayConfig.from_dict(
|
||||
{
|
||||
"address": "http://127.0.0.1:8265",
|
||||
"shared_root": str(tmp_path),
|
||||
"entrypoint_resources": {"worker_node": 1},
|
||||
"runtime_env": {"env_vars": {}},
|
||||
}
|
||||
)
|
||||
spec = mod.AdvancedTaskSpec.from_dict(
|
||||
{
|
||||
"kind": "advanced",
|
||||
"submission_id": "sid3",
|
||||
"nnodes": 2,
|
||||
"n_gpus_per_node": 4,
|
||||
"command": "echo hi",
|
||||
}
|
||||
)
|
||||
|
||||
tool = mod.RayJobTool(cfg)
|
||||
submitted = tool.submit_advanced(spec, no_wait=True, user_id="alice")
|
||||
assert submitted == "sid3"
|
||||
|
||||
job_root = tmp_path / "jobs" / "sid3"
|
||||
assert (job_root / "config" / "ray_config.yaml").exists()
|
||||
assert (job_root / "config" / "taskspec.yaml").exists()
|
||||
assert (job_root / "config" / "ray_submission_id.txt").read_text(encoding="utf-8").strip() == "sid3"
|
||||
|
||||
payload = json.loads((job_root / "config" / "submit_payload.json").read_text(encoding="utf-8"))
|
||||
assert payload["submission_id"] == "sid3"
|
||||
assert "argus.ray.driver_entrypoint" in payload["entrypoint"]
|
||||
assert (job_root / "debug" / "ray_resources_pre.error.txt").exists()
|
||||
25
src/mvp/py/tests/test_ray_resources.py
Normal file
@@ -0,0 +1,25 @@
|
||||
from __future__ import annotations
|
||||
|
||||
def test_get_cluster_available_sums_resources(monkeypatch):
|
||||
from argus.service import ray_resources
|
||||
|
||||
monkeypatch.setattr(ray_resources.ray._private.state, "available_resources_per_node", lambda: { # type: ignore[attr-defined]
|
||||
"n1": {"GPU": 1, "NPU": 2},
|
||||
"n2": {"GPU": 0.5},
|
||||
"bad": "nope",
|
||||
})
|
||||
|
||||
avail = ray_resources.get_cluster_available()
|
||||
assert avail.total_available_gpus == 1.5
|
||||
assert avail.total_available_npus == 2.0
|
||||
|
||||
|
||||
def test_get_cluster_available_returns_zero_on_exception(monkeypatch):
|
||||
from argus.service import ray_resources
|
||||
|
||||
def _boom() -> dict[str, object]:
|
||||
raise RuntimeError("boom")
|
||||
|
||||
monkeypatch.setattr(ray_resources.ray._private.state, "available_resources_per_node", _boom) # type: ignore[attr-defined]
|
||||
avail = ray_resources.get_cluster_available()
|
||||
assert avail.total_available_gpus == 0.0
|
||||
567
src/mvp/py/tests/test_scheduler.py
Normal file
@@ -0,0 +1,567 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import yaml
|
||||
from types import SimpleNamespace
|
||||
|
||||
from argus.ray.models import RayConfig
|
||||
from argus.service.config import V2Config
|
||||
from argus.service.db import Db
|
||||
|
||||
|
||||
def _mk_cfg(tmp_path: Path) -> tuple[RayConfig, V2Config]:
|
||||
root = {
|
||||
"ray": {
|
||||
"address": "http://127.0.0.1:8265",
|
||||
"shared_root": "/private",
|
||||
"entrypoint_resources": {"worker_node": 1},
|
||||
"runtime_env": {"env_vars": {}},
|
||||
},
|
||||
"service": {
|
||||
"sqlite": {"db_path": str(tmp_path / "mvp.sqlite3")},
|
||||
"scheduler": {"tick_s": 1, "retry_interval_s": 1, "max_running_tasks": 1},
|
||||
},
|
||||
}
|
||||
return RayConfig.from_dict(root), V2Config.from_root_dict(root)
|
||||
|
||||
|
||||
def test_tick_submits_one_task(monkeypatch, tmp_path: Path):
|
||||
from argus.service import scheduler as sched_mod
|
||||
|
||||
ray_cfg, v2_cfg = _mk_cfg(tmp_path)
|
||||
db = Db(v2_cfg.sqlite.db_path)
|
||||
db.init()
|
||||
db.create_task_v25(
|
||||
task_id="t1",
|
||||
user_id="alice",
|
||||
workload="ppo",
|
||||
jobspec_yaml=yaml.safe_dump(
|
||||
{
|
||||
"workload": "ppo",
|
||||
"code_path": "/private/common/code/verl",
|
||||
"model_id": "m",
|
||||
"train_file": "/private/common/datasets/t",
|
||||
"val_file": "/private/common/datasets/v",
|
||||
}
|
||||
),
|
||||
nnodes=2,
|
||||
n_gpus_per_node=4,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(sched_mod, "ensure_ray_connected", lambda: None)
|
||||
monkeypatch.setattr(
|
||||
sched_mod,
|
||||
"get_cluster_available",
|
||||
lambda: SimpleNamespace(total_available_gpus=999.0, total_available_npus=0.0),
|
||||
)
|
||||
|
||||
class _Tool:
|
||||
def __init__(self, cfg):
|
||||
self.submitted = []
|
||||
self.job_dirs = []
|
||||
self.extra_env = []
|
||||
self.extra_overrides = []
|
||||
|
||||
def submit(
|
||||
self,
|
||||
spec,
|
||||
no_wait: bool,
|
||||
job_dir: str | None = None,
|
||||
extra_env_vars: dict | None = None,
|
||||
extra_overrides: list[str] | None = None,
|
||||
):
|
||||
self.submitted.append(spec.submission_id)
|
||||
self.job_dirs.append(job_dir)
|
||||
self.extra_env.append(extra_env_vars)
|
||||
self.extra_overrides.append(extra_overrides)
|
||||
return str(spec.submission_id)
|
||||
|
||||
def status(self, submission_id: str):
|
||||
return "RUNNING"
|
||||
|
||||
def logs(self, submission_id: str):
|
||||
return ""
|
||||
|
||||
monkeypatch.setattr(sched_mod, "RayJobTool", _Tool)
|
||||
|
||||
s = sched_mod.Scheduler(db=db, ray_cfg=ray_cfg, v2_cfg=v2_cfg)
|
||||
s.tick()
|
||||
|
||||
row = db.get_task("t1")
|
||||
assert row and row["state"] == "SUBMITTED"
|
||||
attempts = db.list_attempts("t1")
|
||||
assert len(attempts) == 1
|
||||
assert attempts[0]["ray_submission_id"] == "t1--a01"
|
||||
assert s.tool.job_dirs[-1] == "/private/users/alice/jobs/t1--a01"
|
||||
env = s.tool.extra_env[-1]
|
||||
assert env and env["MVP_WANDB_PROJECT"] == "alice_project"
|
||||
assert env["MVP_WANDB_RUN"] == "t1--a01"
|
||||
assert env["MVP_TRAINER_LOGGER"] == "console"
|
||||
assert s.tool.extra_overrides[-1] is None
|
||||
|
||||
|
||||
def test_tick_submits_one_task_advanced(monkeypatch, tmp_path: Path):
|
||||
from argus.service import scheduler as sched_mod
|
||||
|
||||
ray_cfg, v2_cfg = _mk_cfg(tmp_path)
|
||||
db = Db(v2_cfg.sqlite.db_path)
|
||||
db.init()
|
||||
db.create_task_v25(
|
||||
task_id="t1",
|
||||
user_id="alice",
|
||||
workload="advanced",
|
||||
jobspec_yaml=yaml.safe_dump(
|
||||
{
|
||||
"kind": "advanced",
|
||||
"nnodes": 2,
|
||||
"n_gpus_per_node": 4,
|
||||
"command": "python3 -m verl.trainer.main_ppo +ray_kwargs.ray_init.address=auto",
|
||||
}
|
||||
),
|
||||
nnodes=2,
|
||||
n_gpus_per_node=4,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(sched_mod, "ensure_ray_connected", lambda: None)
|
||||
monkeypatch.setattr(
|
||||
sched_mod,
|
||||
"get_cluster_available",
|
||||
lambda: SimpleNamespace(total_available_gpus=999.0, total_available_npus=0.0),
|
||||
)
|
||||
|
||||
class _Tool:
|
||||
def __init__(self, cfg):
|
||||
self.submitted = []
|
||||
self.job_dirs = []
|
||||
self.extra_env = []
|
||||
|
||||
def submit(self, spec, no_wait: bool, job_dir: str | None = None):
|
||||
raise AssertionError("should not call submit() for advanced")
|
||||
|
||||
def submit_advanced(
|
||||
self,
|
||||
spec,
|
||||
no_wait: bool,
|
||||
job_dir: str | None = None,
|
||||
user_id: str | None = None,
|
||||
extra_env_vars: dict | None = None,
|
||||
extra_overrides: list[str] | None = None,
|
||||
):
|
||||
self.submitted.append(spec.submission_id)
|
||||
self.job_dirs.append(job_dir)
|
||||
self.extra_env.append(extra_env_vars)
|
||||
return str(spec.submission_id)
|
||||
|
||||
def status(self, submission_id: str):
|
||||
return "RUNNING"
|
||||
|
||||
def logs(self, submission_id: str):
|
||||
return ""
|
||||
|
||||
monkeypatch.setattr(sched_mod, "RayJobTool", _Tool)
|
||||
|
||||
s = sched_mod.Scheduler(db=db, ray_cfg=ray_cfg, v2_cfg=v2_cfg)
|
||||
s.tick()
|
||||
|
||||
row = db.get_task("t1")
|
||||
assert row and row["state"] == "SUBMITTED"
|
||||
attempts = db.list_attempts("t1")
|
||||
assert len(attempts) == 1
|
||||
assert attempts[0]["ray_submission_id"] == "t1--a01"
|
||||
assert s.tool.job_dirs[-1] == "/private/users/alice/jobs/t1--a01"
|
||||
env = s.tool.extra_env[-1]
|
||||
assert env and env["MVP_WANDB_PROJECT"] == "alice_project"
|
||||
assert env["MVP_WANDB_RUN"] == "t1--a01"
|
||||
assert env["MVP_TRAINER_LOGGER"] == "console"
|
||||
|
||||
|
||||
def test_tick_injects_wandb_env_when_enabled(monkeypatch, tmp_path: Path):
|
||||
from argus.service import scheduler as sched_mod
|
||||
|
||||
root = {
|
||||
"ray": {
|
||||
"address": "http://127.0.0.1:8265",
|
||||
"shared_root": "/private",
|
||||
"entrypoint_resources": {"worker_node": 1},
|
||||
"runtime_env": {"env_vars": {}},
|
||||
},
|
||||
"service": {
|
||||
"sqlite": {"db_path": str(tmp_path / "mvp.sqlite3")},
|
||||
"scheduler": {"tick_s": 1, "retry_interval_s": 1, "max_running_tasks": 1},
|
||||
},
|
||||
"tracking": {"wandb": {"enabled": True, "base_url": "http://172.22.0.1:8090"}},
|
||||
}
|
||||
ray_cfg, v2_cfg = RayConfig.from_dict(root), V2Config.from_root_dict(root)
|
||||
db = Db(v2_cfg.sqlite.db_path)
|
||||
db.init()
|
||||
db.create_task_v25(
|
||||
task_id="t1",
|
||||
user_id="alice",
|
||||
workload="ppo",
|
||||
jobspec_yaml=yaml.safe_dump(
|
||||
{
|
||||
"workload": "ppo",
|
||||
"code_path": "/private/common/code/verl",
|
||||
"model_id": "m",
|
||||
"train_file": "/private/common/datasets/t",
|
||||
"val_file": "/private/common/datasets/v",
|
||||
}
|
||||
),
|
||||
nnodes=2,
|
||||
n_gpus_per_node=4,
|
||||
)
|
||||
|
||||
monkeypatch.setenv("WANDB_API_KEY", "k")
|
||||
monkeypatch.setattr(sched_mod, "ensure_ray_connected", lambda: None)
|
||||
monkeypatch.setattr(
|
||||
sched_mod,
|
||||
"get_cluster_available",
|
||||
lambda: SimpleNamespace(total_available_gpus=999.0, total_available_npus=0.0),
|
||||
)
|
||||
|
||||
class _Tool:
|
||||
def __init__(self, cfg):
|
||||
self.extra_env = []
|
||||
self.job_dirs = []
|
||||
|
||||
def submit(
|
||||
self,
|
||||
spec,
|
||||
no_wait: bool,
|
||||
job_dir: str | None = None,
|
||||
extra_env_vars: dict | None = None,
|
||||
extra_overrides: list[str] | None = None,
|
||||
):
|
||||
self.job_dirs.append(job_dir)
|
||||
self.extra_env.append(extra_env_vars)
|
||||
return str(spec.submission_id)
|
||||
|
||||
def status(self, submission_id: str):
|
||||
return "RUNNING"
|
||||
|
||||
def logs(self, submission_id: str):
|
||||
return ""
|
||||
|
||||
monkeypatch.setattr(sched_mod, "RayJobTool", _Tool)
|
||||
|
||||
s = sched_mod.Scheduler(db=db, ray_cfg=ray_cfg, v2_cfg=v2_cfg)
|
||||
s.tick()
|
||||
|
||||
assert s.tool.job_dirs[-1] == "/private/users/alice/jobs/t1--a01"
|
||||
env = s.tool.extra_env[-1]
|
||||
assert env and env["WANDB_API_KEY"] == "k"
|
||||
assert env["WANDB_BASE_URL"] == "http://172.22.0.1:8090"
|
||||
assert env["WANDB_MODE"] == "online"
|
||||
assert env["WANDB_DIR"] == "/private/users/alice/jobs/t1--a01/wandb"
|
||||
assert env["MVP_WANDB_PROJECT"] == "alice_project"
|
||||
assert env["MVP_WANDB_RUN"] == "t1--a01"
|
||||
assert env["MVP_TRAINER_LOGGER"] == "[console,wandb]"
|
||||
|
||||
|
||||
def test_tick_wandb_enabled_missing_key_degrades(monkeypatch, tmp_path: Path):
|
||||
from argus.service import scheduler as sched_mod
|
||||
|
||||
root = {
|
||||
"ray": {
|
||||
"address": "http://127.0.0.1:8265",
|
||||
"shared_root": "/private",
|
||||
"entrypoint_resources": {"worker_node": 1},
|
||||
"runtime_env": {"env_vars": {}},
|
||||
},
|
||||
"service": {
|
||||
"sqlite": {"db_path": str(tmp_path / "mvp.sqlite3")},
|
||||
"scheduler": {"tick_s": 1, "retry_interval_s": 1, "max_running_tasks": 1},
|
||||
},
|
||||
"tracking": {"wandb": {"enabled": True, "base_url": "http://172.22.0.1:8090"}},
|
||||
}
|
||||
ray_cfg, v2_cfg = RayConfig.from_dict(root), V2Config.from_root_dict(root)
|
||||
db = Db(v2_cfg.sqlite.db_path)
|
||||
db.init()
|
||||
db.create_task_v25(
|
||||
task_id="t1",
|
||||
user_id="alice",
|
||||
workload="ppo",
|
||||
jobspec_yaml=yaml.safe_dump(
|
||||
{
|
||||
"workload": "ppo",
|
||||
"code_path": "/private/common/code/verl",
|
||||
"model_id": "m",
|
||||
"train_file": "/private/common/datasets/t",
|
||||
"val_file": "/private/common/datasets/v",
|
||||
}
|
||||
),
|
||||
nnodes=2,
|
||||
n_gpus_per_node=4,
|
||||
)
|
||||
|
||||
monkeypatch.delenv("WANDB_API_KEY", raising=False)
|
||||
monkeypatch.setattr(sched_mod, "ensure_ray_connected", lambda: None)
|
||||
monkeypatch.setattr(
|
||||
sched_mod,
|
||||
"get_cluster_available",
|
||||
lambda: SimpleNamespace(total_available_gpus=999.0, total_available_npus=0.0),
|
||||
)
|
||||
|
||||
class _Tool:
|
||||
def __init__(self, cfg):
|
||||
self.extra_env = []
|
||||
|
||||
def submit(
|
||||
self,
|
||||
spec,
|
||||
no_wait: bool,
|
||||
job_dir: str | None = None,
|
||||
extra_env_vars: dict | None = None,
|
||||
extra_overrides: list[str] | None = None,
|
||||
):
|
||||
self.extra_env.append(extra_env_vars)
|
||||
return str(spec.submission_id)
|
||||
|
||||
def status(self, submission_id: str):
|
||||
return "RUNNING"
|
||||
|
||||
def logs(self, submission_id: str):
|
||||
return ""
|
||||
|
||||
monkeypatch.setattr(sched_mod, "RayJobTool", _Tool)
|
||||
|
||||
s = sched_mod.Scheduler(db=db, ray_cfg=ray_cfg, v2_cfg=v2_cfg)
|
||||
s.tick()
|
||||
|
||||
env = s.tool.extra_env[-1]
|
||||
assert env and env["MVP_WANDB_PROJECT"] == "alice_project"
|
||||
assert env["MVP_WANDB_RUN"] == "t1--a01"
|
||||
assert env["MVP_TRAINER_LOGGER"] == "console"
|
||||
attempts = db.list_attempts("t1")
|
||||
assert attempts and "WANDB_API_KEY" in str(attempts[0].get("message") or "")
|
||||
|
||||
|
||||
def test_tick_wandb_enabled_updates_jobspec_logger(monkeypatch, tmp_path: Path):
|
||||
from argus.service import scheduler as sched_mod
|
||||
|
||||
root = {
|
||||
"ray": {
|
||||
"address": "http://127.0.0.1:8265",
|
||||
"shared_root": "/private",
|
||||
"entrypoint_resources": {"worker_node": 1},
|
||||
"runtime_env": {"env_vars": {}},
|
||||
},
|
||||
"service": {
|
||||
"sqlite": {"db_path": str(tmp_path / "mvp.sqlite3")},
|
||||
"scheduler": {"tick_s": 1, "retry_interval_s": 1, "max_running_tasks": 1},
|
||||
},
|
||||
"tracking": {"wandb": {"enabled": True, "base_url": "http://172.22.0.1:8090"}},
|
||||
}
|
||||
ray_cfg, v2_cfg = RayConfig.from_dict(root), V2Config.from_root_dict(root)
|
||||
db = Db(v2_cfg.sqlite.db_path)
|
||||
db.init()
|
||||
db.create_task_v25(
|
||||
task_id="t1",
|
||||
user_id="alice",
|
||||
workload="ppo",
|
||||
jobspec_yaml=yaml.safe_dump(
|
||||
{
|
||||
"workload": "ppo",
|
||||
"code_path": "/private/common/code/verl",
|
||||
"model_id": "m",
|
||||
"train_file": "/private/common/datasets/t",
|
||||
"val_file": "/private/common/datasets/v",
|
||||
}
|
||||
),
|
||||
nnodes=2,
|
||||
n_gpus_per_node=4,
|
||||
)
|
||||
|
||||
monkeypatch.setenv("WANDB_API_KEY", "k")
|
||||
monkeypatch.setattr(sched_mod, "ensure_ray_connected", lambda: None)
|
||||
monkeypatch.setattr(
|
||||
sched_mod,
|
||||
"get_cluster_available",
|
||||
lambda: SimpleNamespace(total_available_gpus=999.0, total_available_npus=0.0),
|
||||
)
|
||||
|
||||
class _Tool:
|
||||
def __init__(self, cfg):
|
||||
self.specs = []
|
||||
self.extra_overrides = []
|
||||
|
||||
def submit(
|
||||
self,
|
||||
spec,
|
||||
no_wait: bool,
|
||||
job_dir: str | None = None,
|
||||
extra_env_vars: dict | None = None,
|
||||
extra_overrides: list[str] | None = None,
|
||||
):
|
||||
self.specs.append(spec)
|
||||
self.extra_overrides.append(extra_overrides)
|
||||
return str(spec.submission_id)
|
||||
|
||||
def status(self, submission_id: str):
|
||||
return "RUNNING"
|
||||
|
||||
def logs(self, submission_id: str):
|
||||
return ""
|
||||
|
||||
monkeypatch.setattr(sched_mod, "RayJobTool", _Tool)
|
||||
|
||||
s = sched_mod.Scheduler(db=db, ray_cfg=ray_cfg, v2_cfg=v2_cfg)
|
||||
s.tick()
|
||||
|
||||
spec = s.tool.specs[-1]
|
||||
# v3.6: built-in training should be configured to use wandb when enabled and API key exists.
|
||||
assert spec.workload == "ppo"
|
||||
ov = s.tool.extra_overrides[-1]
|
||||
assert ov and "trainer.logger=[console,wandb]" in ov
|
||||
assert ov and "trainer.project_name=alice_project" in ov
|
||||
assert ov and "trainer.experiment_name=t1--a01" in ov
|
||||
|
||||
|
||||
def test_tick_marks_pending_resources(monkeypatch, tmp_path: Path):
|
||||
from argus.service import scheduler as sched_mod
|
||||
|
||||
ray_cfg, v2_cfg = _mk_cfg(tmp_path)
|
||||
db = Db(v2_cfg.sqlite.db_path)
|
||||
db.init()
|
||||
db.create_task_v25(
|
||||
task_id="t1",
|
||||
user_id="alice",
|
||||
workload="ppo",
|
||||
jobspec_yaml=yaml.safe_dump(
|
||||
{
|
||||
"workload": "ppo",
|
||||
"code_path": "/private/common/code/verl",
|
||||
"model_id": "m",
|
||||
"train_file": "/private/common/datasets/t",
|
||||
"val_file": "/private/common/datasets/v",
|
||||
}
|
||||
),
|
||||
nnodes=2,
|
||||
n_gpus_per_node=4,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(sched_mod, "ensure_ray_connected", lambda: None)
|
||||
monkeypatch.setattr(
|
||||
sched_mod,
|
||||
"get_cluster_available",
|
||||
lambda: SimpleNamespace(total_available_gpus=0.0, total_available_npus=0.0),
|
||||
)
|
||||
monkeypatch.setattr(sched_mod, "RayJobTool", lambda cfg: None)
|
||||
|
||||
s = sched_mod.Scheduler(db=db, ray_cfg=ray_cfg, v2_cfg=v2_cfg)
|
||||
s.tick()
|
||||
row = db.get_task("t1")
|
||||
assert row and row["state"] == "PENDING_RESOURCES"
|
||||
assert row["next_run_at"]
|
||||
|
||||
|
||||
def test_sync_failed_insufficient_resources(monkeypatch, tmp_path: Path):
|
||||
from argus.service import scheduler as sched_mod
|
||||
|
||||
ray_cfg, v2_cfg = _mk_cfg(tmp_path)
|
||||
db = Db(v2_cfg.sqlite.db_path)
|
||||
db.init()
|
||||
db.create_task_v25(
|
||||
task_id="t1",
|
||||
user_id="alice",
|
||||
workload="ppo",
|
||||
jobspec_yaml=yaml.safe_dump(
|
||||
{
|
||||
"workload": "ppo",
|
||||
"code_path": "/private/common/code/verl",
|
||||
"model_id": "m",
|
||||
"train_file": "/private/common/datasets/t",
|
||||
"val_file": "/private/common/datasets/v",
|
||||
}
|
||||
),
|
||||
nnodes=2,
|
||||
n_gpus_per_node=4,
|
||||
)
|
||||
db.create_attempt(task_id="t1", attempt_no=1, ray_submission_id="t1--a01")
|
||||
db.set_task_state(task_id="t1", state="RUNNING", latest_attempt_no=1)
|
||||
|
||||
monkeypatch.setattr(sched_mod, "ensure_ray_connected", lambda: None)
|
||||
monkeypatch.setattr(sched_mod, "RayJobTool", lambda cfg: None)
|
||||
|
||||
s = sched_mod.Scheduler(db=db, ray_cfg=ray_cfg, v2_cfg=v2_cfg)
|
||||
|
||||
class _Tool:
|
||||
def status(self, sid: str):
|
||||
return "FAILED"
|
||||
|
||||
def logs(self, sid: str):
|
||||
# Match the service's regex exactly:
|
||||
# it expects literal backslashes and repeats of 's'/'d' (because of double-escaping).
|
||||
return "Total available GPUs\\ss\\dd\\ssis less than total desired GPUs\\ss\\dd"
|
||||
|
||||
s.tool = _Tool()
|
||||
s.tick()
|
||||
|
||||
row = db.get_task("t1")
|
||||
assert row and row["state"] == "PENDING_RESOURCES"
|
||||
attempts = db.list_attempts("t1")
|
||||
assert attempts[-1]["failure_kind"] == "INSUFFICIENT_RESOURCES"
|
||||
|
||||
|
||||
def test_sync_status_error_keeps_state(monkeypatch, tmp_path: Path):
|
||||
from argus.service import scheduler as sched_mod
|
||||
|
||||
ray_cfg, v2_cfg = _mk_cfg(tmp_path)
|
||||
db = Db(v2_cfg.sqlite.db_path)
|
||||
db.init()
|
||||
db.create_task_v25(
|
||||
task_id="t1",
|
||||
user_id="alice",
|
||||
workload="ppo",
|
||||
jobspec_yaml=yaml.safe_dump({"workload": "ppo", "code_path": "/private/common/code/verl", "model_id": "m", "train_file": "/private/common/datasets/t"}),
|
||||
nnodes=2,
|
||||
n_gpus_per_node=4,
|
||||
)
|
||||
db.create_attempt(task_id="t1", attempt_no=1, ray_submission_id="t1--a01")
|
||||
db.set_task_state(task_id="t1", state="RUNNING", latest_attempt_no=1)
|
||||
|
||||
monkeypatch.setattr(sched_mod, "ensure_ray_connected", lambda: None)
|
||||
monkeypatch.setattr(sched_mod, "RayJobTool", lambda cfg: None)
|
||||
|
||||
s = sched_mod.Scheduler(db=db, ray_cfg=ray_cfg, v2_cfg=v2_cfg)
|
||||
|
||||
class _Tool:
|
||||
def status(self, sid: str):
|
||||
raise RuntimeError("boom")
|
||||
|
||||
s.tool = _Tool()
|
||||
s.tick()
|
||||
row = db.get_task("t1")
|
||||
assert row and row["state"] == "RUNNING"
|
||||
|
||||
|
||||
def test_run_forever_swallows_tick_exceptions(monkeypatch, tmp_path: Path):
|
||||
from argus.service import scheduler as sched_mod
|
||||
|
||||
ray_cfg, v2_cfg = _mk_cfg(tmp_path)
|
||||
db = Db(v2_cfg.sqlite.db_path)
|
||||
db.init()
|
||||
|
||||
monkeypatch.setattr(sched_mod, "RayJobTool", lambda cfg: None)
|
||||
s = sched_mod.Scheduler(db=db, ray_cfg=ray_cfg, v2_cfg=v2_cfg)
|
||||
|
||||
calls = {"n": 0}
|
||||
|
||||
def _tick():
|
||||
calls["n"] += 1
|
||||
raise RuntimeError("boom")
|
||||
|
||||
monkeypatch.setattr(s, "tick", _tick)
|
||||
monkeypatch.setattr(sched_mod.time, "sleep", lambda _: None)
|
||||
|
||||
class _Stop:
|
||||
def __init__(self):
|
||||
self._n = 0
|
||||
|
||||
def is_set(self):
|
||||
self._n += 1
|
||||
return self._n > 1
|
||||
|
||||
s.run_forever(_Stop())
|
||||
assert calls["n"] == 1
|
||||
80
src/mvp/py/tests/test_serve_client.py
Normal file
@@ -0,0 +1,80 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
import types
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
def test_ray_serve_client_calls_start_run_status(monkeypatch):
|
||||
import ray # provided by conftest stub
|
||||
|
||||
calls: list[tuple[str, object]] = []
|
||||
|
||||
def _init(*args, **kwargs):
|
||||
calls.append(("ray.init", {"args": args, "kwargs": kwargs}))
|
||||
|
||||
monkeypatch.setattr(ray, "init", _init, raising=False)
|
||||
|
||||
serve = types.ModuleType("ray.serve")
|
||||
|
||||
def _start(**kwargs):
|
||||
calls.append(("serve.start", kwargs))
|
||||
return None
|
||||
|
||||
def _run(app, name=None, route_prefix=None):
|
||||
calls.append(("serve.run", {"app": app, "name": name, "route_prefix": route_prefix}))
|
||||
return {"deployed": True}
|
||||
|
||||
def _status():
|
||||
calls.append(("serve.status", None))
|
||||
return {"ok": True}
|
||||
|
||||
def _delete(name: str):
|
||||
calls.append(("serve.delete", {"name": name}))
|
||||
return {"deleted": True}
|
||||
|
||||
serve.start = _start # type: ignore[attr-defined]
|
||||
serve.run = _run # type: ignore[attr-defined]
|
||||
serve.status = _status # type: ignore[attr-defined]
|
||||
serve.delete = _delete # type: ignore[attr-defined]
|
||||
|
||||
sys.modules["ray.serve"] = serve
|
||||
ray.serve = serve # type: ignore[attr-defined]
|
||||
|
||||
from argus.service.serve_client import RayServeClient
|
||||
|
||||
client = RayServeClient(http_port=8000, proxy_location="HeadOnly", ray_init_address="auto")
|
||||
client.ensure_started()
|
||||
out = client.apply_app(app="APP", app_name="argus_llm_app", route_prefix="/")
|
||||
st = client.get_status()
|
||||
deleted = client.delete_app(app_name="mk1")
|
||||
|
||||
assert out == {"deployed": True}
|
||||
assert st == {"ok": True}
|
||||
assert deleted == {"deleted": True}
|
||||
|
||||
# Verify call order and key args.
|
||||
assert calls[0][0] == "ray.init"
|
||||
assert calls[0][1]["kwargs"].get("ignore_reinit_error") is True
|
||||
assert calls[1][0] == "serve.start"
|
||||
assert calls[1][1]["http_options"]["port"] == 8000
|
||||
assert calls[2][0] == "serve.run"
|
||||
assert calls[2][1]["name"] == "argus_llm_app"
|
||||
assert calls[3][0] == "serve.status"
|
||||
assert calls[4][0] == "serve.delete"
|
||||
|
||||
|
||||
def test_ray_serve_client_rejects_dict_app(monkeypatch):
|
||||
import ray # provided by conftest stub
|
||||
|
||||
def _init(*args, **kwargs):
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(ray, "init", _init, raising=False)
|
||||
|
||||
from argus.service.serve_client import RayServeClient
|
||||
|
||||
client = RayServeClient(http_port=8000, proxy_location="HeadOnly", ray_init_address="auto")
|
||||
with pytest.raises(ValueError, match="invalid serve app object"):
|
||||
client.apply_app(app={"llm_configs": []}, app_name="mk1", route_prefix="/serve/mk1")
|
||||
38
src/mvp/py/tests/test_server.py
Normal file
@@ -0,0 +1,38 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
def test_server_main_calls_uvicorn(monkeypatch, tmp_path: Path):
|
||||
import server as server_mod
|
||||
|
||||
cfg = tmp_path / "cfg.yaml"
|
||||
cfg.write_text(
|
||||
"ray:\n address: http://127.0.0.1:8265\n shared_root: /private\n entrypoint_resources: {worker_node: 1}\n"
|
||||
"service:\n api: {host: 127.0.0.1, port: 18080}\n sqlite: {db_path: /tmp/x.sqlite3}\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
got = {}
|
||||
|
||||
monkeypatch.setattr(server_mod, "create_app", lambda path: object())
|
||||
monkeypatch.setattr(server_mod.uvicorn, "run", lambda app, host, port, log_level: got.update({"host": host, "port": port}))
|
||||
monkeypatch.setattr(sys, "argv", ["server.py", "--config", str(cfg)])
|
||||
|
||||
assert server_mod.main() == 0
|
||||
assert got["host"] == "127.0.0.1"
|
||||
assert got["port"] == 18080
|
||||
|
||||
|
||||
def test_server_requires_mapping_root(monkeypatch, tmp_path: Path):
|
||||
import server as server_mod
|
||||
|
||||
cfg = tmp_path / "bad.yaml"
|
||||
cfg.write_text("- 1\n- 2\n", encoding="utf-8")
|
||||
monkeypatch.setattr(sys, "argv", ["server.py", "--config", str(cfg)])
|
||||
with pytest.raises(SystemExit, match="config yaml must be a mapping"):
|
||||
server_mod.main()
|
||||
|
||||
125
src/mvp/py/tests/test_service_config.py
Normal file
@@ -0,0 +1,125 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
def test_v2_config_from_root_dict_new_format_defaults():
|
||||
from argus.service.config import V2Config
|
||||
|
||||
cfg = V2Config.from_root_dict(
|
||||
{
|
||||
"ray": {"shared_root": "/private"},
|
||||
"service": {
|
||||
"api": {"host": "127.0.0.1", "port": 9999},
|
||||
"auth": {"token_env": "X"},
|
||||
"sqlite": {"db_path": "/tmp/x.sqlite3"},
|
||||
"scheduler": {"tick_s": 1, "retry_interval_s": 2, "max_running_tasks": 3},
|
||||
},
|
||||
}
|
||||
)
|
||||
assert cfg.api.host == "127.0.0.1"
|
||||
assert cfg.api.port == 9999
|
||||
assert cfg.auth.token_env == "X"
|
||||
assert cfg.sqlite.db_path.endswith(".sqlite3")
|
||||
assert cfg.scheduler.max_running_tasks == 3
|
||||
assert cfg.tracking.wandb.enabled is False
|
||||
assert cfg.serving.enabled is False
|
||||
|
||||
|
||||
def test_v2_config_backward_compat_v2_section_and_default_db_path():
|
||||
from argus.service.config import V2Config
|
||||
|
||||
cfg = V2Config.from_root_dict({"shared_root": "/private", "v2": {"sqlite": {}}})
|
||||
assert cfg.sqlite.db_path == "/private/common/db/mvp.sqlite3"
|
||||
assert cfg.tracking.wandb.enabled is False
|
||||
|
||||
|
||||
def test_v2_config_requires_mappings():
|
||||
from argus.service.config import V2Config
|
||||
|
||||
with pytest.raises(ValueError, match="config\\.service must be a mapping"):
|
||||
V2Config.from_root_dict({"service": ["nope"]})
|
||||
with pytest.raises(ValueError, match="config\\.service\\.\\{api,auth,sqlite,scheduler\\} must be mappings"):
|
||||
V2Config.from_root_dict({"service": {"api": [1], "auth": {}, "sqlite": {}, "scheduler": {}}})
|
||||
|
||||
|
||||
def test_v2_config_requires_data_mappings():
|
||||
from argus.service.config import V2Config
|
||||
|
||||
base = {
|
||||
"ray": {"shared_root": "/private"},
|
||||
"service": {"api": {}, "auth": {}, "sqlite": {}, "scheduler": {}},
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError, match="config\\.data must be a mapping"):
|
||||
V2Config.from_root_dict({**base, "data": ["nope"]})
|
||||
|
||||
with pytest.raises(ValueError, match="config\\.data\\.\\{sftpgo,retention\\} must be mappings"):
|
||||
V2Config.from_root_dict({**base, "data": {"sftpgo": ["x"], "retention": {}}})
|
||||
|
||||
|
||||
def test_v2_config_requires_tracking_and_serving_mappings():
|
||||
from argus.service.config import V2Config
|
||||
|
||||
base = {
|
||||
"ray": {"shared_root": "/private"},
|
||||
"service": {"api": {}, "auth": {}, "sqlite": {}, "scheduler": {}},
|
||||
"data": {"sftpgo": {}, "retention": {}},
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError, match="config\\.tracking must be a mapping"):
|
||||
V2Config.from_root_dict({**base, "tracking": ["nope"]})
|
||||
|
||||
with pytest.raises(ValueError, match="config\\.tracking\\.wandb must be a mapping"):
|
||||
V2Config.from_root_dict({**base, "tracking": {"wandb": ["nope"]}})
|
||||
|
||||
with pytest.raises(ValueError, match="config\\.serving must be a mapping"):
|
||||
V2Config.from_root_dict({**base, "serving": ["nope"]})
|
||||
|
||||
with pytest.raises(ValueError, match="config\\.serving\\.\\{serve,llm\\} must be mappings"):
|
||||
V2Config.from_root_dict({**base, "serving": {"serve": ["x"], "llm": {}}})
|
||||
|
||||
def test_tracking_wandb_defaults_disabled():
|
||||
from argus.service.config import V2Config
|
||||
|
||||
cfg = V2Config.from_root_dict(
|
||||
{
|
||||
"ray": {"shared_root": "/private"},
|
||||
"service": {"api": {}, "auth": {}, "sqlite": {}, "scheduler": {}},
|
||||
"data": {"sftpgo": {}, "retention": {}},
|
||||
}
|
||||
)
|
||||
assert cfg.tracking.wandb.enabled is False
|
||||
assert cfg.tracking.wandb.base_url == ""
|
||||
assert cfg.tracking.wandb.api_key_env == "WANDB_API_KEY"
|
||||
assert cfg.tracking.wandb.project_suffix == "_project"
|
||||
|
||||
|
||||
def test_tracking_wandb_enabled_parses():
|
||||
from argus.service.config import V2Config
|
||||
|
||||
cfg = V2Config.from_root_dict(
|
||||
{
|
||||
"ray": {"shared_root": "/private"},
|
||||
"service": {"api": {}, "auth": {}, "sqlite": {}, "scheduler": {}},
|
||||
"tracking": {"wandb": {"enabled": True, "base_url": "http://127.0.0.1:8090", "api_key_env": "X"}},
|
||||
"data": {"sftpgo": {}, "retention": {}},
|
||||
}
|
||||
)
|
||||
assert cfg.tracking.wandb.enabled is True
|
||||
assert cfg.tracking.wandb.base_url == "http://127.0.0.1:8090"
|
||||
assert cfg.tracking.wandb.api_key_env == "X"
|
||||
|
||||
|
||||
def test_tracking_wandb_enabled_requires_base_url():
|
||||
from argus.service.config import V2Config
|
||||
|
||||
with pytest.raises(ValueError, match="tracking\\.wandb\\.base_url is required"):
|
||||
V2Config.from_root_dict(
|
||||
{
|
||||
"ray": {"shared_root": "/private"},
|
||||
"service": {"api": {}, "auth": {}, "sqlite": {}, "scheduler": {}},
|
||||
"tracking": {"wandb": {"enabled": True, "base_url": ""}},
|
||||
"data": {"sftpgo": {}, "retention": {}},
|
||||
}
|
||||
)
|
||||
23
src/mvp/py/tests/test_serving_model_id_prefix.py
Normal file
@@ -0,0 +1,23 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime, timezone
|
||||
|
||||
import pytest
|
||||
|
||||
from argus.service.serving_spec import make_model_id_prefix
|
||||
|
||||
|
||||
def test_make_model_id_prefix_uses_utc_minutes():
|
||||
dt = datetime(2026, 1, 6, 12, 35, 59, tzinfo=timezone.utc)
|
||||
assert make_model_id_prefix(user_id="alice", now_utc=dt) == "alice-202601061235"
|
||||
|
||||
|
||||
def test_make_model_id_prefix_rejects_empty_user_id():
|
||||
with pytest.raises(ValueError, match="user_id is required"):
|
||||
make_model_id_prefix(user_id="", now_utc=datetime.now(timezone.utc))
|
||||
|
||||
|
||||
def test_make_model_id_prefix_rejects_slash():
|
||||
with pytest.raises(ValueError, match="must not contain"):
|
||||
make_model_id_prefix(user_id="bad/user", now_utc=datetime.now(timezone.utc))
|
||||
|
||||
354
src/mvp/py/tests/test_serving_reconciler.py
Normal file
@@ -0,0 +1,354 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
class _FakeServeClient:
|
||||
def __init__(self):
|
||||
self.started = 0
|
||||
self.applied = []
|
||||
self.deleted = []
|
||||
self.status_calls = 0
|
||||
self.fail_apply = False
|
||||
self.fail_status = False
|
||||
self.fail_delete = False
|
||||
self.fail_start = False
|
||||
|
||||
def ensure_started(self) -> None:
|
||||
if self.fail_start:
|
||||
raise RuntimeError("start boom")
|
||||
self.started += 1
|
||||
|
||||
def apply_app(self, *, app, app_name: str, route_prefix: str = "/"):
|
||||
if self.fail_apply:
|
||||
raise RuntimeError("boom")
|
||||
self.applied.append({"app": app, "app_name": app_name, "route_prefix": route_prefix})
|
||||
return {"ok": True}
|
||||
|
||||
def delete_app(self, *, app_name: str):
|
||||
if self.fail_delete:
|
||||
raise RuntimeError("delete boom")
|
||||
self.deleted.append({"app_name": app_name})
|
||||
return {"ok": True}
|
||||
|
||||
def get_status(self):
|
||||
self.status_calls += 1
|
||||
if self.fail_status:
|
||||
raise RuntimeError("status boom")
|
||||
return {"ok": True}
|
||||
|
||||
|
||||
def _seed_model(db, *, model_key: str, user_id: str, state: str, num_replicas: int = 1, gpus_per_replica: int = 1):
|
||||
spec_yaml = "model_id: x\nmodel_source: $HOME/common/hf/x\n"
|
||||
resolved_yaml = f"user_id: {user_id}\nmodel_id: {user_id}-202601061235-x\n"
|
||||
db.create_serve_model(
|
||||
model_key=model_key,
|
||||
user_id=user_id,
|
||||
model_id_suffix="x",
|
||||
model_id_prefix=f"{user_id}-202601061235",
|
||||
model_id=f"{user_id}-202601061235-x",
|
||||
model_source="/private/hf/x",
|
||||
num_replicas=num_replicas,
|
||||
gpus_per_replica=gpus_per_replica,
|
||||
engine_kwargs_json=json.dumps({"gpu_memory_utilization": 0.9}),
|
||||
spec_yaml=spec_yaml,
|
||||
resolved_spec_yaml=resolved_yaml,
|
||||
)
|
||||
db.set_serve_model_state(model_key=model_key, state=state, event_type="TEST_SEED")
|
||||
|
||||
|
||||
def test_reconciler_skips_when_no_changes(tmp_path: Path):
|
||||
from argus.service.config import V2Config
|
||||
from argus.service.db import Db
|
||||
from argus.service.serving_reconciler import ServingReconciler
|
||||
|
||||
root = {
|
||||
"ray": {"shared_root": "/private"},
|
||||
"service": {"api": {}, "auth": {}, "sqlite": {"db_path": str(tmp_path / "mvp.sqlite3")}, "scheduler": {}},
|
||||
"data": {"sftpgo": {}, "retention": {}},
|
||||
"serving": {"serve": {"http_port": 8000}, "llm": {"accelerator_type": "H20"}},
|
||||
}
|
||||
cfg = V2Config.from_root_dict(root)
|
||||
db = Db(cfg.sqlite.db_path)
|
||||
db.init()
|
||||
|
||||
client = _FakeServeClient()
|
||||
rec = ServingReconciler(db=db, v2_cfg=cfg, ray_runtime_env_env_vars={}, serve_client=client, get_available_fn=lambda: type("A", (), {"total_available_gpus": 8, "total_available_npus": 0})())
|
||||
rec.tick()
|
||||
assert client.started == 0
|
||||
assert client.applied == []
|
||||
|
||||
|
||||
def test_reconciler_pending_resources_no_apply(tmp_path: Path):
|
||||
from argus.service.config import V2Config
|
||||
from argus.service.db import Db
|
||||
from argus.service.serving_reconciler import ServingReconciler
|
||||
|
||||
cfg = V2Config.from_root_dict(
|
||||
{
|
||||
"ray": {"shared_root": "/private"},
|
||||
"service": {"api": {}, "auth": {}, "sqlite": {"db_path": str(tmp_path / "mvp.sqlite3")}, "scheduler": {}},
|
||||
"data": {"sftpgo": {}, "retention": {}},
|
||||
"serving": {"serve": {"http_port": 8000}, "llm": {"accelerator_type": "H20"}},
|
||||
}
|
||||
)
|
||||
db = Db(cfg.sqlite.db_path)
|
||||
db.init()
|
||||
_seed_model(db, model_key="mk1", user_id="alice", state="QUEUED", num_replicas=2, gpus_per_replica=4)
|
||||
|
||||
client = _FakeServeClient()
|
||||
rec = ServingReconciler(
|
||||
db=db,
|
||||
v2_cfg=cfg,
|
||||
ray_runtime_env_env_vars={},
|
||||
serve_client=client,
|
||||
get_available_fn=lambda: type("A", (), {"total_available_gpus": 1, "total_available_npus": 0})(),
|
||||
)
|
||||
rec.tick()
|
||||
# Serve may be started even when resources are insufficient, but apply should not happen.
|
||||
assert client.started == 1
|
||||
assert client.applied == []
|
||||
# State remains QUEUED.
|
||||
row = db.get_serve_model("mk1")
|
||||
assert row and row["state"] == "QUEUED"
|
||||
ev = db.list_serve_events("mk1", limit=50)
|
||||
assert any(e["event_type"] == "SERVE_PENDING_RESOURCES" for e in ev)
|
||||
|
||||
|
||||
def test_reconciler_apply_success_marks_running(tmp_path: Path):
|
||||
from argus.service.config import V2Config
|
||||
from argus.service.db import Db
|
||||
from argus.service.serving_reconciler import ServingReconciler
|
||||
|
||||
cfg = V2Config.from_root_dict(
|
||||
{
|
||||
"ray": {"shared_root": "/private"},
|
||||
"service": {"api": {}, "auth": {}, "sqlite": {"db_path": str(tmp_path / "mvp.sqlite3")}, "scheduler": {}},
|
||||
"data": {"sftpgo": {}, "retention": {}},
|
||||
"serving": {"serve": {"http_port": 8000}, "llm": {"accelerator_type": "H20"}},
|
||||
}
|
||||
)
|
||||
db = Db(cfg.sqlite.db_path)
|
||||
db.init()
|
||||
_seed_model(db, model_key="mk1", user_id="alice", state="QUEUED", num_replicas=1, gpus_per_replica=1)
|
||||
|
||||
client = _FakeServeClient()
|
||||
rec = ServingReconciler(
|
||||
db=db,
|
||||
v2_cfg=cfg,
|
||||
ray_runtime_env_env_vars={"HF_ENDPOINT": "https://hf-mirror.com"},
|
||||
serve_client=client,
|
||||
get_available_fn=lambda: type("A", (), {"total_available_gpus": 8, "total_available_npus": 0})(),
|
||||
)
|
||||
rec.tick()
|
||||
assert client.started == 1
|
||||
assert len(client.applied) == 1
|
||||
assert client.applied[0]["app_name"] == "mk1"
|
||||
assert client.applied[0]["route_prefix"] == "/serve/mk1"
|
||||
applied = client.applied[0]["app"]["llm_configs"]
|
||||
assert applied[0]["engine_kwargs"]["tensor_parallel_size"] == 1
|
||||
assert applied[0]["runtime_env"]["env_vars"]["HF_HUB_OFFLINE"] == "1"
|
||||
row = db.get_serve_model("mk1")
|
||||
assert row and row["state"] == "RUNNING"
|
||||
|
||||
|
||||
def test_reconciler_delete_removes_and_marks_deleted(tmp_path: Path):
|
||||
from argus.service.config import V2Config
|
||||
from argus.service.db import Db
|
||||
from argus.service.serving_reconciler import ServingReconciler
|
||||
|
||||
cfg = V2Config.from_root_dict(
|
||||
{
|
||||
"ray": {"shared_root": "/private"},
|
||||
"service": {"api": {}, "auth": {}, "sqlite": {"db_path": str(tmp_path / "mvp.sqlite3")}, "scheduler": {}},
|
||||
"data": {"sftpgo": {}, "retention": {}},
|
||||
"serving": {"serve": {"http_port": 8000}, "llm": {"accelerator_type": "H20"}},
|
||||
}
|
||||
)
|
||||
db = Db(cfg.sqlite.db_path)
|
||||
db.init()
|
||||
_seed_model(db, model_key="keep", user_id="alice", state="RUNNING", num_replicas=1, gpus_per_replica=1)
|
||||
_seed_model(db, model_key="del", user_id="alice", state="DELETING", num_replicas=1, gpus_per_replica=1)
|
||||
|
||||
client = _FakeServeClient()
|
||||
rec = ServingReconciler(
|
||||
db=db,
|
||||
v2_cfg=cfg,
|
||||
ray_runtime_env_env_vars={},
|
||||
serve_client=client,
|
||||
get_available_fn=lambda: type("A", (), {"total_available_gpus": 8, "total_available_npus": 0})(),
|
||||
)
|
||||
rec.tick()
|
||||
assert client.applied == []
|
||||
assert client.deleted == [{"app_name": "del"}]
|
||||
row = db.get_serve_model("del")
|
||||
assert row and row["state"] == "DELETED"
|
||||
assert row.get("deleted_at")
|
||||
|
||||
|
||||
def test_reconciler_apply_failure_marks_failed(tmp_path: Path):
|
||||
from argus.service.config import V2Config
|
||||
from argus.service.db import Db
|
||||
from argus.service.serving_reconciler import ServingReconciler
|
||||
|
||||
cfg = V2Config.from_root_dict(
|
||||
{
|
||||
"ray": {"shared_root": "/private"},
|
||||
"service": {"api": {}, "auth": {}, "sqlite": {"db_path": str(tmp_path / "mvp.sqlite3")}, "scheduler": {}},
|
||||
"data": {"sftpgo": {}, "retention": {}},
|
||||
"serving": {"serve": {"http_port": 8000}, "llm": {"accelerator_type": "H20"}},
|
||||
}
|
||||
)
|
||||
db = Db(cfg.sqlite.db_path)
|
||||
db.init()
|
||||
_seed_model(db, model_key="mk1", user_id="alice", state="QUEUED")
|
||||
|
||||
client = _FakeServeClient()
|
||||
client.fail_apply = True
|
||||
rec = ServingReconciler(
|
||||
db=db,
|
||||
v2_cfg=cfg,
|
||||
ray_runtime_env_env_vars={},
|
||||
serve_client=client,
|
||||
get_available_fn=lambda: type("A", (), {"total_available_gpus": 8, "total_available_npus": 0})(),
|
||||
)
|
||||
rec.tick()
|
||||
row = db.get_serve_model("mk1")
|
||||
assert row and row["state"] == "FAILED"
|
||||
assert row.get("error_summary")
|
||||
|
||||
|
||||
def test_reconciler_applies_one_model_per_tick(tmp_path: Path):
|
||||
from argus.service.config import V2Config
|
||||
from argus.service.db import Db
|
||||
from argus.service.serving_reconciler import ServingReconciler
|
||||
|
||||
cfg = V2Config.from_root_dict(
|
||||
{
|
||||
"ray": {"shared_root": "/private"},
|
||||
"service": {"api": {}, "auth": {}, "sqlite": {"db_path": str(tmp_path / "mvp.sqlite3")}, "scheduler": {}},
|
||||
"data": {"sftpgo": {}, "retention": {}},
|
||||
"serving": {"serve": {"http_port": 8000}, "llm": {"accelerator_type": "H20"}},
|
||||
}
|
||||
)
|
||||
db = Db(cfg.sqlite.db_path)
|
||||
db.init()
|
||||
_seed_model(db, model_key="mk1", user_id="alice", state="QUEUED", num_replicas=1, gpus_per_replica=1)
|
||||
_seed_model(db, model_key="mk2", user_id="alice", state="QUEUED", num_replicas=1, gpus_per_replica=1)
|
||||
|
||||
client = _FakeServeClient()
|
||||
rec = ServingReconciler(
|
||||
db=db,
|
||||
v2_cfg=cfg,
|
||||
ray_runtime_env_env_vars={},
|
||||
serve_client=client,
|
||||
get_available_fn=lambda: type("A", (), {"total_available_gpus": 8, "total_available_npus": 0})(),
|
||||
)
|
||||
|
||||
rec.tick()
|
||||
assert len(client.applied) == 1
|
||||
assert client.applied[0]["app_name"] == "mk1"
|
||||
assert db.get_serve_model("mk1")["state"] == "RUNNING"
|
||||
assert db.get_serve_model("mk2")["state"] == "QUEUED"
|
||||
|
||||
rec.tick()
|
||||
assert len(client.applied) == 2
|
||||
assert client.applied[1]["app_name"] == "mk2"
|
||||
assert db.get_serve_model("mk2")["state"] == "RUNNING"
|
||||
|
||||
|
||||
def test_reconciler_start_error_is_recorded(tmp_path: Path):
|
||||
from argus.service.config import V2Config
|
||||
from argus.service.db import Db
|
||||
from argus.service.serving_reconciler import ServingReconciler
|
||||
|
||||
cfg = V2Config.from_root_dict(
|
||||
{
|
||||
"ray": {"shared_root": "/private"},
|
||||
"service": {"api": {}, "auth": {}, "sqlite": {"db_path": str(tmp_path / "mvp.sqlite3")}, "scheduler": {}},
|
||||
"data": {"sftpgo": {}, "retention": {}},
|
||||
"serving": {"serve": {"http_port": 8000}, "llm": {"accelerator_type": "H20"}},
|
||||
}
|
||||
)
|
||||
db = Db(cfg.sqlite.db_path)
|
||||
db.init()
|
||||
_seed_model(db, model_key="mk1", user_id="alice", state="QUEUED", num_replicas=1, gpus_per_replica=1)
|
||||
|
||||
client = _FakeServeClient()
|
||||
client.fail_start = True
|
||||
rec = ServingReconciler(db=db, v2_cfg=cfg, ray_runtime_env_env_vars={}, serve_client=client)
|
||||
rec.tick()
|
||||
assert db.get_serve_model("mk1")["state"] == "QUEUED"
|
||||
ev = db.list_serve_events("mk1", limit=50)
|
||||
assert any(e["event_type"] == "SERVE_START_ERROR" for e in ev)
|
||||
|
||||
|
||||
def test_reconciler_delete_failure_marks_failed(tmp_path: Path):
|
||||
from argus.service.config import V2Config
|
||||
from argus.service.db import Db
|
||||
from argus.service.serving_reconciler import ServingReconciler
|
||||
|
||||
cfg = V2Config.from_root_dict(
|
||||
{
|
||||
"ray": {"shared_root": "/private"},
|
||||
"service": {"api": {}, "auth": {}, "sqlite": {"db_path": str(tmp_path / "mvp.sqlite3")}, "scheduler": {}},
|
||||
"data": {"sftpgo": {}, "retention": {}},
|
||||
"serving": {"serve": {"http_port": 8000}, "llm": {"accelerator_type": "H20"}},
|
||||
}
|
||||
)
|
||||
db = Db(cfg.sqlite.db_path)
|
||||
db.init()
|
||||
_seed_model(db, model_key="mk1", user_id="alice", state="DELETING", num_replicas=1, gpus_per_replica=1)
|
||||
|
||||
client = _FakeServeClient()
|
||||
client.fail_delete = True
|
||||
rec = ServingReconciler(db=db, v2_cfg=cfg, ray_runtime_env_env_vars={}, serve_client=client)
|
||||
rec.tick()
|
||||
row = db.get_serve_model("mk1")
|
||||
assert row and row["state"] == "FAILED"
|
||||
assert row.get("error_summary")
|
||||
|
||||
|
||||
def test_reconciler_engine_kwargs_bad_json_is_ignored(tmp_path: Path):
|
||||
from argus.service.config import V2Config
|
||||
from argus.service.db import Db
|
||||
from argus.service.serving_reconciler import ServingReconciler
|
||||
|
||||
cfg = V2Config.from_root_dict(
|
||||
{
|
||||
"ray": {"shared_root": "/private"},
|
||||
"service": {"api": {}, "auth": {}, "sqlite": {"db_path": str(tmp_path / "mvp.sqlite3")}, "scheduler": {}},
|
||||
"data": {"sftpgo": {}, "retention": {}},
|
||||
"serving": {"serve": {"http_port": 8000}, "llm": {"accelerator_type": "H20"}},
|
||||
}
|
||||
)
|
||||
db = Db(cfg.sqlite.db_path)
|
||||
db.init()
|
||||
|
||||
db.create_serve_model(
|
||||
model_key="mk1",
|
||||
user_id="alice",
|
||||
model_id_suffix="x",
|
||||
model_id_prefix="alice-202601061235",
|
||||
model_id="alice-202601061235-x",
|
||||
model_source="/private/hf/x",
|
||||
num_replicas=1,
|
||||
gpus_per_replica=1,
|
||||
engine_kwargs_json="not-json",
|
||||
spec_yaml="model_id: x\nmodel_source: $HOME/common/hf/x\n",
|
||||
resolved_spec_yaml="user_id: alice\nmodel_id: alice-202601061235-x\n",
|
||||
)
|
||||
db.set_serve_model_state(model_key="mk1", state="QUEUED", event_type="TEST_SEED")
|
||||
|
||||
client = _FakeServeClient()
|
||||
rec = ServingReconciler(
|
||||
db=db,
|
||||
v2_cfg=cfg,
|
||||
ray_runtime_env_env_vars={},
|
||||
serve_client=client,
|
||||
get_available_fn=lambda: type("A", (), {"total_available_gpus": 8, "total_available_npus": 0})(),
|
||||
)
|
||||
rec.tick()
|
||||
cfgs = client.applied[0]["app"]["llm_configs"]
|
||||
assert cfgs[0]["engine_kwargs"]["tensor_parallel_size"] == 1
|
||||
47
src/mvp/py/tests/test_serving_spec_paths.py
Normal file
@@ -0,0 +1,47 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime, timezone
|
||||
|
||||
import pytest
|
||||
|
||||
from argus.service.serving_spec import ServingSpec, resolve_serving_spec
|
||||
|
||||
|
||||
def test_expand_home_macro_and_validate_user_path_ok():
|
||||
spec = ServingSpec(
|
||||
model_id="qwen-0.5b",
|
||||
model_source="$HOME/models/my_model",
|
||||
num_replicas=1,
|
||||
gpus_per_replica=1,
|
||||
)
|
||||
r = resolve_serving_spec(spec=spec, user_id="alice", now_utc=datetime(2026, 1, 6, 12, 35, tzinfo=timezone.utc))
|
||||
assert r.model_source == "/private/users/alice/models/my_model"
|
||||
assert r.model_id == "alice-202601061235-qwen-0.5b"
|
||||
|
||||
|
||||
def test_expand_common_hf_macro_ok():
|
||||
spec = ServingSpec(
|
||||
model_id="qwen-0.5b",
|
||||
model_source="$HOME/common/hf/hub/models--Qwen--Qwen2.5/snapshots/abc",
|
||||
num_replicas=1,
|
||||
gpus_per_replica=1,
|
||||
)
|
||||
r = resolve_serving_spec(spec=spec, user_id="alice", now_utc=datetime(2026, 1, 6, 12, 35, tzinfo=timezone.utc))
|
||||
assert r.model_source.startswith("/private/hf/")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"src",
|
||||
[
|
||||
"/etc/passwd",
|
||||
"relative/path",
|
||||
"/private/users/bob/models/x",
|
||||
"/private/users/alice/../bob/x",
|
||||
"/private/common/hf/x",
|
||||
],
|
||||
)
|
||||
def test_model_source_path_rejected(src: str):
|
||||
spec = ServingSpec(model_id="qwen-0.5b", model_source=src, num_replicas=1, gpus_per_replica=1)
|
||||
with pytest.raises((ValueError, PermissionError)):
|
||||
resolve_serving_spec(spec=spec, user_id="alice", now_utc=datetime(2026, 1, 6, 12, 35, tzinfo=timezone.utc))
|
||||
|
||||
72
src/mvp/py/tests/test_serving_spec_validation.py
Normal file
@@ -0,0 +1,72 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from argus.service.serving_spec import ServingSpec, parse_serving_spec, validate_model_id_suffix
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"suffix",
|
||||
[
|
||||
"a",
|
||||
"qwen-0.5b",
|
||||
"Qwen2.5-0.5B",
|
||||
"a_b",
|
||||
"a.b-c",
|
||||
"a" * 64,
|
||||
],
|
||||
)
|
||||
def test_validate_model_id_suffix_accepts(suffix: str):
|
||||
validate_model_id_suffix(suffix)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"suffix",
|
||||
[
|
||||
"",
|
||||
" a",
|
||||
"a ",
|
||||
"-bad",
|
||||
".bad",
|
||||
"bad/",
|
||||
"bad..",
|
||||
"bad\n",
|
||||
"bad\t",
|
||||
"a" * 65,
|
||||
],
|
||||
)
|
||||
def test_validate_model_id_suffix_rejects(suffix: str):
|
||||
with pytest.raises(ValueError):
|
||||
validate_model_id_suffix(suffix)
|
||||
|
||||
|
||||
def test_parse_serving_spec_smoke_defaults():
|
||||
spec = parse_serving_spec(
|
||||
{
|
||||
"model_id": "qwen-0.5b",
|
||||
"model_source": "/private/hf/x",
|
||||
}
|
||||
)
|
||||
assert isinstance(spec, ServingSpec)
|
||||
assert spec.num_replicas == 1
|
||||
assert spec.gpus_per_replica == 1
|
||||
assert spec.engine_kwargs is None
|
||||
|
||||
|
||||
def test_parse_serving_spec_rejects_missing_fields():
|
||||
with pytest.raises(ValueError, match="missing required field: model_id"):
|
||||
parse_serving_spec({"model_source": "/private/hf/x"})
|
||||
with pytest.raises(ValueError, match="missing required field: model_source"):
|
||||
parse_serving_spec({"model_id": "x"})
|
||||
|
||||
|
||||
def test_parse_serving_spec_rejects_bad_types():
|
||||
with pytest.raises(ValueError, match="serving spec must be a mapping"):
|
||||
parse_serving_spec(["nope"])
|
||||
with pytest.raises(ValueError, match="num_replicas"):
|
||||
parse_serving_spec({"model_id": "x", "model_source": "/private/hf/x", "num_replicas": 0})
|
||||
with pytest.raises(ValueError, match="gpus_per_replica"):
|
||||
parse_serving_spec({"model_id": "x", "model_source": "/private/hf/x", "gpus_per_replica": 0})
|
||||
with pytest.raises(ValueError, match="engine_kwargs"):
|
||||
parse_serving_spec({"model_id": "x", "model_source": "/private/hf/x", "engine_kwargs": "nope"})
|
||||
|
||||
322
src/mvp/py/tests/test_sftpgo.py
Normal file
@@ -0,0 +1,322 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import yaml
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
|
||||
def _write_config(tmp_path: Path, *, enabled: bool) -> Path:
|
||||
cfg = {
|
||||
"ray": {
|
||||
"address": "http://127.0.0.1:8265",
|
||||
"shared_root": "/private",
|
||||
"entrypoint_resources": {"worker_node": 1},
|
||||
"runtime_env": {"env_vars": {}},
|
||||
},
|
||||
"data": {
|
||||
"user_root": str(tmp_path / "users"),
|
||||
"sftpgo": {
|
||||
"enabled": bool(enabled),
|
||||
"admin_api_base": "http://127.0.0.1:8081/api/v2",
|
||||
"admin_user": "admin",
|
||||
"admin_password_env": "SFTPGO_ADMIN_PASSWORD",
|
||||
"host": "h1.internal",
|
||||
"sftp_port": 2022,
|
||||
},
|
||||
},
|
||||
"service": {
|
||||
"api": {"host": "127.0.0.1", "port": 0},
|
||||
"auth": {"token_env": "MVP_INTERNAL_TOKEN"},
|
||||
"sqlite": {"db_path": str(tmp_path / "mvp.sqlite3")},
|
||||
"scheduler": {"tick_s": 1, "retry_interval_s": 1, "max_running_tasks": 1},
|
||||
},
|
||||
}
|
||||
p = tmp_path / "cfg.yaml"
|
||||
p.write_text(yaml.safe_dump(cfg), encoding="utf-8")
|
||||
return p
|
||||
|
||||
|
||||
def test_create_user_calls_sftpgo_when_enabled(tmp_path: Path, monkeypatch):
|
||||
from argus.service import app as app_mod
|
||||
|
||||
cfg_path = _write_config(tmp_path, enabled=True)
|
||||
monkeypatch.setenv("MVP_INTERNAL_TOKEN", "adm1")
|
||||
monkeypatch.setenv("SFTPGO_ADMIN_PASSWORD", "pw1")
|
||||
|
||||
calls = {"create": [], "disable": [], "reset": []}
|
||||
|
||||
class _Client:
|
||||
def __init__(self, **kwargs):
|
||||
pass
|
||||
|
||||
def create_user(self, *, username: str, password: str, home_dir: str) -> None:
|
||||
calls["create"].append((username, password, home_dir))
|
||||
|
||||
def enable_user(self, *, username: str, home_dir: str) -> None:
|
||||
# Not used in this test, but required by app for idempotent upsert.
|
||||
return None
|
||||
|
||||
def disable_user(self, *, username: str, home_dir: str) -> None:
|
||||
calls["disable"].append(username)
|
||||
|
||||
def reset_password(self, *, username: str, new_password: str, home_dir: str) -> None:
|
||||
calls["reset"].append((username, new_password))
|
||||
|
||||
monkeypatch.setattr(app_mod, "SFTPGoAdminClient", _Client)
|
||||
|
||||
class _Scheduler:
|
||||
def __init__(self, **kwargs):
|
||||
self.tool = object()
|
||||
|
||||
def run_forever(self, stop_flag):
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(app_mod, "Scheduler", _Scheduler)
|
||||
app = app_mod.create_app(str(cfg_path))
|
||||
|
||||
admin_headers = {"authorization": "Bearer adm1"}
|
||||
with TestClient(app) as c:
|
||||
r = c.post("/api/v2/users", headers=admin_headers, json={"user_id": "alice"})
|
||||
assert r.status_code == 200
|
||||
assert calls["create"]
|
||||
username, password, home_dir = calls["create"][-1]
|
||||
assert username == "alice"
|
||||
assert password
|
||||
assert home_dir.endswith("/users/alice")
|
||||
|
||||
r2 = c.post("/api/v2/users/alice:disable", headers=admin_headers)
|
||||
assert r2.status_code == 200
|
||||
assert calls["disable"] == ["alice"]
|
||||
|
||||
r3 = c.post("/api/v2/users/alice/sftp:reset_password", headers=admin_headers)
|
||||
assert r3.status_code == 200
|
||||
body = r3.json()
|
||||
assert body["user_id"] == "alice"
|
||||
assert body["password"]
|
||||
assert calls["reset"] and calls["reset"][-1][0] == "alice"
|
||||
|
||||
|
||||
def test_create_user_upserts_when_sftpgo_user_already_exists(tmp_path: Path, monkeypatch):
|
||||
from argus.service import app as app_mod
|
||||
from argus.service.sftpgo import SFTPGoError
|
||||
|
||||
cfg_path = _write_config(tmp_path, enabled=True)
|
||||
monkeypatch.setenv("MVP_INTERNAL_TOKEN", "adm1")
|
||||
monkeypatch.setenv("SFTPGO_ADMIN_PASSWORD", "pw1")
|
||||
|
||||
calls = {"create": 0, "reset": [], "enable": []}
|
||||
|
||||
class _Client:
|
||||
def __init__(self, **kwargs):
|
||||
pass
|
||||
|
||||
def create_user(self, *, username: str, password: str, home_dir: str) -> None:
|
||||
calls["create"] += 1
|
||||
raise SFTPGoError("sftpgo http error: 409 Conflict")
|
||||
|
||||
def reset_password(self, *, username: str, new_password: str, home_dir: str) -> None:
|
||||
calls["reset"].append((username, new_password))
|
||||
|
||||
def enable_user(self, *, username: str, home_dir: str) -> None:
|
||||
calls["enable"].append(username)
|
||||
|
||||
def disable_user(self, *, username: str, home_dir: str) -> None:
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(app_mod, "SFTPGoAdminClient", _Client)
|
||||
|
||||
class _Scheduler:
|
||||
def __init__(self, **kwargs):
|
||||
self.tool = object()
|
||||
|
||||
def run_forever(self, stop_flag):
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(app_mod, "Scheduler", _Scheduler)
|
||||
app = app_mod.create_app(str(cfg_path))
|
||||
|
||||
admin_headers = {"authorization": "Bearer adm1"}
|
||||
with TestClient(app) as c:
|
||||
r = c.post("/api/v2/users", headers=admin_headers, json={"user_id": "bob"})
|
||||
assert r.status_code == 200
|
||||
body = r.json()
|
||||
assert body["user_id"] == "bob"
|
||||
assert body.get("sftp", {}).get("password")
|
||||
assert calls["create"] == 1
|
||||
assert calls["reset"] and calls["reset"][-1][0] == "bob"
|
||||
assert calls["enable"] == ["bob"]
|
||||
|
||||
|
||||
def test_sftpgo_enabled_requires_admin_password_env(tmp_path: Path, monkeypatch):
|
||||
from argus.service import app as app_mod
|
||||
|
||||
cfg_path = _write_config(tmp_path, enabled=True)
|
||||
monkeypatch.setenv("MVP_INTERNAL_TOKEN", "adm1")
|
||||
monkeypatch.delenv("SFTPGO_ADMIN_PASSWORD", raising=False)
|
||||
|
||||
class _Scheduler:
|
||||
def __init__(self, **kwargs):
|
||||
self.tool = object()
|
||||
|
||||
def run_forever(self, stop_flag):
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(app_mod, "Scheduler", _Scheduler)
|
||||
app = app_mod.create_app(str(cfg_path))
|
||||
|
||||
admin_headers = {"authorization": "Bearer adm1"}
|
||||
with TestClient(app) as c:
|
||||
r = c.post("/api/v2/users", headers=admin_headers, json={"user_id": "alice"})
|
||||
assert r.status_code == 500
|
||||
assert "SFTPGO_ADMIN_PASSWORD" in r.text
|
||||
|
||||
|
||||
def test_sftpgo_admin_client_builds_requests(monkeypatch):
|
||||
import json
|
||||
from urllib.request import Request
|
||||
|
||||
from argus.service.sftpgo import SFTPGoAdminClient
|
||||
import argus.service.sftpgo as mod
|
||||
|
||||
seen: list[Request] = []
|
||||
|
||||
class _Resp:
|
||||
def __init__(self, body: bytes = b"ok"):
|
||||
self._body = body
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
def read(self):
|
||||
return self._body
|
||||
|
||||
def fake_urlopen(req, timeout=0):
|
||||
seen.append(req)
|
||||
if req.full_url.endswith("/token"):
|
||||
return _Resp(body=b'{"access_token":"t1","expires_at":0}')
|
||||
# allow folder creates to be idempotent in tests
|
||||
if req.full_url.endswith("/folders") and req.get_method() == "POST":
|
||||
return _Resp(body=b'{"message":"created"}')
|
||||
if req.full_url.endswith("/users/alice") and req.get_method() == "GET":
|
||||
return _Resp(
|
||||
body=b'{"username":"alice","status":1,"home_dir":"/private/users/alice","uid":0,"gid":0,"permissions":{"/":["*"]},"virtual_folders":[]}'
|
||||
)
|
||||
return _Resp()
|
||||
|
||||
monkeypatch.setattr(mod, "urlopen", fake_urlopen)
|
||||
|
||||
c = SFTPGoAdminClient(admin_api_base="http://sftpgo.local/api/v2", admin_user="admin", admin_password="pw", common_root="/private/common")
|
||||
c.create_user(username="alice", password="p1", home_dir="/private/users/alice")
|
||||
c.disable_user(username="alice", home_dir="/private/users/alice")
|
||||
c.enable_user(username="alice", home_dir="/private/users/alice")
|
||||
c.reset_password(username="alice", new_password="p2", home_dir="/private/users/alice")
|
||||
|
||||
# Each operation fetches a token, then issues a request.
|
||||
# For create_user: token + ensure folders (2 POSTs) + create user
|
||||
# For disable/enable/reset: token + ensure folders (2 POSTs) + GET user + PUT user
|
||||
assert len(seen) == 1 + 2 + 1 + 3 * (1 + 2 + 1 + 1)
|
||||
assert seen[0].full_url.endswith("/token")
|
||||
assert seen[0].headers.get("Authorization", "").startswith("Basic ")
|
||||
# create_user
|
||||
assert seen[1].full_url.endswith("/folders")
|
||||
assert seen[2].full_url.endswith("/folders")
|
||||
assert seen[3].full_url.endswith("/users")
|
||||
assert seen[3].headers.get("Authorization", "") == "Bearer t1"
|
||||
created = json.loads(seen[3].data.decode("utf-8"))
|
||||
assert created["username"] == "alice"
|
||||
assert "/common" in created.get("permissions", {})
|
||||
assert "/common/datasets" in created.get("permissions", {})
|
||||
assert "/common/hf" in created.get("permissions", {})
|
||||
|
||||
# disable_user
|
||||
assert seen[4].full_url.endswith("/token")
|
||||
assert seen[5].full_url.endswith("/folders")
|
||||
assert seen[6].full_url.endswith("/folders")
|
||||
assert seen[7].full_url.endswith("/users/alice")
|
||||
assert seen[7].get_method() == "GET"
|
||||
assert seen[8].full_url.endswith("/users/alice")
|
||||
assert seen[8].get_method() == "PUT"
|
||||
|
||||
# enable_user
|
||||
assert seen[9].full_url.endswith("/token")
|
||||
assert seen[10].full_url.endswith("/folders")
|
||||
assert seen[11].full_url.endswith("/folders")
|
||||
assert seen[12].full_url.endswith("/users/alice")
|
||||
assert seen[12].get_method() == "GET"
|
||||
assert seen[13].full_url.endswith("/users/alice")
|
||||
assert seen[13].get_method() == "PUT"
|
||||
|
||||
# reset_password
|
||||
assert seen[14].full_url.endswith("/token")
|
||||
assert seen[15].full_url.endswith("/folders")
|
||||
assert seen[16].full_url.endswith("/folders")
|
||||
assert seen[17].full_url.endswith("/users/alice")
|
||||
assert seen[17].get_method() == "GET"
|
||||
assert seen[18].full_url.endswith("/users/alice")
|
||||
assert seen[18].get_method() == "PUT"
|
||||
|
||||
|
||||
def test_sftpgo_admin_client_http_error_raises(monkeypatch):
|
||||
from urllib.error import HTTPError
|
||||
|
||||
from argus.service.sftpgo import SFTPGoAdminClient, SFTPGoError
|
||||
import argus.service.sftpgo as mod
|
||||
|
||||
def fake_urlopen(req, timeout=0):
|
||||
if req.full_url.endswith("/token"):
|
||||
class _Resp:
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
def read(self):
|
||||
return b'{"access_token":"t1","expires_at":0}'
|
||||
|
||||
return _Resp()
|
||||
raise HTTPError(req.full_url, 500, "boom", hdrs=None, fp=None)
|
||||
|
||||
monkeypatch.setattr(mod, "urlopen", fake_urlopen)
|
||||
|
||||
c = SFTPGoAdminClient(admin_api_base="http://sftpgo.local/api/v2", admin_user="admin", admin_password="pw", common_root="/private/common")
|
||||
try:
|
||||
c.create_user(username="alice", password="p1", home_dir="/private/users/alice")
|
||||
assert False, "expected SFTPGoError"
|
||||
except SFTPGoError as e:
|
||||
assert "http error" in str(e)
|
||||
|
||||
|
||||
def test_sftpgo_admin_client_url_error_raises(monkeypatch):
|
||||
from urllib.error import URLError
|
||||
|
||||
from argus.service.sftpgo import SFTPGoAdminClient, SFTPGoError
|
||||
import argus.service.sftpgo as mod
|
||||
|
||||
def fake_urlopen(req, timeout=0):
|
||||
if req.full_url.endswith("/token"):
|
||||
class _Resp:
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
def read(self):
|
||||
return b'{"access_token":"t1","expires_at":0}'
|
||||
|
||||
return _Resp()
|
||||
raise URLError("no route")
|
||||
|
||||
monkeypatch.setattr(mod, "urlopen", fake_urlopen)
|
||||
|
||||
c = SFTPGoAdminClient(admin_api_base="http://sftpgo.local/api/v2", admin_user="admin", admin_password="pw")
|
||||
try:
|
||||
c.disable_user(username="alice", home_dir="/private/users/alice")
|
||||
assert False, "expected SFTPGoError"
|
||||
except SFTPGoError as e:
|
||||
assert "connection error" in str(e)
|
||||
132
src/mvp/py/tests/test_ui.py
Normal file
@@ -0,0 +1,132 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from argus.service.app import create_app
|
||||
|
||||
|
||||
def _write_config(tmp_path: Path) -> Path:
|
||||
p = tmp_path / "cfg.yaml"
|
||||
p.write_text(
|
||||
"""
|
||||
ray:
|
||||
address: "http://127.0.0.1:8265"
|
||||
shared_root: "/private"
|
||||
entrypoint_num_cpus: 1
|
||||
entrypoint_resources: { worker_node: 1 }
|
||||
runtime_env: { env_vars: { PYTHONUNBUFFERED: "1" } }
|
||||
service:
|
||||
api: { host: "127.0.0.1", port: 8080 }
|
||||
auth: { token_env: "MVP_INTERNAL_TOKEN" }
|
||||
sqlite: { db_path: "%(db)s" }
|
||||
data:
|
||||
user_root: "%(users)s"
|
||||
sftpgo: { enabled: false }
|
||||
retention: { jobs_trash_after_days: 3, jobs_purge_after_days: 7, janitor_interval_s: 3600 }
|
||||
"""
|
||||
% {"db": str(tmp_path / "mvp.sqlite3"), "users": str(tmp_path / "users")}
|
||||
)
|
||||
return p
|
||||
|
||||
|
||||
def test_ui_routes_render_200(tmp_path, monkeypatch):
|
||||
cfg = _write_config(tmp_path)
|
||||
monkeypatch.setenv("MVP_INTERNAL_TOKEN", "admin-token")
|
||||
app = create_app(str(cfg))
|
||||
c = TestClient(app)
|
||||
|
||||
for path in (
|
||||
"/ui",
|
||||
"/ui/login",
|
||||
"/ui/tasks",
|
||||
"/ui/tasks/new",
|
||||
"/ui/serving",
|
||||
"/ui/serving/new",
|
||||
"/ui/data",
|
||||
"/ui/admin",
|
||||
"/ui/tasks/any-task-id",
|
||||
"/ui/tasks/any-task-id/logs",
|
||||
"/ui/serving/any-model-key",
|
||||
):
|
||||
r = c.get(path, allow_redirects=True)
|
||||
assert r.status_code == 200
|
||||
assert "<html" in r.text.lower()
|
||||
|
||||
|
||||
def test_ui_contains_sidebar_links(tmp_path, monkeypatch):
|
||||
cfg = _write_config(tmp_path)
|
||||
monkeypatch.setenv("MVP_INTERNAL_TOKEN", "admin-token")
|
||||
app = create_app(str(cfg))
|
||||
c = TestClient(app)
|
||||
|
||||
r = c.get("/ui/tasks")
|
||||
assert r.status_code == 200
|
||||
for link in ("/ui/tasks", "/ui/tasks/new", "/ui/serving", "/ui/data", "/ui/login", "/ui/admin"):
|
||||
assert link in r.text
|
||||
assert "Ray Dashboard" in r.text
|
||||
|
||||
|
||||
def test_ui_task_detail_shows_ids(tmp_path, monkeypatch):
|
||||
cfg = _write_config(tmp_path)
|
||||
monkeypatch.setenv("MVP_INTERNAL_TOKEN", "admin-token")
|
||||
app = create_app(str(cfg))
|
||||
c = TestClient(app)
|
||||
|
||||
task_id = "mvp3-ppo-20250101-000000-aaaa"
|
||||
r = c.get(f"/ui/tasks/{task_id}")
|
||||
assert r.status_code == 200
|
||||
assert task_id in r.text
|
||||
assert f"/ui/tasks/{task_id}/logs" in r.text
|
||||
assert "TaskSpec (YAML)" in r.text
|
||||
assert "/api/v2/tasks/" in r.text
|
||||
|
||||
|
||||
def test_ui_new_task_contains_advanced_example_snippet(tmp_path, monkeypatch):
|
||||
cfg = _write_config(tmp_path)
|
||||
monkeypatch.setenv("MVP_INTERNAL_TOKEN", "admin-token")
|
||||
app = create_app(str(cfg))
|
||||
c = TestClient(app)
|
||||
|
||||
r = c.get("/ui/tasks/new")
|
||||
assert r.status_code == 200
|
||||
# Ensure Advanced example includes required PPO micro batch override (common failure mode if omitted).
|
||||
assert "kind: advanced" in r.text
|
||||
# workload is not needed for advanced in v3.5.
|
||||
assert "# workload:" not in r.text
|
||||
assert "actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu" in r.text
|
||||
# v3.7: default rollout backend switches from sglang -> vllm.
|
||||
assert "actor_rollout_ref.rollout.name=vllm" in r.text
|
||||
# v3.6: Advanced example uses platform-injected env vars so users don't need to edit W&B project/run.
|
||||
assert "trainer.logger=${MVP_TRAINER_LOGGER}" in r.text
|
||||
assert "trainer.project_name=${MVP_WANDB_PROJECT}" in r.text
|
||||
assert "trainer.experiment_name=${MVP_WANDB_RUN}" in r.text
|
||||
assert "Model merge example" in r.text
|
||||
assert "verl.model_merger" in r.text
|
||||
|
||||
|
||||
def test_ui_new_task_contains_evaluation_example(tmp_path, monkeypatch):
|
||||
cfg = _write_config(tmp_path)
|
||||
monkeypatch.setenv("MVP_INTERNAL_TOKEN", "admin-token")
|
||||
app = create_app(str(cfg))
|
||||
c = TestClient(app)
|
||||
|
||||
r = c.get("/ui/tasks/new")
|
||||
assert r.status_code == 200
|
||||
assert "Evaluation example" in r.text
|
||||
assert "verl.trainer.main_eval" in r.text
|
||||
assert "data.path=$HOME/jobs/<RAY_SUBMISSION_ID>/outputs/<EVAL_PARQUET>.parquet" in r.text
|
||||
assert "+ray_kwargs.ray_init.address=auto" in r.text
|
||||
|
||||
|
||||
def test_ui_login_contains_wandb_section(tmp_path, monkeypatch):
|
||||
cfg = _write_config(tmp_path)
|
||||
monkeypatch.setenv("MVP_INTERNAL_TOKEN", "admin-token")
|
||||
app = create_app(str(cfg))
|
||||
c = TestClient(app)
|
||||
|
||||
r = c.get("/ui/login")
|
||||
assert r.status_code == 200
|
||||
assert "W&B" in r.text
|
||||
assert "Open W&B" in r.text
|
||||