forked from lei.chen/myxce
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(
|
||||
"scheduler-agent",
|
||||
"0.1.0",
|
||||
"0.2.0",
|
||||
mcps.WithResourceCapabilities(true, true),
|
||||
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 ──
|
||||
s.AddTool(mcp.NewTool("scheduler_list_jobs",
|
||||
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.WithNumber("limit", mcp.Description("Max results to return")),
|
||||
mcp.WithString("page_token", mcp.Description("Pagination token from a prior response")),
|
||||
mcp.WithReadOnlyHintAnnotation(true),
|
||||
mcp.WithDestructiveHintAnnotation(false),
|
||||
), app.HandleListJobs)
|
||||
|
||||
// ── scheduler_get_job ──
|
||||
|
|
@ -51,16 +56,22 @@ func main() {
|
|||
mcp.Description("Job ID"),
|
||||
mcp.Required(),
|
||||
),
|
||||
mcp.WithReadOnlyHintAnnotation(true),
|
||||
mcp.WithDestructiveHintAnnotation(false),
|
||||
), app.HandleGetJob)
|
||||
|
||||
// ── scheduler_list_queues ──
|
||||
s.AddTool(mcp.NewTool("scheduler_list_queues",
|
||||
mcp.WithDescription("List all queues/partitions and their current state (active, inactive, draining)."),
|
||||
mcp.WithReadOnlyHintAnnotation(true),
|
||||
mcp.WithDestructiveHintAnnotation(false),
|
||||
), app.HandleListQueues)
|
||||
|
||||
// ── scheduler_list_nodes ──
|
||||
s.AddTool(mcp.NewTool("scheduler_list_nodes",
|
||||
mcp.WithDescription("List all compute nodes with their resource usage and status."),
|
||||
mcp.WithReadOnlyHintAnnotation(true),
|
||||
mcp.WithDestructiveHintAnnotation(false),
|
||||
), app.HandleListNodes)
|
||||
|
||||
// ── 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("offset_bytes", mcp.Description("Byte offset to start reading from")),
|
||||
mcp.WithReadOnlyHintAnnotation(true),
|
||||
mcp.WithDestructiveHintAnnotation(false),
|
||||
), app.HandleJobLogs)
|
||||
|
||||
// ── scheduler_cluster_metrics ──
|
||||
s.AddTool(mcp.NewTool("scheduler_cluster_metrics",
|
||||
mcp.WithDescription("Get a snapshot of overall cluster resource usage and utilization."),
|
||||
mcp.WithReadOnlyHintAnnotation(true),
|
||||
mcp.WithDestructiveHintAnnotation(false),
|
||||
), 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)
|
||||
if err := mcps.ServeStdio(s); err != nil {
|
||||
log.Fatalf("server error: %v", err)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
|
@ -4,7 +4,7 @@ import "context"
|
|||
|
||||
// ── 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).
|
||||
//
|
||||
// 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(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
|
||||
// (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
|
||||
|
||||
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