Coverage for gco_mcp / tools / inference.py: 100.00%
177 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-10-28 16:19 +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 framework: str | None = None,
74 extra_args: list[str] | None = None,
75) -> str:
76 """Deploy an inference endpoint across regions.
78 The model to serve is ``MODEL=<id-or-path>`` in ``env_vars``: vLLM and
79 SGLang both take it from there. Servers without a strict adapter, such as
80 Triton, need their command in ``extra_args`` (for Triton,
81 ``["tritonserver", "--model-repository=/models/<name>"]``).
83 Args:
84 name: Endpoint name (e.g. my-llm).
85 image: Container image (e.g. vllm/vllm-openai:v0.32.0).
86 gpu_count: GPUs per replica.
87 replicas: Number of replicas per region.
88 port: Container port.
89 region: Target region(s). Omit for all deployed regions.
90 env_vars: Environment variables as KEY=VALUE strings.
91 framework: Strict serving adapter, ``vllm`` or ``sglang``. Omit to let
92 the image name decide.
93 extra_args: Container arguments, one token per entry, appended in order.
94 """
95 args = [
96 "inference",
97 "deploy",
98 name,
99 "-i",
100 image,
101 "--gpu-count",
102 str(gpu_count),
103 "--replicas",
104 str(replicas),
105 "--port",
106 str(port),
107 ]
108 if region:
109 args += ["-r", region]
110 for env in env_vars or []:
111 args += ["-e", env]
112 if framework:
113 args += ["--framework", framework]
114 # The ``--extra-args=<token>`` form keeps a token that starts with "--"
115 # from being read as an option of its own.
116 args += [f"--extra-args={token}" for token in extra_args or []]
117 return cli_runner._run_cli(*args)
120@mcp.tool(tags={"safe", "inference"})
121@audit_logged
122def list_inference_endpoints(state: str | None = None, region: str | None = None) -> str:
123 """List all inference endpoints.
125 Args:
126 state: Filter by state (deploying, running, stopped, deleted).
127 region: Filter by region.
128 """
129 args = ["inference", "list"]
130 if state:
131 args += ["--state", state]
132 if region:
133 args += ["-r", region]
134 return cli_runner._run_cli(*args)
137@mcp.tool(tags={"safe", "inference"})
138@audit_logged
139def inference_status(name: str) -> str:
140 """Get detailed status of an inference endpoint including per-region breakdown.
142 Args:
143 name: Endpoint name.
144 """
145 return cli_runner._run_cli("inference", "status", name)
148@mcp.tool(tags={"low-risk", "inference"})
149@audit_logged
150def scale_inference(name: str, replicas: int) -> str:
151 """Scale an inference endpoint.
153 Args:
154 name: Endpoint name.
155 replicas: Target replica count.
156 """
157 return cli_runner._run_cli("inference", "scale", name, "--replicas", str(replicas))
160@mcp.tool(tags={"low-risk", "inference"})
161@audit_logged
162def update_inference_image(name: str, image: str) -> str:
163 """Rolling update of an inference endpoint's container image.
165 Args:
166 name: Endpoint name.
167 image: New container image.
168 """
169 return cli_runner._run_cli("inference", "update-image", name, "-i", image)
172@mcp.tool(tags={"low-risk", "inference"})
173@audit_logged
174def stop_inference(name: str) -> str:
175 """Stop an inference endpoint (scales to zero, keeps config).
177 Args:
178 name: Endpoint name.
179 """
180 return cli_runner._run_cli("inference", "stop", name, "-y")
183@mcp.tool(tags={"low-risk", "inference"})
184@audit_logged
185def start_inference(name: str) -> str:
186 """Start a stopped inference endpoint.
188 Args:
189 name: Endpoint name.
190 """
191 return cli_runner._run_cli("inference", "start", name)
194@mcp.tool(tags={"low-risk", "inference"})
195@audit_logged
196def add_inference_region(name: str, region: str) -> str:
197 """Add a Region to an existing inference endpoint.
199 An endpoint keeps the Region list it was deployed with, so a Region added
200 to the deployment later does not run it until it is added here. The
201 inference monitor in that Region then creates the endpoint from its stored
202 spec. The Region must already run a GCO regional stack.
204 Args:
205 name: Endpoint name.
206 region: Deployed GCO Region to add (e.g. ``us-west-2``).
207 """
208 return cli_runner._run_cli("inference", "add-region", name, region)
211@mcp.tool(tags={"low-risk", "inference"})
212@audit_logged
213def remove_inference_region(name: str, region: str) -> str:
214 """Remove a Region from an inference endpoint.
216 The inference monitor in that Region deletes the endpoint's resources
217 there; the endpoint keeps serving from its other Regions and the Region can
218 be added back. Removing the endpoint's last Region is refused.
220 Args:
221 name: Endpoint name.
222 region: Region to remove.
223 """
224 return cli_runner._run_cli("inference", "remove-region", name, region, "-y")
227@mcp.tool(tags={"low-risk", "inference"})
228@audit_logged
229def deploy_disaggregated_inference(
230 name: str,
231 image: str | None = None,
232 prefill: int = 1,
233 decode: int = 1,
234 mooncake_mode: str = "disaggregated",
235 gpu_count: int = 1,
236 port: int = 8000,
237 region: str | None = None,
238 env_vars: list[str] | None = None,
239) -> str:
240 """Deploy a split prefill/decode (XpYd) inference endpoint.
242 Registers an endpoint that serves the prefill and decode phases on
243 separate roles backed by a shared KV-cache store. When no image is
244 supplied the deploy falls back to the default upstream Mooncake-enabled
245 vLLM image.
247 Args:
248 name: Endpoint name (e.g. llama-pd).
249 image: Container image. Omit to use the default upstream
250 Mooncake-enabled vLLM image.
251 prefill: Prefill (X) instance count for the topology.
252 decode: Decode (Y) instance count for the topology.
253 mooncake_mode: Serving mode: disaggregated, store, or both.
254 gpu_count: GPUs per replica.
255 port: Container port.
256 region: Target region(s). Omit for all deployed regions.
257 env_vars: Environment variables as KEY=VALUE strings.
258 """
259 args = [
260 "inference",
261 "deploy",
262 name,
263 "--mooncake-mode",
264 mooncake_mode,
265 "--prefill-replicas",
266 str(prefill),
267 "--decode-replicas",
268 str(decode),
269 "--gpu-count",
270 str(gpu_count),
271 "--port",
272 str(port),
273 ]
274 if image:
275 args += ["-i", image]
276 if region:
277 args += ["-r", region]
278 for env in env_vars or []:
279 args += ["-e", env]
280 return cli_runner._run_cli(*args)
283@mcp.tool(tags={"low-risk", "inference"})
284@audit_logged
285def set_mooncake_topology(name: str, prefill: int, decode: int) -> str:
286 """Resize a disaggregated endpoint's prefill/decode topology.
288 Updates the prefill (X) and decode (Y) instance counts and re-triggers
289 reconciliation so each region adjusts its role replica counts. Both
290 counts must be integers in the range 1..1000.
292 Args:
293 name: Endpoint name.
294 prefill: New prefill (X) instance count.
295 decode: New decode (Y) instance count.
296 """
297 return cli_runner._run_cli(
298 "inference",
299 "set-topology",
300 name,
301 "--prefill",
302 str(prefill),
303 "--decode",
304 str(decode),
305 )
308@mcp.tool(tags={"low-risk", "inference"})
309@audit_logged
310def configure_mooncake_store(
311 endpoint_name: str,
312 cold_tier: bool | None = None,
313 offload: str | None = None,
314 global_segment_size: int | None = None,
315 local_buffer_size: int | None = None,
316 enabled: bool | None = None,
317) -> str:
318 """Update a Mooncake endpoint's shared KV-cache store configuration.
320 At least one optional setting must be supplied. Enabling the cold tier also
321 enables the shared store it extends; the CLI merges supplied fields onto
322 the endpoint's existing store block before reconciliation.
324 Args:
325 endpoint_name: Endpoint whose store configuration should change.
326 cold_tier: Enable or disable the asynchronous per-region S3 cold tier.
327 offload: Spill tier: ``cpu``, ``disk``, or ``none``.
328 global_segment_size: Global KV-cache segment size in bytes.
329 local_buffer_size: Local KV-cache buffer size in bytes.
330 enabled: Enable or disable the shared store itself.
331 """
332 if all(
333 value is None
334 for value in (cold_tier, offload, global_segment_size, local_buffer_size, enabled)
335 ):
336 raise ValueError("At least one Mooncake store setting must be supplied")
338 args = ["inference", "configure-store", endpoint_name]
339 if cold_tier is not None:
340 args.append("--cold-tier" if cold_tier else "--no-cold-tier")
341 if offload is not None:
342 args += ["--offload", offload]
343 if global_segment_size is not None:
344 args += ["--global-segment-size", str(global_segment_size)]
345 if local_buffer_size is not None:
346 args += ["--local-buffer-size", str(local_buffer_size)]
347 if enabled is not None:
348 args.append("--enable-store" if enabled else "--disable-store")
349 return cli_runner._run_cli(*args)
352@mcp.tool(tags={"safe", "inference"})
353@audit_logged
354def mooncake_topology_status(name: str) -> str:
355 """Show a disaggregated endpoint's per-role topology status.
357 Returns the endpoint record including its role-keyed region status, so
358 the prefill and decode roles can be inspected per region.
360 Args:
361 name: Endpoint name.
362 """
363 return cli_runner._run_cli("inference", "status", name)
366@mcp.tool(tags={"low-risk", "inference"})
367@audit_logged
368async def populate_kv_cache(endpoint_name: str, local_path: str, region: str) -> str:
369 """`gco inference populate-kv` — upload data into an endpoint's KV-cache cold tier.
371 Uploads a local file or directory to the region's general-purpose bucket
372 under the cold-tier key prefix the endpoint reads from, pre-warming its
373 Mooncake KV-cache prefix cache. The endpoint must be deployed with the cold
374 tier enabled (deploy with --mooncake-cold-tier, or configure it via
375 configure-store) for its pods to read the uploaded data.
377 Args:
378 endpoint_name: Endpoint whose KV-cache cold tier receives the data.
379 local_path: Local file or directory to upload.
380 region: Region whose general-purpose bucket backs the cold tier.
381 """
382 return await asyncio.to_thread(
383 cli_runner._run_cli,
384 "inference",
385 "populate-kv",
386 endpoint_name,
387 local_path,
388 "-r",
389 region,
390 )
393if is_enabled(FLAG_DESTRUCTIVE_OPERATIONS):
395 @mcp.tool(tags={"destructive", "inference"})
396 @audit_logged
397 async def delete_inference(name: str) -> str:
398 """[gated by GCO_ENABLE_DESTRUCTIVE_OPERATIONS] destructive.
400 Delete an inference endpoint. Cannot be undone — the endpoint, its
401 DynamoDB record, and the underlying Kubernetes resources are removed.
403 Deleting a disaggregated (split prefill/decode) endpoint is refused at
404 invocation when GCO_ENABLE_DESTRUCTIVE_OPERATIONS is not enabled: the
405 request is rejected and no endpoint resource is removed.
407 Args:
408 name: Endpoint name.
409 """
410 if not is_enabled(FLAG_DESTRUCTIVE_OPERATIONS) and await asyncio.to_thread(
411 _endpoint_is_disaggregated, name
412 ):
413 message = (
414 f"Refusing to delete disaggregated endpoint {name!r}: destructive "
415 "operations are disabled. Set GCO_ENABLE_DESTRUCTIVE_OPERATIONS=true "
416 "to allow this deletion."
417 )
418 await _ctx_warning(message)
419 return json.dumps({"error": message, "destructive_operations_disabled": True})
420 await _ctx_warning(f"Deleting inference endpoint {name!r} — this cannot be undone.")
421 return await asyncio.to_thread(cli_runner._run_cli, "inference", "delete", name, "-y")
424@mcp.tool(tags={"low-risk", "inference"})
425@audit_logged
426def canary_deploy(name: str, image: str, weight: int = 10) -> str:
427 """Start a canary deployment (A/B test a new image version).
429 Args:
430 name: Endpoint name.
431 image: New image to canary.
432 weight: Percentage of traffic to send to canary (1-99).
433 """
434 return cli_runner._run_cli("inference", "canary", name, "-i", image, "--weight", str(weight))
437@mcp.tool(tags={"low-risk", "inference"})
438@audit_logged
439def promote_canary(name: str) -> str:
440 """Promote canary to primary (100% traffic to new version).
442 Args:
443 name: Endpoint name.
444 """
445 return cli_runner._run_cli("inference", "promote", name, "-y")
448@mcp.tool(tags={"low-risk", "inference"})
449@audit_logged
450def rollback_canary(name: str) -> str:
451 """Rollback canary (remove canary, 100% traffic to primary).
453 Args:
454 name: Endpoint name.
455 """
456 return cli_runner._run_cli("inference", "rollback", name, "-y")
459@mcp.tool(tags={"safe", "inference"})
460@audit_logged
461def invoke_inference(
462 name: str,
463 prompt: str,
464 max_tokens: int = 100,
465 api_path: str | None = None,
466 region: str | None = None,
467) -> str:
468 """Send a prompt to an inference endpoint and return the generated text.
470 Automatically discovers the endpoint's ingress path, detects the serving
471 framework (vLLM or SGLang), and routes the request through the API
472 Gateway with SigV4 authentication. Triton takes tensors rather than a
473 prompt, so a Triton endpoint is refused with the --path/--data form to use
474 through the CLI.
476 Use this for single-turn text completions. For multi-turn conversations
477 with chat models, use chat_inference instead.
479 Args:
480 name: Endpoint name (e.g. my-llm).
481 prompt: Text prompt to send to the model.
482 max_tokens: Maximum tokens to generate (default: 100).
483 api_path: Override the API sub-path (default: auto-detect from framework).
484 region: Target region for the request (default: nearest via Global Accelerator).
486 This MCP tool intentionally returns one buffered string because the MCP
487 runner has no incremental output contract. The API Gateway/Lambda route
488 supports response streaming; use ``gco inference invoke --stream`` when
489 token-by-token output is required.
490 """
491 args = ["inference", "invoke", name, "-p", prompt, "--max-tokens", str(max_tokens)]
492 if api_path:
493 args += ["--path", api_path]
494 if region:
495 args += ["-r", region]
496 return cli_runner._run_cli(*args)
499@mcp.tool(tags={"safe", "inference"})
500@audit_logged
501def chat_inference(
502 name: str,
503 messages: list[dict[str, str]],
504 max_tokens: int = 256,
505 temperature: float | None = None,
506 region: str | None = None,
507) -> str:
508 """Send a multi-turn chat conversation to an inference endpoint.
510 Sends an OpenAI-compatible /v1/chat/completions request. Works with
511 vLLM, SGLang, and any OpenAI-compatible server.
513 Each message in the list should have 'role' (system/user/assistant) and
514 'content' keys.
516 Args:
517 name: Endpoint name (e.g. my-llm).
518 messages: List of chat messages, e.g. [{"role": "user", "content": "Hello"}].
519 max_tokens: Maximum tokens to generate (default: 256).
520 temperature: Sampling temperature (optional, server default if omitted).
521 region: Target region for the request.
523 This MCP tool intentionally returns one buffered string because the MCP
524 runner has no incremental output contract. The API Gateway/Lambda route
525 supports response streaming; use ``gco inference invoke --stream`` when
526 token-by-token output is required.
527 """
528 body: dict[str, Any] = {"messages": messages, "max_tokens": max_tokens}
529 if temperature is not None:
530 body["temperature"] = temperature
531 data_str = json.dumps(body)
532 args = ["inference", "invoke", name, "-d", data_str, "--path", "/v1/chat/completions"]
533 if region:
534 args += ["-r", region]
535 return cli_runner._run_cli(*args)
538@mcp.tool(tags={"safe", "inference"})
539@audit_logged
540def inference_health(name: str, region: str | None = None) -> str:
541 """Check if an inference endpoint is healthy and ready to serve requests.
543 Hits the endpoint's health check path and returns status and latency.
544 Useful to verify readiness before sending inference requests.
546 Args:
547 name: Endpoint name.
548 region: Target region to check (default: nearest via Global Accelerator).
549 """
550 args = ["inference", "health", name]
551 if region:
552 args += ["-r", region]
553 return cli_runner._run_cli(*args)
556@mcp.tool(tags={"safe", "inference"})
557@audit_logged
558def list_endpoint_models(name: str, region: str | None = None) -> str:
559 """List models loaded on an inference endpoint.
561 Queries the endpoint's /v1/models path (OpenAI-compatible) to discover
562 which models are loaded, their context length, and other metadata.
563 Works with vLLM and other OpenAI-compatible servers; SGLang reports its
564 launcher's model identity, and Triton its /v2/models/stats inventory.
566 Args:
567 name: Endpoint name.
568 region: Target region to query.
569 """
570 args = ["inference", "models", name]
571 if region:
572 args += ["-r", region]
573 return cli_runner._run_cli(*args)