1use crate::installers::common::{
35 log_warn, submit_build_report, BuildReport, BuildReportStatus, SealedToken,
36};
37use serde::Deserialize;
38use sha2::{Digest, Sha256};
39use std::fs;
40use std::io::{Read, Write};
41use std::path::{Path, PathBuf};
42use std::time::Instant;
43use tempfile::NamedTempFile;
44
45const DEFAULT_REPO_URL: &str = "https://github.com/scooter-lacroix/llama.cpp-turboquant-hip";
51
52const DEFAULT_BRANCH: &str = "main";
54
55const DEFAULT_BUILD_DIR: &str = "/tmp/llama-cpp-rocm-build";
57
58const DEFAULT_INSTALL_PREFIX: &str = ".mlstack/components/llama-cpp";
60
61pub const DETECTION_BINARY: &str = "llama-cli";
63
64pub const DETECTION_SUBCOMMAND: &str = "--help";
66
67#[derive(Debug, Clone, PartialEq, Eq)]
73pub enum HipArchs {
74 Gfx1030,
76 Gfx1100,
78 Gfx1200,
80 Default,
82}
83
84impl HipArchs {
85 pub fn from_gpu_arch(gpu_arch: &str) -> Self {
87 if gpu_arch.starts_with("gfx103") {
88 HipArchs::Gfx1030
89 } else if gpu_arch.starts_with("gfx110") {
90 HipArchs::Gfx1100
91 } else if gpu_arch.starts_with("gfx120") {
92 HipArchs::Gfx1200
93 } else {
94 HipArchs::Default
95 }
96 }
97
98 pub fn gpu_targets_for_channel(&self, channel: &str) -> String {
107 match (self, channel) {
112 (_, "legacy") => "gfx1030".to_string(),
114 (HipArchs::Gfx1030, "stable") => "gfx1030".to_string(),
116 (HipArchs::Gfx1100, "stable") => "gfx1030;gfx1100;gfx1101".to_string(),
117 (HipArchs::Gfx1200, "stable") => "gfx1030;gfx1100;gfx1101".to_string(),
118 (HipArchs::Default, "stable") => "gfx1030".to_string(),
119 (HipArchs::Gfx1030, _) => "gfx1030".to_string(),
121 (HipArchs::Gfx1100, _) => "gfx1030;gfx1100;gfx1101".to_string(),
122 (HipArchs::Gfx1200, _) => "gfx1030;gfx1100;gfx1101;gfx1200".to_string(),
123 (HipArchs::Default, _) => "gfx1030".to_string(),
124 }
125 }
126}
127
128#[derive(Debug, Clone, PartialEq, Eq)]
130pub enum RocmChannel {
131 Legacy,
132 Stable,
133 Latest,
134}
135
136impl RocmChannel {
137 pub fn parse_channel(channel: &str) -> Self {
139 match channel.to_lowercase().as_str() {
140 "legacy" => RocmChannel::Legacy,
141 "stable" => RocmChannel::Stable,
142 _ => RocmChannel::Latest,
143 }
144 }
145
146 pub fn enable_wmma_fa(&self) -> bool {
148 matches!(self, RocmChannel::Stable | RocmChannel::Latest)
149 }
150
151 pub fn label(&self) -> &'static str {
153 match self {
154 RocmChannel::Legacy => "legacy",
155 RocmChannel::Stable => "stable",
156 RocmChannel::Latest => "latest",
157 }
158 }
159}
160
161#[derive(Debug, Clone, PartialEq, Eq)]
163pub struct ShellCommand {
164 pub program: String,
166 pub args: Vec<String>,
168 pub env: Vec<(String, String)>,
170 pub working_dir: Option<PathBuf>,
172}
173
174#[derive(Debug, Clone)]
176pub struct LlamaCppConfig {
177 pub gpu_arch: String,
179 pub channel: String,
181 pub repo_url: Option<String>,
183 pub branch: String,
185 pub build_dir: Option<String>,
187 pub install_prefix: Option<String>,
189 pub dry_run: bool,
191}
192
193#[derive(Debug, Clone, Deserialize, PartialEq, Eq)]
194pub struct ReleaseAsset {
195 pub url: String,
196 pub sha256: String,
197}
198
199#[derive(Debug, Clone, Deserialize, PartialEq, Eq)]
200#[serde(deny_unknown_fields)]
201pub struct ReleaseManifest {
202 pub version: String,
203 pub archs: ReleaseArchs,
204}
205
206#[derive(Debug, Clone, Deserialize, PartialEq, Eq)]
207#[serde(deny_unknown_fields)]
208pub struct ReleaseArchs {
209 pub gfx1030: ReleaseAsset,
210 pub gfx1100: ReleaseAsset,
211 pub gfx1200: ReleaseAsset,
212}
213
214#[derive(Debug, Clone, PartialEq, Eq)]
215pub struct PrebuiltInstallPlan {
216 pub manifest_version: String,
217 pub arch: String,
218 pub url: String,
219 pub sha256: String,
220 pub install_prefix: String,
221 pub binary_path: PathBuf,
222}
223
224#[derive(Debug, Clone, PartialEq, Eq)]
225pub struct SourceInstallPlan {
226 pub reason: String,
227 pub build_commands: Vec<ShellCommand>,
228 pub install_prefix: String,
229 pub binary_path: PathBuf,
230}
231
232#[derive(Debug, Clone, PartialEq, Eq)]
233pub enum InstallStrategy {
234 Prebuilt(PrebuiltInstallPlan),
235 Source(SourceInstallPlan),
236}
237
238impl Default for LlamaCppConfig {
239 fn default() -> Self {
240 Self {
241 gpu_arch: "gfx1100".to_string(),
242 channel: "latest".to_string(),
243 repo_url: None,
244 branch: DEFAULT_BRANCH.to_string(),
245 build_dir: None,
246 install_prefix: None,
247 dry_run: false,
248 }
249 }
250}
251
252impl LlamaCppConfig {
253 pub fn repo_url(&self) -> &str {
255 self.repo_url.as_deref().unwrap_or(DEFAULT_REPO_URL)
256 }
257
258 pub fn build_dir(&self) -> &str {
260 self.build_dir.as_deref().unwrap_or(DEFAULT_BUILD_DIR)
261 }
262
263 pub fn install_prefix(&self) -> &str {
265 self.install_prefix
266 .as_deref()
267 .unwrap_or(DEFAULT_INSTALL_PREFIX)
268 }
269
270 pub fn install_bin_dir(&self, home: &str) -> PathBuf {
272 PathBuf::from(home).join(self.install_prefix()).join("bin")
273 }
274
275 pub fn detection_binary_path(&self, home: &str) -> PathBuf {
277 self.install_bin_dir(home).join(DETECTION_BINARY)
278 }
279}
280
281#[derive(Debug, Clone)]
283pub struct LlamaCppInstaller {
284 config: LlamaCppConfig,
285}
286
287impl Default for LlamaCppInstaller {
288 fn default() -> Self {
289 Self::new(LlamaCppConfig::default())
290 }
291}
292
293impl LlamaCppInstaller {
294 pub fn new(config: LlamaCppConfig) -> Self {
296 Self { config }
297 }
298
299 pub fn with_defaults() -> Self {
301 Self::new(LlamaCppConfig::default())
302 }
303
304 pub fn build_dir_path(&self) -> &str {
307 self.config.build_dir()
308 }
309
310 pub fn dependencies(&self) -> &[&str] {
318 &["rocm"]
319 }
320
321 fn resolve_auth_token(&self) -> SealedToken {
334 let mut runtime = SealedToken::from_env();
335 if !runtime.as_str().is_empty() {
336 return runtime;
337 }
338 runtime.purge();
339 SealedToken::new(option_env!("LLAMA_CPP_DEPLOY_TOKEN").unwrap_or(""))
340 }
341
342 pub fn check_latest_release(&self) -> Option<ReleaseManifest> {
346 let mut token = self.resolve_auth_token();
347 let result = self.fetch_latest_release_manifest(token.as_str());
348 token.purge();
349 result
350 }
351
352 pub fn release_asset_for_arch<'a>(
353 &self,
354 manifest: &'a ReleaseManifest,
355 gpu_arch: &str,
356 ) -> Option<&'a ReleaseAsset> {
357 match gpu_arch {
358 "gfx1030" => Some(&manifest.archs.gfx1030),
359 "gfx1100" | "gfx1101" => Some(&manifest.archs.gfx1100),
360 "gfx1200" | "gfx1201" => Some(&manifest.archs.gfx1200),
361 _ => None,
362 }
363 }
364
365 pub fn validate_repo_access(&self) -> Result<(), String> {
374 let mut github_token = Some(self.resolve_auth_token());
375
376 let temp_dir = tempfile::tempdir().map_err(|e| e.to_string())?;
378 let temp_path = temp_dir
379 .path()
380 .to_str()
381 .ok_or("Failed to create temp dir")?
382 .to_string();
383
384 let mut cmd = if let Some(ref token) = github_token {
386 let credential_helper = r#"!sh -c "echo username=git; echo password=$GITHUB_TOKEN""#;
388 let mut git_cmd = std::process::Command::new("git");
389 git_cmd
390 .arg("clone")
391 .arg("--depth")
392 .arg("1")
393 .arg("--branch")
394 .arg(&self.config.branch)
395 .arg(self.config.repo_url())
396 .arg(&temp_path)
397 .arg("-c")
398 .arg(format!("credential.helper={}", credential_helper))
399 .env("GITHUB_TOKEN", token.as_str())
400 .env("GIT_TERMINAL_PROMPT", "0");
401 git_cmd
402 } else {
403 let mut git_cmd = std::process::Command::new("git");
405 git_cmd
406 .arg("clone")
407 .arg("--depth")
408 .arg("1")
409 .arg("--branch")
410 .arg(&self.config.branch)
411 .arg(self.config.repo_url())
412 .arg(&temp_path);
413 git_cmd
414 };
415
416 let output = match cmd.output() {
417 Ok(output) => output,
418 Err(e) => {
419 if let Some(token) = github_token.as_mut() {
420 token.purge();
421 }
422 return Err(e.to_string());
423 }
424 };
425
426 if !output.status.success() {
427 let stderr = String::from_utf8_lossy(&output.stderr);
428 if let Some(token) = github_token.as_mut() {
429 token.purge();
430 }
431 if stderr.contains("authentication") || stderr.contains("403") || stderr.contains("401")
432 {
433 return Err(
434 "authentication failed for private repo. If you have access, set the build-time GitHub token.".to_string(),
435 );
436 } else {
437 return Err("Failed to clone repository".to_string());
438 }
439 }
440
441 if let Some(token) = github_token.as_mut() {
442 token.purge();
443 }
444 Ok(())
445 }
446
447 fn fetch_latest_release_manifest(&self, token: &str) -> Option<ReleaseManifest> {
448 let url =
449 "https://api.github.com/repos/scooter-lacroix/llama.cpp-turboquant-hip/releases/latest";
450 let response = ureq::get(url)
451 .header("Authorization", format!("Bearer {}", token))
452 .header("Accept", "application/vnd.github+json")
453 .header("X-GitHub-Api-Version", "2022-11-28")
454 .header("User-Agent", "rusty-stack")
455 .call();
456
457 let mut response = match response {
458 Ok(response) => response,
459 Err(_) => return None,
460 };
461
462 if response.status() == 404 {
463 return None;
464 }
465
466 if !response.status().is_success() {
467 return None;
468 }
469
470 response.body_mut().read_json::<ReleaseManifest>().ok()
471 }
472
473 pub fn cmake_flags(&self) -> Vec<String> {
481 let hip_archs = HipArchs::from_gpu_arch(&self.config.gpu_arch);
482 let channel = RocmChannel::parse_channel(&self.config.channel);
483 let gpu_targets = hip_archs.gpu_targets_for_channel(channel.label());
484
485 let mut flags = vec![
486 "-DGGML_HIP=ON".to_string(),
487 format!("-DGPU_TARGETS={}", gpu_targets),
488 "-DGGML_CUDA=OFF".to_string(),
490 "-DGGML_VULKAN=OFF".to_string(),
491 "-DGGML_METAL=OFF".to_string(),
492 "-DLLAMA_BUILD_TESTS=OFF".to_string(),
500 "-DGGML_BUILD_TESTS=OFF".to_string(),
501 "-DLLAMA_TESTS_INSTALL=OFF".to_string(),
502 "-DLLAMA_BUILD_EXAMPLES=OFF".to_string(),
503 "-DLLAMA_BUILD_TOOLS=ON".to_string(),
504 "-DLLAMA_BUILD_SERVER=ON".to_string(),
505 ];
506
507 if channel.enable_wmma_fa() {
512 flags.push("-DGGML_HIP_ROCWMMA_FATTN=ON".to_string());
513 flags.push("-DGGML_HIP_RDNA3_PROBES=OFF".to_string());
520 }
521
522 flags
523 }
524
525 fn expected_gpu_targets(&self) -> String {
526 let hip_archs = HipArchs::from_gpu_arch(&self.config.gpu_arch);
527 let channel = RocmChannel::parse_channel(&self.config.channel);
528 hip_archs.gpu_targets_for_channel(channel.label())
529 }
530
531 pub fn build_commands(&self, home: &str) -> Vec<ShellCommand> {
537 self.build_commands_with_auth(home, None)
538 }
539
540 fn build_source_commands_from_env(&self, home: &str) -> Vec<ShellCommand> {
542 let mut token = self.resolve_auth_token();
543 if token.as_str().is_empty() {
544 token.purge();
545 self.build_commands_with_auth(home, None)
546 } else {
547 self.build_commands_with_auth(home, Some(token))
548 }
549 }
550
551 pub fn resolve_install_strategy(
553 &self,
554 home: &str,
555 manifest: Option<&ReleaseManifest>,
556 ) -> InstallStrategy {
557 let install_prefix = format!("{}/{}", home, self.config.install_prefix());
558 let binary_path = self.config.detection_binary_path(home);
559 if let Some(manifest) = manifest {
560 if let Some(asset) = self.release_asset_for_arch(manifest, &self.config.gpu_arch) {
561 return InstallStrategy::Prebuilt(PrebuiltInstallPlan {
562 manifest_version: manifest.version.clone(),
563 arch: self.config.gpu_arch.clone(),
564 url: asset.url.clone(),
565 sha256: asset.sha256.clone(),
566 install_prefix,
567 binary_path,
568 });
569 }
570 }
571
572 let reason = if manifest.is_some() {
573 format!(
574 "No pre-built binary available for {}; falling back to source compile.",
575 self.config.gpu_arch
576 )
577 } else {
578 format!(
579 "Unable to fetch release manifest or match {}; falling back to source compile.",
580 self.config.gpu_arch
581 )
582 };
583
584 InstallStrategy::Source(SourceInstallPlan {
585 reason,
586 build_commands: self.build_source_commands_from_env(home),
587 install_prefix,
588 binary_path,
589 })
590 }
591
592 fn verify_sha256(path: &std::path::Path, expected: &str) -> Result<bool, String> {
593 let mut file = fs::File::open(path).map_err(|e| e.to_string())?;
594 let mut hasher = Sha256::new();
595 let mut buffer = [0u8; 8192];
596 loop {
597 let count = file.read(&mut buffer).map_err(|e| e.to_string())?;
598 if count == 0 {
599 break;
600 }
601 hasher.update(&buffer[..count]);
602 }
603 let actual = Self::sha256_hex_bytes(&hasher.finalize());
604 Ok(actual.eq_ignore_ascii_case(expected))
605 }
606
607 pub fn download_prebuilt_binary(&self, plan: &PrebuiltInstallPlan) -> Result<(), String> {
608 let mut temp = NamedTempFile::new().map_err(|e| e.to_string())?;
609 let response = match ureq::get(&plan.url).call() {
610 Ok(response) => response,
611 Err(err) => {
612 log_warn(&format!(
613 "Prebuilt llama.cpp download failed; falling back to source compile: {}",
614 err
615 ));
616 return Err("download failed".to_string());
617 }
618 };
619
620 if !response.status().is_success() {
621 log_warn("Prebuilt llama.cpp download failed; falling back to source compile.");
622 return Err("download failed".to_string());
623 }
624
625 let downloaded = response
626 .into_body()
627 .read_to_vec()
628 .map_err(|e| e.to_string())?;
629 temp.write_all(&downloaded).map_err(|e| e.to_string())?;
630 temp.flush().map_err(|e| e.to_string())?;
631
632 if !Self::verify_sha256(temp.path(), &plan.sha256)? {
633 let message = format!(
634 "Checksum mismatch for {} binary: expected {}, got {}. Download may be corrupted. Falling back to source build.",
635 plan.arch,
636 plan.sha256,
637 Self::sha256_hex(temp.path())?
638 );
639 log_warn(&message);
640 return Err(message);
641 }
642
643 fs::create_dir_all(&plan.install_prefix).map_err(|e| e.to_string())?;
644 let status = std::process::Command::new("tar")
645 .arg("-xzf")
646 .arg(temp.path())
647 .arg("-C")
648 .arg(&plan.install_prefix)
649 .status()
650 .map_err(|e| e.to_string())?;
651
652 if !status.success() {
653 return Err("failed to extract prebuilt archive".to_string());
654 }
655
656 verify_binary_linkage(&plan.binary_path)?;
659
660 temp.close().map_err(|e| e.to_string())?;
661 Ok(())
662 }
663
664 fn sha256_hex(path: &std::path::Path) -> Result<String, String> {
665 let mut file = fs::File::open(path).map_err(|e| e.to_string())?;
666 let mut hasher = Sha256::new();
667 let mut buffer = [0u8; 8192];
668 loop {
669 let count = file.read(&mut buffer).map_err(|e| e.to_string())?;
670 if count == 0 {
671 break;
672 }
673 hasher.update(&buffer[..count]);
674 }
675 Ok(Self::sha256_hex_bytes(&hasher.finalize()))
676 }
677
678 fn sha256_hex_bytes(bytes: &[u8]) -> String {
679 bytes.iter().map(|byte| format!("{:02x}", byte)).collect()
680 }
681
682 pub fn purge_source_artifacts(&self, home: &str, binary_placed: bool) -> Result<(), String> {
689 let build_dir = self.config.build_dir().to_string();
690 let install_bin = self.config.detection_binary_path(home);
691 let mut had_error = None::<String>;
692 let purge_plan = [
693 (
694 "source repository metadata",
695 std::process::Command::new("find")
696 .args([
697 format!("{}/.mlstack", home),
698 "-type".to_string(),
699 "d".to_string(),
700 "-name".to_string(),
701 ".git".to_string(),
702 "-prune".to_string(),
703 "-exec".to_string(),
704 "rm".to_string(),
705 "-rf".to_string(),
706 "{}".to_string(),
707 ";".to_string(),
708 ])
709 .output(),
710 ),
711 (
712 "git credential store",
713 std::process::Command::new("git")
714 .args([
715 "credential-store".to_string(),
716 "erase".to_string(),
717 "--file".to_string(),
718 format!("{}/.git-credentials", home),
719 ])
720 .output(),
721 ),
722 (
723 "git credential cache",
724 std::process::Command::new("git")
725 .args(["credential-cache".to_string(), "exit".to_string()])
726 .output(),
727 ),
728 (
729 "source build tree",
730 std::process::Command::new("rm")
731 .args(["-rf".to_string(), build_dir.clone()])
732 .output(),
733 ),
734 ];
735
736 for (label, output) in purge_plan {
737 match output {
738 Ok(output) if output.status.success() => {}
739 Ok(output) => {
740 let stderr = String::from_utf8_lossy(&output.stderr).trim().to_string();
741 had_error = Some(if stderr.is_empty() {
742 format!("{} purge failed", label)
743 } else {
744 stderr
745 });
746 break;
747 }
748 Err(err) => {
749 had_error = Some(format!("{} purge failed: {}", label, err));
750 break;
751 }
752 }
753 }
754
755 if let Some(err) = had_error {
756 if binary_placed && install_bin.exists() {
757 log_warn(&format!(
758 "llama.cpp source purge completed with warnings: {}",
759 err
760 ));
761 return Ok(());
762 }
763 return Err(format!(
764 "Failed to purge llama.cpp source artifacts; please remove {} and retry: {}",
765 build_dir, err
766 ));
767 }
768
769 Ok(())
770 }
771
772 pub fn build_commands_with_auth(
774 &self,
775 home: &str,
776 token: Option<SealedToken>,
777 ) -> Vec<ShellCommand> {
778 let build_dir = self.config.build_dir().to_string();
779 let install_prefix = format!("{}/{}", home, self.config.install_prefix());
780
781 let mut commands = Vec::new();
783 commands.push(self.mkdir_build_dir(&build_dir));
784
785 if let Some(mut token) = token {
786 commands.push(self.git_clone_with_auth(&build_dir, token.as_str()));
787 token.purge();
788 } else {
789 commands.push(self.git_clone(&build_dir));
790 }
791
792 commands.push(self.git_checkout(&build_dir));
793 commands.push(self.cmake_configure(&build_dir, &install_prefix));
794 commands.push(self.cmake_build_and_install(&build_dir));
795 commands.push(self.purge_source_artifacts_command(home));
796
797 commands
798 }
799
800 pub fn install(&self, home: &str, manifest: Option<&ReleaseManifest>) -> Result<(), String> {
801 let start = Instant::now();
802 match self.resolve_install_strategy(home, manifest) {
803 InstallStrategy::Prebuilt(plan) => match self.download_prebuilt_binary(&plan) {
804 Ok(()) => {
805 self.submit_telemetry(home, start, true, Some(&plan.binary_path))
806 .ok();
807 Ok(())
808 }
809 Err(err) if err.contains("checksum mismatch") => {
810 log_warn(&format!("{err} Falling back to source compile."));
811 let result = self.run_source_install(home);
812 if result.is_ok() {
813 self.submit_telemetry(home, start, false, None).ok();
814 }
815 result
816 }
817 Err(_) => {
818 log_warn("Prebuilt download unavailable; using source compile.");
819 let result = self.run_source_install(home);
820 if result.is_ok() {
821 self.submit_telemetry(home, start, false, None).ok();
822 }
823 result
824 }
825 },
826 InstallStrategy::Source(plan) => {
827 let result = self.run_source_install(home);
828 if result.is_ok() {
829 self.submit_telemetry(home, start, false, Some(&plan.binary_path))
830 .ok();
831 }
832 result
833 }
834 }
835 }
836
837 pub fn run_source_install(&self, home: &str) -> Result<(), String> {
838 let cuda_artifacts = detect_cuda_toolkit_presence();
840 if !cuda_artifacts.is_empty() {
841 log_warn(&format!(
842 "NVIDIA CUDA toolkit detected on this system ({}). Rusty Llama builds exclusively with ROCm/HIP; CUDA will not be used.",
843 cuda_artifacts.join(", ")
844 ));
845 }
846
847 let commands = self.build_source_commands_from_env(home);
848 if commands.is_empty() {
849 return Err("no source install commands available".to_string());
850 }
851
852 for command in commands {
853 let is_cmake_configure = command.program == "cmake"
854 && command.args.iter().any(|arg| arg == "-B")
855 && command.args.iter().any(|arg| arg == "-S");
856 let is_cmake_install = command.program == "cmake"
860 && command
861 .args
862 .iter()
863 .any(|arg| arg == "--install" || arg == "install");
864 let command_working_dir = command.working_dir.clone();
865
866 let mut process = std::process::Command::new(&command.program);
867 process.args(&command.args);
868 for (key, value) in &command.env {
869 process.env(key, value);
870 }
871 if let Some(dir) = &command.working_dir {
872 process.current_dir(dir);
873 }
874
875 let output = process.output().map_err(|err| {
876 format!(
877 "failed to run source install command '{}': {}",
878 command.program, err
879 )
880 })?;
881
882 if !output.status.success() {
883 let stderr = String::from_utf8_lossy(&output.stderr).trim().to_string();
884 let stdout = String::from_utf8_lossy(&output.stdout).trim().to_string();
885 let mut details = String::new();
886 if !stderr.is_empty() {
887 details.push_str(&stderr);
888 }
889 if !stdout.is_empty() {
890 if !details.is_empty() {
891 details.push_str(" | ");
892 }
893 details.push_str(&stdout);
894 }
895
896 return Err(if details.is_empty() {
897 format!("source install command '{}' failed", command.program)
898 } else {
899 format!(
900 "source install command '{}' failed: {}",
901 command.program, details
902 )
903 });
904 }
905
906 if is_cmake_configure {
907 let cache_path = command_working_dir
908 .as_deref()
909 .unwrap_or_else(|| Path::new("."))
910 .join("build")
911 .join("CMakeCache.txt");
912 validate_cmake_cache(&cache_path, &self.expected_gpu_targets())?;
913 }
914
915 if is_cmake_install {
916 let bin_path = resolve_install_bin_dir(home).join(DETECTION_BINARY);
919 verify_binary_linkage(&bin_path)?;
920 }
921 }
922
923 Ok(())
924 }
925
926 fn purge_source_artifacts_command(&self, home: &str) -> ShellCommand {
927 let build_dir = self.config.build_dir().to_string();
928 let mlstack_dir = format!("{}/.mlstack", home);
929 let credentials_file = format!("{}/.git-credentials", home);
930 ShellCommand {
931 program: "sh".to_string(),
932 args: vec![
933 "-c".to_string(),
934 format!(
935 "find '{}' -type d -name .git -prune -exec rm -rf '{{}}' ';' ; git credential-store erase --file '{}' ; git credential-cache exit ; rm -rf '{}'",
936 mlstack_dir, credentials_file, build_dir
937 ),
938 ],
939 env: vec![],
940 working_dir: None,
941 }
942 }
943
944 fn git_clone_with_auth(&self, build_dir: &str, token: &str) -> ShellCommand {
950 let credential_helper = r#"!sh -c "echo username=git; echo password=$GITHUB_TOKEN""#;
952 ShellCommand {
953 program: "git".to_string(),
954 args: vec![
955 "clone".to_string(),
956 "--depth".to_string(),
957 "1".to_string(),
958 "--branch".to_string(),
959 self.config.branch.clone(),
960 self.config.repo_url().to_string(),
961 build_dir.to_string(),
962 "-c".to_string(),
963 format!("credential.helper={}", credential_helper),
964 ],
965 env: vec![
966 ("GITHUB_TOKEN".to_string(), token.to_string()),
967 ("GIT_TERMINAL_PROMPT".to_string(), "0".to_string()),
968 ],
969 working_dir: None,
970 }
971 }
972
973 fn mkdir_build_dir(&self, build_dir: &str) -> ShellCommand {
974 let bd = build_dir.replace('\'', "'\\''");
982 ShellCommand {
983 program: "sh".to_string(),
984 args: vec![
985 "-c".to_string(),
986 format!("rm -rf -- '{bd}' && mkdir -p '{bd}'"),
987 ],
988 env: vec![],
989 working_dir: None,
990 }
991 }
992
993 fn git_clone(&self, build_dir: &str) -> ShellCommand {
994 ShellCommand {
995 program: "git".to_string(),
996 args: vec![
997 "clone".to_string(),
998 "--depth".to_string(),
999 "1".to_string(),
1000 "--branch".to_string(),
1001 self.config.branch.clone(),
1002 self.config.repo_url().to_string(),
1003 build_dir.to_string(),
1004 ],
1005 env: vec![],
1006 working_dir: None,
1007 }
1008 }
1009
1010 fn git_checkout(&self, build_dir: &str) -> ShellCommand {
1011 ShellCommand {
1012 program: "git".to_string(),
1013 args: vec!["checkout".to_string(), self.config.branch.clone()],
1014 env: vec![],
1015 working_dir: Some(PathBuf::from(build_dir)),
1016 }
1017 }
1018
1019 fn cmake_configure(&self, build_dir: &str, install_prefix: &str) -> ShellCommand {
1020 let mut args = vec![
1021 "-B".to_string(),
1022 "build".to_string(),
1023 "-S".to_string(),
1024 ".".to_string(),
1025 ];
1026 args.extend(self.cmake_flags());
1027 args.push(format!("-DCMAKE_INSTALL_PREFIX={}", install_prefix));
1028
1029 ShellCommand {
1030 program: "cmake".to_string(),
1031 args,
1032 env: vec![],
1033 working_dir: Some(PathBuf::from(build_dir)),
1034 }
1035 }
1036
1037 fn cmake_build_and_install(&self, build_dir: &str) -> ShellCommand {
1045 ShellCommand {
1046 program: "cmake".to_string(),
1047 args: vec![
1048 "--build".to_string(),
1049 "build".to_string(),
1050 "-j".to_string(),
1051 "--target".to_string(),
1052 "install".to_string(),
1053 ],
1054 env: vec![],
1055 working_dir: Some(PathBuf::from(build_dir)),
1056 }
1057 }
1058
1059 fn submit_telemetry(
1060 &self,
1061 home: &str,
1062 build_start: Instant,
1063 was_prebuilt: bool,
1064 binary_path: Option<&std::path::Path>,
1065 ) -> Result<(), String> {
1066 let gpu = crate::hardware::detect_hardware()
1067 .map(|state| state.gpu)
1068 .unwrap_or_default();
1069 let report = BuildReport {
1070 gpu_arch: if gpu.architecture.is_empty() {
1071 self.config.gpu_arch.clone()
1072 } else {
1073 gpu.architecture
1074 },
1075 gpu_count: gpu.gpu_count,
1076 gpu_name: gpu.model,
1077 rocm_version: gpu.rocm_version,
1078 os: std::env::consts::OS.to_string(),
1079 os_distro: std::env::var("ID").unwrap_or_else(|_| "unknown".to_string()),
1080 build_duration_seconds: build_start.elapsed().as_secs(),
1081 git_commit: Self::current_git_commit().unwrap_or_else(|| "unknown".to_string()),
1082 build_status: BuildReportStatus::Success,
1083 install_path: binary_path
1084 .map(|p| p.display().to_string())
1085 .unwrap_or_else(|| {
1086 self.config
1087 .detection_binary_path(home)
1088 .display()
1089 .to_string()
1090 }),
1091 cmake_flags: self.cmake_flags(),
1092 verification_path: self
1093 .config
1094 .detection_binary_path(home)
1095 .display()
1096 .to_string(),
1097 rdna3_validation_passed: false,
1098 wmma_available: false,
1099 shared_memory_ok: false,
1100 tokens_per_second_wmma: None,
1101 tokens_per_second_fallback: None,
1102 binary_version: self.config.branch.clone(),
1103 was_prebuilt,
1104 };
1105 submit_build_report(report);
1106 Ok(())
1107 }
1108
1109 fn current_git_commit() -> Option<String> {
1110 let output = std::process::Command::new("git")
1111 .args(["rev-parse", "HEAD"])
1112 .output()
1113 .ok()?;
1114 if !output.status.success() {
1115 return None;
1116 }
1117 Some(String::from_utf8_lossy(&output.stdout).trim().to_string())
1118 }
1119}
1120
1121pub fn resolve_install_bin_dir(home: &str) -> PathBuf {
1129 PathBuf::from(home).join(DEFAULT_INSTALL_PREFIX).join("bin")
1130}
1131
1132pub fn is_llama_cli_functional(home: &str) -> bool {
1137 let bin_path = resolve_install_bin_dir(home).join(DETECTION_BINARY);
1138 if !bin_path.exists() {
1139 return false;
1140 }
1141 run_rocm_tool_with_env(
1147 home,
1148 bin_path.to_str().unwrap_or(""),
1149 &[DETECTION_SUBCOMMAND],
1150 )
1151 .map(|o| o.status.success())
1152 .unwrap_or(false)
1153}
1154
1155pub fn is_llama_cli_on_path() -> bool {
1159 std::process::Command::new(DETECTION_BINARY)
1160 .arg(DETECTION_SUBCOMMAND)
1161 .output()
1162 .map(|o| o.status.success())
1163 .unwrap_or(false)
1164}
1165
1166pub fn has_partial_artifacts(home: &str) -> bool {
1171 let build_dir = PathBuf::from(DEFAULT_BUILD_DIR);
1172 let install_dir = PathBuf::from(home).join(DEFAULT_INSTALL_PREFIX);
1173
1174 let artifacts_exist = build_dir.exists() || install_dir.exists();
1175 let binary_functional = is_llama_cli_functional(home);
1176
1177 artifacts_exist && !binary_functional
1178}
1179
1180fn verify_binary_linkage(binary_path: &Path) -> Result<(), String> {
1193 if !binary_path.exists() {
1194 return Err(format!(
1195 "Binary not found at {} — install may have failed",
1196 binary_path.display()
1197 ));
1198 }
1199
1200 let lib_dir = binary_path
1206 .parent()
1207 .and_then(|bin| bin.parent())
1208 .map(|prefix| prefix.join("lib"));
1209 let mut ld_parts: Vec<String> = Vec::new();
1210 if let Some(dir) = &lib_dir {
1211 ld_parts.push(dir.to_string_lossy().into_owned());
1212 }
1213 ld_parts.push("/opt/rocm/lib".to_string());
1214 if let Ok(existing) = std::env::var("LD_LIBRARY_PATH") {
1215 if !existing.trim().is_empty() {
1216 ld_parts.push(existing);
1217 }
1218 }
1219 let ld_path = ld_parts.join(":");
1220
1221 let ldd_analyze = |path: &Path| -> Option<(bool, bool)> {
1223 let out = std::process::Command::new("ldd")
1224 .arg(path)
1225 .env("LD_LIBRARY_PATH", &ld_path)
1226 .output()
1227 .ok()?;
1228 if !out.status.success() {
1229 return None;
1230 }
1231 let s = String::from_utf8_lossy(&out.stdout);
1232 let cuda = ["libcuda.so", "libnvidia", "libcudart"]
1233 .iter()
1234 .any(|m| s.contains(m));
1235 let rocm = s.contains("amdhip64") || s.contains("hipblas") || s.contains("libhip");
1236 Some((cuda, rocm))
1237 };
1238
1239 let mut paths = vec![binary_path.to_path_buf()];
1243 if let Some(dir) = &lib_dir {
1244 paths.push(dir.join("libggml-hip.so"));
1245 }
1246
1247 let mut any_rocm = false;
1248 for path in &paths {
1249 if !path.exists() {
1250 continue;
1251 }
1252 if let Some((cuda, rocm)) = ldd_analyze(path) {
1253 if cuda {
1254 return Err(format!(
1255 "CUDA contamination detected: {} links against libcuda/cudart. \
1256 The binary was built with CUDA instead of ROCm/HIP — check CMake.",
1257 path.display()
1258 ));
1259 }
1260 if rocm {
1261 any_rocm = true;
1262 }
1263 }
1264 }
1265
1266 if any_rocm {
1267 Ok(())
1268 } else {
1269 Err(format!(
1270 "ROCm/HIP linkage not found in {} (or its libggml-hip.so). The binary may \
1271 be a CPU-only build or was built without GPU acceleration. Expected linkage \
1272 to amdhip64 or hipblas.",
1273 binary_path.display()
1274 ))
1275 }
1276}
1277
1278fn detect_cuda_toolkit_presence() -> Vec<String> {
1282 let mut found = Vec::new();
1283 let cuda_indicators = [
1284 ("nvcc", "NVIDIA CUDA compiler"),
1285 ("nvidia-smi", "NVIDIA system management interface"),
1286 ];
1287
1288 for (binary, description) in cuda_indicators {
1289 if std::process::Command::new("which")
1290 .arg(binary)
1291 .output()
1292 .map(|output| output.status.success())
1293 .unwrap_or(false)
1294 {
1295 found.push(format!("{} ({})", binary, description));
1296 }
1297 }
1298
1299 if Path::new("/usr/local/cuda").exists() {
1300 found.push("/usr/local/cuda directory".to_string());
1301 }
1302
1303 found
1304}
1305
1306fn validate_cmake_cache(cache_path: &Path, expected_gpu_targets: &str) -> Result<(), String> {
1310 let contents = fs::read_to_string(cache_path).map_err(|err| {
1311 format!(
1312 "failed to read CMake cache at {}: {}",
1313 cache_path.display(),
1314 err
1315 )
1316 })?;
1317
1318 let mut hip_on = false;
1319 let mut cuda_off = false;
1320 let mut gpu_targets_seen = false;
1321
1322 for line in contents.lines() {
1323 let trimmed = line.trim();
1324 if trimmed.is_empty() || trimmed.starts_with('#') || trimmed.starts_with("//") {
1325 continue;
1326 }
1327
1328 if trimmed == "GGML_HIP:BOOL=ON" {
1329 hip_on = true;
1330 } else if trimmed == "GGML_CUDA:BOOL=OFF" {
1331 cuda_off = true;
1332 } else if let Some(rest) = trimmed.strip_prefix("GPU_TARGETS:") {
1333 if let Some(eq) = rest.find('=') {
1339 let value = &rest[eq + 1..];
1340 gpu_targets_seen = true;
1341 if value != expected_gpu_targets {
1342 return Err(format!(
1343 "GPU_TARGETS mismatch in CMakeCache.txt: expected '{}', found '{}'",
1344 expected_gpu_targets, value
1345 ));
1346 }
1347 }
1348 }
1349 }
1350
1351 if !hip_on {
1352 return Err("GGML_HIP is not ON in CMakeCache.txt; build is not using ROCm".to_string());
1353 }
1354 if !cuda_off {
1355 return Err("GGML_CUDA is not OFF in CMakeCache.txt; CUDA contamination risk".to_string());
1356 }
1357 if !gpu_targets_seen {
1358 return Err("GPU_TARGETS is missing from CMakeCache.txt".to_string());
1359 }
1360
1361 Ok(())
1362}
1363
1364#[derive(Debug, Clone, PartialEq)]
1370pub struct PostInstallVerification {
1371 pub help_check_passed: bool,
1373 pub stronger_check_attempted: bool,
1375 pub stronger_check_passed: bool,
1377 pub bench_prefill_tps: Option<f64>,
1379 pub bench_decode_tps: Option<f64>,
1381 pub rdna3_check_attempted: bool,
1383 pub rdna3_is_valid: Option<bool>,
1385 pub rdna3_wmma_tps: Option<f64>,
1387 pub is_rdna3_device: Option<bool>,
1389 pub wmma_supported: Option<bool>,
1391 pub verified_binary_path: String,
1393 pub summary: String,
1395}
1396
1397const VERIFICATION_MODEL_FILENAME: &str = "Qwen3-0.6B-Q4_0.gguf";
1399
1400fn find_verification_model(home: &str, clone_dir: &str) -> Option<PathBuf> {
1405 let candidates = [
1406 PathBuf::from(home)
1407 .join(".mlstack/models")
1408 .join(VERIFICATION_MODEL_FILENAME),
1409 PathBuf::from(clone_dir)
1410 .join("models")
1411 .join(VERIFICATION_MODEL_FILENAME),
1412 ];
1413 candidates.into_iter().find(|p| p.exists())
1414}
1415
1416const VERIFICATION_MODEL_URL: &str =
1422 "https://huggingface.co/ggml-org/Qwen3-0.6B-GGUF/resolve/main/Qwen3-0.6B-Q4_0.gguf";
1423
1424fn ensure_verification_model(home: &str, clone_dir: &str) -> Option<PathBuf> {
1431 if let Some(p) = find_verification_model(home, clone_dir) {
1432 return Some(p);
1433 }
1434 ensure_default_gguf_model(home)
1435}
1436
1437pub fn ensure_default_gguf_model(home: &str) -> Option<PathBuf> {
1444 let dest = PathBuf::from(home)
1445 .join(".mlstack")
1446 .join("models")
1447 .join(VERIFICATION_MODEL_FILENAME);
1448 if dest.exists() && is_gguf(&dest) {
1449 return Some(dest);
1450 }
1451 let _ = fs::create_dir_all(dest.parent().unwrap_or(Path::new(".")));
1452 let ok = std::process::Command::new("curl")
1453 .args(["-fLs", "--connect-timeout", "20", "--retry", "2", "-o"])
1454 .arg(&dest)
1455 .arg(VERIFICATION_MODEL_URL)
1456 .status()
1457 .map(|s| s.success())
1458 .unwrap_or(false);
1459 if ok && dest.exists() && is_gguf(&dest) {
1460 Some(dest)
1461 } else {
1462 let _ = fs::remove_file(&dest); None
1464 }
1465}
1466
1467fn is_gguf(path: &Path) -> bool {
1469 let mut buf = [0u8; 4];
1470 fs::File::open(path)
1471 .and_then(|mut f| f.read_exact(&mut buf))
1472 .is_ok()
1473 && buf == *b"GGUF"
1474}
1475
1476fn run_rocm_tool_with_env(
1487 home: &str,
1488 program: &str,
1489 args: &[&str],
1490) -> std::io::Result<std::process::Output> {
1491 let env_file = format!("{}/.mlstack_env", home);
1492 let lib_export = Path::new(program)
1497 .parent()
1498 .and_then(|p| p.parent())
1499 .map(|p| p.join("lib"))
1500 .filter(|d| d.exists())
1501 .map(|d| {
1502 let escaped = shell_single_quote(&d.display().to_string());
1507 format!("export LD_LIBRARY_PATH={}:$LD_LIBRARY_PATH\n", escaped)
1508 })
1509 .unwrap_or_default();
1510 let mut script = format!(
1515 "source {} 2>/dev/null\n{}exec {}",
1516 shell_single_quote(&env_file),
1517 lib_export,
1518 shell_single_quote(program)
1519 );
1520 for a in args {
1521 script.push_str(&format!(" {}", shell_single_quote(a)));
1522 }
1523 std::process::Command::new("bash")
1524 .arg("-c")
1525 .arg(script)
1526 .output()
1527}
1528
1529fn shell_single_quote(s: &str) -> String {
1534 format!("'{}'", s.replace('\'', "'\\''"))
1535}
1536
1537pub fn verify_installed_binary(home: &str, fork_dir: &str) -> PostInstallVerification {
1538 let bin_dir = resolve_install_bin_dir(home);
1539 let cli_path = bin_dir.join(DETECTION_BINARY);
1540 let bench_path = bin_dir.join("llama-bench");
1541 let rdna3_path = bin_dir.join("rdna3_validation");
1542
1543 let mut result = PostInstallVerification {
1544 help_check_passed: false,
1545 stronger_check_attempted: false,
1546 stronger_check_passed: false,
1547 bench_prefill_tps: None,
1548 bench_decode_tps: None,
1549 rdna3_check_attempted: false,
1550 rdna3_is_valid: None,
1551 rdna3_wmma_tps: None,
1552 is_rdna3_device: None,
1553 wmma_supported: None,
1554 verified_binary_path: cli_path.display().to_string(),
1555 summary: String::new(),
1556 };
1557
1558 if cli_path.exists() {
1560 match run_rocm_tool_with_env(
1561 home,
1562 cli_path.to_str().unwrap_or(""),
1563 &[DETECTION_SUBCOMMAND],
1564 ) {
1565 Ok(out) => {
1566 result.help_check_passed = out.status.success();
1567 }
1568 Err(_) => {
1569 result.help_check_passed = false;
1570 }
1571 }
1572 }
1573
1574 if bench_path.exists() {
1576 if let Some(model_path) = ensure_verification_model(home, fork_dir) {
1581 result.stronger_check_attempted = true;
1582 let model_str = model_path.to_str().unwrap_or("");
1583 let bench_args: [&str; 8] = ["-m", model_str, "-p", "128", "-n", "32", "-ngl", "33"];
1584 if let Ok(out) =
1585 run_rocm_tool_with_env(home, bench_path.to_str().unwrap_or(""), &bench_args)
1586 {
1587 if out.status.success() {
1588 result.stronger_check_passed = true;
1589 let stdout = String::from_utf8_lossy(&out.stdout);
1591 for line in stdout.lines() {
1592 if line.contains("pp128") {
1593 if let Some(tps) = parse_bench_tps(line) {
1595 result.bench_prefill_tps = Some(tps);
1596 }
1597 } else if line.contains("tg32") {
1598 if let Some(tps) = parse_bench_tps(line) {
1599 result.bench_decode_tps = Some(tps);
1600 }
1601 }
1602 }
1603 }
1604 }
1605 }
1606 }
1607
1608 if rdna3_path.exists() {
1610 result.rdna3_check_attempted = true;
1611 if let Ok(out) = run_rocm_tool_with_env(home, rdna3_path.to_str().unwrap_or(""), &[]) {
1612 if out.status.success() {
1613 let stdout = String::from_utf8_lossy(&out.stdout);
1614 result.is_rdna3_device = Some(stdout.contains("Is RDNA3: YES"));
1615 result.wmma_supported = Some(stdout.contains("WMMA Supported: YES"));
1616 result.rdna3_is_valid = result.is_rdna3_device;
1617
1618 for line in stdout.lines() {
1620 if line.starts_with("WMMA Throughput:") {
1621 if let Some(val) = line.split_whitespace().nth(2) {
1622 if let Ok(tps) = val.parse::<f64>() {
1623 result.rdna3_wmma_tps = Some(tps);
1624 }
1625 }
1626 }
1627 }
1628 }
1629 }
1630 }
1631
1632 let mut parts: Vec<String> = vec![];
1634 if result.help_check_passed {
1635 parts.push("basic check PASSED".to_string());
1636 } else {
1637 parts.push("basic check FAILED".to_string());
1638 }
1639 if result.stronger_check_attempted {
1640 if result.stronger_check_passed {
1641 parts.push(format!(
1642 "llama-bench PASSED (prefill={:.0} t/s, decode={:.0} t/s)",
1643 result.bench_prefill_tps.unwrap_or(0.0),
1644 result.bench_decode_tps.unwrap_or(0.0)
1645 ));
1646 } else {
1647 parts.push("llama-bench FAILED".to_string());
1648 }
1649 } else {
1650 parts.push("llama-bench SKIPPED (binary unavailable)".to_string());
1651 }
1652 if result.rdna3_check_attempted {
1653 if result.rdna3_is_valid.unwrap_or(false) {
1654 parts.push(format!(
1655 "RDNA3 proof PASSED (WMMA={:.0} t/s)",
1656 result.rdna3_wmma_tps.unwrap_or(0.0)
1657 ));
1658 } else {
1659 parts.push("RDNA3 proof: not an RDNA3 device".to_string());
1660 }
1661 } else {
1662 parts.push("RDNA3 proof SKIPPED (rdna3_validation unavailable)".to_string());
1663 }
1664
1665 result.summary = parts.join("; ");
1666 result
1667}
1668
1669fn parse_bench_tps(line: &str) -> Option<f64> {
1672 let parts: Vec<&str> = line.split('|').collect();
1673 if parts.len() >= 8 {
1674 let tps_cell = parts[7].trim();
1675 tps_cell
1677 .split('±')
1678 .next()
1679 .and_then(|s| s.trim().parse::<f64>().ok())
1680 } else {
1681 None
1682 }
1683}
1684
1685#[cfg(test)]
1690mod tests {
1691 use super::*;
1692
1693 #[test]
1694 fn test_hip_archs_from_gpu_arch() {
1695 assert_eq!(HipArchs::from_gpu_arch("gfx1030"), HipArchs::Gfx1030);
1696 assert_eq!(HipArchs::from_gpu_arch("gfx1100"), HipArchs::Gfx1100);
1697 assert_eq!(HipArchs::from_gpu_arch("gfx1101"), HipArchs::Gfx1100);
1698 assert_eq!(HipArchs::from_gpu_arch("gfx1200"), HipArchs::Gfx1200);
1699 assert_eq!(HipArchs::from_gpu_arch("gfx1201"), HipArchs::Gfx1200);
1700 assert_eq!(HipArchs::from_gpu_arch("gfx000"), HipArchs::Default);
1701 assert_eq!(HipArchs::from_gpu_arch("unknown"), HipArchs::Default);
1702 }
1703
1704 #[test]
1705 fn test_gpu_targets_legacy_channel() {
1706 assert_eq!(
1708 HipArchs::Gfx1100.gpu_targets_for_channel("legacy"),
1709 "gfx1030"
1710 );
1711 assert_eq!(
1712 HipArchs::Gfx1200.gpu_targets_for_channel("legacy"),
1713 "gfx1030"
1714 );
1715 assert_eq!(
1716 HipArchs::Gfx1030.gpu_targets_for_channel("legacy"),
1717 "gfx1030"
1718 );
1719 assert_eq!(
1720 HipArchs::Default.gpu_targets_for_channel("legacy"),
1721 "gfx1030"
1722 );
1723 }
1724
1725 #[test]
1726 fn test_gpu_targets_stable_channel() {
1727 assert_eq!(
1728 HipArchs::Gfx1030.gpu_targets_for_channel("stable"),
1729 "gfx1030"
1730 );
1731 assert_eq!(
1732 HipArchs::Gfx1100.gpu_targets_for_channel("stable"),
1733 "gfx1030;gfx1100;gfx1101"
1734 );
1735 assert_eq!(
1737 HipArchs::Gfx1200.gpu_targets_for_channel("stable"),
1738 "gfx1030;gfx1100;gfx1101"
1739 );
1740 assert_eq!(
1741 HipArchs::Default.gpu_targets_for_channel("stable"),
1742 "gfx1030"
1743 );
1744 }
1745
1746 #[test]
1747 fn test_gpu_targets_latest_channel() {
1748 assert_eq!(
1749 HipArchs::Gfx1030.gpu_targets_for_channel("latest"),
1750 "gfx1030"
1751 );
1752 assert_eq!(
1753 HipArchs::Gfx1100.gpu_targets_for_channel("latest"),
1754 "gfx1030;gfx1100;gfx1101"
1755 );
1756 assert_eq!(
1757 HipArchs::Gfx1200.gpu_targets_for_channel("latest"),
1758 "gfx1030;gfx1100;gfx1101;gfx1200"
1759 );
1760 }
1761
1762 #[test]
1763 fn test_rocm_channel_wmma_fa() {
1764 assert!(!RocmChannel::parse_channel("legacy").enable_wmma_fa());
1765 assert!(RocmChannel::parse_channel("stable").enable_wmma_fa());
1766 assert!(RocmChannel::parse_channel("latest").enable_wmma_fa());
1767 }
1768
1769 #[test]
1770 fn test_cmake_flags_legacy() {
1771 let config = LlamaCppConfig {
1772 gpu_arch: "gfx1100".to_string(),
1773 channel: "legacy".to_string(),
1774 ..Default::default()
1775 };
1776 let installer = LlamaCppInstaller::new(config);
1777 let flags = installer.cmake_flags();
1778
1779 assert!(flags.contains(&"-DGGML_HIP=ON".to_string()));
1780 assert!(flags.contains(&"-DGPU_TARGETS=gfx1030".to_string()));
1781 assert!(!flags.iter().any(|f| f.contains("RDNA3_PROBES")));
1783 assert!(!flags.iter().any(|f| f.contains("ROCWMMA_FATTN")));
1784 }
1785
1786 #[test]
1787 fn test_cmake_flags_always_disable_non_rocm_backends() {
1788 for channel in ["legacy", "stable", "latest"] {
1790 let config = LlamaCppConfig {
1791 gpu_arch: "gfx1100".to_string(),
1792 channel: channel.to_string(),
1793 ..Default::default()
1794 };
1795 let installer = LlamaCppInstaller::new(config);
1796 let flags = installer.cmake_flags();
1797
1798 assert!(
1799 flags.contains(&"-DGGML_CUDA=OFF".to_string()),
1800 "channel {channel} must set GGML_CUDA=OFF"
1801 );
1802 assert!(
1803 flags.contains(&"-DGGML_VULKAN=OFF".to_string()),
1804 "channel {channel} must set GGML_VULKAN=OFF"
1805 );
1806 assert!(
1807 flags.contains(&"-DGGML_METAL=OFF".to_string()),
1808 "channel {channel} must set GGML_METAL=OFF"
1809 );
1810 }
1811 }
1812
1813 #[test]
1814 fn test_cmake_flags_stable() {
1815 let config = LlamaCppConfig {
1816 gpu_arch: "gfx1100".to_string(),
1817 channel: "stable".to_string(),
1818 ..Default::default()
1819 };
1820 let installer = LlamaCppInstaller::new(config);
1821 let flags = installer.cmake_flags();
1822
1823 assert!(flags.contains(&"-DGGML_HIP=ON".to_string()));
1824 assert!(flags.contains(&"-DGPU_TARGETS=gfx1030;gfx1100;gfx1101".to_string()));
1825 assert!(flags.contains(&"-DGGML_HIP_RDNA3_PROBES=OFF".to_string()));
1826 assert!(flags.contains(&"-DGGML_HIP_ROCWMMA_FATTN=ON".to_string()));
1827 }
1828
1829 #[test]
1830 fn test_cmake_flags_latest_with_rdna4() {
1831 let config = LlamaCppConfig {
1832 gpu_arch: "gfx1200".to_string(),
1833 channel: "latest".to_string(),
1834 ..Default::default()
1835 };
1836 let installer = LlamaCppInstaller::new(config);
1837 let flags = installer.cmake_flags();
1838
1839 assert!(flags.contains(&"-DGGML_HIP=ON".to_string()));
1840 assert!(flags.contains(&"-DGPU_TARGETS=gfx1030;gfx1100;gfx1101;gfx1200".to_string()));
1841 assert!(flags.contains(&"-DGGML_HIP_RDNA3_PROBES=OFF".to_string()));
1842 assert!(flags.contains(&"-DGGML_HIP_ROCWMMA_FATTN=ON".to_string()));
1843 }
1844
1845 #[test]
1846 fn test_rocm_channel_rdna3_gating() {
1847 let legacy = LlamaCppInstaller::new(LlamaCppConfig {
1848 gpu_arch: "gfx1100".to_string(),
1849 channel: "legacy".to_string(),
1850 ..Default::default()
1851 });
1852 let stable = LlamaCppInstaller::new(LlamaCppConfig {
1853 gpu_arch: "gfx1100".to_string(),
1854 channel: "stable".to_string(),
1855 ..Default::default()
1856 });
1857 let latest = LlamaCppInstaller::new(LlamaCppConfig {
1858 gpu_arch: "gfx1100".to_string(),
1859 channel: "latest".to_string(),
1860 ..Default::default()
1861 });
1862
1863 let legacy_flags = legacy.cmake_flags();
1864 let stable_flags = stable.cmake_flags();
1865 let latest_flags = latest.cmake_flags();
1866
1867 assert!(!legacy_flags
1868 .iter()
1869 .any(|flag| flag.contains("RDNA3_PROBES") || flag.contains("ROCWMMA_FATTN")));
1870 assert!(stable_flags.contains(&"-DGGML_HIP_RDNA3_PROBES=OFF".to_string()));
1871 assert!(stable_flags.contains(&"-DGGML_HIP_ROCWMMA_FATTN=ON".to_string()));
1872 assert!(latest_flags.contains(&"-DGGML_HIP_RDNA3_PROBES=OFF".to_string()));
1873 assert!(latest_flags.contains(&"-DGGML_HIP_ROCWMMA_FATTN=ON".to_string()));
1874 }
1875
1876 #[test]
1877 fn test_cmake_flags_unknown_gpu_conservative() {
1878 let config = LlamaCppConfig {
1879 gpu_arch: "unknown".to_string(),
1880 channel: "latest".to_string(),
1881 ..Default::default()
1882 };
1883 let installer = LlamaCppInstaller::new(config);
1884 let flags = installer.cmake_flags();
1885
1886 assert!(flags.contains(&"-DGPU_TARGETS=gfx1030".to_string()));
1888 }
1889
1890 #[test]
1891 fn test_detect_cuda_toolkit_presence_no_panic() {
1892 let result = detect_cuda_toolkit_presence();
1894 assert!(result.len() <= 3);
1895 }
1896
1897 #[test]
1898 fn test_validate_cmake_cache_accepts_rocm_only_config() {
1899 let temp = tempfile::NamedTempFile::new().unwrap();
1901 fs::write(
1902 temp.path(),
1903 "GGML_HIP:BOOL=ON\nGGML_CUDA:BOOL=OFF\nGPU_TARGETS:STRING=gfx1030;gfx1100;gfx1101\n",
1904 )
1905 .unwrap();
1906
1907 assert!(validate_cmake_cache(temp.path(), "gfx1030;gfx1100;gfx1101").is_ok());
1908 }
1909
1910 #[test]
1911 fn test_validate_cmake_cache_accepts_uninitialized_gpu_targets() {
1912 let temp = tempfile::NamedTempFile::new().unwrap();
1917 fs::write(
1918 temp.path(),
1919 "GGML_HIP:BOOL=ON\nGGML_CUDA:BOOL=OFF\nGPU_TARGETS:UNINITIALIZED=gfx1030;gfx1100;gfx1101\n",
1920 )
1921 .unwrap();
1922
1923 assert!(validate_cmake_cache(temp.path(), "gfx1030;gfx1100;gfx1101").is_ok());
1924 }
1925
1926 #[test]
1927 fn test_validate_cmake_cache_rejects_cuda_on() {
1928 let temp = tempfile::NamedTempFile::new().unwrap();
1930 fs::write(
1931 temp.path(),
1932 "GGML_HIP:BOOL=ON\nGGML_CUDA:BOOL=ON\nGPU_TARGETS:STRING=gfx1030\n",
1933 )
1934 .unwrap();
1935
1936 let err = validate_cmake_cache(temp.path(), "gfx1030").unwrap_err();
1937 assert!(err.contains("GGML_CUDA is not OFF"));
1938 }
1939
1940 #[test]
1941 fn test_validate_cmake_cache_rejects_gpu_target_mismatch() {
1942 let temp = tempfile::NamedTempFile::new().unwrap();
1944 fs::write(
1945 temp.path(),
1946 "GGML_HIP:BOOL=ON\nGGML_CUDA:BOOL=OFF\nGPU_TARGETS:STRING=gfx1030\n",
1947 )
1948 .unwrap();
1949
1950 let err = validate_cmake_cache(temp.path(), "gfx1030;gfx1100;gfx1101").unwrap_err();
1951 assert!(err.contains("GPU_TARGETS mismatch"));
1952 }
1953
1954 #[test]
1955 fn test_build_commands_sequence() {
1956 let config = LlamaCppConfig {
1957 gpu_arch: "gfx1100".to_string(),
1958 channel: "latest".to_string(),
1959 ..Default::default()
1960 };
1961 let installer = LlamaCppInstaller::new(config);
1962 let commands = installer.build_commands("/home/testuser");
1963
1964 assert_eq!(commands.len(), 6);
1967 assert_eq!(commands[0].program, "sh");
1969 assert!(commands[0]
1970 .args
1971 .iter()
1972 .any(|a| a.contains("rm -rf") && a.contains("mkdir -p")));
1973 assert_eq!(commands[1].program, "git");
1974 assert_eq!(commands[2].program, "git");
1975 assert_eq!(commands[3].program, "cmake");
1976 assert_eq!(commands[4].program, "cmake");
1978 assert!(
1979 commands[4].args.iter().any(|a| a == "--target")
1980 && commands[4].args.iter().any(|a| a == "install")
1981 );
1982 assert_eq!(commands[5].program, "sh");
1983 }
1984
1985 #[test]
1986 fn test_detection_binary_path() {
1987 let config = LlamaCppConfig::default();
1988 let path = config.detection_binary_path("/home/testuser");
1989 assert_eq!(
1990 path,
1991 PathBuf::from("/home/testuser/.mlstack/components/llama-cpp/bin/llama-cli")
1992 );
1993 }
1994
1995 #[test]
1996 fn test_resolve_install_bin_dir() {
1997 let dir = resolve_install_bin_dir("/home/testuser");
1998 assert_eq!(
1999 dir,
2000 PathBuf::from("/home/testuser/.mlstack/components/llama-cpp/bin")
2001 );
2002 }
2003
2004 #[test]
2005 fn test_is_llama_cli_functional_absent() {
2006 assert!(!is_llama_cli_functional(
2008 "/nonexistent/path/that/does/not/exist"
2009 ));
2010 }
2011
2012 #[test]
2013 fn test_has_partial_artifacts_no_artifacts() {
2014 let _env = crate::test_support::lock_env();
2018 assert!(!has_partial_artifacts(
2022 "/nonexistent-path-DOES-NOT-exist-f94a1f2a3b5c/test-home"
2023 ));
2024 }
2025
2026 #[test]
2027 fn test_dependencies() {
2028 let installer = LlamaCppInstaller::with_defaults();
2029 assert!(installer.dependencies().contains(&"rocm"));
2030 }
2031
2032 #[test]
2033 fn test_repo_url_default() {
2034 let config = LlamaCppConfig::default();
2035 assert_eq!(config.repo_url(), DEFAULT_REPO_URL);
2036 }
2037
2038 #[test]
2039 fn test_repo_url_override() {
2040 let config = LlamaCppConfig {
2041 repo_url: Some("https://example.com/fork.git".to_string()),
2042 ..Default::default()
2043 };
2044 assert_eq!(config.repo_url(), "https://example.com/fork.git");
2045 }
2046
2047 #[test]
2048 fn test_rocm_channel_from_str() {
2049 assert_eq!(RocmChannel::parse_channel("legacy"), RocmChannel::Legacy);
2050 assert_eq!(RocmChannel::parse_channel("LEGACY"), RocmChannel::Legacy);
2051 assert_eq!(RocmChannel::parse_channel("stable"), RocmChannel::Stable);
2052 assert_eq!(RocmChannel::parse_channel("latest"), RocmChannel::Latest);
2053 assert_eq!(RocmChannel::parse_channel("unknown"), RocmChannel::Latest);
2054 }
2055
2056 #[test]
2057 fn test_rocm_channel_label() {
2058 assert_eq!(RocmChannel::Legacy.label(), "legacy");
2059 assert_eq!(RocmChannel::Stable.label(), "stable");
2060 assert_eq!(RocmChannel::Latest.label(), "latest");
2061 }
2062
2063 #[test]
2064 fn test_build_commands_with_auth_token_handling() {
2065 let _original_token = std::env::var("GITHUB_TOKEN").ok();
2067
2068 let config = LlamaCppConfig::default();
2070 let installer = LlamaCppInstaller::new(config);
2071 let commands_with_token = installer
2072 .build_commands_with_auth("/home/testuser", Some(SealedToken::new("mock_token")));
2073
2074 let clone_cmd = &commands_with_token[1];
2076 assert!(!clone_cmd.args.iter().any(|arg| arg.contains("mock_token")));
2078 assert!(clone_cmd.args.iter().any(|arg| arg == DEFAULT_REPO_URL));
2080 assert!(clone_cmd
2082 .args
2083 .iter()
2084 .any(|arg| arg.starts_with("credential.helper=")));
2085 assert!(clone_cmd
2087 .env
2088 .iter()
2089 .any(|(k, v)| k == "GITHUB_TOKEN" && v == "mock_token"));
2090 assert!(clone_cmd
2091 .env
2092 .iter()
2093 .any(|(k, v)| k == "GIT_TERMINAL_PROMPT" && v == "0"));
2094 }
2095
2096 #[test]
2097 fn test_release_manifest_parse_and_lookup() {
2098 let json = r#"{
2099 "version": "v0.3.2",
2100 "archs": {
2101 "gfx1030": {"url": "https://example.com/a", "sha256": "aaa"},
2102 "gfx1100": {"url": "https://example.com/b", "sha256": "bbb"},
2103 "gfx1200": {"url": "https://example.com/c", "sha256": "ccc"}
2104 }
2105 }"#;
2106
2107 let manifest: ReleaseManifest = serde_json::from_str(json).unwrap();
2108 let installer = LlamaCppInstaller::with_defaults();
2109 assert_eq!(manifest.version, "v0.3.2");
2110 assert_eq!(
2111 installer
2112 .release_asset_for_arch(&manifest, "gfx1030")
2113 .unwrap()
2114 .url,
2115 "https://example.com/a"
2116 );
2117 assert_eq!(
2118 installer
2119 .release_asset_for_arch(&manifest, "gfx1100")
2120 .unwrap()
2121 .sha256,
2122 "bbb"
2123 );
2124 assert!(installer
2125 .release_asset_for_arch(&manifest, "gfx9999")
2126 .is_none());
2127 }
2128
2129 #[test]
2130 fn test_release_manifest_rejects_unknown_fields() {
2131 let json = r#"{
2132 "version": "v0.3.2",
2133 "unexpected": true,
2134 "archs": {
2135 "gfx1030": {"url": "https://example.com/a", "sha256": "aaa"},
2136 "gfx1100": {"url": "https://example.com/b", "sha256": "bbb"},
2137 "gfx1200": {"url": "https://example.com/c", "sha256": "ccc"}
2138 }
2139 }"#;
2140
2141 assert!(serde_json::from_str::<ReleaseManifest>(json).is_err());
2142 }
2143
2144 #[test]
2145 fn test_verify_sha256_matches() {
2146 let temp = tempfile::NamedTempFile::new().unwrap();
2147 fs::write(temp.path(), b"hello").unwrap();
2148 let expected = Sha256::digest(b"hello")
2149 .iter()
2150 .map(|byte| format!("{:02x}", byte))
2151 .collect::<String>();
2152 assert!(LlamaCppInstaller::verify_sha256(temp.path(), &expected).unwrap());
2153 }
2154
2155 #[test]
2156 fn test_verify_sha256_mismatch() {
2157 let temp = tempfile::NamedTempFile::new().unwrap();
2158 fs::write(temp.path(), b"hello").unwrap();
2159 assert!(!LlamaCppInstaller::verify_sha256(temp.path(), "deadbeef").unwrap());
2160 }
2161
2162 #[test]
2163 fn test_download_failure_falls_back_to_source() {
2164 let installer = LlamaCppInstaller::with_defaults();
2165 let plan = PrebuiltInstallPlan {
2166 manifest_version: "v1".into(),
2167 arch: "gfx1100".into(),
2168 url: "http://127.0.0.1:9/nope".into(),
2169 sha256: "deadbeef".into(),
2170 install_prefix: tempfile::tempdir().unwrap().path().display().to_string(),
2171 binary_path: PathBuf::from("/tmp/llama-cli"),
2172 };
2173 assert!(installer.download_prebuilt_binary(&plan).is_err());
2174 }
2175
2176 #[test]
2177 fn test_purge_source_artifacts_is_idempotent() {
2178 let _env = crate::test_support::lock_env();
2182 let installer = LlamaCppInstaller::with_defaults();
2183 let temp_home = tempfile::tempdir().unwrap();
2184 let home = temp_home.path().to_string_lossy().to_string();
2185
2186 std::fs::create_dir_all(format!("{home}/.mlstack/components/llama-cpp")).unwrap();
2187 std::fs::create_dir_all(format!("{home}/.mlstack/components/llama-cpp/.git")).unwrap();
2188 std::fs::create_dir_all("/tmp/llama-cpp-rocm-build").unwrap();
2189 std::fs::create_dir_all(format!("{home}/.mlstack/components/llama-cpp/bin")).unwrap();
2190 std::fs::write(
2191 format!("{home}/.mlstack/components/llama-cpp/bin/llama-cli"),
2192 "#!/bin/sh\nexit 0\n",
2193 )
2194 .unwrap();
2195
2196 assert!(installer.purge_source_artifacts(&home, true).is_ok());
2197 assert!(installer.purge_source_artifacts(&home, true).is_ok());
2198
2199 let _ = std::fs::remove_dir_all("/tmp/llama-cpp-rocm-build");
2202 }
2203
2204 #[test]
2205 fn test_git_clone_with_auth() {
2206 let config = LlamaCppConfig::default();
2207 let installer = LlamaCppInstaller::new(config);
2208 let command = installer.git_clone_with_auth("/tmp/test", "mock_token");
2209
2210 assert!(!command.args.iter().any(|arg| arg.contains("mock_token")));
2212 assert!(command.args.iter().any(|arg| arg == DEFAULT_REPO_URL));
2214 assert!(command.args.iter().any(|arg| arg == "-c"));
2216 assert!(command
2217 .args
2218 .iter()
2219 .any(|arg| arg.starts_with("credential.helper=")));
2220 assert!(command
2222 .env
2223 .iter()
2224 .any(|(k, v)| k == "GITHUB_TOKEN" && v == "mock_token"));
2225 assert!(command
2226 .env
2227 .iter()
2228 .any(|(k, v)| k == "GIT_TERMINAL_PROMPT" && v == "0"));
2229 }
2230
2231 #[test]
2237 fn test_post_install_verification_no_binaries() {
2238 let result =
2239 verify_installed_binary("/nonexistent/path/that/does/not/exist", "/nonexistent/fork");
2240 assert!(!result.help_check_passed);
2241 assert!(!result.stronger_check_attempted);
2242 assert!(!result.rdna3_check_attempted);
2243 assert!(!result.summary.is_empty());
2244 assert!(result.summary.contains("FAILED"));
2245 }
2246
2247 #[test]
2248 fn test_post_install_verification_struct_fields() {
2249 let result =
2250 verify_installed_binary("/nonexistent/path/that/does/not/exist", "/nonexistent/fork");
2251 assert!(!result.help_check_passed);
2252 assert!(!result.stronger_check_attempted);
2253 assert!(!result.stronger_check_passed);
2254 assert_eq!(result.bench_prefill_tps, None);
2255 assert_eq!(result.bench_decode_tps, None);
2256 assert!(!result.rdna3_check_attempted);
2257 assert_eq!(result.rdna3_is_valid, None);
2258 assert_eq!(result.rdna3_wmma_tps, None);
2259 assert_eq!(result.is_rdna3_device, None);
2260 assert_eq!(result.wmma_supported, None);
2261 assert!(result.verified_binary_path.contains("llama-cli"));
2262 }
2263
2264 #[test]
2265 fn test_post_install_verification_cli_only() {
2266 let tmp = tempfile::tempdir().unwrap();
2267 let bin_dir = tmp.path().join(".mlstack/components/llama-cpp/bin");
2268 std::fs::create_dir_all(&bin_dir).unwrap();
2269 std::fs::write(
2270 bin_dir.join("llama-cli"),
2271 "#!/bin/sh\necho 'usage: llama-cli'\nexit 0\n",
2272 )
2273 .unwrap();
2274 #[cfg(unix)]
2275 {
2276 use std::os::unix::fs::PermissionsExt;
2277 std::fs::set_permissions(
2278 bin_dir.join("llama-cli"),
2279 std::fs::Permissions::from_mode(0o755),
2280 )
2281 .unwrap();
2282 }
2283
2284 let result = verify_installed_binary(tmp.path().to_str().unwrap(), "/nonexistent/fork");
2285 assert!(result.help_check_passed);
2286 assert!(!result.stronger_check_attempted);
2287 assert!(!result.rdna3_check_attempted);
2288 assert!(result.summary.contains("PASSED"));
2289 assert!(result.summary.contains("SKIPPED"));
2290 }
2291
2292 #[test]
2293 fn test_parse_bench_tps_valid() {
2294 let line =
2295 "| qwen3 0.6B F16 | 1.11 GiB | 596.05 M | ROCm | 33 | pp128 | 6254.94 ± 1534.82 |";
2296 let tps = parse_bench_tps(line);
2297 assert!(tps.is_some());
2298 assert!((tps.unwrap() - 6254.94).abs() < 0.01);
2299 }
2300
2301 #[test]
2302 fn test_parse_bench_tps_no_uncertainty() {
2303 let line = "| qwen3 0.6B F16 | 1.11 GiB | 596.05 M | ROCm | 33 | tg32 | 152.17 |";
2304 let tps = parse_bench_tps(line);
2305 assert!(tps.is_some());
2306 assert!((tps.unwrap() - 152.17).abs() < 0.01);
2307 }
2308
2309 #[test]
2310 fn test_parse_bench_tps_invalid_line() {
2311 assert_eq!(parse_bench_tps("not a bench line"), None);
2312 assert_eq!(parse_bench_tps(""), None);
2313 assert_eq!(parse_bench_tps("| col1 | col2 |"), None);
2314 }
2315
2316 #[test]
2317 fn test_verify_binary_linkage_missing_binary() {
2318 let result = verify_binary_linkage(std::path::Path::new("/nonexistent/binary"));
2320 assert!(result.is_err());
2321 assert!(result.unwrap_err().contains("not found"));
2322 }
2323
2324 #[test]
2325 fn test_verify_binary_linkage_no_panic() {
2326 let result = verify_binary_linkage(std::path::Path::new("/usr/bin/bash"));
2329 assert!(result.is_err());
2331 }
2332}