1use super::{
2 canonical_fingerprint, canonical_runtime_policy_fingerprint, invalid_plan, is_canonical_sha256,
3 joint_candidate_components, joint_partial_precedes, node_weight_requirements,
4 provider_resource_estimator_input_fingerprint, static_contiguous_storage_profile,
5 storage_incompatible_resource_ids, tensor_storage_layout_fingerprint,
6 validate_active_sequence_ceiling, validate_program_bindings, validate_scheduled_token_ceiling,
7 validate_semantic_binding, workspace_base_id, workspace_storage_layout_fingerprint,
8 AliasPolicy, AllocationKind, AllocationLifetime, BTreeMap, BTreeSet, BufferUsage,
9 CanonicalValueBinding, CapabilityCatalog, CapabilityId, DimensionConstraint,
10 DynamicResourceDemand, DynamicResourceDescriptor, DynamicStorageContract,
11 DynamicStorageProfile, DynamicStorageRequirement, ElementType, ExecutionPlanPayload,
12 GlobalValueRange, JointComponentSolution, JointPartialSelection, JointProviderCandidate,
13 JointProviderStorageSelection, JointSelectionObjective, MemoryPlan, NodeId,
14 NodeTokenBindingProjection, NodeWorkContract, OperationDescriptor, OperationRegistryAuthority,
15 PlanBuildRequest, PlanExactAlias, PlanExactAliasKind, PlanHash, PlanHashMaterial, PlanId,
16 PlanNode, PlanNodeResolution, PlanProviderRejectReason, PlanStateEffect, PreparedModelFamily,
17 ProgramNode, ProgramNodeWorkSpec, ProgramValueId, ProviderCompatibilityRequest, ProviderId,
18 ProviderResourcePlan, ProviderSelection, ProviderSelectionReason, ProviderWorkspaceScope,
19 QuantizationFormatId, RejectedProvider, ResolvedValueBinding, ResolvedValueRole,
20 ResourceAllocation, ResourceId, ReusableExecutionMemoryPlan, ReusableExecutionPolicy,
21 RuntimePolicy, Serialize, StateCapacityDemand, StateDependencyTracker, StateInitialization,
22 StateLifetime, TensorAccess, TrustedExecutionWeightPlan, VNextError,
23 ValueAllocationAccumulator, ValueResourceDemand, WeightFormatId, WeightSchema,
24 EXECUTION_PLAN_SCHEMA,
25};
26use super::{resolve_retained_completion_values, CompletionRetentionSpec, RetainedCompletionValue};
27use crate::vnext::{
28 CompletionReadbackRequest, ExecutionDeterminismRequirement, HostTransferLayout,
29 ResourceWorkShape, WeightComponentPayload, WeightComponentSource, WeightComponentSpec,
30};
31
32#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
33pub struct ExecutionPlan {
34 pub(super) payload: ExecutionPlanPayload,
35 pub(super) plan_hash: PlanHash,
36 #[serde(skip)]
37 pub(super) operation_registry_authority: OperationRegistryAuthority,
38 #[serde(skip)]
39 pub(super) trusted_execution_weights: TrustedExecutionWeightPlan,
40}
41
42impl ExecutionPlan {
43 pub fn build<P: RuntimePolicy>(request: PlanBuildRequest<'_, P>) -> Result<Self, VNextError> {
44 request.policy.validate()?;
45 let maximum_active_sequences = request.policy.maximum_active_sequences();
46 validate_active_sequence_ceiling(maximum_active_sequences)?;
47 let maximum_scheduled_tokens = request.policy.maximum_scheduled_tokens();
48 validate_scheduled_token_ceiling(maximum_scheduled_tokens)?;
49 let operation_registry_authority = request
50 .node_resolutions
51 .first()
52 .ok_or_else(|| invalid_plan("plan build request has no node resolutions"))?
53 .operation_registry_authority
54 .clone();
55 if request.node_resolutions.iter().any(|resolution| {
56 resolution.operation_registry_authority != operation_registry_authority
57 }) {
58 return Err(invalid_plan(
59 "node resolutions belong to different operation runtime registries",
60 ));
61 }
62 let policy_capacity = request.policy.memory_capacity_bytes();
63 let memory_reserve = request.policy.memory_reserve_bytes();
64 let device_capacity = request.capabilities.device().total_memory_bytes;
65 if policy_capacity == 0
66 || policy_capacity > device_capacity
67 || memory_reserve >= policy_capacity
68 {
69 return Err(invalid_plan(
70 "runtime policy raw capacity, reserve, or typed admission concurrency is invalid for the device descriptor",
71 ));
72 }
73 let family = request.family;
74 let program = family.program();
75 request
76 .execution_weights
77 .validate_against_catalog(family, request.capabilities)?;
78 let execution_weight_schema = request.execution_weights.plan().schema();
79 let prepared_family_fingerprint = family.fingerprint()?;
80 let program_fingerprint = program.fingerprint()?;
81 let capability_catalog_fingerprint = request.capabilities.fingerprint()?;
82 let device_runtime_implementation_fingerprint = request
83 .capabilities
84 .device()
85 .runtime_implementation_fingerprint
86 .clone();
87 let policy_fingerprint = canonical_runtime_policy_fingerprint(request.policy)?;
88 let weight_format = execution_weight_schema.format_id.clone();
89 let quantization_formats = execution_weight_schema.quantization_formats();
90
91 let mut resolutions = BTreeMap::new();
92 for resolution in request.node_resolutions {
93 let node_id = resolution.node_id.clone();
94 if resolutions.insert(node_id.clone(), resolution).is_some() {
95 return Err(invalid_plan(format!(
96 "node `{node_id}` has duplicate physical resolutions"
97 )));
98 }
99 }
100
101 let program_nodes = program
102 .blocks()
103 .iter()
104 .flat_map(|block| &block.nodes)
105 .collect::<Vec<_>>();
106 let joint_storage = Self::select_joint_provider_storage(
107 &program_nodes,
108 &resolutions,
109 request.capabilities,
110 request.policy,
111 )?;
112 let mut selected_node_resources = joint_storage.node_resources;
113 let selected_resource_profiles = joint_storage.resource_profiles;
114 let mut storage_rejections = joint_storage.storage_rejections;
115 let producers = program_nodes
116 .iter()
117 .flat_map(|node| {
118 node.outputs
119 .iter()
120 .map(move |output| (output.clone(), node.id.clone()))
121 })
122 .collect::<BTreeMap<_, _>>();
123 let mut last_consumers = BTreeMap::<ProgramValueId, usize>::new();
124 for (node_index, node) in program_nodes.iter().enumerate() {
125 for input in &node.inputs {
126 last_consumers.insert(input.clone(), node_index);
127 }
128 }
129 let program_outputs = program.outputs().iter().cloned().collect::<BTreeSet<_>>();
130 let mut canonical_values = BTreeMap::new();
131 let mut bound_values = BTreeSet::new();
132 let mut nodes = Vec::new();
133 let mut state_dependencies = StateDependencyTracker::default();
134 for (node_index, program_node) in program_nodes.into_iter().enumerate() {
135 let resolution = resolutions.remove(&program_node.id).ok_or_else(|| {
136 invalid_plan(format!(
137 "node `{}` has no physical resolution",
138 program_node.id
139 ))
140 })?;
141 let provider_resources = selected_node_resources
142 .remove(&program_node.id)
143 .ok_or_else(|| invalid_plan("joint solver omitted one program node"))?;
144 let storage_rejection = storage_rejections.remove(&program_node.id);
145 let node = Self::build_node(
146 family,
147 execution_weight_schema,
148 &prepared_family_fingerprint,
149 program_node,
150 resolution,
151 provider_resources,
152 storage_rejection,
153 request.capabilities,
154 request.policy.execution_determinism_requirement(),
155 &producers,
156 node_index,
157 &last_consumers,
158 &program_outputs,
159 &mut state_dependencies,
160 &mut canonical_values,
161 &mut bound_values,
162 )?;
163 nodes.push(node);
164 }
165 if !resolutions.is_empty() {
166 return Err(invalid_plan(format!(
167 "physical resolutions contain unknown nodes: {:?}",
168 resolutions.keys().collect::<Vec<_>>()
169 )));
170 }
171 if !selected_node_resources.is_empty() {
172 return Err(invalid_plan(
173 "joint solver returned resources for unknown program nodes",
174 ));
175 }
176 if !storage_rejections.is_empty() {
177 return Err(invalid_plan(
178 "joint solver returned storage rejections for unknown program nodes",
179 ));
180 }
181 Self::validate_semantic_coverage(family, &bound_values)?;
182 Self::validate_global_storage_aliasing(&canonical_values, &nodes)?;
183 let retained_completion_values =
184 resolve_retained_completion_values(&nodes, &request.completion_retention)?;
185 let terminal_output_resources = nodes
186 .iter()
187 .flat_map(|node| node.values())
188 .filter(|binding| program_outputs.contains(binding.value_id()))
189 .flat_map(|binding| binding.storage().components())
190 .map(|component| component.resource_id().clone())
191 .collect::<BTreeSet<_>>()
192 .into_iter()
193 .collect::<Vec<_>>();
194 let retained_completion_resources = retained_completion_values
195 .iter()
196 .map(|value| value.resource_id().clone())
197 .chain(terminal_output_resources.iter().cloned())
198 .collect::<BTreeSet<_>>();
199 let memory = Self::build_memory_plan(
200 family,
201 device_capacity,
202 policy_capacity,
203 memory_reserve,
204 maximum_active_sequences,
205 maximum_scheduled_tokens,
206 &nodes,
207 &selected_resource_profiles,
208 request.policy.reusable_execution_policy(),
209 &retained_completion_resources,
210 )?;
211
212 let mut payload = ExecutionPlanPayload {
213 schema: EXECUTION_PLAN_SCHEMA,
214 plan_id: PlanId::new("plan/unset")?,
215 family_id: family.family_id().clone(),
216 device_id: request.capabilities.device().id.clone(),
217 device_runtime_implementation_fingerprint,
218 prepared_family_fingerprint,
219 program_fingerprint,
220 capability_catalog_fingerprint,
221 policy_version: request.policy.version(),
222 policy_fingerprint,
223 maximum_scheduled_tokens,
224 execution_weights: request.execution_weights.plan().clone(),
225 weight_format,
226 quantization_formats,
227 retained_completion_values,
228 terminal_output_resources,
229 nodes,
230 memory,
231 };
232 let plan_hash = PlanHash::new(canonical_fingerprint(
233 &PlanHashMaterial::from(&payload),
234 "fingerprint execution plan",
235 )?)?;
236 payload.plan_id = Self::plan_id_for_hash(&plan_hash)?;
237 let plan = Self {
238 payload,
239 plan_hash,
240 operation_registry_authority,
241 trusted_execution_weights: request.execution_weights,
242 };
243 plan.validate_internal()?;
244 Ok(plan)
245 }
246
247 #[allow(clippy::too_many_arguments)]
248 pub(super) fn build_node(
249 family: &PreparedModelFamily,
250 execution_weight_schema: &WeightSchema,
251 prepared_family_fingerprint: &str,
252 program_node: &ProgramNode,
253 resolution: PlanNodeResolution,
254 provider_resources: ProviderResourcePlan,
255 storage_rejection: Option<RejectedProvider>,
256 catalog: &CapabilityCatalog,
257 execution_determinism: ExecutionDeterminismRequirement,
258 producers: &BTreeMap<ProgramValueId, NodeId>,
259 node_index: usize,
260 last_consumers: &BTreeMap<ProgramValueId, usize>,
261 program_outputs: &BTreeSet<ProgramValueId>,
262 state_dependencies: &mut StateDependencyTracker,
263 canonical_values: &mut BTreeMap<ProgramValueId, CanonicalValueBinding>,
264 bound_values: &mut BTreeSet<ProgramValueId>,
265 ) -> Result<PlanNode, VNextError> {
266 let operation = catalog.operation_for_node(&program_node.id, &program_node.operation_id)?;
267 if !operation.version.satisfies(program_node.required_version) {
268 return Err(VNextError::IncompatibleOperationVersion {
269 node_id: Some(program_node.id.to_string()),
270 operation_id: program_node.operation_id.to_string(),
271 required_major: program_node.required_version.major,
272 required_minor: program_node.required_version.minor,
273 available_major: operation.version.major,
274 available_minor: operation.version.minor,
275 });
276 }
277 operation.validate_attributes(&program_node.attributes)?;
278 operation.validate_resolved_bindings(&resolution.values)?;
279 validate_program_bindings(program_node, &resolution.values)?;
280 let exact_aliases = Self::extract_exact_aliases(operation, &resolution.values)?;
281 let work = Self::derive_node_work_contract(program_node, operation, &resolution.values)?;
282 Self::validate_alias_liveness(
283 family,
284 program_node,
285 node_index,
286 last_consumers,
287 program_outputs,
288 &exact_aliases,
289 &resolution.values,
290 )?;
291 let state_effects = Self::derive_state_effects(family, &resolution.values)?;
292 for binding in &resolution.values {
293 validate_semantic_binding(family, execution_weight_schema, binding)?;
294 Self::validate_cross_node_value(binding, canonical_values)?;
295 bound_values.insert(binding.value_id().clone());
296 }
297
298 let (required_weight_formats, required_quantization_formats) =
299 node_weight_requirements(family, &resolution.values)?;
300 let selection = Self::select_provider(
301 program_node,
302 operation,
303 catalog,
304 &resolution.required_capabilities,
305 resolution.preferred_provider.as_ref(),
306 &provider_resources.provider_id,
307 storage_rejection,
308 &required_weight_formats,
309 &required_quantization_formats,
310 execution_determinism,
311 )?;
312 provider_resources.validate_shape()?;
313 if provider_resources.provider_id != selection.selected_provider {
314 return Err(invalid_plan(format!(
315 "node `{}` resource estimate belongs to provider `{}` instead of selected provider `{}`",
316 program_node.id,
317 provider_resources.provider_id,
318 selection.selected_provider
319 )));
320 }
321 let selected_provider = catalog
322 .providers_for_node(&program_node.id, &program_node.operation_id)?
323 .iter()
324 .find(|provider| provider.provider_id() == &selection.selected_provider)
325 .ok_or_else(|| {
326 invalid_plan(format!(
327 "node `{}` selected provider is absent from the catalog",
328 program_node.id
329 ))
330 })?;
331 if provider_resources.estimator_id != selected_provider.resource_estimator_id()
332 || provider_resources.estimator_version
333 != selected_provider.resource_estimator_version()
334 || provider_resources.estimator_implementation_fingerprint
335 != selected_provider.resource_estimator_implementation_fingerprint()
336 {
337 return Err(invalid_plan(format!(
338 "node `{}` provider resource estimate is not issued by the selected catalog provider's estimator",
339 program_node.id
340 )));
341 }
342 let minimum_value_alignment = operation.resources.minimum_value_alignment_bytes;
343 if provider_resources.value_alignment_bytes < minimum_value_alignment
344 || provider_resources.value_alignment_bytes % minimum_value_alignment != 0
345 || !operation
346 .resources
347 .scratch
348 .accepts(provider_resources.scratch.is_some())
349 || !operation
350 .resources
351 .binding
352 .accepts(provider_resources.binding.is_some())
353 || !operation
354 .resources
355 .persistent
356 .accepts(provider_resources.persistent.is_some())
357 {
358 return Err(invalid_plan(format!(
359 "node `{}` provider resource estimate violates the operation's alignment or workspace-presence contract",
360 program_node.id
361 )));
362 }
363 let expected_estimator_input = provider_resource_estimator_input_fingerprint(
364 family,
365 prepared_family_fingerprint,
366 operation,
367 program_node,
368 &selection.selected_provider,
369 &resolution.values,
370 &resolution.required_capabilities,
371 )?;
372 if provider_resources.estimator_input_fingerprint != expected_estimator_input {
373 return Err(invalid_plan(format!(
374 "node `{}` provider resource estimate is not bound to its selected provider, shape, attributes, and bindings",
375 program_node.id
376 )));
377 }
378 let mut dependencies = program_node
379 .inputs
380 .iter()
381 .filter_map(|input| producers.get(input))
382 .cloned()
383 .collect::<BTreeSet<_>>();
384 Self::add_state_dependencies(
385 &program_node.id,
386 &state_effects,
387 state_dependencies,
388 &mut dependencies,
389 );
390 let dependencies = dependencies.into_iter().collect::<Vec<_>>();
391 let scratch_resource = provider_resources
392 .scratch
393 .as_ref()
394 .map(|_| {
395 workspace_base_id(
396 &program_node.id,
397 "scratch",
398 &provider_resources.estimate_fingerprint,
399 )
400 })
401 .transpose()?;
402 let binding_resource = provider_resources
403 .binding
404 .as_ref()
405 .map(|_| {
406 workspace_base_id(
407 &program_node.id,
408 "binding",
409 &provider_resources.estimate_fingerprint,
410 )
411 })
412 .transpose()?;
413 let persistent_resource = provider_resources
414 .persistent
415 .as_ref()
416 .map(|_| {
417 workspace_base_id(
418 &program_node.id,
419 "persistent",
420 &provider_resources.estimate_fingerprint,
421 )
422 })
423 .transpose()?;
424 let resources = resolution
425 .values
426 .iter()
427 .flat_map(|binding| binding.storage().components())
428 .map(|component| component.resource_id().clone())
429 .chain(scratch_resource.iter().cloned())
430 .chain(binding_resource.iter().cloned())
431 .chain(persistent_resource.iter().cloned())
432 .collect::<BTreeSet<_>>()
433 .into_iter()
434 .collect();
435 Ok(PlanNode {
436 id: program_node.id.clone(),
437 dependencies,
438 operation_id: program_node.operation_id.clone(),
439 operation_version: program_node.required_version,
440 operation_fingerprint: operation.fingerprint()?,
441 provider_implementation_fingerprint: selected_provider
442 .provider_implementation_fingerprint()
443 .to_owned(),
444 provider_execution_semantics: selected_provider.execution_semantics(),
445 required_capabilities: resolution.required_capabilities,
446 attributes: program_node.attributes.clone(),
447 work,
448 selection,
449 provider_resources,
450 values: resolution.values,
451 exact_aliases,
452 state_effects,
453 scratch_resource,
454 binding_resource,
455 persistent_resource,
456 resources,
457 })
458 }
459
460 pub(super) fn derive_node_work_contract(
461 node: &ProgramNode,
462 operation: &OperationDescriptor,
463 bindings: &[ResolvedValueBinding],
464 ) -> Result<NodeWorkContract, VNextError> {
465 let ProgramNodeWorkSpec::Tokens {
466 value_id,
467 axis: source_axis,
468 } = &node.work
469 else {
470 return Ok(NodeWorkContract::Fixed);
471 };
472 let source_binding = bindings
473 .iter()
474 .find(|binding| binding.value_id() == value_id)
475 .ok_or_else(|| invalid_plan("node token work source has no resolved binding"))?;
476 let source_contract = match source_binding.role() {
477 ResolvedValueRole::Input => operation.inputs.get(source_binding.ordinal() as usize),
478 ResolvedValueRole::Output => operation.outputs.get(source_binding.ordinal() as usize),
479 }
480 .ok_or_else(|| invalid_plan("node token work source ordinal is outside its operation"))?;
481 let source_axis_index = usize::try_from(*source_axis)
482 .map_err(|_| invalid_plan("node token work axis exceeds usize"))?;
483 let source_symbol = match source_contract.dimensions().get(source_axis_index) {
484 Some(DimensionConstraint::Symbol(symbol)) => symbol,
485 _ => {
486 return Err(invalid_plan(
487 "node token work source axis is not one symbolic operation dimension",
488 ))
489 }
490 };
491 if source_binding.usage() != BufferUsage::Activations
492 || source_binding
493 .tensor()
494 .dimensions()
495 .get(source_axis_index)
496 .is_none()
497 {
498 return Err(invalid_plan(
499 "node token work source is not an in-bounds activation axis",
500 ));
501 }
502
503 let mut projections = Vec::new();
504 for binding in bindings {
505 let contract = match binding.role() {
506 ResolvedValueRole::Input => operation.inputs.get(binding.ordinal() as usize),
507 ResolvedValueRole::Output => operation.outputs.get(binding.ordinal() as usize),
508 }
509 .ok_or_else(|| {
510 invalid_plan("resolved work binding ordinal is outside its operation")
511 })?;
512 let matching_axes = contract
513 .dimensions()
514 .iter()
515 .enumerate()
516 .filter(|(_, dimension)| {
517 matches!(dimension, DimensionConstraint::Symbol(symbol) if symbol == source_symbol)
518 })
519 .map(|(axis, _)| axis)
520 .collect::<Vec<_>>();
521 if matching_axes.len() > 1 {
522 return Err(invalid_plan(
523 "one resolved binding repeats the node token work dimension",
524 ));
525 }
526 let Some(axis) = matching_axes.first().copied() else {
527 continue;
528 };
529 let dimensions = binding.tensor().dimensions();
530 if binding.usage() != BufferUsage::Activations || dimensions.get(axis).is_none() {
531 return Err(invalid_plan(
532 "node token work projection is not an in-bounds activation axis",
533 ));
534 }
535 projections.push(NodeTokenBindingProjection {
536 value_id: binding.value_id().clone(),
537 role: binding.role(),
538 ordinal: binding.ordinal(),
539 axis: u32::try_from(axis)
540 .map_err(|_| invalid_plan("node token projection axis exceeds u32"))?,
541 rank: u32::try_from(dimensions.len())
542 .map_err(|_| invalid_plan("node token projection rank exceeds u32"))?,
543 canonical_extent: dimensions[axis],
544 });
545 }
546 projections.sort();
547 if projections.is_empty()
548 || projections
549 .windows(2)
550 .any(|pair| pair[0].role == pair[1].role && pair[0].ordinal == pair[1].ordinal)
551 {
552 return Err(invalid_plan(
553 "node token work projections are empty or non-canonical",
554 ));
555 }
556 let source = projections
557 .iter()
558 .find(|projection| projection.value_id == *value_id && projection.axis == *source_axis)
559 .cloned()
560 .ok_or_else(|| invalid_plan("node token work source did not resolve exactly"))?;
561 if projections
562 .iter()
563 .any(|projection| projection.canonical_extent != source.canonical_extent)
564 {
565 return Err(invalid_plan(
566 "node token work projections disagree on canonical extent",
567 ));
568 }
569 Ok(NodeWorkContract::Tokens {
570 source,
571 projections,
572 })
573 }
574
575 fn validate_node_work_contract(node: &PlanNode) -> Result<(), VNextError> {
576 let NodeWorkContract::Tokens {
577 source,
578 projections,
579 } = &node.work
580 else {
581 return Ok(());
582 };
583 if projections.is_empty()
584 || projections.windows(2).any(|pair| pair[0] >= pair[1])
585 || !projections.contains(source)
586 || projections
587 .iter()
588 .any(|projection| projection.canonical_extent != source.canonical_extent)
589 {
590 return Err(invalid_plan(format!(
591 "node `{}` token work contract is empty or non-canonical",
592 node.id
593 )));
594 }
595 for projection in projections {
596 let binding = node
597 .values
598 .iter()
599 .find(|binding| {
600 binding.role() == projection.role
601 && binding.ordinal() == projection.ordinal
602 && binding.value_id() == &projection.value_id
603 })
604 .ok_or_else(|| {
605 invalid_plan(format!(
606 "node `{}` token projection has no exact value binding",
607 node.id
608 ))
609 })?;
610 let axis = usize::try_from(projection.axis)
611 .map_err(|_| invalid_plan("node token projection axis exceeds usize"))?;
612 if binding.usage() != BufferUsage::Activations
613 || usize::try_from(projection.rank).ok()
614 != Some(binding.tensor().dimensions().len())
615 || binding.tensor().dimensions().get(axis) != Some(&projection.canonical_extent)
616 {
617 return Err(invalid_plan(format!(
618 "node `{}` token projection differs from its resolved tensor",
619 node.id
620 )));
621 }
622 }
623 Ok(())
624 }
625
626 pub(super) fn extract_exact_aliases(
627 operation: &OperationDescriptor,
628 bindings: &[ResolvedValueBinding],
629 ) -> Result<Vec<PlanExactAlias>, VNextError> {
630 let inputs = &bindings[..operation.inputs.len()];
631 let outputs = &bindings[operation.inputs.len()..];
632 let mut aliases = Vec::new();
633 for (output_ordinal, output) in outputs.iter().enumerate() {
634 let (input_ordinal, kind) = match output.alias() {
635 AliasPolicy::NoAlias => continue,
636 AliasPolicy::MayAlias { tensor_index } => {
637 (*tensor_index, PlanExactAliasKind::MayAlias)
638 }
639 AliasPolicy::MustAlias { tensor_index } => {
640 (*tensor_index, PlanExactAliasKind::MustAlias)
641 }
642 };
643 let input = inputs.get(input_ordinal as usize).ok_or_else(|| {
644 invalid_plan(format!(
645 "operation `{}` alias input ordinal is out of range after validation",
646 operation.id
647 ))
648 })?;
649 if output.storage() == input.storage() {
650 aliases.push(PlanExactAlias {
651 output_value_id: output.value_id().clone(),
652 output_ordinal: output_ordinal as u32,
653 input_value_id: input.value_id().clone(),
654 input_ordinal,
655 kind,
656 });
657 } else if kind == PlanExactAliasKind::MustAlias {
658 return Err(invalid_plan(format!(
659 "operation `{}` lost its mandatory exact alias proof",
660 operation.id
661 )));
662 }
663 }
664 Ok(aliases)
665 }
666
667 #[allow(clippy::too_many_arguments)]
668 pub(super) fn validate_alias_liveness(
669 family: &PreparedModelFamily,
670 node: &ProgramNode,
671 node_index: usize,
672 last_consumers: &BTreeMap<ProgramValueId, usize>,
673 program_outputs: &BTreeSet<ProgramValueId>,
674 aliases: &[PlanExactAlias],
675 bindings: &[ResolvedValueBinding],
676 ) -> Result<(), VNextError> {
677 for alias in aliases {
678 let input = bindings
679 .iter()
680 .find(|binding| {
681 binding.role() == ResolvedValueRole::Input
682 && binding.ordinal() == alias.input_ordinal
683 && binding.value_id() == &alias.input_value_id
684 })
685 .ok_or_else(|| invalid_plan("exact alias input proof has no matching binding"))?;
686 if family
687 .program()
688 .states()
689 .iter()
690 .any(|state| state.value_id == alias.input_value_id)
691 {
692 return Err(invalid_plan(format!(
693 "node `{}` output aliases state `{}` without a typed state transition contract",
694 node.id, alias.input_value_id
695 )));
696 }
697 if input.usage() != BufferUsage::Activations
698 || last_consumers.get(&alias.input_value_id) != Some(&node_index)
699 || program_outputs.contains(&alias.input_value_id)
700 {
701 return Err(invalid_plan(format!(
702 "node `{}` aliases activation `{}` before its final legal consumer",
703 node.id, alias.input_value_id
704 )));
705 }
706 }
707 Ok(())
708 }
709
710 pub(super) fn derive_state_effects(
711 family: &PreparedModelFamily,
712 bindings: &[ResolvedValueBinding],
713 ) -> Result<Vec<PlanStateEffect>, VNextError> {
714 let mut effects = Vec::new();
715 for state in family.program().states() {
716 let mut reads = false;
717 let mut writes = false;
718 let state_bindings = bindings
719 .iter()
720 .filter(|binding| binding.value_id() == &state.value_id)
721 .collect::<Vec<_>>();
722 for binding in &state_bindings {
723 match binding.access() {
724 TensorAccess::Read => reads = true,
725 TensorAccess::Write => writes = true,
726 TensorAccess::ReadWrite => {
727 reads = true;
728 writes = true;
729 }
730 }
731 }
732 let access = match (reads, writes) {
733 (false, false) => continue,
734 (true, false) => TensorAccess::Read,
735 (false, true) => TensorAccess::Write,
736 (true, true) => TensorAccess::ReadWrite,
737 };
738 let lifetime = match state.lifetime {
739 StateLifetime::Request => AllocationLifetime::Request,
740 StateLifetime::Sequence => AllocationLifetime::Sequence,
741 StateLifetime::Step => AllocationLifetime::Step,
742 };
743 let resource_ids = state_bindings
744 .iter()
745 .flat_map(|binding| binding.storage().components())
746 .map(|component| component.resource_id().clone())
747 .collect::<BTreeSet<_>>()
748 .into_iter()
749 .collect::<Vec<_>>();
750 if resource_ids.is_empty() {
751 return Err(invalid_plan(format!(
752 "state `{}` effect has no physical resource closure",
753 state.id
754 )));
755 }
756 effects.push(PlanStateEffect {
757 state_id: state.id.clone(),
758 state_value_id: state.value_id.clone(),
759 lifetime,
760 access,
761 resource_ids,
762 });
763 }
764 if effects
765 .windows(2)
766 .any(|pair| pair[0].state_id >= pair[1].state_id)
767 {
768 return Err(invalid_plan("state effects are not canonical"));
769 }
770 Ok(effects)
771 }
772
773 pub(super) fn add_state_dependencies(
774 node_id: &NodeId,
775 effects: &[PlanStateEffect],
776 tracker: &mut StateDependencyTracker,
777 dependencies: &mut BTreeSet<NodeId>,
778 ) {
779 for effect in effects {
780 let state_id = &effect.state_id;
781 match effect.access {
782 TensorAccess::Read => {
783 if let Some(writer) = tracker.last_writer.get(state_id) {
784 dependencies.insert(writer.clone());
785 }
786 tracker
787 .readers_since_write
788 .entry(state_id.clone())
789 .or_default()
790 .insert(node_id.clone());
791 }
792 TensorAccess::Write | TensorAccess::ReadWrite => {
793 if let Some(writer) = tracker.last_writer.get(state_id) {
794 dependencies.insert(writer.clone());
795 }
796 if let Some(readers) = tracker.readers_since_write.remove(state_id) {
797 dependencies.extend(readers);
798 }
799 tracker
800 .last_writer
801 .insert(state_id.clone(), node_id.clone());
802 }
803 }
804 }
805 }
806
807 pub(super) fn validate_cross_node_value(
808 binding: &ResolvedValueBinding,
809 values: &mut BTreeMap<ProgramValueId, CanonicalValueBinding>,
810 ) -> Result<(), VNextError> {
811 let canonical = CanonicalValueBinding {
812 tensor: binding.tensor().clone(),
813 usage: binding.usage(),
814 storage: binding.storage().clone(),
815 };
816 match values.get(binding.value_id()) {
817 Some(previous) if previous != &canonical => Err(invalid_plan(format!(
818 "value `{}` changes tensor or physical storage between nodes",
819 binding.value_id()
820 ))),
821 Some(_) => Ok(()),
822 None => {
823 values.insert(binding.value_id().clone(), canonical);
824 Ok(())
825 }
826 }
827 }
828
829 pub(super) fn validate_global_storage_aliasing(
830 values: &BTreeMap<ProgramValueId, CanonicalValueBinding>,
831 nodes: &[PlanNode],
832 ) -> Result<(), VNextError> {
833 let alias_classes = Self::alias_classes(nodes)?;
834 let mut by_resource = BTreeMap::<ResourceId, Vec<GlobalValueRange>>::new();
835 for (value_id, binding) in values {
836 for component in binding.storage.components() {
837 let end_bytes = component
838 .offset_bytes()
839 .checked_add(component.length_bytes())
840 .ok_or_else(|| invalid_plan("global value storage range overflows u64"))?;
841 let ranges = by_resource
842 .entry(component.resource_id().clone())
843 .or_default();
844 if let Some(previous) = ranges.iter().find(|previous| {
845 previous.value_id != *value_id
846 && previous.offset_bytes < end_bytes
847 && component.offset_bytes() < previous.end_bytes
848 }) {
849 let same_alias_class = alias_classes.get(&previous.value_id).is_some()
850 && alias_classes.get(&previous.value_id) == alias_classes.get(value_id);
851 let previous_binding = values.get(&previous.value_id).ok_or_else(|| {
852 invalid_plan("global alias range has no canonical value binding")
853 })?;
854 if !same_alias_class || previous_binding.storage != binding.storage {
855 return Err(invalid_plan(format!(
856 "values `{}` and `{value_id}` have undeclared, partial, or non-equivalent overlap in physical resource `{}`",
857 previous.value_id,
858 component.resource_id()
859 )));
860 }
861 }
862 ranges.push(GlobalValueRange {
863 value_id: value_id.clone(),
864 offset_bytes: component.offset_bytes(),
865 end_bytes,
866 });
867 }
868 }
869 Ok(())
870 }
871
872 pub(super) fn alias_classes(
873 nodes: &[PlanNode],
874 ) -> Result<BTreeMap<ProgramValueId, ProgramValueId>, VNextError> {
875 let mut graph = BTreeMap::<ProgramValueId, BTreeSet<ProgramValueId>>::new();
876 for node in nodes {
877 let mut previous_output_ordinal = None;
878 for alias in &node.exact_aliases {
879 if previous_output_ordinal.is_some_and(|ordinal| ordinal >= alias.output_ordinal) {
880 return Err(invalid_plan(format!(
881 "node `{}` exact aliases are not canonical",
882 node.id
883 )));
884 }
885 previous_output_ordinal = Some(alias.output_ordinal);
886 let input = node
887 .values
888 .iter()
889 .find(|binding| {
890 binding.role() == ResolvedValueRole::Input
891 && binding.ordinal() == alias.input_ordinal
892 && binding.value_id() == &alias.input_value_id
893 })
894 .ok_or_else(|| invalid_plan("plan exact alias input binding is missing"))?;
895 let output = node
896 .values
897 .iter()
898 .find(|binding| {
899 binding.role() == ResolvedValueRole::Output
900 && binding.ordinal() == alias.output_ordinal
901 && binding.value_id() == &alias.output_value_id
902 })
903 .ok_or_else(|| invalid_plan("plan exact alias output binding is missing"))?;
904 let policy_matches = matches!(
905 (output.alias(), alias.kind),
906 (
907 AliasPolicy::MayAlias { tensor_index },
908 PlanExactAliasKind::MayAlias
909 ) if *tensor_index == alias.input_ordinal
910 ) || matches!(
911 (output.alias(), alias.kind),
912 (
913 AliasPolicy::MustAlias { tensor_index },
914 PlanExactAliasKind::MustAlias
915 ) if *tensor_index == alias.input_ordinal
916 );
917 if !policy_matches
918 || input.storage() != output.storage()
919 || input.usage() != BufferUsage::Activations
920 || output.usage() != BufferUsage::Activations
921 {
922 return Err(invalid_plan(format!(
923 "node `{}` exact alias proof differs from its bindings",
924 node.id
925 )));
926 }
927 graph
928 .entry(alias.input_value_id.clone())
929 .or_default()
930 .insert(alias.output_value_id.clone());
931 graph
932 .entry(alias.output_value_id.clone())
933 .or_default()
934 .insert(alias.input_value_id.clone());
935 }
936 }
937
938 let mut classes = BTreeMap::new();
939 let mut visited = BTreeSet::new();
940 for start in graph.keys() {
941 if visited.contains(start) {
942 continue;
943 }
944 let mut pending = vec![start.clone()];
945 let mut members = BTreeSet::new();
946 while let Some(value) = pending.pop() {
947 if !visited.insert(value.clone()) {
948 continue;
949 }
950 members.insert(value.clone());
951 if let Some(neighbors) = graph.get(&value) {
952 pending.extend(neighbors.iter().cloned());
953 }
954 }
955 let representative = members
956 .first()
957 .cloned()
958 .ok_or_else(|| invalid_plan("empty alias equivalence class"))?;
959 for member in members {
960 classes.insert(member, representative.clone());
961 }
962 }
963 Ok(classes)
964 }
965
966 pub(super) fn validate_semantic_coverage(
967 family: &PreparedModelFamily,
968 bound: &BTreeSet<ProgramValueId>,
969 ) -> Result<(), VNextError> {
970 let required = family
971 .program()
972 .inputs()
973 .iter()
974 .cloned()
975 .chain(
976 family
977 .program()
978 .weights()
979 .iter()
980 .map(|weight| weight.value_id.clone()),
981 )
982 .chain(
983 family
984 .program()
985 .states()
986 .iter()
987 .map(|state| state.value_id.clone()),
988 )
989 .chain(family.program().outputs().iter().cloned())
990 .collect::<BTreeSet<_>>();
991 if !required.is_subset(bound) {
992 return Err(invalid_plan(format!(
993 "semantic values lack physical bindings: {:?}",
994 required.difference(bound).collect::<Vec<_>>()
995 )));
996 }
997 Ok(())
998 }
999
1000 pub(super) fn available_storage_profiles<P: RuntimePolicy>(
1001 requirement: &DynamicStorageRequirement,
1002 catalog: &CapabilityCatalog,
1003 policy: &P,
1004 ) -> BTreeSet<DynamicStorageProfile> {
1005 policy
1006 .dynamic_storage_profile_order()
1007 .iter()
1008 .copied()
1009 .filter(|profile| {
1010 catalog.device().dynamic_storage_profiles.contains(profile)
1011 && requirement.accepts(*profile)
1012 })
1013 .collect()
1014 }
1015
1016 pub(super) fn merge_storage_constraint(
1017 constraints: &mut BTreeMap<ResourceId, BTreeSet<DynamicStorageProfile>>,
1018 resource_id: ResourceId,
1019 accepted: BTreeSet<DynamicStorageProfile>,
1020 ) -> bool {
1021 if accepted.is_empty() {
1022 return false;
1023 }
1024 match constraints.get_mut(&resource_id) {
1025 Some(existing) => {
1026 existing.retain(|profile| accepted.contains(profile));
1027 !existing.is_empty()
1028 }
1029 None => {
1030 constraints.insert(resource_id, accepted);
1031 true
1032 }
1033 }
1034 }
1035
1036 pub(super) fn select_joint_provider_storage<P: RuntimePolicy>(
1037 program_nodes: &[&ProgramNode],
1038 resolutions: &BTreeMap<NodeId, PlanNodeResolution>,
1039 catalog: &CapabilityCatalog,
1040 policy: &P,
1041 ) -> Result<JointProviderStorageSelection, VNextError> {
1042 let mut candidate_sets = Vec::with_capacity(program_nodes.len());
1043 for node in program_nodes {
1044 let resolution = resolutions.get(&node.id).ok_or_else(|| {
1045 invalid_plan(format!("node `{}` has no physical resolution", node.id))
1046 })?;
1047 let providers = catalog.providers_for_node(&node.id, &node.operation_id)?;
1048 let mut candidates = Vec::new();
1049 for resources in &resolution.provider_resource_candidates {
1050 let provider = providers
1051 .iter()
1052 .find(|provider| provider.provider_id() == resources.provider_id())
1053 .ok_or_else(|| {
1054 invalid_plan("provider resource candidate is absent from the catalog")
1055 })?;
1056 let mut constraints = BTreeMap::new();
1057 let mut compatible = true;
1058 for binding in resolution
1059 .values
1060 .iter()
1061 .filter(|binding| binding.usage() != BufferUsage::Weights)
1062 {
1063 let Some(requirement) =
1064 provider.dynamic_storage_for(binding.role(), binding.ordinal())
1065 else {
1066 compatible = false;
1067 break;
1068 };
1069 let accepted = Self::available_storage_profiles(requirement, catalog, policy);
1070 for component in binding.storage().components() {
1071 if !Self::merge_storage_constraint(
1072 &mut constraints,
1073 component.resource_id().clone(),
1074 accepted.clone(),
1075 ) {
1076 compatible = false;
1077 break;
1078 }
1079 }
1080 if !compatible {
1081 break;
1082 }
1083 }
1084 for (kind, workspace) in [
1085 ("scratch", resources.scratch()),
1086 ("binding", resources.binding()),
1087 ("persistent", resources.persistent()),
1088 ] {
1089 let Some(workspace) = workspace else {
1090 continue;
1091 };
1092 let resource_id =
1093 workspace_base_id(&node.id, kind, resources.estimate_fingerprint())?;
1094 if !Self::merge_storage_constraint(
1095 &mut constraints,
1096 resource_id,
1097 Self::available_storage_profiles(workspace.storage(), catalog, policy),
1098 ) {
1099 compatible = false;
1100 break;
1101 }
1102 }
1103 if compatible {
1104 let is_preferred = resolution
1105 .preferred_provider
1106 .as_ref()
1107 .is_some_and(|preferred| preferred == resources.provider_id());
1108 candidates.push(JointProviderCandidate {
1109 resources: resources.clone(),
1110 allowed_profiles: constraints,
1111 is_preferred,
1112 });
1113 }
1114 }
1115 candidates.sort_by(|left, right| {
1116 let left_preferred = resolution
1117 .preferred_provider
1118 .as_ref()
1119 .is_some_and(|preferred| preferred == left.resources.provider_id());
1120 let right_preferred = resolution
1121 .preferred_provider
1122 .as_ref()
1123 .is_some_and(|preferred| preferred == right.resources.provider_id());
1124 right_preferred.cmp(&left_preferred).then(
1125 left.resources
1126 .provider_id()
1127 .cmp(right.resources.provider_id()),
1128 )
1129 });
1130 if candidates.is_empty() {
1131 return Err(invalid_plan(format!(
1132 "node `{}` has no provider candidate with an available storage profile",
1133 node.id
1134 )));
1135 }
1136 candidate_sets.push(candidates);
1137 }
1138
1139 let (chosen, resource_profiles) = Self::solve_joint_provider_candidates(
1140 &candidate_sets,
1141 policy.dynamic_storage_profile_order(),
1142 )?;
1143
1144 let mut storage_rejections = BTreeMap::new();
1145 for (index, node) in program_nodes.iter().enumerate() {
1146 let resolution = resolutions
1147 .get(&node.id)
1148 .ok_or_else(|| invalid_plan("joint storage resolution disappeared"))?;
1149 let Some(preferred) = resolution.preferred_provider.as_ref() else {
1150 continue;
1151 };
1152 if chosen[index].provider_id() == preferred {
1153 continue;
1154 }
1155 let Some(preferred_candidate) = candidate_sets[index]
1156 .iter()
1157 .find(|candidate| candidate.resources.provider_id() == preferred)
1158 else {
1159 if let Some(reason) = resolution.provider_resolution_rejections.get(preferred) {
1160 storage_rejections.insert(
1161 node.id.clone(),
1162 RejectedProvider {
1163 provider_id: preferred.clone(),
1164 reasons: reason.clone(),
1165 },
1166 );
1167 }
1168 continue;
1169 };
1170 let resource_ids =
1171 storage_incompatible_resource_ids(preferred_candidate, &resource_profiles);
1172 if resource_ids.is_empty() {
1173 return Err(invalid_plan(format!(
1174 "preferred provider `{preferred}` was not selected for node `{}` without a storage conflict",
1175 node.id
1176 )));
1177 }
1178 storage_rejections.insert(
1179 node.id.clone(),
1180 RejectedProvider {
1181 provider_id: preferred.clone(),
1182 reasons: PlanProviderRejectReason::StorageIncompatible { resource_ids },
1183 },
1184 );
1185 }
1186
1187 let node_resources = program_nodes
1188 .iter()
1189 .zip(chosen)
1190 .map(|(node, resources)| (node.id.clone(), resources))
1191 .collect();
1192 Ok(JointProviderStorageSelection {
1193 node_resources,
1194 resource_profiles,
1195 storage_rejections,
1196 })
1197 }
1198
1199 pub(super) fn solve_joint_provider_candidates(
1200 candidate_sets: &[Vec<JointProviderCandidate>],
1201 profile_order: &[DynamicStorageProfile],
1202 ) -> Result<
1203 (
1204 Vec<ProviderResourcePlan>,
1205 BTreeMap<ResourceId, DynamicStorageProfile>,
1206 ),
1207 VNextError,
1208 > {
1209 if candidate_sets.is_empty() || candidate_sets.iter().any(Vec::is_empty) {
1210 return Err(invalid_plan(
1211 "joint provider/storage search has an empty candidate set",
1212 ));
1213 }
1214 if profile_order.is_empty() {
1215 return Err(invalid_plan(
1216 "joint provider/storage search has an empty profile order",
1217 ));
1218 }
1219
1220 let components = joint_candidate_components(candidate_sets);
1221 let mut chosen = vec![None; candidate_sets.len()];
1222 let mut resource_profiles = BTreeMap::new();
1223 for component in components {
1224 let solution =
1225 Self::solve_joint_provider_component(&component, candidate_sets, profile_order)?;
1226 for (node_index, resources) in component.iter().copied().zip(solution.chosen) {
1227 if chosen[node_index].replace(resources).is_some() {
1228 return Err(invalid_plan(
1229 "joint storage component assigned one node more than once",
1230 ));
1231 }
1232 }
1233 for (resource_id, profile) in solution.resource_profiles {
1234 if resource_profiles.insert(resource_id, profile).is_some() {
1235 return Err(invalid_plan(
1236 "joint storage components overlap one resource",
1237 ));
1238 }
1239 }
1240 }
1241 let chosen = chosen
1242 .into_iter()
1243 .collect::<Option<Vec<_>>>()
1244 .ok_or_else(|| invalid_plan("joint storage components omitted one node"))?;
1245 Ok((chosen, resource_profiles))
1246 }
1247
1248 pub(super) fn solve_joint_provider_component(
1249 component: &[usize],
1250 candidate_sets: &[Vec<JointProviderCandidate>],
1251 profile_order: &[DynamicStorageProfile],
1252 ) -> Result<JointComponentSolution, VNextError> {
1253 let mut frontier = BTreeMap::from([(
1254 BTreeMap::<ResourceId, BTreeSet<DynamicStorageProfile>>::new(),
1255 JointPartialSelection::default(),
1256 )]);
1257 for node_index in component {
1258 let mut next = BTreeMap::<
1259 BTreeMap<ResourceId, BTreeSet<DynamicStorageProfile>>,
1260 JointPartialSelection,
1261 >::new();
1262 for (constraints, partial) in frontier {
1263 for candidate in &candidate_sets[*node_index] {
1264 let mut next_constraints = constraints.clone();
1265 if candidate
1266 .allowed_profiles
1267 .iter()
1268 .any(|(resource_id, accepted)| {
1269 !Self::merge_storage_constraint(
1270 &mut next_constraints,
1271 resource_id.clone(),
1272 accepted.clone(),
1273 )
1274 })
1275 {
1276 continue;
1277 }
1278 let mut next_partial = partial.clone();
1279 next_partial.chosen.push(candidate.resources.clone());
1280 next_partial.preferred.push(candidate.is_preferred);
1281 match next.entry(next_constraints) {
1282 std::collections::btree_map::Entry::Vacant(entry) => {
1283 entry.insert(next_partial);
1284 }
1285 std::collections::btree_map::Entry::Occupied(mut entry) => {
1286 if joint_partial_precedes(&next_partial, entry.get()) {
1287 entry.insert(next_partial);
1288 }
1289 }
1290 }
1291 }
1292 }
1293 if next.is_empty() {
1294 return Err(invalid_plan(
1295 "no joint provider/storage assignment satisfies shared resource constraints",
1296 ));
1297 }
1298 frontier = next;
1299 }
1300
1301 let mut best: Option<(JointSelectionObjective, JointComponentSolution)> = None;
1302 for (constraints, partial) in frontier {
1303 let resource_profiles = constraints
1304 .into_iter()
1305 .map(|(resource_id, accepted)| {
1306 let (rank, profile) = profile_order
1307 .iter()
1308 .copied()
1309 .enumerate()
1310 .find(|(_, profile)| accepted.contains(profile))
1311 .ok_or_else(|| {
1312 invalid_plan("joint storage solution lost policy-ordered profile")
1313 })?;
1314 Ok((resource_id, (rank, profile)))
1315 })
1316 .collect::<Result<BTreeMap<_, _>, VNextError>>()?;
1317 let objective = JointSelectionObjective::new(
1318 &partial,
1319 resource_profiles.values().map(|(rank, _)| *rank),
1320 profile_order.len(),
1321 )?;
1322 let solution = JointComponentSolution {
1323 chosen: partial.chosen,
1324 resource_profiles: resource_profiles
1325 .into_iter()
1326 .map(|(resource_id, (_, profile))| (resource_id, profile))
1327 .collect(),
1328 };
1329 if best
1330 .as_ref()
1331 .is_none_or(|(current, _)| objective.precedes(current))
1332 {
1333 best = Some((objective, solution));
1334 }
1335 }
1336 best.map(|(_, solution)| solution).ok_or_else(|| {
1337 invalid_plan("no joint provider/storage assignment satisfies one component")
1338 })
1339 }
1340
1341 pub(super) fn select_provider(
1342 node: &ProgramNode,
1343 operation: &OperationDescriptor,
1344 catalog: &CapabilityCatalog,
1345 resolution_required_capabilities: &BTreeSet<CapabilityId>,
1346 preferred_provider: Option<&ProviderId>,
1347 storage_selected_provider: &ProviderId,
1348 storage_rejection: Option<RejectedProvider>,
1349 required_weight_formats: &BTreeSet<WeightFormatId>,
1350 required_quantization_formats: &BTreeSet<QuantizationFormatId>,
1351 execution_determinism: ExecutionDeterminismRequirement,
1352 ) -> Result<ProviderSelection, VNextError> {
1353 let required_capabilities = operation
1354 .provider
1355 .required_capabilities
1356 .union(resolution_required_capabilities)
1357 .cloned()
1358 .collect::<BTreeSet<_>>();
1359 let request = ProviderCompatibilityRequest::new(
1360 node.operation_id.clone(),
1361 node.required_version,
1362 required_capabilities,
1363 required_weight_formats.clone(),
1364 required_quantization_formats.clone(),
1365 execution_determinism,
1366 )?;
1367 let report = catalog.provider_compatibility(request)?;
1368 report.require_compatible_for_node(&catalog.device().id, &node.id)?;
1369 if !report
1370 .compatible_provider_ids()
1371 .contains(storage_selected_provider)
1372 {
1373 return Err(invalid_plan(format!(
1374 "joint storage solver selected incompatible provider `{storage_selected_provider}`"
1375 )));
1376 }
1377 let selected_provider = storage_selected_provider.clone();
1378 let selection_reason = match preferred_provider {
1379 Some(preferred) if preferred == &selected_provider => {
1380 ProviderSelectionReason::PreferredCompatible
1381 }
1382 Some(_) => ProviderSelectionReason::FallbackFromPreferred,
1383 None => ProviderSelectionReason::CanonicalCompatible,
1384 };
1385 let mut rejected_providers = report
1386 .rejected()
1387 .iter()
1388 .map(|rejection| RejectedProvider {
1389 provider_id: rejection.provider_id.clone(),
1390 reasons: PlanProviderRejectReason::Incompatible(rejection.reasons.clone()),
1391 })
1392 .collect::<Vec<_>>();
1393 if let Some(preferred) = preferred_provider {
1394 let registered = catalog
1395 .providers_for_node(&node.id, &node.operation_id)?
1396 .iter()
1397 .any(|provider| provider.provider_id() == preferred);
1398 if !registered {
1399 rejected_providers.push(RejectedProvider {
1400 provider_id: preferred.clone(),
1401 reasons: PlanProviderRejectReason::NotRegistered,
1402 });
1403 }
1404 }
1405 if let Some(rejection) = storage_rejection {
1406 if rejected_providers
1407 .iter()
1408 .any(|existing| existing.provider_id == rejection.provider_id)
1409 {
1410 return Err(invalid_plan(
1411 "provider has duplicate compatibility and storage rejection evidence",
1412 ));
1413 }
1414 rejected_providers.push(rejection);
1415 }
1416 if let Some(preferred) =
1417 preferred_provider.filter(|preferred| *preferred != &selected_provider)
1418 {
1419 if !rejected_providers
1420 .iter()
1421 .any(|rejection| &rejection.provider_id == preferred)
1422 {
1423 return Err(invalid_plan(format!(
1424 "preferred provider `{preferred}` fallback lacks typed rejection evidence"
1425 )));
1426 }
1427 }
1428 rejected_providers.sort_by(|left, right| left.provider_id.cmp(&right.provider_id));
1429 Ok(ProviderSelection {
1430 requested_provider: preferred_provider.cloned(),
1431 selected_provider,
1432 selection_reason,
1433 rejected_providers,
1434 })
1435 }
1436
1437 pub(super) fn validate_provider_selection_evidence(
1438 selection: &ProviderSelection,
1439 ) -> Result<(), VNextError> {
1440 if selection
1441 .rejected_providers
1442 .windows(2)
1443 .any(|pair| pair[0].provider_id >= pair[1].provider_id)
1444 || selection
1445 .rejected_providers
1446 .iter()
1447 .any(|rejection| rejection.provider_id == selection.selected_provider)
1448 || selection.rejected_providers.iter().any(|rejection| {
1449 matches!(
1450 &rejection.reasons,
1451 PlanProviderRejectReason::StorageIncompatible { resource_ids }
1452 if resource_ids.is_empty()
1453 || resource_ids.windows(2).any(|pair| pair[0] >= pair[1])
1454 )
1455 })
1456 {
1457 return Err(invalid_plan(
1458 "provider rejection evidence is duplicate, non-canonical, or rejects the selected provider",
1459 ));
1460 }
1461 match (
1462 selection.requested_provider.as_ref(),
1463 selection.selection_reason,
1464 ) {
1465 (None, ProviderSelectionReason::CanonicalCompatible) => {}
1466 (Some(requested), ProviderSelectionReason::PreferredCompatible)
1467 if requested == &selection.selected_provider => {}
1468 (Some(requested), ProviderSelectionReason::FallbackFromPreferred)
1469 if requested != &selection.selected_provider
1470 && selection
1471 .rejected_providers
1472 .iter()
1473 .any(|rejection| &rejection.provider_id == requested) => {}
1474 _ => return Err(invalid_plan(
1475 "provider selection reason is inconsistent with preference and rejection evidence",
1476 )),
1477 }
1478 Ok(())
1479 }
1480
1481 pub(super) fn build_memory_plan(
1482 family: &PreparedModelFamily,
1483 device_capacity_bytes: u64,
1484 policy_capacity_bytes: u64,
1485 reserve_bytes: u64,
1486 maximum_active_sequences: u32,
1487 maximum_scheduled_tokens: u64,
1488 nodes: &[PlanNode],
1489 selected_resource_profiles: &BTreeMap<ResourceId, DynamicStorageProfile>,
1490 reusable_execution_policy: Option<&ReusableExecutionPolicy>,
1491 retained_completion_resources: &BTreeSet<ResourceId>,
1492 ) -> Result<MemoryPlan, VNextError> {
1493 validate_scheduled_token_ceiling(maximum_scheduled_tokens)?;
1494 let program_inputs = family
1495 .program()
1496 .inputs()
1497 .iter()
1498 .cloned()
1499 .collect::<BTreeSet<_>>();
1500 let program_outputs = family
1501 .program()
1502 .outputs()
1503 .iter()
1504 .cloned()
1505 .collect::<BTreeSet<_>>();
1506 let state_initializations = family
1507 .program()
1508 .states()
1509 .iter()
1510 .map(|state| (state.value_id.clone(), state.initialization))
1511 .collect::<BTreeMap<_, _>>();
1512 let mut values = BTreeMap::<ResourceId, ValueAllocationAccumulator>::new();
1513 let mut static_allocations = Vec::new();
1514 let mut dynamic_descriptors = Vec::new();
1515 let workspace_layout_fingerprint = workspace_storage_layout_fingerprint()?;
1516 for node in nodes {
1517 let value_alignment = node.provider_resources.value_alignment_bytes;
1518 for binding in &node.values {
1519 let logical_layout_fingerprint =
1520 tensor_storage_layout_fingerprint(binding.tensor().layout())?;
1521 for component in binding.storage().components() {
1522 if component.offset_bytes() % value_alignment != 0 {
1523 return Err(invalid_plan(format!(
1524 "resource `{}` offset is not aligned for provider `{}`",
1525 component.resource_id(),
1526 node.provider_resources.provider_id
1527 )));
1528 }
1529 let end = component
1530 .offset_bytes()
1531 .checked_add(component.length_bytes())
1532 .ok_or_else(|| invalid_plan("resource byte range overflows u64"))?;
1533 let token_projection = node
1534 .work
1535 .token_projection(binding.role(), binding.ordinal())
1536 .map(|projection| {
1537 if component.offset_bytes() != 0
1538 || component.length_bytes() % projection.canonical_extent() != 0
1539 {
1540 return Err(invalid_plan(format!(
1541 "token-scaled resource `{}` is not one exact canonical tensor range",
1542 component.resource_id()
1543 )));
1544 }
1545 Ok((
1546 component.length_bytes() / projection.canonical_extent(),
1547 projection.canonical_extent(),
1548 ))
1549 })
1550 .transpose()?;
1551 let demand = Self::value_resource_demand(
1552 family,
1553 binding.value_id(),
1554 binding.usage(),
1555 end,
1556 token_projection,
1557 maximum_active_sequences,
1558 maximum_scheduled_tokens,
1559 &program_inputs,
1560 &program_outputs,
1561 )?;
1562 let initialization = state_initializations
1563 .get(binding.value_id())
1564 .copied()
1565 .unwrap_or(StateInitialization::None);
1566 values
1567 .entry(component.resource_id().clone())
1568 .and_modify(|allocation| {
1569 allocation.merge_result =
1570 allocation.merge_result.take().and_then(|_| {
1571 allocation.merge(
1572 end,
1573 value_alignment,
1574 binding.usage(),
1575 component.element_type(),
1576 demand,
1577 initialization,
1578 logical_layout_fingerprint.clone(),
1579 )
1580 });
1581 })
1582 .or_insert_with(|| ValueAllocationAccumulator {
1583 end_bytes: end,
1584 alignment_bytes: value_alignment,
1585 usage: binding.usage(),
1586 element_type: component.element_type(),
1587 demand,
1588 initialization,
1589 logical_layout_fingerprints: BTreeSet::from([
1590 logical_layout_fingerprint.clone(),
1591 ]),
1592 merge_result: Some(()),
1593 });
1594 }
1595 }
1596 if let Some(workspace) = &node.provider_resources.scratch {
1597 let resource_id = node.scratch_resource.clone().ok_or_else(|| {
1598 invalid_plan(format!(
1599 "node `{}` scratch base identity is missing",
1600 node.id
1601 ))
1602 })?;
1603 if workspace.scope != ProviderWorkspaceScope::Invocation {
1604 return Err(invalid_plan(format!(
1605 "node `{}` scratch workspace is not invocation scoped",
1606 node.id
1607 )));
1608 }
1609 let storage = DynamicStorageContract::new(
1610 *selected_resource_profiles
1611 .get(&resource_id)
1612 .ok_or_else(|| {
1613 invalid_plan(format!(
1614 "scratch resource `{resource_id}` has no selected storage profile"
1615 ))
1616 })?,
1617 workspace_layout_fingerprint.clone(),
1618 )?;
1619 dynamic_descriptors.push(DynamicResourceDescriptor::new(
1620 resource_id,
1621 workspace
1622 .size_formula
1623 .bind_runtime_limits(maximum_active_sequences, maximum_scheduled_tokens)?,
1624 workspace.alignment_bytes,
1625 BufferUsage::Scratch,
1626 ElementType::U8,
1627 AllocationLifetime::Invocation,
1628 AllocationKind::Scratch {
1629 node_id: node.id.clone(),
1630 },
1631 storage,
1632 StateInitialization::None,
1633 maximum_active_sequences,
1634 )?);
1635 } else if node.scratch_resource.is_some() {
1636 return Err(invalid_plan(format!(
1637 "node `{}` has scratch resources without a provider estimate",
1638 node.id
1639 )));
1640 }
1641 if let Some(workspace) = &node.provider_resources.binding {
1642 let resource_id = node.binding_resource.clone().ok_or_else(|| {
1643 invalid_plan(format!(
1644 "node `{}` binding workspace base identity is missing",
1645 node.id
1646 ))
1647 })?;
1648 if workspace.scope != ProviderWorkspaceScope::Invocation {
1649 return Err(invalid_plan(format!(
1650 "node `{}` binding workspace is not invocation scoped",
1651 node.id
1652 )));
1653 }
1654 let storage = DynamicStorageContract::new(
1655 *selected_resource_profiles
1656 .get(&resource_id)
1657 .ok_or_else(|| {
1658 invalid_plan(format!(
1659 "binding resource `{resource_id}` has no selected storage profile"
1660 ))
1661 })?,
1662 workspace_layout_fingerprint.clone(),
1663 )?;
1664 dynamic_descriptors.push(DynamicResourceDescriptor::new(
1665 resource_id,
1666 workspace
1667 .size_formula
1668 .bind_runtime_limits(maximum_active_sequences, maximum_scheduled_tokens)?,
1669 workspace.alignment_bytes,
1670 BufferUsage::Binding,
1671 ElementType::U8,
1672 AllocationLifetime::Invocation,
1673 AllocationKind::Binding {
1674 node_id: node.id.clone(),
1675 },
1676 storage,
1677 StateInitialization::None,
1678 maximum_active_sequences,
1679 )?);
1680 } else if node.binding_resource.is_some() {
1681 return Err(invalid_plan(format!(
1682 "node `{}` has binding resources without a provider estimate",
1683 node.id
1684 )));
1685 }
1686 if let Some(workspace) = &node.provider_resources.persistent {
1687 let resource_id = node.persistent_resource.clone().ok_or_else(|| {
1688 invalid_plan(format!(
1689 "node `{}` persistent base identity is missing",
1690 node.id
1691 ))
1692 })?;
1693 match workspace.scope {
1694 ProviderWorkspaceScope::Plan => {
1695 let bytes = workspace.fixed_bytes().ok_or_else(|| {
1696 invalid_plan(format!(
1697 "node `{}` plan workspace does not have a fixed formula",
1698 node.id
1699 ))
1700 })?;
1701 let storage = DynamicStorageContract::new(
1702 *selected_resource_profiles
1703 .get(&resource_id)
1704 .ok_or_else(|| {
1705 invalid_plan(format!(
1706 "plan workspace `{resource_id}` has no selected storage profile"
1707 ))
1708 })?,
1709 workspace_layout_fingerprint.clone(),
1710 )?;
1711 static_allocations.push(ResourceAllocation::new(
1712 resource_id,
1713 bytes,
1714 workspace.alignment_bytes,
1715 BufferUsage::Persistent,
1716 ElementType::U8,
1717 AllocationKind::Persistent {
1718 node_id: node.id.clone(),
1719 },
1720 storage,
1721 )?);
1722 }
1723 scope @ (ProviderWorkspaceScope::Request
1724 | ProviderWorkspaceScope::Sequence
1725 | ProviderWorkspaceScope::Step) => {
1726 let lifetime = match scope {
1727 ProviderWorkspaceScope::Request => AllocationLifetime::Request,
1728 ProviderWorkspaceScope::Sequence => AllocationLifetime::Sequence,
1729 ProviderWorkspaceScope::Step => AllocationLifetime::Step,
1730 ProviderWorkspaceScope::Plan | ProviderWorkspaceScope::Invocation => {
1731 unreachable!()
1732 }
1733 };
1734 let storage = DynamicStorageContract::new(
1735 *selected_resource_profiles.get(&resource_id).ok_or_else(|| {
1736 invalid_plan(format!(
1737 "persistent resource `{resource_id}` has no selected storage profile"
1738 ))
1739 })?,
1740 workspace_layout_fingerprint.clone(),
1741 )?;
1742 dynamic_descriptors.push(DynamicResourceDescriptor::new(
1743 resource_id,
1744 workspace.size_formula.bind_runtime_limits(
1745 maximum_active_sequences,
1746 maximum_scheduled_tokens,
1747 )?,
1748 workspace.alignment_bytes,
1749 BufferUsage::Persistent,
1750 ElementType::U8,
1751 lifetime,
1752 AllocationKind::Persistent {
1753 node_id: node.id.clone(),
1754 },
1755 storage,
1756 StateInitialization::None,
1757 maximum_active_sequences,
1758 )?);
1759 }
1760 ProviderWorkspaceScope::Invocation => {
1761 return Err(invalid_plan(format!(
1762 "node `{}` persistent workspace cannot be invocation scoped",
1763 node.id
1764 )));
1765 }
1766 }
1767 } else if node.persistent_resource.is_some() {
1768 return Err(invalid_plan(format!(
1769 "node `{}` has persistent resources without a provider estimate",
1770 node.id
1771 )));
1772 }
1773 }
1774 for (resource_id, accumulator) in values {
1775 accumulator.merge_result.ok_or_else(|| {
1776 invalid_plan(format!(
1777 "resource `{resource_id}` has conflicting usage, dtype, lifetime, or demand"
1778 ))
1779 })?;
1780 let logical_layout_fingerprint = canonical_fingerprint(
1781 &accumulator.logical_layout_fingerprints,
1782 "fingerprint dynamic resource tensor layout classes",
1783 )?;
1784 match accumulator.demand {
1785 ValueResourceDemand::PlanStatic => {
1786 let storage = DynamicStorageContract::new(
1787 static_contiguous_storage_profile()?,
1788 logical_layout_fingerprint,
1789 )?;
1790 static_allocations.push(ResourceAllocation::new(
1791 resource_id,
1792 accumulator.end_bytes,
1793 accumulator.alignment_bytes,
1794 accumulator.usage,
1795 accumulator.element_type,
1796 AllocationKind::Value,
1797 storage,
1798 )?);
1799 }
1800 demand => {
1801 let storage = DynamicStorageContract::new(
1802 *selected_resource_profiles.get(&resource_id).ok_or_else(|| {
1803 invalid_plan(format!(
1804 "dynamic value resource `{resource_id}` has no selected storage profile"
1805 ))
1806 })?,
1807 logical_layout_fingerprint,
1808 )?;
1809 dynamic_descriptors.push(DynamicResourceDescriptor::new(
1810 resource_id,
1811 demand
1812 .dynamic_demand(accumulator.end_bytes, accumulator.alignment_bytes)?,
1813 accumulator.alignment_bytes,
1814 accumulator.usage,
1815 accumulator.element_type,
1816 demand.lifetime().ok_or_else(|| {
1817 invalid_plan("dynamic value demand lost its scoped lifetime")
1818 })?,
1819 AllocationKind::Value,
1820 storage,
1821 accumulator.initialization,
1822 maximum_active_sequences,
1823 )?);
1824 }
1825 }
1826 }
1827 MemoryPlan::from_core_with_completion_retention(
1828 device_capacity_bytes,
1829 policy_capacity_bytes,
1830 reserve_bytes,
1831 maximum_active_sequences,
1832 static_allocations,
1833 dynamic_descriptors,
1834 nodes,
1835 reusable_execution_policy,
1836 retained_completion_resources,
1837 )
1838 }
1839
1840 pub(super) fn value_resource_demand(
1841 family: &PreparedModelFamily,
1842 value_id: &ProgramValueId,
1843 usage: BufferUsage,
1844 minimum_bytes: u64,
1845 token_projection: Option<(u64, u64)>,
1846 maximum_active_sequences: u32,
1847 maximum_scheduled_tokens: u64,
1848 program_inputs: &BTreeSet<ProgramValueId>,
1849 program_outputs: &BTreeSet<ProgramValueId>,
1850 ) -> Result<ValueResourceDemand, VNextError> {
1851 if family
1852 .program()
1853 .weights()
1854 .iter()
1855 .any(|weight| &weight.value_id == value_id)
1856 {
1857 return Ok(ValueResourceDemand::PlanStatic);
1858 }
1859 let state = family
1860 .program()
1861 .states()
1862 .iter()
1863 .find(|state| &state.value_id == value_id);
1864 let Some(state) = state else {
1865 if usage != BufferUsage::Activations {
1866 return Err(invalid_plan(format!(
1867 "non-state value `{value_id}` is not backed by activation memory"
1868 )));
1869 }
1870 let is_product_io =
1871 program_inputs.contains(value_id) || program_outputs.contains(value_id);
1872 let lifetime = AllocationLifetime::Step;
1873 if let Some((bytes_per_token, canonical_tokens)) = token_projection {
1874 if bytes_per_token == 0 || canonical_tokens == 0 {
1875 return Err(invalid_plan(
1876 "token-scaled activation has zero bytes or canonical tokens",
1877 ));
1878 }
1879 let maximum_tokens = Self::activation_token_capacity(
1880 lifetime,
1881 canonical_tokens,
1882 maximum_scheduled_tokens,
1883 )?;
1884 return Ok(ValueResourceDemand::TokenScaled {
1885 lifetime,
1886 bytes_per_token,
1887 maximum_tokens,
1888 });
1889 }
1890 if is_product_io {
1891 return Ok(ValueResourceDemand::ParticipantFixed {
1892 lifetime,
1893 maximum_participants: maximum_active_sequences,
1894 });
1895 }
1896 return Ok(ValueResourceDemand::Fixed { lifetime });
1897 };
1898 let lifetime = match state.lifetime {
1899 StateLifetime::Request => AllocationLifetime::Request,
1900 StateLifetime::Sequence => AllocationLifetime::Sequence,
1901 StateLifetime::Step => AllocationLifetime::Step,
1902 };
1903 state.capacity_demand.validate(state.tensor.byte_len()?)?;
1904 match state.capacity_demand {
1905 StateCapacityDemand::FixedPerScope => Ok(ValueResourceDemand::Fixed { lifetime }),
1906 StateCapacityDemand::TokenScaled {
1907 bytes_per_token,
1908 maximum_tokens,
1909 } => {
1910 if bytes_per_token < minimum_bytes {
1911 return Err(invalid_plan(
1912 "token-scaled state demand is smaller than its resolved resource range",
1913 ));
1914 }
1915 Ok(ValueResourceDemand::TokenScaled {
1916 lifetime,
1917 bytes_per_token,
1918 maximum_tokens,
1919 })
1920 }
1921 }
1922 }
1923
1924 pub(super) fn activation_token_capacity(
1925 lifetime: AllocationLifetime,
1926 canonical_tokens: u64,
1927 maximum_scheduled_tokens: u64,
1928 ) -> Result<u64, VNextError> {
1929 if canonical_tokens == 0 {
1930 return Err(invalid_plan(
1931 "token-scaled activation has zero canonical tokens",
1932 ));
1933 }
1934 if lifetime == AllocationLifetime::Request {
1935 Ok(canonical_tokens)
1936 } else {
1937 validate_scheduled_token_ceiling(maximum_scheduled_tokens)?;
1938 Ok(maximum_scheduled_tokens)
1939 }
1940 }
1941
1942 pub(super) fn plan_id_for_hash(hash: &PlanHash) -> Result<PlanId, VNextError> {
1943 PlanId::new(format!("plan/sha256/{}", hash.as_str()))
1944 }
1945
1946 pub(super) fn validate_internal(&self) -> Result<(), VNextError> {
1947 if self.payload.schema != EXECUTION_PLAN_SCHEMA {
1948 return Err(VNextError::UnsupportedPlanSchema {
1949 expected_major: EXECUTION_PLAN_SCHEMA.major,
1950 expected_minor: EXECUTION_PLAN_SCHEMA.minor,
1951 actual_major: self.payload.schema.major,
1952 actual_minor: self.payload.schema.minor,
1953 });
1954 }
1955 let computed = PlanHash::new(canonical_fingerprint(
1956 &PlanHashMaterial::from(&self.payload),
1957 "validate execution plan hash",
1958 )?)?;
1959 if computed != self.plan_hash {
1960 return Err(VNextError::PlanHashMismatch {
1961 expected: computed.to_string(),
1962 actual: self.plan_hash.to_string(),
1963 });
1964 }
1965 if self.payload.plan_id != Self::plan_id_for_hash(&computed)? {
1966 return Err(invalid_plan(
1967 "plan id is not derived from the semantic plan hash",
1968 ));
1969 }
1970 if self.payload.nodes.is_empty()
1971 || !is_canonical_sha256(&self.payload.prepared_family_fingerprint)
1972 || !is_canonical_sha256(&self.payload.program_fingerprint)
1973 || !is_canonical_sha256(&self.payload.capability_catalog_fingerprint)
1974 || !is_canonical_sha256(&self.payload.device_runtime_implementation_fingerprint)
1975 || !is_canonical_sha256(&self.payload.policy_fingerprint)
1976 || self.payload.maximum_scheduled_tokens == 0
1977 {
1978 return Err(invalid_plan("plan provenance or node set is invalid"));
1979 }
1980 self.payload
1981 .execution_weights
1982 .validate_structure(&self.payload.family_id)?;
1983 if &self.payload.execution_weights != self.trusted_execution_weights.plan()
1984 || self.payload.weight_format != self.payload.execution_weights.schema().format_id
1985 || self.payload.quantization_formats
1986 != self
1987 .payload
1988 .execution_weights
1989 .schema()
1990 .quantization_formats()
1991 {
1992 return Err(invalid_plan(
1993 "execution weight summary differs from the execution weight plan",
1994 ));
1995 }
1996 let retention_spec = CompletionRetentionSpec::new(
1997 self.payload
1998 .retained_completion_values
1999 .iter()
2000 .map(|value| value.value_id().clone())
2001 .collect(),
2002 );
2003 let expected_retained_completion_values =
2004 resolve_retained_completion_values(&self.payload.nodes, &retention_spec)?;
2005 if self.payload.retained_completion_values != expected_retained_completion_values {
2006 return Err(invalid_plan(
2007 "retained completion values are not derived from plan outputs",
2008 ));
2009 }
2010 if self.payload.terminal_output_resources.is_empty()
2011 || self
2012 .payload
2013 .terminal_output_resources
2014 .windows(2)
2015 .any(|pair| pair[0] >= pair[1])
2016 {
2017 return Err(invalid_plan(
2018 "terminal output resource evidence is empty or non-canonical",
2019 ));
2020 }
2021 let retained_completion_resources = self
2022 .payload
2023 .retained_completion_values
2024 .iter()
2025 .map(|value| value.resource_id().clone())
2026 .chain(self.payload.terminal_output_resources.iter().cloned())
2027 .collect::<BTreeSet<_>>();
2028 self.payload.memory.validate()?;
2029 let dynamic_capacity_bytes = self
2030 .payload
2031 .memory
2032 .usable_capacity_bytes
2033 .checked_sub(self.payload.memory.static_bytes)
2034 .ok_or_else(|| invalid_plan("static memory exceeds usable capacity"))?;
2035 let base_pools = MemoryPlan::derive_dynamic_pools_with_completion_retention(
2036 &self.payload.memory.dynamic_descriptors,
2037 &self.payload.nodes,
2038 dynamic_capacity_bytes,
2039 &retained_completion_resources,
2040 )?;
2041 let expected_reusable_execution = self
2042 .payload
2043 .memory
2044 .reusable_execution
2045 .as_ref()
2046 .map(|actual| {
2047 MemoryPlan::derive_reusable_execution(
2048 &actual.policy()?,
2049 self.payload.nodes.len(),
2050 &self.payload.memory.dynamic_descriptors,
2051 &base_pools,
2052 )
2053 })
2054 .transpose()?;
2055 if self.payload.memory.reusable_execution != expected_reusable_execution {
2056 return Err(invalid_plan(
2057 "reusable execution budgets are not derived from plan resources",
2058 ));
2059 }
2060 let reusable_workspace_ceilings = expected_reusable_execution
2061 .as_ref()
2062 .map(ReusableExecutionMemoryPlan::pool_workspace_ceilings)
2063 .transpose()?
2064 .unwrap_or_default();
2065 let expected_pools = MemoryPlan::derive_dynamic_pools_with_reusable(
2066 &self.payload.memory.dynamic_descriptors,
2067 &self.payload.nodes,
2068 dynamic_capacity_bytes,
2069 &reusable_workspace_ceilings,
2070 &retained_completion_resources,
2071 )?;
2072 if self.payload.memory.dynamic_pools != expected_pools {
2073 return Err(invalid_plan(
2074 "memory pools or invocation reuse are not derived from plan dependencies",
2075 ));
2076 }
2077 let static_allocations = self
2078 .payload
2079 .memory
2080 .static_allocations
2081 .iter()
2082 .map(|allocation| (allocation.resource_id.clone(), allocation))
2083 .collect::<BTreeMap<_, _>>();
2084 let dynamic_descriptors = self
2085 .payload
2086 .memory
2087 .dynamic_descriptors
2088 .iter()
2089 .map(|descriptor| (descriptor.base_resource_id.clone(), descriptor))
2090 .collect::<BTreeMap<_, _>>();
2091 let mut seen_nodes = BTreeSet::new();
2092 let mut canonical_values = BTreeMap::new();
2093 for node in &self.payload.nodes {
2094 node.provider_resources.validate_shape()?;
2095 Self::validate_provider_selection_evidence(&node.selection)?;
2096 Self::validate_node_work_contract(node)?;
2097 if !seen_nodes.insert(node.id.clone())
2098 || !is_canonical_sha256(&node.provider_implementation_fingerprint)
2099 || node
2100 .dependencies
2101 .iter()
2102 .any(|dependency| dependency == &node.id || !seen_nodes.contains(dependency))
2103 || node.dependencies.windows(2).any(|pair| pair[0] >= pair[1])
2104 || node
2105 .state_effects
2106 .windows(2)
2107 .any(|pair| pair[0].state_id >= pair[1].state_id)
2108 || node.resources.iter().collect::<BTreeSet<_>>().len() != node.resources.len()
2109 || node.resources.iter().any(|resource| {
2110 !static_allocations.contains_key(resource)
2111 && !dynamic_descriptors.contains_key(resource)
2112 })
2113 || node.provider_resources.provider_id != node.selection.selected_provider
2114 {
2115 return Err(invalid_plan(format!(
2116 "node `{}` identity, dependency, or resource closure is invalid",
2117 node.id
2118 )));
2119 }
2120 let expected_resources = node
2121 .values
2122 .iter()
2123 .flat_map(|binding| binding.storage().components())
2124 .map(|component| component.resource_id().clone())
2125 .chain(node.scratch_resource.iter().cloned())
2126 .chain(node.binding_resource.iter().cloned())
2127 .chain(node.persistent_resource.iter().cloned())
2128 .collect::<BTreeSet<_>>()
2129 .into_iter()
2130 .collect::<Vec<_>>();
2131 if node.resources != expected_resources {
2132 return Err(invalid_plan(format!(
2133 "node `{}` resource closure is not canonical",
2134 node.id
2135 )));
2136 }
2137 for effect in &node.state_effects {
2138 if !matches!(
2139 effect.lifetime,
2140 AllocationLifetime::Request
2141 | AllocationLifetime::Sequence
2142 | AllocationLifetime::Step
2143 ) || effect.resource_ids.is_empty()
2144 || effect
2145 .resource_ids
2146 .windows(2)
2147 .any(|pair| pair[0] >= pair[1])
2148 {
2149 return Err(invalid_plan(format!(
2150 "node `{}` state effect has an invalid lifetime or resource closure",
2151 node.id
2152 )));
2153 }
2154 let matching = node
2155 .values
2156 .iter()
2157 .filter(|binding| binding.value_id() == &effect.state_value_id)
2158 .collect::<Vec<_>>();
2159 let expected_effect_resources = matching
2160 .iter()
2161 .flat_map(|binding| binding.storage().components())
2162 .map(|component| component.resource_id().clone())
2163 .collect::<BTreeSet<_>>()
2164 .into_iter()
2165 .collect::<Vec<_>>();
2166 let reads = matching.iter().any(|binding| {
2167 matches!(
2168 binding.access(),
2169 TensorAccess::Read | TensorAccess::ReadWrite
2170 )
2171 });
2172 let writes = matching.iter().any(|binding| {
2173 matches!(
2174 binding.access(),
2175 TensorAccess::Write | TensorAccess::ReadWrite
2176 )
2177 });
2178 let expected_access = match (reads, writes) {
2179 (true, false) => Some(TensorAccess::Read),
2180 (false, true) => Some(TensorAccess::Write),
2181 (true, true) => Some(TensorAccess::ReadWrite),
2182 (false, false) => None,
2183 };
2184 if effect.resource_ids != expected_effect_resources
2185 || expected_access != Some(effect.access)
2186 || effect.resource_ids.iter().any(|resource_id| {
2187 dynamic_descriptors
2188 .get(resource_id)
2189 .is_none_or(|descriptor| descriptor.lifetime != effect.lifetime)
2190 })
2191 {
2192 return Err(invalid_plan(format!(
2193 "node `{}` state effect is not derived from its typed bindings",
2194 node.id
2195 )));
2196 }
2197 }
2198 if node.scratch_resource.is_some() != node.provider_resources.scratch.is_some()
2199 || node.binding_resource.is_some() != node.provider_resources.binding.is_some()
2200 || node.persistent_resource.is_some()
2201 != node.provider_resources.persistent.is_some()
2202 {
2203 return Err(invalid_plan(format!(
2204 "node `{}` workspace base identity presence differs from its provider estimate",
2205 node.id
2206 )));
2207 }
2208 if let Some(resource_id) = &node.scratch_resource {
2209 let descriptor = dynamic_descriptors.get(resource_id).ok_or_else(|| {
2210 invalid_plan(format!("node `{}` scratch descriptor is missing", node.id))
2211 })?;
2212 let workspace = node.provider_resources.scratch.as_ref().ok_or_else(|| {
2213 invalid_plan(format!("node `{}` scratch estimate is missing", node.id))
2214 })?;
2215 if descriptor.demand
2216 != workspace.size_formula.bind_runtime_limits(
2217 self.payload.memory.maximum_active_sequences,
2218 self.payload.maximum_scheduled_tokens,
2219 )?
2220 || descriptor.alignment_bytes != workspace.alignment_bytes
2221 || descriptor.usage != BufferUsage::Scratch
2222 || descriptor.lifetime != AllocationLifetime::Invocation
2223 || descriptor.theoretical_maximum_instances
2224 != self.payload.memory.maximum_active_sequences
2225 || descriptor.kind
2226 != (AllocationKind::Scratch {
2227 node_id: node.id.clone(),
2228 })
2229 {
2230 return Err(invalid_plan(format!(
2231 "node `{}` scratch descriptor differs from its provider estimate",
2232 node.id
2233 )));
2234 }
2235 }
2236 if let Some(resource_id) = &node.binding_resource {
2237 let descriptor = dynamic_descriptors.get(resource_id).ok_or_else(|| {
2238 invalid_plan(format!("node `{}` binding descriptor is missing", node.id))
2239 })?;
2240 let workspace = node.provider_resources.binding.as_ref().ok_or_else(|| {
2241 invalid_plan(format!("node `{}` binding estimate is missing", node.id))
2242 })?;
2243 if descriptor.demand
2244 != workspace.size_formula.bind_runtime_limits(
2245 self.payload.memory.maximum_active_sequences,
2246 self.payload.maximum_scheduled_tokens,
2247 )?
2248 || descriptor.alignment_bytes != workspace.alignment_bytes
2249 || descriptor.usage != BufferUsage::Binding
2250 || descriptor.lifetime != AllocationLifetime::Invocation
2251 || descriptor.theoretical_maximum_instances
2252 != self.payload.memory.maximum_active_sequences
2253 || descriptor.kind
2254 != (AllocationKind::Binding {
2255 node_id: node.id.clone(),
2256 })
2257 {
2258 return Err(invalid_plan(format!(
2259 "node `{}` binding descriptor differs from its provider estimate",
2260 node.id
2261 )));
2262 }
2263 }
2264 if let Some(resource_id) = &node.persistent_resource {
2265 let workspace = node.provider_resources.persistent.as_ref().ok_or_else(|| {
2266 invalid_plan(format!("node `{}` persistent estimate is missing", node.id))
2267 })?;
2268 match workspace.scope {
2269 ProviderWorkspaceScope::Plan => {
2270 let allocation = static_allocations.get(resource_id).ok_or_else(|| {
2271 invalid_plan(format!(
2272 "node `{}` plan-static persistent allocation is missing",
2273 node.id
2274 ))
2275 })?;
2276 if Some(allocation.per_instance_bytes) != workspace.fixed_bytes()
2277 || allocation.alignment_bytes != workspace.alignment_bytes
2278 || allocation.usage != BufferUsage::Persistent
2279 || !workspace.storage.accepts(allocation.storage.profile())
2280 || allocation.storage.logical_layout_fingerprint()
2281 != workspace_storage_layout_fingerprint()?
2282 || allocation.kind
2283 != (AllocationKind::Persistent {
2284 node_id: node.id.clone(),
2285 })
2286 {
2287 return Err(invalid_plan(format!(
2288 "node `{}` plan-static persistent allocation differs from its provider estimate",
2289 node.id
2290 )));
2291 }
2292 }
2293 scope @ (ProviderWorkspaceScope::Request
2294 | ProviderWorkspaceScope::Sequence
2295 | ProviderWorkspaceScope::Step) => {
2296 let expected_lifetime = match scope {
2297 ProviderWorkspaceScope::Request => AllocationLifetime::Request,
2298 ProviderWorkspaceScope::Sequence => AllocationLifetime::Sequence,
2299 ProviderWorkspaceScope::Step => AllocationLifetime::Step,
2300 ProviderWorkspaceScope::Plan | ProviderWorkspaceScope::Invocation => {
2301 unreachable!()
2302 }
2303 };
2304 let descriptor = dynamic_descriptors.get(resource_id).ok_or_else(|| {
2305 invalid_plan(format!(
2306 "node `{}` dynamic persistent descriptor is missing",
2307 node.id
2308 ))
2309 })?;
2310 if descriptor.demand
2311 != workspace.size_formula.bind_runtime_limits(
2312 self.payload.memory.maximum_active_sequences,
2313 self.payload.maximum_scheduled_tokens,
2314 )?
2315 || descriptor.alignment_bytes != workspace.alignment_bytes
2316 || descriptor.usage != BufferUsage::Persistent
2317 || descriptor.lifetime != expected_lifetime
2318 || descriptor.theoretical_maximum_instances
2319 != self.payload.memory.maximum_active_sequences
2320 || descriptor.kind
2321 != (AllocationKind::Persistent {
2322 node_id: node.id.clone(),
2323 })
2324 {
2325 return Err(invalid_plan(format!(
2326 "node `{}` dynamic persistent descriptor differs from its provider estimate",
2327 node.id
2328 )));
2329 }
2330 }
2331 ProviderWorkspaceScope::Invocation => {
2332 return Err(invalid_plan(format!(
2333 "node `{}` persistent workspace cannot be invocation scoped",
2334 node.id
2335 )));
2336 }
2337 }
2338 }
2339 for binding in &node.values {
2340 Self::validate_cross_node_value(binding, &mut canonical_values)?;
2341 }
2342 }
2343 Self::validate_global_storage_aliasing(&canonical_values, &self.payload.nodes)?;
2344 Ok(())
2345 }
2346
2347 pub fn payload(&self) -> &ExecutionPlanPayload {
2348 &self.payload
2349 }
2350
2351 pub fn completion_checkpoint(
2352 &self,
2353 value_id: &ProgramValueId,
2354 ) -> Result<&RetainedCompletionValue, VNextError> {
2355 self.payload
2356 .retained_completion_values
2357 .binary_search_by(|value| value.value_id().cmp(value_id))
2358 .map(|index| &self.payload.retained_completion_values[index])
2359 .map_err(|_| {
2360 invalid_plan(format!(
2361 "semantic value `{value_id}` is not retained for completion readback"
2362 ))
2363 })
2364 }
2365
2366 pub fn completion_checkpoint_readback_for_work(
2370 &self,
2371 value_id: &ProgramValueId,
2372 participant_index: u32,
2373 work: &ResourceWorkShape,
2374 ) -> Result<CompletionReadbackRequest, VNextError> {
2375 let checkpoint = self.completion_checkpoint(value_id)?;
2376 if checkpoint.logical_offset_bytes() != 0 {
2377 return Err(invalid_plan(
2378 "work-shaped completion readback requires a whole-resource activation",
2379 ));
2380 }
2381 let descriptor = self
2382 .payload
2383 .memory
2384 .dynamic_descriptors()
2385 .iter()
2386 .find(|descriptor| descriptor.base_resource_id() == checkpoint.resource_id())
2387 .ok_or_else(|| {
2388 invalid_plan(format!(
2389 "retained completion resource `{}` has no dynamic descriptor",
2390 checkpoint.resource_id()
2391 ))
2392 })?;
2393 if descriptor.element_type() != checkpoint.tensor().element_type() {
2394 return Err(invalid_plan(
2395 "retained completion resource element type differs from its tensor",
2396 ));
2397 }
2398 let byte_len = match descriptor.demand() {
2399 DynamicResourceDemand::ActualSequences { .. } => {
2402 checkpoint.tensor().minimum_storage_bytes()?
2403 }
2404 _ => descriptor.evaluate_logical_request_bytes(work)?,
2405 };
2406 let element_bytes = checkpoint.tensor().element_type().size_bytes();
2407 if byte_len % element_bytes != 0 {
2408 return Err(invalid_plan(
2409 "retained completion byte extent is not element aligned",
2410 ));
2411 }
2412 checkpoint.readback_request(
2413 participant_index,
2414 HostTransferLayout::new(checkpoint.tensor().element_type(), byte_len / element_bytes)?,
2415 )
2416 }
2417
2418 pub fn plan_hash(&self) -> &PlanHash {
2419 &self.plan_hash
2420 }
2421
2422 pub(crate) fn operation_registry_authority(&self) -> &OperationRegistryAuthority {
2423 &self.operation_registry_authority
2424 }
2425
2426 pub(crate) fn materialize_weight_components<'source>(
2427 &self,
2428 family: &PreparedModelFamily,
2429 source: &'source dyn WeightComponentSource,
2430 components: &[&WeightComponentSpec],
2431 ) -> Result<Vec<WeightComponentPayload<'source>>, VNextError> {
2432 self.trusted_execution_weights
2433 .materialize_components(family, source, components)
2434 }
2435}