Coverage for gco/services/inference_monitor.py: 92.66%
1194 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-30 21:22 +0000
1"""
2Inference Monitor — reconciliation controller for inference endpoints.
4Runs in each regional EKS cluster and polls the global DynamoDB table
5(gco-inference-endpoints) to reconcile desired state with actual
6Kubernetes resources. Follows a GitOps-style reconciliation pattern:
8 DynamoDB (desired state) → inference_monitor → Kubernetes (actual state)
10The monitor:
11- Creates and reconciles Deployments, ClusterIP Services, and optional autoscalers
12- Leaves public routing on the shared ``gco-system/gco-gateway`` HTTPRoute:
13 ``/inference`` -> ``gco-system/inference-proxy``
14- Removes legacy endpoint-specific Ingresses so upgrades cannot retain a bypass
15- Updates existing deployments when spec changes
16- Scales deployments up/down
17- Tears down resources when endpoints are deleted
18- Reports per-region status back to DynamoDB
20Environment Variables:
21 CLUSTER_NAME: Name of the EKS cluster
22 REGION: AWS region this monitor runs in
23 INFERENCE_ENDPOINTS_TABLE_NAME: DynamoDB table name
24 RECONCILE_INTERVAL_SECONDS: Seconds between reconciliation loops (default: 15)
25 INFERENCE_NAMESPACE: Namespace for inference workloads (default: gco-inference)
26"""
28import asyncio
29import base64
30import json
31import logging
32import os
33import re
34import secrets
35from dataclasses import dataclass, field
36from datetime import UTC, datetime
37from pathlib import Path
38from typing import Any
39from urllib.parse import urlsplit
41from kubernetes import client, config
42from kubernetes.client.models import V1Deployment
43from kubernetes.client.rest import ApiException
45from gco.services.inference_store import InferenceEndpointStore
46from gco.services.structured_logging import configure_structured_logging
48# <pyflowchart-code-diagram> BEGIN - auto-inserted, do not edit
49# Generated at (UTC): 2026-07-18T01:03:40Z
50# Flowchart(s) generated from this file:
51# * ``InferenceMonitor._reconcile_endpoint`` -> ``diagrams/code_diagrams/gco/services/inference_monitor.InferenceMonitor__reconcile_endpoint.html``
52# (PNG: ``diagrams/code_diagrams/gco/services/inference_monitor.InferenceMonitor__reconcile_endpoint.png``)
53# Regenerate with ``python diagrams/code_diagrams/generate.py``.
54# <pyflowchart-code-diagram> END
57logging.basicConfig(
58 level=logging.INFO,
59 format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
60)
61logger = logging.getLogger(__name__)
64class NetworkPolicyApplyError(Exception):
65 """An intra-namespace allow rule could not be applied.
67 Raised when materialization-time enforcement fails to create or verify one
68 of the allow rules that disaggregated inference depends on. The default-deny
69 posture is left intact — only the widening allow rule failed — and ``rule``
70 names the offending NetworkPolicy so callers can surface exactly which rule
71 could not be applied.
72 """
74 def __init__(self, rule: str, reason: str):
75 self.rule = rule
76 self.reason = reason
77 super().__init__(f"Network policy {rule!r} could not be applied: {reason}")
80class AdminApiKeySecretError(Exception):
81 """A user-named proxy admin API key Secret is missing or empty.
83 Raised before the prefill-decode proxy is materialized when a Secret named
84 by ``proxy.admin_api_key_secret`` is absent or carries no usable
85 ``ADMIN_API_KEY`` value. An endpoint that names no Secret takes the separate
86 auto-managed path and receives a generated ``{name}-admin`` Secret. The
87 proxy never starts without a usable key, so no proxy Deployment or Service
88 is created on this error. ``secret`` records the rejected Secret name.
89 """
91 def __init__(self, secret: str | None, reason: str):
92 self.secret = secret
93 self.reason = reason
94 named = repr(secret) if secret else "<unnamed>"
95 super().__init__(f"Admin API key Secret {named} is unusable: {reason}")
98# Valid TCP port boundaries for KV-transfer bootstrap ports.
99MIN_BOOTSTRAP_PORT = 1024
100MAX_BOOTSTRAP_PORT = 65535
102# vLLM kv_role for each worker role: prefill produces KV, decode consumes it,
103# and a single-instance store node both produces and consumes.
104_KV_ROLE_BY_WORKER_ROLE = {
105 "prefill": "kv_producer",
106 "decode": "kv_consumer",
107 "single": "kv_both",
108}
110# The worker roles each mooncake mode supports. Disaggregated and both split
111# work across prefill/decode; store runs a single kv_both instance.
112_WORKER_ROLES_BY_MODE = {
113 "disaggregated": {"prefill", "decode"},
114 "store": {"single"},
115 "both": {"prefill", "decode"},
116}
118# The EFA RDMA fabric is advertised as a Kubernetes extended resource, gated by
119# a node taint, and selected through a node label. KV cache transfer over
120# RoCE only runs on nodes that carry all three.
121EFA_RESOURCE_NAME = "vpc.amazonaws.com/efa"
122EFA_NODE_SELECTOR_KEY = "efa"
123EFA_NODE_SELECTOR_VALUE = "true"
125# Mooncake KV-transfer role pods are pinned to a dedicated EFA NodePool
126# (mooncake-efa-pool, manifest 46-nodepool-mooncake-efa.yaml) that only offers
127# instance families with >=80GB of GPU memory and FP8-capable Hopper/Blackwell
128# GPUs. The shared training EFA pool (43-nodepool-efa.yaml) also offers p4d
129# (A100 40GB, Ampere, no FP8), which is too small for many disaggregated/store
130# models and can be selected by Karpenter whenever a pod asks only for efa=true.
131# Selecting this extra label keeps role pods off p4d without disturbing the
132# training pool. The value must match the label on the dedicated NodePool.
133MOONCAKE_EFA_NODE_SELECTOR_KEY = "mooncake-efa"
134MOONCAKE_EFA_NODE_SELECTOR_VALUE = "true"
136# The shared per-region Mooncake master exposes its RPC service and the
137# built-in HTTP metadata server on these fixed ports.
138MOONCAKE_MASTER_RPC_PORT = 50051
139MOONCAKE_METADATA_PORT = 8080
140MOONCAKE_MASTER_SERVICE = "mooncake-master"
142# GPU utilization is not a Kubernetes Resource metric, so a native
143# HorizontalPodAutoscaler cannot scale on it (Resource metrics are limited to
144# cpu and memory). The cluster's amazon-cloudwatch-observability agent publishes
145# per-pod GPU utilization to CloudWatch ContainerInsights, so any autoscaler
146# that requests a GPU metric is materialized as a KEDA ScaledObject with an
147# aws-cloudwatch trigger instead. KEDA generates the backing HPA under the hood,
148# and cpu/memory targets ride along as native cpu/memory triggers on the same
149# ScaledObject. KEDA is a mandatory cluster component, so this path is always
150# available.
151KEDA_API_GROUP = "keda.sh"
152KEDA_API_VERSION = "v1alpha1"
153KEDA_SCALEDOBJECT_PLURAL = "scaledobjects"
155# Metric types that can only be served via CloudWatch (KEDA), keyed to the
156# ContainerInsights metric the aws-cloudwatch trigger reads. PodName in the
157# ContainerInsights dimension set is the workload (Deployment) name, so the
158# dimension triple ClusterName/Namespace/PodName yields the average across a
159# Deployment's pods — exactly the signal autoscaling needs.
160GPU_METRIC_NAMESPACE = "ContainerInsights"
161_CLOUDWATCH_METRIC_BY_TYPE = {
162 "gpu": "pod_gpu_utilization",
163 "gpu_memory": "pod_gpu_memory_utilization",
164}
166# Default base port for the KV-transfer bootstrap handshake (VLLM_MOONCAKE_
167# BOOTSTRAP_PORT) and the span of per-worker ports derived from it. vLLM assigns
168# each worker base_port + dp_rank * tp_size + tp_rank, so the intra-namespace
169# allow rule opens a contiguous window starting at the base port. A spec may
170# override the base via mooncake.transfer.bootstrap_base_port.
171MOONCAKE_BOOTSTRAP_BASE_PORT = 8998
172MOONCAKE_BOOTSTRAP_PORT_SPAN = 100
174# Environment variable through which each role pod receives the KV-transfer
175# bootstrap base port. vLLM derives per-worker ports (base + dp_rank * tp_size +
176# tp_rank) from it, so prefill and decode agree on the handshake ports.
177VLLM_MOONCAKE_BOOTSTRAP_PORT_ENV = "VLLM_MOONCAKE_BOOTSTRAP_PORT"
179# Each role pod reads the shared transport settings (metadata-server address,
180# protocol, device) from the rendered mooncake.json. The per-endpoint
181# ``{name}-mooncake`` ConfigMap is mounted read-only at the directory below, and
182# the connector is pointed at the file through MOONCAKE_CONFIG_PATH.
183MOONCAKE_CONFIG_PATH_ENV = "MOONCAKE_CONFIG_PATH"
184MOONCAKE_CONFIG_MOUNT_DIR = "/etc/mooncake"
185MOONCAKE_CONFIG_FILE_PATH = f"{MOONCAKE_CONFIG_MOUNT_DIR}/mooncake.json"
187# Label selector identifying inference workload pods (prefill/decode/proxy and
188# legacy single-Deployment endpoints). Used as both the target and the peer of
189# the intra-namespace allow rules.
190INFERENCE_POD_SELECTOR = {"gco.io/type": "inference"}
192# Names of the intra-namespace allow rules the monitor maintains alongside the
193# default-deny posture in gco-inference. These mirror the manifest names in
194# 03-network-policies.yaml so a failure can point at the same object an operator
195# would inspect with kubectl.
196NETWORK_POLICY_INFERENCE_INTERNAL = "allow-inference-internal"
197NETWORK_POLICY_POD_TO_MASTER = "allow-pod-to-master"
198NETWORK_POLICY_POD_TO_METADATA = "allow-pod-to-metadata"
199NETWORK_POLICY_RDMA_BOOTSTRAP = "allow-rdma-bootstrap"
201# Regional configuration keys through which the in-region deployment supplies
202# the shared master's address. The store cannot be wired without an own-region
203# master address, so a store-bearing endpoint defers configuration when the
204# master address is absent or blank. The metadata server defaults to the master
205# host on the metadata port when not supplied explicitly.
206MOONCAKE_MASTER_ADDRESS_ENV = "MOONCAKE_MASTER_ADDRESS"
207MOONCAKE_METADATA_SERVER_ENV = "MOONCAKE_METADATA_SERVER"
209# Container image for the shared per-region master. Supplied by the in-region
210# deployment so the master tracks the same pinned build the manifests use.
211MOONCAKE_MASTER_IMAGE_ENV = "MOONCAKE_MASTER_IMAGE"
213# Maximum time a store-bearing endpoint keeps deferring role-pod creation while
214# the shared master has not reported a Ready replica. Past this window the
215# monitor keeps deferring and stays in the ``creating`` state, but also surfaces
216# an error so operators can see the master never came up. The master itself is
217# never deleted or modified on account of this timeout.
218MOONCAKE_MASTER_READY_TIMEOUT_SECONDS = 600
220# Object-key prefix under which cold-tier KV objects are written in the
221# general-purpose regional bucket. Mirrors the value the regional stack and the
222# `gco inference populate-kv` upload surface use; kept local so the monitor
223# needs no infrastructure (CDK) imports at runtime.
224MOONCAKE_COLD_TIER_KEY_PREFIX = "mooncake-kv"
226# SSM namespace suffix publishing the always-on general-purpose regional
227# bucket's discovery values for a region. The full namespace is
228# ``/<project_name>/regional-shared-bucket`` (see
229# ``constants.regional_shared_ssm_parameter_prefix``); the monitor builds it at
230# runtime from the injected ``PROJECT_NAME`` env var rather than importing the
231# CDK constant, so it needs no infrastructure imports at runtime. Kept as a
232# suffix constant so the ``/name``, ``/arn``, ``/region`` contract stays in one
233# place. See ``_regional_shared_ssm_parameter_prefix``.
234REGIONAL_SHARED_SSM_PARAMETER_SUFFIX = "regional-shared-bucket"
237def _regional_shared_ssm_parameter_prefix() -> str:
238 """Return this deployment's regional-shared-bucket SSM namespace.
240 Built from the ``PROJECT_NAME`` environment variable (default ``"gco"``)
241 so the monitor reads the same project-scoped path the regional stack
242 writes (``/<project_name>/regional-shared-bucket``). Mirrors
243 ``constants.regional_shared_ssm_parameter_prefix`` without importing CDK.
244 """
245 project_name = os.environ.get("PROJECT_NAME", "gco")
246 return f"/{project_name}/{REGIONAL_SHARED_SSM_PARAMETER_SUFFIX}"
249# Matches an AWS region identifier embedded in an address (host or URI), e.g.
250# ``us-east-1``, ``eu-west-2``, ``ap-southeast-1``, ``us-gov-west-1``. KV
251# transfer over RoCE is intra-region, so any address a topology wires to must
252# resolve to the monitor's own region; an embedded token naming a different
253# region marks the address as out-of-region. An address that carries no token
254# (a bare in-cluster Service name) is region-local by construction.
255_REGION_TOKEN_PATTERN = re.compile(r"\b[a-z]{2}-(?:gov-)?[a-z]+-\d+\b")
257# --- PD proxy behavior -------------------------------------------------------
258#
259# The prefill-decode proxy that fronts a disaggregated endpoint checks whether a
260# prompt's KV blocks already live in the shared store before it sends the prompt
261# to a prefill pod. That check is bounded: it is given this many seconds, and a
262# miss or a check that does not finish in time is treated as "not resident" so
263# the prompt goes to prefill without the request waiting any longer. Holding the
264# bound here keeps the proxy responsive even when the store is slow or
265# unreachable.
266PD_PROXY_RESIDENCY_TIMEOUT_SECONDS = 2
268# Default strategy for spreading requests across the backends of a single role.
269PD_PROXY_DEFAULT_SCHEDULING = "round_robin"
271# Where a residency miss or timed-out lookup is sent. The prompt always goes to
272# a prefill pod in that case; the proxy never stalls the request on the store.
273PD_PROXY_RESIDENCY_MISS_TARGET = "prefill"
275# The residency lookup never blocks the request: a slow or failed store check
276# falls through to prefill rather than holding the client.
277PD_PROXY_RESIDENCY_BLOCKING = "false"
279# Decode-phase requests only ever reach decode pods that report Ready; pods that
280# are still starting are skipped.
281PD_PROXY_DECODE_ROUTING_READY_ONLY = "ready_only"
283# When no decode pod reports Ready, the proxy refuses the request outright
284# instead of streaming a partial generation. The refusal carries a stable
285# status and message so clients can distinguish "no backend yet" from a model
286# error.
287PD_PROXY_NO_DECODE_BACKEND_ACTION_REJECT = "reject"
288PD_PROXY_NO_DECODE_BACKEND_STATUS = "503"
289PD_PROXY_NO_DECODE_BACKEND_MESSAGE = "no available decode backend"
291# Environment variable names the proxy container reads to pick up the behavior
292# above. Surfacing them here keeps the proxy's runtime contract in one place;
293# the reconcile path attaches the values produced by ``build_pd_proxy_config``.
294PD_PROXY_RESIDENCY_TIMEOUT_ENV = "PD_PROXY_RESIDENCY_TIMEOUT_SECONDS"
295PD_PROXY_RESIDENCY_BLOCKING_ENV = "PD_PROXY_RESIDENCY_CHECK_BLOCKING"
296PD_PROXY_RESIDENCY_MISS_TARGET_ENV = "PD_PROXY_RESIDENCY_MISS_TARGET"
297PD_PROXY_DECODE_ROUTING_ENV = "PD_PROXY_DECODE_ROUTING"
298PD_PROXY_NO_DECODE_BACKEND_ACTION_ENV = "PD_PROXY_NO_DECODE_BACKEND_ACTION"
299PD_PROXY_NO_DECODE_BACKEND_STATUS_ENV = "PD_PROXY_NO_DECODE_BACKEND_STATUS"
300PD_PROXY_NO_DECODE_BACKEND_MESSAGE_ENV = "PD_PROXY_NO_DECODE_BACKEND_MESSAGE"
301PD_PROXY_SCHEDULING_ENV = "PD_PROXY_SCHEDULING"
302PD_PROXY_STORE_ADDRESS_ENV = "PD_PROXY_STORE_ADDRESS"
304# Marker label carried by proxy pods so a Service can select the proxy alone,
305# distinct from the prefill/decode role pods (which carry their own role marker).
306PD_PROXY_ROLE_LABEL = "proxy"
308# TCP port the proxy container listens on for authenticated serving requests.
309PD_PROXY_PORT = 8000
311# Environment variable the proxy reads its admin key from, and the data key the
312# backing Kubernetes Secret stores it under. The key value is delivered to the
313# container through a Secret reference at pod start — it is never written to the
314# endpoint spec or passed as a command-line argument.
315PD_PROXY_ADMIN_API_KEY_ENV = "ADMIN_API_KEY"
316ADMIN_API_KEY_SECRET_DATA_KEY = "ADMIN_API_KEY"
318# The proxy program (gco/services/mooncake_pd_proxy.py) is shipped to the proxy
319# pod as a ConfigMap and run from this mount path. The prefill/decode backend
320# URLs and the listen port are passed to it through these env vars; it routes to
321# the role pods through their in-cluster Services so kube-proxy load-balances
322# across only the Ready endpoints of each role.
323PD_PROXY_SCRIPT_FILENAME = "mooncake_pd_proxy.py"
324PD_PROXY_CONFIG_MOUNT_DIR = "/etc/pd-proxy"
325PD_PROXY_SCRIPT_PATH = f"{PD_PROXY_CONFIG_MOUNT_DIR}/{PD_PROXY_SCRIPT_FILENAME}"
326PD_PROXY_PORT_ENV = "PD_PROXY_PORT"
327PD_PROXY_PREFILL_URL_ENV = "PD_PROXY_PREFILL_URL"
328PD_PROXY_DECODE_URL_ENV = "PD_PROXY_DECODE_URL"
331@dataclass
332class RegionServicesResolution:
333 """Outcome of resolving the in-region service addresses an endpoint needs.
335 ``render_mooncake_config`` consumes already-resolved values via a
336 ``region_services`` dict; this carries that dict together with the signals
337 a reconcile pass acts on:
339 - ``region_services`` is the resolved dict to render with, or ``None`` when
340 rendering must be skipped.
341 - ``render_skipped`` is set when the store is enabled but the own-region
342 master address is not configured: the existing endpoint configuration is
343 left untouched and ``store_master_unresolved`` records why.
344 - ``cold_tier_unresolved`` is set when the cold tier was requested but the
345 own-region general-purpose bucket could not be resolved; the cold tier is
346 dropped while the hot-path store keeps operating, and ``error`` explains
347 the condition.
348 """
350 region_services: dict[str, Any] | None = None
351 render_skipped: bool = False
352 store_master_unresolved: bool = False
353 cold_tier_unresolved: bool = False
354 error: str | None = None
357@dataclass
358class MasterReadinessGate:
359 """Outcome of gating dependent role-pod creation on the shared master.
361 A store-bearing endpoint must not materialize its role pods until the
362 single shared ``mooncake-master`` reports a Ready replica. This carries the
363 decision a reconcile pass acts on:
365 - ``proceed`` is ``True`` only when the master reports at least one Ready
366 replica; the caller may then create the dependent role pods and advance
367 out of ``creating``. While it is ``False`` the caller materializes no
368 dependent pods.
369 - ``state`` is the endpoint state to report. It is ``"creating"`` whenever
370 creation is deferred (master not ready, still within the wait window,
371 timed out, or could not be created) and ``None`` when the gate is open.
372 - ``error`` records why creation could not advance: the master did not
373 become Ready within the wait window, or the master could not be created.
374 It is ``None`` while the master is simply still coming up within the
375 window, and ``None`` once the gate is open.
376 """
378 proceed: bool = False
379 state: str | None = None
380 error: str | None = None
383@dataclass
384class RegionalScopeResolution:
385 """Outcome of confirming a disaggregated topology stays inside one region.
387 KV cache transfer over RoCE cannot cross a region boundary, so every
388 ``MooncakeConnector`` peer address and the ``master_server_address`` a
389 topology wires to must resolve to the monitor's own region. An endpoint
390 that targets several regions runs one independent topology per region; each
391 region's monitor reconciles only its own topology and confirms that
392 topology's addresses never escape the region.
394 - ``in_region`` is ``True`` only when every resolved address belongs to the
395 monitor's own region. While it is ``True`` the caller may materialize the
396 topology's role Deployments.
397 - ``peer_addresses`` lists the addresses that were resolved and checked, in
398 a stable order, so callers and logs can show exactly what was wired.
399 - ``state`` is the endpoint state to report. It is ``"failed"`` when an
400 out-of-region address is detected and ``None`` when the topology is
401 wholly in-region.
402 - ``error`` describes the cross-region boundary violation — which addresses
403 resolved to which other regions — when one is found, and is ``None``
404 otherwise. When a violation is reported the caller materializes no role
405 Deployments and leaves any previously materialized resources unchanged.
406 """
408 in_region: bool = True
409 peer_addresses: list[str] = field(default_factory=list)
410 state: str | None = None
411 error: str | None = None
414def _resolved_mooncake_transfer(mooncake: dict[str, Any]) -> tuple[str, str]:
415 """Resolve the persisted transfer intent and optional network device.
417 GCO's spec deliberately uses ``rdma`` as the portable high-performance
418 intent because the mounted Mooncake store configuration accepts
419 ``rdma|tcp``. On AWS, role pods with that intent are placed on EFA nodes,
420 so :func:`build_kv_transfer_config` translates it to vLLM's explicit
421 ``mooncake_protocol=efa`` at the point-to-point connector boundary.
422 """
423 transfer = mooncake.get("transfer", {})
424 if not isinstance(transfer, dict): 424 ↛ 425line 424 didn't jump to line 425 because the condition on line 424 was never true
425 raise ValueError("mooncake.transfer must be a mapping")
427 protocol = transfer.get("protocol", "rdma")
428 if protocol not in {"rdma", "tcp"}: 428 ↛ 429line 428 didn't jump to line 429 because the condition on line 428 was never true
429 raise ValueError(
430 f"mooncake.transfer.protocol must be one of {{rdma, tcp}}, got {protocol!r}"
431 )
432 device_name = transfer.get("device_name", "")
433 if not isinstance(device_name, str): 433 ↛ 434line 433 didn't jump to line 434 because the condition on line 433 was never true
434 raise ValueError(f"mooncake.transfer.device_name must be a string, got {device_name!r}")
435 return protocol, device_name
438def build_kv_transfer_config(mooncake: dict[str, Any], role: str) -> str:
439 """Return the JSON string for vLLM's ``--kv-transfer-config``.
441 Translates a mooncake spec block plus a worker role into the connector
442 configuration vLLM expects:
444 - ``disaggregated`` emits a ``MooncakeConnector``.
445 - ``store`` emits a ``MooncakeStoreConnector``.
446 - ``both`` emits a ``MultiConnector`` wrapping a ``MooncakeConnector``
447 (index 0) followed by a ``MooncakeStoreConnector`` (index 1), both
448 sharing the role's ``kv_role``.
450 Every point-to-point ``MooncakeConnector`` receives explicit
451 ``kv_connector_extra_config``. GCO's default/high-performance ``rdma``
452 intent maps to Mooncake's AWS-specific ``efa`` protocol because the same
453 pod is pinned to the EFA node pool; ``tcp`` remains an explicit fallback.
454 ``device_name`` is forwarded verbatim, with an empty string requesting
455 Mooncake/libfabric auto-detection.
457 The emitted ``kv_role`` is ``kv_producer`` for prefill, ``kv_consumer``
458 for decode, and ``kv_both`` for a single store instance.
460 Args:
461 mooncake: The ``spec["mooncake"]`` block; its ``mode`` selects the
462 connector shape and its optional ``transfer`` block selects the
463 protocol/device.
464 role: One of ``"prefill"``, ``"decode"``, or ``"single"``.
466 Returns:
467 A JSON object string parseable by vLLM.
469 Raises:
470 ValueError: If the ``(mode, role)`` combination or transfer settings
471 are unsupported. No configuration is emitted in that case.
472 """
473 mode = mooncake.get("mode")
474 supported_roles = _WORKER_ROLES_BY_MODE.get(mode) if isinstance(mode, str) else None
475 if supported_roles is None or role not in supported_roles:
476 raise ValueError(f"Unsupported (mode, role) pair: ({mode!r}, {role!r})")
478 kv_role = _KV_ROLE_BY_WORKER_ROLE[role]
480 if mode == "store":
481 return json.dumps({"kv_connector": "MooncakeStoreConnector", "kv_role": kv_role})
483 protocol, device_name = _resolved_mooncake_transfer(mooncake)
484 connector = {
485 "kv_connector": "MooncakeConnector",
486 "kv_role": kv_role,
487 "kv_connector_extra_config": {
488 "mooncake_protocol": "efa" if protocol == "rdma" else "tcp",
489 "device_name": device_name,
490 },
491 }
492 if mode == "disaggregated":
493 return json.dumps(connector)
495 # mode == "both": MultiConnector chains transfer then store.
496 return json.dumps(
497 {
498 "kv_connector": "MultiConnector",
499 "kv_role": kv_role,
500 "kv_connector_extra_config": {
501 "connectors": [
502 connector,
503 {"kv_connector": "MooncakeStoreConnector", "kv_role": kv_role},
504 ]
505 },
506 }
507 )
510def bootstrap_port_for_worker(base_port: int, dp_rank: int, tp_size: int, tp_rank: int) -> int:
511 """Compute the bootstrap port for a ``(dp_rank, tp_rank)`` worker.
513 The port is ``base_port + dp_rank * tp_size + tp_rank``. For a fixed
514 ``base_port`` and ``tp_size`` distinct ``(dp_rank, tp_rank)`` pairs map to
515 distinct ports.
517 Args:
518 base_port: The base bootstrap port for the endpoint.
519 dp_rank: The data-parallel rank of the worker (``>= 0``).
520 tp_size: The tensor-parallel world size (``>= 1``).
521 tp_rank: The tensor-parallel rank within the worker (``0 <= tp_rank < tp_size``).
523 Returns:
524 The TCP port assigned to the worker.
526 Raises:
527 ValueError: If the computed port falls outside the valid range
528 ``1024..65535``. No port is assigned in that case.
529 """
530 port = base_port + dp_rank * tp_size + tp_rank
531 if port < MIN_BOOTSTRAP_PORT or port > MAX_BOOTSTRAP_PORT: 531 ↛ 532line 531 didn't jump to line 532 because the condition on line 531 was never true
532 raise ValueError(
533 f"Computed bootstrap port {port} is outside the valid range "
534 f"{MIN_BOOTSTRAP_PORT}..{MAX_BOOTSTRAP_PORT}"
535 )
536 return port
539def render_mooncake_config(
540 mooncake: dict[str, Any], region_services: dict[str, Any]
541) -> dict[str, Any]:
542 """Render the ``mooncake.json`` contents mounted into each vLLM pod.
544 The returned dict is written verbatim to a ConfigMap and mounted at the
545 path named by ``MOONCAKE_CONFIG_PATH``. It always carries the metadata
546 server and the RDMA/TCP transport settings (``protocol`` and
547 ``device_name``); the key-value store and its optional cold tier are layered
548 on only when requested.
550 The transport block (``protocol``/``device_name``) describes the hot
551 RDMA/RoCE path. The cold tier is an asynchronous object-store backend keyed
552 separately as ``cold_tier_s3_uri``; it is never wired into the transport
553 block, so cold-tier reads and writes stay off the RDMA hot path.
555 Resolution of in-region addresses is the caller's responsibility: this
556 function consumes already-resolved values from ``region_services`` and
557 performs no lookups of its own. In particular, the cold-tier URI is the
558 monitor-resolved general-purpose regional bucket for the monitor's own
559 region; any cold-tier bucket URI in the user spec is ignored.
561 Args:
562 mooncake: The ``spec["mooncake"]`` block.
563 region_services: Resolved in-region addresses, e.g.::
565 {
566 "metadata_server": "http://mooncake-master:8080/metadata",
567 "master_server_address": "mooncake-master:50051",
568 "cold_tier_s3_uri": "s3://gco-regional-shared-<acct>-<region>/...",
569 }
571 ``master_server_address`` is required when the store is enabled and
572 ``cold_tier_s3_uri`` is required when the cold tier is enabled.
574 Returns:
575 The ``mooncake.json`` contents as a dict, where:
577 - ``protocol`` and ``device_name`` are always present.
578 - ``master_server_address`` is present only when the store is enabled.
579 - ``cold_tier_s3_uri`` is present only when the cold tier is enabled,
580 which requires ``cold_tier_enabled`` to be the boolean ``True``; any
581 other value (absent, null, truthy non-bool) leaves the cold tier off.
582 """
583 protocol, device_name = _resolved_mooncake_transfer(mooncake)
584 store = mooncake.get("store", {})
585 cfg: dict[str, Any] = {
586 "metadata_server": region_services["metadata_server"],
587 "protocol": protocol,
588 "device_name": device_name,
589 }
590 if store.get("enabled"):
591 cfg["master_server_address"] = region_services["master_server_address"]
592 # The store runs embedded in each vLLM pod (every rank contributes
593 # `global_segment_size` to the shared pool; GCO's per-region
594 # mooncake-master is only the metadata/master coordinator, not a
595 # standalone store that owns the pool). Embedded mode rejects a zero
596 # segment, so default to 4 GiB (the upstream default) when the spec
597 # does not set one; an operator can tune it via configure-store.
598 cfg["global_segment_size"] = store.get("global_segment_size", "4294967296")
599 cfg["local_buffer_size"] = store.get("local_buffer_size", "2147483648")
600 # Only the boolean True enables the cold tier; any other value leaves it
601 # off. The URI is resolved by the caller for the monitor's own region —
602 # never authored by the user — and is an object-store backend kept off
603 # the RDMA transport block above.
604 if store.get("cold_tier_enabled") is True:
605 cfg["cold_tier_s3_uri"] = region_services["cold_tier_s3_uri"]
606 return cfg
609def apply_efa_scheduling(mooncake: dict[str, Any], pod_spec: client.V1PodSpec) -> None:
610 """Place a role pod on the EFA RDMA fabric when transfer runs over RDMA.
612 KV cache transfer over RoCE only runs on EFA-enabled nodes, which carry a
613 ``vpc.amazonaws.com/efa`` taint, advertise the ``vpc.amazonaws.com/efa``
614 extended resource, and are labelled ``efa=true``. When the transfer
615 protocol is ``rdma`` this mutates ``pod_spec`` in place to:
617 - add a ``vpc.amazonaws.com/efa`` toleration (in addition to any existing
618 tolerations such as the GPU one),
619 - add an ``efa=true`` node selector plus a ``mooncake-efa=true`` node
620 selector (merged with any existing selectors), and
621 - request at least one ``vpc.amazonaws.com/efa`` device on the pod's
622 containers, leaving every existing resource request and limit — including
623 GPU asks — untouched.
625 The ``mooncake-efa=true`` selector pins the pod to the dedicated
626 ``mooncake-efa-pool`` NodePool, which only offers instance families with
627 >=80GB of GPU memory and FP8-capable Hopper/Blackwell GPUs. This keeps role
628 pods off the A100-40GB ``p4d`` family that the shared training EFA pool
629 still offers — that family OOMs on many models and cannot run FP8 KV-cache
630 configs, so Karpenter selecting it for a mooncake pod is a latent failure.
632 When the transfer protocol is explicitly set to anything other than
633 ``rdma`` (for example ``tcp``) the pod is left exactly as it was: no
634 toleration, no node selector, and no device request are added. An unset
635 protocol defaults to ``rdma`` — matching the rest of the Mooncake path — so
636 a disaggregated endpoint lands on EFA by default.
638 Tolerations, selectors, and device requests are applied idempotently, so
639 re-running over an already-scheduled pod produces no duplicates.
641 Args:
642 mooncake: The ``spec["mooncake"]`` block; ``transfer.protocol``
643 (defaulting to ``rdma`` when unset) decides whether EFA scheduling
644 applies.
645 pod_spec: The pod specification to mutate in place.
646 """
647 protocol, _device_name = _resolved_mooncake_transfer(mooncake)
648 if protocol != "rdma":
649 return
651 # Tolerate the EFA taint without disturbing existing tolerations.
652 tolerations = list(pod_spec.tolerations or [])
653 if not any(t.key == EFA_RESOURCE_NAME for t in tolerations): 653 ↛ 662line 653 didn't jump to line 662 because the condition on line 653 was always true
654 tolerations.append(
655 client.V1Toleration(
656 key=EFA_RESOURCE_NAME,
657 operator="Equal",
658 value="true",
659 effect="NoSchedule",
660 )
661 )
662 pod_spec.tolerations = tolerations
664 # Merge the EFA node selectors with any selectors already in place. The
665 # generic efa=true selector lands the pod on EFA fabric; mooncake-efa=true
666 # narrows that to the dedicated mooncake-efa-pool, which excludes the
667 # A100-40GB p4d family that the shared training EFA pool still offers.
668 node_selector = dict(pod_spec.node_selector or {})
669 node_selector[EFA_NODE_SELECTOR_KEY] = EFA_NODE_SELECTOR_VALUE
670 node_selector[MOONCAKE_EFA_NODE_SELECTOR_KEY] = MOONCAKE_EFA_NODE_SELECTOR_VALUE
671 pod_spec.node_selector = node_selector
673 # Request at least one EFA device, preserving existing requests and limits
674 # (notably the GPU asks). Apply to containers that already request an
675 # accelerator; if none do, apply to every container so the pod still asks
676 # for the fabric it needs.
677 containers = pod_spec.containers or []
678 accelerator_keys = ("nvidia.com/gpu", "aws.amazon.com/neuron")
680 def _requests_accelerator(container: client.V1Container) -> bool:
681 reqs = container.resources
682 if reqs is None: 682 ↛ 683line 682 didn't jump to line 683 because the condition on line 682 was never true
683 return False
684 for table in (reqs.requests, reqs.limits): 684 ↛ 687line 684 didn't jump to line 687 because the loop on line 684 didn't complete
685 if table and any(key in table for key in accelerator_keys): 685 ↛ 684line 685 didn't jump to line 684 because the condition on line 685 was always true
686 return True
687 return False
689 targets = [c for c in containers if _requests_accelerator(c)] or list(containers)
690 for container in targets:
691 if container.resources is None: 691 ↛ 692line 691 didn't jump to line 692 because the condition on line 691 was never true
692 container.resources = client.V1ResourceRequirements()
693 if container.resources.requests is None: 693 ↛ 694line 693 didn't jump to line 694 because the condition on line 693 was never true
694 container.resources.requests = {}
695 if container.resources.limits is None: 695 ↛ 696line 695 didn't jump to line 696 because the condition on line 695 was never true
696 container.resources.limits = {}
697 container.resources.requests.setdefault(EFA_RESOURCE_NAME, "1")
698 container.resources.limits.setdefault(EFA_RESOURCE_NAME, "1")
701def build_pd_proxy_config(mooncake: dict[str, Any]) -> dict[str, str]:
702 """Return the environment the prefill-decode proxy runs with.
704 The proxy fronts a disaggregated endpoint and decides, per request, whether
705 to consult the shared store and which backends to dispatch to. Its behavior
706 is fixed by the values returned here so every disaggregated endpoint front
707 behaves identically:
709 - It looks up whether the prompt's KV blocks already reside in the store
710 before sending the prompt to prefill, and that lookup is bounded to
711 ``PD_PROXY_RESIDENCY_TIMEOUT_SECONDS`` seconds.
712 - A miss, or a lookup that does not finish in time, is treated as "not
713 resident": the prompt is sent to a prefill pod and the request is never
714 held waiting on the store.
715 - Decode-phase requests are routed only to decode pods reporting Ready, so a
716 pod that is still starting is skipped.
717 - When no decode pod reports Ready, the proxy refuses the request with a
718 stable status and message rather than streaming any partial output.
720 The residency bound is held constant rather than read from the spec so the
721 responsiveness guarantee cannot be weakened per endpoint. The store address
722 points at the shared in-region master, and the same-role dispatch strategy
723 falls back to round-robin when the spec names none.
725 Args:
726 mooncake: The ``spec["mooncake"]`` block; its optional ``proxy`` section
727 supplies the same-role scheduling strategy.
729 Returns:
730 A mapping of environment variable name to value, ready to attach to the
731 proxy container.
732 """
733 proxy = mooncake.get("proxy", {}) or {}
734 scheduling = proxy.get("scheduling") or PD_PROXY_DEFAULT_SCHEDULING
735 store_address = f"{MOONCAKE_MASTER_SERVICE}:{MOONCAKE_MASTER_RPC_PORT}"
736 return {
737 PD_PROXY_RESIDENCY_TIMEOUT_ENV: str(PD_PROXY_RESIDENCY_TIMEOUT_SECONDS),
738 PD_PROXY_RESIDENCY_BLOCKING_ENV: PD_PROXY_RESIDENCY_BLOCKING,
739 PD_PROXY_RESIDENCY_MISS_TARGET_ENV: PD_PROXY_RESIDENCY_MISS_TARGET,
740 PD_PROXY_DECODE_ROUTING_ENV: PD_PROXY_DECODE_ROUTING_READY_ONLY,
741 PD_PROXY_NO_DECODE_BACKEND_ACTION_ENV: PD_PROXY_NO_DECODE_BACKEND_ACTION_REJECT,
742 PD_PROXY_NO_DECODE_BACKEND_STATUS_ENV: PD_PROXY_NO_DECODE_BACKEND_STATUS,
743 PD_PROXY_NO_DECODE_BACKEND_MESSAGE_ENV: PD_PROXY_NO_DECODE_BACKEND_MESSAGE,
744 PD_PROXY_SCHEDULING_ENV: scheduling,
745 PD_PROXY_STORE_ADDRESS_ENV: store_address,
746 }
749class InferenceMonitor:
750 """
751 Reconciliation controller for inference endpoints.
753 Polls DynamoDB for desired endpoint state and reconciles with
754 the actual Kubernetes resources in the local cluster.
755 """
757 def __init__(
758 self,
759 cluster_id: str,
760 region: str,
761 store: InferenceEndpointStore,
762 namespace: str = "gco-inference",
763 reconcile_interval: int = 15,
764 ):
765 self.cluster_id = cluster_id
766 self.region = region
767 self.store = store
768 self.namespace = namespace
769 self.reconcile_interval = reconcile_interval
770 self._running = False
772 # Initialize Kubernetes clients
773 try:
774 config.load_incluster_config()
775 logger.info("Loaded in-cluster Kubernetes configuration")
776 except config.ConfigException:
777 try:
778 config.load_kube_config()
779 logger.info("Loaded local Kubernetes configuration")
780 except config.ConfigException as e:
781 logger.error("Failed to load Kubernetes configuration: %s", e)
782 raise
784 self.apps_v1 = client.AppsV1Api()
785 self.core_v1 = client.CoreV1Api()
786 self.networking_v1 = client.NetworkingV1Api()
788 # Timeout for Kubernetes API calls (seconds)
789 self._k8s_timeout = int(os.environ.get("K8S_API_TIMEOUT", "30"))
791 # Health watchdog: tracks when each endpoint first became unready.
792 # Inference traffic enters through the shared ``gco-system/gco-gateway``
793 # HTTPRoute at ``/inference`` and then ``gco-system/inference-proxy``, so
794 # model readiness never mutates shared Gateway API resources. Once this
795 # threshold is exceeded, reconciliation emits an explicit degraded-state
796 # warning while the proxy continues returning 503 until a replica is ready.
797 self._unready_since: dict[str, datetime] = {}
798 self._unhealthy_threshold_seconds = int(
799 os.environ.get("INFERENCE_UNHEALTHY_THRESHOLD_SECONDS", "300")
800 ) # 5 minutes default
802 # Master-readiness gate: tracks when each store-bearing endpoint first
803 # deferred its role-pod creation because the shared master was not yet
804 # Ready. The entry is cleared once the master reports a Ready replica so
805 # a later restart of the master restarts the clock cleanly.
806 self._master_deferral_since: dict[str, datetime] = {}
808 # Metrics
809 self._reconcile_count = 0
810 self._errors_count = 0
812 # ------------------------------------------------------------------
813 # Reconciliation loop
814 # ------------------------------------------------------------------
816 async def start(self) -> None:
817 """Start the reconciliation loop with leader election.
819 Uses a Kubernetes Lease object for leader election so that only
820 one replica reconciles at a time. Other replicas stay on standby
821 and take over if the leader dies.
822 """
823 if self._running:
824 logger.warning("Inference monitor already running")
825 return
826 self._running = True
827 logger.info(
828 "Starting inference monitor for %s in %s (interval=%ds)",
829 self.cluster_id,
830 self.region,
831 self.reconcile_interval,
832 )
834 # Namespace and ServiceAccount are pre-created by the kubectl-applier
835 # at deploy time (00-namespaces.yaml, 01-serviceaccounts.yaml). The
836 # inference-monitor SA has namespace-scoped RBAC only — it cannot
837 # read_namespace/create_namespace, so we don't try. If the namespace
838 # is ever missing, deployments below will fail with a clear 404.
840 # Get pod identity for leader election
841 pod_name = os.environ.get("HOSTNAME", f"monitor-{id(self)}")
842 lease_name = "inference-monitor-leader"
844 while self._running:
845 try:
846 if self._try_acquire_lease(lease_name, pod_name): 846 ↛ 849line 846 didn't jump to line 849 because the condition on line 846 was always true
847 await self.reconcile()
848 else:
849 logger.debug("Not the leader, waiting...")
850 except Exception as e:
851 logger.error("Reconciliation error: %s", e, exc_info=True)
852 self._errors_count += 1
853 try:
854 await asyncio.sleep(self.reconcile_interval)
855 except Exception as e:
856 logger.error("Sleep interrupted: %s", e)
857 break
859 def _try_acquire_lease(self, lease_name: str, holder: str) -> bool:
860 """Try to acquire or renew a Kubernetes Lease for leader election.
862 Uses optimistic concurrency via resourceVersion — if two monitors
863 race to update the same lease, K8s returns 409 Conflict for the
864 loser, preventing split-brain.
866 Returns True if this instance is the leader.
867 """
869 coordination_v1 = client.CoordinationV1Api()
870 now = datetime.now(UTC)
872 try:
873 lease = coordination_v1.read_namespaced_lease(lease_name, self.namespace)
874 current_holder = lease.spec.holder_identity
875 renew_time = lease.spec.renew_time
877 # Check if lease is expired (holder hasn't renewed in 3x interval)
878 if renew_time:
879 elapsed = (now - renew_time.replace(tzinfo=UTC)).total_seconds()
880 if elapsed > self.reconcile_interval * 3:
881 # Lease expired — take over
882 logger.info("Lease expired (held by %s), taking over", current_holder)
883 current_holder = None
885 if current_holder == holder:
886 # We're the leader — renew
887 lease.spec.renew_time = now
888 try:
889 coordination_v1.replace_namespaced_lease(lease_name, self.namespace, lease)
890 except ApiException as conflict:
891 if conflict.status == 409:
892 logger.debug("Lease renew conflict (another writer), retrying next cycle")
893 return False
894 raise
895 return True
896 if current_holder is None or current_holder == "":
897 # No leader — claim it
898 lease.spec.holder_identity = holder
899 lease.spec.renew_time = now
900 try:
901 coordination_v1.replace_namespaced_lease(lease_name, self.namespace, lease)
902 except ApiException as conflict:
903 if conflict.status == 409: 903 ↛ 906line 903 didn't jump to line 906 because the condition on line 903 was always true
904 logger.info("Lost lease race to another monitor")
905 return False
906 raise
907 logger.info("Acquired leader lease as %s", holder)
908 return True
909 # Someone else is the leader
910 return False
912 except ApiException as e:
913 if e.status == 404:
914 # Lease doesn't exist — create it
915 lease = client.V1Lease(
916 metadata=client.V1ObjectMeta(
917 name=lease_name,
918 namespace=self.namespace,
919 ),
920 spec=client.V1LeaseSpec(
921 holder_identity=holder,
922 lease_duration_seconds=self.reconcile_interval * 3,
923 renew_time=now,
924 ),
925 )
926 try:
927 coordination_v1.create_namespaced_lease(self.namespace, lease)
928 logger.info("Created leader lease as %s", holder)
929 return True
930 except ApiException:
931 return False
932 logger.warning("Lease check failed: %s", e.reason)
933 return False
935 def stop(self) -> None:
936 """Stop the reconciliation loop."""
937 self._running = False
938 logger.info("Inference monitor stopped")
940 async def reconcile(self) -> list[dict[str, Any]]:
941 """
942 Run one reconciliation cycle.
944 Returns a list of actions taken (for logging/testing).
945 """
946 self._reconcile_count += 1
947 actions: list[dict[str, Any]] = []
949 # Get all endpoints from DynamoDB
950 try:
951 endpoints = self.store.list_endpoints()
952 except Exception as e:
953 logger.error("Failed to list endpoints from DynamoDB: %s", e)
954 return actions
956 for endpoint in endpoints:
957 try:
958 action = await self._reconcile_endpoint(endpoint)
959 if action:
960 actions.append(action)
961 except Exception as e:
962 name = endpoint.get("endpoint_name", "unknown")
963 logger.error("Failed to reconcile endpoint %s: %s", name, e)
964 self._errors_count += 1
965 self.store.update_region_status(
966 name,
967 self.region,
968 "error",
969 error=str(e),
970 )
972 # Purge fully-deleted endpoints from DynamoDB to prevent unbounded growth.
973 # An endpoint is fully deleted when desired_state is "deleted" and all
974 # target regions report "deleted" status.
975 for endpoint in endpoints:
976 if endpoint.get("desired_state") != "deleted":
977 continue
978 region_status = endpoint.get("region_status", {})
979 target_regions = endpoint.get("target_regions", [])
980 if not target_regions:
981 continue
982 all_deleted = all(
983 isinstance(region_status.get(r), dict)
984 and region_status.get(r, {}).get("state") == "deleted"
985 for r in target_regions
986 )
987 if all_deleted:
988 ep_name = endpoint["endpoint_name"]
989 try:
990 self.store.delete_endpoint(ep_name)
991 logger.info("Purged fully-deleted endpoint %s from DynamoDB", ep_name)
992 actions.append({"action": "purge", "endpoint": ep_name})
993 except Exception as e:
994 logger.warning("Failed to purge endpoint %s: %s", ep_name, e)
996 return actions
998 async def _reconcile_endpoint(self, endpoint: dict[str, Any]) -> dict[str, Any] | None:
999 """Reconcile a single endpoint."""
1000 name = endpoint["endpoint_name"]
1001 desired_state = endpoint.get("desired_state", "deploying")
1002 target_regions = endpoint.get("target_regions", [])
1003 spec = endpoint.get("spec", {})
1004 ns = endpoint.get("namespace", self.namespace)
1006 # Am I a target region?
1007 if self.region not in target_regions:
1008 # A classic endpoint uses ``name`` while a split endpoint has only
1009 # role/proxy Deployments. Check every materialized shape so removing
1010 # a region cannot strand Mooncake resources and their GPU nodes.
1011 deployment_names = (
1012 name,
1013 f"{name}-prefill",
1014 f"{name}-decode",
1015 f"{name}-proxy",
1016 )
1017 if any(self._deployment_exists(d, ns) for d in deployment_names):
1018 logger.info(
1019 "Endpoint %s no longer targets %s, cleaning up",
1020 name,
1021 self.region,
1022 )
1023 self._delete_resources(name, ns, spec if isinstance(spec, dict) else None)
1024 self.store.update_region_status(
1025 name,
1026 self.region,
1027 "deleted",
1028 )
1029 return {"action": "cleanup", "endpoint": name, "reason": "region_removed"}
1030 return None
1032 # Reconcile based on desired state
1033 if desired_state in ("deploying", "running"):
1034 if not isinstance(spec, dict): 1034 ↛ 1035line 1034 didn't jump to line 1035 because the condition on line 1034 was never true
1035 error = "endpoint spec must be a mapping"
1036 elif "mooncake" in spec and "canary" in spec:
1037 error = "endpoint spec cannot combine 'mooncake' and 'canary' blocks"
1038 else:
1039 return await self._reconcile_running(name, ns, spec, endpoint)
1041 logger.error("Rejecting invalid endpoint %s: %s", name, error)
1042 self.store.update_region_status(name, self.region, "failed", error=error)
1043 return {"action": "reject", "endpoint": name, "reason": "invalid_spec"}
1044 if desired_state == "stopped":
1045 return self._reconcile_stopped(name, ns)
1046 if desired_state == "deleted":
1047 return self._reconcile_deleted(name, ns, spec if isinstance(spec, dict) else None)
1049 return None
1051 async def _reconcile_running(
1052 self,
1053 name: str,
1054 namespace: str,
1055 spec: dict[str, Any],
1056 endpoint: dict[str, Any],
1057 ) -> dict[str, Any] | None:
1058 """Ensure the endpoint is running with the correct spec."""
1059 # Specs carrying a ``mooncake`` block take the disaggregated path. The
1060 # branch returns ``None`` when no such block is present, so a plain
1061 # endpoint falls through to the single-Deployment path below unchanged.
1062 mooncake_action = await self._reconcile_mooncake(name, namespace, spec, endpoint)
1063 if mooncake_action is not None:
1064 return mooncake_action
1066 deployment = self._get_deployment(name, namespace)
1068 if deployment is None:
1069 # Create everything
1070 logger.info("Creating endpoint %s in %s", name, self.region)
1071 self._create_deployment(name, namespace, spec)
1072 self._create_service(name, namespace, spec)
1073 if spec.get("autoscaling", {}).get("enabled"):
1074 self._create_or_update_hpa(name, namespace, spec)
1075 self.store.update_region_status(
1076 name,
1077 self.region,
1078 "creating",
1079 replicas_desired=spec.get("replicas", 1),
1080 )
1081 return {"action": "create", "endpoint": name}
1083 # Deployment exists — ensure its Service exists. Public traffic follows
1084 # ``gco-system/gco-gateway``'s shared ``/inference`` HTTPRoute to
1085 # ``gco-system/inference-proxy``, which then reaches this endpoint's
1086 # ClusterIP Service.
1087 self._ensure_service(name, namespace, spec)
1089 desired_replicas = spec.get("replicas", 1)
1090 current_replicas = deployment.spec.replicas or 1
1091 ready_replicas = deployment.status.ready_replicas or 0
1093 self._check_health_watchdog(
1094 name, namespace, ready_replicas, desired_replicas, spec, endpoint
1095 )
1097 if current_replicas != desired_replicas:
1098 logger.info(
1099 "Scaling endpoint %s: %d → %d replicas",
1100 name,
1101 current_replicas,
1102 desired_replicas,
1103 )
1104 self._scale_deployment(name, namespace, desired_replicas)
1105 self.store.update_region_status(
1106 name,
1107 self.region,
1108 "updating",
1109 replicas_ready=ready_replicas,
1110 replicas_desired=desired_replicas,
1111 )
1112 return {"action": "scale", "endpoint": name, "replicas": desired_replicas}
1114 # Check if image changed
1115 current_image = self._get_deployment_image(deployment)
1116 desired_image = self._resolve_image_for_region(spec) if spec.get("image") else ""
1117 if current_image and desired_image and current_image != desired_image:
1118 logger.info("Updating endpoint %s image: %s → %s", name, current_image, desired_image)
1119 self._update_deployment_image(name, namespace, desired_image)
1120 self.store.update_region_status(
1121 name,
1122 self.region,
1123 "updating",
1124 replicas_ready=ready_replicas,
1125 replicas_desired=desired_replicas,
1126 )
1127 return {"action": "update_image", "endpoint": name, "image": desired_image}
1129 # Reconcile canary first and publish only observed readiness. The
1130 # authenticated proxy will not sample canary traffic until this exact
1131 # region reports the matching image fully Ready.
1132 canary = spec.get("canary")
1133 canary_status = None
1134 if isinstance(canary, dict): 1134 ↛ 1135line 1134 didn't jump to line 1135 because the condition on line 1134 was never true
1135 canary_status = self._reconcile_canary(name, namespace, spec, canary, endpoint)
1136 else:
1137 self._cleanup_canary(name, namespace)
1139 # Everything is in sync — report status and replace any stale canary
1140 # sub-status in the same region-status write.
1141 state = "running" if ready_replicas >= desired_replicas else "creating"
1142 self.store.update_region_status(
1143 name,
1144 self.region,
1145 state,
1146 replicas_ready=ready_replicas,
1147 replicas_desired=desired_replicas,
1148 extra={"canary": canary_status} if canary_status is not None else None,
1149 )
1151 # Promote desired state only from live local readiness plus explicit
1152 # running observations for every *other* target region. The endpoint
1153 # object may contain a stale local region_status from before this pass.
1154 if state == "running" and endpoint.get("desired_state") == "deploying":
1155 stored_statuses = endpoint.get("region_status", {})
1156 target_regions = endpoint.get("target_regions", [])
1157 all_running = bool(target_regions)
1158 for target_region in target_regions:
1159 if target_region == self.region: 1159 ↛ 1161line 1159 didn't jump to line 1161 because the condition on line 1159 was always true
1160 continue
1161 target_status = (
1162 stored_statuses.get(target_region, {})
1163 if isinstance(stored_statuses, dict)
1164 else {}
1165 )
1166 if not isinstance(target_status, dict) or target_status.get("state") != "running":
1167 all_running = False
1168 break
1169 if all_running: 1169 ↛ 1172line 1169 didn't jump to line 1172 because the condition on line 1169 was always true
1170 self.store.update_desired_state(name, "running")
1172 return None
1174 # ------------------------------------------------------------------
1175 # Mooncake reconciliation branch
1176 # ------------------------------------------------------------------
1178 @staticmethod
1179 def _desired_roles(mode: str | None) -> list[str]:
1180 """Return the worker roles a mode materializes, in a stable order.
1182 Disaggregated and ``both`` modes split work across ``prefill`` then
1183 ``decode``; store mode runs a single ``kv_both`` instance under the
1184 ``single`` role. The order is fixed so role creation and status
1185 reporting are deterministic across passes.
1186 """
1187 roles = _WORKER_ROLES_BY_MODE.get(mode, set()) if isinstance(mode, str) else set()
1188 return [role for role in ("prefill", "decode", "single") if role in roles]
1190 @staticmethod
1191 def _needs_shared_master(mooncake: dict[str, Any]) -> bool:
1192 """Whether the endpoint depends on the shared per-region master.
1194 The store-bearing modes (``store`` and ``both``) always reach the
1195 master for KV metadata, and any endpoint transferring over RDMA reaches
1196 the master's built-in metadata server for the connector handshake. A
1197 disaggregated endpoint transferring over TCP needs no master.
1198 """
1199 mode = mooncake.get("mode")
1200 if mode in ("store", "both"):
1201 return True
1202 transfer = mooncake.get("transfer") or {}
1203 return bool(transfer.get("protocol", "rdma") == "rdma")
1205 def _ensure_mooncake_configmap(self, name: str, ns: str, cfg: dict[str, Any]) -> None:
1206 """Create or update the shared transport ConfigMap for an endpoint.
1208 The rendered transport settings (the dict produced by
1209 :func:`render_mooncake_config`) are written to a ConfigMap named
1210 ``{name}-mooncake`` under the ``mooncake.json`` key, which each role pod
1211 mounts at the configured path. Creation is idempotent: an existing
1212 ConfigMap is patched to the desired contents so a transport change on
1213 the spec propagates on the next pass.
1214 """
1215 cm_name = f"{name}-mooncake"
1216 config_map = client.V1ConfigMap(
1217 metadata=client.V1ObjectMeta(
1218 name=cm_name,
1219 namespace=ns,
1220 labels={"app": name, "project": "gco", "gco.io/type": "inference"},
1221 ),
1222 data={"mooncake.json": json.dumps(cfg, sort_keys=True)},
1223 )
1224 try:
1225 self.core_v1.create_namespaced_config_map(
1226 ns, config_map, _request_timeout=self._k8s_timeout
1227 )
1228 logger.info("Created mooncake config map %s/%s", ns, cm_name)
1229 except ApiException as e:
1230 if e.status == 409: 1230 ↛ 1236line 1230 didn't jump to line 1236 because the condition on line 1230 was always true
1231 self.core_v1.patch_namespaced_config_map(
1232 cm_name, ns, config_map, _request_timeout=self._k8s_timeout
1233 )
1234 logger.info("Updated mooncake config map %s/%s", ns, cm_name)
1235 else:
1236 raise
1238 def _ensure_role_deployment(
1239 self, name: str, ns: str, spec: dict[str, Any], role: str
1240 ) -> tuple[int, int]:
1241 """Ensure one role Deployment exists at its desired replica count.
1243 Creates the role Deployment when absent. When it already exists and an
1244 autoscaler does not own its count, the replica count is reconciled to
1245 the topology-desired value so a topology change on the spec takes
1246 effect. The materialized name is ``{name}`` for the single store role
1247 and ``{name}-{role}`` for prefill and decode.
1249 Returns:
1250 The observed ``(ready, desired)`` replica counts after the pass.
1251 """
1252 mooncake = spec.get("mooncake") or {}
1253 deploy_name = name if role == "single" else f"{name}-{role}"
1254 desired = self._replica_count_for_role(mooncake, role)
1256 deployment = self._get_deployment(deploy_name, ns)
1257 if deployment is None:
1258 self._create_role_deployment(name, ns, spec, role)
1259 return 0, desired
1261 # An autoscaler owns the count for prefill/decode when enabled; leave
1262 # the running count untouched in that case.
1263 autoscaling = mooncake.get("autoscaling") or {}
1264 autoscaled = bool(autoscaling.get("enabled")) and role in ("prefill", "decode")
1265 current = deployment.spec.replicas or 0
1266 if not autoscaled and current != desired: 1266 ↛ 1276line 1266 didn't jump to line 1276 because the condition on line 1266 was always true
1267 logger.info(
1268 "Scaling role deployment %s/%s: %d → %d",
1269 ns,
1270 deploy_name,
1271 current,
1272 desired,
1273 )
1274 self._scale_deployment(deploy_name, ns, desired)
1276 status = getattr(deployment, "status", None)
1277 ready = int(getattr(status, "ready_replicas", 0) or 0) if status else 0
1278 return ready, desired
1280 def _report_role_status(
1281 self,
1282 name: str,
1283 ns: str,
1284 mooncake: dict[str, Any],
1285 region_services: dict[str, Any],
1286 ) -> str:
1287 """Write the role-keyed region status for a Mooncake endpoint.
1289 For split topologies the status carries a ``roles`` map of observed and
1290 desired replica counts per role; for store-bearing endpoints it carries
1291 a ``store`` sub-status with the master's readiness and address. The flat
1292 ``replicas_ready`` / ``replicas_desired`` fields are also populated with
1293 the totals so consumers that only read the flat shape still see motion.
1295 Returns:
1296 The reported endpoint state: ``"running"`` once every desired role
1297 replica (and, when applicable, the master) is Ready, otherwise
1298 ``"creating"``.
1299 """
1300 mode = mooncake.get("mode")
1301 roles = self._desired_roles(mode)
1302 extra: dict[str, Any] = {}
1303 total_ready = 0
1304 total_desired = 0
1305 all_ready = True
1307 if mode in ("disaggregated", "both"):
1308 roles_block: dict[str, Any] = {}
1309 for role in ("prefill", "decode"):
1310 if role not in roles: 1310 ↛ 1311line 1310 didn't jump to line 1311 because the condition on line 1310 was never true
1311 continue
1312 deploy_name = f"{name}-{role}"
1313 desired = self._replica_count_for_role(mooncake, role)
1314 dep = self._get_deployment(deploy_name, ns)
1315 status = getattr(dep, "status", None) if dep else None
1316 ready = int(getattr(status, "ready_replicas", 0) or 0) if status else 0
1317 roles_block[role] = {"ready": ready, "desired": desired}
1318 total_ready += ready
1319 total_desired += desired
1320 if ready < desired:
1321 all_ready = False
1322 extra["roles"] = roles_block
1323 else:
1324 # Store mode runs a single kv_both Deployment under the endpoint name.
1325 desired = self._replica_count_for_role(mooncake, "single")
1326 dep = self._get_deployment(name, ns)
1327 status = getattr(dep, "status", None) if dep else None
1328 ready = int(getattr(status, "ready_replicas", 0) or 0) if status else 0
1329 total_ready += ready
1330 total_desired += desired
1331 if ready < desired: 1331 ↛ 1332line 1331 didn't jump to line 1332 because the condition on line 1331 was never true
1332 all_ready = False
1334 store = mooncake.get("store") or {}
1335 if store.get("enabled"):
1336 master_ready = self._mooncake_master_ready_replicas(ns) >= 1
1337 extra["store"] = {
1338 "ready": master_ready,
1339 "master": region_services.get("master_server_address"),
1340 }
1341 if not master_ready:
1342 all_ready = False
1344 state = "running" if all_ready and total_desired > 0 else "creating"
1345 self.store.update_region_status(
1346 name,
1347 self.region,
1348 state,
1349 replicas_ready=total_ready,
1350 replicas_desired=total_desired,
1351 extra=extra or None,
1352 )
1353 return state
1355 async def _reconcile_mooncake(
1356 self,
1357 name: str,
1358 ns: str,
1359 spec: dict[str, Any],
1360 endpoint: dict[str, Any],
1361 ) -> dict[str, Any] | None:
1362 """Reconcile an endpoint whose spec carries a ``mooncake`` block.
1364 Returns ``None`` when the spec carries no ``mooncake`` block, signalling
1365 the caller to take the single-Deployment path: one Deployment at the
1366 configured replica count and one internal ClusterIP Service, with no role
1367 split, proxy, autoscaler, shared-master dependency, endpoint Ingress,
1368 Gateway, or HTTPRoute. The shared platform route remains unchanged.
1370 With a ``mooncake`` block present, the topology is materialized in
1371 dependency order, and the shared ConfigMap and master are laid down
1372 before any role pod, the roles before the front-end, and the front-end
1373 before status is written:
1375 1. Resolve the in-region addresses (master, metadata, optional cold
1376 tier). When the store is enabled but no own-region master is
1377 configured, nothing further is materialized; the existing
1378 configuration is left unchanged and the endpoint is reported as
1379 still coming up with the unresolved-master reason.
1380 2. Confirm every wired address stays inside the monitor's own region. A
1381 cross-region address fails the endpoint and materializes nothing,
1382 leaving any prior resources in place.
1383 3. Gate dependent pods on the shared per-region master, which also lays
1384 down the intra-namespace allow rules. While the master is not Ready,
1385 or if it could not be created, nothing further is materialized and
1386 the endpoint is reported as still coming up.
1387 4. Render and apply the shared transport ConfigMap.
1388 5. Materialize each role Deployment: prefill and decode for
1389 disaggregated and both modes, a single ``kv_both`` Deployment for
1390 store mode.
1391 6. Materialize each present role's autoscaler when autoscaling is on.
1392 7. Front disaggregated and both modes with the proxy and its internal
1393 ClusterIP Service; give store mode an internal ClusterIP Service.
1394 Public traffic remains on the shared ``gco-system/gco-gateway``
1395 HTTPRoute from ``/inference`` to ``gco-system/inference-proxy``.
1396 8. Write the role-keyed region status.
1398 Returns:
1399 An action record describing what the pass did, or ``None`` when the
1400 spec carries no ``mooncake`` block.
1401 """
1402 mooncake = spec.get("mooncake")
1403 if not mooncake:
1404 return None
1406 mode = mooncake.get("mode")
1408 # Step 1: resolve in-region addresses. A store without an own-region
1409 # master is left untouched and reported as still coming up.
1410 services = self._resolve_region_services(name, mooncake)
1411 if services.render_skipped:
1412 self.store.update_region_status(name, self.region, "creating", error=services.error)
1413 return {
1414 "action": "reconcile_mooncake",
1415 "endpoint": name,
1416 "deferred": "store_master_unresolved",
1417 }
1419 region_services = services.region_services or {}
1421 # Step 2: keep the topology inside its own region.
1422 scope = self._resolve_regional_scope(name, ns, spec, region_services)
1423 if not scope.in_region:
1424 self.store.update_region_status(
1425 name, self.region, scope.state or "failed", error=scope.error
1426 )
1427 return {
1428 "action": "reconcile_mooncake",
1429 "endpoint": name,
1430 "failed": "cross_region_boundary",
1431 }
1433 # Step 3: gate dependent pods on the shared master (and its allow
1434 # rules). This is also where the master itself is created if absent.
1435 if self._needs_shared_master(mooncake):
1436 gate = self._gate_on_mooncake_master(name, ns, spec)
1437 if not gate.proceed:
1438 self.store.update_region_status(
1439 name, self.region, gate.state or "creating", error=gate.error
1440 )
1441 return {
1442 "action": "reconcile_mooncake",
1443 "endpoint": name,
1444 "deferred": "master_not_ready",
1445 }
1447 # Step 4: shared transport ConfigMap, applied once before role pods.
1448 cfg = render_mooncake_config(mooncake, region_services)
1449 self._ensure_mooncake_configmap(name, ns, cfg)
1451 # Step 5: role Deployments, in a stable order.
1452 desired_roles = self._desired_roles(mode)
1453 for role in desired_roles:
1454 self._ensure_role_deployment(name, ns, spec, role)
1456 # Step 6: optional per-role autoscaling.
1457 if (mooncake.get("autoscaling") or {}).get("enabled"):
1458 for role in ("prefill", "decode"):
1459 if role in desired_roles: 1459 ↛ 1458line 1459 didn't jump to line 1458 because the condition on line 1459 was always true
1460 self._create_role_hpa(name, ns, spec, role)
1462 # Step 7: front-end. Disaggregated and both run behind the proxy; store
1463 # exposes its single Deployment directly.
1464 if mode in ("disaggregated", "both"):
1465 # Per-role Services so the proxy can address prefill and decode by
1466 # stable in-cluster DNS. Routing through a Service means kube-proxy
1467 # load-balances across only the Ready pods of each role, which is
1468 # what gives the proxy ready-only decode routing for free.
1469 role_port = spec.get("port", 8000)
1470 for role in desired_roles:
1471 self._create_role_service(name, ns, role, role_port)
1472 try:
1473 self._create_pd_proxy(name, ns, spec, endpoint)
1474 except AdminApiKeySecretError as e:
1475 logger.error("Proxy for endpoint %s in %s not started: %s", name, ns, e)
1476 self.store.update_region_status(name, self.region, "failed", error=str(e))
1477 return {
1478 "action": "reconcile_mooncake",
1479 "endpoint": name,
1480 "failed": "admin_api_key",
1481 }
1482 else:
1483 self._create_service(name, ns, spec)
1485 # Step 8: role-keyed status.
1486 state = self._report_role_status(name, ns, mooncake, region_services)
1487 return {"action": "reconcile_mooncake", "endpoint": name, "state": state}
1489 def _reconcile_stopped(self, name: str, namespace: str) -> dict[str, Any] | None:
1490 """Scale deployment to zero."""
1491 deployment = self._get_deployment(name, namespace)
1492 if deployment is None:
1493 return None
1495 current_replicas = deployment.spec.replicas or 0
1496 if current_replicas > 0:
1497 logger.info("Stopping endpoint %s (scaling to 0)", name)
1498 self._scale_deployment(name, namespace, 0)
1499 self.store.update_region_status(
1500 name,
1501 self.region,
1502 "stopped",
1503 replicas_ready=0,
1504 replicas_desired=0,
1505 )
1506 return {"action": "stop", "endpoint": name}
1508 self.store.update_region_status(
1509 name,
1510 self.region,
1511 "stopped",
1512 replicas_ready=0,
1513 replicas_desired=0,
1514 )
1515 return None
1517 def _reconcile_deleted(
1518 self,
1519 name: str,
1520 namespace: str,
1521 spec: dict[str, Any] | None = None,
1522 ) -> dict[str, Any] | None:
1523 """Delete all resources for the endpoint."""
1524 # Clean up health watchdog tracker
1525 self._unready_since.pop(name, None)
1526 # Clean up the master-readiness deferral tracker
1527 self._master_deferral_since.pop(name, None)
1529 # An endpoint is either a single Deployment named ``name`` or a Mooncake
1530 # role-split topology (``name-prefill``/``name-decode``/``name-proxy``).
1531 # Check all of them so a disaggregated endpoint is actually torn down
1532 # rather than skipped (which would orphan its role Deployments — and the
1533 # GPU nodes they hold).
1534 deployment_names = (
1535 name,
1536 f"{name}-prefill",
1537 f"{name}-decode",
1538 f"{name}-proxy",
1539 )
1540 if any(self._deployment_exists(d, namespace) for d in deployment_names):
1541 logger.info("Deleting endpoint %s from %s", name, self.region)
1542 self._delete_resources(name, namespace, spec)
1543 self.store.update_region_status(name, self.region, "deleted")
1544 return {"action": "delete", "endpoint": name}
1546 self.store.update_region_status(name, self.region, "deleted")
1547 return None
1549 # ------------------------------------------------------------------
1550 # Kubernetes resource management
1551 # ------------------------------------------------------------------
1553 def _deployment_exists(self, name: str, namespace: str) -> bool:
1554 try:
1555 self.apps_v1.read_namespaced_deployment(
1556 name, namespace, _request_timeout=self._k8s_timeout
1557 )
1558 return True
1559 except ApiException as e:
1560 if e.status == 404:
1561 return False
1562 raise
1564 def _get_deployment(self, name: str, namespace: str) -> V1Deployment | None:
1565 try:
1566 return self.apps_v1.read_namespaced_deployment(
1567 name, namespace, _request_timeout=self._k8s_timeout
1568 )
1569 except ApiException as e:
1570 if e.status == 404:
1571 return None
1572 raise
1574 def _get_deployment_image(self, deployment: V1Deployment) -> str | None:
1575 """Get the image of the first container in a deployment."""
1576 containers = deployment.spec.template.spec.containers
1577 if containers:
1578 image: str = containers[0].image
1579 return image
1580 return None
1582 def _resolve_image_for_region(self, spec: dict[str, Any]) -> str:
1583 """Pick the image URI this region should pull from.
1585 ``cli.inference.InferenceManager.deploy`` populates
1586 ``spec["region_image_uris"]`` with a per-region map when the
1587 primary image is an ECR URI, so each cluster can pull from its
1588 local replica instead of crossing the WAN. The map is omitted
1589 for non-ECR refs and for deploys with ``rewrite_image=False``,
1590 in which case we fall back to the flat ``spec["image"]`` URI.
1592 When the map is present but lacks an entry for ``self.region``
1593 (a target region was added after the spec was last written),
1594 the flat URI is also used so the deployment doesn't break — the
1595 next reconcile after a fresh deploy picks up the right URI.
1596 """
1597 region_map = spec.get("region_image_uris")
1598 if isinstance(region_map, dict):
1599 uri = region_map.get(self.region)
1600 if isinstance(uri, str) and uri:
1601 return uri
1602 return str(spec["image"])
1604 # ------------------------------------------------------------------
1605 # In-region service resolution (master address + cold-tier bucket)
1606 # ------------------------------------------------------------------
1608 def _resolve_region_services(
1609 self, name: str, mooncake: dict[str, Any]
1610 ) -> RegionServicesResolution:
1611 """Resolve the in-region addresses an endpoint's ``mooncake.json`` needs.
1613 Everything an endpoint wires to is resolved for the monitor's own
1614 region from regional configuration — never from values typed into the
1615 endpoint spec:
1617 - The shared master's RPC address comes from regional configuration. It
1618 is required whenever the store is enabled. When the store is enabled
1619 and no own-region master address is configured, rendering is skipped
1620 and the existing endpoint configuration is left unchanged; the result
1621 records the unresolved-master condition so the caller can report it.
1622 - The metadata server defaults to the master host on the metadata port
1623 unless regional configuration supplies one explicitly.
1624 - When the cold tier is opted in (``store.cold_tier_enabled`` is the
1625 boolean ``True``), the cold-tier object-store URI is resolved to the
1626 own-region general-purpose regional bucket from that region's
1627 ``/name`` discovery value. Any cold-tier bucket URI in the spec is
1628 ignored. Whether the endpoint writes to the cold tier is governed
1629 solely by the per-endpoint flag, independent of the always-on bucket.
1630 When the bucket cannot be resolved (the region's stack is not yet
1631 deployed), the cold tier is dropped and the condition is recorded,
1632 while the hot-path store stays configured.
1634 Args:
1635 name: The endpoint name, used to scope the cold-tier object key.
1636 mooncake: The ``spec["mooncake"]`` block.
1638 Returns:
1639 A :class:`RegionServicesResolution` carrying the resolved
1640 ``region_services`` dict and any skip/unresolved signals.
1641 """
1642 store = mooncake.get("store", {})
1643 store_enabled = bool(store.get("enabled"))
1645 master_address = os.environ.get(MOONCAKE_MASTER_ADDRESS_ENV, "").strip()
1647 # The store needs an own-region master. It is a fixed in-cluster Service
1648 # the monitor itself provisions per region (mooncake-master:50051), so
1649 # when no override is set in the environment, default to that Service
1650 # rather than deferring — the address is known by construction. An
1651 # operator may still override it via MOONCAKE_MASTER_ADDRESS.
1652 if store_enabled and not master_address:
1653 master_address = f"{MOONCAKE_MASTER_SERVICE}:{MOONCAKE_MASTER_RPC_PORT}"
1655 region_services: dict[str, Any] = {
1656 "metadata_server": self._metadata_server_url(master_address),
1657 }
1658 if store_enabled:
1659 region_services["master_server_address"] = master_address
1661 result = RegionServicesResolution(region_services=region_services)
1663 # Cold tier is opt-in per endpoint. The bucket is always resolved for
1664 # the monitor's own region; any URI supplied in the spec is ignored.
1665 if store_enabled and store.get("cold_tier_enabled") is True:
1666 bucket = self._resolve_regional_shared_bucket()
1667 if bucket:
1668 region_services["cold_tier_s3_uri"] = (
1669 f"s3://{bucket}/{MOONCAKE_COLD_TIER_KEY_PREFIX}/{name}/"
1670 )
1671 else:
1672 # The own-region general-purpose bucket is not resolvable yet.
1673 # Drop the cold tier but keep the hot-path store operating.
1674 result.cold_tier_unresolved = True
1675 result.error = (
1676 "general-purpose regional bucket for region "
1677 f"{self.region} could not be resolved; cold tier disabled, "
1678 "hot-path store still active"
1679 )
1681 return result
1683 def _metadata_server_url(self, master_address: str) -> str:
1684 """Return the metadata server URL, deriving it from the master host.
1686 Regional configuration may supply the metadata server URL directly.
1687 When it does not, the URL defaults to the master host on the metadata
1688 port; if no master host is known the conventional in-cluster service
1689 name is used.
1690 """
1691 configured = os.environ.get(MOONCAKE_METADATA_SERVER_ENV, "").strip()
1692 if configured: 1692 ↛ 1693line 1692 didn't jump to line 1693 because the condition on line 1692 was never true
1693 return configured
1694 host = master_address.rsplit(":", 1)[0] if master_address else MOONCAKE_MASTER_SERVICE
1695 return f"http://{host}:{MOONCAKE_METADATA_PORT}/metadata"
1697 def _resolve_regional_shared_bucket(self) -> str | None:
1698 """Resolve the own-region general-purpose bucket name, or ``None``.
1700 Reads the monitor's own region's ``/name`` discovery value for the
1701 always-on general-purpose regional bucket. Returns ``None`` when the
1702 value is absent (the region's stack is not yet deployed) or cannot be
1703 read, so the caller can drop the cold tier without disturbing the
1704 hot-path store.
1705 """
1706 from gco.services.aws_ssm import get_ssm_parameter_optional
1708 param_name = f"{_regional_shared_ssm_parameter_prefix()}/name"
1709 try:
1710 return get_ssm_parameter_optional(param_name, region=self.region)
1711 except Exception as e: # noqa: BLE001 - any read failure means "unresolved"
1712 logger.warning(
1713 "Failed to resolve general-purpose regional bucket for %s: %s",
1714 self.region,
1715 e,
1716 )
1717 return None
1719 # ------------------------------------------------------------------
1720 # Regional scope boundary (intra-region RDMA enforcement)
1721 # ------------------------------------------------------------------
1723 def _region_of_address(self, address: str) -> str:
1724 """Return the address's explicit AWS region or safe local classification.
1726 AWS region tokens embedded in the host are authoritative. Bare Service
1727 names and Kubernetes ``.svc`` names are local by construction. Any
1728 other host without a region token is external and ambiguous, so it is
1729 classified as ``"unknown"`` and rejected by regional-scope checks.
1730 """
1731 candidate = (address or "").strip()
1732 if not candidate: 1732 ↛ 1733line 1732 didn't jump to line 1733 because the condition on line 1732 was never true
1733 return "unknown"
1735 try:
1736 parsed = urlsplit(candidate if "://" in candidate else f"//{candidate}")
1737 host = (parsed.hostname or "").rstrip(".").lower()
1738 except ValueError:
1739 return "unknown"
1740 if not host: 1740 ↛ 1741line 1740 didn't jump to line 1741 because the condition on line 1740 was never true
1741 return "unknown"
1743 match = _REGION_TOKEN_PATTERN.search(host)
1744 if match:
1745 return match.group(0)
1746 if "." not in host or host.endswith((".svc", ".svc.cluster.local")):
1747 return self.region
1748 return "unknown"
1750 def _resolve_regional_scope(
1751 self,
1752 name: str,
1753 ns: str,
1754 spec: dict[str, Any],
1755 region_services: dict[str, Any] | None,
1756 ) -> RegionalScopeResolution:
1757 """Confirm a disaggregated topology wires only to its own region.
1759 Gathers every address the own-region topology connects to — the
1760 ``MooncakeConnector`` peers (the sibling role Services for prefill and
1761 decode), the shared master's RPC address, and the metadata server — and
1762 confirms each resolves to the monitor's own region. Any explicitly
1763 supplied peer addresses in the spec are checked too, so a misconfigured
1764 endpoint that points a peer or master at another region is caught
1765 before any role pod is materialized.
1767 An endpoint that enumerates two or more target regions runs one
1768 independent topology per region: each region's monitor reconciles only
1769 its own topology (it reconciles only while ``self.region`` is one of the
1770 target regions), and this resolution confirms that topology's addresses
1771 never cross into another region.
1773 Args:
1774 name: The endpoint name; used to derive the in-cluster peer Service
1775 names for the disaggregated roles.
1776 ns: The namespace the topology materializes into.
1777 spec: The endpoint spec being reconciled.
1778 region_services: The resolved in-region addresses (from
1779 :meth:`_resolve_region_services`), or ``None`` when none were
1780 resolved. Supplies the master and metadata addresses to check.
1782 Returns:
1783 A :class:`RegionalScopeResolution`. When every resolved address is
1784 own-region, ``in_region`` is ``True`` and the caller may materialize
1785 the role Deployments. When any address resolves to another region,
1786 ``in_region`` is ``False``, ``state`` is ``"failed"``, and ``error``
1787 names the offending addresses; the caller then materializes no role
1788 Deployments and leaves any prior resources unchanged.
1789 """
1790 mooncake = spec.get("mooncake") or {}
1791 mode = mooncake.get("mode")
1792 roles = _WORKER_ROLES_BY_MODE.get(mode, set()) if isinstance(mode, str) else set()
1794 # MooncakeConnector peers are the sibling role Services within this
1795 # namespace; collect them in a stable order for deterministic reporting.
1796 addresses: list[str] = []
1797 if "prefill" in roles or "decode" in roles:
1798 for role in ("prefill", "decode"):
1799 addresses.append(f"{name}-{role}.{ns}.svc.cluster.local")
1801 # The shared master and metadata server the pods reach. These come from
1802 # the own-region resolution, but are checked here so a foreign address
1803 # supplied through regional configuration is still caught.
1804 if region_services: 1804 ↛ 1814line 1804 didn't jump to line 1814 because the condition on line 1804 was always true
1805 master = region_services.get("master_server_address")
1806 metadata = region_services.get("metadata_server")
1807 if master:
1808 addresses.append(str(master))
1809 if metadata:
1810 addresses.append(str(metadata))
1812 # Defensive: honor any explicit peer/master addresses authored on the
1813 # spec so a hand-edited endpoint cannot smuggle in an out-of-region peer.
1814 store = mooncake.get("store") or {}
1815 transfer = mooncake.get("transfer") or {}
1816 for candidate in (
1817 store.get("master_server_address"),
1818 store.get("metadata_server"),
1819 ):
1820 if candidate: 1820 ↛ 1821line 1820 didn't jump to line 1821 because the condition on line 1820 was never true
1821 addresses.append(str(candidate))
1822 explicit_peers = transfer.get("peer_addresses")
1823 if isinstance(explicit_peers, list):
1824 addresses.extend(str(peer) for peer in explicit_peers if peer)
1826 # De-duplicate while preserving first-seen order.
1827 seen: set[str] = set()
1828 ordered: list[str] = []
1829 for address in addresses:
1830 if address not in seen:
1831 seen.add(address)
1832 ordered.append(address)
1834 resolved_regions = [(address, self._region_of_address(address)) for address in ordered]
1835 out_of_region = [
1836 (address, region) for address, region in resolved_regions if region != self.region
1837 ]
1839 if out_of_region:
1840 detail = ", ".join(
1841 f"{address!r} resolves to region {region}" for address, region in out_of_region
1842 )
1843 logger.error(
1844 "Cross-region boundary violation for endpoint %s in %s: %s",
1845 name,
1846 self.region,
1847 detail,
1848 )
1849 return RegionalScopeResolution(
1850 in_region=False,
1851 peer_addresses=ordered,
1852 state="failed",
1853 error=(f"cross-region boundary violation: {detail}; expected region {self.region}"),
1854 )
1856 return RegionalScopeResolution(in_region=True, peer_addresses=ordered)
1858 def _ensure_mooncake_store(self, ns: str, spec: dict[str, Any]) -> None:
1859 """Maintain the single shared per-region Mooncake master, idempotently.
1861 The master is region-shared, not per-endpoint: every endpoint that
1862 needs the store reaches the same ``mooncake-master`` StatefulSet and the
1863 headless Service that fronts its RPC and metadata ports. This method
1864 uses create-if-absent semantics, so any number of calls within a region
1865 converge on exactly one StatefulSet with a single replica and one
1866 Service. An already-existing master is left untouched — a conflicting
1867 create is treated as success and never overwrites the running master.
1869 The StatefulSet runs the master daemon with its built-in HTTP metadata
1870 server, exposing RPC on :data:`MOONCAKE_MASTER_RPC_PORT` and the
1871 metadata endpoint on :data:`MOONCAKE_METADATA_PORT`. Both ports are
1872 published on the headless Service so in-namespace pods can resolve them.
1874 Args:
1875 ns: The namespace the master shares with the inference workloads.
1876 spec: The endpoint spec being reconciled. Its ``mooncake`` block may
1877 carry a master image override; otherwise the image is taken from
1878 the in-region deployment's environment.
1879 """
1880 mooncake = spec.get("mooncake", {}) or {}
1881 store = mooncake.get("store", {}) or {}
1882 image = store.get("master_image") or os.environ.get(MOONCAKE_MASTER_IMAGE_ENV, "").strip()
1884 labels = {"app": MOONCAKE_MASTER_SERVICE, "project": "gco"}
1886 # Headless Service exposing both the RPC and metadata ports. A None
1887 # cluster IP keeps it headless so the StatefulSet's stable network
1888 # identity resolves directly.
1889 service = client.V1Service(
1890 metadata=client.V1ObjectMeta(
1891 name=MOONCAKE_MASTER_SERVICE,
1892 namespace=ns,
1893 labels=labels,
1894 ),
1895 spec=client.V1ServiceSpec(
1896 cluster_ip="None",
1897 selector={"app": MOONCAKE_MASTER_SERVICE},
1898 ports=[
1899 client.V1ServicePort(
1900 name="rpc",
1901 port=MOONCAKE_MASTER_RPC_PORT,
1902 target_port="rpc",
1903 protocol="TCP",
1904 ),
1905 client.V1ServicePort(
1906 name="metadata",
1907 port=MOONCAKE_METADATA_PORT,
1908 target_port="metadata",
1909 protocol="TCP",
1910 ),
1911 ],
1912 ),
1913 )
1915 container = client.V1Container(
1916 name=MOONCAKE_MASTER_SERVICE,
1917 image=image,
1918 command=["mooncake_master"],
1919 args=[
1920 f"--port={MOONCAKE_MASTER_RPC_PORT}",
1921 "--enable_http_metadata_server=true",
1922 f"--http_metadata_server_port={MOONCAKE_METADATA_PORT}",
1923 ],
1924 ports=[
1925 client.V1ContainerPort(
1926 name="rpc",
1927 container_port=MOONCAKE_MASTER_RPC_PORT,
1928 protocol="TCP",
1929 ),
1930 client.V1ContainerPort(
1931 name="metadata",
1932 container_port=MOONCAKE_METADATA_PORT,
1933 protocol="TCP",
1934 ),
1935 ],
1936 security_context=client.V1SecurityContext(
1937 allow_privilege_escalation=False,
1938 # The upstream mooncake_master launcher chmods its bundled
1939 # binary on startup, which needs a writable root filesystem (a
1940 # read-only root raised OSError: Read-only file system). The pod
1941 # also runs as root so the chmod of the root-owned binary is
1942 # permitted. Privilege escalation stays disabled and all
1943 # capabilities are dropped, so this is constrained root.
1944 read_only_root_filesystem=False,
1945 capabilities=client.V1Capabilities(drop=["ALL"]),
1946 ),
1947 resources=client.V1ResourceRequirements(
1948 requests={"cpu": "250m", "memory": "512Mi"},
1949 limits={"cpu": "1", "memory": "2Gi"},
1950 ),
1951 startup_probe=client.V1Probe(
1952 tcp_socket=client.V1TCPSocketAction(port="rpc"),
1953 initial_delay_seconds=5,
1954 period_seconds=5,
1955 failure_threshold=30,
1956 ),
1957 liveness_probe=client.V1Probe(
1958 tcp_socket=client.V1TCPSocketAction(port="rpc"),
1959 initial_delay_seconds=15,
1960 period_seconds=30,
1961 ),
1962 readiness_probe=client.V1Probe(
1963 # The HTTP metadata server returns 400 for a bare GET /metadata
1964 # (it expects a ?key=), so an HTTP GET readiness probe never
1965 # passes. Confirm readiness by checking the metadata port is
1966 # accepting connections instead.
1967 tcp_socket=client.V1TCPSocketAction(port="metadata"),
1968 initial_delay_seconds=10,
1969 period_seconds=15,
1970 ),
1971 )
1973 stateful_set = client.V1StatefulSet(
1974 metadata=client.V1ObjectMeta(
1975 name=MOONCAKE_MASTER_SERVICE,
1976 namespace=ns,
1977 labels=labels,
1978 ),
1979 spec=client.V1StatefulSetSpec(
1980 service_name=MOONCAKE_MASTER_SERVICE,
1981 replicas=1,
1982 selector=client.V1LabelSelector(match_labels={"app": MOONCAKE_MASTER_SERVICE}),
1983 template=client.V1PodTemplateSpec(
1984 metadata=client.V1ObjectMeta(
1985 labels=labels,
1986 # Keep Karpenter from consolidating the node out from
1987 # under the master: it is a single-replica, stateful
1988 # control-plane daemon holding KV metadata for every
1989 # in-region endpoint, so an eviction drops that state and
1990 # disrupts inference. On lightly-loaded clusters
1991 # consolidation otherwise evicts it mid-image-pull before
1992 # it can even start.
1993 annotations={"karpenter.sh/do-not-disrupt": "true"},
1994 ),
1995 spec=client.V1PodSpec(
1996 service_account_name="gco-service-account",
1997 # The upstream mooncake_master launcher chmods its
1998 # bundled binary (root-owned in the image) on startup, so
1999 # the master must run as root: a non-root uid cannot
2000 # chmod a root-owned file ("Operation not permitted").
2001 # gco-inference does not enforce restricted Pod Security
2002 # and the no-root rule applies only to user-submitted
2003 # jobs, so a root platform daemon is consistent here. The
2004 # container still drops all Linux capabilities and
2005 # disallows privilege escalation (see its securityContext).
2006 security_context=client.V1PodSecurityContext(
2007 run_as_user=0,
2008 run_as_group=0,
2009 ),
2010 containers=[container],
2011 restart_policy="Always",
2012 ),
2013 ),
2014 ),
2015 )
2017 # Create-if-absent: a 409 means the shared master already exists, which
2018 # is the steady state. Leave it untouched and treat it as success.
2019 try:
2020 self.core_v1.create_namespaced_service(ns, service, _request_timeout=self._k8s_timeout)
2021 logger.info("Created shared mooncake master service in %s", ns)
2022 except ApiException as e:
2023 if e.status == 409:
2024 logger.info("Shared mooncake master service already exists in %s", ns)
2025 else:
2026 raise
2028 try:
2029 self.apps_v1.create_namespaced_stateful_set(
2030 ns, stateful_set, _request_timeout=self._k8s_timeout
2031 )
2032 logger.info("Created shared mooncake master statefulset in %s", ns)
2033 except ApiException as e:
2034 if e.status == 409: 2034 ↛ 2037line 2034 didn't jump to line 2037 because the condition on line 2034 was always true
2035 logger.info("Shared mooncake master statefulset already exists in %s", ns)
2036 else:
2037 raise
2039 def _mooncake_master_ready_replicas(self, ns: str) -> int:
2040 """Return the shared master's Ready replica count, 0 when absent.
2042 Reads the ``mooncake-master`` StatefulSet status in ``ns``. A missing
2043 StatefulSet (404) reports zero Ready replicas rather than raising, so a
2044 caller gating on readiness simply keeps deferring until it appears.
2046 Args:
2047 ns: The namespace the shared master lives in.
2049 Returns:
2050 The number of Ready replicas the StatefulSet reports, or 0 when it
2051 does not yet exist or reports no Ready replicas.
2052 """
2053 try:
2054 status = self.apps_v1.read_namespaced_stateful_set_status(
2055 MOONCAKE_MASTER_SERVICE, ns, _request_timeout=self._k8s_timeout
2056 )
2057 except ApiException as e:
2058 if e.status == 404:
2059 return 0
2060 raise
2062 ready = getattr(getattr(status, "status", None), "ready_replicas", 0)
2063 return int(ready or 0)
2065 def _gate_on_mooncake_master(
2066 self, name: str, ns: str, spec: dict[str, Any]
2067 ) -> MasterReadinessGate:
2068 """Gate dependent role-pod creation on the shared master's readiness.
2070 Maintains the single shared master (create-if-absent) and then decides
2071 whether the endpoint's dependent role pods may be materialized:
2073 - If maintaining the master fails, no dependent pods are materialized,
2074 any existing master is left unmodified, and the endpoint stays in
2075 ``creating`` carrying a create-failure error.
2076 - While the master reports fewer than 1 Ready replica, creation is
2077 deferred and the endpoint stays in ``creating``. The first deferral
2078 starts a clock; once it exceeds the wait window the endpoint keeps
2079 deferring and stays in ``creating`` but also surfaces a not-ready
2080 error. The master is never deleted or modified on account of the
2081 timeout.
2082 - Once the master reports at least 1 Ready replica, the gate opens: the
2083 clock is cleared and the caller may create the role pods and advance
2084 out of ``creating``.
2086 Args:
2087 name: The endpoint name, used to track its first deferral.
2088 ns: The namespace the master shares with the workloads.
2089 spec: The endpoint spec being reconciled.
2091 Returns:
2092 A :class:`MasterReadinessGate` describing whether to proceed, the
2093 endpoint state to report, and any error to surface.
2094 """
2095 # Maintain the shared master first. A create failure must not produce
2096 # any dependent pods and must leave an existing master untouched.
2097 try:
2098 self._ensure_mooncake_store(ns, spec)
2099 except ApiException as e:
2100 logger.error(
2101 "Could not create shared mooncake master in %s for endpoint %s: %s",
2102 ns,
2103 name,
2104 e,
2105 )
2106 return MasterReadinessGate(
2107 proceed=False,
2108 state="creating",
2109 error="shared master could not be created",
2110 )
2112 # Apply the intra-namespace allow rules before any role pod is created.
2113 # A failure here must fail pod materialization while leaving the
2114 # default-deny posture intact; surface which rule could not be applied.
2115 try:
2116 self._ensure_intra_namespace_network_policies(ns, spec)
2117 except NetworkPolicyApplyError as e:
2118 logger.error(
2119 "Could not apply network policy %s in %s for endpoint %s: %s",
2120 e.rule,
2121 ns,
2122 name,
2123 e.reason,
2124 )
2125 return MasterReadinessGate(
2126 proceed=False,
2127 state="creating",
2128 error=f"network policy {e.rule} could not be applied",
2129 )
2131 ready_replicas = self._mooncake_master_ready_replicas(ns)
2132 if ready_replicas >= 1:
2133 # Master is Ready: open the gate and reset the deferral clock so a
2134 # later master restart restarts the wait window cleanly.
2135 if name in self._master_deferral_since:
2136 logger.info(
2137 "Shared master Ready in %s, resuming creation for endpoint %s",
2138 ns,
2139 name,
2140 )
2141 del self._master_deferral_since[name]
2142 return MasterReadinessGate(proceed=True, state=None, error=None)
2144 # Master not Ready: defer creation and report creating. Start the clock
2145 # on the first deferral.
2146 now = datetime.now(UTC)
2147 first_deferral = self._master_deferral_since.setdefault(name, now)
2148 deferred_for = (now - first_deferral).total_seconds()
2150 if deferred_for >= MOONCAKE_MASTER_READY_TIMEOUT_SECONDS:
2151 logger.error(
2152 "Shared master not Ready in %s after %.0fs, still deferring endpoint %s",
2153 ns,
2154 deferred_for,
2155 name,
2156 )
2157 return MasterReadinessGate(
2158 proceed=False,
2159 state="creating",
2160 error="shared master did not become Ready",
2161 )
2163 logger.info(
2164 "Deferring creation for endpoint %s in %s until shared master is Ready",
2165 name,
2166 ns,
2167 )
2168 return MasterReadinessGate(proceed=False, state="creating", error=None)
2170 def _ensure_intra_namespace_network_policies(self, ns: str, spec: dict[str, Any]) -> None:
2171 """Apply the intra-namespace allow rules disaggregated inference needs.
2173 Alongside the default-deny posture in ``gco-inference`` (defined in
2174 ``03-network-policies.yaml`` and never touched here), this maintains
2175 four widening allow rules with create-if-absent semantics:
2177 - ``allow-inference-internal`` — managed inference pods exchange TCP
2178 traffic and may reach the shared master's two fixed ports. This
2179 permits proxy-to-role serving and bootstrap traffic while excluding
2180 unselected sources such as the ALB.
2181 - ``allow-pod-to-master`` — inference pods reach the shared master RPC
2182 port (:data:`MOONCAKE_MASTER_RPC_PORT`).
2183 - ``allow-pod-to-metadata`` — inference pods reach the shared metadata
2184 server (:data:`MOONCAKE_METADATA_PORT`).
2185 - ``allow-rdma-bootstrap`` — inference pods reach each other on the
2186 contiguous KV-transfer bootstrap port window starting at the spec's
2187 ``mooncake.transfer.bootstrap_base_port`` (default
2188 :data:`MOONCAKE_BOOTSTRAP_BASE_PORT`).
2190 Each rule is created independently; an already-present rule (409) is the
2191 steady state and counts as success. No deny rule is ever read, modified,
2192 or deleted, so the default-deny policy is preserved regardless of
2193 outcome.
2195 Args:
2196 ns: The inference namespace the rules apply to.
2197 spec: The endpoint spec being reconciled; its
2198 ``mooncake.transfer.bootstrap_base_port`` sizes the bootstrap
2199 port window.
2201 Raises:
2202 NetworkPolicyApplyError: If any single rule cannot be created. The
2203 error names the failing rule; rules created before the failure
2204 remain in place and the default-deny policy is untouched.
2205 """
2206 transfer = (spec.get("mooncake", {}) or {}).get("transfer", {}) or {}
2207 base_port = transfer.get("bootstrap_base_port", MOONCAKE_BOOTSTRAP_BASE_PORT)
2208 try:
2209 base_port = int(base_port)
2210 except TypeError, ValueError:
2211 base_port = MOONCAKE_BOOTSTRAP_BASE_PORT
2212 end_port = min(base_port + MOONCAKE_BOOTSTRAP_PORT_SPAN, MAX_BOOTSTRAP_PORT)
2214 labels = {"project": "gco"}
2215 master_selector = client.V1LabelSelector(match_labels={"app": MOONCAKE_MASTER_SERVICE})
2216 inference_selector = client.V1LabelSelector(match_labels=INFERENCE_POD_SELECTOR)
2217 inference_peer = [client.V1NetworkPolicyPeer(pod_selector=inference_selector)]
2218 master_peer = [client.V1NetworkPolicyPeer(pod_selector=master_selector)]
2219 all_tcp = [client.V1NetworkPolicyPort(protocol="TCP")]
2221 policies = [
2222 (
2223 NETWORK_POLICY_INFERENCE_INTERNAL,
2224 client.V1NetworkPolicy(
2225 metadata=client.V1ObjectMeta(
2226 name=NETWORK_POLICY_INFERENCE_INTERNAL, namespace=ns, labels=labels
2227 ),
2228 spec=client.V1NetworkPolicySpec(
2229 pod_selector=inference_selector,
2230 policy_types=["Ingress", "Egress"],
2231 ingress=[
2232 client.V1NetworkPolicyIngressRule(
2233 _from=inference_peer,
2234 ports=all_tcp,
2235 )
2236 ],
2237 egress=[
2238 client.V1NetworkPolicyEgressRule(
2239 to=inference_peer,
2240 ports=all_tcp,
2241 ),
2242 client.V1NetworkPolicyEgressRule(
2243 to=master_peer,
2244 ports=[
2245 client.V1NetworkPolicyPort(
2246 protocol="TCP", port=MOONCAKE_MASTER_RPC_PORT
2247 ),
2248 client.V1NetworkPolicyPort(
2249 protocol="TCP", port=MOONCAKE_METADATA_PORT
2250 ),
2251 ],
2252 ),
2253 ],
2254 ),
2255 ),
2256 ),
2257 (
2258 NETWORK_POLICY_POD_TO_MASTER,
2259 client.V1NetworkPolicy(
2260 metadata=client.V1ObjectMeta(
2261 name=NETWORK_POLICY_POD_TO_MASTER, namespace=ns, labels=labels
2262 ),
2263 spec=client.V1NetworkPolicySpec(
2264 pod_selector=master_selector,
2265 policy_types=["Ingress"],
2266 ingress=[
2267 client.V1NetworkPolicyIngressRule(
2268 _from=inference_peer,
2269 ports=[
2270 client.V1NetworkPolicyPort(
2271 protocol="TCP", port=MOONCAKE_MASTER_RPC_PORT
2272 )
2273 ],
2274 )
2275 ],
2276 ),
2277 ),
2278 ),
2279 (
2280 NETWORK_POLICY_POD_TO_METADATA,
2281 client.V1NetworkPolicy(
2282 metadata=client.V1ObjectMeta(
2283 name=NETWORK_POLICY_POD_TO_METADATA, namespace=ns, labels=labels
2284 ),
2285 spec=client.V1NetworkPolicySpec(
2286 pod_selector=master_selector,
2287 policy_types=["Ingress"],
2288 ingress=[
2289 client.V1NetworkPolicyIngressRule(
2290 _from=inference_peer,
2291 ports=[
2292 client.V1NetworkPolicyPort(
2293 protocol="TCP", port=MOONCAKE_METADATA_PORT
2294 )
2295 ],
2296 )
2297 ],
2298 ),
2299 ),
2300 ),
2301 (
2302 NETWORK_POLICY_RDMA_BOOTSTRAP,
2303 client.V1NetworkPolicy(
2304 metadata=client.V1ObjectMeta(
2305 name=NETWORK_POLICY_RDMA_BOOTSTRAP, namespace=ns, labels=labels
2306 ),
2307 spec=client.V1NetworkPolicySpec(
2308 pod_selector=inference_selector,
2309 policy_types=["Ingress"],
2310 ingress=[
2311 client.V1NetworkPolicyIngressRule(
2312 _from=inference_peer,
2313 ports=[
2314 client.V1NetworkPolicyPort(
2315 protocol="TCP",
2316 port=base_port,
2317 end_port=end_port,
2318 )
2319 ],
2320 )
2321 ],
2322 ),
2323 ),
2324 ),
2325 ]
2327 for rule_name, policy in policies:
2328 try:
2329 self.networking_v1.create_namespaced_network_policy(
2330 ns, policy, _request_timeout=self._k8s_timeout
2331 )
2332 logger.info("Applied network policy %s in %s", rule_name, ns)
2333 except ApiException as e:
2334 if e.status == 409:
2335 # Already present — the steady state. Leave it untouched.
2336 logger.info("Network policy %s already present in %s", rule_name, ns)
2337 continue
2338 logger.error(
2339 "Could not apply network policy %s in %s: %s",
2340 rule_name,
2341 ns,
2342 e,
2343 )
2344 raise NetworkPolicyApplyError(rule_name, e.reason or str(e)) from e
2346 def _create_deployment(self, name: str, namespace: str, spec: dict[str, Any]) -> None:
2347 """Create a Kubernetes Deployment for an inference endpoint."""
2348 replicas = spec.get("replicas", 1)
2349 deployment = self._build_inference_deployment_object(
2350 name=name,
2351 deploy_name=name,
2352 app_label=name,
2353 namespace=namespace,
2354 spec=spec,
2355 replicas=replicas,
2356 )
2357 self.apps_v1.create_namespaced_deployment(
2358 namespace, deployment, _request_timeout=self._k8s_timeout
2359 )
2360 logger.info("Created deployment %s/%s", namespace, name)
2362 def _build_inference_deployment_object(
2363 self,
2364 name: str,
2365 deploy_name: str,
2366 app_label: str,
2367 namespace: str,
2368 spec: dict[str, Any],
2369 replicas: int,
2370 extra_args: list[str] | None = None,
2371 extra_labels: dict[str, str] | None = None,
2372 ) -> client.V1Deployment:
2373 """Build the ``V1Deployment`` object for an inference workload.
2375 Shared by the single-Deployment path and the role-split prefill/decode/
2376 store paths. ``name`` is the endpoint name (used for the in-cluster
2377 serving prefix and model cache directory), ``deploy_name`` is the
2378 Kubernetes object name, and ``app_label`` is the selector label that
2379 Services and autoscalers target. ``extra_args`` are appended to the
2380 container args (for example the rendered ``--kv-transfer-config``), and
2381 ``extra_labels`` are merged into both the Deployment and pod-template
2382 labels so role pods carry a stable role marker.
2383 """
2384 image = self._resolve_image_for_region(spec)
2385 port = spec.get("port", 8000)
2386 gpu_count = spec.get("gpu_count", 1)
2387 health_path = spec.get("health_check_path", "/health")
2388 env_vars = spec.get("env", {})
2389 # Stable block hashing across data-parallel ranks: identical prompts
2390 # must hash identically so shared prefix-cache hits are not lost
2391 # between pods. The disaggregated serving image (upstream vLLM) no
2392 # longer bakes this in, so default it here; an explicit spec env wins.
2393 env_vars = {"PYTHONHASHSEED": "0", **env_vars}
2394 resources = spec.get("resources", {})
2395 model_path = spec.get("model_path")
2396 command = spec.get("command")
2397 args = spec.get("args")
2399 # Build container
2400 container_env = [client.V1EnvVar(name=k, value=str(v)) for k, v in env_vars.items()]
2402 # Inject --root-path for servers that support it (vLLM, TGI).
2403 # This tells the server to mount its API at /inference/{name} behind the
2404 # shared platform route. We append to existing args (from --extra-args)
2405 # rather than replacing them.
2406 serving_prefix = f"/inference/{name}"
2407 root_path_images = ("vllm", "text-generation-inference", "tgi")
2408 image_lower = image.lower()
2409 if not command and any(tag in image_lower for tag in root_path_images):
2410 if args:
2411 # Append --root-path to user-provided args if not already present
2412 if "--root-path" not in args: 2412 ↛ 2413line 2412 didn't jump to line 2413 because the condition on line 2412 was never true
2413 args = list(args) + ["--root-path", serving_prefix]
2414 else:
2415 args = ["--root-path", serving_prefix]
2417 # Append caller-supplied arguments (for example the rendered
2418 # --kv-transfer-config) after any root-path injection so they survive
2419 # alongside user --extra-args.
2420 if extra_args:
2421 args = (list(args) if args else []) + list(extra_args)
2423 resource_reqs = client.V1ResourceRequirements(
2424 requests=resources.get("requests", {"cpu": "1", "memory": "4Gi"}),
2425 limits=resources.get("limits", {"cpu": "4", "memory": "16Gi"}),
2426 )
2427 # Add accelerator resources (GPU or Neuron)
2428 accelerator = spec.get("accelerator", "nvidia")
2429 if gpu_count > 0:
2430 if accelerator == "neuron":
2431 # AWS Trainium/Inferentia — request Neuron devices
2432 if resource_reqs.limits is None: 2432 ↛ 2433line 2432 didn't jump to line 2433 because the condition on line 2432 was never true
2433 resource_reqs.limits = {}
2434 resource_reqs.limits["aws.amazon.com/neuron"] = str(gpu_count)
2435 if resource_reqs.requests is None: 2435 ↛ 2436line 2435 didn't jump to line 2436 because the condition on line 2435 was never true
2436 resource_reqs.requests = {}
2437 resource_reqs.requests["aws.amazon.com/neuron"] = str(gpu_count)
2438 else:
2439 # NVIDIA GPU (default)
2440 if resource_reqs.limits is None: 2440 ↛ 2441line 2440 didn't jump to line 2441 because the condition on line 2440 was never true
2441 resource_reqs.limits = {}
2442 resource_reqs.limits["nvidia.com/gpu"] = str(gpu_count)
2443 if resource_reqs.requests is None: 2443 ↛ 2444line 2443 didn't jump to line 2444 because the condition on line 2443 was never true
2444 resource_reqs.requests = {}
2445 resource_reqs.requests["nvidia.com/gpu"] = str(gpu_count)
2447 volume_mounts = []
2448 volumes = []
2449 init_containers = []
2450 model_source = spec.get("model_source")
2452 if model_path or model_source:
2453 volume_mounts.append(
2454 client.V1VolumeMount(
2455 name="model-storage",
2456 mount_path="/models",
2457 )
2458 )
2459 volumes.append(
2460 client.V1Volume(
2461 name="model-storage",
2462 persistent_volume_claim=client.V1PersistentVolumeClaimVolumeSource(
2463 claim_name="efs-claim",
2464 ),
2465 )
2466 )
2468 # Mooncake role pods read the shared transport config (metadata-server
2469 # address, protocol, device) from the per-endpoint ``{name}-mooncake``
2470 # ConfigMap mounted read-only at MOONCAKE_CONFIG_MOUNT_DIR, pointed at by
2471 # MOONCAKE_CONFIG_PATH, and learn the KV-transfer bootstrap base port via
2472 # VLLM_MOONCAKE_BOOTSTRAP_PORT. Plain (non-mooncake) endpoints are
2473 # untouched: no volume, mount, or env is added.
2474 mooncake_block = spec.get("mooncake")
2475 if mooncake_block:
2476 volume_mounts.append(
2477 client.V1VolumeMount(
2478 name="mooncake-config",
2479 mount_path=MOONCAKE_CONFIG_MOUNT_DIR,
2480 read_only=True,
2481 )
2482 )
2483 volumes.append(
2484 client.V1Volume(
2485 name="mooncake-config",
2486 config_map=client.V1ConfigMapVolumeSource(name=f"{name}-mooncake"),
2487 )
2488 )
2489 transfer_block = mooncake_block.get("transfer") or {}
2490 base_port = transfer_block.get("bootstrap_base_port", MOONCAKE_BOOTSTRAP_BASE_PORT)
2491 try:
2492 base_port = int(base_port)
2493 except TypeError, ValueError:
2494 base_port = MOONCAKE_BOOTSTRAP_BASE_PORT
2495 container_env.append(
2496 client.V1EnvVar(name=MOONCAKE_CONFIG_PATH_ENV, value=MOONCAKE_CONFIG_FILE_PATH)
2497 )
2498 container_env.append(
2499 client.V1EnvVar(name=VLLM_MOONCAKE_BOOTSTRAP_PORT_ENV, value=str(base_port))
2500 )
2502 # Add init container to sync model from S3 if model_source is set
2503 if model_source and model_source.startswith("s3://"):
2504 model_dest = f"/models/{name}"
2505 init_containers.append(
2506 client.V1Container(
2507 name="model-sync",
2508 image="amazon/aws-cli:latest",
2509 command=["sh", "-c"],
2510 args=[
2511 f"if [ -d '{model_dest}' ] && [ \"$(ls -A '{model_dest}')\" ]; then "
2512 f"echo 'Model already cached at {model_dest}, skipping sync'; "
2513 f"else echo 'Syncing model from {model_source}...'; "
2514 f"aws s3 sync {model_source} {model_dest} --quiet; "
2515 f"echo 'Model sync complete'; fi"
2516 ],
2517 volume_mounts=[
2518 client.V1VolumeMount(
2519 name="model-storage",
2520 mount_path="/models",
2521 )
2522 ],
2523 resources=client.V1ResourceRequirements(
2524 requests={"cpu": "1", "memory": "2Gi"},
2525 limits={"cpu": "4", "memory": "8Gi"},
2526 ),
2527 )
2528 )
2530 # Probe path depends on whether the server handles the prefix
2531 uses_root_path = args is not None and "--root-path" in args
2532 probe_health = f"{serving_prefix}{health_path}" if uses_root_path else health_path
2534 container = client.V1Container(
2535 name="inference",
2536 image=image,
2537 ports=[client.V1ContainerPort(container_port=port)],
2538 env=container_env if container_env else None,
2539 resources=resource_reqs,
2540 volume_mounts=volume_mounts if volume_mounts else None,
2541 command=command,
2542 args=args,
2543 liveness_probe=client.V1Probe(
2544 http_get=client.V1HTTPGetAction(path=probe_health, port=port),
2545 initial_delay_seconds=120,
2546 period_seconds=15,
2547 failure_threshold=5,
2548 ),
2549 readiness_probe=client.V1Probe(
2550 http_get=client.V1HTTPGetAction(path=probe_health, port=port),
2551 initial_delay_seconds=30,
2552 period_seconds=10,
2553 ),
2554 )
2556 # Build tolerations based on accelerator type
2557 if accelerator == "neuron":
2558 tolerations = [
2559 client.V1Toleration(
2560 key="aws.amazon.com/neuron",
2561 operator="Equal",
2562 value="true",
2563 effect="NoSchedule",
2564 )
2565 ]
2566 else:
2567 tolerations = [
2568 client.V1Toleration(
2569 key="nvidia.com/gpu",
2570 operator="Equal",
2571 value="true",
2572 effect="NoSchedule",
2573 )
2574 ]
2576 # Node selector based on accelerator type
2577 node_selector = spec.get("node_selector", {})
2578 if gpu_count > 0 and not node_selector:
2579 if accelerator == "neuron":
2580 node_selector = {"accelerator": "neuron"}
2581 else:
2582 node_selector = {"eks.amazonaws.com/instance-gpu-manufacturer": "nvidia"}
2584 # Apply capacity type preference (spot/on-demand)
2585 capacity_type = spec.get("capacity_type")
2586 if capacity_type in ("spot", "on-demand"):
2587 node_selector["karpenter.sh/capacity-type"] = capacity_type
2589 labels = {
2590 "app": app_label,
2591 "project": "gco",
2592 "gco.io/type": "inference",
2593 }
2594 if extra_labels:
2595 labels.update(extra_labels)
2597 deployment = client.V1Deployment(
2598 metadata=client.V1ObjectMeta(
2599 name=deploy_name,
2600 namespace=namespace,
2601 labels=dict(labels),
2602 ),
2603 spec=client.V1DeploymentSpec(
2604 replicas=replicas,
2605 selector=client.V1LabelSelector(
2606 match_labels={"app": app_label},
2607 ),
2608 template=client.V1PodTemplateSpec(
2609 metadata=client.V1ObjectMeta(
2610 labels=dict(labels),
2611 ),
2612 spec=client.V1PodSpec(
2613 service_account_name="gco-service-account",
2614 containers=[container],
2615 init_containers=init_containers if init_containers else None,
2616 tolerations=tolerations,
2617 node_selector=node_selector if node_selector else None,
2618 volumes=volumes if volumes else None,
2619 ),
2620 ),
2621 ),
2622 )
2624 return deployment
2626 def _replica_count_for_role(self, mooncake: dict[str, Any], role: str) -> int:
2627 """Resolve the materialized replica count for a role.
2629 When per-role autoscaling is enabled and supplies a ``min_replicas``
2630 for the role, the role Deployment is materialized at that lower bound so
2631 the autoscaler owns the count from there. Otherwise the count comes from
2632 the topology: ``topology.prefill`` for prefill and ``topology.decode``
2633 for decode. The single store instance is always one replica.
2634 """
2635 autoscaling = mooncake.get("autoscaling") or {}
2636 if autoscaling.get("enabled") and role in ("prefill", "decode"):
2637 role_cfg = autoscaling.get(role) or {}
2638 min_replicas = role_cfg.get("min_replicas")
2639 if isinstance(min_replicas, int) and not isinstance(min_replicas, bool):
2640 return min_replicas
2642 topology = mooncake.get("topology") or {}
2643 if role == "prefill":
2644 return int(topology.get("prefill", 1))
2645 if role == "decode":
2646 return int(topology.get("decode", 1))
2647 # Single store instance: kv_both runs as one replica.
2648 return 1
2650 def _create_role_deployment(self, name: str, ns: str, spec: dict[str, Any], role: str) -> None:
2651 """Materialize one role Deployment for a Mooncake endpoint.
2653 Disaggregated and ``both`` modes split work across ``{name}-prefill``
2654 and ``{name}-decode``; store mode runs a single ``{name}`` instance with
2655 the ``kv_both`` role. The role's ``--kv-transfer-config`` is attached to
2656 the vLLM container, EFA scheduling is applied when transfer runs over
2657 RDMA, and the replica count is taken from the topology (or the
2658 autoscaling lower bound when that is enabled).
2660 Args:
2661 name: The endpoint name.
2662 ns: The namespace to materialize into.
2663 spec: The endpoint spec; ``spec["mooncake"]`` selects the mode and
2664 topology.
2665 role: One of ``"prefill"``, ``"decode"``, or ``"single"``.
2666 """
2667 mooncake = spec.get("mooncake") or {}
2669 # The store's single instance keeps the endpoint name; prefill and decode
2670 # are suffixed so Services and autoscalers can target each role.
2671 deploy_name = name if role == "single" else f"{name}-{role}"
2673 kv_transfer_config = build_kv_transfer_config(mooncake, role)
2674 replicas = self._replica_count_for_role(mooncake, role)
2676 deployment = self._build_inference_deployment_object(
2677 name=name,
2678 deploy_name=deploy_name,
2679 app_label=deploy_name,
2680 namespace=ns,
2681 spec=spec,
2682 replicas=replicas,
2683 extra_args=["--kv-transfer-config", kv_transfer_config],
2684 extra_labels={"gco.io/role": role},
2685 )
2687 # Land role pods on the EFA fabric when KV transfer runs over RDMA,
2688 # preserving the GPU asks already built into the pod.
2689 apply_efa_scheduling(mooncake, deployment.spec.template.spec)
2691 self.apps_v1.create_namespaced_deployment(
2692 ns, deployment, _request_timeout=self._k8s_timeout
2693 )
2694 logger.info("Created role deployment %s/%s (role=%s)", ns, deploy_name, role)
2696 def _verify_admin_api_key_secret(self, proxy: dict[str, Any], ns: str) -> str:
2697 """Confirm the proxy admin key Secret exists and carries a key value.
2699 The proxy guards a privileged admin path and must never run without a
2700 usable ``ADMIN_API_KEY``. This reads the Secret named by
2701 ``proxy.admin_api_key_secret`` and confirms it holds a non-empty
2702 ``ADMIN_API_KEY`` value, so the value itself never has to be carried on
2703 the spec or a command argument.
2705 Args:
2706 proxy: The ``spec["mooncake"]["proxy"]`` block; ``admin_api_key_secret``
2707 names the backing Secret.
2708 ns: The namespace the Secret lives in.
2710 Returns:
2711 The verified Secret name, suitable for a Secret reference.
2713 Raises:
2714 AdminApiKeySecretError: If the spec names no Secret, the Secret is
2715 absent, or its ``ADMIN_API_KEY`` value is empty or missing.
2716 """
2717 secret_name = proxy.get("admin_api_key_secret")
2718 if not isinstance(secret_name, str) or not secret_name:
2719 raise AdminApiKeySecretError(None, "no admin API key Secret was named")
2721 try:
2722 secret = self.core_v1.read_namespaced_secret(
2723 secret_name, ns, _request_timeout=self._k8s_timeout
2724 )
2725 except ApiException as e:
2726 if e.status == 404: 2726 ↛ 2728line 2726 didn't jump to line 2728 because the condition on line 2726 was always true
2727 raise AdminApiKeySecretError(secret_name, "Secret not found") from e
2728 raise
2730 if not self._secret_has_admin_api_key(secret):
2731 raise AdminApiKeySecretError(
2732 secret_name,
2733 f"{ADMIN_API_KEY_SECRET_DATA_KEY} value is empty or missing",
2734 )
2736 return secret_name
2738 @staticmethod
2739 def _secret_has_admin_api_key(secret: client.V1Secret) -> bool:
2740 """Return whether ``secret`` carries a non-empty ``ADMIN_API_KEY``.
2742 Both the base64 ``data`` and the plaintext ``string_data`` views are
2743 considered, and a value is treated as present only when it decodes to a
2744 non-empty string.
2745 """
2746 string_data = secret.string_data or {}
2747 plain = string_data.get(ADMIN_API_KEY_SECRET_DATA_KEY)
2748 if plain:
2749 return True
2751 data = secret.data or {}
2752 encoded = data.get(ADMIN_API_KEY_SECRET_DATA_KEY)
2753 if not encoded:
2754 return False
2755 try:
2756 return bool(base64.b64decode(encoded))
2757 except ValueError, TypeError:
2758 # A value that cannot be decoded is unusable as an admin key.
2759 return False
2761 def _ensure_admin_api_key_secret(self, name: str, proxy: dict[str, Any], ns: str) -> str:
2762 """Return the proxy admin-key Secret name, provisioning one if needed.
2764 The prefill-decode proxy guards a privileged admin path and must never
2765 run without a usable ``ADMIN_API_KEY``. Two paths satisfy that:
2767 - **Bring-your-own**: when the proxy block names a Secret, that Secret
2768 must already exist and carry a non-empty ``ADMIN_API_KEY``; otherwise
2769 the deployment is rejected, so a typo or a missing pre-created Secret
2770 fails fast. The named Secret is only read, never created or mutated.
2771 - **Auto-managed**: when the proxy names no Secret, a per-endpoint
2772 ``{name}-admin`` Secret is provisioned create-if-absent with a
2773 generated key, so a split deploy needs no manual Secret. The generated
2774 key only ever lives in the cluster — it is never written to the
2775 endpoint spec, a command argument, or a log line.
2777 Args:
2778 name: The endpoint name, used to derive the auto-managed Secret name.
2779 proxy: The ``spec["mooncake"]["proxy"]`` block.
2780 ns: The namespace the Secret lives in.
2782 Returns:
2783 The Secret name to reference from the proxy container.
2785 Raises:
2786 AdminApiKeySecretError: Only on the bring-your-own path, when the
2787 named Secret is absent or its ``ADMIN_API_KEY`` is empty. The
2788 auto-managed path never raises this.
2789 """
2790 named = proxy.get("admin_api_key_secret")
2791 if isinstance(named, str) and named:
2792 return self._verify_admin_api_key_secret(proxy, ns)
2793 return self._provision_admin_api_key_secret(f"{name}-admin", ns)
2795 def _provision_admin_api_key_secret(self, secret_name: str, ns: str) -> str:
2796 """Create the auto-managed proxy admin-key Secret if absent.
2798 Uses create-if-absent semantics so the key stays stable across reconcile
2799 passes: an existing Secret (the steady state, or one a prior pass
2800 created) is left untouched, and a concurrent create (409) is treated as
2801 success. A freshly created Secret carries a cryptographically strong
2802 64-character hex ``ADMIN_API_KEY`` from :func:`secrets.token_hex`, which
2803 reaches the proxy only through a Secret reference — the value is never
2804 logged or written to the spec.
2806 Args:
2807 secret_name: The Secret to ensure exists (``{endpoint}-admin``).
2808 ns: The namespace to create it in.
2810 Returns:
2811 The Secret name, ready for a Secret reference.
2812 """
2813 try:
2814 self.core_v1.read_namespaced_secret(secret_name, ns, _request_timeout=self._k8s_timeout)
2815 # Already present: keep the existing key so proxy pods need no churn.
2816 return secret_name
2817 except ApiException as e:
2818 if e.status != 404: 2818 ↛ 2819line 2818 didn't jump to line 2819 because the condition on line 2818 was never true
2819 raise
2821 secret = client.V1Secret(
2822 metadata=client.V1ObjectMeta(
2823 name=secret_name,
2824 namespace=ns,
2825 labels={"app": secret_name, "project": "gco", "gco.io/type": "inference"},
2826 ),
2827 string_data={ADMIN_API_KEY_SECRET_DATA_KEY: secrets.token_hex(32)},
2828 type="Opaque",
2829 )
2830 # The two logger.info calls below carry a bare `# nosemgrep`: the
2831 # logger-credential-disclosure rule matches the literal word "Secret" in
2832 # the message, but only the Secret's name and namespace (%s/%s) are
2833 # logged here — never the generated key value set above in string_data.
2834 try:
2835 self.core_v1.create_namespaced_secret(ns, secret, _request_timeout=self._k8s_timeout)
2836 logger.info("Provisioned proxy admin-key Secret %s/%s", ns, secret_name) # nosemgrep
2837 except ApiException as e:
2838 if e.status == 409: 2838 ↛ 2841line 2838 didn't jump to line 2841 because the condition on line 2838 was always true
2839 logger.info("Proxy admin-key Secret %s/%s exists", ns, secret_name) # nosemgrep
2840 else:
2841 raise
2842 return secret_name
2844 def _create_pd_proxy(
2845 self, name: str, ns: str, spec: dict[str, Any], endpoint: dict[str, Any]
2846 ) -> None:
2847 """Materialize the prefill-decode proxy front for a disaggregated endpoint.
2849 Disaggregated and ``both`` modes are fronted by a lightweight proxy that
2850 runs the residency check and dispatches each request to the prefill and
2851 decode pods. It materializes a ConfigMap, a proxy Deployment with at
2852 least one replica, and a Service whose selector matches only the proxy
2853 pods. The shared HTTPRoute attached to ``gco-system/gco-gateway`` sends
2854 ``/inference`` to ``gco-system/inference-proxy``; that authenticated
2855 platform proxy then reaches this internal ClusterIP Service.
2856 Endpoint-specific Ingresses are removed as an unsafe legacy path.
2858 Before those resources are created, a user-named
2859 ``mooncake.proxy.admin_api_key_secret`` is verified to contain a usable
2860 ``ADMIN_API_KEY``. When no Secret is named, the monitor auto-provisions
2861 a generated ``{name}-admin`` Secret instead. A missing or empty named
2862 Secret rejects the proxy; the key itself reaches the container only as
2863 a Secret reference at pod start and is never written to the spec or a
2864 command argument.
2866 Creation is idempotent at the API boundary: an already-present Deployment
2867 or ClusterIP Service is left in place, and historical direct Ingresses are
2868 deleted if present. No endpoint Gateway or HTTPRoute is created.
2870 Args:
2871 name: The endpoint name.
2872 ns: The namespace to materialize into.
2873 spec: The endpoint spec; ``spec["mooncake"]`` supplies the proxy
2874 image and behavior.
2875 endpoint: The endpoint record. Legacy per-endpoint routing metadata
2876 is ignored because the shared platform HTTPRoute owns the prefix.
2878 Raises:
2879 AdminApiKeySecretError: If the admin key Secret is missing, names no
2880 Secret, or holds an empty ``ADMIN_API_KEY``. No proxy resource
2881 is created in that case.
2882 """
2883 mooncake = spec.get("mooncake") or {}
2884 proxy = mooncake.get("proxy") or {}
2885 proxy_name = f"{name}-proxy"
2886 del endpoint # Per-endpoint routing metadata is legacy and intentionally ignored.
2888 # The proxy fronts a privileged admin path, so it never starts without a
2889 # usable admin key. When the spec names a Secret it must already exist
2890 # and be non-empty (the deployment is rejected otherwise); when it names
2891 # none, a per-endpoint admin-key Secret is auto-provisioned with a
2892 # generated key. Either way the key reaches the container only by Secret
2893 # reference.
2894 admin_secret_name = self._ensure_admin_api_key_secret(name, proxy, ns)
2896 proxy_env = build_pd_proxy_config(mooncake)
2897 container_env = [client.V1EnvVar(name=k, value=v) for k, v in proxy_env.items()]
2898 # Deliver the admin key by Secret reference only — its value is never
2899 # placed on the spec or a command argument.
2900 container_env.append(
2901 client.V1EnvVar(
2902 name=PD_PROXY_ADMIN_API_KEY_ENV,
2903 value_from=client.V1EnvVarSource(
2904 secret_key_ref=client.V1SecretKeySelector(
2905 name=admin_secret_name,
2906 key=ADMIN_API_KEY_SECRET_DATA_KEY,
2907 )
2908 ),
2909 )
2910 )
2912 # The proxy fronts at least one replica; a spec may ask for more.
2913 replicas = proxy.get("replicas", 1)
2914 if not isinstance(replicas, int) or isinstance(replicas, bool) or replicas < 1: 2914 ↛ 2915line 2914 didn't jump to line 2915 because the condition on line 2914 was never true
2915 replicas = 1
2917 labels = {
2918 "app": proxy_name,
2919 "project": "gco",
2920 "gco.io/type": "inference",
2921 "gco.io/role": PD_PROXY_ROLE_LABEL,
2922 }
2924 # The proxy reaches prefill and decode through their per-role Services
2925 # and listens on PD_PROXY_PORT for requests from the authenticated API
2926 # proxy. Routing via Services means only Ready role pods receive traffic.
2927 port = spec.get("port", 8000)
2928 container_env.extend(
2929 [
2930 client.V1EnvVar(name=PD_PROXY_PORT_ENV, value=str(PD_PROXY_PORT)),
2931 client.V1EnvVar(
2932 name=PD_PROXY_PREFILL_URL_ENV, value=f"http://{name}-prefill:{port}"
2933 ),
2934 client.V1EnvVar(name=PD_PROXY_DECODE_URL_ENV, value=f"http://{name}-decode:{port}"),
2935 ]
2936 )
2938 # Ship the proxy program to the pod as a ConfigMap and run it from there.
2939 self._ensure_pd_proxy_configmap(name, ns)
2940 proxy_volume_name = "pd-proxy-script"
2942 container = client.V1Container(
2943 name="proxy",
2944 image=proxy.get("image"),
2945 command=["python3", PD_PROXY_SCRIPT_PATH],
2946 ports=[client.V1ContainerPort(container_port=PD_PROXY_PORT)],
2947 env=container_env if container_env else None,
2948 resources=client.V1ResourceRequirements(
2949 requests={"cpu": "250m", "memory": "256Mi"},
2950 limits={"cpu": "1", "memory": "1Gi"},
2951 ),
2952 volume_mounts=[
2953 client.V1VolumeMount(
2954 name=proxy_volume_name,
2955 mount_path=PD_PROXY_CONFIG_MOUNT_DIR,
2956 read_only=True,
2957 )
2958 ],
2959 readiness_probe=client.V1Probe(
2960 tcp_socket=client.V1TCPSocketAction(port=PD_PROXY_PORT),
2961 initial_delay_seconds=10,
2962 period_seconds=10,
2963 ),
2964 liveness_probe=client.V1Probe(
2965 tcp_socket=client.V1TCPSocketAction(port=PD_PROXY_PORT),
2966 initial_delay_seconds=30,
2967 period_seconds=15,
2968 failure_threshold=5,
2969 ),
2970 )
2972 deployment = client.V1Deployment(
2973 metadata=client.V1ObjectMeta(
2974 name=proxy_name,
2975 namespace=ns,
2976 labels=dict(labels),
2977 ),
2978 spec=client.V1DeploymentSpec(
2979 replicas=replicas,
2980 selector=client.V1LabelSelector(match_labels={"app": proxy_name}),
2981 template=client.V1PodTemplateSpec(
2982 metadata=client.V1ObjectMeta(labels=dict(labels)),
2983 spec=client.V1PodSpec(
2984 service_account_name="gco-service-account",
2985 containers=[container],
2986 volumes=[
2987 client.V1Volume(
2988 name=proxy_volume_name,
2989 config_map=client.V1ConfigMapVolumeSource(
2990 name=f"{name}-pd-proxy",
2991 default_mode=0o555,
2992 ),
2993 )
2994 ],
2995 ),
2996 ),
2997 ),
2998 )
3000 try:
3001 self.apps_v1.create_namespaced_deployment(
3002 ns, deployment, _request_timeout=self._k8s_timeout
3003 )
3004 logger.info("Created proxy deployment %s/%s", ns, proxy_name)
3005 except ApiException as e:
3006 if e.status == 409: 3006 ↛ 3009line 3006 didn't jump to line 3009 because the condition on line 3006 was always true
3007 logger.info("Proxy deployment %s/%s already exists", ns, proxy_name)
3008 else:
3009 raise
3011 self._create_proxy_service(proxy_name, ns)
3013 def _create_role_service(self, name: str, ns: str, role: str, port: int = 8000) -> None:
3014 """Create the ClusterIP Service that fronts one role's pods.
3016 Named ``{name}-{role}`` and selecting that role Deployment's app label,
3017 so the PD proxy can address prefill or decode by stable in-cluster DNS.
3018 Routing through a Service means kube-proxy load-balances across only the
3019 role's Ready pods, which is what gives the proxy ready-only decode
3020 routing without watching the Kubernetes API. Idempotent at the API
3021 boundary: an already-present Service is left in place.
3022 """
3023 deploy_name = f"{name}-{role}"
3024 service = client.V1Service(
3025 metadata=client.V1ObjectMeta(
3026 name=deploy_name,
3027 namespace=ns,
3028 labels={
3029 "app": deploy_name,
3030 "project": "gco",
3031 "gco.io/type": "inference",
3032 "gco.io/role": role,
3033 },
3034 ),
3035 spec=client.V1ServiceSpec(
3036 selector={"app": deploy_name},
3037 ports=[client.V1ServicePort(port=port, target_port=port, protocol="TCP")],
3038 type="ClusterIP",
3039 ),
3040 )
3041 try:
3042 self.core_v1.create_namespaced_service(ns, service, _request_timeout=self._k8s_timeout)
3043 logger.info("Created role service %s/%s", ns, deploy_name)
3044 except ApiException as e:
3045 if e.status == 409: 3045 ↛ 3048line 3045 didn't jump to line 3048 because the condition on line 3045 was always true
3046 logger.info("Role service %s/%s already exists", ns, deploy_name)
3047 else:
3048 raise
3050 def _ensure_pd_proxy_configmap(self, name: str, ns: str) -> None:
3051 """Publish the PD proxy program to the pod as a ConfigMap.
3053 The proxy program (``mooncake_pd_proxy.py``) ships in this image
3054 alongside the monitor; its source is read here and mounted into the
3055 ``{name}-proxy`` pod, which runs it with ``python3`` from
3056 ``PD_PROXY_SCRIPT_PATH``. The ConfigMap is patched on conflict so the
3057 program tracks the running monitor build.
3058 """
3059 script = (Path(__file__).resolve().parent / PD_PROXY_SCRIPT_FILENAME).read_text(
3060 encoding="utf-8"
3061 )
3062 cm_name = f"{name}-pd-proxy"
3063 body = client.V1ConfigMap(
3064 metadata=client.V1ObjectMeta(
3065 name=cm_name,
3066 namespace=ns,
3067 labels={
3068 "app": f"{name}-proxy",
3069 "project": "gco",
3070 "gco.io/type": "inference",
3071 "gco.io/role": PD_PROXY_ROLE_LABEL,
3072 },
3073 ),
3074 data={PD_PROXY_SCRIPT_FILENAME: script},
3075 )
3076 try:
3077 self.core_v1.create_namespaced_config_map(ns, body, _request_timeout=self._k8s_timeout)
3078 logger.info("Created PD proxy ConfigMap %s/%s", ns, cm_name)
3079 except ApiException as e:
3080 if e.status == 409: 3080 ↛ 3086line 3080 didn't jump to line 3086 because the condition on line 3080 was always true
3081 self.core_v1.patch_namespaced_config_map(
3082 cm_name, ns, body, _request_timeout=self._k8s_timeout
3083 )
3084 logger.info("Updated PD proxy ConfigMap %s/%s", ns, cm_name)
3085 else:
3086 raise
3088 def _create_proxy_service(self, proxy_name: str, namespace: str) -> None:
3089 """Create the Service that fronts only the proxy pods.
3091 The selector is the ``{name}-proxy`` app label together with the proxy
3092 role marker, so the Service resolves exclusively to proxy pods and never
3093 to the prefill or decode role pods that share the namespace.
3094 """
3095 service = client.V1Service(
3096 metadata=client.V1ObjectMeta(
3097 name=proxy_name,
3098 namespace=namespace,
3099 labels={
3100 "app": proxy_name,
3101 "project": "gco",
3102 "gco.io/type": "inference",
3103 "gco.io/role": PD_PROXY_ROLE_LABEL,
3104 },
3105 ),
3106 spec=client.V1ServiceSpec(
3107 selector={"app": proxy_name, "gco.io/role": PD_PROXY_ROLE_LABEL},
3108 ports=[
3109 client.V1ServicePort(
3110 port=80,
3111 target_port=PD_PROXY_PORT,
3112 protocol="TCP",
3113 )
3114 ],
3115 type="ClusterIP",
3116 ),
3117 )
3119 try:
3120 self.core_v1.create_namespaced_service(
3121 namespace, service, _request_timeout=self._k8s_timeout
3122 )
3123 logger.info("Created proxy service %s/%s", namespace, proxy_name)
3124 except ApiException as e:
3125 if e.status == 409: 3125 ↛ 3128line 3125 didn't jump to line 3128 because the condition on line 3125 was always true
3126 logger.info("Proxy service %s/%s already exists", namespace, proxy_name)
3127 else:
3128 raise
3130 def _create_service(self, name: str, namespace: str, spec: dict[str, Any]) -> None:
3131 """Create the internal ClusterIP Service for an inference endpoint."""
3132 port = spec.get("port", 8000)
3134 service = client.V1Service(
3135 metadata=client.V1ObjectMeta(
3136 name=name,
3137 namespace=namespace,
3138 labels={
3139 "app": name,
3140 "project": "gco",
3141 "gco.io/type": "inference",
3142 },
3143 ),
3144 spec=client.V1ServiceSpec(
3145 selector={"app": name},
3146 ports=[
3147 client.V1ServicePort(
3148 port=80,
3149 target_port=port,
3150 protocol="TCP",
3151 )
3152 ],
3153 type="ClusterIP",
3154 ),
3155 )
3157 try:
3158 self.core_v1.create_namespaced_service(
3159 namespace, service, _request_timeout=self._k8s_timeout
3160 )
3161 logger.info("Created service %s/%s", namespace, name)
3162 except ApiException as e:
3163 if e.status == 409: 3163 ↛ 3166line 3163 didn't jump to line 3166 because the condition on line 3163 was always true
3164 logger.info("Service %s/%s already exists", namespace, name)
3165 else:
3166 raise
3168 def _ensure_service(self, name: str, namespace: str, spec: dict[str, Any]) -> None:
3169 """Ensure the endpoint's ClusterIP Service exists, recreating it if missing."""
3170 try:
3171 self.core_v1.read_namespaced_service(
3172 name, namespace, _request_timeout=self._k8s_timeout
3173 )
3174 except ApiException as e:
3175 if e.status == 404:
3176 logger.warning("Service %s/%s missing, recreating", namespace, name)
3177 self._create_service(name, namespace, spec)
3178 else:
3179 raise
3181 def _check_health_watchdog(
3182 self,
3183 name: str,
3184 namespace: str,
3185 ready_replicas: int,
3186 desired_replicas: int,
3187 spec: dict[str, Any],
3188 endpoint: dict[str, Any],
3189 ) -> bool:
3190 """Track prolonged unavailability without changing shared routing.
3192 ``gco-system/gco-gateway`` owns the shared ``/inference`` HTTPRoute to
3193 ``gco-system/inference-proxy``. Individual models are reached through
3194 internal ClusterIP Services, so their readiness never changes the shared
3195 Gateway or HTTPRoute. The threshold still drives degraded-state logging.
3196 """
3197 del namespace, spec, endpoint
3198 if ready_replicas > 0:
3199 if name in self._unready_since:
3200 logger.info("Endpoint %s recovered", name)
3201 del self._unready_since[name]
3202 return False
3204 now = datetime.now(UTC)
3205 if name not in self._unready_since:
3206 self._unready_since[name] = now
3207 logger.warning(
3208 "Endpoint %s has 0/%d ready replicas, starting health watchdog timer",
3209 name,
3210 desired_replicas,
3211 )
3212 return False
3214 unready_duration = (now - self._unready_since[name]).total_seconds()
3215 threshold_exceeded = unready_duration >= self._unhealthy_threshold_seconds
3216 if threshold_exceeded:
3217 logger.warning(
3218 "WATCHDOG: Endpoint %s has been unavailable for %ds (threshold %ds); "
3219 "the authenticated proxy will return 503 until it recovers",
3220 name,
3221 int(unready_duration),
3222 self._unhealthy_threshold_seconds,
3223 )
3224 return threshold_exceeded
3226 def _scale_deployment(self, name: str, namespace: str, replicas: int) -> None:
3227 """Scale a deployment to the desired replica count."""
3228 self.apps_v1.patch_namespaced_deployment(
3229 name,
3230 namespace,
3231 body={"spec": {"replicas": replicas}},
3232 _request_timeout=self._k8s_timeout,
3233 )
3235 def _update_deployment_image(self, name: str, namespace: str, image: str) -> None:
3236 """Update the container image of a deployment."""
3237 self.apps_v1.patch_namespaced_deployment(
3238 name,
3239 namespace,
3240 body={
3241 "spec": {
3242 "template": {"spec": {"containers": [{"name": "inference", "image": image}]}}
3243 }
3244 },
3245 _request_timeout=self._k8s_timeout,
3246 )
3248 def _reconcile_canary(
3249 self,
3250 name: str,
3251 namespace: str,
3252 spec: dict[str, Any],
3253 canary: dict[str, Any],
3254 endpoint: dict[str, Any],
3255 ) -> dict[str, Any]:
3256 """Reconcile a classic canary and return observed readiness for routing."""
3257 canary_image_value = canary.get("image")
3258 if not isinstance(canary_image_value, str) or not canary_image_value.strip():
3259 raise ValueError("canary.image must be a non-empty string")
3260 canary_image = canary_image_value.strip()
3262 canary_replicas = canary.get("replicas", 1)
3263 if (
3264 not isinstance(canary_replicas, int)
3265 or isinstance(canary_replicas, bool)
3266 or canary_replicas < 1
3267 ):
3268 raise ValueError("canary.replicas must be a positive integer")
3270 canary_weight = canary.get("weight", 10)
3271 if (
3272 not isinstance(canary_weight, int)
3273 or isinstance(canary_weight, bool)
3274 or not 1 <= canary_weight <= 99
3275 ):
3276 raise ValueError("canary.weight must be an integer between 1 and 99")
3278 canary_name = f"{name}-canary"
3279 del endpoint # Per-endpoint routing metadata is legacy and intentionally ignored.
3281 canary_spec = dict(spec)
3282 canary_spec["image"] = canary_image
3283 canary_spec["replicas"] = canary_replicas
3284 canary_spec.pop("canary", None)
3285 # A canary image is explicit and global. Retaining the primary's
3286 # region_image_uris would silently deploy the old regional image.
3287 canary_spec.pop("region_image_uris", None)
3289 deployment = self._get_deployment(canary_name, namespace)
3290 state = "creating"
3291 ready_replicas = 0
3292 if deployment is None:
3293 logger.info("Creating canary deployment %s with image %s", canary_name, canary_image)
3294 self._create_deployment(canary_name, namespace, canary_spec)
3295 self._create_service(canary_name, namespace, canary_spec)
3296 else:
3297 self._ensure_service(canary_name, namespace, canary_spec)
3298 current_image = self._get_deployment_image(deployment)
3299 current_replicas = deployment.spec.replicas or 1
3300 ready_replicas = deployment.status.ready_replicas or 0
3301 if current_image != canary_image:
3302 self._update_deployment_image(canary_name, namespace, canary_image)
3303 ready_replicas = 0
3304 state = "updating"
3305 elif current_replicas != canary_replicas:
3306 self._scale_deployment(canary_name, namespace, canary_replicas)
3307 ready_replicas = min(ready_replicas, canary_replicas)
3308 state = "updating"
3309 elif ready_replicas >= canary_replicas: 3309 ↛ 3316line 3309 didn't jump to line 3316 because the condition on line 3309 was always true
3310 state = "running"
3312 # Canary selection happens behind ``gco-system/inference-proxy``; the
3313 # shared ``gco-system/gco-gateway`` HTTPRoute is never changed per
3314 # endpoint. The shared inference proxy consumes the observed canary
3315 # status returned here.
3316 return {
3317 "state": state,
3318 "image": canary_image,
3319 "weight": canary_weight,
3320 "replicas_ready": ready_replicas,
3321 "replicas_desired": canary_replicas,
3322 }
3324 def _cleanup_canary(self, name: str, namespace: str) -> None:
3325 """Remove a canary Deployment and Service."""
3326 canary_name = f"{name}-canary"
3328 # Delete canary deployment
3329 try:
3330 self.apps_v1.delete_namespaced_deployment(
3331 canary_name, namespace, _request_timeout=self._k8s_timeout
3332 )
3333 logger.info("Deleted canary deployment %s", canary_name)
3334 except ApiException as e:
3335 if e.status != 404:
3336 logger.error("Failed to delete canary deployment %s: %s", canary_name, e)
3338 # Delete canary service
3339 try:
3340 self.core_v1.delete_namespaced_service(
3341 canary_name, namespace, _request_timeout=self._k8s_timeout
3342 )
3343 logger.info("Deleted canary service %s", canary_name)
3344 except ApiException as e:
3345 if e.status != 404:
3346 logger.error("Failed to delete canary service %s: %s", canary_name, e)
3348 def _delete_resources(
3349 self,
3350 name: str,
3351 namespace: str,
3352 spec: dict[str, Any] | None = None,
3353 ) -> None:
3354 """Delete all Kubernetes resources owned by an endpoint.
3356 Covers both the single-Deployment endpoint and the Mooncake role-split
3357 topology: role/proxy Deployments and Services, native HPAs or KEDA
3358 ScaledObjects, transport/proxy ConfigMaps, and the generated proxy admin
3359 Secret when the endpoint did not name a user-managed Secret. Each delete
3360 is idempotent: a 404 means that object
3361 is not used by this endpoint's mode and is ignored. The shared regional
3362 ``mooncake-master`` is deliberately NOT deleted because other endpoints
3363 may still depend on it.
3364 """
3365 # Delete canary resources first
3366 self._cleanup_canary(name, namespace)
3368 proxy_name = f"{name}-proxy"
3370 # Deployments: the single-instance endpoint plus the Mooncake prefill/
3371 # decode workers and the PD proxy.
3372 for deployment_name in (name, f"{name}-prefill", f"{name}-decode", proxy_name):
3373 try:
3374 self.apps_v1.delete_namespaced_deployment(
3375 deployment_name, namespace, _request_timeout=self._k8s_timeout
3376 )
3377 logger.info("Deleted deployment %s/%s", namespace, deployment_name)
3378 except ApiException as e:
3379 if e.status != 404:
3380 logger.error("Failed to delete deployment %s: %s", deployment_name, e)
3382 # Services: classic/store, split-role backends, and the PD proxy.
3383 for service_name in (name, f"{name}-prefill", f"{name}-decode", proxy_name):
3384 try:
3385 self.core_v1.delete_namespaced_service(
3386 service_name, namespace, _request_timeout=self._k8s_timeout
3387 )
3388 logger.info("Deleted service %s/%s", namespace, service_name)
3389 except ApiException as e:
3390 if e.status != 404:
3391 logger.error("Failed to delete service %s: %s", service_name, e)
3393 # HPAs: the single-instance endpoint HPA and the per-role Mooncake HPAs.
3394 autoscaling_v2 = client.AutoscalingV2Api()
3395 for hpa_name in (name, f"{name}-prefill", f"{name}-decode"):
3396 try:
3397 autoscaling_v2.delete_namespaced_horizontal_pod_autoscaler(hpa_name, namespace)
3398 logger.info("Deleted HPA %s/%s", namespace, hpa_name)
3399 except ApiException as e:
3400 if e.status != 404:
3401 logger.error("Failed to delete HPA %s: %s", hpa_name, e)
3403 # Per-endpoint Mooncake transport and bundled PD-proxy program. The
3404 # shared regional mooncake-master is intentionally left in place.
3405 for config_map_name in (f"{name}-mooncake", f"{name}-pd-proxy"):
3406 try:
3407 self.core_v1.delete_namespaced_config_map(
3408 config_map_name, namespace, _request_timeout=self._k8s_timeout
3409 )
3410 logger.info("Deleted configmap %s/%s", namespace, config_map_name)
3411 except ApiException as e:
3412 if e.status != 404: 3412 ↛ 3413line 3412 didn't jump to line 3413 because the condition on line 3412 was never true
3413 logger.error("Failed to delete configmap %s: %s", config_map_name, e)
3415 # GPU metrics use KEDA instead of native HPAs. Remove the classic and
3416 # both role-scoped names; absent objects are the common case.
3417 custom_objects = client.CustomObjectsApi()
3418 for scaled_object_name in (name, f"{name}-prefill", f"{name}-decode"):
3419 try:
3420 custom_objects.delete_namespaced_custom_object(
3421 group=KEDA_API_GROUP,
3422 version=KEDA_API_VERSION,
3423 namespace=namespace,
3424 plural=KEDA_SCALEDOBJECT_PLURAL,
3425 name=scaled_object_name,
3426 _request_timeout=self._k8s_timeout,
3427 )
3428 logger.info("Deleted KEDA ScaledObject for %s", scaled_object_name)
3429 except ApiException as e:
3430 if e.status != 404:
3431 logger.error(
3432 "Failed to delete KEDA ScaledObject for %s: %s",
3433 scaled_object_name,
3434 e,
3435 )
3437 # Only the unnamed-secret path is monitor-owned. A Secret explicitly
3438 # named in the endpoint spec is user-managed and must survive deletion.
3439 mooncake = spec.get("mooncake") if isinstance(spec, dict) else None
3440 proxy = mooncake.get("proxy") if isinstance(mooncake, dict) else None
3441 named_secret = proxy.get("admin_api_key_secret") if isinstance(proxy, dict) else None
3442 if (
3443 isinstance(mooncake, dict)
3444 and mooncake.get("mode") in ("disaggregated", "both")
3445 and not (isinstance(named_secret, str) and named_secret)
3446 ):
3447 secret_name = f"{name}-admin"
3448 try:
3449 self.core_v1.delete_namespaced_secret(
3450 secret_name, namespace, _request_timeout=self._k8s_timeout
3451 )
3452 logger.info("Deleted generated proxy access resource %s/%s", namespace, secret_name)
3453 except ApiException as e:
3454 if e.status != 404:
3455 logger.error(
3456 "Failed to delete generated proxy access resource %s: %s",
3457 secret_name,
3458 e,
3459 )
3461 def _build_hpa_metrics(self, metrics_config: list[dict[str, Any]]) -> list[Any]:
3462 """Translate a metrics config list into autoscaler metric specs.
3464 Each entry names a resource (``cpu`` or ``memory``) and a target
3465 average utilization. Unrecognized entries are skipped, and when nothing
3466 recognizable remains the autoscaler falls back to scaling on CPU at 70%
3467 so a Deployment is never left without a scaling signal.
3468 """
3469 hpa_metrics = []
3470 for m in metrics_config:
3471 metric_type = m.get("type", "cpu")
3472 target_value = m.get("target", 70)
3474 if metric_type == "cpu":
3475 hpa_metrics.append(
3476 client.V2MetricSpec(
3477 type="Resource",
3478 resource=client.V2ResourceMetricSource(
3479 name="cpu",
3480 target=client.V2MetricTarget(
3481 type="Utilization",
3482 average_utilization=target_value,
3483 ),
3484 ),
3485 )
3486 )
3487 elif metric_type == "memory":
3488 hpa_metrics.append(
3489 client.V2MetricSpec(
3490 type="Resource",
3491 resource=client.V2ResourceMetricSource(
3492 name="memory",
3493 target=client.V2MetricTarget(
3494 type="Utilization",
3495 average_utilization=target_value,
3496 ),
3497 ),
3498 )
3499 )
3501 if not hpa_metrics:
3502 # Default to CPU if no recognized metrics
3503 hpa_metrics.append(
3504 client.V2MetricSpec(
3505 type="Resource",
3506 resource=client.V2ResourceMetricSource(
3507 name="cpu",
3508 target=client.V2MetricTarget(
3509 type="Utilization",
3510 average_utilization=70,
3511 ),
3512 ),
3513 )
3514 )
3516 return hpa_metrics
3518 @staticmethod
3519 def _metrics_require_keda(metrics_config: list[dict[str, Any]]) -> bool:
3520 """Return True when any metric can only be scaled via KEDA/CloudWatch.
3522 GPU metrics are not Kubernetes Resource metrics, so a native HPA cannot
3523 consume them. Their presence forces the whole autoscaler onto the KEDA
3524 ScaledObject path, where cpu/memory targets become native KEDA triggers
3525 alongside the aws-cloudwatch GPU trigger.
3526 """
3527 return any(m.get("type") in _CLOUDWATCH_METRIC_BY_TYPE for m in metrics_config)
3529 def _build_keda_triggers(
3530 self,
3531 metrics_config: list[dict[str, Any]],
3532 target_name: str,
3533 namespace: str,
3534 ) -> list[dict[str, Any]]:
3535 """Translate a metrics config list into KEDA ScaledObject triggers.
3537 ``cpu`` and ``memory`` map to KEDA's native resource triggers (the same
3538 utilization signal a plain HPA would use). ``gpu``/``gpu_memory`` map to
3539 an ``aws-cloudwatch`` trigger reading the matching ContainerInsights
3540 metric for this Deployment, identified by the
3541 ClusterName/Namespace/PodName dimension triple. Unrecognized entries are
3542 skipped; when nothing recognizable remains the autoscaler falls back to
3543 CPU at 70% so a Deployment is never left without a scaling signal.
3544 """
3545 triggers: list[dict[str, Any]] = []
3546 for m in metrics_config:
3547 metric_type = m.get("type", "cpu")
3548 target_value = m.get("target", 70)
3550 if metric_type in ("cpu", "memory"):
3551 triggers.append(
3552 {
3553 "type": metric_type,
3554 "metricType": "Utilization",
3555 "metadata": {"value": str(target_value)},
3556 }
3557 )
3558 elif metric_type in _CLOUDWATCH_METRIC_BY_TYPE: 3558 ↛ 3546line 3558 didn't jump to line 3546 because the condition on line 3558 was always true
3559 triggers.append(
3560 {
3561 "type": "aws-cloudwatch",
3562 "metadata": {
3563 "namespace": GPU_METRIC_NAMESPACE,
3564 "metricName": _CLOUDWATCH_METRIC_BY_TYPE[metric_type],
3565 "dimensionName": "ClusterName;Namespace;PodName",
3566 "dimensionValue": f"{self.cluster_id};{namespace};{target_name}",
3567 "targetMetricValue": str(target_value),
3568 "minMetricValue": "0",
3569 "metricStat": "Average",
3570 "awsRegion": self.region,
3571 "identityOwner": "operator",
3572 },
3573 }
3574 )
3576 if not triggers: 3576 ↛ 3577line 3576 didn't jump to line 3577 because the condition on line 3576 was never true
3577 triggers.append(
3578 {
3579 "type": "cpu",
3580 "metricType": "Utilization",
3581 "metadata": {"value": "70"},
3582 }
3583 )
3585 return triggers
3587 def _apply_scaled_object(
3588 self,
3589 name: str,
3590 namespace: str,
3591 target_name: str,
3592 min_replicas: int,
3593 max_replicas: int,
3594 metrics_config: list[dict[str, Any]],
3595 ) -> None:
3596 """Create or patch a KEDA ScaledObject targeting one Deployment.
3598 Used whenever the metric set includes a GPU signal (see
3599 :meth:`_metrics_require_keda`). KEDA owns the backing HPA and reads GPU
3600 utilization from CloudWatch via the keda-operator's IRSA role, scaling
3601 ``target_name`` between ``min_replicas`` and ``max_replicas``. An
3602 already-present ScaledObject of the same name is merge-patched rather
3603 than duplicated.
3604 """
3605 body = {
3606 "apiVersion": f"{KEDA_API_GROUP}/{KEDA_API_VERSION}",
3607 "kind": "ScaledObject",
3608 "metadata": {
3609 "name": name,
3610 "namespace": namespace,
3611 "labels": {
3612 "app": name,
3613 "project": "gco",
3614 "gco.io/type": "inference",
3615 },
3616 },
3617 "spec": {
3618 "scaleTargetRef": {"name": target_name},
3619 "minReplicaCount": min_replicas,
3620 "maxReplicaCount": max_replicas,
3621 "triggers": self._build_keda_triggers(metrics_config, target_name, namespace),
3622 },
3623 }
3625 custom = client.CustomObjectsApi()
3626 try:
3627 custom.create_namespaced_custom_object(
3628 group=KEDA_API_GROUP,
3629 version=KEDA_API_VERSION,
3630 namespace=namespace,
3631 plural=KEDA_SCALEDOBJECT_PLURAL,
3632 body=body,
3633 _request_timeout=self._k8s_timeout,
3634 )
3635 logger.info(
3636 "Created KEDA ScaledObject %s targeting %s (min=%d, max=%d)",
3637 name,
3638 target_name,
3639 min_replicas,
3640 max_replicas,
3641 )
3642 except ApiException as e:
3643 if e.status == 409: 3643 ↛ 3655line 3643 didn't jump to line 3655 because the condition on line 3643 was always true
3644 custom.patch_namespaced_custom_object(
3645 group=KEDA_API_GROUP,
3646 version=KEDA_API_VERSION,
3647 namespace=namespace,
3648 plural=KEDA_SCALEDOBJECT_PLURAL,
3649 name=name,
3650 body=body,
3651 _request_timeout=self._k8s_timeout,
3652 )
3653 logger.info("Updated KEDA ScaledObject %s", name)
3654 else:
3655 raise
3657 def _apply_hpa(
3658 self,
3659 hpa_name: str,
3660 namespace: str,
3661 target_name: str,
3662 min_replicas: int,
3663 max_replicas: int,
3664 metrics_config: list[dict[str, Any]],
3665 ) -> None:
3666 """Create or patch a single autoscaler targeting one Deployment.
3668 Builds a V2 autoscaler that scales ``target_name`` between
3669 ``min_replicas`` and ``max_replicas`` on the given metrics, then creates
3670 it. An already-present autoscaler of the same name is patched in place
3671 rather than duplicated. When the metric set includes a GPU signal the
3672 autoscaler is materialized as a KEDA ScaledObject instead (native HPA
3673 Resource metrics cannot read GPU utilization).
3674 """
3675 if self._metrics_require_keda(metrics_config):
3676 self._apply_scaled_object(
3677 name=hpa_name,
3678 namespace=namespace,
3679 target_name=target_name,
3680 min_replicas=min_replicas,
3681 max_replicas=max_replicas,
3682 metrics_config=metrics_config,
3683 )
3684 return
3686 hpa = client.V2HorizontalPodAutoscaler(
3687 metadata=client.V1ObjectMeta(
3688 name=hpa_name,
3689 namespace=namespace,
3690 labels={
3691 "app": hpa_name,
3692 "project": "gco",
3693 "gco.io/type": "inference",
3694 },
3695 ),
3696 spec=client.V2HorizontalPodAutoscalerSpec(
3697 scale_target_ref=client.V2CrossVersionObjectReference(
3698 api_version="apps/v1",
3699 kind="Deployment",
3700 name=target_name,
3701 ),
3702 min_replicas=min_replicas,
3703 max_replicas=max_replicas,
3704 metrics=self._build_hpa_metrics(metrics_config),
3705 ),
3706 )
3708 autoscaling_v2 = client.AutoscalingV2Api()
3709 try:
3710 autoscaling_v2.create_namespaced_horizontal_pod_autoscaler(namespace, hpa)
3711 logger.info(
3712 "Created HPA %s targeting %s (min=%d, max=%d)",
3713 hpa_name,
3714 target_name,
3715 min_replicas,
3716 max_replicas,
3717 )
3718 except ApiException as e:
3719 if e.status == 409: 3719 ↛ 3723line 3719 didn't jump to line 3723 because the condition on line 3719 was always true
3720 autoscaling_v2.patch_namespaced_horizontal_pod_autoscaler(hpa_name, namespace, hpa)
3721 logger.info("Updated HPA %s", hpa_name)
3722 else:
3723 raise
3725 def _create_or_update_hpa(self, name: str, namespace: str, spec: dict[str, Any]) -> None:
3726 """Create or update a Horizontal Pod Autoscaler for an inference endpoint."""
3727 autoscaling_config = spec.get("autoscaling", {})
3728 if not autoscaling_config.get("enabled"):
3729 return
3731 min_replicas = autoscaling_config.get("min_replicas", 1)
3732 max_replicas = autoscaling_config.get("max_replicas", 10)
3733 metrics_config = autoscaling_config.get("metrics", [{"type": "cpu", "target": 70}])
3735 self._apply_hpa(
3736 hpa_name=name,
3737 namespace=namespace,
3738 target_name=name,
3739 min_replicas=min_replicas,
3740 max_replicas=max_replicas,
3741 metrics_config=metrics_config,
3742 )
3744 def _create_role_hpa(self, name: str, ns: str, spec: dict[str, Any], role: str) -> None:
3745 """Create or update one autoscaler for a single Mooncake role.
3747 When the endpoint's ``mooncake.autoscaling`` block is enabled and
3748 carries a config for this role, this materializes exactly one autoscaler
3749 named ``{name}-{role}`` that scales the matching ``{name}-{role}``
3750 Deployment between the role's ``min_replicas`` and ``max_replicas``. The
3751 role Deployment itself is already materialized at ``min_replicas`` by
3752 :meth:`_replica_count_for_role`, so the autoscaler owns the count from
3753 that lower bound. When autoscaling is absent or disabled, or when the
3754 role carries no config, no autoscaler is created and the role's replicas
3755 stay at their topology value.
3757 Args:
3758 name: The endpoint name.
3759 ns: The namespace the role Deployment lives in.
3760 spec: The endpoint spec; ``spec["mooncake"]["autoscaling"]`` drives
3761 the bounds and metrics.
3762 role: One of ``"prefill"`` or ``"decode"``.
3763 """
3764 mooncake = spec.get("mooncake") or {}
3765 autoscaling = mooncake.get("autoscaling") or {}
3766 if not autoscaling.get("enabled"):
3767 return
3769 role_cfg = autoscaling.get(role)
3770 if not role_cfg:
3771 return
3773 min_replicas = role_cfg.get("min_replicas", 1)
3774 max_replicas = role_cfg.get("max_replicas", 10)
3775 metrics_config = role_cfg.get("metrics", [{"type": "cpu", "target": 70}])
3777 target_name = f"{name}-{role}"
3778 self._apply_hpa(
3779 hpa_name=target_name,
3780 namespace=ns,
3781 target_name=target_name,
3782 min_replicas=min_replicas,
3783 max_replicas=max_replicas,
3784 metrics_config=metrics_config,
3785 )
3787 # ------------------------------------------------------------------
3788 # Metrics
3789 # ------------------------------------------------------------------
3791 def get_metrics(self) -> dict[str, Any]:
3792 return {
3793 "cluster_id": self.cluster_id,
3794 "region": self.region,
3795 "running": self._running,
3796 "reconcile_count": self._reconcile_count,
3797 "errors_count": self._errors_count,
3798 }
3801def create_inference_monitor_from_env() -> InferenceMonitor:
3802 """Create an InferenceMonitor from environment variables."""
3803 cluster_id = os.getenv("CLUSTER_NAME", "unknown-cluster")
3804 region = os.getenv("REGION", "unknown-region")
3805 namespace = os.getenv("INFERENCE_NAMESPACE", "gco-inference")
3806 interval = int(os.getenv("RECONCILE_INTERVAL_SECONDS", "15"))
3808 # Enable structured JSON logging for CloudWatch Insights
3809 configure_structured_logging(
3810 service_name="inference-monitor",
3811 cluster_id=cluster_id,
3812 region=region,
3813 )
3815 store = InferenceEndpointStore() # Uses DYNAMODB_REGION env var, falls back to REGION
3817 return InferenceMonitor(
3818 cluster_id=cluster_id,
3819 region=region,
3820 store=store,
3821 namespace=namespace,
3822 reconcile_interval=interval,
3823 )
3826async def main() -> None:
3827 """Entry point for the inference monitor."""
3828 monitor = create_inference_monitor_from_env()
3829 logger.info("Inference monitor initialized: %s", monitor.get_metrics())
3831 # Expose Prometheus metrics on a dedicated port for the in-cluster
3832 # observability scrape. A scrape-time collector reflects the monitor's live
3833 # counters (reconcile_count, errors_count, running), so no push from the
3834 # reconcile loop is needed.
3835 from gco.services.service_metrics import start_metrics_server
3837 metrics_port = int(os.getenv("METRICS_PORT", "9090"))
3838 start_metrics_server(metrics_port, "inference-monitor", monitor.get_metrics)
3840 while True:
3841 try:
3842 await monitor.start()
3843 except KeyboardInterrupt:
3844 logger.info("Shutting down inference monitor")
3845 monitor.stop()
3846 break
3847 except Exception as e:
3848 logger.error("Monitor crashed, restarting in 10s: %s", e, exc_info=True)
3849 monitor.stop()
3850 monitor._running = False
3851 await asyncio.sleep(10)
3854if __name__ == "__main__":
3855 asyncio.run(main())