1use ferrum_types::{
2 NativeOperatorBackend, NativeOperatorContractVersion, NativeOperatorProviderCatalog,
3 NativeOperatorProviderCatalogRow, NATIVE_OPERATOR_PROVIDER_CATALOG_SCHEMA_VERSION,
4};
5use serde::{Deserialize, Deserializer, Serialize};
6use sha2::{Digest, Sha256};
7use std::collections::{BTreeMap, BTreeSet};
8
9use super::super::{
10 ContractVersion, DeviceDescriptor, NodeId, OperationId, ProviderId, VNextError,
11 WeightMaterializerDescriptor, WeightMaterializerId, MAX_WEIGHT_MATERIALIZERS,
12};
13use super::foundation::{invalid_operation, operation_error_for_node};
14use super::{
15 EngineProviderDescriptor, OperationDescriptor, OperationProviderDescriptor, OracleSpec,
16 ProviderCompatibilityRejectReason, ProviderCompatibilityRejection, ProviderCompatibilityReport,
17 ProviderCompatibilityRequest,
18};
19
20pub const MAX_OPERATION_CATALOG_ROWS: usize = 4096;
21pub const MAX_OPERATION_PROVIDER_ROWS: usize = 16384;
22pub const MAX_ENGINE_PROVIDER_ROWS: usize = 4096;
23pub const MAX_REFERENCE_ORACLE_DEPTH: usize = 64;
24
25#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
27pub struct CapabilityCatalog {
28 device: DeviceDescriptor,
29 operations: BTreeMap<OperationId, OperationDescriptor>,
30 providers: BTreeMap<OperationId, Vec<OperationProviderDescriptor>>,
31 engine_providers: BTreeMap<ProviderId, EngineProviderDescriptor>,
32 weight_materializers: BTreeMap<WeightMaterializerId, WeightMaterializerDescriptor>,
33}
34
35#[derive(Deserialize)]
36#[serde(deny_unknown_fields)]
37struct CapabilityCatalogWire {
38 device: DeviceDescriptor,
39 operations: BTreeMap<OperationId, OperationDescriptor>,
40 providers: BTreeMap<OperationId, Vec<OperationProviderDescriptor>>,
41 engine_providers: BTreeMap<ProviderId, EngineProviderDescriptor>,
42 #[serde(default = "identity_weight_materializer_descriptors")]
43 weight_materializers: BTreeMap<WeightMaterializerId, WeightMaterializerDescriptor>,
44}
45
46impl<'de> Deserialize<'de> for CapabilityCatalog {
47 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
48 where
49 D: Deserializer<'de>,
50 {
51 let wire = CapabilityCatalogWire::deserialize(deserializer)?;
52 Self::from_maps(
53 wire.device,
54 wire.operations,
55 wire.providers,
56 wire.engine_providers,
57 wire.weight_materializers,
58 )
59 .map_err(serde::de::Error::custom)
60 }
61}
62
63impl CapabilityCatalog {
64 pub fn new(
65 device: DeviceDescriptor,
66 operations: Vec<OperationDescriptor>,
67 providers: BTreeMap<OperationId, Vec<OperationProviderDescriptor>>,
68 engine_providers: Vec<EngineProviderDescriptor>,
69 ) -> Result<Self, VNextError> {
70 let mut operation_map = BTreeMap::new();
71 for operation in operations {
72 let operation_id = operation.id.clone();
73 if operation_map
74 .insert(operation_id.clone(), operation)
75 .is_some()
76 {
77 return Err(invalid_operation(format!(
78 "duplicate operation descriptor `{operation_id}`"
79 )));
80 }
81 }
82 let mut engine_map = BTreeMap::new();
83 for engine in engine_providers {
84 let provider_id = engine.provider_id().clone();
85 if engine_map.insert(provider_id.clone(), engine).is_some() {
86 return Err(invalid_operation(format!(
87 "duplicate engine provider `{provider_id}`"
88 )));
89 }
90 }
91 let weight_materializers = identity_weight_materializer_descriptors();
92 Self::from_maps(
93 device,
94 operation_map,
95 providers,
96 engine_map,
97 weight_materializers,
98 )
99 }
100
101 fn from_maps(
102 device: DeviceDescriptor,
103 operations: BTreeMap<OperationId, OperationDescriptor>,
104 mut providers: BTreeMap<OperationId, Vec<OperationProviderDescriptor>>,
105 engine_providers: BTreeMap<ProviderId, EngineProviderDescriptor>,
106 weight_materializers: BTreeMap<WeightMaterializerId, WeightMaterializerDescriptor>,
107 ) -> Result<Self, VNextError> {
108 device.validate()?;
109 validate_weight_materializer_descriptors(&device, &weight_materializers)?;
110 let provider_row_count = providers.values().try_fold(0_usize, |total, entries| {
111 total.checked_add(entries.len()).ok_or_else(|| {
112 invalid_operation("capability catalog provider row count overflows usize")
113 })
114 })?;
115 if operations.is_empty()
116 || providers.is_empty()
117 || engine_providers.is_empty()
118 || operations.len() > MAX_OPERATION_CATALOG_ROWS
119 || provider_row_count > MAX_OPERATION_PROVIDER_ROWS
120 || engine_providers.len() > MAX_ENGINE_PROVIDER_ROWS
121 {
122 return Err(invalid_operation(
123 "capability catalog is empty or exceeds its operation/provider/engine row budget",
124 ));
125 }
126 if operations.keys().collect::<BTreeSet<_>>() != providers.keys().collect::<BTreeSet<_>>() {
127 return Err(invalid_operation(
128 "capability catalog operation and provider rows do not match",
129 ));
130 }
131 for (operation_id, operation) in &operations {
132 if operation_id != &operation.id {
133 return Err(invalid_operation(format!(
134 "operation descriptor `{}` is stored under `{operation_id}`",
135 operation.id
136 )));
137 }
138 operation.validate()?;
139 if !operation
140 .provider
141 .required_capabilities
142 .is_subset(&device.capabilities)
143 {
144 return Err(VNextError::UnsupportedOperation {
145 node_id: None,
146 operation_id: operation_id.to_string(),
147 device_id: device.id.to_string(),
148 reason: "device does not advertise the operation's required capabilities"
149 .to_owned(),
150 });
151 }
152 }
153 validate_reference_oracle_graph(&operations)?;
154 for (operation_id, entries) in &mut providers {
155 if entries.is_empty() {
156 return Err(VNextError::UnsupportedOperation {
157 node_id: None,
158 operation_id: operation_id.to_string(),
159 device_id: device.id.to_string(),
160 reason: "provider row is empty".to_owned(),
161 });
162 }
163 let operation =
164 operations
165 .get(operation_id)
166 .ok_or_else(|| VNextError::UnsupportedOperation {
167 node_id: None,
168 operation_id: operation_id.to_string(),
169 device_id: device.id.to_string(),
170 reason: "provider row has no operation descriptor".to_owned(),
171 })?;
172 let operation_fingerprint = operation.fingerprint()?;
173 for entry in entries.iter() {
174 if entry.operation_id() != operation_id
175 || entry.operation_fingerprint() != operation_fingerprint
176 {
177 return Err(VNextError::UnsupportedOperation {
178 node_id: None,
179 operation_id: operation_id.to_string(),
180 device_id: device.id.to_string(),
181 reason: format!(
182 "provider `{}` is bound to a different operation descriptor",
183 entry.provider_id()
184 ),
185 });
186 }
187 if entry.device_id() != &device.id {
188 return Err(VNextError::UnsupportedOperation {
189 node_id: None,
190 operation_id: operation_id.to_string(),
191 device_id: device.id.to_string(),
192 reason: format!(
193 "provider `{}` belongs to device `{}`",
194 entry.provider_id(),
195 entry.device_id()
196 ),
197 });
198 }
199 if !entry.version().satisfies(operation.version)
200 || !entry
201 .version()
202 .satisfies(operation.provider.minimum_version)
203 {
204 return Err(VNextError::UnsupportedOperation {
205 node_id: None,
206 operation_id: operation_id.to_string(),
207 device_id: device.id.to_string(),
208 reason: format!(
209 "provider `{}` does not satisfy the operation version",
210 entry.provider_id()
211 ),
212 });
213 }
214 if !entry.capabilities().is_subset(&device.capabilities)
215 || !operation
216 .provider
217 .required_capabilities
218 .is_subset(entry.capabilities())
219 {
220 return Err(VNextError::UnsupportedOperation {
221 node_id: None,
222 operation_id: operation_id.to_string(),
223 device_id: device.id.to_string(),
224 reason: format!(
225 "provider `{}` capabilities are incompatible with the device or operation",
226 entry.provider_id()
227 ),
228 });
229 }
230 }
231 entries.sort_by(|left, right| {
232 left.provider_id()
233 .cmp(right.provider_id())
234 .then(left.version().cmp(&right.version()))
235 });
236 let mut seen = BTreeSet::new();
237 if entries
238 .iter()
239 .any(|entry| !seen.insert(entry.provider_id().clone()))
240 {
241 return Err(VNextError::UnsupportedOperation {
242 node_id: None,
243 operation_id: operation_id.to_string(),
244 device_id: device.id.to_string(),
245 reason: "duplicate provider identity".to_owned(),
246 });
247 }
248 }
249 for (provider_id, engine) in &engine_providers {
250 if provider_id != engine.provider_id()
251 || engine.device_id() != &device.id
252 || !engine.capabilities().is_subset(&device.capabilities)
253 {
254 return Err(invalid_operation(format!(
255 "engine provider `{provider_id}` identity, device, or capabilities are invalid"
256 )));
257 }
258 }
259 Ok(Self {
260 device,
261 operations,
262 providers,
263 engine_providers,
264 weight_materializers,
265 })
266 }
267
268 pub(crate) fn with_weight_materializer_descriptors(
269 mut self,
270 weight_materializers: BTreeMap<WeightMaterializerId, WeightMaterializerDescriptor>,
271 ) -> Result<Self, VNextError> {
272 validate_weight_materializer_descriptors(&self.device, &weight_materializers)?;
273 self.weight_materializers = weight_materializers;
274 Ok(self)
275 }
276
277 pub fn device(&self) -> &DeviceDescriptor {
278 &self.device
279 }
280
281 pub fn providers_for(
282 &self,
283 operation_id: &OperationId,
284 ) -> Result<&[OperationProviderDescriptor], VNextError> {
285 self.providers
286 .get(operation_id)
287 .map(Vec::as_slice)
288 .ok_or_else(|| VNextError::UnsupportedOperation {
289 node_id: None,
290 operation_id: operation_id.to_string(),
291 device_id: self.device.id.to_string(),
292 reason: "no provider is registered".to_owned(),
293 })
294 }
295
296 pub fn providers_for_node(
298 &self,
299 node_id: &NodeId,
300 operation_id: &OperationId,
301 ) -> Result<&[OperationProviderDescriptor], VNextError> {
302 self.providers_for(operation_id)
303 .map_err(|error| operation_error_for_node(error, node_id))
304 }
305
306 pub fn operation(
307 &self,
308 operation_id: &OperationId,
309 ) -> Result<&OperationDescriptor, VNextError> {
310 self.operations
311 .get(operation_id)
312 .ok_or_else(|| VNextError::UnsupportedOperation {
313 node_id: None,
314 operation_id: operation_id.to_string(),
315 device_id: self.device.id.to_string(),
316 reason: "operation descriptor is not registered".to_owned(),
317 })
318 }
319
320 pub fn operation_for_node(
322 &self,
323 node_id: &NodeId,
324 operation_id: &OperationId,
325 ) -> Result<&OperationDescriptor, VNextError> {
326 self.operation(operation_id)
327 .map_err(|error| operation_error_for_node(error, node_id))
328 }
329
330 pub fn provider_compatibility(
331 &self,
332 mut request: ProviderCompatibilityRequest,
333 ) -> Result<ProviderCompatibilityReport, VNextError> {
334 let operation = self.operation(request.operation_id())?;
335 request
336 .extend_required_capabilities(operation.provider.required_capabilities.iter().cloned());
337 let mut compatible_provider_ids = Vec::new();
338 let mut rejected = Vec::new();
339 for provider in self.providers_for(request.operation_id())? {
340 let mut reasons = Vec::new();
341 if !operation.version.satisfies(request.required_version()) {
342 reasons.push(
343 ProviderCompatibilityRejectReason::OperationVersionMismatch {
344 required: request.required_version(),
345 available: operation.version,
346 },
347 );
348 }
349 if !provider.version().satisfies(request.required_version()) {
350 reasons.push(ProviderCompatibilityRejectReason::ProviderVersionMismatch {
351 required: request.required_version(),
352 available: provider.version(),
353 });
354 }
355 let missing_capabilities = request
356 .required_capabilities()
357 .difference(provider.capabilities())
358 .cloned()
359 .collect::<BTreeSet<_>>();
360 if !missing_capabilities.is_empty() {
361 reasons.push(ProviderCompatibilityRejectReason::MissingCapabilities {
362 capabilities: missing_capabilities,
363 });
364 }
365 let missing_weight_formats = request
366 .required_weight_formats()
367 .difference(provider.accepted_weight_formats())
368 .cloned()
369 .collect::<BTreeSet<_>>();
370 if !missing_weight_formats.is_empty() {
371 reasons.push(
372 ProviderCompatibilityRejectReason::UnsupportedWeightFormats {
373 formats: missing_weight_formats,
374 },
375 );
376 }
377 let missing_quantization_formats = request
378 .required_quantization_formats()
379 .difference(provider.accepted_quantization_formats())
380 .cloned()
381 .collect::<BTreeSet<_>>();
382 if !missing_quantization_formats.is_empty() {
383 reasons.push(
384 ProviderCompatibilityRejectReason::UnsupportedQuantizationFormats {
385 formats: missing_quantization_formats,
386 },
387 );
388 }
389 if !request
390 .execution_determinism()
391 .accepts(provider.execution_semantics())
392 {
393 reasons.push(
394 ProviderCompatibilityRejectReason::InsufficientExecutionDeterminism {
395 required: request.execution_determinism(),
396 available: provider.execution_semantics(),
397 },
398 );
399 }
400 if reasons.is_empty() {
401 compatible_provider_ids.push(provider.provider_id().clone());
402 } else {
403 rejected.push(ProviderCompatibilityRejection {
404 provider_id: provider.provider_id().clone(),
405 reasons,
406 });
407 }
408 }
409 ProviderCompatibilityReport::from_classification(request, compatible_provider_ids, rejected)
410 }
411
412 pub fn operations(&self) -> &BTreeMap<OperationId, OperationDescriptor> {
413 &self.operations
414 }
415
416 pub fn providers(&self) -> &BTreeMap<OperationId, Vec<OperationProviderDescriptor>> {
417 &self.providers
418 }
419
420 pub fn engine_providers(&self) -> &BTreeMap<ProviderId, EngineProviderDescriptor> {
421 &self.engine_providers
422 }
423
424 pub fn weight_materializers(
425 &self,
426 ) -> &BTreeMap<WeightMaterializerId, WeightMaterializerDescriptor> {
427 &self.weight_materializers
428 }
429
430 pub fn weight_materializer(
431 &self,
432 materializer_id: &WeightMaterializerId,
433 ) -> Result<&WeightMaterializerDescriptor, VNextError> {
434 self.weight_materializers
435 .get(materializer_id)
436 .ok_or_else(|| {
437 invalid_operation(format!(
438 "weight materializer `{materializer_id}` is absent from the capability catalog"
439 ))
440 })
441 }
442
443 pub fn engine_provider(
444 &self,
445 provider_id: &ProviderId,
446 required_version: ContractVersion,
447 ) -> Result<&EngineProviderDescriptor, VNextError> {
448 let provider = self.engine_providers.get(provider_id).ok_or_else(|| {
449 invalid_operation(format!("engine provider `{provider_id}` is not registered"))
450 })?;
451 if !provider.contract_version().satisfies(required_version) {
452 return Err(invalid_operation(format!(
453 "engine provider `{provider_id}` version {} does not satisfy {required_version}",
454 provider.contract_version()
455 )));
456 }
457 Ok(provider)
458 }
459
460 pub fn native_operator_provider_catalog(
464 &self,
465 backend: NativeOperatorBackend,
466 ) -> Result<NativeOperatorProviderCatalog, VNextError> {
467 let mut providers = Vec::new();
468 for (operation_id, descriptors) in &self.providers {
469 let operation = self.operations.get(operation_id).ok_or_else(|| {
470 invalid_operation(
471 "capability catalog provider row lacks its operation while exporting native identities",
472 )
473 })?;
474 let operation_fingerprint = operation.fingerprint()?;
475 for provider in descriptors {
476 providers.push(NativeOperatorProviderCatalogRow {
477 operation_id: operation_id.to_string(),
478 operation_contract_version: NativeOperatorContractVersion::new(
479 operation.version.major,
480 operation.version.minor,
481 ),
482 operation_fingerprint: operation_fingerprint.clone(),
483 provider_id: provider.provider_id().to_string(),
484 provider_version: NativeOperatorContractVersion::new(
485 provider.version().major,
486 provider.version().minor,
487 ),
488 provider_implementation_fingerprint: provider
489 .provider_implementation_fingerprint()
490 .to_owned(),
491 });
492 }
493 }
494 providers.sort();
495 let catalog = NativeOperatorProviderCatalog {
496 schema_version: NATIVE_OPERATOR_PROVIDER_CATALOG_SCHEMA_VERSION,
497 backend,
498 providers,
499 };
500 catalog.validate().map_err(invalid_operation)?;
501 Ok(catalog)
502 }
503
504 pub fn fingerprint(&self) -> Result<String, VNextError> {
505 let bytes = serde_json::to_vec(self).map_err(|error| VNextError::Serialization {
506 context: "serialize capability catalog",
507 message: error.to_string(),
508 })?;
509 Ok(format!("{:x}", Sha256::digest(bytes)))
510 }
511}
512
513fn validate_weight_materializer_descriptors(
514 device: &DeviceDescriptor,
515 descriptors: &BTreeMap<WeightMaterializerId, WeightMaterializerDescriptor>,
516) -> Result<(), VNextError> {
517 if descriptors.is_empty() || descriptors.len() > MAX_WEIGHT_MATERIALIZERS {
518 return Err(invalid_operation(
519 "capability catalog weight materializers are empty or exceed their row budget",
520 ));
521 }
522 for (id, descriptor) in descriptors {
523 if id != descriptor.id() {
524 return Err(invalid_operation(format!(
525 "weight materializer `{}` is stored under `{id}`",
526 descriptor.id()
527 )));
528 }
529 descriptor.validate_for_device(device)?;
530 }
531 let identity = WeightMaterializerDescriptor::identity()?;
532 if descriptors.get(identity.id()) != Some(&identity) {
533 return Err(invalid_operation(
534 "capability catalog lacks the canonical identity weight materializer",
535 ));
536 }
537 Ok(())
538}
539
540fn identity_weight_materializer_descriptors(
541) -> BTreeMap<WeightMaterializerId, WeightMaterializerDescriptor> {
542 let identity = WeightMaterializerDescriptor::identity()
543 .expect("the built-in identity weight materializer descriptor is valid");
544 BTreeMap::from([(identity.id().clone(), identity)])
545}
546
547fn validate_reference_oracle_graph(
548 operations: &BTreeMap<OperationId, OperationDescriptor>,
549) -> Result<(), VNextError> {
550 for (operation_id, operation) in operations {
551 if let OracleSpec::ReferenceOperation {
552 operation_id: reference_id,
553 version,
554 } = &operation.oracle
555 {
556 let reference = operations.get(reference_id).ok_or_else(|| {
557 invalid_operation(format!(
558 "operation `{operation_id}` references missing oracle `{reference_id}`"
559 ))
560 })?;
561 if !reference.version.satisfies(*version) {
562 return Err(invalid_operation(format!(
563 "operation `{operation_id}` oracle `{reference_id}` version {} does not satisfy {version}",
564 reference.version
565 )));
566 }
567 if operation.inputs != reference.inputs
568 || operation.outputs != reference.outputs
569 || operation.attributes != reference.attributes
570 {
571 return Err(invalid_operation(format!(
572 "operation `{operation_id}` oracle `{reference_id}` has an incompatible input/output/attribute contract"
573 )));
574 }
575 }
576 }
577
578 #[derive(Clone, Copy, PartialEq, Eq)]
579 enum VisitState {
580 Visiting,
581 Visited,
582 }
583
584 let mut states = BTreeMap::<OperationId, VisitState>::new();
585 for root in operations.keys() {
586 if states.get(root) == Some(&VisitState::Visited) {
587 continue;
588 }
589 let mut path = Vec::<OperationId>::new();
590 let mut current = root.clone();
591 loop {
592 match states.get(¤t) {
593 Some(VisitState::Visited) => break,
594 Some(VisitState::Visiting) => {
595 return Err(invalid_operation(format!(
596 "reference-oracle graph contains a cycle at `{current}`"
597 )));
598 }
599 None => {}
600 }
601 if path.len() >= MAX_REFERENCE_ORACLE_DEPTH {
602 return Err(invalid_operation(format!(
603 "reference-oracle chain from `{root}` exceeds depth {MAX_REFERENCE_ORACLE_DEPTH}"
604 )));
605 }
606 states.insert(current.clone(), VisitState::Visiting);
607 path.push(current.clone());
608 let Some(OperationDescriptor {
609 oracle:
610 OracleSpec::ReferenceOperation {
611 operation_id: reference_id,
612 ..
613 },
614 ..
615 }) = operations.get(¤t)
616 else {
617 break;
618 };
619 current = reference_id.clone();
620 }
621 for operation_id in path {
622 states.insert(operation_id, VisitState::Visited);
623 }
624 }
625 Ok(())
626}