feat(scheduler-agent): Phase 1 write operations — submit/cancel/hold/release + templates
Add write operations to the scheduler-agent MCP server (Phase 1): - schedulerclient: SubmitJob/CancelJob/HoldJob/ReleaseJob + JobSpec/ SubmitJobResult/JobControlRequest/ControlResult types + doPost helper (HTTP 200 + envelope-code dual check). TDD: 12 new tests. - mcp-server: 4 write tools (scheduler_submit/cancel/hold/release) with destructive annotations; scheduler_submit has dry_run review mode + general_serial (通用串行脚本) job template registry. Fix the 6 read-only tools' misleading destructive annotations. TDD: 14 new tests. Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
lei.chen
parent
2e299f4a76
commit
35416d5cf4
9 changed files with 1447 additions and 6 deletions
227
src/frontend/mcp-server/internal/templates.go
Normal file
227
src/frontend/mcp-server/internal/templates.go
Normal file
|
|
@ -0,0 +1,227 @@
|
|||
package internal
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"regexp"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
sclient "schedulerclient"
|
||||
)
|
||||
|
||||
// ── job templates ────────────────────────────────────────────────────────────
|
||||
|
||||
// TemplateGeneralSerial is the "通用串行脚本" template — submit a shell
|
||||
// script/command to a queue with optional resource limits.
|
||||
const TemplateGeneralSerial = "general_serial"
|
||||
|
||||
// JobTemplate resolves raw LLM tool arguments into a validated JobSpec.
|
||||
// Implementations read only the keys they need, so future templates
|
||||
// (GPU training, job arrays, dependency chains) add new params without
|
||||
// touching tool registration in cmd/main.go — they register here instead.
|
||||
type JobTemplate interface {
|
||||
Name() string
|
||||
Description() string
|
||||
Resolve(args map[string]any) (*sclient.JobSpec, error)
|
||||
}
|
||||
|
||||
// TemplateRegistry holds all registered job templates.
|
||||
type TemplateRegistry struct {
|
||||
m map[string]JobTemplate
|
||||
}
|
||||
|
||||
// NewTemplateRegistry returns a registry seeded with the built-in templates.
|
||||
func NewTemplateRegistry() *TemplateRegistry {
|
||||
r := &TemplateRegistry{m: make(map[string]JobTemplate)}
|
||||
r.Register(generalSerialTemplate{})
|
||||
return r
|
||||
}
|
||||
|
||||
// Register adds a template, overwriting any existing template with the same name.
|
||||
func (r *TemplateRegistry) Register(t JobTemplate) {
|
||||
r.m[t.Name()] = t
|
||||
}
|
||||
|
||||
// Get returns the template with the given name.
|
||||
func (r *TemplateRegistry) Get(name string) (JobTemplate, bool) {
|
||||
t, ok := r.m[name]
|
||||
return t, ok
|
||||
}
|
||||
|
||||
// Has reports whether a template with the given name is registered.
|
||||
func (r *TemplateRegistry) Has(name string) bool {
|
||||
_, ok := r.m[name]
|
||||
return ok
|
||||
}
|
||||
|
||||
// Names returns the sorted list of registered template names.
|
||||
func (r *TemplateRegistry) Names() []string {
|
||||
names := make([]string, 0, len(r.m))
|
||||
for n := range r.m {
|
||||
names = append(names, n)
|
||||
}
|
||||
// deterministic order for error messages and tool enums
|
||||
sortStrings(names)
|
||||
return names
|
||||
}
|
||||
|
||||
func sortStrings(s []string) {
|
||||
for i := 1; i < len(s); i++ {
|
||||
for j := i; j > 0 && s[j-1] > s[j]; j-- {
|
||||
s[j-1], s[j] = s[j], s[j-1]
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── general_serial ───────────────────────────────────────────────────────────
|
||||
|
||||
// generalSerialTemplate implements the "通用串行脚本" template: a single shell
|
||||
// command/script run on one or more processors in a queue.
|
||||
type generalSerialTemplate struct{}
|
||||
|
||||
func (generalSerialTemplate) Name() string { return TemplateGeneralSerial }
|
||||
|
||||
func (generalSerialTemplate) Description() string {
|
||||
return "提交一个 shell 脚本/命令（通用串行作业）：command 必填，可选 queue/cores/walltime/memory/cwd/env 等"
|
||||
}
|
||||
|
||||
func (generalSerialTemplate) Resolve(args map[string]any) (*sclient.JobSpec, error) {
|
||||
command := argString(args, "command")
|
||||
if strings.TrimSpace(command) == "" {
|
||||
return nil, errors.New("scheduler_submit: command is required")
|
||||
}
|
||||
|
||||
spec := &sclient.JobSpec{
|
||||
Name: argString(args, "name"),
|
||||
Command: command,
|
||||
Queue: firstNonEmpty(argString(args, "queue"), "normal"),
|
||||
Project: argString(args, "project"),
|
||||
CWD: argString(args, "cwd"),
|
||||
OutFile: argString(args, "out_file"),
|
||||
ErrFile: argString(args, "err_file"),
|
||||
JobDescription: argString(args, "description"),
|
||||
}
|
||||
|
||||
if cores := argNumber(args, "cores"); cores > 0 {
|
||||
spec.NumProcessors = int(cores)
|
||||
}
|
||||
|
||||
if mem := argString(args, "memory"); mem != "" {
|
||||
bytes, err := parseMemoryToBytes(mem)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("scheduler_submit: invalid memory %q: %w", mem, err)
|
||||
}
|
||||
spec.MemoryBytes = bytes
|
||||
}
|
||||
|
||||
if wt := argString(args, "walltime"); wt != "" {
|
||||
secs, err := parseWalltimeSeconds(wt)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("scheduler_submit: invalid walltime %q: %w", wt, err)
|
||||
}
|
||||
spec.WalltimeSeconds = secs
|
||||
}
|
||||
|
||||
if env := argString(args, "env"); env != "" {
|
||||
for _, kv := range strings.Split(env, ",") {
|
||||
kv = strings.TrimSpace(kv)
|
||||
if kv != "" {
|
||||
spec.Env = append(spec.Env, kv)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return spec, nil
|
||||
}
|
||||
|
||||
// ── argument helpers ─────────────────────────────────────────────────────────
|
||||
|
||||
func argString(args map[string]any, key string) string {
|
||||
raw, ok := args[key]
|
||||
if !ok {
|
||||
return ""
|
||||
}
|
||||
s, _ := raw.(string)
|
||||
return s
|
||||
}
|
||||
|
||||
func argNumber(args map[string]any, key string) float64 {
|
||||
raw, ok := args[key]
|
||||
if !ok {
|
||||
return 0
|
||||
}
|
||||
n, _ := raw.(float64)
|
||||
return n
|
||||
}
|
||||
|
||||
func firstNonEmpty(vals ...string) string {
|
||||
for _, v := range vals {
|
||||
if v != "" {
|
||||
return v
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// ── unit parsers (pure functions, unit-testable) ─────────────────────────────
|
||||
|
||||
// memoryRe matches "<number> <unit>", "<number><unit>", and "<number> <unit>"
|
||||
// with optional whitespace between value and unit.
|
||||
var memoryRe = regexp.MustCompile(`^([0-9]*\.?[0-9]+)\s*([A-Za-z]+)?$`)
|
||||
|
||||
// parseMemoryToBytes parses "512 MB" / "512M" / "4 GB" / "1.5 G" into bytes
|
||||
// using the package's decimal convention (1 MB = 1_000_000 bytes, matching
|
||||
// schedulerclient.mbToBytes) so "4 GB" round-trips to "4000 MB" on the wire.
|
||||
// A bare number without a unit is rejected (unit is required).
|
||||
func parseMemoryToBytes(s string) (int64, error) {
|
||||
s = strings.TrimSpace(s)
|
||||
if s == "" {
|
||||
return 0, errors.New("empty memory value")
|
||||
}
|
||||
m := memoryRe.FindStringSubmatch(s)
|
||||
if m == nil {
|
||||
return 0, fmt.Errorf("expected '<number> <unit>'")
|
||||
}
|
||||
value, err := strconv.ParseFloat(m[1], 64)
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("invalid number %q", m[1])
|
||||
}
|
||||
|
||||
unit := strings.ToUpper(m[2])
|
||||
var mult float64
|
||||
switch unit {
|
||||
case "MB", "M":
|
||||
mult = 1
|
||||
case "GB", "G":
|
||||
mult = 1_000
|
||||
case "TB", "T":
|
||||
mult = 1_000_000
|
||||
default:
|
||||
return 0, fmt.Errorf("unsupported unit %q", m[2])
|
||||
}
|
||||
return int64(value * mult * 1_000_000), nil
|
||||
}
|
||||
|
||||
// parseWalltimeSeconds parses "HH:MM" or "H:MM" into seconds. Bare integers
|
||||
// are rejected — the apiserver's run_limit treats bare ints as MINUTES, which
|
||||
// would silently change semantics, so the template requires explicit H:MM.
|
||||
func parseWalltimeSeconds(s string) (int, error) {
|
||||
s = strings.TrimSpace(s)
|
||||
parts := strings.Split(s, ":")
|
||||
if len(parts) != 2 {
|
||||
return 0, fmt.Errorf("expected 'HH:MM', got %q", s)
|
||||
}
|
||||
hours, err := strconv.Atoi(parts[0])
|
||||
if err != nil || hours < 0 || hours > 24 {
|
||||
return 0, fmt.Errorf("hour must be between 0 and 24, got %q", parts[0])
|
||||
}
|
||||
minutes, err := strconv.Atoi(parts[1])
|
||||
if err != nil || minutes < 0 || minutes > 59 {
|
||||
return 0, fmt.Errorf("minute must be between 0 and 59, got %q", parts[1])
|
||||
}
|
||||
if hours == 24 && minutes != 0 {
|
||||
return 0, fmt.Errorf("24:00 is the max; got %q", s)
|
||||
}
|
||||
return hours*3600 + minutes*60, nil
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue