rlx_cli/
device.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3//
4// This program is free software: you can redistribute it and/or modify
5// it under the terms of the GNU General Public License as published by
6// the Free Software Foundation, version 3.
7//
8// This program is distributed in the hope that it will be useful,
9// but WITHOUT ANY WARRANTY; without even the implied warranty of
10// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
11// GNU General Public License for more details.
12//
13// You should have received a copy of the GNU General Public License
14// along with this program. If not, see <https://www.gnu.org/licenses/>.
15
16use anyhow::Result;
17use rlx_core::{validate_sam_device, validate_standard_device};
18use rlx_runtime::Device;
19use std::str::FromStr;
20
21/// Parse a device alias. Delegates to the upstream `FromStr for Device`
22/// impl in `rlx-driver`, which knows every alias (`mps`, `wgpu`, `hip`,
23/// …) and returns a uniform error string. Kept as a wrapper that
24/// surfaces `anyhow::Result` so existing callers don't change.
25pub fn parse_device(s: &str) -> Result<Device> {
26    Device::from_str(s).map_err(|e| anyhow::anyhow!("{e}"))
27}
28
29/// Parse a CLI device name and enforce the workspace-wide standard backend set.
30pub fn parse_standard_device(family: &str, s: &str) -> Result<Device> {
31    let d = parse_device(s)?;
32    validate_standard_device(family, d)?;
33    Ok(d)
34}
35
36/// Parse a causal-LM CLI device name (standard backends + CoreML / ANE + `auto`).
37pub fn parse_lm_device(family: &str, s: &str) -> Result<Device> {
38    rlx_core::resolve_lm_device_str(family, s)
39}
40
41pub fn parse_llama32_device(s: &str) -> Result<Device> {
42    parse_lm_device("llama32", s)
43}
44
45pub fn parse_gemma_device(s: &str) -> Result<Device> {
46    parse_lm_device("gemma", s)
47}
48
49pub fn parse_qwen35_device(s: &str) -> Result<Device> {
50    parse_standard_device("qwen35", s)
51}
52
53pub fn parse_llada2_device(s: &str) -> Result<Device> {
54    parse_standard_device("llada2", s)
55}
56
57/// Parse a CLI device name and enforce the SAM backend set (standard + `tpu`).
58pub fn parse_sam_device(family: &str, s: &str) -> Result<Device> {
59    let d = match s {
60        "tpu" => Device::Tpu,
61        other => parse_device(other)?,
62    };
63    validate_sam_device(family, d)?;
64    Ok(d)
65}
66
67#[cfg(test)]
68mod tests {
69    use super::*;
70
71    #[cfg(feature = "coreml")]
72    #[test]
73    fn parse_gemma_coreml_device() {
74        assert_eq!(parse_gemma_device("coreml").unwrap(), Device::Ane);
75        assert_eq!(parse_gemma_device("ane").unwrap(), Device::Ane);
76    }
77
78    #[test]
79    fn parse_gemma_auto_device() {
80        assert_eq!(
81            parse_gemma_device("auto").unwrap(),
82            rlx_core::pick_lm_device()
83        );
84    }
85}
rlx_cli/device.rs

rlx_cli/
device.rs
