Coverage for gco_mcp/tools/inference.py: 96.73%
166 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-30 21:22 +0000
1"""Inference endpoint management MCP tools."""
3import asyncio
4import contextlib
5import json
6from typing import Any
8import cli_runner
9from audit import audit_logged
10from feature_flags import FLAG_DESTRUCTIVE_OPERATIONS, is_enabled
11from server import mcp
14async def _ctx_warning(message: str) -> None:
15 """Emit ``ctx.warning(...)`` from inside a tool body, no-op when no Context.
17 The destructive ``delete_inference`` tool runs short — we don't need the
18 full long-task progress stack, just an audited warning back to the
19 operator (and the audit log via the middleware spy).
20 """
21 try:
22 from fastmcp.server.dependencies import get_context
24 ctx = get_context()
25 except Exception:
26 return
27 with contextlib.suppress(Exception):
28 await ctx.warning(message)
31#: Serving modes that split an endpoint into separate prefill and decode roles.
32#: These endpoints carry extra Kubernetes resources (per-role Deployments, a
33#: proxy, and a shared store dependency), so removing one is held behind the
34#: destructive-operations flag.
35_DISAGGREGATED_MODES = frozenset({"disaggregated", "both"})
38def _endpoint_is_disaggregated(name: str) -> bool:
39 """Return True when the named endpoint runs a split prefill/decode topology.
41 Reads the endpoint's persisted spec (read-only) and inspects its
42 ``mooncake`` block. A ``mode`` of ``disaggregated`` or ``both`` runs
43 separate prefill and decode roles. Anything else — a plain
44 single-Deployment endpoint, an endpoint without a ``mooncake`` block, or a
45 lookup that cannot be parsed — is treated as not disaggregated.
46 """
47 raw = cli_runner._run_cli("inference", "status", name)
48 try:
49 data = json.loads(raw)
50 except TypeError, ValueError:
51 return False
52 if not isinstance(data, dict):
53 return False
54 spec = data.get("spec")
55 if not isinstance(spec, dict):
56 return False
57 mooncake = spec.get("mooncake")
58 if not isinstance(mooncake, dict):
59 return False
60 return mooncake.get("mode") in _DISAGGREGATED_MODES
63@mcp.tool(tags={"low-risk", "inference"})
64@audit_logged
65def deploy_inference(
66 name: str,
67 image: str,
68 gpu_count: int = 1,
69 replicas: int = 1,
70 port: int = 8000,
71 region: str | None = None,
72 env_vars: list[str] | None = None,
73) -> str:
74 """Deploy an inference endpoint across regions.
76 Args:
77 name: Endpoint name (e.g. my-llm).
78 image: Container image (e.g. vllm/vllm-openai:v0.24.0).
79 gpu_count: GPUs per replica.
80 replicas: Number of replicas per region.
81 port: Container port.
82 region: Target region(s). Omit for all deployed regions.
83 env_vars: Environment variables as KEY=VALUE strings.
84 """
85 args = [
86 "inference",
87 "deploy",
88 name,
89 "-i",
90 image,
91 "--gpu-count",
92 str(gpu_count),
93 "--replicas",
94 str(replicas),
95 "--port",
96 str(port),
97 ]
98 if region:
99 args += ["-r", region]
100 for env in env_vars or []:
101 args += ["-e", env]
102 return cli_runner._run_cli(*args)
105@mcp.tool(tags={"safe", "inference"})
106@audit_logged
107def list_inference_endpoints(state: str | None = None, region: str | None = None) -> str:
108 """List all inference endpoints.
110 Args:
111 state: Filter by state (deploying, running, stopped, deleted).
112 region: Filter by region.
113 """
114 args = ["inference", "list"]
115 if state: 115 ↛ 117line 115 didn't jump to line 117 because the condition on line 115 was always true
116 args += ["--state", state]
117 if region: 117 ↛ 118line 117 didn't jump to line 118 because the condition on line 117 was never true
118 args += ["-r", region]
119 return cli_runner._run_cli(*args)
122@mcp.tool(tags={"safe", "inference"})
123@audit_logged
124def inference_status(name: str) -> str:
125 """Get detailed status of an inference endpoint including per-region breakdown.
127 Args:
128 name: Endpoint name.
129 """
130 return cli_runner._run_cli("inference", "status", name)
133@mcp.tool(tags={"low-risk", "inference"})
134@audit_logged
135def scale_inference(name: str, replicas: int) -> str:
136 """Scale an inference endpoint.
138 Args:
139 name: Endpoint name.
140 replicas: Target replica count.
141 """
142 return cli_runner._run_cli("inference", "scale", name, "--replicas", str(replicas))
145@mcp.tool(tags={"low-risk", "inference"})
146@audit_logged
147def update_inference_image(name: str, image: str) -> str:
148 """Rolling update of an inference endpoint's container image.
150 Args:
151 name: Endpoint name.
152 image: New container image.
153 """
154 return cli_runner._run_cli("inference", "update-image", name, "-i", image)
157@mcp.tool(tags={"low-risk", "inference"})
158@audit_logged
159def stop_inference(name: str) -> str:
160 """Stop an inference endpoint (scales to zero, keeps config).
162 Args:
163 name: Endpoint name.
164 """
165 return cli_runner._run_cli("inference", "stop", name, "-y")
168@mcp.tool(tags={"low-risk", "inference"})
169@audit_logged
170def start_inference(name: str) -> str:
171 """Start a stopped inference endpoint.
173 Args:
174 name: Endpoint name.
175 """
176 return cli_runner._run_cli("inference", "start", name)
179@mcp.tool(tags={"low-risk", "inference"})
180@audit_logged
181def deploy_disaggregated_inference(
182 name: str,
183 image: str | None = None,
184 prefill: int = 1,
185 decode: int = 1,
186 mooncake_mode: str = "disaggregated",
187 gpu_count: int = 1,
188 port: int = 8000,
189 region: str | None = None,
190 env_vars: list[str] | None = None,
191) -> str:
192 """Deploy a split prefill/decode (XpYd) inference endpoint.
194 Registers an endpoint that serves the prefill and decode phases on
195 separate roles backed by a shared KV-cache store. When no image is
196 supplied the deploy falls back to the default upstream Mooncake-enabled
197 vLLM image.
199 Args:
200 name: Endpoint name (e.g. llama-pd).
201 image: Container image. Omit to use the default upstream
202 Mooncake-enabled vLLM image.
203 prefill: Prefill (X) instance count for the topology.
204 decode: Decode (Y) instance count for the topology.
205 mooncake_mode: Serving mode: disaggregated, store, or both.
206 gpu_count: GPUs per replica.
207 port: Container port.
208 region: Target region(s). Omit for all deployed regions.
209 env_vars: Environment variables as KEY=VALUE strings.
210 """
211 args = [
212 "inference",
213 "deploy",
214 name,
215 "--mooncake-mode",
216 mooncake_mode,
217 "--prefill-replicas",
218 str(prefill),
219 "--decode-replicas",
220 str(decode),
221 "--gpu-count",
222 str(gpu_count),
223 "--port",
224 str(port),
225 ]
226 if image:
227 args += ["-i", image]
228 if region:
229 args += ["-r", region]
230 for env in env_vars or []:
231 args += ["-e", env]
232 return cli_runner._run_cli(*args)
235@mcp.tool(tags={"low-risk", "inference"})
236@audit_logged
237def set_mooncake_topology(name: str, prefill: int, decode: int) -> str:
238 """Resize a disaggregated endpoint's prefill/decode topology.
240 Updates the prefill (X) and decode (Y) instance counts and re-triggers
241 reconciliation so each region adjusts its role replica counts. Both
242 counts must be integers in the range 1..1000.
244 Args:
245 name: Endpoint name.
246 prefill: New prefill (X) instance count.
247 decode: New decode (Y) instance count.
248 """
249 return cli_runner._run_cli(
250 "inference",
251 "set-topology",
252 name,
253 "--prefill",
254 str(prefill),
255 "--decode",
256 str(decode),
257 )
260@mcp.tool(tags={"low-risk", "inference"})
261@audit_logged
262def configure_mooncake_store(
263 endpoint_name: str,
264 cold_tier: bool | None = None,
265 offload: str | None = None,
266 global_segment_size: int | None = None,
267 local_buffer_size: int | None = None,
268 enabled: bool | None = None,
269) -> str:
270 """Update a Mooncake endpoint's shared KV-cache store configuration.
272 At least one optional setting must be supplied. Enabling the cold tier also
273 enables the shared store it extends; the CLI merges supplied fields onto
274 the endpoint's existing store block before reconciliation.
276 Args:
277 endpoint_name: Endpoint whose store configuration should change.
278 cold_tier: Enable or disable the asynchronous per-region S3 cold tier.
279 offload: Spill tier: ``cpu``, ``disk``, or ``none``.
280 global_segment_size: Global KV-cache segment size in bytes.
281 local_buffer_size: Local KV-cache buffer size in bytes.
282 enabled: Enable or disable the shared store itself.
283 """
284 if all(
285 value is None
286 for value in (cold_tier, offload, global_segment_size, local_buffer_size, enabled)
287 ):
288 raise ValueError("At least one Mooncake store setting must be supplied")
290 args = ["inference", "configure-store", endpoint_name]
291 if cold_tier is not None: 291 ↛ 293line 291 didn't jump to line 293 because the condition on line 291 was always true
292 args.append("--cold-tier" if cold_tier else "--no-cold-tier")
293 if offload is not None:
294 args += ["--offload", offload]
295 if global_segment_size is not None:
296 args += ["--global-segment-size", str(global_segment_size)]
297 if local_buffer_size is not None:
298 args += ["--local-buffer-size", str(local_buffer_size)]
299 if enabled is not None: 299 ↛ 301line 299 didn't jump to line 301 because the condition on line 299 was always true
300 args.append("--enable-store" if enabled else "--disable-store")
301 return cli_runner._run_cli(*args)
304@mcp.tool(tags={"safe", "inference"})
305@audit_logged
306def mooncake_topology_status(name: str) -> str:
307 """Show a disaggregated endpoint's per-role topology status.
309 Returns the endpoint record including its role-keyed region status, so
310 the prefill and decode roles can be inspected per region.
312 Args:
313 name: Endpoint name.
314 """
315 return cli_runner._run_cli("inference", "status", name)
318@mcp.tool(tags={"low-risk", "inference"})
319@audit_logged
320async def populate_kv_cache(endpoint_name: str, local_path: str, region: str) -> str:
321 """`gco inference populate-kv` — upload data into an endpoint's KV-cache cold tier.
323 Uploads a local file or directory to the region's general-purpose bucket
324 under the cold-tier key prefix the endpoint reads from, pre-warming its
325 Mooncake KV-cache prefix cache. The endpoint must be deployed with the cold
326 tier enabled (deploy with --mooncake-cold-tier, or configure it via
327 configure-store) for its pods to read the uploaded data.
329 Args:
330 endpoint_name: Endpoint whose KV-cache cold tier receives the data.
331 local_path: Local file or directory to upload.
332 region: Region whose general-purpose bucket backs the cold tier.
333 """
334 return await asyncio.to_thread(
335 cli_runner._run_cli,
336 "inference",
337 "populate-kv",
338 endpoint_name,
339 local_path,
340 "-r",
341 region,
342 )
345if is_enabled(FLAG_DESTRUCTIVE_OPERATIONS):
347 @mcp.tool(tags={"destructive", "inference"})
348 @audit_logged
349 async def delete_inference(name: str) -> str:
350 """[gated by GCO_ENABLE_DESTRUCTIVE_OPERATIONS] destructive.
352 Delete an inference endpoint. Cannot be undone — the endpoint, its
353 DynamoDB record, and the underlying Kubernetes resources are removed.
355 Deleting a disaggregated (split prefill/decode) endpoint is refused at
356 invocation when GCO_ENABLE_DESTRUCTIVE_OPERATIONS is not enabled: the
357 request is rejected and no endpoint resource is removed.
359 Args:
360 name: Endpoint name.
361 """
362 if not is_enabled(FLAG_DESTRUCTIVE_OPERATIONS) and await asyncio.to_thread(
363 _endpoint_is_disaggregated, name
364 ):
365 message = (
366 f"Refusing to delete disaggregated endpoint {name!r}: destructive "
367 "operations are disabled. Set GCO_ENABLE_DESTRUCTIVE_OPERATIONS=true "
368 "to allow this deletion."
369 )
370 await _ctx_warning(message)
371 return json.dumps({"error": message, "destructive_operations_disabled": True})
372 await _ctx_warning(f"Deleting inference endpoint {name!r} — this cannot be undone.")
373 return await asyncio.to_thread(cli_runner._run_cli, "inference", "delete", name, "-y")
376@mcp.tool(tags={"low-risk", "inference"})
377@audit_logged
378def canary_deploy(name: str, image: str, weight: int = 10) -> str:
379 """Start a canary deployment (A/B test a new image version).
381 Args:
382 name: Endpoint name.
383 image: New image to canary.
384 weight: Percentage of traffic to send to canary (1-99).
385 """
386 return cli_runner._run_cli("inference", "canary", name, "-i", image, "--weight", str(weight))
389@mcp.tool(tags={"low-risk", "inference"})
390@audit_logged
391def promote_canary(name: str) -> str:
392 """Promote canary to primary (100% traffic to new version).
394 Args:
395 name: Endpoint name.
396 """
397 return cli_runner._run_cli("inference", "promote", name, "-y")
400@mcp.tool(tags={"low-risk", "inference"})
401@audit_logged
402def rollback_canary(name: str) -> str:
403 """Rollback canary (remove canary, 100% traffic to primary).
405 Args:
406 name: Endpoint name.
407 """
408 return cli_runner._run_cli("inference", "rollback", name, "-y")
411@mcp.tool(tags={"safe", "inference"})
412@audit_logged
413def invoke_inference(
414 name: str,
415 prompt: str,
416 max_tokens: int = 100,
417 api_path: str | None = None,
418 region: str | None = None,
419) -> str:
420 """Send a prompt to an inference endpoint and return the generated text.
422 Automatically discovers the endpoint's ingress path, detects the serving
423 framework (vLLM, TGI, Triton), and routes the request through the API
424 Gateway with SigV4 authentication.
426 Use this for single-turn text completions. For multi-turn conversations
427 with chat models, use chat_inference instead.
429 Args:
430 name: Endpoint name (e.g. my-llm).
431 prompt: Text prompt to send to the model.
432 max_tokens: Maximum tokens to generate (default: 100).
433 api_path: Override the API sub-path (default: auto-detect from framework).
434 region: Target region for the request (default: nearest via Global Accelerator).
436 This MCP tool intentionally returns one buffered string because the MCP
437 runner has no incremental output contract. The API Gateway/Lambda route
438 supports response streaming; use ``gco inference invoke --stream`` when
439 token-by-token output is required.
440 """
441 args = ["inference", "invoke", name, "-p", prompt, "--max-tokens", str(max_tokens)]
442 if api_path:
443 args += ["--path", api_path]
444 if region:
445 args += ["-r", region]
446 return cli_runner._run_cli(*args)
449@mcp.tool(tags={"safe", "inference"})
450@audit_logged
451def chat_inference(
452 name: str,
453 messages: list[dict[str, str]],
454 max_tokens: int = 256,
455 temperature: float | None = None,
456 region: str | None = None,
457) -> str:
458 """Send a multi-turn chat conversation to an inference endpoint.
460 Sends an OpenAI-compatible /v1/chat/completions request. Works with
461 vLLM, TGI (with --api-protocol openai), and any OpenAI-compatible server.
463 Each message in the list should have 'role' (system/user/assistant) and
464 'content' keys.
466 Args:
467 name: Endpoint name (e.g. my-llm).
468 messages: List of chat messages, e.g. [{"role": "user", "content": "Hello"}].
469 max_tokens: Maximum tokens to generate (default: 256).
470 temperature: Sampling temperature (optional, server default if omitted).
471 region: Target region for the request.
473 This MCP tool intentionally returns one buffered string because the MCP
474 runner has no incremental output contract. The API Gateway/Lambda route
475 supports response streaming; use ``gco inference invoke --stream`` when
476 token-by-token output is required.
477 """
478 body: dict[str, Any] = {"messages": messages, "max_tokens": max_tokens}
479 if temperature is not None:
480 body["temperature"] = temperature
481 data_str = json.dumps(body)
482 args = ["inference", "invoke", name, "-d", data_str, "--path", "/v1/chat/completions"]
483 if region:
484 args += ["-r", region]
485 return cli_runner._run_cli(*args)
488@mcp.tool(tags={"safe", "inference"})
489@audit_logged
490def inference_health(name: str, region: str | None = None) -> str:
491 """Check if an inference endpoint is healthy and ready to serve requests.
493 Hits the endpoint's health check path and returns status and latency.
494 Useful to verify readiness before sending inference requests.
496 Args:
497 name: Endpoint name.
498 region: Target region to check (default: nearest via Global Accelerator).
499 """
500 args = ["inference", "health", name]
501 if region:
502 args += ["-r", region]
503 return cli_runner._run_cli(*args)
506@mcp.tool(tags={"safe", "inference"})
507@audit_logged
508def list_endpoint_models(name: str, region: str | None = None) -> str:
509 """List models loaded on an inference endpoint.
511 Queries the endpoint's /v1/models path (OpenAI-compatible) to discover
512 which models are loaded, their context length, and other metadata.
513 Works with vLLM and other OpenAI-compatible servers.
515 Args:
516 name: Endpoint name.
517 region: Target region to query.
518 """
519 args = ["inference", "models", name]
520 if region:
521 args += ["-r", region]
522 return cli_runner._run_cli(*args)