diff --git a/.github/codeowners/areas.yaml b/.github/codeowners/areas.yaml index 1194e6730eec..61efd2cf27b0 100644 --- a/.github/codeowners/areas.yaml +++ b/.github/codeowners/areas.yaml @@ -19,6 +19,7 @@ areas: - lib/llm/src/grpc/protos/ - lib/llm/src/grpc/service.rs - lib/llm/src/grpc/service/ + - lib/http-server/ - lib/llm/src/http.rs - lib/llm/src/http/ - lib/llm/src/http/service.rs @@ -73,10 +74,14 @@ areas: path_globs: - lib/global-view-broadcaster/ - components/src/dynamo/pool_relay/ + - container/global-discovery-service.yaml - lib/global-view-receiver/ - docs/fern/pages/developer-guide/knowledge-base/modular-components/router/multi-dc-kv-routing.md + - docs/fern/pages/developer-guide/**/multi-cluster-services-registry.md - lib/global-*/** + - docs/fern/pages/kubernetes/kv-aware-routing/multi-cluster-services.md - docs/fern/pages/kubernetes/kv-aware-routing/pool-relay.md + - docs/fern/pages/reference/components/global-discovery-mcs-configuration.md - docs/fern/pages/reference/components/pool-relay-configuration.md - lib/bindings/python/rust/llm/pool_relay.rs - lib/llm/src/pool_relay.rs @@ -1085,6 +1090,10 @@ shared: owners: - frontend - observability +- glob: lib/http-server/src/metrics.rs + owners: + - frontend + - observability - glob: lib/llm/src/kv_router/shared_cache.rs owners: - router @@ -1284,6 +1293,10 @@ shared: owners: - operator - backend-vllm +- glob: container/Dockerfile.global-discovery-service + owners: + - global-router + - ops - glob: deploy/operator/Dockerfile owners: - operator diff --git a/CODEOWNERS b/CODEOWNERS index 50d852d90af4..9ad8af2848ba 100644 --- a/CODEOWNERS +++ b/CODEOWNERS @@ -67,6 +67,7 @@ * @ai-dynamo/dynamo-runtime-codeowners # === frontend (@ai-dynamo/dynamo-frontend-codeowners) === +/lib/http-server/ @ai-dynamo/dynamo-frontend-codeowners /examples/chat_templates/ @ai-dynamo/dynamo-frontend-codeowners /components/src/dynamo/frontend/ @ai-dynamo/dynamo-frontend-codeowners /components/src/dynamo/frontend/tests/ @ai-dynamo/dynamo-frontend-codeowners @@ -94,6 +95,7 @@ /lib/global-view-receiver/ @ai-dynamo/dynamo-global-router-codeowners /lib/global-view-broadcaster/ @ai-dynamo/dynamo-global-router-codeowners /components/src/dynamo/pool_relay/ @ai-dynamo/dynamo-global-router-codeowners +/container/global-discovery-service.yaml @ai-dynamo/dynamo-global-router-codeowners # === backend-vllm (@ai-dynamo/dynamo-backend-vllm-codeowners) === /lib/sidecar/vllm/ @ai-dynamo/dynamo-backend-vllm-codeowners @@ -600,7 +602,10 @@ requirements.txt /components/src/dynamo/sglang/request_handlers/image_diffusion/ @ai-dynamo/dynamo-diffusion-codeowners /recipes/kustomize/components/dynamo-openapi/dynamo-openapi.json @ai-dynamo/dynamo-operator-codeowners /docs/fern/pages/reference/components/pool-relay-configuration.md @ai-dynamo/dynamo-global-router-codeowners +/docs/fern/pages/developer-guide/**/multi-cluster-services-registry.md @ai-dynamo/dynamo-global-router-codeowners +/docs/fern/pages/kubernetes/kv-aware-routing/multi-cluster-services.md @ai-dynamo/dynamo-global-router-codeowners /recipes/templates/kustomize/components/dynamo-openapi/dynamo-openapi.json @ai-dynamo/dynamo-operator-codeowners +/docs/fern/pages/reference/components/global-discovery-mcs-configuration.md @ai-dynamo/dynamo-global-router-codeowners /docs/fern/pages/developer-guide/knowledge-base/modular-components/backends/triton/ @ai-dynamo/dynamo-backend-triton-codeowners /docs/fern/pages/developer-guide/knowledge-base/modular-components/router/multi-dc-kv-routing.md @ai-dynamo/dynamo-global-router-codeowners @@ -663,6 +668,7 @@ requirements.txt /examples/diffusers/Dockerfile @ai-dynamo/dynamo-diffusion-codeowners @ai-dynamo/dynamo-ops-codeowners /lib/kvbm-logical/src/metrics/ @ai-dynamo/dynamo-kv-memory-codeowners @ai-dynamo/dynamo-observability-codeowners /tests/serve/test_vllm_omni.py @ai-dynamo/dynamo-runtime-codeowners @ai-dynamo/dynamo-backend-vllm-codeowners +/lib/http-server/src/metrics.rs @ai-dynamo/dynamo-frontend-codeowners @ai-dynamo/dynamo-observability-codeowners /tests/test_kustomize_matrix.py @ai-dynamo/dynamo-runtime-codeowners @ai-dynamo/dynamo-ops-codeowners /.github/workflows/fern-docs.yml @ai-dynamo/dynamo-ops-codeowners @ai-dynamo/dynamo-docs-codeowners /examples/deployments/**/vllm/** @ai-dynamo/dynamo-operator-codeowners @ai-dynamo/dynamo-backend-vllm-codeowners @@ -758,6 +764,7 @@ requirements.txt /components/src/dynamo/vllm/multimodal_utils/ @ai-dynamo/dynamo-backend-vllm-codeowners @ai-dynamo/dynamo-multimodal-codeowners /deploy/inference-gateway/ext-proc/Dockerfile @ai-dynamo/dynamo-epp-codeowners @ai-dynamo/dynamo-router-codeowners @ai-dynamo/dynamo-ops-codeowners /components/src/dynamo/trtllm/encode_helper.py @ai-dynamo/dynamo-backend-trtllm-codeowners @ai-dynamo/dynamo-multimodal-codeowners +/container/Dockerfile.global-discovery-service @ai-dynamo/dynamo-global-router-codeowners @ai-dynamo/dynamo-ops-codeowners /container/templates/sglang_runtime.Dockerfile @ai-dynamo/dynamo-ops-codeowners @ai-dynamo/dynamo-backend-sglang-codeowners /container/templates/trtllm_runtime.Dockerfile @ai-dynamo/dynamo-ops-codeowners @ai-dynamo/dynamo-backend-trtllm-codeowners /tests/router/test_router_e2e_with_vllm_xpu.py @ai-dynamo/dynamo-router-codeowners @ai-dynamo/dynamo-xpu-codeowners @ai-dynamo/dynamo-backend-vllm-codeowners diff --git a/Cargo.lock b/Cargo.lock index 3c99a69128f7..f52cc079ff75 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2504,6 +2504,26 @@ dependencies = [ [[package]] name = "dynamo-global-discovery-service" version = "1.7.0" +dependencies = [ + "anyhow", + "async-stream", + "async-trait", + "axum 0.8.4", + "clap", + "dynamo-global-common", + "dynamo-http-server", + "dynamo-runtime", + "futures", + "k8s-openapi", + "kube", + "reqwest 0.12.28", + "serde_json", + "thiserror 2.0.18", + "tokio", + "tokio-util", + "tracing", + "url", +] [[package]] name = "dynamo-global-frontend" @@ -2530,6 +2550,28 @@ dependencies = [ "thiserror 2.0.18", ] +[[package]] +name = "dynamo-http-server" +version = "1.7.0" +dependencies = [ + "anyhow", + "axum 0.8.4", + "dynamo-runtime", + "futures", + "hyper-util", + "parking_lot", + "prometheus", + "rand 0.9.4", + "rcgen", + "rustls", + "serde", + "tokio", + "tokio-rustls", + "tokio-util", + "tower 0.5.3", + "tracing", +] + [[package]] name = "dynamo-kv-hashing" version = "1.7.0" diff --git a/Cargo.toml b/Cargo.toml index 4e254e8d1d7f..0106d09e337c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -12,6 +12,7 @@ members = [ "lib/kv-hashing", "lib/mocker", "lib/kv-router", + "lib/http-server", "lib/global-common", "lib/global-discovery-client", "lib/global-discovery-service", @@ -73,6 +74,7 @@ dynamo-truthy = { path = "lib/truthy", version = "1.7.0" } dynamo-data-gen = { path = "lib/data-gen", version = "1.7.0" } dynamo-kv-hashing = { path = "lib/kv-hashing", version = "1.7.0" } dynamo-global-common = { path = "lib/global-common", version = "1.7.0" } +dynamo-http-server = { path = "lib/http-server", version = "1.7.0" } dynamo-memory = { path = "lib/memory", version = "1.7.0" } dynamo-mocker = { path = "lib/mocker", version = "1.7.0" } dynamo-kv-router = { path = "lib/kv-router", version = "1.7.0" } diff --git a/container/Dockerfile.global-discovery-service b/container/Dockerfile.global-discovery-service new file mode 100644 index 000000000000..dfecb17ffc2e --- /dev/null +++ b/container/Dockerfile.global-discovery-service @@ -0,0 +1,26 @@ +# syntax=docker/dockerfile:1.10.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# CPU-only image for the Global Discovery Service, which serves pool discovery +# for the Global Router over HTTP; see global-discovery-service.yaml. Build it +# from the repository root: +# docker build -f container/Dockerfile.global-discovery-service -t global-discovery-service . +FROM rust:1.96.1-trixie AS builder +RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ + build-essential cmake pkg-config protobuf-compiler libclang-dev libssl-dev curl \ + && rm -rf /var/lib/apt/lists/* +WORKDIR /workspace +COPY . . +RUN --mount=type=cache,target=/usr/local/cargo/registry,sharing=locked \ + --mount=type=cache,target=/usr/local/cargo/git,sharing=locked \ + cargo build --locked --release -p dynamo-global-discovery-service --bin global-discovery-service + +FROM debian:trixie-slim AS global-discovery-service +RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ + ca-certificates libstdc++6 \ + && rm -rf /var/lib/apt/lists/* +COPY --from=builder /workspace/target/release/global-discovery-service /usr/local/bin/global-discovery-service +USER 10001:0 +EXPOSE 8080 +ENTRYPOINT ["/usr/local/bin/global-discovery-service"] diff --git a/container/global-discovery-service.yaml b/container/global-discovery-service.yaml new file mode 100644 index 000000000000..75332dd33c2a --- /dev/null +++ b/container/global-discovery-service.yaml @@ -0,0 +1,116 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Global Discovery Service with pools from Karmada, applied directly to the hub +# cluster (not through Karmada). Replace IMAGE with the pushed image, for example: +# sed "s#IMAGE#/global-discovery-service:#" global-discovery-service.yaml | kubectl apply -f - +# It needs only read access to EndpointSlices. It lists the pools whose +# `-pool-relay` Services Karmada imports, and the replicas of the +# Global Router Service named by the --router-* arguments: set those to the +# router's namespace, its Service and the port PoolRelays connect to. +# +# It reads the imported slices with the pod's credentials, so this cluster +# needs a ServiceImport for each pool's PoolRelay Service and for the router +# Service. To read the Karmada control plane instead, which holds every +# exported Service's slices, mount its kubeconfig from a Secret, point +# KUBECONFIG at it, and grant the same read access there. +apiVersion: v1 +kind: Namespace +metadata: + name: global-discovery +--- +apiVersion: v1 +kind: ServiceAccount +metadata: + name: global-discovery-service + namespace: global-discovery +--- +apiVersion: rbac.authorization.k8s.io/v1 +kind: ClusterRole +metadata: + name: global-discovery-service +rules: + - apiGroups: ["discovery.k8s.io"] + resources: ["endpointslices"] + verbs: ["get", "list", "watch"] +--- +apiVersion: rbac.authorization.k8s.io/v1 +kind: ClusterRoleBinding +metadata: + name: global-discovery-service +roleRef: + apiGroup: rbac.authorization.k8s.io + kind: ClusterRole + name: global-discovery-service +subjects: + - kind: ServiceAccount + name: global-discovery-service + namespace: global-discovery +--- +apiVersion: apps/v1 +kind: Deployment +metadata: + name: global-discovery-service + namespace: global-discovery +spec: + replicas: 1 + selector: + matchLabels: + app.kubernetes.io/name: global-discovery-service + template: + metadata: + labels: + app.kubernetes.io/name: global-discovery-service + spec: + serviceAccountName: global-discovery-service + securityContext: + runAsNonRoot: true + containers: + - name: global-discovery + image: IMAGE + args: + - --listen=0.0.0.0:8080 + - --pool-discovery=karmada + - --router-namespace=global-discovery + - --router-service=global-router + - --router-port-name=http + env: + - name: DYN_LOG + value: info + ports: + - name: http + containerPort: 8080 + readinessProbe: + httpGet: + path: /healthz + port: http + periodSeconds: 5 + lifecycle: + # Keep serving while the Pod leaves the Service's endpoints. + preStop: + exec: + command: ["sleep", "5"] + resources: + requests: + cpu: 50m + memory: 64Mi + limits: + memory: 256Mi + securityContext: + allowPrivilegeEscalation: false + readOnlyRootFilesystem: true + capabilities: + drop: ["ALL"] +--- +apiVersion: v1 +kind: Service +metadata: + name: global-discovery-service + namespace: global-discovery +spec: + selector: + app.kubernetes.io/name: global-discovery-service + ports: + - name: http + port: 8080 + targetPort: http diff --git a/docs/fern/index.yml b/docs/fern/index.yml index 82624693c617..fba37800f2cc 100644 --- a/docs/fern/index.yml +++ b/docs/fern/index.yml @@ -171,6 +171,9 @@ navigation: path: pages/kubernetes/kv-aware-routing/dynamo-frontend.md - page: Deploy the Pool Relay (Experimental) path: pages/kubernetes/kv-aware-routing/pool-relay.md + - page: Deploy Multi-Cluster Services (Experimental) + path: pages/kubernetes/kv-aware-routing/multi-cluster-services.md + slug: multi-cluster-services - page: Using GAIE with Dynamo path: pages/kubernetes/kv-aware-routing/gateway-api.mdx - page: Using GAIE with vanilla vLLM @@ -478,6 +481,9 @@ navigation: - page: Pool Relay Concepts (Experimental) slug: multi-dc-kv-routing path: pages/developer-guide/knowledge-base/modular-components/router/multi-dc-kv-routing.md + - page: Multi-Cluster Services Registry (Experimental) + slug: multi-cluster-services-registry + path: pages/developer-guide/knowledge-base/modular-components/router/multi-cluster-services-registry.md - section: Standalone Services skip-slug: true collapsed: true @@ -1272,6 +1278,9 @@ navigation: path: pages/reference/components/dynosim-sweep-reference.mdx - page: Global Router Configuration path: pages/reference/components/global-router-configuration.mdx + - page: Multi-Cluster Services Configuration + path: pages/reference/components/global-discovery-mcs-configuration.md + slug: global-discovery-mcs-configuration - page: Gateway API Routing path: pages/reference/components/gateway-api-routing.mdx - page: ThunderAgent Scheduler diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/multi-cluster-services-registry.md b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/multi-cluster-services-registry.md new file mode 100644 index 000000000000..36989fe8599d --- /dev/null +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/multi-cluster-services-registry.md @@ -0,0 +1,398 @@ +--- +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +title: Multi-Cluster Services Registry +subtitle: How the registry and controller implement the Multi-Cluster Services API, and how the Global Discovery Service reads the imports +--- + +**Experimental.** The module +[`deploy/global-discovery-mcs`](https://github.com/ai-dynamo/dynamo/tree/main/deploy/global-discovery-mcs) +implements the Kubernetes Multi-Cluster Services (MCS) API (KEP-1645) with a central registry +instead of a Karmada control plane. A controller in every member cluster writes the cluster's +exports to the registry and materialises the merged imports of the whole clusterset, and the Global +Discovery Service reads those imports to find pools and Global Router replicas in every cluster. +This page describes the data model, the merge and consistency rules, the storage layout, the +controller, and the boundary with Karmada. + +For deployment, see [Deploy Multi-Cluster Services](../../../../kubernetes/kv-aware-routing/multi-cluster-services.md). +For flags, conditions and metrics, see +[Multi-Cluster Services Configuration](../../../../reference/components/global-discovery-mcs-configuration.md). + +## Background + +KEP-1645 defines two resources in `multicluster.x-k8s.io`. A ServiceExport, created with the name +and in the namespace of a Service, makes that Service visible to the clusterset. An implementation +then creates a ServiceImport of the same name in every cluster of the clusterset where the +namespace exists ("namespace sameness"), with the union of the exported ports and either a +clusterset virtual IP (ClusterSetIP) or none (Headless). Endpoints arrive as EndpointSlices +labelled `multicluster.kubernetes.io/service-name` and `multicluster.kubernetes.io/source-cluster`, +one set per source cluster. When the exports of one Service disagree, the oldest export wins and +the others report a `Conflict` condition. Cluster identity comes from the KEP-2149 ClusterProperty +objects `cluster.clusterset.k8s.io` and `clusterset.k8s.io`. + +The module follows the `sigs.k8s.io/mcs-api` v0.5.2 types, CRDs and derived Service naming. Its +merge and conflict rules implement the conflict resolution of KEP-1645 as expressed by the +`sigs.k8s.io/mcs-api` condition reasons. They run once, in the registry, so every cluster sees the +same result. + +## Components + +```mermaid +flowchart LR + subgraph A["Member cluster A"] + SA["Services and EndpointSlices"] --> CA["mcs-controller"] + CA --> IA["ServiceImports, derived Services, imported EndpointSlices"] + IA --> GD["Global Discovery Service"] + end + subgraph B["Member cluster B"] + SB["Services and EndpointSlices"] --> CB["mcs-controller"] + CB --> IB["ServiceImports, derived Services, imported EndpointSlices"] + end + subgraph R1["Region 1"] + G1["mcs-registry replicas"] --> D1[("DynamoDB table replica")] + end + subgraph R2["Region 2"] + G2["mcs-registry replicas"] --> D2[("DynamoDB table replica")] + end + CA -->|"gRPC writes and watch"| G1 + CB -->|"gRPC writes and watch"| G2 + D1 ---|"Global Table replication"| D2 +``` + +| Component | Role | +|---|---| +| `mcs-registry` | Stateless server. Keeps an in-memory view of each clusterset, merges exports into ServiceImports, evaluates leases, and serves the `dynamo.mcs.v1` API over gRPC and REST. Every replica runs every loop; no replica is elected. | +| Store | `memory` (one process) or `dynamodb` (one table, replicated between regions as a Global Table). | +| `mcs-controller` | One per member cluster, leader-elected. Registers the cluster, writes its exports and endpoints, and applies the clusterset's imports to its cluster. | +| `mcs-admin` | Inspects a store, deletes a cluster, creates the DynamoDB table. | +| `pkg/client` | Go client of the registry: ordered URL failover with DNS re-resolution, writer generations, and a reconnecting watch loop. | + +The registry does not authenticate callers in this version. It belongs on a network that only the +member clusters reach. + +## Data Model + +### What the Registry Stores + +The registry stores facts, each written by one writer: the `mcs-controller` of the cluster the +record belongs to. + +| Record | Key | Content | +|---|---|---| +| Cluster | clusterset, cluster | Region, zones, provider, labels, the lease (epoch, TTL, last renewal) and a generation | +| ServiceExport | clusterset, cluster, namespace, name | Ports, type, session affinity and its client IP timeout, IP families, internal traffic policy, traffic distribution, exported labels and annotations, the export creation timestamp, and a writer generation | +| Endpoint set | clusterset, cluster, namespace, name | Every endpoint of the export in that cluster: addresses, address type, `ready`, `serving` and `terminating`, hostname, node name, zone, ports and attributes, with a writer generation. At most 1000 endpoints. | + +An endpoint set is only ever replaced as a whole. The controller's source of truth is the set of +local EndpointSlices, so a full replace is idempotent across restarts, needs no tombstones, and +cannot leave a stale endpoint behind. + +### What the Registry Computes + +The registry derives, and never stores: + +- the lease state of each cluster, at read time; +- one ServiceImport per (namespace, name) in a clusterset: the merged spec, the conflict if any, + a status per exporting cluster (lease state, participation, its own conflict, endpoint counts, + and `Valid`, `Ready` and `Conflict` conditions), and the endpoints grouped by source cluster. + +Merge rules: + +1. Only exports of clusters whose lease is active take part. They are ordered oldest first by + export creation timestamp, then by cluster name; the first one wins. +2. Ports are the union by (port, protocol) in that order. A name or app protocol disagreement + resolves to the oldest export, and a port whose name another (port, protocol) already uses is + skipped. +3. IP families are intersected. An export with no family in common with the running intersection + is skipped; an export that lists no family does not constrain the result. +4. Type, session affinity and its configuration, annotations, labels, internal traffic policy and + traffic distribution come from the winner. + +A conflict is set when a participating export disagrees with the winner. The registry reports one +reason, the first that applies in the order ports, IP families, type, session affinity, its client +IP timeout, annotations, labels, internal traffic policy and traffic distribution, using the +`sigs.k8s.io/mcs-api` reasons plus `IPFamilyConflict`. As KEP-1645 asks, the `Conflict` condition is +set on every participating cluster; the per-cluster `conflict` field additionally names the +clusters whose own export differs. + +A cluster whose lease expired stops taking part: its endpoints leave every import, and its exports +no longer count in conflict resolution, so a dead cluster cannot pin the winner. An import whose +exporting clusters have all expired keeps its type and ports, with no endpoints, so derived +Services keep their virtual IPs while a controller restarts. The records stay for `--retention` +(24 h), so a returning controller keeps its export creation timestamps, then compaction removes +them. `DeregisterCluster` removes everything of a cluster at once. + +## Writes, Generations and Leases + +Writes are conditional in the store: + +- Exports and endpoint sets carry a writer generation that must exceed the stored one. Writers use + `max(previous + 1, unix milliseconds)`, which stays monotonic across restarts and across a + failover to another registry region without remembering earlier values. A stale write fails with + `STALE_GENERATION` and the stored generation; the controller retries once above it. +- `RegisterCluster` mints a lease epoch, `max(stored + 1, unix milliseconds)`. Heartbeats must + carry the current epoch; any other epoch fails with `LEASE_EPOCH_MISMATCH`, and the controller + registers again and resends every export. + +A handler writes to the store, then applies the stored record to its own replica's view before it +responds (write-through), so a writer and every client of that replica read the write at once. +Sibling replicas pick it up on their next refresh. A record written through a replica wins over a +refresh snapshot for one refresh interval, unless the snapshot holds a newer version, so a refresh +that started before the write does not undo it. + +Leases are evaluated at read time from server timestamps: a lease is expired once +`renewed_at + ttl + skew` is in the past. Nothing writes expiry to the store. Each replica +re-evaluates every lease every `--lease-eval-interval` (1 s) without a store read and publishes the +changes. Replicas may disagree on a lease by the skew allowance (2 s by default); the only effect +is an endpoint set that appears or disappears a moment earlier on one replica than on another. + +## Watches and Revisions + +Each replica refreshes every clusterset every `--refresh-interval` (2 s): it loads the clusterset +from the store, merges, diffs against its previous view, and publishes the changes to the +clusterset's watchers. It lists the store's clustersets every 30 s, and loads a clusterset at once +when a request names one it does not know. + +```mermaid +sequenceDiagram + participant S as Service and EndpointSlices in cluster A + participant CA as mcs-controller in A + participant R as mcs-registry + participant DB as Store + participant CB as mcs-controller in B + participant K as ServiceImport, derived Service and EndpointSlices in B + S->>CA: ServiceExport, created or from the export label + CA->>R: PutServiceExport with spec and generation + CA->>R: ReplaceEndpoints with the full endpoint set + R->>DB: conditional writes + R->>R: apply to own view, write-through + R->>R: refresh every 2 s, merge, evaluate leases + R-->>CB: WatchServiceImports events + CB->>K: server-side apply + R-->>CA: WatchServiceImports, status of cluster A + CA->>S: ServiceExport conditions Valid, Ready and Conflict + loop every TTL/3 + CA->>R: HeartbeatCluster with the lease epoch + end +``` + +A watch sends `SNAPSHOT_START`, one `UPSERT` per object, `SNAPSHOT_END` with a revision, then +deltas (`UPSERT` with the full object, `DELETE` with the last known one), and a `BOOKMARK` with the +current revision every 5 s, which keeps load balancer idle timeouts from closing quiet streams. + +- A revision is `view_epoch:seq`. The view epoch is minted when the replica starts; the sequence + counts the changes of one clusterset on that replica. Revisions of different replicas, or of one + replica before and after a restart, are unrelated. +- Each replica keeps the last 10000 events per clusterset. A watch that resumes with a revision of + the same replica whose later events are all still kept receives only those events. Any other + revision yields a fresh snapshot, so a resume never fails and never skips an event; a client that + fails over to another replica or region loses only the replay shortcut. +- Filters (namespace, label selector, endpoints or not) apply per watcher. An object that stops + matching a watcher's filter arrives as `DELETE` with its new state. A bookmark's revision covers + changes the filter skipped. +- A watcher that falls 1024 events behind is sent a new snapshot instead of being disconnected. A + replica serves at most 5000 watch streams; beyond that a watch fails with `TOO_MANY_WATCHERS`. + +## Storage + +### Memory Store + +The memory store keeps the same conditional semantics in the registry process. It suits tests and +evaluation with one replica. A restart loses every record: each controller's next heartbeat fails +with `CLUSTER_NOT_REGISTERED`, and the controller registers again and resends its exports. In +between, the importers receive a snapshot without those imports and delete them, then re-create +them as the exports return. + +### DynamoDB Single-Table Layout + +The DynamoDB store keeps every record in one table with a string partition key `pk` and sort key +`sk`. Record bodies are serialized protobuf; the values that write conditions need are top-level +attributes (`gen`, `lease_epoch`, `lease_renewed_at`, `updated_at`, `expires_at`). + +| Item | `pk` | `sk` | +|---|---|---| +| Clusterset index | `CSX` | `CS#` | +| Cluster | `CS#` | `CL#` | +| ServiceExport | `CS#` | `EX###` | +| Endpoint partition marker | `CS#` | `EPX#` | +| Endpoint set | `EP##` | `#` | + +- The clusterset partition holds the small, heartbeat-heavy records. Endpoint sets, which are large + and change often, live in one partition per cluster. Loading a clusterset or deleting a cluster + needs prefix queries only; the table has no global secondary index. +- Registration writes on `attribute_not_exists(pk) OR (lease_epoch < :new AND gen = :seen)`, + heartbeats on `lease_epoch = :epoch`, and export and endpoint set puts on + `attribute_not_exists(pk) OR gen < :gen`, returning the stored item on failure. +- A refresh reads a clusterset with one query of its partition and one query per endpoint + partition, eventually consistent and paginated. +- Cluster items carry a 7-day TTL (`expires_at`) as a garbage-collection safety net. Leases never + depend on DynamoDB TTL timing. + +The module's +[storage document](https://github.com/ai-dynamo/dynamo/blob/main/deploy/global-discovery-mcs/docs/storage.md) +has the full item catalogue, conditions, capacity estimates and table definition. + +### Global Tables and Multi-Region Reads + +With a Global Table, every registry replica reads and writes the table replica of its own region. +DynamoDB replicates each write to the other regions, usually within a second, and resolves +concurrent writes to one item by last writer wins. Conditions are evaluated only in the region +that accepts a write. The registry stays correct without cross-region conditions because of three +properties of its records: + +1. **One logical writer per item.** A cluster's records are written only by that cluster's + controller, which has one leader at a time. +2. **Monotonic generations.** The generations of one writer never decrease: a generation of + `max(previous + 1, unix milliseconds)` stays above everything that writer wrote before, + including across a restart and a failover to another region. In practice generation order also + follows wall-clock order, provided the controllers' clocks are roughly right. DynamoDB chooses + the winner of concurrent cross-region writes by its own internal timestamps, which the + registry's generations do not set, so the registry does not rely on a generation to win that + race. It relies on the next property to repair the outcome. +3. **Full-state writes.** Every put replaces a whole record, so a writer's next write of that + record repairs whatever a conflict resolution kept. The controller writes a record again when + the local object changes and after a resync trigger. It has no timer today; see + [Guarantees and Typical Timing](#guarantees-and-typical-timing). + +When the store does not answer, a replica keeps serving reads from its last view, fails writes +with `STORE_UNAVAILABLE`, and reports not ready once its view is older than three refresh +intervals, so load balancer and DNS health checks move clients to healthy replicas or regions. + +### Guarantees and Typical Timing + +This section separates what the mechanisms above guarantee from what is typical. Every duration is +a default; the flag that sets it is named where it applies. + +**Guaranteed by the implementation:** + +- **Lease expiry.** A cluster whose heartbeats stop leaves every import once `--lease-ttl` plus + `--lease-skew` has passed since its last renewal (30 s and 2 s by default), evaluated by each + registry replica every `--lease-eval-interval` (1 s). Leases never depend on DynamoDB TTL timing. + Sibling replicas learn the last renewal at their next refresh, so allow roughly 32 to 35 s with + the defaults, plus the time to deliver the change to the controllers and apply it. +- **Resync on loss.** At startup, after each reconnect of the registry client, and after each + re-registration, the leader deletes this cluster's registry exports that have no local + ServiceExport and writes every export again. A failed resync is retried with a backoff from 1 s + to 30 s until it succeeds. A reconnect includes a failover to another registry URL or region. +- **Writes follow local changes.** When a ServiceExport, its Service or its EndpointSlices + change, the controller writes the spec or the endpoint set whose hash changed, with a + generation above the stored one, and retries a transient failure with backoff. +- **No stale overwrite, no lost event.** A write below the stored generation is refused. A watch + that resumes never fails and never skips an event, and the importer deletes imports missing + from the registry only after `SNAPSHOT_END`. + +**Typical, not guaranteed:** + +- **Propagation.** A write is visible on the replica that accepted it at once, on its sibling + replicas within one `--refresh-interval` (2 s), and in other regions after DynamoDB's + replication plus one refresh interval. Replication usually takes about a second, and the + `ReplicationLatency` metric shows the actual delay; DynamoDB gives no bound. An endpoint change + in one cluster typically reaches the slices of another cluster in 3 to 6 s, which includes the + write, replication, the refresh, the watch event and the apply. Treat 5 s as the usual target, + not as a limit. +- **Recovery after a registry restart with the memory store.** The controller's next heartbeat + fails, which re-registers the cluster and resends every export. That takes up to one heartbeat + interval, `--lease-ttl` divided by three with 10 percent jitter (about 10 s by default), or less + when a write fails first and wakes the lease. Until then, every cluster has deleted the imports + the registry no longer holds. + +**Not covered today:** + +- **Repair without a trigger.** The controller remembers the hash of each spec and endpoint set + it wrote and writes again only when the local object changes. A resync is triggered by startup, + a reconnect and a re-registration, not by a timer. If a registry record differs from what the + controller last wrote while the connection stays up and the cluster stays registered, for + example after concurrent cross-region writes kept an older generation, the record stays as it is + until the local object changes or a trigger fires. There is no upper bound on that time. The + registry does not detect such drift. + +A periodic resync of exports and imports is planned; this page will describe its intervals when it +exists. + +## Controller + +All parts run on one controller-runtime manager. Everything that writes to the registry or applies +imports runs on the leader only. + +| Part | Role | +|---|---| +| Identity | Resolves the clusterset and cluster ids (flags, else the ClusterProperty objects), the region (flag, else node labels), zones and provider at startup. | +| Lease | Registers the cluster, then heartbeats every TTL/3. A refused heartbeat registers again and triggers a full resync. | +| ServiceExport reconciler | For each ServiceExport, and again on each change of the ServiceExport, its Service or its EndpointSlices: adds the finalizer, validates the Service, writes the spec and the endpoint set when their hashes change, and writes the conditions. On deletion it deletes the registry record before removing the finalizer. | +| Resyncer | At startup, after each reconnect and after each re-registration, and at no other time: deletes this cluster's registry exports that have no local ServiceExport, then writes every export again. It has no timer. | +| Auto-export | Creates a ServiceExport, owned by the Service, for each Service labelled `dynamo.nvidia.com/mcs-export: "true"`, and deletes it when the label goes. It never touches a ServiceExport it did not create. | +| Importer | Holds one `WatchServiceImports` stream with endpoints and applies each import with server-side apply. | + +The export side builds the endpoint set from the Service's EndpointSlices, skipping the slices it +imported itself. For each endpoint it records the addresses, the slice's address type, `ready` +(unset means ready), `serving` (unset follows ready), `terminating`, hostname, node, zone, the +slice's ports, and the endpoint's target reference as the attributes `targetRef.kind`, +`targetRef.name` and `targetRef.uid`. The Service's `dynamo.nvidia.com/mcs-*` labels and annotations +travel with the export as exported labels and annotations. + +The import side writes, for each import in a namespace that exists locally: + +| Object | Name | Content | +|---|---|---| +| ServiceImport | The exported Service's name | Merged type, ports, session affinity and IP families; the derived Service's ClusterIPs as `spec.ips`; the participating clusters and the `Ready` and `Conflict` conditions | +| Derived Service | `derived-` | No selector; ClusterIP, or `clusterIP: None` for Headless; merged ports; owned by the ServiceImport | +| EndpointSlices | `--` | One set per source cluster, the local cluster included, at most 100 endpoints each; ports filtered to the merged port names; conditions as exported; `targetRef` from the exported attributes; labels `kubernetes.io/service-name`, `multicluster.kubernetes.io/service-name`, `multicluster.kubernetes.io/source-cluster` and `endpointslice.kubernetes.io/managed-by` | + +The importer deletes imports that are missing from the registry only after `SNAPSHOT_END`, so a +reconnect in the middle of a snapshot never removes an import. An import whose namespace does not +exist waits until the namespace appears. A ServiceImport without the controller's managed-by label +belongs to another implementation and is skipped. + +### Derived Service Naming + +The derived Service of the import `/` is `derived-` followed by the first ten +characters of the lowercase, unpadded base32hex encoding of the SHA-256 of `/`. +This is the `sigs.k8s.io/mcs-api` algorithm, so the name is the same in every cluster and in every +implementation that follows mcs-api, and a client can compute it to dial +`derived-..svc.cluster.local`. Names under `clusterset.local` need the CoreDNS +`multicluster` plugin, which reads the ServiceImports and slices the controller writes. + +## How the Global Discovery Service Reads the Imports + +The Global Discovery Service, with `--pool-discovery mcs`, never calls the registry. It runs in a +member cluster and watches the EndpointSlices that `mcs-controller` imports there: + +- It selects every slice labelled `multicluster.kubernetes.io/source-cluster`, in every namespace, + and groups the slices by `multicluster.kubernetes.io/service-name`. The source-cluster label tells + apart the endpoints that each cluster exported for the same Service. +- A pool is a Service named `<--pool-service-suffix>` (default `-pool-relay`) in one source + cluster and namespace. It is live while one of its endpoints is Ready. +- The Global Router replicas are the Ready endpoints of the Service named by `--router-namespace` + and `--router-service`, reached on the TCP port named `--router-port-name`. The controller keeps + the port names of the merged import on every slice. +- A replica is named `/`, after the imported endpoint's `targetRef.name`, + which the controller carries across clusters in the `targetRef.name` attribute; an endpoint + without one is named after its address. + +The contract between the module and the Global Discovery Service is these two labels, the +`targetRef.name` attribute, the copied endpoint conditions and the port names. An endpoint change +in one cluster typically reaches the other clusters' slices in 3 to 6 seconds. A cluster whose +lease expires (30 s plus the 2 s skew by default) loses its slices in every cluster, so its pools +and replicas leave the lists. See [Guarantees and Typical Timing](#guarantees-and-typical-timing) +for what is guaranteed. + +## Compatibility with Karmada + +The Global Discovery Service also reads Karmada's imports, with `--pool-discovery karmada`. The two +implementations share the export side and differ on the import side: + +| Concern | Dynamo MCS | Karmada | +|---|---|---| +| Export | ServiceExport in the cluster that runs the Service | The same object | +| Where imports appear | Every member cluster with the controller | The Karmada control plane, and member clusters with a ServiceImport for the Service | +| Imported slice identification | Labels `multicluster.kubernetes.io/service-name` and `multicluster.kubernetes.io/source-cluster` | Name `imported--`, label `kubernetes.io/service-name: derived-` | +| Derived Service | `derived-` | `derived-`, which limits Service names to 55 characters | +| Control plane | `mcs-registry`, with the memory store or a DynamoDB Global Table | The Karmada control plane | + +Because the export side is the same, a deployment moves between the two by replacing the MCS +implementation and changing `--pool-discovery`; the exported Services and their ServiceExports +stay the same, subject to Karmada's 55-character limit on Service names. The ServiceExport and +ServiceImport CRDs from mcs-api v0.5.2 serve `v1alpha1` and `v1beta1`, so clients of either version +share them. Within one cluster, keep the +implementations apart by namespace: the controller adds its finalizer to, and writes conditions +on, every ServiceExport in the namespaces it watches, and it skips ServiceImports it did not create. diff --git a/docs/fern/pages/kubernetes/kv-aware-routing/multi-cluster-services.md b/docs/fern/pages/kubernetes/kv-aware-routing/multi-cluster-services.md new file mode 100644 index 000000000000..b063bb164f5b --- /dev/null +++ b/docs/fern/pages/kubernetes/kv-aware-routing/multi-cluster-services.md @@ -0,0 +1,492 @@ +--- +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +title: Deploy Multi-Cluster Services +subtitle: Share Services between Kubernetes clusters for the Global Discovery Service without Karmada +--- + +**Experimental.** NVIDIA Dynamo includes an implementation of the Kubernetes Multi-Cluster Services +(MCS) API ([KEP-1645](https://github.com/kubernetes/enhancements/tree/master/keps/sig-multicluster/1645-multi-cluster-services-api), +`multicluster.x-k8s.io`) in +[`deploy/global-discovery-mcs`](https://github.com/ai-dynamo/dynamo/tree/main/deploy/global-discovery-mcs). +It has two parts: + +- `mcs-registry`, a registry server that holds the clusters of a clusterset, the Services each + cluster exports and their endpoints, and merges the exports of each Service into one + ServiceImport. +- `mcs-controller`, which runs in every member cluster. It sends the cluster's ServiceExports and + their endpoints to the registry, and creates the ServiceImports, derived Services and + EndpointSlices of every exported Service in its own cluster. + +The Global Discovery Service, run with `--pool-discovery mcs`, reads those EndpointSlices to find +the Global Router pools and replicas of every cluster. + +For every flag, condition and metric, see +[Multi-Cluster Services Configuration](../../reference/components/global-discovery-mcs-configuration.md). +For the data model and the consistency rules, see +[Multi-Cluster Services Registry](../../developer-guide/knowledge-base/modular-components/router/multi-cluster-services-registry.md). + +## Choose an Implementation + +The Global Discovery Service reads imports from this implementation (`--pool-discovery mcs`) or +from Karmada (`--pool-discovery karmada`). Both read the same ServiceExport in the cluster that runs +the Service, and both need pod addresses that are routable between the clusters. + +| | Dynamo MCS (`mcs`) | Karmada (`karmada`) | +|---|---|---| +| Control plane | `mcs-registry` in one cluster, or in several regions over one DynamoDB Global Table | The Karmada control plane | +| Where imports appear | Every member cluster that runs `mcs-controller`, in the namespace of the exported Service | The Karmada control plane, and member clusters that have a ServiceImport for the Service | +| Longest exported Service name | 63 characters | 55 characters (Karmada names the imported Service `derived-`) | +| Imported EndpointSlices | Labelled `multicluster.kubernetes.io/service-name` and `multicluster.kubernetes.io/source-cluster` | Named `imported--` | + +Use the Dynamo implementation when all of the following hold: + +- **One owner and a trusted network.** The registry does not authenticate callers: any client + that reaches its ports can register a cluster and write exports for it. Run it on a private + network that only the member clusters reach, and use TLS on its listeners where that network + is shared. +- **A pod network across clusters.** Imported EndpointSlices hold the pod addresses of the source + cluster. Consumers, the Global Discovery Service and the Global Router included, dial those + addresses directly, for example over peered VPCs with non-overlapping pod CIDRs. Neither + implementation checks that the addresses are reachable. +- **No single control-plane cluster, or several regions.** The registry runs in each region over + one DynamoDB Global Table, and each controller fails over between registry URLs. + +Use Karmada when the clusters are already Karmada members and you manage them through its control +plane. + +## Prerequisites + +- Two or more Kubernetes clusters, and kubectl and Helm access to each of them. +- A clusterset name shared by every cluster and a cluster id unique to each cluster, both DNS + labels, for example clusterset `prod` and cluster ids `use1-a` and `usw2-a`. The cluster id is + the `multicluster.kubernetes.io/source-cluster` label of every EndpointSlice imported from that + cluster. +- Pod addresses routable between the clusters, and network rules that allow each cluster's pod + CIDR to reach the exported ports. +- A network path from every member cluster to the registry's gRPC port, `9090`. +- For the DynamoDB store: an AWS account, permission to create a DynamoDB table, and an IAM role + for the registry. The module's + [IAM document](https://github.com/ai-dynamo/dynamo/blob/main/deploy/global-discovery-mcs/docs/iam.md) + lists the policies. + +Build the `global-discovery-mcs` image, which holds `mcs-registry`, `mcs-controller`, `mcs-admin` and the +CRDs, and push it to a registry your clusters can pull from: + +```bash +cd deploy/global-discovery-mcs +make docker-build docker-push IMG=my-registry.example.com/global-discovery-mcs:dev +``` + +The commands below install the Helm chart +[`deploy/helm/charts/global-discovery-mcs`](https://github.com/ai-dynamo/dynamo/tree/main/deploy/helm/charts/global-discovery-mcs) +from a checkout of the repository. They set only the values each step needs; the +[chart README](https://github.com/ai-dynamo/dynamo/blob/main/deploy/helm/charts/global-discovery-mcs/README.md) +lists every value, including the Service, TLS and ServiceAccount settings. + +## Install the Registry + +Run one registry for the clusterset, in any cluster that every member cluster can reach: a cluster +of its own, or one of the members. Choose a store: + +| Store | Replicas | Use | +|---|---|---| +| `memory` | Exactly one | Evaluation. A restart loses every record. The controllers register again on their next heartbeat (every 10 s with the default 30 s lease) and send their exports again; in between, every cluster deletes the imports the registry no longer has. | +| `dynamodb` | Any number, in one or more regions | Production. Records survive restarts, and a DynamoDB Global Table replicates them between regions. | + +The registry and every controller must agree on the transport. Choose one path before you install +either of them: + +| Path | Registry | Controllers | Use | +|---|---|---|---| +| TLS | `registry.tls.secretName` names a Secret that holds the registry certificate | An `https://` URL, and `controller.registryTlsCa.configMapName` when a private CA signed the certificate | Production | +| Plaintext | `registry.tls.secretName` empty (the default) | An `http://` URL and `controller.registryAllowInsecureHttp=true` | Evaluation on a private network | + +The chart refuses an `http://` controller URL without `controller.registryAllowInsecureHttp=true`. +It cannot check a registry in another cluster, so an `https://` URL in front of a plaintext +registry, or the reverse, shows up as a controller that cannot register its cluster. + +The two choices are independent: the commands below pair the memory store with the plaintext path +and the DynamoDB store with the TLS path. To run the memory store over TLS, add +`registry.tls.secretName` to the first command and follow the TLS steps. To run the DynamoDB store +without TLS, leave it out of the second command and follow the plaintext steps. + +### Evaluation: Memory Store without TLS + +```bash +helm install mcs-registry deploy/helm/charts/global-discovery-mcs \ + --namespace mcs-system --create-namespace \ + --set registry.enabled=true \ + --set controller.enabled=false \ + --set image.repository=my-registry.example.com/global-discovery-mcs \ + --set image.tag=dev +``` + +The registry serves plaintext (h2c) on the gRPC and REST ports. Install the controllers with an +`http://` URL and `controller.registryAllowInsecureHttp=true`, as the next section shows. + +### Production: DynamoDB Store with TLS + +Create the table once with `mcs-admin` (`/mcs-admin` in the image, or `bin/mcs-admin` after +`make build`), using the credentials of an administrator role: + +```bash +mcs-admin create-table --table=global-discovery-mcs --region=us-east-1 +``` + +Create the registry's certificate as a `kubernetes.io/tls` Secret in the registry's namespace, for +example with cert-manager or from a certificate that your CA issued. The certificate must list, as +a subject alternative name (SAN), the host of every URL in `controller.registryUrls`: the +controller checks the URL's host against the certificate. This page uses +`mcs-registry.example.internal`. When the registry's cluster also runs a controller from the same +release, that controller dials the registry Service and the certificate also needs the name +`mcs-registry.mcs-system.svc`. + +```bash +kubectl create namespace mcs-system +kubectl -n mcs-system create secret tls mcs-registry-tls \ + --cert=registry.crt --key=registry.key +``` + +Install the registry with the store, the table name and the Secret: + +```bash +helm install mcs-registry deploy/helm/charts/global-discovery-mcs \ + --namespace mcs-system \ + --set registry.enabled=true \ + --set controller.enabled=false \ + --set registry.store=dynamodb \ + --set registry.dynamodb.table=global-discovery-mcs \ + --set registry.tls.secretName=mcs-registry-tls \ + --set image.repository=my-registry.example.com/global-discovery-mcs \ + --set image.tag=dev +``` + +With `registry.tls.secretName`, the registry serves TLS on the gRPC and REST ports +(`--tls-cert-file` and `--tls-key-file`). The admin port stays plaintext. The registry reads the key +pair at startup, so restart it after you rotate the certificate. The controllers need the CA +certificate that signed it, as the next section shows. + +The registry takes its credentials from the default AWS chain, and its region from `--region` or +from `AWS_REGION`, which IAM roles for service accounts (IRSA) and EKS Pod Identity set. Bind its +ServiceAccount to the registry role. + +To run the registry in several regions, add a Global Table replica of the table in each region, +as the module's +[storage document](https://github.com/ai-dynamo/dynamo/blob/main/deploy/global-discovery-mcs/docs/storage.md#global-table-replicas) +describes, and install the registry in each region with the same table name. Each registry reads +and writes the replica of its own region. A write usually reaches the other regions within a few +seconds, with no guaranteed bound; see +[Guarantees and Typical Timing](../../developer-guide/knowledge-base/modular-components/router/multi-cluster-services-registry.md#guarantees-and-typical-timing). +Give every regional hostname that a controller dials a SAN in the certificate of the registries +behind it. + +### Expose the Registry + +Expose the gRPC port, `9090`, to the member clusters, for example through an internal load +balancer. The REST port, `8080`, serves the same API as JSON and is useful for inspection; the +admin port, `9091`, serves `/metrics`, `/healthz` and `/readyz`. + +To terminate TLS at a load balancer instead of in the registry, leave `registry.tls.secretName` +empty, give the load balancer a certificate for the host in `controller.registryUrls`, and have it +forward HTTP/2 (gRPC) to port `9090`, for example an NLB TLS listener. The controllers then follow +the TLS path: an `https://` URL, `controller.registryTlsCa` when a private CA signed the load +balancer's certificate, and `controller.registryAllowInsecureHttp` left `false`. The hop from the +load balancer to the registry is plaintext, so keep it on the private network. In the registry's +own cluster, set `controller.registryUrls` to the same `https://` URL, because the default URL of +a release without `registry.tls.secretName` is plaintext. + +### Check the Registry + +Check that the registry is ready. Forward its admin port: + +```bash +kubectl -n mcs-system port-forward deployment/mcs-registry 9091:9091 +``` + +In another terminal: + +```bash +curl -s localhost:9091/readyz +``` + +## Install the Controller in Every Member Cluster + +Install the controller in each member cluster with that cluster's id and region. The clusterset +name is the same everywhere. Use the sequence that matches the registry's transport. + +### Production: TLS + +Give the controller the CA certificate that signed the registry certificate, as a ConfigMap in the +controller's namespace, and set the registry URL to an `https://` URL whose host is one of the +certificate's SANs: + +```bash +kubectl create namespace mcs-system +kubectl -n mcs-system create configmap mcs-registry-ca --from-file=ca.crt=ca.crt +helm install mcs-controller deploy/helm/charts/global-discovery-mcs \ + --namespace mcs-system \ + --set registry.enabled=false \ + --set controller.enabled=true \ + --set 'controller.registryUrls={https://mcs-registry.example.internal:9090}' \ + --set controller.registryTlsCa.configMapName=mcs-registry-ca \ + --set controller.clustersetId=prod \ + --set controller.clusterId=use1-a \ + --set controller.region=us-east-1 \ + --set image.repository=my-registry.example.com/global-discovery-mcs \ + --set image.tag=dev +``` + +When a CA in the system trust roots of the image signed the certificate, skip the ConfigMap and +`controller.registryTlsCa.configMapName`. To read the bundle from a Secret instead, set +`controller.registryTlsCa.secretName`; the bundle's key is `ca.crt` unless you set +`controller.registryTlsCa.key`. + +### Evaluation: plaintext + +Set an `http://` URL and opt in to plaintext explicitly: + +```bash +helm install mcs-controller deploy/helm/charts/global-discovery-mcs \ + --namespace mcs-system --create-namespace \ + --set registry.enabled=false \ + --set controller.enabled=true \ + --set 'controller.registryUrls={http://mcs-registry.example.internal:9090}' \ + --set controller.registryAllowInsecureHttp=true \ + --set controller.clustersetId=prod \ + --set controller.clusterId=use1-a \ + --set controller.region=us-east-1 \ + --set image.repository=my-registry.example.com/global-discovery-mcs \ + --set image.tag=dev +``` + +### In the Registry's Cluster + +In the cluster that runs the registry, install one release with both `registry.enabled=true` and +`controller.enabled=true` instead. Without `controller.registryUrls`, its controller dials the +registry Service of the release, `https://mcs-registry..svc:9090` with +`registry.tls.secretName` and `http://mcs-registry..svc:9090` without it. On the TLS +path, also set `controller.registryTlsCa.configMapName` and give the certificate the name +`mcs-registry..svc`. On the plaintext path, set +`controller.registryAllowInsecureHttp=true`, because the chart refuses the plaintext URL otherwise. + +The values map to `mcs-controller` flags: + +- `controller.registryUrls` (`--registry-url`): the registry's gRPC URLs, in failover order. Put + a global DNS name first and per-region names after it. A URL without a port uses `443` for + `https` and `80` for `http`, so name the port when the registry serves on `9090`. Every URL + uses the same scheme. +- `controller.registryTlsCa` (`--registry-tls-ca-file`): the PEM bundle of the CAs that sign the + registry certificates. When unset, the controller verifies the registry against the system roots. +- `controller.registryAllowInsecureHttp` (`--registry-allow-insecure-http`): permits `http://` + URLs. It defaults to `false`; the chart and the controller both refuse an `http://` URL without + it. +- `controller.clustersetId` (`--clusterset-id`) and `controller.clusterId` (`--cluster-id`): when + unset, the controller reads them from the KEP-2149 ClusterProperty objects `clusterset.k8s.io` + and `cluster.clusterset.k8s.io`. +- `controller.region` (`--region`): when unset, the controller uses the most common + `topology.kubernetes.io/region` node label. + +The controller uses leader election: only the leader writes to the registry and applies imports, +and a second replica takes over when the leader stops. The leader registers the cluster with a +30 s lease and renews it every 10 s. If the renewals stop, the registry removes the cluster's +endpoints from every import once the lease expires. + +Check that every cluster registered. In the registry's cluster, forward its REST port: + +```bash +kubectl -n mcs-system port-forward deployment/mcs-registry 8080:8080 +``` + +In another terminal, list the clusters of the clusterset. Every cluster id appears with +`"state":"LEASE_STATE_ACTIVE"` in its lease: + +```bash +curl -s localhost:8080/v1/clustersets/prod/clusters +``` + +When the registry serves TLS, its REST port does too. Verify the certificate against the CA and +resolve the certificate's host to the forwarded port: + +```bash +curl -s --cacert ca.crt --resolve mcs-registry.example.internal:8080:127.0.0.1 \ + https://mcs-registry.example.internal:8080/v1/clustersets/prod/clusters +``` + +## Export Services + +An exported Service appears as a ServiceImport, with a derived Service and EndpointSlices, in every +member cluster where its namespace exists. Export a Service in one of three ways. + +Label the Service. The controller creates a ServiceExport of the same name, owned by the Service, +and deletes it when the label is removed: + +```bash +kubectl -n dynamo label service my-dgd-pool-relay dynamo.nvidia.com/mcs-export=true +``` + +Or create the ServiceExport yourself, in the namespace and with the name of the Service: + +```yaml +apiVersion: multicluster.x-k8s.io/v1beta1 +kind: ServiceExport +metadata: + name: my-dgd-pool-relay + namespace: dynamo +``` + +Or let the Dynamo operator create the ServiceExports. Set `mcsImplementation` to `mcs` in the +operator configuration (`dynamo-operator.dynamo.multiCluster.mcsImplementation` in the platform +chart), and annotate each DynamoGraphDeployment (DGD) that serves as a Global Router pool: + +```bash +kubectl -n dynamo annotate dynamographdeployment my-dgd nvidia.com/global-router-pool=true +``` + +The operator enables its `multiClusterExport` feature gate when the cluster serves the ServiceExport +API, and then creates ServiceExports for `-pool-relay` and `-frontend` of every +annotated DGD. Removing the annotation deletes them. + +For the Global Discovery Service, export in every cluster: + +- the PoolRelay Service of each pool, `-pool-relay`; +- the Global Router Service, in every cluster that runs Global Router replicas. + +Keep the port names and numbers of an exported Service the same in every cluster. When exports of +one Service disagree, the oldest export wins and every participating cluster's ServiceExport reports +`Conflict=True`. + +A ServiceImport lands in the namespace of the exported Service. In a cluster without that namespace, +the controller holds the import until the namespace is created; `--create-namespaces` creates it +instead. + +## Verify the Imports + +In the exporting cluster, check the ServiceExport's conditions: + +```bash +kubectl -n dynamo get serviceexport my-dgd-pool-relay \ + -o jsonpath='{range .status.conditions[*]}{.type}={.status} {.reason}{"\n"}{end}' +``` + +A working export reports: + +```text +Valid=True Valid +Ready=True Ready +Conflict=False NoConflicts +``` + +In any member cluster, list the ServiceImports, then the derived Services and the imported +EndpointSlices with their source clusters: + +```bash +kubectl -n dynamo get serviceimports +kubectl -n dynamo get services,endpointslices -l multicluster.kubernetes.io/service-name \ + -L multicluster.kubernetes.io/service-name,multicluster.kubernetes.io/source-cluster +``` + +Each exported Service has one derived Service named `derived-` and one set of EndpointSlices +per exporting cluster, the local cluster included. The ServiceExport and ServiceImport objects do +not carry the `multicluster.kubernetes.io/service-name` label, so list them without the selector. + +## Run the Global Discovery Service with MCS + +The Global Discovery Service does not call the registry. Run it in any member cluster where +`mcs-controller` imports the clusterset; it watches the imported EndpointSlices with its pod's +credentials and needs `get`, `list` and `watch` on `endpointslices` in every namespace. + +Start from the example manifest +[`container/global-discovery-service.yaml`](https://github.com/ai-dynamo/dynamo/blob/main/container/global-discovery-service.yaml), +which holds the RBAC, Deployment and Service, and set `--pool-discovery=mcs`: + +```yaml +args: + - --listen=0.0.0.0:8080 + - --pool-discovery=mcs + - --router-namespace=global-discovery + - --router-service=global-router + - --router-port-name=http +``` + +- `--router-namespace` and `--router-service` name the exported Global Router Service, and + `--router-port-name` its TCP port that PoolRelays connect to. All three are required with + `--pool-discovery mcs`. +- `--pool-service-suffix` (default `-pool-relay`) names each pool's PoolRelay Service, + ``. + +Check the lists it serves. Forward its port: + +```bash +kubectl -n global-discovery port-forward service/global-discovery-service 8080:8080 +``` + +In another terminal: + +```bash +curl -s localhost:8080/pools +curl -s localhost:8080/routers +``` + +`/pools` lists one pool per source cluster, namespace and DGD whose PoolRelay Service has a Ready +endpoint; `key.cluster_id` is the cluster id of the exporting cluster. `/routers` lists every Ready +Global Router replica as `/` with its URL. + +The module's end-to-end test, `make -C deploy/global-discovery-mcs e2e`, runs this setup on three local kind clusters. + +## Troubleshooting + +| Symptom | Check | +|---|---| +| ServiceExport `Valid=False`, reason `NoService` | Create the Service with the ServiceExport's name in the same namespace. | +| ServiceExport `Valid=False`, reason `InvalidServiceType` | ExternalName Services cannot be exported. | +| ServiceExport `Ready=Unknown`, reason `Pending`, for more than a few seconds | The controller has not registered the cluster. Check the controller logs, the registry URL and port, that the URL scheme matches the registry's TLS setting, that the certificate names the URL's host and the controller trusts its CA, and network access to the registry. | +| ServiceExport `Ready=False`, reason `Failed` | The registry rejected the export, for example above 64 ports, labels or 32 annotations. The `ExportFailed` Event names the reason: `kubectl -n dynamo events --for serviceexport/my-dgd-pool-relay`. | +| ServiceExport `Ready=False`, reason `ClusterLeaseExpired` | The cluster's heartbeats stopped reaching the registry. Check the controller leader and its logs. | +| ServiceExport `Conflict=True` | The exports of the Service differ between clusters; the reason names the field (`PortConflict`, `TypeConflict` and so on). Align the Services; the oldest export wins in the meantime. | +| No ServiceImport in a cluster | The namespace is missing (the controller holds the import), `--watch-namespaces` excludes it, or a ServiceImport of another MCS implementation exists with that name. The controller skips ServiceImports it did not create unless `--adopt-unmanaged-imports` is set. | +| A cluster's endpoints are missing from an import | Its lease expired or it did not export the Service. List the clusters and imports through the registry's REST port, or run `mcs-admin describe` against the DynamoDB store. | +| The Global Discovery Service lists no pools | Check that the imported EndpointSlices exist in its cluster, that the PoolRelay Service names end with `--pool-service-suffix`, that the endpoints are Ready, and that its ServiceAccount can list EndpointSlices in every namespace. | +| The Global Discovery Service lists no routers | Check `--router-namespace`, `--router-service` and that `--router-port-name` names a TCP port of the exported router Service. | + +To read the controller logs: + +```bash +kubectl -n mcs-system logs deployment/mcs-controller +``` + +To inspect what the registry holds, list the imports through its REST port (forwarded as above, with the TLS options when the registry serves TLS): + +```bash +curl -s 'localhost:8080/v1/clustersets/prod/imports?includeEndpoints=true' +``` + +With the DynamoDB store, `mcs-admin` reads the table directly, with the registry role's +permissions. `describe` lists the clusters with their lease states, the exports and the endpoint +sets: + +```bash +mcs-admin --store=dynamodb --region=us-east-1 describe --clusterset=prod +``` + +To remove a cluster that will not return, together with its exports, run +`mcs-admin --store=dynamodb --region=us-east-1 delete-cluster --clusterset=prod --cluster=use1-a --yes`. + +## CRDs and Other MCS Implementations + +The chart installs three CRDs when `crds.install=true`: ServiceExport and ServiceImport, copied from +[`sigs.k8s.io/mcs-api`](https://github.com/kubernetes-sigs/mcs-api) v0.5.2, and the KEP-2149 +ClusterProperty (`about.k8s.io/v1alpha1`). The mcs-api v0.5.2 CRDs serve both +`multicluster.x-k8s.io/v1alpha1` and `v1beta1`, and store `v1beta1`. The controller reads and +writes `v1beta1`; clients that write `v1alpha1`, such as the Dynamo operator and Karmada, keep +working against the same CRDs. + +Set `crds.install=false` when another installation already owns the CRDs in a cluster, for example +Karmada or another release of this chart. The cluster must still serve +`multicluster.x-k8s.io/v1beta1`; replace a CRD that serves only `v1alpha1` with the mcs-api v0.5.2 +version first. + +Within one cluster, run one MCS implementation per namespace. The controller acts on every +ServiceExport in the namespaces it watches: it adds its finalizer, writes the export to the +registry and sets the export's conditions. Limit it with `--watch-namespaces` where another +implementation exports Services from the same cluster. diff --git a/docs/fern/pages/kubernetes/kv-aware-routing/pool-relay.md b/docs/fern/pages/kubernetes/kv-aware-routing/pool-relay.md index ab64ba1221e5..10804d360cff 100644 --- a/docs/fern/pages/kubernetes/kv-aware-routing/pool-relay.md +++ b/docs/fern/pages/kubernetes/kv-aware-routing/pool-relay.md @@ -317,6 +317,9 @@ protected listener through your network ingress. Keep its upstream connection HT and allow long-lived server streams. See the [gRPC contract](https://github.com/ai-dynamo/dynamo/blob/main/lib/llm/src/pool_relay/docs/grpc-contract.md) for message sizes, reconnect behavior, and error reasons. +Within one trusted network whose clusters route each other's pod addresses, the Kubernetes +Multi-Cluster Services API can make a Relay Service visible in the other clusters instead; see +[Deploy Multi-Cluster Services](multi-cluster-services.md). ## Optional mTLS Sidecar diff --git a/docs/fern/pages/reference/components/global-discovery-mcs-configuration.md b/docs/fern/pages/reference/components/global-discovery-mcs-configuration.md new file mode 100644 index 000000000000..c7438201925f --- /dev/null +++ b/docs/fern/pages/reference/components/global-discovery-mcs-configuration.md @@ -0,0 +1,450 @@ +--- +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +title: Multi-Cluster Services Configuration +subtitle: Flags, labels, conditions, metrics, ports and RBAC of mcs-registry, mcs-controller and mcs-admin +--- + +**Experimental.** This page is the reference for the Kubernetes Multi-Cluster Services (MCS) +implementation in +[`deploy/global-discovery-mcs`](https://github.com/ai-dynamo/dynamo/tree/main/deploy/global-discovery-mcs): +the registry server `mcs-registry`, the per-cluster `mcs-controller`, the `mcs-admin` tool, and the +Global Discovery Service flags that read its imports. The `global-discovery-mcs` image holds the three +binaries (`/mcs-registry`, the entrypoint, `/mcs-controller` and `/mcs-admin`), runs as user +`65532`, and ships the CRDs in `/opt/global-discovery-mcs/crds`. + +For deployment, see [Deploy Multi-Cluster Services](../../kubernetes/kv-aware-routing/multi-cluster-services.md). +For the design, see +[Multi-Cluster Services Registry](../../developer-guide/knowledge-base/modular-components/router/multi-cluster-services-registry.md). +The Helm chart's values are listed in the +[chart README](https://github.com/ai-dynamo/dynamo/blob/main/deploy/helm/charts/global-discovery-mcs/README.md). + +## Ports + +| Component | Port | Serves | +|---|---|---| +| `mcs-registry` | `9090` | gRPC API, package `dynamo.mcs.v1` (h2c, or TLS with `--tls-cert-file`), and the gRPC health service | +| `mcs-registry` | `8080` | REST/JSON gateway of the same API, under `/v1/` | +| `mcs-registry` | `9091` | `/metrics`, `/healthz`, `/readyz`, `/debug/pprof/` | +| `mcs-controller` | `8080` | `/metrics` | +| `mcs-controller` | `8081` | `/healthz`, `/readyz` | + +## mcs-registry Flags + +Every flag has an environment fallback named `MCS_` plus the flag name in upper case with dashes +replaced by underscores. A flag on the command line wins over its environment variable. An invalid +configuration exits with status 2 and lists every invalid setting. + +| Flag | Environment variable | Default | Description | +|---|---|---|---| +| `--grpc-listen-address` | `MCS_GRPC_LISTEN_ADDRESS` | `:9090` | gRPC listener. `host:port`, different from the other two listeners. | +| `--http-listen-address` | `MCS_HTTP_LISTEN_ADDRESS` | `:8080` | REST gateway listener. The gateway dials the gRPC listener over loopback. `host:port`. | +| `--admin-listen-address` | `MCS_ADMIN_LISTEN_ADDRESS` | `:9091` | Admin listener for metrics, health and pprof. `host:port`. | +| `--region` | `MCS_REGION`, then `AWS_REGION` | none | Region returned by `GetServerInfo` and used by the DynamoDB store. | +| `--store` | `MCS_STORE` | `memory` | Store backend: `memory` or `dynamodb`. | +| `--dynamodb-table` | `MCS_DYNAMODB_TABLE` | `global-discovery-mcs` | DynamoDB table. Must not be empty with `--store=dynamodb`. | +| `--dynamodb-endpoint` | `MCS_DYNAMODB_ENDPOINT` | none | DynamoDB endpoint override, for dynamodb-local. An `http` or `https` URL with a host. | +| `--refresh-interval` | `MCS_REFRESH_INTERVAL` | `2s` | Interval of the full reload of each clusterset from the store. At least `100ms`. | +| `--lease-eval-interval` | `MCS_LEASE_EVAL_INTERVAL` | `1s` | Interval of the lease evaluation between refreshes, without a store read. At least `100ms`. | +| `--clusterset-discovery-interval` | `MCS_CLUSTERSET_DISCOVERY_INTERVAL` | `30s` | Interval of the listing of the store's clustersets. A request that names an unknown clusterset loads it at once. At least `1s`. | +| `--lease-skew` | `MCS_LEASE_SKEW` | `2s` | Clock skew allowance before a lease expires. `0s` to `1m`. | +| `--retention` | `MCS_RETENTION` | `24h` | How long the records of an expired cluster are kept before compaction removes them. At least `5m`. | +| `--compaction-interval` | `MCS_COMPACTION_INTERVAL` | `10m` | Interval of the retention cleanup. At least `1s`. | +| `--bookmark-interval` | `MCS_BOOKMARK_INTERVAL` | `5s` | Interval of watch `BOOKMARK` events. `1s` to `5m`. | +| `--watch-history` | `MCS_WATCH_HISTORY` | `10000` | Watch events kept per clusterset for resume. At least `1`. | +| `--watch-queue-size` | `MCS_WATCH_QUEUE_SIZE` | `1024` | Per-watcher event queue. A watcher that overflows it is sent a new snapshot. At least `1`. | +| `--max-watchers` | `MCS_MAX_WATCHERS` | `5000` | Concurrent watch streams per replica. Above it a watch fails with `TOO_MANY_WATCHERS`. At least `1`. | +| `--max-message-bytes` | `MCS_MAX_MESSAGE_BYTES` | `4194304` | gRPC send and receive message size limit. At least `1048576`. | +| `--advertise-endpoints` | `MCS_ADVERTISE_ENDPOINTS` | none | JSON array of `{"region","grpcUrl","httpUrl","primary"}` returned by `GetServerInfo`; clients append the `grpcUrl` values to their URL list as fallbacks. Each entry needs `region` and `grpcUrl` or `httpUrl` (`http` or `https` URLs); at most one entry is `primary`. | +| `--tls-cert-file` | `MCS_TLS_CERT_FILE` | none | TLS certificate of the gRPC and REST listeners. Set together with `--tls-key-file`. | +| `--tls-key-file` | `MCS_TLS_KEY_FILE` | none | TLS private key of the gRPC and REST listeners. | +| `--shutdown-drain` | `MCS_SHUTDOWN_DRAIN` | `5s` | Time the replica reports not ready on SIGTERM before it stops serving. At least `0s`. | +| `--shutdown-timeout` | `MCS_SHUTDOWN_TIMEOUT` | `25s` | Graceful stop of the gRPC and REST servers. At least `1s`. | +| `--otel-enabled` | `MCS_OTEL_ENABLED` | `false` | OpenTelemetry tracing over OTLP, configured by the standard `OTEL_*` environment variables. | +| `--log-level` | `MCS_LOG_LEVEL` | `info` | `debug`, `info`, `warn` or `error`. Logs are JSON. | + +The DynamoDB store takes its credentials from the default AWS credential chain. The gRPC server +pings idle connections after 60 s (timeout 20 s) and accepts client pings every 10 s or more, also +without active streams. + +### Registry Health + +| Check | Passes when | +|---|---| +| `/healthz` | The process serves the admin port. | +| `/readyz` and the gRPC health service | The store answers a ping, every clusterset the store knows was loaded at least once, and no clusterset view is older than three refresh intervals. Fails from the start of shutdown. | + +On SIGTERM a replica reports not ready, waits `--shutdown-drain`, ends open watches with +`UNAVAILABLE` and reason `SHUTTING_DOWN`, stops the REST and gRPC servers within +`--shutdown-timeout`, and flushes traces. Set the pod's `terminationGracePeriodSeconds` above the +sum of the two. + +## mcs-controller Flags + +The controller reads two environment variables: `MCS_REGISTRY_URLS` and `POD_NAMESPACE`, as the +defaults of `--registry-url` and `--operator-namespace`. It has no other environment fallbacks. An +invalid configuration exits with status 2 and lists every invalid setting. + +| Flag | Default | Description | +|---|---|---| +| `--registry-url` | `$MCS_REGISTRY_URLS` | Registry gRPC URLs in failover order, repeatable or comma-separated. Required. Every URL uses the same scheme, `https` or `http`; `http` needs `--registry-allow-insecure-http`. A URL without a port uses `443` (`https`) or `80` (`http`). Any `--registry-url` on the command line replaces the environment value. | +| `--registry-tls-ca-file` | system roots | PEM bundle of the CAs that sign the registry certificates. The file must exist. | +| `--registry-allow-insecure-http` | `false` | Permit plaintext `http://` registry URLs. | +| `--clusterset-id` | none | Clusterset name, a DNS label. Without it, the ClusterProperty `clusterset.k8s.io` provides it. | +| `--cluster-id` | none | Cluster id, a DNS label unique in the clusterset. It is the `multicluster.kubernetes.io/source-cluster` value of this cluster's endpoints in every importing cluster. Without it, the ClusterProperty `cluster.clusterset.k8s.io` provides it. | +| `--cluster-identity-source` | `auto` | Where the two ids come from. `auto`: a flag wins, an id without its flag comes from its ClusterProperty, and a missing ClusterProperty CRD counts as no property. `flags`: both flags are required and ClusterProperty objects are not read. `clusterproperty`: both ClusterProperty objects are required, and flags, when set, must agree with them. | +| `--region` | none | Region of the registration. Empty uses the most common `topology.kubernetes.io/region` node label. | +| `--cluster-labels` | none | Comma-separated `key=value` labels of the registration. Keys are qualified names and values valid label values; at most 64. | +| `--lease-ttl` | `30s` | Cluster lease TTL requested from the registry, `10s` to `5m`. The leader heartbeats every TTL/3, with 10 percent jitter. | +| `--watch-namespaces` | all namespaces | Comma-separated namespaces to export from and import into. Restricts the caches of Services, EndpointSlices, ServiceExports and ServiceImports. | +| `--create-namespaces` | `false` | Create a missing namespace of an import instead of holding the import pending. The controller never deletes namespaces. | +| `--adopt-unmanaged-imports` | `false` | Take over ServiceImports that the controller did not create. | +| `--write-cluster-properties` | `false` | Create the ClusterProperty objects `clusterset.k8s.io` and `cluster.clusterset.k8s.io` from the resolved ids. An existing property with another value stops the controller. | +| `--deregister-on-shutdown` | `false` | Deregister the cluster when the controller stops, which deletes its exports in the registry. | +| `--endpoints-per-slice` | `100` | Endpoints per imported EndpointSlice, `1` to `1000`. | +| `--max-concurrent-reconciles` | `4` | Workers of each controller. At least `1`. | +| `--metrics-bind-address` | `:8080` | Address of `/metrics`. `0` disables it. | +| `--health-probe-bind-address` | `:8081` | Address of `/healthz` and `/readyz`. Required. | +| `--leader-elect` | `true` | Leader election. Only the leader writes to the registry and applies imports; required for more than one replica. | +| `--operator-namespace` | `$POD_NAMESPACE` | Namespace of the leader election Lease, a namespace name. Required with `--leader-elect`. | +| `--otel-enabled` | `false` | OpenTelemetry tracing over OTLP, configured by the standard `OTEL_*` environment variables. | +| `--kubeconfig` | in-cluster configuration | Kubeconfig path, for running out of cluster. | +| `--zap-devel` | `false` | Development logging: console encoder, debug level, stack traces from warn. | +| `--zap-encoder` | `json` | Log encoding, `json` or `console`. | +| `--zap-log-level` | `info` | `debug`, `info`, `error`, `panic`, or an integer for more verbose debug levels. | +| `--zap-stacktrace-level` | `error` | Level at and above which stack traces are logged: `info`, `error` or `panic`. | +| `--zap-time-encoding` | `epoch` | `epoch`, `millis`, `nano`, `iso8601`, `rfc3339` or `rfc3339nano`. | + +### Controller Health + +| Check | Passes when | +|---|---| +| `/healthz` | The process answers. | +| `/readyz` | The registry client connected at least once and, on the leader, the cluster lease is registered. | + +## mcs-admin + +`mcs-admin` reads and repairs a store directly, without a running registry: + +```text +mcs-admin [store flags] describe [--clusterset NAME] [--output text|json] +mcs-admin [store flags] delete-cluster --clusterset NAME --cluster NAME --yes +mcs-admin [store flags] create-table [create-table flags] +``` + +`describe` and `delete-cluster` act on the store the store flags open. The memory store lives in +the registry process, so use them with `--store=dynamodb`; with the default `--store=memory` they +print a warning to stderr and see an empty store. Exit status: `0` on success, `1` on a +failure, `2` on a usage error. + +### Store Flags + +Given before the command. Every store flag has the same `MCS_` environment fallback as the +`mcs-registry` flag of the same name. + +| Flag | Environment variable | Default | Description | +|---|---|---|---| +| `--store` | `MCS_STORE` | `memory` | `memory` or `dynamodb`. | +| `--dynamodb-table` | `MCS_DYNAMODB_TABLE` | `global-discovery-mcs` | DynamoDB table. | +| `--dynamodb-endpoint` | `MCS_DYNAMODB_ENDPOINT` | none | DynamoDB endpoint override, for dynamodb-local. | +| `--region` | `MCS_REGION`, then `AWS_REGION` | none | AWS region. | +| `--lease-skew` | `MCS_LEASE_SKEW` | `2s` | Clock skew allowance when `describe` evaluates lease states. | + +### Commands + +| Command | Flags | Description | +|---|---|---| +| `describe` | `--clusterset` (default: every clusterset), `--output` (`text` or `json`, default `text`) | Lists each clusterset's clusters with region, lease state, epoch, renewal, expiry and generation; its exports with type, generation and export creation time; and its endpoint sets with endpoint and ready counts. | +| `delete-cluster` | `--clusterset`, `--cluster` (DNS labels, required), `--yes` | Deletes the cluster with its exports and endpoint sets. Without `--yes` it prints what it would delete and exits with status 2. | +| `create-table` | See below | Creates the DynamoDB table, waits until it is `ACTIVE`, then enables TTL on `expires_at` and point-in-time recovery. On an existing table it checks the key schema and enables only what is missing. | + +`create-table` takes its own flags after the command. `--table`, `--region` and `--endpoint` default +to the store flags `--dynamodb-table`, `--region` and `--dynamodb-endpoint`, so the order of +precedence is the `create-table` flag, then the store flag, then its `MCS_` environment variable, +then the built-in default: + +| Flag | Default | Description | +|---|---|---| +| `--table` | `--dynamodb-table`, else `global-discovery-mcs` | Table name. | +| `--region` | `--region` store flag, else `AWS_REGION` or the shared AWS configuration | AWS region. | +| `--endpoint` | `--dynamodb-endpoint` store flag, else none | DynamoDB endpoint override, for dynamodb-local. | +| `--deletion-protection` | `true` | Create the table with deletion protection. | +| `--pitr` | `true` | Enable point-in-time recovery. dynamodb-local does not support it; pass `--pitr=false`. | +| `--ttl` | `true` | Enable TTL on `expires_at`. | +| `--wait-timeout` | `5m` | How long to wait for the table to become `ACTIVE`. | + +## Global Discovery Service Flags for MCS + +The Global Discovery Service reads the EndpointSlices that `mcs-controller` imports with these +flags. It does not connect to the registry. + +| Flag | Default | Description | +|---|---|---| +| `--pool-discovery` | `registrations` | `mcs` reads imported EndpointSlices labelled `multicluster.kubernetes.io/source-cluster` and `multicluster.kubernetes.io/service-name`. `karmada` reads Karmada's imports. | +| `--pool-service-suffix` | `-pool-relay` | A pool's PoolRelay Service is named ``. Must not be empty. | +| `--router-namespace` | none | Namespace of the exported Global Router Service. Required with `mcs`. | +| `--router-service` | none | The exported Global Router Service. Required with `mcs`. | +| `--router-port-name` | none | Name of the router Service's TCP port that PoolRelays connect to. Required with `mcs`. | + +With `--pool-discovery mcs`, the flags of the `registrations` mode (`--registration-storage`, +`--heartbeat-interval-secs`, `--lease-duration-secs` and `--router`) fail at startup, and +`POST /pools` answers `501`. The service needs `get`, `list` and `watch` on +`discovery.k8s.io` `endpointslices` in every namespace. + +## Labels, Annotations and Names + +| Key or name | On | Value | +|---|---|---| +| `dynamo.nvidia.com/mcs-export` | Service | `"true"` makes the controller create and own a ServiceExport of the same name. | +| Other `dynamo.nvidia.com/mcs-*` labels and annotations | Service | Forwarded to the registry as exported labels and annotations, with `spec.exportedLabels` and `spec.exportedAnnotations` of the ServiceExport, which win on the same key. | +| `dynamo.nvidia.com/mcs-managed-by` | ServiceExports, ServiceImports, derived Services, EndpointSlices and namespaces the controller creates | `mcs-controller`. Not forwarded. | +| `dynamo.nvidia.com/mcs-clusterset` | The same objects | The clusterset. Not forwarded. | +| `dynamo.nvidia.com/mcs-service-export` | ServiceExport finalizer | The controller deletes the registry record before removing it. | +| `multicluster.kubernetes.io/service-name` | Derived Services and imported EndpointSlices | Name of the exported Service. | +| `multicluster.kubernetes.io/source-cluster` | Imported EndpointSlices | Cluster id of the exporting cluster. | +| `kubernetes.io/service-name` | Imported EndpointSlices | Name of the derived Service. | +| `endpointslice.kubernetes.io/managed-by` | Imported EndpointSlices | `mcs-controller.dynamo.nvidia.com`. The controller never exports slices with this value. | +| `multicluster.kubernetes.io/derived-service` | ServiceImport annotation | Name of the derived Service. | +| Derived Service name | Importing clusters | `derived-` plus the first ten characters of the lowercase, unpadded base32hex encoding of the SHA-256 of `/`, as in `sigs.k8s.io/mcs-api`. | +| Imported EndpointSlice name | Importing clusters | `--`. | +| `targetRef.kind`, `targetRef.name`, `targetRef.uid` | Endpoint attributes in the registry | The exported endpoint's target reference. The importer sets the imported endpoint's `targetRef` from the name, the kind (`Pod` when absent) and `targetRef.namespace` when present. | +| `clusterset.k8s.io`, `cluster.clusterset.k8s.io` | ClusterProperty objects (`about.k8s.io/v1alpha1`) | Clusterset name and cluster id. | +| `mcs-controller` | Server-side apply field manager | | +| `mcs-controller.dynamo.nvidia.com` | Leader election Lease, in `--operator-namespace` | | +| `mcs.dynamo.nvidia.com` | Domain of every registry error's `ErrorInfo` | | + +## Conditions + +### ServiceExport + +The controller sets `Valid` itself and copies `Ready` and `Conflict` from the status the registry +reports for its cluster. + +| Type | Status | Reason | Meaning | +|---|---|---|---| +| `Valid` | `True` | `Valid` | The Service exists and is exportable. | +| `Valid` | `False` | `NoService` | No Service of the same name exists. The registry record is deleted. | +| `Valid` | `False` | `InvalidServiceType` | The Service is of type ExternalName. The registry record is deleted. | +| `Ready` | `True` | `Ready` | The registry holds the export and the cluster's lease is active. | +| `Ready` | `Unknown` | `Pending` | The registry has not reported the export yet, or the cluster is not registered. | +| `Ready` | `False` | `Failed` | The registry rejected the export, or the export is not valid. | +| `Ready` | `False` | `ClusterLeaseExpired` | The cluster's lease expired; its endpoints left every import. | +| `Conflict` | `False` | `NoConflicts` | The exports of the Service agree. | +| `Conflict` | `True` | A conflict reason | The exports of the Service disagree. Set on every participating cluster, whichever cluster's export differs. | +| `Conflict` | `Unknown` | `ClusterLeaseExpired` | The cluster's lease expired. | + +An invalid ServiceExport has no `Conflict` condition. Conflict reasons, in the order the registry +checks them; it reports the first that applies: + +| Reason | The exports differ in | +|---|---| +| `PortConflict` | Ports: a name or app protocol of the same (port, protocol), or a name used by another (port, protocol) | +| `IPFamilyConflict` | IP families, with no family in common | +| `TypeConflict` | Type: ClusterSetIP or Headless | +| `SessionAffinityConflict` | Session affinity | +| `SessionAffinityConfigConflict` | Session affinity client IP timeout | +| `AnnotationsConflict` | Exported annotations | +| `LabelsConflict` | Exported labels | +| `InternalTrafficPolicyConflict` | Internal traffic policy | +| `TrafficDistributionConflict` | Traffic distribution | + +`IPFamilyConflict` and `ClusterLeaseExpired` are registry additions; the other reasons are the +`sigs.k8s.io/mcs-api` constants. + +### ServiceImport + +| Type | Status | Reason | Meaning | +|---|---|---|---| +| `Ready` | `True` | `Ready` | At least one exporting cluster participates. The message counts them. | +| `Ready` | `False` | `Pending` | No exporting cluster participates: their leases expired. | +| `Conflict` | `False` | `NoConflicts` | The exports agree. | +| `Conflict` | `True` | A conflict reason | The exports disagree. | + +`status.clusters` lists the participating clusters. + +### Events + +The controller writes Events on ServiceExports through the `events.k8s.io` API, with reporting +controller `mcs-controller` and action `Export`. + +| Reason | Type | When | +|---|---|---| +| `InvalidExport` | Warning | The ServiceExport became invalid (`NoService` or `InvalidServiceType`). | +| `EndpointsTruncated` | Warning | The Service has more than 1000 endpoints. Ready endpoints are kept first. | +| `ExportFailed` | Warning | The registry rejected a write as invalid. | + +## Registry API + +The gRPC services are `ClusterService`, `ServiceExportService`, `ServiceImportService` and +`ServerInfoService` in package `dynamo.mcs.v1`. The REST gateway maps them as follows, where `{cs}` +is the clusterset; every path segment is a DNS-1123 label: + +| RPC | Method and path | +|---|---| +| `RegisterCluster` | `PUT /v1/clustersets/{cs}/clusters/{cluster}` | +| `HeartbeatCluster` | `POST /v1/clustersets/{cs}/clusters/{cluster}:heartbeat` | +| `DeregisterCluster` | `DELETE /v1/clustersets/{cs}/clusters/{cluster}` | +| `GetCluster` | `GET /v1/clustersets/{cs}/clusters/{cluster}` | +| `ListClusters` | `GET /v1/clustersets/{cs}/clusters` | +| `WatchClusters` | `GET /v1/clustersets/{cs}/clusters:watch` | +| `ListClusterSets` | `GET /v1/clustersets` | +| `PutServiceExport` | `PUT /v1/clustersets/{cs}/clusters/{cluster}/exports/{namespace}/{name}` | +| `DeleteServiceExport` | `DELETE /v1/clustersets/{cs}/clusters/{cluster}/exports/{namespace}/{name}` | +| `ListServiceExports` | `GET /v1/clustersets/{cs}/clusters/{cluster}/exports` | +| `ReplaceEndpoints` | `PUT /v1/clustersets/{cs}/clusters/{cluster}/exports/{namespace}/{name}/endpoints` | +| `GetServiceImport` | `GET /v1/clustersets/{cs}/imports/{namespace}/{name}` | +| `ListServiceImports` | `GET /v1/clustersets/{cs}/imports` | +| `WatchServiceImports` | `GET /v1/clustersets/{cs}/imports:watch` | +| `GetServerInfo` | `GET /v1/serverinfo` | + +JSON follows the proto3 JSON mapping: lowerCamelCase field names, enums as their full names (for +example `SERVICE_TYPE_CLUSTER_SET_IP`), 64-bit integers as strings. GET filters are query parameters, +for example `?namespace=dynamo&includeEndpoints=true`. A REST watch is a chunked response of +newline-delimited JSON objects, each wrapped in `result`. + +### Errors + +Every error is a `google.rpc.Status` with a `google.rpc.ErrorInfo` detail in domain +`mcs.dynamo.nvidia.com`: + +| Reason | gRPC code | HTTP | Meaning | +|---|---|---|---| +| `STALE_GENERATION` | `FAILED_PRECONDITION` | 400 | The write's generation does not exceed the stored one; `metadata.current_generation` holds it. | +| `LEASE_EPOCH_MISMATCH` | `ABORTED` | 409 | The heartbeat's lease epoch is not current. Register again. | +| `CLUSTER_NOT_REGISTERED` | `FAILED_PRECONDITION` or `NOT_FOUND` | 400 or 404 | The cluster has no registration. | +| `TOO_MANY_ENDPOINTS` | `INVALID_ARGUMENT` | 400 | The endpoint set exceeds 1000 endpoints. | +| `TOO_MANY_WATCHERS` | `RESOURCE_EXHAUSTED` | 429 | The replica is at `--max-watchers`. | +| `STORE_UNAVAILABLE` | `UNAVAILABLE` | 503 | The store did not answer. | +| `SHUTTING_DOWN` | `UNAVAILABLE` | 503 | The replica is draining. | +| `INVALID_ARGUMENT` | `INVALID_ARGUMENT` | 400 | Invalid input; a `google.rpc.BadRequest` detail names every invalid field. | +| `NOT_FOUND` | `NOT_FOUND` | 404 | No such ServiceImport. | +| `INTERNAL` | `INTERNAL` | 500 | An unexpected error. | + +### Limits + +| Item | Limit | +|---|---| +| Clusterset, cluster, namespace and name | DNS-1123 label | +| Labels per object | 64 | +| Annotations per object | 32, at most 16 KiB in total | +| Ports per export | 64 | +| Endpoints per export per cluster | 1000 | +| Attributes per endpoint | 16 | +| Message size | 4 MiB (`--max-message-bytes`) | +| Cluster lease TTL | 10 s to 5 m, default 30 s | + +## Metrics + +Counters and labelled gauges are created with every known label value, so they appear on the first +scrape. + +### mcs-registry + +Served on the admin port, `9091`. + +| Metric | Type | Labels | Meaning | +|---|---|---|---| +| `dynamo_mcs_registry_clusters` | gauge | `clusterset`, `state` (`active`, `expired`) | Registered clusters by lease state | +| `dynamo_mcs_registry_exports` | gauge | `clusterset` | ServiceExport records | +| `dynamo_mcs_registry_endpoints` | gauge | `clusterset`, `condition` (`ready`, `not_ready`) | Endpoints in stored endpoint sets | +| `dynamo_mcs_registry_service_imports` | gauge | `clusterset`, `conflict` (`true`, `false`) | Merged ServiceImports | +| `dynamo_mcs_registry_lease_expirations_total` | counter | `kind` (`cluster`) | Leases this replica saw expire | +| `dynamo_mcs_registry_watchers` | gauge | `kind` (`cluster`, `service_import`) | Open watch streams | +| `dynamo_mcs_registry_watch_events_sent_total` | counter | `kind`, `type` (`snapshot_start`, `upsert`, `delete`, `snapshot_end`, `bookmark`) | Watch events sent | +| `dynamo_mcs_registry_watch_resyncs_total` | counter | `kind`, `reason` (`lagged`, `revision_unknown`, `revision_expired`) | Snapshots sent instead of deltas | +| `dynamo_mcs_registry_watch_rejected_total` | counter | `reason` (`too_many_watchers`, `shutting_down`) | Watch streams refused | +| `dynamo_mcs_registry_view_refresh_duration_seconds` | histogram | | Duration of one clusterset refresh from the store | +| `dynamo_mcs_registry_view_refresh_failures_total` | counter | | Failed clusterset refreshes and discoveries | +| `dynamo_mcs_registry_view_age_seconds` | gauge | | Age of the oldest clusterset view since its last successful refresh | +| `dynamo_mcs_registry_store_request_duration_seconds` | histogram | `operation`, `result` (`ok`, `conditional`, `error`) | Store request duration | +| `dynamo_mcs_registry_store_conditional_failures_total` | counter | `operation`, `reason` | Store writes refused by their condition: `renew_cluster_lease` (`lease_epoch_mismatch`, `not_found`), `put_export` (`stale_generation`), `put_endpoint_set` (`stale_generation`, `too_large`) | +| `dynamo_mcs_registry_compaction_deleted_total` | counter | `kind` (`cluster`, `export`, `endpoint_set`) | Records removed by compaction | +| `dynamo_mcs_registry_build_info` | gauge | `version`, `git_hash` | Build metadata; always 1 | + +Store operations: `put_cluster`, `renew_cluster_lease`, `get_cluster`, `delete_cluster`, +`put_export`, `delete_export`, `put_endpoint_set`, `delete_endpoint_set`, `list_cluster_sets`, +`load_cluster_set` and `ping`. + +The gRPC server adds `dynamo_mcs_registry_grpc_server_started_total`, +`dynamo_mcs_registry_grpc_server_handled_total` (with `grpc_code`), +`dynamo_mcs_registry_grpc_server_msg_received_total`, +`dynamo_mcs_registry_grpc_server_msg_sent_total` and the histogram +`dynamo_mcs_registry_grpc_server_handling_seconds`, labelled `grpc_type`, `grpc_service` and +`grpc_method`. + +### mcs-controller + +Served on `--metrics-bind-address`, `:8080`. + +| Metric | Type | Labels | Meaning | +|---|---|---|---| +| `dynamo_mcs_controller_exports` | gauge | `state` (`exported`, `pending`, `invalid`, `failed`) | ServiceExports by state | +| `dynamo_mcs_controller_reconcile_duration_seconds` | histogram | `controller` (`serviceexport`, `autoexport`, `serviceimport`) | Reconcile duration | +| `dynamo_mcs_controller_reconciles_total` | counter | `controller`, `result` (`success`, `error`) | Reconciles | +| `dynamo_mcs_controller_export_writes_total` | counter | `op` (`put_export`, `delete_export`), `result` (`success`, `stale_generation`, `not_registered`, `error`) | `PutServiceExport` and `DeleteServiceExport` calls | +| `dynamo_mcs_controller_endpoint_replacements_total` | counter | `result` (`success`, `stale_generation`, `not_registered`, `error`) | `ReplaceEndpoints` calls | +| `dynamo_mcs_controller_exported_endpoints` | gauge | | Endpoints of all exports as last written | +| `dynamo_mcs_controller_imports` | gauge | | ServiceImports in the controller's view of the registry | +| `dynamo_mcs_controller_import_events_total` | counter | `type` (`snapshot_start`, `upsert`, `delete`, `snapshot_end`, `bookmark`) | `WatchServiceImports` events received | +| `dynamo_mcs_controller_imports_pending_namespace` | gauge | | Imports waiting for their namespace | +| `dynamo_mcs_controller_stream_reconnects_total` | counter | | Registry client reconnects after the first connection | +| `dynamo_mcs_controller_resyncs_total` | counter | `reason` (`startup`, `reconnect`, `reregister`, `snapshot`) | Full resyncs of the exports for `startup`, `reconnect` and `reregister`; restarts of the import watch with a new snapshot for `snapshot` | +| `dynamo_mcs_controller_gc_deletions_total` | counter | `kind` (`registry_export`, `service_import`, `service`, `endpoint_slice`) | Objects deleted because their source is gone | +| `dynamo_mcs_controller_heartbeat_failures_total` | counter | `reason` (`epoch_mismatch`, `not_registered`, `error`) | Failed cluster heartbeats | +| `dynamo_mcs_controller_registrations_total` | counter | `result` (`success`, `error`) | `RegisterCluster` calls | +| `dynamo_mcs_controller_cluster_registered` | gauge | | 1 while this replica holds a registered cluster lease | + +controller-runtime adds its `controller_runtime_*` and `workqueue_*` metrics on the same endpoint. + +### Registry Client + +The controller's registry client adds these metrics on the controller's endpoint: + +| Metric | Type | Labels | Meaning | +|---|---|---|---| +| `dynamo_mcs_client_requests_total` | counter | `rpc`, `code` | Calls by RPC and gRPC status code; a watch counts once, when it ends | +| `dynamo_mcs_client_request_duration_seconds` | histogram | `rpc` | Unary call duration, retries included | +| `dynamo_mcs_client_endpoint_failovers_total` | counter | | Connections to a different registry URL than the previous one, the return to the primary included | +| `dynamo_mcs_client_connected` | gauge | `url` | 1 for the registry URL in use, 0 for the other configured and advertised URLs | +| `dynamo_mcs_client_watch_reconnects_total` | counter | `watch` (`clusters`, `service_imports`) | Watch streams reopened | + +## RBAC + +`mcs-registry` does not use the Kubernetes API. With the DynamoDB store, its IAM role needs +`dynamodb:GetItem`, `Query`, `PutItem`, `UpdateItem`, `DeleteItem` and `BatchWriteItem` on the +table in each region; the module's +[IAM document](https://github.com/ai-dynamo/dynamo/blob/main/deploy/global-discovery-mcs/docs/iam.md) +has the policies for the registry and for `mcs-admin create-table`. + +`mcs-controller` needs the ClusterRole below, generated from the code into +`deploy/global-discovery-mcs/config/rbac/role.yaml`, and a Role in its own namespace for leader +election on `coordination.k8s.io` `leases` (`get`, `list`, `watch`, `create`, `update`, `patch`, +`delete`). + +| API group | Resources | Verbs | Used for | +|---|---|---|---| +| `multicluster.x-k8s.io` | `serviceexports`, `serviceimports` | `get`, `list`, `watch`, `create`, `update`, `patch`, `delete` | Exports, label-driven exports and imports | +| `multicluster.x-k8s.io` | `serviceexports/status`, `serviceimports/status` | `get`, `update`, `patch` | Conditions and import status | +| `multicluster.x-k8s.io` | `serviceexports/finalizers` | `update` | The ServiceExport finalizer | +| core | `services` | `get`, `list`, `watch`, `create`, `update`, `patch`, `delete` | Exported Services and derived Services | +| `discovery.k8s.io` | `endpointslices` | `get`, `list`, `watch`, `create`, `update`, `patch`, `delete` | Exported endpoints and imported EndpointSlices | +| core | `namespaces` | `get`, `list`, `watch`, `create` | Pending imports; `create` only for `--create-namespaces` | +| core | `nodes` | `get`, `list`, `watch` | Region and zones | +| `about.k8s.io` | `clusterproperties` | `get`, `list`, `watch`, `create`, `update`, `patch` | Identity; writes only for `--write-cluster-properties` | +| core, `events.k8s.io` | `events` | `create`, `patch` | Events on ServiceExports | + +The Global Discovery Service needs `get`, `list` and `watch` on `discovery.k8s.io` `endpointslices` +in every namespace. + +## CRDs + +The controller needs three CRDs, which the image ships in `/opt/global-discovery-mcs/crds`: + +| CRD | Versions | Source | +|---|---|---| +| `serviceexports.multicluster.x-k8s.io` | `v1alpha1` and `v1beta1` served, `v1beta1` stored | `sigs.k8s.io/mcs-api` v0.5.2 | +| `serviceimports.multicluster.x-k8s.io` | `v1alpha1` and `v1beta1` served, `v1beta1` stored | `sigs.k8s.io/mcs-api` v0.5.2 | +| `clusterproperties.about.k8s.io` | `v1alpha1` | KEP-2149 | + +The controller reads and writes `multicluster.x-k8s.io/v1beta1`, so a cluster where another +installation owns these CRDs must serve that version. diff --git a/lib/global-discovery-service/Cargo.toml b/lib/global-discovery-service/Cargo.toml index dab58d34a06a..1add7153bb04 100644 --- a/lib/global-discovery-service/Cargo.toml +++ b/lib/global-discovery-service/Cargo.toml @@ -9,3 +9,31 @@ edition.workspace = true authors.workspace = true license.workspace = true repository.workspace = true + +[[bin]] +name = "global-discovery-service" +path = "src/main.rs" + +[dependencies] +dynamo-global-common = { workspace = true } +dynamo-http-server = { workspace = true } +dynamo-runtime = { workspace = true } + +anyhow = { workspace = true } +async-stream = { workspace = true } +async-trait = { workspace = true } +axum = { workspace = true } +clap = { version = "4.5.49", features = ["derive"] } +futures = { workspace = true } +k8s-openapi = { workspace = true } +kube = { workspace = true } +thiserror = { workspace = true } +tokio = { workspace = true } +tokio-util = { workspace = true } +tracing = { workspace = true } +url = { workspace = true } + +[dev-dependencies] +reqwest = { workspace = true } +serde_json = { workspace = true } +tokio = { workspace = true, features = ["test-util"] } diff --git a/lib/global-discovery-service/src/backends.rs b/lib/global-discovery-service/src/backends.rs new file mode 100644 index 000000000000..085ddbfaf94e --- /dev/null +++ b/lib/global-discovery-service/src/backends.rs @@ -0,0 +1,9 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Implementations of [`crate::pool_discovery`], one module per +//! backend. + +pub mod kubernetes; +pub mod memory; +mod stream; diff --git a/lib/global-discovery-service/src/backends/kubernetes.rs b/lib/global-discovery-service/src/backends/kubernetes.rs new file mode 100644 index 000000000000..bbf3a1a0dd44 --- /dev/null +++ b/lib/global-discovery-service/src/backends/kubernetes.rs @@ -0,0 +1,49 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! The MCS backend, using SIG Multicluster: pools ([`McsPoolDirectory`]) and +//! Global Router replicas ([`McsReplicaDirectory`]) from the EndpointSlices +//! that MCS imports, named as [`McsImplementation`] describes. +//! +//! It assumes a trusted network where pod addresses are reachable between +//! clusters (for example Azure CNI with peered VNets, or the AWS VPC CNI with +//! peered VPCs), and an MCS implementation, Karmada first. +//! +//! Each pool's PoolRelay Service, named `-pool-relay` by default, +//! needs a `ServiceExport`, and so does the router's Service in the hub, so +//! its replicas appear in every importing cluster. A pool is live while a +//! PoolRelay endpoint is Ready. Its `relay_identities` are that PoolRelay +//! pod's addresses, for routers on a trusted network to admit it. +//! +//! With Karmada, a Service name may have at most 55 characters: DGD names up +//! to 44 with the default `-pool-relay` suffix. Karmada imports a Service as +//! `derived-`, and that name may have at most 63 characters, so +//! Karmada imports nothing for a longer name, and the backend never sees that +//! Service's endpoints. [`McsReplicaDirectory::new`] logs an error for a +//! Service whose name is too long. +//! +//! Create the PoolRelay Service and its `ServiceExport` in the workload +//! cluster, not through Karmada propagation. Kubernetes copies a Service's +//! labels onto its EndpointSlices, so a Karmada-propagated Service gives its +//! slices Karmada's labels. In a test with Karmada v1.19, deleting such a +//! `ServiceExport` made karmada-agent delete the Service's EndpointSlice +//! too. Exporting the Service again then found no slice until the Service +//! changed. +//! +//! ```ignore +//! let backend = PoolDiscoveryBackend { +//! pools: Arc::new(McsPoolDirectory::new(client.clone(), mcs, suffix)), +//! routers: Arc::new(McsReplicaDirectory::new(client, &router, mcs)), +//! registrar: None, +//! }; +//! service::serve(listener, backend, None, shutdown).await?; +//! ``` + +mod cache; +mod endpoints; +mod hub; +mod mcs; + +pub use endpoints::{ExportedService, McsReplicaDirectory}; +pub use hub::McsPoolDirectory; +pub use mcs::McsImplementation; diff --git a/lib/global-discovery-service/src/backends/kubernetes/cache.rs b/lib/global-discovery-service/src/backends/kubernetes/cache.rs new file mode 100644 index 000000000000..e1dd113b65be --- /dev/null +++ b/lib/global-discovery-service/src/backends/kubernetes/cache.rs @@ -0,0 +1,553 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! One EndpointSlice watch feeding a derived index, shared by the pool and +//! router replica directories. Each slice event updates only what that slice +//! affects; a relist rebuilds the index once, at `InitDone`. + +use std::collections::HashMap; +use std::fmt; +use std::sync::{Arc, OnceLock, PoisonError, RwLock, RwLockReadGuard}; +use std::time::Duration; + +use crate::pool_discovery::{ + Completeness, Directory, DiscoveryError, Snapshot, WatchEvent, WatchStream, +}; +use async_trait::async_trait; +use futures::{Stream, StreamExt}; +use k8s_openapi::api::discovery::v1::{Endpoint, EndpointSlice}; +use kube::Api; +use kube::api::ListParams; +use kube::runtime::{WatchStreamExt, watcher}; +use tokio::sync::broadcast; +use tokio::time::Instant; +use tokio_util::sync::CancellationToken; + +use super::super::stream::{EVENT_BUFFER, wake_at, watch_stream}; + +/// How long the watch may keep failing before the cache reports its last view +/// as `Estimated`. It keeps serving that view, and the service stays ready. +pub(super) const STALE_AFTER: Duration = Duration::from_secs(60); + +/// How long without an error before a failing watch counts as working again. +/// kube resumes a watch without an event, so silence is the only sign of it. +/// While the watch fails, errors come at least every 30 s backoff plus the +/// time an attempt takes to fail: within the client's connect and read +/// timeouts, which together should be under 90 s (the service's are 10 s and +/// 75 s), since a watch asks the server to end it after +/// [`WATCH_TIMEOUT_SECS`]. +pub(super) const RECOVERED_AFTER: Duration = Duration::from_secs(120); + +/// Each watch asks the server to end it after this long and starts again, so +/// a server that stops answering fails within the client's read timeout +/// rather than kube's default of nearly five minutes. +const WATCH_TIMEOUT_SECS: u32 = 50; + +/// Slices are keyed by (namespace, name). +pub(super) type SliceKey = (String, String); + +pub(super) fn slice_key(slice: &EndpointSlice) -> SliceKey { + ( + slice.metadata.namespace.clone().unwrap_or_default(), + slice.metadata.name.clone().unwrap_or_default(), + ) +} + +/// A slice's named TCP port, unless it is 0. One slice lists one set of ports +/// for all its endpoints. +pub(super) fn tcp_port(slice: &EndpointSlice, name: &str) -> Option { + slice + .ports + .iter() + .flatten() + .find(|port| { + port.name.as_deref() == Some(name) + && port + .protocol + .as_deref() + .is_none_or(|protocol| protocol == "TCP") + }) + .and_then(|port| port.port) + .and_then(|port| u16::try_from(port).ok()) + .filter(|port| *port != 0) +} + +/// The slice's Ready endpoints. The EndpointSlice API reads a missing `ready` +/// as ready. +pub(super) fn ready_endpoints(slice: &EndpointSlice) -> impl Iterator { + slice.endpoints.iter().filter(|endpoint| { + endpoint + .conditions + .as_ref() + .and_then(|conditions| conditions.ready) + .unwrap_or(true) + }) +} + +/// What a directory derives from slices. +pub(super) trait SliceIndex: Default + Send + Sync + 'static { + type Item: Clone + Send + Sync + 'static; + type Key: Clone + Send + Sync + 'static; + /// Turns a slice into an entry. + type Rules: Send + Sync + 'static; + /// What one slice contributes. + type Entry: Send + Sync + 'static; + + /// `None` for a slice this index does not use. + fn entry(rules: &Self::Rules, slice: &EndpointSlice) -> Option; + + /// One slice changed, or went away (`None`). Touches only what that slice + /// affects and returns the events to publish. + fn upsert( + &mut self, + rules: &Self::Rules, + key: SliceKey, + entry: Option, + ) -> Vec>; + + /// A relist: every slice at once. The cache publishes one `Resync`. + fn replace(&mut self, rules: &Self::Rules, entries: HashMap); + + fn snapshot(&self) -> Snapshot; +} + +pub(super) struct CacheState { + pub(super) synced: bool, + /// The current run of watch errors. `None` while the watch works. + failures: Option, + /// Slices listed since the current relist began; they replace the index + /// at `InitDone`. + relist: Option>, + pub(super) index: I, +} + +impl Default for CacheState { + fn default() -> Self { + Self { + synced: false, + failures: None, + relist: None, + index: I::default(), + } + } +} + +impl CacheState { + /// `None` until the first list completes. + fn snapshot(&self) -> Option> { + let mut snapshot = self.synced.then(|| self.index.snapshot())?; + if self + .failures + .is_some_and(|failures| failures.stale(Instant::now())) + { + snapshot.completeness = Completeness::Estimated; + } + Some(snapshot) + } + + /// The index with the given completeness, for a `Resync` that tells + /// watchers it changed. `None` until the first list completes. + fn resync(&self, completeness: Completeness) -> Option> { + let mut snapshot = self.synced.then(|| self.index.snapshot())?; + snapshot.completeness = completeness; + Some(WatchEvent::Resync(snapshot)) + } +} + +/// Watch errors no more than `RECOVERED_AFTER` apart. +#[derive(Clone, Copy)] +struct Failures { + first: Instant, + last: Instant, +} + +impl Failures { + /// The errors span `STALE_AFTER`, so the view turned `Estimated`. + fn long(&self) -> bool { + self.last.saturating_duration_since(self.first) >= STALE_AFTER + } + + /// A long run whose last error is recent. One error followed by silence + /// is a watch that resumed. + fn stale(&self, now: Instant) -> bool { + now.saturating_duration_since(self.last) < RECOVERED_AFTER && self.long() + } +} + +/// The index kept current from slice events. Shared by the watch task, `list` +/// and every `list_and_watch` stream. +pub(super) struct SliceCache { + rules: I::Rules, + state: RwLock>, + /// Sent while holding the `state` write lock. A subscriber that takes a + /// snapshot and subscribes under the read lock therefore misses no event + /// and sees none twice. + pub(super) events: broadcast::Sender>, +} + +impl SliceCache { + pub(super) fn new(rules: I::Rules) -> Self { + Self { + rules, + state: RwLock::default(), + events: broadcast::channel(EVENT_BUFFER).0, + } + } + + pub(super) fn read(&self) -> RwLockReadGuard<'_, CacheState> { + self.state.read().unwrap_or_else(PoisonError::into_inner) + } + + /// `None` until the first list completes. + pub(super) fn snapshot(&self) -> Option> { + self.read().snapshot() + } + + /// Applies one watch event. Returns `true` for a relist that found + /// nothing. + pub(super) fn apply(&self, event: watcher::Event) -> bool { + let mut state = self.state.write().unwrap_or_else(PoisonError::into_inner); + // `Apply`, `Delete` and `InitDone` show that the watch works and the + // index is current. A relist's slices replace the index only at + // `InitDone`, so `InitApply` does not end a run of errors, and kube + // sends `Init` before each relist attempt, before any request. + let ended_long_run = !matches!(event, watcher::Event::Init | watcher::Event::InitApply(_)) + && state + .failures + .take() + .is_some_and(|failures| failures.long()); + let mut listed_nothing = false; + let mut events = match event { + watcher::Event::Init => { + state.relist = Some(HashMap::new()); + Vec::new() + } + watcher::Event::InitApply(slice) => { + // The watch drops `Init` (see `retried`), and a watcher that + // streams its lists sends none, so the first slice starts the + // relist. + if let Some(entry) = I::entry(&self.rules, &slice) { + state + .relist + .get_or_insert_with(HashMap::new) + .insert(slice_key(&slice), entry); + } + Vec::new() + } + watcher::Event::InitDone => { + let entries = state.relist.take().unwrap_or_default(); + state.index.replace(&self.rules, entries); + state.synced = true; + let snapshot = state.index.snapshot(); + listed_nothing = snapshot.items.is_empty(); + vec![WatchEvent::Resync(snapshot)] + } + watcher::Event::Apply(slice) => { + let entry = I::entry(&self.rules, &slice); + state.index.upsert(&self.rules, slice_key(&slice), entry) + } + watcher::Event::Delete(slice) => { + state.index.upsert(&self.rules, slice_key(&slice), None) + } + }; + // Watchers learn that the view is complete again, as a relist would + // tell them. + if ended_long_run && !matches!(events.last(), Some(WatchEvent::Resync(_))) { + events.extend(state.resync(Completeness::Exact)); + } + for event in events { + // No subscribers is not an error. + let _ = self.events.send(event); + } + listed_nothing + } + + /// The watch failed. Returns how long it has been failing. + pub(super) fn failed(&self) -> Duration { + let now = Instant::now(); + let mut state = self.state.write().unwrap_or_else(PoisonError::into_inner); + // A relist that failed starts over, so its slices so far are dropped. + state.relist = None; + // A run that has been quiet for `RECOVERED_AFTER` is over, though + // `quiet` may not have run yet; this error starts a new one. + self.end_quiet_run(&mut state, now); + let (failures, was_long) = match state.failures { + Some(failures) => ( + Failures { + last: now, + ..failures + }, + failures.long(), + ), + None => ( + Failures { + first: now, + last: now, + }, + false, + ), + }; + state.failures = Some(failures); + // Watchers learn that the view turned `Estimated`. + if failures.long() + && !was_long + && let Some(event) = state.resync(Completeness::Estimated) + { + // No subscribers is not an error. + let _ = self.events.send(event); + } + now.saturating_duration_since(failures.first) + } + + /// Ends a run of errors that has been quiet for `RECOVERED_AFTER`: kube + /// resumed the watch without an event. Watchers learn that the view is + /// complete again. + pub(super) fn quiet(&self) { + let mut state = self.state.write().unwrap_or_else(PoisonError::into_inner); + self.end_quiet_run(&mut state, Instant::now()); + } + + /// Ends the run of errors if it has been quiet for `RECOVERED_AFTER`, and + /// tells watchers if it had made the view `Estimated`. + fn end_quiet_run(&self, state: &mut CacheState, now: Instant) { + let Some(failures) = state.failures else { + return; + }; + if now.saturating_duration_since(failures.last) < RECOVERED_AFTER { + return; + } + state.failures = None; + if failures.long() + && let Some(event) = state.resync(Completeness::Exact) + { + // No subscribers is not an error. + let _ = self.events.send(event); + } + } + + /// A watch stream that starts at the first complete list and ends when + /// `cancel` (the caller's) or `stop` (the directory's) is cancelled. + pub(super) fn subscribe( + self: &Arc, + cancel: CancellationToken, + stop: CancellationToken, + ) -> WatchStream { + let cache = Arc::clone(self); + watch_stream( + move || { + let state = cache.read(); + let receiver = cache.events.subscribe(); + (receiver, state.snapshot()) + }, + cancel, + stop, + ) + } +} + +/// Single consumer of the slice watch, so the cache applies events in order. +async fn watch_slices( + slices: Api, + selector: String, + cache: Arc>, + stop: CancellationToken, +) { + // Lists come from the API server's watch cache (`resourceVersion=0`), as + // client-go's reflector does, rather than paged quorum reads from etcd: + // the watch catches up from there, and discovery needs no more. kube sends + // `resourceVersion=0` only without a page size, so paging is off. + let config = watcher::Config { + page_size: None, + ..watcher::Config::default() + .labels(&selector) + .any_semantic() + .timeout(WATCH_TIMEOUT_SECS) + }; + follow(retried(watcher(slices, config)), &selector, &cache, &stop).await; +} + +/// Applies `events` to `cache` until they end or `stop` is cancelled. +pub(super) async fn follow( + events: impl Stream>>, + selector: &str, + cache: &SliceCache, + stop: &CancellationToken, +) { + let mut events = std::pin::pin!(events); + // When a run of errors counts as over, unless another error comes first. + let mut quiet_at = None; + loop { + let event = tokio::select! { + biased; + () = stop.cancelled() => return, + () = wake_at(quiet_at) => { + quiet_at = None; + cache.quiet(); + continue; + } + event = events.next() => event, + }; + match event { + Some(Ok(event)) => { + quiet_at = None; + if cache.apply(event) { + tracing::warn!( + %selector, + "the EndpointSlices list nothing: check that MCS imports the expected Services into this cluster, under the expected names" + ); + } + } + // The cache keeps the last good view, and the watcher retries with + // backoff. The service serves that view and stays ready; after + // `STALE_AFTER`, the view is `Estimated`. + Some(Err(error)) => { + let failing_for = cache.failed(); + quiet_at = Some(Instant::now() + RECOVERED_AFTER); + if failing_for < STALE_AFTER { + tracing::warn!(%error, %selector, "EndpointSlice watch error"); + } else { + tracing::error!( + %error, + %selector, + ?failing_for, + "EndpointSlice watch still failing" + ); + } + } + None => return, + } + } +} + +/// The watch with kube's backoff. kube resets its backoff on every item, and +/// sends `Init` before each relist attempt, before any request, so a list that +/// keeps failing would be retried every 800 ms. `Init` is dropped: the cache +/// starts a relist at its first `InitApply` instead, which works because, with +/// paging off, a list's slices arrive together. +fn retried( + events: impl Stream>>, +) -> impl Stream>> { + events + .filter(|event| std::future::ready(!matches!(event, Ok(watcher::Event::Init)))) + .default_backoff() +} + +/// A [`Directory`] served from one EndpointSlice watch. The watch starts on +/// first use and stops when the directory is dropped. +pub(super) struct SliceDirectory { + slices: Api, + selector: String, + cache: Arc>, + watching: OnceLock<()>, + stop: CancellationToken, +} + +impl SliceDirectory { + pub(super) fn new(slices: Api, selector: String, rules: I::Rules) -> Self { + Self { + slices, + selector, + cache: Arc::new(SliceCache::new(rules)), + watching: OnceLock::new(), + stop: CancellationToken::new(), + } + } + + /// Every `Directory` method runs in a Tokio runtime, so the spawn is safe. + fn ensure_watching(&self) { + self.watching.get_or_init(|| { + tokio::spawn(watch_slices( + self.slices.clone(), + self.selector.clone(), + Arc::clone(&self.cache), + self.stop.clone(), + )); + }); + } +} + +impl fmt::Debug for SliceDirectory { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("SliceDirectory") + .field("selector", &self.selector) + .finish_non_exhaustive() + } +} + +impl Drop for SliceDirectory { + fn drop(&mut self) { + self.stop.cancel(); + } +} + +#[async_trait] +impl Directory for SliceDirectory { + /// From the cache once it has synced. Before that, one list from the API + /// server's cache, so a caller never waits for the watch to start. + async fn list(&self) -> Result, DiscoveryError> { + self.ensure_watching(); + if let Some(snapshot) = self.cache.snapshot() { + return Ok(snapshot); + } + let listed = self + .slices + .list(&ListParams::default().labels(&self.selector).match_any()) + .await + .map_err(unavailable)?; + let entries = listed + .items + .iter() + .filter_map(|slice| Some((slice_key(slice), I::entry(&self.cache.rules, slice)?))) + .collect(); + let mut index = I::default(); + index.replace(&self.cache.rules, entries); + Ok(index.snapshot()) + } + + /// The stream starts at the first complete list. Until then it waits while + /// the watch retries and logs its failures, as `dynamo_runtime`'s + /// Kubernetes discovery waits to be ready; `cancel_token` ends the wait. + async fn list_and_watch( + &self, + cancel_token: Option, + ) -> Result, DiscoveryError> { + self.ensure_watching(); + Ok(self + .cache + .subscribe(cancel_token.unwrap_or_default(), self.stop.clone())) + } +} + +fn unavailable(error: impl Into) -> DiscoveryError { + DiscoveryError::Unavailable(error.into()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test(start_paused = true)] + async fn a_relist_that_keeps_failing_backs_off_without_its_init() { + let attempts = (0..4).flat_map(|_| { + [ + Ok(watcher::Event::Init), + Err(watcher::Error::NoResourceVersion), + ] + }); + let mut events = std::pin::pin!(retried(futures::stream::iter(attempts))); + let mut last = Instant::now(); + let mut gaps = Vec::new(); + while let Some(event) = events.next().await { + assert!(event.is_err(), "`Init` is dropped"); + gaps.push(last.elapsed()); + last = Instant::now(); + } + // Each retry waits longer, as kube's backoff grows from 800 ms. + assert_eq!(gaps.len(), 4); + assert_eq!(gaps[0], Duration::ZERO); + assert!( + gaps.windows(2).skip(1).all(|pair| pair[1] > pair[0]), + "{gaps:?}" + ); + } +} diff --git a/lib/global-discovery-service/src/backends/kubernetes/endpoints.rs b/lib/global-discovery-service/src/backends/kubernetes/endpoints.rs new file mode 100644 index 000000000000..54a14ffe75cb --- /dev/null +++ b/lib/global-discovery-service/src/backends/kubernetes/endpoints.rs @@ -0,0 +1,558 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Router replicas from the EndpointSlices that MCS imports. The Global +//! Discovery Service serves them to PoolRelays as `GET /routers`. Router +//! replicas find each other through the Dynamo runtime, not through this. + +use std::collections::{BTreeMap, HashMap, HashSet}; +use std::net::{IpAddr, SocketAddr}; + +use crate::pool_discovery::{ + Completeness, Directory, DiscoveryError, RouterReplica, RouterReplicaEvent, RouterReplicaId, + Snapshot, WatchEvent, WatchStream, +}; +use async_trait::async_trait; +use k8s_openapi::api::discovery::v1::EndpointSlice; +use kube::{Api, Client}; +use tokio_util::sync::CancellationToken; +use url::Url; + +use super::cache::{SliceDirectory, SliceIndex, SliceKey, ready_endpoints, tcp_port}; +use super::mcs::{self, McsImplementation}; + +/// A Service exported through MCS, such as the Global Router's. It need not +/// be headless: imported slices list every pod, and readers use slices, not +/// DNS. +#[derive(Clone, Debug)] +pub struct ExportedService { + pub namespace: String, + pub name: String, + /// The named port its replicas are reached on, for example the router's + /// state stream. + pub port_name: String, +} + +/// Ready replicas from a Service's imported slices, named +/// `/`, or `/
` when a slice +/// does not name the pod. Lists are `Exact` unless the watch keeps failing. +#[derive(Debug)] +pub struct McsReplicaDirectory(SliceDirectory); + +impl McsReplicaDirectory { + /// The Service's slices as MCS imports them into the cluster `client` + /// reaches. The client's connect and read timeouts should add up to under + /// 90 s, so a server that stops answering shows up as failures. Logs an + /// error for a Service that `mcs` cannot import, since none of its + /// replicas would be listed. + pub fn new(client: Client, service: &ExportedService, mcs: McsImplementation) -> Self { + if !mcs::can_import(mcs, &service.name) { + tracing::error!( + service = %service.name, + karmada_max_service_name = mcs::KARMADA_MAX_SERVICE_NAME, + "the Service's name is too long for Karmada to import it, so none of its replicas will be listed" + ); + } + Self(SliceDirectory::new( + Api::namespaced(client, &service.namespace), + mcs::service_selector(mcs, &service.name), + ReplicaNaming { + mcs, + service: service.name.clone(), + port_name: service.port_name.clone(), + }, + )) + } +} + +#[async_trait] +impl Directory for McsReplicaDirectory { + async fn list(&self) -> Result, DiscoveryError> { + self.0.list().await + } + + async fn list_and_watch( + &self, + cancel_token: Option, + ) -> Result, DiscoveryError> { + self.0.list_and_watch(cancel_token).await + } +} + +/// Turns one imported slice into replicas. +struct ReplicaNaming { + mcs: McsImplementation, + service: String, + port_name: String, +} + +impl ReplicaNaming { + /// Ready endpoints only. The URL comes from the endpoint's address and + /// the named port, formatted through `SocketAddr` so IPv6 gets brackets. + /// Per-pod DNS names are not used: they exist only for pods with + /// hostnames. `None` for a slice of another Service, a slice whose source + /// cluster is unknown, or one without the named port. + fn replicas(&self, slice: &EndpointSlice) -> Option> { + if mcs::service_name(self.mcs, slice) != Some(self.service.as_str()) { + return None; + } + let Some(cluster) = mcs::source_cluster(self.mcs, slice) else { + mcs::warn_unknown_cluster(slice, &self.service); + return None; + }; + let Some(port) = tcp_port(slice, &self.port_name) else { + // The slice is the router Service's, so the name is misconfigured. + tracing::warn!( + slice = slice.metadata.name.as_deref().unwrap_or_default(), + port_name = %self.port_name, + "router EndpointSlice has no TCP port of that name; its replicas are not listed" + ); + return None; + }; + let replicas = ready_endpoints(slice) + .filter_map(|endpoint| { + // An endpoint's addresses are interchangeable, so the first is + // enough. + let address: IpAddr = endpoint.addresses.first()?.parse().ok()?; + let url = Url::parse(&format!("http://{}", SocketAddr::new(address, port))).ok()?; + let name = endpoint + .target_ref + .as_ref() + .and_then(|target| target.name.clone()) + .unwrap_or_else(|| address.to_string()); + Some(RouterReplica { + replica: RouterReplicaId(format!("{cluster}/{name}")), + url, + }) + }) + .collect(); + Some(replicas) + } +} + +#[derive(Default)] +struct ReplicaIndex { + /// The replicas each slice lists. + slices: HashMap>, + /// Every listed replica, with the slices that list it. + replicas: HashMap>, +} + +impl ReplicaIndex { + /// A replica in several slices is served at its IPv4 address if it has + /// one, as on a dual-stack cluster, so its URL does not depend on how the + /// slices are named. Otherwise from the first slice in key order. + fn served(&self, id: &RouterReplicaId) -> Option<&RouterReplica> { + let listed = self.replicas.get(id)?; + listed + .values() + .find(|endpoint| matches!(endpoint.url.host(), Some(url::Host::Ipv4(_)))) + .or_else(|| listed.values().next()) + } +} + +impl SliceIndex for ReplicaIndex { + type Item = RouterReplica; + type Key = RouterReplicaId; + type Rules = ReplicaNaming; + type Entry = Vec; + + fn entry(naming: &ReplicaNaming, slice: &EndpointSlice) -> Option> { + naming.replicas(slice) + } + + /// Touches only the replicas this slice lists now or listed before. + fn upsert( + &mut self, + _naming: &ReplicaNaming, + key: SliceKey, + entry: Option>, + ) -> Vec { + let old = self.slices.remove(&key).unwrap_or_default(); + let new = entry.unwrap_or_default(); + let touched: HashSet = old + .iter() + .chain(&new) + .map(|endpoint| endpoint.replica.clone()) + .collect(); + let before: Vec<(RouterReplicaId, Option)> = touched + .into_iter() + .map(|id| { + let served = self.served(&id).cloned(); + (id, served) + }) + .collect(); + for endpoint in &old { + if let Some(listed) = self.replicas.get_mut(&endpoint.replica) { + listed.remove(&key); + if listed.is_empty() { + self.replicas.remove(&endpoint.replica); + } + } + } + for endpoint in &new { + self.replicas + .entry(endpoint.replica.clone()) + .or_default() + .insert(key.clone(), endpoint.clone()); + } + if !new.is_empty() { + self.slices.insert(key, new); + } + before + .into_iter() + .filter_map(|(id, before)| match (before, self.served(&id)) { + (Some(_), None) => Some(WatchEvent::Removed(id)), + (before, Some(after)) if before.as_ref() != Some(after) => { + Some(WatchEvent::Upserted(after.clone())) + } + _ => None, + }) + .collect() + } + + /// A relist: every slice at once. Rebuilds every replica; the cache + /// publishes one `Resync` instead of per-replica events. + fn replace(&mut self, _naming: &ReplicaNaming, entries: HashMap>) { + self.replicas.clear(); + for (key, endpoints) in &entries { + for endpoint in endpoints { + self.replicas + .entry(endpoint.replica.clone()) + .or_default() + .insert(key.clone(), endpoint.clone()); + } + } + self.slices = entries + .into_iter() + .filter(|(_, endpoints)| !endpoints.is_empty()) + .collect(); + } + + fn snapshot(&self) -> Snapshot { + Snapshot { + items: self + .replicas + .keys() + .filter_map(|id| self.served(id).cloned()) + .collect(), + completeness: Completeness::Exact, + } + } +} + +#[cfg(test)] +mod tests { + use std::collections::BTreeMap; + use std::sync::Arc; + + use k8s_openapi::api::core::v1::ObjectReference; + use k8s_openapi::api::discovery::v1::{Endpoint, EndpointConditions, EndpointPort}; + use k8s_openapi::apimachinery::pkg::apis::meta::v1::ObjectMeta; + use kube::runtime::watcher; + + use super::super::cache::SliceCache; + use super::super::mcs::{ + KARMADA_SLICE_MANAGER, KUBERNETES_SERVICE_NAME_LABEL, MANAGED_BY_LABEL, + }; + use super::*; + + const ROUTER: &str = "global-router"; + + fn naming() -> ReplicaNaming { + ReplicaNaming { + mcs: McsImplementation::Karmada, + service: ROUTER.to_owned(), + port_name: "state".to_owned(), + } + } + + /// A Karmada-imported slice of the router Service from `cluster`, with + /// endpoints as (address, Ready, pod). + fn slice( + cluster: &str, + suffix: &str, + endpoints: &[(&str, Option, Option<&str>)], + ) -> EndpointSlice { + EndpointSlice { + address_type: "IPv4".to_owned(), + metadata: ObjectMeta { + name: Some(format!("imported-{cluster}-{ROUTER}-{suffix}")), + namespace: Some("dynamo".to_owned()), + labels: Some(BTreeMap::from([ + ( + MANAGED_BY_LABEL.to_owned(), + KARMADA_SLICE_MANAGER.to_owned(), + ), + ( + KUBERNETES_SERVICE_NAME_LABEL.to_owned(), + format!("derived-{ROUTER}"), + ), + ])), + ..Default::default() + }, + ports: Some(vec![EndpointPort { + name: Some("state".to_owned()), + port: Some(50051), + ..Default::default() + }]), + endpoints: endpoints + .iter() + .map(|(address, ready, pod)| Endpoint { + addresses: vec![(*address).to_owned()], + conditions: Some(EndpointConditions { + ready: *ready, + ..Default::default() + }), + target_ref: pod.map(|pod| ObjectReference { + kind: Some("Pod".to_owned()), + name: Some(pod.to_owned()), + ..Default::default() + }), + ..Default::default() + }) + .collect(), + } + } + + fn cache() -> Arc> { + Arc::new(SliceCache::new(naming())) + } + + fn sync(cache: &SliceCache, slices: Vec) { + cache.apply(watcher::Event::Init); + for slice in slices { + cache.apply(watcher::Event::InitApply(slice)); + } + cache.apply(watcher::Event::InitDone); + } + + fn id(name: &str) -> RouterReplicaId { + RouterReplicaId(name.to_owned()) + } + + /// Served replicas as (id, URL), sorted. + fn served(cache: &SliceCache) -> Vec<(String, String)> { + let mut replicas: Vec<(String, String)> = cache + .snapshot() + .expect("synced") + .items + .into_iter() + .map(|endpoint| (endpoint.replica.0, endpoint.url.to_string())) + .collect(); + replicas.sort(); + replicas + } + + #[test] + fn ready_endpoints_become_replicas() { + let cache = cache(); + sync( + &cache, + vec![slice( + "ohio", + "a1b2c", + &[ + ("10.42.0.5", Some(true), Some("global-router-0")), + // The EndpointSlice API reads a missing `ready` as ready. + ("10.42.0.6", None, None), + ("10.42.0.7", Some(false), Some("global-router-2")), + ], + )], + ); + assert_eq!( + served(&cache), + [ + ( + "ohio/10.42.0.6".to_owned(), + "http://10.42.0.6:50051/".to_owned() + ), + ( + "ohio/global-router-0".to_owned(), + "http://10.42.0.5:50051/".to_owned() + ), + ] + ); + } + + #[test] + fn brackets_ipv6_addresses() { + let cache = cache(); + let mut ipv6 = slice( + "ohio", + "a1b2c", + &[("fd00::5", Some(true), Some("global-router-0"))], + ); + ipv6.address_type = "IPv6".to_owned(); + sync(&cache, vec![ipv6]); + assert_eq!( + served(&cache), + [( + "ohio/global-router-0".to_owned(), + "http://[fd00::5]:50051/".to_owned() + )] + ); + } + + #[test] + fn skips_slices_without_the_port_or_of_another_service() { + let mut no_port = slice("ohio", "a1b2c", &[("10.42.0.5", Some(true), None)]); + no_port.ports = Some(Vec::new()); + let mut other = slice("ohio", "d3e4f", &[("10.42.0.6", Some(true), None)]); + other.metadata.labels.as_mut().unwrap().insert( + KUBERNETES_SERVICE_NAME_LABEL.to_owned(), + "derived-echo".to_owned(), + ); + // Only the named TCP port counts, and not port 0. + let mut udp = slice("ohio", "g5h6i", &[("10.42.0.7", Some(true), None)]); + udp.ports.as_mut().unwrap()[0].protocol = Some("UDP".to_owned()); + let mut zero = slice("ohio", "j7k8l", &[("10.42.0.8", Some(true), None)]); + zero.ports.as_mut().unwrap()[0].port = Some(0); + let cache = cache(); + sync(&cache, vec![no_port, other, udp, zero]); + assert!(served(&cache).is_empty()); + } + + #[test] + fn changes_publish_added_and_removed() { + let cache = cache(); + sync(&cache, Vec::new()); + let mut events = cache.events.subscribe(); + + let one = slice( + "ohio", + "a1b2c", + &[("10.42.0.5", Some(true), Some("global-router-0"))], + ); + cache.apply(watcher::Event::Apply(one.clone())); + assert!( + matches!(events.try_recv().unwrap(), WatchEvent::Upserted(e) if e.replica == id("ohio/global-router-0")) + ); + cache.apply(watcher::Event::Apply(one.clone())); + assert!(events.try_recv().is_err()); + // The pod moves to a new address: the replica changes. + cache.apply(watcher::Event::Apply(slice( + "ohio", + "a1b2c", + &[("10.42.0.9", Some(true), Some("global-router-0"))], + ))); + assert!( + matches!(events.try_recv().unwrap(), WatchEvent::Upserted(e) if e.url.as_str() == "http://10.42.0.9:50051/") + ); + // It stops being Ready. + cache.apply(watcher::Event::Apply(slice( + "ohio", + "a1b2c", + &[("10.42.0.9", Some(false), Some("global-router-0"))], + ))); + assert!( + matches!(events.try_recv().unwrap(), WatchEvent::Removed(r) if r == id("ohio/global-router-0")) + ); + assert!(cache.read().index.replicas.is_empty()); + } + + #[test] + fn replica_in_two_slices_stays_until_both_drop_it() { + let cache = cache(); + let first = slice( + "ohio", + "a1b2c", + &[("10.42.0.5", Some(true), Some("global-router-0"))], + ); + let second = slice( + "ohio", + "d3e4f", + &[("10.42.0.5", Some(true), Some("global-router-0"))], + ); + sync(&cache, vec![first.clone(), second]); + assert_eq!(served(&cache).len(), 1); + let mut events = cache.events.subscribe(); + cache.apply(watcher::Event::Delete(first)); + assert!(events.try_recv().is_err()); + assert_eq!(served(&cache).len(), 1); + cache.apply(watcher::Event::Delete(slice("ohio", "d3e4f", &[]))); + assert!(matches!(events.try_recv().unwrap(), WatchEvent::Removed(_))); + assert!(served(&cache).is_empty()); + } + + #[test] + fn a_dual_stack_replica_is_served_at_its_ipv4_address() { + let cache = cache(); + // The IPv6 slice's name sorts first. + let mut ipv6 = slice( + "ohio", + "a1b2c", + &[("fd00::5", Some(true), Some("global-router-0"))], + ); + ipv6.address_type = "IPv6".to_owned(); + let ipv4 = slice( + "ohio", + "d3e4f", + &[("10.42.0.5", Some(true), Some("global-router-0"))], + ); + sync(&cache, vec![ipv6, ipv4]); + assert_eq!( + served(&cache), + [( + "ohio/global-router-0".to_owned(), + "http://10.42.0.5:50051/".to_owned() + )] + ); + } + + #[test] + fn replicas_from_two_clusters_keep_their_names_apart() { + let cache = cache(); + sync( + &cache, + vec![ + slice( + "ohio", + "a1b2c", + &[("10.42.0.5", Some(true), Some("global-router-0"))], + ), + slice( + "west", + "d3e4f", + &[("10.43.0.5", Some(true), Some("global-router-0"))], + ), + ], + ); + let ids: Vec = served(&cache).into_iter().map(|(id, _)| id).collect(); + assert_eq!(ids, ["ohio/global-router-0", "west/global-router-0"]); + } + + #[test] + fn relist_replaces_everything() { + let cache = cache(); + sync( + &cache, + vec![slice( + "ohio", + "a1b2c", + &[("10.42.0.5", Some(true), Some("global-router-0"))], + )], + ); + let mut events = cache.events.subscribe(); + sync( + &cache, + vec![slice( + "west", + "d3e4f", + &[("10.43.0.5", Some(true), Some("global-router-1"))], + )], + ); + let WatchEvent::Resync(snapshot) = events.try_recv().unwrap() else { + panic!("expected Resync"); + }; + let ids: Vec = snapshot.items.into_iter().map(|e| e.replica).collect(); + assert_eq!(ids, [id("west/global-router-1")]); + assert_eq!( + served(&cache), + [( + "west/global-router-1".to_owned(), + "http://10.43.0.5:50051/".to_owned() + )] + ); + } +} diff --git a/lib/global-discovery-service/src/backends/kubernetes/hub.rs b/lib/global-discovery-service/src/backends/kubernetes/hub.rs new file mode 100644 index 000000000000..eb07507a84fe --- /dev/null +++ b/lib/global-discovery-service/src/backends/kubernetes/hub.rs @@ -0,0 +1,805 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Pools from imported EndpointSlices: the Global Discovery Service's MCS +//! backend. + +use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet}; +use std::net::IpAddr; + +use crate::pool_discovery::{ + Completeness, Directory, DiscoveryError, PoolEvent, PoolListing, Snapshot, WatchEvent, + WatchStream, +}; +use async_trait::async_trait; +use dynamo_global_common::{PoolKey, PoolLocation, address_identity}; +use k8s_openapi::api::discovery::v1::EndpointSlice; +use kube::{Api, Client}; +use tokio_util::sync::CancellationToken; + +use super::cache::{SliceDirectory, SliceIndex, SliceKey, ready_endpoints}; +use super::mcs::{self, McsImplementation}; + +/// One pool per (source cluster, namespace, DGD name), from imported +/// EndpointSlices. MCS merges same-named exports from several clusters into +/// one import, so pools are split by source cluster. Lists are `Exact` unless +/// the watch keeps failing. Each slice event updates one pool. +/// +/// One watch per directory feeds a cache that serves `list` and every +/// `list_and_watch` stream. The watch starts on first use and stops when the +/// directory is dropped. +#[derive(Debug)] +pub struct McsPoolDirectory(SliceDirectory); + +impl McsPoolDirectory { + /// `imports` reaches the cluster where the imported pool slices are: the + /// hub for standard MCS. For Karmada, the Karmada control plane, which + /// collects every exported Service's slices, so no per-pool + /// `ServiceImport` is needed. A member cluster with a `ServiceImport` per + /// pool also works. Each PoolRelay Service is named + /// ``. The client's connect and read + /// timeouts should add up to under 90 s, so a server that stops answering + /// shows up as failures. + pub fn new(imports: Client, mcs: McsImplementation, pool_service_suffix: String) -> Self { + Self(SliceDirectory::new( + Api::all(imports), + mcs::imported_selector(mcs), + PoolNaming { + mcs, + pool_service_suffix, + }, + )) + } +} + +#[async_trait] +impl Directory for McsPoolDirectory { + async fn list(&self) -> Result, DiscoveryError> { + self.0.list().await + } + + async fn list_and_watch( + &self, + cancel_token: Option, + ) -> Result, DiscoveryError> { + self.0.list_and_watch(cancel_token).await + } +} + +/// Turns imported slices into pool listings. +struct PoolNaming { + mcs: McsImplementation, + pool_service_suffix: String, +} + +impl PoolNaming { + /// Source cluster from [`mcs::source_cluster`]; DGD name from the + /// Service name without `pool_service_suffix`. `None` for imports that + /// are not pools, and for a pool slice from an unknown cluster, which is + /// logged. + fn pool_key(&self, slice: &EndpointSlice) -> Option { + let service = mcs::service_name(self.mcs, slice)?; + let dgd_name = service + .strip_suffix(self.pool_service_suffix.as_str()) + .filter(|name| !name.is_empty())?; + let Some(cluster) = mcs::source_cluster(self.mcs, slice) else { + mcs::warn_unknown_cluster(slice, service); + return None; + }; + let namespace = slice.metadata.namespace.as_deref()?; + PoolKey::new(cluster, namespace, dgd_name).ok() + } + + fn slice(&self, slice: &EndpointSlice) -> Option { + Some(SliceEntry { + pool: self.pool_key(slice)?, + relays: ready_relays(slice), + }) + } +} + +/// MCS knows only the pool's cluster, so its location is named after it. +fn location(cluster: &str) -> PoolLocation { + PoolLocation { + region: cluster.to_owned(), + availability_zone: None, + datacenter: None, + } +} + +/// A pool is live while a PoolRelay endpoint is Ready. Ready can be stale +/// when a cluster loses contact with the MCS implementation, so the runtime +/// also drops a pool whose Relay connection is gone. +/// +/// Each Ready address comes with its pod: the endpoint's `targetRef`, or the +/// address itself when the endpoint has none. A dual-stack pod is in two +/// slices, one per address family. +fn ready_relays(slice: &EndpointSlice) -> Vec<(String, IpAddr)> { + ready_endpoints(slice) + .flat_map(|endpoint| { + let pod = endpoint + .target_ref + .as_ref() + .and_then(|target| target.name.as_deref()); + endpoint.addresses.iter().filter_map(move |address| { + let address: IpAddr = address.parse().ok()?; + let pod = pod.map_or_else(|| address.to_string(), str::to_owned); + Some((pod, address)) + }) + }) + .collect() +} + +struct SliceEntry { + pool: PoolKey, + relays: Vec<(String, IpAddr)>, +} + +#[derive(Default)] +struct PoolState { + slices: HashSet, + /// `Some` while the pool is live. + listing: Option, +} + +#[derive(Default)] +struct PoolIndex { + slices: HashMap, + pools: HashMap, +} + +impl SliceIndex for PoolIndex { + type Item = PoolListing; + type Key = PoolKey; + type Rules = PoolNaming; + type Entry = SliceEntry; + + fn entry(naming: &PoolNaming, slice: &EndpointSlice) -> Option { + naming.slice(slice) + } + + /// One slice changed or went away. Touches only that slice's pool. + fn upsert( + &mut self, + _naming: &PoolNaming, + key: SliceKey, + entry: Option, + ) -> Vec { + let old_pool = self.slices.remove(&key).map(|old| old.pool); + let new_pool = entry.as_ref().map(|entry| entry.pool.clone()); + if let Some(entry) = entry { + self.slices.insert(key.clone(), entry); + } + let mut events = Vec::new(); + if let Some(pool) = old_pool.filter(|pool| new_pool.as_ref() != Some(pool)) { + if let Some(state) = self.pools.get_mut(&pool) { + state.slices.remove(&key); + } + events.extend(self.rebuild(&pool)); + } + if let Some(pool) = new_pool { + self.pools + .entry(pool.clone()) + .or_default() + .slices + .insert(key); + events.extend(self.rebuild(&pool)); + } + events + } + + /// A relist: every slice at once. Rebuilds every pool; the cache + /// publishes one `Resync` instead of per-pool events. + fn replace(&mut self, _naming: &PoolNaming, slices: HashMap) { + self.slices = slices; + for pool in self.pools.values_mut() { + pool.slices.clear(); + } + for (key, entry) in &self.slices { + self.pools + .entry(entry.pool.clone()) + .or_default() + .slices + .insert(key.clone()); + } + let pools: Vec = self.pools.keys().cloned().collect(); + for pool in &pools { + self.rebuild(pool); + } + } + + fn snapshot(&self) -> Snapshot { + Snapshot { + items: self + .pools + .values() + .filter_map(|pool| pool.listing.clone()) + .collect(), + completeness: Completeness::Exact, + } + } +} + +impl PoolIndex { + /// Recomputes one pool from its slices. Returns the event to publish if + /// the listing changed. + fn rebuild(&mut self, key: &PoolKey) -> Option { + let state = self.pools.get(key)?; + let mut relays: BTreeMap<&str, BTreeSet> = BTreeMap::new(); + for entry in state + .slices + .iter() + .filter_map(|slice| self.slices.get(slice)) + { + for (pod, address) in &entry.relays { + relays.entry(pod).or_default().insert(*address); + } + } + // Bind the Ready PoolRelay pod with the lowest address. The choice + // depends only on the current slices, so every service instance binds + // the same pod, and routers that ask different instances admit the + // same PoolRelay. A rolling update can move the binding to the new + // pod while the old one is still Ready. + let relay_identities: Option> = relays + .values() + .min_by_key(|addresses| addresses.first()) + .map(|addresses| addresses.iter().copied().map(address_identity).collect()); + + let state = self.pools.get_mut(key)?; + let Some(relay_identities) = relay_identities else { + let was_live = state.listing.take().is_some(); + if state.slices.is_empty() { + self.pools.remove(key); + } + return was_live.then(|| WatchEvent::Removed(key.clone())); + }; + let listing = PoolListing { + key: key.clone(), + location: location(key.cluster_id()), + model: None, + capabilities: Vec::new(), + frontends: Vec::new(), + capacity: Vec::new(), + relay_identities, + }; + if state.listing.as_ref() == Some(&listing) { + return None; + } + state.listing = Some(listing.clone()); + Some(WatchEvent::Upserted(listing)) + } +} + +#[cfg(test)] +mod tests { + use std::collections::BTreeMap; + use std::sync::Arc; + + use futures::{FutureExt, StreamExt}; + use std::time::Duration; + + use k8s_openapi::api::core::v1::ObjectReference; + use k8s_openapi::api::discovery::v1::{Endpoint, EndpointConditions}; + use k8s_openapi::apimachinery::pkg::apis::meta::v1::ObjectMeta; + use kube::runtime::watcher; + + use super::super::cache::{RECOVERED_AFTER, STALE_AFTER, SliceCache, follow}; + use super::mcs::{KARMADA_SLICE_MANAGER, KUBERNETES_SERVICE_NAME_LABEL, MANAGED_BY_LABEL}; + use super::*; + + const SUFFIX: &str = "-pool-relay"; + + fn naming() -> PoolNaming { + PoolNaming { + mcs: McsImplementation::Karmada, + pool_service_suffix: SUFFIX.to_owned(), + } + } + + /// A Karmada-imported slice of `-pool-relay` from `cluster`. + fn slice(cluster: &str, dgd: &str, suffix: &str, endpoints: &[(&str, bool)]) -> EndpointSlice { + let service = format!("{dgd}{SUFFIX}"); + EndpointSlice { + address_type: "IPv4".to_owned(), + metadata: ObjectMeta { + name: Some(format!("imported-{cluster}-{service}-{suffix}")), + namespace: Some("dynamo".to_owned()), + labels: Some(BTreeMap::from([ + ( + MANAGED_BY_LABEL.to_owned(), + KARMADA_SLICE_MANAGER.to_owned(), + ), + ( + KUBERNETES_SERVICE_NAME_LABEL.to_owned(), + format!("derived-{service}"), + ), + ])), + ..Default::default() + }, + endpoints: endpoints + .iter() + .map(|(address, ready)| Endpoint { + addresses: vec![(*address).to_owned()], + conditions: Some(EndpointConditions { + ready: Some(*ready), + ..Default::default() + }), + ..Default::default() + }) + .collect(), + ports: None, + } + } + + /// Like [`slice`], with each endpoint's pod named. + fn pod_slice( + cluster: &str, + suffix: &str, + address_type: &str, + endpoints: &[(&str, &str)], + ) -> EndpointSlice { + let mut slice = slice(cluster, "mocker", suffix, &[]); + slice.address_type = address_type.to_owned(); + slice.endpoints = endpoints + .iter() + .map(|(pod, address)| Endpoint { + addresses: vec![(*address).to_owned()], + target_ref: Some(ObjectReference { + kind: Some("Pod".to_owned()), + name: Some((*pod).to_owned()), + ..Default::default() + }), + ..Default::default() + }) + .collect(); + slice + } + + fn key(cluster: &str, dgd: &str) -> PoolKey { + PoolKey::new(cluster, "dynamo", dgd).unwrap() + } + + fn cache() -> Arc> { + Arc::new(SliceCache::new(naming())) + } + + fn sync(cache: &SliceCache, slices: Vec) { + cache.apply(watcher::Event::Init); + for slice in slices { + cache.apply(watcher::Event::InitApply(slice)); + } + cache.apply(watcher::Event::InitDone); + } + + fn completeness(cache: &SliceCache) -> Completeness { + cache.snapshot().unwrap().completeness + } + + #[tokio::test(start_paused = true)] + async fn a_watch_that_keeps_failing_marks_its_last_list_incomplete() { + let cache = cache(); + let ready = || slice("west", "mocker", "a1b2c", &[("10.43.0.1", true)]); + sync(&cache, vec![ready()]); + // A relist that keeps failing. kube sends `Init` before each attempt; + // the watch drops it (see `retried`), and the cache does not count it + // as a working watch either. + for _ in 0..3 { + cache.apply(watcher::Event::Init); + cache.failed(); + tokio::time::advance(STALE_AFTER / 3).await; + } + assert_eq!(completeness(&cache), Completeness::Exact); + cache.apply(watcher::Event::Init); + assert_eq!(cache.failed(), STALE_AFTER); + let stale = cache.snapshot().unwrap(); + assert_eq!(stale.items.len(), 1); + assert_eq!(stale.completeness, Completeness::Estimated); + + // A relist that completes means the watch works again. Until then the + // old view stays `Estimated`. + cache.apply(watcher::Event::InitApply(ready())); + assert_eq!(completeness(&cache), Completeness::Estimated); + cache.apply(watcher::Event::InitDone); + assert_eq!(completeness(&cache), Completeness::Exact); + } + + #[tokio::test(start_paused = true)] + async fn one_error_then_silence_is_a_resumed_watch() { + let cache = cache(); + sync( + &cache, + vec![slice("west", "mocker", "a1b2c", &[("10.43.0.1", true)])], + ); + // kube resumes the watch without an event. + cache.failed(); + tokio::time::advance(STALE_AFTER + Duration::from_secs(1)).await; + assert_eq!(completeness(&cache), Completeness::Exact); + } + + #[tokio::test(start_paused = true)] + async fn a_failing_watch_that_falls_silent_counts_as_recovered() { + let cache = cache(); + sync( + &cache, + vec![slice("west", "mocker", "a1b2c", &[("10.43.0.1", true)])], + ); + for _ in 0..3 { + cache.failed(); + tokio::time::advance(STALE_AFTER / 2).await; + } + assert_eq!(completeness(&cache), Completeness::Estimated); + tokio::time::advance(RECOVERED_AFTER).await; + assert_eq!(completeness(&cache), Completeness::Exact); + // A later error starts a new run. + assert_eq!(cache.failed(), Duration::ZERO); + assert_eq!(completeness(&cache), Completeness::Exact); + } + + fn listing(cache: &SliceCache, key: &PoolKey) -> Option { + cache + .read() + .index + .pools + .get(key) + .and_then(|pool| pool.listing.clone()) + } + + #[test] + fn pool_listing_from_imported_slice() { + let cache = cache(); + sync( + &cache, + vec![slice("west", "mocker", "a1b2c", &[("10.43.66.175", true)])], + ); + let listing = listing(&cache, &key("west", "mocker")).expect("live pool"); + assert_eq!(listing.location.region, "west"); + assert_eq!(listing.relay_identities, ["ip:10.43.66.175"]); + assert_eq!(listing.model, None); + } + + #[test] + fn a_dual_stack_relay_is_one_pod_with_both_addresses() { + let cache = cache(); + sync( + &cache, + vec![ + pod_slice( + "west", + "a1b2c", + "IPv4", + &[("relay-a", "10.43.0.7"), ("relay-b", "10.43.0.9")], + ), + pod_slice( + "west", + "d3e4f", + "IPv6", + &[("relay-a", "fd00::7"), ("relay-b", "fd00::1")], + ), + ], + ); + let listing = listing(&cache, &key("west", "mocker")).expect("live pool"); + assert_eq!(listing.relay_identities, ["ip:10.43.0.7", "ip:fd00::7"]); + } + + #[test] + fn skips_imports_that_are_not_pools() { + let mut other = slice("west", "mocker", "a1b2c", &[("10.43.0.9", true)]); + other.metadata.labels.as_mut().unwrap().insert( + KUBERNETES_SERVICE_NAME_LABEL.to_owned(), + "derived-echo".to_owned(), + ); + let cache = cache(); + sync(&cache, vec![other]); + assert!(cache.snapshot().unwrap().items.is_empty()); + } + + #[test] + fn same_export_from_two_clusters_is_two_pools() { + let cache = cache(); + sync( + &cache, + vec![ + slice("west", "mocker", "a1b2c", &[("10.43.0.1", true)]), + slice("ohio", "mocker", "d3e4f", &[("10.42.0.1", true)]), + ], + ); + assert!(listing(&cache, &key("west", "mocker")).is_some()); + assert!(listing(&cache, &key("ohio", "mocker")).is_some()); + } + + #[test] + fn pool_without_ready_relay_is_not_live() { + let cache = cache(); + sync( + &cache, + vec![slice("west", "mocker", "a1b2c", &[("10.43.0.1", false)])], + ); + assert!(cache.snapshot().unwrap().items.is_empty()); + } + + #[test] + fn changes_publish_added_and_removed() { + let cache = cache(); + sync(&cache, Vec::new()); + let mut events = cache.events.subscribe(); + + let live = slice("west", "mocker", "a1b2c", &[("10.43.0.1", true)]); + cache.apply(watcher::Event::Apply(live.clone())); + assert!(matches!( + events.try_recv().unwrap(), + WatchEvent::Upserted(_) + )); + cache.apply(watcher::Event::Apply(live.clone())); + assert!(events.try_recv().is_err()); + // The Relay stops being Ready. + cache.apply(watcher::Event::Apply(slice( + "west", + "mocker", + "a1b2c", + &[("10.43.0.1", false)], + ))); + assert!( + matches!(events.try_recv().unwrap(), WatchEvent::Removed(k) if k == key("west", "mocker")) + ); + // Ready again. + cache.apply(watcher::Event::Apply(live.clone())); + assert!(matches!( + events.try_recv().unwrap(), + WatchEvent::Upserted(_) + )); + cache.apply(watcher::Event::Delete(live)); + assert!(matches!(events.try_recv().unwrap(), WatchEvent::Removed(_))); + assert!(cache.read().index.pools.is_empty()); + } + + #[test] + fn pool_spans_several_slices() { + let cache = cache(); + sync( + &cache, + vec![ + slice("west", "mocker", "a1b2c", &[("10.43.0.1", false)]), + slice("west", "mocker", "d3e4f", &[("10.43.0.2", true)]), + ], + ); + assert!(listing(&cache, &key("west", "mocker")).is_some()); + // Losing the slice with the Ready Relay ends the pool. + cache.apply(watcher::Event::Delete(slice( + "west", + "mocker", + "d3e4f", + &[], + ))); + assert!(listing(&cache, &key("west", "mocker")).is_none()); + } + + #[test] + fn every_instance_binds_the_same_relay() { + let pool = key("west", "mocker"); + let both = || { + slice( + "west", + "mocker", + "a1b2c", + &[("10.43.0.9", true), ("10.43.0.2", true)], + ) + }; + // One instance saw the old PoolRelay alone first. + let first = cache(); + sync( + &first, + vec![slice("west", "mocker", "a1b2c", &[("10.43.0.9", true)])], + ); + first.apply(watcher::Event::Apply(both())); + // Another started later. + let second = cache(); + sync(&second, vec![both()]); + let bound = |cache: &SliceCache| listing(cache, &pool).unwrap().relay_identities; + assert_eq!(bound(&first), bound(&second)); + assert_eq!( + bound(&first), + [address_identity("10.43.0.2".parse().unwrap())] + ); + } + + #[tokio::test(start_paused = true)] + async fn watchers_learn_when_the_view_turns_stale_and_back() { + let cache = cache(); + let ready = || slice("west", "mocker", "a1b2c", &[("10.43.0.1", true)]); + sync(&cache, vec![ready()]); + let mut events = cache.events.subscribe(); + let mut next_resync = || match events.try_recv() { + Ok(WatchEvent::Resync(snapshot)) => Some(snapshot.completeness), + Ok(_) => panic!("expected a Resync"), + Err(_) => None, + }; + for _ in 0..3 { + cache.failed(); + tokio::time::advance(STALE_AFTER / 2).await; + } + assert_eq!(next_resync(), Some(Completeness::Estimated)); + // Further errors change nothing. + cache.failed(); + assert_eq!(next_resync(), None); + + // Recovered by silence. + tokio::time::advance(RECOVERED_AFTER).await; + cache.quiet(); + assert_eq!(next_resync(), Some(Completeness::Exact)); + assert_eq!(next_resync(), None); + + // Recovered by an event from the API server. + for _ in 0..3 { + cache.failed(); + tokio::time::advance(STALE_AFTER / 2).await; + } + assert_eq!(next_resync(), Some(Completeness::Estimated)); + cache.apply(watcher::Event::Apply(ready())); + assert_eq!(next_resync(), Some(Completeness::Exact)); + cache.quiet(); + assert_eq!(next_resync(), None); + + // Recovered by a relist, which announces only its own view. + for _ in 0..3 { + cache.failed(); + tokio::time::advance(STALE_AFTER / 2).await; + } + assert_eq!(next_resync(), Some(Completeness::Estimated)); + cache.apply(watcher::Event::InitApply(ready())); + assert_eq!(next_resync(), None); + cache.apply(watcher::Event::InitDone); + assert_eq!(next_resync(), Some(Completeness::Exact)); + assert_eq!(next_resync(), None); + + // An error after a quiet spell ends the old run first, though + // `quiet` has not run. + for _ in 0..3 { + cache.failed(); + tokio::time::advance(STALE_AFTER / 2).await; + } + assert_eq!(next_resync(), Some(Completeness::Estimated)); + tokio::time::advance(RECOVERED_AFTER).await; + assert_eq!(cache.failed(), Duration::ZERO); + assert_eq!(next_resync(), Some(Completeness::Exact)); + assert_eq!(next_resync(), None); + } + + #[tokio::test(start_paused = true)] + async fn the_watch_ends_a_run_of_errors_that_falls_silent() { + let cache = cache(); + sync( + &cache, + vec![slice("west", "mocker", "a1b2c", &[("10.43.0.1", true)])], + ); + let mut events = cache.events.subscribe(); + let (errors, watch) = futures::channel::mpsc::unbounded(); + let stop = CancellationToken::new(); + let following = tokio::spawn({ + let (cache, stop) = (Arc::clone(&cache), stop.clone()); + async move { follow(watch, "test", &cache, &stop).await } + }); + for _ in 0..3 { + errors + .unbounded_send(Err(watcher::Error::NoResourceVersion)) + .unwrap(); + tokio::task::yield_now().await; + tokio::time::advance(STALE_AFTER / 2).await; + } + let resync = tokio::time::timeout(Duration::from_secs(1), events.recv()).await; + assert!( + matches!(resync, Ok(Ok(WatchEvent::Resync(s))) if s.completeness == Completeness::Estimated) + ); + + // kube resumed the watch without an event, so only the loop's timer + // can end the run. + tokio::time::advance(RECOVERED_AFTER).await; + let resync = tokio::time::timeout(Duration::from_secs(1), events.recv()).await; + assert!( + matches!(resync, Ok(Ok(WatchEvent::Resync(s))) if s.completeness == Completeness::Exact) + ); + stop.cancel(); + following.await.unwrap(); + } + + #[test] + fn relist_replaces_everything() { + let cache = cache(); + sync( + &cache, + vec![slice("west", "old", "a1b2c", &[("10.43.0.1", true)])], + ); + let mut events = cache.events.subscribe(); + cache.apply(watcher::Event::Init); + cache.apply(watcher::Event::InitApply(slice( + "west", + "new", + "d3e4f", + &[("10.43.0.2", true)], + ))); + // The old view stays until the relist completes. + assert!(listing(&cache, &key("west", "old")).is_some()); + cache.apply(watcher::Event::InitDone); + let WatchEvent::Resync(snapshot) = events.try_recv().unwrap() else { + panic!("expected Resync"); + }; + let keys: Vec = snapshot + .items + .into_iter() + .map(|listing| listing.key) + .collect(); + assert_eq!(keys, [key("west", "new")]); + assert!(events.try_recv().is_err()); + } + + #[test] + fn a_failed_relist_starts_over() { + let cache = cache(); + sync(&cache, Vec::new()); + // A relist fails after its first slices, as a paged one could. + cache.apply(watcher::Event::InitApply(slice( + "west", + "gone", + "a1b2c", + &[("10.43.0.1", true)], + ))); + cache.failed(); + cache.apply(watcher::Event::InitApply(slice( + "west", + "kept", + "d3e4f", + &[("10.43.0.2", true)], + ))); + cache.apply(watcher::Event::InitDone); + assert!(listing(&cache, &key("west", "gone")).is_none()); + assert!(listing(&cache, &key("west", "kept")).is_some()); + } + + #[tokio::test] + async fn watch_starts_with_added_then_resync() { + let cache = cache(); + let cancel = CancellationToken::new(); + let mut stream = cache.subscribe(cancel.clone(), CancellationToken::new()); + // Polled once, so it subscribes before the first list completes. + assert!(stream.next().now_or_never().is_none()); + sync( + &cache, + vec![slice("west", "mocker", "a1b2c", &[("10.43.0.1", true)])], + ); + assert!(matches!( + stream.next().await, + Some(Ok(WatchEvent::Upserted(_))) + )); + assert!( + matches!(stream.next().await, Some(Ok(WatchEvent::Resync(s))) if s.items.len() == 1) + ); + cache.apply(watcher::Event::Delete(slice( + "west", + "mocker", + "a1b2c", + &[], + ))); + assert!(matches!( + stream.next().await, + Some(Ok(WatchEvent::Removed(_))) + )); + cancel.cancel(); + assert!(stream.next().await.is_none()); + } + + #[tokio::test] + async fn late_subscriber_gets_the_current_pools() { + let cache = cache(); + sync( + &cache, + vec![slice("west", "mocker", "a1b2c", &[("10.43.0.1", true)])], + ); + let mut stream = cache.subscribe(CancellationToken::new(), CancellationToken::new()); + assert!( + matches!(stream.next().await, Some(Ok(WatchEvent::Upserted(r))) if r.key == key("west", "mocker")) + ); + assert!(matches!( + stream.next().await, + Some(Ok(WatchEvent::Resync(_))) + )); + } +} diff --git a/lib/global-discovery-service/src/backends/kubernetes/mcs.rs b/lib/global-discovery-service/src/backends/kubernetes/mcs.rs new file mode 100644 index 000000000000..4d810f7abd90 --- /dev/null +++ b/lib/global-discovery-service/src/backends/kubernetes/mcs.rs @@ -0,0 +1,270 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! How each MCS implementation names imported EndpointSlices. Only the +//! slices are read: they hold the endpoints, their readiness, and the source +//! cluster. + +use k8s_openapi::api::discovery::v1::EndpointSlice; + +/// How the MCS implementation names imported EndpointSlices. Not +/// `non_exhaustive`: the MCS backend handles every variant. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum McsImplementation { + /// The `multicluster.kubernetes.io` labels (Submariner, Cilium, GKE). + Standard, + /// Karmada's own names: a `derived-` Service, and slices named + /// `imported--`. + Karmada, +} + +/// Standard MCS: About API cluster ID of the cluster a slice came from. +pub const SOURCE_CLUSTER_LABEL: &str = "multicluster.kubernetes.io/source-cluster"; + +/// Standard MCS: name of the imported Service a slice belongs to. +pub const SERVICE_NAME_LABEL: &str = "multicluster.kubernetes.io/service-name"; + +/// Kubernetes: name of the Service a slice belongs to. +pub const KUBERNETES_SERVICE_NAME_LABEL: &str = "kubernetes.io/service-name"; + +/// Kubernetes: the controller that manages a slice. +pub const MANAGED_BY_LABEL: &str = "endpointslice.kubernetes.io/managed-by"; + +/// Karmada: imported slices belong to a Service named `derived-`, +/// through the usual `kubernetes.io/service-name` label. +pub const KARMADA_DERIVED_PREFIX: &str = "derived-"; + +/// Karmada names each imported slice `imported--`. +pub const KARMADA_IMPORTED_PREFIX: &str = "imported-"; + +/// Karmada's `endpointslice.kubernetes.io/managed-by` value on the slices it +/// collects. +pub const KARMADA_SLICE_MANAGER: &str = "endpointslice-controller.karmada.io"; + +/// Label selector for every imported slice. Narrowed further by name, since +/// not every implementation copies `exportedLabels`. +pub fn imported_selector(mcs: McsImplementation) -> String { + match mcs { + McsImplementation::Standard => SOURCE_CLUSTER_LABEL.to_owned(), + McsImplementation::Karmada => format!("{MANAGED_BY_LABEL}={KARMADA_SLICE_MANAGER}"), + } +} + +/// Label selector for one imported Service's slices. +pub fn service_selector(mcs: McsImplementation, service: &str) -> String { + match mcs { + McsImplementation::Standard => format!("{SERVICE_NAME_LABEL}={service}"), + McsImplementation::Karmada => { + format!("{KUBERNETES_SERVICE_NAME_LABEL}={KARMADA_DERIVED_PREFIX}{service}") + } + } +} + +/// The imported Service's name as exported, without Karmada's prefix. +pub fn service_name(mcs: McsImplementation, slice: &EndpointSlice) -> Option<&str> { + let labels = slice.metadata.labels.as_ref()?; + match mcs { + McsImplementation::Standard => labels.get(SERVICE_NAME_LABEL).map(String::as_str), + McsImplementation::Karmada => labels + .get(KUBERNETES_SERVICE_NAME_LABEL)? + .strip_prefix(KARMADA_DERIVED_PREFIX), + } +} + +/// The longest a Service name, or a label value, may be. +const MAX_NAME_LENGTH: usize = 63; + +/// The longest Service name Karmada can import. Karmada names the imported +/// Service `derived-`, and puts that name in each imported slice's +/// `kubernetes.io/service-name` label, so a longer name gets neither. +pub const KARMADA_MAX_SERVICE_NAME: usize = MAX_NAME_LENGTH - KARMADA_DERIVED_PREFIX.len(); + +/// Whether `mcs` can import a Service of this name. A directory lists +/// nothing for one it cannot. +pub fn can_import(mcs: McsImplementation, service: &str) -> bool { + match mcs { + McsImplementation::Standard => true, + McsImplementation::Karmada => service.len() <= KARMADA_MAX_SERVICE_NAME, + } +} + +/// Logs a slice of a Service the directory wants, but from an unknown +/// cluster. Such a slice is skipped, so its pool or router replica is not +/// listed. +pub fn warn_unknown_cluster(slice: &EndpointSlice, service: &str) { + tracing::warn!( + slice = slice.metadata.name.as_deref().unwrap_or_default(), + service, + "skipping an imported EndpointSlice that does not say which cluster it came from; with Karmada, its name must be imported---" + ); +} + +/// The cluster the slice came from. With Karmada this is the Karmada member +/// name, which must equal the About API cluster ID. +/// +/// Karmada has no source-cluster label, so the cluster comes from the slice +/// name: `imported--`, where the source slice name +/// starts with `-`. The name is the same on the Karmada control +/// plane and in member clusters. The `work.karmada.io/namespace` annotation +/// is not used: it names the source cluster (`karmada-es-`) only on +/// the control plane, and the destination cluster on copies that Karmada +/// propagates into members. +pub fn source_cluster(mcs: McsImplementation, slice: &EndpointSlice) -> Option<&str> { + match mcs { + McsImplementation::Standard => slice + .metadata + .labels + .as_ref()? + .get(SOURCE_CLUSTER_LABEL) + .map(String::as_str), + McsImplementation::Karmada => { + let service = service_name(mcs, slice)?; + let rest = slice + .metadata + .name + .as_deref()? + .strip_prefix(KARMADA_IMPORTED_PREFIX)?; + // The last `--` ends the cluster name. + rest.rmatch_indices(service) + .find_map(|(start, _)| { + let cluster = rest[..start].strip_suffix('-')?; + rest[start + service.len()..] + .starts_with('-') + .then_some(cluster) + }) + .filter(|cluster| !cluster.is_empty()) + } + } +} + +#[cfg(test)] +mod tests { + use std::collections::BTreeMap; + + use k8s_openapi::apimachinery::pkg::apis::meta::v1::ObjectMeta; + + use super::*; + + fn slice(name: &str, labels: &[(&str, &str)]) -> EndpointSlice { + EndpointSlice { + address_type: "IPv4".to_owned(), + endpoints: Vec::new(), + metadata: ObjectMeta { + name: Some(name.to_owned()), + labels: Some( + labels + .iter() + .map(|(key, value)| ((*key).to_owned(), (*value).to_owned())) + .collect::>(), + ), + ..Default::default() + }, + ports: None, + } + } + + fn karmada_slice(name: &str, service: &str) -> EndpointSlice { + slice( + name, + &[ + (MANAGED_BY_LABEL, KARMADA_SLICE_MANAGER), + ( + KUBERNETES_SERVICE_NAME_LABEL, + &format!("{KARMADA_DERIVED_PREFIX}{service}"), + ), + ], + ) + } + + #[test] + fn selectors() { + assert_eq!( + imported_selector(McsImplementation::Standard), + "multicluster.kubernetes.io/source-cluster" + ); + assert_eq!( + imported_selector(McsImplementation::Karmada), + "endpointslice.kubernetes.io/managed-by=endpointslice-controller.karmada.io" + ); + assert_eq!( + service_selector(McsImplementation::Standard, "pool-relay"), + "multicluster.kubernetes.io/service-name=pool-relay" + ); + assert_eq!( + service_selector(McsImplementation::Karmada, "pool-relay"), + "kubernetes.io/service-name=derived-pool-relay" + ); + } + + #[test] + fn standard_reads_labels() { + let slice = slice( + "pool-relay-x7k2p", + &[ + (SOURCE_CLUSTER_LABEL, "west"), + (SERVICE_NAME_LABEL, "pool-relay"), + ], + ); + assert_eq!( + service_name(McsImplementation::Standard, &slice), + Some("pool-relay") + ); + assert_eq!( + source_cluster(McsImplementation::Standard, &slice), + Some("west") + ); + } + + // Names as Karmada v1.19 created them in a two-cluster test. + #[test] + fn karmada_reads_service_and_cluster() { + let slice = karmada_slice( + "imported-aws-west-relay-mocker-west-1-28g2h", + "relay-mocker-west-1", + ); + assert_eq!( + service_name(McsImplementation::Karmada, &slice), + Some("relay-mocker-west-1") + ); + assert_eq!( + source_cluster(McsImplementation::Karmada, &slice), + Some("aws-west") + ); + } + + #[test] + fn karmada_cluster_name_may_contain_the_service_name() { + let slice = karmada_slice("imported-a-svc-b-svc-x1y2z", "svc"); + assert_eq!( + source_cluster(McsImplementation::Karmada, &slice), + Some("a-svc-b") + ); + } + + #[test] + fn karmada_rejects_unexpected_names() { + // Not imported by Karmada. + let local = karmada_slice("relay-mocker-west-1-28g2h", "relay-mocker-west-1"); + assert_eq!(source_cluster(McsImplementation::Karmada, &local), None); + // The source slice is not named after the Service. + let misnamed = karmada_slice("imported-west-relay-mock-28g2h", "relay-mocker-west-1"); + assert_eq!(source_cluster(McsImplementation::Karmada, &misnamed), None); + // No cluster before the service name. + let empty = karmada_slice("imported--svc-x1y2z", "svc"); + assert_eq!(source_cluster(McsImplementation::Karmada, &empty), None); + // Not a derived Service. + let plain = slice( + "imported-west-svc-x1y2z", + &[(KUBERNETES_SERVICE_NAME_LABEL, "svc")], + ); + assert_eq!(service_name(McsImplementation::Karmada, &plain), None); + } + + #[test] + fn karmada_imports_names_of_up_to_55_characters() { + assert_eq!(KARMADA_MAX_SERVICE_NAME, 55); + assert!(can_import(McsImplementation::Karmada, &"a".repeat(55))); + assert!(!can_import(McsImplementation::Karmada, &"a".repeat(56))); + assert!(can_import(McsImplementation::Standard, &"a".repeat(56))); + } +} diff --git a/lib/global-discovery-service/src/backends/memory.rs b/lib/global-discovery-service/src/backends/memory.rs new file mode 100644 index 000000000000..ee93dbe85bd6 --- /dev/null +++ b/lib/global-discovery-service/src/backends/memory.rs @@ -0,0 +1,484 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! The memory backend: one process keeps every registration. It needs no +//! Kubernetes, so it suits tests and a first deployment. A restarted process +//! starts empty and fills as PoolRelays register again, so its lists are +//! `Estimated` until one lease has passed: by then every live pool has +//! registered again. Router replicas come from configuration until routers +//! register too. Run one replica: each keeps only the registrations it +//! receives. + +use std::collections::HashMap; +use std::fmt; +use std::sync::{Arc, Mutex, MutexGuard, OnceLock, PoisonError}; +use std::time::Duration; + +use crate::pool_discovery::{ + Completeness, Directory, DiscoveryError, Lease, PoolEvent, PoolListing, PoolRegistrar, + PoolRegistration, RouterReplica, RouterReplicaEvent, RouterReplicaId, Snapshot, WatchEvent, + WatchStream, +}; +use async_trait::async_trait; +use dynamo_global_common::PoolKey; +use tokio::sync::broadcast; +use tokio::time::{Instant, MissedTickBehavior}; +use tokio_util::sync::CancellationToken; + +use super::stream::{EVENT_BUFFER, wake_at, watch_stream}; + +/// Pool registrations in memory. A pool expires the lease duration after its +/// last registration. Reads skip an expired pool at once; a sweep removes it, +/// with a `Removed` event, within one heartbeat interval. +/// +/// Any caller may register any pool: V1 assumes a trusted network, and +/// several PoolRelay replicas may register the same pool. +pub struct MemoryPoolStore { + lease: Lease, + started: Instant, + pools: Arc>, + /// Sent while holding the `pools` lock, so a subscriber that snapshots + /// and subscribes under it misses no event and sees none twice. + events: broadcast::Sender, + sweeping: OnceLock<()>, + stop: CancellationToken, +} + +#[derive(Default)] +struct Pools { + registrations: HashMap, +} + +struct Registered { + /// As the PoolRelay sent it, with its lists sorted. + registration: PoolRegistration, + expires_at: Instant, +} + +impl Pools { + fn snapshot(&self, now: Instant, completeness: Completeness) -> Snapshot { + Snapshot { + items: self + .registrations + .values() + .filter(|registered| registered.expires_at > now) + .map(|registered| listing(®istered.registration)) + .collect(), + completeness, + } + } +} + +impl MemoryPoolStore { + /// The service bounds the lease, since the store adds it to the current + /// time. + pub fn new(lease: Lease) -> Self { + Self { + lease, + started: Instant::now(), + pools: Arc::default(), + events: broadcast::channel(EVENT_BUFFER).0, + sweeping: OnceLock::new(), + stop: CancellationToken::new(), + } + } + + fn lock(&self) -> MutexGuard<'_, Pools> { + self.pools.lock().unwrap_or_else(PoisonError::into_inner) + } + + /// Every trait method runs in a Tokio runtime, so the spawn is safe. + fn ensure_sweeping(&self) { + self.sweeping.get_or_init(|| { + tokio::spawn(sweep( + Arc::clone(&self.pools), + self.events.clone(), + self.lease.heartbeat_interval(), + self.started.checked_add(self.lease.duration()), + self.stop.clone(), + )); + }); + } +} + +impl fmt::Debug for MemoryPoolStore { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("MemoryPoolStore") + .field("lease", &self.lease) + .finish_non_exhaustive() + } +} + +impl Drop for MemoryPoolStore { + fn drop(&mut self) { + self.stop.cancel(); + } +} + +/// `Exact` once a lease has passed since the store started. +fn completeness(started: Instant, lease: Lease, now: Instant) -> Completeness { + if now.saturating_duration_since(started) >= lease.duration() { + Completeness::Exact + } else { + Completeness::Estimated + } +} + +/// Removes expired pools and publishes `Removed` for each. At `complete_at`, +/// once a lease has passed since the store started, publishes a `Resync`, so +/// watchers learn that its lists are complete. +async fn sweep( + pools: Arc>, + events: broadcast::Sender, + every: Duration, + mut complete_at: Option, + stop: CancellationToken, +) { + let mut ticks = tokio::time::interval(every); + ticks.set_missed_tick_behavior(MissedTickBehavior::Delay); + loop { + tokio::select! { + biased; + () = stop.cancelled() => return, + () = wake_at(complete_at) => { + complete_at = None; + let pools = pools.lock().unwrap_or_else(PoisonError::into_inner); + let complete = pools.snapshot(Instant::now(), Completeness::Exact); + // No subscribers is not an error. + let _ = events.send(WatchEvent::Resync(complete)); + continue; + } + _ = ticks.tick() => {} + } + let now = Instant::now(); + let mut pools = pools.lock().unwrap_or_else(PoisonError::into_inner); + pools.registrations.retain(|key, registered| { + let live = registered.expires_at > now; + if !live { + // No subscribers is not an error. + let _ = events.send(WatchEvent::Removed(key.clone())); + } + live + }); + } +} + +/// Sorts a registration's lists and drops repeats, so the same lists in +/// another order are not a change. +fn sorted(mut registration: PoolRegistration) -> PoolRegistration { + registration.capabilities.sort_unstable(); + registration.capabilities.dedup(); + registration.frontends.sort_unstable(); + registration.frontends.dedup(); + registration.capacity.sort_unstable(); + registration.capacity.dedup(); + registration +} + +/// A registration as listed. The store does not know the PoolRelay's +/// addresses. +fn listing(registration: &PoolRegistration) -> PoolListing { + PoolListing { + key: registration.key.clone(), + location: registration.location.clone(), + model: Some(registration.model.clone()), + capabilities: registration.capabilities.clone(), + frontends: registration.frontends.clone(), + capacity: registration.capacity.clone(), + relay_identities: Vec::new(), + } +} + +#[async_trait] +impl PoolRegistrar for MemoryPoolStore { + /// An unchanged pool only renews its lease. A new or changed one replaces + /// the old listing, with an `Upserted` event. Lists in another order, for + /// example from another PoolRelay replica, are not a change. + async fn register(&self, registration: PoolRegistration) -> Result { + let registration = sorted(registration); + self.ensure_sweeping(); + let now = Instant::now(); + let expires_at = now + self.lease.duration(); + let mut pools = self.lock(); + if let Some(registered) = pools.registrations.get_mut(®istration.key) + && registered.expires_at > now + && registered.registration == registration + { + registered.expires_at = expires_at; + return Ok(self.lease); + } + let added = listing(®istration); + pools.registrations.insert( + registration.key.clone(), + Registered { + registration, + expires_at, + }, + ); + // No subscribers is not an error. + let _ = self.events.send(WatchEvent::Upserted(added)); + Ok(self.lease) + } +} + +#[async_trait] +impl Directory for MemoryPoolStore { + async fn list(&self) -> Result, DiscoveryError> { + let now = Instant::now(); + let completeness = completeness(self.started, self.lease, now); + Ok(self.lock().snapshot(now, completeness)) + } + + /// Starts from the view `list` returns, then follows registrations and + /// expired leases, and sends a `Resync` when the lists become `Exact`. + async fn list_and_watch( + &self, + cancel_token: Option, + ) -> Result, DiscoveryError> { + self.ensure_sweeping(); + let pools = Arc::clone(&self.pools); + let events = self.events.clone(); + let (started, lease) = (self.started, self.lease); + Ok(watch_stream( + move || { + let pools = pools.lock().unwrap_or_else(PoisonError::into_inner); + let now = Instant::now(); + let completeness = completeness(started, lease, now); + (events.subscribe(), Some(pools.snapshot(now, completeness))) + }, + cancel_token.unwrap_or_default(), + self.stop.clone(), + )) + } +} + +/// Router replicas from configuration. The set never changes, and lists are +/// `Exact`. +#[derive(Debug)] +pub struct StaticReplicaDirectory { + replicas: Vec, + /// Never sent on. Keeps watch streams open until the directory is + /// dropped. + events: broadcast::Sender, + stop: CancellationToken, +} + +impl StaticReplicaDirectory { + pub fn new(replicas: Vec) -> Self { + Self { + replicas, + events: broadcast::channel(1).0, + stop: CancellationToken::new(), + } + } + + fn snapshot(&self) -> Snapshot { + Snapshot { + items: self.replicas.clone(), + completeness: Completeness::Exact, + } + } +} + +impl Drop for StaticReplicaDirectory { + fn drop(&mut self) { + self.stop.cancel(); + } +} + +#[async_trait] +impl Directory for StaticReplicaDirectory { + async fn list(&self) -> Result, DiscoveryError> { + Ok(self.snapshot()) + } + + /// The configured replicas. The stream then stays open, with no events, + /// until the caller cancels it or the directory is dropped. + async fn list_and_watch( + &self, + cancel_token: Option, + ) -> Result, DiscoveryError> { + let events = self.events.clone(); + let snapshot = self.snapshot(); + Ok(watch_stream( + move || (events.subscribe(), Some(snapshot.clone())), + cancel_token.unwrap_or_default(), + self.stop.clone(), + )) + } +} + +#[cfg(test)] +mod tests { + use crate::pool_discovery::Capacity; + use dynamo_global_common::PoolLocation; + use futures::StreamExt; + use url::Url; + + use super::*; + + fn lease() -> Lease { + Lease::new(5, 15).unwrap() + } + + fn key(dgd: &str) -> PoolKey { + PoolKey::new("west", "dynamo", dgd).unwrap() + } + + fn capacity(sku: &str, workers: u32) -> Capacity { + Capacity { + sku: sku.to_owned(), + workers, + } + } + + fn registration(dgd: &str, model: &str) -> PoolRegistration { + PoolRegistration { + key: key(dgd), + location: PoolLocation { + region: "us-west-2".to_owned(), + availability_zone: None, + datacenter: None, + }, + model: model.to_owned(), + capabilities: vec!["prefill".to_owned(), "decode".to_owned()], + frontends: vec![ + Url::parse("http://10.43.1.3:8000/").unwrap(), + Url::parse("http://10.43.1.2:8000/").unwrap(), + ], + capacity: vec![capacity("H100", 8), capacity("B200", 4)], + } + } + + async fn find(store: &MemoryPoolStore, dgd: &str) -> Option { + let snapshot = store.list().await.unwrap(); + snapshot + .items + .into_iter() + .find(|listing| listing.key == key(dgd)) + } + + #[tokio::test(start_paused = true)] + async fn registering_again_only_renews() { + let store = MemoryPoolStore::new(lease()); + assert_eq!( + store.register(registration("mocker", "m1")).await.unwrap(), + lease() + ); + let mut watch = store.list_and_watch(None).await.unwrap(); + let Some(Ok(WatchEvent::Upserted(first))) = watch.next().await else { + panic!("expected the pool"); + }; + assert_eq!(first.capabilities, ["decode", "prefill"]); + assert!(matches!( + watch.next().await, + Some(Ok(WatchEvent::Resync(_))) + )); + + // Heartbeats every 10 s keep the 15 s lease alive past 15 s. Lists in + // another order, or with an item twice, are no change. + for _ in 0..3 { + tokio::time::advance(Duration::from_secs(10)).await; + let mut heartbeat = registration("mocker", "m1"); + heartbeat.capabilities.reverse(); + heartbeat.frontends.reverse(); + heartbeat.capacity.push(capacity("H100", 8)); + let _ = store.register(heartbeat).await.unwrap(); + } + assert_eq!(find(&store, "mocker").await, Some(first)); + // A lease has passed, so the store's lists are complete; the renewals + // published nothing. + assert!( + matches!(watch.next().await, Some(Ok(WatchEvent::Resync(s))) if s.completeness == Completeness::Exact) + ); + + // A changed pool is published. + let _ = store.register(registration("mocker", "m2")).await.unwrap(); + let Some(Ok(WatchEvent::Upserted(changed))) = watch.next().await else { + panic!("expected the change"); + }; + assert_eq!(changed.model.as_deref(), Some("m2")); + } + + #[tokio::test(start_paused = true)] + async fn a_missed_lease_hides_then_removes_the_pool() { + let store = MemoryPoolStore::new(lease()); + let _ = store.register(registration("mocker", "m1")).await.unwrap(); + let mut watch = store.list_and_watch(None).await.unwrap(); + assert!(matches!( + watch.next().await, + Some(Ok(WatchEvent::Upserted(_))) + )); + assert!(matches!( + watch.next().await, + Some(Ok(WatchEvent::Resync(_))) + )); + + tokio::time::advance(Duration::from_secs(16)).await; + // Reads skip it before the sweep runs. + assert!(find(&store, "mocker").await.is_none()); + // A lease has passed, so the store's lists are complete. Then the + // sweep removes the pool and tells watchers. + assert!(matches!( + watch.next().await, + Some(Ok(WatchEvent::Resync(_))) + )); + assert!( + matches!(watch.next().await, Some(Ok(WatchEvent::Removed(k))) if k == key("mocker")) + ); + assert!(store.lock().registrations.is_empty()); + + // Registering again brings it back. + let _ = store.register(registration("mocker", "m1")).await.unwrap(); + assert!( + matches!(watch.next().await, Some(Ok(WatchEvent::Upserted(listing))) if listing.key == key("mocker")) + ); + } + + #[tokio::test(start_paused = true)] + async fn lists_are_exact_once_a_lease_has_passed() { + let store = MemoryPoolStore::new(lease()); + let completeness = || async { store.list().await.unwrap().completeness }; + let mut watch = store.list_and_watch(None).await.unwrap(); + let Some(Ok(WatchEvent::Resync(first))) = watch.next().await else { + panic!("expected the first Resync"); + }; + assert_eq!(first.completeness, Completeness::Estimated); + assert_eq!(completeness().await, Completeness::Estimated); + + tokio::time::advance(lease().duration()).await; + assert_eq!(completeness().await, Completeness::Exact); + // Watchers learn it too. With no pools nothing else is published, so + // the wait is bounded: a missing Resync fails the test, not hangs it. + let next = tokio::time::timeout(Duration::from_secs(1), watch.next()).await; + let Ok(Some(Ok(WatchEvent::Resync(complete)))) = next else { + panic!("expected a Resync"); + }; + assert_eq!(complete.completeness, Completeness::Exact); + } + + #[tokio::test] + async fn static_replicas_list_and_watch() { + let replica = RouterReplica { + replica: RouterReplicaId("ohio/global-router-0".to_owned()), + url: Url::parse("http://10.42.0.5:50051").unwrap(), + }; + let directory = StaticReplicaDirectory::new(vec![replica.clone()]); + let snapshot = directory.list().await.unwrap(); + assert_eq!(snapshot.items, std::slice::from_ref(&replica)); + assert_eq!(snapshot.completeness, Completeness::Exact); + + let cancel = CancellationToken::new(); + let mut watch = directory + .list_and_watch(Some(cancel.clone())) + .await + .unwrap(); + assert!(matches!(watch.next().await, Some(Ok(WatchEvent::Upserted(r))) if r == replica)); + assert!(matches!( + watch.next().await, + Some(Ok(WatchEvent::Resync(_))) + )); + cancel.cancel(); + assert!(watch.next().await.is_none()); + } +} diff --git a/lib/global-discovery-service/src/backends/stream.rs b/lib/global-discovery-service/src/backends/stream.rs new file mode 100644 index 000000000000..0184b5076ffc --- /dev/null +++ b/lib/global-discovery-service/src/backends/stream.rs @@ -0,0 +1,141 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! The list-and-watch stream every directory here serves: one `Upserted` per +//! current item, then a `Resync` with the same set, then changes. A +//! subscriber that falls behind gets a fresh `Resync`. + +use crate::pool_discovery::{Snapshot, WatchEvent, WatchStream}; +use tokio::sync::broadcast; +use tokio::time::Instant; +use tokio_util::sync::CancellationToken; + +/// Events a slow subscriber may fall behind by before it gets a `Resync`. +pub(super) const EVENT_BUFFER: usize = 1024; + +/// `subscribe_now` must subscribe and take its snapshot under the lock that +/// senders hold while they send, so the stream misses no event and sees none +/// twice. The snapshot is `None` until the source has a complete view; the +/// stream then starts at the first `Resync` it receives. +pub(super) fn watch_stream( + subscribe_now: F, + cancel: CancellationToken, + stop: CancellationToken, +) -> WatchStream +where + T: Clone + Send + 'static, + K: Clone + Send + 'static, + F: Fn() -> (broadcast::Receiver>, Option>) + Send + 'static, +{ + Box::pin(async_stream::stream! { + let (mut receiver, snapshot) = subscribe_now(); + let mut started = snapshot.is_some(); + if let Some(snapshot) = snapshot { + for event in initial_events(snapshot) { + yield Ok(event); + } + } + loop { + let received = tokio::select! { + biased; + () = cancel.cancelled() => break, + () = stop.cancelled() => break, + received = receiver.recv() => received, + }; + match received { + // The first complete view arrived after this subscription. + Ok(WatchEvent::Resync(snapshot)) if !started => { + started = true; + for event in initial_events(snapshot) { + yield Ok(event); + } + } + Ok(event) => { + if started { + yield Ok(event); + } + } + Err(broadcast::error::RecvError::Lagged(missed)) => { + tracing::warn!(missed, "discovery watch fell behind; resyncing"); + let (fresh, snapshot) = subscribe_now(); + receiver = fresh; + if let Some(snapshot) = snapshot { + started = true; + yield Ok(WatchEvent::Resync(snapshot)); + } + } + Err(broadcast::error::RecvError::Closed) => break, + } + } + }) +} + +/// Sleeps until `at`, or forever without one: a timer branch in `select!` +/// that is off until it is armed. +pub(super) async fn wake_at(at: Option) { + match at { + Some(at) => tokio::time::sleep_until(at).await, + None => std::future::pending().await, + } +} + +fn initial_events(snapshot: Snapshot) -> Vec> { + let mut events: Vec> = snapshot + .items + .iter() + .cloned() + .map(WatchEvent::Upserted) + .collect(); + events.push(WatchEvent::Resync(snapshot)); + events +} + +#[cfg(test)] +mod tests { + use std::sync::{Arc, Mutex}; + + use futures::StreamExt; + + use super::*; + use crate::pool_discovery::Completeness; + + #[tokio::test] + async fn a_subscriber_that_falls_behind_gets_a_fresh_resync() { + let (events, _) = broadcast::channel::>(2); + let items = Arc::new(Mutex::new(vec![1])); + let subscribe_now = { + let events = events.clone(); + let items = Arc::clone(&items); + move || { + let snapshot = Snapshot { + items: items.lock().unwrap().clone(), + completeness: Completeness::Exact, + }; + (events.subscribe(), Some(snapshot)) + } + }; + let mut stream = watch_stream( + subscribe_now, + CancellationToken::new(), + CancellationToken::new(), + ); + assert!(matches!( + stream.next().await, + Some(Ok(WatchEvent::Upserted(1))) + )); + assert!(matches!( + stream.next().await, + Some(Ok(WatchEvent::Resync(_))) + )); + + // Three changes overflow the two-event buffer. + for item in 2..5 { + items.lock().unwrap().push(item); + events.send(WatchEvent::Upserted(item)).unwrap(); + } + let Some(Ok(WatchEvent::Resync(snapshot))) = stream.next().await else { + panic!("expected a fresh Resync"); + }; + assert_eq!(snapshot.items, [1, 2, 3, 4]); + } +} diff --git a/lib/global-discovery-service/src/lib.rs b/lib/global-discovery-service/src/lib.rs index ad1857d36fe1..6778fc5b74aa 100644 --- a/lib/global-discovery-service/src/lib.rs +++ b/lib/global-discovery-service/src/lib.rs @@ -2,3 +2,17 @@ // SPDX-License-Identifier: Apache-2.0 //! Pool discovery for the Global Router, served over HTTP. +//! +//! The contract comes first: +//! - [`service`]: the HTTP API, whose JSON types are in `dynamo_global_common`, +//! and [`PoolDiscoveryBackend`], the backend behind it. +//! - [`pool_discovery`]: the traits that backend is built from. +//! +//! [`backends`] implements those traits: registrations kept in memory, and +//! the EndpointSlices that SIG Multicluster (MCS) imports. + +pub mod backends; +pub mod pool_discovery; +pub mod service; + +pub use service::{PoolDiscoveryBackend, Tls, routes, serve}; diff --git a/lib/global-discovery-service/src/main.rs b/lib/global-discovery-service/src/main.rs new file mode 100644 index 000000000000..a3e29b9b291a --- /dev/null +++ b/lib/global-discovery-service/src/main.rs @@ -0,0 +1,444 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Global Discovery Service. See `dynamo_global_discovery_service::service`. +//! +//! Two flags choose the backend behind the API: +//! - `--pool-discovery`: where the pool and router lists come from. Either +//! the registrations PoolRelays send (`registrations`), or the +//! EndpointSlices that MCS imports, named the Karmada way (`karmada`) or +//! the standard way (`mcs`). +//! - `--registration-storage`: where registrations are kept. Only `memory` +//! for now. +//! +//! ```bash +//! cargo run -p dynamo-global-discovery-service -- \ +//! --pool-discovery registrations --router aws-ohio/global-router-0=http://10.42.0.5:50051 +//! ``` + +use std::collections::HashSet; +use std::ffi::OsString; +use std::net::SocketAddr; +use std::path::PathBuf; +use std::sync::Arc; +use std::time::Duration; + +use anyhow::{Context, Result, bail}; +use clap::error::ErrorKind; +use clap::parser::ValueSource; +use clap::{CommandFactory, FromArgMatches, Parser, ValueEnum}; +use dynamo_global_common::Lease; +use dynamo_global_discovery_service::backends::kubernetes::{ + ExportedService, McsImplementation, McsPoolDirectory, McsReplicaDirectory, +}; +use dynamo_global_discovery_service::backends::memory::{MemoryPoolStore, StaticReplicaDirectory}; +use dynamo_global_discovery_service::pool_discovery::{ + PoolDirectory, PoolRegistrar, RouterReplica, RouterReplicaId, +}; +use dynamo_global_discovery_service::service::{self, PoolDiscoveryBackend, Tls}; +use kube::Client; +use tokio::net::TcpListener; +use tokio_util::sync::CancellationToken; +use url::Url; + +/// The longest lease the service grants. `Lease` itself has no upper bound, +/// and the registration storage adds the lease to the current time, which a +/// lease near `u64::MAX` seconds would overflow. +const MAX_LEASE_DURATION_SECS: u64 = 86_400; + +/// How long the Kubernetes client waits to connect to the API server, TLS +/// included. kube's default is 30 s. +const KUBE_CONNECT_TIMEOUT: Duration = Duration::from_secs(10); + +/// How long the Kubernetes client waits for the API server to answer. Pool +/// discovery from MCS needs this and [`KUBE_CONNECT_TIMEOUT`] to add up to +/// under 90 s, so a server that stops answering shows up as failures; its +/// watches end after 50 s, so a working watch never reaches it. kube's +/// default is 295 s. +const KUBE_READ_TIMEOUT: Duration = Duration::from_secs(75); +const _: () = assert!(KUBE_CONNECT_TIMEOUT.as_secs() + KUBE_READ_TIMEOUT.as_secs() < 90); + +/// `--help` headings for the flags that work with only some `--pool-discovery` +/// values. A flag under the heading of another value fails at startup, rather +/// than being ignored. +const WITH_REGISTRATIONS: &str = "With --pool-discovery registrations"; +const WITH_MCS: &str = "With --pool-discovery karmada or mcs"; + +#[derive(Debug, Parser)] +#[command(about = "Serve Global Router pool discovery over HTTP")] +struct Args { + /// Address to serve on. Use `[::]:8080` on clusters with IPv6 pods. + #[arg(long, default_value = "0.0.0.0:8080")] + listen: SocketAddr, + /// Serve HTTPS with this certificate. Needs `--tls-key-path`. + #[arg(long, requires = "tls_key_path")] + tls_cert_path: Option, + /// The private key for `--tls-cert-path`. + #[arg(long, requires = "tls_cert_path")] + tls_key_path: Option, + /// Also require callers to present a certificate that this CA signed + /// (mTLS). Probes need one too, so probe readiness with TCP or a command. + #[arg(long, requires = "tls_cert_path")] + tls_client_ca_cert_path: Option, + /// Where the pool and router lists come from. + #[arg(long, value_enum, default_value_t = PoolDiscovery::Registrations)] + pool_discovery: PoolDiscovery, + /// Where registrations are kept. + #[arg( + long, + value_enum, + default_value_t = RegistrationStorage::Memory, + help_heading = WITH_REGISTRATIONS + )] + registration_storage: RegistrationStorage, + /// How often a PoolRelay registers again. + #[arg(long, default_value_t = 5, help_heading = WITH_REGISTRATIONS)] + heartbeat_interval_secs: u64, + /// How long a registration lasts. Longer than the heartbeat interval, and + /// at most a day. + #[arg( + long, + default_value_t = 15, + value_parser = clap::value_parser!(u64).range(..=MAX_LEASE_DURATION_SECS), + help_heading = WITH_REGISTRATIONS + )] + lease_duration_secs: u64, + /// A Global Router replica as `=`. Repeat it for each replica. + #[arg(long = "router", value_parser = parse_router, help_heading = WITH_REGISTRATIONS)] + routers: Vec, + /// Each PoolRelay Service is named ``. + #[arg( + long, + default_value = "-pool-relay", + allow_hyphen_values = true, + value_parser = clap::builder::NonEmptyStringValueParser::new(), + help_heading = WITH_MCS + )] + pool_service_suffix: String, + /// Namespace of the Global Router Service. + #[arg(long, help_heading = WITH_MCS, required_if_eq_any = MCS_VALUES)] + router_namespace: Option, + /// The Global Router Service, exported through MCS. + #[arg(long, help_heading = WITH_MCS, required_if_eq_any = MCS_VALUES)] + router_service: Option, + /// The router Service's named port for state streams. + #[arg(long, help_heading = WITH_MCS, required_if_eq_any = MCS_VALUES)] + router_port_name: Option, +} + +/// The `--pool-discovery` values that read MCS, for clap's +/// `required_if_eq_any`. +const MCS_VALUES: [(&str, &str); 2] = [("pool_discovery", "karmada"), ("pool_discovery", "mcs")]; + +/// Where the pool and router lists come from. +#[derive(Clone, Copy, Debug, PartialEq, Eq, ValueEnum)] +enum PoolDiscovery { + /// The pools PoolRelays register (`POST /pools`), while their leases + /// last. Router replicas come from `--router`. + Registrations, + /// The pools whose PoolRelay Services Karmada imports, and the replicas + /// of the router Service, from EndpointSlices named the Karmada way. + /// Takes no registrations. + Karmada, + /// As `karmada`, from EndpointSlices with the standard + /// `multicluster.kubernetes.io` labels. + Mcs, +} + +/// Where registrations are kept. +#[derive(Clone, Copy, Debug, PartialEq, Eq, ValueEnum)] +enum RegistrationStorage { + /// In this process. Run one replica: each keeps only the registrations it + /// receives. + Memory, +} + +fn parse_router(value: &str) -> Result { + let (id, url) = value + .split_once('=') + .ok_or("expected =".to_owned())?; + if id.is_empty() { + return Err("the replica id is empty".to_owned()); + } + let url = Url::parse(url).map_err(|error| error.to_string())?; + Ok(RouterReplica { + replica: RouterReplicaId(id.to_owned()), + url, + }) +} + +/// Parses the flags. A flag set under the `--help` heading of another +/// `--pool-discovery` value fails, as clap's own errors do. A flag left at +/// its default is not refused. +fn parse_args(args: I) -> Result +where + I: IntoIterator, + T: Into + Clone, +{ + let mut command = Args::command(); + let matches = command.try_get_matches_from_mut(args)?; + let parsed = Args::from_arg_matches(&matches)?; + let refused = match parsed.pool_discovery { + PoolDiscovery::Registrations => WITH_MCS, + PoolDiscovery::Karmada | PoolDiscovery::Mcs => WITH_REGISTRATIONS, + }; + let given = command + .get_arguments() + .find(|arg| { + arg.get_help_heading() == Some(refused) + && matches + .value_source(arg.get_id().as_str()) + .is_some_and(|source| source != ValueSource::DefaultValue) + }) + .map(|arg| arg.get_long().unwrap_or(arg.get_id().as_str()).to_owned()); + match given { + Some(flag) => Err(command.error( + ErrorKind::ArgumentConflict, + format!("--{flag} works only {}", refused.to_lowercase()), + )), + None => Ok(parsed), + } +} + +#[tokio::main] +async fn main() -> Result<()> { + dynamo_runtime::logging::init(); + let args = parse_args(std::env::args_os()).unwrap_or_else(|error| error.exit()); + let backend = match args.pool_discovery { + PoolDiscovery::Registrations => registrations_backend(&args)?, + PoolDiscovery::Karmada => mcs_backend(&args, McsImplementation::Karmada).await?, + PoolDiscovery::Mcs => mcs_backend(&args, McsImplementation::Standard).await?, + }; + let listener = TcpListener::bind(args.listen) + .await + .with_context(|| format!("bind the Global Discovery listener at {}", args.listen))?; + tracing::info!(address = %args.listen, "Global Discovery Service listening"); + + let shutdown = CancellationToken::new(); + cancel_on_signal(shutdown.clone())?; + let tls = args + .tls_cert_path + .clone() + .zip(args.tls_key_path.clone()) + .map(|(cert, key)| Tls { + cert, + key, + client_ca: args.tls_client_ca_cert_path.clone(), + }); + service::serve(listener, backend, tls, shutdown) + .await + .context("serve the Global Discovery Service") +} + +/// Pools from registrations; router replicas from `--router`. +fn registrations_backend(args: &Args) -> Result { + let lease = Lease::new(args.heartbeat_interval_secs, args.lease_duration_secs) + .context("check --heartbeat-interval-secs and --lease-duration-secs")?; + let registrations = match args.registration_storage { + RegistrationStorage::Memory => Arc::new(MemoryPoolStore::new(lease)), + }; + let mut ids = HashSet::new(); + if let Some(repeated) = args + .routers + .iter() + .find(|router| !ids.insert(&router.replica)) + { + bail!("--router {} is given twice", repeated.replica.0); + } + let pools: Arc = registrations.clone(); + let registrar: Arc = registrations; + Ok(PoolDiscoveryBackend { + pools, + routers: Arc::new(StaticReplicaDirectory::new(args.routers.clone())), + registrar: Some(registrar), + }) +} + +/// Pools from the PoolRelay Services that MCS imports, and router replicas +/// from the router Service's. +async fn mcs_backend(args: &Args, mcs: McsImplementation) -> Result { + let (Some(namespace), Some(name), Some(port_name)) = ( + &args.router_namespace, + &args.router_service, + &args.router_port_name, + ) else { + bail!( + "--pool-discovery karmada and mcs need --router-namespace, --router-service and \ + --router-port-name" + ); + }; + let client = kubernetes_client() + .await + .context("create the Kubernetes client")?; + let router = ExportedService { + namespace: namespace.clone(), + name: name.clone(), + port_name: port_name.clone(), + }; + Ok(PoolDiscoveryBackend { + pools: Arc::new(McsPoolDirectory::new( + client.clone(), + mcs, + args.pool_service_suffix.clone(), + )), + routers: Arc::new(McsReplicaDirectory::new(client, &router, mcs)), + registrar: None, + }) +} + +/// A client configured as `Client::try_default` does, from the pod's service +/// account or the local kubeconfig, but with [`KUBE_CONNECT_TIMEOUT`] and +/// [`KUBE_READ_TIMEOUT`]. +async fn kubernetes_client() -> Result { + let mut config = kube::Config::infer().await?; + config.connect_timeout = Some(KUBE_CONNECT_TIMEOUT); + config.read_timeout = Some(KUBE_READ_TIMEOUT); + Ok(Client::try_from(config)?) +} + +/// Cancels `shutdown` on SIGTERM or Ctrl-C. Fails at startup if it cannot +/// listen for SIGTERM, so the service never runs without a clean way to stop. +fn cancel_on_signal(shutdown: CancellationToken) -> Result<()> { + #[cfg(unix)] + let mut terminate = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) + .context("listen for SIGTERM")?; + tokio::spawn(async move { + #[cfg(unix)] + let terminated = terminate.recv(); + #[cfg(not(unix))] + let terminated = std::future::pending::>(); + tokio::select! { + // If Ctrl-C cannot be watched, SIGTERM still stops the service. + Ok(()) = tokio::signal::ctrl_c() => {} + _ = terminated => {} + } + shutdown.cancel(); + }); + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + const ROUTER: &str = "--router=aws-ohio/global-router-0=http://10.42.0.5:50051"; + + /// `--pool-discovery karmada` with the router flags it needs, then `extra`. + fn karmada<'a>(extra: &[&'a str]) -> Vec<&'a str> { + let mut args = vec![ + "gds", + "--pool-discovery", + "karmada", + "--router-namespace", + "global-discovery", + "--router-service", + "global-router", + "--router-port-name", + "http", + ]; + args.extend_from_slice(extra); + args + } + + /// `args` fail because `--{flag}` works only with `--pool-discovery {values}`. + fn assert_refused(args: &[&str], flag: &str, values: &str) { + let error = parse_args(args.iter().copied()).unwrap_err(); + assert_eq!(error.kind(), ErrorKind::ArgumentConflict, "{error}"); + let reason = format!("--{flag} works only with --pool-discovery {values}"); + assert!(error.to_string().contains(&reason), "{error}"); + } + + #[test] + fn pools_come_from_registrations_kept_in_memory_by_default() { + let args = parse_args(["gds"]).unwrap(); + assert_eq!(args.pool_discovery, PoolDiscovery::Registrations); + assert_eq!(args.registration_storage, RegistrationStorage::Memory); + } + + #[test] + fn every_flag_says_which_pool_discovery_it_works_with() { + // The flags that work with every --pool-discovery value. Any other + // flag needs one of the two headings, which `parse_args` enforces. + const EVERY_VALUE: &[&str] = &[ + "listen", + "tls_cert_path", + "tls_key_path", + "tls_client_ca_cert_path", + "pool_discovery", + "help", + ]; + for arg in Args::command().get_arguments() { + let id = arg.get_id().as_str(); + let headed = matches!(arg.get_help_heading(), Some(WITH_REGISTRATIONS | WITH_MCS)); + assert_ne!( + EVERY_VALUE.contains(&id), + headed, + "--{id}: list it in EVERY_VALUE or give it a --pool-discovery heading" + ); + } + } + + #[test] + fn mcs_needs_the_router_service() { + let error = parse_args(["gds", "--pool-discovery", "mcs"]).unwrap_err(); + assert_eq!(error.kind(), ErrorKind::MissingRequiredArgument, "{error}"); + let all = karmada(&[]); + assert!(parse_args(all.clone()).is_ok()); + for flag in [ + "--router-namespace", + "--router-service", + "--router-port-name", + ] { + let mut args = all.clone(); + let at = args.iter().position(|arg| *arg == flag).unwrap(); + args.drain(at..at + 2); + let error = parse_args(args).unwrap_err(); + assert_eq!(error.kind(), ErrorKind::MissingRequiredArgument, "{error}"); + assert!(error.to_string().contains(flag), "{error}"); + } + } + + #[test] + fn a_flag_for_another_pool_discovery_fails() { + assert!(parse_args(["gds", ROUTER]).is_ok()); + assert_refused(&karmada(&[ROUTER]), "router", "registrations"); + assert_refused( + &karmada(&["--lease-duration-secs", "30"]), + "lease-duration-secs", + "registrations", + ); + assert_refused( + &karmada(&["--registration-storage", "memory"]), + "registration-storage", + "registrations", + ); + assert_refused( + &["gds", "--router-service", "global-router"], + "router-service", + "karmada or mcs", + ); + assert_refused( + &["gds", "--pool-service-suffix=-relay"], + "pool-service-suffix", + "karmada or mcs", + ); + // Flags left at their defaults are not refused, whatever their heading. + assert!(parse_args(karmada(&[])).is_ok()); + } + + #[test] + fn pool_service_suffix_can_start_with_a_hyphen_but_not_be_empty() { + let args = parse_args(karmada(&["--pool-service-suffix", "-pool-gateway"])).unwrap(); + assert_eq!(args.pool_service_suffix, "-pool-gateway"); + // An empty suffix would make every imported Service a pool. + assert!(parse_args(karmada(&["--pool-service-suffix="])).is_err()); + } + + #[test] + fn a_lease_is_at_most_a_day() { + let day = MAX_LEASE_DURATION_SECS.to_string(); + assert!(parse_args(["gds", "--lease-duration-secs", &day]).is_ok()); + let longer = (MAX_LEASE_DURATION_SECS + 1).to_string(); + assert!(parse_args(["gds", "--lease-duration-secs", &longer]).is_err()); + } +} diff --git a/lib/global-discovery-service/src/pool_discovery.rs b/lib/global-discovery-service/src/pool_discovery.rs new file mode 100644 index 000000000000..0c83205d9d42 --- /dev/null +++ b/lib/global-discovery-service/src/pool_discovery.rs @@ -0,0 +1,36 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! The contract between the Global Discovery Service and its backends. +//! Design: . +//! +//! The service reads and writes only through this contract, so a backend +//! can change without changing PoolRelays or routers: +//! - [`PoolDirectory`]: which pools exist. +//! - [`PoolRegistrar`]: a PoolRelay registers its pool. Each registration is +//! also the pool's heartbeat. +//! - [`ReplicaDirectory`]: the Global Router replicas a PoolRelay connects to. +//! +//! The two directories are aliases of one trait, [`Directory`], over pools and +//! over router replicas. What they list, and what a PoolRelay registers, are +//! the service's JSON types from `dynamo_global_common`. + +mod pools; +mod registrar; +mod replicas; +mod watch; + +pub use pools::{Capacity, PoolDirectory, PoolEvent, PoolListing}; +pub use registrar::{Lease, PoolRegistrar, PoolRegistration}; +pub use replicas::{ReplicaDirectory, RouterReplica, RouterReplicaEvent, RouterReplicaId}; +pub use watch::{Completeness, Directory, Snapshot, WatchEvent, WatchStream}; + +use thiserror::Error; + +/// Shared by every backend, so the service maps each case to one status code. +#[derive(Debug, Error)] +#[non_exhaustive] +pub enum DiscoveryError { + #[error("discovery backend is unavailable")] + Unavailable(#[source] anyhow::Error), +} diff --git a/lib/global-discovery-service/src/pool_discovery/pools.rs b/lib/global-discovery-service/src/pool_discovery/pools.rs new file mode 100644 index 000000000000..be053dda1025 --- /dev/null +++ b/lib/global-discovery-service/src/pool_discovery/pools.rs @@ -0,0 +1,13 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use dynamo_global_common::PoolKey; +pub use dynamo_global_common::{Capacity, PoolListing}; + +use super::watch::{Directory, WatchEvent}; + +/// The pools discovery knows. A pool's state (load, KV, readiness) is not part +/// of its listing. Router replicas derive the routing `PoolId` from each +/// listing's `key` with the shared `PoolIdDeriver`. +pub type PoolDirectory = dyn Directory; +pub type PoolEvent = WatchEvent; diff --git a/lib/global-discovery-service/src/pool_discovery/registrar.rs b/lib/global-discovery-service/src/pool_discovery/registrar.rs new file mode 100644 index 000000000000..2f720c60b5c1 --- /dev/null +++ b/lib/global-discovery-service/src/pool_discovery/registrar.rs @@ -0,0 +1,21 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use async_trait::async_trait; +pub use dynamo_global_common::{Lease, PoolRegistration}; + +use super::DiscoveryError; + +/// The write side. A PoolRelay registers its pool on a timer, and each call +/// is also the pool's heartbeat: registering an unchanged pool again only +/// renews it. The [`Lease`] says when to register again, and how long the +/// pool stays listed after its last registration, measured on the backend's +/// clock. PoolRelays get router replicas from a [`super::ReplicaDirectory`], +/// not from this reply. +/// +/// V1 assumes a trusted network, where any caller may register any pool, so +/// `register` does not take the caller. +#[async_trait] +pub trait PoolRegistrar: Send + Sync { + async fn register(&self, registration: PoolRegistration) -> Result; +} diff --git a/lib/global-discovery-service/src/pool_discovery/replicas.rs b/lib/global-discovery-service/src/pool_discovery/replicas.rs new file mode 100644 index 000000000000..637c7c9cdc9a --- /dev/null +++ b/lib/global-discovery-service/src/pool_discovery/replicas.rs @@ -0,0 +1,16 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Global Router replicas a pool must connect to. Every replica needs every +//! pool's state, so a PoolRelay opens one connection per replica, and gets +//! the list from the service. +//! +//! A replica's readiness must not wait for its pools: Relays dial only +//! listed (Ready) replicas, so otherwise no pool could ever reach it. Pool +//! completeness is a separate signal ([`super::Completeness`]). + +pub use dynamo_global_common::{RouterReplica, RouterReplicaEvent, RouterReplicaId}; + +use super::watch::Directory; + +pub type ReplicaDirectory = dyn Directory; diff --git a/lib/global-discovery-service/src/pool_discovery/watch.rs b/lib/global-discovery-service/src/pool_discovery/watch.rs new file mode 100644 index 000000000000..e478430a9e7c --- /dev/null +++ b/lib/global-discovery-service/src/pool_discovery/watch.rs @@ -0,0 +1,30 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! The list-and-watch contract of `dynamo_runtime::discovery::Discovery`, +//! made generic so pools and router replicas share it. Its events and +//! snapshots are the ones in `dynamo_global_common`. + +use std::pin::Pin; + +use async_trait::async_trait; +use futures::Stream; +use tokio_util::sync::CancellationToken; + +pub use dynamo_global_common::{Completeness, Snapshot, WatchEvent}; + +use super::DiscoveryError; + +pub type WatchStream = + Pin, DiscoveryError>> + Send>>; + +/// Read side. Method names follow `dynamo_runtime::discovery::Discovery`. +#[async_trait] +pub trait Directory: Send + Sync { + async fn list(&self) -> Result, DiscoveryError>; + + async fn list_and_watch( + &self, + cancel_token: Option, + ) -> Result, DiscoveryError>; +} diff --git a/lib/global-discovery-service/src/service.rs b/lib/global-discovery-service/src/service.rs new file mode 100644 index 000000000000..6f6b0966f25a --- /dev/null +++ b/lib/global-discovery-service/src/service.rs @@ -0,0 +1,358 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! The Global Discovery Service's HTTP API. Its JSON types are in +//! `dynamo_global_common`, shared with the service's clients. +//! +//! - `POST /pools`: a PoolRelay registers its pool. It calls this on a timer, +//! and each call is also the pool's heartbeat. +//! - `GET /routers`: the Global Router replicas a PoolRelay connects to. +//! - `GET /pools`: the pools a Global Router routes to. +//! - `GET /healthz`: OK while the service can answer from its backend. Use it +//! for readiness; liveness needs only the TCP port. With mTLS, a probe +//! needs a client certificate too, so probe with TCP or a command. +//! +//! The service reads and writes only through the [`crate::pool_discovery`] +//! traits. It runs on `dynamo-http-server`, so it gets TLS and mTLS the way +//! the Dynamo frontend does. + +use std::fmt; +use std::path::PathBuf; +use std::sync::Arc; + +use crate::pool_discovery::{ + Completeness, DiscoveryError, PoolDirectory, PoolRegistrar, PoolRegistration, ReplicaDirectory, +}; +use axum::extract::State; +use axum::http::StatusCode; +use axum::response::{IntoResponse, Response}; +use axum::routing::get; +use axum::{Json, Router}; +use dynamo_global_common::{PoolsResponse, RoutersResponse}; +use dynamo_http_server::lifecycle::ServiceObserver; +use dynamo_http_server::server::{HttpServer, HttpServerConfig}; +use tokio::net::TcpListener; +use tokio_util::sync::CancellationToken; + +/// The backend behind the API: where the pool and router lists come from, +/// and where registrations go. +pub struct PoolDiscoveryBackend { + pub pools: Arc, + pub routers: Arc, + /// `None` when the backend takes no registrations: `POST /pools` then + /// answers `501`. + pub registrar: Option>, +} + +impl fmt::Debug for PoolDiscoveryBackend { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("PoolDiscoveryBackend") + .field("takes_registrations", &self.registrar.is_some()) + .finish_non_exhaustive() + } +} + +/// TLS for the listener. With `client_ca`, callers must also present a +/// certificate that it signed (mTLS). +#[derive(Clone, Debug)] +pub struct Tls { + pub cert: PathBuf, + pub key: PathBuf, + pub client_ca: Option, +} + +/// The service's HTTP routes. +pub fn routes(backend: PoolDiscoveryBackend) -> Router { + Router::new() + .route("/pools", get(list_pools).post(register_pool)) + .route("/routers", get(list_routers)) + .route("/healthz", get(healthz)) + .with_state(Arc::new(backend)) +} + +/// Serves until `shutdown` is cancelled. Then it stops taking connections and +/// gives requests in flight up to `DYN_HTTP_GRACEFUL_SHUTDOWN_TIMEOUT_SECS` +/// (5 s by default) to finish. +pub async fn serve( + listener: TcpListener, + backend: PoolDiscoveryBackend, + tls: Option, + shutdown: CancellationToken, +) -> anyhow::Result<()> { + let address = listener.local_addr()?; + let config = HttpServerConfig { + host: address.ip().to_string(), + port: address.port(), + enable_tls: tls.is_some(), + tls_cert_path: tls.as_ref().map(|tls| tls.cert.clone()), + tls_key_path: tls.as_ref().map(|tls| tls.key.clone()), + tls_client_ca_cert_path: tls.and_then(|tls| tls.client_ca), + }; + // The observer counts inference requests, and this service has none, so + // shutdown goes straight to letting open requests finish. + let observer = Arc::new(ServiceObserver::default()); + HttpServer::prepare(routes(backend), config, Some(listener)) + .await? + .serve(observer, CancellationToken::new(), shutdown) + .await +} + +async fn register_pool( + State(backend): State>, + Json(registration): Json, +) -> Response { + let Some(registrar) = &backend.registrar else { + return ( + StatusCode::NOT_IMPLEMENTED, + "this service takes no registrations", + ) + .into_response(); + }; + match registrar.register(registration).await { + Ok(lease) => Json(lease).into_response(), + Err(error) => error_response(error), + } +} + +async fn list_pools(State(backend): State>) -> Response { + match backend.pools.list().await { + Ok(snapshot) => Json(PoolsResponse { + complete: snapshot.completeness == Completeness::Exact, + pools: snapshot.items, + }) + .into_response(), + Err(error) => error_response(error), + } +} + +async fn list_routers(State(backend): State>) -> Response { + match backend.routers.list().await { + Ok(snapshot) => Json(RoutersResponse { + routers: snapshot.items, + }) + .into_response(), + Err(error) => error_response(error), + } +} + +/// Lists what `GET /pools` and `GET /routers` would serve. With pools from +/// MCS, the backend answers from its cache once that has synced, so an API +/// server outage does not turn every replica unready at once. The first probe +/// also starts its watch. +async fn healthz(State(backend): State>) -> StatusCode { + let (pools, routers) = tokio::join!(backend.pools.list(), backend.routers.list()); + match pools.and(routers) { + Ok(_) => StatusCode::OK, + Err(error) => { + tracing::warn!(?error, "discovery backend is unreachable"); + StatusCode::SERVICE_UNAVAILABLE + } + } +} + +/// Maps a backend error to a status. The cause is logged, not returned. +fn error_response(error: DiscoveryError) -> Response { + let status = match &error { + DiscoveryError::Unavailable(_) => StatusCode::SERVICE_UNAVAILABLE, + }; + tracing::warn!(?error, %status, "discovery request failed"); + (status, error.to_string()).into_response() +} + +#[cfg(test)] +mod tests { + use crate::pool_discovery::{Directory, RouterReplica, RouterReplicaId, Snapshot, WatchStream}; + use async_trait::async_trait; + use dynamo_global_common::Lease; + use serde_json::{Value, json}; + use url::Url; + + use super::*; + use crate::backends::memory::{MemoryPoolStore, StaticReplicaDirectory}; + + /// The body `dynamo-global-discovery-client` sends. + fn registration() -> Value { + json!({ + "key": {"cluster_id": "aws-west", "namespace": "dynamo", "dgd_name": "mocker-west-1"}, + "location": { + "region": "us-west-2", + "availability_zone": null, + "datacenter": null + }, + "model": "Qwen/Qwen3-0.6B", + "capabilities": ["aggregated"], + "frontends": ["http://10.43.1.2:8000/"], + "capacity": [{"sku": "H100", "workers": 8}] + }) + } + + fn router_replica() -> RouterReplica { + RouterReplica { + replica: RouterReplicaId("aws-ohio/global-router-0".to_owned()), + url: Url::parse("http://10.42.0.5:50051").unwrap(), + } + } + + fn registrations_backend() -> PoolDiscoveryBackend { + let pools = Arc::new(MemoryPoolStore::new(Lease::new(5, 15).unwrap())); + PoolDiscoveryBackend { + pools: pools.clone(), + routers: Arc::new(StaticReplicaDirectory::new(vec![router_replica()])), + registrar: Some(pools), + } + } + + /// Serves `backend` on a free local port and returns its base URL. + async fn start(backend: PoolDiscoveryBackend) -> (String, CancellationToken) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base = format!("http://{}", listener.local_addr().unwrap()); + let shutdown = CancellationToken::new(); + tokio::spawn(serve(listener, backend, None, shutdown.clone())); + (base, shutdown) + } + + async fn get_json(base: &str, path: &str) -> (reqwest::StatusCode, Value) { + let response = reqwest::get(format!("{base}{path}")).await.unwrap(); + let status = response.status(); + let body = response.bytes().await.unwrap(); + (status, serde_json::from_slice(&body).unwrap_or(Value::Null)) + } + + async fn post_pool(base: &str, body: &Value) -> (reqwest::StatusCode, Value) { + let response = reqwest::Client::new() + .post(format!("{base}/pools")) + .header("content-type", "application/json") + .body(body.to_string()) + .send() + .await + .unwrap(); + let status = response.status(); + let body = response.bytes().await.unwrap(); + (status, serde_json::from_slice(&body).unwrap_or(Value::Null)) + } + + #[tokio::test] + async fn register_then_list_pools() { + let (base, shutdown) = start(registrations_backend()).await; + let (status, lease) = post_pool(&base, ®istration()).await; + assert_eq!(status, reqwest::StatusCode::OK); + assert_eq!( + lease, + json!({"heartbeat_interval_s": 5, "lease_duration_s": 15}) + ); + + let (status, pools) = get_json(&base, "/pools").await; + assert_eq!(status, reqwest::StatusCode::OK); + assert_eq!( + pools, + json!({ + "pools": [{ + "key": {"cluster_id": "aws-west", "namespace": "dynamo", "dgd_name": "mocker-west-1"}, + "location": { + "region": "us-west-2", + "availability_zone": null, + "datacenter": null + }, + "model": "Qwen/Qwen3-0.6B", + "capabilities": ["aggregated"], + "frontends": ["http://10.43.1.2:8000/"], + "capacity": [{"sku": "H100", "workers": 8}], + "relay_identities": [] + }], + "complete": false + }) + ); + shutdown.cancel(); + } + + #[tokio::test] + async fn routers_use_the_client_shape() { + let (base, shutdown) = start(registrations_backend()).await; + let (status, routers) = get_json(&base, "/routers").await; + assert_eq!(status, reqwest::StatusCode::OK); + assert_eq!( + routers, + json!({"routers": [{"replica": "aws-ohio/global-router-0", "url": "http://10.42.0.5:50051/"}]}) + ); + shutdown.cancel(); + } + + #[tokio::test] + async fn invalid_key_is_rejected() { + let (base, shutdown) = start(registrations_backend()).await; + let mut body = registration(); + body["key"]["cluster_id"] = json!(" "); + let (status, _) = post_pool(&base, &body).await; + assert_eq!(status, reqwest::StatusCode::UNPROCESSABLE_ENTITY); + let (_, pools) = get_json(&base, "/pools").await; + assert_eq!(pools["pools"], json!([])); + shutdown.cancel(); + } + + #[tokio::test] + async fn backend_without_registrar_refuses_registrations() { + let mut backend = registrations_backend(); + backend.registrar = None; + let (base, shutdown) = start(backend).await; + let (status, _) = post_pool(&base, ®istration()).await; + assert_eq!(status, reqwest::StatusCode::NOT_IMPLEMENTED); + shutdown.cancel(); + } + + /// A directory that is down. + struct Down; + + #[async_trait] + impl Directory for Down { + async fn list(&self) -> Result, DiscoveryError> { + Err(DiscoveryError::Unavailable(anyhow::anyhow!("down"))) + } + + async fn list_and_watch( + &self, + _cancel_token: Option, + ) -> Result, DiscoveryError> { + Err(DiscoveryError::Unavailable(anyhow::anyhow!("down"))) + } + } + + #[tokio::test] + async fn unreachable_backend_is_unavailable() { + let (base, shutdown) = start(PoolDiscoveryBackend { + pools: Arc::new(Down), + routers: Arc::new(Down), + registrar: None, + }) + .await; + for path in ["/pools", "/routers", "/healthz"] { + let response = reqwest::get(format!("{base}{path}")).await.unwrap(); + assert_eq!( + response.status(), + reqwest::StatusCode::SERVICE_UNAVAILABLE, + "{path}" + ); + // The cause stays in the log. + assert!(!response.text().await.unwrap().contains("down"), "{path}"); + } + shutdown.cancel(); + } + + #[tokio::test] + async fn stops_when_shut_down() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let healthz = format!("http://{}/healthz", listener.local_addr().unwrap()); + let shutdown = CancellationToken::new(); + let server = tokio::spawn(serve( + listener, + registrations_backend(), + None, + shutdown.clone(), + )); + let response = reqwest::get(&healthz).await.unwrap(); + assert_eq!(response.status(), reqwest::StatusCode::OK); + + shutdown.cancel(); + server.await.unwrap().unwrap(); + assert!(reqwest::get(&healthz).await.is_err()); + } +} diff --git a/lib/http-server/Cargo.toml b/lib/http-server/Cargo.toml new file mode 100644 index 000000000000..829c94c7aac8 --- /dev/null +++ b/lib/http-server/Cargo.toml @@ -0,0 +1,33 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +[package] +name = "dynamo-http-server" +version.workspace = true +edition.workspace = true +authors.workspace = true +license.workspace = true +homepage.workspace = true +repository.workspace = true +description = "Plane-independent production HTTP server for Dynamo services" + +[dependencies] +anyhow = { workspace = true } +axum = { workspace = true } +hyper-util = { version = "0.1.17", features = ["server-auto", "service", "tokio"] } +tokio-rustls = { workspace = true } +tower = { workspace = true } +parking_lot = { workspace = true } +rand = { workspace = true } +serde = { workspace = true } +dynamo-runtime = { workspace = true } +futures = { workspace = true } +prometheus = { workspace = true } +tokio = { workspace = true } +tokio-util = { workspace = true } +tracing = { workspace = true } + +[dev-dependencies] +rcgen = { workspace = true } +rustls = { workspace = true } +tokio = { workspace = true, features = ["test-util"] } diff --git a/lib/http-server/src/disconnect.rs b/lib/http-server/src/disconnect.rs new file mode 100644 index 000000000000..cfee7a84977b --- /dev/null +++ b/lib/http-server/src/disconnect.rs @@ -0,0 +1,138 @@ +// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Connection monitoring for HTTP and gRPC services. + +use crate::metrics::{CancellationLabels, Metrics}; +use dynamo_runtime::engine::AsyncEngineContext; +use std::sync::Arc; + +#[derive(Clone, Copy)] +pub enum ConnectionStatus { + Disabled, + ClosedUnexpectedly, + ClosedGracefully, +} + +pub struct ConnectionHandle { + sender: Option>, + on_drop: ConnectionStatus, +} + +impl ConnectionHandle { + /// Handle which by default will issue a [`ConnectionStatus::ClosedGracefully`] signal when dropped. + pub fn create_disarmed(sender: tokio::sync::oneshot::Sender) -> Self { + Self { + sender: Some(sender), + on_drop: ConnectionStatus::ClosedGracefully, + } + } + + /// Handle which will issue a [`ConnectionStatus::ClosedUnexpectedly`] signal when dropped. + pub fn create_armed(sender: tokio::sync::oneshot::Sender) -> Self { + Self { + sender: Some(sender), + on_drop: ConnectionStatus::ClosedUnexpectedly, + } + } + + /// Handle which will not issue a signal when dropped. + pub fn create_disabled(sender: tokio::sync::oneshot::Sender) -> Self { + Self { + sender: Some(sender), + on_drop: ConnectionStatus::Disabled, + } + } + + /// Handle which will issue a [`ConnectionStatus::ClosedGracefully`] signal when dropped. + pub fn disarm(&mut self) { + self.on_drop = ConnectionStatus::ClosedGracefully; + } + + /// Handle which will issue a [`ConnectionStatus::ClosedUnexpectedly`] signal when dropped. + pub fn arm(&mut self) { + self.on_drop = ConnectionStatus::ClosedUnexpectedly; + } +} + +impl Drop for ConnectionHandle { + fn drop(&mut self) { + if let Some(sender) = self.sender.take() { + let _ = sender.send(self.on_drop); + } + } +} + +/// Creates a pair of handles which will monitor for disconnects from the client. +/// +/// The first handle is armed and will issue a [`ConnectionStatus::ClosedUnexpectedly`] signal when dropped. +/// The second handle is disarmed and will issue a [`ConnectionStatus::ClosedGracefully`] signal when dropped. +/// +/// The handles are returned in the order of the first being armed and the second being disarmed. +pub async fn create_connection_monitor( + engine_context: Arc, + metrics: Option>, + cancellation_labels: CancellationLabels, +) -> (ConnectionHandle, ConnectionHandle) { + // these oneshot channels monitor possible disconnects from the client in two different scopes: + // - the local task (connection_handle) + // - an optionally streaming response (stream_handle) + let (connection_tx, connection_rx) = tokio::sync::oneshot::channel(); + let (stream_tx, stream_rx) = tokio::sync::oneshot::channel(); + + // detached task that will naturally close when both handles are dropped + tokio::spawn(connection_monitor( + engine_context.clone(), + connection_rx, + stream_rx, + metrics, + cancellation_labels, + )); + + // Two handles, the first is armed, the second is disarmed + ( + ConnectionHandle::create_armed(connection_tx), + ConnectionHandle::create_disabled(stream_tx), + ) +} + +/// Monitor handler and response-body completion and cancel on unexpected drops. +#[tracing::instrument(level = "trace", skip_all, fields(request_id = %engine_context.id()))] +pub async fn connection_monitor( + engine_context: Arc, + connection_rx: tokio::sync::oneshot::Receiver, + stream_rx: tokio::sync::oneshot::Receiver, + metrics: Option>, + cancellation_labels: CancellationLabels, +) { + match connection_rx.await { + Err(_) | Ok(ConnectionStatus::ClosedUnexpectedly) => { + // the client has disconnected, no need to gracefully cancel, just kill the context + tracing::warn!("Connection closed unexpectedly; issuing cancellation"); + if let Some(metrics) = &metrics { + metrics.inc_client_disconnect(); + metrics.inc_cancellation(&cancellation_labels); + } + engine_context.kill(); + } + Ok(ConnectionStatus::ClosedGracefully) => { + tracing::trace!("Connection closed gracefully"); + } + Ok(ConnectionStatus::Disabled) => {} + } + + match stream_rx.await { + Err(_) | Ok(ConnectionStatus::ClosedUnexpectedly) => { + tracing::warn!("Stream closed unexpectedly; issuing cancellation"); + if let Some(metrics) = &metrics { + metrics.inc_client_disconnect(); + metrics.inc_cancellation(&cancellation_labels); + } + engine_context.kill(); + } + Ok(ConnectionStatus::ClosedGracefully) => { + tracing::trace!("Stream closed gracefully"); + } + Ok(ConnectionStatus::Disabled) => {} + } +} diff --git a/lib/http-server/src/lib.rs b/lib/http-server/src/lib.rs new file mode 100644 index 000000000000..018750aa9871 --- /dev/null +++ b/lib/http-server/src/lib.rs @@ -0,0 +1,10 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +pub mod disconnect; +pub mod lifecycle; +pub mod load; +pub mod metrics; +pub mod request; +pub mod server; +pub mod write_timeout; diff --git a/lib/http-server/src/lifecycle.rs b/lib/http-server/src/lifecycle.rs new file mode 100644 index 000000000000..77e7a5bcc643 --- /dev/null +++ b/lib/http-server/src/lifecycle.rs @@ -0,0 +1,198 @@ +// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use axum::body::Body; +use axum::http::Response; +use futures::StreamExt; +use std::sync::Arc; +use std::sync::atomic::{AtomicU8, AtomicU64, Ordering}; +use std::time::Duration; +use tokio::sync::Notify; + +/// Middleware that echoes `x-request-id` from request to response headers. +pub async fn echo_request_id_header( + request: axum::extract::Request, + next: axum::middleware::Next, +) -> axum::response::Response { + let x_request_id = request.headers().get("x-request-id").cloned(); + let mut response = next.run(request).await; + if let Some(value) = x_request_id { + response.headers_mut().insert("x-request-id", value); + } + response +} + +/// Lifecycle stage for the HTTP frontend. +/// +/// The stage gates readiness and request admission separately from the runtime +/// cancellation token so the frontend can stop accepting new requests before +/// tearing down discovery and transport state. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum ServiceStage { + /// The frontend is ready to admit new inference requests. + Ready = 0, + /// The frontend is rejecting new requests while admitted responses drain. + Draining = 1, + /// The frontend is cancelling runtime state and shutting down. + Stopping = 2, +} + +impl ServiceStage { + fn as_u8(self) -> u8 { + self as u8 + } + + fn from_u8(value: u8) -> Self { + match value { + 0 => Self::Ready, + 1 => Self::Draining, + _ => Self::Stopping, + } + } +} + +impl std::fmt::Display for ServiceStage { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Ready => f.write_str("ready"), + Self::Draining => f.write_str("draining"), + Self::Stopping => f.write_str("stopping"), + } + } +} + +/// Shared HTTP frontend lifecycle and inflight request tracker. +/// +/// `ServiceObserver` is shared by HTTP handlers, health endpoints, and the +/// shutdown path. It lets shutdown first mark the frontend as draining, then +/// wait for admitted inference response bodies to complete before cancelling +/// runtime state. +#[derive(Debug)] +pub struct ServiceObserver { + stage: AtomicU8, + inflight_inference: AtomicU64, + inflight_zero: Notify, +} + +impl Default for ServiceObserver { + fn default() -> Self { + Self { + stage: AtomicU8::new(ServiceStage::Ready.as_u8()), + inflight_inference: AtomicU64::new(0), + inflight_zero: Notify::new(), + } + } +} + +impl ServiceObserver { + /// Return the current frontend lifecycle stage. + pub fn stage(&self) -> ServiceStage { + ServiceStage::from_u8(self.stage.load(Ordering::Acquire)) + } + + /// Return true when the frontend should admit new inference requests. + pub fn is_ready(&self) -> bool { + self.stage() == ServiceStage::Ready + } + + /// Mark the frontend as draining. + /// + /// Draining makes readiness fail and causes request admission checks to + /// reject new inference requests while existing response bodies continue. + pub fn start_draining(&self) { + tracing::info!( + previous_stage = ?self.stage(), + inflight_requests = self.inflight_count(), + "frontend service entering draining stage" + ); + self.stage + .store(ServiceStage::Draining.as_u8(), Ordering::Release); + } + + /// Mark the frontend as stopping. + /// + /// Stopping is entered after inflight requests drain or the graceful + /// shutdown timeout expires. + pub fn start_stopping(&self) { + tracing::info!( + previous_stage = ?self.stage(), + inflight_requests = self.inflight_count(), + "frontend service entering stopping stage" + ); + self.stage + .store(ServiceStage::Stopping.as_u8(), Ordering::Release); + } + + /// Track one admitted inference response body. + /// + /// The returned permit must live for the full HTTP response body lifetime, + /// including streaming responses. Dropping the permit decrements the + /// inflight count and wakes shutdown waiters when the count reaches zero. + pub fn acquire_inflight(self: &Arc) -> InflightPermit { + self.inflight_inference.fetch_add(1, Ordering::Relaxed); + InflightPermit { + observer: self.clone(), + } + } + + /// Return the number of admitted inference requests still in flight. + pub fn inflight_count(&self) -> u64 { + self.inflight_inference.load(Ordering::Acquire) + } + + /// Wait until all admitted inference requests drain or `timeout` expires. + /// + /// Returns `true` when inflight work drained before the timeout and `false` + /// when shutdown should proceed because the timeout expired. + pub async fn wait_inflight_zero_or_timeout(&self, timeout: Duration) -> bool { + tokio::time::timeout(timeout, async { + loop { + let notified = self.inflight_zero.notified(); + tokio::pin!(notified); + // Register before reading the count so a final permit drop + // cannot notify between the count check and the await. + notified.as_mut().enable(); + if self.inflight_count() == 0 { + break; + } + notified.as_mut().await; + } + }) + .await + .is_ok() + } +} + +/// RAII guard for one admitted inference response. +/// +/// This permit is held by a response-body wrapper so it is released only when +/// the client response body finishes or is dropped. +pub struct InflightPermit { + observer: Arc, +} + +impl Drop for InflightPermit { + fn drop(&mut self) { + if self + .observer + .inflight_inference + .fetch_sub(1, Ordering::AcqRel) + == 1 + && self.observer.stage() != ServiceStage::Ready + { + self.observer.inflight_zero.notify_waiters(); + } + } +} + +/// Hold an inflight permit until the complete HTTP body is sent or dropped. +pub fn track_inflight_response(response: Response, permit: InflightPermit) -> Response { + let (parts, body) = response.into_parts(); + // Keep the permit alive until the full response body, including streams, + // finishes or is dropped. + let stream = body.into_data_stream().map(move |result| { + let _permit = &permit; + result + }); + Response::from_parts(parts, Body::from_stream(stream)) +} diff --git a/lib/http-server/src/load.rs b/lib/http-server/src/load.rs new file mode 100644 index 000000000000..a4d9ceec23b3 --- /dev/null +++ b/lib/http-server/src/load.rs @@ -0,0 +1,110 @@ +// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Metric snapshot values, independent of model discovery and event publication. + +use serde::{Deserialize, Serialize}; + +/// One frontend's complete load report. Each frame replaces the previous frame +/// from the same frontend incarnation; the Relay never merges partial frames. +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct FrontendLoadFrame { + /// Stable across event-publisher reconnects. + pub frontend_instance_id: u64, + /// Identifies the lifetime of the counters; reconnecting does not reset it. + pub incarnation: u64, + /// Starts at 1 and strictly increases within one incarnation. + pub sequence: u64, + pub is_serving_ready: bool, + pub models: Vec, +} + +/// Load for one canonical model on one frontend. +#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] +pub struct FrontendModelLoad { + pub model: String, + pub is_ready: bool, + pub gauges: RequestGauges, + pub totals: RequestTotals, +} + +/// In-flight requests at publish time. +/// +/// Prompt token counts are known once a request has been tokenized; a request +/// still being tokenized contributes zero input tokens. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] +pub struct RequestGauges { + /// Admitted requests that have not produced their first output token. + pub requests_awaiting_first_token: u64, + pub requests_generating: u64, + /// Prompt tokens of `requests_awaiting_first_token`. + pub awaiting_first_token_input_tokens: u64, + /// Prompt tokens of every in-flight request. + pub inflight_input_tokens: u64, +} + +impl RequestGauges { + pub fn add(&mut self, other: &Self) { + self.requests_awaiting_first_token = self + .requests_awaiting_first_token + .saturating_add(other.requests_awaiting_first_token); + self.requests_generating = self + .requests_generating + .saturating_add(other.requests_generating); + self.awaiting_first_token_input_tokens = self + .awaiting_first_token_input_tokens + .saturating_add(other.awaiting_first_token_input_tokens); + self.inflight_input_tokens = self + .inflight_input_tokens + .saturating_add(other.inflight_input_tokens); + } +} + +/// Cumulative request counters since the frontend incarnation started. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] +pub struct RequestTotals { + pub requests_started: u64, + pub requests_completed: u64, + pub requests_failed: u64, + pub requests_cancelled: u64, + pub input_tokens: u64, + pub output_tokens: u64, +} + +impl RequestTotals { + pub fn add(&mut self, other: &Self) { + self.requests_started = self.requests_started.saturating_add(other.requests_started); + self.requests_completed = self + .requests_completed + .saturating_add(other.requests_completed); + self.requests_failed = self.requests_failed.saturating_add(other.requests_failed); + self.requests_cancelled = self + .requests_cancelled + .saturating_add(other.requests_cancelled); + self.input_tokens = self.input_tokens.saturating_add(other.input_tokens); + self.output_tokens = self.output_tokens.saturating_add(other.output_tokens); + } + + /// Growth since `previous`, from the same frontend incarnation. A total that + /// went backwards (an alias unregistered, say) grows by nothing; a restart + /// arrives as a new incarnation and starts from a fresh baseline instead. + pub fn growth_since(&self, previous: &Self) -> Self { + let grow = |current: u64, previous: u64| current.saturating_sub(previous); + Self { + requests_started: grow(self.requests_started, previous.requests_started), + requests_completed: grow(self.requests_completed, previous.requests_completed), + requests_failed: grow(self.requests_failed, previous.requests_failed), + requests_cancelled: grow(self.requests_cancelled, previous.requests_cancelled), + input_tokens: grow(self.input_tokens, previous.input_tokens), + output_tokens: grow(self.output_tokens, previous.output_tokens), + } + } +} + +/// Model identity and readiness supplied by the service composition. +#[derive(Clone, Debug)] +pub struct ModelLoadView { + pub name: String, + pub is_ready: bool, + pub aliases: Vec, +} diff --git a/lib/http-server/src/metrics.rs b/lib/http-server/src/metrics.rs new file mode 100644 index 000000000000..8e3b550a6c1c --- /dev/null +++ b/lib/http-server/src/metrics.rs @@ -0,0 +1,4223 @@ +// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use crate::load::{ + FrontendLoadFrame, FrontendModelLoad, ModelLoadView, RequestGauges, RequestTotals, +}; +use dynamo_runtime::metrics::prometheus_names::{ + clamp_u64_to_i64, frontend_service, name_prefix, sanitize_frontend_prometheus_prefix, +}; +use dynamo_runtime::{ + config::environment_names::llm::metrics as env_metrics, + engine::Data, + error::{DynamoError, ErrorType as DynamoErrorType}, + pipeline::Context, +}; +use prometheus::{ + GaugeVec, HistogramOpts, HistogramVec, IntCounterVec, IntGaugeVec, Opts, + core::{Collector, Desc}, + proto::{Metric, MetricFamily}, +}; +use std::{ + collections::HashMap, + sync::{ + Arc, LazyLock, OnceLock, + atomic::{AtomicU64, Ordering::Relaxed}, + }, + time::{Duration, Instant}, +}; + +fn new_failure_counter(metrics_prefix: Option<&str>) -> IntCounterVec { + let prefix = + sanitize_frontend_prometheus_prefix(metrics_prefix.unwrap_or(name_prefix::FRONTEND)); + IntCounterVec::new( + Opts::new( + format!("{}_{}", prefix, frontend_service::FAILURES_TOTAL), + "Total number of terminal semantic request failures", + ), + &["class", "reason"], + ) + .expect("the fixed failure metric name and labels are valid") +} + +/// Process-wide because protocol error renderers can run without a request-scoped `Metrics`. +static FAILURE_METRICS_PREFIX: OnceLock = OnceLock::new(); +pub static DYNAM_FAILURES_TOTAL: LazyLock = LazyLock::new(|| { + let prefix = FAILURE_METRICS_PREFIX.get_or_init(|| { + let raw = std::env::var(env_metrics::DYN_METRICS_PREFIX) + .unwrap_or_else(|_| name_prefix::FRONTEND.to_string()); + sanitize_frontend_prometheus_prefix(&raw) + }); + new_failure_counter(Some(prefix.as_str())) +}); + +fn record_failure_into(counter: &IntCounterVec, error: &DynamoError) { + counter + .with_label_values(&[error.class().as_str(), error.reason().as_str()]) + .inc(); +} + +pub fn record_failure(error: &DynamoError) { + record_failure_into(&DYNAM_FAILURES_TOTAL, error); +} + +/// Check whether an error chain indicates the request was rejected. +pub fn request_was_rejected(err: &(dyn std::error::Error + 'static)) -> bool { + // Both overload flavors are client-visible rejections (HTTP 529). They differ + // only in whether migration may retry elsewhere. + const REJECTION: &[DynamoErrorType] = &[ + DynamoErrorType::ResourceExhausted, + DynamoErrorType::WorkerOverloaded, + ]; + const NON_REJECTION: &[DynamoErrorType] = &[]; + dynamo_runtime::error::match_error_chain(err, REJECTION, NON_REJECTION) +} + +/// User-facing body for a request whose caller-supplied deadline elapsed +/// before dispatch; shared by every HTTP surface so the wording stays uniform. +pub const REQUEST_DEADLINE_EXCEEDED_MESSAGE: &str = "request deadline exceeded"; + +/// Identify a deadline that expired while waiting in the router's queue. +pub fn request_deadline_exceeded(err: &(dyn std::error::Error + 'static)) -> bool { + queue_deadline_error(err).is_some() +} + +pub fn queue_deadline_error<'a>( + err: &'a (dyn std::error::Error + 'static), +) -> Option<&'a DynamoError> { + let mut current = Some(err); + while let Some(error) = current { + if let Some(error) = error.downcast_ref::() + && error.reason().as_str() == "router.queue_deadline_exceeded" + { + return Some(error); + } + current = error.source(); + } + None +} + +/// Check whether an error chain indicates that no backend worker is available +/// to this request. Both flavors are HTTP 503; they differ only in whether +/// migration may retry elsewhere. +pub fn request_was_unavailable(err: &(dyn std::error::Error + 'static)) -> bool { + const UNAVAILABLE: &[DynamoErrorType] = &[ + DynamoErrorType::Unavailable, + DynamoErrorType::WorkerUnavailable, + ]; + const AVAILABLE: &[DynamoErrorType] = &[]; + dynamo_runtime::error::match_error_chain(err, UNAVAILABLE, AVAILABLE) +} + +/// Check whether an error chain indicates the request was cancelled. +pub fn request_was_cancelled(err: &(dyn std::error::Error + 'static)) -> bool { + const CANCELLATION: &[DynamoErrorType] = &[ + DynamoErrorType::Cancelled, + DynamoErrorType::Backend(dynamo_runtime::error::BackendError::Cancelled), + ]; + const NON_CANCELLATION: &[DynamoErrorType] = &[]; + dynamo_runtime::error::match_error_chain(err, CANCELLATION, NON_CANCELLATION) +} + +pub fn request_was_timed_out(err: &(dyn std::error::Error + 'static)) -> bool { + use dynamo_runtime::error::BackendError; + + const TIMEOUT: &[DynamoErrorType] = &[ + DynamoErrorType::ResponseTimeout, + DynamoErrorType::ConnectionTimeout, + DynamoErrorType::Backend(BackendError::ResponseTimeout), + DynamoErrorType::Backend(BackendError::ConnectionTimeout), + ]; + const NON_TIMEOUT: &[DynamoErrorType] = &[]; + dynamo_runtime::error::match_error_chain(err, TIMEOUT, NON_TIMEOUT) +} + +pub use prometheus::Registry; + +pub const WORKER_TYPE_DECODE: &str = "decode"; +pub const WORKER_TYPE_PREFILL: &str = "prefill"; +const UNSET_DP_RANK_LABEL: &str = "none"; +const ITL_LOCAL_FLUSH_TOKENS: u64 = 64; + +pub struct ServingLoadCollector { + awaiting_requests: IntGaugeVec, + input_tokens: IntGaugeVec, + input_tokens_total: IntCounterVec, +} + +fn initialize_model_load( + awaiting_requests: &IntGaugeVec, + input_tokens: &IntGaugeVec, + input_tokens_total: &IntCounterVec, + model: &str, +) { + awaiting_requests.with_label_values(&[model]); + input_tokens_total.with_label_values(&[model]); + for state in [STATE_AWAITING_FIRST_TOKEN, STATE_GENERATING] { + input_tokens.with_label_values(&[model, state]); + } +} + +impl Collector for ServingLoadCollector { + fn desc(&self) -> Vec<&Desc> { + let mut descriptions = self.awaiting_requests.desc(); + descriptions.extend(self.input_tokens.desc()); + descriptions.extend(self.input_tokens_total.desc()); + descriptions + } + + fn collect(&self) -> Vec { + let mut families = self.awaiting_requests.collect(); + families.extend(self.input_tokens.collect()); + families.extend(self.input_tokens_total.collect()); + families + } +} + +/// Register the scrape-time model readiness collector with the frontend registry. +/// Global Prometheus gauge for last observed TTFT per worker (in seconds) +/// Labels: worker_id, dp_rank, worker_type +pub static WORKER_LAST_TIME_TO_FIRST_TOKEN_GAUGE: LazyLock = LazyLock::new(|| { + GaugeVec::new( + Opts::new( + format!( + "{}_{}", + name_prefix::FRONTEND, + frontend_service::WORKER_LAST_TIME_TO_FIRST_TOKEN_SECONDS + ), + "Last observed time to first token per worker (seconds)", + ), + &["worker_id", "dp_rank", "worker_type"], + ) + .expect("Failed to create worker_last_time_to_first_token gauge") +}); + +/// Global Prometheus gauge for last observed input sequence tokens per worker +/// Labels: worker_id, dp_rank, worker_type +/// Updated atomically with TTFT - represents the input token count from the same request +pub static WORKER_LAST_INPUT_SEQUENCE_TOKENS_GAUGE: LazyLock = LazyLock::new(|| { + IntGaugeVec::new( + Opts::new( + format!( + "{}_{}", + name_prefix::FRONTEND, + frontend_service::WORKER_LAST_INPUT_SEQUENCE_TOKENS + ), + "Last observed input sequence tokens per worker", + ), + &["worker_id", "dp_rank", "worker_type"], + ) + .expect("Failed to create worker_last_input_sequence_tokens gauge") +}); + +/// Global Prometheus gauge for last observed ITL per worker (in seconds) +/// Labels: worker_id, dp_rank, worker_type +pub static WORKER_LAST_INTER_TOKEN_LATENCY_GAUGE: LazyLock = LazyLock::new(|| { + GaugeVec::new( + Opts::new( + format!( + "{}_{}", + name_prefix::FRONTEND, + frontend_service::WORKER_LAST_INTER_TOKEN_LATENCY_SECONDS + ), + "Last observed inter-token latency per worker (seconds)", + ), + &["worker_id", "dp_rank", "worker_type"], + ) + .expect("Failed to create worker_last_inter_token_latency gauge") +}); + +/// Generate log-spaced histogram buckets with values rounded to 2 significant figures. +/// +/// # Arguments +/// * `min` - Minimum value for the buckets (must be > 0 for log spacing) +/// * `max` - Maximum value for the buckets (must be > min) +/// * `count` - Number of buckets to generate +/// +/// # Returns +/// A vector of log-spaced values, always starting with 0.0 and ending with the rounded max value. +/// Duplicates created by rounding are removed, so the final count may be less than requested. +/// +/// # Note +/// With 2 significant figures, there are roughly 90 unique values per order of magnitude. +/// Requesting more buckets than can be uniquely represented will result in deduplication. +pub fn generate_log_buckets(min: f64, max: f64, count: usize) -> Vec { + if count == 0 { + return vec![]; + } + if count == 1 { + return vec![0.0]; + } + + let requested_count = count; + let mut buckets = Vec::with_capacity(count); + buckets.push(0.0); + + // Generate log-spaced values from min to max + for i in 1..count { + let log_min = min.ln(); + let log_max = max.ln(); + let log_value = log_min + (log_max - log_min) * (i as f64) / ((count - 1) as f64); + let value = log_value.exp(); + buckets.push(round_to_sig_figs(value, 2)); + } + + // Remove consecutive duplicates (buckets are already sorted) + let original_len = buckets.len(); + buckets.dedup(); + + // Warn if significant deduplication occurred + if buckets.len() < original_len && (original_len - buckets.len()) > original_len / 10 { + tracing::warn!( + requested = requested_count, + unique = buckets.len(), + duplicates = original_len - buckets.len(), + min = min, + max = max, + "Histogram bucket generation: Significant duplicate values after rounding to 2 sig figs. \ + Consider reducing bucket count or increasing range." + ); + } + + buckets +} + +/// Round a number to a specified number of significant figures +pub fn round_to_sig_figs(value: f64, sig_figs: u32) -> f64 { + if value == 0.0 { + return 0.0; + } + + let magnitude = value.abs().log10().floor(); + let scale = 10_f64.powf(sig_figs as f64 - 1.0 - magnitude); + (value * scale).round() / scale +} + +const MAX_BUCKET_COUNT: usize = 512; +const MAX_BUCKET_LIST_BYTES: usize = MAX_BUCKET_COUNT * 32; + +fn validate_bucket_config(min: f64, max: f64, count: usize) -> bool { + min.is_finite() + && max.is_finite() + && min > 0.0 + && min < max + && count > 0 + && count <= MAX_BUCKET_COUNT +} + +/// How bucket settings are looked up. Injected rather than calling `std::env::var` +/// directly so tests can supply values without mutating the process environment: +/// concurrent `setenv`/`getenv` is undefined behavior on Unix, and many other tests in +/// this module construct `Metrics` and read these same variables. +type EnvLookup<'a> = &'a dyn Fn(&str) -> Option; + +/// The real lookup, used everywhere outside tests. +fn system_env(name: &str) -> Option { + std::env::var(name).ok() +} + +/// Read one histogram bucket setting, e.g. `DYN_METRICS_ITL_MAX`, from the environment. +/// +/// Falls back to the deprecated doubled form (`DYN_HISTOGRAM_DYN_METRICS_ITL_MAX`), +/// warning so operators can migrate off it. +/// +/// Returns `default` when neither name is set. A name that *is* set but cannot be +/// parsed also returns `default`, but warns first: silently ignoring a typo makes +/// the whole knob look like it does not exist. +fn bucket_env_var( + env: EnvLookup<'_>, + prefix: &str, + suffix: &str, + default: T, +) -> T { + let name = format!("{prefix}_{suffix}"); + let found = match env(&name) { + Some(value) => Some((name, value)), + None => { + let deprecated = format!( + "{}{prefix}_{suffix}", + env_metrics::DEPRECATED_HISTOGRAM_PREFIX + ); + env(&deprecated).map(|value| { + tracing::warn!( + deprecated = %deprecated, + replacement = %name, + "Deprecated histogram bucket environment variable; rename it, \ + support for the old name will be removed in a future release" + ); + (deprecated, value) + }) + } + }; + + match found { + None => default, + Some((name, value)) => value.parse::().unwrap_or_else(|_| { + tracing::warn!( + env_var = %name, + value = %value, + "Could not parse histogram bucket environment variable, using default" + ); + default + }), + } +} + +/// Parse histogram bucket configuration from environment variables +/// Returns (min, max, count) with defaults if not specified +fn parse_bucket_config( + env: EnvLookup<'_>, + env_prefix: &str, + default_min: f64, + default_max: f64, + default_count: usize, +) -> (f64, f64, usize) { + if !validate_bucket_config(default_min, default_max, default_count) { + tracing::error!( + default_min, + default_max, + default_count, + "Invalid default histogram configuration" + ); + return (1.0, 10.0, 10); + } + let mut min = bucket_env_var(env, env_prefix, "MIN", default_min); + let mut max = bucket_env_var(env, env_prefix, "MAX", default_max); + let mut count = bucket_env_var(env, env_prefix, "COUNT", default_count); + + if !validate_bucket_config(min, max, count) { + tracing::warn!( + min=%min, + max=%max, + count=%count, + "Invalid histogram configuration given, using defaults" + ); + min = default_min; + max = default_max; + count = default_count; + } + + (min, max, count) +} + +fn parse_bucket_list(value: &str) -> Option> { + if value.len() > MAX_BUCKET_LIST_BYTES { + return None; + } + + let buckets = value + .split(',') + .map(|item| item.trim().parse::()) + .collect::, _>>() + .ok()?; + + if buckets.len() > MAX_BUCKET_COUNT + || buckets + .iter() + .any(|bucket| !bucket.is_finite() || *bucket < 0.0) + || buckets.windows(2).any(|pair| pair[0] >= pair[1]) + { + return None; + } + + Some(buckets) +} + +fn inter_token_latency_buckets(env: EnvLookup<'_>) -> Vec { + if let Some(value) = env(env_metrics::DYN_METRICS_ITL_BUCKETS) { + return parse_bucket_list(&value).unwrap_or_else(|| { + tracing::warn!( + env_var = env_metrics::DYN_METRICS_ITL_BUCKETS, + value_len = value.len(), + max_value_bytes = MAX_BUCKET_LIST_BYTES, + "Invalid explicit ITL histogram buckets, using defaults" + ); + generate_log_buckets(0.001, 80.0, 20) + }); + } + + let (min, max, count) = parse_bucket_config(env, env_metrics::DYN_METRICS_ITL, 0.001, 80.0, 20); + generate_log_buckets(min, max, count) +} + +/// State for metrics handler. +/// Optionally holds a reference to the DRT's [`MetricsRegistry`] so that +/// metrics created via `metrics().create*()` anywhere in the process are +/// automatically included in the `/metrics` response. +pub struct Metrics { + load_incarnation: u64, + load_sequence: AtomicU64, + awaiting_first_token_requests: IntGaugeVec, + inflight_input_tokens: IntGaugeVec, + input_tokens_counter: IntCounterVec, + request_started_counter: IntCounterVec, + request_counter: IntCounterVec, + /// Deprecated: use `active_requests_gauge`. Kept for backwards compatibility. + inflight_gauge: IntGaugeVec, + active_requests_gauge: IntGaugeVec, + client_disconnect_gauge: prometheus::IntGauge, + http_queue_gauge: IntGaugeVec, + request_duration: HistogramVec, + input_sequence_length: HistogramVec, + output_sequence_length: HistogramVec, + cached_tokens: HistogramVec, + tokenizer_latency: HistogramVec, + output_tokens_counter: IntCounterVec, + time_to_first_token: HistogramVec, + inter_token_latency: HistogramVec, + /// End-to-end latency of an OpenAI `/v1/embeddings` request. Distinct from + /// `request_duration` so the buckets can be sized for sub-second pooling + /// inference rather than the 1-512s LLM-generation range. Labeled by + /// `model`. + embedding_latency: HistogramVec, + + // Per-request multimodal content-part count histograms (labeled by `model`). + images_per_request: HistogramVec, + videos_per_request: HistogramVec, + audio_per_request: HistogramVec, + image_tokens_per_request: HistogramVec, + + // Runtime configuration metrics. Note: Some of these metrics represent counter-like values from + // source systems, but are implemented as gauges because they are copied/synchronized from upstream + // counter values rather than being directly incremented. + model_total_kv_blocks: IntGaugeVec, + model_max_num_seqs: IntGaugeVec, + model_max_num_batched_tokens: IntGaugeVec, + model_context_length: IntGaugeVec, + model_kv_cache_block_size: IntGaugeVec, + model_migration_limit: IntGaugeVec, + model_migration_total: IntCounterVec, + model_migration_duration_seconds: HistogramVec, + model_migration_max_seq_len_exceeded_total: IntCounterVec, + model_cancellation_total: IntCounterVec, + model_rejection_total: IntCounterVec, +} + +// Inflight tracks requests from HTTP handler start until complete response is finished. +// HTTP queue tracks requests from HTTP handler start until first token generation begins (including prefill time). +// HTTP queue time is a subset of inflight time. For detailed explanation, see: +// deploy/metrics/README.md - "Request Processing Flow" section + +/// RAII object for HTTP queue gauge +/// Tracks requests from HTTP handler start until metrics processing begins +pub struct HttpQueueGuard { + metrics: Arc, + model: String, +} + +/// RAII object for inflight gauge and request counters +/// If this object is dropped without calling `mark_ok`, then the request will increment +/// the request counter with the `status` label with [`frontend_service::status::ERROR`]; otherwise, it will increment +/// the counter with `status` label [`frontend_service::status::SUCCESS`] +pub struct InflightGuard { + metrics: Arc, + request: Arc, + model: String, + endpoint: Endpoint, + request_type: RequestType, + status: Status, + error_type: ErrorType, + timer: Instant, + request_id: String, + span: tracing::Span, +} + +// Shared by the guard, tokenization, and response collector. The metric handles +// point to the same counters exposed by Prometheus; only request progress is local. +pub struct RequestMetrics { + // Prompt tokens and phase flags transition together; gauge deltas may arrive + // out of order across pipeline tasks, so gauges are signed and samples clamp. + state: AtomicU64, + awaiting_first_token_requests: prometheus::IntGauge, + awaiting_first_token_input_tokens: prometheus::IntGauge, + generating_input_tokens: prometheus::IntGauge, + input_tokens: prometheus::IntCounter, + output_tokens: prometheus::IntCounter, +} + +const STATE_AWAITING_FIRST_TOKEN: &str = "awaiting_first_token"; +const STATE_GENERATING: &str = "generating"; + +const REQUEST_METRICS_KEY: &str = "frontend.request_metrics"; +const FINISHED: u64 = 1 << 63; +const GENERATING: u64 = 1 << 62; +const INPUT_TOKENS: u64 = GENERATING - 1; + +impl RequestMetrics { + pub fn new(metrics: &Metrics, model: &str) -> Self { + let labels = &[model]; + let awaiting_first_token_requests = metrics + .awaiting_first_token_requests + .with_label_values(labels); + awaiting_first_token_requests.inc(); + Self { + state: AtomicU64::new(0), + awaiting_first_token_requests, + awaiting_first_token_input_tokens: metrics + .inflight_input_tokens + .with_label_values(&[model, STATE_AWAITING_FIRST_TOKEN]), + generating_input_tokens: metrics + .inflight_input_tokens + .with_label_values(&[model, STATE_GENERATING]), + input_tokens: metrics.input_tokens_counter.with_label_values(labels), + output_tokens: metrics.output_tokens_counter.with_label_values(labels), + } + } + + pub fn attach_to(self: &Arc, context: &mut Context) { + context.insert_arc(REQUEST_METRICS_KEY, self.clone()); + } + + pub fn from_context(context: &Context) -> Option> { + context + .get_optional(REQUEST_METRICS_KEY) + .expect("request metrics have a fixed registry type") + } + + /// Count a prompt of `tokens` tokens as soon as it is tokenized. A batched + /// completion request reports each of its prompts. + pub fn add_input_tokens(&self, tokens: usize) { + self.update_input_tokens(|known| known.saturating_add(tokens as u64)); + } + + /// Report the request's whole prompt size from a response annotation, which + /// repeats it on every chunk. Never lowers the known size, so it neither + /// double counts nor undoes [`add_input_tokens`](Self::add_input_tokens). + pub fn observe_input_tokens(&self, tokens: usize) { + self.update_input_tokens(|known| known.max(tokens as u64)); + } + + fn update_input_tokens(&self, input_tokens: impl Fn(u64) -> u64) { + let mut added = 0; + let Ok(previous) = self.state.fetch_update(Relaxed, Relaxed, |state| { + if state & FINISHED != 0 { + return None; + } + let known = state & INPUT_TOKENS; + added = input_tokens(known).min(INPUT_TOKENS).saturating_sub(known); + (added > 0).then_some(state + added) + }) else { + return; + }; + self.input_tokens.inc_by(added); + let added = added as i64; + if previous & GENERATING == 0 { + self.awaiting_first_token_input_tokens.add(added); + } else { + self.generating_input_tokens.add(added); + } + } + + /// Count `tokens` newly generated output tokens (a delta, not a total). The + /// first output token moves the request from awaiting to generating. + pub fn add_output_tokens(&self, tokens: usize) { + let state = self.state.load(Relaxed); + if tokens == 0 || state & FINISHED != 0 { + return; + } + self.output_tokens.inc_by(tokens as u64); + if state & GENERATING != 0 { + return; + } + let previous = self.state.fetch_or(GENERATING, Relaxed); + if previous & (GENERATING | FINISHED) != 0 { + return; + } + let input = (previous & INPUT_TOKENS) as i64; + self.awaiting_first_token_requests.dec(); + self.awaiting_first_token_input_tokens.sub(input); + self.generating_input_tokens.add(input); + } + + /// Stop tracking the request. Only the first call on any clone counts. + fn finish(&self) { + let previous = self.state.fetch_or(FINISHED, Relaxed); + if previous & FINISHED != 0 { + return; + } + let input = (previous & INPUT_TOKENS) as i64; + if previous & GENERATING != 0 { + self.generating_input_tokens.sub(input); + } else { + self.awaiting_first_token_requests.dec(); + self.awaiting_first_token_input_tokens.sub(input); + } + } +} + +fn metric_label<'a>(metric: &'a Metric, name: &str) -> &'a str { + metric + .get_label() + .iter() + .find(|label| label.name() == name) + .expect("frontend metrics have fixed labels") + .value() +} + +// Read the integer counter handles, not the exported f64 samples, so load totals +// remain exact even after a counter exceeds 2^53. +// Only re-resolve labels returned by collect(), so no new series are created. +fn model_counter_values(family: &IntCounterVec) -> impl Iterator + '_ { + family + .collect() + .into_iter() + .flat_map(|family| family.metric) + .map(|metric| { + let model = metric_label(&metric, "model"); + (model.to_owned(), family.with_label_values(&[model]).get()) + }) +} + +fn model_gauge_values(family: &IntGaugeVec) -> impl Iterator { + family + .collect() + .into_iter() + .flat_map(|family| family.metric) + .map(|metric| { + ( + metric_label(&metric, "model").to_owned(), + metric.get_gauge().value().max(0.0) as u64, + ) + }) +} + +/// Requests will be logged by the type of endpoint hit +/// This will include llamastack in the future +#[derive(Clone, Copy)] +pub enum Endpoint { + /// OAI Completions + Completions, + + /// OAI Chat Completions + ChatCompletions, + + /// OAI Embeddings + Embeddings, + + /// Classification (sequence classification / cross-encoder pooling) + Classify, + + /// Cross-encoder reranking + Rerank, + + /// Pooling (raw pooler output) + Pooling, + + /// OAI Images + Images, + + /// OAI Videos + Videos, + + /// OAI Audio Speech + Audios, + + /// OAI Responses + Responses, + + /// Anthropic Messages + AnthropicMessages, + + /// Tensor + Tensor, + + /// Generate (token-in/token-out) + Generate, +} + +/// Metrics for the HTTP service +pub enum RequestType { + /// SingleIn / SingleOut + Unary, + + /// SingleIn / ManyOut + Stream, +} + +/// Labels for cancellation metrics +pub struct CancellationLabels { + pub model: String, + pub endpoint: String, + pub request_type: String, +} + +/// Status +#[derive(PartialEq)] +pub enum Status { + Success, + Error, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum RequestCompletion { + Success, + Cancelled, + Error, +} + +impl RequestCompletion { + fn from_labels(status: &str, error: &str) -> Self { + if status == Status::Success.as_str() { + Self::Success + } else if error == ErrorType::Cancelled.as_str() { + Self::Cancelled + } else { + Self::Error + } + } + + fn as_str(self) -> &'static str { + match self { + Self::Success => "success", + Self::Cancelled => "cancelled", + Self::Error => "error", + } + } +} + +/// Error type classification for fine-grained observability +#[derive(PartialEq, Clone, Debug)] +pub enum ErrorType { + /// No error (for successful requests) + None, + /// Client validation error (4xx with "Validation:" prefix) + Validation, + /// Model or resource not found (404) + NotFound, + /// Service overloaded or rate limited (429 or 529) + Overload, + /// Service unavailable because no backend worker can serve the request + Unavailable, + /// Request cancelled by client or timeout + Cancelled, + /// Backend accepted the request but stopped responding (response inactivity timeout) + ResponseTimeout, + /// Internal server error (500 and other unexpected errors) + Internal, + /// Feature not implemented (501) + NotImplemented, +} + +/// Track response-specific metrics +pub struct ResponseMetricCollector { + metrics: Arc, + request: Arc, + model: String, + // Per-model metric handles cached for the request. Most are resolved at construction; + // ITL is resolved lazily on its first observation so requests that never produce ITL + // do not allocate a local histogram. Caching avoids re-hashing the `model` label on + // every chunk or output token. Each handle shares the underlying metric with its vec, + // so observations are equivalent. + time_to_first_token: prometheus::Histogram, + inter_token_latency: Option, + itl_pending_tokens: u64, + input_sequence_length: prometheus::Histogram, + cached_tokens: prometheus::Histogram, + // Per-model multimodal count histogram handles, resolved once at construction. + images_per_request: prometheus::Histogram, + videos_per_request: prometheus::Histogram, + audio_per_request: prometheus::Histogram, + image_tokens_per_request: prometheus::Histogram, + // Latched per-request counts (for the tracing span fields recorded in `Drop`). + image_count_val: usize, + video_count_val: usize, + audio_count_val: usize, + multimodal_counts_observed: bool, + start_time: Instant, + // we use is_first_token to distinguish TTFT from ITL. It is true by default and + // flipped to false when the first token is returned and TTFT is published. + is_first_token: bool, + // we track the last response time so that ITL for the newly returned tokens can + // be computed. + last_response_time: Option, + osl: usize, + isl: usize, + ttft_ms: Option, + itl_sum_secs: f64, + itl_count: u64, + // we track if cached_tokens has been observed to ensure we only increment once per request + cached_tokens_observed: bool, + // we track if tokenize latency has been observed to ensure we only increment once per request + tokenize_latency_observed: bool, + // latest accumulated detokenize latency and sample count reported by tracker + detokenize_latency_total: Duration, + detokenize_count_total: u64, + // Prefill worker info for TTFT attribution (set from LLMMetricAnnotation) + prefill_worker_id: Option, + prefill_dp_rank: Option, + // Prefill worker type for Prometheus labeling - stored at routing time to avoid MDC lookup + prefill_worker_type: Option, + // Decode worker info for ITL attribution (set from LLMMetricAnnotation) + decode_worker_id: Option, + decode_dp_rank: Option, + // Decode worker type for Prometheus labeling - stored at routing time to avoid MDC lookup + decode_worker_type: Option, + // Cached per-worker ITL gauge handle. The decode-worker labels are latched once at + // routing time (`set_worker_info` only sets when unset), so this GaugeVec handle is + // resolved a single time and reused — the per-token observe path then does no label + // formatting (`worker_id.to_string()`) or label hashing, only `set`. + decode_itl_gauge: Option, +} + +impl Default for Metrics { + fn default() -> Self { + Self::new() + } +} + +impl Metrics { + /// Create Metrics with the standard prefix defined by [`name_prefix::FRONTEND`] or specify custom prefix via the following environment variable: + /// - `DYN_METRICS_PREFIX`: Override the default metrics prefix + /// + /// The following metrics will be created with the configured prefix: + /// - `{prefix}_requests_total` - IntCounterVec for the total number of requests processed + /// - `{prefix}_inflight_requests` - IntGaugeVec for the number of inflight/concurrent requests + /// - `{prefix}_disconnected_clients` - IntGauge for the number of disconnected clients + /// - `{prefix}_request_duration_seconds` - HistogramVec for the duration of requests + /// - `{prefix}_input_sequence_tokens` - HistogramVec for input sequence length in tokens + /// - `{prefix}_output_sequence_tokens` - HistogramVec for output sequence length in tokens + /// - `{prefix}_tokenizer_latency_ms` - HistogramVec for tokenizer latency in milliseconds + /// - `{prefix}_output_tokens_total` - IntCounterVec for total output tokens generated (real-time updates) + /// - `{prefix}_time_to_first_token_seconds` - HistogramVec for time to first token in seconds + /// - `{prefix}_inter_token_latency_seconds` - HistogramVec for inter-token latency in seconds + /// + /// ## Histogram Bucket Configuration + /// + /// Histograms use log-spaced buckets rounded to 2 significant figures unless noted otherwise. + /// Bucket configuration can be customized via environment variables (MIN: minimum value, + /// MAX: maximum value, COUNT: number of buckets): + /// + /// - `DYN_METRICS_REQUEST_DURATION_{MIN,MAX,COUNT}` - Request duration histogram (defaults: 1.0, 512.0, 10) + /// - `DYN_METRICS_INPUT_SEQUENCE_{MIN,MAX,COUNT}` - Input sequence length histogram (defaults: 50.0, 128000.0, 12) + /// - `DYN_METRICS_OUTPUT_SEQUENCE_{MIN,MAX,COUNT}` - Output sequence length histogram (defaults: 50.0, 32000.0, 10) + /// - `DYN_METRICS_TTFT_{MIN,MAX,COUNT}` - Time to first token histogram (defaults: 0.001, 480.0, 18) + /// - `DYN_METRICS_ITL_BUCKETS` - Comma-separated inter-token latency bucket boundaries + /// - `DYN_METRICS_ITL_{MIN,MAX,COUNT}` - Log-spaced inter-token latency configuration (defaults: 0.001, 80.0, 20) + /// - `DYN_METRICS_EMBEDDING_LATENCY_{MIN,MAX,COUNT}` - End-to-end `/v1/embeddings` latency histogram (defaults: 0.001, 10.0, 14) + /// + /// ## Model Configuration Metrics + /// + /// Runtime config metrics (from ModelRuntimeConfig): + /// - `{prefix}_model_total_kv_blocks` - IntGaugeVec for total KV cache blocks available for a worker serving the model + /// - `{prefix}_model_max_num_seqs` - IntGaugeVec for maximum sequences for a worker serving the model + /// - `{prefix}_model_max_num_batched_tokens` - IntGaugeVec for maximum batched tokens for a worker serving the model + /// + /// MDC metrics (from ModelDeploymentCard): + /// - `{prefix}_model_context_length` - IntGaugeVec for maximum context length for a worker serving the model + /// - `{prefix}_model_kv_cache_block_size` - IntGaugeVec for KV cache block size for a worker serving the model + /// - `{prefix}_model_migration_limit` - IntGaugeVec for request migration limit for a worker serving the model + /// + /// ## Runtime Config Polling Configuration + /// + /// The polling behavior can be configured via environment variables: + /// - `DYN_HTTP_SVC_CONFIG_METRICS_POLL_INTERVAL_SECS`: Poll interval in seconds (must be > 0, supports fractional seconds, defaults to 8) + /// + /// Metrics are never removed to preserve historical data. Runtime config and MDC + /// metrics are updated when models are discovered and their configurations are available. + pub fn new() -> Self { + Self::new_with_prefix(std::env::var(env_metrics::DYN_METRICS_PREFIX).ok()) + } + + /// Create Metrics with an explicit optional prefix. `None` uses the standard + /// frontend prefix and does not read environment variables. + pub fn new_with_prefix(metrics_prefix: Option) -> Self { + Self::build(metrics_prefix, &system_env) + } + + /// Real constructor. `env` is injected so tests can pin bucket configuration + /// without touching the process environment. + fn build(metrics_prefix: Option, env: EnvLookup<'_>) -> Self { + // TODO: Remove DYN_METRICS_PREFIX env-var override (added in PR #2432 for + // NIM compatibility with the old "nv_llm_http_service_" prefix). No longer + // needed — hardcode name_prefix::FRONTEND and drop the sanitize function. + let raw_prefix = metrics_prefix.unwrap_or_else(|| name_prefix::FRONTEND.to_string()); + let prefix = sanitize_frontend_prometheus_prefix(&raw_prefix); + if let Err(configured_prefix) = FAILURE_METRICS_PREFIX.set(prefix.clone()) + && configured_prefix != prefix + { + tracing::warn!( + configured=%configured_prefix, + requested=%prefix, + "Semantic failure metrics already use a different process-wide prefix" + ); + } + if prefix != raw_prefix { + tracing::warn!( + raw=%raw_prefix, + sanitized=%prefix, + env=%frontend_service::METRICS_PREFIX_ENV, + "Sanitized HTTP metrics prefix" + ); + } + let frontend_metric_name = |suffix: &str| format!("{}_{}", &prefix, suffix); + + let awaiting_first_token_requests = IntGaugeVec::new( + Opts::new( + frontend_metric_name(frontend_service::AWAITING_FIRST_TOKEN_REQUESTS), + "Admitted requests awaiting their first output token; endpoints that do not emit tokens remain here until completion", + ), + &["model"], + ) + .unwrap(); + + let inflight_input_tokens = IntGaugeVec::new( + Opts::new( + frontend_metric_name(frontend_service::INFLIGHT_INPUT_TOKENS), + "Known prompt tokens in inflight requests by state; non-token endpoints remain awaiting_first_token until completion", + ), + &["model", "state"], + ) + .unwrap(); + + let input_tokens_counter = IntCounterVec::new( + Opts::new( + frontend_metric_name(frontend_service::INPUT_TOKENS_TOTAL), + "Total known input tokens of admitted requests", + ), + &["model"], + ) + .unwrap(); + + let request_counter = IntCounterVec::new( + Opts::new( + frontend_metric_name(frontend_service::REQUESTS_TOTAL), + "Total number of LLM requests processed", + ), + &["model", "endpoint", "request_type", "status", "error_type"], + ) + .unwrap(); + + let request_started_counter = IntCounterVec::new( + Opts::new( + frontend_metric_name(frontend_service::REQUESTS_STARTED_TOTAL), + "Total number of LLM requests accepted by the frontend handler", + ), + &["model", "endpoint", "request_type"], + ) + .unwrap(); + + let inflight_gauge = IntGaugeVec::new( + Opts::new( + frontend_metric_name(frontend_service::INFLIGHT_REQUESTS), + "Number of inflight requests", + ), + &["model"], + ) + .unwrap(); + + let active_requests_gauge = IntGaugeVec::new( + Opts::new( + frontend_metric_name(frontend_service::ACTIVE_REQUESTS), + "Number of requests currently being handled by the frontend, from HTTP handler entry to response completion", + ), + &["model"], + ) + .unwrap(); + + let client_disconnect_gauge = prometheus::IntGauge::new( + frontend_metric_name(frontend_service::DISCONNECTED_CLIENTS), + "Number of disconnected clients", + ) + .unwrap(); + + let http_queue_gauge = IntGaugeVec::new( + Opts::new( + frontend_metric_name(frontend_service::QUEUED_REQUESTS), + "Number of requests in HTTP processing queue", + ), + &["model"], + ) + .unwrap(); + + // Request duration buckets: configurable via DYN_METRICS_REQUEST_DURATION_{MIN,MAX,COUNT} + let (req_dur_min, req_dur_max, req_dur_count) = parse_bucket_config( + env, + env_metrics::DYN_METRICS_REQUEST_DURATION, + 1.0, + 512.0, + 10, + ); + let request_duration_buckets = + generate_log_buckets(req_dur_min, req_dur_max, req_dur_count); + + let request_duration = HistogramVec::new( + HistogramOpts::new( + frontend_metric_name(frontend_service::REQUEST_DURATION_SECONDS), + "Duration of LLM requests", + ) + .buckets(request_duration_buckets), + &["model"], + ) + .unwrap(); + + // Input sequence length buckets: configurable via DYN_METRICS_INPUT_SEQUENCE_{MIN,MAX,COUNT} + let (isl_min, isl_max, isl_count) = parse_bucket_config( + env, + env_metrics::DYN_METRICS_INPUT_SEQUENCE, + 50.0, + 128000.0, + 12, + ); + let input_sequence_buckets = generate_log_buckets(isl_min, isl_max, isl_count); + + let input_sequence_length = HistogramVec::new( + HistogramOpts::new( + frontend_metric_name(frontend_service::INPUT_SEQUENCE_TOKENS), + "Input sequence length in tokens", + ) + .buckets(input_sequence_buckets.clone()), + &["model"], + ) + .unwrap(); + + // Output sequence length buckets: configurable via DYN_METRICS_OUTPUT_SEQUENCE_{MIN,MAX,COUNT} + let (osl_min, osl_max, osl_count) = parse_bucket_config( + env, + env_metrics::DYN_METRICS_OUTPUT_SEQUENCE, + 50.0, + 32000.0, + 10, + ); + let output_sequence_buckets = generate_log_buckets(osl_min, osl_max, osl_count); + + let output_sequence_length = HistogramVec::new( + HistogramOpts::new( + frontend_metric_name(frontend_service::OUTPUT_SEQUENCE_TOKENS), + "Output sequence length in tokens", + ) + .buckets(output_sequence_buckets), + &["model"], + ) + .unwrap(); + + let output_tokens_counter = IntCounterVec::new( + Opts::new( + frontend_metric_name(frontend_service::OUTPUT_TOKENS_TOTAL), + "Total number of output tokens generated (updates in real-time)", + ), + &["model"], + ) + .unwrap(); + + // Time to first token buckets: configurable via DYN_METRICS_TTFT_{MIN,MAX,COUNT} + let (ttft_min, ttft_max, ttft_count) = + parse_bucket_config(env, env_metrics::DYN_METRICS_TTFT, 0.001, 480.0, 18); + let time_to_first_token_buckets = generate_log_buckets(ttft_min, ttft_max, ttft_count); + + let time_to_first_token = HistogramVec::new( + HistogramOpts::new( + frontend_metric_name(frontend_service::TIME_TO_FIRST_TOKEN_SECONDS), + "Time to first token in seconds", + ) + .buckets(time_to_first_token_buckets), + &["model"], + ) + .unwrap(); + + let inter_token_latency_buckets = inter_token_latency_buckets(env); + + let inter_token_latency = HistogramVec::new( + HistogramOpts::new( + frontend_metric_name(frontend_service::INTER_TOKEN_LATENCY_SECONDS), + "Inter-token latency in seconds", + ) + .buckets(inter_token_latency_buckets), + &["model"], + ) + .unwrap(); + + // Embedding latency buckets: pooling-model inference is typically + // sub-second (60-200 token inputs on L4-class GPUs land in the 5-50ms + // range), so 1ms..10s on a log scale gives p50/p99 resolution that + // the 1..512s `request_duration` buckets cannot. + let (emb_min, emb_max, emb_count) = parse_bucket_config( + env, + env_metrics::DYN_METRICS_EMBEDDING_LATENCY, + 0.001, + 10.0, + 14, + ); + let embedding_latency_buckets = generate_log_buckets(emb_min, emb_max, emb_count); + + let embedding_latency = HistogramVec::new( + HistogramOpts::new( + frontend_metric_name(frontend_service::EMBEDDING_LATENCY_SECONDS), + "End-to-end latency of /v1/embeddings requests, in seconds. \ + Distinct from request_duration_seconds so the bucket range can \ + cover sub-second pooling-model inference.", + ) + .buckets(embedding_latency_buckets), + &["model"], + ) + .unwrap(); + + // Per-request media-count buckets: every integer through 10 (where most + // requests land), then coarse steps for the large-batch tail. + let multimodal_count_buckets = vec![ + 0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 20.0, 30.0, 40.0, 50.0, 100.0, + ]; + + let images_per_request = HistogramVec::new( + HistogramOpts::new( + frontend_metric_name(frontend_service::IMAGES_PER_REQUEST), + "Number of image_url content parts per request", + ) + .buckets(multimodal_count_buckets.clone()), + &["model"], + ) + .unwrap(); + + let videos_per_request = HistogramVec::new( + HistogramOpts::new( + frontend_metric_name(frontend_service::VIDEOS_PER_REQUEST), + "Number of video_url content parts per request", + ) + .buckets(multimodal_count_buckets.clone()), + &["model"], + ) + .unwrap(); + + let audio_per_request = HistogramVec::new( + HistogramOpts::new( + frontend_metric_name(frontend_service::AUDIO_PER_REQUEST), + "Number of audio_url content parts per request", + ) + .buckets(multimodal_count_buckets), + &["model"], + ) + .unwrap(); + + // Image placeholder-token buckets: powers of two cover the common + // single-image range and multi-image request totals. + let image_token_buckets = vec![ + 4.0, 8.0, 16.0, 32.0, 64.0, 128.0, 256.0, 512.0, 1024.0, 2048.0, 4096.0, 8192.0, + 16384.0, 32768.0, 65536.0, + ]; + let image_tokens_per_request = HistogramVec::new( + HistogramOpts::new( + frontend_metric_name(frontend_service::IMAGE_TOKENS_PER_REQUEST), + "Calculated image-placeholder token count per image-bearing request; \ + recorded from the response path only when every image resolves and \ + processor overrides are absent", + ) + .buckets(image_token_buckets), + &["model"], + ) + .unwrap(); + + let cached_tokens = HistogramVec::new( + HistogramOpts::new( + frontend_metric_name(frontend_service::CACHED_TOKENS), + "Number of cached tokens (prefix cache hits) per request", + ) + .buckets(input_sequence_buckets.clone()), + &["model"], + ) + .unwrap(); + + let tokenizer_latency = HistogramVec::new( + HistogramOpts::new( + frontend_metric_name(frontend_service::TOKENIZER_LATENCY_MS), + "Tokenizer latency in milliseconds", + ) + .buckets(vec![ + 0.5, 1.0, 2.0, 4.0, 8.0, 16.0, 32.0, 64.0, 128.0, 256.0, 512.0, + ]), + &[frontend_service::OPERATION_LABEL], + ) + .unwrap(); + + // Runtime configuration metrics + // Note: Some of these metrics represent counter-like values from source systems, + // but are implemented as gauges because they are copied/synchronized from upstream + // counter values rather than being directly incremented. + let model_total_kv_blocks = IntGaugeVec::new( + Opts::new( + frontend_metric_name(frontend_service::MODEL_TOTAL_KV_BLOCKS), + "Total KV cache blocks available for a worker serving the model", + ), + &["model"], + ) + .unwrap(); + + let model_max_num_seqs = IntGaugeVec::new( + Opts::new( + frontend_metric_name(frontend_service::MODEL_MAX_NUM_SEQS), + "Maximum number of sequences for a worker serving the model", + ), + &["model"], + ) + .unwrap(); + + let model_max_num_batched_tokens = IntGaugeVec::new( + Opts::new( + frontend_metric_name(frontend_service::MODEL_MAX_NUM_BATCHED_TOKENS), + "Maximum number of batched tokens for a worker serving the model", + ), + &["model"], + ) + .unwrap(); + + let model_context_length = IntGaugeVec::new( + Opts::new( + frontend_metric_name(frontend_service::MODEL_CONTEXT_LENGTH), + "Maximum context length in tokens for a worker serving the model", + ), + &["model"], + ) + .unwrap(); + + let model_kv_cache_block_size = IntGaugeVec::new( + Opts::new( + frontend_metric_name(frontend_service::MODEL_KV_CACHE_BLOCK_SIZE), + "KV cache block size in tokens for a worker serving the model", + ), + &["model"], + ) + .unwrap(); + + let model_migration_limit = IntGaugeVec::new( + Opts::new( + frontend_metric_name(frontend_service::MODEL_MIGRATION_LIMIT), + "Maximum number of request migrations allowed for the model", + ), + &["model"], + ) + .unwrap(); + + let model_migration_total = IntCounterVec::new( + Opts::new( + frontend_metric_name(frontend_service::MODEL_MIGRATION_TOTAL), + "Total number of request migrations due to worker unavailability", + ), + &["model", frontend_service::MIGRATION_TYPE_LABEL], + ) + .unwrap(); + + let model_migration_duration_seconds = HistogramVec::new( + HistogramOpts::new( + frontend_metric_name(frontend_service::MODEL_MIGRATION_DURATION_SECONDS), + "Time from detecting a migratable failure until recovery, terminal failure, or cancellation", + ) + .buckets(vec![ + 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, 20.0, 40.0, 60.0, + ]), + &[ + "model", + frontend_service::MIGRATION_TYPE_LABEL, + frontend_service::MIGRATION_OUTCOME_LABEL, + ], + ) + .unwrap(); + + let model_migration_max_seq_len_exceeded_total = IntCounterVec::new( + Opts::new( + frontend_metric_name(frontend_service::MODEL_MIGRATION_MAX_SEQ_LEN_EXCEEDED_TOTAL), + "Total number of times migration was disabled by max_seq_len limit", + ), + &["model"], + ) + .unwrap(); + + let model_cancellation_total = IntCounterVec::new( + Opts::new( + frontend_metric_name(frontend_service::MODEL_CANCELLATION_TOTAL), + "Total number of request cancellations", + ), + &["model", "endpoint", "request_type"], + ) + .unwrap(); + + let model_rejection_total = IntCounterVec::new( + Opts::new( + frontend_metric_name(frontend_service::MODEL_REJECTION_TOTAL), + "Total number of requests rejected due to resource exhaustion", + ), + &["model", "endpoint"], + ) + .unwrap(); + + let metrics = Metrics { + load_incarnation: rand::random(), + load_sequence: AtomicU64::new(0), + awaiting_first_token_requests, + inflight_input_tokens, + input_tokens_counter, + request_started_counter, + request_counter, + inflight_gauge, + active_requests_gauge, + client_disconnect_gauge, + http_queue_gauge, + request_duration, + input_sequence_length, + output_sequence_length, + cached_tokens, + tokenizer_latency, + output_tokens_counter, + time_to_first_token, + inter_token_latency, + embedding_latency, + images_per_request, + videos_per_request, + audio_per_request, + image_tokens_per_request, + model_total_kv_blocks, + model_max_num_seqs, + model_max_num_batched_tokens, + model_context_length, + model_kv_cache_block_size, + model_migration_limit, + model_migration_total, + model_migration_duration_seconds, + model_migration_max_seq_len_exceeded_total, + model_cancellation_total, + model_rejection_total, + }; + initialize_model_load( + &metrics.awaiting_first_token_requests, + &metrics.inflight_input_tokens, + &metrics.input_tokens_counter, + "unknown_model", + ); + metrics + } + + /// Get the number of successful requests for the given dimensions: + /// - model + /// - endpoint (completions/chat_completions) + /// - request type (unary/stream) + /// - status (success/error) + pub fn get_request_counter( + &self, + model: &str, + endpoint: &Endpoint, + request_type: &RequestType, + status: &Status, + error_type: &ErrorType, + ) -> u64 { + self.request_counter + .with_label_values(&[ + model, + endpoint.as_str(), + request_type.as_str(), + status.as_str(), + error_type.as_str(), + ]) + .get() + } + + /// Increment the counter for requests for the given dimensions: + /// - model + /// - endpoint (completions/chat_completions) + /// - request type (unary/stream) + /// - status (success/error) + fn inc_request_counter( + &self, + model: &str, + endpoint: &Endpoint, + request_type: &RequestType, + status: &Status, + error_type: &ErrorType, + ) { + self.request_counter + .with_label_values(&[ + model, + endpoint.as_str(), + request_type.as_str(), + status.as_str(), + error_type.as_str(), + ]) + .inc() + } + + /// Increment the counter for requests accepted by the frontend handler. + fn inc_request_started_counter( + &self, + model: &str, + endpoint: &Endpoint, + request_type: &RequestType, + ) { + self.request_started_counter + .with_label_values(&[model, endpoint.as_str(), request_type.as_str()]) + .inc() + } + + /// Get the number if inflight requests for the given model + pub fn get_inflight_count(&self, model: &str) -> i64 { + self.inflight_gauge.with_label_values(&[model]).get() + } + + /// Record one observation of the end-to-end latency of an + /// `/v1/embeddings` request, in seconds. Labeled by model. + /// + /// Negative values are dropped (they can only happen from a clock skew or + /// programming bug; recording them would corrupt percentile estimates). + pub fn observe_embedding_latency(&self, model: &str, seconds: f64) { + if !seconds.is_finite() || seconds < 0.0 { + return; + } + self.embedding_latency + .with_label_values(&[model]) + .observe(seconds); + } + + /// Get the cumulative count of embedding latency observations for the given + /// model. Test helper. + #[cfg(test)] + pub fn embedding_latency_count(&self, model: &str) -> u64 { + self.embedding_latency + .with_label_values(&[model]) + .get_sample_count() + } + + /// Get the cumulative sum (seconds) of embedding latency observations for + /// the given model. Test helper. + #[cfg(test)] + pub fn embedding_latency_sum(&self, model: &str) -> f64 { + self.embedding_latency + .with_label_values(&[model]) + .get_sample_sum() + } + + fn inc_inflight_gauge(&self, model: &str) { + self.inflight_gauge.with_label_values(&[model]).inc(); + self.active_requests_gauge.with_label_values(&[model]).inc(); + } + + fn dec_inflight_gauge(&self, model: &str) { + self.inflight_gauge.with_label_values(&[model]).dec(); + self.active_requests_gauge.with_label_values(&[model]).dec(); + } + + /// Increment the gauge for client disconnections + pub fn inc_client_disconnect(&self) { + self.client_disconnect_gauge.inc(); + } + + /// Get the count of client disconnections + pub fn get_client_disconnect_count(&self) -> i64 { + self.client_disconnect_gauge.get() + } + + fn inc_http_queue_gauge(&self, model: &str) { + self.http_queue_gauge.with_label_values(&[model]).inc() + } + + fn dec_http_queue_gauge(&self, model: &str) { + self.http_queue_gauge.with_label_values(&[model]).dec() + } + + pub fn register(&self, registry: &Registry) -> Result<(), prometheus::Error> { + self.register_with_serving_load_collector(registry, Box::new(self.serving_load_collector())) + } + + /// Metric handles used by a composition's scrape-time model initializer. + pub fn serving_load_collector(&self) -> ServingLoadCollector { + ServingLoadCollector { + awaiting_requests: self.awaiting_first_token_requests.clone(), + input_tokens: self.inflight_input_tokens.clone(), + input_tokens_total: self.input_tokens_counter.clone(), + } + } + + pub fn register_with_serving_load_collector( + &self, + registry: &Registry, + collector: Box, + ) -> Result<(), prometheus::Error> { + registry.register(collector)?; + registry.register(Box::new(self.request_started_counter.clone()))?; + registry.register(Box::new(self.request_counter.clone()))?; + registry.register(Box::new(self.inflight_gauge.clone()))?; + registry.register(Box::new(self.active_requests_gauge.clone()))?; + registry.register(Box::new(self.client_disconnect_gauge.clone()))?; + registry.register(Box::new(self.http_queue_gauge.clone()))?; + registry.register(Box::new(self.request_duration.clone()))?; + registry.register(Box::new(self.input_sequence_length.clone()))?; + registry.register(Box::new(self.output_sequence_length.clone()))?; + registry.register(Box::new(self.cached_tokens.clone()))?; + registry.register(Box::new(self.tokenizer_latency.clone()))?; + registry.register(Box::new(self.output_tokens_counter.clone()))?; + registry.register(Box::new(self.time_to_first_token.clone()))?; + registry.register(Box::new(self.inter_token_latency.clone()))?; + registry.register(Box::new(self.embedding_latency.clone()))?; + registry.register(Box::new(self.images_per_request.clone()))?; + registry.register(Box::new(self.videos_per_request.clone()))?; + registry.register(Box::new(DYNAM_FAILURES_TOTAL.clone()))?; + registry.register(Box::new(self.audio_per_request.clone()))?; + registry.register(Box::new(self.image_tokens_per_request.clone()))?; + + // Register runtime configuration metrics + registry.register(Box::new(self.model_total_kv_blocks.clone()))?; + registry.register(Box::new(self.model_max_num_seqs.clone()))?; + registry.register(Box::new(self.model_max_num_batched_tokens.clone()))?; + registry.register(Box::new(self.model_context_length.clone()))?; + registry.register(Box::new(self.model_kv_cache_block_size.clone()))?; + registry.register(Box::new(self.model_migration_limit.clone()))?; + registry.register(Box::new(self.model_migration_total.clone()))?; + registry.register(Box::new(self.model_migration_duration_seconds.clone()))?; + registry.register(Box::new( + self.model_migration_max_seq_len_exceeded_total.clone(), + ))?; + registry.register(Box::new(self.model_cancellation_total.clone()))?; + registry.register(Box::new(self.model_rejection_total.clone()))?; + + Ok(()) + } + + /// Update runtime configuration metrics for a model + /// This should be called when model runtime configuration is available or updated + pub fn update_runtime_config_metrics_raw( + &self, + model_name: &str, + total_kv_blocks: Option, + max_num_seqs: Option, + max_num_batched_tokens: Option, + ) { + if let Some(total_kv_blocks) = total_kv_blocks { + self.model_total_kv_blocks + .with_label_values(&[model_name]) + .set(clamp_u64_to_i64(total_kv_blocks)); + } + + if let Some(max_num_seqs) = max_num_seqs { + self.model_max_num_seqs + .with_label_values(&[model_name]) + .set(clamp_u64_to_i64(max_num_seqs)); + } + + if let Some(max_batched_tokens) = max_num_batched_tokens { + self.model_max_num_batched_tokens + .with_label_values(&[model_name]) + .set(clamp_u64_to_i64(max_batched_tokens)); + } + } + + pub fn update_deployment_metrics( + &self, + model_name: &str, + context_length: u32, + kv_cache_block_size: u32, + migration_limit: u32, + ) { + self.model_context_length + .with_label_values(&[model_name]) + .set(context_length as i64); + self.model_kv_cache_block_size + .with_label_values(&[model_name]) + .set(kv_cache_block_size as i64); + self.model_migration_limit + .with_label_values(&[model_name]) + .set(migration_limit as i64); + } + + /// Increment the migration counter for a new request migration + pub fn inc_migration_new_request(&self, model: &str) { + self.model_migration_total + .with_label_values(&[model, frontend_service::migration_type::NEW_REQUEST]) + .inc(); + } + + /// Increment the migration counter for an ongoing request migration + pub fn inc_migration_ongoing_request(&self, model: &str) { + self.model_migration_total + .with_label_values(&[model, frontend_service::migration_type::ONGOING_REQUEST]) + .inc(); + } + + /// Get the current count of new request migrations for a model + pub fn get_migration_new_request_count(&self, model: &str) -> u64 { + self.model_migration_total + .with_label_values(&[model, frontend_service::migration_type::NEW_REQUEST]) + .get() + } + + /// Get the current count of ongoing request migrations for a model + pub fn get_migration_ongoing_request_count(&self, model: &str) -> u64 { + self.model_migration_total + .with_label_values(&[model, frontend_service::migration_type::ONGOING_REQUEST]) + .get() + } + + /// Observe the elapsed time for a completed migration event. + pub fn observe_migration_duration( + &self, + model: &str, + migration_type: &str, + outcome: &str, + duration: Duration, + ) { + self.model_migration_duration_seconds + .with_label_values(&[model, migration_type, outcome]) + .observe(duration.as_secs_f64()); + } + + /// Get the number of observed migration durations for the given dimensions. + pub fn get_migration_duration_sample_count( + &self, + model: &str, + migration_type: &str, + outcome: &str, + ) -> u64 { + self.model_migration_duration_seconds + .with_label_values(&[model, migration_type, outcome]) + .get_sample_count() + } + + /// Increment the counter for migrations disabled by max_seq_len being exceeded + pub fn inc_migration_max_seq_len_exceeded(&self, model: &str) { + self.model_migration_max_seq_len_exceeded_total + .with_label_values(&[model]) + .inc(); + } + + /// Get the current count of migrations disabled by max_seq_len being exceeded + pub fn get_migration_max_seq_len_exceeded_count(&self, model: &str) -> u64 { + self.model_migration_max_seq_len_exceeded_total + .with_label_values(&[model]) + .get() + } + + /// Increment the cancellation counter + pub fn inc_cancellation(&self, labels: &CancellationLabels) { + self.model_cancellation_total + .with_label_values(&[&labels.model, &labels.endpoint, &labels.request_type]) + .inc(); + } + + /// Get the current cancellation count + pub fn get_cancellation_count(&self, labels: &CancellationLabels) -> u64 { + self.model_cancellation_total + .with_label_values(&[&labels.model, &labels.endpoint, &labels.request_type]) + .get() + } + + /// Increment the rejection counter for a request rejected due to resource exhaustion + pub fn inc_rejection(&self, model: &str, endpoint: Endpoint) { + self.model_rejection_total + .with_label_values(&[model, &endpoint.to_string()]) + .inc(); + } + + /// Get the current rejection count for a model and endpoint + pub fn get_rejection_count(&self, model: &str, endpoint: Endpoint) -> u64 { + self.model_rejection_total + .with_label_values(&[model, &endpoint.to_string()]) + .get() + } + + /// Create a new [`InflightGuard`] for the given model and annotate if its a streaming request, + /// and the kind of endpoint that was hit + /// + /// The [`InflightGuard`] is an RAII object will handle incrementing the inflight gauge and + /// request counters. + /// + /// # Metrics Distinction + /// + /// This method creates an inflight guard t tracks requests actively being processed by the LLM engine. + /// This is distinct from [`HttpQueueGuard`] which tracks requests from HTTP handler start until + /// first token generation (including prefill time). The separation allows monitoring both HTTP processing queue time + /// and actual LLM processing time. + pub fn create_inflight_guard( + self: Arc, + model: &str, + endpoint: Endpoint, + streaming: bool, + request_id: &str, + ) -> InflightGuard { + let request_type = if streaming { + RequestType::Stream + } else { + RequestType::Unary + }; + + InflightGuard::new( + self.clone(), + model.to_string(), + endpoint, + request_type, + request_id.to_string(), + ) + } + + /// Snapshot existing metric handles; counters are never reset by publication. + pub fn frontend_load( + &self, + frontend_instance_id: u64, + is_serving_ready: bool, + models: Vec, + ) -> FrontendLoadFrame { + let mut totals = HashMap::::new(); + for family in self.request_started_counter.collect() { + for metric in family.get_metric() { + let model = metric_label(metric, "model"); + let count = self + .request_started_counter + .with_label_values(&[ + model, + metric_label(metric, "endpoint"), + metric_label(metric, "request_type"), + ]) + .get(); + let total = &mut totals.entry(model.to_owned()).or_default().requests_started; + *total = total.saturating_add(count); + } + } + for family in self.request_counter.collect() { + for metric in family.get_metric() { + let model = metric_label(metric, "model"); + let status = metric_label(metric, "status"); + let error = metric_label(metric, "error_type"); + let count = self + .request_counter + .with_label_values(&[ + model, + metric_label(metric, "endpoint"), + metric_label(metric, "request_type"), + status, + error, + ]) + .get(); + let totals = totals.entry(model.to_owned()).or_default(); + let total = match RequestCompletion::from_labels(status, error) { + RequestCompletion::Success => &mut totals.requests_completed, + RequestCompletion::Cancelled => &mut totals.requests_cancelled, + RequestCompletion::Error => &mut totals.requests_failed, + }; + *total = total.saturating_add(count); + } + } + for (model, value) in model_counter_values(&self.input_tokens_counter) { + totals.entry(model).or_default().input_tokens = value; + } + for (model, value) in model_counter_values(&self.output_tokens_counter) { + totals.entry(model).or_default().output_tokens = value; + } + let mut gauges = HashMap::::new(); + for (model, value) in model_gauge_values(&self.awaiting_first_token_requests) { + gauges + .entry(model) + .or_default() + .requests_awaiting_first_token = value; + } + for (model, value) in model_gauge_values(&self.active_requests_gauge) { + let gauge = gauges.entry(model).or_default(); + gauge.requests_generating = value.saturating_sub(gauge.requests_awaiting_first_token); + } + for family in self.inflight_input_tokens.collect() { + for metric in family.get_metric() { + let gauge = gauges + .entry(metric_label(metric, "model").to_owned()) + .or_default(); + let value = metric.get_gauge().value().max(0.0) as u64; + gauge.inflight_input_tokens = gauge.inflight_input_tokens.saturating_add(value); + if metric_label(metric, "state") == STATE_AWAITING_FIRST_TOKEN { + gauge.awaiting_first_token_input_tokens = value; + } + } + } + let models = models + .into_iter() + .map(|view| { + let mut load = FrontendModelLoad { + model: view.name, + is_ready: view.is_ready, + ..Default::default() + }; + for name in std::iter::once(&load.model).chain(&view.aliases) { + if let Some(gauge) = gauges.get(name) { + load.gauges.add(gauge); + } + if let Some(total) = totals.get(name) { + load.totals.add(total); + } + } + load + }) + .collect(); + FrontendLoadFrame { + frontend_instance_id, + incarnation: self.load_incarnation, + sequence: self.load_sequence.fetch_add(1, Relaxed) + 1, + is_serving_ready, + models, + } + } + + /// Create a new [`HttpQueueGuard`] for tracking HTTP processing queue + /// + /// This guard tracks requests from HTTP handler start until first token generation, + /// providing visibility into HTTP processing queue time before actual LLM processing begins. + pub fn create_http_queue_guard(self: Arc, model: &str) -> HttpQueueGuard { + HttpQueueGuard::new(self, model.to_string()) + } +} + +impl HttpQueueGuard { + fn new(metrics: Arc, model: String) -> Self { + // Increment the HTTP queue gauge when the guard is created + metrics.inc_http_queue_gauge(&model); + + HttpQueueGuard { metrics, model } + } +} + +impl Drop for HttpQueueGuard { + fn drop(&mut self) { + // Decrement the HTTP queue gauge when the guard is dropped + self.metrics.dec_http_queue_gauge(&self.model); + } +} + +impl InflightGuard { + fn new( + metrics: Arc, + model: String, + endpoint: Endpoint, + request_type: RequestType, + request_id: String, + ) -> Self { + let timer = Instant::now(); + let request = Arc::new(RequestMetrics::new(&metrics, &model)); + metrics.inc_inflight_gauge(&model); + metrics.inc_request_started_counter(&model, &endpoint, &request_type); + + tracing::Span::current().record("model", model.as_str()); + + tracing::info!( + request_id = %request_id, + model = %model, + endpoint = %endpoint, + request_type = %request_type, + "request received" + ); + + InflightGuard { + metrics, + request, + model, + endpoint, + request_type, + status: Status::Error, + error_type: ErrorType::Internal, + timer, + request_id, + span: tracing::Span::current(), + } + } + + pub fn request_id(&self) -> &str { + &self.request_id + } + pub fn model(&self) -> &str { + &self.model + } + pub fn endpoint(&self) -> &Endpoint { + &self.endpoint + } + pub fn request_type(&self) -> &RequestType { + &self.request_type + } + pub fn error_type(&self) -> &ErrorType { + &self.error_type + } + pub fn elapsed_ms(&self) -> u128 { + self.timer.elapsed().as_millis() + } + + pub fn request_metrics(&self) -> &Arc { + &self.request + } + + /// Create the [`ResponseMetricCollector`] for this request's responses. + pub fn response_collector(&self) -> ResponseMetricCollector { + ResponseMetricCollector::new( + self.metrics.clone(), + self.model.clone(), + self.request.clone(), + ) + } + + pub fn mark_ok(&mut self) { + self.status = Status::Success; + self.error_type = ErrorType::None; + } + + pub fn mark_error(&mut self, error_type: ErrorType) { + self.status = Status::Error; + self.error_type = error_type; + } + + fn request_completion(&self) -> RequestCompletion { + RequestCompletion::from_labels(self.status.as_str(), self.error_type.as_str()) + } +} + +impl Drop for InflightGuard { + fn drop(&mut self) { + let _enter = self.span.enter(); + let duration = self.timer.elapsed().as_secs_f64(); + self.metrics.dec_inflight_gauge(&self.model); + self.metrics.inc_request_counter( + &self.model, + &self.endpoint, + &self.request_type, + &self.status, + &self.error_type, + ); + self.metrics + .request_duration + .with_label_values(&[&self.model]) + .observe(duration); + + let completion = self.request_completion(); + self.request.finish(); + self.span.record("request.outcome", completion.as_str()); + + let elapsed_ms = (duration * 1000.0) as u64; + let status_str = self.status.as_str(); + match completion { + RequestCompletion::Error => { + let detail = match self.error_type { + ErrorType::ResponseTimeout => "backend stream inactivity timeout", + ErrorType::Internal => "internal server error during processing", + ErrorType::Validation => "invalid request parameters", + ErrorType::NotFound => "model or resource not found", + ErrorType::Overload => "service overloaded or rate limited", + ErrorType::Unavailable => "no backend worker available", + ErrorType::NotImplemented => "requested feature not implemented", + // `request_completion()` routes `(Error, Cancelled)` to + // `RequestCompletion::Cancelled`, so only `None` reaches + // here. `Cancelled` is listed to keep the match total + // without a panic in a `Drop` impl. + ErrorType::None | ErrorType::Cancelled => "unknown error", + }; + tracing::error!( + request_id = %self.request_id, + model = %self.model, + endpoint = %self.endpoint, + request_type = %self.request_type, + status = %status_str, + error_type = %self.error_type, + error_detail = %detail, + elapsed_ms = %elapsed_ms, + "request completed" + ); + } + RequestCompletion::Cancelled => { + tracing::info!( + request_id = %self.request_id, + model = %self.model, + endpoint = %self.endpoint, + request_type = %self.request_type, + status = %completion.as_str(), + error_type = %self.error_type, + error_detail = "cancelled before completion", + elapsed_ms = %elapsed_ms, + "request completed" + ); + } + RequestCompletion::Success => { + tracing::info!( + request_id = %self.request_id, + model = %self.model, + endpoint = %self.endpoint, + request_type = %self.request_type, + status = %status_str, + elapsed_ms = %elapsed_ms, + "request completed" + ); + } + } + } +} + +impl std::fmt::Display for Endpoint { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Endpoint::Completions => write!(f, "completions"), + Endpoint::ChatCompletions => write!(f, "chat_completions"), + Endpoint::Embeddings => write!(f, "embeddings"), + Endpoint::Classify => write!(f, "classify"), + Endpoint::Rerank => write!(f, "rerank"), + Endpoint::Pooling => write!(f, "pooling"), + Endpoint::Images => write!(f, "images"), + Endpoint::Videos => write!(f, "videos"), + Endpoint::Audios => write!(f, "audios"), + Endpoint::Responses => write!(f, "responses"), + Endpoint::AnthropicMessages => write!(f, "anthropic_messages"), + Endpoint::Tensor => write!(f, "tensor"), + Endpoint::Generate => write!(f, "generate"), + } + } +} + +impl Endpoint { + pub fn as_str(&self) -> &'static str { + match self { + Endpoint::Completions => "completions", + Endpoint::ChatCompletions => "chat_completions", + Endpoint::Embeddings => "embeddings", + Endpoint::Classify => "classify", + Endpoint::Rerank => "rerank", + Endpoint::Pooling => "pooling", + Endpoint::Images => "images", + Endpoint::Videos => "videos", + Endpoint::Audios => "audios", + Endpoint::Responses => "responses", + Endpoint::AnthropicMessages => "anthropic_messages", + Endpoint::Tensor => "tensor", + Endpoint::Generate => "generate", + } + } +} + +impl RequestType { + pub fn as_str(&self) -> &'static str { + match self { + RequestType::Unary => frontend_service::request_type::UNARY, + RequestType::Stream => frontend_service::request_type::STREAM, + } + } +} + +impl std::fmt::Display for RequestType { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(self.as_str()) + } +} + +impl Status { + pub fn as_str(&self) -> &'static str { + match self { + Status::Success => frontend_service::status::SUCCESS, + Status::Error => frontend_service::status::ERROR, + } + } +} + +impl ErrorType { + pub fn as_str(&self) -> &'static str { + match self { + ErrorType::None => frontend_service::error_type::NONE, + ErrorType::Validation => frontend_service::error_type::VALIDATION, + ErrorType::NotFound => frontend_service::error_type::NOT_FOUND, + ErrorType::Overload => frontend_service::error_type::OVERLOAD, + ErrorType::Unavailable => frontend_service::error_type::UNAVAILABLE, + ErrorType::Cancelled => frontend_service::error_type::CANCELLED, + ErrorType::ResponseTimeout => frontend_service::error_type::RESPONSE_TIMEOUT, + ErrorType::Internal => frontend_service::error_type::INTERNAL, + ErrorType::NotImplemented => frontend_service::error_type::NOT_IMPLEMENTED, + } + } +} + +impl std::fmt::Display for ErrorType { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(self.as_str()) + } +} + +impl ResponseMetricCollector { + pub fn new(metrics: Arc, model: String, request: Arc) -> Self { + // Resolve the per-model handles once (cheap clones of the vec entries) so the + // per-chunk / per-token hot path in `observe_response` does no label hashing. + let time_to_first_token = metrics.time_to_first_token.with_label_values(&[&model]); + let input_sequence_length = metrics.input_sequence_length.with_label_values(&[&model]); + let cached_tokens = metrics.cached_tokens.with_label_values(&[&model]); + let images_per_request = metrics.images_per_request.with_label_values(&[&model]); + let videos_per_request = metrics.videos_per_request.with_label_values(&[&model]); + let audio_per_request = metrics.audio_per_request.with_label_values(&[&model]); + let image_tokens_per_request = metrics + .image_tokens_per_request + .with_label_values(&[&model]); + ResponseMetricCollector { + metrics, + request, + model, + time_to_first_token, + inter_token_latency: None, + itl_pending_tokens: 0, + input_sequence_length, + cached_tokens, + images_per_request, + videos_per_request, + audio_per_request, + image_tokens_per_request, + image_count_val: 0, + video_count_val: 0, + audio_count_val: 0, + multimodal_counts_observed: false, + is_first_token: true, + last_response_time: None, + start_time: Instant::now(), + osl: 0, + isl: 0, + ttft_ms: None, + itl_sum_secs: 0.0, + itl_count: 0, + cached_tokens_observed: false, + tokenize_latency_observed: false, + detokenize_latency_total: Duration::ZERO, + detokenize_count_total: 0, + prefill_worker_id: None, + prefill_dp_rank: None, + prefill_worker_type: None, + decode_worker_id: None, + decode_dp_rank: None, + decode_worker_type: None, + decode_itl_gauge: None, + } + } + + /// Set the worker info for per-worker TTFT/ITL metrics. + /// In disaggregated mode, TTFT is attributed to prefill worker, ITL to decode worker. + /// Worker types are stored at routing time to avoid expensive MDC lookup when updating metrics. + pub fn set_worker_info( + &mut self, + prefill_worker_id: Option, + prefill_dp_rank: Option, + prefill_worker_type: Option, + decode_worker_id: Option, + decode_dp_rank: Option, + decode_worker_type: Option, + ) { + if self.prefill_worker_id.is_none() { + self.prefill_worker_id = prefill_worker_id; + } + if self.prefill_dp_rank.is_none() { + self.prefill_dp_rank = prefill_dp_rank; + } + if self.prefill_worker_type.is_none() { + self.prefill_worker_type = prefill_worker_type; + } + if self.decode_worker_id.is_none() { + self.decode_worker_id = decode_worker_id; + } + if self.decode_dp_rank.is_none() { + self.decode_dp_rank = decode_dp_rank; + } + if self.decode_worker_type.is_none() { + self.decode_worker_type = decode_worker_type; + } + } + + pub fn set_worker_info_from_parts( + &mut self, + prefill_worker_id: Option, + prefill_dp_rank: Option, + prefill_worker_type: Option<&str>, + decode_worker_id: Option, + decode_dp_rank: Option, + decode_worker_type: Option<&str>, + ) { + let prefill_worker_type = self + .prefill_worker_type + .is_none() + .then(|| prefill_worker_type.map(str::to_owned)) + .flatten(); + let decode_worker_type = self + .decode_worker_type + .is_none() + .then(|| decode_worker_type.map(str::to_owned)) + .flatten(); + + self.set_worker_info( + prefill_worker_id, + prefill_dp_rank, + prefill_worker_type, + decode_worker_id, + decode_dp_rank, + decode_worker_type, + ); + } + + fn local_inter_token_latency(&mut self) -> &mut prometheus::local::LocalHistogram { + let metrics = &self.metrics; + let model = self.model.as_str(); + self.inter_token_latency.get_or_insert_with(|| { + metrics + .inter_token_latency + .with_label_values(&[model]) + .local() + }) + } + + /// Observe the current output sequence length + pub fn observe_current_osl(&mut self, osl: usize) { + self.osl = osl; + } + + /// Check if this will be the first token (before calling observe_response) + pub fn is_first_token(&self) -> bool { + self.is_first_token + } + + /// Observe cached tokens (prefix cache hits), observing only once per request when value is available + pub fn observe_cached_tokens(&mut self, cached_tokens: Option) { + if let Some(tokens) = cached_tokens + && !self.cached_tokens_observed + { + self.cached_tokens_observed = true; + self.cached_tokens.observe(tokens as f64); + } + } + + /// Observe per-request multimodal content-part counts, latched to run + /// exactly once per request (the counts are constant across chunks). + pub fn observe_multimodal_counts( + &mut self, + image_count: usize, + video_count: usize, + audio_count: usize, + ) { + self.observe_multimodal_metrics(image_count, video_count, audio_count, None); + } + + pub fn observe_multimodal_metrics( + &mut self, + image_count: usize, + video_count: usize, + audio_count: usize, + image_tokens: Option, + ) { + if self.multimodal_counts_observed { + return; + } + self.multimodal_counts_observed = true; + + self.image_count_val = image_count; + self.video_count_val = video_count; + self.audio_count_val = audio_count; + + self.images_per_request.observe(image_count as f64); + self.videos_per_request.observe(video_count as f64); + self.audio_per_request.observe(audio_count as f64); + if let Some(image_tokens) = image_tokens { + self.image_tokens_per_request.observe(image_tokens as f64); + } + } + + /// Observe tokenize/detokenize latencies in milliseconds. + /// Tokenize is observed once per request; detokenize is accumulated and observed at request end. + pub fn observe_tokenize_latencies( + &mut self, + tokenize_latency: Option, + detokenize_latency: Option, + detokenize_count: Option, + ) { + if let Some(latency) = tokenize_latency + && !self.tokenize_latency_observed + { + self.tokenize_latency_observed = true; + self.metrics + .tokenizer_latency + .with_label_values(&[frontend_service::operation::TOKENIZE]) + .observe(latency.as_secs_f64() * 1000.0); + } + + if let Some(latency) = detokenize_latency { + self.detokenize_latency_total = latency; + } + if let Some(count) = detokenize_count { + self.detokenize_count_total = count; + } + } + + /// Observe a response with input sequence length and number of new tokens + pub fn observe_response(&mut self, isl: usize, num_tokens: usize) { + self.request.observe_input_tokens(isl); + self.request.add_output_tokens(num_tokens); + if num_tokens == 0 { + return; + } + + // Store ISL for span recording on drop + self.isl = isl; + + if self.is_first_token { + // NOTE: when there are multiple tokens in the first response, + // we use the full response time as TTFT and ignore the ITL + self.is_first_token = false; + + // Publish TTFT and store for span recording + let ttft = self.start_time.elapsed().as_secs_f64(); + self.ttft_ms = Some(ttft * 1000.0); + self.time_to_first_token.observe(ttft); + + // Update per-worker TTFT and input sequence tokens gauges - attributed to prefill worker. + // Both gauges are updated atomically from the same request to correlate latency with input size. + // Use stored worker_type (from routing time) to avoid MDC lookup. + // Falls back to WORKER_TYPE_PREFILL if not available. + if let Some(worker_id) = self.prefill_worker_id { + let worker_id_str = worker_id.to_string(); + let dp_rank_str = self + .prefill_dp_rank + .map_or(UNSET_DP_RANK_LABEL.to_string(), |r| r.to_string()); + let worker_type = self + .prefill_worker_type + .as_deref() + .unwrap_or(WORKER_TYPE_PREFILL); + let labels = &[worker_id_str.as_str(), dp_rank_str.as_str(), worker_type]; + WORKER_LAST_TIME_TO_FIRST_TOKEN_GAUGE + .with_label_values(labels) + .set(ttft); + WORKER_LAST_INPUT_SEQUENCE_TOKENS_GAUGE + .with_label_values(labels) + .set(isl as i64); + } + + // Publish ISL + // TODO: publish ISL as soon as the tokenization process completes + self.input_sequence_length.observe(isl as f64); + } + + let current_duration = self.start_time.elapsed(); + + if let Some(last_response_time) = self.last_response_time { + let response_duration = current_duration - last_response_time; + let itl = response_duration.as_secs_f64() / num_tokens as f64; + self.itl_sum_secs += itl * num_tokens as f64; + self.itl_count += num_tokens as u64; + self.itl_pending_tokens = self.itl_pending_tokens.saturating_add(num_tokens as u64); + let should_flush = self.itl_pending_tokens >= ITL_LOCAL_FLUSH_TOKENS; + { + // Resolve the request-local histogram on the first ITL only, then reuse + // it without re-hashing the model label for each output token. + let histogram = self.local_inter_token_latency(); + for _ in 0..num_tokens { + histogram.observe(itl); + } + if should_flush { + histogram.flush(); + } + } + if should_flush { + self.itl_pending_tokens = 0; + } + + // Update per-worker ITL gauge - attributed to decode worker. + // Use stored worker_type (from routing time) to avoid MDC lookup. + // Falls back to WORKER_TYPE_DECODE if not available. + // + // The worker labels are fixed for the request's lifetime (latched in + // `set_worker_info`), so resolve the GaugeVec handle once and cache it. + // This keeps the per-token path free of `to_string()` allocations and + // label hashing; subsequent tokens only call `set`. + if self.decode_itl_gauge.is_none() + && let Some(worker_id) = self.decode_worker_id + { + let worker_id_str = worker_id.to_string(); + let dp_rank_str = self + .decode_dp_rank + .map_or(UNSET_DP_RANK_LABEL.to_string(), |r| r.to_string()); + let worker_type = self + .decode_worker_type + .as_deref() + .unwrap_or(WORKER_TYPE_DECODE); + self.decode_itl_gauge = + Some(WORKER_LAST_INTER_TOKEN_LATENCY_GAUGE.with_label_values(&[ + worker_id_str.as_str(), + dp_rank_str.as_str(), + worker_type, + ])); + } + if let Some(gauge) = &self.decode_itl_gauge { + gauge.set(itl); + } + } + + self.last_response_time = Some(current_duration); + } +} + +impl Drop for ResponseMetricCollector { + fn drop(&mut self) { + if let Some(histogram) = &self.inter_token_latency { + histogram.flush(); + } + self.itl_pending_tokens = 0; + + if !self.detokenize_latency_total.is_zero() && self.detokenize_count_total > 0 { + let avg_detokenize_latency_ms = (self.detokenize_latency_total.as_secs_f64() * 1000.0) + / self.detokenize_count_total as f64; + self.metrics + .tokenizer_latency + .with_label_values(&[frontend_service::operation::DETOKENIZE]) + .observe(avg_detokenize_latency_ms); + } + + // Publish final OSL when the collector is dropped, but only for + // requests that actually produced output tokens. Recording zero for + // failed/cancelled requests would corrupt histogram averages (sum/count) + // and percentiles. Failures are already tracked by requests_total with + // status and error_type labels. + if self.osl > 0 { + self.metrics + .output_sequence_length + .with_label_values(&[&self.model]) + .observe(self.osl as f64); + } + + // Record request summary on the enclosing span. + // InflightGuard::Drop and on_response logs will inherit these. + let span = tracing::Span::current(); + span.record("input_tokens", self.isl as u32); + span.record("output_tokens", self.osl as u32); + // Only record once observed, so requests that never carried a metrics + // annotation (e.g. errored before the first chunk) don't stamp a false zero. + if self.multimodal_counts_observed { + span.record("image_count", self.image_count_val as u32); + span.record("video_count", self.video_count_val as u32); + span.record("audio_count", self.audio_count_val as u32); + } + if let Some(ttft_ms) = self.ttft_ms { + span.record("ttft_ms", format!("{:.2}", ttft_ms).as_str()); + } + if self.itl_count > 0 { + let avg_ms = (self.itl_sum_secs / self.itl_count as f64) * 1000.0; + span.record("avg_itl_ms", format!("{:.2}", avg_ms).as_str()); + } + if let Some(worker_id) = self.prefill_worker_id { + span.record("prefill_worker_id", worker_id); + } + if let Some(worker_id) = self.decode_worker_id { + span.record("decode_worker_id", worker_id); + } + } +} + +impl ServingLoadCollector { + pub fn initialize_model(&self, model: &str) { + initialize_model_load( + &self.awaiting_requests, + &self.input_tokens, + &self.input_tokens_total, + model, + ); + } +} + +#[cfg(test)] +mod tests { + use super::*; + fn response_collector(metrics: &Arc, model: &str) -> ResponseMetricCollector { + ResponseMetricCollector::new( + metrics.clone(), + model.to_string(), + Arc::new(RequestMetrics::new(metrics, model)), + ) + } + + #[test] + fn test_round_to_sig_figs() { + // Test rounding to 2 significant figures + assert_eq!(round_to_sig_figs(0.0026, 2), 0.0026); + assert_eq!(round_to_sig_figs(0.26, 2), 0.26); + assert_eq!(round_to_sig_figs(0.2356, 2), 0.24); + assert_eq!(round_to_sig_figs(1.234, 2), 1.2); + assert_eq!(round_to_sig_figs(12.34, 2), 12.0); + assert_eq!(round_to_sig_figs(123.4, 2), 120.0); + assert_eq!(round_to_sig_figs(1234.0, 2), 1200.0); + assert_eq!(round_to_sig_figs(0.0, 2), 0.0); + + // Test edge cases + assert_eq!(round_to_sig_figs(0.999, 2), 1.0); + assert_eq!(round_to_sig_figs(9.99, 2), 10.0); + assert_eq!(round_to_sig_figs(99.9, 2), 100.0); + } + + #[test] + fn test_generate_log_buckets_basic() { + // Test basic properties + let buckets = generate_log_buckets(1.0, 100.0, 5); + + // Check length + assert_eq!(buckets.len(), 5); + + // Check first value is 0 + assert_eq!(buckets[0], 0.0); + + // Check last value is approximately max (rounded to 2 sig figs) + assert_eq!(buckets[buckets.len() - 1], 100.0); + + // Check values are increasing + for i in 1..buckets.len() { + assert!( + buckets[i] > buckets[i - 1], + "Bucket values should be increasing: {} <= {}", + buckets[i - 1], + buckets[i] + ); + } + } + + /// A stand-in for the process environment. Tests inject this instead of calling + /// `temp_env`: mutating the real environment races the many other tests in this + /// module that construct `Metrics` and read these same variables, which on Unix is + /// an unsafe `setenv`/`getenv` race, not merely an ordering hazard. + fn fake_env<'a>(pairs: &'a [(&'a str, &'a str)]) -> impl Fn(&str) -> Option + 'a { + move |name| { + pairs + .iter() + .find(|(key, _)| *key == name) + .map(|(_, value)| (*value).to_string()) + } + } + + #[test] + fn bucket_config_reads_the_documented_env_var_names() { + // These names were once read under an extra DYN_HISTOGRAM_ prefix, which + // silently stopped them working; they must resolve as documented. + let env = fake_env(&[ + ("DYN_METRICS_ITL_MIN", "0.002"), + ("DYN_METRICS_ITL_MAX", "80"), + ("DYN_METRICS_ITL_COUNT", "20"), + ]); + let cfg = parse_bucket_config(&env, env_metrics::DYN_METRICS_ITL, 0.001, 2.0, 13); + assert_eq!(cfg, (0.002, 80.0, 20)); + } + + #[test] + fn bucket_config_falls_back_to_the_deprecated_doubled_name() { + // The doubled form was the only name that worked for several releases, so it + // stays supported for one more rather than silently reverting to defaults. + let env = fake_env(&[("DYN_HISTOGRAM_DYN_METRICS_ITL_MAX", "80")]); + let (_, max, _) = parse_bucket_config(&env, env_metrics::DYN_METRICS_ITL, 0.001, 2.0, 13); + assert_eq!(max, 80.0); + } + + #[test] + fn bucket_config_prefers_the_new_name_over_the_deprecated_one() { + let env = fake_env(&[ + ("DYN_METRICS_ITL_MAX", "80"), + ("DYN_HISTOGRAM_DYN_METRICS_ITL_MAX", "30"), + ]); + let (_, max, _) = parse_bucket_config(&env, env_metrics::DYN_METRICS_ITL, 0.001, 2.0, 13); + assert_eq!(max, 80.0); + } + + #[test] + fn bucket_config_falls_back_to_defaults_for_unset_and_unparseable_values() { + let env = fake_env(&[ + ("DYN_METRICS_ITL_MIN", "not-a-number"), + ("DYN_METRICS_ITL_COUNT", "12.5"), + ]); + let cfg = parse_bucket_config(&env, env_metrics::DYN_METRICS_ITL, 0.001, 2.0, 13); + assert_eq!(cfg, (0.001, 2.0, 13)); + } + + #[test] + fn bucket_config_rejects_a_parseable_but_invalid_combination() { + // min >= max parses fine but cannot produce buckets, so the whole set reverts — + // including a COUNT that was valid on its own. + let env = fake_env(&[ + ("DYN_METRICS_ITL_MIN", "5"), + ("DYN_METRICS_ITL_MAX", "1"), + ("DYN_METRICS_ITL_COUNT", "20"), + ]); + let cfg = parse_bucket_config(&env, env_metrics::DYN_METRICS_ITL, 0.001, 2.0, 13); + assert_eq!(cfg, (0.001, 2.0, 13)); + } + + /// Upper bounds of a histogram's buckets, i.e. the `le` label values that a + /// dashboard's `histogram_quantile` query sees on `/metrics`. + fn bucket_upper_bounds(registry: &Registry, metric_name: &str) -> Vec { + registry + .gather() + .into_iter() + .find(|family| family.name() == metric_name) + .expect("histogram not registered") + .get_metric()[0] + .get_histogram() + .get_bucket() + .iter() + .map(|b| b.upper_bound()) + .collect() + } + + #[test] + fn itl_ceiling_env_var_reaches_the_exported_le_labels() { + // Deliberately not the default (80.0, 20): a test that sets the env to the + // default value would pass even if the variable were ignored entirely. + let env = fake_env(&[ + ("DYN_METRICS_ITL_MAX", "30"), + ("DYN_METRICS_ITL_COUNT", "8"), + ]); + let registry = Registry::new(); + let metrics = Metrics::build(None, &env); + metrics.register(®istry).unwrap(); + metrics + .inter_token_latency + .with_label_values(&["m"]) + .observe(0.5); + + let bounds = bucket_upper_bounds( + ®istry, + &format!( + "{}_{}", + name_prefix::FRONTEND, + frontend_service::INTER_TOKEN_LATENCY_SECONDS + ), + ); + assert_eq!( + bounds.last().copied(), + Some(30.0), + "top finite bucket should follow DYN_METRICS_ITL_MAX, got {bounds:?}" + ); + assert_eq!( + bounds.len(), + 8, + "bucket count should follow DYN_METRICS_ITL_COUNT" + ); + } + + #[test] + fn itl_default_buckets_reach_80_seconds() { + // The ceiling was 2.0s, so any inter-token latency above it landed in `+Inf` and + // `histogram_quantile` pinned p99 at exactly 2.0 -- indistinguishable from a real + // 2s measurement. Guard the shipped default, not just the env-var override. + let registry = Registry::new(); + let metrics = Metrics::build(None, &fake_env(&[])); + metrics.register(®istry).unwrap(); + metrics + .inter_token_latency + .with_label_values(&["m"]) + .observe(0.5); + + let bounds = bucket_upper_bounds( + ®istry, + &format!( + "{}_{}", + name_prefix::FRONTEND, + frontend_service::INTER_TOKEN_LATENCY_SECONDS + ), + ); + // The exact set is the contract. Note the bottom edge of 0.0018: vLLM's equivalent + // starts at 0.01 and cannot resolve anything faster than 10ms per token, so raising + // the ceiling must not cost the low-end resolution Dynamo has and vLLM does not. + assert_eq!( + bounds, + vec![ + 0.0, 0.0018, 0.0033, 0.0059, 0.011, 0.02, 0.035, 0.064, 0.12, 0.21, 0.38, 0.69, + 1.2, 2.3, 4.1, 7.4, 13.0, 24.0, 44.0, 80.0 + ], + ); + } + + #[test] + fn explicit_itl_buckets_take_precedence_over_log_config() { + let env = fake_env(&[ + ("DYN_METRICS_ITL_BUCKETS", "0.01, 0.02, 0.05"), + ("DYN_METRICS_ITL_MIN", "1"), + ("DYN_METRICS_ITL_MAX", "10"), + ("DYN_METRICS_ITL_COUNT", "3"), + ]); + let registry = Registry::new(); + let metrics = Metrics::build(None, &env); + metrics.register(®istry).unwrap(); + metrics + .inter_token_latency + .with_label_values(&["m"]) + .observe(0.03); + + let bounds = bucket_upper_bounds( + ®istry, + &format!( + "{}_{}", + name_prefix::FRONTEND, + frontend_service::INTER_TOKEN_LATENCY_SECONDS + ), + ); + assert_eq!(bounds, vec![0.01, 0.02, 0.05]); + } + + #[test] + fn invalid_explicit_itl_buckets_fall_back_to_defaults() { + for value in [ + "0.01,nope", + "0.01,NaN", + "-0.01,0.02", + "0.02,0.01", + "0.01,0.01", + ] { + let pairs = [ + ("DYN_METRICS_ITL_BUCKETS", value), + ("DYN_METRICS_ITL_MAX", "30"), + ("DYN_METRICS_ITL_COUNT", "8"), + ]; + let env = fake_env(&pairs); + assert_eq!( + inter_token_latency_buckets(&env), + generate_log_buckets(0.001, 80.0, 20) + ); + } + + let oversized_value = format!("{:>width$}", "0", width = MAX_BUCKET_LIST_BYTES + 1); + let pairs = [("DYN_METRICS_ITL_BUCKETS", oversized_value.as_str())]; + let env = fake_env(&pairs); + assert_eq!( + inter_token_latency_buckets(&env), + generate_log_buckets(0.001, 80.0, 20) + ); + } + + #[test] + fn test_generate_log_buckets_edge_cases() { + // Test empty buckets + let buckets = generate_log_buckets(1.0, 100.0, 0); + assert_eq!(buckets.len(), 0); + + // Test single bucket + let buckets = generate_log_buckets(1.0, 100.0, 1); + assert_eq!(buckets.len(), 1); + assert_eq!(buckets[0], 0.0); + + // Test two buckets + let buckets = generate_log_buckets(1.0, 100.0, 2); + assert_eq!(buckets.len(), 2); + assert_eq!(buckets[0], 0.0); + assert_eq!(buckets[1], 100.0); + } + + #[test] + fn test_generate_log_buckets_always_includes_zero() { + // Test various configurations + for count in 1..=20 { + let buckets = generate_log_buckets(0.1, 1000.0, count); + assert_eq!( + buckets[0], 0.0, + "First bucket should always be 0.0 for count={}", + count + ); + } + } + + #[test] + fn test_all_buckets_are_two_sig_figs() { + let test_cases = vec![ + (1.0, 512.0, 10), + (50.0, 128000.0, 12), + (50.0, 32000.0, 10), + (0.001, 480.0, 18), + (0.001, 80.0, 20), + (0.001, 10.0, 14), + ]; + + for (min, max, count) in test_cases { + let buckets = generate_log_buckets(min, max, count); + for &value in buckets.iter().skip(1) { + let rounded = round_to_sig_figs(value, 2); + assert_eq!( + value, rounded, + "Value {} should be rounded to 2 sig figs (min={}, max={}, count={})", + value, min, max, count + ); + } + } + } + + #[test] + fn test_sig_fig_limitation_with_many_buckets() { + // This test demonstrates that 2 sig figs limits the number of unique bucket values + // With 1000 requested buckets but only 2 sig figs, we'll get automatic deduplication + let buckets = generate_log_buckets(0.0001, 1.0, 1000); + + println!( + "Requested 1000 buckets, got {} total values (including 0.0)", + buckets.len() + ); + + // With 2 sig figs across 4 orders of magnitude (0.0001 to 1.0), + // we can have roughly 90 unique values per order of magnitude + // So we expect around 360 unique values maximum + assert!( + buckets.len() < 500, + "Expected fewer than 500 unique buckets due to 2 sig fig limitation, got {}", + buckets.len() + ); + + // Verify all values are unique (no duplicates remain after deduplication) + let mut sorted_buckets = buckets.clone(); + sorted_buckets.sort_by(|a, b| a.partial_cmp(b).unwrap()); + sorted_buckets.dedup(); + assert_eq!( + buckets.len(), + sorted_buckets.len(), + "All buckets should be unique after deduplication" + ); + + // Verify first is still 0.0 + assert_eq!(buckets[0], 0.0); + + // Verify values are still in increasing order + for i in 1..buckets.len() { + assert!( + buckets[i] > buckets[i - 1], + "Buckets should be in increasing order" + ); + } + } + + #[test] + fn test_deduplication_preserves_order() { + // Test that deduplication maintains increasing order + let buckets = generate_log_buckets(0.01, 1.0, 50); + + // Verify all values are unique + let mut unique_check = std::collections::HashSet::new(); + for &bucket in &buckets { + assert!( + unique_check.insert(bucket.to_bits()), + "Duplicate value {} found after deduplication", + bucket + ); + } + + // Verify order is maintained + for i in 1..buckets.len() { + assert!( + buckets[i] > buckets[i - 1], + "Bucket values should be in increasing order after deduplication" + ); + } + } + + #[test] + fn test_output_tokens_counter_increments() { + let metrics = Arc::new(Metrics::new()); + let registry = prometheus::Registry::new(); + metrics.register(®istry).unwrap(); + + let model = "test-model"; + + // Create response collector + let mut collector = response_collector(&metrics, model); + + // Simulate first chunk (5 tokens) + collector.observe_response(100, 5); + + // Verify counter incremented by 5 + let counter_value = metrics + .output_tokens_counter + .with_label_values(&[model]) + .get(); + assert_eq!(counter_value, 5); + + // Simulate second chunk (10 tokens) + collector.observe_response(100, 10); + + // Verify counter incremented to 15 + let counter_value = metrics + .output_tokens_counter + .with_label_values(&[model]) + .get(); + assert_eq!(counter_value, 15); + + // Simulate third chunk (7 tokens) + collector.observe_response(100, 7); + + // Verify counter incremented to 22 + let counter_value = metrics + .output_tokens_counter + .with_label_values(&[model]) + .get(); + assert_eq!(counter_value, 22); + } + + #[test] + fn test_inflight_guard_feeds_frontend_load() { + let metrics = Arc::new(Metrics::new()); + let model = "frontend-load-model"; + let load = || { + let view = crate::load::ModelLoadView { + name: model.to_string(), + is_ready: true, + aliases: Vec::new(), + }; + metrics.frontend_load(0, true, vec![view]).models.remove(0) + }; + let mut inflight = + metrics + .clone() + .create_inflight_guard(model, Endpoint::Completions, true, ""); + let mut collector = inflight.response_collector(); + + let awaiting_requests = metrics + .awaiting_first_token_requests + .with_label_values(&[model]); + let awaiting_tokens = metrics + .inflight_input_tokens + .with_label_values(&[model, STATE_AWAITING_FIRST_TOKEN]); + let generating_tokens = metrics + .inflight_input_tokens + .with_label_values(&[model, STATE_GENERATING]); + collector.observe_response(10, 0); + assert_eq!( + ( + awaiting_requests.get(), + awaiting_tokens.get(), + generating_tokens.get() + ), + (1, 10, 0) + ); + let awaiting = load(); + assert_eq!(awaiting.gauges.requests_awaiting_first_token, 1); + assert_eq!(awaiting.gauges.awaiting_first_token_input_tokens, 10); + + // Output chunks are deltas; the repeated input count is not re-added. + collector.observe_response(10, 3); + collector.observe_response(10, 2); + assert_eq!( + ( + awaiting_requests.get(), + awaiting_tokens.get(), + generating_tokens.get() + ), + (0, 0, 10) + ); + inflight.request_metrics().add_input_tokens(4); + assert_eq!(generating_tokens.get(), 14); + let generating = load(); + assert_eq!(generating.gauges.requests_generating, 1); + assert_eq!(generating.gauges.inflight_input_tokens, 14); + + inflight.mark_ok(); + drop(inflight); + assert_eq!( + ( + awaiting_requests.get(), + awaiting_tokens.get(), + generating_tokens.get() + ), + (0, 0, 0) + ); + let finished = load(); + assert_eq!(finished.gauges, Default::default()); + assert_eq!(finished.totals.requests_completed, 1); + assert_eq!(finished.totals.input_tokens, 14); + assert_eq!(finished.totals.output_tokens, 5); + } + + #[test] + fn frontend_frames_keep_the_metrics_incarnation_across_publications() { + let metrics = Metrics::new(); + let first = metrics.frontend_load(7, true, Vec::new()); + let second = metrics.frontend_load(7, false, Vec::new()); + assert_eq!(first.sequence, 1); + assert_eq!(second.sequence, 2); + assert_eq!(first.incarnation, second.incarnation); + } + + /// Publishing load only reads /metrics: a snapshot of a model with an + /// alias adds no series to any family. + #[test] + fn frontend_load_creates_no_prometheus_series() { + let metrics = Metrics::new(); + let registry = Registry::new(); + metrics.register(®istry).unwrap(); + let series = || { + registry + .gather() + .iter() + .flat_map(|family| { + family + .get_metric() + .iter() + .filter(|metric| { + metric + .get_label() + .iter() + .any(|label| label.name() == "model") + }) + .map(move |metric| { + let labels = metric + .get_label() + .iter() + .map(|label| format!("{}={}", label.name(), label.value())) + .collect::>() + .join(","); + format!("{}{{{labels}}}", family.name()) + }) + }) + .collect::>() + }; + + let before = series(); + metrics.frontend_load( + 0, + true, + vec![crate::load::ModelLoadView { + name: "idle-model".into(), + aliases: vec!["idle-alias".into()], + is_ready: true, + }], + ); + assert_eq!(series(), before); + } + + #[test] + fn frontend_load_reads_existing_counters_and_keeps_integer_precision() { + let metrics = Metrics::new(); + let large_count = (1_u64 << 53) + 1; + metrics + .request_started_counter + .with_label_values(&["uncommitted", "completions", "stream"]) + .inc(); + metrics + .request_started_counter + .with_label_values(&["primary", "completions", "stream"]) + .inc_by(large_count); + metrics + .request_counter + .with_label_values(&["alias", "embeddings", "unary", "error", "cancelled"]) + .inc_by(3); + metrics + .request_counter + .with_label_values(&["primary", "completions", "stream", "success", "none"]) + .inc_by(7); + metrics + .request_counter + .with_label_values(&["primary", "completions", "stream", "error", "internal"]) + .inc_by(2); + metrics + .output_tokens_counter + .with_label_values(&["alias"]) + .inc_by(large_count); + metrics + .active_requests_gauge + .with_label_values(&["primary"]) + .set(4); + metrics + .awaiting_first_token_requests + .with_label_values(&["primary"]) + .set(1); + metrics + .inflight_input_tokens + .with_label_values(&["primary", STATE_AWAITING_FIRST_TOKEN]) + .set(5); + metrics + .inflight_input_tokens + .with_label_values(&["primary", STATE_GENERATING]) + .set(8); + let frame = metrics.frontend_load( + 0, + true, + vec![crate::load::ModelLoadView { + name: "primary".into(), + is_ready: true, + aliases: vec!["alias".into()], + }], + ); + assert_eq!( + frame.models[0].gauges, + RequestGauges { + requests_awaiting_first_token: 1, + requests_generating: 3, + awaiting_first_token_input_tokens: 5, + inflight_input_tokens: 13, + } + ); + assert_eq!( + frame.models[0].totals, + RequestTotals { + requests_started: large_count, + requests_completed: 7, + requests_failed: 2, + requests_cancelled: 3, + output_tokens: large_count, + ..Default::default() + } + ); + } + + #[test] + fn frontend_request_metrics_ignore_updates_after_guard_finishes() { + let metrics = Arc::new(Metrics::new()); + let mut guard = + metrics + .clone() + .create_inflight_guard("model", Endpoint::Completions, true, ""); + let request = guard.request_metrics().clone(); + let mut context = Context::new(()); + request.attach_to(&mut context); + assert!(Arc::ptr_eq( + &request, + &RequestMetrics::from_context(&context).unwrap() + )); + request.add_input_tokens(3); + request.add_input_tokens(5); + request.observe_input_tokens(5); + request.observe_input_tokens(3); + assert_eq!( + metrics + .input_tokens_counter + .with_label_values(&["model"]) + .get(), + 8 + ); + assert_eq!( + metrics + .input_sequence_length + .with_label_values(&["model"]) + .get_sample_sum(), + 0.0 + ); + guard.mark_error(ErrorType::Cancelled); + drop(guard); + request.add_input_tokens(9); + request.add_output_tokens(1); + let frame = metrics.frontend_load( + 0, + true, + vec![crate::load::ModelLoadView { + name: "model".into(), + is_ready: true, + aliases: vec![], + }], + ); + assert_eq!(frame.models[0].gauges, RequestGauges::default()); + assert_eq!( + frame.models[0].totals, + RequestTotals { + requests_started: 1, + requests_cancelled: 1, + input_tokens: 8, + ..Default::default() + } + ); + } + + #[test] + fn frontend_gauges_settle_after_concurrent_request_transitions() { + let metrics = Arc::new(Metrics::new()); + for _ in 0..32 { + let guard = + metrics + .clone() + .create_inflight_guard("model", Endpoint::Completions, true, ""); + let request = guard.request_metrics().clone(); + std::thread::scope(|scope| { + scope.spawn(|| { + request.add_input_tokens(7); + }); + scope.spawn(|| { + request.add_output_tokens(1); + }); + scope.spawn(move || drop(guard)); + }); + } + assert_eq!( + metrics.get_request_counter( + "model", + &Endpoint::Completions, + &RequestType::Stream, + &Status::Error, + &ErrorType::Internal, + ), + 32 + ); + assert_eq!( + metrics + .awaiting_first_token_requests + .with_label_values(&["model"]) + .get(), + 0 + ); + assert_eq!( + metrics + .active_requests_gauge + .with_label_values(&["model"]) + .get(), + 0 + ); + for state in [STATE_AWAITING_FIRST_TOKEN, STATE_GENERATING] { + assert_eq!( + metrics + .inflight_input_tokens + .with_label_values(&["model", state]) + .get(), + 0 + ); + } + } + + #[test] + fn test_local_itl_histogram_is_initialized_lazily() { + let metrics = Arc::new(Metrics::new()); + let model = "lazy-local-itl-model"; + let global = metrics.inter_token_latency.with_label_values(&[model]); + let mut collector = response_collector(&metrics, model); + + assert!(collector.inter_token_latency.is_none()); + + collector.observe_response(10, 0); + assert!( + collector.inter_token_latency.is_none(), + "zero-token responses must not allocate a local histogram" + ); + + collector.observe_response(10, 1); + assert!( + collector.inter_token_latency.is_none(), + "the TTFT-only response must not allocate a local histogram" + ); + + collector.observe_response(10, 1); + assert!( + collector.inter_token_latency.is_some(), + "the first ITL observation must allocate a local histogram" + ); + assert_eq!( + global.get_sample_count(), + 0, + "the first ITL observation must remain local until flush or drop" + ); + + drop(collector); + assert_eq!(global.get_sample_count(), 1); + } + + #[test] + fn test_cached_handles_record_through_vec() { + // The collector caches per-model handles; observing through them must update + // the same metric the vec exposes. Regression guard for the handle cache + // (including the lazily resolved handle used by the per-token ITL loop). + let metrics = Arc::new(Metrics::new()); + let registry = prometheus::Registry::new(); + metrics.register(®istry).unwrap(); + let model = "cached-handle-model"; + + let mut collector = response_collector(&metrics, model); + // First chunk (3 tokens): TTFT + ISL observed once; no ITL yet. + collector.observe_response(42, 3); + // Second chunk (4 tokens): ITL observed once per token via the cached handle. + collector.observe_response(42, 4); + + // ITL observations are request-local until the batch threshold or drop. + assert_eq!( + metrics + .inter_token_latency + .with_label_values(&[model]) + .get_sample_count(), + 0 + ); + drop(collector); + + let ttft = metrics.time_to_first_token.with_label_values(&[model]); + assert_eq!(ttft.get_sample_count(), 1, "TTFT observed once"); + + let isl = metrics.input_sequence_length.with_label_values(&[model]); + assert_eq!(isl.get_sample_count(), 1, "ISL observed once"); + assert_eq!(isl.get_sample_sum(), 42.0); + + let itl = metrics.inter_token_latency.with_label_values(&[model]); + assert_eq!( + itl.get_sample_count(), + 4, + "ITL observed once per token of 2nd chunk" + ); + + let out = metrics + .output_tokens_counter + .with_label_values(&[model]) + .get(); + assert_eq!(out, 7, "output tokens = 3 + 4 via cached counter handle"); + } + + #[test] + fn test_local_itl_histogram_flushes_at_threshold_and_drop() { + use prometheus::core::Collector; + + let metrics = Arc::new(Metrics::new()); + let registry = prometheus::Registry::new(); + metrics.register(®istry).unwrap(); + let model = "local-itl-flush-model"; + let global = metrics.inter_token_latency.with_label_values(&[model]); + + let mut collector = response_collector(&metrics, model); + collector.observe_response(10, 1); // TTFT only. + assert!(collector.inter_token_latency.is_none()); + collector.observe_response(10, 63); + assert!(collector.inter_token_latency.is_some()); + assert_eq!( + global.get_sample_count(), + 0, + "a partial local batch must remain unpublished" + ); + + collector.observe_response(10, 1); + assert_eq!( + global.get_sample_count(), + ITL_LOCAL_FLUSH_TOKENS, + "the 64th pending ITL sample must flush the local histogram" + ); + + collector.observe_response(10, 7); + assert_eq!( + global.get_sample_count(), + ITL_LOCAL_FLUSH_TOKENS, + "the next partial batch must remain local" + ); + drop(collector); + + assert_eq!(global.get_sample_count(), 71); + assert!(global.get_sample_sum() > 0.0); + + let families = global.collect(); + let histogram = families[0].get_metric()[0].get_histogram(); + let last_bucket = histogram.get_bucket().last().unwrap(); + assert_eq!( + last_bucket.cumulative_count(), + global.get_sample_count(), + "flushing must update histogram buckets as well as count and sum" + ); + } + + #[test] + fn test_multimodal_counts_observed_once() { + // Counts are request-level constants carried on every chunk's annotation. + // The collector must latch them on the first observation: each histogram + // records exactly one per-request sample, and later (differing) values are + // ignored. The histogram `_sum` doubles as the cumulative volume. + let metrics = Arc::new(Metrics::new()); + let registry = prometheus::Registry::new(); + metrics.register(®istry).unwrap(); + let model = "mm-counts-model"; + + let mut collector = response_collector(&metrics, model); + // Repeated calls simulate the same counts arriving on each streamed chunk. + collector.observe_multimodal_metrics(2, 1, 0, Some(1290)); + collector.observe_multimodal_metrics(2, 1, 0, Some(1290)); + // A later, differing value must not override the latched counts. + collector.observe_multimodal_metrics(99, 99, 99, Some(9999)); + + let img = metrics.images_per_request.with_label_values(&[model]); + assert_eq!(img.get_sample_count(), 1, "image histogram observed once"); + assert_eq!(img.get_sample_sum(), 2.0, "image _sum == cumulative volume"); + let vid = metrics.videos_per_request.with_label_values(&[model]); + assert_eq!(vid.get_sample_count(), 1, "video histogram observed once"); + assert_eq!(vid.get_sample_sum(), 1.0); + let aud = metrics.audio_per_request.with_label_values(&[model]); + assert_eq!( + aud.get_sample_count(), + 1, + "audio histogram observes even a 0 (text-only distribution)" + ); + assert_eq!(aud.get_sample_sum(), 0.0); + let image_tokens = metrics.image_tokens_per_request.with_label_values(&[model]); + assert_eq!( + image_tokens.get_sample_count(), + 1, + "image-token histogram observed once" + ); + assert_eq!(image_tokens.get_sample_sum(), 1290.0); + } + + #[test] + fn test_multimodal_counts_text_only_zeros() { + // A text-only request carries zero counts; the histograms still get one + // sample each (at 0), so `_count` is 1 and `_sum` is 0. + let metrics = Arc::new(Metrics::new()); + let registry = prometheus::Registry::new(); + metrics.register(®istry).unwrap(); + let model = "text-only-model"; + + let mut collector = response_collector(&metrics, model); + collector.observe_multimodal_counts(0, 0, 0); + + let img = metrics.images_per_request.with_label_values(&[model]); + assert_eq!(img.get_sample_count(), 1); + assert_eq!(img.get_sample_sum(), 0.0); + let image_tokens = metrics.image_tokens_per_request.with_label_values(&[model]); + assert_eq!( + image_tokens.get_sample_count(), + 0, + "unavailable image-token count must not emit a sample" + ); + } + + #[test] + fn test_cached_decode_itl_gauge_records_for_worker() { + // The per-worker ITL gauge handle is resolved once (worker labels are latched + // at routing time via `set_worker_info`) and reused on every output token. + // Verify observations through the cached handle land on the same GaugeVec + // series the labels address. Regression guard for the per-token handle cache. + let metrics = Arc::new(Metrics::new()); + let registry = prometheus::Registry::new(); + metrics.register(®istry).ok(); + let model = "itl-gauge-model"; + + // Distinctive worker id so this never collides with other tests on the + // process-global WORKER_LAST_INTER_TOKEN_LATENCY_GAUGE. + let worker_id: u64 = 9_876_543_210; + let mut collector = response_collector(&metrics, model); + collector.set_worker_info( + None, + None, + None, + Some(worker_id), + Some(0), + Some("decode".to_string()), + ); + + // First chunk: TTFT only, no ITL yet (no prior response time) -> gauge unset. + collector.observe_response(10, 1); + // Elapse measurable time so ITL > 0, then a second chunk drives the cached + // per-worker ITL gauge. + std::thread::sleep(std::time::Duration::from_millis(2)); + collector.observe_response(10, 1); + + let worker_id_str = worker_id.to_string(); + let gauge = WORKER_LAST_INTER_TOKEN_LATENCY_GAUGE.with_label_values(&[ + worker_id_str.as_str(), + "0", + "decode", + ]); + assert!( + gauge.get() > 0.0, + "cached per-worker ITL gauge must record a positive ITL through the vec, got {}", + gauge.get() + ); + } + + #[test] + fn test_output_tokens_counter_zero_tokens() { + let metrics = Arc::new(Metrics::new()); + let registry = prometheus::Registry::new(); + metrics.register(®istry).unwrap(); + + let model = "test-model"; + let mut collector = response_collector(&metrics, model); + + // Simulate chunk with zero tokens (should not increment) + collector.observe_response(100, 0); + + // Verify counter remains 0 + let counter_value = metrics + .output_tokens_counter + .with_label_values(&[model]) + .get(); + assert_eq!(counter_value, 0); + + // Add some tokens + collector.observe_response(100, 5); + assert_eq!( + metrics + .output_tokens_counter + .with_label_values(&[model]) + .get(), + 5 + ); + + // Try zero tokens again (should not change counter) + collector.observe_response(100, 0); + assert_eq!( + metrics + .output_tokens_counter + .with_label_values(&[model]) + .get(), + 5 + ); + } + + #[test] + fn test_output_tokens_counter_multiple_models() { + let metrics = Arc::new(Metrics::new()); + let registry = prometheus::Registry::new(); + metrics.register(®istry).unwrap(); + + let model1 = "model-1"; + let model2 = "model-2"; + + // Create collectors for different models + let mut collector1 = response_collector(&metrics, model1); + let mut collector2 = response_collector(&metrics, model2); + + // Increment model1 + collector1.observe_response(100, 10); + assert_eq!( + metrics + .output_tokens_counter + .with_label_values(&[model1]) + .get(), + 10 + ); + assert_eq!( + metrics + .output_tokens_counter + .with_label_values(&[model2]) + .get(), + 0 + ); + + // Increment model2 + collector2.observe_response(200, 20); + assert_eq!( + metrics + .output_tokens_counter + .with_label_values(&[model1]) + .get(), + 10 + ); + assert_eq!( + metrics + .output_tokens_counter + .with_label_values(&[model2]) + .get(), + 20 + ); + + // Increment model1 again + collector1.observe_response(100, 5); + assert_eq!( + metrics + .output_tokens_counter + .with_label_values(&[model1]) + .get(), + 15 + ); + assert_eq!( + metrics + .output_tokens_counter + .with_label_values(&[model2]) + .get(), + 20 + ); + } + + #[test] + fn test_cached_tokens_once_per_request() { + let metrics = Arc::new(Metrics::new()); + let registry = prometheus::Registry::new(); + metrics.register(®istry).unwrap(); + + let model = "test-model"; + let expected_metric_name = "dynamo_frontend_cached_tokens"; + let mut collector = response_collector(&metrics, model); + + // Create histogram handle first + let _histogram = metrics.cached_tokens.with_label_values(&[model]); + + // First call should observe and record 1 sample + collector.observe_cached_tokens(Some(100)); + let metric_families = registry.gather(); + let histogram_family = metric_families + .iter() + .find(|mf| mf.name() == expected_metric_name) + .expect("histogram should be registered"); + assert_eq!( + histogram_family.get_metric()[0] + .get_histogram() + .get_sample_count(), + 1 + ); + + // Second call with same collector should not observe again (idempotent) + collector.observe_cached_tokens(Some(50)); + let metric_families = registry.gather(); + let histogram_family = metric_families + .iter() + .find(|mf| mf.name() == expected_metric_name) + .expect("histogram should be registered"); + assert_eq!( + histogram_family.get_metric()[0] + .get_histogram() + .get_sample_count(), + 1 + ); + + // Third call with different value should still be idempotent + collector.observe_cached_tokens(Some(75)); + let metric_families = registry.gather(); + let histogram_family = metric_families + .iter() + .find(|mf| mf.name() == expected_metric_name) + .expect("histogram should be registered"); + assert_eq!( + histogram_family.get_metric()[0] + .get_histogram() + .get_sample_count(), + 1 + ); + } + + #[test] + fn test_error_type_as_str() { + assert_eq!(ErrorType::None.as_str(), ""); + assert_eq!(ErrorType::Validation.as_str(), "validation"); + assert_eq!(ErrorType::NotFound.as_str(), "not_found"); + assert_eq!(ErrorType::Overload.as_str(), "overload"); + assert_eq!(ErrorType::Unavailable.as_str(), "unavailable"); + assert_eq!(ErrorType::Cancelled.as_str(), "cancelled"); + assert_eq!(ErrorType::ResponseTimeout.as_str(), "response_timeout"); + assert_eq!(ErrorType::Internal.as_str(), "internal"); + assert_eq!(ErrorType::NotImplemented.as_str(), "not_implemented"); + } + + #[test] + fn test_inflight_guard_marks_success_with_correct_error_type() { + let metrics = Arc::new(Metrics::new()); + let registry = prometheus::Registry::new(); + metrics.register(®istry).unwrap(); + + let model = "test-model"; + + { + let mut guard = + metrics + .clone() + .create_inflight_guard(model, Endpoint::ChatCompletions, false, ""); + guard.mark_ok(); + } // guard drops here + + // Verify counter incremented with status=success, error_type="" + let counter_value = metrics + .request_counter + .with_label_values(&[ + model, + Endpoint::ChatCompletions.as_str(), + RequestType::Unary.as_str(), + Status::Success.as_str(), + ErrorType::None.as_str(), + ]) + .get(); + assert_eq!(counter_value, 1); + + let started_counter_value = metrics + .request_started_counter + .with_label_values(&[ + model, + Endpoint::ChatCompletions.as_str(), + RequestType::Unary.as_str(), + ]) + .get(); + assert_eq!(started_counter_value, 1); + } + + #[test] + fn test_inflight_guard_marks_validation_error() { + let metrics = Arc::new(Metrics::new()); + let registry = prometheus::Registry::new(); + metrics.register(®istry).unwrap(); + + let model = "test-model"; + + { + let mut guard = + metrics + .clone() + .create_inflight_guard(model, Endpoint::ChatCompletions, false, ""); + guard.mark_error(ErrorType::Validation); + } // guard drops here + + // Verify counter incremented with status=error, error_type=validation + let counter_value = metrics + .request_counter + .with_label_values(&[ + model, + Endpoint::ChatCompletions.as_str(), + RequestType::Unary.as_str(), + Status::Error.as_str(), + ErrorType::Validation.as_str(), + ]) + .get(); + assert_eq!(counter_value, 1); + } + + #[test] + fn test_inflight_guard_classifies_cancellation_separately_for_tracing() { + let metrics = Arc::new(Metrics::new()); + let registry = prometheus::Registry::new(); + metrics.register(®istry).unwrap(); + + let model = "test-model"; + let mut guard = metrics.clone().create_inflight_guard( + model, + Endpoint::ChatCompletions, + true, + "cancelled-request", + ); + guard.mark_error(ErrorType::Cancelled); + + assert_eq!(guard.request_completion(), RequestCompletion::Cancelled); + drop(guard); + + // Keep the existing Prometheus contract while tracing reports cancellation + // as an expected request outcome rather than a span error. + let counter_value = metrics + .request_counter + .with_label_values(&[ + model, + Endpoint::ChatCompletions.as_str(), + RequestType::Stream.as_str(), + Status::Error.as_str(), + ErrorType::Cancelled.as_str(), + ]) + .get(); + assert_eq!(counter_value, 1); + } + + #[test] + fn test_inflight_guard_defaults_to_internal_error_on_drop() { + let metrics = Arc::new(Metrics::new()); + let registry = prometheus::Registry::new(); + metrics.register(®istry).unwrap(); + + let model = "test-model"; + + { + let _guard = + metrics + .clone() + .create_inflight_guard(model, Endpoint::ChatCompletions, false, ""); + // Don't call mark_ok() or mark_error() - simulate panic/unhandled error + } // guard drops with default error_type=Internal + + // Verify counter incremented with status=error, error_type=internal + let counter_value = metrics + .request_counter + .with_label_values(&[ + model, + Endpoint::ChatCompletions.as_str(), + RequestType::Unary.as_str(), + Status::Error.as_str(), + ErrorType::Internal.as_str(), + ]) + .get(); + assert_eq!(counter_value, 1); + } + + #[test] + fn test_active_requests_tracks_inflight_guard() { + let metrics = Arc::new(Metrics::new()); + let registry = prometheus::Registry::new(); + metrics.register(®istry).unwrap(); + + let model = "test-model-active"; + + // Both gauges start at 0 + assert_eq!(metrics.inflight_gauge.with_label_values(&[model]).get(), 0); + assert_eq!( + metrics + .active_requests_gauge + .with_label_values(&[model]) + .get(), + 0 + ); + + { + let _guard = metrics.clone().create_inflight_guard( + model, + Endpoint::ChatCompletions, + true, + "req-1", + ); + + // Both gauges increment together + assert_eq!(metrics.inflight_gauge.with_label_values(&[model]).get(), 1); + assert_eq!( + metrics + .active_requests_gauge + .with_label_values(&[model]) + .get(), + 1 + ); + + { + let _guard2 = metrics.clone().create_inflight_guard( + model, + Endpoint::ChatCompletions, + true, + "req-2", + ); + assert_eq!(metrics.inflight_gauge.with_label_values(&[model]).get(), 2); + assert_eq!( + metrics + .active_requests_gauge + .with_label_values(&[model]) + .get(), + 2 + ); + } + // guard2 dropped + assert_eq!(metrics.inflight_gauge.with_label_values(&[model]).get(), 1); + assert_eq!( + metrics + .active_requests_gauge + .with_label_values(&[model]) + .get(), + 1 + ); + } + // guard dropped — both back to 0 + assert_eq!(metrics.inflight_gauge.with_label_values(&[model]).get(), 0); + assert_eq!( + metrics + .active_requests_gauge + .with_label_values(&[model]) + .get(), + 0 + ); + } + + #[test] + fn test_lifecycle_constructors_preserve_model_label_casing() { + let metrics = Arc::new(Metrics::new()); + let registry = prometheus::Registry::new(); + metrics.register(®istry).unwrap(); + + let model = "Llama-3.1-8B-Instruct"; + + let _inflight = metrics.clone().create_inflight_guard( + model, + Endpoint::ChatCompletions, + true, + "req-case", + ); + let mut collector = _inflight.response_collector(); + collector.observe_response(100, 1); + let _queue = metrics.clone().create_http_queue_guard(model); + + // Drop the lifecycle guards so their Drop impls record + // completion-time metrics (request_counter, request_duration, + // detokenize observations) before we gather. Without this the + // regression only covers constructor-time labels. + drop(_inflight); + drop(collector); + drop(_queue); + + let metric_families = registry.gather(); + let observed: Vec = metric_families + .iter() + .flat_map(|mf| mf.get_metric()) + .flat_map(|m| m.get_label()) + .filter(|l| l.name() == "model" && l.value() != "unknown_model") + .map(|l| l.value().to_string()) + .collect(); + + assert!( + !observed.is_empty(), + "expected at least one metric to carry a model label" + ); + for value in &observed { + assert_eq!( + value, model, + "model label was modified; expected original casing to be preserved" + ); + } + } + + #[test] + fn test_all_error_types_recorded_correctly() { + let metrics = Arc::new(Metrics::new()); + let registry = prometheus::Registry::new(); + metrics.register(®istry).unwrap(); + + let model = "test-model"; + let endpoint = Endpoint::ChatCompletions; + + // Test each error type + let error_types = vec![ + ErrorType::Validation, + ErrorType::NotFound, + ErrorType::Overload, + ErrorType::Unavailable, + ErrorType::Cancelled, + ErrorType::ResponseTimeout, + ErrorType::Internal, + ErrorType::NotImplemented, + ]; + + for error_type in &error_types { + let mut guard = metrics + .clone() + .create_inflight_guard(model, endpoint, false, ""); + guard.mark_error(error_type.clone()); + drop(guard); + } + + // Verify each error type recorded correctly + for error_type in &error_types { + let counter_value = metrics + .request_counter + .with_label_values(&[ + model, + endpoint.as_str(), + RequestType::Unary.as_str(), + Status::Error.as_str(), + error_type.as_str(), + ]) + .get(); + assert_eq!( + counter_value, + 1, + "Should have 1 request for error_type={}", + error_type.as_str() + ); + } + } + + #[test] + fn semantic_failure_metric_honors_custom_prefix() { + let registry = prometheus::Registry::new(); + let counter = new_failure_counter(Some("nv_llm_http_service")); + counter + .with_label_values(&["Internal", "runtime.internal"]) + .inc(); + registry.register(Box::new(counter)).unwrap(); + + assert!( + registry + .gather() + .iter() + .any(|family| family.name() == "nv_llm_http_service_failures_total") + ); + } + + #[test] + fn semantic_failure_metric_has_only_class_and_reason_labels() { + use dynamo_runtime::error::{DynamoError, ErrorClass, ErrorReason}; + + let registry = prometheus::Registry::new(); + let counter = new_failure_counter(None); + registry.register(Box::new(counter.clone())).unwrap(); + let error = DynamoError::builder() + .class(ErrorClass::InvalidRequest) + .reason(ErrorReason::new("request.invalid").unwrap()) + .build(); + + record_failure_into(&counter, &error); + + let family = registry + .gather() + .into_iter() + .find(|family| family.name() == "dynamo_frontend_failures_total") + .expect("dynamo_frontend_failures_total metric"); + let metric = family + .get_metric() + .iter() + .find(|metric| { + metric + .get_label() + .iter() + .any(|label| label.name() == "reason" && label.value() == "request.invalid") + }) + .expect("request.invalid sample"); + let labels: Vec<_> = metric + .get_label() + .iter() + .map(|label| (label.name(), label.value())) + .collect(); + assert_eq!( + labels, + vec![("class", "InvalidRequest"), ("reason", "request.invalid")] + ); + assert_eq!(metric.get_counter().value(), 1.0); + } + + #[test] + fn test_multiple_requests_different_error_types() { + let metrics = Arc::new(Metrics::new()); + let registry = prometheus::Registry::new(); + metrics.register(®istry).unwrap(); + + let model = "test-model"; + + // Record 2 validation errors, 3 internal errors, 1 success + for _ in 0..2 { + let mut guard = + metrics + .clone() + .create_inflight_guard(model, Endpoint::ChatCompletions, false, ""); + guard.mark_error(ErrorType::Validation); + drop(guard); + } + + for _ in 0..3 { + let mut guard = + metrics + .clone() + .create_inflight_guard(model, Endpoint::Completions, false, ""); + guard.mark_error(ErrorType::Internal); + drop(guard); + } + + { + let mut guard = + metrics + .clone() + .create_inflight_guard(model, Endpoint::Embeddings, false, ""); + guard.mark_ok(); + drop(guard); + } + + // Check validation errors (2 from ChatCompletions) + let validation_count = metrics + .request_counter + .with_label_values(&[ + model, + Endpoint::ChatCompletions.as_str(), + RequestType::Unary.as_str(), + Status::Error.as_str(), + ErrorType::Validation.as_str(), + ]) + .get(); + assert_eq!(validation_count, 2); + + // Check internal errors (3 from Completions) + let internal_count = metrics + .request_counter + .with_label_values(&[ + model, + Endpoint::Completions.as_str(), + RequestType::Unary.as_str(), + Status::Error.as_str(), + ErrorType::Internal.as_str(), + ]) + .get(); + assert_eq!(internal_count, 3); + + // Check success (1 from Embeddings) + let success_count = metrics + .request_counter + .with_label_values(&[ + model, + Endpoint::Embeddings.as_str(), + RequestType::Unary.as_str(), + Status::Success.as_str(), + ErrorType::None.as_str(), + ]) + .get(); + assert_eq!(success_count, 1); + } + + #[test] + fn test_observe_embedding_latency_records_observation() { + let metrics = Arc::new(Metrics::new()); + let registry = prometheus::Registry::new(); + metrics.register(®istry).unwrap(); + + let model = "embed-test"; + assert_eq!(metrics.embedding_latency_count(model), 0); + assert_eq!(metrics.embedding_latency_sum(model), 0.0); + + metrics.observe_embedding_latency(model, 0.005); + metrics.observe_embedding_latency(model, 0.020); + + assert_eq!(metrics.embedding_latency_count(model), 2); + assert!((metrics.embedding_latency_sum(model) - 0.025).abs() < 1e-9); + } + + #[test] + fn test_observe_embedding_latency_drops_invalid_values() { + let metrics = Arc::new(Metrics::new()); + let registry = prometheus::Registry::new(); + metrics.register(®istry).unwrap(); + + let model = "embed-test"; + + metrics.observe_embedding_latency(model, -0.5); + metrics.observe_embedding_latency(model, f64::NAN); + metrics.observe_embedding_latency(model, f64::INFINITY); + metrics.observe_embedding_latency(model, f64::NEG_INFINITY); + + assert_eq!(metrics.embedding_latency_count(model), 0); + assert_eq!(metrics.embedding_latency_sum(model), 0.0); + + // Zero is allowed (boundary case for finite, non-negative). + metrics.observe_embedding_latency(model, 0.0); + assert_eq!(metrics.embedding_latency_count(model), 1); + } + + #[test] + fn test_observe_embedding_latency_partitions_by_model() { + let metrics = Arc::new(Metrics::new()); + let registry = prometheus::Registry::new(); + metrics.register(®istry).unwrap(); + + let m1 = "embed-a"; + let m2 = "embed-b"; + + metrics.observe_embedding_latency(m1, 0.010); + metrics.observe_embedding_latency(m1, 0.030); + metrics.observe_embedding_latency(m2, 0.005); + + assert_eq!(metrics.embedding_latency_count(m1), 2); + assert_eq!(metrics.embedding_latency_count(m2), 1); + assert!((metrics.embedding_latency_sum(m1) - 0.040).abs() < 1e-9); + assert!((metrics.embedding_latency_sum(m2) - 0.005).abs() < 1e-9); + } + + #[test] + fn test_embedding_latency_registered_in_registry() { + let metrics = Arc::new(Metrics::new()); + let registry = prometheus::Registry::new(); + metrics.register(®istry).unwrap(); + + metrics.observe_embedding_latency("embed-test", 0.012); + + let metric_families = registry.gather(); + let found = metric_families.iter().any(|mf| { + mf.name() + .ends_with(frontend_service::EMBEDDING_LATENCY_SECONDS) + }); + assert!( + found, + "embedding_latency_seconds histogram must be registered with the registry" + ); + } + + #[test] + fn test_worker_metadata_is_latched_without_replacement() { + let metrics = Arc::new(Metrics::new()); + let mut collector = response_collector(&metrics, "worker-latch-model"); + let first = ( + Some(11), + Some(1), + Some("prefill-first".to_string()), + Some(22), + Some(2), + Some("decode-first".to_string()), + ); + collector.set_worker_info_from_parts( + first.0, + first.1, + first.2.as_deref(), + first.3, + first.4, + first.5.as_deref(), + ); + + let later = ( + Some(33), + Some(3), + Some("prefill-later".to_string()), + Some(44), + Some(4), + Some("decode-later".to_string()), + ); + collector.set_worker_info_from_parts( + later.0, + later.1, + later.2.as_deref(), + later.3, + later.4, + later.5.as_deref(), + ); + + assert_eq!(collector.prefill_worker_id, Some(11)); + assert_eq!(collector.prefill_dp_rank, Some(1)); + assert_eq!( + collector.prefill_worker_type.as_deref(), + Some("prefill-first") + ); + assert_eq!(collector.decode_worker_id, Some(22)); + assert_eq!(collector.decode_dp_rank, Some(2)); + assert_eq!( + collector.decode_worker_type.as_deref(), + Some("decode-first") + ); + } +} diff --git a/lib/http-server/src/request.rs b/lib/http-server/src/request.rs new file mode 100644 index 000000000000..50d8646eefe1 --- /dev/null +++ b/lib/http-server/src/request.rs @@ -0,0 +1,44 @@ +// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use dynamo_runtime::config::environment_names::llm::DYN_HTTP_BODY_LIMIT_MB; + +// Default axum max body limit without configuring is 2MB: https://docs.rs/axum/latest/axum/extract/struct.DefaultBodyLimit.html +/// Default body limit in bytes (45MB) to support 500k+ token payloads. +/// Can be configured at runtime using the DYN_HTTP_BODY_LIMIT_MB environment variable. +pub fn get_body_limit() -> usize { + std::env::var(DYN_HTTP_BODY_LIMIT_MB) + .ok() + .and_then(|s| s.parse::().ok()) + .map(|mb| mb * 1024 * 1024) + .unwrap_or(45 * 1024 * 1024) +} + +pub fn is_json_content_type(content_type: &str) -> bool { + let media_type = content_type.split(';').next().unwrap_or_default().trim(); + let Some((media_type, subtype)) = media_type.split_once('/') else { + return false; + }; + + media_type.eq_ignore_ascii_case("application") + && (subtype.eq_ignore_ascii_case("json") + || subtype + .to_ascii_lowercase() + .rsplit_once('+') + .is_some_and(|(_, suffix)| suffix == "json")) +} +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_is_json_content_type() { + assert!(is_json_content_type("application/json")); + assert!(is_json_content_type("application/json; charset=utf-8")); + assert!(is_json_content_type("Application/JSON")); + assert!(is_json_content_type("application/vnd.dynamo+json")); + assert!(!is_json_content_type("text/plain")); + assert!(!is_json_content_type("application/json-patch")); + assert!(!is_json_content_type("application")); + } +} diff --git a/lib/http-server/src/server.rs b/lib/http-server/src/server.rs new file mode 100644 index 000000000000..216b3aaf8ee8 --- /dev/null +++ b/lib/http-server/src/server.rs @@ -0,0 +1,398 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! The frontend's HTTP(S) accept loop, with graceful shutdown bounded by a grace period. +//! +//! A connection task is polled on every wakeup, i.e. for every streamed response chunk, and +//! each poll also polls its shutdown signals. Every connection waits on its own child +//! [`CancellationToken`]s, whose lock and `Notify` no other connection touches. A shared +//! `tokio::sync::Notify` would lock its waiter list on every poll of a pending waiter. + +use crate::lifecycle::ServiceObserver; +use anyhow::{Context, Result}; +use dynamo_runtime::config::environment_names::llm as env_llm; +use std::net::{IpAddr, SocketAddr}; +use std::path::PathBuf; +use std::sync::Arc; + +use std::{future::Future, pin::pin, time::Duration}; + +use axum::Router; +use hyper_util::rt::{TokioExecutor, TokioIo}; +use hyper_util::server::conn::auto::Builder; +use hyper_util::service::TowerToHyperService; +use tokio::io::{AsyncRead, AsyncWrite}; +use tokio::net::{TcpListener, TcpStream}; +use tokio_rustls::TlsAcceptor; +use tokio_util::sync::CancellationToken; +use tokio_util::task::TaskTracker; + +use crate::write_timeout::{ConnectionDeadline, ConnectionService, WriteTimeoutIo}; + +/// The handshake timeout of axum-server's rustls acceptor, which served the frontend before. +const TLS_HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(10); + +/// Serves `router` on `listener` until `shutdown` resolves with a grace period. +/// +/// Then it stops accepting connections and lets open ones finish their requests; any still +/// open after the grace period are closed, dropping their response bodies. Returning, or +/// being dropped, closes every connection that was not upgraded. An upgraded connection, +/// such as a WebSocket session, runs on its own task and stays open after `serve` returns. +/// +/// Every connection is wrapped in [`WriteTimeoutIo`], and every request on it carries the +/// connection's [`ConnectionDeadline`]. +pub async fn serve( + listener: TcpListener, + tls: Option, + router: Router, + write_idle_timeout: Duration, + shutdown: impl Future, +) { + // Build each route's handler once, as `Router::into_make_service` does, not per request. + let router = router.with_state(()); + let graceful = CancellationToken::new(); + let hard = CancellationToken::new(); + let _close_connections = hard.clone().drop_guard(); + let connections = TaskTracker::new(); + + let mut shutdown = pin!(shutdown); + let grace_period = loop { + let stream = tokio::select! { + biased; + grace_period = &mut shutdown => break grace_period, + stream = accept(&listener) => stream, + }; + let connection = serve_connection( + stream, + tls.clone(), + router.clone(), + write_idle_timeout, + graceful.child_token(), + ); + connections.spawn(hard.child_token().run_until_cancelled_owned(connection)); + }; + + // Refuse new connections at once instead of leaving them in the backlog. + drop(listener); + graceful.cancel(); + connections.close(); + let _ = tokio::time::timeout(grace_period, connections.wait()).await; +} + +async fn accept(listener: &TcpListener) -> TcpStream { + loop { + match listener.accept().await { + Ok((stream, _)) => { + // Streamed tokens are small writes. With Nagle's algorithm on, one can wait + // for the client's delayed ACK of the previous write, 40 ms or more. Set on + // the TCP socket, this also covers TLS connections. + if let Err(error) = stream.set_nodelay(true) { + tracing::debug!(%error, "Failed to set TCP_NODELAY"); + } + return stream; + } + // Errors such as EMFILE are transient: back off and retry, as axum-server did. + Err(error) => { + tracing::debug!(%error, "Failed to accept HTTP connection"); + tokio::time::sleep(Duration::from_millis(50)).await; + } + } + } +} + +async fn serve_connection( + stream: TcpStream, + tls: Option, + router: Router, + write_idle_timeout: Duration, + graceful: CancellationToken, +) { + let Some(tls) = tls else { + return serve_io(stream, router, write_idle_timeout, graceful).await; + }; + match tokio::time::timeout(TLS_HANDSHAKE_TIMEOUT, tls.accept(stream)).await { + Ok(Ok(stream)) => serve_io(stream, router, write_idle_timeout, graceful).await, + Ok(Err(error)) => tracing::debug!(%error, "TLS handshake failed"), + Err(_) => tracing::debug!("TLS handshake timed out"), + } +} + +async fn serve_io( + io: I, + router: Router, + write_idle_timeout: Duration, + graceful: CancellationToken, +) where + I: AsyncRead + AsyncWrite + Unpin + Send + 'static, +{ + let deadline = ConnectionDeadline::default(); + let io = TokioIo::new(WriteTimeoutIo::new( + io, + write_idle_timeout, + deadline.clone(), + )); + let service = TowerToHyperService::new(ConnectionService::new(router, deadline)); + let builder = Builder::new(TokioExecutor::new()); + let mut connection = pin!(builder.serve_connection_with_upgrades(io, service)); + if graceful + .run_until_cancelled(connection.as_mut()) + .await + .is_none() + { + connection.as_mut().graceful_shutdown(); + let _ = connection.await; + } +} + +pub struct HttpServerConfig { + pub host: String, + pub port: u16, + pub enable_tls: bool, + pub tls_cert_path: Option, + pub tls_key_path: Option, + pub tls_client_ca_cert_path: Option, +} + +pub struct HttpServer { + router: Router, + listener: TcpListener, + tls: Option, + protocol: &'static str, +} + +impl HttpServer { + /// Bind and prepare TLS before the caller starts background tasks. + pub async fn prepare( + router: Router, + config: HttpServerConfig, + listener: Option, + ) -> Result { + let address = format!("{}:{}", config.host, config.port); + let protocol = if config.enable_tls { "HTTPS" } else { "HTTP" }; + tracing::info!(protocol, address, "Starting HTTP(S) service"); + + let tls = if config.enable_tls { + let cert_path = config + .tls_cert_path + .as_ref() + .ok_or_else(|| anyhow::anyhow!("TLS certificate path not provided"))?; + let key_path = config + .tls_key_path + .as_ref() + .ok_or_else(|| anyhow::anyhow!("TLS private key path not provided"))?; + let mut server_config = dynamo_runtime::tls_utils::server_tls_config( + cert_path, + key_path, + config.tls_client_ca_cert_path.as_deref(), + ) + .context("Failed to create TLS config")?; + server_config.alpn_protocols = vec![b"h2".to_vec(), b"http/1.1".to_vec()]; + Some(tokio_rustls::TlsAcceptor::from(Arc::new(server_config))) + } else { + None + }; + + let listener = match listener { + Some(l) => l, + None => { + let addr = listen_socket_addr(&config.host, config.port)?; + bind_listener(addr).map_err(|e| { + tracing::error!( + protocol = %protocol, + address = %address, + error = %e, + "Failed to bind server to address" + ); + match e.kind() { + std::io::ErrorKind::AddrInUse => anyhow::anyhow!( + "Failed to start {} server: port {} already in use. Use --http-port to specify a different port.", + protocol, + config.port + ), + _ => anyhow::anyhow!( + "Failed to start {} server on {}: {}", + protocol, + address, + e + ), + } + })? + } + }; + tracing::info!(address = %listener.local_addr()?, "{protocol} server listening"); + + Ok(Self { + router, + listener, + tls, + protocol, + }) + } + + /// Serve while admitted response bodies drain before cancelling service state. + pub async fn serve( + self, + state: Arc, + state_cancel: CancellationToken, + cancel_token: CancellationToken, + ) -> Result<()> { + let observer = cancel_token.child_token(); + let protocol = self.protocol; + // The server keeps accepting while inflight requests drain, then gives open + // connections what is left of the shutdown budget. + let shutdown = async { + observer.cancelled().await; + state.start_draining(); + tracing::info!("{protocol} server shutdown requested"); + let shutdown_timeout = Duration::from_secs(get_graceful_shutdown_timeout() as u64); + let drain_started = tokio::time::Instant::now(); + if !state.wait_inflight_zero_or_timeout(shutdown_timeout).await { + tracing::warn!( + inflight_requests = state.inflight_count(), + "Timed out waiting for inflight inference requests to drain" + ); + } + shutdown_timeout.saturating_sub(drain_started.elapsed()) + }; + serve( + self.listener, + self.tls, + self.router, + write_idle_timeout(), + shutdown, + ) + .await; + state.start_stopping(); + state_cancel.cancel(); + Ok(()) + } +} + +fn get_graceful_shutdown_timeout() -> usize { + std::env::var(env_llm::DYN_HTTP_GRACEFUL_SHUTDOWN_TIMEOUT_SECS) + .ok() + .and_then(|s| s.parse::().ok()) + .unwrap_or(5) +} + +fn write_idle_timeout() -> Duration { + std::env::var(env_llm::DYN_HTTP_WRITE_IDLE_TIMEOUT_SECS) + .ok() + .and_then(|value| value.parse::().ok()) + .filter(|value| *value > 0) + .map(Duration::from_secs) + .filter(|duration| tokio::time::Instant::now().checked_add(*duration).is_some()) + .unwrap_or(Duration::from_secs(120)) +} + +const DEFAULT_LISTEN_BACKLOG: u32 = 4096; + +// `listen(2)` takes an `int`; anything above `i32::MAX` would go negative. +fn parse_listen_backlog(value: Result) -> u32 { + value + .ok() + .and_then(|s| s.trim().parse::().ok()) + .filter(|n| (1..=i32::MAX as u32).contains(n)) + .unwrap_or(DEFAULT_LISTEN_BACKLOG) +} + +fn listen_backlog() -> u32 { + parse_listen_backlog(std::env::var(env_llm::DYN_HTTP_LISTEN_BACKLOG)) +} + +/// Build the listen address from an IP literal host. IPv6 brackets are optional, +/// so `::` and `[::]` are equivalent. +pub fn listen_socket_addr(host: &str, port: u16) -> Result { + let literal = host.trim(); + let literal = literal + .strip_prefix('[') + .and_then(|inner| inner.strip_suffix(']')) + .unwrap_or(literal); + let ip: IpAddr = literal + .parse() + .map_err(|e| anyhow::anyhow!("Invalid address '{host}:{port}': {e}"))?; + Ok(SocketAddr::new(ip, port)) +} + +/// `tokio::net::TcpListener::bind` listens with a backlog of 128. A few thousand +/// clients connecting within seconds overflow that, and overflowed connections +/// are delayed or, depending on host TCP settings, fail. +fn bind_listener(addr: SocketAddr) -> std::io::Result { + let socket = if addr.is_ipv4() { + tokio::net::TcpSocket::new_v4()? + } else { + tokio::net::TcpSocket::new_v6()? + }; + socket.set_reuseaddr(true)?; + socket.bind(addr)?; + socket.listen(listen_backlog()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn accepted_connections_turn_nagle_off() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let (_client, stream) = tokio::join!(TcpStream::connect(addr), accept(&listener)); + assert!(stream.nodelay().unwrap()); + } + #[test] + fn test_listen_backlog_env_var() { + assert_eq!( + parse_listen_backlog(Err(std::env::VarError::NotPresent)), + DEFAULT_LISTEN_BACKLOG + ); + assert_eq!( + parse_listen_backlog(Ok("0".to_string())), + DEFAULT_LISTEN_BACKLOG + ); + assert_eq!( + parse_listen_backlog(Ok("invalid".to_string())), + DEFAULT_LISTEN_BACKLOG + ); + assert_eq!( + parse_listen_backlog(Ok((i32::MAX as u32 + 1).to_string())), + DEFAULT_LISTEN_BACKLOG + ); + assert_eq!( + parse_listen_backlog(Ok(i32::MAX.to_string())), + i32::MAX as u32 + ); + assert_eq!(parse_listen_backlog(Ok(" 8192 ".to_string())), 8192); + } + #[tokio::test] + async fn test_bind_listener_accepts_connections() { + let listener = bind_listener("127.0.0.1:0".parse().unwrap()).unwrap(); + let addr = listener.local_addr().unwrap(); + assert_ne!(addr.port(), 0); + let client = tokio::net::TcpStream::connect(addr).await.unwrap(); + let (_server_side, peer) = listener.accept().await.unwrap(); + assert_eq!(peer, client.local_addr().unwrap()); + } + #[test] + fn test_listen_socket_addr_accepts_ip_literals() { + for (host, expected) in [ + ("0.0.0.0", "0.0.0.0:8000"), + ("127.0.0.1", "127.0.0.1:8000"), + ("::", "[::]:8000"), + ("[::]", "[::]:8000"), + (" [::1] ", "[::1]:8000"), + ] { + assert_eq!( + listen_socket_addr(host, 8000).unwrap(), + expected.parse::().unwrap(), + "{host}" + ); + } + for host in ["localhost", "[::", "::1]", ""] { + let error = listen_socket_addr(host, 8000).unwrap_err().to_string(); + assert!( + error.starts_with(&format!("Invalid address '{host}:8000'")), + "{host}: {error}" + ); + } + } +} diff --git a/lib/http-server/src/write_timeout.rs b/lib/http-server/src/write_timeout.rs new file mode 100644 index 000000000000..adf929b3ca06 --- /dev/null +++ b/lib/http-server/src/write_timeout.rs @@ -0,0 +1,381 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Connection-level write and batch-transfer bounds, independent of body polling. + +use std::sync::{ + Arc, + atomic::{AtomicBool, Ordering}, +}; +use std::task::{Context, Poll}; +use std::time::Duration; +use std::{collections::BTreeMap, future::Future, io, pin::Pin}; + +use futures::task::AtomicWaker; +use parking_lot::Mutex; +use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; +use tokio::time::Instant; +use tokio::time::Sleep; + +#[derive(Clone, Default)] +pub struct ConnectionDeadline(Arc); + +#[derive(Default)] +struct Deadlines { + transfers: Mutex>, + is_active: AtomicBool, + waker: AtomicWaker, +} + +impl ConnectionDeadline { + pub fn start(&self, timeout: Duration) -> TransferDeadline { + let now = Instant::now(); + let deadline = now + .checked_add(timeout) + .unwrap_or_else(|| now + Duration::from_secs(600)); + *self.0.transfers.lock().entry(deadline).or_default() += 1; + self.0.is_active.store(true, Ordering::Release); + self.0.waker.wake(); + TransferDeadline { + connection: self.clone(), + deadline, + } + } + + fn next(&self, cx: &mut Context<'_>) -> Option { + self.0.waker.register(cx.waker()); + if !self.0.is_active.load(Ordering::Acquire) { + return None; + } + self.0 + .transfers + .lock() + .first_key_value() + .map(|(deadline, _)| *deadline) + } +} + +pub struct TransferDeadline { + connection: ConnectionDeadline, + deadline: Instant, +} + +impl Drop for TransferDeadline { + fn drop(&mut self) { + let mut transfers = self.connection.0.transfers.lock(); + if let Some(count) = transfers.get_mut(&self.deadline) { + *count -= 1; + if *count == 0 { + transfers.remove(&self.deadline); + } + } + self.connection + .0 + .is_active + .store(!transfers.is_empty(), Ordering::Release); + drop(transfers); + self.connection.0.waker.wake(); + } +} + +#[derive(Clone)] +pub(crate) struct ConnectionService { + inner: S, + deadline: ConnectionDeadline, +} + +impl ConnectionService { + pub(crate) fn new(inner: S, deadline: ConnectionDeadline) -> Self { + Self { inner, deadline } + } +} + +impl tower::Service> for ConnectionService +where + S: tower::Service>, +{ + type Response = S::Response; + type Error = S::Error; + type Future = S::Future; + + fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { + self.inner.poll_ready(cx) + } + + fn call(&mut self, mut request: axum::http::Request) -> Self::Future { + request.extensions_mut().insert(self.deadline.clone()); + self.inner.call(request) + } +} + +pub(crate) struct WriteTimeoutIo { + inner: I, + timeout: Duration, + stalled: Option>>, + is_write_pending: bool, + expired: bool, + connection_deadline: ConnectionDeadline, + transfer_timer: Option>>, +} + +impl WriteTimeoutIo { + pub(crate) fn new( + inner: I, + timeout: Duration, + connection_deadline: ConnectionDeadline, + ) -> Self { + Self { + inner, + timeout, + stalled: None, + is_write_pending: false, + expired: false, + connection_deadline, + transfer_timer: None, + } + } + + fn check_timeout(&mut self, cx: &mut Context<'_>) -> io::Result<()> { + if let Some(deadline) = self.connection_deadline.next(cx) { + let timer = self + .transfer_timer + .get_or_insert_with(|| Box::pin(tokio::time::sleep_until(deadline))); + if timer.deadline() != deadline { + timer.as_mut().reset(deadline); + } + if timer.as_mut().poll(cx).is_ready() { + self.expired = true; + } + } else { + self.transfer_timer = None; + } + if self + .stalled + .as_mut() + .is_some_and(|timer| timer.as_mut().poll(cx).is_ready()) + { + self.expired = true; + } + if self.expired { + return Err(io::Error::new( + io::ErrorKind::TimedOut, + "Frontend connection transfer deadline expired", + )); + } + Ok(()) + } + + fn start_timer(&mut self, cx: &mut Context<'_>) { + let timer = self + .stalled + .get_or_insert_with(|| Box::pin(tokio::time::sleep(self.timeout))); + // Wake the connection task even when Hyper stops polling the body. + let _ = timer.as_mut().poll(cx); + } + + fn track_write(&mut self, cx: &mut Context<'_>, result: &Poll>) { + match result { + Poll::Pending => { + self.is_write_pending = true; + self.start_timer(cx); + } + Poll::Ready(Ok(written)) if *written > 0 => { + self.is_write_pending = false; + self.stalled = None; + } + _ => {} + } + } + + fn track_flush(&mut self, cx: &mut Context<'_>, result: &Poll>) { + match result { + Poll::Pending => self.start_timer(cx), + // A no-op flush must not reset a blocked write's deadline. + Poll::Ready(Ok(())) if !self.is_write_pending => self.stalled = None, + _ => {} + } + } +} + +impl AsyncRead for WriteTimeoutIo { + fn poll_read( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + self.check_timeout(cx)?; + Pin::new(&mut self.inner).poll_read(cx, buf) + } +} + +impl AsyncWrite for WriteTimeoutIo { + fn poll_write( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + self.check_timeout(cx)?; + let result = Pin::new(&mut self.inner).poll_write(cx, buf); + self.track_write(cx, &result); + result + } + + fn poll_write_vectored( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + bufs: &[io::IoSlice<'_>], + ) -> Poll> { + self.check_timeout(cx)?; + let result = Pin::new(&mut self.inner).poll_write_vectored(cx, bufs); + self.track_write(cx, &result); + result + } + + fn is_write_vectored(&self) -> bool { + self.inner.is_write_vectored() + } + + fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + self.check_timeout(cx)?; + let result = Pin::new(&mut self.inner).poll_flush(cx); + self.track_flush(cx, &result); + result + } + + fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + self.check_timeout(cx)?; + let result = Pin::new(&mut self.inner).poll_shutdown(cx); + self.track_flush(cx, &result); + result + } +} + +#[cfg(test)] +mod tests { + use super::super::server::serve; + use super::*; + use axum::Router; + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + use tokio_rustls::TlsAcceptor; + use tokio_util::task::AbortOnDropHandle; + + #[tokio::test(start_paused = true)] + async fn completed_transfers_remove_their_connection_deadline() { + let (inner, _) = tokio::io::duplex(1); + let connection = ConnectionDeadline::default(); + let first = connection.start(Duration::from_secs(1)); + let second = connection.start(Duration::from_secs(5)); + let mut writer = WriteTimeoutIo::new(inner, Duration::from_secs(30), connection.clone()); + let waker = futures::task::noop_waker(); + let mut cx = Context::from_waker(&waker); + writer.check_timeout(&mut cx).unwrap(); + assert_eq!( + writer.transfer_timer.as_ref().unwrap().deadline(), + first.deadline + ); + drop(first); + writer.check_timeout(&mut cx).unwrap(); + assert_eq!( + writer.transfer_timer.as_ref().unwrap().deadline(), + second.deadline + ); + tokio::time::advance(Duration::from_secs(2)).await; + writer.check_timeout(&mut cx).unwrap(); + drop(second); + tokio::time::advance(Duration::from_secs(10)).await; + writer.check_timeout(&mut cx).unwrap(); + assert!(writer.transfer_timer.is_none()); + let _expired = connection.start(Duration::from_secs(1)); + writer.check_timeout(&mut cx).unwrap(); + tokio::time::advance(Duration::from_secs(2)).await; + assert_eq!( + writer.check_timeout(&mut cx).unwrap_err().kind(), + io::ErrorKind::TimedOut + ); + } + + #[tokio::test] + async fn write_progress_resets_deadline_and_stalled_write_expires() { + tokio::time::timeout(Duration::from_secs(2), async { + let (inner, mut reader) = tokio::io::duplex(1); + let mut writer = WriteTimeoutIo::new( + inner, + Duration::from_millis(100), + ConnectionDeadline::default(), + ); + writer.write_all(b"a").await.unwrap(); + let write = writer.write_all(b"b"); + tokio::pin!(write); + // Start a stalled write, then free capacity before its idle deadline. + assert!( + tokio::time::timeout(Duration::from_millis(5), &mut write) + .await + .is_err() + ); + let mut byte = [0]; + reader.read_exact(&mut byte).await.unwrap(); + write.await.unwrap(); + assert_eq!(byte, *b"a"); + { + let bufs = [io::IoSlice::new(b"c")]; + let write = writer.write_vectored(&bufs); + tokio::pin!(write); + assert!( + tokio::time::timeout(Duration::from_millis(5), &mut write) + .await + .is_err() + ); + } + // DuplexStream's flush is a no-op: it cannot renew a blocked write. + writer.flush().await.unwrap(); + let error = writer.write_all(b"c").await.unwrap_err(); + assert_eq!(error.kind(), io::ErrorKind::TimedOut); + // An expired connection must not become usable again. + reader.read_exact(&mut byte).await.unwrap(); + assert_eq!( + writer.write_all(b"d").await.unwrap_err().kind(), + io::ErrorKind::TimedOut + ); + }) + .await + .unwrap(); + } + + fn tls_acceptor(cert: &rcgen::CertifiedKey) -> TlsAcceptor { + let provider = Arc::new(rustls::crypto::ring::default_provider()); + let key = rustls::pki_types::PrivatePkcs8KeyDer::from(cert.key_pair.serialize_der()); + let mut config = rustls::ServerConfig::builder_with_provider(provider) + .with_safe_default_protocol_versions() + .unwrap() + .with_no_client_auth() + .with_single_cert(vec![cert.cert.der().clone()], key.into()) + .unwrap(); + config.alpn_protocols = vec![b"h2".to_vec(), b"http/1.1".to_vec()]; + TlsAcceptor::from(Arc::new(config)) + } + + // A client that connects and never sends a ClientHello is closed once the TLS + // handshake timeout passes, instead of holding its connection task open. + #[tokio::test(start_paused = true)] + async fn silent_tls_client_is_closed_after_handshake_timeout() { + let cert = rcgen::generate_simple_self_signed(vec!["localhost".into()]).unwrap(); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let server = serve( + listener, + Some(tls_acceptor(&cert)), + Router::new(), + Duration::from_secs(30), + std::future::pending(), + ); + let _server = AbortOnDropHandle::new(tokio::spawn(server)); + let mut client = tokio::net::TcpStream::connect(addr).await.unwrap(); + let started = tokio::time::Instant::now(); + let mut buf = [0; 1]; + let read = tokio::time::timeout(Duration::from_secs(60), client.read(&mut buf)) + .await + .expect("silent TLS client was not closed"); + assert_eq!(read.unwrap(), 0); + assert!(started.elapsed() >= Duration::from_secs(10)); + } +}