V3.5 支持Advanced模式,用户提供python命令

This commit is contained in:
2026-01-04 15:28:41 +08:00
parent e3dcfe526f
commit 5d0d849c28
22 changed files with 1355 additions and 34 deletions

View File

@@ -9,8 +9,9 @@ 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 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
@@ -316,23 +317,52 @@ def create_app(config_path: str) -> FastAPI:
async def submit_task(req: Request) -> dict[str, Any]:
subject = _auth(req)
body = (await req.body()).decode("utf-8")
obj = yaml.safe_load(body) or {}
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 = JobSpec.from_dict(obj)
spec = parse_taskspec(obj)
except Exception as e:
raise HTTPException(status_code=400, detail=f"invalid jobspec: {e!r}")
# v3.0 path policy:
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/...
root = ray_cfg.shared_root.rstrip("/")
common_prefix = f"{root}/common/"
user_prefix = f"{root}/users/{str(subject['user_id']).strip()}/"
user_prefix = f"{root}/users/{user_id}/"
common_datasets_prefix = f"{common_prefix}datasets/"
user_datasets_prefix = f"{user_prefix}datasets/"
@@ -360,15 +390,17 @@ def create_app(config_path: str) -> FastAPI:
detail=f"model_id local path must start with {common_models_prefix} or {user_models_prefix}",
)
task_id = new_task_id(spec.workload, user_id=str(subject["user_id"]))
db.create_task_v25(
task_id=task_id,
user_id=str(subject["user_id"]),
workload=spec.workload,
jobspec_yaml=body,
nnodes=spec.nnodes,
n_gpus_per_node=spec.n_gpus_per_node,
)
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}")
@@ -428,7 +460,7 @@ def create_app(config_path: str) -> FastAPI:
obj = yaml.safe_load(raw) or {}
if not isinstance(obj, dict):
raise ValueError("jobspec must be a YAML mapping")
spec = JobSpec.from_dict(obj)
spec = parse_taskspec(obj)
resolved = spec.to_public_dict()
attempts = db.list_attempts(task_id)
if attempts: