Coverage for gco_mcp / tools / jobs.py: 100.00%
110 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-10-28 16:19 +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 ``status`` is the Job's verdict and stays ``running`` while an unplaced
102 pod exists, so read the pod state next to it: ``pod_phase`` (the phase all
103 pods share; ``null`` when they disagree), ``pod_phases`` (count per
104 phase), ``scheduled_pods`` / ``unscheduled_pods``, and ``unscheduled`` —
105 one entry per pod with no node, carrying the ``PodScheduled`` condition's
106 ``reason`` and ``message`` ("0/6 nodes are available ...", a taint nobody
107 tolerates, Karpenter still nominating capacity). A Job whose only pod sits
108 in ``unscheduled`` is waiting on capacity, not training, and needs no
109 ``get_job_events`` call to learn why. Counts are ``null`` (unknown), not
110 ``0``, when the region runs a bridge that predates them.
112 Args:
113 job_name: Name of the job.
114 region: AWS region where the job is running.
115 namespace: Kubernetes namespace.
116 """
117 return cli_runner._run_cli("jobs", "get", job_name, "-r", region, "-n", namespace)
120@mcp.tool(tags={"safe", "jobs"})
121@audit_logged
122def get_job_logs(job_name: str, region: str, namespace: str = "gco-jobs", tail: int = 100) -> str:
123 """Get logs from a job.
125 Args:
126 job_name: Name of the job.
127 region: AWS region.
128 namespace: Kubernetes namespace.
129 tail: Number of log lines to return.
130 """
131 return cli_runner._run_cli(
132 "jobs", "logs", job_name, "-r", region, "-n", namespace, "--tail", str(tail)
133 )
136if is_enabled(FLAG_DESTRUCTIVE_OPERATIONS):
138 @mcp.tool(tags={"destructive", "jobs"})
139 @audit_logged
140 async def delete_job(job_name: str, region: str, namespace: str = "gco-jobs") -> str:
141 """[gated by GCO_ENABLE_DESTRUCTIVE_OPERATIONS] destructive.
143 Delete a job. Cannot be undone — the Kubernetes Job and its pods
144 are removed and any pod logs not yet shipped to CloudWatch are lost.
146 Args:
147 job_name: Name of the job to delete.
148 region: AWS region.
149 namespace: Kubernetes namespace.
150 """
151 await _ctx_warning(
152 f"Deleting job {job_name!r} in {region}/{namespace} — this cannot be undone."
153 )
154 return await asyncio.to_thread(
155 cli_runner._run_cli, "jobs", "delete", job_name, "-r", region, "-n", namespace, "-y"
156 )
159@mcp.tool(tags={"safe", "jobs"})
160@audit_logged
161def get_job_events(job_name: str, region: str, namespace: str = "gco-jobs") -> str:
162 """Get Kubernetes events for a job (useful for debugging).
164 Args:
165 job_name: Name of the job.
166 region: AWS region.
167 namespace: Kubernetes namespace.
168 """
169 return cli_runner._run_cli("jobs", "events", job_name, "-r", region, "-n", namespace)
172@mcp.tool(tags={"safe", "jobs"})
173@audit_logged
174def get_job_pods(job_name: str, region: str, namespace: str = "gco-jobs") -> str:
175 """Get pod details, placement, and container status for a job.
177 Each pod carries a ``node`` block with the instance type and spot/on-demand
178 capacity type of the node it landed on.
180 Args:
181 job_name: Name of the owning Kubernetes Job.
182 region: AWS region where the job is running.
183 namespace: Kubernetes namespace.
184 """
185 return cli_runner._run_cli("jobs", "pods", job_name, "-r", region, "-n", namespace)
188@mcp.tool(tags={"safe", "jobs"})
189@audit_logged
190def get_pod_logs(
191 job_name: str,
192 pod_name: str,
193 region: str,
194 namespace: str = "gco-jobs",
195 tail: int = 100,
196 container: str | None = None,
197) -> str:
198 """Get a bounded log tail from one specific pod belonging to a job.
200 Args:
201 job_name: Name of the owning Kubernetes Job.
202 pod_name: Exact pod name returned by ``get_job_pods``.
203 region: AWS region where the pod is running.
204 namespace: Kubernetes namespace.
205 tail: Maximum number of log lines to return.
206 container: Container name for a multi-container pod.
207 """
208 args = [
209 "jobs",
210 "pod-logs",
211 job_name,
212 pod_name,
213 "-r",
214 region,
215 "-n",
216 namespace,
217 "--tail",
218 str(tail),
219 ]
220 if container:
221 args += ["--container", container]
222 return cli_runner._run_cli(*args)
225@mcp.tool(tags={"safe", "jobs"})
226@audit_logged
227def get_job_metrics(job_name: str, region: str, namespace: str = "gco-jobs") -> str:
228 """Get CPU and memory usage for all pods in a job.
230 Requires metrics-server in the target cluster.
232 Args:
233 job_name: Name of the Kubernetes Job.
234 region: AWS region where the job is running.
235 namespace: Kubernetes namespace.
236 """
237 return cli_runner._run_cli("jobs", "metrics", job_name, "-r", region, "-n", namespace)
240@mcp.tool(tags={"low-risk", "jobs"})
241@audit_logged
242def retry_job(job_name: str, region: str, namespace: str = "gco-jobs") -> str:
243 """Retry a failed job by creating a new Job while preserving the original.
245 Args:
246 job_name: Failed Kubernetes Job to retry.
247 region: AWS region where the job ran.
248 namespace: Kubernetes namespace.
249 """
250 return cli_runner._run_cli("jobs", "retry", job_name, "-r", region, "-n", namespace, "--yes")
253@mcp.tool(tags={"safe", "jobs"})
254@audit_logged
255def cluster_health(region: str | None = None) -> str:
256 """Get health status of GCO clusters.
258 Args:
259 region: Specific region, or omit for all regions.
260 """
261 args = ["jobs", "health"]
262 if region:
263 args += ["-r", region]
264 else:
265 args += ["--all-regions"]
266 return cli_runner._run_cli(*args)
269@mcp.tool(tags={"safe", "jobs"})
270@audit_logged
271def get_job_validation_policy(region: str) -> str:
272 """Get the job validation policy a region actually enforces, as deployed.
274 Use this before submitting to check whether a manifest will be admitted,
275 rather than paying to provision a region and discovering the conflict at
276 submit time. Returns the per-manifest cpu/memory/gpu caps,
277 allowed_namespaces, allowed_kinds, trusted_registries, the pod-security
278 block_* flags, and the namespace's live ResourceQuota / LimitRange
279 ceilings.
281 This reads the deployed cluster, not a local cdk.json. The two diverge
282 whenever a stack was deployed from a different checkout, and CDK augments
283 trusted_registries with the project's own ECR hostnames at synth time, so
284 the effective allowlist is strictly larger than the configured one.
286 A manifest must clear all three layers: the front-door policy, the
287 per-container LimitRange, and the aggregate ResourceQuota.
289 Args:
290 region: AWS region (e.g. us-east-1).
291 """
292 return cli_runner._run_cli("jobs", "policy", "-r", region)
295@mcp.tool(tags={"safe", "jobs"})
296@audit_logged
297def check_job_policy(
298 manifest_path: str,
299 regions: list[str] | None = None,
300 namespace: str | None = None,
301 offline: bool = False,
302) -> str:
303 """Check which regions would admit a manifest, and whether regions agree.
305 Answers two questions get_job_validation_policy leaves to the caller.
307 Which regions would take this job: the same manifest is evaluated against
308 each region's deployed policy using the code the manifest processor runs,
309 so a job that is admissible in one region and over-cap in another is
310 reported as such instead of being discovered by submitting.
312 Whether the regions still agree: there are no per-region policy overrides,
313 so any field that differs across regions means a region was deployed from a
314 different checkout of cdk.json. That is invisible until a manifest that
315 worked yesterday is refused. trusted_registries is compared with ECR
316 hostnames stripped, since CDK adds those per deployment.
318 Advisory. The cluster is the authoritative gate and this reads a snapshot
319 of its policy, so a reject here is a strong signal, not a verdict.
321 Args:
322 manifest_path: Path to a manifest file or a directory of them.
323 regions: Regions to check. Omit for every configured region.
324 namespace: Namespace to assume for manifests that don't declare one.
325 offline: Read cdk.json instead of calling AWS. Needs no credentials,
326 but reports the CONFIGURED policy rather than the deployed one, and
327 a deployed region trusts ECR registries cdk.json never mentions --
328 so an image rejection may be a false positive.
329 """
330 args = ["jobs", "check-policy", manifest_path]
331 for region in regions or []:
332 args += ["-r", region]
333 if namespace:
334 args += ["-n", namespace]
335 if offline:
336 args.append("--offline")
337 return cli_runner._run_cli(*args)
340@mcp.tool(tags={"safe", "jobs"})
341@audit_logged
342def queue_status(region: str | None = None) -> str:
343 """View SQS queue status (pending, in-flight, DLQ counts).
345 Args:
346 region: Specific region, or omit for all regions.
347 """
348 args = ["jobs", "queue-status"]
349 if region:
350 args += ["-r", region]
351 else:
352 args += ["--all-regions"]
353 return cli_runner._run_cli(*args)