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