1use super::{
2 align_up, invalid_plan, Deserialize, Deserializer, DynamicResourceDemand, DynamicResourceShape,
3 DynamicResourceShapeBucket, DynamicStorageRequirement, ResourceWorkShape, Serialize,
4 VNextError, MAX_PROVIDER_WORKSPACE_SHAPE_BUCKETS,
5};
6
7#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
12#[serde(rename_all = "snake_case")]
13pub enum ProviderWorkspaceSizeFormula {
14 Fixed {
15 bytes: u64,
16 },
17 ActualSequences {
18 bytes_per_sequence: u64,
19 },
20 Tokens {
21 bytes_per_token: u64,
22 },
23 Affine {
24 fixed_bytes: u64,
25 bytes_per_sequence: u64,
26 bytes_per_token: u64,
27 },
28 Pages {
29 bytes_per_page: u64,
30 maximum_pages: u64,
31 },
32 BoundedShapeBuckets {
33 buckets: Vec<DynamicResourceShapeBucket>,
34 },
35}
36
37#[derive(Deserialize)]
38#[serde(rename_all = "snake_case", deny_unknown_fields)]
39enum ProviderWorkspaceSizeFormulaWire {
40 Fixed {
41 bytes: u64,
42 },
43 ActualSequences {
44 bytes_per_sequence: u64,
45 },
46 Tokens {
47 bytes_per_token: u64,
48 },
49 Affine {
50 fixed_bytes: u64,
51 bytes_per_sequence: u64,
52 bytes_per_token: u64,
53 },
54 Pages {
55 bytes_per_page: u64,
56 maximum_pages: u64,
57 },
58 BoundedShapeBuckets {
59 buckets: Vec<DynamicResourceShapeBucket>,
60 },
61}
62
63impl ProviderWorkspaceSizeFormula {
64 pub fn fixed(bytes: u64) -> Result<Self, VNextError> {
65 Self::validated(Self::Fixed { bytes })
66 }
67
68 pub fn actual_sequences(bytes_per_sequence: u64) -> Result<Self, VNextError> {
69 Self::validated(Self::ActualSequences { bytes_per_sequence })
70 }
71
72 pub fn tokens(bytes_per_token: u64) -> Result<Self, VNextError> {
73 Self::validated(Self::Tokens { bytes_per_token })
74 }
75
76 pub fn affine(
77 fixed_bytes: u64,
78 bytes_per_sequence: u64,
79 bytes_per_token: u64,
80 ) -> Result<Self, VNextError> {
81 Self::validated(Self::Affine {
82 fixed_bytes,
83 bytes_per_sequence,
84 bytes_per_token,
85 })
86 }
87
88 pub fn pages(bytes_per_page: u64, maximum_pages: u64) -> Result<Self, VNextError> {
89 Self::validated(Self::Pages {
90 bytes_per_page,
91 maximum_pages,
92 })
93 }
94
95 pub fn bounded_shape_buckets(
96 buckets: Vec<DynamicResourceShapeBucket>,
97 ) -> Result<Self, VNextError> {
98 Self::validated(Self::BoundedShapeBuckets { buckets })
99 }
100
101 fn validated(formula: Self) -> Result<Self, VNextError> {
102 formula.validate()?;
103 Ok(formula)
104 }
105
106 fn validate(&self) -> Result<(), VNextError> {
107 let valid = match self {
108 Self::Fixed { bytes } => *bytes > 0,
109 Self::ActualSequences { bytes_per_sequence }
110 | Self::Tokens {
111 bytes_per_token: bytes_per_sequence,
112 } => *bytes_per_sequence > 0,
113 Self::Affine {
114 fixed_bytes,
115 bytes_per_sequence,
116 bytes_per_token,
117 } => {
118 (*bytes_per_sequence > 0 || *bytes_per_token > 0)
119 && fixed_bytes
120 .checked_add(*bytes_per_sequence)
121 .and_then(|bytes| bytes.checked_add(*bytes_per_token))
122 .is_some()
123 }
124 Self::Pages {
125 bytes_per_page,
126 maximum_pages,
127 } => {
128 *bytes_per_page > 0
129 && *maximum_pages > 0
130 && bytes_per_page.checked_mul(*maximum_pages).is_some()
131 }
132 Self::BoundedShapeBuckets { buckets } => {
133 !buckets.is_empty()
134 && buckets.len() <= MAX_PROVIDER_WORKSPACE_SHAPE_BUCKETS
135 && buckets.windows(2).all(|pair| {
136 let previous = &pair[0];
137 let next = &pair[1];
138 next.maximum_sequences() >= previous.maximum_sequences()
139 && next.maximum_tokens() >= previous.maximum_tokens()
140 && next.maximum_pages() >= previous.maximum_pages()
141 && (next.maximum_sequences() > previous.maximum_sequences()
142 || next.maximum_tokens() > previous.maximum_tokens()
143 || next.maximum_pages() > previous.maximum_pages())
144 && next.bytes() >= previous.bytes()
145 })
146 }
147 };
148 if !valid {
149 return Err(invalid_plan(
150 "provider workspace formula is zero, overflowing, or non-canonical",
151 ));
152 }
153 Ok(())
154 }
155
156 fn evaluate_shape_bytes(&self, shape: DynamicResourceShape) -> Result<u64, VNextError> {
157 self.validate()?;
158 let bytes = match self {
159 Self::Fixed { bytes } => *bytes,
160 Self::ActualSequences { bytes_per_sequence } => bytes_per_sequence
161 .checked_mul(u64::from(shape.sequences))
162 .ok_or_else(|| invalid_plan("provider sequence workspace overflows u64"))?,
163 Self::Tokens { bytes_per_token } => bytes_per_token
164 .checked_mul(shape.tokens)
165 .ok_or_else(|| invalid_plan("provider token workspace overflows u64"))?,
166 Self::Affine {
167 fixed_bytes,
168 bytes_per_sequence,
169 bytes_per_token,
170 } => fixed_bytes
171 .checked_add(
172 bytes_per_sequence
173 .checked_mul(u64::from(shape.sequences))
174 .ok_or_else(|| {
175 invalid_plan("provider affine sequence workspace overflows u64")
176 })?,
177 )
178 .and_then(|bytes| {
179 bytes_per_token
180 .checked_mul(shape.tokens)
181 .and_then(|token_bytes| bytes.checked_add(token_bytes))
182 })
183 .ok_or_else(|| invalid_plan("provider affine token workspace overflows u64"))?,
184 Self::Pages {
185 bytes_per_page,
186 maximum_pages,
187 } if shape.pages <= *maximum_pages => bytes_per_page
188 .checked_mul(shape.pages)
189 .ok_or_else(|| invalid_plan("provider page workspace overflows u64"))?,
190 Self::BoundedShapeBuckets { buckets } => buckets
191 .iter()
192 .find(|bucket| bucket.covers(shape))
193 .map(DynamicResourceShapeBucket::bytes)
194 .ok_or_else(|| invalid_plan("actual invocation shape exceeds provider buckets"))?,
195 Self::Pages { .. } => {
196 return Err(invalid_plan(
197 "actual invocation pages exceed the provider implementation bound",
198 ))
199 }
200 };
201 if bytes == 0 {
202 return Err(invalid_plan("provider workspace evaluates to zero bytes"));
203 }
204 Ok(bytes)
205 }
206
207 pub(super) fn bind_runtime_limits(
208 &self,
209 maximum_sequences: u32,
210 maximum_tokens: u64,
211 ) -> Result<DynamicResourceDemand, VNextError> {
212 self.validate()?;
213 match self {
214 Self::Fixed { bytes } => DynamicResourceDemand::fixed(*bytes),
215 Self::ActualSequences { bytes_per_sequence } => {
216 DynamicResourceDemand::actual_sequences(*bytes_per_sequence, maximum_sequences)
217 }
218 Self::Tokens { bytes_per_token } => {
219 DynamicResourceDemand::tokens(*bytes_per_token, maximum_tokens)
220 }
221 Self::Affine {
222 fixed_bytes,
223 bytes_per_sequence,
224 bytes_per_token,
225 } => DynamicResourceDemand::affine(
226 *fixed_bytes,
227 *bytes_per_sequence,
228 maximum_sequences,
229 *bytes_per_token,
230 maximum_tokens,
231 ),
232 Self::Pages {
233 bytes_per_page,
234 maximum_pages,
235 } => DynamicResourceDemand::pages(*bytes_per_page, *maximum_pages),
236 Self::BoundedShapeBuckets { buckets } => {
237 DynamicResourceDemand::bounded_shape_buckets(buckets.clone())
238 }
239 }
240 }
241
242 fn is_fixed(&self) -> bool {
243 matches!(self, Self::Fixed { .. })
244 }
245
246 fn is_valid_for_sequence_scope(&self) -> bool {
247 match self {
248 Self::Fixed { .. } | Self::Tokens { .. } | Self::Pages { .. } => true,
249 Self::Affine {
250 bytes_per_sequence, ..
251 } => *bytes_per_sequence == 0,
252 Self::BoundedShapeBuckets { buckets } => {
253 buckets.iter().all(|bucket| bucket.maximum_sequences() == 1)
254 }
255 Self::ActualSequences { .. } => false,
256 }
257 }
258}
259
260impl<'de> Deserialize<'de> for ProviderWorkspaceSizeFormula {
261 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
262 where
263 D: serde::Deserializer<'de>,
264 {
265 let formula = match ProviderWorkspaceSizeFormulaWire::deserialize(deserializer)? {
266 ProviderWorkspaceSizeFormulaWire::Fixed { bytes } => Self::Fixed { bytes },
267 ProviderWorkspaceSizeFormulaWire::ActualSequences { bytes_per_sequence } => {
268 Self::ActualSequences { bytes_per_sequence }
269 }
270 ProviderWorkspaceSizeFormulaWire::Tokens { bytes_per_token } => {
271 Self::Tokens { bytes_per_token }
272 }
273 ProviderWorkspaceSizeFormulaWire::Affine {
274 fixed_bytes,
275 bytes_per_sequence,
276 bytes_per_token,
277 } => Self::Affine {
278 fixed_bytes,
279 bytes_per_sequence,
280 bytes_per_token,
281 },
282 ProviderWorkspaceSizeFormulaWire::Pages {
283 bytes_per_page,
284 maximum_pages,
285 } => Self::Pages {
286 bytes_per_page,
287 maximum_pages,
288 },
289 ProviderWorkspaceSizeFormulaWire::BoundedShapeBuckets { buckets } => {
290 Self::BoundedShapeBuckets { buckets }
291 }
292 };
293 Self::validated(formula).map_err(serde::de::Error::custom)
294 }
295}
296
297#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
298#[serde(rename_all = "snake_case")]
299pub enum ProviderWorkspaceScope {
300 Plan,
301 Request,
302 Sequence,
303 Step,
304 Invocation,
305}
306
307#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
314#[serde(rename_all = "snake_case")]
315pub enum ProviderWorkspaceReusePolicy {
316 OverwriteBeforeRead,
318 ZeroBeforeUse,
320 Preserve,
322}
323
324#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
325#[serde(deny_unknown_fields)]
326pub struct ProviderWorkspaceRequirement {
327 pub(super) size_formula: ProviderWorkspaceSizeFormula,
328 pub(super) alignment_bytes: u64,
329 pub(super) scope: ProviderWorkspaceScope,
330 pub(super) reuse_policy: ProviderWorkspaceReusePolicy,
331 pub(super) storage: DynamicStorageRequirement,
332}
333
334impl ProviderWorkspaceRequirement {
335 pub fn new(
338 fixed_bytes: u64,
339 alignment_bytes: u64,
340 scope: ProviderWorkspaceScope,
341 reuse_policy: ProviderWorkspaceReusePolicy,
342 storage: DynamicStorageRequirement,
343 ) -> Result<Self, VNextError> {
344 Self::from_formula(
345 ProviderWorkspaceSizeFormula::fixed(fixed_bytes)?,
346 alignment_bytes,
347 scope,
348 reuse_policy,
349 storage,
350 )
351 }
352
353 pub fn from_formula(
354 size_formula: ProviderWorkspaceSizeFormula,
355 alignment_bytes: u64,
356 scope: ProviderWorkspaceScope,
357 reuse_policy: ProviderWorkspaceReusePolicy,
358 storage: DynamicStorageRequirement,
359 ) -> Result<Self, VNextError> {
360 size_formula.validate()?;
361 if alignment_bytes == 0
362 || !alignment_bytes.is_power_of_two()
363 || (scope == ProviderWorkspaceScope::Plan && !size_formula.is_fixed())
364 || (scope == ProviderWorkspaceScope::Sequence
365 && !size_formula.is_valid_for_sequence_scope())
366 {
367 return Err(invalid_plan(
368 "provider workspace has invalid formula, alignment, or scope",
369 ));
370 }
371 let requirement = Self {
372 size_formula,
373 alignment_bytes,
374 scope,
375 reuse_policy,
376 storage,
377 };
378 requirement.minimum_bytes()?;
379 Ok(requirement)
380 }
381
382 pub fn size_formula(&self) -> &ProviderWorkspaceSizeFormula {
383 &self.size_formula
384 }
385
386 pub fn evaluate_bytes(&self, work: &ResourceWorkShape) -> Result<u64, VNextError> {
387 self.evaluate_shape_bytes(work.immediate_shape())
388 }
389
390 pub fn evaluate_fit_bytes(&self, work: &ResourceWorkShape) -> Result<u64, VNextError> {
391 self.evaluate_shape_bytes(work.fit_shape())
392 }
393
394 pub(crate) fn evaluate_shape_bytes(
395 &self,
396 shape: DynamicResourceShape,
397 ) -> Result<u64, VNextError> {
398 align_up(
399 self.size_formula.evaluate_shape_bytes(shape)?,
400 self.alignment_bytes,
401 )
402 }
403
404 pub fn minimum_bytes(&self) -> Result<u64, VNextError> {
405 self.evaluate_shape_bytes(DynamicResourceShape::from_validated(1, 1, 1))
406 }
407
408 pub fn fixed_bytes(&self) -> Option<u64> {
409 match &self.size_formula {
410 ProviderWorkspaceSizeFormula::Fixed { bytes } => Some(*bytes),
411 _ => None,
412 }
413 }
414
415 pub const fn alignment_bytes(&self) -> u64 {
416 self.alignment_bytes
417 }
418
419 pub const fn scope(&self) -> ProviderWorkspaceScope {
420 self.scope
421 }
422
423 pub const fn reuse_policy(&self) -> ProviderWorkspaceReusePolicy {
424 self.reuse_policy
425 }
426
427 pub fn storage(&self) -> &DynamicStorageRequirement {
428 &self.storage
429 }
430}
431
432#[derive(Deserialize)]
433#[serde(deny_unknown_fields)]
434pub(super) struct ProviderWorkspaceRequirementWire {
435 pub(super) size_formula: ProviderWorkspaceSizeFormula,
436 pub(super) alignment_bytes: u64,
437 pub(super) scope: ProviderWorkspaceScope,
438 pub(super) reuse_policy: ProviderWorkspaceReusePolicy,
439 pub(super) storage: DynamicStorageRequirement,
440}
441
442impl<'de> Deserialize<'de> for ProviderWorkspaceRequirement {
443 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
444 where
445 D: Deserializer<'de>,
446 {
447 let wire = ProviderWorkspaceRequirementWire::deserialize(deserializer)?;
448 Self::from_formula(
449 wire.size_formula,
450 wire.alignment_bytes,
451 wire.scope,
452 wire.reuse_policy,
453 wire.storage,
454 )
455 .map_err(serde::de::Error::custom)
456 }
457}