ts_l3/layout.rs
1// Copyright (c) 2026 LightSeek Foundation
2//
3// Permission is hereby granted, free of charge, to any person obtaining a copy
4// of this software and associated documentation files (the "Software"), to deal
5// in the Software without restriction, including without limitation the rights
6// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
7// copies of the Software, and to permit persons to whom the Software is
8// furnished to do so, subject to the following conditions:
9//
10// The above copyright notice and this permission notice shall be included in
11// all copies or substantial portions of the Software.
12//
13// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
14// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
15// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
16// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
17// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
18// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
19// SOFTWARE.
20
21//! Where a scheduler Host block lives in the shared pinned Host allocation,
22//! and whether this rank owns it.
23//!
24//! The data plane allocates the Host pool (`cache/host_storage.py`) and the
25//! L2 tier copies through it; the same bytes are what an L3 GET fills and a
26//! PUT reads. One LCM block holds `packing` cache blocks of each group, laid
27//! out field after field. A group with `shard_count` D deals its blocks
28//! cyclically to D owners: block `v > 0` belongs to rank `(v - 1) % D` and is
29//! that rank's local block `(v - 1) / D + 1`
30//! (`cache/transfer/ownership.py`).
31
32/// One cache group's blocks in the Host pool.
33#[derive(Debug, Clone, Copy, PartialEq, Eq)]
34pub struct GroupLayout {
35 /// Cache blocks per LCM block.
36 pub packing: usize,
37 /// Bytes of one packed cache block.
38 pub block_bytes: usize,
39 /// Owners the group's blocks are dealt to; 1 is replicated.
40 pub shard_count: usize,
41}
42
43/// The Host pool's geometry and this rank's place among the owners.
44#[derive(Debug, Clone, PartialEq, Eq)]
45pub struct HostLayout {
46 /// Each group, in the scheduler's group order.
47 pub groups: Vec<GroupLayout>,
48 /// Bytes of one LCM block.
49 pub lcm_block_bytes: usize,
50 /// LCM blocks in the pool, the null block excluded.
51 pub num_lcm_blocks: usize,
52 /// This rank among a sharded group's owners (the KV-parallel rank).
53 pub owner_rank: usize,
54 /// This rank's attention-TP rank.
55 pub tp_rank: usize,
56}
57
58/// A block or group the layout cannot place.
59#[derive(Debug, thiserror::Error, PartialEq, Eq)]
60pub enum LayoutError {
61 /// No such group.
62 #[error("cache transfer names unknown group {0}")]
63 Group(u32),
64 /// The null block.
65 #[error("a cache transfer cannot name the null block 0")]
66 Null,
67 /// Past the pool.
68 #[error("Host block {block} of group {group} is out of range")]
69 Block {
70 /// The group.
71 group: u32,
72 /// The block.
73 block: i64,
74 },
75}
76
77/// A byte range of the Host pool.
78#[derive(Debug, Clone, Copy, PartialEq, Eq)]
79pub struct Range {
80 /// From the pool's start.
81 pub offset: usize,
82 /// Its length.
83 pub size: usize,
84}
85
86impl HostLayout {
87 /// The pool's size.
88 pub fn bytes(&self) -> usize {
89 self.num_lcm_blocks * self.lcm_block_bytes
90 }
91
92 fn group(&self, group: u32) -> Result<&GroupLayout, LayoutError> {
93 self.groups
94 .get(group as usize)
95 .ok_or(LayoutError::Group(group))
96 }
97
98 /// The bytes of `group`'s scheduler block `block` on this rank, or None
99 /// when another rank owns it.
100 pub fn owned_range(&self, group: u32, block: i32) -> Result<Option<Range>, LayoutError> {
101 let layout = self.group(group)?;
102 if block == 0 {
103 return Err(LayoutError::Null);
104 }
105 let virtual_count = 1 + self.num_lcm_blocks * layout.packing * layout.shard_count;
106 let Ok(virtual_block) = usize::try_from(block) else {
107 return Err(LayoutError::Block {
108 group,
109 block: i64::from(block),
110 });
111 };
112 if virtual_block >= virtual_count {
113 return Err(LayoutError::Block {
114 group,
115 block: i64::from(block),
116 });
117 }
118 let dealt = virtual_block - 1;
119 if layout.shard_count > 1 && dealt % layout.shard_count != self.owner_rank {
120 return Ok(None);
121 }
122 let local = dealt / layout.shard_count + 1;
123 Ok(Some(self.local_range(group, layout, local)?))
124 }
125
126 /// The bytes of this rank's local block `local` of `group`
127 /// (`HostCacheStorage.host_block_range`).
128 fn local_range(
129 &self,
130 group: u32,
131 layout: &GroupLayout,
132 local: usize,
133 ) -> Result<Range, LayoutError> {
134 if local == 0 || local > self.num_lcm_blocks * layout.packing {
135 return Err(LayoutError::Block {
136 group,
137 block: local as i64,
138 });
139 }
140 let zero_based = local - 1;
141 let (parent, child) = (zero_based / layout.packing, zero_based % layout.packing);
142 Ok(Range {
143 offset: parent * self.lcm_block_bytes + child * layout.block_bytes,
144 size: layout.block_bytes,
145 })
146 }
147
148 /// The bytes of one of `group`'s blocks, owned here or not.
149 pub fn block_bytes(&self, group: u32) -> Result<usize, LayoutError> {
150 Ok(self.group(group)?.block_bytes)
151 }
152
153 /// The rank an object key of `group` names: the owners of a sharded
154 /// group share one name per shard, so a later allocation may move a
155 /// block to another owner of it.
156 pub fn key_rank(&self, group: u32) -> Result<i64, LayoutError> {
157 Ok((self.tp_rank / self.group(group)?.shard_count) as i64)
158 }
159}
160
161#[cfg(test)]
162mod tests {
163 use super::*;
164
165 fn layout(shard_count: usize, owner_rank: usize) -> HostLayout {
166 HostLayout {
167 groups: vec![
168 GroupLayout {
169 packing: 2,
170 block_bytes: 10,
171 shard_count: 1,
172 },
173 GroupLayout {
174 packing: 1,
175 block_bytes: 4,
176 shard_count,
177 },
178 ],
179 lcm_block_bytes: 20,
180 num_lcm_blocks: 3,
181 owner_rank,
182 tp_rank: 5,
183 }
184 }
185
186 #[test]
187 fn packed_blocks_sit_inside_their_lcm_block() {
188 let host = layout(1, 0);
189 assert_eq!(host.bytes(), 60);
190 let range = |block| host.owned_range(0, block).unwrap().unwrap();
191 assert_eq!(
192 range(1),
193 Range {
194 offset: 0,
195 size: 10
196 }
197 );
198 assert_eq!(
199 range(2),
200 Range {
201 offset: 10,
202 size: 10
203 }
204 );
205 assert_eq!(
206 range(3),
207 Range {
208 offset: 20,
209 size: 10
210 }
211 );
212 assert_eq!(
213 range(6),
214 Range {
215 offset: 50,
216 size: 10
217 }
218 );
219 assert_eq!(
220 host.owned_range(0, 7),
221 Err(LayoutError::Block { group: 0, block: 7 })
222 );
223 assert_eq!(host.owned_range(0, 0), Err(LayoutError::Null));
224 assert_eq!(host.owned_range(2, 1), Err(LayoutError::Group(2)));
225 }
226
227 #[test]
228 fn a_sharded_group_deals_its_blocks_to_the_owners() {
229 // Blocks 1, 3, 5 belong to rank 0 as its local 1, 2, 3; 2, 4, 6 to
230 // rank 1.
231 let rank0 = layout(2, 0);
232 let rank1 = layout(2, 1);
233 assert_eq!(
234 rank0.owned_range(1, 3).unwrap(),
235 Some(Range {
236 offset: 20,
237 size: 4
238 })
239 );
240 assert_eq!(rank0.owned_range(1, 4).unwrap(), None);
241 assert_eq!(
242 rank1.owned_range(1, 4).unwrap(),
243 Some(Range {
244 offset: 20,
245 size: 4
246 })
247 );
248 // The bound counts every owner's blocks.
249 assert!(rank1.owned_range(1, 6).unwrap().is_some());
250 assert!(rank1.owned_range(1, 7).is_err());
251 assert_eq!(rank0.key_rank(1).unwrap(), 2);
252 assert_eq!(rank0.key_rank(0).unwrap(), 5);
253 }
254}