v2.0 代码重构，跑通run_all_api.sh
This commit is contained in:
190
src/mvp/py/argus/service/app.py
Normal file
190
src/mvp/py/argus/service/app.py
Normal file
@@ -0,0 +1,190 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import threading
|
||||
from typing import Any
|
||||
|
||||
import yaml
|
||||
from fastapi import FastAPI, HTTPException, Request, Response
|
||||
|
||||
from argus.core.ids import new_task_id
|
||||
from argus.ray.models import JobSpec, RayConfig
|
||||
|
||||
from .config import V2Config
|
||||
from .db import Db
|
||||
from .scheduler import Scheduler
|
||||
|
||||
|
||||
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)
|
||||
stop_flag = threading.Event()
|
||||
tool = scheduler.tool
|
||||
|
||||
app = FastAPI(title="mvp-v2", version="2.0")
|
||||
|
||||
def _require_token(req: Request) -> None:
|
||||
token_env = v2_cfg.auth.token_env
|
||||
expected = os.environ.get(token_env, "")
|
||||
if not expected:
|
||||
# Misconfigured service; treat as server error.
|
||||
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 != expected:
|
||||
raise HTTPException(status_code=401, detail="invalid token")
|
||||
|
||||
@app.on_event("startup")
|
||||
def _startup() -> None:
|
||||
t = threading.Thread(target=scheduler.run_forever, args=(stop_flag,), daemon=True)
|
||||
t.start()
|
||||
|
||||
@app.on_event("shutdown")
|
||||
def _shutdown() -> None:
|
||||
stop_flag.set()
|
||||
|
||||
@app.post("/api/v2/tasks")
|
||||
async def submit_task(req: Request) -> dict[str, Any]:
|
||||
_require_token(req)
|
||||
body = (await req.body()).decode("utf-8")
|
||||
obj = yaml.safe_load(body) or {}
|
||||
if not isinstance(obj, dict):
|
||||
raise HTTPException(status_code=400, detail="jobspec must be a YAML mapping")
|
||||
|
||||
try:
|
||||
spec = JobSpec.from_dict(obj)
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=400, detail=f"invalid jobspec: {e!r}")
|
||||
|
||||
task_id = new_task_id(spec.workload)
|
||||
db.create_task(
|
||||
task_id=task_id,
|
||||
workload=spec.workload,
|
||||
jobspec_yaml=body,
|
||||
nnodes=spec.nnodes,
|
||||
n_gpus_per_node=spec.n_gpus_per_node,
|
||||
)
|
||||
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]:
|
||||
_require_token(req)
|
||||
row = db.get_task(task_id)
|
||||
if not row:
|
||||
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:
|
||||
out["latest_attempt"] = {
|
||||
"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"),
|
||||
}
|
||||
return out
|
||||
|
||||
@app.get("/api/v2/tasks/{task_id}/attempts")
|
||||
async def get_attempts(task_id: str, req: Request) -> dict[str, Any]:
|
||||
_require_token(req)
|
||||
row = db.get_task(task_id)
|
||||
if not row:
|
||||
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]:
|
||||
_require_token(req)
|
||||
row = db.get_task(task_id)
|
||||
if not row:
|
||||
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:
|
||||
_require_token(req)
|
||||
row = db.get_task(task_id)
|
||||
if not row:
|
||||
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]:
|
||||
_require_token(req)
|
||||
return db.list_queue()
|
||||
|
||||
return app
|
||||
Reference in New Issue
Block a user