Coverage for cli / commands / inference_cmd.py: 100.00%
729 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-09-14 22:07 +0000
1"""Inference endpoint commands."""
3import codecs
4import sys
5from email.message import Message
6from typing import Any
8import click
10from ..config import GCOConfig
11from ..output import (
12 confirm,
13 emit_structured_document,
14 get_output_formatter,
15 interactive_echo,
16)
18pass_config = click.make_pass_decorator(GCOConfig, ensure=True)
21@click.group()
22@pass_config
23def inference(config: Any) -> None:
24 """Manage multi-region inference endpoints."""
25 pass
28@inference.command("deploy")
29@click.argument("endpoint_name")
30@click.option(
31 "--image",
32 "-i",
33 default=None,
34 help="Container image (e.g. vllm/vllm-openai:v0.29.0). Optional with "
35 "--mooncake-mode: falls back to the default upstream Mooncake-enabled vLLM image.",
36)
37@click.option(
38 "--framework",
39 type=click.Choice(["vllm", "tgi"]),
40 default=None,
41 help="Explicit serving runtime contract; persisted for renderer/probe behavior.",
42)
43@click.option(
44 "--region",
45 "-r",
46 multiple=True,
47 help="Target region(s). Repeatable. Default: all deployed regions",
48)
49@click.option("--replicas", default=1, help="Replicas per region (default: 1)")
50@click.option("--gpu-count", default=1, help="GPUs per replica (default: 1)")
51@click.option("--gpu-type", help="GPU instance type hint (e.g. g5.xlarge)")
52@click.option("--port", default=8000, help="Container port (default: 8000)")
53@click.option("--model-path", help="EFS path for model weights")
54@click.option(
55 "--model-source",
56 help="S3 URI for model weights (e.g. s3://bucket/models/llama3). "
57 "Auto-synced to each region via init container.",
58)
59@click.option("--health-path", default="/health", help="Health check path (default: /health)")
60@click.option("--env", "-e", multiple=True, help="Environment variable (KEY=VALUE). Repeatable")
61@click.option("--namespace", "-n", default="gco-inference", help="Kubernetes namespace")
62@click.option("--label", "-l", multiple=True, help="Label (key=value). Repeatable")
63@click.option("--min-replicas", type=int, default=None, help="Autoscaling: minimum replicas")
64@click.option("--max-replicas", type=int, default=None, help="Autoscaling: maximum replicas")
65@click.option(
66 "--autoscale-metric",
67 multiple=True,
68 help="Autoscaling metric (cpu:70, memory:80, gpu:60). Repeatable. Enables "
69 "autoscaling. CPU/memory scale via the native HPA; gpu (and gpu_memory) "
70 "scale on CloudWatch GPU utilization via KEDA.",
71)
72@click.option(
73 "--capacity-type",
74 type=click.Choice(["on-demand", "spot"]),
75 default=None,
76 help="Node capacity type. 'spot' uses cheaper preemptible instances.",
77)
78@click.option(
79 "--extra-args",
80 multiple=True,
81 help="Extra arguments passed to the container (e.g. '--kv-transfer-config {...}'). Repeatable.",
82)
83@click.option(
84 "--accelerator",
85 type=click.Choice(["nvidia", "neuron"]),
86 default="nvidia",
87 help="Accelerator type: 'nvidia' for GPU instances (default), 'neuron' for Trainium/Inferentia.",
88)
89@click.option(
90 "--node-selector",
91 multiple=True,
92 help="Node selector (key=value). Repeatable. E.g. --node-selector eks.amazonaws.com/instance-family=inf2",
93)
94@click.option(
95 "--no-rewrite-image",
96 is_flag=True,
97 default=False,
98 help="Skip the per-region ECR URI rewrite. The image URI is sent verbatim "
99 "to every target region (operator owns cross-region pulls).",
100)
101@click.option(
102 "--mooncake-mode",
103 type=click.Choice(["disaggregated", "store", "both"]),
104 default=None,
105 help="Enable Mooncake serving: 'disaggregated' splits prefill/decode, "
106 "'store' runs a shared KV-cache store, 'both' composes the two. When set "
107 "and -i is omitted, the default upstream Mooncake-enabled vLLM image is used.",
108)
109@click.option(
110 "--prefill-replicas",
111 type=int,
112 default=1,
113 help="Prefill instance count (X in an XpYd topology) for split modes.",
114)
115@click.option(
116 "--decode-replicas",
117 type=int,
118 default=1,
119 help="Decode instance count (Y in an XpYd topology) for split modes.",
120)
121@click.option(
122 "--mooncake-protocol",
123 type=click.Choice(["rdma", "tcp"]),
124 default=None,
125 help="Mooncake transfer intent. 'rdma' (the default) schedules role pods "
126 "on EFA and configures vLLM's connector protocol as 'efa'; 'tcp' is the "
127 "non-EFA fallback. Requires --mooncake-mode.",
128)
129@click.option(
130 "--mooncake-device-name",
131 default=None,
132 help="Network device passed to Mooncake (for example efa_0 or eth0). "
133 "Omit or pass an empty value for auto-detection. Requires --mooncake-mode.",
134)
135@click.option(
136 "--mooncake-autoscale",
137 multiple=True,
138 help="Per-role Mooncake autoscaling as ROLE:MIN:MAX[:METRIC:TARGET ...], "
139 "e.g. 'prefill:1:8' or 'decode:2:16:cpu:70:gpu:60'. Repeatable (one per "
140 "role); append additional METRIC:TARGET pairs to scale a role on multiple "
141 "metrics (cpu/memory via HPA, gpu/gpu_memory via KEDA CloudWatch). Requires "
142 "--mooncake-mode disaggregated|both; populates spec.mooncake.autoscaling "
143 "(distinct from the legacy --autoscale-metric/--min-replicas flags).",
144)
145@click.option(
146 "--mooncake-cold-tier",
147 is_flag=True,
148 default=False,
149 help="Enable the asynchronous per-region S3 cold tier for the shared "
150 "KV-cache store (the cold tier extends the store). Pre-warm it with "
151 "'gco inference populate-kv'. Requires --mooncake-mode store or both.",
152)
153@click.option(
154 "--mooncake-proxy-image",
155 default=None,
156 help="Container image for the prefill-decode proxy (disaggregated/both). "
157 "Defaults to the endpoint image, which bundles the reference proxy.",
158)
159@click.option(
160 "--mooncake-admin-key-secret",
161 default=None,
162 help="Name of an existing Kubernetes Secret holding the prefill-decode "
163 "proxy ADMIN_API_KEY. Optional: when omitted, each region's monitor "
164 "auto-provisions a {name}-admin Secret with a generated key.",
165)
166@pass_config
167def inference_deploy(
168 config: Any,
169 endpoint_name: Any,
170 image: Any,
171 framework: Any,
172 region: Any,
173 replicas: Any,
174 gpu_count: Any,
175 gpu_type: Any,
176 port: Any,
177 model_path: Any,
178 model_source: Any,
179 health_path: Any,
180 env: Any,
181 namespace: Any,
182 label: Any,
183 min_replicas: Any,
184 max_replicas: Any,
185 autoscale_metric: Any,
186 capacity_type: Any,
187 extra_args: Any,
188 accelerator: Any,
189 node_selector: Any,
190 no_rewrite_image: Any,
191 mooncake_mode: Any,
192 prefill_replicas: Any,
193 decode_replicas: Any,
194 mooncake_protocol: Any,
195 mooncake_device_name: Any,
196 mooncake_autoscale: Any,
197 mooncake_cold_tier: Any,
198 mooncake_proxy_image: Any,
199 mooncake_admin_key_secret: Any,
200) -> None:
201 """Deploy an inference endpoint to one or more regions.
203 The endpoint is registered in DynamoDB and the inference_monitor
204 in each target region creates the Kubernetes resources automatically.
206 Examples:
207 gco inference deploy my-llm -i vllm/vllm-openai:v0.29.0
209 gco inference deploy llama3-70b \\
210 -i vllm/vllm-openai:v0.29.0 \\
211 -r us-east-1 -r eu-west-1 \\
212 --replicas 2 --gpu-count 4 \\
213 --model-path /mnt/gco/models/llama3-70b \\
214 -e MODEL_NAME=meta-llama/Llama-3-70B
215 """
216 from ..inference import get_inference_manager
218 formatter = get_output_formatter(config)
220 # Parse env vars and labels
221 env_dict = {}
222 for e_var in env:
223 if "=" in e_var:
224 k, v = e_var.split("=", 1)
225 env_dict[k] = v
227 labels_dict = {}
228 for lbl in label:
229 if "=" in lbl:
230 k, v = lbl.split("=", 1)
231 labels_dict[k] = v
233 node_selector_dict = {}
234 for ns in node_selector:
235 if "=" in ns:
236 k, v = ns.split("=", 1)
237 node_selector_dict[k] = v
239 # Build autoscaling config
240 autoscaling_config = None
241 if autoscale_metric:
242 metrics = []
243 for m in autoscale_metric:
244 if ":" in m:
245 mtype, mtarget = m.split(":", 1)
246 metrics.append({"type": mtype, "target": int(mtarget)})
247 else:
248 metrics.append({"type": m, "target": 70})
249 autoscaling_config = {
250 "enabled": True,
251 "min_replicas": min_replicas or 1,
252 "max_replicas": max_replicas or 10,
253 "metrics": metrics,
254 }
256 # Transfer overrides are meaningful only when a Mooncake block is being
257 # authored. With no override, the monitor resolves the default RDMA intent
258 # to vLLM's explicit EFA connector protocol and auto-detects the device.
259 if (mooncake_protocol is not None or mooncake_device_name is not None) and not mooncake_mode:
260 formatter.print_error(
261 "--mooncake-protocol and --mooncake-device-name require --mooncake-mode."
262 )
263 sys.exit(1)
265 mooncake_transfer_config: dict[str, Any] | None = None
266 if mooncake_protocol is not None or mooncake_device_name is not None:
267 mooncake_transfer_config = {}
268 if mooncake_protocol is not None:
269 mooncake_transfer_config["protocol"] = mooncake_protocol
270 if mooncake_device_name is not None:
271 mooncake_transfer_config["device_name"] = mooncake_device_name
273 # Build per-role Mooncake autoscaling config (spec.mooncake.autoscaling).
274 # This is distinct from the legacy single-Deployment autoscaling above:
275 # each ROLE:MIN:MAX token sets a role's bounds, and any number of trailing
276 # METRIC:TARGET pairs add scaling signals for that role. Bounds and metrics
277 # are validated fail-fast in the deploy path before anything is persisted.
278 mooncake_autoscaling_config: dict[str, Any] | None = None
279 if mooncake_autoscale:
280 if not mooncake_mode:
281 formatter.print_error(
282 "--mooncake-autoscale requires --mooncake-mode (disaggregated or both)."
283 )
284 sys.exit(1)
285 mooncake_autoscaling_config = {"enabled": True}
286 for entry in mooncake_autoscale:
287 parts = entry.split(":")
288 # ROLE:MIN:MAX, then zero or more METRIC:TARGET pairs.
289 if len(parts) < 3 or (len(parts) - 3) % 2 != 0:
290 formatter.print_error(
291 f"Invalid --mooncake-autoscale value '{entry}'. Expected "
292 "ROLE:MIN:MAX optionally followed by METRIC:TARGET pairs."
293 )
294 sys.exit(1)
295 role = parts[0]
296 if role not in ("prefill", "decode"):
297 formatter.print_error(
298 f"Invalid --mooncake-autoscale role '{role}'. Expected 'prefill' or 'decode'."
299 )
300 sys.exit(1)
301 try:
302 role_block: dict[str, Any] = {
303 "min_replicas": int(parts[1]),
304 "max_replicas": int(parts[2]),
305 }
306 metric_tokens = parts[3:]
307 metrics = [
308 {"type": metric_tokens[i], "target": int(metric_tokens[i + 1])}
309 for i in range(0, len(metric_tokens), 2)
310 ]
311 if metrics:
312 role_block["metrics"] = metrics
313 except ValueError:
314 formatter.print_error(
315 f"Invalid --mooncake-autoscale numbers in '{entry}'. MIN, MAX, "
316 "and each TARGET must be integers."
317 )
318 sys.exit(1)
319 mooncake_autoscaling_config[role] = role_block
321 # --mooncake-cold-tier opts into the async per-region S3 cold tier, which
322 # extends the shared store, so it only applies to store/both modes.
323 if mooncake_cold_tier and mooncake_mode not in ("store", "both"):
324 formatter.print_error(
325 "--mooncake-cold-tier requires --mooncake-mode store or both "
326 "(the cold tier extends the shared KV-cache store)."
327 )
328 sys.exit(1)
330 mooncake_store_config: dict[str, Any] | None = None
331 if mooncake_cold_tier:
332 mooncake_store_config = {"enabled": True, "cold_tier_enabled": True}
334 # Configure the prefill-decode proxy that fronts split modes: an explicit
335 # image (otherwise it defaults to the endpoint image) and the name of the
336 # Kubernetes Secret holding its ADMIN_API_KEY.
337 mooncake_proxy_config: dict[str, Any] | None = None
338 if mooncake_proxy_image or mooncake_admin_key_secret:
339 mooncake_proxy_config = {}
340 if mooncake_proxy_image:
341 mooncake_proxy_config["image"] = mooncake_proxy_image
342 if mooncake_admin_key_secret:
343 mooncake_proxy_config["admin_api_key_secret"] = mooncake_admin_key_secret
345 # When no admin-key Secret is named, each region's monitor auto-provisions a
346 # {name}-admin Secret with a generated key, so no manual step is needed.
347 if mooncake_mode in ("disaggregated", "both") and not mooncake_admin_key_secret:
348 formatter.print_info(
349 "No --mooncake-admin-key-secret given; each region's inference "
350 "monitor will auto-provision a '{name}-admin' Secret with a "
351 "generated ADMIN_API_KEY. Pass --mooncake-admin-key-secret to use "
352 "your own Secret instead."
353 )
355 try:
356 manager = get_inference_manager(config)
357 result = manager.deploy(
358 endpoint_name=endpoint_name,
359 image=image,
360 framework=framework,
361 target_regions=list(region) if region else None,
362 replicas=replicas,
363 gpu_count=gpu_count,
364 gpu_type=gpu_type,
365 port=port,
366 model_path=model_path,
367 model_source=model_source,
368 health_check_path=health_path,
369 env=env_dict if env_dict else None,
370 namespace=namespace,
371 labels=labels_dict if labels_dict else None,
372 autoscaling=autoscaling_config,
373 capacity_type=capacity_type,
374 extra_args=list(extra_args) if extra_args else None,
375 accelerator=accelerator,
376 node_selector=node_selector_dict if node_selector_dict else None,
377 rewrite_image=not no_rewrite_image,
378 mooncake_mode=mooncake_mode,
379 prefill_replicas=prefill_replicas,
380 decode_replicas=decode_replicas,
381 mooncake_store=mooncake_store_config,
382 mooncake_transfer=mooncake_transfer_config,
383 mooncake_proxy=mooncake_proxy_config,
384 mooncake_autoscaling=mooncake_autoscaling_config,
385 )
387 formatter.print_success(f"Endpoint '{endpoint_name}' registered for deployment")
388 regions_str = ", ".join(result.get("target_regions", []))
389 formatter.print_info(f"Target regions: {regions_str}")
390 formatter.print_info(f"Ingress path: {result.get('ingress_path', '')}")
391 formatter.print_info(
392 "The inference_monitor in each region will create the resources. "
393 "Use 'gco inference status' to track progress."
394 )
396 # Warn if deploying to a subset of regions
397 if region:
398 from ..aws_client import get_aws_client as _get_client
400 all_stacks = _get_client(config).discover_regional_stacks()
401 all_regions = set(all_stacks.keys())
402 target_set = set(result.get("target_regions", []))
403 missing = all_regions - target_set
404 if missing:
405 formatter.print_warning(
406 f"Endpoint is NOT deployed to: {', '.join(sorted(missing))}. "
407 "Global Accelerator may route users to those regions where "
408 "the endpoint won't exist. Consider deploying to all regions "
409 "(omit -r) for consistent global routing."
410 )
412 if config.output_format != "table":
413 formatter.print(result)
415 except ValueError as e:
416 formatter.print_error(str(e))
417 sys.exit(1)
418 except Exception as e:
419 formatter.print_error(f"Failed to deploy endpoint: {e}")
420 sys.exit(1)
423@inference.command("list")
424@click.option("--state", "-s", help="Filter by state (deploying, running, stopped, deleted)")
425@click.option("--region", "-r", help="Filter by target region")
426@pass_config
427def inference_list(config: Any, state: Any, region: Any) -> None:
428 """List inference endpoints.
430 Examples:
431 gco inference list
432 gco inference list --state running
433 gco inference list -r us-east-1
434 """
435 from ..inference import get_inference_manager
437 formatter = get_output_formatter(config)
439 try:
440 manager = get_inference_manager(config)
441 endpoints = manager.list_endpoints(desired_state=state, region=region)
443 if config.output_format != "table":
444 formatter.print(endpoints)
445 return
447 if not endpoints:
448 formatter.print_info("No inference endpoints found")
449 return
451 print(f"\n Inference Endpoints ({len(endpoints)} found)")
452 print(" " + "-" * 85)
453 print(f" {'NAME':<25} {'STATE':<12} {'REGIONS':<25} {'REPLICAS':>8} {'IMAGE'}")
454 print(" " + "-" * 85)
455 for ep in endpoints:
456 name = ep.get("endpoint_name", "")[:24]
457 ep_state = ep.get("desired_state", "unknown")
458 regions = ", ".join(ep.get("target_regions", []))[:24]
459 spec = ep.get("spec", {})
460 replicas = spec.get("replicas", 1) if isinstance(spec, dict) else 1
461 image = spec.get("image", "")[:40] if isinstance(spec, dict) else ""
462 print(f" {name:<25} {ep_state:<12} {regions:<25} {replicas:>8} {image}")
464 print()
466 except Exception as e:
467 formatter.print_error(f"Failed to list endpoints: {e}")
468 sys.exit(1)
471@inference.command("status")
472@click.argument("endpoint_name")
473@pass_config
474def inference_status(config: Any, endpoint_name: Any) -> None:
475 """Show detailed status of an inference endpoint.
477 Examples:
478 gco inference status my-llm
479 """
480 from ..inference import get_inference_manager
482 formatter = get_output_formatter(config)
484 try:
485 manager = get_inference_manager(config)
486 endpoint = manager.get_endpoint(endpoint_name)
488 if not endpoint:
489 formatter.print_error(f"Endpoint '{endpoint_name}' not found")
490 sys.exit(1)
492 if config.output_format != "table":
493 formatter.print(endpoint)
494 return
496 spec = endpoint.get("spec", {})
497 print(f"\n Endpoint: {endpoint_name}")
498 print(" " + "-" * 60)
499 print(f" State: {endpoint.get('desired_state', 'unknown')}")
500 print(f" Image: {spec.get('image', 'N/A')}")
501 print(f" Replicas: {spec.get('replicas', 1)}")
502 print(f" GPUs: {spec.get('gpu_count', 0)}")
503 print(f" Port: {spec.get('port', 8000)}")
504 print(f" Path: {endpoint.get('ingress_path', 'N/A')}")
505 print(f" Namespace: {endpoint.get('namespace', 'N/A')}")
506 print(f" Created: {endpoint.get('created_at', 'N/A')}")
508 # Region status
509 region_status = endpoint.get("region_status", {})
510 if region_status:
511 print("\n Region Status:")
512 print(f" {'REGION':<18} {'STATE':<12} {'READY':>5} {'DESIRED':>7} {'LAST SYNC'}")
513 print(" " + "-" * 65)
514 for r, status in region_status.items():
515 if isinstance(status, dict):
516 r_state = status.get("state", "unknown")
517 ready = status.get("replicas_ready", 0)
518 desired = status.get("replicas_desired", 0)
519 last_sync = status.get("last_sync", "N/A")
520 if last_sync and len(last_sync) > 19:
521 last_sync = last_sync[:19]
522 print(f" {r:<18} {r_state:<12} {ready:>5} {desired:>7} {last_sync}")
523 else:
524 target_regions = endpoint.get("target_regions", [])
525 print(f"\n Target regions: {', '.join(target_regions)}")
526 print(" (Waiting for inference_monitor to sync)")
528 print()
530 except Exception as e:
531 formatter.print_error(f"Failed to get endpoint status: {e}")
532 sys.exit(1)
535@inference.command("scale")
536@click.argument("endpoint_name")
537@click.option("--replicas", "-r", required=True, type=int, help="New replica count")
538@pass_config
539def inference_scale(config: Any, endpoint_name: Any, replicas: Any) -> None:
540 """Scale an inference endpoint.
542 Examples:
543 gco inference scale my-llm --replicas 4
544 """
545 from ..inference import get_inference_manager
547 formatter = get_output_formatter(config)
549 try:
550 manager = get_inference_manager(config)
551 result = manager.scale(endpoint_name, replicas)
553 if result:
554 formatter.print_success(f"Endpoint '{endpoint_name}' scaled to {replicas} replicas")
555 else:
556 formatter.print_error(f"Endpoint '{endpoint_name}' not found")
557 sys.exit(1)
559 except Exception as e:
560 formatter.print_error(f"Failed to scale endpoint: {e}")
561 sys.exit(1)
564@inference.command("stop")
565@click.argument("endpoint_name")
566@click.option("--yes", "-y", is_flag=True, help="Skip confirmation")
567@pass_config
568def inference_stop(config: Any, endpoint_name: Any, yes: Any) -> None:
569 """Stop an inference endpoint (scale to zero, keep config).
571 Examples:
572 gco inference stop my-llm -y
573 """
574 from ..inference import get_inference_manager
576 formatter = get_output_formatter(config)
578 if not yes:
579 confirm(f"Stop endpoint '{endpoint_name}'?", abort=True)
581 try:
582 manager = get_inference_manager(config)
583 result = manager.stop(endpoint_name)
585 if result:
586 formatter.print_success(f"Endpoint '{endpoint_name}' marked for stop")
587 else:
588 formatter.print_error(f"Endpoint '{endpoint_name}' not found")
589 sys.exit(1)
591 except Exception as e:
592 formatter.print_error(f"Failed to stop endpoint: {e}")
593 sys.exit(1)
596@inference.command("start")
597@click.argument("endpoint_name")
598@pass_config
599def inference_start(config: Any, endpoint_name: Any) -> None:
600 """Start a stopped inference endpoint.
602 Examples:
603 gco inference start my-llm
604 """
605 from ..inference import get_inference_manager
607 formatter = get_output_formatter(config)
609 try:
610 manager = get_inference_manager(config)
611 result = manager.start(endpoint_name)
613 if result:
614 formatter.print_success(f"Endpoint '{endpoint_name}' marked for start")
615 else:
616 formatter.print_error(
617 f"Endpoint '{endpoint_name}' is not a stopped endpoint. If it was deleted "
618 "or purged, redeploy it with 'gco inference deploy'."
619 )
620 sys.exit(1)
622 except Exception as e:
623 formatter.print_error(f"Failed to start endpoint: {e}")
624 sys.exit(1)
627@inference.command("delete")
628@click.argument("endpoint_name")
629@click.option(
630 "--expected-owner-label",
631 default=None,
632 metavar="KEY=VALUE",
633 hidden=True,
634 help="Require an exact stored owner label in the same atomic delete-state update",
635)
636@click.option(
637 "--expected-lifecycle-id",
638 default=None,
639 hidden=True,
640 help="Require the immutable endpoint lifecycle observed by an automation client",
641)
642@click.option("--yes", "-y", is_flag=True, help="Skip confirmation")
643@pass_config
644def inference_delete(
645 config: Any,
646 endpoint_name: Any,
647 expected_owner_label: Any,
648 expected_lifecycle_id: Any,
649 yes: Any,
650) -> None:
651 """Delete an inference endpoint from all regions.
653 The inference_monitor in each region will clean up the K8s resources.
655 Examples:
656 gco inference delete my-llm -y
657 """
658 from ..inference import get_inference_manager
660 formatter = get_output_formatter(config)
662 owner_condition: tuple[str, str] | None = None
663 if expected_owner_label is not None:
664 label_name, separator, label_value = str(expected_owner_label).partition("=")
665 if not separator or not label_name or not label_value:
666 formatter.print_error("--expected-owner-label must be KEY=VALUE")
667 sys.exit(1)
668 owner_condition = (label_name, label_value)
669 lifecycle_condition = (
670 str(expected_lifecycle_id).strip() if expected_lifecycle_id is not None else None
671 )
672 if lifecycle_condition == "":
673 formatter.print_error("--expected-lifecycle-id must be non-empty")
674 sys.exit(1)
675 if (owner_condition is None) != (lifecycle_condition is None):
676 formatter.print_error(
677 "--expected-owner-label and --expected-lifecycle-id must be supplied together"
678 )
679 sys.exit(1)
681 if not yes:
682 confirm(f"Delete endpoint '{endpoint_name}' from all regions?", abort=True)
684 try:
685 manager = get_inference_manager(config)
686 result = manager.delete(
687 endpoint_name,
688 expected_owner_label=owner_condition,
689 expected_lifecycle_id=lifecycle_condition,
690 )
692 if result:
693 formatter.print_success(
694 f"Endpoint '{endpoint_name}' marked for deletion. "
695 "The inference_monitor will clean up resources in each region."
696 )
697 else:
698 suffix = " or ownership/lifecycle condition failed" if owner_condition else ""
699 formatter.print_error(f"Endpoint '{endpoint_name}' not found{suffix}")
700 sys.exit(1)
702 except Exception as e:
703 formatter.print_error(f"Failed to delete endpoint: {e}")
704 sys.exit(1)
707@inference.command("update-image")
708@click.argument("endpoint_name")
709@click.option("--image", "-i", required=True, help="New container image")
710@pass_config
711def inference_update_image(config: Any, endpoint_name: Any, image: Any) -> None:
712 """Update the container image for an inference endpoint.
714 Triggers a rolling update across all target regions.
716 Examples:
717 gco inference update-image my-llm -i vllm/vllm-openai:v0.29.0
718 """
719 from ..inference import get_inference_manager
721 formatter = get_output_formatter(config)
723 try:
724 manager = get_inference_manager(config)
725 result = manager.update_image(endpoint_name, image)
727 if result:
728 formatter.print_success(f"Endpoint '{endpoint_name}' image updated to {image}")
729 formatter.print_info("Rolling update will be applied by inference_monitor")
730 else:
731 formatter.print_error(f"Endpoint '{endpoint_name}' not found")
732 sys.exit(1)
734 except Exception as e:
735 formatter.print_error(f"Failed to update image: {e}")
736 sys.exit(1)
739@inference.command("invoke")
740@click.argument("endpoint_name")
741@click.option("--prompt", "-p", help="Text prompt to send")
742@click.option("--data", "-d", help="Raw JSON body to send")
743@click.option(
744 "--path", "api_path", default=None, help="API sub-path (default: auto-detect from framework)"
745)
746@click.option("--region", "-r", help="Target region for the request")
747@click.option(
748 "--max-tokens", type=int, default=100, help="Maximum tokens to generate (default: 100)"
749)
750@click.option(
751 "--stream/--no-stream",
752 default=None,
753 help="Enable or disable incremental response streaming. Raw JSON with "
754 "'stream': true enables streaming automatically.",
755)
756@pass_config
757def inference_invoke(
758 config: Any,
759 endpoint_name: Any,
760 prompt: Any,
761 data: Any,
762 api_path: Any,
763 region: Any,
764 max_tokens: Any,
765 stream: Any,
766) -> None:
767 """Send a request to an inference endpoint and print the response.
769 Automatically discovers the endpoint's stored API path (the legacy
770 ``ingress_path`` record field) and routes the request through API Gateway
771 with SigV4 authentication.
773 Examples:
774 gco inference invoke my-llm -p "What is GPU orchestration?"
776 gco inference invoke my-llm -d '{"prompt": "Hello", "max_tokens": 50}'
778 gco inference invoke my-llm -p "Explain K8s" --path /v1/completions
779 """
780 import json as _json
782 from ..aws_client import get_aws_client
783 from ..inference import get_inference_manager
785 formatter = get_output_formatter(config)
787 if not prompt and not data:
788 formatter.print_error("Provide --prompt (-p) or --data (-d)")
789 sys.exit(1)
791 try:
792 # Look up the endpoint's stored API prefix and serving spec. The record
793 # retains the historical ``ingress_path`` field name for compatibility;
794 # requests still traverse only the shared authenticated Ingress.
795 manager = get_inference_manager(config)
796 endpoint = manager.get_endpoint(endpoint_name)
797 if not endpoint:
798 formatter.print_error(f"Endpoint '{endpoint_name}' not found")
799 sys.exit(1)
801 endpoint_path = endpoint.get("ingress_path", f"/inference/{endpoint_name}")
802 spec = endpoint.get("spec", {})
803 image = spec.get("image", "") if isinstance(spec, dict) else ""
804 image_lower = image.lower() if isinstance(image, str) else ""
805 persisted_framework = spec.get("framework") if isinstance(spec, dict) else None
807 parsed_data: dict[str, Any] | None = None
808 if data:
809 parsed_json = _json.loads(data)
810 if not isinstance(parsed_json, dict):
811 raise ValueError("--data must contain a JSON object")
812 parsed_data = parsed_json
814 # An explicit flag wins over the body. Without a flag, raw OpenAI JSON
815 # can opt into streamed transport by carrying its normal stream field.
816 if stream is None:
817 stream_response = parsed_data is not None and parsed_data.get("stream") is True
818 else:
819 stream_response = bool(stream)
820 if parsed_data is not None and stream is not None:
821 parsed_data["stream"] = stream_response
823 # Persisted framework is authoritative for neutral/private image names;
824 # image heuristics remain only for legacy records created before that
825 # field existed.
826 if api_path is None:
827 if persisted_framework == "tgi":
828 api_path = "/generate_stream" if stream_response else "/generate"
829 elif persisted_framework == "vllm" or "vllm" in image_lower:
830 api_path = "/v1/completions"
831 elif "text-generation-inference" in image_lower or "tgi" in image_lower:
832 api_path = "/generate_stream" if stream_response else "/generate"
833 elif "tritonserver" in image_lower or "triton" in image_lower:
834 api_path = "/v2/models"
835 else:
836 api_path = "/v1/completions"
838 full_path = f"{endpoint_path}{api_path}"
840 # Build the request body.
841 body: dict[str, Any]
842 if parsed_data is not None:
843 body = parsed_data
844 else:
845 assert prompt is not None
846 if "generate" in api_path:
847 # TGI format; /generate_stream controls response streaming.
848 body = {"inputs": prompt, "parameters": {"max_new_tokens": max_tokens}}
849 elif "/v2/" in api_path:
850 # Triton — just list models, prompt not used for this path.
851 body = {}
852 else:
853 # OpenAI-compatible (vLLM, etc.)
854 # Determine model name for OpenAI-compatible request
855 model_name = endpoint_name
856 if isinstance(spec, dict):
857 # Check env vars first
858 model_name = spec.get("env", {}).get("MODEL", model_name)
859 # Check container args for --model (vLLM, etc.)
860 args_list = spec.get("args") or []
861 for i, arg in enumerate(args_list):
862 if arg == "--model" and i + 1 < len(args_list):
863 model_name = args_list[i + 1]
864 break
865 # Default for vLLM with no explicit model — auto-detect
866 # by querying /v1/models on the running endpoint
867 if model_name == endpoint_name and (
868 persisted_framework == "vllm" or "vllm" in image_lower
869 ):
870 try:
871 detect_client = get_aws_client(config)
872 models_path = f"/inference/{endpoint_name}/v1/models"
873 models_resp = detect_client.make_authenticated_request(
874 method="GET",
875 path=models_path,
876 target_region=region,
877 )
878 if models_resp.ok:
879 models_data = models_resp.json().get("data", [])
880 if models_data:
881 model_name = models_data[0]["id"]
882 except Exception:
883 pass # Fall through to endpoint_name as model
884 body = {
885 "model": model_name,
886 "prompt": prompt,
887 "max_tokens": max_tokens,
888 "stream": stream_response,
889 }
891 if stream_response:
892 # Keep streamed stdout byte-for-byte pipeline-friendly; request
893 # metadata belongs on stderr when the response itself is streamed.
894 print(f"ℹ POST {full_path}", file=sys.stderr)
895 else:
896 formatter.print_info(f"POST {full_path}")
898 # Make the authenticated request. ``stream=True`` prevents requests
899 # from preloading the body so chunks can reach stdout as they arrive.
900 client = get_aws_client(config)
901 response = client.make_authenticated_request(
902 method="POST",
903 path=full_path,
904 body=body,
905 target_region=region,
906 stream=stream_response,
907 )
909 if stream_response:
910 try:
911 if not response.ok:
912 formatter.print_error(f"HTTP {response.status_code}: {response.text[:500]}")
913 sys.exit(1)
915 # Requests assumes ISO-8859-1 for text/* without a declared
916 # charset. Model token streams are UTF-8 in practice, so honor
917 # only an explicit response charset and otherwise use UTF-8.
918 content_type = response.headers.get("content-type", "")
919 encoding = "utf-8"
920 if isinstance(content_type, str):
921 parsed_content_type = Message()
922 parsed_content_type["content-type"] = content_type
923 declared_charset = parsed_content_type.get_content_charset()
924 if declared_charset is not None:
925 try:
926 codecs.lookup(declared_charset)
927 except LookupError:
928 pass
929 else:
930 encoding = declared_charset
932 decoder = codecs.getincrementaldecoder(encoding)(errors="replace")
933 for chunk in response.iter_content(chunk_size=8192, decode_unicode=False):
934 if not chunk:
935 continue
936 output = chunk if isinstance(chunk, str) else decoder.decode(chunk)
937 if output:
938 sys.stdout.write(output)
939 sys.stdout.flush()
940 remainder = decoder.decode(b"", final=True)
941 if remainder:
942 sys.stdout.write(remainder)
943 sys.stdout.flush()
944 finally:
945 response.close()
946 return
948 # Buffered responses retain the friendly extraction used by the CLI.
949 if response.ok:
950 try:
951 resp_json = response.json()
952 # Extract the generated text for common formats
953 text = None
954 if "choices" in resp_json:
955 # OpenAI format
956 choices = resp_json["choices"]
957 if choices:
958 text = choices[0].get("text") or choices[0].get("message", {}).get(
959 "content"
960 )
961 elif "generated_text" in resp_json:
962 # TGI format
963 text = resp_json["generated_text"]
964 elif isinstance(resp_json, list) and resp_json and "generated_text" in resp_json[0]:
965 text = resp_json[0]["generated_text"]
967 if text and config.output_format == "table":
968 print(f"\n{text.strip()}\n")
969 else:
970 emit_structured_document(
971 resp_json,
972 output_format="json",
973 rendered=_json.dumps(resp_json, indent=2),
974 )
975 except _json.JSONDecodeError:
976 print(response.text)
977 else:
978 formatter.print_error(f"HTTP {response.status_code}: {response.text[:500]}")
979 sys.exit(1)
981 except Exception as e:
982 formatter.print_error(f"Failed to invoke endpoint: {e}")
983 sys.exit(1)
986@inference.command("canary")
987@click.argument("endpoint_name")
988@click.option("--image", "-i", required=True, help="New container image for canary")
989@click.option(
990 "--weight",
991 "-w",
992 default=10,
993 type=int,
994 help="Percentage of traffic to canary (1-99, default: 10)",
995)
996@click.option(
997 "--replicas", "-r", default=1, type=int, help="Number of canary replicas (default: 1)"
998)
999@pass_config
1000def inference_canary(
1001 config: Any, endpoint_name: Any, image: Any, weight: Any, replicas: Any
1002) -> None:
1003 """Start a canary deployment with a new image.
1005 Routes a percentage of traffic to the canary while the primary
1006 continues serving the rest. Use 'promote' to make the canary
1007 the new primary, or 'rollback' to remove it.
1009 Examples:
1010 gco inference canary my-llm -i vllm/vllm-openai:v0.29.0 --weight 10
1011 gco inference canary my-llm -i new-image:latest -w 25 -r 2
1012 """
1013 from ..inference import get_inference_manager
1015 formatter = get_output_formatter(config)
1017 try:
1018 manager = get_inference_manager(config)
1019 result = manager.canary_deploy(endpoint_name, image, weight=weight, replicas=replicas)
1021 if not result:
1022 formatter.print_error(f"Endpoint '{endpoint_name}' not found")
1023 sys.exit(1)
1025 formatter.print_success(
1026 f"Canary started: {weight}% traffic → {image} ({replicas} replica(s))"
1027 )
1028 formatter.print_info(f"Monitor with: gco inference status {endpoint_name}")
1029 formatter.print_info(f"Promote with: gco inference promote {endpoint_name}")
1030 formatter.print_info(f"Rollback with: gco inference rollback {endpoint_name}")
1032 except ValueError as e:
1033 formatter.print_error(str(e))
1034 sys.exit(1)
1035 except Exception as e:
1036 formatter.print_error(f"Failed to start canary: {e}")
1037 sys.exit(1)
1040@inference.command("promote")
1041@click.argument("endpoint_name")
1042@click.option("--yes", "-y", is_flag=True, help="Skip confirmation")
1043@pass_config
1044def inference_promote(config: Any, endpoint_name: Any, yes: Any) -> None:
1045 """Promote the canary to primary.
1047 Replaces the primary image with the canary image and removes
1048 the canary deployment. All traffic goes to the new image.
1050 Examples:
1051 gco inference promote my-llm -y
1052 """
1053 from ..inference import get_inference_manager
1055 formatter = get_output_formatter(config)
1057 try:
1058 manager = get_inference_manager(config)
1059 endpoint = manager.get_endpoint(endpoint_name)
1061 if not endpoint:
1062 formatter.print_error(f"Endpoint '{endpoint_name}' not found")
1063 sys.exit(1)
1065 canary = endpoint.get("spec", {}).get("canary")
1066 if not canary:
1067 formatter.print_error(f"Endpoint '{endpoint_name}' has no active canary")
1068 sys.exit(1)
1070 if not yes:
1071 current_image = endpoint.get("spec", {}).get("image", "unknown")
1072 interactive_echo(f" Current primary: {current_image}")
1073 interactive_echo(f" Canary image: {canary.get('image', 'unknown')}")
1074 interactive_echo(f" Canary weight: {canary.get('weight', 0)}%")
1075 if not confirm(" Promote canary to primary?"):
1076 formatter.print_info("Cancelled")
1077 return
1079 result = manager.promote_canary(endpoint_name)
1080 if result:
1081 new_image = result.get("spec", {}).get("image", "unknown")
1082 formatter.print_success(f"Promoted: all traffic now serving {new_image}")
1083 else:
1084 formatter.print_error("Promotion failed")
1085 sys.exit(1)
1087 except ValueError as e:
1088 formatter.print_error(str(e))
1089 sys.exit(1)
1090 except Exception as e:
1091 formatter.print_error(f"Failed to promote canary: {e}")
1092 sys.exit(1)
1095@inference.command("rollback")
1096@click.argument("endpoint_name")
1097@click.option("--yes", "-y", is_flag=True, help="Skip confirmation")
1098@pass_config
1099def inference_rollback(config: Any, endpoint_name: Any, yes: Any) -> None:
1100 """Remove the canary deployment, keeping the primary unchanged.
1102 All traffic returns to the primary deployment.
1104 Examples:
1105 gco inference rollback my-llm -y
1106 """
1107 from ..inference import get_inference_manager
1109 formatter = get_output_formatter(config)
1111 try:
1112 manager = get_inference_manager(config)
1113 endpoint = manager.get_endpoint(endpoint_name)
1115 if not endpoint:
1116 formatter.print_error(f"Endpoint '{endpoint_name}' not found")
1117 sys.exit(1)
1119 canary = endpoint.get("spec", {}).get("canary")
1120 if not canary:
1121 formatter.print_error(f"Endpoint '{endpoint_name}' has no active canary")
1122 sys.exit(1)
1124 if not yes:
1125 interactive_echo(f" Canary image: {canary.get('image', 'unknown')}")
1126 interactive_echo(f" Canary weight: {canary.get('weight', 0)}%")
1127 if not confirm(" Remove canary and restore full traffic to primary?"):
1128 formatter.print_info("Cancelled")
1129 return
1131 result = manager.rollback_canary(endpoint_name)
1132 if result:
1133 primary_image = result.get("spec", {}).get("image", "unknown")
1134 formatter.print_success(f"Rolled back: all traffic now serving {primary_image}")
1135 else:
1136 formatter.print_error("Rollback failed")
1137 sys.exit(1)
1139 except ValueError as e:
1140 formatter.print_error(str(e))
1141 sys.exit(1)
1142 except Exception as e:
1143 formatter.print_error(f"Failed to rollback canary: {e}")
1144 sys.exit(1)
1147@inference.command("health")
1148@click.argument("endpoint_name")
1149@click.option("--region", "-r", help="Target region to check")
1150@pass_config
1151def inference_health(config: Any, endpoint_name: Any, region: Any) -> None:
1152 """Check if an inference endpoint is healthy and ready to serve.
1154 Hits the endpoint's health check path and reports status and latency.
1156 Examples:
1157 gco inference health my-llm
1159 gco inference health my-llm -r us-east-1
1160 """
1161 import time as _time
1163 from ..aws_client import get_aws_client
1164 from ..inference import get_inference_manager
1166 formatter = get_output_formatter(config)
1168 try:
1169 manager = get_inference_manager(config)
1170 endpoint = manager.get_endpoint(endpoint_name)
1171 if not endpoint:
1172 formatter.print_error(f"Endpoint '{endpoint_name}' not found")
1173 sys.exit(1)
1175 endpoint_path = endpoint.get("ingress_path", f"/inference/{endpoint_name}")
1176 spec = endpoint.get("spec", {})
1177 health_path = (
1178 spec.get("health_check_path", "/health") if isinstance(spec, dict) else "/health"
1179 )
1180 full_path = f"{endpoint_path}{health_path}"
1182 client = get_aws_client(config)
1183 start = _time.monotonic()
1184 response = client.make_authenticated_request(
1185 method="GET",
1186 path=full_path,
1187 target_region=region,
1188 )
1189 latency_ms = (_time.monotonic() - start) * 1000
1191 result = {
1192 "endpoint": endpoint_name,
1193 "status": "healthy" if response.ok else "unhealthy",
1194 "http_status": response.status_code,
1195 "latency_ms": round(latency_ms, 1),
1196 "path": full_path,
1197 }
1199 try:
1200 result["body"] = response.json()
1201 except Exception:
1202 result["body"] = response.text[:200] if response.text else None
1204 if config.output_format == "table":
1205 status_icon = "✓" if response.ok else "✗"
1206 formatter.print_info(
1207 f"{status_icon} {endpoint_name}: {result['status']} "
1208 f"(HTTP {response.status_code}, {result['latency_ms']}ms)"
1209 )
1210 else:
1211 formatter.print(result)
1213 except Exception as e:
1214 formatter.print_error(f"Health check failed: {e}")
1215 sys.exit(1)
1218@inference.command("models")
1219@click.argument("endpoint_name")
1220@click.option(
1221 "--framework",
1222 type=click.Choice(["vllm", "tgi"]),
1223 default=None,
1224 help="Runtime metadata contract (default: persisted endpoint framework or vLLM).",
1225)
1226@click.option("--region", "-r", help="Target region to query")
1227@pass_config
1228def inference_models(
1229 config: Any,
1230 endpoint_name: Any,
1231 framework: Any,
1232 region: Any,
1233) -> None:
1234 """Read exact model identity from vLLM /v1/models or TGI /info."""
1235 import json as _json
1237 from ..aws_client import get_aws_client
1238 from ..inference import get_inference_manager
1240 formatter = get_output_formatter(config)
1242 try:
1243 manager = get_inference_manager(config)
1244 endpoint = manager.get_endpoint(endpoint_name)
1245 if not endpoint:
1246 formatter.print_error(f"Endpoint '{endpoint_name}' not found")
1247 sys.exit(1)
1249 ingress_path = endpoint.get("ingress_path", f"/inference/{endpoint_name}")
1250 spec = endpoint.get("spec")
1251 persisted_framework = spec.get("framework") if isinstance(spec, dict) else None
1252 image = spec.get("image") if isinstance(spec, dict) else None
1253 inferred_framework = (
1254 "tgi"
1255 if isinstance(image, str)
1256 and ("text-generation-inference" in image.lower() or "/tgi" in image.lower())
1257 else "vllm"
1258 )
1259 selected_framework = framework or persisted_framework or inferred_framework
1260 if selected_framework not in ("vllm", "tgi"):
1261 raise ValueError("endpoint has an unsupported persisted inference framework")
1262 model_path = "info" if selected_framework == "tgi" else "v1/models"
1263 full_path = f"{ingress_path}/{model_path}"
1265 client = get_aws_client(config)
1266 response = client.make_authenticated_request(
1267 method="GET",
1268 path=full_path,
1269 target_region=region,
1270 )
1272 if response.ok:
1273 try:
1274 resp_json = response.json()
1275 emit_structured_document(
1276 resp_json,
1277 output_format="json",
1278 rendered=_json.dumps(resp_json, indent=2),
1279 )
1280 except _json.JSONDecodeError:
1281 print(response.text)
1282 else:
1283 formatter.print_error(f"HTTP {response.status_code}: {response.text[:500]}")
1284 sys.exit(1)
1286 except Exception as e:
1287 formatter.print_error(f"Failed to list models: {e}")
1288 sys.exit(1)
1291@inference.command("set-topology")
1292@click.argument("endpoint_name")
1293@click.option(
1294 "--prefill",
1295 required=True,
1296 type=int,
1297 help="Prefill (X) instance count for the XpYd topology.",
1298)
1299@click.option(
1300 "--decode",
1301 required=True,
1302 type=int,
1303 help="Decode (Y) instance count for the XpYd topology.",
1304)
1305@pass_config
1306def inference_set_topology(config: Any, endpoint_name: Any, prefill: Any, decode: Any) -> None:
1307 """Resize a disaggregated endpoint's prefill/decode topology.
1309 Updates the endpoint's prefill (X) and decode (Y) instance counts and
1310 re-triggers reconciliation so each region's monitor adjusts the role
1311 replica counts. Both counts must be integers in the range 1..1000.
1313 Examples:
1314 gco inference set-topology llama-pd --prefill 3 --decode 2
1315 """
1316 from ..inference import get_inference_manager
1318 formatter = get_output_formatter(config)
1320 try:
1321 manager = get_inference_manager(config)
1322 result = manager.set_topology(endpoint_name, prefill, decode)
1324 if result:
1325 formatter.print_success(
1326 f"Endpoint '{endpoint_name}' topology set to {prefill}p{decode}d"
1327 )
1328 formatter.print_info(
1329 "The inference_monitor will adjust prefill and decode "
1330 "replica counts in each region."
1331 )
1332 if config.output_format != "table":
1333 formatter.print(result)
1334 else:
1335 formatter.print_error(f"Endpoint '{endpoint_name}' not found")
1336 sys.exit(1)
1338 except ValueError as e:
1339 formatter.print_error(str(e))
1340 sys.exit(1)
1341 except Exception as e:
1342 formatter.print_error(f"Failed to set topology: {e}")
1343 sys.exit(1)
1346@inference.command("configure-store")
1347@click.argument("endpoint_name")
1348@click.option(
1349 "--cold-tier/--no-cold-tier",
1350 "cold_tier",
1351 default=None,
1352 help="Opt the endpoint into (or out of) the asynchronous S3 cold tier. "
1353 "Enabling it also enables the shared store it extends.",
1354)
1355@click.option(
1356 "--offload",
1357 type=click.Choice(["cpu", "disk", "none"]),
1358 default=None,
1359 help="KV-store offload tier for spilling cache beyond GPU memory.",
1360)
1361@click.option(
1362 "--global-segment-size",
1363 type=int,
1364 default=None,
1365 help="Global segment size in bytes for the KV-cache store.",
1366)
1367@click.option(
1368 "--local-buffer-size",
1369 type=int,
1370 default=None,
1371 help="Local buffer size in bytes for the KV-cache store.",
1372)
1373@click.option(
1374 "--enable-store/--disable-store",
1375 "enabled",
1376 default=None,
1377 help="Enable or disable the shared KV-cache store.",
1378)
1379@pass_config
1380def inference_configure_store(
1381 config: Any,
1382 endpoint_name: Any,
1383 cold_tier: Any,
1384 offload: Any,
1385 global_segment_size: Any,
1386 local_buffer_size: Any,
1387 enabled: Any,
1388) -> None:
1389 """Update the shared KV-cache store on a Mooncake endpoint.
1391 Merges the given settings into the endpoint's existing KV-cache store
1392 configuration and re-triggers reconciliation so each region's monitor picks
1393 up the change. Enabling the cold tier also enables the shared store it
1394 extends. Use 'gco inference populate-kv' to pre-warm the cold tier.
1396 Examples:
1397 gco inference configure-store my-llm --cold-tier
1398 gco inference configure-store my-llm --offload cpu --local-buffer-size 2147483648
1399 """
1400 from ..inference import get_inference_manager
1402 formatter = get_output_formatter(config)
1404 try:
1405 manager = get_inference_manager(config)
1406 endpoint = manager.get_endpoint(endpoint_name)
1407 if not endpoint:
1408 formatter.print_error(f"Endpoint '{endpoint_name}' not found")
1409 sys.exit(1)
1411 # Merge onto the endpoint's current store block so changing one field
1412 # does not drop the others.
1413 spec = endpoint.get("spec", {}) if isinstance(endpoint, dict) else {}
1414 mooncake = spec.get("mooncake", {}) if isinstance(spec, dict) else {}
1415 store_config = dict(mooncake.get("store") or {})
1417 if enabled is not None:
1418 store_config["enabled"] = enabled
1419 if cold_tier is not None:
1420 store_config["cold_tier_enabled"] = cold_tier
1421 if cold_tier:
1422 # The cold tier extends the shared store, so enabling it enables
1423 # the store too.
1424 store_config["enabled"] = True
1425 if offload is not None:
1426 store_config["offload"] = offload
1427 if global_segment_size is not None:
1428 store_config["global_segment_size"] = global_segment_size
1429 if local_buffer_size is not None:
1430 store_config["local_buffer_size"] = local_buffer_size
1432 if not store_config:
1433 formatter.print_error(
1434 "No store settings given. Pass --cold-tier, --offload, "
1435 "--global-segment-size, --local-buffer-size, or --enable-store."
1436 )
1437 sys.exit(1)
1439 result = manager.configure_store(endpoint_name, store_config)
1440 if result:
1441 formatter.print_success(f"Endpoint '{endpoint_name}' store configuration updated")
1442 formatter.print_info(
1443 "The inference_monitor will re-render the KV-cache store "
1444 "configuration in each region."
1445 )
1446 if config.output_format != "table":
1447 formatter.print(result)
1448 else:
1449 formatter.print_error(f"Endpoint '{endpoint_name}' not found")
1450 sys.exit(1)
1452 except ValueError as e:
1453 formatter.print_error(str(e))
1454 sys.exit(1)
1455 except Exception as e:
1456 formatter.print_error(f"Failed to configure store: {e}")
1457 sys.exit(1)
1460@inference.command("populate-kv")
1461@click.argument("endpoint_name")
1462@click.argument("local_path")
1463@click.option(
1464 "--region",
1465 "-r",
1466 required=True,
1467 help="Region whose general-purpose bucket backs the endpoint's KV-cache cold tier.",
1468)
1469@pass_config
1470def inference_populate_kv(config: Any, endpoint_name: Any, local_path: Any, region: Any) -> None:
1471 """Upload data into an endpoint's Mooncake KV-cache cold tier.
1473 Uploads a local file or directory to the region's general-purpose bucket
1474 under the cold-tier key prefix the endpoint reads from
1475 (mooncake-kv/<endpoint>/). The endpoint must be deployed with the cold tier
1476 enabled (deploy with --mooncake-cold-tier, or run
1477 'gco inference configure-store <name> --cold-tier') for its pods to read the
1478 uploaded data.
1480 Examples:
1481 gco inference populate-kv my-llm ./kv-warm-set/ --region us-east-1
1482 """
1483 from ..models import get_regional_bucket_manager
1485 formatter = get_output_formatter(config)
1487 try:
1488 manager = get_regional_bucket_manager(config)
1489 formatter.print_info(
1490 f"Uploading {local_path} into the KV-cache cold tier for "
1491 f"'{endpoint_name}' in '{region}'..."
1492 )
1493 result = manager.populate_kv_cache(local_path, region, endpoint_name)
1495 formatter.print_success(
1496 f"Uploaded {result['files_uploaded']} file(s) to {result['s3_uri']}"
1497 )
1498 formatter.print_info(
1499 "Pods for this endpoint read the cold tier when it is enabled "
1500 "(deploy with --mooncake-cold-tier or 'gco inference configure-store')."
1501 )
1503 if config.output_format != "table":
1504 formatter.print(result)
1506 except Exception as e:
1507 formatter.print_error(f"Failed to populate KV cache: {e}")
1508 sys.exit(1)