Coverage for cli/inference.py: 97.89%

1""" 

2Inference endpoint management for GCO CLI. 

3 

4Provides functionality to deploy, manage, and monitor inference endpoints 

5across multi-region EKS clusters via the DynamoDB-backed reconciliation 

6pattern (inference_monitor). 

7""" 

8 

9from __future__ import annotations 

10 

11import logging 

12from copy import deepcopy 

13from typing import TYPE_CHECKING, Any, Literal, TypedDict, TypeGuard 

14 

15from .aws_client import get_aws_client 

16from .config import GCOConfig, get_config 

17 

18# <pyflowchart-code-diagram> BEGIN - auto-inserted, do not edit 

19# Generated at (UTC): 2026-07-18T01:03:40Z 

20# Flowchart(s) generated from this file: 

21# * ``InferenceManager.deploy`` -> ``diagrams/code_diagrams/cli/inference.InferenceManager_deploy.html`` 

22# (PNG: ``diagrams/code_diagrams/cli/inference.InferenceManager_deploy.png``) 

23# * ``InferenceManager.canary_deploy`` -> ``diagrams/code_diagrams/cli/inference.InferenceManager_canary_deploy.html`` 

24# (PNG: ``diagrams/code_diagrams/cli/inference.InferenceManager_canary_deploy.png``) 

25# Regenerate with ``python diagrams/code_diagrams/generate.py``. 

26# <pyflowchart-code-diagram> END 

27 

28 

29if TYPE_CHECKING: 

30 from gco.services.inference_store import InferenceEndpointStore 

31 

32logger = logging.getLogger(__name__) 

33 

34 

35# --------------------------------------------------------------------------- 

36# Mooncake topology — optional endpoint-spec extension 

37# --------------------------------------------------------------------------- 

38# 

39# An endpoint spec may carry an optional ``mooncake`` block describing 

40# disaggregated prefill/decode (PD) serving and/or a shared KV-cache store. 

41# The block is entirely additive: when it is absent the endpoint reconciles 

42# exactly as it does today — one Deployment and one internal ClusterIP Service 

43# behind the shared authenticated inference route. 

44# 

45# The definitions below describe the shape of that block (the dict written to 

46# DynamoDB and read back by the per-region monitor) and the constant 

47# vocabularies its enumerated fields draw from. Byte-size fields are authored 

48# as base-10 integer decimal strings (see :func:`author_byte_size`) so they 

49# round-trip through DynamoDB without being coerced to ``Decimal`` via a float 

50# literal. 

51 

52#: Serving modes a ``mooncake`` block may declare. 

53#: ``disaggregated`` splits prefill and decode; ``store`` runs a single 

54#: KV-store instance; ``both`` composes the two. 

55MOONCAKE_MODES: frozenset[str] = frozenset({"disaggregated", "store", "both"}) 

56 

57#: KV transfer / store intents. ``rdma`` is the default high-performance 

58#: intent: GCO schedules the pod on EFA and renders vLLM's point-to-point 

59#: ``mooncake_protocol`` as ``efa``. ``tcp`` is the non-EFA fallback. 

60MOONCAKE_TRANSFER_PROTOCOLS: frozenset[str] = frozenset({"rdma", "tcp"}) 

61 

62#: KV-store offload tiers for spilling cache beyond GPU memory. 

63MOONCAKE_OFFLOAD_TIERS: frozenset[str] = frozenset({"cpu", "disk", "none"}) 

64 

65#: PD proxy request-scheduling strategies supported today. 

66MOONCAKE_PROXY_SCHEDULING: frozenset[str] = frozenset({"round_robin"}) 

67 

68#: Inclusive bounds for per-role replica counts in an XpYd topology. 

69MOONCAKE_TOPOLOGY_MIN: int = 1 

70MOONCAKE_TOPOLOGY_MAX: int = 1000 

71 

72#: Inclusive bounds for byte-size fields. The ceiling is the signed 64-bit 

73#: maximum; authoring sizes as decimal strings in ``[MIN, MAX]`` keeps them out 

74#: of float/Decimal coercion when they round-trip through DynamoDB. 

75MOONCAKE_BYTE_SIZE_MIN: int = 0 

76MOONCAKE_BYTE_SIZE_MAX: int = 9223372036854775807 

77 

78#: Transfer-engine defaults mirroring Mooncake's reference configuration. 

79MOONCAKE_DEFAULT_BOOTSTRAP_BASE_PORT: int = 8998 

80MOONCAKE_DEFAULT_NUM_WORKERS: int = 10 

81MOONCAKE_DEFAULT_ABORT_REQUEST_TIMEOUT: int = 480 

82 

83 

84class MooncakeTopology(TypedDict): 

85 """An XpYd topology: ``prefill`` (X) and ``decode`` (Y) instance counts.""" 

86 

87 prefill: int 

88 decode: int 

89 

90 

91class MooncakeStoreConfig(TypedDict, total=False): 

92 """KV-cache store pool configuration. 

93 

94 ``global_segment_size`` and ``local_buffer_size`` are byte counts authored 

95 as base-10 integer decimal strings. ``cold_tier_enabled`` opts this 

96 endpoint into the asynchronous, per-region object-store cold tier; the 

97 cold-tier bucket is resolved by the monitor from regional configuration and 

98 is never a user-typed URI. 

99 """ 

100 

101 enabled: bool 

102 metadata_server: str 

103 master_server_address: str 

104 protocol: Literal["rdma", "tcp"] 

105 device_name: str 

106 global_segment_size: str 

107 local_buffer_size: str 

108 offload: Literal["cpu", "disk", "none"] 

109 cold_tier_enabled: bool 

110 

111 

112class MooncakeTransferConfig(TypedDict, total=False): 

113 """RDMA/TCP transfer-engine configuration for KV cache movement.""" 

114 

115 protocol: Literal["rdma", "tcp"] 

116 device_name: str 

117 num_workers: int 

118 bootstrap_base_port: int 

119 abort_request_timeout: int 

120 

121 

122class MooncakeProxyConfig(TypedDict, total=False): 

123 """PD proxy configuration. ``admin_api_key_secret`` names the Kubernetes 

124 Secret holding the proxy admin key; the key value is never carried on the 

125 endpoint spec.""" 

126 

127 image: str 

128 scheduling: Literal["round_robin"] 

129 admin_api_key_secret: str 

130 

131 

132class MooncakeRoleAutoscaling(TypedDict, total=False): 

133 """Per-role autoscaling bounds and metrics for one of prefill/decode.""" 

134 

135 min_replicas: int 

136 max_replicas: int 

137 metrics: list[dict[str, Any]] 

138 

139 

140class MooncakeAutoscalingConfig(TypedDict, total=False): 

141 """Optional per-role pod autoscaling. When absent the topology is static.""" 

142 

143 enabled: bool 

144 prefill: MooncakeRoleAutoscaling 

145 decode: MooncakeRoleAutoscaling 

146 

147 

148class MooncakeSpec(TypedDict, total=False): 

149 """The optional ``mooncake`` block carried on an endpoint spec dict.""" 

150 

151 mode: Literal["disaggregated", "store", "both"] 

152 topology: MooncakeTopology 

153 store: MooncakeStoreConfig 

154 transfer: MooncakeTransferConfig 

155 proxy: MooncakeProxyConfig 

156 autoscaling: MooncakeAutoscalingConfig 

157 

158 

159def author_byte_size(value: int | str) -> str: 

160 """Render a byte-size value as a canonical base-10 integer decimal string. 

161 

162 Mooncake store/transfer sizes (segment size, local buffer) are carried on 

163 the endpoint spec as digit-only strings so they survive the DynamoDB 

164 round-trip without being coerced to ``Decimal`` through a float literal. 

165 

166 Accepts a non-negative ``int`` or a string of base-10 ASCII digits and 

167 returns the same whole number as ``str``. The value must fall in 

168 ``[MOONCAKE_BYTE_SIZE_MIN, MOONCAKE_BYTE_SIZE_MAX]``. Signs, decimal 

169 points, exponents, floats, booleans, and any non-digit text are not 

170 accepted. 

171 

172 Raises: 

173 ValueError: when ``value`` cannot be authored as an in-range base-10 

174 integer. 

175 """ 

176 # ``bool`` is a subclass of ``int``; reject it explicitly so ``True``/``False`` 

177 # never masquerade as 1/0 byte sizes. 

178 if isinstance(value, bool): 

179 raise ValueError(f"byte-size value must be an integer, got bool: {value!r}") 

180 

181 if isinstance(value, int): 

182 size = value 

183 elif isinstance(value, str): 

184 text = value.strip() 

185 if not text or any(ch not in "0123456789" for ch in text): 

186 raise ValueError( 

187 "byte-size value must be a base-10 integer string " 

188 f"(ASCII digits only, no sign, point, or exponent), got {value!r}" 

189 ) 

190 size = int(text) 

191 else: 

192 raise ValueError( 

193 f"byte-size value must be an int or a base-10 digit string, got {type(value).__name__}" 

194 ) 

195 

196 if not MOONCAKE_BYTE_SIZE_MIN <= size <= MOONCAKE_BYTE_SIZE_MAX: 

197 raise ValueError( 

198 "byte-size value out of range " 

199 f"[{MOONCAKE_BYTE_SIZE_MIN}, {MOONCAKE_BYTE_SIZE_MAX}]: {size}" 

200 ) 

201 

202 return str(size) 

203 

204 

205#: Byte-size fields a ``mooncake`` store block may carry. Each is authored as a 

206#: base-10 integer decimal string via :func:`author_byte_size`. 

207_MOONCAKE_STORE_BYTE_SIZE_FIELDS: tuple[str, ...] = ( 

208 "global_segment_size", 

209 "local_buffer_size", 

210) 

211 

212#: Modes that run a split prefill/decode topology and therefore require a 

213#: valid ``topology`` and may carry per-role autoscaling. 

214_MOONCAKE_DISAGGREGATED_MODES: frozenset[str] = frozenset({"disaggregated", "both"}) 

215 

216 

217def _is_plain_int(value: Any) -> TypeGuard[int]: 

218 """True when ``value`` is an ``int`` and not a ``bool``. 

219 

220 ``bool`` is a subclass of ``int``; counts and replica bounds must be real 

221 integers, so ``True``/``False`` are not accepted as 1/0. 

222 """ 

223 return isinstance(value, int) and not isinstance(value, bool) 

224 

225 

226def _validate_role_autoscaling_bounds(role: str, role_block: dict[str, Any]) -> None: 

227 """Validate one role's ``min_replicas``/``max_replicas`` bounds. 

228 

229 Raises :class:`ValueError` naming the violated bound. ``min_replicas`` must 

230 be an integer ``>= 1`` and ``max_replicas`` an integer no smaller than the 

231 effective minimum (which defaults to 1 when ``min_replicas`` is absent). 

232 """ 

233 min_replicas = role_block.get("min_replicas") 

234 max_replicas = role_block.get("max_replicas") 

235 

236 if min_replicas is not None: 

237 if not _is_plain_int(min_replicas): 

238 raise ValueError( 

239 f"mooncake.autoscaling.{role}.min_replicas must be an integer, got {min_replicas!r}" 

240 ) 

241 if min_replicas < 1: 

242 raise ValueError( 

243 f"mooncake.autoscaling.{role}.min_replicas must be >= 1, got {min_replicas}" 

244 ) 

245 

246 if max_replicas is not None: 246 ↛ exitline 246 didn't return from function '_validate_role_autoscaling_bounds' because the condition on line 246 was always true

247 if not _is_plain_int(max_replicas): 

248 raise ValueError( 

249 f"mooncake.autoscaling.{role}.max_replicas must be an integer, got {max_replicas!r}" 

250 ) 

251 effective_min = min_replicas if _is_plain_int(min_replicas) else 1 

252 if max_replicas < effective_min: 

253 raise ValueError( 

254 f"mooncake.autoscaling.{role}.max_replicas ({max_replicas}) " 

255 f"must be >= min_replicas ({effective_min})" 

256 ) 

257 

258 

259def validate_mooncake_spec(mooncake: dict[str, Any]) -> None: 

260 """Validate a ``mooncake`` endpoint-spec block, failing fast. 

261 

262 Raises :class:`ValueError` on the first rejected field, naming the 

263 offending field so the caller can correct it. The check is pure — it reads 

264 nothing and writes nothing — so a caller that validates before persisting 

265 leaves any previously stored spec untouched when a block is rejected. 

266 

267 The rules enforced here are: 

268 

269 * ``mode`` must be one of the supported serving modes 

270 (:data:`MOONCAKE_MODES`). 

271 * ``transfer`` must be a mapping when present; ``protocol`` must be one of 

272 :data:`MOONCAKE_TRANSFER_PROTOCOLS` and ``device_name`` must be a string 

273 (the empty string requests automatic interface detection). 

274 * Store byte-size fields must author as in-range base-10 integers. 

275 * ``disaggregated``/``both`` modes require integer ``topology.prefill`` and 

276 ``topology.decode`` in 

277 ``[MOONCAKE_TOPOLOGY_MIN, MOONCAKE_TOPOLOGY_MAX]``. 

278 * ``store.cold_tier_enabled`` may be true only while ``store.enabled`` is 

279 true (the cold tier extends the hot store). 

280 * Autoscaling may be enabled only for ``disaggregated``/``both`` modes, and 

281 each present role's ``min_replicas``/``max_replicas`` must satisfy 

282 ``min_replicas >= 1`` and ``max_replicas >= min_replicas``. 

283 """ 

284 if not isinstance(mooncake, dict): 

285 raise ValueError("mooncake block must be a mapping") 

286 

287 mode = mooncake.get("mode") 

288 if mode not in MOONCAKE_MODES: 

289 allowed = ", ".join(sorted(MOONCAKE_MODES)) 

290 raise ValueError(f"mooncake.mode must be one of {{{allowed}}}, got {mode!r}") 

291 

292 store = mooncake.get("store") 

293 if store is not None and not isinstance(store, dict): 

294 raise ValueError("mooncake.store must be a mapping") 

295 

296 transfer = mooncake.get("transfer") 

297 if transfer is not None and not isinstance(transfer, dict): 

298 raise ValueError("mooncake.transfer must be a mapping") 

299 if isinstance(transfer, dict): 

300 protocol = transfer.get("protocol", "rdma") 

301 if protocol not in MOONCAKE_TRANSFER_PROTOCOLS: 

302 allowed = ", ".join(sorted(MOONCAKE_TRANSFER_PROTOCOLS)) 

303 raise ValueError( 

304 f"mooncake.transfer.protocol must be one of {{{allowed}}}, got {protocol!r}" 

305 ) 

306 device_name = transfer.get("device_name", "") 

307 if not isinstance(device_name, str): 

308 raise ValueError(f"mooncake.transfer.device_name must be a string, got {device_name!r}") 

309 

310 # Byte-size fields must author cleanly; surface the offending field name. 

311 if isinstance(store, dict): 

312 for field in _MOONCAKE_STORE_BYTE_SIZE_FIELDS: 

313 if field in store: 

314 try: 

315 author_byte_size(store[field]) 

316 except ValueError as exc: 

317 raise ValueError(f"mooncake.store.{field}: {exc}") from exc 

318 

319 # Split topologies need integer prefill/decode counts in range. 

320 if mode in _MOONCAKE_DISAGGREGATED_MODES: 

321 topology = mooncake.get("topology") 

322 if not isinstance(topology, dict): 

323 raise ValueError( 

324 f"mooncake.topology is required for mode {mode!r} with integer " 

325 "'prefill' and 'decode' counts" 

326 ) 

327 for field in ("prefill", "decode"): 

328 count = topology.get(field) 

329 if not _is_plain_int(count): 

330 raise ValueError( 

331 f"mooncake.topology.{field} must be an integer in " 

332 f"[{MOONCAKE_TOPOLOGY_MIN}, {MOONCAKE_TOPOLOGY_MAX}], " 

333 f"got {count!r}" 

334 ) 

335 if not MOONCAKE_TOPOLOGY_MIN <= count <= MOONCAKE_TOPOLOGY_MAX: 

336 raise ValueError( 

337 f"mooncake.topology.{field} out of range " 

338 f"[{MOONCAKE_TOPOLOGY_MIN}, {MOONCAKE_TOPOLOGY_MAX}]: {count}" 

339 ) 

340 

341 # The cold tier extends the hot store; it cannot be enabled on its own. 

342 if ( 

343 isinstance(store, dict) 

344 and store.get("cold_tier_enabled") is True 

345 and store.get("enabled") is not True 

346 ): 

347 raise ValueError( 

348 "mooncake.store.cold_tier_enabled requires mooncake.store.enabled to be true" 

349 ) 

350 

351 autoscaling = mooncake.get("autoscaling") 

352 if autoscaling is not None: 

353 if not isinstance(autoscaling, dict): 

354 raise ValueError("mooncake.autoscaling must be a mapping") 

355 if autoscaling.get("enabled") is True and mode not in _MOONCAKE_DISAGGREGATED_MODES: 

356 raise ValueError( 

357 "mooncake.autoscaling.enabled requires a 'disaggregated' or " 

358 f"'both' mode, got {mode!r}" 

359 ) 

360 for role in ("prefill", "decode"): 

361 role_block = autoscaling.get(role) 

362 if role_block is None: 

363 continue 

364 if not isinstance(role_block, dict): 

365 raise ValueError(f"mooncake.autoscaling.{role} must be a mapping") 

366 _validate_role_autoscaling_bounds(role, role_block) 

367 

368 

369class InferenceManager: 

370 """Manages inference endpoints via the DynamoDB store.""" 

371 

372 def __init__(self, config: GCOConfig | None = None): 

373 self.config = config or get_config() 

374 self._aws_client = get_aws_client(config) 

375 

376 def _get_store(self, region: str | None = None) -> InferenceEndpointStore: 

377 """Get an InferenceEndpointStore for the global region.""" 

378 from gco.services.inference_store import InferenceEndpointStore 

379 

380 # Use the global region for DynamoDB (same as job store) 

381 store_region = region or self.config.global_region 

382 return InferenceEndpointStore(region=store_region) 

383 

384 def _build_mooncake_block( 

385 self, 

386 *, 

387 mode: str, 

388 prefill_replicas: int, 

389 decode_replicas: int, 

390 store: dict[str, Any] | None, 

391 transfer: dict[str, Any] | None, 

392 proxy: dict[str, Any] | None, 

393 autoscaling: dict[str, Any] | None, 

394 default_proxy_image: str | None = None, 

395 ) -> dict[str, Any]: 

396 """Assemble and validate an optional ``spec.mooncake`` block. 

397 

398 Composes the topology and any supplied store/transfer/proxy/autoscaling 

399 sub-blocks into a single mapping, authoring store byte-size fields as 

400 base-10 integer decimal strings so they round-trip through DynamoDB, 

401 then validates the result. Validation is pure and runs before the 

402 caller persists anything, so a rejected block leaves any previously 

403 stored spec untouched. Raises :class:`ValueError` — naming the offending 

404 field — when the mode is unsupported or any field is invalid. 

405 

406 The store-bearing modes (``store`` and ``both``) default the store to 

407 enabled so the shared master address is wired in (the ``both``-mode 

408 MultiConnector's store half depends on it), and split modes 

409 (``disaggregated`` and ``both``) default the prefill-decode proxy image 

410 to ``default_proxy_image`` when the caller supplies no explicit proxy 

411 image. 

412 """ 

413 block: dict[str, Any] = {"mode": mode} 

414 

415 # Split modes carry an XpYd topology; a single-instance store does not. 

416 if mode in _MOONCAKE_DISAGGREGATED_MODES: 

417 block["topology"] = { 

418 "prefill": prefill_replicas, 

419 "decode": decode_replicas, 

420 } 

421 

422 # The store-bearing modes (store and both) only function with the KV 

423 # store enabled: the both-mode MultiConnector's store half is wired to 

424 # the shared master address, which the monitor renders only for an 

425 # enabled store. So a store block is always present for those modes, 

426 # defaulting enabled to True; an explicit store block still tunes 

427 # offload, sizes, and the cold tier. 

428 store_block = dict(store) if store is not None else None 

429 if mode in ("store", "both"): 

430 store_block = dict(store_block or {}) 

431 store_block.setdefault("enabled", True) 

432 if store_block is not None: 

433 # Author byte-size fields as canonical decimal strings up front so 

434 # the persisted spec round-trips through DynamoDB without float or 

435 # Decimal coercion. Authoring also fails fast on bad inputs. 

436 for field in _MOONCAKE_STORE_BYTE_SIZE_FIELDS: 

437 if field in store_block: 

438 try: 

439 store_block[field] = author_byte_size(store_block[field]) 

440 except ValueError as exc: 

441 raise ValueError(f"mooncake.store.{field}: {exc}") from exc 

442 block["store"] = store_block 

443 

444 if transfer is not None: 

445 block["transfer"] = dict(transfer) 

446 

447 # Split modes are fronted by the prefill-decode proxy, which needs a 

448 # container image. Default it to the same image the role pods serve from 

449 # (the upstream vLLM image bundles the reference proxy) so a split deploy 

450 # stands up without a separate proxy image; an explicit proxy image 

451 # still wins. 

452 proxy_block = dict(proxy) if proxy is not None else None 

453 if mode in _MOONCAKE_DISAGGREGATED_MODES and default_proxy_image: 

454 proxy_block = dict(proxy_block or {}) 

455 proxy_block.setdefault("image", default_proxy_image) 

456 if proxy_block is not None: 

457 block["proxy"] = proxy_block 

458 

459 if autoscaling is not None: 

460 block["autoscaling"] = dict(autoscaling) 

461 

462 # Fail fast before persisting: rejects unsupported modes (naming the 

463 # allowed values) and every other invalid field. 

464 validate_mooncake_spec(block) 

465 return block 

466 

467 def deploy( 

468 self, 

469 endpoint_name: str, 

470 image: str | None = None, 

471 target_regions: list[str] | None = None, 

472 replicas: int = 1, 

473 gpu_count: int = 1, 

474 gpu_type: str | None = None, 

475 port: int = 8000, 

476 model_path: str | None = None, 

477 model_source: str | None = None, 

478 health_check_path: str = "/health", 

479 env: dict[str, str] | None = None, 

480 namespace: str = "gco-inference", 

481 labels: dict[str, str] | None = None, 

482 autoscaling: dict[str, Any] | None = None, 

483 capacity_type: str | None = None, 

484 extra_args: list[str] | None = None, 

485 accelerator: str = "nvidia", 

486 node_selector: dict[str, str] | None = None, 

487 rewrite_image: bool = True, 

488 *, 

489 mooncake_mode: str | None = None, 

490 prefill_replicas: int = 1, 

491 decode_replicas: int = 1, 

492 mooncake_store: dict[str, Any] | None = None, 

493 mooncake_transfer: dict[str, Any] | None = None, 

494 mooncake_proxy: dict[str, Any] | None = None, 

495 mooncake_autoscaling: dict[str, Any] | None = None, 

496 ) -> dict[str, Any]: 

497 """ 

498 Deploy an inference endpoint to one or more regions. 

499 

500 The endpoint spec is written to DynamoDB. The inference_monitor 

501 in each target region picks it up and creates the K8s resources. 

502 

503 Args: 

504 endpoint_name: Unique name for the endpoint 

505 image: Container image (e.g. vllm/vllm-openai:v0.26.0). Optional 

506 when ``mooncake_mode`` is set: a disaggregated/store deploy 

507 with no image falls back to the default upstream 

508 Mooncake-enabled vLLM image. A plain deploy still requires an 

509 image. 

510 target_regions: Regions to deploy to (default: all deployed regions) 

511 replicas: Number of replicas per region 

512 gpu_count: GPUs per replica 

513 gpu_type: GPU instance type hint for node selector 

514 port: Container port 

515 model_path: EFS path for model weights 

516 health_check_path: Health check endpoint path 

517 env: Environment variables 

518 namespace: Kubernetes namespace 

519 labels: Labels for the endpoint 

520 rewrite_image: When True (the default), rewrite ECR URIs in 

521 ``image`` to target each region's local replica. Non-ECR 

522 refs (Docker Hub, GHCR, etc.) are left unchanged. When 

523 False, the URI is written verbatim to every region's 

524 spec — the operator is responsible for cross-region 

525 pulls. Per-region rewrites are stored under a 

526 ``region_overrides`` map on the spec keyed by region. 

527 mooncake_mode: When set to one of ``disaggregated``, ``store``, 

528 or ``both``, build and persist a ``spec.mooncake`` block for 

529 disaggregated prefill/decode serving and/or a shared KV-cache 

530 store. An unsupported value is rejected before anything is 

531 persisted. 

532 prefill_replicas: X in an XpYd topology — prefill instance count 

533 for split (``disaggregated``/``both``) modes. 

534 decode_replicas: Y in an XpYd topology — decode instance count for 

535 split modes. 

536 mooncake_store: Optional KV-store pool configuration merged into 

537 ``spec.mooncake.store``. Byte-size fields are authored as 

538 base-10 integer decimal strings so they round-trip through 

539 DynamoDB. 

540 mooncake_transfer: Optional Mooncake transfer intent and network 

541 device. ``protocol`` accepts ``rdma`` (the default; scheduled 

542 on EFA and rendered to vLLM as ``mooncake_protocol=efa``) or 

543 ``tcp`` (no EFA placement). ``device_name`` is forwarded to 

544 both the connector and mounted Mooncake configuration; an 

545 empty string lets Mooncake auto-detect it. 

546 mooncake_proxy: Optional PD proxy configuration merged into 

547 ``spec.mooncake.proxy``. 

548 mooncake_autoscaling: Optional per-role autoscaling configuration 

549 merged into ``spec.mooncake.autoscaling``. 

550 

551 Returns: 

552 Created endpoint record 

553 """ 

554 # Build the optional mooncake block first and validate it before any 

555 # persistence so a rejected block leaves any stored spec untouched. A 

556 # disaggregated/store deploy without an explicit image falls back to 

557 # the default upstream Mooncake-enabled vLLM image. 

558 mooncake_block: dict[str, Any] | None = None 

559 if mooncake_mode is not None: 

560 # Resolve the image before building the block so a split mode's 

561 # prefill-decode proxy can default to the same image the role pods 

562 # serve from (the upstream vLLM image bundles the reference proxy). 

563 if image is None: 

564 from .images import default_disaggregated_image 

565 

566 image = default_disaggregated_image(config=self.config) 

567 mooncake_block = self._build_mooncake_block( 

568 mode=mooncake_mode, 

569 prefill_replicas=prefill_replicas, 

570 decode_replicas=decode_replicas, 

571 store=mooncake_store, 

572 transfer=mooncake_transfer, 

573 proxy=mooncake_proxy, 

574 autoscaling=mooncake_autoscaling, 

575 default_proxy_image=image, 

576 ) 

577 

578 if image is None: 

579 raise ValueError( 

580 "an image is required (pass image, or set mooncake_mode to use " 

581 "the default upstream Mooncake-enabled vLLM image)" 

582 ) 

583 

584 if not target_regions: 

585 stacks = self._aws_client.discover_regional_stacks() 

586 target_regions = list(stacks.keys()) 

587 if not target_regions: 

588 raise ValueError("No deployed regions found. Deploy infrastructure first.") 

589 

590 # Per-region image-URI rewrites for ECR refs. Each target region 

591 # gets the local replica's URI on its own spec, so the 

592 # inference_monitor's pod-spec materialiser pulls in-region 

593 # rather than across the WAN. Non-ECR URIs come back unchanged 

594 # from the helper, so this is a no-op for Docker Hub / GHCR refs. 

595 # 

596 # The helper lives in ``cli._image_uri`` rather than ``cli.images`` 

597 # so this import doesn't create a module-level cycle: 

598 # ``cli.images`` itself imports the same helper. ``cli._image_uri`` 

599 # is a leaf module with no project-side dependencies. 

600 region_image_map: dict[str, str] = {} 

601 if rewrite_image: 

602 from ._image_uri import rewrite_image_uri_for_region 

603 

604 for region in target_regions: 

605 region_image_map[region] = rewrite_image_uri_for_region(image, region) 

606 

607 spec = { 

608 "image": image, 

609 "port": port, 

610 "replicas": replicas, 

611 "gpu_count": gpu_count, 

612 "health_check_path": health_check_path, 

613 } 

614 # Preserve the rewrite map on the spec so the inference_monitor 

615 # service can pick the right URI per region when materialising 

616 # pods. When ``rewrite_image=False`` no map is set and the flat 

617 # ``image`` field is the only source. 

618 if region_image_map and any(uri != image for uri in region_image_map.values()): 

619 spec["region_image_uris"] = region_image_map 

620 if gpu_type: 

621 spec["gpu_type"] = gpu_type 

622 if model_path: 

623 spec["model_path"] = model_path 

624 if model_source: 

625 spec["model_source"] = model_source 

626 if env: 

627 spec["env"] = env 

628 if autoscaling: 

629 spec["autoscaling"] = autoscaling 

630 if capacity_type: 

631 spec["capacity_type"] = capacity_type 

632 if extra_args: 

633 spec["args"] = extra_args 

634 if accelerator != "nvidia": 

635 spec["accelerator"] = accelerator 

636 if node_selector: 

637 spec["node_selector"] = node_selector 

638 if mooncake_block is not None: 

639 spec["mooncake"] = mooncake_block 

640 

641 store = self._get_store() 

642 result: dict[str, Any] = store.create_endpoint( 

643 endpoint_name=endpoint_name, 

644 spec=spec, 

645 target_regions=target_regions, 

646 namespace=namespace, 

647 labels=labels, 

648 ) 

649 return result 

650 

651 def list_endpoints( 

652 self, 

653 desired_state: str | None = None, 

654 region: str | None = None, 

655 ) -> list[dict[str, Any]]: 

656 """List all inference endpoints.""" 

657 store = self._get_store() 

658 result: list[dict[str, Any]] = store.list_endpoints( 

659 desired_state=desired_state, 

660 target_region=region, 

661 ) 

662 return result 

663 

664 def get_endpoint(self, endpoint_name: str) -> dict[str, Any] | None: 

665 """Get details of a specific endpoint.""" 

666 store = self._get_store() 

667 result: dict[str, Any] | None = store.get_endpoint(endpoint_name) 

668 return result 

669 

670 def scale(self, endpoint_name: str, replicas: int) -> dict[str, Any] | None: 

671 """Scale an endpoint to a new replica count.""" 

672 store = self._get_store() 

673 result: dict[str, Any] | None = store.scale_endpoint(endpoint_name, replicas) 

674 return result 

675 

676 def set_topology( 

677 self, 

678 endpoint_name: str, 

679 prefill: int, 

680 decode: int, 

681 ) -> dict[str, Any] | None: 

682 """Resize a disaggregated endpoint's prefill/decode topology. 

683 

684 Updates ``spec.mooncake.topology`` to the new XpYd counts and 

685 re-triggers reconciliation (via :meth:`InferenceEndpointStore.update_spec`, 

686 which flips ``desired_state`` to ``deploying``) so the per-region 

687 monitor adjusts the prefill and decode role replica counts. 

688 

689 Both counts must be integers in the inclusive range 

690 ``[MOONCAKE_TOPOLOGY_MIN, MOONCAKE_TOPOLOGY_MAX]``. The counts are 

691 validated before anything is read or written, so a rejected request 

692 names the offending count and leaves the stored topology and 

693 ``desired_state`` untouched. 

694 

695 Args: 

696 endpoint_name: Name of the disaggregated endpoint to resize. 

697 prefill: New prefill (X) instance count. 

698 decode: New decode (Y) instance count. 

699 

700 Returns: 

701 The updated endpoint record, or ``None`` when no endpoint with 

702 ``endpoint_name`` exists. 

703 

704 Raises: 

705 ValueError: when ``prefill`` or ``decode`` is not an integer in 

706 ``[MOONCAKE_TOPOLOGY_MIN, MOONCAKE_TOPOLOGY_MAX]``. 

707 """ 

708 # Validate before any read or write so a bad count names the offending 

709 # field and leaves the stored topology and desired_state unchanged. 

710 for field, count in (("prefill", prefill), ("decode", decode)): 

711 if not _is_plain_int(count): 

712 raise ValueError( 

713 f"topology {field} count must be an integer in " 

714 f"[{MOONCAKE_TOPOLOGY_MIN}, {MOONCAKE_TOPOLOGY_MAX}], " 

715 f"got {count!r}" 

716 ) 

717 if not MOONCAKE_TOPOLOGY_MIN <= count <= MOONCAKE_TOPOLOGY_MAX: 

718 raise ValueError( 

719 f"topology {field} count out of range " 

720 f"[{MOONCAKE_TOPOLOGY_MIN}, {MOONCAKE_TOPOLOGY_MAX}]: {count}" 

721 ) 

722 

723 store = self._get_store() 

724 endpoint = store.get_endpoint(endpoint_name) 

725 if not endpoint: 

726 return None 

727 

728 spec = endpoint.get("spec", {}) 

729 # Preserve any existing mooncake sub-fields and replace only the 

730 # topology counts. 

731 mooncake = dict(spec.get("mooncake") or {}) 

732 mooncake["topology"] = {"prefill": prefill, "decode": decode} 

733 spec["mooncake"] = mooncake 

734 

735 result: dict[str, Any] | None = store.update_spec(endpoint_name, spec) 

736 return result 

737 

738 def configure_store( 

739 self, 

740 endpoint_name: str, 

741 store_config: dict[str, Any], 

742 ) -> dict[str, Any] | None: 

743 """Update an endpoint's KV-cache store configuration. 

744 

745 Merges ``store_config`` into ``spec.mooncake.store`` and re-triggers 

746 reconciliation (via :meth:`InferenceEndpointStore.update_spec`, which 

747 flips ``desired_state`` to ``deploying``) so the per-region monitor 

748 picks up the new store settings. 

749 

750 Store byte-size fields are authored as base-10 integer decimal strings 

751 (so they round-trip through DynamoDB without float/Decimal coercion) 

752 and the resulting ``mooncake`` block is validated before anything is 

753 written. A rejected configuration names the offending field and leaves 

754 the stored spec untouched. 

755 

756 Args: 

757 endpoint_name: Name of the endpoint to reconfigure. 

758 store_config: KV-store pool settings merged into 

759 ``spec.mooncake.store``. 

760 

761 Returns: 

762 The updated endpoint record, or ``None`` when no endpoint with 

763 ``endpoint_name`` exists. 

764 

765 Raises: 

766 ValueError: when the resulting ``mooncake`` block is invalid (for 

767 example an out-of-range byte-size field). 

768 """ 

769 store = self._get_store() 

770 endpoint = store.get_endpoint(endpoint_name) 

771 if not endpoint: 

772 return None 

773 

774 spec = endpoint.get("spec", {}) 

775 # Preserve any existing mooncake sub-fields and replace only the store 

776 # block, authoring byte-size fields as canonical decimal strings. 

777 mooncake = dict(spec.get("mooncake") or {}) 

778 store_block = dict(store_config) 

779 for field in _MOONCAKE_STORE_BYTE_SIZE_FIELDS: 

780 if field in store_block: 

781 try: 

782 store_block[field] = author_byte_size(store_block[field]) 

783 except ValueError as exc: 

784 raise ValueError(f"mooncake.store.{field}: {exc}") from exc 

785 mooncake["store"] = store_block 

786 spec["mooncake"] = mooncake 

787 

788 # Fail fast before persisting so a rejected block leaves the stored 

789 # spec untouched. 

790 validate_mooncake_spec(mooncake) 

791 

792 result: dict[str, Any] | None = store.update_spec(endpoint_name, spec) 

793 return result 

794 

795 def stop(self, endpoint_name: str) -> dict[str, Any] | None: 

796 """Stop an endpoint (scale to zero, keep resources).""" 

797 store = self._get_store() 

798 result: dict[str, Any] | None = store.update_desired_state(endpoint_name, "stopped") 

799 return result 

800 

801 def start(self, endpoint_name: str) -> dict[str, Any] | None: 

802 """Start a stopped endpoint.""" 

803 store = self._get_store() 

804 result: dict[str, Any] | None = store.update_desired_state(endpoint_name, "running") 

805 return result 

806 

807 def delete(self, endpoint_name: str) -> dict[str, Any] | None: 

808 """Mark an endpoint for deletion (inference_monitor cleans up).""" 

809 store = self._get_store() 

810 result: dict[str, Any] | None = store.update_desired_state(endpoint_name, "deleted") 

811 return result 

812 

813 def update_image(self, endpoint_name: str, image: str) -> dict[str, Any] | None: 

814 """Update the container image for an endpoint.""" 

815 if not isinstance(image, str) or not image.strip(): 

816 raise ValueError("Image must be a non-empty string") 

817 

818 store = self._get_store() 

819 endpoint = store.get_endpoint(endpoint_name) 

820 if not endpoint: 

821 return None 

822 raw_spec = endpoint.get("spec") 

823 if not isinstance(raw_spec, dict): 823 ↛ 824line 823 didn't jump to line 824 because the condition on line 823 was never true

824 raise ValueError(f"Endpoint '{endpoint_name}' has an invalid spec") 

825 spec = deepcopy(raw_spec) 

826 spec["image"] = image.strip() 

827 # A direct image update is global. Stale regional rewrites would take 

828 # precedence in the monitor and silently keep serving the old image. 

829 spec.pop("region_image_uris", None) 

830 result: dict[str, Any] | None = store.update_spec(endpoint_name, spec) 

831 return result 

832 

833 def add_region(self, endpoint_name: str, region: str) -> dict[str, Any] | None: 

834 """Add a region to an existing endpoint.""" 

835 from datetime import UTC, datetime 

836 

837 store = self._get_store() 

838 endpoint = store.get_endpoint(endpoint_name) 

839 if not endpoint: 

840 return None 

841 regions = endpoint.get("target_regions", []) 

842 if region not in regions: 

843 regions.append(region) 

844 # Update via raw DynamoDB update 

845 try: 

846 response = store._table.update_item( 

847 Key={"endpoint_name": endpoint_name}, 

848 UpdateExpression="SET target_regions = :r, updated_at = :u", 

849 ExpressionAttributeValues={ 

850 ":r": regions, 

851 ":u": datetime.now(UTC).isoformat(), 

852 }, 

853 ReturnValues="ALL_NEW", 

854 ) 

855 result: dict[str, Any] | None = response.get("Attributes") 

856 return result 

857 except Exception as e: 

858 logger.error("Failed to add region: %s", e) 

859 return None 

860 

861 def remove_region(self, endpoint_name: str, region: str) -> dict[str, Any] | None: 

862 """Remove a region from an existing endpoint.""" 

863 store = self._get_store() 

864 endpoint = store.get_endpoint(endpoint_name) 

865 if not endpoint: 

866 return None 

867 regions = endpoint.get("target_regions", []) 

868 if region in regions: 

869 regions.remove(region) 

870 try: 

871 from datetime import UTC, datetime 

872 

873 response = store._table.update_item( 

874 Key={"endpoint_name": endpoint_name}, 

875 UpdateExpression="SET target_regions = :r, updated_at = :u", 

876 ExpressionAttributeValues={ 

877 ":r": regions, 

878 ":u": datetime.now(UTC).isoformat(), 

879 }, 

880 ReturnValues="ALL_NEW", 

881 ) 

882 result: dict[str, Any] | None = response.get("Attributes") 

883 return result 

884 except Exception as e: 

885 logger.error("Failed to remove region: %s", e)