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
|
|
@ -29,11 +29,14 @@ func main() {
|
||||||
|
|
||||||
s := mcps.NewMCPServer(
|
s := mcps.NewMCPServer(
|
||||||
"scheduler-agent",
|
"scheduler-agent",
|
||||||
"0.1.0",
|
"0.2.0",
|
||||||
mcps.WithResourceCapabilities(true, true),
|
mcps.WithResourceCapabilities(true, true),
|
||||||
mcps.WithLogging(),
|
mcps.WithLogging(),
|
||||||
)
|
)
|
||||||
|
|
||||||
|
// Read-only tools are annotated as such so the MCP client does not prompt
|
||||||
|
// for confirmation on them (NewTool defaults to destructive=true).
|
||||||
|
|
||||||
// ── scheduler_list_jobs ──
|
// ── scheduler_list_jobs ──
|
||||||
s.AddTool(mcp.NewTool("scheduler_list_jobs",
|
s.AddTool(mcp.NewTool("scheduler_list_jobs",
|
||||||
mcp.WithDescription("List jobs in the cluster. Supports filtering by queue, owner, status, and pagination."),
|
mcp.WithDescription("List jobs in the cluster. Supports filtering by queue, owner, status, and pagination."),
|
||||||
|
|
@ -42,6 +45,8 @@ func main() {
|
||||||
mcp.WithString("status", mcp.Description("Filter by status: pending|running|succeeded|failed|cancelled|suspended|unknown")),
|
mcp.WithString("status", mcp.Description("Filter by status: pending|running|succeeded|failed|cancelled|suspended|unknown")),
|
||||||
mcp.WithNumber("limit", mcp.Description("Max results to return")),
|
mcp.WithNumber("limit", mcp.Description("Max results to return")),
|
||||||
mcp.WithString("page_token", mcp.Description("Pagination token from a prior response")),
|
mcp.WithString("page_token", mcp.Description("Pagination token from a prior response")),
|
||||||
|
mcp.WithReadOnlyHintAnnotation(true),
|
||||||
|
mcp.WithDestructiveHintAnnotation(false),
|
||||||
), app.HandleListJobs)
|
), app.HandleListJobs)
|
||||||
|
|
||||||
// ── scheduler_get_job ──
|
// ── scheduler_get_job ──
|
||||||
|
|
@ -51,16 +56,22 @@ func main() {
|
||||||
mcp.Description("Job ID"),
|
mcp.Description("Job ID"),
|
||||||
mcp.Required(),
|
mcp.Required(),
|
||||||
),
|
),
|
||||||
|
mcp.WithReadOnlyHintAnnotation(true),
|
||||||
|
mcp.WithDestructiveHintAnnotation(false),
|
||||||
), app.HandleGetJob)
|
), app.HandleGetJob)
|
||||||
|
|
||||||
// ── scheduler_list_queues ──
|
// ── scheduler_list_queues ──
|
||||||
s.AddTool(mcp.NewTool("scheduler_list_queues",
|
s.AddTool(mcp.NewTool("scheduler_list_queues",
|
||||||
mcp.WithDescription("List all queues/partitions and their current state (active, inactive, draining)."),
|
mcp.WithDescription("List all queues/partitions and their current state (active, inactive, draining)."),
|
||||||
|
mcp.WithReadOnlyHintAnnotation(true),
|
||||||
|
mcp.WithDestructiveHintAnnotation(false),
|
||||||
), app.HandleListQueues)
|
), app.HandleListQueues)
|
||||||
|
|
||||||
// ── scheduler_list_nodes ──
|
// ── scheduler_list_nodes ──
|
||||||
s.AddTool(mcp.NewTool("scheduler_list_nodes",
|
s.AddTool(mcp.NewTool("scheduler_list_nodes",
|
||||||
mcp.WithDescription("List all compute nodes with their resource usage and status."),
|
mcp.WithDescription("List all compute nodes with their resource usage and status."),
|
||||||
|
mcp.WithReadOnlyHintAnnotation(true),
|
||||||
|
mcp.WithDestructiveHintAnnotation(false),
|
||||||
), app.HandleListNodes)
|
), app.HandleListNodes)
|
||||||
|
|
||||||
// ── scheduler_job_logs ──
|
// ── scheduler_job_logs ──
|
||||||
|
|
@ -72,13 +83,68 @@ func main() {
|
||||||
),
|
),
|
||||||
mcp.WithNumber("tail", mcp.Description("Number of lines from the end to return (0 = all)")),
|
mcp.WithNumber("tail", mcp.Description("Number of lines from the end to return (0 = all)")),
|
||||||
mcp.WithNumber("offset_bytes", mcp.Description("Byte offset to start reading from")),
|
mcp.WithNumber("offset_bytes", mcp.Description("Byte offset to start reading from")),
|
||||||
|
mcp.WithReadOnlyHintAnnotation(true),
|
||||||
|
mcp.WithDestructiveHintAnnotation(false),
|
||||||
), app.HandleJobLogs)
|
), app.HandleJobLogs)
|
||||||
|
|
||||||
// ── scheduler_cluster_metrics ──
|
// ── scheduler_cluster_metrics ──
|
||||||
s.AddTool(mcp.NewTool("scheduler_cluster_metrics",
|
s.AddTool(mcp.NewTool("scheduler_cluster_metrics",
|
||||||
mcp.WithDescription("Get a snapshot of overall cluster resource usage and utilization."),
|
mcp.WithDescription("Get a snapshot of overall cluster resource usage and utilization."),
|
||||||
|
mcp.WithReadOnlyHintAnnotation(true),
|
||||||
|
mcp.WithDestructiveHintAnnotation(false),
|
||||||
), app.HandleClusterMetrics)
|
), app.HandleClusterMetrics)
|
||||||
|
|
||||||
|
// ── write tools (Phase 1) ──
|
||||||
|
// Write tools are annotated destructive so the MCP client (e.g. Claude
|
||||||
|
// Code) prompts the user for confirmation before executing them — the
|
||||||
|
// design doc's "写操作默认 require confirm" at the layer MCP supports it.
|
||||||
|
|
||||||
|
// ── scheduler_submit ──
|
||||||
|
s.AddTool(mcp.NewTool("scheduler_submit",
|
||||||
|
mcp.WithDescription("Submit a job. REVIEW REQUIRED: call with dry_run=true first to inspect the resolved job spec, then re-call with dry_run=false to actually submit."),
|
||||||
|
mcp.WithString("template", mcp.Description("Job template (default 'general_serial' 通用串行脚本)."),
|
||||||
|
mcp.Enum(internal.TemplateGeneralSerial), mcp.DefaultString(internal.TemplateGeneralSerial)),
|
||||||
|
mcp.WithBoolean("dry_run", mcp.Description("Review mode: validate and return the resolved job spec without submitting."), mcp.DefaultBool(false)),
|
||||||
|
mcp.WithString("command", mcp.Description("Shell command to run (required)."), mcp.Required()),
|
||||||
|
mcp.WithString("name", mcp.Description("Job name; array syntax 'myjob[1-10]' supported.")),
|
||||||
|
mcp.WithString("queue", mcp.Description("Queue (default 'normal').")),
|
||||||
|
mcp.WithString("project", mcp.Description("Project name.")),
|
||||||
|
mcp.WithNumber("cores", mcp.Description("Number of processors.")),
|
||||||
|
mcp.WithString("walltime", mcp.Description("Walltime limit 'HH:MM', e.g. '02:30'.")),
|
||||||
|
mcp.WithString("memory", mcp.Description("Memory limit, e.g. '4 GB', '2048 MB'.")),
|
||||||
|
mcp.WithString("cwd", mcp.Description("Working directory.")),
|
||||||
|
mcp.WithString("env", mcp.Description("Comma-separated KEY=VALUE environment variables.")),
|
||||||
|
mcp.WithString("out_file", mcp.Description("Standard output file.")),
|
||||||
|
mcp.WithString("err_file", mcp.Description("Standard error file.")),
|
||||||
|
mcp.WithString("description", mcp.Description("Job description.")),
|
||||||
|
mcp.WithReadOnlyHintAnnotation(false),
|
||||||
|
mcp.WithDestructiveHintAnnotation(true),
|
||||||
|
), app.HandleSubmitJob)
|
||||||
|
|
||||||
|
// ── scheduler_cancel ──
|
||||||
|
s.AddTool(mcp.NewTool("scheduler_cancel",
|
||||||
|
mcp.WithDescription("Cancel (kill) a job by ID. Destructive — you will be asked to confirm before execution."),
|
||||||
|
mcp.WithString("id", mcp.Description("Job ID to cancel"), mcp.Required()),
|
||||||
|
mcp.WithReadOnlyHintAnnotation(false),
|
||||||
|
mcp.WithDestructiveHintAnnotation(true),
|
||||||
|
), app.HandleCancelJob)
|
||||||
|
|
||||||
|
// ── scheduler_hold ──
|
||||||
|
s.AddTool(mcp.NewTool("scheduler_hold",
|
||||||
|
mcp.WithDescription("Hold (suspend/stop) a job by ID. Destructive — you will be asked to confirm before execution."),
|
||||||
|
mcp.WithString("id", mcp.Description("Job ID to hold"), mcp.Required()),
|
||||||
|
mcp.WithReadOnlyHintAnnotation(false),
|
||||||
|
mcp.WithDestructiveHintAnnotation(true),
|
||||||
|
), app.HandleHoldJob)
|
||||||
|
|
||||||
|
// ── scheduler_release ──
|
||||||
|
s.AddTool(mcp.NewTool("scheduler_release",
|
||||||
|
mcp.WithDescription("Release (resume) a held job by ID. Destructive — you will be asked to confirm before execution."),
|
||||||
|
mcp.WithString("id", mcp.Description("Job ID to release"), mcp.Required()),
|
||||||
|
mcp.WithReadOnlyHintAnnotation(false),
|
||||||
|
mcp.WithDestructiveHintAnnotation(true),
|
||||||
|
), app.HandleReleaseJob)
|
||||||
|
|
||||||
log.Printf("Starting scheduler-agent MCP server (base-url: %s)", baseURL)
|
log.Printf("Starting scheduler-agent MCP server (base-url: %s)", baseURL)
|
||||||
if err := mcps.ServeStdio(s); err != nil {
|
if err := mcps.ServeStdio(s); err != nil {
|
||||||
log.Fatalf("server error: %v", err)
|
log.Fatalf("server error: %v", err)
|
||||||
|
|
|
||||||
|
|
@ -13,13 +13,14 @@ import (
|
||||||
|
|
||||||
// App wires the schedulerclient to the MCP tool handlers.
|
// App wires the schedulerclient to the MCP tool handlers.
|
||||||
type App struct {
|
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.
|
// NewApp creates a new App with a schedulerclient connected to the given base URL.
|
||||||
func NewApp(baseURL, token string) (*App, error) {
|
func NewApp(baseURL, token string) (*App, error) {
|
||||||
hc := sclient.NewClient(baseURL, token)
|
hc := sclient.NewClient(baseURL, token)
|
||||||
return &App{client: hc}, nil
|
return &App{client: hc, templates: NewTemplateRegistry()}, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
// ── helpers ─────────────────────────────────────────────────────────────────
|
// ── helpers ─────────────────────────────────────────────────────────────────
|
||||||
|
|
@ -60,6 +61,17 @@ func getNumber(req mcp.CallToolRequest, key string) float64 {
|
||||||
return n
|
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 ───────────────────────────────────────────────────────────
|
// ── tool handlers ───────────────────────────────────────────────────────────
|
||||||
|
|
||||||
func (a *App) HandleListJobs(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
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)
|
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 ─────────────────────────────────────────────────────────────────
|
// ── helpers ─────────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
// parseStatuses parses a comma-separated list of URM status strings.
|
// 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
|
||||||
|
}
|
||||||
|
|
@ -4,7 +4,7 @@ import "context"
|
||||||
|
|
||||||
// ── Client interface ────────────────────────────────────────────────────────
|
// ── Client interface ────────────────────────────────────────────────────────
|
||||||
|
|
||||||
// Client provides read-only access to the scheduler cluster through a
|
// Client provides read and write access to the scheduler cluster through a
|
||||||
// scheduler-agnostic unified resource model (URM).
|
// scheduler-agnostic unified resource model (URM).
|
||||||
//
|
//
|
||||||
// Implementations talk to the apiserver's REST API and authenticate via
|
// Implementations talk to the apiserver's REST API and authenticate via
|
||||||
|
|
@ -36,4 +36,18 @@ type Client interface {
|
||||||
|
|
||||||
// GetClusterMetrics returns a snapshot of overall cluster resource usage.
|
// GetClusterMetrics returns a snapshot of overall cluster resource usage.
|
||||||
GetClusterMetrics(ctx context.Context) (*Metrics, error)
|
GetClusterMetrics(ctx context.Context) (*Metrics, error)
|
||||||
|
|
||||||
|
// ── Write operations (Phase 1) ──
|
||||||
|
|
||||||
|
// SubmitJob submits a job and returns its assigned ID.
|
||||||
|
SubmitJob(ctx context.Context, spec JobSpec) (*SubmitJobResult, error)
|
||||||
|
|
||||||
|
// CancelJob kills a job (action "kill", default SIGTERM).
|
||||||
|
CancelJob(ctx context.Context, id string) error
|
||||||
|
|
||||||
|
// HoldJob suspends a job (action "stop", default SIGTSTP).
|
||||||
|
HoldJob(ctx context.Context, id string) error
|
||||||
|
|
||||||
|
// ReleaseJob resumes a suspended job (action "resume", default SIGCONT).
|
||||||
|
ReleaseJob(ctx context.Context, id string) error
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -1,8 +1,9 @@
|
||||||
// Package schedulerclient provides a read-only Go client for the XCE scheduler API,
|
// Package schedulerclient provides a Go client for the XCE scheduler API,
|
||||||
// exposing a unified resource model (URM) that maps the cluster's native concepts
|
// exposing a unified resource model (URM) that maps the cluster's native concepts
|
||||||
// (jobs, queues, nodes) into a scheduler-agnostic representation.
|
// (jobs, queues, nodes) into a scheduler-agnostic representation.
|
||||||
//
|
//
|
||||||
// This is the library that the Phase 0 MCP Server will use to talk to the apiserver.
|
// This is the library the scheduler-agent MCP Server uses to talk to the
|
||||||
|
// apiserver: Phase 0 read-only operations, Phase 1 write operations.
|
||||||
package schedulerclient
|
package schedulerclient
|
||||||
|
|
||||||
import "time"
|
import "time"
|
||||||
|
|
|
||||||
224
src/frontend/schedulerclient/write.go
Normal file
224
src/frontend/schedulerclient/write.go
Normal file
|
|
@ -0,0 +1,224 @@
|
||||||
|
package schedulerclient
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"net/http"
|
||||||
|
"strconv"
|
||||||
|
"strings"
|
||||||
|
)
|
||||||
|
|
||||||
|
// ── Write types (Phase 1) ────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
// JobSpec is a unified job submission specification. Canonical units are
|
||||||
|
// WalltimeSeconds (seconds) and MemoryBytes (bytes); the HTTP layer converts
|
||||||
|
// them to the apiserver's wire format ("HH:MM" and "<N> MB").
|
||||||
|
type JobSpec struct {
|
||||||
|
Name string `json:"name,omitempty"` // job_name; array syntax "myjob[1-10]" is passed through
|
||||||
|
Command string `json:"command"` // required
|
||||||
|
Queue string `json:"queue,omitempty"` // empty → apiserver defaults to "normal"
|
||||||
|
Project string `json:"project,omitempty"` // project_name
|
||||||
|
CWD string `json:"cwd,omitempty"` // working directory
|
||||||
|
OutFile string `json:"out_file,omitempty"` // standard output file
|
||||||
|
ErrFile string `json:"err_file,omitempty"` // standard error file
|
||||||
|
Env []string `json:"env,omitempty"` // "KEY=VALUE"
|
||||||
|
NumProcessors int `json:"num_processors,omitempty"` // number of processors
|
||||||
|
MaxNumProcessors int `json:"maxnum_processors,omitempty"`
|
||||||
|
GpuReq string `json:"gpureq,omitempty"` // gpureq resource string
|
||||||
|
WalltimeSeconds int `json:"walltime_seconds,omitempty"` // run_limit → "HH:MM"
|
||||||
|
MemoryBytes int64 `json:"memory_bytes,omitempty"` // mem_limit → "<MB> MB"
|
||||||
|
Exclusive bool `json:"exclusive,omitempty"`
|
||||||
|
Rerunable bool `json:"rerunable,omitempty"`
|
||||||
|
DependCond string `json:"depend_cond,omitempty"` // dependency condition string
|
||||||
|
JobDescription string `json:"description,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// SubmitJobResult is the apiserver's response to a successful job submission.
|
||||||
|
type SubmitJobResult struct {
|
||||||
|
JobID int64 `json:"jobid"` // NB: lowercase field name in the apiserver response
|
||||||
|
QueueName string `json:"queue_name"`
|
||||||
|
Message string `json:"message"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// JobControlRequest is the body sent to POST /xce/v1/jobs/:jobids/control.
|
||||||
|
type JobControlRequest struct {
|
||||||
|
Action string `json:"action"` // kill | stop | resume | checkpoint
|
||||||
|
Signal string `json:"signal,omitempty"`
|
||||||
|
Comment string `json:"comment,omitempty"`
|
||||||
|
ForceKill bool `json:"force_kill,omitempty"`
|
||||||
|
Method string `json:"method,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// ControlResult is the apiserver's per-item control response.
|
||||||
|
type ControlResult struct {
|
||||||
|
Status string `json:"status"` // accepted | failed | partially_accepted | no_items_processed
|
||||||
|
Accepted []ControlItem `json:"accepted"`
|
||||||
|
Failed []ControlItem `json:"failed"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// ControlItem reports the outcome for one job (or array element).
|
||||||
|
type ControlItem struct {
|
||||||
|
UniqueJobID int64 `json:"unique_job_id"`
|
||||||
|
JobID int64 `json:"job_id"`
|
||||||
|
ArrayIdx int `json:"array_idx"`
|
||||||
|
Msg string `json:"msg"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// apiJobSubmitRequest mirrors the subset of apiserver models.JobSubmitRequest
|
||||||
|
// JSON tags we send. Zero-valued fields are omitted (omitempty) so the
|
||||||
|
// apiserver applies its own defaults (e.g. queue_name → "normal").
|
||||||
|
type apiJobSubmitRequest struct {
|
||||||
|
JobName string `json:"job_name,omitempty"`
|
||||||
|
Command string `json:"command"`
|
||||||
|
QueueName string `json:"queue_name,omitempty"`
|
||||||
|
ProjectName string `json:"project_name,omitempty"`
|
||||||
|
CWD string `json:"cwd,omitempty"`
|
||||||
|
OutFile string `json:"out_file,omitempty"`
|
||||||
|
ErrFile string `json:"err_file,omitempty"`
|
||||||
|
Env []string `json:"env,omitempty"`
|
||||||
|
NumProcessors int `json:"num_processors,omitempty"`
|
||||||
|
MaxNumProcessors int `json:"maxnum_processors,omitempty"`
|
||||||
|
Gpureq string `json:"gpureq,omitempty"`
|
||||||
|
Exclusive bool `json:"exclusive,omitempty"`
|
||||||
|
Rerunable bool `json:"rerunable,omitempty"`
|
||||||
|
DependCond string `json:"depend_cond,omitempty"`
|
||||||
|
JobDescription string `json:"job_description,omitempty"`
|
||||||
|
RunLimit string `json:"run_limit,omitempty"` // "HH:MM" (bare ints are MINUTES — never sent)
|
||||||
|
MemLimit string `json:"mem_limit,omitempty"` // "<N> MB"
|
||||||
|
}
|
||||||
|
|
||||||
|
// jobSpecToAPI converts a JobSpec into the wire request. run_limit is rendered
|
||||||
|
// as "H:MM" and mem_limit as "<MB> MB" using the package's decimal
|
||||||
|
// mbToBytes convention (1 MB = 1_000_000 bytes), so "4 GB" → 4e9 bytes →
|
||||||
|
// "4000 MB" → CustomLimit parses cleanly.
|
||||||
|
func jobSpecToAPI(s JobSpec) apiJobSubmitRequest {
|
||||||
|
req := apiJobSubmitRequest{
|
||||||
|
JobName: s.Name,
|
||||||
|
Command: s.Command,
|
||||||
|
QueueName: s.Queue,
|
||||||
|
ProjectName: s.Project,
|
||||||
|
CWD: s.CWD,
|
||||||
|
OutFile: s.OutFile,
|
||||||
|
ErrFile: s.ErrFile,
|
||||||
|
Env: s.Env,
|
||||||
|
NumProcessors: s.NumProcessors,
|
||||||
|
MaxNumProcessors: s.MaxNumProcessors,
|
||||||
|
Gpureq: s.GpuReq,
|
||||||
|
Exclusive: s.Exclusive,
|
||||||
|
Rerunable: s.Rerunable,
|
||||||
|
DependCond: s.DependCond,
|
||||||
|
JobDescription: s.JobDescription,
|
||||||
|
}
|
||||||
|
if s.WalltimeSeconds > 0 {
|
||||||
|
req.RunLimit = formatWalltime(s.WalltimeSeconds)
|
||||||
|
}
|
||||||
|
if s.MemoryBytes > 0 {
|
||||||
|
req.MemLimit = formatMemory(s.MemoryBytes)
|
||||||
|
}
|
||||||
|
return req
|
||||||
|
}
|
||||||
|
|
||||||
|
// formatWalltime renders seconds as "H:MM", matching the apiserver's
|
||||||
|
// CustomInt.String() output. Sub-minute seconds are truncated (walltime
|
||||||
|
// granularity is minutes).
|
||||||
|
func formatWalltime(secs int) string {
|
||||||
|
return fmt.Sprintf("%d:%02d", secs/3600, (secs%3600)/60)
|
||||||
|
}
|
||||||
|
|
||||||
|
// formatMemory renders bytes as "<N> MB" using the package's decimal
|
||||||
|
// mbToBytes convention. Rounds up so sub-MB amounts don't become "0 MB".
|
||||||
|
func formatMemory(bytes int64) string {
|
||||||
|
mb := bytes / mbToBytes
|
||||||
|
if bytes%mbToBytes != 0 {
|
||||||
|
mb++
|
||||||
|
}
|
||||||
|
return strconv.Itoa(int(mb)) + " MB"
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── shared POST helper ───────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
// doPost performs an authenticated POST with a JSON body and decodes the
|
||||||
|
// envelope. expectedCode is the envelope code on success (201 create, 202
|
||||||
|
// control). NB: ReturnSuccess always uses HTTP 200; the logical status is the
|
||||||
|
// envelope's code field. ReturnError sets both HTTP status and envelope code
|
||||||
|
// to the error code.
|
||||||
|
func (hc *HTTPClient) doPost(ctx context.Context, url string, body any, expectedCode int) (*apiEnvelope, error) {
|
||||||
|
req := hc.rc.R().SetContext(ctx).
|
||||||
|
SetHeader("Authorization", hc.token).
|
||||||
|
SetHeader("Content-Type", "application/json")
|
||||||
|
|
||||||
|
var env apiEnvelope
|
||||||
|
resp, err := req.SetBody(body).SetResult(&env).Post(url)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("post %s: %w", url, err)
|
||||||
|
}
|
||||||
|
if resp.StatusCode() != http.StatusOK {
|
||||||
|
// resty does not decode non-2xx bodies into SetResult, so read the raw
|
||||||
|
// body (the apiserver's error envelope) for a useful message.
|
||||||
|
return nil, fmt.Errorf("post %s: HTTP %d: %s", url, resp.StatusCode(), resp.String())
|
||||||
|
}
|
||||||
|
if env.Code != expectedCode {
|
||||||
|
return nil, fmt.Errorf("post %s: api error %d: %s", url, env.Code, env.Msg)
|
||||||
|
}
|
||||||
|
return &env, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── write method implementations ─────────────────────────────────────────────
|
||||||
|
|
||||||
|
func (hc *HTTPClient) SubmitJob(ctx context.Context, spec JobSpec) (*SubmitJobResult, error) {
|
||||||
|
if strings.TrimSpace(spec.Command) == "" {
|
||||||
|
return nil, fmt.Errorf("submit job: command is required")
|
||||||
|
}
|
||||||
|
|
||||||
|
url := hc.baseURL + "/xce/v1/jobs/create"
|
||||||
|
env, err := hc.doPost(ctx, url, jobSpecToAPI(spec), http.StatusCreated)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("submit job: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
var res SubmitJobResult
|
||||||
|
if err := json.Unmarshal(env.Data, &res); err != nil {
|
||||||
|
return nil, fmt.Errorf("submit job: parse result: %w", err)
|
||||||
|
}
|
||||||
|
return &res, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// controlJob sends a job-control action and returns a descriptive error unless
|
||||||
|
// every requested item was accepted.
|
||||||
|
func (hc *HTTPClient) controlJob(ctx context.Context, id, action string) error {
|
||||||
|
if strings.TrimSpace(id) == "" {
|
||||||
|
return fmt.Errorf("%s job: job id is required", action)
|
||||||
|
}
|
||||||
|
|
||||||
|
url := hc.baseURL + "/xce/v1/jobs/" + id + "/control"
|
||||||
|
env, err := hc.doPost(ctx, url, JobControlRequest{Action: action}, http.StatusAccepted)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("%s job %s: %w", action, id, err)
|
||||||
|
}
|
||||||
|
|
||||||
|
var res ControlResult
|
||||||
|
if err := json.Unmarshal(env.Data, &res); err != nil {
|
||||||
|
return fmt.Errorf("%s job %s: parse result: %w", action, id, err)
|
||||||
|
}
|
||||||
|
switch res.Status {
|
||||||
|
case "accepted":
|
||||||
|
return nil
|
||||||
|
case "no_items_processed":
|
||||||
|
return fmt.Errorf("%s job %s: no matching job found", action, id)
|
||||||
|
default: // "failed" or "partially_accepted"
|
||||||
|
return fmt.Errorf("%s job %s: status=%s failed=%+v", action, id, res.Status, res.Failed)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (hc *HTTPClient) CancelJob(ctx context.Context, id string) error {
|
||||||
|
return hc.controlJob(ctx, id, "kill")
|
||||||
|
}
|
||||||
|
|
||||||
|
func (hc *HTTPClient) HoldJob(ctx context.Context, id string) error {
|
||||||
|
return hc.controlJob(ctx, id, "stop")
|
||||||
|
}
|
||||||
|
|
||||||
|
func (hc *HTTPClient) ReleaseJob(ctx context.Context, id string) error {
|
||||||
|
return hc.controlJob(ctx, id, "resume")
|
||||||
|
}
|
||||||
318
src/frontend/schedulerclient/write_test.go
Normal file
318
src/frontend/schedulerclient/write_test.go
Normal file
|
|
@ -0,0 +1,318 @@
|
||||||
|
package schedulerclient
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"io"
|
||||||
|
"net/http"
|
||||||
|
"net/http/httptest"
|
||||||
|
"strings"
|
||||||
|
"sync"
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
|
// ── test helper ──────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
// newTestServer spins up an httptest server that records the last request and
|
||||||
|
// returns the fixed JSON body. The returned *HTTPClient points at it.
|
||||||
|
func newTestServer(t *testing.T, handler http.HandlerFunc) (*HTTPClient, func() (method, path, auth, body string)) {
|
||||||
|
t.Helper()
|
||||||
|
|
||||||
|
var (
|
||||||
|
mu sync.Mutex
|
||||||
|
lastReq struct{ method, path, auth, body string }
|
||||||
|
)
|
||||||
|
|
||||||
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
b, _ := io.ReadAll(r.Body)
|
||||||
|
mu.Lock()
|
||||||
|
lastReq = struct{ method, path, auth, body string }{r.Method, r.URL.Path, r.Header.Get("Authorization"), string(b)}
|
||||||
|
mu.Unlock()
|
||||||
|
handler(w, r)
|
||||||
|
}))
|
||||||
|
t.Cleanup(srv.Close)
|
||||||
|
|
||||||
|
capture := func() (method, path, auth, body string) {
|
||||||
|
mu.Lock()
|
||||||
|
defer mu.Unlock()
|
||||||
|
return lastReq.method, lastReq.path, lastReq.auth, lastReq.body
|
||||||
|
}
|
||||||
|
|
||||||
|
return NewClient(srv.URL, "test-token"), capture
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── SubmitJob ────────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
func TestSubmitJob_Success(t *testing.T) {
|
||||||
|
hc, capture := newTestServer(t, func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
io.WriteString(w, `{"code":201,"msg":"success","data":{"jobid":42,"queue_name":"normal","message":"ok"}}`)
|
||||||
|
})
|
||||||
|
|
||||||
|
res, err := hc.SubmitJob(context.Background(), JobSpec{Command: "bash run.sh"})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("SubmitJob returned error: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
method, path, auth, body := capture()
|
||||||
|
if method != http.MethodPost {
|
||||||
|
t.Errorf("method = %q, want POST", method)
|
||||||
|
}
|
||||||
|
if path != "/xce/v1/jobs/create" {
|
||||||
|
t.Errorf("path = %q, want /xce/v1/jobs/create", path)
|
||||||
|
}
|
||||||
|
if auth != "test-token" {
|
||||||
|
t.Errorf("authorization = %q, want test-token", auth)
|
||||||
|
}
|
||||||
|
if !strings.Contains(body, `"command":"bash run.sh"`) {
|
||||||
|
t.Errorf("body = %s, want command field", body)
|
||||||
|
}
|
||||||
|
if res.JobID != 42 {
|
||||||
|
t.Errorf("JobID = %d, want 42", res.JobID)
|
||||||
|
}
|
||||||
|
if res.QueueName != "normal" {
|
||||||
|
t.Errorf("QueueName = %q, want normal", res.QueueName)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSubmitJob_ServerRejectsValidation(t *testing.T) {
|
||||||
|
hc, _ := newTestServer(t, func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
w.WriteHeader(http.StatusBadRequest)
|
||||||
|
io.WriteString(w, `{"code":400,"msg":"invalid queue name","details":null}`)
|
||||||
|
})
|
||||||
|
|
||||||
|
_, err := hc.SubmitJob(context.Background(), JobSpec{Command: "bash run.sh"})
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected error for 400 response")
|
||||||
|
}
|
||||||
|
if !strings.Contains(err.Error(), "invalid queue name") {
|
||||||
|
t.Errorf("error = %q, want to mention 'invalid queue name'", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSubmitJob_Forbidden(t *testing.T) {
|
||||||
|
hc, _ := newTestServer(t, func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
w.WriteHeader(http.StatusForbidden)
|
||||||
|
io.WriteString(w, `{"code":403,"msg":"access denied","details":null}`)
|
||||||
|
})
|
||||||
|
|
||||||
|
_, err := hc.SubmitJob(context.Background(), JobSpec{Command: "bash run.sh"})
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected error for 403 response")
|
||||||
|
}
|
||||||
|
if !strings.Contains(err.Error(), "access denied") {
|
||||||
|
t.Errorf("error = %q, want to mention 'access denied'", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSubmitJob_ServerError(t *testing.T) {
|
||||||
|
hc, _ := newTestServer(t, func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
w.WriteHeader(http.StatusInternalServerError)
|
||||||
|
io.WriteString(w, `{"code":500,"msg":"internal error","details":null}`)
|
||||||
|
})
|
||||||
|
|
||||||
|
_, err := hc.SubmitJob(context.Background(), JobSpec{Command: "bash run.sh"})
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected error for 500 response")
|
||||||
|
}
|
||||||
|
if !strings.Contains(err.Error(), "internal error") {
|
||||||
|
t.Errorf("error = %q, want to mention 'internal error'", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSubmitJob_EmptyCommand(t *testing.T) {
|
||||||
|
// No server needed: validation must happen before any network call.
|
||||||
|
hc, _ := newTestServer(t, func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
t.Error("server should not be called when command is empty")
|
||||||
|
})
|
||||||
|
|
||||||
|
_, err := hc.SubmitJob(context.Background(), JobSpec{Command: " "})
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected error for empty command")
|
||||||
|
}
|
||||||
|
if !strings.Contains(err.Error(), "command is required") {
|
||||||
|
t.Errorf("error = %q, want 'command is required'", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSubmitJob_WireBodyMapping(t *testing.T) {
|
||||||
|
hc, capture := newTestServer(t, func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
io.WriteString(w, `{"code":201,"msg":"success","data":{"jobid":7,"queue_name":"normal","message":"ok"}}`)
|
||||||
|
})
|
||||||
|
|
||||||
|
spec := JobSpec{
|
||||||
|
Name: "myjob",
|
||||||
|
Command: "bash run.sh",
|
||||||
|
Queue: "normal",
|
||||||
|
Project: "default",
|
||||||
|
CWD: "/home/u/work",
|
||||||
|
OutFile: "out.log",
|
||||||
|
ErrFile: "err.log",
|
||||||
|
Env: []string{"FOO=bar", "BAZ=qux"},
|
||||||
|
NumProcessors: 4,
|
||||||
|
MaxNumProcessors: 8,
|
||||||
|
GpuReq: "",
|
||||||
|
WalltimeSeconds: 9000, // → "2:30"
|
||||||
|
MemoryBytes: 2_000_000_000, // → "2000 MB"
|
||||||
|
Exclusive: true,
|
||||||
|
Rerunable: false,
|
||||||
|
DependCond: "",
|
||||||
|
JobDescription: "my desc",
|
||||||
|
}
|
||||||
|
|
||||||
|
if _, err := hc.SubmitJob(context.Background(), spec); err != nil {
|
||||||
|
t.Fatalf("SubmitJob returned error: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
_, _, _, body := capture()
|
||||||
|
var got map[string]any
|
||||||
|
if err := json.Unmarshal([]byte(body), &got); err != nil {
|
||||||
|
t.Fatalf("captured body is not valid JSON: %v\nbody=%s", err, body)
|
||||||
|
}
|
||||||
|
|
||||||
|
want := map[string]any{
|
||||||
|
"job_name": "myjob",
|
||||||
|
"command": "bash run.sh",
|
||||||
|
"queue_name": "normal",
|
||||||
|
"project_name": "default",
|
||||||
|
"cwd": "/home/u/work",
|
||||||
|
"out_file": "out.log",
|
||||||
|
"err_file": "err.log",
|
||||||
|
"env": []any{"FOO=bar", "BAZ=qux"},
|
||||||
|
"num_processors": float64(4),
|
||||||
|
"maxnum_processors": float64(8),
|
||||||
|
"exclusive": true,
|
||||||
|
"run_limit": "2:30",
|
||||||
|
"mem_limit": "2000 MB",
|
||||||
|
"job_description": "my desc",
|
||||||
|
}
|
||||||
|
for k, wantVal := range want {
|
||||||
|
if got[k] == nil {
|
||||||
|
t.Errorf("body missing key %q", k)
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
// compare as JSON-normalized values
|
||||||
|
gotJSON, _ := json.Marshal(got[k])
|
||||||
|
wantJSON, _ := json.Marshal(wantVal)
|
||||||
|
if string(gotJSON) != string(wantJSON) {
|
||||||
|
t.Errorf("body[%q] = %s, want %s", k, gotJSON, wantJSON)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// gpureq and rerunable are zero-valued → omitted via omitempty.
|
||||||
|
for _, absent := range []string{"gpureq", "rerunable", "depend_cond"} {
|
||||||
|
if _, ok := got[absent]; ok {
|
||||||
|
t.Errorf("body should omit empty field %q, got %s", absent, got[absent])
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Cancel / Hold / Release ──────────────────────────────────────────────────
|
||||||
|
|
||||||
|
func TestCancelJob_Success(t *testing.T) {
|
||||||
|
hc, capture := newTestServer(t, func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
io.WriteString(w, `{"code":202,"msg":"accepted","data":{"status":"accepted","accepted":[{"unique_job_id":123,"job_id":123,"array_idx":0,"msg":"kill signal TERM sent"}],"failed":[]}}`)
|
||||||
|
})
|
||||||
|
|
||||||
|
if err := hc.CancelJob(context.Background(), "123"); err != nil {
|
||||||
|
t.Fatalf("CancelJob returned error: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
method, path, _, body := capture()
|
||||||
|
if method != http.MethodPost {
|
||||||
|
t.Errorf("method = %q, want POST", method)
|
||||||
|
}
|
||||||
|
if path != "/xce/v1/jobs/123/control" {
|
||||||
|
t.Errorf("path = %q, want /xce/v1/jobs/123/control", path)
|
||||||
|
}
|
||||||
|
if strings.TrimSpace(body) != `{"action":"kill"}` {
|
||||||
|
t.Errorf("body = %s, want {\"action\":\"kill\"}", body)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestCancelJob_StatusFailed(t *testing.T) {
|
||||||
|
hc, _ := newTestServer(t, func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
io.WriteString(w, `{"code":202,"msg":"failed","data":{"status":"failed","accepted":[],"failed":[{"unique_job_id":123,"job_id":123,"array_idx":0,"msg":"job 123 is in state DONE, cannot be killed"}]}}`)
|
||||||
|
})
|
||||||
|
|
||||||
|
err := hc.CancelJob(context.Background(), "123")
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected error when control status is failed")
|
||||||
|
}
|
||||||
|
if !strings.Contains(err.Error(), "123") {
|
||||||
|
t.Errorf("error = %q, want to mention job id 123", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestHoldJob_SendsStop(t *testing.T) {
|
||||||
|
hc, capture := newTestServer(t, func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
io.WriteString(w, `{"code":202,"msg":"accepted","data":{"status":"accepted","accepted":[],"failed":[]}}`)
|
||||||
|
})
|
||||||
|
|
||||||
|
if err := hc.HoldJob(context.Background(), "456"); err != nil {
|
||||||
|
t.Fatalf("HoldJob returned error: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
_, _, _, body := capture()
|
||||||
|
if strings.TrimSpace(body) != `{"action":"stop"}` {
|
||||||
|
t.Errorf("body = %s, want {\"action\":\"stop\"}", body)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestReleaseJob_SendsResume(t *testing.T) {
|
||||||
|
hc, capture := newTestServer(t, func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
io.WriteString(w, `{"code":202,"msg":"accepted","data":{"status":"accepted","accepted":[],"failed":[]}}`)
|
||||||
|
})
|
||||||
|
|
||||||
|
if err := hc.ReleaseJob(context.Background(), "789"); err != nil {
|
||||||
|
t.Fatalf("ReleaseJob returned error: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
_, _, _, body := capture()
|
||||||
|
if strings.TrimSpace(body) != `{"action":"resume"}` {
|
||||||
|
t.Errorf("body = %s, want {\"action\":\"resume\"}", body)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── conversion helpers ───────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
func TestFormatWalltimeSeconds(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
in int
|
||||||
|
want string
|
||||||
|
}{
|
||||||
|
{0, "0:00"},
|
||||||
|
{3600, "1:00"},
|
||||||
|
{9000, "2:30"},
|
||||||
|
{3661, "1:01"},
|
||||||
|
}
|
||||||
|
for _, tt := range tests {
|
||||||
|
if got := formatWalltime(tt.in); got != tt.want {
|
||||||
|
t.Errorf("formatWalltime(%d) = %q, want %q", tt.in, got, tt.want)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFormatMemoryBytes(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
in int64
|
||||||
|
want string
|
||||||
|
}{
|
||||||
|
{1_000_000, "1 MB"},
|
||||||
|
{2_000_000_000, "2000 MB"},
|
||||||
|
{1_500_000, "2 MB"}, // rounds up to avoid 1.5 → 1 MB truncation
|
||||||
|
}
|
||||||
|
for _, tt := range tests {
|
||||||
|
if got := formatMemory(tt.in); got != tt.want {
|
||||||
|
t.Errorf("formatMemory(%d) = %q, want %q", tt.in, got, tt.want)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
Loading…
Add table
Add a link
Reference in a new issue