dynamo_runtime/routing_policy/
types.rs

1// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
5pub struct RouteTarget {
6    pub worker_id: u64,
7    pub dp_rank: Option<u32>,
8}
9
10impl RouteTarget {
11    pub const fn worker(worker_id: u64) -> Self {
12        Self {
13            worker_id,
14            dp_rank: None,
15        }
16    }
17
18    pub const fn new(worker_id: u64, dp_rank: Option<u32>) -> Self {
19        Self { worker_id, dp_rank }
20    }
21}
22
23#[derive(Clone, Copy, Debug, PartialEq, Eq)]
24pub(crate) enum RoutePolicy {
25    RoundRobin,
26    Random,
27    PowerOfTwoChoices,
28    LeastLoaded,
29    DeviceAwareWeighted,
30}
31
32#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
33pub(crate) enum RouteDevice {
34    Cpu,
35    #[default]
36    Accelerator,
37}
38
39#[derive(Clone, Copy, Debug, PartialEq, Eq)]
40pub(crate) struct RouteCandidate {
41    pub(crate) target: RouteTarget,
42    pub(crate) device: RouteDevice,
43    pub(crate) cache_hits: usize,
44}
45
46impl RouteCandidate {
47    pub(crate) const fn worker(worker_id: u64) -> Self {
48        Self {
49            target: RouteTarget::worker(worker_id),
50            device: RouteDevice::Accelerator,
51            cache_hits: 0,
52        }
53    }
54}
55
56#[derive(Clone, Copy, Debug)]
57pub(crate) enum CandidateView<'a> {
58    Workers(&'a [u64]),
59    DeviceAware(&'a [RouteCandidate]),
60}
61
62impl CandidateView<'_> {
63    #[inline(always)]
64    pub(super) fn len(&self) -> usize {
65        match self {
66            Self::Workers(workers) => workers.len(),
67            Self::DeviceAware(candidates) => candidates.len(),
68        }
69    }
70
71    #[inline(always)]
72    pub(super) fn is_empty(&self) -> bool {
73        self.len() == 0
74    }
75
76    #[inline(always)]
77    pub(super) fn target(&self, index: usize) -> RouteTarget {
78        match self {
79            Self::Workers(workers) => RouteTarget::worker(workers[index]),
80            Self::DeviceAware(candidates) => candidates[index].target,
81        }
82    }
83}
84
85#[derive(Clone, Copy, Debug)]
86pub(crate) struct RouteContext {
87    pub(crate) required_cache_hits: usize,
88    pub(crate) non_cpu_to_cpu_ratio: usize,
89}
90
91impl Default for RouteContext {
92    fn default() -> Self {
93        Self {
94            required_cache_hits: 0,
95            non_cpu_to_cpu_ratio: 8,
96        }
97    }
98}
99
100#[derive(Clone, Copy, Debug, PartialEq, Eq)]
101pub(crate) enum AdmissionKind {
102    None,
103    Occupancy,
104}
105
106#[derive(Clone, Copy, Debug, PartialEq, Eq)]
107pub(crate) struct RouteDecision {
108    pub(crate) target: RouteTarget,
109    pub(crate) admission: AdmissionKind,
110}
dynamo_runtime/routing_policy/types.rs

dynamo_runtime/routing_policy/
types.rs
