787 lines
34 KiB
Python
787 lines
34 KiB
Python
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
|