feat(mcp-server): add scheduler_checkpoint + scheduler_restart tools
Expose two backend capabilities MCP previously omitted: - scheduler_checkpoint: control action "checkpoint" (CRIU) - scheduler_restart: POST /xce/v1/jobs/:id/restart (re-queues as a new job, resuming from the last checkpoint if one exists; returns new jobid) schedulerclient gains CheckpointJob/RestartJob (RestartJobResult, JobRestartRequest); Client interface updated. Both tools marked destructive (confirm required). tools/list now reports 17 tools. Unit tests added/updated in both modules — all pass. Backend endpoints are existing apiserver routes; no backend changes in this commit. Also syncs left_scheduler_work.md §八.4 (checkpoint/restart done; the only remaining MCP tool gap is scheduler_modify, blocked on a controller job-field-update opcode that does not exist yet). Co-Authored-By: Claude Code <noreply@anthropic.com>
This commit is contained in:
lei.chen
parent
05ac57b183
commit
6020015fa7
7 changed files with 279 additions and 10 deletions
|
|
@ -25,7 +25,7 @@ const serverInstructions = "XCE scheduler agent — cluster status, not an incid
|
|||
"root causes, do NOT reference \"known bugs\", and do NOT propose diagnostic or repair actions (checking logs, " +
|
||||
"restarting daemons, opening tickets, ssh'ing to nodes, etc.) unless the user explicitly asks for help or " +
|
||||
"remediation. If you believe action may be warranted, say so in at most one line and let the user decide.\n\n" +
|
||||
"Write tools (scheduler_submit/hold/release/cancel) are destructive and require the user's explicit intent."
|
||||
"Write tools (scheduler_submit/cancel/hold/release/checkpoint/restart) are destructive and require the user's explicit intent."
|
||||
|
||||
func main() {
|
||||
baseURL := os.Getenv("XCE_API_URL")
|
||||
|
|
@ -221,6 +221,22 @@ func main() {
|
|||
mcp.WithDestructiveHintAnnotation(true),
|
||||
), app.HandleReleaseJob)
|
||||
|
||||
// ── scheduler_checkpoint ──
|
||||
s.AddTool(mcp.NewTool("scheduler_checkpoint",
|
||||
mcp.WithDescription("Checkpoint a running job by ID (action 'checkpoint', CRIU). Destructive — you will be asked to confirm before execution."),
|
||||
mcp.WithString("id", mcp.Description("Job ID to checkpoint"), mcp.Required()),
|
||||
mcp.WithReadOnlyHintAnnotation(false),
|
||||
mcp.WithDestructiveHintAnnotation(true),
|
||||
), app.HandleCheckpointJob)
|
||||
|
||||
// ── scheduler_restart ──
|
||||
s.AddTool(mcp.NewTool("scheduler_restart",
|
||||
mcp.WithDescription("Restart (re-queue) a finished job as a new job by ID, resuming from its last checkpoint if one exists. Returns the new job ID. Destructive — you will be asked to confirm before execution."),
|
||||
mcp.WithString("id", mcp.Description("Job ID to restart"), mcp.Required()),
|
||||
mcp.WithReadOnlyHintAnnotation(false),
|
||||
mcp.WithDestructiveHintAnnotation(true),
|
||||
), app.HandleRestartJob)
|
||||
|
||||
log.Printf("Starting scheduler-agent MCP server (base-url: %s)", baseURL)
|
||||
if err := mcps.ServeStdio(s); err != nil {
|
||||
log.Fatalf("server error: %v", err)
|
||||
|
|
|
|||
|
|
@ -288,6 +288,29 @@ func (a *App) HandleReleaseJob(ctx context.Context, req mcp.CallToolRequest) (*m
|
|||
return textResultStr(fmt.Sprintf("job %s released (resume accepted)", id)), nil
|
||||
}
|
||||
|
||||
func (a *App) HandleCheckpointJob(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.CheckpointJob(ctx, id); err != nil {
|
||||
return mcp.NewToolResultError(fmt.Sprintf("checkpoint job failed: %v", err)), nil
|
||||
}
|
||||
return textResultStr(fmt.Sprintf("job %s checkpoint accepted", id)), nil
|
||||
}
|
||||
|
||||
func (a *App) HandleRestartJob(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
||||
id := getString(req, "id")
|
||||
if id == "" {
|
||||
return mcp.NewToolResultError("job id is required"), nil
|
||||
}
|
||||
res, err := a.client.RestartJob(ctx, id)
|
||||
if err != nil {
|
||||
return mcp.NewToolResultError(fmt.Sprintf("restart job failed: %v", err)), nil
|
||||
}
|
||||
return textResult(res)
|
||||
}
|
||||
|
||||
// ── helpers ─────────────────────────────────────────────────────────────────
|
||||
|
||||
// parseStatuses parses a comma-separated list of URM status strings.
|
||||
|
|
|
|||
|
|
@ -18,12 +18,16 @@ type fakeClient struct {
|
|||
submitSpec *sclient.JobSpec
|
||||
submitErr error
|
||||
|
||||
cancelID string
|
||||
cancelErr error
|
||||
holdID string
|
||||
holdErr error
|
||||
releaseID string
|
||||
releaseErr error
|
||||
cancelID string
|
||||
cancelErr error
|
||||
holdID string
|
||||
holdErr error
|
||||
releaseID string
|
||||
releaseErr error
|
||||
checkpointID string
|
||||
checkpointErr error
|
||||
restartID string
|
||||
restartErr error
|
||||
}
|
||||
|
||||
func (f *fakeClient) SubmitJob(ctx context.Context, spec sclient.JobSpec) (*sclient.SubmitJobResult, error) {
|
||||
|
|
@ -46,6 +50,19 @@ func (f *fakeClient) ReleaseJob(ctx context.Context, id string) error {
|
|||
return f.releaseErr
|
||||
}
|
||||
|
||||
func (f *fakeClient) CheckpointJob(ctx context.Context, id string) error {
|
||||
f.checkpointID = id
|
||||
return f.checkpointErr
|
||||
}
|
||||
|
||||
func (f *fakeClient) RestartJob(ctx context.Context, id string) (*sclient.RestartJobResult, error) {
|
||||
f.restartID = id
|
||||
if f.restartErr != nil {
|
||||
return nil, f.restartErr
|
||||
}
|
||||
return &sclient.RestartJobResult{JobID: 999, QueueName: "normal", Message: "ok"}, nil
|
||||
}
|
||||
|
||||
func newTestApp(client sclient.Client) *App {
|
||||
a := &App{client: client, templates: NewTemplateRegistry()}
|
||||
return a
|
||||
|
|
@ -288,6 +305,92 @@ func TestHandleReleaseJob(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
func TestHandleCheckpointJob(t *testing.T) {
|
||||
fake := &fakeClient{}
|
||||
a := newTestApp(fake)
|
||||
|
||||
res, err := a.HandleCheckpointJob(context.Background(), idRequest("scheduler_checkpoint", "101"))
|
||||
if err != nil {
|
||||
t.Fatalf("HandleCheckpointJob error: %v", err)
|
||||
}
|
||||
if res.IsError {
|
||||
t.Fatalf("expected success, got error: %s", resultText(t, res))
|
||||
}
|
||||
if fake.checkpointID != "101" {
|
||||
t.Errorf("checkpointID = %q, want 101", fake.checkpointID)
|
||||
}
|
||||
if !strings.Contains(resultText(t, res), "checkpoint accepted") {
|
||||
t.Errorf("result should mention checkpoint accepted, got: %s", resultText(t, res))
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleCheckpointJob_MissingID(t *testing.T) {
|
||||
fake := &fakeClient{}
|
||||
a := newTestApp(fake)
|
||||
|
||||
res, err := a.HandleCheckpointJob(context.Background(), idRequest("scheduler_checkpoint", ""))
|
||||
if err != nil {
|
||||
t.Fatalf("HandleCheckpointJob error: %v", err)
|
||||
}
|
||||
if !res.IsError {
|
||||
t.Fatal("expected error for missing id")
|
||||
}
|
||||
if fake.checkpointID != "" {
|
||||
t.Errorf("client should not be called, got checkpointID = %q", fake.checkpointID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleRestartJob(t *testing.T) {
|
||||
fake := &fakeClient{}
|
||||
a := newTestApp(fake)
|
||||
|
||||
res, err := a.HandleRestartJob(context.Background(), idRequest("scheduler_restart", "123"))
|
||||
if err != nil {
|
||||
t.Fatalf("HandleRestartJob error: %v", err)
|
||||
}
|
||||
if res.IsError {
|
||||
t.Fatalf("expected success, got error: %s", resultText(t, res))
|
||||
}
|
||||
if fake.restartID != "123" {
|
||||
t.Errorf("restartID = %q, want 123", fake.restartID)
|
||||
}
|
||||
if !strings.Contains(resultText(t, res), `"jobid": 999`) {
|
||||
t.Errorf("result should include the new job id 999, got: %s", resultText(t, res))
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleRestartJob_ClientError(t *testing.T) {
|
||||
fake := &fakeClient{restartErr: &fakeError{msg: "restart failed"}}
|
||||
a := newTestApp(fake)
|
||||
|
||||
res, err := a.HandleRestartJob(context.Background(), idRequest("scheduler_restart", "123"))
|
||||
if err != nil {
|
||||
t.Fatalf("HandleRestartJob error: %v", err)
|
||||
}
|
||||
if !res.IsError {
|
||||
t.Fatal("expected error result when client fails")
|
||||
}
|
||||
if !strings.Contains(resultText(t, res), "restart failed") {
|
||||
t.Errorf("result should mention the client error, got: %s", resultText(t, res))
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleRestartJob_MissingID(t *testing.T) {
|
||||
fake := &fakeClient{}
|
||||
a := newTestApp(fake)
|
||||
|
||||
res, err := a.HandleRestartJob(context.Background(), idRequest("scheduler_restart", ""))
|
||||
if err != nil {
|
||||
t.Fatalf("HandleRestartJob error: %v", err)
|
||||
}
|
||||
if !res.IsError {
|
||||
t.Fatal("expected error for missing id")
|
||||
}
|
||||
if fake.restartID != "" {
|
||||
t.Errorf("client should not be called, got restartID = %q", fake.restartID)
|
||||
}
|
||||
}
|
||||
|
||||
// fakeError is a minimal error for tests.
|
||||
type fakeError struct{ msg string }
|
||||
|
||||
|
|
|
|||
|
|
@ -51,6 +51,13 @@ type Client interface {
|
|||
// ReleaseJob resumes a suspended job (action "resume", default SIGCONT).
|
||||
ReleaseJob(ctx context.Context, id string) error
|
||||
|
||||
// CheckpointJob checkpoints a running job (action "checkpoint", CRIU).
|
||||
CheckpointJob(ctx context.Context, id string) error
|
||||
|
||||
// RestartJob re-queues a finished job as a new job (resuming from its last
|
||||
// checkpoint when one exists) and returns the new job's identity.
|
||||
RestartJob(ctx context.Context, id string) (*RestartJobResult, error)
|
||||
|
||||
// ── Event stream (Phase 2) ──
|
||||
|
||||
// WatchJobs subscribes to the live job-event SSE stream and returns a
|
||||
|
|
|
|||
|
|
@ -51,6 +51,12 @@ type JobControlRequest struct {
|
|||
Method string `json:"method,omitempty"`
|
||||
}
|
||||
|
||||
// JobRestartRequest is the body sent to POST /xce/v1/jobs/:jobids/restart.
|
||||
// All fields are optional; an empty request restarts to any eligible host.
|
||||
type JobRestartRequest struct {
|
||||
AskedHosts string `json:"asked_hosts,omitempty"`
|
||||
}
|
||||
|
||||
// ControlResult is the apiserver's per-item control response.
|
||||
type ControlResult struct {
|
||||
Status string `json:"status"` // accepted | failed | partially_accepted | no_items_processed
|
||||
|
|
@ -58,6 +64,15 @@ type ControlResult struct {
|
|||
Failed []ControlItem `json:"failed"`
|
||||
}
|
||||
|
||||
// RestartJobResult is the apiserver's response to a job restart — a new job
|
||||
// is created (optionally resumed from the old job's last checkpoint).
|
||||
type RestartJobResult struct {
|
||||
JobID int64 `json:"jobid"`
|
||||
JobIDHuman string `json:"job_id_human"`
|
||||
QueueName string `json:"queue_name"`
|
||||
Message string `json:"message"`
|
||||
}
|
||||
|
||||
// ControlItem reports the outcome for one job (or array element).
|
||||
type ControlItem struct {
|
||||
UniqueJobID int64 `json:"unique_job_id"`
|
||||
|
|
@ -223,3 +238,29 @@ func (hc *HTTPClient) HoldJob(ctx context.Context, id string) error {
|
|||
func (hc *HTTPClient) ReleaseJob(ctx context.Context, id string) error {
|
||||
return hc.controlJob(ctx, id, "resume")
|
||||
}
|
||||
|
||||
// CheckpointJob checkpoints a running job (action "checkpoint", CRIU).
|
||||
func (hc *HTTPClient) CheckpointJob(ctx context.Context, id string) error {
|
||||
return hc.controlJob(ctx, id, "checkpoint")
|
||||
}
|
||||
|
||||
// RestartJob re-queues a finished job as a new job, resuming from its last
|
||||
// checkpoint when one exists. Returns the new job's identity. The apiserver
|
||||
// backs this with the same submit path, so a restart is accepted with 201.
|
||||
func (hc *HTTPClient) RestartJob(ctx context.Context, id string) (*RestartJobResult, error) {
|
||||
if strings.TrimSpace(id) == "" {
|
||||
return nil, fmt.Errorf("restart job: job id is required")
|
||||
}
|
||||
|
||||
url := hc.baseURL + "/xce/v1/jobs/" + id + "/restart"
|
||||
env, err := hc.doPost(ctx, url, JobRestartRequest{}, http.StatusCreated)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("restart job %s: %w", id, err)
|
||||
}
|
||||
|
||||
var res RestartJobResult
|
||||
if err := json.Unmarshal(env.Data, &res); err != nil {
|
||||
return nil, fmt.Errorf("restart job %s: parse result: %w", id, err)
|
||||
}
|
||||
return &res, nil
|
||||
}
|
||||
|
|
|
|||
|
|
@ -282,6 +282,85 @@ func TestReleaseJob_SendsResume(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
// ── Checkpoint / Restart ──────────────────────────────────────────────────────
|
||||
|
||||
func TestCheckpointJob_SendsCheckpoint(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.CheckpointJob(context.Background(), "101"); err != nil {
|
||||
t.Fatalf("CheckpointJob returned error: %v", err)
|
||||
}
|
||||
|
||||
_, path, _, body := capture()
|
||||
if path != "/xce/v1/jobs/101/control" {
|
||||
t.Errorf("path = %q, want /xce/v1/jobs/101/control", path)
|
||||
}
|
||||
if strings.TrimSpace(body) != `{"action":"checkpoint"}` {
|
||||
t.Errorf("body = %s, want {\"action\":\"checkpoint\"}", body)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCheckpointJob_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":101,"job_id":101,"array_idx":0,"msg":"job 101 is in state DONE, cannot be checkpointed"}]}}`)
|
||||
})
|
||||
|
||||
err := hc.CheckpointJob(context.Background(), "101")
|
||||
if err == nil {
|
||||
t.Fatal("expected error when control status is failed")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "101") {
|
||||
t.Errorf("error = %q, want to mention job id 101", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRestartJob_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":999,"queue_name":"normal","message":"Job <123> will be restarted with new job <999> to default queue <normal>."}}`)
|
||||
})
|
||||
|
||||
res, err := hc.RestartJob(context.Background(), "123")
|
||||
if err != nil {
|
||||
t.Fatalf("RestartJob 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/restart" {
|
||||
t.Errorf("path = %q, want /xce/v1/jobs/123/restart", path)
|
||||
}
|
||||
if strings.TrimSpace(body) != "{}" {
|
||||
t.Errorf("body = %s, want {}", body)
|
||||
}
|
||||
if res.JobID != 999 {
|
||||
t.Errorf("JobID = %d, want 999", res.JobID)
|
||||
}
|
||||
if res.QueueName != "normal" {
|
||||
t.Errorf("QueueName = %q, want normal", res.QueueName)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRestartJob_EmptyID(t *testing.T) {
|
||||
hc, _ := newTestServer(t, func(w http.ResponseWriter, r *http.Request) {
|
||||
t.Error("server should not be called when id is empty")
|
||||
})
|
||||
|
||||
_, err := hc.RestartJob(context.Background(), " ")
|
||||
if err == nil {
|
||||
t.Fatal("expected error for empty id")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "job id is required") {
|
||||
t.Errorf("error = %q, want 'job id is required'", err)
|
||||
}
|
||||
}
|
||||
|
||||
// ── conversion helpers ───────────────────────────────────────────────────────
|
||||
|
||||
func TestFormatWalltimeSeconds(t *testing.T) {
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue