rusty_stack/installers/components/
onnxruntime.rs

1//! ONNX Runtime installer — ports `scripts/build_onnxruntime_multi.sh`.
2//!
3//! Constructs correct cmake commands with ROCm integration and Python bindings,
4//! builds from source using ONNX Runtime's build.sh, and pip installs from
5//! the built wheel. Falls back to prebuilt wheel if available.
6//!
7//! ONNX Runtime depends on ROCm.
8//!
9//! # Validation Assertions
10//!
11//! - **VAL-INSTALL-010**: ONNX Runtime correct cmake command
12//! - **VAL-INSTALL-011**: ONNX Runtime pip install from build output
13//! - **VAL-INSTALL-014**: Source builds parse build output for errors
14//! - **VAL-INSTALL-045**: ONNX Runtime declares dependency on ROCm
15
16use crate::installers::common::RocmEnv;
17use std::path::PathBuf;
18
19// ===========================================================================
20// Types
21// ===========================================================================
22
23/// GPU architecture specification for HIP builds.
24#[derive(Debug, Clone, PartialEq, Eq)]
25pub enum HipArchs {
26    /// gfx1030 family.
27    Gfx1030,
28    /// gfx1100 family (gfx1100-gfx1103).
29    Gfx1100,
30    /// gfx1200 family (gfx1200-gfx1201).
31    Gfx1200,
32    /// Default fallback.
33    Default,
34}
35
36impl HipArchs {
37    /// Detect HIP architectures from a GPU arch string.
38    pub fn from_gpu_arch(gpu_arch: &str) -> Self {
39        if gpu_arch.starts_with("gfx103") {
40            HipArchs::Gfx1030
41        } else if gpu_arch.starts_with("gfx110") {
42            HipArchs::Gfx1100
43        } else if gpu_arch.starts_with("gfx120") {
44            HipArchs::Gfx1200
45        } else {
46            HipArchs::Default
47        }
48    }
49
50    /// Get the CMAKE_HIP_ARCHITECTURES string.
51    pub fn cmake_hip_architectures(&self) -> &str {
52        match self {
53            HipArchs::Gfx1030 => "gfx1030",
54            HipArchs::Gfx1100 => "gfx1100;gfx1101;gfx1102;gfx1103",
55            HipArchs::Gfx1200 => "gfx1200;gfx1201",
56            HipArchs::Default => "gfx1100",
57        }
58    }
59}
60
61/// ONNX Runtime install method.
62///
63/// **Option ordering:** [`MigraphxWheel`] is the primary/default (PyPI
64/// `onnxruntime-migraphx`, current release); [`PrebuiltWheel`] is the legacy
65/// fallback; [`SourceBuild`] is the last resort (custom/when no wheel exists).
66#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
67pub enum OnnxInstallMethod {
68    /// **Option 1 (primary, default):** install the pinned
69    /// `onnxruntime-migraphx` wheel from PyPI ([`PREBUILT_MIGRAPHX_VERSION`],
70    /// currently 1.27.1). PyPI's Microsoft-published wheel ships the
71    /// MIGraphXExecutionProvider built against ROCm and tracks the current
72    /// release. The AMD manylinux repo (`repo.radeon.com`) is NOT used — it lags
73    /// far behind (1.23.2 for ROCm 7.2.4) and fabricates 404 URLs for newer
74    /// versions.
75    #[default]
76    MigraphxWheel,
77    /// **Option 3 (last resort):** build from source with ROCm + MIGraphX EP
78    /// (includes ROCMExecutionProvider), targeting [`DEFAULT_ONNXRUNTIME_VERSION`].
79    /// Use only when the PyPI wheel is unavailable or a custom build is required.
80    SourceBuild,
81    /// **Option 2 (legacy):** install prebuilt `onnxruntime-rocm` from PyPI.
82    /// May be ABI-incompatible with the installed ROCm; prefer [`MigraphxWheel`].
83    PrebuiltWheel,
84}
85
86/// Current ROCm-compatible ONNX Runtime release: the source-build target and
87/// the version the bundled manifest proposes.
88pub const DEFAULT_ONNXRUNTIME_VERSION: &str = "1.27.1";
89
90/// Pinned prebuilt `onnxruntime-migraphx` wheel version installed from PyPI by
91/// the default ([`OnnxInstallMethod::MigraphxWheel`]) path — the primary,
92/// production-proven ROCm build that ships the MIGraphXExecutionProvider.
93///
94/// Tracks the current PyPI release line (1.27.1) and is kept in lock-step with
95/// [`DEFAULT_ONNXRUNTIME_VERSION`] and the manifest's `onnx` target version
96/// (`baseline_manifest.json`). That lock-step is load-bearing: the default
97/// MIGraphX-wheel install path must land on exactly the version the post-install
98/// honesty guard verifies against — otherwise an explicit
99/// `rusty-stack update onnx` would install this pinned wheel but be reported as
100/// failed because the detected version mismatches the manifest target.
101///
102/// The Rust `ort` crate's MIGraphX execution-provider builder exposes no
103/// `with_model_cache_dir` / model-cache API in ANY released version (newest is
104/// `ort 2.0.0-rc.12`), and the v2.0 "arbitrarily configurable" EP change left
105/// MIGraphX out — so the version choice is driven by manifest/honesty-guard
106/// alignment rather than an `ort` API benefit. The source-build path
107/// ([`OnnxInstallMethod::SourceBuild`]) targets
108/// [`DEFAULT_ONNXRUNTIME_VERSION`] when a custom build is actually required.
109/// The AMD manylinux repo (`repo.radeon.com/rocm-rel-<release>`) is NOT used —
110/// it lags far behind (still 1.23.2 for ROCm 7.2.4) and fabricates 404 URLs for
111/// newer versions. Bump only when a newer wheel is verified to add real value
112/// (and bump all three in lock-step: this constant,
113/// [`DEFAULT_ONNXRUNTIME_VERSION`], and the manifest target).
114pub const PREBUILT_MIGRAPHX_VERSION: &str = "1.27.1";
115
116/// Env var selecting the ONNX install method
117/// (`migraphx`/`source`/`prebuilt`). Default: [`OnnxInstallMethod::MigraphxWheel`].
118pub const ONNX_INSTALL_METHOD_ENV: &str = "MLSTACK_ONNX_INSTALL_METHOD";
119
120/// Env var overriding the prebuilt `onnxruntime-migraphx` version (else the
121/// pinned [`PREBUILT_MIGRAPHX_VERSION`] is used).
122pub const ONNX_VERSION_ENV: &str = "MLSTACK_ONNX_VERSION";
123
124/// Resolve the ONNX install method from [`ONNX_INSTALL_METHOD_ENV`].
125pub fn install_method_from_env() -> OnnxInstallMethod {
126    match std::env::var(ONNX_INSTALL_METHOD_ENV)
127        .unwrap_or_default()
128        .trim()
129        .to_ascii_lowercase()
130        .as_str()
131    {
132        "source" | "source-build" | "src" => OnnxInstallMethod::SourceBuild,
133        "prebuilt" | "rocm" | "onnxruntime-rocm" => OnnxInstallMethod::PrebuiltWheel,
134        _ => OnnxInstallMethod::MigraphxWheel,
135    }
136}
137
138/// Resolve an overridden prebuilt version from [`ONNX_VERSION_ENV`], else `None`.
139pub fn prebuilt_version_from_env() -> Option<String> {
140    std::env::var(ONNX_VERSION_ENV)
141        .ok()
142        .map(|s| s.trim().to_string())
143        .filter(|s| !s.is_empty())
144}
145
146/// Configuration for the ONNX Runtime installer.
147#[derive(Debug, Clone)]
148pub struct OnnxRuntimeConfig {
149    /// Python binary to use.
150    pub python_bin: String,
151    /// Whether to run in dry-run mode.
152    pub dry_run: bool,
153    /// ROCm version string (e.g., "7.2.0").
154    pub rocm_version: Option<String>,
155    /// ROCm release for AMD repo URL (e.g., "7.2.4").
156    pub rocm_release: Option<String>,
157    /// GPU architecture string (e.g., "gfx1100").
158    pub gpu_arch: Option<String>,
159    /// Working directory for the build (defaults to /tmp/onnxruntime-rocm).
160    pub workdir: Option<PathBuf>,
161    /// Preinstalled Eigen path.
162    pub eigen_path: Option<PathBuf>,
163    /// Whether to use preinstalled Eigen.
164    pub use_preinstalled_eigen: bool,
165    /// Install method (default: MIGraphX wheel from PyPI).
166    pub install_method: OnnxInstallMethod,
167    /// ONNX Runtime wheel version. `None` uses the current bundled default.
168    pub runtime_version: Option<String>,
169    /// Pinned prebuilt `onnxruntime-migraphx` version for the default path.
170    /// `None` uses [`PREBUILT_MIGRAPHX_VERSION`].
171    pub prebuilt_version: Option<String>,
172}
173
174impl Default for OnnxRuntimeConfig {
175    fn default() -> Self {
176        Self {
177            python_bin: "python3".to_string(),
178            dry_run: false,
179            rocm_version: None,
180            rocm_release: None,
181            gpu_arch: None,
182            workdir: None,
183            eigen_path: None,
184            use_preinstalled_eigen: false,
185            install_method: OnnxInstallMethod::default(),
186            runtime_version: None,
187            prebuilt_version: None,
188        }
189    }
190}
191
192impl OnnxRuntimeConfig {
193    /// Get the working directory for the build.
194    pub fn workdir(&self) -> PathBuf {
195        self.workdir
196            .clone()
197            .unwrap_or_else(|| std::env::temp_dir().join("onnxruntime-rocm"))
198    }
199
200    /// Get the effective GPU arch string.
201    pub fn gpu_arch(&self) -> &str {
202        self.gpu_arch.as_deref().unwrap_or("gfx1100")
203    }
204
205    /// Get the effective ROCm version string.
206    pub fn rocm_version(&self) -> &str {
207        self.rocm_version.as_deref().unwrap_or("7.2")
208    }
209
210    /// Get the effective ROCm release string (for AMD repo URL).
211    pub fn rocm_release(&self) -> &str {
212        self.rocm_release.as_deref().unwrap_or("7.2.4")
213    }
214
215    /// Get the effective ONNX Runtime wheel version.
216    ///
217    /// Filters an empty/whitespace `Some("")` to the default: a stale or
218    /// untrimmed `ctx.target_version` would otherwise produce an invalid
219    /// `git checkout v` (no tag suffix) on the source-build path. Mirrors
220    /// [`prebuilt_version_from_env`]'s empty-filtering.
221    pub fn runtime_version(&self) -> &str {
222        match self.runtime_version.as_deref() {
223            Some(v) if !v.trim().is_empty() => v,
224            _ => DEFAULT_ONNXRUNTIME_VERSION,
225        }
226    }
227
228    /// Get the pinned prebuilt `onnxruntime-migraphx` version for the default
229    /// PyPI install path.
230    pub fn prebuilt_version(&self) -> &str {
231        self.prebuilt_version
232            .as_deref()
233            .unwrap_or(PREBUILT_MIGRAPHX_VERSION)
234    }
235}
236
237/// A constructed shell command.
238#[derive(Debug, Clone, PartialEq, Eq)]
239pub struct ShellCommand {
240    /// The program to run.
241    pub program: String,
242    /// Arguments to pass.
243    pub args: Vec<String>,
244    /// Environment variables to set.
245    pub env: Vec<(String, String)>,
246    /// Working directory for the command.
247    pub working_dir: Option<PathBuf>,
248}
249
250/// The ONNX Runtime installer.
251pub struct OnnxRuntimeInstaller {
252    config: OnnxRuntimeConfig,
253}
254
255impl OnnxRuntimeInstaller {
256    /// Create a new ONNX Runtime installer with the given config.
257    pub fn new(config: OnnxRuntimeConfig) -> Self {
258        Self { config }
259    }
260
261    /// Create with default config.
262    pub fn with_defaults() -> Self {
263        Self::new(OnnxRuntimeConfig::default())
264    }
265
266    /// Get the requested/effective runtime version for status reporting.
267    pub fn runtime_version(&self) -> &str {
268        self.config.runtime_version()
269    }
270
271    /// Get the pinned prebuilt `onnxruntime-migraphx` version installed by the
272    /// default path.
273    pub fn prebuilt_version(&self) -> &str {
274        self.config.prebuilt_version()
275    }
276
277    // -----------------------------------------------------------------------
278    // Dependencies (VAL-INSTALL-045)
279    // -----------------------------------------------------------------------
280
281    /// Get the list of required dependencies.
282    ///
283    /// ONNX Runtime depends on ROCm.
284    pub fn dependencies(&self) -> &[&str] {
285        &["rocm"]
286    }
287
288    // -----------------------------------------------------------------------
289    // Prebuilt wheel install command
290    // -----------------------------------------------------------------------
291
292    /// Construct the pip install command for prebuilt onnxruntime-rocm wheel.
293    ///
294    /// The original script tries this first:
295    /// `pip install --upgrade --prefer-binary onnxruntime-rocm`
296    pub fn build_prebuilt_install_command(&self) -> ShellCommand {
297        ShellCommand {
298            program: self.config.python_bin.clone(),
299            args: vec![
300                "-m".to_string(),
301                "pip".to_string(),
302                "install".to_string(),
303                "--upgrade".to_string(),
304                "--force-reinstall".to_string(),
305                "--no-deps".to_string(),
306                "--prefer-binary".to_string(),
307                "onnxruntime-rocm".to_string(),
308            ],
309            env: vec![],
310            working_dir: None,
311        }
312    }
313
314    // -----------------------------------------------------------------------
315    // MIGraphX wheel install from PyPI (default)
316    // -----------------------------------------------------------------------
317
318    /// Construct the pip install command for the pinned `onnxruntime-migraphx`
319    /// wheel from PyPI.
320    ///
321    /// This is the default install method. The Microsoft-published PyPI wheel
322    /// ships the MIGraphX execution provider built against ROCm. `--no-deps`
323    /// honors the No-CUDA hard-prime (prevents pulling nvidia/cuda runtime
324    /// transitive deps); `--no-cache-dir` forces a fresh fetch.
325    pub fn build_migraphx_install_command(&self) -> ShellCommand {
326        let spec = format!("onnxruntime-migraphx=={}", self.config.prebuilt_version());
327        ShellCommand {
328            program: self.config.python_bin.clone(),
329            args: vec![
330                "-m".to_string(),
331                "pip".to_string(),
332                "install".to_string(),
333                "--upgrade".to_string(),
334                "--force-reinstall".to_string(),
335                "--no-deps".to_string(),
336                "--no-cache-dir".to_string(),
337                spec,
338            ],
339            env: vec![],
340            working_dir: None,
341        }
342    }
343
344    /// Construct a provider validation command for AMD ONNX Runtime.
345    pub fn build_provider_validation_command(&self) -> ShellCommand {
346        let script = r#"
347import ctypes
348import os
349import onnxruntime as ort
350
351if not getattr(ort, "__version__", None) or not hasattr(ort, "get_available_providers"):
352    raise SystemExit(f"ONNX Runtime import incomplete: {getattr(ort, '__file__', '<unknown>')}")
353
354priority = ["MIGraphXExecutionProvider", "ROCMExecutionProvider"]
355available = ort.get_available_providers()
356selected = next((p for p in priority if p in available), None)
357if selected is None:
358    capi = os.path.join(os.path.dirname(ort.__file__), "capi")
359    loader_errors = []
360    for lib in ("libonnxruntime_providers_migraphx.so", "libonnxruntime_providers_rocm.so"):
361        path = os.path.join(capi, lib)
362        if os.path.exists(path):
363            try:
364                ctypes.CDLL(path)
365            except OSError as exc:
366                loader_errors.append(f"{lib}: {exc}")
367    raise SystemExit(
368        "ONNX Runtime AMD provider unavailable; expected MIGraphXExecutionProvider "
369        f"or legacy ROCMExecutionProvider; available={available}; "
370        f"loader_errors={loader_errors or 'none'}"
371    )
372print(f"ONNX Runtime AMD provider ready: {selected}; available={available}")
373"#;
374
375        ShellCommand {
376            program: self.config.python_bin.clone(),
377            args: vec!["-c".to_string(), script.trim().to_string()],
378            env: vec![],
379            working_dir: None,
380        }
381    }
382
383    // -----------------------------------------------------------------------
384    // Model optimizer (ORT graph optimization for quantized models)
385    // -----------------------------------------------------------------------
386
387    /// Construct the Python command to run ORT graph optimization on an ONNX model.
388    ///
389    /// Applies `ORT_ENABLE_ALL` optimization level which fuses quantized ops
390    /// (DynamicQuantizeLinear, MatMulInteger, etc.) into custom ops that bypass
391    /// MIGraphX's broken kernels. The optimized model is saved alongside the original.
392    pub fn build_model_optimizer_command(&self, model_path: &str) -> ShellCommand {
393        let optimized_path = format!("{}.optimized", model_path);
394
395        let script = format!(
396            "import onnxruntime as ort; \
397             providers = [p for p in ['MIGraphXExecutionProvider', 'ROCMExecutionProvider'] if p in ort.get_available_providers()]; \
398             assert providers, f'No AMD ONNX Runtime provider available: {{ort.get_available_providers()}}'; \
399             opts = ort.SessionOptions(); \
400             opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL; \
401             opts.optimized_model_filepath = '{optimized_path}'; \
402             ort.InferenceSession('{model_path}', opts, providers=providers); \
403             print('Optimized: {model_path} -> {optimized_path} using ' + providers[0])"
404        );
405
406        ShellCommand {
407            program: self.config.python_bin.clone(),
408            args: vec!["-c".to_string(), script],
409            env: vec![],
410            working_dir: None,
411        }
412    }
413
414    // -----------------------------------------------------------------------
415    // Git clone command
416    // -----------------------------------------------------------------------
417
418    /// Construct the git clone command for ONNX Runtime.
419    ///
420    /// The original script clones:
421    /// `git clone --recursive https://github.com/microsoft/onnxruntime.git`
422    pub fn build_git_clone_command(&self) -> ShellCommand {
423        ShellCommand {
424            program: "git".to_string(),
425            args: vec![
426                "clone".to_string(),
427                "--recursive".to_string(),
428                "https://github.com/microsoft/onnxruntime.git".to_string(),
429            ],
430            env: vec![],
431            working_dir: Some(self.config.workdir()),
432        }
433    }
434
435    /// Construct the git checkout command for the release tag matching the
436    /// configured runtime version (e.g. `v1.27.1`).
437    pub fn build_git_checkout_command(&self) -> ShellCommand {
438        ShellCommand {
439            program: "git".to_string(),
440            args: vec![
441                "checkout".to_string(),
442                format!("v{}", self.config.runtime_version()),
443            ],
444            env: vec![],
445            working_dir: Some(self.config.workdir().join("onnxruntime")),
446        }
447    }
448
449    /// Construct the git submodule update command.
450    pub fn build_git_submodule_command(&self) -> ShellCommand {
451        ShellCommand {
452            program: "git".to_string(),
453            args: vec![
454                "submodule".to_string(),
455                "update".to_string(),
456                "--init".to_string(),
457                "--recursive".to_string(),
458            ],
459            env: vec![],
460            working_dir: Some(self.config.workdir().join("onnxruntime")),
461        }
462    }
463
464    // -----------------------------------------------------------------------
465    // Build command (VAL-INSTALL-010)
466    // -----------------------------------------------------------------------
467
468    /// Construct the ONNX Runtime build.sh command with ROCm flags.
469    ///
470    /// The original script runs `./build.sh` with many flags:
471    /// - `--config Release`
472    /// - `--build_wheel`
473    /// - `--parallel $(nproc)-1`
474    /// - `--use_rocm --rocm_home /opt/rocm`
475    /// - `--rocm_version 70200`
476    /// - `--use_migraphx --migraphx_home /opt/rocm`
477    /// - `--cmake_extra_defines CMAKE_HIP_ARCHITECTURES=...`
478    /// - Various CMAKE isolation defines
479    /// - `--allow_running_as_root`
480    pub fn build_build_command(&self, rocm_env: &RocmEnv) -> ShellCommand {
481        let rocm_path = rocm_env
482            .path()
483            .map(|p| p.to_string_lossy().to_string())
484            .unwrap_or_else(|| "/opt/rocm".to_string());
485
486        let hip_archs = HipArchs::from_gpu_arch(self.config.gpu_arch());
487        let nproc = std::thread::available_parallelism()
488            .map(|n| n.get().saturating_sub(1))
489            .unwrap_or(3);
490
491        // ONNX Runtime v1.23+ dropped `--use_rocm`/`--rocm_home`/`--rocm_version`
492        // (verified against v1.27.1's build.py — zero `use_rocm` references).
493        // The MIGraphX EP is built with `--use_migraphx --migraphx_home`; CMake
494        // locates ROCm via CMAKE_PREFIX_PATH (below) and targets the GPU arch
495        // via CMAKE_HIP_ARCHITECTURES. Passing the removed flags would make
496        // argparse reject the build outright.
497        let mut args = vec![
498            "--config".to_string(),
499            "Release".to_string(),
500            "--build_wheel".to_string(),
501            "--parallel".to_string(),
502            nproc.to_string(),
503            "--skip_tests".to_string(),
504            "--use_migraphx".to_string(),
505            "--migraphx_home".to_string(),
506            rocm_path.clone(),
507        ];
508
509        if self.config.use_preinstalled_eigen {
510            if let Some(ref eigen_path) = self.config.eigen_path {
511                args.push("--use_preinstalled_eigen".to_string());
512                args.push("--eigen_path".to_string());
513                args.push(eigen_path.to_string_lossy().to_string());
514            }
515        }
516
517        // CMAKE extra defines
518        let cmake_extra_defines = vec![
519            format!(
520                "CMAKE_HIP_ARCHITECTURES={}",
521                hip_archs.cmake_hip_architectures()
522            ),
523            "CMAKE_CXX_STANDARD=20".to_string(),
524            "onnxruntime_USE_EXTERNAL_ABSEIL=OFF".to_string(),
525            "CMAKE_DISABLE_FIND_PACKAGE_re2=ON".to_string(),
526            "CMAKE_FIND_PACKAGE_NO_PACKAGE_REGISTRY=ON".to_string(),
527            "CMAKE_FIND_PACKAGE_NO_SYSTEM_PACKAGE_REGISTRY=ON".to_string(),
528            "CMAKE_FIND_USE_PACKAGE_REGISTRY=OFF".to_string(),
529            "CMAKE_FIND_USE_SYSTEM_PACKAGE_REGISTRY=OFF".to_string(),
530            format!("CMAKE_PREFIX_PATH={}", rocm_path),
531            "re2_DIR=RE2_DIR-NOTFOUND".to_string(),
532            "CMAKE_POLICY_VERSION_MINIMUM=3.5".to_string(),
533        ];
534
535        for define in cmake_extra_defines {
536            args.push("--cmake_extra_defines".to_string());
537            args.push(define);
538        }
539
540        if self.config.use_preinstalled_eigen {
541            if let Some(ref eigen_path) = self.config.eigen_path {
542                args.push("--cmake_extra_defines".to_string());
543                args.push(format!(
544                    "FETCHCONTENT_SOURCE_DIR_EIGEN={}",
545                    eigen_path.to_string_lossy()
546                ));
547                args.push("--cmake_extra_defines".to_string());
548                args.push("FETCHCONTENT_TRY_FIND_PACKAGE_MODE=NEVER".to_string());
549            }
550        }
551
552        args.push("--allow_running_as_root".to_string());
553
554        let env = vec![];
555
556        ShellCommand {
557            program: "./build.sh".to_string(),
558            args,
559            env,
560            working_dir: Some(self.config.workdir().join("onnxruntime")),
561        }
562    }
563
564    // -----------------------------------------------------------------------
565    // Pip install from build output (VAL-INSTALL-011)
566    // -----------------------------------------------------------------------
567
568    /// Construct the pip uninstall command to remove all onnxruntime variants.
569    pub fn build_uninstall_command(&self) -> ShellCommand {
570        ShellCommand {
571            program: self.config.python_bin.clone(),
572            args: vec![
573                "-m".to_string(),
574                "pip".to_string(),
575                "uninstall".to_string(),
576                "-y".to_string(),
577                "onnxruntime".to_string(),
578                "onnxruntime-rocm".to_string(),
579                "onnxruntime-gpu".to_string(),
580                "onnxruntime-migraphx".to_string(),
581            ],
582            env: vec![],
583            working_dir: None,
584        }
585    }
586
587    /// Construct the pip install command for the built wheel.
588    ///
589    /// The original script installs from:
590    /// `build/Linux/Release/dist/*.whl`
591    pub fn build_wheel_install_command(&self) -> ShellCommand {
592        let wheel_path = self
593            .config
594            .workdir()
595            .join("onnxruntime")
596            .join("build/Linux/Release/dist/*.whl");
597
598        ShellCommand {
599            program: self.config.python_bin.clone(),
600            args: vec![
601                "-m".to_string(),
602                "pip".to_string(),
603                "install".to_string(),
604                wheel_path.to_string_lossy().to_string(),
605            ],
606            env: vec![],
607            working_dir: None,
608        }
609    }
610
611    // -----------------------------------------------------------------------
612    // Build error detection (VAL-INSTALL-014)
613    // -----------------------------------------------------------------------
614
615    /// Check build output for errors.
616    ///
617    /// Detects common build failure patterns and propagates as an error.
618    pub fn check_build_output(
619        &self,
620        stdout: &str,
621        stderr: &str,
622        exit_code: i32,
623    ) -> Result<(), String> {
624        if exit_code != 0 {
625            let error_patterns = [
626                "CMake Error",
627                "error:",
628                "FAILED:",
629                "fatal error:",
630                "Build failed",
631                "ninja: build stopped",
632                "Could NOT find",
633            ];
634
635            let combined = format!("{}\n{}", stdout, stderr);
636            let detected_errors: Vec<&str> = error_patterns
637                .iter()
638                .filter(|p| combined.contains(*p))
639                .copied()
640                .collect();
641
642            if detected_errors.is_empty() {
643                return Err(format!(
644                    "ONNX Runtime build failed with exit code {} (no specific error pattern detected)",
645                    exit_code
646                ));
647            }
648
649            return Err(format!(
650                "ONNX Runtime build failed with exit code {}. Detected errors: {}",
651                exit_code,
652                detected_errors.join(", ")
653            ));
654        }
655
656        Ok(())
657    }
658}
659
660// ===========================================================================
661// Tests
662// ===========================================================================
663
664#[cfg(test)]
665mod tests {
666    use super::*;
667    use std::path::PathBuf;
668
669    // --- VAL-INSTALL-010: ONNX Runtime correct cmake command ---
670
671    #[test]
672    fn test_build_command_uses_migraphx_recipe_not_use_rocm() {
673        // ONNX Runtime v1.23+ dropped --use_rocm/--rocm_home/--rocm_version
674        // (verified: v1.27.1 build.py has zero `use_rocm`). The source build
675        // uses --use_migraphx --migraphx_home + cmake defines (CMAKE_PREFIX_PATH
676        // + CMAKE_HIP_ARCHITECTURES).
677        let installer = OnnxRuntimeInstaller::new(OnnxRuntimeConfig {
678            rocm_version: Some("7.2.0".to_string()),
679            gpu_arch: Some("gfx1100".to_string()),
680            ..Default::default()
681        });
682        let rocm_env = RocmEnv::from_known(Some(PathBuf::from("/opt/rocm")), "7.2.0".to_string());
683        let cmd = installer.build_build_command(&rocm_env);
684
685        assert_eq!(cmd.program, "./build.sh");
686        assert!(cmd.args.contains(&"--use_migraphx".to_string()));
687        assert!(cmd.args.contains(&"--migraphx_home".to_string()));
688        assert!(cmd.args.contains(&"--build_wheel".to_string()));
689        assert!(cmd.args.contains(&"--skip_tests".to_string()));
690        assert!(cmd.args.contains(&"--allow_running_as_root".to_string()));
691        assert!(cmd.args.iter().any(|a| a == "/opt/rocm"));
692        // Removed flags must NOT appear (argparse would reject them outright).
693        for removed in ["--use_rocm", "--rocm_home", "--rocm_version"] {
694            assert!(
695                !cmd.args.contains(&removed.to_string()),
696                "removed flag {removed} must not appear: {:?}",
697                cmd.args
698            );
699        }
700    }
701
702    #[test]
703    fn test_build_command_has_hip_architectures() {
704        let installer = OnnxRuntimeInstaller::new(OnnxRuntimeConfig {
705            gpu_arch: Some("gfx1100".to_string()),
706            ..Default::default()
707        });
708        let rocm_env = RocmEnv::from_known(Some(PathBuf::from("/opt/rocm")), "7.2.0".to_string());
709        let cmd = installer.build_build_command(&rocm_env);
710
711        assert!(cmd
712            .args
713            .iter()
714            .any(|a| a.contains("CMAKE_HIP_ARCHITECTURES=gfx1100;gfx1101;gfx1102;gfx1103")));
715    }
716
717    #[test]
718    fn test_build_command_has_cmake_extra_defines() {
719        let installer = OnnxRuntimeInstaller::with_defaults();
720        let rocm_env = RocmEnv::from_known(Some(PathBuf::from("/opt/rocm")), "7.2.0".to_string());
721        let cmd = installer.build_build_command(&rocm_env);
722
723        assert!(cmd
724            .args
725            .iter()
726            .any(|a| a.contains("onnxruntime_USE_EXTERNAL_ABSEIL=OFF")));
727        assert!(cmd
728            .args
729            .iter()
730            .any(|a| a.contains("CMAKE_DISABLE_FIND_PACKAGE_re2=ON")));
731        assert!(cmd
732            .args
733            .iter()
734            .any(|a| a.contains("CMAKE_POLICY_VERSION_MINIMUM=3.5")));
735    }
736
737    // --- C++20 standard fix (fix-onnx-cmake-cxx20) ---
738
739    #[test]
740    fn test_build_command_includes_cmake_cxx_standard_20() {
741        let installer = OnnxRuntimeInstaller::with_defaults();
742        let rocm_env = RocmEnv::from_known(Some(PathBuf::from("/opt/rocm")), "7.2.0".to_string());
743        let cmd = installer.build_build_command(&rocm_env);
744
745        // The cmake command must include CMAKE_CXX_STANDARD=20 to fix
746        // C++20 standard library test failures during cmake configuration.
747        assert!(
748            cmd.args.iter().any(|a| a.contains("CMAKE_CXX_STANDARD=20")),
749            "build command must include CMAKE_CXX_STANDARD=20 cmake extra define, got args: {:?}",
750            cmd.args
751        );
752    }
753
754    #[test]
755    fn test_build_command_cmake_cxx_standard_is_cmake_extra_define() {
756        let installer = OnnxRuntimeInstaller::with_defaults();
757        let rocm_env = RocmEnv::from_known(Some(PathBuf::from("/opt/rocm")), "7.2.0".to_string());
758        let cmd = installer.build_build_command(&rocm_env);
759
760        // Find the CMAKE_CXX_STANDARD=20 arg and verify it's preceded by --cmake_extra_defines
761        let define_idx = cmd
762            .args
763            .iter()
764            .position(|a| a.contains("CMAKE_CXX_STANDARD=20"));
765        assert!(
766            define_idx.is_some(),
767            "CMAKE_CXX_STANDARD=20 not found in args"
768        );
769
770        let idx = define_idx.unwrap();
771        assert!(
772            idx > 0 && cmd.args[idx - 1] == "--cmake_extra_defines",
773            "CMAKE_CXX_STANDARD=20 must be preceded by --cmake_extra_defines"
774        );
775    }
776
777    #[test]
778    fn test_build_command_does_not_override_env_vars() {
779        let installer = OnnxRuntimeInstaller::with_defaults();
780        let rocm_env = RocmEnv::from_known(Some(PathBuf::from("/opt/rocm")), "7.2.0".to_string());
781        let cmd = installer.build_build_command(&rocm_env);
782
783        assert!(
784            cmd.env.is_empty(),
785            "ONNX build must inherit ~/.mlstack_env, got env: {:?}",
786            cmd.env
787        );
788    }
789
790    #[test]
791    fn test_build_command_has_parallel_flag() {
792        let installer = OnnxRuntimeInstaller::with_defaults();
793        let rocm_env = RocmEnv::from_known(Some(PathBuf::from("/opt/rocm")), "7.2.0".to_string());
794        let cmd = installer.build_build_command(&rocm_env);
795
796        assert!(cmd.args.contains(&"--parallel".to_string()));
797    }
798
799    #[test]
800    fn test_build_command_env_vars() {
801        let installer = OnnxRuntimeInstaller::with_defaults();
802        let rocm_env = RocmEnv::from_known(Some(PathBuf::from("/opt/rocm")), "7.2.0".to_string());
803        let cmd = installer.build_build_command(&rocm_env);
804
805        assert!(cmd.env.is_empty());
806    }
807
808    #[test]
809    fn test_build_command_with_eigen() {
810        let installer = OnnxRuntimeInstaller::new(OnnxRuntimeConfig {
811            use_preinstalled_eigen: true,
812            eigen_path: Some(PathBuf::from("/usr/include/eigen3")),
813            ..Default::default()
814        });
815        let rocm_env = RocmEnv::from_known(Some(PathBuf::from("/opt/rocm")), "7.2.0".to_string());
816        let cmd = installer.build_build_command(&rocm_env);
817
818        assert!(cmd.args.contains(&"--use_preinstalled_eigen".to_string()));
819        assert!(cmd.args.iter().any(|a| a.contains("/usr/include/eigen3")));
820        assert!(cmd
821            .args
822            .iter()
823            .any(|a| a.contains("FETCHCONTENT_SOURCE_DIR_EIGEN")));
824    }
825
826    // --- VAL-INSTALL-011: ONNX Runtime pip install from build output ---
827
828    #[test]
829    fn test_wheel_install_command() {
830        let installer = OnnxRuntimeInstaller::with_defaults();
831        let cmd = installer.build_wheel_install_command();
832
833        assert_eq!(cmd.program, "python3");
834        assert!(cmd.args.contains(&"-m".to_string()));
835        assert!(cmd.args.contains(&"pip".to_string()));
836        assert!(cmd.args.contains(&"install".to_string()));
837        assert!(cmd
838            .args
839            .iter()
840            .any(|a| a.contains("build/Linux/Release/dist/*.whl")));
841    }
842
843    #[test]
844    fn test_uninstall_command() {
845        let installer = OnnxRuntimeInstaller::with_defaults();
846        let cmd = installer.build_uninstall_command();
847
848        assert!(cmd.args.contains(&"uninstall".to_string()));
849        assert!(cmd.args.contains(&"-y".to_string()));
850        assert!(cmd.args.contains(&"onnxruntime".to_string()));
851        assert!(cmd.args.contains(&"onnxruntime-rocm".to_string()));
852        assert!(cmd.args.contains(&"onnxruntime-gpu".to_string()));
853        assert!(cmd.args.contains(&"onnxruntime-migraphx".to_string()));
854    }
855
856    // --- Prebuilt wheel install ---
857
858    #[test]
859    fn test_prebuilt_install_command() {
860        let installer = OnnxRuntimeInstaller::with_defaults();
861        let cmd = installer.build_prebuilt_install_command();
862
863        assert_eq!(cmd.program, "python3");
864        assert!(cmd.args.contains(&"-m".to_string()));
865        assert!(cmd.args.contains(&"pip".to_string()));
866        assert!(cmd.args.contains(&"--upgrade".to_string()));
867        assert!(cmd.args.contains(&"--force-reinstall".to_string()));
868        assert!(cmd.args.contains(&"--no-deps".to_string()));
869        assert!(cmd.args.contains(&"--prefer-binary".to_string()));
870        assert!(cmd.args.contains(&"onnxruntime-rocm".to_string()));
871        assert!(cmd.env.is_empty());
872    }
873
874    // --- Git commands ---
875
876    #[test]
877    fn test_git_clone_command() {
878        let installer = OnnxRuntimeInstaller::with_defaults();
879        let cmd = installer.build_git_clone_command();
880        assert_eq!(cmd.program, "git");
881        assert!(cmd.args.contains(&"clone".to_string()));
882        assert!(cmd.args.contains(&"--recursive".to_string()));
883        assert!(cmd.args.iter().any(|a| a.contains("microsoft/onnxruntime")));
884    }
885
886    #[test]
887    fn test_git_checkout_command() {
888        let installer = OnnxRuntimeInstaller::with_defaults();
889        let cmd = installer.build_git_checkout_command();
890        // Checks out the tag matching the configured runtime version (1.27.1),
891        // not a stale hardcoded tag.
892        assert!(cmd
893            .args
894            .contains(&format!("v{}", DEFAULT_ONNXRUNTIME_VERSION)));
895        assert!(cmd.args.contains(&"v1.27.1".to_string()));
896    }
897
898    #[test]
899    fn test_git_checkout_command_follows_runtime_version() {
900        let installer = OnnxRuntimeInstaller::new(OnnxRuntimeConfig {
901            runtime_version: Some("1.25.0".to_string()),
902            ..Default::default()
903        });
904        let cmd = installer.build_git_checkout_command();
905        assert!(cmd.args.contains(&"v1.25.0".to_string()));
906    }
907
908    #[test]
909    fn test_runtime_version_filters_empty_to_default() {
910        // An untrimmed/stale ctx.target_version of Some("") must NOT fall through
911        // as the tag — it would produce an invalid `git checkout v` (no suffix).
912        // Empty/whitespace values resolve to DEFAULT_ONNXRUNTIME_VERSION.
913        let empty = OnnxRuntimeInstaller::new(OnnxRuntimeConfig {
914            runtime_version: Some(String::new()),
915            ..Default::default()
916        });
917        assert_eq!(empty.config.runtime_version(), DEFAULT_ONNXRUNTIME_VERSION);
918        assert_ne!(empty.config.runtime_version(), "");
919        let ws = OnnxRuntimeInstaller::new(OnnxRuntimeConfig {
920            runtime_version: Some("   ".to_string()),
921            ..Default::default()
922        });
923        assert_eq!(ws.config.runtime_version(), DEFAULT_ONNXRUNTIME_VERSION);
924        // And a real value is preserved.
925        let real = OnnxRuntimeInstaller::new(OnnxRuntimeConfig {
926            runtime_version: Some("1.25.0".to_string()),
927            ..Default::default()
928        });
929        assert_eq!(real.config.runtime_version(), "1.25.0");
930    }
931
932    // --- HipArchs ---
933
934    #[test]
935    fn test_hip_archs_gfx1030() {
936        let archs = HipArchs::from_gpu_arch("gfx1030");
937        assert_eq!(archs.cmake_hip_architectures(), "gfx1030");
938    }
939
940    #[test]
941    fn test_hip_archs_gfx1100() {
942        let archs = HipArchs::from_gpu_arch("gfx1100");
943        assert_eq!(
944            archs.cmake_hip_architectures(),
945            "gfx1100;gfx1101;gfx1102;gfx1103"
946        );
947    }
948
949    #[test]
950    fn test_hip_archs_gfx1200() {
951        let archs = HipArchs::from_gpu_arch("gfx1200");
952        assert_eq!(archs.cmake_hip_architectures(), "gfx1200;gfx1201");
953    }
954
955    #[test]
956    fn test_hip_archs_default() {
957        let archs = HipArchs::from_gpu_arch("gfx900");
958        assert_eq!(archs.cmake_hip_architectures(), "gfx1100");
959    }
960
961    // --- VAL-INSTALL-045: ONNX Runtime declares dependency on ROCm ---
962
963    #[test]
964    fn test_dependencies() {
965        let installer = OnnxRuntimeInstaller::with_defaults();
966        let deps = installer.dependencies();
967        assert!(deps.contains(&"rocm"));
968    }
969
970    // --- VAL-INSTALL-014: Build error detection ---
971
972    #[test]
973    fn test_build_error_cmake() {
974        let installer = OnnxRuntimeInstaller::with_defaults();
975        let result = installer.check_build_output("", "CMake Error at cmake/CMakeLists.txt:42", 1);
976        assert!(result.is_err());
977        assert!(result.unwrap_err().contains("CMake Error"));
978    }
979
980    #[test]
981    fn test_build_error_ninja() {
982        let installer = OnnxRuntimeInstaller::with_defaults();
983        let result =
984            installer.check_build_output("ninja: build stopped: subcommand failed.", "", 1);
985        assert!(result.is_err());
986        assert!(result.unwrap_err().contains("ninja: build stopped"));
987    }
988
989    #[test]
990    fn test_build_error_success() {
991        let installer = OnnxRuntimeInstaller::with_defaults();
992        let result = installer.check_build_output("Build succeeded", "", 0);
993        assert!(result.is_ok());
994    }
995
996    // --- Config defaults ---
997
998    #[test]
999    fn test_config_defaults() {
1000        let config = OnnxRuntimeConfig::default();
1001        assert_eq!(config.python_bin, "python3");
1002        assert!(!config.dry_run);
1003        assert!(config.rocm_version.is_none());
1004        assert!(config.rocm_release.is_none());
1005        assert!(config.gpu_arch.is_none());
1006        assert!(config.workdir.is_none());
1007        assert!(!config.use_preinstalled_eigen);
1008        assert!(config.eigen_path.is_none());
1009        assert_eq!(config.install_method, OnnxInstallMethod::MigraphxWheel);
1010        assert_eq!(config.runtime_version(), DEFAULT_ONNXRUNTIME_VERSION);
1011        assert_eq!(config.prebuilt_version(), PREBUILT_MIGRAPHX_VERSION);
1012    }
1013
1014    // --- MIGraphX wheel install (PyPI prebuilt, default path) ---
1015
1016    #[test]
1017    fn test_prebuilt_version_defaults_to_pinned() {
1018        let installer = OnnxRuntimeInstaller::with_defaults();
1019        assert_eq!(installer.prebuilt_version(), PREBUILT_MIGRAPHX_VERSION);
1020        // Lock-step invariant: the default wheel version matches the manifest
1021        // target + DEFAULT_ONNXRUNTIME_VERSION so the honesty guard verifies the
1022        // version the default path actually installs.
1023        assert_eq!(installer.prebuilt_version(), DEFAULT_ONNXRUNTIME_VERSION);
1024        assert_eq!(installer.prebuilt_version(), "1.27.1");
1025    }
1026
1027    #[test]
1028    fn test_install_method_from_env() {
1029        let _env = crate::test_support::lock_env();
1030        std::env::remove_var(ONNX_INSTALL_METHOD_ENV);
1031        assert_eq!(install_method_from_env(), OnnxInstallMethod::MigraphxWheel);
1032        for (v, expected) in [
1033            ("source", OnnxInstallMethod::SourceBuild),
1034            ("SOURCE-BUILD", OnnxInstallMethod::SourceBuild),
1035            ("prebuilt", OnnxInstallMethod::PrebuiltWheel),
1036            ("rocm", OnnxInstallMethod::PrebuiltWheel),
1037            ("migraphx", OnnxInstallMethod::MigraphxWheel),
1038            ("", OnnxInstallMethod::MigraphxWheel),
1039        ] {
1040            if v.is_empty() {
1041                std::env::remove_var(ONNX_INSTALL_METHOD_ENV);
1042            } else {
1043                std::env::set_var(ONNX_INSTALL_METHOD_ENV, v);
1044            }
1045            assert_eq!(install_method_from_env(), expected, "for {v:?}");
1046        }
1047    }
1048
1049    #[test]
1050    fn test_prebuilt_version_from_env_override() {
1051        let _env = crate::test_support::lock_env();
1052        std::env::remove_var(ONNX_VERSION_ENV);
1053        assert_eq!(prebuilt_version_from_env(), None);
1054        std::env::set_var(ONNX_VERSION_ENV, "  1.30.0  ");
1055        assert_eq!(prebuilt_version_from_env(), Some("1.30.0".to_string()));
1056        std::env::set_var(ONNX_VERSION_ENV, "   ");
1057        assert_eq!(prebuilt_version_from_env(), None);
1058    }
1059
1060    #[test]
1061    fn test_prebuilt_version_overridable() {
1062        let installer = OnnxRuntimeInstaller::new(OnnxRuntimeConfig {
1063            prebuilt_version: Some("1.23.2".to_string()),
1064            ..Default::default()
1065        });
1066        assert_eq!(installer.prebuilt_version(), "1.23.2");
1067    }
1068
1069    #[test]
1070    fn test_migraphx_install_command_pypi() {
1071        let installer = OnnxRuntimeInstaller::with_defaults();
1072        let cmd = installer.build_migraphx_install_command();
1073        assert_eq!(cmd.program, "python3");
1074        assert!(cmd.args.contains(&"-m".to_string()));
1075        assert!(cmd.args.contains(&"pip".to_string()));
1076        assert!(cmd.args.contains(&"install".to_string()));
1077        assert!(cmd.args.contains(&"--upgrade".to_string()));
1078        assert!(cmd.args.contains(&"--force-reinstall".to_string()));
1079        // --no-deps honors the No-CUDA hard-prime (no nvidia/cuda transitive deps)
1080        assert!(cmd.args.contains(&"--no-deps".to_string()));
1081        assert!(cmd.args.contains(&"--no-cache-dir".to_string()));
1082        // Pins the prebuilt version from PyPI — NOT an AMD repo URL.
1083        assert!(cmd
1084            .args
1085            .iter()
1086            .any(|a| a == &format!("onnxruntime-migraphx=={}", PREBUILT_MIGRAPHX_VERSION)));
1087        assert!(
1088            !cmd.args.iter().any(|a| a.contains("repo.radeon.com")),
1089            "must not use the AMD manylinux repo (it lags and 404s): {:?}",
1090            cmd.args
1091        );
1092        assert!(cmd.env.is_empty());
1093    }
1094
1095    // --- Provider validation / model optimizer ---
1096
1097    #[test]
1098    fn test_provider_validation_command() {
1099        let installer = OnnxRuntimeInstaller::with_defaults();
1100        let cmd = installer.build_provider_validation_command();
1101        assert_eq!(cmd.program, "python3");
1102        assert!(cmd.args.contains(&"-c".to_string()));
1103        let script = &cmd.args[1];
1104        assert!(script.contains("MIGraphXExecutionProvider"));
1105        assert!(script.contains("ROCMExecutionProvider"));
1106        assert!(script.contains("ctypes.CDLL"));
1107        assert!(script.contains("loader_errors"));
1108        assert!(cmd.env.is_empty());
1109    }
1110
1111    #[test]
1112    fn test_model_optimizer_command() {
1113        let installer = OnnxRuntimeInstaller::with_defaults();
1114        let cmd = installer.build_model_optimizer_command("/path/to/model.onnx");
1115        assert_eq!(cmd.program, "python3");
1116        assert!(cmd.args.contains(&"-c".to_string()));
1117        let script = &cmd.args[1];
1118        assert!(script.contains("MIGraphXExecutionProvider"));
1119        assert!(script.contains("ROCMExecutionProvider"));
1120        assert!(!script.contains("CPUExecutionProvider"));
1121        assert!(script.contains("ORT_ENABLE_ALL"));
1122        assert!(script.contains("/path/to/model.onnx"));
1123        assert!(script.contains("/path/to/model.onnx.optimized"));
1124        assert!(cmd.env.is_empty());
1125    }
1126
1127    // --- OnnxInstallMethod ---
1128
1129    #[test]
1130    fn test_install_method_default() {
1131        assert_eq!(
1132            OnnxInstallMethod::default(),
1133            OnnxInstallMethod::MigraphxWheel
1134        );
1135    }
1136}
rusty_stack/installers/components/onnxruntime.rs

rusty_stack/installers/components/
onnxruntime.rs
