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 if existing
1814 .engines
1815 .contains_key(&(family_id, engine_id.clone()))
1816 {
1817 return Ok(self);
1818 }
1819 }
1820
1821 let record = configure_module(Arc::clone(&value))
1822 .map_err(|source| RuntimeConfigError::ExtensionModule { source })?;
1823 if !record.engines.contains_key(&(family_id, engine_id.clone())) {
1824 return Err(RuntimeConfigError::MissingExtensionEngine {
1825 module_id,
1826 family_id,
1827 engine_id: engine_id.clone(),
1828 });
1829 }
1830 self.candidate.modules.insert(module_id, record);
1831 *self.changed = true;
1832 Ok(self)
1833 }
1834
1835 #[doc(hidden)]
1895 pub fn replace_extension_module_for_engine(
1896 &mut self,
1897 value: Arc<dyn ExtensionModule>,
1898 family_id: &'static str,
1899 engine_id: &EngineId,
1900 ) -> Result<&mut Self, RuntimeConfigError> {
1901 let module_id = value.module_id().clone();
1902 if let Some(existing) = self.candidate.modules.get(&module_id) {
1903 if existing.module_identical(&value) {
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
1918 let record = configure_module(Arc::clone(&value))
1919 .map_err(|source| RuntimeConfigError::ExtensionModule { source })?;
1920 if !record.engines.contains_key(&(family_id, engine_id.clone())) {
1921 return Err(RuntimeConfigError::MissingExtensionEngine {
1922 module_id,
1923 family_id,
1924 engine_id: engine_id.clone(),
1925 });
1926 }
1927 self.candidate.modules.insert(module_id, record);
1928 *self.changed = true;
1929 Ok(self)
1930 }
1931
1932 pub fn remove_extension_module(
1939 &mut self,
1940 id: &ExtensionModuleId,
1941 ) -> Result<&mut Self, RuntimeConfigError> {
1942 remove_extension_module_candidate(self.candidate, id, self.changed)?;
1943 Ok(self)
1944 }
1945}
1946
1947impl fmt::Debug for RuntimeReconfiguration<'_> {
1948 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
1949 formatter
1950 .debug_struct("RuntimeReconfiguration")
1951 .field("engine_count", &self.candidate.engines.len())
1952 .field("extension_module_count", &self.candidate.modules.len())
1953 .field("transfer_provider_count", &self.candidate.transfers.len())
1954 .field("changed", &*self.changed)
1955 .finish_non_exhaustive()
1956 }
1957}
1958
1959#[derive(Clone, Copy)]
1971pub struct EngineSnapshotView<'a> {
1972 slot: &'a FrozenEngineSlot,
1973}
1974
1975impl<'a> EngineSnapshotView<'a> {
1976 pub fn engine_id(&self) -> &'a EngineId {
1978 self.slot.engine_id()
1979 }
1980
1981 pub fn provider_device_identity(&self) -> &'a super::ProviderDeviceIdentity {
1991 self.slot.provider_device_identity()
1992 }
1993
1994 pub fn registration_identity(&self) -> RegistrationIdentity {
1996 self.slot.metadata().identity
1997 }
1998
1999 pub fn context_identity(&self) -> ExecutionContextIdentity {
2001 self.slot.context_identity()
2002 }
2003
2004 pub fn event_domain_id(&self) -> EventDomainId {
2006 self.slot.metadata().event_domain_id
2007 }
2008
2009 pub(super) fn executable_witness(&self) -> Option<&'a Arc<ExecutableEngineSnapshot>> {
2010 self.slot.executable()
2011 }
2012
2013 pub fn hardware_class(&self) -> &'a HardwareClassId {
2015 self.slot.hardware_class()
2016 }
2017
2018 pub fn capabilities(&self) -> &'a CoreCapabilityBundle {
2020 self.slot.capabilities()
2021 }
2022
2023 pub(super) fn storage_classes(&self) -> &'a [StorageClass] {
2024 self.slot.storage_classes()
2025 }
2026
2027 pub(super) fn default_storage_class(&self) -> &'a StorageClass {
2028 self.slot.default_storage_class()
2029 }
2030
2031 #[doc(hidden)]
2058 pub fn accepts_input_signature(&self, input: &super::InputSignatureEntry) -> bool {
2059 self.slot
2060 .storage_classes()
2061 .iter()
2062 .any(|storage_class| self.accepts_input_signature_for_storage(input, storage_class))
2063 }
2064
2065 pub(super) fn accepts_input_signature_for_storage(
2066 &self,
2067 input: &super::InputSignatureEntry,
2068 storage_class: &StorageClass,
2069 ) -> bool {
2070 self.slot
2071 .executable()
2072 .is_some_and(|snapshot| snapshot.accepts_input_signature(input, storage_class))
2073 }
2074
2075 #[cfg(test)]
2076 pub(crate) fn has_execution_engine_for_test(&self) -> bool {
2077 self.slot.executable().is_some()
2078 }
2079}
2080
2081impl fmt::Debug for EngineSnapshotView<'_> {
2082 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
2083 formatter
2084 .debug_struct("EngineSnapshotView")
2085 .field("engine_id", self.engine_id())
2086 .field("registration_identity", &self.registration_identity())
2087 .field("context_identity", &self.context_identity())
2088 .field("hardware_class", self.hardware_class())
2089 .field("capabilities", self.capabilities())
2090 .finish()
2091 }
2092}
2093
2094fn default_execution_policy() -> ExecutionPolicy {
2095 ExecutionPolicy::new(super::Determinism::Fast, None, 0)
2096}
2097
2098fn register_engine_candidate(
2099 candidate: &mut CandidateConfig,
2100 registration: EngineRegistration,
2101 changed: &mut bool,
2102) -> Result<(), RuntimeConfigError> {
2103 let engine_id = registration.engine_id().clone();
2104 match candidate.engines.get(&engine_id) {
2105 Some(existing) if existing.registration.candidate_identical(®istration) => Ok(()),
2106 Some(_) => Err(RuntimeConfigError::DuplicateEngine { engine_id }),
2107 None => {
2108 ensure_unique_provider_device_target(candidate, ®istration)?;
2109 candidate.engines.insert(
2110 engine_id,
2111 CandidateEngineRecord {
2112 registration,
2113 identity: CandidateRegistrationIdentity::New,
2114 },
2115 );
2116 *changed = true;
2117 Ok(())
2118 }
2119 }
2120}
2121
2122fn replace_engine_candidate(
2123 candidate: &mut CandidateConfig,
2124 registration: EngineRegistration,
2125 changed: &mut bool,
2126) -> Result<(), RuntimeConfigError> {
2127 let engine_id = registration.engine_id().clone();
2128 let Some(existing) = candidate.engines.get(&engine_id) else {
2129 return Err(RuntimeConfigError::MissingEngine { engine_id });
2130 };
2131 if existing.registration.candidate_identical(®istration) {
2132 return Ok(());
2133 }
2134 if existing.registration.provider_device_identity() != registration.provider_device_identity() {
2135 return Err(RuntimeConfigError::EngineTargetRebind {
2136 engine_id,
2137 current: existing.registration.provider_device_identity().clone(),
2138 replacement: registration.provider_device_identity().clone(),
2139 });
2140 }
2141 ensure_unique_provider_device_target_except(candidate, ®istration, &engine_id)?;
2142 candidate.engines.insert(
2143 engine_id,
2144 CandidateEngineRecord {
2145 registration,
2146 identity: CandidateRegistrationIdentity::New,
2147 },
2148 );
2149 *changed = true;
2150 Ok(())
2151}
2152
2153fn remove_engine_candidate(
2154 candidate: &mut CandidateConfig,
2155 id: &EngineId,
2156 changed: &mut bool,
2157) -> Result<(), RuntimeConfigError> {
2158 match candidate.engines.remove(id) {
2159 Some(_) => {
2160 *changed = true;
2161 Ok(())
2162 }
2163 None => Err(RuntimeConfigError::MissingEngine {
2164 engine_id: id.clone(),
2165 }),
2166 }
2167}
2168
2169fn install_extension_module_candidate(
2170 candidate: &mut CandidateConfig,
2171 module: Arc<dyn ExtensionModule>,
2172 changed: &mut bool,
2173) -> Result<(), RuntimeConfigError> {
2174 let module_id = module.module_id().clone();
2175 match candidate.modules.get(&module_id) {
2176 Some(existing) if existing.module_identical(&module) => Ok(()),
2177 Some(_) => Err(RuntimeConfigError::ExtensionModule {
2178 source: ExtensionModuleError::ConflictingModule { module_id },
2179 }),
2180 None => {
2181 let record = configure_module(module)
2182 .map_err(|source| RuntimeConfigError::ExtensionModule { source })?;
2183 candidate.modules.insert(module_id, record);
2184 *changed = true;
2185 Ok(())
2186 }
2187 }
2188}
2189
2190fn replace_extension_module_candidate(
2191 candidate: &mut CandidateConfig,
2192 module: Arc<dyn ExtensionModule>,
2193 changed: &mut bool,
2194) -> Result<(), RuntimeConfigError> {
2195 let module_id = module.module_id().clone();
2196 match candidate.modules.get(&module_id) {
2197 Some(existing) if existing.module_identical(&module) => Ok(()),
2198 _ => {
2199 let record = configure_module(module)
2200 .map_err(|source| RuntimeConfigError::ExtensionModule { source })?;
2201 candidate.modules.insert(module_id, record);
2202 *changed = true;
2203 Ok(())
2204 }
2205 }
2206}
2207
2208fn remove_extension_module_candidate(
2209 candidate: &mut CandidateConfig,
2210 id: &ExtensionModuleId,
2211 changed: &mut bool,
2212) -> Result<(), RuntimeConfigError> {
2213 if candidate.modules.remove(id).is_some() {
2214 *changed = true;
2215 }
2216 Ok(())
2217}
2218
2219fn register_transfer_provider_candidate(
2220 candidate: &mut CandidateConfig,
2221 source: TransferEndpoint,
2222 destination: TransferEndpoint,
2223 provider: Arc<dyn TransferProvider>,
2224 changed: &mut bool,
2225) -> Result<(), RuntimeConfigError> {
2226 let key = TransferRoute::new(source, destination);
2227 match candidate.transfers.get(&key) {
2228 Some(existing) if Arc::ptr_eq(&existing.provider, &provider) => Ok(()),
2229 Some(_) => Err(RuntimeConfigError::ConflictingRegistration {
2230 key: RegistrationKey::TransferProvider {
2231 source: key.source().clone(),
2232 destination: key.destination().clone(),
2233 },
2234 }),
2235 None => {
2236 candidate.transfers.insert(
2237 key,
2238 CandidateTransferRecord {
2239 provider,
2240 binding: CandidateTransferBinding::New,
2241 },
2242 );
2243 *changed = true;
2244 Ok(())
2245 }
2246 }
2247}
2248
2249fn remove_transfer_provider_candidate(
2250 candidate: &mut CandidateConfig,
2251 source: TransferEndpoint,
2252 destination: TransferEndpoint,
2253 changed: &mut bool,
2254) -> Result<(), RuntimeConfigError> {
2255 let key = TransferRoute::new(source, destination);
2256 if candidate.transfers.remove(&key).is_none() {
2257 return Err(RuntimeConfigError::MissingTransferProvider {
2258 source_endpoint: key.source().clone(),
2259 destination: key.destination().clone(),
2260 });
2261 }
2262 *changed = true;
2263 Ok(())
2264}
2265
2266fn ensure_unique_provider_device_target(
2267 candidate: &CandidateConfig,
2268 registration: &EngineRegistration,
2269) -> Result<(), RuntimeConfigError> {
2270 ensure_unique_provider_device_target_except(candidate, registration, registration.engine_id())
2271}
2272
2273fn ensure_unique_provider_device_target_except(
2274 candidate: &CandidateConfig,
2275 registration: &EngineRegistration,
2276 ignored_engine_id: &EngineId,
2277) -> Result<(), RuntimeConfigError> {
2278 if let Some((first_engine_id, _)) = candidate.engines.iter().find(|(engine_id, record)| {
2279 *engine_id != ignored_engine_id
2280 && record.registration.provider_device_identity()
2281 == registration.provider_device_identity()
2282 }) {
2283 return Err(RuntimeConfigError::DuplicateProviderDeviceTarget {
2284 provider_device_identity: registration.provider_device_identity().clone(),
2285 first_engine_id: first_engine_id.clone(),
2286 duplicate_engine_id: registration.engine_id().clone(),
2287 });
2288 }
2289 Ok(())
2290}
2291
2292fn validate_candidate(
2293 candidate: CandidateConfig,
2294 issuer: NonZeroU64,
2295 next_ordinal: NonZeroU64,
2296) -> Result<(BoundCandidateConfig, NonZeroU64), RuntimeConfigError> {
2297 let mut seen_targets = BTreeMap::<ProviderDeviceIdentity, EngineId>::new();
2298 for (engine_id, record) in &candidate.engines {
2299 if let Some(first_engine_id) = seen_targets.insert(
2300 record.registration.provider_device_identity().clone(),
2301 engine_id.clone(),
2302 ) {
2303 return Err(RuntimeConfigError::DuplicateProviderDeviceTarget {
2304 provider_device_identity: record.registration.provider_device_identity().clone(),
2305 first_engine_id,
2306 duplicate_engine_id: engine_id.clone(),
2307 });
2308 }
2309 }
2310
2311 let mut bound_transfers = BTreeMap::new();
2312 for (route, record) in &candidate.transfers {
2313 let source_binding = validate_transfer_endpoint(&candidate, route.source())?;
2314 let destination_binding = validate_transfer_endpoint(&candidate, route.destination())?;
2315 let preserved = match &record.binding {
2316 CandidateTransferBinding::New => None,
2317 CandidateTransferBinding::Preserved {
2318 source,
2319 destination,
2320 } => Some((source, destination)),
2321 };
2322 if let Some((registered_source, registered_destination)) = preserved {
2323 if registered_source != &source_binding {
2324 return Err(RuntimeConfigError::StaleTransferRoute {
2325 source_endpoint: route.source().clone(),
2326 destination: route.destination().clone(),
2327 endpoint: route.source().clone(),
2328 registered: Box::new(registered_source.clone()),
2329 current: Box::new(source_binding.clone()),
2330 });
2331 }
2332 if registered_destination != &destination_binding {
2333 return Err(RuntimeConfigError::StaleTransferRoute {
2334 source_endpoint: route.source().clone(),
2335 destination: route.destination().clone(),
2336 endpoint: route.destination().clone(),
2337 registered: Box::new(registered_destination.clone()),
2338 current: Box::new(destination_binding.clone()),
2339 });
2340 }
2341 }
2342 bound_transfers.insert(
2343 route.clone(),
2344 BoundCandidateTransferRecord {
2345 provider: Arc::clone(&record.provider),
2346 source: source_binding,
2347 destination: destination_binding,
2348 },
2349 );
2350 }
2351 let mut seen = BTreeMap::<(ExtensionFamilyId, EngineId), ExtensionModuleId>::new();
2352 for (module_id, module) in &candidate.modules {
2353 for family_engine in module.engines.keys() {
2354 if seen
2355 .insert(
2356 (family_engine.0, family_engine.1.clone()),
2357 module_id.clone(),
2358 )
2359 .is_some()
2360 {
2361 return Err(RuntimeConfigError::ConflictingRegistration {
2362 key: RegistrationKey::ExtensionEngine {
2363 family: family_engine.0,
2364 engine: family_engine.1.clone(),
2365 },
2366 });
2367 }
2368 }
2369 }
2370
2371 let CandidateConfig {
2372 policy,
2373 engines,
2374 modules,
2375 transfers: _,
2376 } = candidate;
2377 let mut allocator = RegistrationIdentityAllocator::new(issuer, next_ordinal);
2378 let engines = engines
2379 .into_iter()
2380 .map(|(engine_id, record)| {
2381 let identity = match record.identity {
2382 CandidateRegistrationIdentity::New => allocator.allocate()?,
2383 CandidateRegistrationIdentity::Preserved(identity) => identity,
2384 };
2385 Ok((
2386 engine_id,
2387 BoundCandidateEngineRecord {
2388 registration: record.registration,
2389 identity,
2390 },
2391 ))
2392 })
2393 .collect::<Result<BTreeMap<_, _>, RuntimeConfigError>>()?;
2394 let modules = modules
2395 .into_iter()
2396 .map(|(module_id, module)| {
2397 let mut allocate = || allocator.allocate();
2398 Ok((module_id, bind_candidate_module(module, &mut allocate)?))
2399 })
2400 .collect::<Result<BTreeMap<_, _>, RuntimeConfigError>>()?;
2401 Ok((
2402 BoundCandidateConfig {
2403 policy,
2404 engines,
2405 modules,
2406 transfers: bound_transfers,
2407 },
2408 allocator.next_ordinal(),
2409 ))
2410}
2411
2412fn validate_transfer_endpoint(
2413 candidate: &CandidateConfig,
2414 endpoint: &TransferEndpoint,
2415) -> Result<ProviderDeviceIdentity, RuntimeConfigError> {
2416 let Some(engine) = candidate.engines.get(endpoint.engine_id()) else {
2417 return Err(RuntimeConfigError::UnknownTransferEndpointEngine {
2418 endpoint: endpoint.clone(),
2419 });
2420 };
2421 if !engine
2422 .registration
2423 .storage_classes()
2424 .contains(endpoint.storage_class())
2425 {
2426 return Err(RuntimeConfigError::UnsupportedTransferEndpointStorage {
2427 endpoint: endpoint.clone(),
2428 });
2429 }
2430 Ok(engine.registration.provider_device_identity().clone())
2431}
2432
2433struct RegistrationIdentityAllocator {
2434 issuer: NonZeroU64,
2435 next: NonZeroU64,
2436}
2437
2438impl RegistrationIdentityAllocator {
2439 fn new(issuer: NonZeroU64, next: NonZeroU64) -> Self {
2440 Self { issuer, next }
2441 }
2442
2443 fn allocate(&mut self) -> Result<RegistrationIdentity, RuntimeConfigError> {
2444 let identity = RegistrationIdentity::new(self.issuer, self.next);
2445 let next = self
2446 .next
2447 .get()
2448 .checked_add(1)
2449 .and_then(NonZeroU64::new)
2450 .ok_or(RuntimeConfigError::IdentityExhausted)?;
2451 self.next = next;
2452 Ok(identity)
2453 }
2454
2455 fn next_ordinal(&self) -> NonZeroU64 {
2456 self.next
2457 }
2458}
2459
2460fn freeze_candidate(
2461 runtime_id: RuntimeId,
2462 epoch: RuntimeEpoch,
2463 candidate: BoundCandidateConfig,
2464) -> Result<RuntimeConfigSnapshot, RuntimeConfigError> {
2465 let mut engines = Vec::with_capacity(candidate.engines.len());
2466 let mut engine_indices = BTreeMap::new();
2467 let mut engine_locations = BTreeMap::new();
2468 let mut cache_owners = Vec::new();
2469 for (index, (engine_id, record)) in candidate.engines.into_iter().enumerate() {
2470 let BoundCandidateEngineRecord {
2471 registration,
2472 identity,
2473 } = record;
2474 let event_domain_id = EventDomainId::new(runtime_id, epoch, identity);
2475 let (state, candidate_token) = registration.into_state_and_token();
2476 let provider_device_identity = state.provider_device_identity().clone();
2477 let metadata = FrozenEngineMetadata {
2478 candidate_token,
2479 identity,
2480 event_domain_id,
2481 };
2482 let frozen = match state {
2483 EngineRegistrationState::PreparationOnly { binding } => {
2484 FrozenEngineSlot::PreparationOnly(Arc::new(PreparationOnlyEngineSnapshot {
2485 metadata,
2486 binding,
2487 }))
2488 }
2489 EngineRegistrationState::Executable(binding) => {
2490 if let Some(owner) = binding.contract().cache_owner().cloned() {
2491 cache_owners.push(FrozenCacheOwner {
2492 id: engine_cache_owner_id(&engine_id),
2493 kind: FrozenCacheOwnerKind::Engine,
2494 owner,
2495 });
2496 }
2497 cache_owners.push(FrozenCacheOwner {
2498 id: engine_extension_cache_owner_id(&engine_id),
2499 kind: FrozenCacheOwnerKind::Extension,
2500 owner: execution::extension_cache_owner(binding.contract().executor().clone()),
2501 });
2502 FrozenEngineSlot::Executable(Arc::new(ExecutableEngineSnapshot {
2503 metadata,
2504 binding,
2505 }))
2506 }
2507 };
2508 engine_locations.insert(
2509 engine_id.clone(),
2510 (provider_device_identity, event_domain_id),
2511 );
2512 engine_indices.insert(engine_id, index);
2513 engines.push(frozen);
2514 }
2515 let extensions = freeze_extension_slots(candidate.modules)?;
2516 for (id, owner) in extensions.cache_owner_records() {
2517 cache_owners.push(FrozenCacheOwner {
2518 id,
2519 kind: FrozenCacheOwnerKind::Extension,
2520 owner,
2521 });
2522 }
2523 let mut transfers = BTreeMap::new();
2524 for (route, record) in candidate.transfers {
2525 let BoundCandidateTransferRecord {
2526 provider,
2527 source: source_binding,
2528 destination: destination_binding,
2529 } = record;
2530 let (_, source_event_domain_id) = bound_engine_location(&engine_locations, route.source())?;
2531 let (_, destination_event_domain_id) =
2532 bound_engine_location(&engine_locations, route.destination())?;
2533 let resolved_route = ResolvedTransferRoute::new(
2534 ResolvedTransferEndpoint::new(
2535 route.source().clone(),
2536 source_binding,
2537 *source_event_domain_id,
2538 ),
2539 ResolvedTransferEndpoint::new(
2540 route.destination().clone(),
2541 destination_binding,
2542 *destination_event_domain_id,
2543 ),
2544 );
2545 transfers.insert(resolved_route, provider);
2546 }
2547 Ok(RuntimeConfigSnapshot {
2548 runtime_id,
2549 epoch,
2550 policy: candidate.policy,
2551 engines: engines.into(),
2552 engine_indices,
2553 extensions,
2554 transfers: FrozenTransferRegistry::new(transfers),
2555 cache_owners: cache_owners.into(),
2556 })
2557}
2558
2559fn bound_engine_location<'a>(
2560 locations: &'a BTreeMap<EngineId, (ProviderDeviceIdentity, EventDomainId)>,
2561 endpoint: &TransferEndpoint,
2562) -> Result<&'a (ProviderDeviceIdentity, EventDomainId), RuntimeConfigError> {
2563 locations
2564 .get(endpoint.engine_id())
2565 .ok_or_else(|| RuntimeConfigError::BoundCandidateInvariant {
2566 endpoint: endpoint.clone(),
2567 })
2568}
2569
2570fn engine_cache_owner_id(engine_id: &EngineId) -> CacheOwnerId {
2571 let id = engine_id.as_str();
2572 CacheOwnerId::from_canonical_owner_id(Arc::<str>::from(format!("engine[{}]:{id}", id.len())))
2573}
2574
2575fn engine_extension_cache_owner_id(engine_id: &EngineId) -> CacheOwnerId {
2576 let id = engine_id.as_str();
2577 CacheOwnerId::from_canonical_owner_id(Arc::<str>::from(format!(
2578 "extension-executor[{}]:{id}",
2579 id.len()
2580 )))
2581}
2582
2583fn allocate_nonzero(counter: &AtomicU64) -> Result<NonZeroU64, RuntimeConfigError> {
2584 let value = counter
2585 .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |next| {
2586 next.checked_add(1)
2587 })
2588 .map_err(|_| RuntimeConfigError::IdentityExhausted)?;
2589 NonZeroU64::new(value).ok_or(RuntimeConfigError::IdentityExhausted)
2590}
2591
2592#[cfg(test)]
2593mod freeze_tests {
2594 use crate::{ProviderId, TransferRequest};
2595
2596 use super::*;
2597
2598 #[derive(Debug)]
2599 struct FreezeTestContext;
2600
2601 #[derive(Debug)]
2602 struct FreezeTestProvider;
2603
2604 impl TransferProvider for FreezeTestProvider {
2605 fn transfer_blocking(
2606 &self,
2607 _request: TransferRequest<'_>,
2608 ) -> crate::Result<tenferro_tensor::Tensor> {
2609 Err(crate::Error::Internal("freeze test provider".into()))
2610 }
2611 }
2612
2613 fn registration(
2614 engine_id: &str,
2615 target: &str,
2616 ) -> Result<EngineRegistration, RuntimeConfigError> {
2617 let engine_id = EngineId::new(engine_id).map_err(RuntimeConfigError::from)?;
2618 let storage =
2619 StorageClass::new("tenferro.test.freeze.storage").map_err(RuntimeConfigError::from)?;
2620 Ok(EngineRegistration::preparation_only(
2621 super::super::ProviderPreparationBinding::new(
2622 engine_id,
2623 ProviderDeviceIdentity::new(
2624 ProviderId::new("tenferro.test.freeze.provider")
2625 .map_err(RuntimeConfigError::from)?,
2626 target,
2627 )
2628 .map_err(RuntimeConfigError::from)?,
2629 ExecutionContextIdentity::of::<FreezeTestContext>(),
2630 HardwareClassId::new("tenferro.test.freeze.hardware")
2631 .map_err(RuntimeConfigError::from)?,
2632 Arc::from(vec![storage.clone()]),
2633 storage,
2634 CoreCapabilityBundle::default(),
2635 )?,
2636 ))
2637 }
2638
2639 fn candidate(binding: CandidateTransferBinding) -> Result<CandidateConfig, RuntimeConfigError> {
2640 let source_id =
2641 EngineId::new("tenferro.test.freeze.source").map_err(RuntimeConfigError::from)?;
2642 let destination_id =
2643 EngineId::new("tenferro.test.freeze.destination").map_err(RuntimeConfigError::from)?;
2644 let storage =
2645 StorageClass::new("tenferro.test.freeze.storage").map_err(RuntimeConfigError::from)?;
2646 let source_endpoint = TransferEndpoint::new(source_id.clone(), storage.clone());
2647 let destination_endpoint = TransferEndpoint::new(destination_id.clone(), storage);
2648 let mut candidate = CandidateConfig::empty();
2649 let mut changed = false;
2650 register_engine_candidate(
2651 &mut candidate,
2652 registration(source_id.as_str(), "freeze-source")?,
2653 &mut changed,
2654 )?;
2655 register_engine_candidate(
2656 &mut candidate,
2657 registration(destination_id.as_str(), "freeze-destination")?,
2658 &mut changed,
2659 )?;
2660 register_transfer_provider_candidate(
2661 &mut candidate,
2662 source_endpoint.clone(),
2663 destination_endpoint.clone(),
2664 Arc::new(FreezeTestProvider),
2665 &mut changed,
2666 )?;
2667 candidate
2668 .transfers
2669 .get_mut(&TransferRoute::new(source_endpoint, destination_endpoint))
2670 .expect("registered route")
2671 .binding = binding;
2672 Ok(candidate)
2673 }
2674
2675 #[test]
2676 fn validation_owns_stale_route_rejection_and_bound_freeze_is_total() {
2677 let wrong_source = ProviderDeviceIdentity::new(
2678 ProviderId::new("tenferro.test.freeze.provider").unwrap(),
2679 "different-source",
2680 )
2681 .unwrap();
2682 let preserved = CandidateTransferBinding::Preserved {
2683 source: wrong_source,
2684 destination: ProviderDeviceIdentity::new(
2685 ProviderId::new("tenferro.test.freeze.provider").unwrap(),
2686 "freeze-destination",
2687 )
2688 .unwrap(),
2689 };
2690 let result = validate_candidate(
2691 candidate(preserved).unwrap(),
2692 NonZeroU64::new(1).unwrap(),
2693 NonZeroU64::new(1).unwrap(),
2694 );
2695 let error = match result {
2696 Ok(_) => panic!("candidate validation must reject stale preserved bindings"),
2697 Err(error) => error,
2698 };
2699 assert!(matches!(
2700 error,
2701 RuntimeConfigError::StaleTransferRoute { .. }
2702 ));
2703
2704 let (bound, _) = validate_candidate(
2705 candidate(CandidateTransferBinding::New).unwrap(),
2706 NonZeroU64::new(1).unwrap(),
2707 NonZeroU64::new(1).unwrap(),
2708 )
2709 .expect("validation must produce a complete bound candidate");
2710 freeze_candidate(
2711 RuntimeId::from_nonzero(NonZeroU64::new(1).unwrap()),
2712 RuntimeEpoch::one(),
2713 bound,
2714 )
2715 .expect("a bound candidate must freeze without semantic route revalidation");
2716 }
2717
2718 #[test]
2719 fn frozen_engine_slots_are_arc_sized() {
2720 let slot_size = std::mem::size_of::<FrozenEngineSlot>();
2721 let arc_size = std::mem::size_of::<Arc<()>>();
2722
2723 assert!(
2724 slot_size <= 2 * arc_size,
2725 "frozen engine slots should keep immutable snapshot payloads behind Arc: slot_size={slot_size}, arc_size={arc_size}",
2726 );
2727 }
2728}