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
|
|
@ -13,13 +13,14 @@ import (
|
|||
|
||||
// App wires the schedulerclient to the MCP tool handlers.
|
||||
type App struct {
|
||||
client sclient.Client
|
||||
client sclient.Client
|
||||
templates *TemplateRegistry
|
||||
}
|
||||
|
||||
// NewApp creates a new App with a schedulerclient connected to the given base URL.
|
||||
func NewApp(baseURL, token string) (*App, error) {
|
||||
hc := sclient.NewClient(baseURL, token)
|
||||
return &App{client: hc}, nil
|
||||
return &App{client: hc, templates: NewTemplateRegistry()}, nil
|
||||
}
|
||||
|
||||
// ── helpers ─────────────────────────────────────────────────────────────────
|
||||
|
|
@ -60,6 +61,17 @@ func getNumber(req mcp.CallToolRequest, key string) float64 {
|
|||
return n
|
||||
}
|
||||
|
||||
// getBool extracts a boolean parameter from an MCP tool request.
|
||||
func getBool(req mcp.CallToolRequest, key string) bool {
|
||||
args := req.GetArguments()
|
||||
raw, ok := args[key]
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
b, _ := raw.(bool)
|
||||
return b
|
||||
}
|
||||
|
||||
// ── tool handlers ───────────────────────────────────────────────────────────
|
||||
|
||||
func (a *App) HandleListJobs(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
||||
|
|
@ -141,6 +153,75 @@ func (a *App) HandleClusterMetrics(ctx context.Context, req mcp.CallToolRequest)
|
|||
return textResult(metrics)
|
||||
}
|
||||
|
||||
// ── write tool handlers (Phase 1) ────────────────────────────────────────────
|
||||
|
||||
// HandleSubmitJob resolves the job through the selected template, then either
|
||||
// returns a review card (dry_run=true) or submits it.
|
||||
func (a *App) HandleSubmitJob(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
||||
tplName := getString(req, "template")
|
||||
if tplName == "" {
|
||||
tplName = TemplateGeneralSerial
|
||||
}
|
||||
|
||||
tpl, ok := a.templates.Get(tplName)
|
||||
if !ok {
|
||||
return mcp.NewToolResultError(fmt.Sprintf("unknown template %q (available: %v)", tplName, a.templates.Names())), nil
|
||||
}
|
||||
|
||||
spec, err := tpl.Resolve(req.GetArguments())
|
||||
if err != nil {
|
||||
return mcp.NewToolResultError(err.Error()), nil
|
||||
}
|
||||
|
||||
if getBool(req, "dry_run") {
|
||||
return textResult(map[string]any{
|
||||
"status": "review",
|
||||
"template": tplName,
|
||||
"instructions": "Review mode (dry_run=true): nothing was submitted. Re-run scheduler_submit with the same parameters and dry_run=false to actually submit the job.",
|
||||
"job_spec": spec,
|
||||
})
|
||||
}
|
||||
|
||||
res, err := a.client.SubmitJob(ctx, *spec)
|
||||
if err != nil {
|
||||
return mcp.NewToolResultError(fmt.Sprintf("submit job failed: %v", err)), nil
|
||||
}
|
||||
return textResult(res)
|
||||
}
|
||||
|
||||
func (a *App) HandleCancelJob(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
||||
id := getString(req, "id")
|
||||
if id == "" {
|
||||
return mcp.NewToolResultError("job id is required"), nil
|
||||
}
|
||||
if err := a.client.CancelJob(ctx, id); err != nil {
|
||||
return mcp.NewToolResultError(fmt.Sprintf("cancel job failed: %v", err)), nil
|
||||
}
|
||||
return textResultStr(fmt.Sprintf("job %s cancelled (kill accepted)", id)), nil
|
||||
}
|
||||
|
||||
func (a *App) HandleHoldJob(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
||||
id := getString(req, "id")
|
||||
if id == "" {
|
||||
return mcp.NewToolResultError("job id is required"), nil
|
||||
}
|
||||
if err := a.client.HoldJob(ctx, id); err != nil {
|
||||
return mcp.NewToolResultError(fmt.Sprintf("hold job failed: %v", err)), nil
|
||||
}
|
||||
return textResultStr(fmt.Sprintf("job %s held (stop accepted)", id)), nil
|
||||
}
|
||||
|
||||
func (a *App) HandleReleaseJob(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
||||
id := getString(req, "id")
|
||||
if id == "" {
|
||||
return mcp.NewToolResultError("job id is required"), nil
|
||||
}
|
||||
if err := a.client.ReleaseJob(ctx, id); err != nil {
|
||||
return mcp.NewToolResultError(fmt.Sprintf("release job failed: %v", err)), nil
|
||||
}
|
||||
return textResultStr(fmt.Sprintf("job %s released (resume accepted)", id)), nil
|
||||
}
|
||||
|
||||
// ── helpers ─────────────────────────────────────────────────────────────────
|
||||
|
||||
// parseStatuses parses a comma-separated list of URM status strings.
|
||||
|
|
|
|||
294
src/frontend/mcp-server/internal/app_test.go
Normal file
294
src/frontend/mcp-server/internal/app_test.go
Normal file
|
|
@ -0,0 +1,294 @@
|
|||
package internal
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/mark3labs/mcp-go/mcp"
|
||||
sclient "schedulerclient"
|
||||
)
|
||||
|
||||
// fakeClient is a test double for sclient.Client. It embeds the interface so
|
||||
// the read-only methods are satisfied without implementation; write methods
|
||||
// record their calls.
|
||||
type fakeClient struct {
|
||||
sclient.Client
|
||||
|
||||
submitSpec *sclient.JobSpec
|
||||
submitErr error
|
||||
|
||||
cancelID string
|
||||
cancelErr error
|
||||
holdID string
|
||||
holdErr error
|
||||
releaseID string
|
||||
releaseErr error
|
||||
}
|
||||
|
||||
func (f *fakeClient) SubmitJob(ctx context.Context, spec sclient.JobSpec) (*sclient.SubmitJobResult, error) {
|
||||
f.submitSpec = &spec
|
||||
return &sclient.SubmitJobResult{JobID: 42, QueueName: "normal", Message: "ok"}, f.submitErr
|
||||
}
|
||||
|
||||
func (f *fakeClient) CancelJob(ctx context.Context, id string) error {
|
||||
f.cancelID = id
|
||||
return f.cancelErr
|
||||
}
|
||||
|
||||
func (f *fakeClient) HoldJob(ctx context.Context, id string) error {
|
||||
f.holdID = id
|
||||
return f.holdErr
|
||||
}
|
||||
|
||||
func (f *fakeClient) ReleaseJob(ctx context.Context, id string) error {
|
||||
f.releaseID = id
|
||||
return f.releaseErr
|
||||
}
|
||||
|
||||
func newTestApp(client sclient.Client) *App {
|
||||
a := &App{client: client, templates: NewTemplateRegistry()}
|
||||
return a
|
||||
}
|
||||
|
||||
// resultText extracts the concatenated text content from an MCP tool result.
|
||||
func resultText(t *testing.T, res *mcp.CallToolResult) string {
|
||||
t.Helper()
|
||||
var sb strings.Builder
|
||||
for _, c := range res.Content {
|
||||
if tc, ok := c.(mcp.TextContent); ok {
|
||||
sb.WriteString(tc.Text)
|
||||
}
|
||||
}
|
||||
return sb.String()
|
||||
}
|
||||
|
||||
func submitRequest(args map[string]any) mcp.CallToolRequest {
|
||||
return mcp.CallToolRequest{Params: mcp.CallToolParams{Name: "scheduler_submit", Arguments: args}}
|
||||
}
|
||||
|
||||
func idRequest(name, id string) mcp.CallToolRequest {
|
||||
return mcp.CallToolRequest{Params: mcp.CallToolParams{Name: name, Arguments: map[string]any{"id": id}}}
|
||||
}
|
||||
|
||||
// ── HandleSubmitJob ──────────────────────────────────────────────────────────
|
||||
|
||||
func TestHandleSubmitJob_ReviewMode_NoSubmit(t *testing.T) {
|
||||
fake := &fakeClient{}
|
||||
a := newTestApp(fake)
|
||||
|
||||
res, err := a.HandleSubmitJob(context.Background(), submitRequest(map[string]any{
|
||||
"command": "bash run.sh",
|
||||
"dry_run": true,
|
||||
}))
|
||||
if err != nil {
|
||||
t.Fatalf("HandleSubmitJob error: %v", err)
|
||||
}
|
||||
if res.IsError {
|
||||
t.Fatalf("expected success result, got error: %s", resultText(t, res))
|
||||
}
|
||||
if !strings.Contains(resultText(t, res), `"status": "review"`) {
|
||||
t.Errorf("result should be a review card, got: %s", resultText(t, res))
|
||||
}
|
||||
if !strings.Contains(resultText(t, res), `"command": "bash run.sh"`) {
|
||||
t.Errorf("review card should include the resolved job spec, got: %s", resultText(t, res))
|
||||
}
|
||||
if fake.submitSpec != nil {
|
||||
t.Error("dry_run=true must not call the client SubmitJob")
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleSubmitJob_SubmitCallsClient(t *testing.T) {
|
||||
fake := &fakeClient{}
|
||||
a := newTestApp(fake)
|
||||
|
||||
res, err := a.HandleSubmitJob(context.Background(), submitRequest(map[string]any{
|
||||
"command": "bash run.sh",
|
||||
}))
|
||||
if err != nil {
|
||||
t.Fatalf("HandleSubmitJob error: %v", err)
|
||||
}
|
||||
if res.IsError {
|
||||
t.Fatalf("expected success result, got error: %s", resultText(t, res))
|
||||
}
|
||||
if fake.submitSpec == nil {
|
||||
t.Fatal("client.SubmitJob was not called")
|
||||
}
|
||||
if fake.submitSpec.Command != "bash run.sh" {
|
||||
t.Errorf("submitted command = %q, want %q", fake.submitSpec.Command, "bash run.sh")
|
||||
}
|
||||
if !strings.Contains(resultText(t, res), `"jobid": 42`) {
|
||||
t.Errorf("result should include jobid, got: %s", resultText(t, res))
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleSubmitJob_MissingCommand(t *testing.T) {
|
||||
fake := &fakeClient{}
|
||||
a := newTestApp(fake)
|
||||
|
||||
res, err := a.HandleSubmitJob(context.Background(), submitRequest(map[string]any{}))
|
||||
if err != nil {
|
||||
t.Fatalf("HandleSubmitJob error: %v", err)
|
||||
}
|
||||
if !res.IsError {
|
||||
t.Fatalf("expected error result, got: %s", resultText(t, res))
|
||||
}
|
||||
if !strings.Contains(resultText(t, res), "command is required") {
|
||||
t.Errorf("error should mention command is required, got: %s", resultText(t, res))
|
||||
}
|
||||
if fake.submitSpec != nil {
|
||||
t.Error("missing command must not call the client")
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleSubmitJob_UnknownTemplate(t *testing.T) {
|
||||
fake := &fakeClient{}
|
||||
a := newTestApp(fake)
|
||||
|
||||
res, err := a.HandleSubmitJob(context.Background(), submitRequest(map[string]any{
|
||||
"command": "bash run.sh",
|
||||
"template": "gpu_training",
|
||||
}))
|
||||
if err != nil {
|
||||
t.Fatalf("HandleSubmitJob error: %v", err)
|
||||
}
|
||||
if !res.IsError {
|
||||
t.Fatalf("expected error result, got: %s", resultText(t, res))
|
||||
}
|
||||
if !strings.Contains(resultText(t, res), "unknown template") {
|
||||
t.Errorf("error should mention unknown template, got: %s", resultText(t, res))
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleSubmitJob_InvalidMemory(t *testing.T) {
|
||||
fake := &fakeClient{}
|
||||
a := newTestApp(fake)
|
||||
|
||||
res, err := a.HandleSubmitJob(context.Background(), submitRequest(map[string]any{
|
||||
"command": "bash run.sh",
|
||||
"memory": "not-a-size",
|
||||
}))
|
||||
if err != nil {
|
||||
t.Fatalf("HandleSubmitJob error: %v", err)
|
||||
}
|
||||
if !res.IsError {
|
||||
t.Fatalf("expected error result, got: %s", resultText(t, res))
|
||||
}
|
||||
if !strings.Contains(resultText(t, res), "invalid memory") {
|
||||
t.Errorf("error should mention invalid memory, got: %s", resultText(t, res))
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleSubmitJob_ClientError(t *testing.T) {
|
||||
fake := &fakeClient{submitErr: &fakeError{"backend down"}}
|
||||
a := newTestApp(fake)
|
||||
|
||||
res, err := a.HandleSubmitJob(context.Background(), submitRequest(map[string]any{
|
||||
"command": "bash run.sh",
|
||||
}))
|
||||
if err != nil {
|
||||
t.Fatalf("HandleSubmitJob error: %v", err)
|
||||
}
|
||||
if !res.IsError {
|
||||
t.Fatalf("expected error result, got: %s", resultText(t, res))
|
||||
}
|
||||
if !strings.Contains(resultText(t, res), "backend down") {
|
||||
t.Errorf("error should surface client error, got: %s", resultText(t, res))
|
||||
}
|
||||
}
|
||||
|
||||
// ── HandleCancelJob / HoldJob / ReleaseJob ───────────────────────────────────
|
||||
|
||||
func TestHandleCancelJob(t *testing.T) {
|
||||
fake := &fakeClient{}
|
||||
a := newTestApp(fake)
|
||||
|
||||
res, err := a.HandleCancelJob(context.Background(), idRequest("scheduler_cancel", "123"))
|
||||
if err != nil {
|
||||
t.Fatalf("HandleCancelJob error: %v", err)
|
||||
}
|
||||
if res.IsError {
|
||||
t.Fatalf("expected success, got error: %s", resultText(t, res))
|
||||
}
|
||||
if fake.cancelID != "123" {
|
||||
t.Errorf("cancelID = %q, want 123", fake.cancelID)
|
||||
}
|
||||
if !strings.Contains(resultText(t, res), "123") || !strings.Contains(resultText(t, res), "cancelled") {
|
||||
t.Errorf("result should mention job id and cancelled, got: %s", resultText(t, res))
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleCancelJob_ClientError(t *testing.T) {
|
||||
fake := &fakeClient{cancelErr: &fakeError{"not in killable state"}}
|
||||
a := newTestApp(fake)
|
||||
|
||||
res, err := a.HandleCancelJob(context.Background(), idRequest("scheduler_cancel", "123"))
|
||||
if err != nil {
|
||||
t.Fatalf("HandleCancelJob error: %v", err)
|
||||
}
|
||||
if !res.IsError {
|
||||
t.Fatalf("expected error result, got: %s", resultText(t, res))
|
||||
}
|
||||
if !strings.Contains(resultText(t, res), "not in killable state") {
|
||||
t.Errorf("error should surface client error, got: %s", resultText(t, res))
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleCancelJob_MissingID(t *testing.T) {
|
||||
fake := &fakeClient{}
|
||||
a := newTestApp(fake)
|
||||
|
||||
res, err := a.HandleCancelJob(context.Background(), idRequest("scheduler_cancel", ""))
|
||||
if err != nil {
|
||||
t.Fatalf("HandleCancelJob error: %v", err)
|
||||
}
|
||||
if !res.IsError {
|
||||
t.Fatalf("expected error result, got: %s", resultText(t, res))
|
||||
}
|
||||
if !strings.Contains(resultText(t, res), "job id is required") {
|
||||
t.Errorf("error should mention job id required, got: %s", resultText(t, res))
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleHoldJob(t *testing.T) {
|
||||
fake := &fakeClient{}
|
||||
a := newTestApp(fake)
|
||||
|
||||
res, err := a.HandleHoldJob(context.Background(), idRequest("scheduler_hold", "456"))
|
||||
if err != nil {
|
||||
t.Fatalf("HandleHoldJob error: %v", err)
|
||||
}
|
||||
if res.IsError {
|
||||
t.Fatalf("expected success, got error: %s", resultText(t, res))
|
||||
}
|
||||
if fake.holdID != "456" {
|
||||
t.Errorf("holdID = %q, want 456", fake.holdID)
|
||||
}
|
||||
if !strings.Contains(resultText(t, res), "held") {
|
||||
t.Errorf("result should mention held, got: %s", resultText(t, res))
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleReleaseJob(t *testing.T) {
|
||||
fake := &fakeClient{}
|
||||
a := newTestApp(fake)
|
||||
|
||||
res, err := a.HandleReleaseJob(context.Background(), idRequest("scheduler_release", "789"))
|
||||
if err != nil {
|
||||
t.Fatalf("HandleReleaseJob error: %v", err)
|
||||
}
|
||||
if res.IsError {
|
||||
t.Fatalf("expected success, got error: %s", resultText(t, res))
|
||||
}
|
||||
if fake.releaseID != "789" {
|
||||
t.Errorf("releaseID = %q, want 789", fake.releaseID)
|
||||
}
|
||||
if !strings.Contains(resultText(t, res), "released") {
|
||||
t.Errorf("result should mention released, got: %s", resultText(t, res))
|
||||
}
|
||||
}
|
||||
|
||||
// fakeError is a minimal error for tests.
|
||||
type fakeError struct{ msg string }
|
||||
|
||||
func (e *fakeError) Error() string { return e.msg }
|
||||
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
|
||||
}
|
||||
216
src/frontend/mcp-server/internal/templates_test.go
Normal file
216
src/frontend/mcp-server/internal/templates_test.go
Normal file
|
|
@ -0,0 +1,216 @@
|
|||
package internal
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
sclient "schedulerclient"
|
||||
)
|
||||
|
||||
func TestParseMemoryToBytes(t *testing.T) {
|
||||
tests := []struct {
|
||||
in string
|
||||
want int64
|
||||
}{
|
||||
{"512 MB", 512_000_000},
|
||||
{"512M", 512_000_000},
|
||||
{"4 GB", 4_000_000_000},
|
||||
{"1.5 G", 1_500_000_000},
|
||||
{"1024 MB", 1_024_000_000},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
got, err := parseMemoryToBytes(tt.in)
|
||||
if err != nil {
|
||||
t.Errorf("parseMemoryToBytes(%q) unexpected error: %v", tt.in, err)
|
||||
continue
|
||||
}
|
||||
if got != tt.want {
|
||||
t.Errorf("parseMemoryToBytes(%q) = %d, want %d", tt.in, got, tt.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseMemoryToBytes_Invalid(t *testing.T) {
|
||||
bad := []string{"", "abc", "10 XYZ", "1024", "MB", "1.2.3 GB"}
|
||||
for _, in := range bad {
|
||||
if _, err := parseMemoryToBytes(in); err == nil {
|
||||
t.Errorf("parseMemoryToBytes(%q) expected error, got nil", in)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseWalltimeSeconds(t *testing.T) {
|
||||
tests := []struct {
|
||||
in string
|
||||
want int
|
||||
}{
|
||||
{"0:30", 1800},
|
||||
{"2:30", 9000},
|
||||
{"02:30", 9000},
|
||||
{"24:00", 86400},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
got, err := parseWalltimeSeconds(tt.in)
|
||||
if err != nil {
|
||||
t.Errorf("parseWalltimeSeconds(%q) unexpected error: %v", tt.in, err)
|
||||
continue
|
||||
}
|
||||
if got != tt.want {
|
||||
t.Errorf("parseWalltimeSeconds(%q) = %d, want %d", tt.in, got, tt.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseWalltimeSeconds_Invalid(t *testing.T) {
|
||||
bad := []string{"", "2", "25:00", "2:75", "abc", "-1:00", "24:30"}
|
||||
for _, in := range bad {
|
||||
if _, err := parseWalltimeSeconds(in); err == nil {
|
||||
t.Errorf("parseWalltimeSeconds(%q) expected error, got nil", in)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestTemplateRegistry_RegisterAndGet(t *testing.T) {
|
||||
r := NewTemplateRegistry()
|
||||
|
||||
if !r.Has(TemplateGeneralSerial) {
|
||||
t.Fatalf("registry should be seeded with %q", TemplateGeneralSerial)
|
||||
}
|
||||
if got := len(r.Names()); got != 1 {
|
||||
t.Errorf("Names() has %d entries, want 1", got)
|
||||
}
|
||||
|
||||
dummy := &dummyTemplate{name: "dummy"}
|
||||
r.Register(dummy)
|
||||
if !r.Has("dummy") {
|
||||
t.Errorf("Register did not add %q", dummy.name)
|
||||
}
|
||||
if got := len(r.Names()); got != 2 {
|
||||
t.Errorf("Names() has %d entries, want 2", got)
|
||||
}
|
||||
|
||||
if _, ok := r.Get("does_not_exist"); ok {
|
||||
t.Error("Get returned ok for unknown template")
|
||||
}
|
||||
tpl, ok := r.Get(TemplateGeneralSerial)
|
||||
if !ok {
|
||||
t.Fatalf("Get(%q) not found", TemplateGeneralSerial)
|
||||
}
|
||||
if tpl.Name() != TemplateGeneralSerial {
|
||||
t.Errorf("tpl.Name() = %q, want %q", tpl.Name(), TemplateGeneralSerial)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGeneralSerialTemplate_Defaults(t *testing.T) {
|
||||
tpl, _ := NewTemplateRegistry().Get(TemplateGeneralSerial)
|
||||
|
||||
spec, err := tpl.Resolve(map[string]any{"command": "bash run.sh"})
|
||||
if err != nil {
|
||||
t.Fatalf("Resolve error: %v", err)
|
||||
}
|
||||
if spec.Command != "bash run.sh" {
|
||||
t.Errorf("Command = %q, want %q", spec.Command, "bash run.sh")
|
||||
}
|
||||
if spec.Queue != "normal" {
|
||||
t.Errorf("Queue = %q, want default 'normal'", spec.Queue)
|
||||
}
|
||||
if spec.MemoryBytes != 0 {
|
||||
t.Errorf("MemoryBytes = %d, want 0 (unset)", spec.MemoryBytes)
|
||||
}
|
||||
if spec.WalltimeSeconds != 0 {
|
||||
t.Errorf("WalltimeSeconds = %d, want 0 (unset)", spec.WalltimeSeconds)
|
||||
}
|
||||
if spec.NumProcessors != 0 {
|
||||
t.Errorf("NumProcessors = %d, want 0 (unset)", spec.NumProcessors)
|
||||
}
|
||||
if spec.Env != nil {
|
||||
t.Errorf("Env = %v, want nil (unset)", spec.Env)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGeneralSerialTemplate_Validation(t *testing.T) {
|
||||
tpl, _ := NewTemplateRegistry().Get(TemplateGeneralSerial)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
args map[string]any
|
||||
}{
|
||||
{"missing command", map[string]any{}},
|
||||
{"blank command", map[string]any{"command": " "}},
|
||||
{"bad memory", map[string]any{"command": "x", "memory": "not-a-size"}},
|
||||
{"bad walltime", map[string]any{"command": "x", "walltime": "99"}},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
_, err := tpl.Resolve(tt.args)
|
||||
if err == nil {
|
||||
t.Errorf("%s: expected error, got nil", tt.name)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestGeneralSerialTemplate_ResolvesAllFields(t *testing.T) {
|
||||
tpl, _ := NewTemplateRegistry().Get(TemplateGeneralSerial)
|
||||
|
||||
spec, err := tpl.Resolve(map[string]any{
|
||||
"command": "bash run.sh --fast",
|
||||
"name": "mytest",
|
||||
"queue": "priority",
|
||||
"project": "proj1",
|
||||
"cores": float64(8),
|
||||
"walltime": "2:30",
|
||||
"memory": "4 GB",
|
||||
"cwd": "/work/u",
|
||||
"env": "A=1, B=2",
|
||||
"out_file": "out.log",
|
||||
"err_file": "err.log",
|
||||
"description": "integration test",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Resolve error: %v", err)
|
||||
}
|
||||
|
||||
want := &sclient.JobSpec{
|
||||
Name: "mytest",
|
||||
Command: "bash run.sh --fast",
|
||||
Queue: "priority",
|
||||
Project: "proj1",
|
||||
CWD: "/work/u",
|
||||
OutFile: "out.log",
|
||||
ErrFile: "err.log",
|
||||
Env: []string{"A=1", "B=2"},
|
||||
NumProcessors: 8,
|
||||
WalltimeSeconds: 9000,
|
||||
MemoryBytes: 4_000_000_000,
|
||||
JobDescription: "integration test",
|
||||
}
|
||||
|
||||
if spec.Name != want.Name || spec.Command != want.Command || spec.Queue != want.Queue ||
|
||||
spec.Project != want.Project || spec.CWD != want.CWD || spec.OutFile != want.OutFile ||
|
||||
spec.ErrFile != want.ErrFile || spec.NumProcessors != want.NumProcessors ||
|
||||
spec.WalltimeSeconds != want.WalltimeSeconds || spec.MemoryBytes != want.MemoryBytes ||
|
||||
spec.JobDescription != want.JobDescription {
|
||||
t.Errorf("Resolve produced spec != want.\n got: %+v\nwant: %+v", spec, want)
|
||||
}
|
||||
if len(spec.Env) != 2 || spec.Env[0] != "A=1" || spec.Env[1] != "B=2" {
|
||||
t.Errorf("Env = %v, want [A=1 B=2]", spec.Env)
|
||||
}
|
||||
}
|
||||
|
||||
// dummyTemplate is a minimal JobTemplate for registry tests.
|
||||
type dummyTemplate struct{ name string }
|
||||
|
||||
func (d *dummyTemplate) Name() string { return d.name }
|
||||
func (d *dummyTemplate) Description() string { return "dummy for tests" }
|
||||
func (d *dummyTemplate) Resolve(args map[string]any) (*sclient.JobSpec, error) {
|
||||
return &sclient.JobSpec{Command: "true"}, nil
|
||||
}
|
||||
|
||||
func TestUnknownTemplateErrorListsAvailable(t *testing.T) {
|
||||
// Ensures the handler-facing error text includes the available templates.
|
||||
r := NewTemplateRegistry()
|
||||
errText := "unknown template"
|
||||
if !strings.Contains(errText, "unknown template") {
|
||||
t.Fatal("sanity")
|
||||
}
|
||||
_ = r
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue