1use std::collections::{BTreeMap, BTreeSet};
10
11use serde::{Deserialize, Serialize};
12use sha2::{Digest, Sha256};
13
14use crate::campaign::stable_json_fingerprint;
15use crate::canonical::parse_typed_json;
16use crate::fold::FoldSet;
17use crate::metrics::ScoreSet;
18use crate::plan::CampaignSpec;
19use crate::runtime::InMemoryLineageRecorder;
20use crate::selection::{RefitStrategy, SelectionPolicy};
21use crate::training::{TrainingInfluenceKind, TrainingInfluenceManifest};
22
23pub const HPO_MANIFEST_SCHEMA_VERSION: u32 = 1;
24pub const N4MOPT_CHECKPOINT_SCHEMA_VERSION: u32 = 1;
25pub const N4MOPT_ARTIFACT_KIND: &str = "n4m_optimizer_checkpoint";
26pub const N4MOPT_FORMAT: &str = "N4MOPT";
27pub const MAX_N4MOPT_CHECKPOINT_BYTES: usize = 64 * 1024 * 1024;
30
31pub(crate) fn campaign_provenance_fingerprint(campaign: &CampaignSpec) -> crate::Result<String> {
37 let mut canonical = campaign.clone();
38 if let Some(serde_json::Value::Object(operation)) =
39 canonical.metadata.get_mut("methods_hpo_operation")
40 {
41 operation.remove("resume_package_json");
42 operation.remove("trials");
45 }
46 stable_json_fingerprint(&canonical)
47}
48
49pub type HpoResult<T> = std::result::Result<T, HpoError>;
50
51#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
54#[serde(deny_unknown_fields)]
55pub struct HpoNativeError {
56 pub status: i32,
57 pub kind: String,
58 pub message: String,
59 pub retryable: bool,
60}
61
62#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
63#[serde(tag = "code", rename_all = "snake_case")]
64pub enum HpoError {
65 MethodsOptimizerFeatureDisabled,
66 RuntimeConfiguration {
67 reason: String,
68 },
69 InvalidManifest {
70 reason: String,
71 },
72 InvalidSearchSpace {
73 reason: String,
74 },
75 InvalidTrial {
76 reason: String,
77 },
78 InvalidCheckpoint {
79 reason: String,
80 },
81 CheckpointBindingMismatch {
82 reason: String,
83 },
84 Native {
85 operation: String,
86 error: HpoNativeError,
87 },
88 PartialBatch {
89 committed: Vec<HpoTrial>,
90 error: HpoNativeError,
91 },
92 Evaluation {
93 reason: String,
94 },
95}
96
97impl std::fmt::Display for HpoError {
98 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
99 match self {
100 Self::MethodsOptimizerFeatureDisabled => f.write_str(
101 "Methods optimizer support is disabled; enable the published `methods-optimizer` feature",
102 ),
103 Self::RuntimeConfiguration { reason } => {
104 write!(f, "invalid Methods runtime configuration: {reason}")
105 }
106 Self::InvalidManifest { reason } => write!(f, "invalid HPO manifest: {reason}"),
107 Self::InvalidSearchSpace { reason } => write!(f, "invalid HPO search space: {reason}"),
108 Self::InvalidTrial { reason } => write!(f, "invalid native HPO trial: {reason}"),
109 Self::InvalidCheckpoint { reason } => write!(f, "invalid N4MOPT checkpoint: {reason}"),
110 Self::CheckpointBindingMismatch { reason } => write!(f, "checkpoint binding mismatch: {reason}"),
111 Self::Native { operation, error } => write!(f, "n4m {operation} failed ({}/{}): {}", error.kind, error.status, error.message),
112 Self::PartialBatch { committed, error } => write!(f, "n4m ask_batch committed {} trial(s), then failed ({}/{}): {}", committed.len(), error.kind, error.status, error.message),
113 Self::Evaluation { reason } => write!(f, "HPO evaluation failed: {reason}"),
114 }
115 }
116}
117impl std::error::Error for HpoError {}
118
119#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
120#[serde(deny_unknown_fields)]
121pub struct HpoStudyBinding {
122 pub controller_id: String,
123 pub study_id: String,
124 pub search_space_fingerprint: String,
125 pub optimizer_fingerprint: String,
126}
127
128impl HpoStudyBinding {
129 pub fn validate(&self) -> HpoResult<()> {
130 for (field, value) in [
131 ("controller_id", &self.controller_id),
132 ("study_id", &self.study_id),
133 ("search_space_fingerprint", &self.search_space_fingerprint),
134 ("optimizer_fingerprint", &self.optimizer_fingerprint),
135 ] {
136 if value.trim().is_empty() {
137 return Err(HpoError::InvalidManifest {
138 reason: format!("{field} must not be empty"),
139 });
140 }
141 }
142 Ok(())
143 }
144}
145
146#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
147#[serde(deny_unknown_fields)]
148pub struct MethodsHpoControllerManifest {
149 pub schema_version: u32,
150 pub binding: HpoStudyBinding,
151}
152
153impl MethodsHpoControllerManifest {
154 pub fn validate(&self) -> HpoResult<()> {
155 if self.schema_version != HPO_MANIFEST_SCHEMA_VERSION {
156 return Err(HpoError::InvalidManifest {
157 reason: format!(
158 "unsupported schema_version {}; expected {HPO_MANIFEST_SCHEMA_VERSION}",
159 self.schema_version
160 ),
161 });
162 }
163 self.binding.validate()
164 }
165}
166
167#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
170#[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)]
171pub enum HpoParameter {
172 Int {
173 name: String,
174 low: i64,
175 high: i64,
176 step: i64,
177 log: bool,
178 },
179 Float {
180 name: String,
181 low: f64,
182 high: f64,
183 step: f64,
184 log: bool,
185 },
186 Categorical {
187 name: String,
188 values: Vec<HpoCategory>,
189 },
190 Ordinal {
191 name: String,
192 values: Vec<f64>,
193 },
194 SortedTuple {
195 name: String,
196 length: i32,
197 low: f64,
198 high: f64,
199 integer: bool,
200 },
201}
202
203impl HpoParameter {
204 fn name(&self) -> &str {
205 match self {
206 Self::Int { name, .. }
207 | Self::Float { name, .. }
208 | Self::Categorical { name, .. }
209 | Self::Ordinal { name, .. }
210 | Self::SortedTuple { name, .. } => name,
211 }
212 }
213 fn output_names(&self) -> Vec<String> {
214 match self {
215 Self::SortedTuple { name, length, .. } => (0..*length)
216 .map(|index| format!("{name}#{index}"))
217 .collect(),
218 _ => vec![self.name().to_string()],
219 }
220 }
221 fn validate(&self) -> HpoResult<()> {
222 if self.name().trim().is_empty() {
223 return Err(HpoError::InvalidSearchSpace {
224 reason: "parameter name must not be empty".to_string(),
225 });
226 }
227 match self {
228 Self::Int {
229 low, high, step, ..
230 } if low > high || *step <= 0 => Err(HpoError::InvalidSearchSpace {
231 reason: format!(
232 "integer parameter `{}` has invalid bounds or step",
233 self.name()
234 ),
235 }),
236 Self::Float {
237 low, high, step, ..
238 } if !low.is_finite()
239 || !high.is_finite()
240 || !step.is_finite()
241 || low > high
242 || *step < 0.0 =>
243 {
244 Err(HpoError::InvalidSearchSpace {
245 reason: format!(
246 "float parameter `{}` has invalid bounds or step",
247 self.name()
248 ),
249 })
250 }
251 Self::Categorical { values, .. } if values.is_empty() => {
252 Err(HpoError::InvalidSearchSpace {
253 reason: format!("categorical parameter `{}` has no values", self.name()),
254 })
255 }
256 Self::Ordinal { values, .. }
257 if values.is_empty() || values.iter().any(|value| !value.is_finite()) =>
258 {
259 Err(HpoError::InvalidSearchSpace {
260 reason: format!("ordinal parameter `{}` is invalid", self.name()),
261 })
262 }
263 Self::SortedTuple {
264 length, low, high, ..
265 } if *length <= 0 || !low.is_finite() || !high.is_finite() || low > high => {
266 Err(HpoError::InvalidSearchSpace {
267 reason: format!("sorted tuple parameter `{}` is invalid", self.name()),
268 })
269 }
270 _ => Ok(()),
271 }
272 }
273}
274
275#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
276#[serde(rename_all = "snake_case", untagged)]
277pub enum HpoCategory {
278 String(String),
279 Integer(i64),
280 Float(f64),
281 Boolean(bool),
282}
283
284#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
285#[serde(deny_unknown_fields)]
286pub struct HpoSearchSpace {
287 pub parameters: Vec<HpoParameter>,
288}
289
290impl HpoSearchSpace {
291 pub fn validate(&self) -> HpoResult<()> {
292 if self.parameters.is_empty() {
293 return Err(HpoError::InvalidSearchSpace {
294 reason: "search space has no parameters".to_string(),
295 });
296 }
297 let mut names = BTreeSet::new();
298 for parameter in &self.parameters {
299 parameter.validate()?;
300 for name in parameter.output_names() {
301 if !names.insert(name.clone()) {
302 return Err(HpoError::InvalidSearchSpace {
303 reason: format!("duplicate emitted parameter `{name}`"),
304 });
305 }
306 }
307 }
308 Ok(())
309 }
310 pub fn fingerprint(&self) -> HpoResult<String> {
313 self.validate()?;
314 let json = serde_json::to_string(self).map_err(|error| HpoError::InvalidSearchSpace {
315 reason: error.to_string(),
316 })?;
317 parse_typed_json(&json)
318 .and_then(|value| value.fingerprint())
319 .map_err(|error| HpoError::InvalidSearchSpace {
320 reason: format!("cannot canonically fingerprint search space: {error}"),
321 })
322 }
323}
324
325#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
329#[serde(deny_unknown_fields)]
330pub struct MethodsHpoStudyConfig {
331 pub controller_id: String,
332 pub study_id: String,
333 pub methods_abi: String,
337 pub search_space: HpoSearchSpace,
338 pub optimizer: HpoOptimizerConfig,
339}
340
341impl MethodsHpoStudyConfig {
342 #[cfg(feature = "methods-optimizer")]
343 fn methods_abi_identity(&self) -> HpoResult<String> {
344 if self.methods_abi.trim().is_empty() {
345 return Err(HpoError::InvalidManifest {
346 reason: "Methods ABI identity must be supplied by the native controller"
347 .to_string(),
348 });
349 }
350 Ok(self.methods_abi.clone())
351 }
352}
353
354#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
355#[serde(deny_unknown_fields)]
356pub struct HpoOptimizerConfig {
357 pub sampler: HpoSampler,
358 pub pruner: HpoPruner,
359 pub direction: HpoDirection,
360 pub metric: HpoMetric,
361 pub seed: u64,
362 pub n_startup_trials: i32,
363 pub max_resource: i32,
364 pub reduction_factor: i32,
365}
366
367impl HpoOptimizerConfig {
368 #[cfg(feature = "methods-optimizer")]
369 fn fingerprint(&self) -> HpoResult<String> {
370 let json = serde_json::to_string(self).map_err(|error| HpoError::InvalidManifest {
371 reason: error.to_string(),
372 })?;
373 parse_typed_json(&json)
374 .and_then(|value| value.fingerprint())
375 .map_err(|error| HpoError::InvalidManifest {
376 reason: format!("cannot canonically fingerprint optimizer configuration: {error}"),
377 })
378 }
379}
380
381#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)]
382#[serde(rename_all = "snake_case")]
383pub enum HpoSampler {
384 Random,
385 Sobol,
386 Lhs,
387 Ternary,
388 Ga,
389 Pso,
390 Cmaes,
391 Tpe,
392 GpEi,
393}
394#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)]
395#[serde(rename_all = "snake_case")]
396pub enum HpoPruner {
397 None,
398 Median,
399 Asha,
400 Hyperband,
401 Racing,
402}
403#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)]
404#[serde(rename_all = "snake_case")]
405pub enum HpoDirection {
406 Auto,
407 Minimize,
408 Maximize,
409}
410#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)]
411#[serde(rename_all = "snake_case")]
412pub enum HpoMetric {
413 Rmse,
414 Mse,
415 Mae,
416 R2,
417 Accuracy,
418 BalancedAccuracy,
419 F1,
420 Logloss,
421}
422
423#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
424#[serde(deny_unknown_fields)]
425pub struct HpoTrialParameter {
426 pub name: String,
427 pub value: f64,
428 #[serde(default)]
431 pub native_kind: Option<HpoNativeParameterKind>,
432 #[serde(default)]
433 pub category_type: Option<HpoCategoryType>,
434 #[serde(default)]
435 pub integer: bool,
436 pub active: bool,
437 pub category_index: Option<i32>,
438 pub category_label: Option<String>,
439}
440
441#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)]
442#[serde(rename_all = "snake_case")]
443pub enum HpoNativeParameterKind {
444 Int,
445 Float,
446 LogInt,
447 LogFloat,
448 Categorical,
449 Ordinal,
450 SortedTuple,
451}
452
453#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)]
454#[serde(rename_all = "snake_case")]
455pub enum HpoCategoryType {
456 String,
457 Integer,
458 Float,
459 Boolean,
460}
461
462#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)]
463#[serde(rename_all = "snake_case")]
464pub enum HpoTrialStatus {
465 Running,
466 Completed,
467 Pruned,
468 Failed,
469 Cancelled,
470}
471
472#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
473#[serde(deny_unknown_fields)]
474pub struct HpoFailure {
475 pub code: String,
476 pub message: String,
477 pub retryable: bool,
478}
479
480#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
481#[serde(deny_unknown_fields)]
482pub struct HpoIntermediate {
483 pub sequence: i64,
484 pub step: i32,
485 pub score: f64,
486 pub should_prune: bool,
487}
488
489#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
490#[serde(deny_unknown_fields)]
491pub struct HpoTrial {
492 pub id: i64,
493 pub ask_sequence: i64,
494 pub terminal_sequence: Option<i64>,
495 pub parameters: BTreeMap<String, HpoTrialParameter>,
496 pub parameter_order: Vec<String>,
497 pub status: HpoTrialStatus,
498 pub score: Option<f64>,
499 pub rung: i32,
500 pub duration: f64,
501 pub intermediates: Vec<HpoIntermediate>,
502 pub failure: Option<HpoFailure>,
503}
504
505#[cfg(any(test, feature = "methods-optimizer"))]
512fn canonical_hpo_terminal_ledger(trials: Vec<HpoTrial>) -> crate::Result<Vec<HpoTrial>> {
513 let json = serde_json::to_string(&trials)?;
514 parse_typed_json(&json).map_err(|error| {
515 crate::DagMlError::RuntimeValidation(format!(
516 "native Methods HPO terminal ledger has no strict TCV1 JSON preimage: {error}"
517 ))
518 })?;
519 Ok(serde_json::from_str(&json)?)
520}
521
522#[cfg(any(test, feature = "methods-optimizer"))]
523fn scores_within_one_ulp(left: f64, right: f64) -> bool {
524 if left == right {
525 return true;
526 }
527 if !left.is_finite() || !right.is_finite() {
528 return false;
529 }
530 let ordered = |value: f64| {
531 let bits = value.to_bits();
532 if bits & (1_u64 << 63) != 0 {
533 (!bits) as i128
534 } else {
535 (bits | (1_u64 << 63)) as i128
536 }
537 };
538 (ordered(left) - ordered(right)).abs() <= 1
539}
540
541#[cfg(any(test, feature = "methods-optimizer"))]
542fn optional_scores_within_one_ulp(left: Option<f64>, right: Option<f64>) -> bool {
543 match (left, right) {
544 (Some(left), Some(right)) => scores_within_one_ulp(left, right),
545 (None, None) => true,
546 _ => false,
547 }
548}
549
550#[cfg(any(test, feature = "methods-optimizer"))]
551fn hpo_terminal_trials_match(native: &[HpoTrial], persisted: &[HpoTrial]) -> bool {
552 native.len() == persisted.len()
553 && native.iter().zip(persisted).all(|(native, persisted)| {
554 native.id == persisted.id
555 && native.ask_sequence == persisted.ask_sequence
556 && native.terminal_sequence == persisted.terminal_sequence
557 && native.parameters == persisted.parameters
558 && native.parameter_order == persisted.parameter_order
559 && native.status == persisted.status
560 && optional_scores_within_one_ulp(native.score, persisted.score)
561 && native.rung == persisted.rung
562 && native.duration == persisted.duration
563 && native.failure == persisted.failure
564 && native.intermediates.len() == persisted.intermediates.len()
565 && native
566 .intermediates
567 .iter()
568 .zip(&persisted.intermediates)
569 .all(|(native, persisted)| {
570 native.sequence == persisted.sequence
571 && native.step == persisted.step
572 && native.should_prune == persisted.should_prune
573 && scores_within_one_ulp(native.score, persisted.score)
574 })
575 })
576}
577
578#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
581#[serde(deny_unknown_fields)]
582pub struct HpoBestTrial {
583 pub trial: HpoTrial,
584 pub score: f64,
585}
586
587#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
588#[serde(tag = "event", rename_all = "snake_case", deny_unknown_fields)]
589pub enum HpoEvent {
590 Asked {
591 trial_id: i64,
592 },
593 Intermediate {
594 trial_id: i64,
595 step: i32,
596 score: f64,
597 should_prune: bool,
598 },
599 Terminal {
600 trial_id: i64,
601 status: HpoTrialStatus,
602 score: Option<f64>,
603 failure: Option<HpoFailure>,
604 },
605}
606
607pub struct HpoEvaluationBoundary<'a> {
611 pub folds: &'a FoldSet,
612 pub influence: &'a TrainingInfluenceManifest,
613 pub lineage: &'a mut InMemoryLineageRecorder,
614 pub scores: &'a mut ScoreSet,
615 pub selection: &'a SelectionPolicy,
616 pub refit_strategy: Option<RefitStrategy>,
617}
618
619impl HpoEvaluationBoundary<'_> {
620 pub fn validate(&self) -> HpoResult<()> {
621 self.folds
622 .validate()
623 .map_err(|error| HpoError::Evaluation {
624 reason: error.to_string(),
625 })?;
626 self.scores
627 .validate()
628 .map_err(|error| HpoError::Evaluation {
629 reason: error.to_string(),
630 })?;
631 self.selection
632 .validate()
633 .map_err(|error| HpoError::Evaluation {
634 reason: error.to_string(),
635 })?;
636 self.influence
637 .validate()
638 .map_err(|error| HpoError::Evaluation {
639 reason: error.to_string(),
640 })?;
641 if !self
642 .influence
643 .entries
644 .iter()
645 .any(|entry| entry.kind == TrainingInfluenceKind::HpoSelection)
646 {
647 return Err(HpoError::Evaluation {
648 reason: "training influence manifest has no hpo_selection entry".to_string(),
649 });
650 }
651 Ok(())
652 }
653}
654
655pub trait HpoEvaluator {
656 fn evaluate(
657 &mut self,
658 trial: &HpoTrial,
659 boundary: &mut HpoEvaluationBoundary<'_>,
660 ) -> HpoResult<HpoTerminal>;
661
662 fn evaluate_with_reporter(
666 &mut self,
667 trial: &HpoTrial,
668 boundary: &mut HpoEvaluationBoundary<'_>,
669 _reporter: &mut dyn HpoIntermediateReporter,
670 ) -> HpoResult<HpoTerminal> {
671 self.evaluate(trial, boundary)
672 }
673}
674
675pub trait HpoIntermediateReporter {
676 fn report(&mut self, step: i32, score: f64) -> HpoResult<HpoReportOutcome>;
677}
678
679#[derive(Clone, Debug, PartialEq)]
680pub enum HpoReportOutcome {
681 Continue,
682 Pruned(HpoTrial),
683}
684
685#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
686#[serde(tag = "state", rename_all = "snake_case")]
687pub enum HpoTerminal {
688 Completed { score: f64 },
689 Failed { failure: HpoFailure },
690 Pruned { failure: HpoFailure },
691 Cancelled { failure: HpoFailure },
692}
693
694#[derive(Clone, Debug, PartialEq)]
695pub struct HpoBatch {
696 pub trials: Vec<HpoTrial>,
697 pub native_error: Option<HpoNativeError>,
698}
699
700#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
701#[serde(deny_unknown_fields)]
702pub struct N4moptCheckpointArtifact {
703 pub schema_version: u32,
704 pub artifact_kind: String,
705 pub format: String,
706 pub binding: HpoStudyBinding,
707 pub methods_abi: String,
708 pub opaque_payload: Vec<u8>,
709 pub payload_sha256: String,
710}
711
712impl N4moptCheckpointArtifact {
713 #[cfg(feature = "methods-optimizer")]
714 fn new(
715 binding: HpoStudyBinding,
716 methods_abi: String,
717 opaque_payload: Vec<u8>,
718 ) -> HpoResult<Self> {
719 let value = Self {
720 schema_version: N4MOPT_CHECKPOINT_SCHEMA_VERSION,
721 artifact_kind: N4MOPT_ARTIFACT_KIND.to_string(),
722 format: N4MOPT_FORMAT.to_string(),
723 binding,
724 methods_abi,
725 payload_sha256: payload_sha256(&opaque_payload),
726 opaque_payload,
727 };
728 value.validate()?;
729 Ok(value)
730 }
731 pub fn validate(&self) -> HpoResult<()> {
732 if self.schema_version != N4MOPT_CHECKPOINT_SCHEMA_VERSION
733 || self.artifact_kind != N4MOPT_ARTIFACT_KIND
734 || self.format != N4MOPT_FORMAT
735 {
736 return Err(HpoError::InvalidCheckpoint {
737 reason: "checkpoint schema, kind, or format is invalid".to_string(),
738 });
739 }
740 self.binding.validate()?;
741 if self.methods_abi.trim().is_empty()
742 || self.opaque_payload.is_empty()
743 || self.opaque_payload.len() > MAX_N4MOPT_CHECKPOINT_BYTES
744 {
745 return Err(HpoError::InvalidCheckpoint {
746 reason: "checkpoint ABI/payload is invalid or exceeds the maximum size".to_string(),
747 });
748 }
749 if self.payload_sha256 != payload_sha256(&self.opaque_payload) {
750 return Err(HpoError::InvalidCheckpoint {
751 reason: "checkpoint payload SHA-256 differs from envelope".to_string(),
752 });
753 }
754 Ok(())
755 }
756}
757
758#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
764#[serde(deny_unknown_fields)]
765pub struct N4moptCheckpointReference {
766 pub artifact: crate::runtime::ArtifactRef,
767 pub binding: HpoStudyBinding,
768 pub methods_abi: String,
769}
770
771impl N4moptCheckpointReference {
772 pub fn validate(&self) -> HpoResult<()> {
773 self.binding.validate()?;
774 if self.methods_abi.trim().is_empty() {
775 return Err(HpoError::InvalidCheckpoint {
776 reason: "checkpoint reference has no Methods ABI identity".to_string(),
777 });
778 }
779 self.artifact
780 .validate_portable()
781 .map_err(|error| HpoError::InvalidCheckpoint {
782 reason: format!("checkpoint archive reference is invalid: {error}"),
783 })?;
784 if self.artifact.kind != N4MOPT_ARTIFACT_KIND
785 || self.artifact.controller_id.as_str() != self.binding.controller_id
786 || self.artifact.backend != Some(crate::runtime::ArtifactBackend::Raw)
787 {
788 return Err(HpoError::InvalidCheckpoint {
789 reason: "checkpoint reference must be a raw Methods-owned N4MOPT artifact"
790 .to_string(),
791 });
792 }
793 Ok(())
794 }
795}
796
797fn payload_sha256(payload: &[u8]) -> String {
798 format!("{:x}", Sha256::digest(payload))
799}
800
801pub fn methods_optimizer_preflight() -> HpoResult<()> {
804 #[cfg(feature = "methods-optimizer")]
805 {
806 Ok(())
807 }
808 #[cfg(not(feature = "methods-optimizer"))]
809 {
810 Err(HpoError::MethodsOptimizerFeatureDisabled)
811 }
812}
813
814pub const METHODS_PLS_CONTROLLER_ID: &str = "controller:methods.pls";
819
820pub const METHODS_RIDGE_CONTROLLER_ID: &str = "controller:methods.ridge";
826
827#[cfg(feature = "methods-optimizer")]
833#[derive(Clone, Debug, Eq, PartialEq)]
834pub struct MethodsRuntime {
835 library_path: std::path::PathBuf,
836}
837
838#[cfg(feature = "methods-optimizer")]
839impl MethodsRuntime {
840 pub fn configure(library_path: impl AsRef<std::path::Path>) -> HpoResult<Self> {
841 let library_path = library_path.as_ref();
842 if !library_path.is_absolute() {
843 return Err(HpoError::RuntimeConfiguration {
844 reason: "libn4m path must be absolute".to_string(),
845 });
846 }
847 let canonical = std::fs::canonicalize(library_path).map_err(|error| {
848 HpoError::RuntimeConfiguration {
849 reason: format!(
850 "cannot resolve libn4m path `{}`: {error}",
851 library_path.display()
852 ),
853 }
854 })?;
855 let metadata =
856 std::fs::metadata(&canonical).map_err(|error| HpoError::RuntimeConfiguration {
857 reason: format!(
858 "cannot inspect libn4m path `{}`: {error}",
859 canonical.display()
860 ),
861 })?;
862 if !metadata.is_file() {
863 return Err(HpoError::RuntimeConfiguration {
864 reason: format!(
865 "libn4m path `{}` is not a regular file",
866 canonical.display()
867 ),
868 });
869 }
870 n4m::configure_library(&canonical).map_err(|error| HpoError::RuntimeConfiguration {
871 reason: format!("cannot load libn4m `{}`: {error}", canonical.display()),
872 })?;
873 Ok(Self {
874 library_path: canonical,
875 })
876 }
877
878 pub fn library_path(&self) -> &std::path::Path {
879 &self.library_path
880 }
881}
882
883#[cfg(feature = "methods-optimizer")]
887pub struct MethodsHpoController {
888 id: crate::ControllerId,
889 _runtime: MethodsRuntime,
890}
891
892#[cfg(feature = "methods-optimizer")]
893impl MethodsHpoController {
894 pub fn new(id: crate::ControllerId, runtime: MethodsRuntime) -> Self {
895 Self {
896 id,
897 _runtime: runtime,
898 }
899 }
900}
901
902#[cfg(feature = "methods-optimizer")]
903impl crate::runtime::RuntimeController for MethodsHpoController {
904 fn controller_id(&self) -> &crate::ControllerId {
905 &self.id
906 }
907
908 fn invoke(&self, task: &crate::runtime::NodeTask) -> crate::Result<crate::runtime::NodeResult> {
909 Err(crate::DagMlError::RuntimeValidation(format!(
910 "Methods HPO controller `{}` is training-owned and cannot execute graph task `{}` directly",
911 self.id, task.node_plan.node_id
912 )))
913 }
914
915 fn create_tuner_session(
916 &self,
917 task: &crate::runtime::RuntimeHpoCampaignTask,
918 context: &crate::runtime::RuntimeHpoExecutionContext,
919 ) -> crate::Result<Box<dyn crate::runtime::RuntimeTunerSession>> {
920 if task.operation_id != context.operation_id
921 || task.controller_id != context.controller_id
922 || task.target_node_id != context.target_node_id
923 || context.study.controller_id != self.id.as_str()
924 {
925 return Err(crate::DagMlError::RuntimeValidation(
926 "Methods HPO tuner task/context identity mismatch".to_string(),
927 ));
928 }
929 let study = if let Some(checkpoint) = &context.resume_checkpoint {
930 MethodsHpoStudy::restore(context.study.clone(), checkpoint)
931 } else {
932 MethodsHpoStudy::create(context.study.clone())
933 }
934 .map_err(|error| {
935 crate::DagMlError::RuntimeValidation(format!(
936 "cannot create controller-owned native Methods HPO study: {error}"
937 ))
938 })?;
939 if context.resume_checkpoint.is_some() {
940 let mut native = study.trials().map_err(|error| {
941 crate::DagMlError::RuntimeValidation(format!(
942 "cannot attest restored native Methods HPO ledger: {error}"
943 ))
944 })?;
945 native.sort_by_key(|trial| trial.id);
946 let native = canonical_hpo_terminal_ledger(native)?;
947 let persisted = canonical_hpo_terminal_ledger(
948 context
949 .resume_terminal_trials
950 .iter()
951 .map(|snapshot| snapshot.trial.clone())
952 .collect(),
953 )?;
954 if !hpo_terminal_trials_match(&native, &persisted) {
955 return Err(crate::DagMlError::RuntimeValidation(
956 "restored native Methods HPO ledger does not exactly match persisted terminal evidence"
957 .to_string(),
958 ));
959 }
960 }
961 Ok(Box::new(MethodsHpoSession {
962 study,
963 context: context.clone(),
964 controller_id: self.id.clone(),
965 }))
966 }
967}
968
969#[cfg(feature = "methods-optimizer")]
970struct MethodsHpoSession {
971 study: MethodsHpoStudy,
972 context: crate::runtime::RuntimeHpoExecutionContext,
973 controller_id: crate::ControllerId,
974}
975
976#[cfg(feature = "methods-optimizer")]
977impl crate::runtime::RuntimeTunerSession for MethodsHpoSession {
978 fn trial_history_len(&self) -> crate::Result<u32> {
979 let count = self
980 .study
981 .trials()
982 .map_err(|error| crate::DagMlError::RuntimeValidation(error.to_string()))?
983 .len();
984 u32::try_from(count).map_err(|_| {
985 crate::DagMlError::RuntimeValidation(
986 "native Methods HPO trial history exceeds u32 budget".to_string(),
987 )
988 })
989 }
990
991 fn ask(&mut self) -> crate::Result<Option<crate::runtime::RuntimeHpoProposal>> {
992 let trial = self
993 .study
994 .ask()
995 .map_err(|error| crate::DagMlError::RuntimeValidation(error.to_string()))?;
996 if self.context.study.search_space.parameters.len() != 1
997 || self.context.parameter_paths.len() != 1
998 || self.context.parameter_paths.get("n_components") != Some(&"n_components".to_string())
999 || !matches!(self.context.study.search_space.parameters.first(), Some(HpoParameter::Int { name, low: 1, high: 3, step: 1, log: false }) if name == "n_components")
1000 {
1001 return Err(crate::DagMlError::RuntimeValidation(
1002 "Methods HPO v1 accepts only active integer n_components=1..3 mapped directly to the target model".to_string(),
1003 ));
1004 }
1005 let parameter = trial.parameters.get("n_components").ok_or_else(|| {
1006 crate::DagMlError::RuntimeValidation(
1007 "native Methods HPO trial omitted active n_components".to_string(),
1008 )
1009 })?;
1010 if !parameter.active
1011 || !parameter.integer
1012 || parameter.value.fract() != 0.0
1013 || !(1.0..=3.0).contains(¶meter.value)
1014 {
1015 return Err(crate::DagMlError::RuntimeValidation(
1016 "native Methods HPO emitted invalid n_components outside V1 integer bounds"
1017 .to_string(),
1018 ));
1019 }
1020 let mut variant = self.context.base_variant.clone();
1021 variant.choices.insert(
1022 "native_methods_hpo".to_string(),
1023 crate::generation::GenerationChoice {
1024 label: format!("trial:{}", trial.id),
1025 value: serde_json::json!({"trial_id": trial.id}),
1026 param_overrides: vec![crate::generation::GenerationParamOverride {
1027 node_id: self.context.target_node_id.clone(),
1028 params: BTreeMap::from([(
1029 "n_components".to_string(),
1030 serde_json::json!(parameter.value as i64),
1031 )]),
1032 }],
1033 active_subsequence: None,
1034 },
1035 );
1036 variant.variant_id = crate::VariantId::new(format!("hpo:trial:{}", trial.id))
1037 .map_err(|error| crate::DagMlError::RuntimeValidation(error.to_string()))?;
1038 variant.fingerprint = crate::campaign::stable_json_fingerprint(&(
1039 self.context.base_variant.fingerprint.as_str(),
1040 &variant.choices,
1041 trial.id,
1042 ))?;
1043 Ok(Some(crate::runtime::RuntimeHpoProposal {
1044 trial_id: trial.id,
1045 variant,
1046 }))
1047 }
1048
1049 fn report_intermediate(
1050 &mut self,
1051 value: crate::runtime::RuntimeHpoIntermediate,
1052 ) -> crate::Result<crate::runtime::RuntimeHpoIntermediateOutcome> {
1053 let pruned = self
1054 .study
1055 .report_intermediate(value.trial_id, value.step, value.score)
1056 .map_err(|error| crate::DagMlError::RuntimeValidation(error.to_string()))?;
1057 Ok(if pruned {
1058 crate::runtime::RuntimeHpoIntermediateOutcome::Pruned
1059 } else {
1060 crate::runtime::RuntimeHpoIntermediateOutcome::Continue
1061 })
1062 }
1063
1064 fn tell(
1065 &mut self,
1066 trial_id: i64,
1067 terminal: crate::runtime::RuntimeHpoTerminal,
1068 ) -> crate::Result<()> {
1069 let terminal = match terminal {
1070 crate::runtime::RuntimeHpoTerminal::Completed { score } => {
1071 HpoTerminal::Completed { score }
1072 }
1073 crate::runtime::RuntimeHpoTerminal::Failed { failure } => HpoTerminal::Failed {
1074 failure: HpoFailure {
1075 code: failure.code,
1076 message: failure.message,
1077 retryable: failure.retryable,
1078 },
1079 },
1080 };
1081 self.study
1082 .tell(trial_id, terminal)
1083 .map_err(|error| crate::DagMlError::RuntimeValidation(error.to_string()))?;
1084 Ok(())
1085 }
1086
1087 fn checkpoint(&self) -> crate::Result<N4moptCheckpointArtifact> {
1088 let checkpoint = self.study.save_checkpoint().map_err(|error| {
1089 crate::DagMlError::RuntimeValidation(format!(
1090 "cannot save native Methods HPO checkpoint: {error}"
1091 ))
1092 })?;
1093 checkpoint.validate().map_err(|error| {
1094 crate::DagMlError::RuntimeValidation(format!(
1095 "invalid native Methods HPO checkpoint: {error}"
1096 ))
1097 })?;
1098 if checkpoint.binding.controller_id != self.controller_id.as_str()
1099 || checkpoint.binding.controller_id != self.context.study.controller_id
1100 || checkpoint.binding.study_id != self.context.study.study_id
1101 || checkpoint.methods_abi != self.context.study.methods_abi
1102 {
1103 return Err(crate::DagMlError::RuntimeValidation(
1104 "native Methods HPO checkpoint binding/ABI does not match its scheduler context"
1105 .to_string(),
1106 ));
1107 }
1108 Ok(checkpoint)
1109 }
1110
1111 fn incumbent(
1112 &self,
1113 variants: &BTreeMap<i64, crate::VariantId>,
1114 ) -> crate::Result<Option<crate::runtime::RuntimeHpoIncumbent>> {
1115 let Some(best) = self.study.best().map_err(|error| {
1116 crate::DagMlError::RuntimeValidation(format!(
1117 "cannot read native Methods HPO incumbent: {error}"
1118 ))
1119 })?
1120 else {
1121 return Ok(None);
1122 };
1123 let score = if let Some(persisted) = self
1124 .context
1125 .resume_terminal_trials
1126 .iter()
1127 .find(|snapshot| snapshot.trial.id == best.trial.id)
1128 {
1129 let native = canonical_hpo_terminal_ledger(vec![best.trial.clone()])?;
1130 let prior = canonical_hpo_terminal_ledger(vec![persisted.trial.clone()])?;
1131 if !hpo_terminal_trials_match(&native, &prior) {
1132 return Err(crate::DagMlError::RuntimeValidation(
1133 "native Methods HPO incumbent does not match persisted terminal evidence"
1134 .to_string(),
1135 ));
1136 }
1137 persisted.trial.score.ok_or_else(|| {
1138 crate::DagMlError::RuntimeValidation(
1139 "persisted Methods HPO incumbent has no terminal score".to_string(),
1140 )
1141 })?
1142 } else {
1143 best.score
1144 };
1145 let variant_id = variants.get(&best.trial.id).cloned().ok_or_else(|| {
1146 crate::DagMlError::RuntimeValidation(
1147 "native Methods HPO best() returned a trial without scheduler variant identity"
1148 .to_string(),
1149 )
1150 })?;
1151 Ok(Some(crate::runtime::RuntimeHpoIncumbent {
1152 trial_id: best.trial.id,
1153 score,
1154 metric: self.context.selection.metric.name().to_string(),
1155 direction: self.context.selection.direction,
1156 variant_id,
1157 }))
1158 }
1159
1160 fn terminal_trial_snapshots(
1161 &self,
1162 variants: &BTreeMap<i64, crate::VariantId>,
1163 ) -> crate::Result<Vec<crate::runtime::RuntimeHpoTerminalSnapshot>> {
1164 let mut trials = self.study.trials().map_err(|error| {
1165 crate::DagMlError::RuntimeValidation(format!(
1166 "cannot read native Methods HPO terminal ledger: {error}"
1167 ))
1168 })?;
1169 trials.sort_by_key(|trial| trial.id);
1170 if trials.iter().any(|trial| {
1171 !matches!(
1172 trial.status,
1173 HpoTrialStatus::Completed | HpoTrialStatus::Pruned | HpoTrialStatus::Failed
1174 )
1175 }) {
1176 return Err(crate::DagMlError::RuntimeValidation(
1177 "native Methods HPO trial ledger contains a non-terminal trial".to_string(),
1178 ));
1179 }
1180 let persisted_by_id = self
1181 .context
1182 .resume_terminal_trials
1183 .iter()
1184 .map(|snapshot| (snapshot.trial.id, snapshot))
1185 .collect::<BTreeMap<_, _>>();
1186 trials
1187 .into_iter()
1188 .map(|trial| {
1189 if let Some(persisted) = persisted_by_id.get(&trial.id) {
1190 let native = canonical_hpo_terminal_ledger(vec![trial])?;
1191 let prior = canonical_hpo_terminal_ledger(vec![persisted.trial.clone()])?;
1192 if !hpo_terminal_trials_match(&native, &prior) {
1193 return Err(crate::DagMlError::RuntimeValidation(
1194 "restored native Methods HPO trial does not match persisted terminal evidence"
1195 .to_string(),
1196 ));
1197 }
1198 return Ok((*persisted).clone());
1199 }
1200 Ok(crate::runtime::RuntimeHpoTerminalSnapshot {
1201 variant_id: variants.get(&trial.id).cloned(),
1202 trial,
1203 })
1204 })
1205 .collect()
1206 }
1207}
1208
1209#[cfg(feature = "methods-optimizer")]
1210mod pls_controller {
1211 use std::collections::{BTreeMap, BTreeSet};
1212 use std::sync::atomic::{AtomicU64, Ordering};
1213 use std::sync::Mutex;
1214
1215 use super::*;
1216 use crate::runtime::{
1217 ArtifactBackend, ArtifactRef, HandleKind, HandleRef, LineageRecord, MethodsPlsData,
1218 MethodsPlsDataRequest, NodeResult, NodeTask, PredictionBlock, PredictionInputSpec,
1219 PredictionPartition, RegressionTargetBlock, RuntimeController, RuntimeDataProvider,
1220 };
1221 use crate::{
1222 ArtifactId, ControllerId, DagMlError, LineageId, Phase, PredictionLevel, PredictionUnitId,
1223 Result,
1224 };
1225 use n4m::{Config, Context, MatrixRef, Model};
1226
1227 pub struct MethodsPlsController {
1232 id: ControllerId,
1233 _runtime: MethodsRuntime,
1234 next_handle: AtomicU64,
1235 exported_n4mm_by_artifact: Mutex<BTreeMap<ArtifactId, Vec<u8>>>,
1238 hydrated_n4mm_by_handle: Mutex<BTreeMap<u64, Vec<u8>>>,
1242 }
1243
1244 impl MethodsPlsController {
1245 pub fn new(runtime: MethodsRuntime) -> Self {
1246 Self {
1247 id: ControllerId::new(METHODS_PLS_CONTROLLER_ID)
1248 .expect("Methods PLS controller id is valid"),
1249 _runtime: runtime,
1250 next_handle: AtomicU64::new(0),
1251 exported_n4mm_by_artifact: Mutex::new(BTreeMap::new()),
1252 hydrated_n4mm_by_handle: Mutex::new(BTreeMap::new()),
1253 }
1254 }
1255
1256 #[doc(hidden)]
1258 pub fn hydrated_payload_count(&self) -> Result<usize> {
1259 self.hydrated_n4mm_by_handle
1260 .lock()
1261 .map(|payloads| payloads.len())
1262 .map_err(|_| {
1263 DagMlError::RuntimeValidation(
1264 "portable Methods PLS hydrated N4MM lock poisoned".to_string(),
1265 )
1266 })
1267 }
1268
1269 fn handle(&self, kind: HandleKind) -> HandleRef {
1270 HandleRef {
1271 handle: self.next_handle.fetch_add(1, Ordering::SeqCst) + 1,
1272 kind,
1273 owner_controller: self.id.clone(),
1274 }
1275 }
1276
1277 fn request(
1278 task: &NodeTask,
1279 provider: &dyn RuntimeDataProvider,
1280 data_port: &str,
1281 ) -> Result<MethodsPlsDataRequest> {
1282 let bindings = task
1283 .node_plan
1284 .data_bindings
1285 .iter()
1286 .filter(|binding| binding.input_name == data_port)
1287 .collect::<Vec<_>>();
1288 let [binding] = bindings.as_slice() else {
1289 return Err(DagMlError::RuntimeValidation(format!(
1290 "portable Methods node `{}` requires exactly one `{data_port}` DataBinding",
1291 task.node_plan.node_id,
1292 )));
1293 };
1294 let identity = provider.training_data_identity(binding)?;
1295 if task.phase != Phase::Predict && identity.is_none() {
1296 return Err(DagMlError::RuntimeValidation(format!(
1297 "portable Methods PLS provider did not attest target-bound DataBinding `{}.{}` for {:?}",
1298 binding.node_id, binding.input_name, task.phase
1299 )));
1300 }
1301 let data_view_key = format!("data:{data_port}");
1302 let fit_view = task.data_views.get(data_port).or_else(|| task.data_views.get(&data_view_key)).cloned().ok_or_else(|| {
1303 DagMlError::RuntimeValidation(format!(
1304 "portable Methods node `{}` requires its scheduler-created `{data_port}` data view (available: {:?})",
1305 task.node_plan.node_id,
1306 task.data_views.keys().collect::<Vec<_>>(),
1307 ))
1308 })?;
1309 let prediction_view = if task.phase == Phase::FitCv {
1310 let validation_view_key = format!("data:{data_port}:validation");
1311 let validation_key = format!("{data_port}:validation");
1312 Some(task.data_views.get(&validation_key).or_else(|| task.data_views.get(&validation_view_key)).cloned().ok_or_else(|| {
1313 DagMlError::RuntimeValidation(format!(
1314 "portable Methods node `{}` requires its scheduler-created `{data_port}` validation view",
1315 task.node_plan.node_id,
1316 ))
1317 })?)
1318 } else {
1319 None
1320 };
1321 let request = MethodsPlsDataRequest {
1322 node_id: task.node_plan.node_id.clone(),
1323 phase: task.phase,
1324 variant_id: task.variant_id.clone(),
1325 fold_id: task.fold_id.clone(),
1326 binding: (*binding).clone(),
1327 identity,
1328 fit_view,
1329 prediction_view,
1330 };
1331 request.validate()?;
1332 Ok(request)
1333 }
1334
1335 fn components(task: &NodeTask) -> Result<i32> {
1336 let value = task.node_plan.params.get("n_components").ok_or_else(|| {
1337 DagMlError::RuntimeValidation(format!(
1338 "portable Methods PLS node `{}` requires integer `n_components`",
1339 task.node_plan.node_id
1340 ))
1341 })?;
1342 let value = value.as_i64().ok_or_else(|| {
1343 DagMlError::RuntimeValidation(
1344 "portable Methods PLS `n_components` must be an integer".to_string(),
1345 )
1346 })?;
1347 i32::try_from(value)
1348 .ok()
1349 .filter(|value| *value > 0)
1350 .ok_or_else(|| {
1351 DagMlError::RuntimeValidation(
1352 "portable Methods PLS `n_components` must be a positive i32".to_string(),
1353 )
1354 })
1355 }
1356
1357 fn native_error(operation: &str, error: n4m::Error) -> DagMlError {
1358 DagMlError::RuntimeValidation(format!(
1359 "portable Methods PLS {operation} failed: {error}"
1360 ))
1361 }
1362
1363 fn fit(task: &NodeTask, data: &MethodsPlsData) -> Result<(Context, Model)> {
1364 let context =
1365 Context::new().map_err(|error| Self::native_error("context_create", error))?;
1366 let mut config =
1367 Config::new().map_err(|error| Self::native_error("config_create", error))?;
1368 config
1369 .set_n_components(Self::components(task)?)
1370 .map_err(|error| Self::native_error("config_set_n_components", error))?;
1371 let x = MatrixRef::row_major(&data.fit.x.values, data.fit.x.rows, data.fit.x.cols)
1372 .map_err(|error| Self::native_error("fit_x_matrix", error))?;
1373 let targets = data.fit.y.as_ref().ok_or_else(|| {
1374 DagMlError::RuntimeValidation(
1375 "portable Methods PLS fit requires targets".to_string(),
1376 )
1377 })?;
1378 let y = MatrixRef::row_major(&targets.values, targets.rows, targets.cols)
1379 .map_err(|error| Self::native_error("fit_y_matrix", error))?;
1380 let model = Model::fit(&context, &config, x, y)
1381 .map_err(|error| Self::native_error("fit", error))?;
1382 Ok((context, model))
1383 }
1384
1385 fn predict(
1386 context: &Context,
1387 model: &Model,
1388 data: &crate::runtime::MethodsPlsDataset,
1389 ) -> Result<Vec<Vec<f64>>> {
1390 let x = MatrixRef::row_major(&data.x.values, data.x.rows, data.x.cols)
1391 .map_err(|error| Self::native_error("predict_x_matrix", error))?;
1392 let prediction = model
1393 .predict(context, x)
1394 .map_err(|error| Self::native_error("predict", error))?;
1395 Ok(prediction
1396 .data
1397 .chunks(prediction.cols)
1398 .map(|row| row.to_vec())
1399 .collect())
1400 }
1401
1402 fn result(
1403 &self,
1404 task: &NodeTask,
1405 dataset: &crate::runtime::MethodsPlsDataset,
1406 values: Vec<Vec<f64>>,
1407 artifact: Option<(ArtifactRef, HandleRef)>,
1408 partition: PredictionPartition,
1409 ) -> Result<NodeResult> {
1410 let prediction = PredictionBlock {
1411 prediction_id: Some(format!(
1412 "methods-pls:{}:{}:{}",
1413 task.node_plan.node_id,
1414 task.phase.as_str(),
1415 task.fold_id
1416 .as_ref()
1417 .map(|id| id.as_str())
1418 .unwrap_or("full")
1419 )),
1420 producer_node: task.node_plan.node_id.clone(),
1421 producer_port: Some("oof".to_string()),
1422 partition,
1423 fold_id: (task.phase == Phase::FitCv)
1424 .then(|| task.fold_id.clone())
1425 .flatten(),
1426 sample_ids: dataset.sample_ids.clone(),
1427 values,
1428 target_names: dataset.target_names.clone(),
1429 };
1430 let regression_targets = if task.phase == Phase::FitCv {
1431 let targets = dataset.y.as_ref().ok_or_else(|| {
1432 DagMlError::RuntimeValidation(
1433 "portable Methods PLS FIT_CV requires validation targets".to_string(),
1434 )
1435 })?;
1436 vec![RegressionTargetBlock {
1437 level: PredictionLevel::Sample,
1438 unit_ids: dataset
1439 .sample_ids
1440 .iter()
1441 .cloned()
1442 .map(PredictionUnitId::Sample)
1443 .collect(),
1444 values: targets
1445 .values
1446 .chunks(targets.cols)
1447 .map(|row| row.to_vec())
1448 .collect(),
1449 target_names: dataset.target_names.clone(),
1450 }]
1451 } else {
1452 Vec::new()
1453 };
1454 let (artifacts, artifact_handles) = artifact
1455 .map(|(artifact, handle)| {
1456 (
1457 vec![artifact.clone()],
1458 BTreeMap::from([(artifact.id, handle)]),
1459 )
1460 })
1461 .unwrap_or_default();
1462 let artifact_refs = artifacts.clone();
1463 Ok(NodeResult {
1464 schema_version: None,
1465 node_id: task.node_plan.node_id.clone(),
1466 outputs: BTreeMap::from([("oof".to_string(), self.handle(HandleKind::Prediction))]),
1467 predictions: vec![prediction],
1468 observation_predictions: Vec::new(),
1469 aggregated_predictions: Vec::new(),
1470 explanations: Vec::new(),
1471 shape_deltas: Vec::new(),
1472 artifacts,
1473 artifact_handles,
1474 fit_influence_diagnostics: Vec::new(),
1475 regression_targets,
1476 lineage: LineageRecord {
1477 record_id: LineageId::new(format!(
1478 "lineage:methods-pls:{}:{}:{}:{}",
1479 task.node_plan.node_id,
1480 task.phase.as_str(),
1481 task.variant_id
1482 .as_ref()
1483 .map(|id| id.as_str())
1484 .unwrap_or("base"),
1485 task.fold_id
1486 .as_ref()
1487 .map(|id| id.as_str())
1488 .unwrap_or("full")
1489 ))
1490 .expect("valid native PLS lineage id"),
1491 run_id: task.run_id.clone(),
1492 node_id: task.node_plan.node_id.clone(),
1493 phase: task.phase,
1494 controller_id: self.id.clone(),
1495 controller_version: task.node_plan.controller_version.clone(),
1496 variant_id: task.variant_id.clone(),
1497 fold_id: task.fold_id.clone(),
1498 branch_path: task.branch_path.clone(),
1499 input_lineage: Vec::new(),
1500 artifact_refs,
1501 params_fingerprint: task.node_plan.params_fingerprint.clone(),
1502 data_model_shape_fingerprint: None,
1503 aggregation_policy_fingerprint: None,
1504 seed: task.seed,
1505 unsafe_flags: BTreeSet::new(),
1506 metrics: BTreeMap::new(),
1507 loss_attestations: Vec::new(),
1508 early_stopping_records: Vec::new(),
1509 },
1510 })
1511 }
1512 }
1513
1514 impl RuntimeController for MethodsPlsController {
1515 fn controller_id(&self) -> &ControllerId {
1516 &self.id
1517 }
1518
1519 fn export_artifact_payload(&self, artifact_id: &ArtifactId) -> Result<Option<Vec<u8>>> {
1520 Ok(self
1521 .exported_n4mm_by_artifact
1522 .lock()
1523 .map_err(|_| {
1524 DagMlError::RuntimeValidation(
1525 "portable Methods PLS N4MM sidecar lock poisoned".to_string(),
1526 )
1527 })?
1528 .remove(artifact_id))
1529 }
1530
1531 fn hydrate_artifact_payload(
1532 &self,
1533 request: &crate::runtime::ArtifactMaterializationRequest,
1534 payload: &[u8],
1535 ) -> Result<HandleRef> {
1536 if request.artifact.kind != "n4m_model"
1537 || request.artifact.backend != Some(ArtifactBackend::Raw)
1538 {
1539 return Err(DagMlError::RuntimeValidation(format!(
1540 "portable Methods PLS cannot hydrate non-N4MM artifact `{}`",
1541 request.artifact.id
1542 )));
1543 }
1544 if format!("{:x}", Sha256::digest(payload))
1545 != request
1546 .artifact
1547 .content_fingerprint
1548 .as_deref()
1549 .unwrap_or_default()
1550 || request.artifact.size_bytes != Some(payload.len() as u64)
1551 {
1552 return Err(DagMlError::RuntimeValidation(format!(
1553 "portable Methods PLS payload `{}` does not match its artifact reference",
1554 request.artifact.id
1555 )));
1556 }
1557 let context = Context::new()
1561 .map_err(|error| Self::native_error("hydrate_context_create", error))?;
1562 Model::import_n4mm(&context, payload)
1563 .map_err(|error| Self::native_error("hydrate_import_n4mm", error))?;
1564 let handle = self.handle(HandleKind::Model);
1565 self.hydrated_n4mm_by_handle
1566 .lock()
1567 .map_err(|_| {
1568 DagMlError::RuntimeValidation(
1569 "portable Methods PLS hydrated N4MM lock poisoned".to_string(),
1570 )
1571 })?
1572 .insert(handle.handle, payload.to_vec());
1573 Ok(handle)
1574 }
1575
1576 fn release_hydrated_artifact_payload(&self, handle: &HandleRef) -> Result<()> {
1577 if handle.kind != HandleKind::Model || handle.owner_controller != self.id {
1578 return Err(DagMlError::RuntimeValidation(format!(
1579 "portable Methods PLS cannot release foreign hydrated handle {}",
1580 handle.handle
1581 )));
1582 }
1583 self.hydrated_n4mm_by_handle
1588 .lock()
1589 .map_err(|_| {
1590 DagMlError::RuntimeValidation(
1591 "portable Methods PLS hydrated N4MM lock poisoned".to_string(),
1592 )
1593 })?
1594 .remove(&handle.handle);
1595 Ok(())
1596 }
1597
1598 fn invoke(&self, task: &NodeTask) -> Result<NodeResult> {
1599 Err(DagMlError::RuntimeValidation(format!(
1600 "portable Methods PLS node `{}` requires a RuntimeDataProvider numeric view",
1601 task.node_plan.node_id
1602 )))
1603 }
1604
1605 fn invoke_with_data_provider(
1606 &self,
1607 task: &NodeTask,
1608 provider: &dyn RuntimeDataProvider,
1609 ) -> Result<NodeResult> {
1610 if task.node_plan.kind != crate::graph::NodeKind::Model {
1611 return Err(DagMlError::RuntimeValidation(
1612 "portable Methods PLS controller only serves model nodes".to_string(),
1613 ));
1614 }
1615 let request = Self::request(task, provider, "x")?;
1616 provider.preflight_methods_pls(&request)?;
1617 let data = provider.methods_pls_data(&request)?;
1618 data.validate_for(&request)?;
1619 match task.phase {
1620 Phase::FitCv | Phase::Refit => {
1621 let (context, model) = Self::fit(task, &data)?;
1622 let prediction_data = data.prediction.as_ref().unwrap_or(&data.fit);
1623 let values = Self::predict(&context, &model, prediction_data)?;
1624 let artifact = if task.phase == Phase::Refit {
1625 let bytes = model
1626 .export_n4mm()
1627 .map_err(|error| Self::native_error("export_n4mm", error))?;
1628 let handle = self.handle(HandleKind::Model);
1629 let fingerprint = format!("{:x}", Sha256::digest(&bytes));
1630 let id = ArtifactId::new(format!(
1631 "artifact:methods-pls:{}:refit",
1632 task.node_plan.node_id
1633 ))
1634 .map_err(|error| DagMlError::RuntimeValidation(error.to_string()))?;
1635 self.exported_n4mm_by_artifact
1636 .lock()
1637 .map_err(|_| {
1638 DagMlError::RuntimeValidation(
1639 "portable Methods PLS N4MM sidecar lock poisoned".to_string(),
1640 )
1641 })?
1642 .insert(id.clone(), bytes.clone());
1643 Some((
1644 ArtifactRef {
1645 id,
1646 kind: "n4m_model".to_string(),
1647 controller_id: self.id.clone(),
1648 backend: Some(ArtifactBackend::Raw),
1649 uri: Some(format!(
1654 "methods/{}.n4mm",
1655 task.node_plan.node_id.as_str().replace(':', "_")
1656 )),
1657 content_fingerprint: Some(fingerprint),
1658 size_bytes: Some(bytes.len() as u64),
1659 plugin: None,
1660 plugin_version: None,
1661 },
1662 handle,
1663 ))
1664 } else {
1665 None
1666 };
1667 let partition = match task.phase {
1668 Phase::FitCv => PredictionPartition::Validation,
1669 Phase::Refit if data.prediction.is_some() => PredictionPartition::Test,
1674 Phase::Refit | Phase::Predict => PredictionPartition::Final,
1675 _ => unreachable!("match arm admits only FIT_CV/REFIT"),
1676 };
1677 self.result(task, prediction_data, values, artifact, partition)
1678 }
1679 Phase::Predict => {
1680 let artifact = task.artifact_inputs.values().find(|artifact| artifact.controller_id == self.id).ok_or_else(|| DagMlError::RuntimeValidation("portable Methods PLS PREDICT requires its retained N4MM artifact reference".to_string()))?;
1681 let handle = task
1682 .input_handles
1683 .get(&crate::runtime::refit_artifact_input_key(&artifact.artifact.id))
1684 .ok_or_else(|| {
1685 DagMlError::RuntimeValidation(
1686 "portable Methods PLS PREDICT requires a hydrated N4MM runtime handle"
1687 .to_string(),
1688 )
1689 })?;
1690 let bytes = self.hydrated_n4mm_by_handle.lock().map_err(|_| DagMlError::RuntimeValidation("portable Methods PLS hydrated N4MM lock poisoned".to_string()))?.remove(&handle.handle).ok_or_else(|| DagMlError::RuntimeValidation("portable Methods PLS PREDICT requires N4MM bytes hydrated from the execution bundle in this controller instance".to_string()))?;
1694 let context = Context::new()
1695 .map_err(|error| Self::native_error("context_create", error))?;
1696 let model = Model::import_n4mm(&context, &bytes)
1697 .map_err(|error| Self::native_error("import_n4mm", error))?;
1698 let values = Self::predict(&context, &model, &data.fit)?;
1699 self.result(task, &data.fit, values, None, PredictionPartition::Final)
1700 }
1701 _ => Err(DagMlError::RuntimeValidation(
1702 "portable Methods PLS supports FIT_CV, REFIT, and PREDICT only".to_string(),
1703 )),
1704 }
1705 }
1706 }
1707
1708 pub struct MethodsRidgeController {
1717 id: ControllerId,
1718 _runtime: MethodsRuntime,
1719 next_handle: AtomicU64,
1720 exported_n4mm_by_artifact: Mutex<BTreeMap<ArtifactId, Vec<u8>>>,
1721 hydrated_n4mm_by_handle: Mutex<BTreeMap<u64, Vec<u8>>>,
1722 }
1723
1724 type RidgePredictionInputs<'a> = BTreeMap<String, &'a PredictionInputSpec>;
1725
1726 struct RidgePredictionOutput {
1727 sample_ids: Vec<crate::SampleId>,
1728 target_names: Vec<String>,
1729 values: Vec<Vec<f64>>,
1730 partition: PredictionPartition,
1731 }
1732
1733 impl MethodsRidgeController {
1734 pub fn new(runtime: MethodsRuntime) -> Self {
1735 Self {
1736 id: ControllerId::new(METHODS_RIDGE_CONTROLLER_ID)
1737 .expect("Methods Ridge controller id is valid"),
1738 _runtime: runtime,
1739 next_handle: AtomicU64::new(0),
1740 exported_n4mm_by_artifact: Mutex::new(BTreeMap::new()),
1741 hydrated_n4mm_by_handle: Mutex::new(BTreeMap::new()),
1742 }
1743 }
1744
1745 fn handle(&self, kind: HandleKind) -> HandleRef {
1746 HandleRef {
1747 handle: self.next_handle.fetch_add(1, Ordering::SeqCst) + 1,
1748 kind,
1749 owner_controller: self.id.clone(),
1750 }
1751 }
1752
1753 fn lambda(task: &NodeTask) -> Result<f64> {
1754 let lambda = task
1755 .node_plan
1756 .params
1757 .get("ridge_lambda")
1758 .and_then(serde_json::Value::as_f64)
1759 .ok_or_else(|| {
1760 DagMlError::RuntimeValidation(format!(
1761 "portable Methods Ridge node `{}` requires finite numeric `ridge_lambda`",
1762 task.node_plan.node_id
1763 ))
1764 })?;
1765 if !lambda.is_finite() || lambda < 0.0 {
1766 return Err(DagMlError::RuntimeValidation(
1767 "portable Methods Ridge `ridge_lambda` must be finite and non-negative"
1768 .to_string(),
1769 ));
1770 }
1771 Ok(lambda)
1772 }
1773
1774 fn feature_matrix(
1775 specs: &BTreeMap<String, &PredictionInputSpec>,
1776 sample_ids: &[crate::SampleId],
1777 expected_partition: PredictionPartition,
1778 label: &str,
1779 ) -> Result<crate::runtime::MethodsPlsMatrix> {
1780 if specs.len() < 2 {
1781 return Err(DagMlError::RuntimeValidation(format!(
1782 "portable Methods Ridge {label} requires OOF predictions from at least two base producers"
1783 )));
1784 }
1785 let mut cols = 0usize;
1786 for (key, spec) in specs {
1787 if spec.partition != expected_partition
1788 || spec.prediction_level != PredictionLevel::Sample
1789 || spec.sample_ids != sample_ids
1790 || spec.prediction_width == 0
1791 || spec.values.len() != sample_ids.len()
1792 || spec.values.iter().any(|row| {
1793 row.len() != spec.prediction_width
1794 || row.iter().any(|value| !value.is_finite())
1795 })
1796 {
1797 return Err(DagMlError::RuntimeValidation(format!(
1798 "portable Methods Ridge {label} input `{key}` is not an exact finite sample-level prediction matrix for the scheduler scope"
1799 )));
1800 }
1801 cols = cols.checked_add(spec.prediction_width).ok_or_else(|| {
1802 DagMlError::RuntimeValidation(
1803 "portable Methods Ridge feature width overflows usize".to_string(),
1804 )
1805 })?;
1806 }
1807 let capacity = sample_ids.len().checked_mul(cols).ok_or_else(|| {
1808 DagMlError::RuntimeValidation(
1809 "portable Methods Ridge feature matrix size overflows usize".to_string(),
1810 )
1811 })?;
1812 let mut values = Vec::with_capacity(capacity);
1813 for row in 0..sample_ids.len() {
1814 for spec in specs.values() {
1815 values.extend_from_slice(&spec.values[row]);
1816 }
1817 }
1818 let matrix = crate::runtime::MethodsPlsMatrix {
1819 values,
1820 rows: sample_ids.len(),
1821 cols,
1822 };
1823 matrix.validate(&format!("Ridge {label} OOF"))?;
1824 Ok(matrix)
1825 }
1826
1827 fn split_prediction_inputs<'a>(
1828 task: &'a NodeTask,
1829 suffix: &str,
1830 output_required: bool,
1831 ) -> Result<(RidgePredictionInputs<'a>, RidgePredictionInputs<'a>)> {
1832 let mut fit = BTreeMap::new();
1833 let mut output = BTreeMap::new();
1834 for (key, spec) in &task.prediction_inputs {
1835 if let Some(base) = key.strip_suffix(suffix) {
1836 if base.is_empty() || output.insert(base.to_string(), spec).is_some() {
1837 return Err(DagMlError::RuntimeValidation(format!(
1838 "portable Methods Ridge received duplicate or malformed output OOF input `{key}`"
1839 )));
1840 }
1841 } else if fit.insert(key.clone(), spec).is_some() {
1845 return Err(DagMlError::RuntimeValidation(format!(
1846 "portable Methods Ridge received unsupported prediction input `{key}`; expected base keys and `{suffix}` counterparts"
1847 )));
1848 }
1849 }
1850 if fit.is_empty()
1851 || (output_required
1852 && fit.keys().collect::<Vec<_>>() != output.keys().collect::<Vec<_>>())
1853 || (!output.is_empty()
1854 && fit.keys().collect::<Vec<_>>() != output.keys().collect::<Vec<_>>())
1855 {
1856 return Err(DagMlError::RuntimeValidation(format!(
1857 "portable Methods Ridge requires base OOF inputs and, when present, exactly paired `{suffix}` prediction inputs"
1858 )));
1859 }
1860 Ok((fit, output))
1861 }
1862
1863 fn predict_only_inputs(task: &NodeTask) -> Result<RidgePredictionInputs<'_>> {
1864 let mut inputs = BTreeMap::new();
1865 for (key, spec) in &task.prediction_inputs {
1866 let Some(base) = key.strip_suffix(":predict") else {
1867 return Err(DagMlError::RuntimeValidation(format!(
1868 "portable Methods Ridge PREDICT accepts only `:predict` OOF inputs, received `{key}`"
1869 )));
1870 };
1871 if base.is_empty() || inputs.insert(base.to_string(), spec).is_some() {
1872 return Err(DagMlError::RuntimeValidation(format!(
1873 "portable Methods Ridge PREDICT received duplicate or malformed OOF input `{key}`"
1874 )));
1875 }
1876 }
1877 if inputs.len() < 2 {
1878 return Err(DagMlError::RuntimeValidation(
1879 "portable Methods Ridge PREDICT requires at least two `:predict` OOF inputs"
1880 .to_string(),
1881 ));
1882 }
1883 Ok(inputs)
1884 }
1885
1886 fn fit(
1887 task: &NodeTask,
1888 features: &crate::runtime::MethodsPlsMatrix,
1889 targets: &crate::runtime::MethodsPlsMatrix,
1890 ) -> Result<(Context, Model)> {
1891 let context = Context::new().map_err(|error| {
1892 MethodsPlsController::native_error("ridge_context_create", error)
1893 })?;
1894 let config = Config::new().map_err(|error| {
1895 MethodsPlsController::native_error("ridge_config_create", error)
1896 })?;
1897 let x = MatrixRef::row_major(&features.values, features.rows, features.cols)
1898 .map_err(|error| MethodsPlsController::native_error("ridge_fit_features", error))?;
1899 let y = MatrixRef::row_major(&targets.values, targets.rows, targets.cols)
1900 .map_err(|error| MethodsPlsController::native_error("ridge_fit_targets", error))?;
1901 let model = Model::fit_ridge(&context, &config, x, y, Self::lambda(task)?)
1902 .map_err(|error| MethodsPlsController::native_error("ridge_fit", error))?;
1903 Ok((context, model))
1904 }
1905
1906 fn predict(
1907 context: &Context,
1908 model: &Model,
1909 features: &crate::runtime::MethodsPlsMatrix,
1910 ) -> Result<Vec<Vec<f64>>> {
1911 let x = MatrixRef::row_major(&features.values, features.rows, features.cols).map_err(
1912 |error| MethodsPlsController::native_error("ridge_predict_features", error),
1913 )?;
1914 let prediction = model
1915 .predict(context, x)
1916 .map_err(|error| MethodsPlsController::native_error("ridge_predict", error))?;
1917 Ok(prediction
1918 .data
1919 .chunks(prediction.cols)
1920 .map(|row| row.to_vec())
1921 .collect())
1922 }
1923
1924 fn result(
1925 &self,
1926 task: &NodeTask,
1927 output: RidgePredictionOutput,
1928 targets: Option<&crate::runtime::MethodsPlsMatrix>,
1929 artifact: Option<(ArtifactRef, HandleRef)>,
1930 ) -> Result<NodeResult> {
1931 let regression_targets = if task.phase == Phase::FitCv {
1932 let targets = targets.ok_or_else(|| {
1933 DagMlError::RuntimeValidation(
1934 "portable Methods Ridge FIT_CV requires validation targets".to_string(),
1935 )
1936 })?;
1937 vec![RegressionTargetBlock {
1938 level: PredictionLevel::Sample,
1939 unit_ids: output
1940 .sample_ids
1941 .iter()
1942 .cloned()
1943 .map(PredictionUnitId::Sample)
1944 .collect(),
1945 values: targets
1946 .values
1947 .chunks(targets.cols)
1948 .map(|row| row.to_vec())
1949 .collect(),
1950 target_names: output.target_names.clone(),
1951 }]
1952 } else {
1953 Vec::new()
1954 };
1955 let (artifacts, artifact_handles) = artifact
1956 .map(|(artifact, handle)| {
1957 (
1958 vec![artifact.clone()],
1959 BTreeMap::from([(artifact.id, handle)]),
1960 )
1961 })
1962 .unwrap_or_default();
1963 let artifact_refs = artifacts.clone();
1964 Ok(NodeResult {
1965 schema_version: None,
1966 node_id: task.node_plan.node_id.clone(),
1967 outputs: BTreeMap::from([("oof".to_string(), self.handle(HandleKind::Prediction))]),
1968 predictions: vec![PredictionBlock {
1969 prediction_id: Some(format!(
1970 "methods-ridge:{}:{}:{}",
1971 task.node_plan.node_id,
1972 task.phase.as_str(),
1973 task.fold_id
1974 .as_ref()
1975 .map(|id| id.as_str())
1976 .unwrap_or("full")
1977 )),
1978 producer_node: task.node_plan.node_id.clone(),
1979 producer_port: Some("oof".to_string()),
1980 partition: output.partition,
1981 fold_id: (task.phase == Phase::FitCv)
1982 .then(|| task.fold_id.clone())
1983 .flatten(),
1984 sample_ids: output.sample_ids,
1985 values: output.values,
1986 target_names: output.target_names,
1987 }],
1988 observation_predictions: Vec::new(),
1989 aggregated_predictions: Vec::new(),
1990 explanations: Vec::new(),
1991 shape_deltas: Vec::new(),
1992 artifacts,
1993 artifact_handles,
1994 fit_influence_diagnostics: Vec::new(),
1995 regression_targets,
1996 lineage: LineageRecord {
1997 record_id: LineageId::new(format!(
1998 "lineage:methods-ridge:{}:{}:{}:{}",
1999 task.node_plan.node_id,
2000 task.phase.as_str(),
2001 task.variant_id
2002 .as_ref()
2003 .map(|id| id.as_str())
2004 .unwrap_or("base"),
2005 task.fold_id
2006 .as_ref()
2007 .map(|id| id.as_str())
2008 .unwrap_or("full")
2009 ))
2010 .expect("valid native Ridge lineage id"),
2011 run_id: task.run_id.clone(),
2012 node_id: task.node_plan.node_id.clone(),
2013 phase: task.phase,
2014 controller_id: self.id.clone(),
2015 controller_version: task.node_plan.controller_version.clone(),
2016 variant_id: task.variant_id.clone(),
2017 fold_id: task.fold_id.clone(),
2018 branch_path: task.branch_path.clone(),
2019 input_lineage: Vec::new(),
2020 artifact_refs,
2021 params_fingerprint: task.node_plan.params_fingerprint.clone(),
2022 data_model_shape_fingerprint: None,
2023 aggregation_policy_fingerprint: None,
2024 seed: task.seed,
2025 unsafe_flags: BTreeSet::new(),
2026 metrics: BTreeMap::new(),
2027 loss_attestations: Vec::new(),
2028 early_stopping_records: Vec::new(),
2029 },
2030 })
2031 }
2032 }
2033
2034 impl RuntimeController for MethodsRidgeController {
2035 fn controller_id(&self) -> &ControllerId {
2036 &self.id
2037 }
2038
2039 fn export_artifact_payload(&self, artifact_id: &ArtifactId) -> Result<Option<Vec<u8>>> {
2040 Ok(self
2041 .exported_n4mm_by_artifact
2042 .lock()
2043 .map_err(|_| {
2044 DagMlError::RuntimeValidation(
2045 "portable Methods Ridge N4MM sidecar lock poisoned".to_string(),
2046 )
2047 })?
2048 .remove(artifact_id))
2049 }
2050
2051 fn hydrate_artifact_payload(
2052 &self,
2053 request: &crate::runtime::ArtifactMaterializationRequest,
2054 payload: &[u8],
2055 ) -> Result<HandleRef> {
2056 if request.artifact.kind != "n4m_model"
2057 || request.artifact.backend != Some(ArtifactBackend::Raw)
2058 || format!("{:x}", Sha256::digest(payload))
2059 != request
2060 .artifact
2061 .content_fingerprint
2062 .as_deref()
2063 .unwrap_or_default()
2064 || request.artifact.size_bytes != Some(payload.len() as u64)
2065 {
2066 return Err(DagMlError::RuntimeValidation(format!(
2067 "portable Methods Ridge payload `{}` does not match its N4MM artifact reference",
2068 request.artifact.id
2069 )));
2070 }
2071 let context = Context::new().map_err(|error| {
2072 MethodsPlsController::native_error("ridge_hydrate_context_create", error)
2073 })?;
2074 Model::import_n4mm(&context, payload).map_err(|error| {
2075 MethodsPlsController::native_error("ridge_hydrate_import_n4mm", error)
2076 })?;
2077 let handle = self.handle(HandleKind::Model);
2078 self.hydrated_n4mm_by_handle
2079 .lock()
2080 .map_err(|_| {
2081 DagMlError::RuntimeValidation(
2082 "portable Methods Ridge hydrated N4MM lock poisoned".to_string(),
2083 )
2084 })?
2085 .insert(handle.handle, payload.to_vec());
2086 Ok(handle)
2087 }
2088
2089 fn release_hydrated_artifact_payload(&self, handle: &HandleRef) -> Result<()> {
2090 if handle.kind != HandleKind::Model || handle.owner_controller != self.id {
2091 return Err(DagMlError::RuntimeValidation(format!(
2092 "portable Methods Ridge cannot release foreign hydrated handle {}",
2093 handle.handle
2094 )));
2095 }
2096 self.hydrated_n4mm_by_handle
2097 .lock()
2098 .map_err(|_| {
2099 DagMlError::RuntimeValidation(
2100 "portable Methods Ridge hydrated N4MM lock poisoned".to_string(),
2101 )
2102 })?
2103 .remove(&handle.handle);
2104 Ok(())
2105 }
2106
2107 fn invoke(&self, task: &NodeTask) -> Result<NodeResult> {
2108 Err(DagMlError::RuntimeValidation(format!(
2109 "portable Methods Ridge node `{}` requires a RuntimeDataProvider numeric view",
2110 task.node_plan.node_id
2111 )))
2112 }
2113
2114 fn invoke_with_data_provider(
2115 &self,
2116 task: &NodeTask,
2117 provider: &dyn RuntimeDataProvider,
2118 ) -> Result<NodeResult> {
2119 if task.node_plan.kind != crate::graph::NodeKind::Model {
2120 return Err(DagMlError::RuntimeValidation(
2121 "portable Methods Ridge controller only serves model nodes".to_string(),
2122 ));
2123 }
2124 let request = MethodsPlsController::request(task, provider, "x_original")?;
2128 provider.preflight_methods_pls(&request)?;
2129 let data = provider.methods_pls_data(&request)?;
2130 data.validate_for(&request)?;
2131 match task.phase {
2132 Phase::FitCv => {
2133 let (fit_specs, validation_specs) =
2134 Self::split_prediction_inputs(task, ":outer", true)?;
2135 let prediction = data.prediction.as_ref().ok_or_else(|| {
2136 DagMlError::RuntimeValidation(
2137 "portable Methods Ridge FIT_CV requires a validation data view"
2138 .to_string(),
2139 )
2140 })?;
2141 let fit_features = Self::feature_matrix(
2142 &fit_specs,
2143 &data.fit.sample_ids,
2144 PredictionPartition::Validation,
2145 "inner FIT_CV",
2146 )?;
2147 let validation_features = Self::feature_matrix(
2148 &validation_specs,
2149 &prediction.sample_ids,
2150 PredictionPartition::Validation,
2151 "outer FIT_CV",
2152 )?;
2153 let targets = data.fit.y.as_ref().ok_or_else(|| {
2154 DagMlError::RuntimeValidation(
2155 "portable Methods Ridge FIT_CV requires fitting targets".to_string(),
2156 )
2157 })?;
2158 let validation_targets = prediction.y.as_ref().ok_or_else(|| {
2159 DagMlError::RuntimeValidation(
2160 "portable Methods Ridge FIT_CV requires validation targets".to_string(),
2161 )
2162 })?;
2163 let (context, model) = Self::fit(task, &fit_features, targets)?;
2164 let values = Self::predict(&context, &model, &validation_features)?;
2165 Self::result(
2166 self,
2167 task,
2168 RidgePredictionOutput {
2169 sample_ids: prediction.sample_ids.clone(),
2170 target_names: prediction.target_names.clone(),
2171 values,
2172 partition: PredictionPartition::Validation,
2173 },
2174 Some(validation_targets),
2175 None,
2176 )
2177 }
2178 Phase::Refit => {
2179 let (fit_specs, refit_specs) =
2180 Self::split_prediction_inputs(task, ":refit", false)?;
2181 let fit_features = Self::feature_matrix(
2182 &fit_specs,
2183 &data.fit.sample_ids,
2184 PredictionPartition::Validation,
2185 "REFIT",
2186 )?;
2187 let targets = data.fit.y.as_ref().ok_or_else(|| {
2188 DagMlError::RuntimeValidation(
2189 "portable Methods Ridge REFIT requires targets".to_string(),
2190 )
2191 })?;
2192 let (context, model) = Self::fit(task, &fit_features, targets)?;
2193 let (output_ids, values, partition) = if refit_specs.is_empty() {
2194 (
2199 data.fit.sample_ids.clone(),
2200 Self::predict(&context, &model, &fit_features)?,
2201 PredictionPartition::Final,
2202 )
2203 } else {
2204 let output_ids = refit_specs
2205 .values()
2206 .next()
2207 .expect("paired inputs checked")
2208 .sample_ids
2209 .clone();
2210 let refit_features = Self::feature_matrix(
2211 &refit_specs,
2212 &output_ids,
2213 PredictionPartition::Test,
2214 "REFIT output",
2215 )?;
2216 (
2217 output_ids,
2218 Self::predict(&context, &model, &refit_features)?,
2219 PredictionPartition::Test,
2220 )
2221 };
2222 let bytes = model.export_n4mm().map_err(|error| {
2223 MethodsPlsController::native_error("ridge_export_n4mm", error)
2224 })?;
2225 let id = ArtifactId::new(format!(
2226 "artifact:methods-ridge:{}:refit",
2227 task.node_plan.node_id
2228 ))
2229 .map_err(|error| DagMlError::RuntimeValidation(error.to_string()))?;
2230 let handle = self.handle(HandleKind::Model);
2231 self.exported_n4mm_by_artifact
2232 .lock()
2233 .map_err(|_| {
2234 DagMlError::RuntimeValidation(
2235 "portable Methods Ridge N4MM sidecar lock poisoned".to_string(),
2236 )
2237 })?
2238 .insert(id.clone(), bytes.clone());
2239 let artifact = ArtifactRef {
2240 id,
2241 kind: "n4m_model".to_string(),
2242 controller_id: self.id.clone(),
2243 backend: Some(ArtifactBackend::Raw),
2244 uri: Some(format!(
2245 "methods/{}.n4mm",
2246 task.node_plan.node_id.as_str().replace(':', "_")
2247 )),
2248 content_fingerprint: Some(format!("{:x}", Sha256::digest(&bytes))),
2249 size_bytes: Some(bytes.len() as u64),
2250 plugin: None,
2251 plugin_version: None,
2252 };
2253 Self::result(
2254 self,
2255 task,
2256 RidgePredictionOutput {
2257 sample_ids: output_ids,
2258 target_names: data.fit.target_names.clone(),
2259 values,
2260 partition,
2261 },
2262 None,
2263 Some((artifact, handle)),
2264 )
2265 }
2266 Phase::Predict => {
2267 let predict_specs = Self::predict_only_inputs(task)?;
2268 let features = Self::feature_matrix(
2269 &predict_specs,
2270 &data.fit.sample_ids,
2271 PredictionPartition::Final,
2272 "PREDICT",
2273 )?;
2274 let artifact = task.artifact_inputs.values().find(|artifact| artifact.controller_id == self.id).ok_or_else(|| DagMlError::RuntimeValidation("portable Methods Ridge PREDICT requires its retained N4MM artifact reference".to_string()))?;
2275 let handle = task.input_handles.get(&crate::runtime::refit_artifact_input_key(&artifact.artifact.id)).ok_or_else(|| DagMlError::RuntimeValidation("portable Methods Ridge PREDICT requires a hydrated N4MM runtime handle".to_string()))?;
2276 let bytes = self.hydrated_n4mm_by_handle.lock().map_err(|_| DagMlError::RuntimeValidation("portable Methods Ridge hydrated N4MM lock poisoned".to_string()))?.remove(&handle.handle).ok_or_else(|| DagMlError::RuntimeValidation("portable Methods Ridge PREDICT requires N4MM bytes hydrated from the execution bundle in this controller instance".to_string()))?;
2277 let context = Context::new().map_err(|error| {
2278 MethodsPlsController::native_error("ridge_predict_context_create", error)
2279 })?;
2280 let model = Model::import_n4mm(&context, &bytes).map_err(|error| {
2281 MethodsPlsController::native_error("ridge_predict_import_n4mm", error)
2282 })?;
2283 let values = Self::predict(&context, &model, &features)?;
2284 Self::result(
2285 self,
2286 task,
2287 RidgePredictionOutput {
2288 sample_ids: data.fit.sample_ids.clone(),
2289 target_names: data.fit.target_names.clone(),
2290 values,
2291 partition: PredictionPartition::Final,
2292 },
2293 None,
2294 None,
2295 )
2296 }
2297 _ => Err(DagMlError::RuntimeValidation(
2298 "portable Methods Ridge supports FIT_CV, REFIT, and PREDICT only".to_string(),
2299 )),
2300 }
2301 }
2302 }
2303
2304 #[cfg(test)]
2305 mod ridge_tests {
2306 use std::collections::BTreeMap;
2307
2308 use super::*;
2309
2310 fn sample_ids() -> Vec<crate::SampleId> {
2311 vec![
2312 crate::SampleId::new("sample:1").unwrap(),
2313 crate::SampleId::new("sample:2").unwrap(),
2314 ]
2315 }
2316
2317 fn prediction(values: Vec<Vec<f64>>) -> PredictionInputSpec {
2318 PredictionInputSpec {
2319 producer_node: crate::NodeId::new("model:base").unwrap(),
2320 source_port: "oof".to_string(),
2321 target_port: "oof".to_string(),
2322 partition: PredictionPartition::Validation,
2323 prediction_level: PredictionLevel::Sample,
2324 fold_id: None,
2325 fold_ids: Vec::new(),
2326 unit_ids: Vec::new(),
2327 sample_ids: sample_ids(),
2328 values,
2329 prediction_width: 1,
2330 target_names: vec!["y".to_string()],
2331 }
2332 }
2333
2334 #[test]
2335 fn ridge_features_are_stably_ordered_and_identity_aligned() {
2336 let mut inputs = BTreeMap::new();
2337 inputs.insert(
2340 "model:z.oof".to_string(),
2341 prediction(vec![vec![30.0], vec![40.0]]),
2342 );
2343 inputs.insert(
2344 "model:a.oof".to_string(),
2345 prediction(vec![vec![10.0], vec![20.0]]),
2346 );
2347
2348 let refs = inputs
2349 .iter()
2350 .map(|(key, spec)| (key.clone(), spec))
2351 .collect();
2352 let matrix = MethodsRidgeController::feature_matrix(
2353 &refs,
2354 &sample_ids(),
2355 PredictionPartition::Validation,
2356 "test",
2357 )
2358 .unwrap();
2359 assert_eq!(matrix.rows, 2);
2360 assert_eq!(matrix.cols, 2);
2361 assert_eq!(matrix.values, vec![10.0, 30.0, 20.0, 40.0]);
2362
2363 let mut misaligned = inputs["model:a.oof"].clone();
2364 misaligned.sample_ids.reverse();
2365 let refs = BTreeMap::from([
2366 ("model:a.oof".to_string(), &misaligned),
2367 ("model:z.oof".to_string(), &inputs["model:z.oof"]),
2368 ]);
2369 assert!(MethodsRidgeController::feature_matrix(
2370 &refs,
2371 &sample_ids(),
2372 PredictionPartition::Validation,
2373 "test",
2374 )
2375 .is_err());
2376 }
2377 }
2378}
2379
2380#[cfg(feature = "methods-optimizer")]
2381pub use pls_controller::{MethodsPlsController, MethodsRidgeController};
2382
2383#[cfg(feature = "methods-optimizer")]
2391pub fn register_methods_runtime_controllers(
2392 registry: &mut crate::runtime::RuntimeControllerRegistry,
2393 hpo_controller_id: crate::ControllerId,
2394 runtime: MethodsRuntime,
2395) -> crate::Result<()> {
2396 let pls_controller_id = crate::ControllerId::new(METHODS_PLS_CONTROLLER_ID)
2397 .expect("the fixed Methods PLS controller id is valid");
2398 let ridge_controller_id = crate::ControllerId::new(METHODS_RIDGE_CONTROLLER_ID)
2399 .expect("the fixed Methods Ridge controller id is valid");
2400 if hpo_controller_id == pls_controller_id || hpo_controller_id == ridge_controller_id {
2401 return Err(crate::DagMlError::RuntimeValidation(
2402 "Methods HPO controller id must differ from the Methods PLS and Ridge controller ids"
2403 .to_string(),
2404 ));
2405 }
2406 for controller_id in [&pls_controller_id, &ridge_controller_id, &hpo_controller_id] {
2407 if registry.get(controller_id).is_some() {
2408 return Err(crate::DagMlError::RuntimeValidation(format!(
2409 "duplicate runtime controller `{controller_id}`"
2410 )));
2411 }
2412 }
2413 registry.register(Box::new(MethodsPlsController::new(runtime.clone())))?;
2414 registry.register(Box::new(MethodsRidgeController::new(runtime.clone())))?;
2415 registry.register(Box::new(MethodsHpoController::new(
2416 hpo_controller_id,
2417 runtime,
2418 )))?;
2419 Ok(())
2420}
2421
2422#[cfg(feature = "methods-optimizer")]
2423mod native {
2424 use super::*;
2425 use n4m::{
2426 Category, Direction, Error, ErrorKind, Optimizer, OptimizerOptions, Pruner, Sampler,
2427 SearchSpace, TrialError, TrialSnapshot, TrialStatus,
2428 };
2429
2430 pub(super) struct MethodsHpoStudy {
2431 manifest: MethodsHpoControllerManifest,
2432 methods_abi: String,
2433 _context: n4m::Context,
2434 optimizer: Optimizer,
2435 events: Vec<HpoEvent>,
2436 }
2437
2438 #[allow(dead_code)]
2442 struct MethodsHpoReporter<'a> {
2443 study: &'a mut MethodsHpoStudy,
2444 trial_id: i64,
2445 pruned: Option<HpoTrial>,
2446 }
2447
2448 impl HpoIntermediateReporter for MethodsHpoReporter<'_> {
2449 fn report(&mut self, step: i32, score: f64) -> HpoResult<HpoReportOutcome> {
2450 if self.pruned.is_some() {
2451 return Err(HpoError::InvalidTrial {
2452 reason: "cannot report after native pruning terminalized the trial".to_string(),
2453 });
2454 }
2455 if self.study.report_intermediate(self.trial_id, step, score)? {
2456 let snapshot = self.study.snapshot_for(self.trial_id)?;
2457 if snapshot.status != HpoTrialStatus::Pruned {
2458 return Err(HpoError::InvalidTrial {
2459 reason: "native pruner returned true without a PRUNED snapshot".to_string(),
2460 });
2461 }
2462 self.pruned = Some(snapshot.clone());
2463 return Ok(HpoReportOutcome::Pruned(snapshot));
2464 }
2465 Ok(HpoReportOutcome::Continue)
2466 }
2467 }
2468
2469 impl MethodsHpoStudy {
2470 pub(super) fn create(config: MethodsHpoStudyConfig) -> HpoResult<Self> {
2471 methods_optimizer_preflight()?;
2472 let methods_abi = config.methods_abi_identity()?;
2473 let fingerprint = config.search_space.fingerprint()?;
2474 let optimizer_fingerprint = config.optimizer.fingerprint()?;
2475 let manifest = MethodsHpoControllerManifest {
2476 schema_version: HPO_MANIFEST_SCHEMA_VERSION,
2477 binding: HpoStudyBinding {
2478 controller_id: config.controller_id,
2479 study_id: config.study_id,
2480 search_space_fingerprint: fingerprint,
2481 optimizer_fingerprint,
2482 },
2483 };
2484 manifest.validate()?;
2485 let context =
2486 n4m::Context::new().map_err(|error| native_error("context_create", error))?;
2487 let native_space = create_space(&config.search_space)?;
2488 let options = create_options(&config.optimizer);
2489 let optimizer = Optimizer::new(&context, &native_space, &options)
2490 .map_err(|error| native_error("optimizer_create", error))?;
2491 Ok(Self {
2492 manifest,
2493 methods_abi,
2494 _context: context,
2495 optimizer,
2496 events: Vec::new(),
2497 })
2498 }
2499 pub(super) fn restore(
2500 config: MethodsHpoStudyConfig,
2501 checkpoint: &N4moptCheckpointArtifact,
2502 ) -> HpoResult<Self> {
2503 methods_optimizer_preflight()?;
2504 let methods_abi = config.methods_abi_identity()?;
2505 let fingerprint = config.search_space.fingerprint()?;
2506 let optimizer_fingerprint = config.optimizer.fingerprint()?;
2507 let binding = HpoStudyBinding {
2508 controller_id: config.controller_id,
2509 study_id: config.study_id,
2510 search_space_fingerprint: fingerprint,
2511 optimizer_fingerprint,
2512 };
2513 checkpoint.validate()?;
2514 if checkpoint.binding != binding || checkpoint.methods_abi != methods_abi {
2515 return Err(HpoError::CheckpointBindingMismatch {
2516 reason: "study/search-space or Methods ABI differs from checkpoint".to_string(),
2517 });
2518 }
2519 let context =
2522 n4m::Context::new().map_err(|error| native_error("context_create", error))?;
2523 let optimizer = Optimizer::load_n4mopt(&context, &checkpoint.opaque_payload)
2524 .map_err(|error| native_error("load_n4mopt", error))?;
2525 Ok(Self {
2526 manifest: MethodsHpoControllerManifest {
2527 schema_version: HPO_MANIFEST_SCHEMA_VERSION,
2528 binding,
2529 },
2530 methods_abi,
2531 _context: context,
2532 optimizer,
2533 events: Vec::new(),
2534 })
2535 }
2536 #[allow(dead_code)]
2537 pub fn manifest(&self) -> &MethodsHpoControllerManifest {
2538 &self.manifest
2539 }
2540 #[allow(dead_code)]
2541 pub fn events(&self) -> &[HpoEvent] {
2542 &self.events
2543 }
2544 pub fn ask(&mut self) -> HpoResult<HpoTrial> {
2545 let id = self
2546 .optimizer
2547 .ask()
2548 .and_then(|trial| trial.id())
2549 .map_err(|error| native_error("ask", error))?;
2550 let trial = self.snapshot_for(id)?;
2551 self.events.push(HpoEvent::Asked { trial_id: trial.id });
2552 Ok(trial)
2553 }
2554 #[allow(dead_code)]
2555 pub fn ask_batch(&mut self, count: i32) -> HpoResult<HpoBatch> {
2556 match self.optimizer.ask_batch(count) {
2557 Ok(native_trials) => {
2558 let ids = native_trials
2559 .iter()
2560 .map(|trial| {
2561 trial
2562 .id()
2563 .map_err(|error| native_error("trial_get_id", error))
2564 })
2565 .collect::<HpoResult<Vec<_>>>()?;
2566 let trials = ids
2567 .into_iter()
2568 .map(|id| self.snapshot_for(id))
2569 .collect::<HpoResult<Vec<_>>>()?;
2570 for trial in &trials {
2571 self.events.push(HpoEvent::Asked { trial_id: trial.id });
2572 }
2573 Ok(HpoBatch {
2574 trials,
2575 native_error: None,
2576 })
2577 }
2578 Err(n4m::AskBatchError::Partial {
2579 error,
2580 trials: native_trials,
2581 }) => {
2582 let ids = native_trials
2583 .iter()
2584 .map(|trial| {
2585 trial
2586 .id()
2587 .map_err(|error| native_error("trial_get_id", error))
2588 })
2589 .collect::<HpoResult<Vec<_>>>()?;
2590 let trials = ids
2591 .into_iter()
2592 .map(|id| self.snapshot_for(id))
2593 .collect::<HpoResult<Vec<_>>>()?;
2594 for trial in &trials {
2595 self.events.push(HpoEvent::Asked { trial_id: trial.id });
2596 }
2597 Ok(HpoBatch {
2598 trials,
2599 native_error: Some(to_native_error(error)),
2600 })
2601 }
2602 Err(n4m::AskBatchError::Error(error)) => Err(native_error("ask_batch", error)),
2603 }
2604 }
2605 pub fn report_intermediate(
2606 &mut self,
2607 trial_id: i64,
2608 step: i32,
2609 score: f64,
2610 ) -> HpoResult<bool> {
2611 if !score.is_finite() {
2612 return Err(HpoError::InvalidTrial {
2613 reason: "intermediate score must be finite".to_string(),
2614 });
2615 }
2616 let should_prune = self
2617 .optimizer
2618 .tell_intermediate(trial_id, step, score)
2619 .map_err(|error| native_error("tell_intermediate", error))?;
2620 self.events.push(HpoEvent::Intermediate {
2621 trial_id,
2622 step,
2623 score,
2624 should_prune,
2625 });
2626 if should_prune {
2629 self.events.push(HpoEvent::Terminal {
2630 trial_id,
2631 status: HpoTrialStatus::Pruned,
2632 score: None,
2633 failure: None,
2634 });
2635 }
2636 Ok(should_prune)
2637 }
2638 pub fn tell(&mut self, trial_id: i64, terminal: HpoTerminal) -> HpoResult<HpoTrial> {
2642 let (status, score, failure) = match terminal {
2643 HpoTerminal::Completed { score } if score.is_finite() => {
2644 (TrialStatus::Completed, score, None)
2645 }
2646 HpoTerminal::Completed { .. } => {
2647 return Err(HpoError::InvalidTrial {
2648 reason: "terminal score must be finite".to_string(),
2649 })
2650 }
2651 HpoTerminal::Failed { failure } => (TrialStatus::Failed, 0.0, Some(failure)),
2652 HpoTerminal::Pruned { failure } => (TrialStatus::Pruned, 0.0, Some(failure)),
2655 HpoTerminal::Cancelled { failure } => (TrialStatus::Cancelled, 0.0, Some(failure)),
2656 };
2657 let native_failure = matches!(status, TrialStatus::Failed | TrialStatus::Cancelled)
2658 .then_some(failure.as_ref())
2659 .flatten()
2660 .map(|failure| TrialError {
2661 code: failure.code.clone(),
2662 message: failure.message.clone(),
2663 retryable: failure.retryable,
2664 });
2665 self.optimizer
2666 .tell_result(trial_id, status, score, native_failure.as_ref())
2667 .map_err(|error| native_error("tell_result", error))?;
2668 let snapshot = self.snapshot_for(trial_id)?;
2669 self.events.push(HpoEvent::Terminal {
2670 trial_id,
2671 status: snapshot.status,
2672 score: snapshot.score,
2673 failure: snapshot.failure.clone(),
2674 });
2675 Ok(snapshot)
2676 }
2677 #[allow(dead_code)]
2678 pub fn evaluate_one<E: HpoEvaluator>(
2679 &mut self,
2680 evaluator: &mut E,
2681 boundary: &mut HpoEvaluationBoundary<'_>,
2682 ) -> HpoResult<HpoTrial> {
2683 boundary.validate()?;
2684 let trial = self.ask()?;
2685 let (terminal, pruned) = {
2686 let mut reporter = MethodsHpoReporter {
2687 study: self,
2688 trial_id: trial.id,
2689 pruned: None,
2690 };
2691 let terminal = evaluator.evaluate_with_reporter(&trial, boundary, &mut reporter);
2692 (terminal, reporter.pruned.take())
2693 };
2694 let terminal = match terminal {
2695 Ok(terminal) => terminal,
2696 Err(error) => {
2697 if pruned.is_some() {
2698 return Err(error);
2701 }
2702 let failure = HpoFailure {
2703 code: "HPO_EVALUATION".to_string(),
2704 message: error.to_string(),
2705 retryable: true,
2706 };
2707 let _ = self.tell(trial.id, HpoTerminal::Failed { failure });
2710 return Err(error);
2711 }
2712 };
2713 if let Some(snapshot) = pruned {
2714 return Ok(snapshot);
2717 }
2718 self.tell(trial.id, terminal)
2719 }
2720 pub fn best(&self) -> HpoResult<Option<HpoBestTrial>> {
2721 let Some((trial, score)) = self
2722 .optimizer
2723 .best()
2724 .map_err(|error| native_error("best", error))?
2725 else {
2726 return Ok(None);
2727 };
2728 let id = trial
2729 .id()
2730 .map_err(|error| native_error("trial_get_id", error))?;
2731 self.snapshot_for(id)
2732 .map(|trial| Some(HpoBestTrial { trial, score }))
2733 }
2734 pub fn trials(&self) -> HpoResult<Vec<HpoTrial>> {
2735 self.optimizer
2736 .trials(0)
2737 .map_err(|error| native_error("trials", error))?
2738 .iter()
2739 .map(snapshot_trial)
2740 .collect()
2741 }
2742 pub fn save_checkpoint(&self) -> HpoResult<N4moptCheckpointArtifact> {
2743 let payload = self
2744 .optimizer
2745 .save_n4mopt()
2746 .map_err(|error| native_error("save_n4mopt", error))?;
2747 if payload.len() > MAX_N4MOPT_CHECKPOINT_BYTES {
2748 return Err(HpoError::InvalidCheckpoint {
2749 reason: "native checkpoint exceeds configured limit".to_string(),
2750 });
2751 }
2752 N4moptCheckpointArtifact::new(
2753 self.manifest.binding.clone(),
2754 self.methods_abi.clone(),
2755 payload,
2756 )
2757 }
2758 fn snapshot_for(&self, id: i64) -> HpoResult<HpoTrial> {
2759 self.optimizer
2760 .trials(id)
2761 .map_err(|error| native_error("trials", error))?
2762 .into_iter()
2763 .find(|trial| trial.id == id)
2764 .ok_or_else(|| HpoError::InvalidTrial {
2765 reason: format!("native trial history omitted committed trial `{id}`"),
2766 })
2767 .and_then(|trial| snapshot_trial(&trial))
2768 }
2769 }
2770
2771 fn create_space(space: &HpoSearchSpace) -> HpoResult<SearchSpace> {
2772 let mut result =
2773 SearchSpace::new().map_err(|error| native_error("search_space_create", error))?;
2774 for parameter in &space.parameters {
2775 let call = match parameter {
2776 HpoParameter::Int {
2777 name,
2778 low,
2779 high,
2780 step,
2781 log,
2782 } => result.add_int(name, *low, *high, *step, *log),
2783 HpoParameter::Float {
2784 name,
2785 low,
2786 high,
2787 step,
2788 log,
2789 } => result.add_float(name, *low, *high, *step, *log),
2790 HpoParameter::Categorical { name, values } => result
2791 .add_categorical(name, &values.iter().map(map_category).collect::<Vec<_>>()),
2792 HpoParameter::Ordinal { name, values } => result.add_ordinal(name, values),
2793 HpoParameter::SortedTuple {
2794 name,
2795 length,
2796 low,
2797 high,
2798 integer,
2799 } => result.add_sorted_tuple(name, *length, *low, *high, *integer),
2800 };
2801 call.map_err(|error| native_error("search_space_add", error))?;
2802 }
2803 Ok(result)
2804 }
2805 fn map_category(value: &HpoCategory) -> Category {
2806 match value {
2807 HpoCategory::String(value) => Category::Str(value.clone()),
2808 HpoCategory::Integer(value) => Category::Int(*value),
2809 HpoCategory::Float(value) => Category::Float(*value),
2810 HpoCategory::Boolean(value) => Category::Bool(*value),
2811 }
2812 }
2813 fn create_options(config: &HpoOptimizerConfig) -> OptimizerOptions {
2814 OptimizerOptions {
2815 sampler: match config.sampler {
2816 HpoSampler::Random => Sampler::Random,
2817 HpoSampler::Sobol => Sampler::Sobol,
2818 HpoSampler::Lhs => Sampler::Lhs,
2819 HpoSampler::Ternary => Sampler::Ternary,
2820 HpoSampler::Ga => Sampler::Ga,
2821 HpoSampler::Pso => Sampler::Pso,
2822 HpoSampler::Cmaes => Sampler::Cmaes,
2823 HpoSampler::Tpe => Sampler::Tpe,
2824 HpoSampler::GpEi => Sampler::GpEi,
2825 },
2826 pruner: match config.pruner {
2827 HpoPruner::None => Pruner::None,
2828 HpoPruner::Median => Pruner::Median,
2829 HpoPruner::Asha => Pruner::Asha,
2830 HpoPruner::Hyperband => Pruner::Hyperband,
2831 HpoPruner::Racing => Pruner::Racing,
2832 },
2833 direction: match config.direction {
2834 HpoDirection::Auto => Direction::Auto,
2835 HpoDirection::Minimize => Direction::Minimize,
2836 HpoDirection::Maximize => Direction::Maximize,
2837 },
2838 metric: match config.metric {
2839 HpoMetric::Rmse => n4m::Metric::Rmse,
2840 HpoMetric::Mse => n4m::Metric::Mse,
2841 HpoMetric::Mae => n4m::Metric::Mae,
2842 HpoMetric::R2 => n4m::Metric::R2,
2843 HpoMetric::Accuracy => n4m::Metric::Accuracy,
2844 HpoMetric::BalancedAccuracy => n4m::Metric::BalancedAccuracy,
2845 HpoMetric::F1 => n4m::Metric::F1,
2846 HpoMetric::Logloss => n4m::Metric::Logloss,
2847 },
2848 seed: config.seed,
2849 n_startup_trials: config.n_startup_trials,
2850 max_resource: config.max_resource,
2851 reduction_factor: config.reduction_factor,
2852 ..OptimizerOptions::default()
2853 }
2854 }
2855 fn map_status(value: TrialStatus) -> HpoTrialStatus {
2856 match value {
2857 TrialStatus::Running => HpoTrialStatus::Running,
2858 TrialStatus::Completed => HpoTrialStatus::Completed,
2859 TrialStatus::Pruned => HpoTrialStatus::Pruned,
2860 TrialStatus::Failed => HpoTrialStatus::Failed,
2861 TrialStatus::Cancelled => HpoTrialStatus::Cancelled,
2862 }
2863 }
2864 fn snapshot_trial(value: &TrialSnapshot) -> HpoResult<HpoTrial> {
2865 let mut parameters = BTreeMap::new();
2866 for (name, parameter) in &value.parameters {
2867 parameters.insert(
2868 name.clone(),
2869 HpoTrialParameter {
2870 name: name.clone(),
2871 value: parameter.value,
2872 native_kind: Some(map_parameter_kind(parameter.kind)),
2873 category_type: parameter.category_type.map(map_category_type),
2874 integer: parameter.integer,
2875 active: parameter.active,
2876 category_index: parameter.category_index,
2877 category_label: parameter.category_label.clone(),
2878 },
2879 );
2880 }
2881 Ok(HpoTrial {
2882 id: value.id,
2883 ask_sequence: value.ask_sequence,
2884 terminal_sequence: value.terminal_sequence,
2885 parameters,
2886 parameter_order: value.parameter_order.clone(),
2887 status: map_status(value.status),
2888 score: value.score,
2889 rung: value.rung,
2890 duration: value.duration,
2891 intermediates: value
2892 .intermediates
2893 .iter()
2894 .map(|item| HpoIntermediate {
2895 sequence: item.sequence,
2896 step: item.step,
2897 score: item.score,
2898 should_prune: item.should_prune,
2899 })
2900 .collect(),
2901 failure: value.error.as_ref().map(|item| HpoFailure {
2902 code: item.code.clone(),
2903 message: item.message.clone(),
2904 retryable: item.retryable,
2905 }),
2906 })
2907 }
2908 fn map_parameter_kind(value: n4m::ParameterKind) -> HpoNativeParameterKind {
2909 match value {
2910 n4m::ParameterKind::Int => HpoNativeParameterKind::Int,
2911 n4m::ParameterKind::Float => HpoNativeParameterKind::Float,
2912 n4m::ParameterKind::LogInt => HpoNativeParameterKind::LogInt,
2913 n4m::ParameterKind::LogFloat => HpoNativeParameterKind::LogFloat,
2914 n4m::ParameterKind::Categorical => HpoNativeParameterKind::Categorical,
2915 n4m::ParameterKind::Ordinal => HpoNativeParameterKind::Ordinal,
2916 n4m::ParameterKind::SortedTuple => HpoNativeParameterKind::SortedTuple,
2917 }
2918 }
2919 fn map_category_type(value: n4m::CategoryType) -> HpoCategoryType {
2920 match value {
2921 n4m::CategoryType::Str => HpoCategoryType::String,
2922 n4m::CategoryType::Int => HpoCategoryType::Integer,
2923 n4m::CategoryType::Float => HpoCategoryType::Float,
2924 n4m::CategoryType::Bool => HpoCategoryType::Boolean,
2925 }
2926 }
2927 fn native_error(operation: &str, error: Error) -> HpoError {
2928 HpoError::Native {
2929 operation: operation.to_string(),
2930 error: to_native_error(error),
2931 }
2932 }
2933 fn to_native_error(error: Error) -> HpoNativeError {
2934 HpoNativeError {
2935 status: error.status,
2936 kind: format!("{:?}", error.kind).to_lowercase(),
2937 retryable: matches!(
2938 error.kind,
2939 ErrorKind::OutOfMemory
2940 | ErrorKind::BackendUnavailable
2941 | ErrorKind::Cancelled
2942 | ErrorKind::Io
2943 ),
2944 message: error.message,
2945 }
2946 }
2947}
2948
2949#[cfg(feature = "methods-optimizer")]
2950use native::MethodsHpoStudy;
2951
2952#[cfg(test)]
2953mod tests {
2954 use super::*;
2955 #[cfg(feature = "methods-optimizer-local")]
2956 use crate::controller::{
2957 ArtifactPolicy, ControllerCapability, ControllerFitScope, ControllerManifest,
2958 ControllerRegistry, RngPolicy,
2959 };
2960 #[cfg(feature = "methods-optimizer-local")]
2961 use crate::graph::{GraphInterface, GraphSpec, NodeKind, NodeSpec, PortSchema};
2962 #[cfg(feature = "methods-optimizer-local")]
2963 use crate::metrics::RegressionMetricKind;
2964 #[cfg(feature = "methods-optimizer-local")]
2965 use crate::phase::Phase;
2966 #[cfg(feature = "methods-optimizer-local")]
2967 use crate::plan::{build_execution_plan, CampaignSpec, ExecutionPlan};
2968 #[cfg(feature = "methods-optimizer-local")]
2969 use crate::runtime::{
2970 ArtifactBackend, ArtifactMaterializationRequest, ArtifactRef, RuntimeController,
2971 RuntimeControllerRegistry, RuntimeHpoExecutionContext, RuntimeHpoIntermediate,
2972 RuntimeHpoIntermediateOutcome, RuntimeHpoProvenance, RuntimeHpoSelectionTarget,
2973 RuntimeHpoTerminal,
2974 };
2975
2976 #[cfg(feature = "methods-optimizer-local")]
2977 fn native_runtime() -> MethodsRuntime {
2978 let library_path = std::env::var_os("N4M_LIBRARY_PATH")
2979 .expect("methods native tests require an explicit N4M_LIBRARY_PATH");
2980 MethodsRuntime::configure(library_path).expect("configure explicit Methods test runtime")
2981 }
2982
2983 #[test]
2984 fn default_build_refuses_before_any_native_or_host_work() {
2985 assert_eq!(
2986 methods_optimizer_preflight(),
2987 if cfg!(feature = "methods-optimizer") {
2988 Ok(())
2989 } else {
2990 Err(HpoError::MethodsOptimizerFeatureDisabled)
2991 }
2992 );
2993 }
2994
2995 #[cfg(feature = "methods-optimizer")]
2996 #[test]
2997 fn methods_runtime_refuses_relative_library_paths_before_native_loading() {
2998 assert!(matches!(
2999 MethodsRuntime::configure("libn4m.so"),
3000 Err(HpoError::RuntimeConfiguration { reason })
3001 if reason == "libn4m path must be absolute"
3002 ));
3003 }
3004
3005 #[cfg(feature = "methods-optimizer-local")]
3006 #[test]
3007 fn methods_runtime_registers_all_controllers_atomically() {
3008 let runtime = native_runtime();
3009 let hpo_id = crate::ControllerId::new("controller:tuner.methods").unwrap();
3010 let mut registry = RuntimeControllerRegistry::new();
3011
3012 register_methods_runtime_controllers(&mut registry, hpo_id.clone(), runtime.clone())
3013 .unwrap();
3014 let pls_id = crate::ControllerId::new(METHODS_PLS_CONTROLLER_ID).unwrap();
3015 let ridge_id = crate::ControllerId::new(METHODS_RIDGE_CONTROLLER_ID).unwrap();
3016 assert!(registry.get(&pls_id).is_some());
3017 assert!(registry.get(&ridge_id).is_some());
3018 assert!(registry.get(&hpo_id).is_some());
3019
3020 let error = register_methods_runtime_controllers(&mut registry, hpo_id.clone(), runtime)
3021 .unwrap_err();
3022 assert!(error.to_string().contains("duplicate runtime controller"));
3023 assert!(registry.get(&pls_id).is_some());
3024 assert!(registry.get(&ridge_id).is_some());
3025 assert!(registry.get(&hpo_id).is_some());
3026 }
3027
3028 #[cfg(feature = "methods-optimizer-local")]
3029 #[test]
3030 fn methods_pls_hydrated_payload_release_is_idempotent_and_handle_local() {
3031 let runtime = native_runtime();
3032 let context = n4m::Context::new().unwrap();
3033 let mut config = n4m::Config::new().unwrap();
3034 config.set_n_components(1).unwrap();
3035 let x_values = [1.0, 1.0, 2.0, 4.0, 3.0, 9.0, 4.0, 16.0];
3036 let y_values = [1.0, 2.0, 3.0, 4.0];
3037 let x = n4m::MatrixRef::row_major(&x_values, 4, 2).unwrap();
3038 let y = n4m::MatrixRef::row_major(&y_values, 4, 1).unwrap();
3039 let payload = n4m::Model::fit(&context, &config, x, y)
3040 .unwrap()
3041 .export_n4mm()
3042 .unwrap();
3043 let controller = MethodsPlsController::new(runtime);
3044 let controller_id = controller.controller_id().clone();
3045 let request = ArtifactMaterializationRequest {
3046 run_id: crate::RunId::new("run:methods-pls.release").unwrap(),
3047 bundle_id: crate::BundleId::new("bundle:methods-pls.release").unwrap(),
3048 node_id: crate::NodeId::new("model:methods-pls").unwrap(),
3049 phase: Phase::Predict,
3050 variant_id: None,
3051 controller_id: controller_id.clone(),
3052 artifact: ArtifactRef {
3053 id: crate::ArtifactId::new("artifact:methods-pls.release").unwrap(),
3054 kind: "n4m_model".to_string(),
3055 controller_id,
3056 backend: Some(ArtifactBackend::Raw),
3057 uri: Some("methods/release.n4mm".to_string()),
3058 content_fingerprint: Some(format!("{:x}", Sha256::digest(&payload))),
3059 size_bytes: Some(payload.len() as u64),
3060 plugin: None,
3061 plugin_version: None,
3062 },
3063 params_fingerprint: "params:methods-pls.release".to_string(),
3064 training_loss_fingerprint: None,
3065 };
3066
3067 let first = controller
3068 .hydrate_artifact_payload(&request, &payload)
3069 .unwrap();
3070 assert_eq!(controller.hydrated_payload_count().unwrap(), 1);
3071 controller
3072 .release_hydrated_artifact_payload(&first)
3073 .unwrap();
3074 controller
3075 .release_hydrated_artifact_payload(&first)
3076 .unwrap();
3077 assert_eq!(controller.hydrated_payload_count().unwrap(), 0);
3078
3079 let second = controller
3080 .hydrate_artifact_payload(&request, &payload)
3081 .unwrap();
3082 assert_ne!(first.handle, second.handle);
3083 controller
3084 .release_hydrated_artifact_payload(&second)
3085 .unwrap();
3086 assert_eq!(controller.hydrated_payload_count().unwrap(), 0);
3087 }
3088
3089 fn ledger_trial(score: f64) -> HpoTrial {
3090 HpoTrial {
3091 id: 7,
3092 ask_sequence: 3,
3093 terminal_sequence: Some(4),
3094 parameters: BTreeMap::new(),
3095 parameter_order: Vec::new(),
3096 status: HpoTrialStatus::Completed,
3097 score: Some(score),
3098 rung: 0,
3099 duration: 0.25,
3100 intermediates: vec![HpoIntermediate {
3101 sequence: 2,
3102 step: 0,
3103 score,
3104 should_prune: false,
3105 }],
3106 failure: None,
3107 }
3108 }
3109
3110 #[test]
3111 fn restored_ledger_accepts_only_one_ulp_score_drift_after_tcv1_projection() {
3112 let native = canonical_hpo_terminal_ledger(vec![ledger_trial(1.0)]).unwrap();
3113 let mut one_ulp = native.clone();
3114 one_ulp[0].score = Some(f64::from_bits(1.0_f64.to_bits() + 1));
3115 one_ulp[0].intermediates[0].score = f64::from_bits(1.0_f64.to_bits() + 1);
3116 let one_ulp = canonical_hpo_terminal_ledger(one_ulp).unwrap();
3117 assert!(hpo_terminal_trials_match(&native, &one_ulp));
3118
3119 let mut two_ulps = native.clone();
3120 two_ulps[0].score = Some(f64::from_bits(1.0_f64.to_bits() + 2));
3121 assert!(!hpo_terminal_trials_match(&native, &two_ulps));
3122
3123 let mut tampered = native.clone();
3124 tampered[0].score = Some(99.0);
3125 assert!(!hpo_terminal_trials_match(&native, &tampered));
3126 }
3127
3128 #[test]
3129 fn search_space_digest_is_canonical_and_order_sensitive() {
3130 let space = HpoSearchSpace {
3131 parameters: vec![HpoParameter::Int {
3132 name: "depth".into(),
3133 low: 1,
3134 high: 5,
3135 step: 1,
3136 log: false,
3137 }],
3138 };
3139 assert_eq!(space.fingerprint().unwrap(), space.fingerprint().unwrap());
3140 let swapped = HpoSearchSpace {
3141 parameters: vec![
3142 HpoParameter::Float {
3143 name: "rate".into(),
3144 low: 0.1,
3145 high: 1.0,
3146 step: 0.1,
3147 log: false,
3148 },
3149 space.parameters[0].clone(),
3150 ],
3151 };
3152 assert_ne!(space.fingerprint().unwrap(), swapped.fingerprint().unwrap());
3153 }
3154 #[test]
3155 fn checkpoint_rejects_oversize_before_native_decoder() {
3156 let binding = HpoStudyBinding {
3157 controller_id: "controller:hpo".into(),
3158 study_id: "study:one".into(),
3159 search_space_fingerprint: "a".into(),
3160 optimizer_fingerprint: "b".into(),
3161 };
3162 let checkpoint = N4moptCheckpointArtifact {
3163 schema_version: 1,
3164 artifact_kind: N4MOPT_ARTIFACT_KIND.into(),
3165 format: N4MOPT_FORMAT.into(),
3166 binding,
3167 methods_abi: "n4m-abi-2.2".into(),
3168 opaque_payload: vec![0; MAX_N4MOPT_CHECKPOINT_BYTES + 1],
3169 payload_sha256: "x".into(),
3170 };
3171 assert!(matches!(
3172 checkpoint.validate(),
3173 Err(HpoError::InvalidCheckpoint { .. })
3174 ));
3175 }
3176
3177 #[cfg(feature = "methods-optimizer-local")]
3178 fn native_config() -> MethodsHpoStudyConfig {
3179 MethodsHpoStudyConfig {
3180 controller_id: "controller:methods-hpo".into(),
3181 study_id: "study:native-lifecycle".into(),
3182 methods_abi: "n4m-abi-2.2".into(),
3183 search_space: HpoSearchSpace {
3184 parameters: vec![HpoParameter::Int {
3185 name: "n_components".into(),
3186 low: 1,
3187 high: 3,
3188 step: 1,
3189 log: false,
3190 }],
3191 },
3192 optimizer: HpoOptimizerConfig {
3193 sampler: HpoSampler::Random,
3194 pruner: HpoPruner::None,
3195 direction: HpoDirection::Minimize,
3196 metric: HpoMetric::Rmse,
3197 seed: 7,
3198 n_startup_trials: 1,
3199 max_resource: 0,
3200 reduction_factor: 0,
3201 },
3202 }
3203 }
3204
3205 #[cfg(feature = "methods-optimizer-local")]
3206 fn native_hpo_manifest(id: &str, kind: NodeKind) -> ControllerManifest {
3207 ControllerManifest {
3208 controller_id: crate::ControllerId::new(id).unwrap(),
3209 controller_version: "native-hpo-test".to_string(),
3210 operator_kind: kind,
3211 priority: 0,
3212 supported_phases: BTreeSet::from([Phase::FitCv]),
3213 input_ports: Vec::new(),
3214 output_ports: Vec::new(),
3215 data_requirements: None,
3216 capabilities: BTreeSet::from([ControllerCapability::Deterministic]),
3217 operator_selectors: Vec::new(),
3218 fit_scope: ControllerFitScope::FoldTrain,
3219 rng_policy: RngPolicy::UsesCoreSeed,
3220 artifact_policy: ArtifactPolicy::Serializable,
3221 }
3222 }
3223
3224 #[cfg(feature = "methods-optimizer-local")]
3225 fn native_hpo_node(id: &str, kind: NodeKind) -> NodeSpec {
3226 NodeSpec {
3227 id: crate::NodeId::new(id).unwrap(),
3228 kind,
3229 operator: None,
3230 params: BTreeMap::new(),
3231 ports: PortSchema {
3232 inputs: Vec::new(),
3233 outputs: Vec::new(),
3234 },
3235 metadata: BTreeMap::new(),
3236 seed_label: None,
3237 }
3238 }
3239
3240 #[cfg(feature = "methods-optimizer-local")]
3241 fn attested_native_hpo_context() -> (
3242 ExecutionPlan,
3243 RuntimeHpoExecutionContext,
3244 crate::runtime::RuntimeHpoCampaignTask,
3245 ) {
3246 let target_node_id = crate::NodeId::new("model:methods-pls").unwrap();
3247 let controller_id = crate::ControllerId::new("controller:methods-hpo").unwrap();
3248 let mut registry = ControllerRegistry::new();
3249 registry
3250 .register(native_hpo_manifest(
3251 "controller:methods-pls",
3252 NodeKind::Model,
3253 ))
3254 .unwrap();
3255 let plan = build_execution_plan(
3256 "plan:methods-hpo-checkpoint",
3257 GraphSpec {
3258 id: "graph:methods-hpo-checkpoint".to_string(),
3259 interface: GraphInterface::default(),
3260 nodes: vec![native_hpo_node("model:methods-pls", NodeKind::Model)],
3261 edges: Vec::new(),
3262 search_space_fingerprint: None,
3263 metadata: BTreeMap::new(),
3264 },
3265 CampaignSpec {
3266 inner_cv: None,
3267 id: "campaign:methods-hpo-checkpoint".to_string(),
3268 root_seed: Some(17),
3269 leakage_policy: Default::default(),
3270 aggregation_policy: Default::default(),
3271 split_invocation: None,
3272 generation: Default::default(),
3273 shape_plans: BTreeMap::new(),
3274 data_bindings: BTreeMap::new(),
3275 branch_view_plans: Vec::new(),
3276 metadata: BTreeMap::new(),
3277 },
3278 ®istry,
3279 )
3280 .unwrap();
3281 let context = RuntimeHpoExecutionContext {
3282 operation_id: "hpo:methods".to_string(),
3283 controller_id: controller_id.clone(),
3284 target_node_id: target_node_id.clone(),
3285 base_variant: plan.variants[0].clone(),
3286 trial_budget_total: 2,
3287 study: native_config(),
3288 parameter_paths: BTreeMap::from([(
3289 "n_components".to_string(),
3290 "n_components".to_string(),
3291 )]),
3292 resume_checkpoint: None,
3293 resume_variants: BTreeMap::new(),
3294 resume_terminal_trials: Vec::new(),
3295 selection: RuntimeHpoSelectionTarget {
3296 producer_node: target_node_id.clone(),
3297 producer_port: "prediction".to_string(),
3298 metric: RegressionMetricKind::Rmse,
3299 direction: HpoDirection::Minimize,
3300 },
3301 provenance: RuntimeHpoProvenance {
3302 graph_fingerprint: plan.graph_fingerprint.clone(),
3303 campaign_fingerprint: plan.campaign_fingerprint.clone(),
3304 controller_fingerprint: plan.controller_fingerprint.clone(),
3305 data_identities_fingerprint: "data-identities:methods-hpo".to_string(),
3306 fold_set_fingerprint: None,
3307 training_influence_fingerprint: "influence:methods-hpo".to_string(),
3308 relation_fingerprint: "relations:methods-hpo".to_string(),
3309 },
3310 };
3311 context.validate_for_plan(&plan).unwrap();
3312 let task = crate::runtime::RuntimeHpoCampaignTask {
3313 run_id: crate::RunId::new("run:methods-hpo-checkpoint").unwrap(),
3314 operation_id: "hpo:methods".to_string(),
3315 controller_id: controller_id.clone(),
3316 target_node_id,
3317 seed: Some(17),
3318 };
3319 assert_eq!(task.controller_id, controller_id);
3320 (plan, context, task)
3321 }
3322
3323 #[cfg(feature = "methods-optimizer-local")]
3324 fn proposal_components(proposal: &crate::runtime::RuntimeHpoProposal) -> i64 {
3325 let choice = proposal.variant.choices.get("native_methods_hpo").unwrap();
3326 let override_ = choice.param_overrides.first().unwrap();
3327 assert_eq!(override_.params.len(), 1);
3328 override_.params["n_components"].as_i64().unwrap()
3329 }
3330
3331 #[cfg(feature = "methods-optimizer-local")]
3332 fn assert_runtime_refusal(error: crate::DagMlError) {
3333 assert!(matches!(error, crate::DagMlError::RuntimeValidation(_)));
3334 }
3335
3336 #[cfg(feature = "methods-optimizer-local")]
3337 #[test]
3338 fn registered_methods_session_checkpoints_restores_and_refuses_tampering() {
3339 let runtime = native_runtime();
3340 let (plan, context, task) = attested_native_hpo_context();
3341 let controller_id = task.controller_id.clone();
3342 let mut controllers = RuntimeControllerRegistry::new();
3343 controllers
3344 .register(Box::new(MethodsHpoController::new(
3345 controller_id.clone(),
3346 runtime,
3347 )))
3348 .unwrap();
3349 let controller = controllers.get(&controller_id).unwrap();
3350
3351 let mut session = controller.create_tuner_session(&task, &context).unwrap();
3352 let first = session.ask().unwrap().unwrap();
3353 assert!((1..=3).contains(&proposal_components(&first)));
3354 assert_eq!(
3355 session
3356 .report_intermediate(RuntimeHpoIntermediate {
3357 trial_id: first.trial_id,
3358 step: 0,
3359 score: 1.5,
3360 })
3361 .unwrap(),
3362 RuntimeHpoIntermediateOutcome::Continue
3363 );
3364 session
3365 .tell(first.trial_id, RuntimeHpoTerminal::Completed { score: 1.0 })
3366 .unwrap();
3367 let checkpoint = session.checkpoint().unwrap();
3368 checkpoint.validate().unwrap();
3369 assert_eq!(checkpoint.binding.controller_id, controller_id.as_str());
3370 assert_eq!(checkpoint.binding.study_id, context.study.study_id);
3371 assert_eq!(checkpoint.methods_abi, context.study.methods_abi);
3372
3373 let checkpoint_trace = MethodsHpoStudy::restore(context.study.clone(), &checkpoint)
3376 .unwrap()
3377 .trials()
3378 .unwrap();
3379 assert_eq!(checkpoint_trace.len(), 1);
3380 assert_eq!(checkpoint_trace[0].id, first.trial_id);
3381 assert_eq!(checkpoint_trace[0].status, HpoTrialStatus::Completed);
3382 assert_eq!(checkpoint_trace[0].score, Some(1.0));
3383 assert_eq!(
3384 MethodsHpoStudy::restore(context.study.clone(), &checkpoint)
3385 .unwrap()
3386 .best()
3387 .unwrap()
3388 .unwrap()
3389 .trial
3390 .id,
3391 first.trial_id
3392 );
3393
3394 let mut expected = MethodsHpoStudy::restore(context.study.clone(), &checkpoint).unwrap();
3395 assert_eq!(expected.trials().unwrap(), checkpoint_trace);
3396 assert_eq!(expected.best().unwrap().unwrap().trial.id, first.trial_id);
3397 let expected_next = expected.ask().unwrap();
3398 let mut resumed_context = context.clone();
3399 resumed_context.resume_checkpoint = Some(checkpoint.clone());
3400 resumed_context.resume_terminal_trials = vec![crate::runtime::RuntimeHpoTerminalSnapshot {
3401 trial: checkpoint_trace[0].clone(),
3402 variant_id: Some(first.variant.variant_id.clone()),
3403 }];
3404 resumed_context.validate_for_plan(&plan).unwrap();
3405
3406 let mut tampered_ledger = resumed_context.clone();
3411 tampered_ledger.resume_terminal_trials[0].trial.score = Some(99.0);
3412 let error = match controller.create_tuner_session(&task, &tampered_ledger) {
3413 Err(error) => error,
3414 Ok(_) => panic!("tampered restored terminal ledger unexpectedly created a session"),
3415 };
3416 assert_runtime_refusal(error);
3417
3418 let mut resumed = controller
3419 .create_tuner_session(&task, &resumed_context)
3420 .unwrap();
3421 let resumed_next = resumed.ask().unwrap().unwrap();
3422 assert_eq!(resumed_next.trial_id, expected_next.id);
3423 assert_eq!(
3424 proposal_components(&resumed_next),
3425 expected_next.parameters["n_components"].value as i64
3426 );
3427 resumed
3428 .report_intermediate(RuntimeHpoIntermediate {
3429 trial_id: resumed_next.trial_id,
3430 step: 0,
3431 score: 0.5,
3432 })
3433 .unwrap();
3434 resumed
3435 .tell(
3436 resumed_next.trial_id,
3437 RuntimeHpoTerminal::Completed { score: 0.25 },
3438 )
3439 .unwrap();
3440 let resumed_checkpoint = resumed.checkpoint().unwrap();
3441 let resumed_trace = MethodsHpoStudy::restore(context.study.clone(), &resumed_checkpoint)
3442 .unwrap()
3443 .trials()
3444 .unwrap();
3445 assert_eq!(resumed_trace.len(), 2);
3446 assert_eq!(resumed_trace[1].id, resumed_next.trial_id);
3447 assert_eq!(resumed_trace[1].status, HpoTrialStatus::Completed);
3448 assert_eq!(resumed_trace[1].score, Some(0.25));
3449 assert_eq!(
3450 MethodsHpoStudy::restore(context.study.clone(), &resumed_checkpoint)
3451 .unwrap()
3452 .best()
3453 .unwrap()
3454 .unwrap()
3455 .trial
3456 .id,
3457 resumed_next.trial_id
3458 );
3459
3460 let mut wrong_abi = resumed_context.clone();
3461 wrong_abi.study.methods_abi = "n4m-abi-wrong".to_string();
3462 wrong_abi.validate_for_plan(&plan).unwrap();
3463 let error = match controller.create_tuner_session(&task, &wrong_abi) {
3464 Err(error) => error,
3465 Ok(_) => panic!("mismatched Methods ABI unexpectedly restored a session"),
3466 };
3467 assert_runtime_refusal(error);
3468
3469 let mut wrong_binding = resumed_context.clone();
3470 wrong_binding
3471 .resume_checkpoint
3472 .as_mut()
3473 .unwrap()
3474 .binding
3475 .study_id = "study:wrong-binding".to_string();
3476 wrong_binding.validate_for_plan(&plan).unwrap();
3477 let error = match controller.create_tuner_session(&task, &wrong_binding) {
3478 Err(error) => error,
3479 Ok(_) => panic!("mismatched checkpoint binding unexpectedly restored a session"),
3480 };
3481 assert_runtime_refusal(error);
3482
3483 let mut wrong_checksum = resumed_context;
3484 wrong_checksum
3485 .resume_checkpoint
3486 .as_mut()
3487 .unwrap()
3488 .opaque_payload[0] ^= 1;
3489 assert_runtime_refusal(wrong_checksum.validate_for_plan(&plan).unwrap_err());
3490 let error = match controller.create_tuner_session(&task, &wrong_checksum) {
3491 Err(error) => error,
3492 Ok(_) => panic!("bad checkpoint checksum unexpectedly restored a session"),
3493 };
3494 assert_runtime_refusal(error);
3495 }
3496
3497 #[cfg(feature = "methods-optimizer-local")]
3498 #[test]
3499 fn real_n4m_lifecycle_batch_trials_best_and_checkpoint() {
3500 let _runtime = native_runtime();
3501 let config = native_config();
3502 let mut study = MethodsHpoStudy::create(config.clone()).unwrap();
3503 let batch = study.ask_batch(2).unwrap();
3504 assert_eq!(batch.trials.len(), 2);
3505 assert!(batch.native_error.is_none());
3506 assert!(batch.trials.iter().all(|trial| trial.id >= 0));
3507 assert_eq!(batch.trials[0].parameter_order, vec!["n_components"]);
3508
3509 study
3510 .report_intermediate(batch.trials[0].id, 0, 2.0)
3511 .unwrap();
3512 study
3513 .tell(batch.trials[0].id, HpoTerminal::Completed { score: 1.0 })
3514 .unwrap();
3515 study
3516 .tell(
3517 batch.trials[1].id,
3518 HpoTerminal::Failed {
3519 failure: HpoFailure {
3520 code: "EVALUATION_FAILED".into(),
3521 message: "controlled test failure".into(),
3522 retryable: true,
3523 },
3524 },
3525 )
3526 .unwrap();
3527
3528 let trials = study.trials().unwrap();
3529 assert_eq!(trials.len(), 2);
3530 assert_eq!(trials[0].status, HpoTrialStatus::Completed);
3531 assert_eq!(trials[0].score, Some(1.0));
3532 assert_eq!(trials[1].status, HpoTrialStatus::Failed);
3533 assert!(trials[1].failure.as_ref().unwrap().retryable);
3534 assert_eq!(study.best().unwrap().unwrap().score, 1.0);
3535 assert!(study
3536 .events()
3537 .iter()
3538 .any(|event| matches!(event, HpoEvent::Intermediate { step: 0, .. })));
3539
3540 let checkpoint = study.save_checkpoint().unwrap();
3541 let restored = MethodsHpoStudy::restore(config, &checkpoint).unwrap();
3542 assert_eq!(restored.trials().unwrap().len(), 2);
3543 }
3544
3545 #[cfg(feature = "methods-optimizer-local")]
3546 #[test]
3547 fn real_tpe_pruner_failure_trace_and_checkpoint_resume_are_native() {
3548 let _runtime = native_runtime();
3549 let mut config = native_config();
3550 config.optimizer.sampler = HpoSampler::Tpe;
3551 config.optimizer.pruner = HpoPruner::Median;
3552 config.optimizer.n_startup_trials = 2;
3553 config.optimizer.seed = 51;
3554 let mut study = MethodsHpoStudy::create(config.clone()).unwrap();
3555
3556 let failed = study.ask().unwrap();
3560 let failed = study
3561 .tell(
3562 failed.id,
3563 HpoTerminal::Failed {
3564 failure: HpoFailure {
3565 code: "CV_PROVIDER_FAILURE".into(),
3566 message: "controlled fold materialization failure".into(),
3567 retryable: false,
3568 },
3569 },
3570 )
3571 .unwrap();
3572 assert_eq!(failed.status, HpoTrialStatus::Failed);
3573 assert_eq!(failed.failure.unwrap().code, "CV_PROVIDER_FAILURE");
3574
3575 let first = study.ask().unwrap();
3580 assert!(!study.report_intermediate(first.id, 0, 1.0).unwrap());
3581 let second = study.ask().unwrap();
3582 assert!(!study.report_intermediate(second.id, 0, 2.0).unwrap());
3583 let third = study.ask().unwrap();
3584 assert!(study.report_intermediate(third.id, 0, 9.0).unwrap());
3585
3586 let trials = study.trials().unwrap();
3587 let pruned = trials.iter().find(|trial| trial.id == third.id).unwrap();
3588 assert_eq!(pruned.status, HpoTrialStatus::Pruned);
3589 assert!(pruned.terminal_sequence.is_some());
3590 assert!(pruned
3591 .intermediates
3592 .iter()
3593 .any(|item| item.step == 0 && item.score == 9.0 && item.should_prune));
3594 assert!(study.events().iter().any(|event| {
3595 matches!(event, HpoEvent::Terminal { trial_id, status: HpoTrialStatus::Failed, .. } if *trial_id == failed.id)
3596 }));
3597 assert!(study.events().iter().any(|event| {
3598 matches!(event, HpoEvent::Terminal { trial_id, status: HpoTrialStatus::Pruned, .. } if *trial_id == third.id)
3599 }));
3600
3601 let checkpoint = study.save_checkpoint().unwrap();
3602 let checkpoint: N4moptCheckpointArtifact =
3606 serde_json::from_str(&serde_json::to_string(&checkpoint).unwrap()).unwrap();
3607 let mut resumed = MethodsHpoStudy::restore(config, &checkpoint).unwrap();
3608 for _ in 0..4 {
3609 let uninterrupted = study.ask().unwrap();
3610 let restored = resumed.ask().unwrap();
3611 assert_eq!(uninterrupted.id, restored.id);
3612 assert_eq!(uninterrupted.parameter_order, restored.parameter_order);
3613 assert_eq!(uninterrupted.parameters, restored.parameters);
3614 }
3615 }
3616
3617 #[cfg(feature = "methods-optimizer-local")]
3618 #[test]
3619 fn malformed_checkpoint_reaches_native_n4mopt_decoder_as_typed_error() {
3620 let _runtime = native_runtime();
3621 let config = native_config();
3622 let study = MethodsHpoStudy::create(config.clone()).unwrap();
3623 let mut checkpoint = study.save_checkpoint().unwrap();
3624 checkpoint.opaque_payload[0] ^= 1;
3625 checkpoint.payload_sha256 = payload_sha256(&checkpoint.opaque_payload);
3626 assert!(matches!(
3627 MethodsHpoStudy::restore(config, &checkpoint),
3628 Err(HpoError::Native { operation, .. }) if operation == "load_n4mopt"
3629 ));
3630 }
3631}