dynamo_runtime/discovery/kube/
utils.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4use anyhow::Result;
5use k8s_openapi::api::discovery::v1::EndpointSlice;
6use std::collections::hash_map::DefaultHasher;
7use std::hash::{Hash, Hasher};
8
9/// Hash a pod name to get a consistent instance ID
10pub fn hash_pod_name(pod_name: &str) -> u64 {
11    // Clear top 11 bits to ensure it can be safely rounded to IEEE-754 f64
12    const INSTANCE_ID_MASK: u64 = 0x001F_FFFF_FFFF_FFFFu64;
13    let mut hasher = DefaultHasher::new();
14    pod_name.hash(&mut hasher);
15    hasher.finish() & INSTANCE_ID_MASK
16}
17
18/// Extract endpoint information from an EndpointSlice
19/// Returns (instance_id, pod_name) tuples for ready endpoints
20pub(super) fn extract_endpoint_info(slice: &EndpointSlice) -> Vec<(u64, String)> {
21    let mut result = Vec::new();
22
23    for endpoint in &slice.endpoints {
24        let is_ready = endpoint
25            .conditions
26            .as_ref()
27            .and_then(|c| c.ready)
28            .unwrap_or(false);
29
30        if !is_ready {
31            continue;
32        }
33
34        let pod_name = match endpoint.target_ref.as_ref() {
35            Some(target_ref) => target_ref.name.as_deref().unwrap_or(""),
36            None => continue,
37        };
38
39        if pod_name.is_empty() {
40            continue;
41        }
42
43        let instance_id = hash_pod_name(pod_name);
44
45        result.push((instance_id, pod_name.to_string()));
46    }
47
48    result
49}
50
51/// Pod information extracted from environment
52#[derive(Debug, Clone)]
53pub(super) struct PodInfo {
54    pub pod_name: String,
55    pub pod_namespace: String,
56    pub pod_uid: String,
57    pub system_port: u16,
58}
59
60impl PodInfo {
61    /// Discover pod information from environment variables
62    ///
63    /// Required environment variables:
64    /// - `POD_NAME`: Name of the pod (required)
65    /// - `POD_UID`: UID of the pod (required for CR owner reference)
66    /// - `POD_NAMESPACE`: Namespace of the pod (defaults to "default")
67    pub fn from_env() -> Result<Self> {
68        let pod_name = std::env::var("POD_NAME")
69            .map_err(|_| anyhow::anyhow!("POD_NAME environment variable not set"))?;
70
71        let pod_uid = std::env::var("POD_UID")
72            .map_err(|_| anyhow::anyhow!("POD_UID environment variable not set"))?;
73
74        let pod_namespace = std::env::var("POD_NAMESPACE").unwrap_or_else(|_| {
75            tracing::warn!("POD_NAMESPACE not set, defaulting to 'default'");
76            "default".to_string()
77        });
78
79        // Read system server port from config
80        let config = crate::config::RuntimeConfig::from_settings().unwrap_or_default();
81        let system_port = config.system_port as u16;
82
83        Ok(Self {
84            pod_name,
85            pod_namespace,
86            pod_uid,
87            system_port,
88        })
89    }
90}
91
92#[cfg(test)]
93mod tests {
94    use super::*;
95
96    #[test]
97    fn test_hash_json_serialization_roundtrip() {
98        // Verify that JSON serialization/deserialization preserves exact values
99        let pod_names = [
100            "worker-0",
101            "worker-99999",
102            "deployment-with-hash-suffix-a1b2c3d4e5f6",
103            "fake-name-1-0-worker-nrdfv",
104        ];
105
106        for pod_name in &pod_names {
107            let original_hash = hash_pod_name(pod_name);
108            let json = serde_json::to_string(&original_hash).unwrap();
109            let deserialized_hash: u64 = serde_json::from_str(&json).unwrap();
110
111            assert_eq!(
112                original_hash, deserialized_hash,
113                "JSON roundtrip changed hash value for pod_name={:?}: {} -> {} (json: {})",
114                pod_name, original_hash, deserialized_hash, json
115            );
116        }
117    }
118
119    #[test]
120    fn test_hash_in_struct_serialization() {
121        // Test serialization when the hash is embedded in a struct
122        #[derive(serde::Serialize, serde::Deserialize, Debug, PartialEq)]
123        struct WorkerInfo {
124            instance_id: u64,
125            name: String,
126        }
127
128        let pod_name = "fake-name-1-0-worker-nrdfv";
129        let info = WorkerInfo {
130            instance_id: hash_pod_name(pod_name),
131            name: pod_name.to_string(),
132        };
133
134        let json = serde_json::to_string(&info).unwrap();
135        let deserialized: WorkerInfo = serde_json::from_str(&json).unwrap();
136
137        assert_eq!(info, deserialized);
138    }
139}
dynamo_runtime/discovery/kube/utils.rs

dynamo_runtime/discovery/kube/
utils.rs
