agora_agentkit/scheduler/
grouping.rs

1//! Batch grouping algorithm.
2//!
3//! Groups [`WorkItem`]s into [`BatchGroup`]s optimized for cache reuse.
4//! Sort priority: model > prefix hash > context length bucket.
5
6use std::collections::BTreeMap;
7
8use super::WorkItem;
9
10/// Configuration for the grouping algorithm.
11#[derive(Debug, Clone)]
12pub struct GroupingConfig {
13    /// Context length bucket size in tokens. Items whose token counts
14    /// fall in the same bucket (token_count / bucket_size) are grouped
15    /// together to minimize Ollama memory reallocation.
16    pub context_length_bucket: u32,
17}
18
19impl Default for GroupingConfig {
20    fn default() -> Self {
21        Self {
22            context_length_bucket: 4096,
23        }
24    }
25}
26
27/// A group of work items that should be submitted together.
28///
29/// All items in a group share the same model and prefix hash,
30/// maximizing cache reuse.
31#[derive(Debug)]
32pub struct BatchGroup<P> {
33    /// Model identifier shared by all items in this group.
34    pub model: String,
35    /// Prefix hash shared by all items in this group.
36    pub prefix_hash: u64,
37    /// Context length bucket (token_count / bucket_size).
38    pub context_bucket: u32,
39    /// The work items in this group.
40    pub items: Vec<WorkItem<P>>,
41}
42
43/// Composite key for grouping: (model, prefix_hash, context_bucket).
44///
45/// Uses `BTreeMap` ordering: model (String) > prefix_hash (u64) >
46/// context_bucket (u32), which matches our desired sort priority.
47type GroupKey = (String, u64, u32);
48
49/// Group work items into [`BatchGroup`]s by (model, prefix_hash, context_bucket).
50///
51/// Items are grouped by the composite key. Within each group, items are
52/// in their original order. The groups themselves are sorted by key
53/// (model first, then prefix_hash, then context_bucket).
54pub fn group_work_items<P>(
55    items: Vec<WorkItem<P>>,
56    config: &GroupingConfig,
57) -> Vec<BatchGroup<P>> {
58    let bucket_size = config.context_length_bucket.max(1); // avoid div by zero
59
60    let mut groups: BTreeMap<GroupKey, Vec<WorkItem<P>>> = BTreeMap::new();
61
62    for item in items {
63        let bucket = item.token_count / bucket_size;
64        let key = (item.model.clone(), item.prefix_hash, bucket);
65        groups.entry(key).or_default().push(item);
66    }
67
68    groups
69        .into_iter()
70        .map(|((model, prefix_hash, context_bucket), items)| BatchGroup {
71            model,
72            prefix_hash,
73            context_bucket,
74            items,
75        })
76        .collect()
77}
78
79#[cfg(test)]
80mod tests {
81    use std::time::Instant;
82
83    use crate::ids::AgentId;
84
85    use super::super::CycleStep;
86    use super::*;
87
88    fn item(model: &str, prefix_hash: u64, token_count: u32) -> WorkItem<()> {
89        WorkItem {
90            agent_id: AgentId::new(),
91            prompt: (),
92            step: CycleStep::Think,
93            prefix_hash,
94            model: model.to_string(),
95            queued_at: Instant::now(),
96            token_count,
97        }
98    }
99
100    #[test]
101    fn groups_by_model() {
102        let items = vec![
103            item("claude", 100, 5000),
104            item("cogito", 100, 5000),
105            item("claude", 100, 5000),
106        ];
107
108        let groups = group_work_items(items, &GroupingConfig::default());
109        assert_eq!(groups.len(), 2);
110        assert_eq!(groups[0].model, "claude");
111        assert_eq!(groups[0].items.len(), 2);
112        assert_eq!(groups[1].model, "cogito");
113        assert_eq!(groups[1].items.len(), 1);
114    }
115
116    #[test]
117    fn groups_by_prefix_hash() {
118        let items = vec![
119            item("claude", 100, 5000),
120            item("claude", 200, 5000),
121            item("claude", 100, 5000),
122        ];
123
124        let groups = group_work_items(items, &GroupingConfig::default());
125        assert_eq!(groups.len(), 2);
126        assert_eq!(groups[0].prefix_hash, 100);
127        assert_eq!(groups[0].items.len(), 2);
128        assert_eq!(groups[1].prefix_hash, 200);
129    }
130
131    #[test]
132    fn groups_by_context_bucket() {
133        let config = GroupingConfig {
134            context_length_bucket: 4096,
135        };
136
137        let items = vec![
138            item("claude", 100, 4000),  // bucket 0
139            item("claude", 100, 5000),  // bucket 1
140            item("claude", 100, 4500),  // bucket 1
141            item("claude", 100, 12000), // bucket 2
142        ];
143
144        let groups = group_work_items(items, &config);
145        assert_eq!(groups.len(), 3);
146        assert_eq!(groups[0].context_bucket, 0);
147        assert_eq!(groups[0].items.len(), 1);
148        assert_eq!(groups[1].context_bucket, 1);
149        assert_eq!(groups[1].items.len(), 2);
150        assert_eq!(groups[2].context_bucket, 2);
151        assert_eq!(groups[2].items.len(), 1);
152    }
153
154    #[test]
155    fn empty_input() {
156        let groups = group_work_items::<()>(vec![], &GroupingConfig::default());
157        assert!(groups.is_empty());
158    }
159}
agora_agentkit/scheduler/grouping.rs

agora_agentkit/scheduler/
grouping.rs
