Coverage for gco_mcp / tools / jobs.py: 100.00%
110 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-09-14 22:07 +0000
1"""Job management MCP tools."""
3import asyncio
4import contextlib
6import cli_runner
7from audit import audit_logged
8from feature_flags import FLAG_DESTRUCTIVE_OPERATIONS, is_enabled
9from server import mcp
12async def _ctx_warning(message: str) -> None:
13 """Emit ``ctx.warning(...)`` from inside a tool body, no-op when no Context.
15 The destructive ``delete_job`` tool runs short — we don't need the
16 full long-task progress stack, just an audited warning back to the
17 operator (and the audit log via the middleware spy).
18 """
19 try:
20 from fastmcp.server.dependencies import get_context
22 ctx = get_context()
23 except Exception:
24 return
25 with contextlib.suppress(Exception):
26 await ctx.warning(message)
29@mcp.tool(tags={"safe", "jobs"})
30@audit_logged
31def list_jobs(
32 region: str | None = None, namespace: str | None = None, status: str | None = None
33) -> str:
34 """List jobs across GCO clusters.
36 Args:
37 region: AWS region (e.g. us-east-1). If omitted, lists across all regions.
38 namespace: Filter by Kubernetes namespace.
39 status: Filter by job status (pending, running, completed, succeeded, failed).
40 """
41 args = ["jobs", "list"]
42 if region:
43 args += ["-r", region]
44 else:
45 args += ["--all-regions"]
46 if namespace:
47 args += ["-n", namespace]
48 if status:
49 args += ["-s", status]
50 return cli_runner._run_cli(*args)
53@mcp.tool(tags={"low-risk", "jobs"})
54@audit_logged
55def submit_job_sqs(
56 manifest_path: str, region: str, namespace: str | None = None, priority: int | None = None
57) -> str:
58 """Submit a job via SQS queue (recommended for production).
60 Args:
61 manifest_path: Path to the YAML manifest file (relative to project root).
62 region: Target AWS region for the SQS queue.
63 namespace: Override the namespace in the manifest.
64 priority: Job priority (0-100, higher = more important).
65 """
66 args = ["jobs", "submit-sqs", manifest_path, "-r", region]
67 if namespace:
68 args += ["-n", namespace]
69 if priority is not None:
70 args += ["--priority", str(priority)]
71 return cli_runner._run_cli(*args)
74@mcp.tool(tags={"low-risk", "jobs"})
75@audit_logged
76def submit_job_api(manifest_path: str, namespace: str | None = None) -> str:
77 """Submit a job via the authenticated API Gateway (SigV4).
79 Args:
80 manifest_path: Path to the YAML manifest file.
81 namespace: Override the namespace in the manifest.
82 """
83 args = ["jobs", "submit", manifest_path]
84 if namespace:
85 args += ["-n", namespace]
86 return cli_runner._run_cli(*args)
89@mcp.tool(tags={"safe", "jobs"})
90@audit_logged
91def get_job(job_name: str, region: str, namespace: str = "gco-jobs") -> str:
92 """Get details of a specific job, including where its pods were scheduled.
94 Reports ``node_name``, ``node_instance_type`` and ``node_capacity_type``
95 (spot vs on-demand) for the node the job's pod landed on, plus a ``nodes``
96 list covering every node involved. A job authorized to run on a set of
97 interchangeable instance types therefore reports the one it actually used,
98 not just what the manifest permitted. The fields are unset — never guessed
99 — when nothing is scheduled yet or the pods have been garbage-collected.
101 Args:
102 job_name: Name of the job.
103 region: AWS region where the job is running.
104 namespace: Kubernetes namespace.
105 """
106 return cli_runner._run_cli("jobs", "get", job_name, "-r", region, "-n", namespace)
109@mcp.tool(tags={"safe", "jobs"})
110@audit_logged
111def get_job_logs(job_name: str, region: str, namespace: str = "gco-jobs", tail: int = 100) -> str:
112 """Get logs from a job.
114 Args:
115 job_name: Name of the job.
116 region: AWS region.
117 namespace: Kubernetes namespace.
118 tail: Number of log lines to return.
119 """
120 return cli_runner._run_cli(
121 "jobs", "logs", job_name, "-r", region, "-n", namespace, "--tail", str(tail)
122 )
125if is_enabled(FLAG_DESTRUCTIVE_OPERATIONS):
127 @mcp.tool(tags={"destructive", "jobs"})
128 @audit_logged
129 async def delete_job(job_name: str, region: str, namespace: str = "gco-jobs") -> str:
130 """[gated by GCO_ENABLE_DESTRUCTIVE_OPERATIONS] destructive.
132 Delete a job. Cannot be undone — the Kubernetes Job and its pods
133 are removed and any pod logs not yet shipped to CloudWatch are lost.
135 Args:
136 job_name: Name of the job to delete.
137 region: AWS region.
138 namespace: Kubernetes namespace.
139 """
140 await _ctx_warning(
141 f"Deleting job {job_name!r} in {region}/{namespace} — this cannot be undone."
142 )
143 return await asyncio.to_thread(
144 cli_runner._run_cli, "jobs", "delete", job_name, "-r", region, "-n", namespace, "-y"
145 )
148@mcp.tool(tags={"safe", "jobs"})
149@audit_logged
150def get_job_events(job_name: str, region: str, namespace: str = "gco-jobs") -> str:
151 """Get Kubernetes events for a job (useful for debugging).
153 Args:
154 job_name: Name of the job.
155 region: AWS region.
156 namespace: Kubernetes namespace.
157 """
158 return cli_runner._run_cli("jobs", "events", job_name, "-r", region, "-n", namespace)
161@mcp.tool(tags={"safe", "jobs"})
162@audit_logged
163def get_job_pods(job_name: str, region: str, namespace: str = "gco-jobs") -> str:
164 """Get pod details, placement, and container status for a job.
166 Each pod carries a ``node`` block with the instance type and spot/on-demand
167 capacity type of the node it landed on.
169 Args:
170 job_name: Name of the owning Kubernetes Job.
171 region: AWS region where the job is running.
172 namespace: Kubernetes namespace.
173 """
174 return cli_runner._run_cli("jobs", "pods", job_name, "-r", region, "-n", namespace)
177@mcp.tool(tags={"safe", "jobs"})
178@audit_logged
179def get_pod_logs(
180 job_name: str,
181 pod_name: str,
182 region: str,
183 namespace: str = "gco-jobs",
184 tail: int = 100,
185 container: str | None = None,
186) -> str:
187 """Get a bounded log tail from one specific pod belonging to a job.
189 Args:
190 job_name: Name of the owning Kubernetes Job.
191 pod_name: Exact pod name returned by ``get_job_pods``.
192 region: AWS region where the pod is running.
193 namespace: Kubernetes namespace.
194 tail: Maximum number of log lines to return.
195 container: Container name for a multi-container pod.
196 """
197 args = [
198 "jobs",
199 "pod-logs",
200 job_name,
201 pod_name,
202 "-r",
203 region,
204 "-n",
205 namespace,
206 "--tail",
207 str(tail),
208 ]
209 if container:
210 args += ["--container", container]
211 return cli_runner._run_cli(*args)
214@mcp.tool(tags={"safe", "jobs"})
215@audit_logged
216def get_job_metrics(job_name: str, region: str, namespace: str = "gco-jobs") -> str:
217 """Get CPU and memory usage for all pods in a job.
219 Requires metrics-server in the target cluster.
221 Args:
222 job_name: Name of the Kubernetes Job.
223 region: AWS region where the job is running.
224 namespace: Kubernetes namespace.
225 """
226 return cli_runner._run_cli("jobs", "metrics", job_name, "-r", region, "-n", namespace)
229@mcp.tool(tags={"low-risk", "jobs"})
230@audit_logged
231def retry_job(job_name: str, region: str, namespace: str = "gco-jobs") -> str:
232 """Retry a failed job by creating a new Job while preserving the original.
234 Args:
235 job_name: Failed Kubernetes Job to retry.
236 region: AWS region where the job ran.
237 namespace: Kubernetes namespace.
238 """
239 return cli_runner._run_cli("jobs", "retry", job_name, "-r", region, "-n", namespace, "--yes")
242@mcp.tool(tags={"safe", "jobs"})
243@audit_logged
244def cluster_health(region: str | None = None) -> str:
245 """Get health status of GCO clusters.
247 Args:
248 region: Specific region, or omit for all regions.
249 """
250 args = ["jobs", "health"]
251 if region:
252 args += ["-r", region]
253 else:
254 args += ["--all-regions"]
255 return cli_runner._run_cli(*args)
258@mcp.tool(tags={"safe", "jobs"})
259@audit_logged
260def get_job_validation_policy(region: str) -> str:
261 """Get the job validation policy a region actually enforces, as deployed.
263 Use this before submitting to check whether a manifest will be admitted,
264 rather than paying to provision a region and discovering the conflict at
265 submit time. Returns the per-manifest cpu/memory/gpu caps,
266 allowed_namespaces, allowed_kinds, trusted_registries, the pod-security
267 block_* flags, and the namespace's live ResourceQuota / LimitRange
268 ceilings.
270 This reads the deployed cluster, not a local cdk.json. The two diverge
271 whenever a stack was deployed from a different checkout, and CDK augments
272 trusted_registries with the project's own ECR hostnames at synth time, so
273 the effective allowlist is strictly larger than the configured one.
275 A manifest must clear all three layers: the front-door policy, the
276 per-container LimitRange, and the aggregate ResourceQuota.
278 Args:
279 region: AWS region (e.g. us-east-1).
280 """
281 return cli_runner._run_cli("jobs", "policy", "-r", region)
284@mcp.tool(tags={"safe", "jobs"})
285@audit_logged
286def check_job_policy(
287 manifest_path: str,
288 regions: list[str] | None = None,
289 namespace: str | None = None,
290 offline: bool = False,
291) -> str:
292 """Check which regions would admit a manifest, and whether regions agree.
294 Answers two questions get_job_validation_policy leaves to the caller.
296 Which regions would take this job: the same manifest is evaluated against
297 each region's deployed policy using the code the manifest processor runs,
298 so a job that is admissible in one region and over-cap in another is
299 reported as such instead of being discovered by submitting.
301 Whether the regions still agree: there are no per-region policy overrides,
302 so any field that differs across regions means a region was deployed from a
303 different checkout of cdk.json. That is invisible until a manifest that
304 worked yesterday is refused. trusted_registries is compared with ECR
305 hostnames stripped, since CDK adds those per deployment.
307 Advisory. The cluster is the authoritative gate and this reads a snapshot
308 of its policy, so a reject here is a strong signal, not a verdict.
310 Args:
311 manifest_path: Path to a manifest file or a directory of them.
312 regions: Regions to check. Omit for every configured region.
313 namespace: Namespace to assume for manifests that don't declare one.
314 offline: Read cdk.json instead of calling AWS. Needs no credentials,
315 but reports the CONFIGURED policy rather than the deployed one, and
316 a deployed region trusts ECR registries cdk.json never mentions --
317 so an image rejection may be a false positive.
318 """
319 args = ["jobs", "check-policy", manifest_path]
320 for region in regions or []:
321 args += ["-r", region]
322 if namespace:
323 args += ["-n", namespace]
324 if offline:
325 args.append("--offline")
326 return cli_runner._run_cli(*args)
329@mcp.tool(tags={"safe", "jobs"})
330@audit_logged
331def queue_status(region: str | None = None) -> str:
332 """View SQS queue status (pending, in-flight, DLQ counts).
334 Args:
335 region: Specific region, or omit for all regions.
336 """
337 args = ["jobs", "queue-status"]
338 if region:
339 args += ["-r", region]
340 else:
341 args += ["--all-regions"]
342 return cli_runner._run_cli(*args)