Coverage for gco_mcp / tools / slurm.py: 100.00%
82 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-10-28 16:19 +0000
1"""Slurm MCP tools (wrappers over ``gco slurm``).
3The Slinky Slurm cluster GCO deploys when ``cdk.json`` ``helm.slurm.enabled``
4is on, reached through the GCO API's ``/api/v1/slurm/*`` routes: every tool
5is pinned to one region because each region runs its own cluster. The
6interactive ``gco slurm shell`` and the ``--follow`` form of ``gco slurm
7output`` have no tool on purpose (they hold a terminal open over the cluster
8tunnel); ``slurm_cancel`` is gated like every other destructive operation.
9"""
11import asyncio
12import contextlib
14import cli_runner
15from audit import audit_logged
16from feature_flags import FLAG_DESTRUCTIVE_OPERATIONS, is_enabled
17from server import mcp
20def _args(*parts: str | None, region: str) -> list[str]:
21 """``gco slurm <parts...> -r REGION``, dropping omitted optional parts."""
22 return ["slurm", *[part for part in parts if part is not None], "-r", region]
25# =============================================================================
26# Read-only tools
27# =============================================================================
30@mcp.tool(tags={"safe", "slurm"})
31@audit_logged
32async def slurm_status(region: str) -> str:
33 """`gco slurm status` — controller reachability, partitions, node and job counts.
35 Args:
36 region: Region whose Slurm cluster to read (one cluster per region).
37 """
38 return await asyncio.to_thread(cli_runner._run_cli, *_args("status", region=region))
41@mcp.tool(tags={"safe", "slurm"})
42@audit_logged
43async def slurm_partitions(region: str) -> str:
44 """`gco slurm partitions` — partitions with state, time limit and node counts (sinfo).
46 An autoscaled partition reads zero nodes while idle: the first pending job
47 makes KEDA scale its NodeSet up.
49 Args:
50 region: Region whose Slurm cluster to read.
51 """
52 return await asyncio.to_thread(cli_runner._run_cli, *_args("partitions", region=region))
55@mcp.tool(tags={"safe", "slurm"})
56@audit_logged
57async def slurm_nodes(region: str, partition: str | None = None) -> str:
58 """`gco slurm nodes` — Slurm nodes (slurmd worker pods) with state, resources and gres.
60 Args:
61 region: Region whose Slurm cluster to read.
62 partition: Only nodes in this partition.
63 """
64 args = _args("nodes", region=region)
65 if partition:
66 args += ["--partition", partition]
67 return await asyncio.to_thread(cli_runner._run_cli, *args)
70@mcp.tool(tags={"safe", "slurm"})
71@audit_logged
72async def slurm_queue(
73 region: str,
74 state: str | None = None,
75 partition: str | None = None,
76 name: str | None = None,
77 limit: int | None = None,
78) -> str:
79 """`gco slurm queue` — the live queue as slurmctld sees it (squeue).
81 Args:
82 region: Region whose Slurm cluster to read.
83 state: Only jobs in this state (PENDING, RUNNING, COMPLETING, ...).
84 partition: Only jobs in this partition.
85 name: Only jobs with exactly this name.
86 limit: Maximum jobs returned (default 200).
87 """
88 args = _args("queue", region=region)
89 if state:
90 args += ["--state", state]
91 if partition:
92 args += ["--partition", partition]
93 if name:
94 args += ["--name", name]
95 if limit is not None:
96 args += ["--limit", str(limit)]
97 return await asyncio.to_thread(cli_runner._run_cli, *args)
100@mcp.tool(tags={"safe", "slurm"})
101@audit_logged
102async def slurm_job(region: str, job_id: int) -> str:
103 """`gco slurm job` — one job in full (state, TRES, times, exit code, output paths).
105 Args:
106 region: Region whose Slurm cluster to read.
107 job_id: The Slurm job id.
108 """
109 return await asyncio.to_thread(cli_runner._run_cli, *_args("job", str(job_id), region=region))
112@mcp.tool(tags={"safe", "slurm"})
113@audit_logged
114async def slurm_accounting(
115 region: str,
116 since: str | None = None,
117 start: str | None = None,
118 end: str | None = None,
119 user: str | None = None,
120 partition: str | None = None,
121 state: str | None = None,
122 name: str | None = None,
123 limit: int | None = None,
124) -> str:
125 """`gco slurm accounting` — job history from slurmdbd (sacct).
127 Args:
128 region: Region whose Slurm cluster to read.
129 since: Relative window such as 30m, 24h, 7d or 2w (default 1d unless start is given).
130 start: Window start as a Slurm time (2026-10-01, 2026-10-01T08:00, now-2days).
131 end: Window end as a Slurm time (default now).
132 user: Only this Slurm user (API submissions run as ``slurm``; the
133 submitting caller is in each row's comment).
134 partition: Only this partition.
135 state: Only these states, comma-separated (COMPLETED,FAILED,TIMEOUT).
136 name: Only jobs with exactly this name.
137 limit: Maximum rows (default 200).
138 """
139 args = _args("accounting", region=region)
140 for flag, value in (
141 ("--since", since),
142 ("--start", start),
143 ("--end", end),
144 ("--user", user),
145 ("--partition", partition),
146 ("--state", state),
147 ("--name", name),
148 ):
149 if value:
150 args += [flag, value]
151 if limit is not None:
152 args += ["--limit", str(limit)]
153 return await asyncio.to_thread(cli_runner._run_cli, *args)
156# =============================================================================
157# Mutating tools (low-risk)
158# =============================================================================
161@mcp.tool(tags={"low-risk", "slurm"})
162@audit_logged
163async def slurm_submit(
164 region: str,
165 script_path: str,
166 name: str | None = None,
167 partition: str | None = None,
168 nodes: int | None = None,
169 tasks: int | None = None,
170 cpus_per_task: int | None = None,
171 gpus_per_node: int | None = None,
172 memory: str | None = None,
173 time_limit: str | None = None,
174 working_directory: str | None = None,
175 environment: list[str] | None = None,
176 comment: str | None = None,
177 exclusive: bool = False,
178) -> str:
179 """`gco slurm submit` — submit a local batch script to Slurm (sbatch).
181 The script's leading ``#SBATCH`` directives are honoured; the arguments
182 here override them. The job runs as Slurm's own user with the caller
183 identity in its comment, and allocates compute (an idle autoscaled
184 partition scales up for it).
186 Args:
187 region: Region whose Slurm cluster runs the job.
188 script_path: Local path of the batch script (an interpreter line first).
189 name: Job name (default: the script's #SBATCH --job-name, else gco-job).
190 partition: Partition (default: the cluster's default partition).
191 nodes: Node count.
192 tasks: Task count (--ntasks).
193 cpus_per_task: CPUs per task.
194 gpus_per_node: GPUs per node.
195 memory: Memory per node: MiB or a K/M/G/T suffix such as 64G.
196 time_limit: Wall time: minutes, HH:MM:SS or D-HH:MM:SS.
197 working_directory: Working directory on the worker (default /home/slurm).
198 environment: Extra environment variables as NAME=VALUE strings.
199 comment: Free-text comment appended after the caller identity.
200 exclusive: Whole-node allocation.
201 """
202 args = _args("submit", script_path, region=region)
203 for flag, value in (
204 ("--name", name),
205 ("--partition", partition),
206 ("--mem", memory),
207 ("--time", time_limit),
208 ("--chdir", working_directory),
209 ("--comment", comment),
210 ):
211 if value:
212 args += [flag, value]
213 for flag, number in (
214 ("--nodes", nodes),
215 ("--tasks", tasks),
216 ("--cpus-per-task", cpus_per_task),
217 ("--gpus-per-node", gpus_per_node),
218 ):
219 if number is not None:
220 args += [flag, str(number)]
221 for item in environment or []:
222 args += ["--env", item]
223 if exclusive:
224 args.append("--exclusive")
225 return await asyncio.to_thread(cli_runner._run_cli, *args)
228# =============================================================================
229# Destructive tools — gated by GCO_ENABLE_DESTRUCTIVE_OPERATIONS
230# =============================================================================
233async def _ctx_warning(message: str) -> None:
234 """Emit ``ctx.warning(...)`` from inside a tool body, no-op when no Context."""
235 try:
236 from fastmcp.server.dependencies import get_context
238 ctx = get_context()
239 except Exception:
240 return
241 with contextlib.suppress(Exception):
242 await ctx.warning(message)
245if is_enabled(FLAG_DESTRUCTIVE_OPERATIONS):
247 @mcp.tool(tags={"destructive", "slurm"})
248 @audit_logged
249 async def slurm_cancel(region: str, job_id: int, signal: str | None = None) -> str:
250 """[gated by GCO_ENABLE_DESTRUCTIVE_OPERATIONS] destructive.
252 `gco slurm cancel` — cancel a Slurm job (scancel). A running job's
253 work is lost; a pending job leaves the queue.
255 Args:
256 region: Region whose Slurm cluster runs the job.
257 job_id: The Slurm job id.
258 signal: Signal to send instead of the default (e.g. SIGUSR1).
259 """
260 await _ctx_warning(f"Cancelling Slurm job {job_id} in {region} — its work is lost.")
261 args = [*_args("cancel", str(job_id), region=region), "-y"]
262 if signal:
263 args += ["--signal", signal]
264 return await asyncio.to_thread(cli_runner._run_cli, *args)