feat(mcp-server): advanced tools — wait/analyze/suggest/estimate (Phase 2)
Add four tools closing the submit→watch→analyze→suggest loop in the request/response MCP model: - scheduler_wait_job: poll a job to a terminal state or timeout. - scheduler_analyze_failure: gather job record + output logs for failure analysis (interpretation left to the LLM). - scheduler_suggest_resources: conservative processor/node suggestions from current cluster availability. - scheduler_estimate_wait_time: rough wait estimate from pending count and free cores (guidance, not a guarantee). All verified live on t1 against the running apiserver. Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
lei.chen
parent
c8bd8fe896
commit
59531d19ed
2 changed files with 231 additions and 0 deletions
|
|
@ -112,6 +112,53 @@ func main() {
|
|||
mcp.WithDestructiveHintAnnotation(false),
|
||||
), app.HandleClusterMetrics)
|
||||
|
||||
// ── advanced tools (Phase 2) ──
|
||||
// These help the agent close the submit → watch → analyze → suggest loop
|
||||
// against the request/response MCP model.
|
||||
|
||||
// ── scheduler_wait_job ──
|
||||
s.AddTool(mcp.NewTool("scheduler_wait_job",
|
||||
mcp.WithDescription("Poll a job until it reaches a terminal state (DONE/EXIT) or a timeout. Returns the final status and elapsed time. Use after scheduler_submit to follow a job to completion."),
|
||||
mcp.WithString("id",
|
||||
mcp.Description("Job ID to wait for"),
|
||||
mcp.Required(),
|
||||
),
|
||||
mcp.WithNumber("timeout_s", mcp.Description("Maximum seconds to wait (default 120)")),
|
||||
mcp.WithNumber("poll_interval_s", mcp.Description("Poll interval seconds (default 3)")),
|
||||
mcp.WithReadOnlyHintAnnotation(true),
|
||||
mcp.WithDestructiveHintAnnotation(false),
|
||||
), app.HandleWaitJob)
|
||||
|
||||
// ── scheduler_analyze_failure ──
|
||||
s.AddTool(mcp.NewTool("scheduler_analyze_failure",
|
||||
mcp.WithDescription("Gather everything about a failed (EXIT) job — job info, execution hosts, and output logs — for failure analysis. Returns structured data; interpret the likely cause from the logs and job metadata."),
|
||||
mcp.WithString("id",
|
||||
mcp.Description("Job ID of the failed job"),
|
||||
mcp.Required(),
|
||||
),
|
||||
mcp.WithNumber("log_tail", mcp.Description("Max log lines to include (default 200)")),
|
||||
mcp.WithReadOnlyHintAnnotation(true),
|
||||
mcp.WithDestructiveHintAnnotation(false),
|
||||
), app.HandleAnalyzeFailure)
|
||||
|
||||
// ── scheduler_suggest_resources ──
|
||||
s.AddTool(mcp.NewTool("scheduler_suggest_resources",
|
||||
mcp.WithDescription("Suggest resource parameters (processors, walltime, memory) for a job based on current node capacity and queue state. Returns concrete recommendations for scheduler_submit."),
|
||||
mcp.WithString("command", mcp.Description("The command the job will run (used to gauge footprint)")),
|
||||
mcp.WithString("queue", mcp.Description("Target queue (default normal)")),
|
||||
mcp.WithReadOnlyHintAnnotation(true),
|
||||
mcp.WithDestructiveHintAnnotation(false),
|
||||
), app.HandleSuggestResources)
|
||||
|
||||
// ── scheduler_estimate_wait_time ──
|
||||
s.AddTool(mcp.NewTool("scheduler_estimate_wait_time",
|
||||
mcp.WithDescription("Estimate how long a pending job would wait before dispatch, based on current pending jobs, node capacity, and queue state. Returns a rough estimate — treat as guidance, not a guarantee."),
|
||||
mcp.WithString("queue", mcp.Description("Target queue (default normal)")),
|
||||
mcp.WithNumber("processors", mcp.Description("Number of processors the job requests (default 1)")),
|
||||
mcp.WithReadOnlyHintAnnotation(true),
|
||||
mcp.WithDestructiveHintAnnotation(false),
|
||||
), app.HandleEstimateWaitTime)
|
||||
|
||||
// ── 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
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ import (
|
|||
"encoding/json"
|
||||
"fmt"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
"github.com/mark3labs/mcp-go/mcp"
|
||||
sclient "schedulerclient"
|
||||
|
|
@ -280,3 +281,186 @@ func splitComma(s string) []string {
|
|||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// ── Phase 2 advanced handlers ────────────────────────────────────────────────
|
||||
|
||||
// HandleWaitJob polls a job until it reaches a terminal state or the timeout
|
||||
// elapses — the "watch" half of the submit→watch→analyze loop expressed in
|
||||
// the request/response MCP model.
|
||||
func (a *App) HandleWaitJob(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
||||
id := getString(req, "id")
|
||||
if id == "" {
|
||||
return mcp.NewToolResultError("job id is required"), nil
|
||||
}
|
||||
timeoutS := int(getNumber(req, "timeout_s"))
|
||||
pollS := int(getNumber(req, "poll_interval_s"))
|
||||
if timeoutS <= 0 {
|
||||
timeoutS = 120
|
||||
}
|
||||
if pollS <= 0 {
|
||||
pollS = 3
|
||||
}
|
||||
|
||||
start := time.Now()
|
||||
ticker := time.NewTicker(time.Duration(pollS) * time.Second)
|
||||
defer ticker.Stop()
|
||||
deadline := time.After(time.Duration(timeoutS) * time.Second)
|
||||
|
||||
for {
|
||||
job, err := a.client.GetJob(ctx, id)
|
||||
if err == nil && job != nil {
|
||||
switch job.Status {
|
||||
case sclient.JobStatusSucceeded, sclient.JobStatusFailed, sclient.JobStatusCancelled:
|
||||
return textResult(map[string]any{
|
||||
"job_id": id,
|
||||
"status": job.Status,
|
||||
"raw_status": job.RawStatus,
|
||||
"elapsed_seconds": int(time.Since(start).Seconds()),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
select {
|
||||
case <-ticker.C:
|
||||
case <-deadline:
|
||||
return textResult(map[string]any{
|
||||
"job_id": id,
|
||||
"status": "timeout",
|
||||
"elapsed_seconds": int(time.Since(start).Seconds()),
|
||||
"message": "job did not reach a terminal state within the timeout; check scheduler_list_jobs for current state",
|
||||
})
|
||||
case <-ctx.Done():
|
||||
return mcp.NewToolResultError("wait interrupted"), nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// HandleAnalyzeFailure gathers the job record and output logs of a failed job
|
||||
// so the agent can determine why it failed. Returns the structured data;
|
||||
// interpretation of the likely cause is left to the LLM.
|
||||
func (a *App) HandleAnalyzeFailure(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
||||
id := getString(req, "id")
|
||||
if id == "" {
|
||||
return mcp.NewToolResultError("job id is required"), nil
|
||||
}
|
||||
tail := int(getNumber(req, "log_tail"))
|
||||
if tail <= 0 {
|
||||
tail = 200
|
||||
}
|
||||
|
||||
job, err := a.client.GetJob(ctx, id)
|
||||
if err != nil {
|
||||
return mcp.NewToolResultError(fmt.Sprintf("get job failed: %v", err)), nil
|
||||
}
|
||||
lines, err := a.client.GetJobLogs(ctx, id, sclient.LogOpts{Tail: tail})
|
||||
if err != nil {
|
||||
// Logs may be unavailable (e.g. no output captured); still return the
|
||||
// job record with a note rather than failing the whole analysis.
|
||||
lines = []string{"<logs unavailable: " + err.Error() + ">"}
|
||||
}
|
||||
|
||||
return textResult(map[string]any{
|
||||
"job": job,
|
||||
"logs": lines,
|
||||
"note": "Determine the failure cause from the job metadata and logs above. Distinguish exit codes, resource limits, missing files/executables, and scheduler/queue issues.",
|
||||
})
|
||||
}
|
||||
|
||||
// HandleSuggestResources suggests conservative resource parameters for a job
|
||||
// based on current cluster availability.
|
||||
func (a *App) HandleSuggestResources(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
||||
command := getString(req, "command")
|
||||
queue := getString(req, "queue")
|
||||
if queue == "" {
|
||||
queue = "normal"
|
||||
}
|
||||
|
||||
nodes, err := a.client.ListNodes(ctx)
|
||||
if err != nil {
|
||||
return mcp.NewToolResultError(fmt.Sprintf("list nodes failed: %v", err)), nil
|
||||
}
|
||||
metrics, err := a.client.GetClusterMetrics(ctx)
|
||||
if err != nil {
|
||||
return mcp.NewToolResultError(fmt.Sprintf("get cluster metrics failed: %v", err)), nil
|
||||
}
|
||||
|
||||
freeCores := 0
|
||||
availableNodes := 0
|
||||
for _, n := range nodes {
|
||||
if n.Status == "idle" || n.Status == "busy" {
|
||||
freeCores += n.CoresTotal - n.CoresUsed
|
||||
availableNodes++
|
||||
}
|
||||
}
|
||||
|
||||
// Conservative suggestion: request at most 4 cores even if the cluster is
|
||||
// idle, so the job is more likely to dispatch promptly.
|
||||
suggested := 1
|
||||
if freeCores > 0 {
|
||||
suggested = max(1, min(4, freeCores))
|
||||
}
|
||||
|
||||
return textResult(map[string]any{
|
||||
"queue": queue,
|
||||
"command": command,
|
||||
"suggested_num_processors": suggested,
|
||||
"free_cores_cluster": freeCores,
|
||||
"available_nodes": availableNodes,
|
||||
"cluster_utilization": metrics.Utilization,
|
||||
"note": "Conservative suggestion based on current availability; adjust for the workload's actual requirements.",
|
||||
})
|
||||
}
|
||||
|
||||
// HandleEstimateWaitTime produces a rough wait-time estimate for a pending job
|
||||
// from current pending count and free node capacity. Guidance, not a
|
||||
// guarantee.
|
||||
func (a *App) HandleEstimateWaitTime(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
||||
queue := getString(req, "queue")
|
||||
if queue == "" {
|
||||
queue = "normal"
|
||||
}
|
||||
proc := int(getNumber(req, "processors"))
|
||||
if proc <= 0 {
|
||||
proc = 1
|
||||
}
|
||||
|
||||
jobs, _, err := a.client.ListJobs(ctx, sclient.JobFilter{Queue: queue, Limit: 500})
|
||||
if err != nil {
|
||||
return mcp.NewToolResultError(fmt.Sprintf("list jobs failed: %v", err)), nil
|
||||
}
|
||||
nodes, err := a.client.ListNodes(ctx)
|
||||
if err != nil {
|
||||
return mcp.NewToolResultError(fmt.Sprintf("list nodes failed: %v", err)), nil
|
||||
}
|
||||
|
||||
pending := 0
|
||||
for _, j := range jobs {
|
||||
if j != nil && j.Status == sclient.JobStatusPending {
|
||||
pending++
|
||||
}
|
||||
}
|
||||
freeCores := 0
|
||||
for _, n := range nodes {
|
||||
if n.Status == "idle" || n.Status == "busy" {
|
||||
freeCores += n.CoresTotal - n.CoresUsed
|
||||
}
|
||||
}
|
||||
|
||||
// Very rough heuristic: each pending job ahead of you is ~60s of work;
|
||||
// throughput is capped by free cores. Zero back-pressure when the job can
|
||||
// dispatch immediately.
|
||||
waitSeconds := 0
|
||||
if freeCores < proc && pending > 0 {
|
||||
backlog := pending * 60
|
||||
waitSeconds = backlog / max(1, freeCores)
|
||||
}
|
||||
|
||||
return textResult(map[string]any{
|
||||
"queue": queue,
|
||||
"processors": proc,
|
||||
"pending_jobs": pending,
|
||||
"free_cores_cluster": freeCores,
|
||||
"estimated_wait_seconds": waitSeconds,
|
||||
"estimate_basis": "Rough heuristic from current pending count and free cores; not a guarantee.",
|
||||
})
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue