Coverage for cli / inference.py: 100.00%
506 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-10-28 16:19 +0000
1"""
2Inference endpoint management for GCO CLI.
4Provides functionality to deploy, manage, and monitor inference endpoints
5across multi-region EKS clusters via the DynamoDB-backed reconciliation
6pattern (inference_monitor).
7"""
9from __future__ import annotations
11import logging
12import secrets
13from copy import deepcopy
14from typing import TYPE_CHECKING, Any, Literal, TypedDict, TypeGuard
16from gco.models.inference_models import (
17 GPU_NAME_NODE_LABEL,
18 INSTANCE_FAMILY_NODE_LABEL,
19 SGLANG_SUPPORTED_GPU_EXAMPLES,
20 SGLANG_UNSUPPORTED_GPU_FAMILIES,
21 SGLANG_UNSUPPORTED_GPU_NAMES,
22)
24from .aws_client import get_aws_client
25from .config import GCOConfig, get_config
27# <pyflowchart-code-diagram> BEGIN - auto-inserted, do not edit
28# Generated at (UTC): 2026-10-25T13:00:05Z
29# Generated from Git commit: de508a874e7ac267180e26a4c6fc56c1cd0878c2
30# Flowchart(s) generated from this file:
31# * ``InferenceManager.deploy`` -> ``diagrams/code_diagrams/cli/inference.InferenceManager_deploy.html``
32# (PNG: ``diagrams/code_diagrams/cli/inference.InferenceManager_deploy.png``)
33# * ``InferenceManager.canary_deploy`` -> ``diagrams/code_diagrams/cli/inference.InferenceManager_canary_deploy.html``
34# (PNG: ``diagrams/code_diagrams/cli/inference.InferenceManager_canary_deploy.png``)
35# Regenerate with ``SOURCE_DATE_EPOCH=<unix-seconds> GCO_DIAGRAM_SOURCE_COMMIT=<40-char-sha> python diagrams/generate.py --code-only``.
36# <pyflowchart-code-diagram> END
39if TYPE_CHECKING:
40 from gco.services.inference_store import InferenceEndpointStore
42logger = logging.getLogger(__name__)
45# ---------------------------------------------------------------------------
46# Mooncake topology — optional endpoint-spec extension
47# ---------------------------------------------------------------------------
48#
49# An endpoint spec may carry an optional ``mooncake`` block describing
50# disaggregated prefill/decode (PD) serving and/or a shared KV-cache store.
51# The block is entirely additive: when it is absent the endpoint reconciles
52# exactly as it does today — one Deployment and one internal ClusterIP Service
53# behind the shared authenticated inference route.
54#
55# The definitions below describe the shape of that block (the dict written to
56# DynamoDB and read back by the per-region monitor) and the constant
57# vocabularies its enumerated fields draw from. Byte-size fields are authored
58# as base-10 integer decimal strings (see :func:`author_byte_size`) so they
59# round-trip through DynamoDB without being coerced to ``Decimal`` via a float
60# literal.
62#: Serving modes a ``mooncake`` block may declare.
63#: ``disaggregated`` splits prefill and decode; ``store`` runs a single
64#: KV-store instance; ``both`` composes the two.
65MOONCAKE_MODES: frozenset[str] = frozenset({"disaggregated", "store", "both"})
67#: KV transfer / store intents. ``rdma`` is the default high-performance
68#: intent: GCO schedules the pod on EFA and renders vLLM's point-to-point
69#: ``mooncake_protocol`` as ``efa``. ``tcp`` is the non-EFA fallback.
70MOONCAKE_TRANSFER_PROTOCOLS: frozenset[str] = frozenset({"rdma", "tcp"})
72#: KV-store offload tiers for spilling cache beyond GPU memory.
73MOONCAKE_OFFLOAD_TIERS: frozenset[str] = frozenset({"cpu", "disk", "none"})
75#: PD proxy request-scheduling strategies supported today.
76MOONCAKE_PROXY_SCHEDULING: frozenset[str] = frozenset({"round_robin"})
78#: Inclusive bounds for per-role replica counts in an XpYd topology.
79MOONCAKE_TOPOLOGY_MIN: int = 1
80MOONCAKE_TOPOLOGY_MAX: int = 1000
82#: Inclusive bounds for byte-size fields. The ceiling is the signed 64-bit
83#: maximum; authoring sizes as decimal strings in ``[MIN, MAX]`` keeps them out
84#: of float/Decimal coercion when they round-trip through DynamoDB.
85MOONCAKE_BYTE_SIZE_MIN: int = 0
86MOONCAKE_BYTE_SIZE_MAX: int = 9223372036854775807
88#: Transfer-engine defaults mirroring Mooncake's reference configuration.
89MOONCAKE_DEFAULT_BOOTSTRAP_BASE_PORT: int = 8998
90MOONCAKE_DEFAULT_NUM_WORKERS: int = 10
91MOONCAKE_DEFAULT_ABORT_REQUEST_TIMEOUT: int = 480
94class MooncakeTopology(TypedDict):
95 """An XpYd topology: ``prefill`` (X) and ``decode`` (Y) instance counts."""
97 prefill: int
98 decode: int
101class MooncakeStoreConfig(TypedDict, total=False):
102 """KV-cache store pool configuration.
104 ``global_segment_size`` and ``local_buffer_size`` are byte counts authored
105 as base-10 integer decimal strings. ``cold_tier_enabled`` opts this
106 endpoint into the asynchronous, per-region object-store cold tier; the
107 cold-tier bucket is resolved by the monitor from regional configuration and
108 is never a user-typed URI.
109 """
111 enabled: bool
112 metadata_server: str
113 master_server_address: str
114 protocol: Literal["rdma", "tcp"]
115 device_name: str
116 global_segment_size: str
117 local_buffer_size: str
118 offload: Literal["cpu", "disk", "none"]
119 cold_tier_enabled: bool
122class MooncakeTransferConfig(TypedDict, total=False):
123 """RDMA/TCP transfer-engine configuration for KV cache movement."""
125 protocol: Literal["rdma", "tcp"]
126 device_name: str
127 num_workers: int
128 bootstrap_base_port: int
129 abort_request_timeout: int
132class MooncakeProxyConfig(TypedDict, total=False):
133 """PD proxy configuration. ``admin_api_key_secret`` names the Kubernetes
134 Secret holding the proxy admin key; the key value is never carried on the
135 endpoint spec."""
137 image: str
138 scheduling: Literal["round_robin"]
139 admin_api_key_secret: str
142class MooncakeRoleAutoscaling(TypedDict, total=False):
143 """Per-role autoscaling bounds and metrics for one of prefill/decode."""
145 min_replicas: int
146 max_replicas: int
147 metrics: list[dict[str, Any]]
150class MooncakeAutoscalingConfig(TypedDict, total=False):
151 """Optional per-role pod autoscaling. When absent the topology is static."""
153 enabled: bool
154 prefill: MooncakeRoleAutoscaling
155 decode: MooncakeRoleAutoscaling
158class MooncakeSpec(TypedDict, total=False):
159 """The optional ``mooncake`` block carried on an endpoint spec dict."""
161 mode: Literal["disaggregated", "store", "both"]
162 topology: MooncakeTopology
163 store: MooncakeStoreConfig
164 transfer: MooncakeTransferConfig
165 proxy: MooncakeProxyConfig
166 autoscaling: MooncakeAutoscalingConfig
169def author_byte_size(value: int | str) -> str:
170 """Render a byte-size value as a canonical base-10 integer decimal string.
172 Mooncake store/transfer sizes (segment size, local buffer) are carried on
173 the endpoint spec as digit-only strings so they survive the DynamoDB
174 round-trip without being coerced to ``Decimal`` through a float literal.
176 Accepts a non-negative ``int`` or a string of base-10 ASCII digits and
177 returns the same whole number as ``str``. The value must fall in
178 ``[MOONCAKE_BYTE_SIZE_MIN, MOONCAKE_BYTE_SIZE_MAX]``. Signs, decimal
179 points, exponents, floats, booleans, and any non-digit text are not
180 accepted.
182 Raises:
183 ValueError: when ``value`` cannot be authored as an in-range base-10
184 integer.
185 """
186 # ``bool`` is a subclass of ``int``; reject it explicitly so ``True``/``False``
187 # never masquerade as 1/0 byte sizes.
188 if isinstance(value, bool):
189 raise ValueError(f"byte-size value must be an integer, got bool: {value!r}")
191 if isinstance(value, int):
192 size = value
193 elif isinstance(value, str):
194 text = value.strip()
195 if not text or any(ch not in "0123456789" for ch in text):
196 raise ValueError(
197 "byte-size value must be a base-10 integer string "
198 f"(ASCII digits only, no sign, point, or exponent), got {value!r}"
199 )
200 size = int(text)
201 else:
202 raise ValueError(
203 f"byte-size value must be an int or a base-10 digit string, got {type(value).__name__}"
204 )
206 if not MOONCAKE_BYTE_SIZE_MIN <= size <= MOONCAKE_BYTE_SIZE_MAX:
207 raise ValueError(
208 "byte-size value out of range "
209 f"[{MOONCAKE_BYTE_SIZE_MIN}, {MOONCAKE_BYTE_SIZE_MAX}]: {size}"
210 )
212 return str(size)
215#: Byte-size fields a ``mooncake`` store block may carry. Each is authored as a
216#: base-10 integer decimal string via :func:`author_byte_size`.
217_MOONCAKE_STORE_BYTE_SIZE_FIELDS: tuple[str, ...] = (
218 "global_segment_size",
219 "local_buffer_size",
220)
222#: Modes that run a split prefill/decode topology and therefore require a
223#: valid ``topology`` and may carry per-role autoscaling.
224_MOONCAKE_DISAGGREGATED_MODES: frozenset[str] = frozenset({"disaggregated", "both"})
227def _is_plain_int(value: Any) -> TypeGuard[int]:
228 """True when ``value`` is an ``int`` and not a ``bool``.
230 ``bool`` is a subclass of ``int``; counts and replica bounds must be real
231 integers, so ``True``/``False`` are not accepted as 1/0.
232 """
233 return isinstance(value, int) and not isinstance(value, bool)
236def _validate_role_autoscaling_bounds(role: str, role_block: dict[str, Any]) -> None:
237 """Validate one role's ``min_replicas``/``max_replicas`` bounds.
239 Raises :class:`ValueError` naming the violated bound. ``min_replicas`` must
240 be an integer ``>= 1`` and ``max_replicas`` an integer no smaller than the
241 effective minimum (which defaults to 1 when ``min_replicas`` is absent).
242 """
243 min_replicas = role_block.get("min_replicas")
244 max_replicas = role_block.get("max_replicas")
246 if min_replicas is not None:
247 if not _is_plain_int(min_replicas):
248 raise ValueError(
249 f"mooncake.autoscaling.{role}.min_replicas must be an integer, got {min_replicas!r}"
250 )
251 if min_replicas < 1:
252 raise ValueError(
253 f"mooncake.autoscaling.{role}.min_replicas must be >= 1, got {min_replicas}"
254 )
256 if max_replicas is not None:
257 if not _is_plain_int(max_replicas):
258 raise ValueError(
259 f"mooncake.autoscaling.{role}.max_replicas must be an integer, got {max_replicas!r}"
260 )
261 effective_min = min_replicas if _is_plain_int(min_replicas) else 1
262 if max_replicas < effective_min:
263 raise ValueError(
264 f"mooncake.autoscaling.{role}.max_replicas ({max_replicas}) "
265 f"must be >= min_replicas ({effective_min})"
266 )
269def validate_mooncake_spec(mooncake: dict[str, Any]) -> None:
270 """Validate a ``mooncake`` endpoint-spec block, failing fast.
272 Raises :class:`ValueError` on the first rejected field, naming the
273 offending field so the caller can correct it. The check is pure — it reads
274 nothing and writes nothing — so a caller that validates before persisting
275 leaves any previously stored spec untouched when a block is rejected.
277 The rules enforced here are:
279 * ``mode`` must be one of the supported serving modes
280 (:data:`MOONCAKE_MODES`).
281 * ``transfer`` must be a mapping when present; ``protocol`` must be one of
282 :data:`MOONCAKE_TRANSFER_PROTOCOLS` and ``device_name`` must be a string
283 (the empty string requests automatic interface detection).
284 * Store byte-size fields must author as in-range base-10 integers.
285 * ``disaggregated``/``both`` modes require integer ``topology.prefill`` and
286 ``topology.decode`` in
287 ``[MOONCAKE_TOPOLOGY_MIN, MOONCAKE_TOPOLOGY_MAX]``.
288 * ``store.cold_tier_enabled`` may be true only while ``store.enabled`` is
289 true (the cold tier extends the hot store).
290 * Autoscaling may be enabled only for ``disaggregated``/``both`` modes, and
291 each present role's ``min_replicas``/``max_replicas`` must satisfy
292 ``min_replicas >= 1`` and ``max_replicas >= min_replicas``.
293 """
294 if not isinstance(mooncake, dict):
295 raise ValueError("mooncake block must be a mapping")
297 mode = mooncake.get("mode")
298 if mode not in MOONCAKE_MODES:
299 allowed = ", ".join(sorted(MOONCAKE_MODES))
300 raise ValueError(f"mooncake.mode must be one of {{{allowed}}}, got {mode!r}")
302 store = mooncake.get("store")
303 if store is not None and not isinstance(store, dict):
304 raise ValueError("mooncake.store must be a mapping")
306 transfer = mooncake.get("transfer")
307 if transfer is not None and not isinstance(transfer, dict):
308 raise ValueError("mooncake.transfer must be a mapping")
309 if isinstance(transfer, dict):
310 protocol = transfer.get("protocol", "rdma")
311 if protocol not in MOONCAKE_TRANSFER_PROTOCOLS:
312 allowed = ", ".join(sorted(MOONCAKE_TRANSFER_PROTOCOLS))
313 raise ValueError(
314 f"mooncake.transfer.protocol must be one of {{{allowed}}}, got {protocol!r}"
315 )
316 device_name = transfer.get("device_name", "")
317 if not isinstance(device_name, str):
318 raise ValueError(f"mooncake.transfer.device_name must be a string, got {device_name!r}")
320 # Byte-size fields must author cleanly; surface the offending field name.
321 if isinstance(store, dict):
322 for field in _MOONCAKE_STORE_BYTE_SIZE_FIELDS:
323 if field in store:
324 try:
325 author_byte_size(store[field])
326 except ValueError as exc:
327 raise ValueError(f"mooncake.store.{field}: {exc}") from exc
329 # Split topologies need integer prefill/decode counts in range.
330 if mode in _MOONCAKE_DISAGGREGATED_MODES:
331 topology = mooncake.get("topology")
332 if not isinstance(topology, dict):
333 raise ValueError(
334 f"mooncake.topology is required for mode {mode!r} with integer "
335 "'prefill' and 'decode' counts"
336 )
337 for field in ("prefill", "decode"):
338 count = topology.get(field)
339 if not _is_plain_int(count):
340 raise ValueError(
341 f"mooncake.topology.{field} must be an integer in "
342 f"[{MOONCAKE_TOPOLOGY_MIN}, {MOONCAKE_TOPOLOGY_MAX}], "
343 f"got {count!r}"
344 )
345 if not MOONCAKE_TOPOLOGY_MIN <= count <= MOONCAKE_TOPOLOGY_MAX:
346 raise ValueError(
347 f"mooncake.topology.{field} out of range "
348 f"[{MOONCAKE_TOPOLOGY_MIN}, {MOONCAKE_TOPOLOGY_MAX}]: {count}"
349 )
351 # The cold tier extends the hot store; it cannot be enabled on its own.
352 if (
353 isinstance(store, dict)
354 and store.get("cold_tier_enabled") is True
355 and store.get("enabled") is not True
356 ):
357 raise ValueError(
358 "mooncake.store.cold_tier_enabled requires mooncake.store.enabled to be true"
359 )
361 autoscaling = mooncake.get("autoscaling")
362 if autoscaling is not None:
363 if not isinstance(autoscaling, dict):
364 raise ValueError("mooncake.autoscaling must be a mapping")
365 if autoscaling.get("enabled") is True and mode not in _MOONCAKE_DISAGGREGATED_MODES:
366 raise ValueError(
367 "mooncake.autoscaling.enabled requires a 'disaggregated' or "
368 f"'both' mode, got {mode!r}"
369 )
370 for role in ("prefill", "decode"):
371 role_block = autoscaling.get(role)
372 if role_block is None:
373 continue
374 if not isinstance(role_block, dict):
375 raise ValueError(f"mooncake.autoscaling.{role} must be a mapping")
376 _validate_role_autoscaling_bounds(role, role_block)
379class InferenceManager:
380 """Manages inference endpoints via the DynamoDB store."""
382 def __init__(self, config: GCOConfig | None = None):
383 self.config = config or get_config()
384 self._aws_client = get_aws_client(config)
386 def _get_store(self, region: str | None = None) -> InferenceEndpointStore:
387 """Get an InferenceEndpointStore for the global region."""
388 from gco.services.inference_store import InferenceEndpointStore
390 # Use the global region for DynamoDB (same as job store)
391 store_region = region or self.config.global_region
392 return InferenceEndpointStore(region=store_region)
394 def _build_mooncake_block(
395 self,
396 *,
397 mode: str,
398 prefill_replicas: int,
399 decode_replicas: int,
400 store: dict[str, Any] | None,
401 transfer: dict[str, Any] | None,
402 proxy: dict[str, Any] | None,
403 autoscaling: dict[str, Any] | None,
404 default_proxy_image: str | None = None,
405 ) -> dict[str, Any]:
406 """Assemble and validate an optional ``spec.mooncake`` block.
408 Composes the topology and any supplied store/transfer/proxy/autoscaling
409 sub-blocks into a single mapping, authoring store byte-size fields as
410 base-10 integer decimal strings so they round-trip through DynamoDB,
411 then validates the result. Validation is pure and runs before the
412 caller persists anything, so a rejected block leaves any previously
413 stored spec untouched. Raises :class:`ValueError` — naming the offending
414 field — when the mode is unsupported or any field is invalid.
416 The store-bearing modes (``store`` and ``both``) default the store to
417 enabled so the shared master address is wired in (the ``both``-mode
418 MultiConnector's store half depends on it), and split modes
419 (``disaggregated`` and ``both``) default the prefill-decode proxy image
420 to ``default_proxy_image`` when the caller supplies no explicit proxy
421 image.
422 """
423 block: dict[str, Any] = {"mode": mode}
425 # Split modes carry an XpYd topology; a single-instance store does not.
426 if mode in _MOONCAKE_DISAGGREGATED_MODES:
427 block["topology"] = {
428 "prefill": prefill_replicas,
429 "decode": decode_replicas,
430 }
432 # The store-bearing modes (store and both) only function with the KV
433 # store enabled: the both-mode MultiConnector's store half is wired to
434 # the shared master address, which the monitor renders only for an
435 # enabled store. So a store block is always present for those modes,
436 # defaulting enabled to True; an explicit store block still tunes
437 # offload, sizes, and the cold tier.
438 store_block = dict(store) if store is not None else None
439 if mode in ("store", "both"):
440 store_block = dict(store_block or {})
441 store_block.setdefault("enabled", True)
442 if store_block is not None:
443 # Author byte-size fields as canonical decimal strings up front so
444 # the persisted spec round-trips through DynamoDB without float or
445 # Decimal coercion. Authoring also fails fast on bad inputs.
446 for field in _MOONCAKE_STORE_BYTE_SIZE_FIELDS:
447 if field in store_block:
448 try:
449 store_block[field] = author_byte_size(store_block[field])
450 except ValueError as exc:
451 raise ValueError(f"mooncake.store.{field}: {exc}") from exc
452 block["store"] = store_block
454 if transfer is not None:
455 block["transfer"] = dict(transfer)
457 # Split modes are fronted by the prefill-decode proxy, which needs a
458 # container image. Default it to the same image the role pods serve from
459 # (the upstream vLLM image bundles the reference proxy) so a split deploy
460 # stands up without a separate proxy image; an explicit proxy image
461 # still wins.
462 proxy_block = dict(proxy) if proxy is not None else None
463 if mode in _MOONCAKE_DISAGGREGATED_MODES and default_proxy_image:
464 proxy_block = dict(proxy_block or {})
465 proxy_block.setdefault("image", default_proxy_image)
466 if proxy_block is not None:
467 block["proxy"] = proxy_block
469 if autoscaling is not None:
470 block["autoscaling"] = dict(autoscaling)
472 # Fail fast before persisting: rejects unsupported modes (naming the
473 # allowed values) and every other invalid field.
474 validate_mooncake_spec(block)
475 return block
477 def deploy(
478 self,
479 endpoint_name: str,
480 image: str | None = None,
481 target_regions: list[str] | None = None,
482 replicas: int = 1,
483 gpu_count: int = 1,
484 gpu_type: str | None = None,
485 port: int = 8000,
486 model_path: str | None = None,
487 model_source: str | None = None,
488 health_check_path: str = "/health",
489 env: dict[str, str] | None = None,
490 namespace: str = "gco-inference",
491 labels: dict[str, str] | None = None,
492 autoscaling: dict[str, Any] | None = None,
493 capacity_type: str | None = None,
494 extra_args: list[str] | None = None,
495 accelerator: str = "nvidia",
496 node_selector: dict[str, str] | None = None,
497 rewrite_image: bool = True,
498 *,
499 framework: str | None = None,
500 mooncake_mode: str | None = None,
501 prefill_replicas: int = 1,
502 decode_replicas: int = 1,
503 mooncake_store: dict[str, Any] | None = None,
504 mooncake_transfer: dict[str, Any] | None = None,
505 mooncake_proxy: dict[str, Any] | None = None,
506 mooncake_autoscaling: dict[str, Any] | None = None,
507 ) -> dict[str, Any]:
508 """
509 Deploy an inference endpoint to one or more regions.
511 The endpoint spec is written to DynamoDB. The inference_monitor
512 in each target region picks it up and creates the K8s resources.
514 Args:
515 endpoint_name: Unique name for the endpoint
516 image: Container image (e.g. vllm/vllm-openai:v0.32.0). Optional
517 when ``mooncake_mode`` is set: a disaggregated/store deploy
518 with no image falls back to the default upstream
519 Mooncake-enabled vLLM image. A plain deploy still requires an
520 image.
521 target_regions: Regions to deploy to (default: all deployed regions)
522 replicas: Number of replicas per region
523 gpu_count: GPUs per replica
524 gpu_type: Preferred instance type, rendered as a preferred node
525 affinity (a hint, not a requirement)
526 port: Container port
527 model_path: Mounts the inference namespace's EFS volume at
528 ``/models`` without syncing anything
529 health_check_path: Health check endpoint path
530 env: Environment variables
531 namespace: Kubernetes namespace
532 labels: Labels for the endpoint
533 rewrite_image: When True (the default), rewrite ECR URIs in
534 ``image`` to target each region's local replica. Non-ECR
535 refs (Docker Hub, GHCR, etc.) are left unchanged. When
536 False, the URI is written verbatim to every region's
537 spec — the operator is responsible for cross-region
538 pulls. Per-region rewrites are stored under a
539 ``region_overrides`` map on the spec keyed by region.
540 mooncake_mode: When set to one of ``disaggregated``, ``store``,
541 or ``both``, build and persist a ``spec.mooncake`` block for
542 disaggregated prefill/decode serving and/or a shared KV-cache
543 store. An unsupported value is rejected before anything is
544 persisted.
545 prefill_replicas: X in an XpYd topology — prefill instance count
546 for split (``disaggregated``/``both``) modes.
547 decode_replicas: Y in an XpYd topology — decode instance count for
548 split modes.
549 mooncake_store: Optional KV-store pool configuration merged into
550 ``spec.mooncake.store``. Byte-size fields are authored as
551 base-10 integer decimal strings so they round-trip through
552 DynamoDB.
553 mooncake_transfer: Optional Mooncake transfer intent and network
554 device. ``protocol`` accepts ``rdma`` (the default; scheduled
555 on EFA and rendered to vLLM as ``mooncake_protocol=efa``) or
556 ``tcp`` (no EFA placement). ``device_name`` is forwarded to
557 both the connector and mounted Mooncake configuration; an
558 empty string lets Mooncake auto-detect it.
559 mooncake_proxy: Optional PD proxy configuration merged into
560 ``spec.mooncake.proxy``.
561 mooncake_autoscaling: Optional per-role autoscaling configuration
562 merged into ``spec.mooncake.autoscaling``.
564 Returns:
565 Created endpoint record
566 """
567 if framework not in (None, "vllm", "sglang"):
568 raise ValueError("framework must be 'vllm' or 'sglang'")
569 if mooncake_mode is not None and framework == "sglang":
570 raise ValueError("Mooncake serving requires the vllm framework")
571 if (
572 framework == "sglang"
573 and not (env or {}).get("MODEL")
574 and not {"--model-path", "--model"} & set(extra_args or [])
575 ):
576 # The renderer supplies the launcher and takes the model from the
577 # ``MODEL`` convention or an explicit launcher flag; without either
578 # the pod would only crash-loop on a missing --model-path.
579 raise ValueError(
580 "framework 'sglang' needs the model to serve: pass -e MODEL=<id-or-path> "
581 "or --extra-args=--model-path --extra-args <id-or-path>"
582 )
583 if framework == "sglang" and node_selector:
584 # The renderer keeps SGLang off pre-Ampere GPUs on its own; a
585 # selector that pins one would only make the pod unschedulable
586 # forever, so refuse it here with the reason instead.
587 pinned_gpu = str(node_selector.get(GPU_NAME_NODE_LABEL, "")).lower()
588 pinned_family = str(node_selector.get(INSTANCE_FAMILY_NODE_LABEL, "")).lower()
589 pinned = (
590 pinned_gpu
591 if pinned_gpu in SGLANG_UNSUPPORTED_GPU_NAMES
592 else pinned_family
593 if pinned_family in SGLANG_UNSUPPORTED_GPU_FAMILIES
594 else None
595 )
596 if pinned is not None:
597 raise ValueError(
598 "framework 'sglang' needs an NVIDIA GPU with compute capability 8.0 or "
599 f"newer ({SGLANG_SUPPORTED_GPU_EXAMPLES}); the node selector pins "
600 f"'{pinned}', for which SGLang ships no prebuilt kernels"
601 )
602 if mooncake_mode is not None and framework is None:
603 framework = "vllm"
605 # Build the optional mooncake block first and validate it before any
606 # persistence so a rejected block leaves any stored spec untouched. A
607 # disaggregated/store deploy without an explicit image falls back to
608 # the default upstream Mooncake-enabled vLLM image.
609 mooncake_block: dict[str, Any] | None = None
610 if mooncake_mode is not None:
611 # Resolve the image before building the block so a split mode's
612 # prefill-decode proxy can default to the same image the role pods
613 # serve from (the upstream vLLM image bundles the reference proxy).
614 if image is None:
615 from .images import default_disaggregated_image
617 image = default_disaggregated_image(config=self.config)
618 mooncake_block = self._build_mooncake_block(
619 mode=mooncake_mode,
620 prefill_replicas=prefill_replicas,
621 decode_replicas=decode_replicas,
622 store=mooncake_store,
623 transfer=mooncake_transfer,
624 proxy=mooncake_proxy,
625 autoscaling=mooncake_autoscaling,
626 default_proxy_image=image,
627 )
629 if image is None:
630 raise ValueError(
631 "an image is required (pass image, or set mooncake_mode to use "
632 "the default upstream Mooncake-enabled vLLM image)"
633 )
635 if not target_regions:
636 stacks = self._aws_client.discover_regional_stacks()
637 target_regions = list(stacks.keys())
638 if not target_regions:
639 raise ValueError("No deployed regions found. Deploy infrastructure first.")
641 # Per-region image-URI rewrites for ECR refs. Each target region
642 # gets the local replica's URI on its own spec, so the
643 # inference_monitor's pod-spec materialiser pulls in-region
644 # rather than across the WAN. Non-ECR URIs come back unchanged
645 # from the helper, so this is a no-op for Docker Hub / GHCR refs.
646 #
647 # The helper lives in ``cli._image_uri`` rather than ``cli.images``
648 # so this import doesn't create a module-level cycle:
649 # ``cli.images`` itself imports the same helper. ``cli._image_uri``
650 # is a leaf module with no project-side dependencies.
651 region_image_map: dict[str, str] = {}
652 if rewrite_image:
653 from ._image_uri import rewrite_image_uri_for_region
655 for region in target_regions:
656 region_image_map[region] = rewrite_image_uri_for_region(image, region)
658 spec = {
659 "image": image,
660 "port": port,
661 "replicas": replicas,
662 "gpu_count": gpu_count,
663 "health_check_path": health_check_path,
664 }
665 if framework:
666 spec["framework"] = framework
667 # Preserve the rewrite map on the spec so the inference_monitor
668 # service can pick the right URI per region when materialising
669 # pods. When ``rewrite_image=False`` no map is set and the flat
670 # ``image`` field is the only source.
671 if region_image_map and any(uri != image for uri in region_image_map.values()):
672 spec["region_image_uris"] = region_image_map
673 if gpu_type:
674 spec["gpu_type"] = gpu_type
675 if model_path:
676 spec["model_path"] = model_path
677 if model_source:
678 spec["model_source"] = model_source
679 if env:
680 spec["env"] = env
681 if autoscaling:
682 spec["autoscaling"] = autoscaling
683 if capacity_type:
684 spec["capacity_type"] = capacity_type
685 if extra_args:
686 spec["args"] = extra_args
687 if accelerator != "nvidia":
688 spec["accelerator"] = accelerator
689 if node_selector:
690 spec["node_selector"] = node_selector
691 if mooncake_block is not None:
692 spec["mooncake"] = mooncake_block
694 store = self._get_store()
695 result: dict[str, Any] = store.create_endpoint(
696 endpoint_name=endpoint_name,
697 spec=spec,
698 target_regions=target_regions,
699 namespace=namespace,
700 labels=labels,
701 )
702 return result
704 def list_endpoints(
705 self,
706 desired_state: str | None = None,
707 region: str | None = None,
708 ) -> list[dict[str, Any]]:
709 """List all inference endpoints."""
710 store = self._get_store()
711 result: list[dict[str, Any]] = store.list_endpoints(
712 desired_state=desired_state,
713 target_region=region,
714 )
715 return result
717 def get_endpoint(self, endpoint_name: str) -> dict[str, Any] | None:
718 """Get details of a specific endpoint."""
719 store = self._get_store()
720 result: dict[str, Any] | None = store.get_endpoint(endpoint_name)
721 return result
723 @staticmethod
724 def _load_mutable_endpoint(
725 store: InferenceEndpointStore,
726 endpoint_name: str,
727 operation: str,
728 ) -> tuple[dict[str, Any], str, str | None] | None:
729 """Read and lifecycle-fence an endpoint before an ordinary mutation."""
730 endpoint = store.get_endpoint(endpoint_name, consistent_read=True)
731 if endpoint is None:
732 return None
733 if endpoint.get("desired_state") == "deleted":
734 raise ValueError(
735 f"Endpoint '{endpoint_name}' is deleted and cannot be {operation}; "
736 "redeploy it after deletion completes."
737 )
738 lifecycle_id = endpoint.get("lifecycle_id")
739 if not isinstance(lifecycle_id, str) or not lifecycle_id:
740 migrated = store.ensure_lifecycle_metadata(endpoint)
741 if migrated is None:
742 raise ValueError(
743 f"Endpoint '{endpoint_name}' changed while initializing lifecycle identity"
744 )
745 endpoint = migrated
746 lifecycle_id = endpoint.get("lifecycle_id")
747 if not isinstance(lifecycle_id, str) or not lifecycle_id:
748 raise ValueError(f"Endpoint '{endpoint_name}' has no lifecycle identity")
749 updated_at = endpoint.get("updated_at")
750 return endpoint, lifecycle_id, updated_at if isinstance(updated_at, str) else None
752 @staticmethod
753 def _raise_write_conflict(endpoint_name: str, operation: str) -> None:
754 raise ValueError(
755 f"Endpoint '{endpoint_name}' changed while being {operation}; retry after reading status"
756 )
758 def scale(self, endpoint_name: str, replicas: int) -> dict[str, Any] | None:
759 """Scale a classic static endpoint only."""
760 store = self._get_store()
761 loaded = self._load_mutable_endpoint(store, endpoint_name, "scaled")
762 if loaded is None:
763 return None
764 endpoint, lifecycle_id, _updated_at = loaded
765 spec = endpoint.get("spec")
766 if isinstance(spec, dict) and "mooncake" in spec:
767 raise ValueError(
768 f"Endpoint '{endpoint_name}' uses Mooncake topology; use "
769 "'gco inference set-topology' instead of 'gco inference scale'."
770 )
771 autoscaling = spec.get("autoscaling") if isinstance(spec, dict) else None
772 if isinstance(autoscaling, dict) and autoscaling.get("enabled") is True:
773 raise ValueError(
774 f"Endpoint '{endpoint_name}' is autoscaled; update its min/max autoscaling "
775 "bounds or disable autoscaling before using 'gco inference scale'."
776 )
777 result: dict[str, Any] | None = store.scale_endpoint(
778 endpoint_name,
779 replicas,
780 expected_lifecycle_id=lifecycle_id,
781 )
782 if result is None:
783 self._raise_write_conflict(endpoint_name, "scaled")
784 return result
786 def set_topology(
787 self,
788 endpoint_name: str,
789 prefill: int,
790 decode: int,
791 ) -> dict[str, Any] | None:
792 """Resize a disaggregated endpoint's prefill/decode topology.
794 Updates ``spec.mooncake.topology`` to the new XpYd counts and
795 re-triggers reconciliation (via :meth:`InferenceEndpointStore.update_spec`,
796 which flips ``desired_state`` to ``deploying``) so the per-region
797 monitor adjusts the prefill and decode role replica counts.
799 Both counts must be integers in the inclusive range
800 ``[MOONCAKE_TOPOLOGY_MIN, MOONCAKE_TOPOLOGY_MAX]``. The counts are
801 validated before anything is read or written, so a rejected request
802 names the offending count and leaves the stored topology and
803 ``desired_state`` untouched.
805 Args:
806 endpoint_name: Name of the disaggregated endpoint to resize.
807 prefill: New prefill (X) instance count.
808 decode: New decode (Y) instance count.
810 Returns:
811 The updated endpoint record, or ``None`` when no endpoint with
812 ``endpoint_name`` exists.
814 Raises:
815 ValueError: when ``prefill`` or ``decode`` is not an integer in
816 ``[MOONCAKE_TOPOLOGY_MIN, MOONCAKE_TOPOLOGY_MAX]``.
817 """
818 # Validate before any read or write so a bad count names the offending
819 # field and leaves the stored topology and desired_state unchanged.
820 for field, count in (("prefill", prefill), ("decode", decode)):
821 if not _is_plain_int(count):
822 raise ValueError(
823 f"topology {field} count must be an integer in "
824 f"[{MOONCAKE_TOPOLOGY_MIN}, {MOONCAKE_TOPOLOGY_MAX}], "
825 f"got {count!r}"
826 )
827 if not MOONCAKE_TOPOLOGY_MIN <= count <= MOONCAKE_TOPOLOGY_MAX:
828 raise ValueError(
829 f"topology {field} count out of range "
830 f"[{MOONCAKE_TOPOLOGY_MIN}, {MOONCAKE_TOPOLOGY_MAX}]: {count}"
831 )
833 store = self._get_store()
834 loaded = self._load_mutable_endpoint(store, endpoint_name, "updated")
835 if loaded is None:
836 return None
837 endpoint, lifecycle_id, updated_at = loaded
839 raw_spec = endpoint.get("spec")
840 if not isinstance(raw_spec, dict):
841 raise ValueError(f"Endpoint '{endpoint_name}' has an invalid spec")
842 spec = deepcopy(raw_spec)
843 # Preserve any existing mooncake sub-fields and replace only the
844 # topology counts.
845 mooncake = dict(spec.get("mooncake") or {})
846 mooncake["topology"] = {"prefill": prefill, "decode": decode}
847 spec["mooncake"] = mooncake
849 result: dict[str, Any] | None = store.update_spec(
850 endpoint_name,
851 spec,
852 expected_lifecycle_id=lifecycle_id,
853 expected_updated_at=updated_at,
854 )
855 if result is None:
856 self._raise_write_conflict(endpoint_name, "updated")
857 return result
859 def configure_store(
860 self,
861 endpoint_name: str,
862 store_config: dict[str, Any],
863 ) -> dict[str, Any] | None:
864 """Update an endpoint's KV-cache store configuration.
866 Merges ``store_config`` into ``spec.mooncake.store`` and re-triggers
867 reconciliation (via :meth:`InferenceEndpointStore.update_spec`, which
868 flips ``desired_state`` to ``deploying``) so the per-region monitor
869 picks up the new store settings.
871 Store byte-size fields are authored as base-10 integer decimal strings
872 (so they round-trip through DynamoDB without float/Decimal coercion)
873 and the resulting ``mooncake`` block is validated before anything is
874 written. A rejected configuration names the offending field and leaves
875 the stored spec untouched.
877 Args:
878 endpoint_name: Name of the endpoint to reconfigure.
879 store_config: KV-store pool settings merged into
880 ``spec.mooncake.store``.
882 Returns:
883 The updated endpoint record, or ``None`` when no endpoint with
884 ``endpoint_name`` exists.
886 Raises:
887 ValueError: when the resulting ``mooncake`` block is invalid (for
888 example an out-of-range byte-size field).
889 """
890 store = self._get_store()
891 loaded = self._load_mutable_endpoint(store, endpoint_name, "configured")
892 if loaded is None:
893 return None
894 endpoint, lifecycle_id, updated_at = loaded
896 raw_spec = endpoint.get("spec")
897 if not isinstance(raw_spec, dict):
898 raise ValueError(f"Endpoint '{endpoint_name}' has an invalid spec")
899 spec = deepcopy(raw_spec)
900 # Preserve any existing mooncake sub-fields and replace only the store
901 # block, authoring byte-size fields as canonical decimal strings.
902 mooncake = dict(spec.get("mooncake") or {})
903 store_block = dict(store_config)
904 for field in _MOONCAKE_STORE_BYTE_SIZE_FIELDS:
905 if field in store_block:
906 try:
907 store_block[field] = author_byte_size(store_block[field])
908 except ValueError as exc:
909 raise ValueError(f"mooncake.store.{field}: {exc}") from exc
910 mooncake["store"] = store_block
911 spec["mooncake"] = mooncake
913 # Fail fast before persisting so a rejected block leaves the stored
914 # spec untouched.
915 validate_mooncake_spec(mooncake)
917 result: dict[str, Any] | None = store.update_spec(
918 endpoint_name,
919 spec,
920 expected_lifecycle_id=lifecycle_id,
921 expected_updated_at=updated_at,
922 )
923 if result is None:
924 self._raise_write_conflict(endpoint_name, "configured")
925 return result
927 def stop(self, endpoint_name: str) -> dict[str, Any] | None:
928 """Stop a live endpoint without crossing a deletion boundary."""
929 store = self._get_store()
930 loaded = self._load_mutable_endpoint(store, endpoint_name, "stopped")
931 if loaded is None:
932 return None
933 _endpoint, lifecycle_id, _updated_at = loaded
934 result: dict[str, Any] | None = store.update_desired_state(
935 endpoint_name,
936 "stopped",
937 expected_lifecycle_id=lifecycle_id,
938 )
939 if result is None:
940 self._raise_write_conflict(endpoint_name, "stopped")
941 return result
943 def start(self, endpoint_name: str) -> dict[str, Any] | None:
944 """Start a stopped endpoint; deleted endpoints must be redeployed."""
945 store = self._get_store()
946 result: dict[str, Any] | None = store.start_endpoint(endpoint_name)
947 return result
949 def delete(
950 self,
951 endpoint_name: str,
952 *,
953 expected_owner_label: tuple[str, str] | None = None,
954 expected_lifecycle_id: str | None = None,
955 ) -> dict[str, Any] | None:
956 """Mark an endpoint deleted with optional exact-incarnation ownership."""
957 if expected_owner_label is not None and not expected_lifecycle_id:
958 raise ValueError("Owner-conditioned deletion also requires an immutable lifecycle id")
959 store = self._get_store()
960 if expected_owner_label is None:
961 endpoint = store.get_endpoint(endpoint_name, consistent_read=True)
962 if endpoint is None:
963 return None
964 migrated = store.ensure_lifecycle_metadata(endpoint)
965 if migrated is None:
966 raise ValueError(
967 f"Endpoint '{endpoint_name}' changed while initializing deletion identity"
968 )
969 lifecycle_id = migrated.get("lifecycle_id")
970 if not isinstance(lifecycle_id, str) or not lifecycle_id:
971 raise ValueError(f"Endpoint '{endpoint_name}' has no lifecycle identity")
972 result: dict[str, Any] | None = store.update_desired_state(
973 endpoint_name,
974 "deleted",
975 expected_lifecycle_id=lifecycle_id,
976 )
977 else:
978 result = store.update_desired_state(
979 endpoint_name,
980 "deleted",
981 expected_label=expected_owner_label,
982 expected_lifecycle_id=expected_lifecycle_id,
983 )
984 return result
986 def update_image(self, endpoint_name: str, image: str) -> dict[str, Any] | None:
987 """Update the container image for an endpoint."""
988 if not isinstance(image, str) or not image.strip():
989 raise ValueError("Image must be a non-empty string")
991 store = self._get_store()
992 loaded = self._load_mutable_endpoint(store, endpoint_name, "updated")
993 if loaded is None:
994 return None
995 endpoint, lifecycle_id, updated_at = loaded
996 raw_spec = endpoint.get("spec")
997 if not isinstance(raw_spec, dict):
998 raise ValueError(f"Endpoint '{endpoint_name}' has an invalid spec")
999 spec = deepcopy(raw_spec)
1000 spec["image"] = image.strip()
1001 # A direct image update is global. Stale regional rewrites would take
1002 # precedence in the monitor and silently keep serving the old image.
1003 spec.pop("region_image_uris", None)
1004 result: dict[str, Any] | None = store.update_spec(
1005 endpoint_name,
1006 spec,
1007 expected_lifecycle_id=lifecycle_id,
1008 expected_updated_at=updated_at,
1009 )
1010 if result is None:
1011 self._raise_write_conflict(endpoint_name, "updated")
1012 return result
1014 def add_region(self, endpoint_name: str, region: str) -> dict[str, Any] | None:
1015 """Add a target and append it to this lifecycle's cleanup history."""
1016 store = self._get_store()
1017 loaded = self._load_mutable_endpoint(store, endpoint_name, "added to a Region")
1018 if loaded is None:
1019 return None
1020 endpoint, lifecycle_id, updated_at = loaded
1021 if updated_at is None:
1022 raise ValueError(f"Endpoint '{endpoint_name}' has no conditional update timestamp")
1023 regions = list(endpoint.get("target_regions") or [])
1024 cleanup_regions = list(endpoint.get("cleanup_regions") or regions)
1025 if region in regions:
1026 return endpoint
1027 regions.append(region)
1028 if region not in cleanup_regions:
1029 cleanup_regions.append(region)
1030 raw_generations = endpoint.get("region_generations")
1031 stored_generations = raw_generations if isinstance(raw_generations, dict) else {}
1032 region_generations = {
1033 cleanup_region: (
1034 stored_generations[cleanup_region]
1035 if isinstance(stored_generations.get(cleanup_region), str)
1036 and stored_generations[cleanup_region]
1037 else secrets.token_hex(32)
1038 )
1039 for cleanup_region in cleanup_regions
1040 }
1041 # Every target membership transition gets a fresh token, even when the
1042 # Region appeared earlier in cleanup history. This invalidates a prior
1043 # terminal removal acknowledgement before resources can be recreated.
1044 region_generations[region] = secrets.token_hex(32)
1045 try:
1046 result = store.update_target_regions(
1047 endpoint_name,
1048 regions,
1049 cleanup_regions,
1050 region_generations,
1051 expected_lifecycle_id=lifecycle_id,
1052 expected_updated_at=updated_at,
1053 )
1054 except Exception as e:
1055 logger.error("Failed to add region: %s", e)
1056 return None
1057 if result is None:
1058 self._raise_write_conflict(endpoint_name, "added to a Region")
1059 return result
1061 def remove_region(self, endpoint_name: str, region: str) -> dict[str, Any] | None:
1062 """Remove a target without erasing its authoritative cleanup history."""
1063 store = self._get_store()
1064 loaded = self._load_mutable_endpoint(store, endpoint_name, "removed from a Region")
1065 if loaded is None:
1066 return None
1067 endpoint, lifecycle_id, updated_at = loaded
1068 if updated_at is None:
1069 raise ValueError(f"Endpoint '{endpoint_name}' has no conditional update timestamp")
1070 regions = list(endpoint.get("target_regions") or [])
1071 cleanup_regions = list(endpoint.get("cleanup_regions") or regions)
1072 if region not in regions:
1073 return endpoint
1074 regions.remove(region)
1075 if region not in cleanup_regions:
1076 cleanup_regions.append(region)
1077 raw_generations = endpoint.get("region_generations")
1078 stored_generations = raw_generations if isinstance(raw_generations, dict) else {}
1079 region_generations = {
1080 cleanup_region: (
1081 stored_generations[cleanup_region]
1082 if isinstance(stored_generations.get(cleanup_region), str)
1083 and stored_generations[cleanup_region]
1084 else secrets.token_hex(32)
1085 )
1086 for cleanup_region in cleanup_regions
1087 }
1088 # The cleanup pass must acknowledge this removal, not any earlier
1089 # remove/re-add cycle for the same Region.
1090 region_generations[region] = secrets.token_hex(32)
1091 try:
1092 result = store.update_target_regions(
1093 endpoint_name,
1094 regions,
1095 cleanup_regions,
1096 region_generations,
1097 expected_lifecycle_id=lifecycle_id,
1098 expected_updated_at=updated_at,
1099 )
1100 except Exception as e:
1101 logger.error("Failed to remove region: %s", e)
1102 return None
1103 if result is None:
1104 self._raise_write_conflict(endpoint_name, "removed from a Region")
1105 return result
1107 def canary_deploy(
1108 self,
1109 endpoint_name: str,
1110 image: str,
1111 weight: int = 10,
1112 replicas: int = 1,
1113 ) -> dict[str, Any] | None:
1114 """Start a canary deployment for an existing classic endpoint.
1116 Creates a canary variant with the new image receiving ``weight``%
1117 of traffic. Mooncake endpoints are excluded because their split-role
1118 topology cannot be represented by the classic canary Deployment.
1120 Args:
1121 endpoint_name: Existing endpoint to canary
1122 image: New container image for the canary
1123 weight: Percentage of traffic to route to canary (1-99)
1124 replicas: Positive number of canary replicas
1126 Returns:
1127 Updated endpoint record, or None if endpoint not found
1128 """
1129 if not isinstance(image, str) or not image.strip():
1130 raise ValueError("Canary image must be a non-empty string")
1131 if not _is_plain_int(weight) or not 1 <= weight <= 99:
1132 raise ValueError("Canary weight must be an integer between 1 and 99")
1133 if not _is_plain_int(replicas) or replicas < 1:
1134 raise ValueError("Canary replicas must be a positive integer")
1136 store = self._get_store()
1137 loaded = self._load_mutable_endpoint(store, endpoint_name, "updated")
1138 if loaded is None:
1139 return None
1140 endpoint, lifecycle_id, updated_at = loaded
1142 if endpoint.get("desired_state") not in ("running", "deploying"):
1143 raise ValueError(
1144 f"Cannot canary an endpoint in '{endpoint.get('desired_state')}' state. "
1145 "Endpoint must be running or deploying."
1146 )
1148 raw_spec = endpoint.get("spec")
1149 if not isinstance(raw_spec, dict):
1150 raise ValueError(f"Endpoint '{endpoint_name}' has an invalid spec")
1151 if "mooncake" in raw_spec:
1152 raise ValueError("Canary deployments are not supported for Mooncake endpoints")
1154 # Never mutate the object returned by the store; callers and test
1155 # doubles may retain it as shared state.
1156 spec = deepcopy(raw_spec)
1157 spec["canary"] = {
1158 "image": image.strip(),
1159 "weight": weight,
1160 "replicas": replicas,
1161 }
1163 result: dict[str, Any] | None = store.update_spec(
1164 endpoint_name,
1165 spec,
1166 expected_lifecycle_id=lifecycle_id,
1167 expected_updated_at=updated_at,
1168 )
1169 if result is None:
1170 self._raise_write_conflict(endpoint_name, "updated")
1171 return result
1173 def promote_canary(self, endpoint_name: str) -> dict[str, Any] | None:
1174 """Promote a classic canary to primary and remove its deployment."""
1175 store = self._get_store()
1176 loaded = self._load_mutable_endpoint(store, endpoint_name, "updated")
1177 if loaded is None:
1178 return None
1179 endpoint, lifecycle_id, updated_at = loaded
1181 raw_spec = endpoint.get("spec")
1182 if not isinstance(raw_spec, dict):
1183 raise ValueError(f"Endpoint '{endpoint_name}' has an invalid spec")
1184 if "mooncake" in raw_spec:
1185 raise ValueError("Canary promotion is not supported for Mooncake endpoints")
1187 canary = raw_spec.get("canary")
1188 if not isinstance(canary, dict):
1189 raise ValueError(f"Endpoint '{endpoint_name}' has no active canary deployment")
1190 if "image" not in canary:
1191 raise ValueError(
1192 f"Canary deployment for '{endpoint_name}' is missing the 'image' field"
1193 )
1194 canary_image = canary["image"]
1195 if not isinstance(canary_image, str) or not canary_image.strip():
1196 raise ValueError(
1197 f"Canary deployment for '{endpoint_name}' has an invalid 'image' field"
1198 )
1200 spec = deepcopy(raw_spec)
1201 spec["image"] = canary_image.strip()
1202 spec.pop("canary", None)
1203 # The canary image is explicit and global. Existing per-region primary
1204 # rewrites point at the superseded image and must not take precedence.
1205 spec.pop("region_image_uris", None)
1207 result: dict[str, Any] | None = store.update_spec(
1208 endpoint_name,
1209 spec,
1210 expected_lifecycle_id=lifecycle_id,
1211 expected_updated_at=updated_at,
1212 )
1213 if result is None:
1214 self._raise_write_conflict(endpoint_name, "updated")
1215 return result
1217 def rollback_canary(self, endpoint_name: str) -> dict[str, Any] | None:
1218 """Remove the canary deployment, keeping the primary unchanged."""
1219 store = self._get_store()
1220 loaded = self._load_mutable_endpoint(store, endpoint_name, "updated")
1221 if loaded is None:
1222 return None
1223 endpoint, lifecycle_id, updated_at = loaded
1225 raw_spec = endpoint.get("spec")
1226 if not isinstance(raw_spec, dict):
1227 raise ValueError(f"Endpoint '{endpoint_name}' has an invalid spec")
1228 if "canary" not in raw_spec:
1229 raise ValueError(f"Endpoint '{endpoint_name}' has no active canary deployment")
1231 # Rollback is deliberately allowed for a legacy invalid
1232 # Mooncake-plus-canary record so an operator can repair it.
1233 spec = deepcopy(raw_spec)
1234 spec.pop("canary", None)
1235 result: dict[str, Any] | None = store.update_spec(
1236 endpoint_name,
1237 spec,
1238 expected_lifecycle_id=lifecycle_id,
1239 expected_updated_at=updated_at,
1240 )
1241 if result is None:
1242 self._raise_write_conflict(endpoint_name, "updated")
1243 return result
1246def get_inference_manager(config: GCOConfig | None = None) -> InferenceManager:
1247 """Factory function for InferenceManager."""
1248 return InferenceManager(config)