protocols.rs 9.03 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
Ryan Olson's avatar
Ryan Olson committed
15
16

use serde::{Deserialize, Serialize};
17
18
19
use std::str::FromStr;

use crate::pipeline::PipelineError;
Ryan Olson's avatar
Ryan Olson committed
20
21
22

pub mod annotated;

23
24
pub type LeaseId = i64;

25
26
27
28
29
30
31
/// Default namespace if user does not provide one
const DEFAULT_NAMESPACE: &str = "NS";

const DEFAULT_COMPONENT: &str = "C";

const DEFAULT_ENDPOINT: &str = "E";

Ryan Olson's avatar
Ryan Olson committed
32
33
34
35
36
37
#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq)]
pub struct Component {
    pub name: String,
    pub namespace: String,
}

38
39
/// Represents an endpoint with a namespace, component, and name.
///
40
/// An `Endpoint` is defined by a three-part string separated by `/` or a '.':
41
42
43
44
45
/// - **namespace**
/// - **component**
/// - **name**
///
/// Example format: `"namespace/component/endpoint"`
Ryan Olson's avatar
Ryan Olson committed
46
47
#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq)]
pub struct Endpoint {
48
49
    pub namespace: String,
    pub component: String,
Ryan Olson's avatar
Ryan Olson committed
50
    pub name: String,
51
}
52

53
54
55
56
57
impl PartialEq<Vec<&str>> for Endpoint {
    fn eq(&self, other: &Vec<&str>) -> bool {
        if other.len() != 3 {
            return false;
        }
58

59
60
61
62
63
64
65
66
        self.namespace == other[0] && self.component == other[1] && self.name == other[2]
    }
}

impl PartialEq<Endpoint> for Vec<&str> {
    fn eq(&self, other: &Endpoint) -> bool {
        other == self
    }
Ryan Olson's avatar
Ryan Olson committed
67
68
}

69
70
71
72
73
74
75
76
77
impl Default for Endpoint {
    fn default() -> Self {
        Endpoint {
            namespace: DEFAULT_NAMESPACE.to_string(),
            component: DEFAULT_COMPONENT.to_string(),
            name: DEFAULT_ENDPOINT.to_string(),
        }
    }
}
78

79
80
impl From<&str> for Endpoint {
    /// Creates an `Endpoint` from a string.
81
82
83
84
    ///
    /// # Arguments
    /// - `path`: A string in the format `"namespace/component/endpoint"`.
    ///
85
86
87
88
89
90
91
92
93
94
    /// The first two parts become the first two elements of the vector.
    /// The third and subsequent parts are joined with '_' and become the third element.
    /// Default values are used for missing parts.
    ///
    /// # Examples:
    /// - "component" -> ["DEFAULT_NS", "component", "DEFAULT_E"]
    /// - "namespace.component" -> ["namespace", "component", "DEFAULT_E"]
    /// - "namespace.component.endpoint" -> ["namespace", "component", "endpoint"]
    /// - "namespace/component" -> ["namespace", "component", "DEFAULT_E"]
    /// - "namespace.component.endpoint.other.parts" -> ["namespace", "component", "endpoint_other_parts"]
95
96
97
    ///
    /// # Examples
    /// ```ignore
98
    /// use dynemo_runtime:protocols::Endpoint;
99
    ///
100
    /// let endpoint = Endpoint::from("namespace/component/endpoint");
101
102
103
104
    /// assert_eq!(endpoint.namespace, "namespace");
    /// assert_eq!(endpoint.component, "component");
    /// assert_eq!(endpoint.name, "endpoint");
    /// ```
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
    fn from(input: &str) -> Self {
        let mut result = Endpoint::default();

        // Split the input string on either '.' or '/'
        let elements: Vec<&str> = input
            .trim_matches([' ', '/', '.'])
            .split(['.', '/'])
            .filter(|x| !x.is_empty())
            .collect();

        match elements.len() {
            0 => {}
            1 => {
                result.component = elements[0].to_string();
            }
            2 => {
                result.namespace = elements[0].to_string();
                result.component = elements[1].to_string();
            }
            3 => {
                result.namespace = elements[0].to_string();
                result.component = elements[1].to_string();
                result.name = elements[2].to_string();
            }
            x if x > 3 => {
                result.namespace = elements[0].to_string();
                result.component = elements[1].to_string();
                result.name = elements[2..].join("_");
            }
            _ => unreachable!(),
135
        }
136
        result
137
138
139
140
141
142
143
144
    }
}

impl FromStr for Endpoint {
    type Err = PipelineError;

    /// Parses an `Endpoint` from a string using the standard Rust `.parse::<T>()` pattern.
    ///
145
    /// This is implemented in terms of [`From<&str>`].
146
147
    ///
    /// # Errors
148
    /// Does not fail
149
150
151
152
    ///
    /// # Examples
    /// ```ignore
    /// use std::str::FromStr;
153
    /// use dynemo_runtime:protocols::Endpoint;
154
155
156
157
158
159
160
    ///
    /// let endpoint: Endpoint = "namespace/component/endpoint".parse().unwrap();
    /// assert_eq!(endpoint.namespace, "namespace");
    /// assert_eq!(endpoint.component, "component");
    /// assert_eq!(endpoint.name, "endpoint");
    /// ```
    fn from_str(s: &str) -> Result<Self, Self::Err> {
161
        Ok(Endpoint::from(s))
162
163
164
    }
}

Ryan Olson's avatar
Ryan Olson committed
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq)]
#[serde(rename_all = "snake_case")]
pub enum RouterType {
    PushRoundRobin,
    PushRandom,
}

impl Default for RouterType {
    fn default() -> Self {
        Self::PushRandom
    }
}

#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq)]
pub struct ModelMetaData {
    pub name: String,
    pub component: Component,
    pub router_type: RouterType,
}
184
185
186
187

#[cfg(test)]
mod tests {
    use super::*;
188
189
    use std::convert::TryFrom;
    use std::str::FromStr;
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218

    #[test]
    fn test_router_type_default() {
        let default_router = RouterType::default();
        assert_eq!(default_router, RouterType::PushRandom);
    }

    #[test]
    fn test_router_type_serialization() {
        let router_round_robin = RouterType::PushRoundRobin;
        let router_random = RouterType::PushRandom;

        let serialized_round_robin = serde_json::to_string(&router_round_robin).unwrap();
        let serialized_random = serde_json::to_string(&router_random).unwrap();

        assert_eq!(serialized_round_robin, "\"push_round_robin\"");
        assert_eq!(serialized_random, "\"push_random\"");
    }

    #[test]
    fn test_router_type_deserialization() {
        let round_robin: RouterType = serde_json::from_str("\"push_round_robin\"").unwrap();
        let random: RouterType = serde_json::from_str("\"push_random\"").unwrap();

        assert_eq!(round_robin, RouterType::PushRoundRobin);
        assert_eq!(random, RouterType::PushRandom);
    }

    #[test]
219
    fn test_valid_endpoint_from() {
220
        let input = "namespace1/component1/endpoint1";
221
        let endpoint = Endpoint::from(input);
222
223
224
225
226
227
228
229
230

        assert_eq!(endpoint.namespace, "namespace1");
        assert_eq!(endpoint.component, "component1");
        assert_eq!(endpoint.name, "endpoint1");
    }

    #[test]
    fn test_valid_endpoint_from_str() {
        let input = "namespace2/component2/endpoint2";
231
        let endpoint = Endpoint::from_str(input).unwrap();
232
233
234
235
236
237
238
239
240

        assert_eq!(endpoint.namespace, "namespace2");
        assert_eq!(endpoint.component, "component2");
        assert_eq!(endpoint.name, "endpoint2");
    }

    #[test]
    fn test_valid_endpoint_parse() {
        let input = "namespace3/component3/endpoint3";
241
        let endpoint: Endpoint = input.parse().unwrap();
242
243
244
245
246
247
248

        assert_eq!(endpoint.namespace, "namespace3");
        assert_eq!(endpoint.component, "component3");
        assert_eq!(endpoint.name, "endpoint3");
    }

    #[test]
249
250
    fn test_endpoint_from() {
        let result = Endpoint::from("component");
251
        assert_eq!(
252
253
            result,
            vec![DEFAULT_NAMESPACE, "component", DEFAULT_ENDPOINT]
254
255
256
257
        );
    }

    #[test]
258
259
260
    fn test_namespace_component_endpoint() {
        let result = Endpoint::from("namespace.component.endpoint");
        assert_eq!(result, vec!["namespace", "component", "endpoint"]);
261
262
263
    }

    #[test]
264
265
266
    fn test_forward_slash_separator() {
        let result = Endpoint::from("namespace/component");
        assert_eq!(result, vec!["namespace", "component", DEFAULT_ENDPOINT]);
267
268
269
    }

    #[test]
270
271
    fn test_multiple_parts() {
        let result = Endpoint::from("namespace.component.endpoint.other.parts");
272
        assert_eq!(
273
274
            result,
            vec!["namespace", "component", "endpoint_other_parts"]
275
276
277
278
        );
    }

    #[test]
279
280
281
282
    fn test_mixed_separators() {
        // Do it the .into way for variety and documentation
        let result: Endpoint = "namespace/component.endpoint".into();
        assert_eq!(result, vec!["namespace", "component", "endpoint"]);
283
284
285
    }

    #[test]
286
287
288
289
290
291
    fn test_empty_string() {
        let result = Endpoint::from("");
        assert_eq!(
            result,
            vec![DEFAULT_NAMESPACE, DEFAULT_COMPONENT, DEFAULT_ENDPOINT]
        );
292

293
294
295
296
297
        // White space is equivalent to an empty string
        let result = Endpoint::from("   ");
        assert_eq!(
            result,
            vec![DEFAULT_NAMESPACE, DEFAULT_COMPONENT, DEFAULT_ENDPOINT]
298
299
        );
    }
300
}