1use crate::core::manifest::{Manifest, ManifestComponent};
31use crate::core::plan::PlanItem;
32use crate::core::types::{ExecutorKind, ValidationTier};
33use serde::{Deserialize, Serialize};
34use std::collections::{HashMap, HashSet};
35use std::fmt;
36
37#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
43#[serde(rename_all = "lowercase")]
44pub enum UpdateClassification {
45 Safe,
47 Guarded,
49 Blocked,
51 Candidate,
53 Experimental,
55}
56
57impl UpdateClassification {
58 pub fn label(self) -> &'static str {
60 match self {
61 UpdateClassification::Safe => "safe",
62 UpdateClassification::Guarded => "guarded",
63 UpdateClassification::Blocked => "blocked",
64 UpdateClassification::Candidate => "candidate",
65 UpdateClassification::Experimental => "experimental",
66 }
67 }
68
69 pub fn is_visible(self) -> bool {
71 !matches!(
72 self,
73 UpdateClassification::Blocked | UpdateClassification::Experimental
74 )
75 }
76
77 pub fn is_preselected(self) -> bool {
79 matches!(self, UpdateClassification::Safe)
80 }
81}
82
83impl fmt::Display for UpdateClassification {
84 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
85 f.write_str(self.label())
86 }
87}
88
89#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
98pub struct PlannerItem {
99 #[serde(flatten)]
101 pub plan_item: PlanItem,
102 pub classification: UpdateClassification,
104 pub visible: bool,
106 pub selected: bool,
108 pub classification_reason: String,
110 pub requires_hardware_check: bool,
112 pub min_rocm_version: String,
114}
115
116#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)]
122pub struct CompatibilityContext {
123 pub rocm_version: String,
125 pub rocm_channel: String,
127 pub available_executors: HashSet<ExecutorKind>,
129 pub installed_components: HashSet<String>,
131 pub installed_versions: HashMap<String, String>,
133 pub gpu_architecture: String,
135 pub runtime_version: String,
137}
138
139impl CompatibilityContext {
140 pub fn new() -> Self {
142 Self::default()
143 }
144
145 pub fn is_installed(&self, component_id: &str) -> bool {
147 self.installed_components.contains(component_id)
148 }
149
150 pub fn installed_version(&self, component_id: &str) -> Option<&str> {
152 self.installed_versions
153 .get(component_id)
154 .map(|s| s.as_str())
155 }
156}
157
158#[derive(Debug, Clone, Default)]
164pub struct PlannerOptions {
165 pub scan_only: bool,
167 pub all_safe: bool,
169 pub include_experimental: bool,
171 pub target_components: Vec<String>,
173 pub json_output: bool,
175}
176
177#[derive(Debug, Clone, PartialEq)]
183pub enum PlannerError {
184 ComponentBlocked {
186 component_id: String,
187 reason: String,
188 },
189 ComponentUnknown { component_id: String },
191 RuntimeTooOld { current: String, required: String },
193 DependencyConflict {
195 component_id: String,
196 dependency_id: String,
197 reason: String,
198 },
199}
200
201impl fmt::Display for PlannerError {
202 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
203 match self {
204 PlannerError::ComponentBlocked {
205 component_id,
206 reason,
207 } => {
208 write!(f, "component '{component_id}' is blocked: {reason}")
209 }
210 PlannerError::ComponentUnknown { component_id } => {
211 write!(f, "unknown component: '{component_id}'")
212 }
213 PlannerError::RuntimeTooOld { current, required } => {
214 write!(
215 f,
216 "runtime version {current} is too old; requires {required}. Run `rusty upgrade` first."
217 )
218 }
219 PlannerError::DependencyConflict {
220 component_id,
221 dependency_id,
222 reason,
223 } => {
224 write!(
225 f,
226 "dependency conflict for '{component_id}' on '{dependency_id}': {reason}"
227 )
228 }
229 }
230 }
231}
232
233impl std::error::Error for PlannerError {}
234
235pub struct UpdatePlanner;
249
250impl UpdatePlanner {
251 pub fn new() -> Self {
253 Self
254 }
255
256 pub fn build_plan(
261 &self,
262 manifest: &Manifest,
263 context: &CompatibilityContext,
264 options: &PlannerOptions,
265 ) -> Result<Vec<PlannerItem>, PlannerError> {
266 self.check_runtime_version(manifest, context)?;
268
269 if !options.target_components.is_empty() {
270 for target in &options.target_components {
271 if !manifest.components.iter().any(|c| c.id == *target) {
272 return Err(PlannerError::ComponentUnknown {
273 component_id: target.clone(),
274 });
275 }
276 }
277 }
278
279 let target_set: HashSet<&str> = options
280 .target_components
281 .iter()
282 .map(|s| s.as_str())
283 .collect();
284 let required_deps = self.collect_required_dependencies(manifest, &target_set);
285 let rocm_available_or_planned = !context.rocm_version.is_empty()
286 || target_set.contains("rocm")
287 || required_deps.contains("rocm");
288 let planned_rocm_version = manifest
289 .components
290 .iter()
291 .find(|c| {
292 c.id == "rocm" && (target_set.contains("rocm") || required_deps.contains("rocm"))
293 })
294 .map(|c| c.version.as_str());
295
296 let mut items = Vec::new();
297
298 for component in &manifest.components {
299 if !options.target_components.is_empty() {
301 if !target_set.contains(component.id.as_str())
302 && !required_deps.contains(component.id.as_str())
303 {
304 continue;
305 }
306 }
307
308 let classification = self.classify_update(
310 component,
311 context,
312 rocm_available_or_planned,
313 planned_rocm_version,
314 );
315
316 if classification == UpdateClassification::Blocked {
318 if target_set.contains(component.id.as_str())
320 || required_deps.contains(component.id.as_str())
321 {
322 let reason = self.block_reason(
323 component,
324 context,
325 rocm_available_or_planned,
326 planned_rocm_version,
327 );
328 return Err(PlannerError::ComponentBlocked {
329 component_id: component.id.clone(),
330 reason,
331 });
332 }
333 items.push(self.build_planner_item(
335 component,
336 classification,
337 context,
338 rocm_available_or_planned,
339 planned_rocm_version,
340 ));
341 continue;
342 }
343
344 if classification == UpdateClassification::Experimental && !options.include_experimental
346 {
347 continue;
348 }
349
350 items.push(self.build_planner_item(
351 component,
352 classification,
353 context,
354 rocm_available_or_planned,
355 planned_rocm_version,
356 ));
357 }
358
359 if options.all_safe {
362 if options.include_experimental {
363 items.retain(|i| {
364 if target_set.contains(i.plan_item.component_id.as_str()) {
365 return true;
366 }
367 matches!(
368 i.classification,
369 UpdateClassification::Safe | UpdateClassification::Experimental
370 )
371 });
372 } else {
373 items.retain(|i| {
374 target_set.contains(i.plan_item.component_id.as_str())
375 || i.classification == UpdateClassification::Safe
376 });
377 }
378 }
379
380 if !target_set.is_empty() {
382 for item in &mut items {
383 if target_set.contains(item.plan_item.component_id.as_str()) {
384 item.selected = true;
385 item.plan_item.selected = true;
386 }
387 }
388 }
389
390 self.enforce_dependency_rules(&mut items)?;
392
393 Ok(items)
394 }
395
396 fn classify_update(
398 &self,
399 component: &ManifestComponent,
400 context: &CompatibilityContext,
401 rocm_available_or_planned: bool,
402 planned_rocm_version: Option<&str>,
403 ) -> UpdateClassification {
404 if component.validation_tier == ValidationTier::Experimental {
406 return UpdateClassification::Experimental;
407 }
408
409 if component.validation_tier == ValidationTier::Blocked {
411 return UpdateClassification::Blocked;
412 }
413
414 if !self.is_hardware_compatible(component, context) {
416 return UpdateClassification::Blocked;
417 }
418
419 if !self.is_executor_available(component, context) {
421 return UpdateClassification::Blocked;
422 }
423
424 if !self.is_rocm_compatible(
426 component,
427 context,
428 rocm_available_or_planned,
429 planned_rocm_version,
430 ) {
431 return UpdateClassification::Blocked;
432 }
433
434 let current_version = context.installed_version(&component.id).unwrap_or("");
436
437 if current_version.is_empty() {
439 return UpdateClassification::Candidate;
440 }
441
442 if current_version == component.version {
444 return UpdateClassification::Safe;
445 }
446
447 match self.version_bump_level(current_version, &component.version) {
449 BumpLevel::Patch => {
450 if component.validation_tier == ValidationTier::Validated {
452 UpdateClassification::Safe
453 } else {
454 UpdateClassification::Guarded
455 }
456 }
457 BumpLevel::Minor => UpdateClassification::Guarded,
458 BumpLevel::Major => UpdateClassification::Candidate,
459 BumpLevel::Unknown => {
460 UpdateClassification::Guarded
462 }
463 }
464 }
465
466 fn build_planner_item(
468 &self,
469 component: &ManifestComponent,
470 classification: UpdateClassification,
471 context: &CompatibilityContext,
472 rocm_available_or_planned: bool,
473 planned_rocm_version: Option<&str>,
474 ) -> PlannerItem {
475 let current_version = context
476 .installed_version(&component.id)
477 .unwrap_or("")
478 .to_string();
479
480 let visible = classification.is_visible();
481 let selected = classification.is_preselected() && current_version != component.version;
482 let isolation_safe = matches!(classification, UpdateClassification::Safe);
483
484 let classification_reason = self.classification_reason(
485 component,
486 classification,
487 context,
488 rocm_available_or_planned,
489 planned_rocm_version,
490 );
491
492 let dependencies = self.dependencies_for_component(component);
493
494 PlannerItem {
495 plan_item: PlanItem::new(
496 &component.id,
497 ¤t_version,
498 &component.version,
499 component.validation_tier,
500 selected,
501 &classification_reason,
502 dependencies,
503 isolation_safe,
504 ),
505 classification,
506 visible,
507 selected,
508 classification_reason,
509 requires_hardware_check: self.requires_hardware_check(&component.id),
510 min_rocm_version: component.min_rocm_version.clone(),
511 }
512 }
513
514 fn is_hardware_compatible(
516 &self,
517 component: &ManifestComponent,
518 context: &CompatibilityContext,
519 ) -> bool {
520 if !component.compatible_channels.is_empty()
522 && self.requires_rocm(&component.id)
523 && !component
524 .compatible_channels
525 .iter()
526 .any(|ch| ch.eq_ignore_ascii_case(&context.rocm_channel))
527 {
528 return false;
529 }
530 true
531 }
532
533 fn is_executor_available(
535 &self,
536 _component: &ManifestComponent,
537 context: &CompatibilityContext,
538 ) -> bool {
539 context
541 .available_executors
542 .contains(&ExecutorKind::LegacyScript)
543 || context.available_executors.contains(&ExecutorKind::Rust)
544 }
545
546 fn is_rocm_compatible(
548 &self,
549 component: &ManifestComponent,
550 context: &CompatibilityContext,
551 rocm_available_or_planned: bool,
552 planned_rocm_version: Option<&str>,
553 ) -> bool {
554 if !self.requires_rocm(&component.id) {
556 return true;
557 }
558
559 if component.id == "rocm" {
561 return true;
562 }
563
564 if context.rocm_version.is_empty() && !rocm_available_or_planned {
566 return false;
567 }
568
569 let effective_rocm_version = self.effective_rocm_version(context, planned_rocm_version);
571 if !component.min_rocm_version.is_empty()
572 && !effective_rocm_version.is_empty()
573 && !version_gte(&effective_rocm_version, &component.min_rocm_version)
574 {
575 return false;
576 }
577
578 if !component.compatible_channels.is_empty()
580 && !component
581 .compatible_channels
582 .iter()
583 .any(|ch| ch.eq_ignore_ascii_case(&context.rocm_channel))
584 {
585 return false;
586 }
587
588 true
589 }
590
591 fn requires_rocm(&self, component_id: &str) -> bool {
593 matches!(
594 component_id,
595 "rocm"
596 | "rocm-smi"
597 | "pytorch"
598 | "triton"
599 | "onnx"
600 | "migraphx"
601 | "flash-attn"
602 | "rccl"
603 | "vllm"
604 | "aiter"
605 | "deepspeed"
606 | "bitsandbytes"
607 | "llama-cpp"
608 )
609 }
610
611 fn requires_hardware_check(&self, component_id: &str) -> bool {
613 self.requires_rocm(component_id)
614 }
615
616 fn block_reason(
618 &self,
619 component: &ManifestComponent,
620 context: &CompatibilityContext,
621 rocm_available_or_planned: bool,
622 planned_rocm_version: Option<&str>,
623 ) -> String {
624 if component.validation_tier == ValidationTier::Blocked {
625 return "component is marked as blocked in manifest".to_string();
626 }
627
628 if !self.is_hardware_compatible(component, context) {
629 return format!(
630 "incompatible with ROCm channel '{}' (requires: {})",
631 context.rocm_channel,
632 component.compatible_channels.join(", ")
633 );
634 }
635
636 if !self.is_executor_available(component, context) {
637 return "no available executor for this platform".to_string();
638 }
639
640 if !self.is_rocm_compatible(
641 component,
642 context,
643 rocm_available_or_planned,
644 planned_rocm_version,
645 ) {
646 if context.rocm_version.is_empty() && !rocm_available_or_planned {
647 return "requires ROCm but ROCm is not installed".to_string();
648 }
649 let effective_rocm_version = self.effective_rocm_version(context, planned_rocm_version);
650 if !component.min_rocm_version.is_empty() {
651 if effective_rocm_version.is_empty() {
652 return format!(
653 "requires ROCm >= {} and will rely on the planned ROCm install",
654 component.min_rocm_version
655 );
656 }
657 return format!(
658 "requires ROCm >= {} but effective plan version is {}",
659 component.min_rocm_version, effective_rocm_version
660 );
661 }
662 return format!(
663 "incompatible with effective ROCm version ({})",
664 effective_rocm_version
665 );
666 }
667
668 "unknown reason".to_string()
669 }
670
671 fn classification_reason(
673 &self,
674 component: &ManifestComponent,
675 classification: UpdateClassification,
676 context: &CompatibilityContext,
677 rocm_available_or_planned: bool,
678 planned_rocm_version: Option<&str>,
679 ) -> String {
680 let current = context.installed_version(&component.id).unwrap_or("");
681
682 match classification {
683 UpdateClassification::Safe => {
684 if current == component.version {
685 format!("already up to date (v{})", component.version)
686 } else {
687 format!(
688 "patch update {} → {} (validated, no new deps)",
689 current, component.version
690 )
691 }
692 }
693 UpdateClassification::Guarded => {
694 format!(
695 "minor update {} → {} (new dependencies or config changes)",
696 current, component.version
697 )
698 }
699 UpdateClassification::Blocked => self.block_reason(
700 component,
701 context,
702 rocm_available_or_planned,
703 planned_rocm_version,
704 ),
705 UpdateClassification::Candidate => {
706 if current.is_empty() {
707 format!("new install v{}", component.version)
708 } else {
709 format!(
710 "major update {} → {} (compatible)",
711 current, component.version
712 )
713 }
714 }
715 UpdateClassification::Experimental => {
716 format!("experimental component v{}", component.version)
717 }
718 }
719 }
720
721 fn dependencies_for_component(&self, component: &ManifestComponent) -> Vec<String> {
723 if !component.dependencies.is_empty() {
724 return component.dependencies.clone();
725 }
726 self.derive_dependencies(&component.id)
727 }
728
729 fn derive_dependencies(&self, component_id: &str) -> Vec<String> {
731 match component_id {
732 "pytorch" => vec!["rocm".to_string()],
733 "triton" => vec!["pytorch".to_string()],
734 "flash-attn" => vec!["pytorch".to_string()],
735 "deepspeed" => vec!["pytorch".to_string()],
736 "vllm" => vec!["pytorch".to_string()],
737 "megatron" => vec!["pytorch".to_string(), "mpi4py".to_string()],
738 "onnx" => vec!["rocm".to_string()],
739 "migraphx" => vec!["rocm".to_string()],
740 "bitsandbytes" => vec!["pytorch".to_string()],
741 "aiter" => vec!["pytorch".to_string()],
742 "mpi4py" => vec![],
743 "wandb" => vec![],
744 "comfyui" => vec!["pytorch".to_string()],
745 "vllm-studio" => vec!["vllm".to_string()],
746 "textgen" => vec!["pytorch".to_string()],
747 "rocm" => vec![],
748 "rocm-smi" => vec!["rocm".to_string()],
749 "llama-cpp" => vec!["rocm".to_string()],
750 "permanent-env" => vec![],
751 _ => vec![],
752 }
753 }
754
755 fn collect_required_dependencies<'a>(
756 &self,
757 manifest: &Manifest,
758 targets: &HashSet<&'a str>,
759 ) -> HashSet<String> {
760 let mut required = HashSet::new();
761 let mut stack: Vec<String> = targets.iter().map(|s| (*s).to_string()).collect();
762
763 while let Some(component) = stack.pop() {
764 let deps = manifest
765 .components
766 .iter()
767 .find(|c| c.id == component)
768 .map(|c| self.dependencies_for_component(c))
769 .unwrap_or_else(|| self.derive_dependencies(&component));
770 for dep in deps {
771 if required.insert(dep.clone()) {
772 stack.push(dep);
773 }
774 }
775 }
776
777 for target in targets {
778 required.remove(*target);
779 }
780
781 required
782 }
783
784 fn effective_rocm_version(
785 &self,
786 context: &CompatibilityContext,
787 planned_rocm_version: Option<&str>,
788 ) -> String {
789 match planned_rocm_version {
790 Some(planned) if context.rocm_version.is_empty() => planned.to_string(),
791 Some(planned) if version_gte(planned, &context.rocm_version) => planned.to_string(),
792 _ => context.rocm_version.clone(),
793 }
794 }
795
796 fn version_bump_level(&self, current: &str, proposed: &str) -> BumpLevel {
798 let cur_parts = parse_version_parts(current);
799 let pro_parts = parse_version_parts(proposed);
800
801 match (cur_parts, pro_parts) {
802 (Some(cur), Some(pro)) if !cur.is_empty() && !pro.is_empty() => {
803 let len = cur.len().max(pro.len());
804 let mut cmp = std::cmp::Ordering::Equal;
805 for idx in 0..len {
806 let cur_part = *cur.get(idx).unwrap_or(&0);
807 let pro_part = *pro.get(idx).unwrap_or(&0);
808 cmp = pro_part.cmp(&cur_part);
809 if cmp != std::cmp::Ordering::Equal {
810 break;
811 }
812 }
813
814 match cmp {
815 std::cmp::Ordering::Less => BumpLevel::Unknown,
816 std::cmp::Ordering::Equal => BumpLevel::Patch,
817 std::cmp::Ordering::Greater => {
818 if pro.get(0).copied().unwrap_or(0) > cur.get(0).copied().unwrap_or(0) {
819 BumpLevel::Major
820 } else if pro.get(1).copied().unwrap_or(0)
821 > cur.get(1).copied().unwrap_or(0)
822 {
823 BumpLevel::Minor
824 } else {
825 BumpLevel::Patch
826 }
827 }
828 }
829 }
830 _ => BumpLevel::Unknown,
831 }
832 }
833
834 fn check_runtime_version(
836 &self,
837 manifest: &Manifest,
838 context: &CompatibilityContext,
839 ) -> Result<(), PlannerError> {
840 if manifest.min_runtime_version.is_empty() {
841 return Ok(());
842 }
843
844 if context.runtime_version.is_empty() {
845 return Ok(());
846 }
847
848 if !version_gte(&context.runtime_version, &manifest.min_runtime_version) {
849 return Err(PlannerError::RuntimeTooOld {
850 current: context.runtime_version.clone(),
851 required: manifest.min_runtime_version.clone(),
852 });
853 }
854
855 Ok(())
856 }
857
858 fn enforce_dependency_rules(&self, items: &mut [PlannerItem]) -> Result<(), PlannerError> {
863 let mut changed = true;
866 while changed {
867 changed = false;
868
869 let deps_to_select: Vec<String> = items
871 .iter()
872 .filter(|i| i.selected)
873 .flat_map(|i| i.plan_item.dependencies.clone())
874 .collect();
875
876 for dep_id in &deps_to_select {
877 if let Some(dep_item) = items
878 .iter_mut()
879 .find(|i| i.plan_item.component_id == *dep_id)
880 {
881 if !dep_item.selected
882 && dep_item.classification != UpdateClassification::Blocked
883 {
884 dep_item.selected = true;
885 changed = true;
886 }
887 }
888 }
889 }
890
891 Ok(())
892 }
893}
894
895impl Default for UpdatePlanner {
896 fn default() -> Self {
897 Self::new()
898 }
899}
900
901#[derive(Debug, Clone, Copy, PartialEq, Eq)]
907enum BumpLevel {
908 Patch,
909 Minor,
910 Major,
911 Unknown,
912}
913
914fn parse_version_parts(version: &str) -> Option<Vec<u32>> {
916 let base = version.split('-').next()?;
917 base.split('.')
918 .map(|s| s.parse::<u32>().ok())
919 .collect::<Option<Vec<_>>>()
920}
921
922fn version_gte(actual: &str, required: &str) -> bool {
926 let act = parse_version_parts(actual).unwrap_or_default();
927 let req = parse_version_parts(required).unwrap_or_default();
928
929 let max_len = act.len().max(req.len());
930 for i in 0..max_len {
931 let a = act.get(i).copied().unwrap_or(0);
932 let r = req.get(i).copied().unwrap_or(0);
933 if a > r {
934 return true;
935 }
936 if a < r {
937 return false;
938 }
939 }
940 true }
942
943#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
949pub struct ScanOutput {
950 pub installed: Vec<InstalledComponent>,
952 pub manifest_source: String,
954 pub rocm_channel: String,
956 pub rocm_version: String,
958 pub gpu_architecture: String,
960}
961
962#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
964pub struct InstalledComponent {
965 pub id: String,
966 pub version: String,
967 pub status: String,
968}
969
970#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
976pub struct PlanOutput {
977 pub plan: Vec<PlannerItemOutput>,
979 pub summary: PlanSummary,
981}
982
983#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
985pub struct PlannerItemOutput {
986 pub component_id: String,
987 pub current_version: String,
988 pub proposed_version: String,
989 pub classification: String,
990 pub risk_tier: String,
991 pub selected: bool,
992 pub visible: bool,
993 pub rationale: String,
994 pub dependencies: Vec<String>,
995}
996
997impl From<&PlannerItem> for PlannerItemOutput {
998 fn from(item: &PlannerItem) -> Self {
999 Self {
1000 component_id: item.plan_item.component_id.clone(),
1001 current_version: item.plan_item.current_version.clone(),
1002 proposed_version: item.plan_item.proposed_version.clone(),
1003 classification: item.classification.label().to_string(),
1004 risk_tier: item.plan_item.risk_tier.label().to_string(),
1005 selected: item.selected,
1006 visible: item.visible,
1007 rationale: item.plan_item.rationale.clone(),
1008 dependencies: item.plan_item.dependencies.clone(),
1009 }
1010 }
1011}
1012
1013#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
1015pub struct PlanSummary {
1016 pub total: usize,
1017 pub safe: usize,
1018 pub guarded: usize,
1019 pub candidate: usize,
1020 pub experimental: usize,
1021 pub blocked: usize,
1022 pub selected: usize,
1023}
1024
1025impl PlanSummary {
1026 pub fn from_items(items: &[PlannerItem]) -> Self {
1028 Self {
1029 total: items.len(),
1030 safe: items
1031 .iter()
1032 .filter(|i| i.classification == UpdateClassification::Safe)
1033 .count(),
1034 guarded: items
1035 .iter()
1036 .filter(|i| i.classification == UpdateClassification::Guarded)
1037 .count(),
1038 candidate: items
1039 .iter()
1040 .filter(|i| i.classification == UpdateClassification::Candidate)
1041 .count(),
1042 experimental: items
1043 .iter()
1044 .filter(|i| i.classification == UpdateClassification::Experimental)
1045 .count(),
1046 blocked: items
1047 .iter()
1048 .filter(|i| i.classification == UpdateClassification::Blocked)
1049 .count(),
1050 selected: items.iter().filter(|i| i.selected).count(),
1051 }
1052 }
1053}
1054
1055#[cfg(test)]
1060mod tests {
1061 use super::*;
1062 use crate::core::manifest::{Manifest, ManifestComponent};
1063 use crate::core::types::{Category, ValidationTier};
1064
1065 fn make_component(id: &str, version: &str, tier: ValidationTier) -> ManifestComponent {
1070 ManifestComponent {
1071 id: id.to_string(),
1072 version: version.to_string(),
1073 script: format!("install_{id}.sh"),
1074 category: Category::Core,
1075 validation_tier: tier,
1076 min_rocm_version: String::new(),
1077 compatible_channels: vec![],
1078 dependencies: vec![],
1079 }
1080 }
1081
1082 fn make_manifest(components: Vec<ManifestComponent>) -> Manifest {
1083 Manifest {
1084 schema_version: 2,
1085 sequence: 1,
1086 generated_at: String::new(),
1087 expires_at: None,
1088 min_runtime_version: String::new(),
1089 components,
1090 signature: None,
1091 }
1092 }
1093
1094 fn make_context() -> CompatibilityContext {
1095 let mut ctx = CompatibilityContext::new();
1096 ctx.rocm_version = "7.2.1".to_string();
1097 ctx.rocm_channel = "latest".to_string();
1098 ctx.gpu_architecture = "gfx1100".to_string();
1099 ctx.available_executors = HashSet::from([
1100 ExecutorKind::LegacyScript,
1101 ExecutorKind::Rust,
1102 ExecutorKind::ExternalPackageManager,
1103 ]);
1104 ctx
1105 }
1106
1107 fn planner() -> UpdatePlanner {
1108 UpdatePlanner::new()
1109 }
1110
1111 #[test]
1116 fn test_safe_classification_patch_bump_validated() {
1117 let mut context = make_context();
1118 context
1119 .installed_versions
1120 .insert("pytorch".to_string(), "2.4.0".to_string());
1121 context.installed_components.insert("pytorch".to_string());
1122
1123 let component = make_component("pytorch", "2.4.1", ValidationTier::Validated);
1124 let classification =
1125 planner().classify_update(&component, &context, !context.rocm_version.is_empty(), None);
1126 assert_eq!(classification, UpdateClassification::Safe);
1127 }
1128
1129 #[test]
1130 fn test_safe_classification_same_version_reinstall() {
1131 let mut context = make_context();
1132 context
1133 .installed_versions
1134 .insert("pytorch".to_string(), "2.4.0".to_string());
1135 context.installed_components.insert("pytorch".to_string());
1136
1137 let component = make_component("pytorch", "2.4.0", ValidationTier::Validated);
1138 let classification =
1139 planner().classify_update(&component, &context, !context.rocm_version.is_empty(), None);
1140 assert_eq!(classification, UpdateClassification::Safe);
1141 }
1142
1143 #[test]
1148 fn test_guarded_classification_minor_bump() {
1149 let mut context = make_context();
1150 context
1151 .installed_versions
1152 .insert("pytorch".to_string(), "2.4.0".to_string());
1153 context.installed_components.insert("pytorch".to_string());
1154
1155 let component = make_component("pytorch", "2.5.0", ValidationTier::Validated);
1156 let classification =
1157 planner().classify_update(&component, &context, !context.rocm_version.is_empty(), None);
1158 assert_eq!(classification, UpdateClassification::Guarded);
1159 }
1160
1161 #[test]
1162 fn test_guarded_classification_patch_bump_candidate_tier() {
1163 let mut context = make_context();
1164 context
1165 .installed_versions
1166 .insert("triton".to_string(), "3.0.0".to_string());
1167 context.installed_components.insert("triton".to_string());
1168
1169 let component = make_component("triton", "3.0.1", ValidationTier::Candidate);
1171 let classification =
1172 planner().classify_update(&component, &context, !context.rocm_version.is_empty(), None);
1173 assert_eq!(classification, UpdateClassification::Guarded);
1174 }
1175
1176 #[test]
1181 fn test_blocked_classification_missing_executor() {
1182 let mut context = CompatibilityContext::new();
1183 context.rocm_version = "7.2.1".to_string();
1185
1186 let component = make_component("pytorch", "2.4.0", ValidationTier::Validated);
1187 let classification =
1188 planner().classify_update(&component, &context, !context.rocm_version.is_empty(), None);
1189 assert_eq!(classification, UpdateClassification::Blocked);
1190 }
1191
1192 #[test]
1193 fn test_blocked_classification_blocked_tier() {
1194 let context = make_context();
1195 let component = make_component("bad-pkg", "1.0.0", ValidationTier::Blocked);
1196 let classification =
1197 planner().classify_update(&component, &context, !context.rocm_version.is_empty(), None);
1198 assert_eq!(classification, UpdateClassification::Blocked);
1199 }
1200
1201 #[test]
1206 fn test_candidate_classification_major_bump() {
1207 let mut context = make_context();
1208 context
1209 .installed_versions
1210 .insert("pytorch".to_string(), "2.4.0".to_string());
1211 context.installed_components.insert("pytorch".to_string());
1212
1213 let component = make_component("pytorch", "3.0.0", ValidationTier::Validated);
1214 let classification =
1215 planner().classify_update(&component, &context, !context.rocm_version.is_empty(), None);
1216 assert_eq!(classification, UpdateClassification::Candidate);
1217 }
1218
1219 #[test]
1220 fn test_candidate_classification_new_install() {
1221 let context = make_context();
1222 let component = make_component("triton", "3.1.0", ValidationTier::Validated);
1224 let classification =
1225 planner().classify_update(&component, &context, !context.rocm_version.is_empty(), None);
1226 assert_eq!(classification, UpdateClassification::Candidate);
1227 }
1228
1229 #[test]
1234 fn test_experimental_classification() {
1235 let context = make_context();
1236 let component = make_component("deepspeed", "0.15.0", ValidationTier::Experimental);
1237 let classification =
1238 planner().classify_update(&component, &context, !context.rocm_version.is_empty(), None);
1239 assert_eq!(classification, UpdateClassification::Experimental);
1240 }
1241
1242 #[test]
1243 fn test_experimental_hidden_without_flag() {
1244 let context = make_context();
1245 let manifest = make_manifest(vec![make_component(
1246 "deepspeed",
1247 "0.15.0",
1248 ValidationTier::Experimental,
1249 )]);
1250
1251 let options = PlannerOptions {
1252 include_experimental: false,
1253 ..Default::default()
1254 };
1255
1256 let items = planner().build_plan(&manifest, &context, &options).unwrap();
1257 assert!(
1258 items.is_empty(),
1259 "Experimental items should be hidden without flag"
1260 );
1261 }
1262
1263 #[test]
1264 fn test_experimental_visible_with_flag() {
1265 let context = make_context();
1266 let manifest = make_manifest(vec![make_component(
1267 "deepspeed",
1268 "0.15.0",
1269 ValidationTier::Experimental,
1270 )]);
1271
1272 let options = PlannerOptions {
1273 include_experimental: true,
1274 ..Default::default()
1275 };
1276
1277 let items = planner().build_plan(&manifest, &context, &options).unwrap();
1278 assert_eq!(items.len(), 1);
1279 assert_eq!(items[0].classification, UpdateClassification::Experimental);
1280 }
1281
1282 #[test]
1287 fn test_safe_updates_preselected() {
1288 let mut context = make_context();
1289 context
1290 .installed_versions
1291 .insert("pytorch".to_string(), "2.4.0".to_string());
1292 context.installed_components.insert("pytorch".to_string());
1293
1294 let manifest = make_manifest(vec![make_component(
1295 "pytorch",
1296 "2.4.1",
1297 ValidationTier::Validated,
1298 )]);
1299
1300 let options = PlannerOptions::default();
1301 let items = planner().build_plan(&manifest, &context, &options).unwrap();
1302
1303 assert_eq!(items.len(), 1);
1304 assert!(items[0].selected, "Safe updates should be preselected");
1305 }
1306
1307 #[test]
1312 fn test_guarded_updates_not_preselected() {
1313 let mut context = make_context();
1314 context
1315 .installed_versions
1316 .insert("pytorch".to_string(), "2.4.0".to_string());
1317 context.installed_components.insert("pytorch".to_string());
1318
1319 let manifest = make_manifest(vec![make_component(
1320 "pytorch",
1321 "2.5.0",
1322 ValidationTier::Validated,
1323 )]);
1324
1325 let options = PlannerOptions::default();
1326 let items = planner().build_plan(&manifest, &context, &options).unwrap();
1327
1328 assert_eq!(items.len(), 1);
1329 assert!(
1330 !items[0].selected,
1331 "Guarded updates should NOT be preselected"
1332 );
1333 }
1334
1335 #[test]
1340 fn test_candidate_updates_visible_not_preselected() {
1341 let mut context = make_context();
1342 context
1343 .installed_versions
1344 .insert("pytorch".to_string(), "2.4.0".to_string());
1345 context.installed_components.insert("pytorch".to_string());
1346
1347 let manifest = make_manifest(vec![make_component(
1348 "pytorch",
1349 "3.0.0",
1350 ValidationTier::Validated,
1351 )]);
1352
1353 let options = PlannerOptions::default();
1354 let items = planner().build_plan(&manifest, &context, &options).unwrap();
1355
1356 assert_eq!(items.len(), 1);
1357 assert!(items[0].visible, "Candidate updates should be visible");
1358 assert!(
1359 !items[0].selected,
1360 "Candidate updates should NOT be preselected"
1361 );
1362 }
1363
1364 #[test]
1374 fn test_blocked_not_in_default_plan() {
1375 let context = make_context();
1376 let manifest = make_manifest(vec![make_component(
1377 "bad-pkg",
1378 "1.0.0",
1379 ValidationTier::Blocked,
1380 )]);
1381
1382 let options = PlannerOptions::default();
1383 let items = planner().build_plan(&manifest, &context, &options).unwrap();
1384
1385 assert_eq!(items.len(), 1);
1387 assert!(!items[0].visible);
1388 assert!(!items[0].selected);
1389 }
1390
1391 #[test]
1392 fn test_blocked_explicit_target_fails() {
1393 let context = make_context();
1394 let manifest = make_manifest(vec![make_component(
1395 "bad-pkg",
1396 "1.0.0",
1397 ValidationTier::Blocked,
1398 )]);
1399
1400 let options = PlannerOptions {
1401 target_components: vec!["bad-pkg".to_string()],
1402 ..Default::default()
1403 };
1404
1405 let result = planner().build_plan(&manifest, &context, &options);
1406 assert!(result.is_err());
1407 match result.unwrap_err() {
1408 PlannerError::ComponentBlocked {
1409 component_id,
1410 reason,
1411 } => {
1412 assert_eq!(component_id, "bad-pkg");
1413 assert!(!reason.is_empty());
1414 }
1415 other => panic!("Expected ComponentBlocked, got: {other:?}"),
1416 }
1417 }
1418
1419 #[test]
1424 fn test_dependency_auto_selected() {
1425 let mut context = make_context();
1426 context
1427 .installed_versions
1428 .insert("rocm".to_string(), "7.2.0".to_string());
1429 context.installed_components.insert("rocm".to_string());
1430 context
1431 .installed_versions
1432 .insert("pytorch".to_string(), "2.4.0".to_string());
1433 context.installed_components.insert("pytorch".to_string());
1434
1435 let manifest = make_manifest(vec![
1436 make_component("rocm", "7.2.1", ValidationTier::Validated), make_component("pytorch", "2.4.1", ValidationTier::Validated), ]);
1439
1440 let options = PlannerOptions::default();
1441 let items = planner().build_plan(&manifest, &context, &options).unwrap();
1442
1443 let rocm_item = items
1445 .iter()
1446 .find(|i| i.plan_item.component_id == "rocm")
1447 .unwrap();
1448 let pytorch_item = items
1449 .iter()
1450 .find(|i| i.plan_item.component_id == "pytorch")
1451 .unwrap();
1452
1453 assert!(rocm_item.selected, "ROCm should be selected (safe)");
1454 assert!(pytorch_item.selected, "PyTorch should be selected (safe)");
1455 }
1456
1457 #[test]
1462 fn test_scan_only_produces_plan() {
1463 let mut context = make_context();
1464 context
1465 .installed_versions
1466 .insert("pytorch".to_string(), "2.4.0".to_string());
1467 context.installed_components.insert("pytorch".to_string());
1468
1469 let manifest = make_manifest(vec![make_component(
1470 "pytorch",
1471 "2.4.1",
1472 ValidationTier::Validated,
1473 )]);
1474
1475 let options = PlannerOptions {
1476 scan_only: true,
1477 ..Default::default()
1478 };
1479
1480 let items = planner().build_plan(&manifest, &context, &options).unwrap();
1481 assert_eq!(items.len(), 1);
1482 }
1484
1485 #[test]
1490 fn test_all_safe_filters_to_safe_only() {
1491 let mut context = make_context();
1492 context
1493 .installed_versions
1494 .insert("pytorch".to_string(), "2.4.0".to_string());
1495 context.installed_components.insert("pytorch".to_string());
1496 context
1497 .installed_versions
1498 .insert("triton".to_string(), "3.0.0".to_string());
1499 context.installed_components.insert("triton".to_string());
1500
1501 let manifest = make_manifest(vec![
1502 make_component("pytorch", "2.4.1", ValidationTier::Validated), make_component("triton", "3.1.0", ValidationTier::Validated), ]);
1505
1506 let options = PlannerOptions {
1507 all_safe: true,
1508 ..Default::default()
1509 };
1510
1511 let items = planner().build_plan(&manifest, &context, &options).unwrap();
1512
1513 assert_eq!(items.len(), 1, "Only safe items should remain");
1514 assert_eq!(items[0].classification, UpdateClassification::Safe);
1515 assert_eq!(items[0].plan_item.component_id, "pytorch");
1516 }
1517
1518 #[test]
1519 fn test_targeted_candidate_is_selected() {
1520 let mut context = make_context();
1521 context
1522 .installed_versions
1523 .insert("pytorch".to_string(), "2.4.0".to_string());
1524 context.installed_components.insert("pytorch".to_string());
1525
1526 let manifest = make_manifest(vec![make_component(
1527 "pytorch",
1528 "3.0.0",
1529 ValidationTier::Validated,
1530 )]);
1531
1532 let options = PlannerOptions {
1533 target_components: vec!["pytorch".to_string()],
1534 ..Default::default()
1535 };
1536
1537 let items = planner().build_plan(&manifest, &context, &options).unwrap();
1538 assert_eq!(items.len(), 1);
1539 assert_eq!(items[0].classification, UpdateClassification::Candidate);
1540 assert!(items[0].selected);
1541 assert!(items[0].plan_item.selected);
1542 }
1543
1544 #[test]
1545 fn test_all_safe_keeps_explicit_targeted_non_safe_component() {
1546 let mut context = make_context();
1547 context
1548 .installed_versions
1549 .insert("pytorch".to_string(), "2.4.0".to_string());
1550 context.installed_components.insert("pytorch".to_string());
1551 context
1552 .installed_versions
1553 .insert("triton".to_string(), "3.0.0".to_string());
1554 context.installed_components.insert("triton".to_string());
1555
1556 let manifest = make_manifest(vec![
1557 make_component("pytorch", "2.4.1", ValidationTier::Validated), make_component("triton", "3.1.0", ValidationTier::Validated), ]);
1560
1561 let options = PlannerOptions {
1562 all_safe: true,
1563 target_components: vec!["triton".to_string()],
1564 ..Default::default()
1565 };
1566
1567 let items = planner().build_plan(&manifest, &context, &options).unwrap();
1568
1569 assert_eq!(
1570 items.len(),
1571 2,
1572 "Targeted component and dependencies are in scope"
1573 );
1574
1575 let triton = items
1576 .iter()
1577 .find(|i| i.plan_item.component_id == "triton")
1578 .unwrap();
1579 assert_eq!(triton.classification, UpdateClassification::Guarded);
1580 assert!(triton.selected);
1581 assert!(triton.plan_item.selected);
1582
1583 let pytorch = items
1584 .iter()
1585 .find(|i| i.plan_item.component_id == "pytorch")
1586 .unwrap();
1587 assert_eq!(pytorch.classification, UpdateClassification::Safe);
1588 assert!(pytorch.selected);
1589 }
1590
1591 #[test]
1596 fn test_targeted_blocked_component_fails() {
1597 let context = make_context();
1598 let manifest = make_manifest(vec![
1599 make_component("pytorch", "2.5.0", ValidationTier::Validated),
1600 make_component("bad-pkg", "1.0.0", ValidationTier::Blocked),
1601 ]);
1602
1603 let options = PlannerOptions {
1604 target_components: vec!["bad-pkg".to_string()],
1605 ..Default::default()
1606 };
1607
1608 let result = planner().build_plan(&manifest, &context, &options);
1609 assert!(result.is_err(), "Targeting a blocked component should fail");
1610 }
1611
1612 #[test]
1613 fn test_targeted_rocm_component_allows_planned_rocm_prereq() {
1614 let context = make_context();
1615 let manifest = make_manifest(vec![
1616 make_component("rocm", "7.2.4", ValidationTier::Validated),
1617 make_component("pytorch", "2.5.0", ValidationTier::Validated),
1618 ]);
1619
1620 let options = PlannerOptions {
1621 target_components: vec!["pytorch".to_string()],
1622 ..Default::default()
1623 };
1624
1625 let items = planner().build_plan(&manifest, &context, &options).unwrap();
1626 let ids: HashSet<_> = items
1627 .iter()
1628 .map(|i| i.plan_item.component_id.as_str())
1629 .collect();
1630 assert!(ids.contains("rocm"));
1631 assert!(ids.contains("pytorch"));
1632 }
1633
1634 #[test]
1635 fn test_targeted_component_uses_manifest_declared_dependencies() {
1636 let context = make_context();
1637 let mut custom = make_component("custom-runtime", "1.2.0", ValidationTier::Validated);
1638 custom.dependencies = vec!["rocm".to_string(), "mpi4py".to_string()];
1639
1640 let manifest = make_manifest(vec![
1641 make_component("rocm", "7.2.4", ValidationTier::Validated),
1642 make_component("mpi4py", "4.0.0", ValidationTier::Validated),
1643 custom,
1644 ]);
1645
1646 let options = PlannerOptions {
1647 target_components: vec!["custom-runtime".to_string()],
1648 ..Default::default()
1649 };
1650
1651 let items = planner().build_plan(&manifest, &context, &options).unwrap();
1652 let ids: HashSet<_> = items
1653 .iter()
1654 .map(|i| i.plan_item.component_id.as_str())
1655 .collect();
1656 assert!(ids.contains("custom-runtime"));
1657 assert!(ids.contains("rocm"));
1658 assert!(ids.contains("mpi4py"));
1659 }
1660
1661 #[test]
1662 fn test_targeted_component_with_planned_rocm_not_blocked_by_min_version_on_fresh_host() {
1663 let context = make_context();
1664 let manifest = make_manifest(vec![
1665 make_component("rocm", "7.2.4", ValidationTier::Validated),
1666 ManifestComponent {
1667 min_rocm_version: "7.2.4".to_string(),
1668 ..make_component("llama-cpp", "0.2.0", ValidationTier::Validated)
1669 },
1670 ]);
1671
1672 let options = PlannerOptions {
1673 target_components: vec!["llama-cpp".to_string()],
1674 ..Default::default()
1675 };
1676
1677 let items = planner().build_plan(&manifest, &context, &options).unwrap();
1678 let ids: HashSet<_> = items
1679 .iter()
1680 .map(|i| i.plan_item.component_id.as_str())
1681 .collect();
1682 assert!(ids.contains("rocm"));
1683 assert!(ids.contains("llama-cpp"));
1684 }
1685
1686 #[test]
1687 fn test_targeted_unknown_component_fails() {
1688 let context = make_context();
1689 let manifest = make_manifest(vec![make_component(
1690 "pytorch",
1691 "2.5.0",
1692 ValidationTier::Validated,
1693 )]);
1694
1695 let options = PlannerOptions {
1696 target_components: vec!["nonexistent".to_string()],
1697 ..Default::default()
1698 };
1699
1700 let result = planner().build_plan(&manifest, &context, &options);
1701 assert!(
1702 result.is_err(),
1703 "Targeting a nonexistent component should fail"
1704 );
1705 }
1706
1707 #[test]
1713 fn test_planner_works_with_baseline_manifest() {
1714 let context = make_context();
1715 let manifest = make_manifest(vec![make_component(
1716 "pytorch",
1717 "2.4.0",
1718 ValidationTier::Validated,
1719 )]);
1720
1721 let options = PlannerOptions::default();
1722 let items = planner().build_plan(&manifest, &context, &options).unwrap();
1723 assert!(!items.is_empty());
1724 }
1725
1726 #[test]
1731 fn test_runtime_too_old_error() {
1732 let error = PlannerError::RuntimeTooOld {
1733 current: "0.1.0".to_string(),
1734 required: "0.2.0".to_string(),
1735 };
1736 let msg = error.to_string();
1737 assert!(msg.contains("0.1.0"));
1738 assert!(msg.contains("0.2.0"));
1739 assert!(msg.contains("rusty upgrade"));
1740 }
1741
1742 #[test]
1747 fn test_rocm_channel_in_context() {
1748 let mut context = make_context();
1749 context.rocm_channel = "legacy".to_string();
1750
1751 let manifest = make_manifest(vec![make_component(
1752 "pytorch",
1753 "2.4.0",
1754 ValidationTier::Validated,
1755 )]);
1756
1757 let options = PlannerOptions::default();
1758 let items = planner().build_plan(&manifest, &context, &options).unwrap();
1759 assert!(!items.is_empty());
1760 }
1761
1762 #[test]
1772 fn test_classification_serde_roundtrip() {
1773 let classifications = [
1774 UpdateClassification::Safe,
1775 UpdateClassification::Guarded,
1776 UpdateClassification::Blocked,
1777 UpdateClassification::Candidate,
1778 UpdateClassification::Experimental,
1779 ];
1780 for c in &classifications {
1781 let json = serde_json::to_string(c).unwrap();
1782 let back: UpdateClassification = serde_json::from_str(&json).unwrap();
1783 assert_eq!(*c, back);
1784 }
1785 }
1786
1787 #[test]
1792 fn test_plan_summary_from_items() {
1793 let mut context = make_context();
1794 context
1795 .installed_versions
1796 .insert("pytorch".to_string(), "2.4.0".to_string());
1797 context.installed_components.insert("pytorch".to_string());
1798 context
1799 .installed_versions
1800 .insert("triton".to_string(), "3.0.0".to_string());
1801 context.installed_components.insert("triton".to_string());
1802
1803 let manifest = make_manifest(vec![
1804 make_component("pytorch", "2.4.1", ValidationTier::Validated), make_component("triton", "3.1.0", ValidationTier::Validated), make_component("new-pkg", "1.0.0", ValidationTier::Validated), make_component("bad-pkg", "1.0.0", ValidationTier::Blocked), ]);
1809
1810 let options = PlannerOptions::default();
1811 let items = planner().build_plan(&manifest, &context, &options).unwrap();
1812
1813 let summary = PlanSummary::from_items(&items);
1814 assert_eq!(summary.total, 4);
1815 assert_eq!(summary.safe, 1);
1816 assert_eq!(summary.guarded, 1);
1817 assert_eq!(summary.candidate, 1);
1818 assert_eq!(summary.blocked, 1);
1819 assert_eq!(summary.selected, 1); }
1821
1822 #[test]
1827 fn test_planner_item_output_conversion() {
1828 let mut context = make_context();
1829 context
1830 .installed_versions
1831 .insert("pytorch".to_string(), "2.4.0".to_string());
1832 context.installed_components.insert("pytorch".to_string());
1833
1834 let manifest = make_manifest(vec![make_component(
1835 "pytorch",
1836 "2.4.1",
1837 ValidationTier::Validated,
1838 )]);
1839
1840 let options = PlannerOptions::default();
1841 let items = planner().build_plan(&manifest, &context, &options).unwrap();
1842
1843 let output: Vec<PlannerItemOutput> = items.iter().map(PlannerItemOutput::from).collect();
1844 assert_eq!(output.len(), 1);
1845 assert_eq!(output[0].component_id, "pytorch");
1846 assert_eq!(output[0].classification, "safe");
1847 assert!(output[0].selected);
1848 }
1849
1850 #[test]
1855 fn test_full_plan_build_mixed_classifications() {
1856 let mut context = make_context();
1857 context
1858 .installed_versions
1859 .insert("rocm".to_string(), "7.2.0".to_string());
1860 context.installed_components.insert("rocm".to_string());
1861 context
1862 .installed_versions
1863 .insert("pytorch".to_string(), "2.4.0".to_string());
1864 context.installed_components.insert("pytorch".to_string());
1865 context
1866 .installed_versions
1867 .insert("triton".to_string(), "3.0.0".to_string());
1868 context.installed_components.insert("triton".to_string());
1869
1870 let manifest = make_manifest(vec![
1871 make_component("rocm", "7.2.1", ValidationTier::Validated), make_component("pytorch", "2.5.0", ValidationTier::Validated), make_component("triton", "4.0.0", ValidationTier::Validated), make_component("wandb", "0.17.0", ValidationTier::Validated), make_component("deepspeed", "0.15.0", ValidationTier::Experimental), ]);
1877
1878 let options = PlannerOptions {
1879 include_experimental: true,
1880 ..Default::default()
1881 };
1882
1883 let items = planner().build_plan(&manifest, &context, &options).unwrap();
1884
1885 assert_eq!(items.len(), 5);
1886
1887 let rocm = items
1888 .iter()
1889 .find(|i| i.plan_item.component_id == "rocm")
1890 .unwrap();
1891 assert_eq!(rocm.classification, UpdateClassification::Safe);
1892 assert!(rocm.selected);
1893
1894 let pytorch = items
1895 .iter()
1896 .find(|i| i.plan_item.component_id == "pytorch")
1897 .unwrap();
1898 assert_eq!(pytorch.classification, UpdateClassification::Guarded);
1899 assert!(!pytorch.selected);
1900
1901 let triton = items
1902 .iter()
1903 .find(|i| i.plan_item.component_id == "triton")
1904 .unwrap();
1905 assert_eq!(triton.classification, UpdateClassification::Candidate);
1906 assert!(!triton.selected);
1907
1908 let wandb = items
1909 .iter()
1910 .find(|i| i.plan_item.component_id == "wandb")
1911 .unwrap();
1912 assert_eq!(wandb.classification, UpdateClassification::Candidate);
1913 assert!(!wandb.selected);
1914
1915 let deepspeed = items
1916 .iter()
1917 .find(|i| i.plan_item.component_id == "deepspeed")
1918 .unwrap();
1919 assert_eq!(deepspeed.classification, UpdateClassification::Experimental);
1920 assert!(!deepspeed.selected);
1921 }
1922
1923 #[test]
1928 fn test_compatibility_context_default() {
1929 let ctx = CompatibilityContext::default();
1930 assert!(ctx.rocm_version.is_empty());
1931 assert!(ctx.rocm_channel.is_empty());
1932 assert!(ctx.available_executors.is_empty());
1933 assert!(ctx.installed_components.is_empty());
1934 assert!(ctx.installed_versions.is_empty());
1935 assert!(ctx.gpu_architecture.is_empty());
1936 assert!(ctx.runtime_version.is_empty());
1937 }
1938
1939 #[test]
1940 fn test_compatibility_context_is_installed() {
1941 let mut ctx = CompatibilityContext::new();
1942 assert!(!ctx.is_installed("pytorch"));
1943 ctx.installed_components.insert("pytorch".to_string());
1944 assert!(ctx.is_installed("pytorch"));
1945 }
1946
1947 #[test]
1948 fn test_compatibility_context_installed_version() {
1949 let mut ctx = CompatibilityContext::new();
1950 assert!(ctx.installed_version("pytorch").is_none());
1951 ctx.installed_versions
1952 .insert("pytorch".to_string(), "2.4.0".to_string());
1953 assert_eq!(ctx.installed_version("pytorch"), Some("2.4.0"));
1954 }
1955
1956 #[test]
1961 fn test_planner_error_display() {
1962 let err = PlannerError::ComponentBlocked {
1963 component_id: "bad-pkg".to_string(),
1964 reason: "incompatible hardware".to_string(),
1965 };
1966 assert!(err.to_string().contains("bad-pkg"));
1967 assert!(err.to_string().contains("incompatible hardware"));
1968
1969 let err = PlannerError::ComponentUnknown {
1970 component_id: "nonexistent".to_string(),
1971 };
1972 assert!(err.to_string().contains("nonexistent"));
1973
1974 let err = PlannerError::RuntimeTooOld {
1975 current: "0.1.0".to_string(),
1976 required: "0.2.0".to_string(),
1977 };
1978 assert!(err.to_string().contains("0.1.0"));
1979 assert!(err.to_string().contains("rusty upgrade"));
1980 }
1981
1982 #[test]
1983 fn test_planner_error_is_std_error() {
1984 let err = PlannerError::ComponentUnknown {
1985 component_id: "test".to_string(),
1986 };
1987 let _: &dyn std::error::Error = &err;
1988 }
1989
1990 #[test]
1995 fn test_scan_output_serde_roundtrip() {
1996 let output = ScanOutput {
1997 installed: vec![InstalledComponent {
1998 id: "pytorch".to_string(),
1999 version: "2.4.0".to_string(),
2000 status: "installed".to_string(),
2001 }],
2002 manifest_source: "baseline".to_string(),
2003 rocm_channel: "latest".to_string(),
2004 rocm_version: "7.2.1".to_string(),
2005 gpu_architecture: "gfx1100".to_string(),
2006 };
2007 let json = serde_json::to_string(&output).unwrap();
2008 let back: ScanOutput = serde_json::from_str(&json).unwrap();
2009 assert_eq!(output, back);
2010 }
2011
2012 #[test]
2013 fn test_plan_output_serde_roundtrip() {
2014 let output = PlanOutput {
2015 plan: vec![PlannerItemOutput {
2016 component_id: "pytorch".to_string(),
2017 current_version: "2.4.0".to_string(),
2018 proposed_version: "2.4.1".to_string(),
2019 classification: "safe".to_string(),
2020 risk_tier: "low".to_string(),
2021 selected: true,
2022 visible: true,
2023 rationale: "patch update".to_string(),
2024 dependencies: vec!["rocm".to_string()],
2025 }],
2026 summary: PlanSummary {
2027 total: 1,
2028 safe: 1,
2029 guarded: 0,
2030 candidate: 0,
2031 experimental: 0,
2032 blocked: 0,
2033 selected: 1,
2034 },
2035 };
2036 let json = serde_json::to_string(&output).unwrap();
2037 let back: PlanOutput = serde_json::from_str(&json).unwrap();
2038 assert_eq!(output, back);
2039 }
2040
2041 #[test]
2046 fn test_classification_visibility() {
2047 assert!(UpdateClassification::Safe.is_visible());
2048 assert!(UpdateClassification::Guarded.is_visible());
2049 assert!(!UpdateClassification::Blocked.is_visible());
2050 assert!(UpdateClassification::Candidate.is_visible());
2051 assert!(!UpdateClassification::Experimental.is_visible());
2052 }
2053
2054 #[test]
2055 fn test_classification_preselection() {
2056 assert!(UpdateClassification::Safe.is_preselected());
2057 assert!(!UpdateClassification::Guarded.is_preselected());
2058 assert!(!UpdateClassification::Blocked.is_preselected());
2059 assert!(!UpdateClassification::Candidate.is_preselected());
2060 assert!(!UpdateClassification::Experimental.is_preselected());
2061 }
2062
2063 #[test]
2064 fn test_classification_labels() {
2065 assert_eq!(UpdateClassification::Safe.label(), "safe");
2066 assert_eq!(UpdateClassification::Guarded.label(), "guarded");
2067 assert_eq!(UpdateClassification::Blocked.label(), "blocked");
2068 assert_eq!(UpdateClassification::Candidate.label(), "candidate");
2069 assert_eq!(UpdateClassification::Experimental.label(), "experimental");
2070 }
2071
2072 #[test]
2077 fn test_derive_dependencies_pytorch() {
2078 let deps = planner().derive_dependencies("pytorch");
2079 assert!(deps.contains(&"rocm".to_string()));
2080 }
2081
2082 #[test]
2083 fn test_derive_dependencies_triton() {
2084 let deps = planner().derive_dependencies("triton");
2085 assert!(deps.contains(&"pytorch".to_string()));
2086 }
2087
2088 #[test]
2089 fn test_derive_dependencies_rocm() {
2090 let deps = planner().derive_dependencies("rocm");
2091 assert!(deps.is_empty());
2092 }
2093
2094 #[test]
2095 fn test_derive_dependencies_unknown() {
2096 let deps = planner().derive_dependencies("unknown-component");
2097 assert!(deps.is_empty());
2098 }
2099
2100 #[test]
2105 fn test_version_bump_patch() {
2106 assert_eq!(
2107 planner().version_bump_level("1.0.0", "1.0.1"),
2108 BumpLevel::Patch
2109 );
2110 }
2111
2112 #[test]
2113 fn test_version_bump_minor() {
2114 assert_eq!(
2115 planner().version_bump_level("1.0.0", "1.1.0"),
2116 BumpLevel::Minor
2117 );
2118 }
2119
2120 #[test]
2121 fn test_version_bump_major() {
2122 assert_eq!(
2123 planner().version_bump_level("1.0.0", "2.0.0"),
2124 BumpLevel::Major
2125 );
2126 }
2127
2128 #[test]
2129 fn test_version_bump_unknown() {
2130 assert_eq!(
2131 planner().version_bump_level("abc", "def"),
2132 BumpLevel::Unknown
2133 );
2134 }
2135
2136 #[test]
2137 fn test_version_bump_downgrade_is_unknown() {
2138 assert_eq!(
2139 planner().version_bump_level("7.3.0", "7.2.4"),
2140 BumpLevel::Unknown
2141 );
2142 assert_eq!(
2143 planner().version_bump_level("2.5.1", "2.5.0"),
2144 BumpLevel::Unknown
2145 );
2146 }
2147
2148 #[test]
2153 fn test_version_gte_equal() {
2154 assert!(version_gte("7.2.1", "7.2.1"));
2155 }
2156
2157 #[test]
2158 fn test_version_gte_greater() {
2159 assert!(version_gte("7.3.0", "7.2.1"));
2160 assert!(version_gte("8.0.0", "7.2.1"));
2161 assert!(version_gte("7.2.2", "7.2.1"));
2162 }
2163
2164 #[test]
2165 fn test_version_gte_less() {
2166 assert!(!version_gte("7.2.0", "7.2.1"));
2167 assert!(!version_gte("6.0.0", "7.0.0"));
2168 }
2169
2170 #[test]
2171 fn test_version_gte_different_lengths() {
2172 assert!(version_gte("7.2", "7.2.0"));
2173 assert!(version_gte("7.2.1", "7.2"));
2174 assert!(!version_gte("7.1", "7.2.0"));
2175 }
2176
2177 #[test]
2182 fn test_blocked_by_rocm_version_requirement() {
2183 let context = make_context(); let component = ManifestComponent {
2186 id: "pytorch".to_string(),
2187 version: "3.0.0".to_string(),
2188 script: "install_pytorch.sh".to_string(),
2189 category: crate::core::types::Category::Core,
2190 validation_tier: ValidationTier::Validated,
2191 min_rocm_version: "99.0.0".to_string(), compatible_channels: vec![],
2193 dependencies: vec![],
2194 };
2195
2196 let classification =
2197 planner().classify_update(&component, &context, !context.rocm_version.is_empty(), None);
2198 assert_eq!(classification, UpdateClassification::Blocked);
2199 }
2200
2201 #[test]
2202 fn test_not_blocked_by_satisfied_rocm_requirement() {
2203 let context = make_context(); let component = ManifestComponent {
2206 id: "pytorch".to_string(),
2207 version: "3.0.0".to_string(),
2208 script: "install_pytorch.sh".to_string(),
2209 category: crate::core::types::Category::Core,
2210 validation_tier: ValidationTier::Validated,
2211 min_rocm_version: "7.0.0".to_string(),
2212 compatible_channels: vec![],
2213 dependencies: vec![],
2214 };
2215
2216 let classification =
2217 planner().classify_update(&component, &context, !context.rocm_version.is_empty(), None);
2218 assert_ne!(classification, UpdateClassification::Blocked);
2219 }
2220
2221 #[test]
2226 fn test_blocked_by_incompatible_channel() {
2227 let mut context = make_context();
2228 context.rocm_channel = "legacy".to_string();
2229
2230 let component = ManifestComponent {
2231 id: "pytorch".to_string(),
2232 version: "3.0.0".to_string(),
2233 script: "install_pytorch.sh".to_string(),
2234 category: crate::core::types::Category::Core,
2235 validation_tier: ValidationTier::Validated,
2236 min_rocm_version: String::new(),
2237 compatible_channels: vec!["latest".to_string(), "stable".to_string()],
2238 dependencies: vec![],
2239 };
2240
2241 let classification =
2242 planner().classify_update(&component, &context, !context.rocm_version.is_empty(), None);
2243 assert_eq!(classification, UpdateClassification::Blocked);
2244 }
2245
2246 #[test]
2247 fn test_allowed_by_compatible_channel() {
2248 let mut context = make_context();
2249 context.rocm_channel = "latest".to_string();
2250
2251 let component = ManifestComponent {
2252 id: "pytorch".to_string(),
2253 version: "3.0.0".to_string(),
2254 script: "install_pytorch.sh".to_string(),
2255 category: crate::core::types::Category::Core,
2256 validation_tier: ValidationTier::Validated,
2257 min_rocm_version: String::new(),
2258 compatible_channels: vec!["latest".to_string(), "stable".to_string()],
2259 dependencies: vec![],
2260 };
2261
2262 let classification =
2263 planner().classify_update(&component, &context, !context.rocm_version.is_empty(), None);
2264 assert_ne!(classification, UpdateClassification::Blocked);
2265 }
2266
2267 #[test]
2268 fn test_channel_filter_legacy_excludes_latest_only() {
2269 let mut context = make_context();
2270 context.rocm_channel = "legacy".to_string();
2271
2272 let manifest = make_manifest(vec![ManifestComponent {
2273 id: "pytorch".to_string(),
2274 version: "3.0.0".to_string(),
2275 script: "install_pytorch.sh".to_string(),
2276 category: crate::core::types::Category::Core,
2277 validation_tier: ValidationTier::Validated,
2278 min_rocm_version: String::new(),
2279 compatible_channels: vec!["latest".to_string()],
2280 dependencies: vec![],
2281 }]);
2282
2283 let options = PlannerOptions::default();
2284 let items = planner().build_plan(&manifest, &context, &options).unwrap();
2285
2286 assert_eq!(items.len(), 1);
2288 assert_eq!(items[0].classification, UpdateClassification::Blocked);
2289 assert!(!items[0].visible);
2290 }
2291
2292 #[test]
2297 fn test_runtime_too_old_blocks_plan() {
2298 let context = make_context(); let mut manifest = make_manifest(vec![make_component(
2301 "pytorch",
2302 "3.0.0",
2303 ValidationTier::Validated,
2304 )]);
2305 manifest.min_runtime_version = "99.0.0".to_string();
2306
2307 let mut ctx = context;
2308 ctx.runtime_version = "0.1.0".to_string();
2309
2310 let options = PlannerOptions::default();
2311 let result = planner().build_plan(&manifest, &ctx, &options);
2312 assert!(result.is_err());
2313 match result.unwrap_err() {
2314 PlannerError::RuntimeTooOld { current, required } => {
2315 assert_eq!(current, "0.1.0");
2316 assert_eq!(required, "99.0.0");
2317 }
2318 other => panic!("Expected RuntimeTooOld, got: {other:?}"),
2319 }
2320 }
2321
2322 #[test]
2323 fn test_runtime_version_satisfied() {
2324 let mut context = make_context();
2325 context.runtime_version = "1.0.0".to_string();
2326
2327 let mut manifest = make_manifest(vec![make_component(
2328 "pytorch",
2329 "3.0.0",
2330 ValidationTier::Validated,
2331 )]);
2332 manifest.min_runtime_version = "0.5.0".to_string();
2333
2334 let options = PlannerOptions::default();
2335 let result = planner().build_plan(&manifest, &context, &options);
2336 assert!(result.is_ok());
2337 }
2338
2339 #[test]
2344 fn test_blocked_when_rocm_not_installed_but_required() {
2345 let mut context = CompatibilityContext::new();
2346 context.rocm_version = String::new(); context.available_executors =
2348 HashSet::from([ExecutorKind::LegacyScript, ExecutorKind::Rust]);
2349
2350 let component = make_component("pytorch", "3.0.0", ValidationTier::Validated);
2351 let classification =
2352 planner().classify_update(&component, &context, !context.rocm_version.is_empty(), None);
2353 assert_eq!(classification, UpdateClassification::Blocked);
2354 }
2355
2356 #[test]
2357 fn test_rocm_component_allowed_when_rocm_not_installed() {
2358 let mut context = CompatibilityContext::new();
2359 context.rocm_version = String::new();
2360 context.available_executors =
2361 HashSet::from([ExecutorKind::LegacyScript, ExecutorKind::Rust]);
2362
2363 let component = make_component("rocm", "7.2.4", ValidationTier::Validated);
2364 let classification =
2365 planner().classify_update(&component, &context, !context.rocm_version.is_empty(), None);
2366 assert_ne!(classification, UpdateClassification::Blocked);
2367 }
2368
2369 #[test]
2370 fn test_isolation_safe_tracks_classification() {
2371 let mut context = make_context();
2372 context
2373 .installed_versions
2374 .insert("pytorch".to_string(), "2.4.0".to_string());
2375 context.installed_components.insert("pytorch".to_string());
2376 context
2377 .installed_versions
2378 .insert("triton".to_string(), "3.0.0".to_string());
2379 context.installed_components.insert("triton".to_string());
2380
2381 let manifest = make_manifest(vec![
2382 make_component("pytorch", "2.4.1", ValidationTier::Validated), make_component("triton", "3.1.0", ValidationTier::Validated), ]);
2385 let items = planner()
2386 .build_plan(&manifest, &context, &PlannerOptions::default())
2387 .unwrap();
2388
2389 let pytorch = items
2390 .iter()
2391 .find(|i| i.plan_item.component_id == "pytorch")
2392 .unwrap();
2393 let triton = items
2394 .iter()
2395 .find(|i| i.plan_item.component_id == "triton")
2396 .unwrap();
2397
2398 assert!(pytorch.plan_item.isolation_safe);
2399 assert!(!triton.plan_item.isolation_safe);
2400 }
2401
2402 #[test]
2407 fn test_block_reason_includes_rocm_version() {
2408 let context = make_context(); let component = ManifestComponent {
2411 id: "pytorch".to_string(),
2412 version: "3.0.0".to_string(),
2413 script: "install_pytorch.sh".to_string(),
2414 category: crate::core::types::Category::Core,
2415 validation_tier: ValidationTier::Validated,
2416 min_rocm_version: "99.0.0".to_string(),
2417 compatible_channels: vec![],
2418 dependencies: vec![],
2419 };
2420
2421 let reason =
2422 planner().block_reason(&component, &context, !context.rocm_version.is_empty(), None);
2423 assert!(reason.contains("99.0.0"));
2424 assert!(reason.contains("7.2.1"));
2425 }
2426
2427 #[test]
2428 fn test_block_reason_channel_incompatibility() {
2429 let mut context = make_context();
2430 context.rocm_channel = "legacy".to_string();
2431
2432 let component = ManifestComponent {
2433 id: "pytorch".to_string(),
2434 version: "3.0.0".to_string(),
2435 script: "install_pytorch.sh".to_string(),
2436 category: crate::core::types::Category::Core,
2437 validation_tier: ValidationTier::Validated,
2438 min_rocm_version: String::new(),
2439 compatible_channels: vec!["latest".to_string()],
2440 dependencies: vec![],
2441 };
2442
2443 let reason =
2444 planner().block_reason(&component, &context, !context.rocm_version.is_empty(), None);
2445 assert!(reason.contains("legacy"));
2446 }
2447}