1use std::collections::BTreeMap;
2use std::fmt;
3use std::num::NonZeroU64;
4use std::sync::atomic::{AtomicU64, Ordering};
5use std::sync::{Arc, RwLock};
6
7use tenferro_tensor::{Tensor, TensorValue};
8
9use crate::graph::CompiledGraph;
10use crate::program::FrozenProgram;
11
12use super::cache::{PreparedPlanCacheLimits, RuntimeCacheSet};
13use super::cache_owner::{FrozenCacheOwner, FrozenCacheOwnerKind};
14use super::engine_registration::{CandidateRegistrationToken, EngineRegistrationState};
15use super::execution;
16#[cfg(test)]
17use super::extension::ExtensionSlotFullForTest;
18use super::extension::{
19 bind_candidate_module, configure_module, freeze_extension_slots, BoundCandidateModuleRecord,
20 CandidateModuleRecord, CandidateRegistrationIdentity, ExtensionEngineSnapshotView,
21 ExtensionFamilyId, FrozenExtensionSlots,
22};
23use super::preparation::{PreparedEntryKey, PreparedProgram, PreparedProgramResult};
24use super::schedule::EventDomainId;
25use super::{
26 CacheOwnerId, CoreCapabilityBundle, EngineId, EngineRegistration, ExecutionContextIdentity,
27 ExecutionPolicy, ExtensionModule, ExtensionModuleError, ExtensionModuleId,
28 FrozenTransferRegistry, HardwareClassId, InputSignature, PrepareCapability, PrepareOptions,
29 ProviderDeviceIdentity, RegistrationIdentity, RegistrationKey, ResolvedTransferEndpoint,
30 ResolvedTransferRoute, RuntimeCacheError, RuntimeCacheStats, RuntimeConfigError, RuntimeEpoch,
31 RuntimeId, RuntimeReconfigureError, RuntimeStateError, StorageClass, TransferEndpoint,
32 TransferProvider, TransferRoute,
33};
34use crate::{Error, ErrorPhase};
35
36static NEXT_RUNTIME_ID: AtomicU64 = AtomicU64::new(1);
37static NEXT_REGISTRATION_ISSUER: AtomicU64 = AtomicU64::new(1);
38const INITIAL_REGISTRATION_ORDINAL: NonZeroU64 = NonZeroU64::MIN;
39
40#[derive(Clone, Debug)]
41struct CandidateEngineRecord {
42 registration: EngineRegistration,
43 identity: CandidateRegistrationIdentity,
44}
45
46#[derive(Clone, Debug)]
47struct BoundCandidateEngineRecord {
48 registration: EngineRegistration,
49 identity: RegistrationIdentity,
50}
51
52#[derive(Clone, Debug)]
53enum CandidateTransferBinding {
54 New,
56 Preserved {
58 source: ProviderDeviceIdentity,
59 destination: ProviderDeviceIdentity,
60 },
61}
62
63#[derive(Clone, Debug)]
64struct CandidateTransferRecord {
65 provider: Arc<dyn TransferProvider>,
66 binding: CandidateTransferBinding,
67}
68
69struct BoundCandidateTransferRecord {
70 provider: Arc<dyn TransferProvider>,
71 source: ProviderDeviceIdentity,
72 destination: ProviderDeviceIdentity,
73}
74
75#[derive(Clone, Debug)]
76struct CandidateConfig {
77 policy: ExecutionPolicy,
78 engines: BTreeMap<EngineId, CandidateEngineRecord>,
79 modules: BTreeMap<ExtensionModuleId, CandidateModuleRecord>,
80 transfers: BTreeMap<TransferRoute, CandidateTransferRecord>,
81}
82
83struct BoundCandidateConfig {
84 policy: ExecutionPolicy,
85 engines: BTreeMap<EngineId, BoundCandidateEngineRecord>,
86 modules: BTreeMap<ExtensionModuleId, BoundCandidateModuleRecord>,
87 transfers: BTreeMap<TransferRoute, BoundCandidateTransferRecord>,
88}
89
90impl CandidateConfig {
91 fn empty() -> Self {
92 Self {
93 policy: default_execution_policy(),
94 engines: BTreeMap::new(),
95 modules: BTreeMap::new(),
96 transfers: BTreeMap::new(),
97 }
98 }
99
100 fn from_snapshot(snapshot: &RuntimeConfigSnapshot) -> Result<Self, RuntimeConfigError> {
101 let engines = snapshot
102 .engines
103 .iter()
104 .map(|slot| {
105 let registration = slot.to_registration()?;
106 Ok((
107 registration.engine_id().clone(),
108 CandidateEngineRecord {
109 registration,
110 identity: CandidateRegistrationIdentity::Preserved(
111 slot.metadata().identity,
112 ),
113 },
114 ))
115 })
116 .collect::<Result<BTreeMap<_, _>, RuntimeConfigError>>()?;
117 Ok(Self {
118 policy: snapshot.policy.clone(),
119 engines,
120 modules: snapshot.extensions.to_candidate_modules(),
121 transfers: snapshot
122 .transfers
123 .iter()
124 .map(|(resolved_route, provider)| {
125 (
126 TransferRoute::new(
127 resolved_route.source().logical().clone(),
128 resolved_route.destination().logical().clone(),
129 ),
130 CandidateTransferRecord {
131 provider: Arc::clone(provider),
132 binding: CandidateTransferBinding::Preserved {
133 source: resolved_route.source().provider_device_identity().clone(),
134 destination: resolved_route
135 .destination()
136 .provider_device_identity()
137 .clone(),
138 },
139 },
140 )
141 })
142 .collect(),
143 })
144 }
145}
146
147#[derive(Clone, Debug)]
148struct FrozenEngineMetadata {
149 candidate_token: Arc<CandidateRegistrationToken>,
150 identity: RegistrationIdentity,
151 event_domain_id: EventDomainId,
152}
153
154#[derive(Clone)]
155struct PreparationOnlyEngineSnapshot {
156 metadata: FrozenEngineMetadata,
157 binding: super::ProviderPreparationBinding,
158}
159
160#[derive(Clone, Debug)]
161pub(super) struct ExecutableEngineSnapshot {
162 metadata: FrozenEngineMetadata,
163 binding: super::ProviderExecutableBinding,
164}
165
166#[derive(Clone)]
167enum FrozenEngineSlot {
168 PreparationOnly(Arc<PreparationOnlyEngineSnapshot>),
169 Executable(Arc<ExecutableEngineSnapshot>),
170}
171
172impl FrozenEngineSlot {
173 fn metadata(&self) -> &FrozenEngineMetadata {
174 match self {
175 Self::PreparationOnly(snapshot) => &snapshot.metadata,
176 Self::Executable(snapshot) => &snapshot.metadata,
177 }
178 }
179
180 fn provider_device_identity(&self) -> &ProviderDeviceIdentity {
181 match self {
182 Self::PreparationOnly(snapshot) => snapshot.binding.provider_device_identity(),
183 Self::Executable(snapshot) => snapshot.binding.contract().provider_device_identity(),
184 }
185 }
186
187 fn engine_id(&self) -> &EngineId {
188 match self {
189 Self::PreparationOnly(snapshot) => snapshot.binding.engine_id(),
190 Self::Executable(snapshot) => snapshot.binding.engine_id(),
191 }
192 }
193
194 fn hardware_class(&self) -> &HardwareClassId {
195 match self {
196 Self::PreparationOnly(snapshot) => snapshot.binding.hardware_class(),
197 Self::Executable(snapshot) => snapshot.binding.hardware_class(),
198 }
199 }
200
201 fn storage_classes(&self) -> &[StorageClass] {
202 match self {
203 Self::PreparationOnly(snapshot) => snapshot.binding.storage_classes(),
204 Self::Executable(snapshot) => snapshot.binding.storage_classes(),
205 }
206 }
207
208 fn default_storage_class(&self) -> &StorageClass {
209 match self {
210 Self::PreparationOnly(snapshot) => snapshot.binding.default_storage_class(),
211 Self::Executable(snapshot) => snapshot.binding.default_storage_class(),
212 }
213 }
214
215 fn context_identity(&self) -> ExecutionContextIdentity {
216 match self {
217 Self::PreparationOnly(snapshot) => snapshot.binding.context_identity(),
218 Self::Executable(snapshot) => snapshot.binding.contract().context_identity(),
219 }
220 }
221
222 fn capabilities(&self) -> &CoreCapabilityBundle {
223 match self {
224 Self::PreparationOnly(snapshot) => snapshot.binding.capabilities(),
225 Self::Executable(snapshot) => snapshot.binding.contract().capabilities(),
226 }
227 }
228
229 fn executable(&self) -> Option<&Arc<ExecutableEngineSnapshot>> {
230 match self {
231 Self::PreparationOnly(_) => None,
232 Self::Executable(snapshot) => Some(snapshot),
233 }
234 }
235}
236
237impl ExecutableEngineSnapshot {
238 pub(super) fn engine_id(&self) -> &EngineId {
239 self.binding.engine_id()
240 }
241
242 pub(super) fn event_domain_id(&self) -> EventDomainId {
243 self.metadata.event_domain_id
244 }
245
246 pub(super) fn provider_device_identity(&self) -> &ProviderDeviceIdentity {
247 self.binding.contract().provider_device_identity()
248 }
249
250 #[cfg(test)]
251 pub(super) fn context_identity(&self) -> ExecutionContextIdentity {
252 self.binding.contract().context_identity()
253 }
254
255 pub(super) fn executor(&self) -> &Arc<dyn super::execution::ErasedTensorBackendExecutor> {
256 self.binding.contract().executor()
257 }
258
259 pub(super) fn event_domain_driver(&self) -> &Arc<dyn super::EventDomainDriver> {
260 self.binding.contract().event_domain_driver()
261 }
262
263 #[cfg(test)]
264 pub(super) fn has_executor(&self) -> bool {
265 true
266 }
267
268 #[cfg(test)]
269 pub(super) fn has_event_domain_driver(&self) -> bool {
270 true
271 }
272
273 pub(super) fn accepts_input_placement(
274 &self,
275 placement: &tenferro_tensor::Placement,
276 storage_class: &StorageClass,
277 ) -> bool {
278 self.binding.storage_classes().contains(storage_class)
279 && self
280 .binding
281 .contract()
282 .accepts_input_placement(placement, storage_class)
283 }
284
285 pub(super) fn accepts_input_signature(
286 &self,
287 input: &super::InputSignatureEntry,
288 storage_class: &StorageClass,
289 ) -> bool {
290 self.binding.storage_classes().contains(storage_class)
291 && self
292 .binding
293 .contract()
294 .accepts_input_signature(input, storage_class)
295 }
296
297 pub(super) fn accepts_runtime_input(
298 &self,
299 input: &tenferro_tensor::TensorRead<'_>,
300 storage_class: &StorageClass,
301 ) -> bool {
302 self.binding.storage_classes().contains(storage_class)
303 && self
304 .binding
305 .contract()
306 .accepts_runtime_input(input, storage_class)
307 }
308
309 pub(super) fn owns_resident_tensor(
310 &self,
311 input: &tenferro_tensor::TensorRead<'_>,
312 storage_class: &StorageClass,
313 ) -> bool {
314 self.binding.storage_classes().contains(storage_class)
315 && self
316 .binding
317 .contract()
318 .owns_resident_tensor(input, storage_class)
319 }
320
321 #[cfg(test)]
322 pub(super) fn for_test(
323 engine_id: EngineId,
324 provider_device_identity: ProviderDeviceIdentity,
325 event_domain_id: EventDomainId,
326 storage_class: StorageClass,
327 ) -> Arc<Self> {
328 Self::for_test_with_driver(
329 engine_id,
330 provider_device_identity,
331 event_domain_id,
332 storage_class,
333 Arc::new(super::ImmediateEventDomainDriver::new()),
334 )
335 }
336
337 #[cfg(test)]
338 pub(super) fn for_test_with_driver(
339 engine_id: EngineId,
340 provider_device_identity: ProviderDeviceIdentity,
341 event_domain_id: EventDomainId,
342 storage_class: StorageClass,
343 event_domain_driver: Arc<dyn super::EventDomainDriver>,
344 ) -> Arc<Self> {
345 let ingress = super::InputIngressContract::new(
346 super::InputPlacementContract::new(|_, _| true),
347 super::InputSignatureContract::new(|_, _, _, _| true),
348 super::RuntimeInputContract::new(|_, _| true),
349 super::ResidentOutputContract::new(|_, _| true),
350 );
351 let contract = super::ExecutableEngineContract::new(
352 provider_device_identity,
353 CoreCapabilityBundle::default(),
354 tenferro_cpu::CpuBackend::new(),
355 event_domain_driver,
356 ingress,
357 None,
358 );
359 let binding = super::ProviderExecutableBinding::new(
360 engine_id,
361 HardwareClassId::new("tenferro.test.schedule.hardware").expect("test hardware class"),
362 Arc::from(vec![storage_class.clone()]),
363 storage_class,
364 contract,
365 )
366 .expect("test executable binding");
367 Arc::new(Self {
368 metadata: FrozenEngineMetadata {
369 candidate_token: Arc::new(CandidateRegistrationToken),
370 identity: event_domain_id.registration_identity(),
371 event_domain_id,
372 },
373 binding,
374 })
375 }
376}
377
378impl FrozenEngineSlot {
379 fn to_registration(&self) -> Result<EngineRegistration, RuntimeConfigError> {
380 let metadata = self.metadata();
381 let registration = match self {
382 Self::PreparationOnly(snapshot) => {
383 EngineRegistration::from_state(EngineRegistrationState::PreparationOnly {
384 binding: snapshot.binding.clone(),
385 })
386 }
387 Self::Executable(snapshot) => EngineRegistration::from_state(
388 EngineRegistrationState::Executable(snapshot.binding.clone()),
389 ),
390 };
391 Ok(registration.with_candidate_token(Arc::clone(&metadata.candidate_token)))
392 }
393}
394
395impl fmt::Debug for FrozenEngineSlot {
396 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
397 let metadata = self.metadata();
398 formatter
399 .debug_struct("FrozenEngineSlot")
400 .field("engine_id", self.engine_id())
401 .field("registration_identity", &metadata.identity)
402 .field("event_domain_id", &metadata.event_domain_id)
403 .field("context_identity", &self.context_identity())
404 .field("hardware_class", self.hardware_class())
405 .field(
406 "state",
407 &match self {
408 Self::PreparationOnly(_) => "preparation-only",
409 Self::Executable(_) => "executable",
410 },
411 )
412 .finish_non_exhaustive()
413 }
414}
415
416#[derive(Clone)]
434pub struct RuntimeConfigSnapshot {
435 runtime_id: RuntimeId,
436 epoch: RuntimeEpoch,
437 policy: ExecutionPolicy,
438 engines: Arc<[FrozenEngineSlot]>,
439 engine_indices: BTreeMap<EngineId, usize>,
440 extensions: FrozenExtensionSlots,
441 transfers: FrozenTransferRegistry,
442 cache_owners: Arc<[FrozenCacheOwner]>,
443}
444
445impl RuntimeConfigSnapshot {
446 pub fn runtime_id(&self) -> RuntimeId {
448 self.runtime_id
449 }
450
451 pub fn epoch(&self) -> RuntimeEpoch {
453 self.epoch
454 }
455
456 pub fn execution_policy(&self) -> &ExecutionPolicy {
458 &self.policy
459 }
460
461 pub fn engine_count(&self) -> usize {
463 self.engines.len()
464 }
465
466 pub fn extension_module_count(&self) -> usize {
468 self.extensions.module_count()
469 }
470
471 pub fn transfer_provider_count(&self) -> usize {
473 self.transfers.len()
474 }
475
476 #[doc(hidden)]
478 pub fn has_extension_family(&self, family_id: &'static str) -> bool {
479 self.extensions.has_family(family_id)
480 }
481
482 #[doc(hidden)]
501 pub fn has_extension_engine(&self, family_id: &'static str, engine_id: &EngineId) -> bool {
502 self.extensions.has_engine(family_id, engine_id)
503 }
504
505 #[doc(hidden)]
512 pub fn has_extension_module_identical(&self, module: &Arc<dyn ExtensionModule>) -> bool {
513 self.extensions.has_module_identical(module)
514 }
515
516 #[doc(hidden)]
524 pub fn has_extension_module_engine(
525 &self,
526 module_id: &ExtensionModuleId,
527 family_id: &'static str,
528 engine_id: &EngineId,
529 ) -> bool {
530 self.extensions
531 .has_module_engine(module_id, family_id, engine_id)
532 }
533
534 pub fn engine(&self, id: &EngineId) -> Option<EngineSnapshotView<'_>> {
536 self.engine_indices
537 .get(id)
538 .map(|&index| EngineSnapshotView {
539 slot: &self.engines[index],
540 })
541 }
542
543 #[cfg(test)]
544 pub(crate) fn engine_ids_for_test(&self) -> impl Iterator<Item = &EngineId> {
545 self.engines.iter().map(FrozenEngineSlot::engine_id)
546 }
547
548 #[cfg(test)]
549 pub(crate) fn transfer_routes_for_test(&self) -> impl Iterator<Item = &ResolvedTransferRoute> {
550 self.transfers.iter().map(|(route, _)| route)
551 }
552
553 pub(super) fn engine_views_for_preparation(
554 &self,
555 ) -> impl Iterator<Item = EngineSnapshotView<'_>> + '_ {
556 self.engines.iter().map(|slot| EngineSnapshotView { slot })
557 }
558
559 pub(super) fn extension_slot_for_preparation(
560 &self,
561 family_id: ExtensionFamilyId,
562 engine_id: &EngineId,
563 ) -> Option<ExtensionEngineSnapshotView<'_>> {
564 self.extensions.slot_for_preparation(family_id, engine_id)
565 }
566
567 pub(super) fn transfer_registry_for_preparation(&self) -> FrozenTransferRegistry {
568 self.transfers.clone()
569 }
570
571 #[cfg(test)]
572 pub(crate) fn extension_slots_for_test(
573 &self,
574 ) -> impl Iterator<
575 Item = (
576 &ExtensionModuleId,
577 ExtensionFamilyId,
578 &EngineId,
579 RegistrationIdentity,
580 ),
581 > {
582 self.extensions.slots_for_test()
583 }
584
585 #[cfg(test)]
586 pub(crate) fn extension_slot_identity_for_test(
587 &self,
588 family_id: ExtensionFamilyId,
589 engine_id: &EngineId,
590 ) -> Option<RegistrationIdentity> {
591 self.extensions.slot_identity_for_test(family_id, engine_id)
592 }
593
594 #[cfg(test)]
595 pub(crate) fn extension_slot_full_for_test(
596 &self,
597 family_id: ExtensionFamilyId,
598 engine_id: &EngineId,
599 ) -> Option<ExtensionSlotFullForTest<'_>> {
600 self.extensions.slot_full_for_test(family_id, engine_id)
601 }
602
603 #[cfg(test)]
604 pub(super) fn cache_owners_for_test(&self) -> &[FrozenCacheOwner] {
605 &self.cache_owners
606 }
607
608 pub(super) fn cache_owners_for_runtime(&self) -> &[FrozenCacheOwner] {
609 &self.cache_owners
610 }
611}
612
613impl fmt::Debug for RuntimeConfigSnapshot {
614 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
615 formatter
616 .debug_struct("RuntimeConfigSnapshot")
617 .field("runtime_id", &self.runtime_id)
618 .field("epoch", &self.epoch)
619 .field("execution_policy", &self.policy)
620 .field("engine_count", &self.engines.len())
621 .field("extension_module_count", &self.extensions.module_count())
622 .field("extension_engine_count", &self.extensions.engine_count())
623 .field("transfer_provider_count", &self.transfers.len())
624 .field("cache_owner_count", &self.cache_owners.len())
625 .finish_non_exhaustive()
626 }
627}
628
629struct RuntimeState {
630 runtime_id: RuntimeId,
631 issuer: NonZeroU64,
632 next_registration_ordinal: AtomicU64,
633 active: RwLock<Arc<RuntimeConfigSnapshot>>,
634 published_epoch: AtomicU64,
635 caches: RuntimeCacheSet<PreparedEntryKey, PreparedProgram>,
636}
637
638#[derive(Clone)]
652pub struct Runtime(Arc<RuntimeState>);
653
654impl Runtime {
655 pub fn builder() -> RuntimeConfigBuilder {
657 RuntimeConfigBuilder::new()
658 }
659
660 pub fn id(&self) -> RuntimeId {
662 self.0.runtime_id
663 }
664
665 pub fn snapshot(&self) -> Result<Arc<RuntimeConfigSnapshot>, RuntimeStateError> {
672 self.0
673 .active
674 .read()
675 .map(|snapshot| Arc::clone(&snapshot))
676 .map_err(|_| RuntimeStateError::Poisoned {
677 lock: "runtime.active",
678 })
679 }
680
681 pub fn epoch(&self) -> Result<RuntimeEpoch, RuntimeStateError> {
689 match NonZeroU64::new(self.0.published_epoch.load(Ordering::Acquire)) {
690 Some(value) => Ok(RuntimeEpoch::from_nonzero(value)),
691 None => Err(RuntimeStateError::Poisoned {
692 lock: "runtime.published_epoch",
693 }),
694 }
695 }
696
697 pub fn prepared_cache_limits(&self) -> Result<PreparedPlanCacheLimits, RuntimeStateError> {
716 self.0.caches.prepared().limits()
717 }
718
719 pub fn set_prepared_cache_limits(
746 &self,
747 limits: PreparedPlanCacheLimits,
748 ) -> Result<(), RuntimeStateError> {
749 self.0.caches.prepared().set_limits(limits)
750 }
751
752 pub fn clear_prepared_cache(&self) -> Result<(), RuntimeStateError> {
771 self.0.caches.prepared().clear()
772 }
773
774 pub fn cache_stats(&self) -> Result<RuntimeCacheStats, RuntimeCacheError> {
794 super::preparation::cache_stats(self, &self.0.caches)
795 }
796
797 pub fn clear_caches(&self) -> Result<(), RuntimeCacheError> {
817 super::preparation::clear_caches(self, &self.0.caches)
818 }
819
820 #[allow(
821 dead_code,
822 reason = "Phase 5 graph execution consumes crate-private prepared programs"
823 )]
824 pub(crate) fn prepare_for(
825 &self,
826 frozen: &FrozenProgram,
827 signature: &InputSignature,
828 options: &PrepareOptions,
829 ) -> PreparedProgramResult<Arc<PreparedProgram>> {
830 super::preparation::prepare_for(self, &self.0.caches, frozen, signature, options)
831 }
832
833 pub(crate) fn prepare_compiled_for(
834 &self,
835 program: &CompiledGraph,
836 signature: &InputSignature,
837 options: &PrepareOptions,
838 ) -> PreparedProgramResult<Arc<PreparedProgram>> {
839 super::preparation::prepare_compiled_for(self, &self.0.caches, program, signature, options)
840 }
841
842 pub fn prepare_extension_immediate(
947 &self,
948 engine_id: &EngineId,
949 op: &dyn tenferro_ops::ext_op::ExtensionOp,
950 signature: &InputSignature,
951 ) -> crate::Result<PrepareCapability> {
952 super::preparation::prepare_extension_immediate(self, engine_id, op, signature).map_err(
953 |source| {
954 Error::runtime_state_source(
955 "Runtime::prepare_extension_immediate",
956 ErrorPhase::Execution,
957 source,
958 )
959 },
960 )
961 }
962
963 pub fn run_compiled(
1005 &self,
1006 program: &CompiledGraph,
1007 inputs: &[&Tensor],
1008 ) -> crate::Result<Vec<Tensor>> {
1009 super::execution::run_compiled(self, program, inputs)
1010 }
1011
1012 pub fn execute_scoped_read_only<'env>(
1026 &self,
1027 program: &CompiledGraph,
1028 inputs: super::execution::ScopedReadInputs<'env>,
1029 ) -> std::result::Result<
1030 super::execution::ScopedExecutionOutcome<'env>,
1031 super::execution::ScopedSubmitRejected<'env>,
1032 > {
1033 super::execution::execute_scoped_read_only(self, program, inputs)
1034 }
1035
1036 pub fn prepare_compiled(
1061 &self,
1062 program: &CompiledGraph,
1063 inputs: &[&Tensor],
1064 ) -> crate::Result<super::execution::PreparedCompiledGraph> {
1065 super::execution::prepare_compiled(self, program, inputs)
1066 }
1067
1068 pub fn run_prepared(
1078 &self,
1079 prepared: &super::execution::PreparedCompiledGraph,
1080 inputs: &[&Tensor],
1081 ) -> crate::Result<Vec<Tensor>> {
1082 super::execution::run_prepared(self, prepared, inputs)
1083 }
1084
1085 pub fn submit(
1101 &self,
1102 program: &CompiledGraph,
1103 inputs: super::execution::ExecutionInputs,
1104 ) -> std::result::Result<super::execution::ExecutionHandle, super::execution::SubmitError> {
1105 super::execution::submit(self, program, inputs)
1106 }
1107
1108 pub fn run_compiled_values(
1146 &self,
1147 program: &CompiledGraph,
1148 inputs: &[&Tensor],
1149 ) -> crate::Result<Vec<TensorValue>> {
1150 super::execution::run_compiled_values(self, program, inputs)
1151 }
1152
1153 pub fn reconfigure(
1169 &self,
1170 edit: impl FnOnce(&mut RuntimeReconfiguration<'_>) -> Result<(), RuntimeConfigError>,
1171 ) -> Result<RuntimeEpoch, RuntimeReconfigureError> {
1172 let base = self
1173 .snapshot()
1174 .map_err(|source| RuntimeReconfigureError::State { source })?;
1175 let mut candidate = CandidateConfig::from_snapshot(&base)
1176 .map_err(|source| RuntimeReconfigureError::Edit { source })?;
1177 let mut changed = false;
1178 {
1179 let mut reconfiguration = RuntimeReconfiguration {
1180 candidate: &mut candidate,
1181 changed: &mut changed,
1182 };
1183 edit(&mut reconfiguration)
1184 .map_err(|source| RuntimeReconfigureError::Edit { source })?;
1185 }
1186
1187 if !changed {
1188 return Ok(base.epoch());
1189 }
1190 let next_identity_ordinal = NonZeroU64::new(
1191 self.0.next_registration_ordinal.load(Ordering::SeqCst),
1192 )
1193 .ok_or(RuntimeReconfigureError::Edit {
1194 source: RuntimeConfigError::IdentityExhausted,
1195 })?;
1196 let (bound_candidate, post_ordinal) =
1197 validate_candidate(candidate, self.0.issuer, next_identity_ordinal)
1198 .map_err(|source| RuntimeReconfigureError::Edit { source })?;
1199
1200 let next_epoch =
1201 base.epoch()
1202 .checked_next()
1203 .ok_or(RuntimeReconfigureError::EpochExhausted {
1204 current: base.epoch(),
1205 })?;
1206
1207 let mut guard = self
1208 .0
1209 .active
1210 .write()
1211 .map_err(|_| RuntimeReconfigureError::State {
1212 source: RuntimeStateError::Poisoned {
1213 lock: "runtime.active",
1214 },
1215 })?;
1216 if !Arc::ptr_eq(&*guard, &base) {
1217 return Err(RuntimeReconfigureError::ConcurrentReconfiguration {
1218 base: base.epoch(),
1219 current: guard.epoch(),
1220 });
1221 }
1222
1223 let next_snapshot = Arc::new(
1224 freeze_candidate(self.0.runtime_id, next_epoch, bound_candidate)
1225 .map_err(|source| RuntimeReconfigureError::Edit { source })?,
1226 );
1227
1228 self.0
1229 .next_registration_ordinal
1230 .store(post_ordinal.get(), Ordering::SeqCst);
1231 *guard = next_snapshot;
1232 self.0
1233 .published_epoch
1234 .store(next_epoch.get().get(), Ordering::Release);
1235 Ok(next_epoch)
1236 }
1237
1238 #[cfg(test)]
1239 pub(crate) fn force_epoch_for_test(&self, epoch: RuntimeEpoch) {
1240 let mut guard = self.0.active.write().expect("test runtime lock");
1241 let mut replacement = (**guard).clone();
1242 replacement.epoch = epoch;
1243 *guard = Arc::new(replacement);
1244 self.0
1245 .published_epoch
1246 .store(epoch.get().get(), Ordering::Release);
1247 }
1248
1249 #[cfg(test)]
1250 pub(crate) fn force_next_registration_ordinal_for_test(&self, next: NonZeroU64) {
1251 self.0
1252 .next_registration_ordinal
1253 .store(next.get(), Ordering::SeqCst);
1254 }
1255
1256 #[cfg(test)]
1257 pub(crate) fn poison_active_lock_for_test(&self) {
1258 let state = Arc::clone(&self.0);
1259 let _ = std::panic::catch_unwind(std::panic::AssertUnwindSafe(move || {
1260 let _guard = state.active.write().expect("test runtime lock");
1261 panic!("poison runtime.active for test");
1262 }));
1263 }
1264}
1265
1266impl fmt::Debug for Runtime {
1267 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
1268 formatter
1269 .debug_struct("Runtime")
1270 .field("runtime_id", &self.0.runtime_id)
1271 .field("published_epoch", &self.epoch().ok())
1272 .finish_non_exhaustive()
1273 }
1274}
1275
1276pub struct RuntimeConfigBuilder {
1290 candidate: CandidateConfig,
1291}
1292
1293impl RuntimeConfigBuilder {
1294 pub fn new() -> Self {
1296 Self {
1297 candidate: CandidateConfig::empty(),
1298 }
1299 }
1300
1301 pub fn execution_policy(&mut self, value: ExecutionPolicy) -> &mut Self {
1303 self.candidate.policy = value;
1304 self
1305 }
1306
1307 pub fn register_engine(
1314 &mut self,
1315 value: EngineRegistration,
1316 ) -> Result<&mut Self, RuntimeConfigError> {
1317 let mut changed = false;
1318 register_engine_candidate(&mut self.candidate, value, &mut changed)?;
1319 Ok(self)
1320 }
1321
1322 pub fn replace_engine(
1328 &mut self,
1329 value: EngineRegistration,
1330 ) -> Result<&mut Self, RuntimeConfigError> {
1331 let mut changed = false;
1332 replace_engine_candidate(&mut self.candidate, value, &mut changed)?;
1333 Ok(self)
1334 }
1335
1336 pub fn remove_engine(&mut self, id: &EngineId) -> Result<&mut Self, RuntimeConfigError> {
1342 let mut changed = false;
1343 remove_engine_candidate(&mut self.candidate, id, &mut changed)?;
1344 Ok(self)
1345 }
1346
1347 pub fn install_extension_module(
1354 &mut self,
1355 value: Arc<dyn ExtensionModule>,
1356 ) -> Result<&mut Self, RuntimeConfigError> {
1357 let mut changed = false;
1358 install_extension_module_candidate(&mut self.candidate, value, &mut changed)?;
1359 Ok(self)
1360 }
1361
1362 pub fn register_transfer_provider(
1442 &mut self,
1443 source: TransferEndpoint,
1444 destination: TransferEndpoint,
1445 provider: Arc<dyn TransferProvider>,
1446 ) -> Result<&mut Self, RuntimeConfigError> {
1447 let mut changed = false;
1448 register_transfer_provider_candidate(
1449 &mut self.candidate,
1450 source,
1451 destination,
1452 provider,
1453 &mut changed,
1454 )?;
1455 Ok(self)
1456 }
1457
1458 pub fn remove_transfer_provider(
1469 &mut self,
1470 source: TransferEndpoint,
1471 destination: TransferEndpoint,
1472 ) -> Result<&mut Self, RuntimeConfigError> {
1473 let mut changed = false;
1474 remove_transfer_provider_candidate(&mut self.candidate, source, destination, &mut changed)?;
1475 Ok(self)
1476 }
1477
1478 pub fn replace_extension_module(
1485 &mut self,
1486 value: Arc<dyn ExtensionModule>,
1487 ) -> Result<&mut Self, RuntimeConfigError> {
1488 let mut changed = false;
1489 replace_extension_module_candidate(&mut self.candidate, value, &mut changed)?;
1490 Ok(self)
1491 }
1492
1493 pub fn remove_extension_module(
1500 &mut self,
1501 id: &ExtensionModuleId,
1502 ) -> Result<&mut Self, RuntimeConfigError> {
1503 let mut changed = false;
1504 remove_extension_module_candidate(&mut self.candidate, id, &mut changed)?;
1505 Ok(self)
1506 }
1507
1508 pub fn build(self) -> Result<Runtime, RuntimeConfigError> {
1518 let runtime_id = RuntimeId::from_nonzero(allocate_nonzero(&NEXT_RUNTIME_ID)?);
1519 let issuer = allocate_nonzero(&NEXT_REGISTRATION_ISSUER)?;
1520 let (bound_candidate, post_ordinal) =
1521 validate_candidate(self.candidate, issuer, INITIAL_REGISTRATION_ORDINAL)?;
1522 let epoch = RuntimeEpoch::one();
1523 let snapshot = Arc::new(freeze_candidate(runtime_id, epoch, bound_candidate)?);
1524 let state = RuntimeState {
1525 runtime_id,
1526 issuer,
1527 next_registration_ordinal: AtomicU64::new(post_ordinal.get()),
1528 active: RwLock::new(snapshot),
1529 published_epoch: AtomicU64::new(epoch.get().get()),
1530 caches: RuntimeCacheSet::new(PreparedPlanCacheLimits::default()),
1531 };
1532 Ok(Runtime(Arc::new(state)))
1533 }
1534}
1535
1536impl Default for RuntimeConfigBuilder {
1537 fn default() -> Self {
1538 Self::new()
1539 }
1540}
1541
1542impl fmt::Debug for RuntimeConfigBuilder {
1543 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
1544 formatter
1545 .debug_struct("RuntimeConfigBuilder")
1546 .field("execution_policy", &self.candidate.policy)
1547 .field("engine_count", &self.candidate.engines.len())
1548 .field("extension_module_count", &self.candidate.modules.len())
1549 .field("transfer_provider_count", &self.candidate.transfers.len())
1550 .finish_non_exhaustive()
1551 }
1552}
1553
1554pub struct RuntimeReconfiguration<'a> {
1556 candidate: &'a mut CandidateConfig,
1557 changed: &'a mut bool,
1558}
1559
1560impl RuntimeReconfiguration<'_> {
1561 pub fn execution_policy(&mut self, policy: ExecutionPolicy) -> &mut Self {
1563 if self.candidate.policy != policy {
1564 self.candidate.policy = policy;
1565 *self.changed = true;
1566 }
1567 self
1568 }
1569
1570 pub fn register_engine(
1577 &mut self,
1578 value: EngineRegistration,
1579 ) -> Result<&mut Self, RuntimeConfigError> {
1580 register_engine_candidate(self.candidate, value, self.changed)?;
1581 Ok(self)
1582 }
1583
1584 pub fn replace_engine(
1590 &mut self,
1591 value: EngineRegistration,
1592 ) -> Result<&mut Self, RuntimeConfigError> {
1593 replace_engine_candidate(self.candidate, value, self.changed)?;
1594 Ok(self)
1595 }
1596
1597 pub fn remove_engine(&mut self, id: &EngineId) -> Result<&mut Self, RuntimeConfigError> {
1603 remove_engine_candidate(self.candidate, id, self.changed)?;
1604 Ok(self)
1605 }
1606
1607 pub fn install_extension_module(
1614 &mut self,
1615 value: Arc<dyn ExtensionModule>,
1616 ) -> Result<&mut Self, RuntimeConfigError> {
1617 install_extension_module_candidate(self.candidate, value, self.changed)?;
1618 Ok(self)
1619 }
1620
1621 pub fn register_transfer_provider(
1696 &mut self,
1697 source: TransferEndpoint,
1698 destination: TransferEndpoint,
1699 provider: Arc<dyn TransferProvider>,
1700 ) -> Result<&mut Self, RuntimeConfigError> {
1701 register_transfer_provider_candidate(
1702 self.candidate,
1703 source,
1704 destination,
1705 provider,
1706 self.changed,
1707 )?;
1708 Ok(self)
1709 }
1710
1711 pub fn remove_transfer_provider(
1719 &mut self,
1720 source: TransferEndpoint,
1721 destination: TransferEndpoint,
1722 ) -> Result<&mut Self, RuntimeConfigError> {
1723 remove_transfer_provider_candidate(self.candidate, source, destination, self.changed)?;
1724 Ok(self)
1725 }
1726
1727 pub fn replace_extension_module(
1735 &mut self,
1736 value: Arc<dyn ExtensionModule>,
1737 ) -> Result<&mut Self, RuntimeConfigError> {
1738 replace_extension_module_candidate(self.candidate, value, self.changed)?;
1739 Ok(self)
1740 }
1741
1742 #[doc(hidden)]
1805 pub fn ensure_extension_module_for_engine(
1806 &mut self,
1807 value: Arc<dyn ExtensionModule>,
1808 family_id: &'static str,
1809 engine_id: &EngineId,
1810 ) -> Result<&mut Self, RuntimeConfigError> {
1811 let module_id = value.module_id().clone();
1812 if let Some(existing) = self.candidate.modules.get(&module_id)
1813 && existing
1814 .engines
1815 .contains_key(&(family_id, engine_id.clone()))
1816 {
1817 return Ok(self);
1818 }
1819
1820 let record = configure_module(Arc::clone(&value))
1821 .map_err(|source| RuntimeConfigError::ExtensionModule { source })?;
1822 if !record.engines.contains_key(&(family_id, engine_id.clone())) {
1823 return Err(RuntimeConfigError::MissingExtensionEngine {
1824 module_id,
1825 family_id,
1826 engine_id: engine_id.clone(),
1827 });
1828 }
1829 self.candidate.modules.insert(module_id, record);
1830 *self.changed = true;
1831 Ok(self)
1832 }
1833
1834 #[doc(hidden)]
1894 pub fn replace_extension_module_for_engine(
1895 &mut self,
1896 value: Arc<dyn ExtensionModule>,
1897 family_id: &'static str,
1898 engine_id: &EngineId,
1899 ) -> Result<&mut Self, RuntimeConfigError> {
1900 let module_id = value.module_id().clone();
1901 if let Some(existing) = self.candidate.modules.get(&module_id)
1902 && existing.module_identical(&value)
1903 {
1904 if existing
1905 .engines
1906 .contains_key(&(family_id, engine_id.clone()))
1907 {
1908 return Ok(self);
1909 }
1910 return Err(RuntimeConfigError::MissingExtensionEngine {
1911 module_id,
1912 family_id,
1913 engine_id: engine_id.clone(),
1914 });
1915 }
1916
1917 let record = configure_module(Arc::clone(&value))
1918 .map_err(|source| RuntimeConfigError::ExtensionModule { source })?;
1919 if !record.engines.contains_key(&(family_id, engine_id.clone())) {
1920 return Err(RuntimeConfigError::MissingExtensionEngine {
1921 module_id,
1922 family_id,
1923 engine_id: engine_id.clone(),
1924 });
1925 }
1926 self.candidate.modules.insert(module_id, record);
1927 *self.changed = true;
1928 Ok(self)
1929 }
1930
1931 pub fn remove_extension_module(
1938 &mut self,
1939 id: &ExtensionModuleId,
1940 ) -> Result<&mut Self, RuntimeConfigError> {
1941 remove_extension_module_candidate(self.candidate, id, self.changed)?;
1942 Ok(self)
1943 }
1944}
1945
1946impl fmt::Debug for RuntimeReconfiguration<'_> {
1947 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
1948 formatter
1949 .debug_struct("RuntimeReconfiguration")
1950 .field("engine_count", &self.candidate.engines.len())
1951 .field("extension_module_count", &self.candidate.modules.len())
1952 .field("transfer_provider_count", &self.candidate.transfers.len())
1953 .field("changed", &*self.changed)
1954 .finish_non_exhaustive()
1955 }
1956}
1957
1958#[derive(Clone, Copy)]
1970pub struct EngineSnapshotView<'a> {
1971 slot: &'a FrozenEngineSlot,
1972}
1973
1974impl<'a> EngineSnapshotView<'a> {
1975 pub fn engine_id(&self) -> &'a EngineId {
1977 self.slot.engine_id()
1978 }
1979
1980 pub fn provider_device_identity(&self) -> &'a super::ProviderDeviceIdentity {
1990 self.slot.provider_device_identity()
1991 }
1992
1993 pub fn registration_identity(&self) -> RegistrationIdentity {
1995 self.slot.metadata().identity
1996 }
1997
1998 pub fn context_identity(&self) -> ExecutionContextIdentity {
2000 self.slot.context_identity()
2001 }
2002
2003 pub fn event_domain_id(&self) -> EventDomainId {
2005 self.slot.metadata().event_domain_id
2006 }
2007
2008 pub(super) fn executable_witness(&self) -> Option<&'a Arc<ExecutableEngineSnapshot>> {
2009 self.slot.executable()
2010 }
2011
2012 pub fn hardware_class(&self) -> &'a HardwareClassId {
2014 self.slot.hardware_class()
2015 }
2016
2017 pub fn capabilities(&self) -> &'a CoreCapabilityBundle {
2019 self.slot.capabilities()
2020 }
2021
2022 pub(super) fn storage_classes(&self) -> &'a [StorageClass] {
2023 self.slot.storage_classes()
2024 }
2025
2026 pub(super) fn default_storage_class(&self) -> &'a StorageClass {
2027 self.slot.default_storage_class()
2028 }
2029
2030 #[doc(hidden)]
2057 pub fn accepts_input_signature(&self, input: &super::InputSignatureEntry) -> bool {
2058 self.slot
2059 .storage_classes()
2060 .iter()
2061 .any(|storage_class| self.accepts_input_signature_for_storage(input, storage_class))
2062 }
2063
2064 pub(super) fn accepts_input_signature_for_storage(
2065 &self,
2066 input: &super::InputSignatureEntry,
2067 storage_class: &StorageClass,
2068 ) -> bool {
2069 self.slot
2070 .executable()
2071 .is_some_and(|snapshot| snapshot.accepts_input_signature(input, storage_class))
2072 }
2073
2074 #[cfg(test)]
2075 pub(crate) fn has_execution_engine_for_test(&self) -> bool {
2076 self.slot.executable().is_some()
2077 }
2078}
2079
2080impl fmt::Debug for EngineSnapshotView<'_> {
2081 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
2082 formatter
2083 .debug_struct("EngineSnapshotView")
2084 .field("engine_id", self.engine_id())
2085 .field("registration_identity", &self.registration_identity())
2086 .field("context_identity", &self.context_identity())
2087 .field("hardware_class", self.hardware_class())
2088 .field("capabilities", self.capabilities())
2089 .finish()
2090 }
2091}
2092
2093fn default_execution_policy() -> ExecutionPolicy {
2094 ExecutionPolicy::new(super::Determinism::Fast, None, 0)
2095}
2096
2097fn register_engine_candidate(
2098 candidate: &mut CandidateConfig,
2099 registration: EngineRegistration,
2100 changed: &mut bool,
2101) -> Result<(), RuntimeConfigError> {
2102 let engine_id = registration.engine_id().clone();
2103 match candidate.engines.get(&engine_id) {
2104 Some(existing) if existing.registration.candidate_identical(®istration) => Ok(()),
2105 Some(_) => Err(RuntimeConfigError::DuplicateEngine { engine_id }),
2106 None => {
2107 ensure_unique_provider_device_target(candidate, ®istration)?;
2108 candidate.engines.insert(
2109 engine_id,
2110 CandidateEngineRecord {
2111 registration,
2112 identity: CandidateRegistrationIdentity::New,
2113 },
2114 );
2115 *changed = true;
2116 Ok(())
2117 }
2118 }
2119}
2120
2121fn replace_engine_candidate(
2122 candidate: &mut CandidateConfig,
2123 registration: EngineRegistration,
2124 changed: &mut bool,
2125) -> Result<(), RuntimeConfigError> {
2126 let engine_id = registration.engine_id().clone();
2127 let Some(existing) = candidate.engines.get(&engine_id) else {
2128 return Err(RuntimeConfigError::MissingEngine { engine_id });
2129 };
2130 if existing.registration.candidate_identical(®istration) {
2131 return Ok(());
2132 }
2133 if existing.registration.provider_device_identity() != registration.provider_device_identity() {
2134 return Err(RuntimeConfigError::EngineTargetRebind {
2135 engine_id,
2136 current: existing.registration.provider_device_identity().clone(),
2137 replacement: registration.provider_device_identity().clone(),
2138 });
2139 }
2140 ensure_unique_provider_device_target_except(candidate, ®istration, &engine_id)?;
2141 candidate.engines.insert(
2142 engine_id,
2143 CandidateEngineRecord {
2144 registration,
2145 identity: CandidateRegistrationIdentity::New,
2146 },
2147 );
2148 *changed = true;
2149 Ok(())
2150}
2151
2152fn remove_engine_candidate(
2153 candidate: &mut CandidateConfig,
2154 id: &EngineId,
2155 changed: &mut bool,
2156) -> Result<(), RuntimeConfigError> {
2157 match candidate.engines.remove(id) {
2158 Some(_) => {
2159 *changed = true;
2160 Ok(())
2161 }
2162 None => Err(RuntimeConfigError::MissingEngine {
2163 engine_id: id.clone(),
2164 }),
2165 }
2166}
2167
2168fn install_extension_module_candidate(
2169 candidate: &mut CandidateConfig,
2170 module: Arc<dyn ExtensionModule>,
2171 changed: &mut bool,
2172) -> Result<(), RuntimeConfigError> {
2173 let module_id = module.module_id().clone();
2174 match candidate.modules.get(&module_id) {
2175 Some(existing) if existing.module_identical(&module) => Ok(()),
2176 Some(_) => Err(RuntimeConfigError::ExtensionModule {
2177 source: ExtensionModuleError::ConflictingModule { module_id },
2178 }),
2179 None => {
2180 let record = configure_module(module)
2181 .map_err(|source| RuntimeConfigError::ExtensionModule { source })?;
2182 candidate.modules.insert(module_id, record);
2183 *changed = true;
2184 Ok(())
2185 }
2186 }
2187}
2188
2189fn replace_extension_module_candidate(
2190 candidate: &mut CandidateConfig,
2191 module: Arc<dyn ExtensionModule>,
2192 changed: &mut bool,
2193) -> Result<(), RuntimeConfigError> {
2194 let module_id = module.module_id().clone();
2195 match candidate.modules.get(&module_id) {
2196 Some(existing) if existing.module_identical(&module) => Ok(()),
2197 _ => {
2198 let record = configure_module(module)
2199 .map_err(|source| RuntimeConfigError::ExtensionModule { source })?;
2200 candidate.modules.insert(module_id, record);
2201 *changed = true;
2202 Ok(())
2203 }
2204 }
2205}
2206
2207fn remove_extension_module_candidate(
2208 candidate: &mut CandidateConfig,
2209 id: &ExtensionModuleId,
2210 changed: &mut bool,
2211) -> Result<(), RuntimeConfigError> {
2212 if candidate.modules.remove(id).is_some() {
2213 *changed = true;
2214 }
2215 Ok(())
2216}
2217
2218fn register_transfer_provider_candidate(
2219 candidate: &mut CandidateConfig,
2220 source: TransferEndpoint,
2221 destination: TransferEndpoint,
2222 provider: Arc<dyn TransferProvider>,
2223 changed: &mut bool,
2224) -> Result<(), RuntimeConfigError> {
2225 let key = TransferRoute::new(source, destination);
2226 match candidate.transfers.get(&key) {
2227 Some(existing) if Arc::ptr_eq(&existing.provider, &provider) => Ok(()),
2228 Some(_) => Err(RuntimeConfigError::ConflictingRegistration {
2229 key: RegistrationKey::TransferProvider {
2230 source: key.source().clone(),
2231 destination: key.destination().clone(),
2232 },
2233 }),
2234 None => {
2235 candidate.transfers.insert(
2236 key,
2237 CandidateTransferRecord {
2238 provider,
2239 binding: CandidateTransferBinding::New,
2240 },
2241 );
2242 *changed = true;
2243 Ok(())
2244 }
2245 }
2246}
2247
2248fn remove_transfer_provider_candidate(
2249 candidate: &mut CandidateConfig,
2250 source: TransferEndpoint,
2251 destination: TransferEndpoint,
2252 changed: &mut bool,
2253) -> Result<(), RuntimeConfigError> {
2254 let key = TransferRoute::new(source, destination);
2255 if candidate.transfers.remove(&key).is_none() {
2256 return Err(RuntimeConfigError::MissingTransferProvider {
2257 source_endpoint: key.source().clone(),
2258 destination: key.destination().clone(),
2259 });
2260 }
2261 *changed = true;
2262 Ok(())
2263}
2264
2265fn ensure_unique_provider_device_target(
2266 candidate: &CandidateConfig,
2267 registration: &EngineRegistration,
2268) -> Result<(), RuntimeConfigError> {
2269 ensure_unique_provider_device_target_except(candidate, registration, registration.engine_id())
2270}
2271
2272fn ensure_unique_provider_device_target_except(
2273 candidate: &CandidateConfig,
2274 registration: &EngineRegistration,
2275 ignored_engine_id: &EngineId,
2276) -> Result<(), RuntimeConfigError> {
2277 if let Some((first_engine_id, _)) = candidate.engines.iter().find(|(engine_id, record)| {
2278 *engine_id != ignored_engine_id
2279 && record.registration.provider_device_identity()
2280 == registration.provider_device_identity()
2281 }) {
2282 return Err(RuntimeConfigError::DuplicateProviderDeviceTarget {
2283 provider_device_identity: registration.provider_device_identity().clone(),
2284 first_engine_id: first_engine_id.clone(),
2285 duplicate_engine_id: registration.engine_id().clone(),
2286 });
2287 }
2288 Ok(())
2289}
2290
2291fn validate_candidate(
2292 candidate: CandidateConfig,
2293 issuer: NonZeroU64,
2294 next_ordinal: NonZeroU64,
2295) -> Result<(BoundCandidateConfig, NonZeroU64), RuntimeConfigError> {
2296 let mut seen_targets = BTreeMap::<ProviderDeviceIdentity, EngineId>::new();
2297 for (engine_id, record) in &candidate.engines {
2298 if let Some(first_engine_id) = seen_targets.insert(
2299 record.registration.provider_device_identity().clone(),
2300 engine_id.clone(),
2301 ) {
2302 return Err(RuntimeConfigError::DuplicateProviderDeviceTarget {
2303 provider_device_identity: record.registration.provider_device_identity().clone(),
2304 first_engine_id,
2305 duplicate_engine_id: engine_id.clone(),
2306 });
2307 }
2308 }
2309
2310 let mut bound_transfers = BTreeMap::new();
2311 for (route, record) in &candidate.transfers {
2312 let source_binding = validate_transfer_endpoint(&candidate, route.source())?;
2313 let destination_binding = validate_transfer_endpoint(&candidate, route.destination())?;
2314 let preserved = match &record.binding {
2315 CandidateTransferBinding::New => None,
2316 CandidateTransferBinding::Preserved {
2317 source,
2318 destination,
2319 } => Some((source, destination)),
2320 };
2321 if let Some((registered_source, registered_destination)) = preserved {
2322 if registered_source != &source_binding {
2323 return Err(RuntimeConfigError::StaleTransferRoute {
2324 source_endpoint: route.source().clone(),
2325 destination: route.destination().clone(),
2326 endpoint: route.source().clone(),
2327 registered: Box::new(registered_source.clone()),
2328 current: Box::new(source_binding.clone()),
2329 });
2330 }
2331 if registered_destination != &destination_binding {
2332 return Err(RuntimeConfigError::StaleTransferRoute {
2333 source_endpoint: route.source().clone(),
2334 destination: route.destination().clone(),
2335 endpoint: route.destination().clone(),
2336 registered: Box::new(registered_destination.clone()),
2337 current: Box::new(destination_binding.clone()),
2338 });
2339 }
2340 }
2341 bound_transfers.insert(
2342 route.clone(),
2343 BoundCandidateTransferRecord {
2344 provider: Arc::clone(&record.provider),
2345 source: source_binding,
2346 destination: destination_binding,
2347 },
2348 );
2349 }
2350 let mut seen = BTreeMap::<(ExtensionFamilyId, EngineId), ExtensionModuleId>::new();
2351 for (module_id, module) in &candidate.modules {
2352 for family_engine in module.engines.keys() {
2353 if seen
2354 .insert(
2355 (family_engine.0, family_engine.1.clone()),
2356 module_id.clone(),
2357 )
2358 .is_some()
2359 {
2360 return Err(RuntimeConfigError::ConflictingRegistration {
2361 key: RegistrationKey::ExtensionEngine {
2362 family: family_engine.0,
2363 engine: family_engine.1.clone(),
2364 },
2365 });
2366 }
2367 }
2368 }
2369
2370 let CandidateConfig {
2371 policy,
2372 engines,
2373 modules,
2374 transfers: _,
2375 } = candidate;
2376 let mut allocator = RegistrationIdentityAllocator::new(issuer, next_ordinal);
2377 let engines = engines
2378 .into_iter()
2379 .map(|(engine_id, record)| {
2380 let identity = match record.identity {
2381 CandidateRegistrationIdentity::New => allocator.allocate()?,
2382 CandidateRegistrationIdentity::Preserved(identity) => identity,
2383 };
2384 Ok((
2385 engine_id,
2386 BoundCandidateEngineRecord {
2387 registration: record.registration,
2388 identity,
2389 },
2390 ))
2391 })
2392 .collect::<Result<BTreeMap<_, _>, RuntimeConfigError>>()?;
2393 let modules = modules
2394 .into_iter()
2395 .map(|(module_id, module)| {
2396 let mut allocate = || allocator.allocate();
2397 Ok((module_id, bind_candidate_module(module, &mut allocate)?))
2398 })
2399 .collect::<Result<BTreeMap<_, _>, RuntimeConfigError>>()?;
2400 Ok((
2401 BoundCandidateConfig {
2402 policy,
2403 engines,
2404 modules,
2405 transfers: bound_transfers,
2406 },
2407 allocator.next_ordinal(),
2408 ))
2409}
2410
2411fn validate_transfer_endpoint(
2412 candidate: &CandidateConfig,
2413 endpoint: &TransferEndpoint,
2414) -> Result<ProviderDeviceIdentity, RuntimeConfigError> {
2415 let Some(engine) = candidate.engines.get(endpoint.engine_id()) else {
2416 return Err(RuntimeConfigError::UnknownTransferEndpointEngine {
2417 endpoint: endpoint.clone(),
2418 });
2419 };
2420 if !engine
2421 .registration
2422 .storage_classes()
2423 .contains(endpoint.storage_class())
2424 {
2425 return Err(RuntimeConfigError::UnsupportedTransferEndpointStorage {
2426 endpoint: endpoint.clone(),
2427 });
2428 }
2429 Ok(engine.registration.provider_device_identity().clone())
2430}
2431
2432struct RegistrationIdentityAllocator {
2433 issuer: NonZeroU64,
2434 next: NonZeroU64,
2435}
2436
2437impl RegistrationIdentityAllocator {
2438 fn new(issuer: NonZeroU64, next: NonZeroU64) -> Self {
2439 Self { issuer, next }
2440 }
2441
2442 fn allocate(&mut self) -> Result<RegistrationIdentity, RuntimeConfigError> {
2443 let identity = RegistrationIdentity::new(self.issuer, self.next);
2444 let next = self
2445 .next
2446 .get()
2447 .checked_add(1)
2448 .and_then(NonZeroU64::new)
2449 .ok_or(RuntimeConfigError::IdentityExhausted)?;
2450 self.next = next;
2451 Ok(identity)
2452 }
2453
2454 fn next_ordinal(&self) -> NonZeroU64 {
2455 self.next
2456 }
2457}
2458
2459fn freeze_candidate(
2460 runtime_id: RuntimeId,
2461 epoch: RuntimeEpoch,
2462 candidate: BoundCandidateConfig,
2463) -> Result<RuntimeConfigSnapshot, RuntimeConfigError> {
2464 let mut engines = Vec::with_capacity(candidate.engines.len());
2465 let mut engine_indices = BTreeMap::new();
2466 let mut engine_locations = BTreeMap::new();
2467 let mut cache_owners = Vec::new();
2468 for (index, (engine_id, record)) in candidate.engines.into_iter().enumerate() {
2469 let BoundCandidateEngineRecord {
2470 registration,
2471 identity,
2472 } = record;
2473 let event_domain_id = EventDomainId::new(runtime_id, epoch, identity);
2474 let (state, candidate_token) = registration.into_state_and_token();
2475 let provider_device_identity = state.provider_device_identity().clone();
2476 let metadata = FrozenEngineMetadata {
2477 candidate_token,
2478 identity,
2479 event_domain_id,
2480 };
2481 let frozen = match state {
2482 EngineRegistrationState::PreparationOnly { binding } => {
2483 FrozenEngineSlot::PreparationOnly(Arc::new(PreparationOnlyEngineSnapshot {
2484 metadata,
2485 binding,
2486 }))
2487 }
2488 EngineRegistrationState::Executable(binding) => {
2489 if let Some(owner) = binding.contract().cache_owner().cloned() {
2490 cache_owners.push(FrozenCacheOwner {
2491 id: engine_cache_owner_id(&engine_id),
2492 kind: FrozenCacheOwnerKind::Engine,
2493 owner,
2494 });
2495 }
2496 cache_owners.push(FrozenCacheOwner {
2497 id: engine_extension_cache_owner_id(&engine_id),
2498 kind: FrozenCacheOwnerKind::Extension,
2499 owner: execution::extension_cache_owner(binding.contract().executor().clone()),
2500 });
2501 FrozenEngineSlot::Executable(Arc::new(ExecutableEngineSnapshot {
2502 metadata,
2503 binding,
2504 }))
2505 }
2506 };
2507 engine_locations.insert(
2508 engine_id.clone(),
2509 (provider_device_identity, event_domain_id),
2510 );
2511 engine_indices.insert(engine_id, index);
2512 engines.push(frozen);
2513 }
2514 let extensions = freeze_extension_slots(candidate.modules)?;
2515 for (id, owner) in extensions.cache_owner_records() {
2516 cache_owners.push(FrozenCacheOwner {
2517 id,
2518 kind: FrozenCacheOwnerKind::Extension,
2519 owner,
2520 });
2521 }
2522 let mut transfers = BTreeMap::new();
2523 for (route, record) in candidate.transfers {
2524 let BoundCandidateTransferRecord {
2525 provider,
2526 source: source_binding,
2527 destination: destination_binding,
2528 } = record;
2529 let (_, source_event_domain_id) = bound_engine_location(&engine_locations, route.source())?;
2530 let (_, destination_event_domain_id) =
2531 bound_engine_location(&engine_locations, route.destination())?;
2532 let resolved_route = ResolvedTransferRoute::new(
2533 ResolvedTransferEndpoint::new(
2534 route.source().clone(),
2535 source_binding,
2536 *source_event_domain_id,
2537 ),
2538 ResolvedTransferEndpoint::new(
2539 route.destination().clone(),
2540 destination_binding,
2541 *destination_event_domain_id,
2542 ),
2543 );
2544 transfers.insert(resolved_route, provider);
2545 }
2546 Ok(RuntimeConfigSnapshot {
2547 runtime_id,
2548 epoch,
2549 policy: candidate.policy,
2550 engines: engines.into(),
2551 engine_indices,
2552 extensions,
2553 transfers: FrozenTransferRegistry::new(transfers),
2554 cache_owners: cache_owners.into(),
2555 })
2556}
2557
2558fn bound_engine_location<'a>(
2559 locations: &'a BTreeMap<EngineId, (ProviderDeviceIdentity, EventDomainId)>,
2560 endpoint: &TransferEndpoint,
2561) -> Result<&'a (ProviderDeviceIdentity, EventDomainId), RuntimeConfigError> {
2562 locations
2563 .get(endpoint.engine_id())
2564 .ok_or_else(|| RuntimeConfigError::BoundCandidateInvariant {
2565 endpoint: endpoint.clone(),
2566 })
2567}
2568
2569fn engine_cache_owner_id(engine_id: &EngineId) -> CacheOwnerId {
2570 let id = engine_id.as_str();
2571 CacheOwnerId::from_canonical_owner_id(Arc::<str>::from(format!("engine[{}]:{id}", id.len())))
2572}
2573
2574fn engine_extension_cache_owner_id(engine_id: &EngineId) -> CacheOwnerId {
2575 let id = engine_id.as_str();
2576 CacheOwnerId::from_canonical_owner_id(Arc::<str>::from(format!(
2577 "extension-executor[{}]:{id}",
2578 id.len()
2579 )))
2580}
2581
2582fn allocate_nonzero(counter: &AtomicU64) -> Result<NonZeroU64, RuntimeConfigError> {
2583 let value = counter
2584 .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |next| {
2585 next.checked_add(1)
2586 })
2587 .map_err(|_| RuntimeConfigError::IdentityExhausted)?;
2588 NonZeroU64::new(value).ok_or(RuntimeConfigError::IdentityExhausted)
2589}
2590
2591#[cfg(test)]
2592mod freeze_tests {
2593 use crate::{ProviderId, TransferRequest};
2594
2595 use super::*;
2596
2597 #[derive(Debug)]
2598 struct FreezeTestContext;
2599
2600 #[derive(Debug)]
2601 struct FreezeTestProvider;
2602
2603 impl TransferProvider for FreezeTestProvider {
2604 fn transfer_blocking(
2605 &self,
2606 _request: TransferRequest<'_>,
2607 ) -> crate::Result<tenferro_tensor::Tensor> {
2608 Err(crate::Error::Internal("freeze test provider".into()))
2609 }
2610 }
2611
2612 fn registration(
2613 engine_id: &str,
2614 target: &str,
2615 ) -> Result<EngineRegistration, RuntimeConfigError> {
2616 let engine_id = EngineId::new(engine_id).map_err(RuntimeConfigError::from)?;
2617 let storage =
2618 StorageClass::new("tenferro.test.freeze.storage").map_err(RuntimeConfigError::from)?;
2619 Ok(EngineRegistration::preparation_only(
2620 super::super::ProviderPreparationBinding::new(
2621 engine_id,
2622 ProviderDeviceIdentity::new(
2623 ProviderId::new("tenferro.test.freeze.provider")
2624 .map_err(RuntimeConfigError::from)?,
2625 target,
2626 )
2627 .map_err(RuntimeConfigError::from)?,
2628 ExecutionContextIdentity::of::<FreezeTestContext>(),
2629 HardwareClassId::new("tenferro.test.freeze.hardware")
2630 .map_err(RuntimeConfigError::from)?,
2631 Arc::from(vec![storage.clone()]),
2632 storage,
2633 CoreCapabilityBundle::default(),
2634 )?,
2635 ))
2636 }
2637
2638 fn candidate(binding: CandidateTransferBinding) -> Result<CandidateConfig, RuntimeConfigError> {
2639 let source_id =
2640 EngineId::new("tenferro.test.freeze.source").map_err(RuntimeConfigError::from)?;
2641 let destination_id =
2642 EngineId::new("tenferro.test.freeze.destination").map_err(RuntimeConfigError::from)?;
2643 let storage =
2644 StorageClass::new("tenferro.test.freeze.storage").map_err(RuntimeConfigError::from)?;
2645 let source_endpoint = TransferEndpoint::new(source_id.clone(), storage.clone());
2646 let destination_endpoint = TransferEndpoint::new(destination_id.clone(), storage);
2647 let mut candidate = CandidateConfig::empty();
2648 let mut changed = false;
2649 register_engine_candidate(
2650 &mut candidate,
2651 registration(source_id.as_str(), "freeze-source")?,
2652 &mut changed,
2653 )?;
2654 register_engine_candidate(
2655 &mut candidate,
2656 registration(destination_id.as_str(), "freeze-destination")?,
2657 &mut changed,
2658 )?;
2659 register_transfer_provider_candidate(
2660 &mut candidate,
2661 source_endpoint.clone(),
2662 destination_endpoint.clone(),
2663 Arc::new(FreezeTestProvider),
2664 &mut changed,
2665 )?;
2666 candidate
2667 .transfers
2668 .get_mut(&TransferRoute::new(source_endpoint, destination_endpoint))
2669 .expect("registered route")
2670 .binding = binding;
2671 Ok(candidate)
2672 }
2673
2674 #[test]
2675 fn validation_owns_stale_route_rejection_and_bound_freeze_is_total() {
2676 let wrong_source = ProviderDeviceIdentity::new(
2677 ProviderId::new("tenferro.test.freeze.provider").unwrap(),
2678 "different-source",
2679 )
2680 .unwrap();
2681 let preserved = CandidateTransferBinding::Preserved {
2682 source: wrong_source,
2683 destination: ProviderDeviceIdentity::new(
2684 ProviderId::new("tenferro.test.freeze.provider").unwrap(),
2685 "freeze-destination",
2686 )
2687 .unwrap(),
2688 };
2689 let result = validate_candidate(
2690 candidate(preserved).unwrap(),
2691 NonZeroU64::new(1).unwrap(),
2692 NonZeroU64::new(1).unwrap(),
2693 );
2694 let error = match result {
2695 Ok(_) => panic!("candidate validation must reject stale preserved bindings"),
2696 Err(error) => error,
2697 };
2698 assert!(matches!(
2699 error,
2700 RuntimeConfigError::StaleTransferRoute { .. }
2701 ));
2702
2703 let (bound, _) = validate_candidate(
2704 candidate(CandidateTransferBinding::New).unwrap(),
2705 NonZeroU64::new(1).unwrap(),
2706 NonZeroU64::new(1).unwrap(),
2707 )
2708 .expect("validation must produce a complete bound candidate");
2709 freeze_candidate(
2710 RuntimeId::from_nonzero(NonZeroU64::new(1).unwrap()),
2711 RuntimeEpoch::one(),
2712 bound,
2713 )
2714 .expect("a bound candidate must freeze without semantic route revalidation");
2715 }
2716
2717 #[test]
2718 fn frozen_engine_slots_are_arc_sized() {
2719 let slot_size = std::mem::size_of::<FrozenEngineSlot>();
2720 let arc_size = std::mem::size_of::<Arc<()>>();
2721
2722 assert!(
2723 slot_size <= 2 * arc_size,
2724 "frozen engine slots should keep immutable snapshot payloads behind Arc: slot_size={slot_size}, arc_size={arc_size}",
2725 );
2726 }
2727}