furiosa_device/config/
mod.rs

1mod builder;
2mod env;
3pub(crate) mod find;
4mod inner;
5
6use std::fmt::Display;
7use std::str::FromStr;
8
9pub use builder::DeviceConfigBuilder;
10pub(crate) use find::{expand_status, find_device_files_in};
11use serde::{Deserialize, Serialize};
12
13pub use self::builder::NotDetermined;
14pub use self::env::EnvBuilder;
15use self::inner::DeviceConfigInner;
16use crate::{Arch, DeviceError};
17
18/// Describes a required set of devices for [`find_device_files`][crate::find_device_files].
19///
20/// # Examples
21/// ```rust
22/// use furiosa_device::DeviceConfig;
23///
24/// // 1 core
25/// DeviceConfig::warboy().build();
26///
27/// // 1 core x 2
28/// DeviceConfig::warboy().count(2);
29///
30/// // Fused 2 cores x 2
31/// DeviceConfig::warboy().fused().count(2);
32/// ```
33///
34/// # Textual Representation
35///
36/// DeviceConfig supports textual representation, which is its equivalent string representation.
37/// One can obtain the corresponding DeviceConfig from the textual representation
38/// by using the FromStr trait, or by calling [`from_env`][`DeviceConfig::from_env`]
39/// after setting an environment variable.
40///
41/// ```rust
42/// use std::str::FromStr;
43///
44/// use furiosa_device::DeviceConfig;
45///
46/// let config = DeviceConfig::from_env("SOME_OTHER_ENV_KEY").build();
47/// let config = DeviceConfig::from_str("npu:0:0,npu:0:1").unwrap(); // get config directly from a string literal
48/// ```
49///
50/// The rules for textual representation are as follows:
51///
52/// ```rust
53/// use std::str::FromStr;
54///
55/// use furiosa_device::DeviceConfig;
56///
57/// // Using specific device names
58/// DeviceConfig::from_str("npu:0:0").unwrap(); // npu0pe0
59/// DeviceConfig::from_str("npu:0:0-1").unwrap(); // npu0pe0-1
60///
61/// // Using device configs
62/// DeviceConfig::from_str("warboy*2").unwrap(); // single pe x 2 (equivalent to "warboy(1)*2")
63/// DeviceConfig::from_str("warboy(1)*2").unwrap(); // single pe x 2
64/// DeviceConfig::from_str("warboy(2)*2").unwrap(); // 2-pe fusioned x 2
65///
66/// // Combine multiple representations separated by commas
67/// DeviceConfig::from_str("npu:0:0-1,npu:1:0-1").unwrap(); // npu0pe0-1, npu1pe0-1
68/// ```
69#[derive(Clone, Debug, Serialize, Deserialize)]
70#[serde(into = "String", try_from = "&str")]
71pub struct DeviceConfig {
72    pub(crate) inner: DeviceConfigInner,
73}
74
75impl DeviceConfig {
76    /// Returns a builder associated with Warboy NPUs.
77    pub fn warboy() -> DeviceConfigBuilder<Arch, NotDetermined, NotDetermined> {
78        DeviceConfigBuilder {
79            arch: Arch::WarboyB0,
80            mode: NotDetermined { _priv: () },
81            count: NotDetermined { _priv: () },
82        }
83    }
84
85    pub fn warboy_a0() -> DeviceConfigBuilder<Arch, NotDetermined, NotDetermined> {
86        DeviceConfigBuilder {
87            arch: Arch::WarboyA0,
88            mode: NotDetermined { _priv: () },
89            count: NotDetermined { _priv: () },
90        }
91    }
92
93    /// Returns a builder struct to read config saved in an environment variable.
94    /// You can provide fallback options to the builder in case the envrionment variable is empty.
95    pub fn from_env<K: ToString>(key: K) -> EnvBuilder<NotDetermined> {
96        EnvBuilder::<NotDetermined>::from_env(key)
97    }
98}
99
100impl Default for DeviceConfig {
101    fn default() -> Self {
102        DeviceConfig::warboy().fused().count(1)
103    }
104}
105
106impl FromStr for DeviceConfig {
107    type Err = DeviceError;
108
109    fn from_str(s: &str) -> Result<Self, Self::Err> {
110        Ok(Self {
111            inner: DeviceConfigInner::from_str(s)?,
112        })
113    }
114}
115
116impl<'a> TryFrom<&'a str> for DeviceConfig {
117    type Error = DeviceError;
118
119    fn try_from(value: &'a str) -> Result<Self, Self::Error> {
120        DeviceConfig::from_str(value)
121    }
122}
123
124impl Display for DeviceConfig {
125    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
126        self.inner.fmt(f)
127    }
128}
129
130impl From<DeviceConfig> for String {
131    fn from(config: DeviceConfig) -> Self {
132        config.to_string()
133    }
134}
135
136#[cfg(test)]
137mod tests {
138    use super::*;
139    use crate::list::list_devices_with;
140
141    #[tokio::test]
142    async fn test_find_device_files() -> eyre::Result<()> {
143        // test directory contains 2 warboy NPUs
144        let devices =
145            list_devices_with("../test_data/test-0/dev", "../test_data/test-0/sys").await?;
146        let devices_with_statuses = expand_status(devices).await?;
147
148        // try lookup 4 different single cores
149        let config = DeviceConfig::warboy().single().count(4);
150        let found_device_files = find_device_files_in(&config, &devices_with_statuses)?;
151        let mut found_device_file_names: Vec<&str> =
152            found_device_files.iter().map(|f| f.filename()).collect();
153        found_device_file_names.sort();
154        assert_eq!(
155            found_device_file_names,
156            &["npu0pe0", "npu0pe1", "npu1pe0", "npu1pe1"],
157        );
158
159        // try lookup all single cores
160        let config = DeviceConfig::warboy().single().all();
161        let found_device_files = find_device_files_in(&config, &devices_with_statuses)?;
162        let mut found_device_file_names: Vec<&str> =
163            found_device_files.iter().map(|f| f.filename()).collect();
164        found_device_file_names.sort();
165        assert_eq!(
166            found_device_file_names,
167            &["npu0pe0", "npu0pe1", "npu1pe0", "npu1pe1"],
168        );
169
170        // // looking for 5 different cores should fail
171        let config = DeviceConfig::warboy().single().count(5);
172        let found = find_device_files_in(&config, &devices_with_statuses);
173        match found {
174            Ok(_) => panic!("looking for 5 different cores should fail"),
175            Err(e) => assert!(matches!(e, DeviceError::DeviceNotFound { .. })),
176        }
177
178        // // try lookup 2 different fused cores
179        let config = DeviceConfig::warboy().fused().count(2);
180        let found_device_files = find_device_files_in(&config, &devices_with_statuses)?;
181        let mut found_device_file_names: Vec<&str> =
182            found_device_files.iter().map(|f| f.filename()).collect();
183        found_device_file_names.sort();
184        assert_eq!(found_device_file_names, &["npu0pe0-1", "npu1pe0-1"],);
185
186        // // try lookup all fused cores
187        let config = DeviceConfig::warboy().fused().all();
188        let found_device_files = find_device_files_in(&config, &devices_with_statuses)?;
189        let mut found_device_file_names: Vec<&str> =
190            found_device_files.iter().map(|f| f.filename()).collect();
191        found_device_file_names.sort();
192        assert_eq!(found_device_file_names, &["npu0pe0-1", "npu1pe0-1"],);
193
194        // // looking for 3 different fused cores should fail
195        let config = DeviceConfig::warboy().fused().count(3);
196        let found = find_device_files_in(&config, &devices_with_statuses);
197        match found {
198            Ok(_) => panic!("looking for 3 different fused cores should fail"),
199            Err(e) => assert!(matches!(e, DeviceError::DeviceNotFound { .. })),
200        }
201
202        Ok(())
203    }
204
205    #[test]
206    fn test_config_symmetric_display() -> eyre::Result<()> {
207        assert_eq!("npu:0".parse::<DeviceConfig>()?.to_string(), "npu:0");
208        assert_eq!("npu:1".parse::<DeviceConfig>()?.to_string(), "npu:1");
209        assert_eq!("npu:0:0".parse::<DeviceConfig>()?.to_string(), "npu:0:0");
210        assert_eq!("npu:0:1".parse::<DeviceConfig>()?.to_string(), "npu:0:1");
211        assert_eq!("npu:1:0".parse::<DeviceConfig>()?.to_string(), "npu:1:0");
212        assert_eq!(
213            "npu:0:0-1".parse::<DeviceConfig>()?.to_string(),
214            "npu:0:0-1"
215        );
216
217        assert_eq!(
218            "warboy(1)*2".parse::<DeviceConfig>()?.to_string(),
219            "warboy(1)*2"
220        );
221        assert_eq!(
222            "warboy(2)*4".parse::<DeviceConfig>()?.to_string(),
223            "warboy(2)*4"
224        );
225
226        Ok(())
227    }
228
229    #[test]
230    fn test_config_comma_separated() -> eyre::Result<()> {
231        let config =
232            "npu:0:0,npu:0:1,npu:0:0-1,warboy(1)*1,warboy(2)*2,npu0pe0".parse::<DeviceConfig>()?;
233
234        assert_eq!(
235            config.inner.cfgs,
236            vec![
237                "npu:0:0".parse::<crate::config::inner::Config>()?,
238                "npu:0:1".parse::<crate::config::inner::Config>()?,
239                "npu:0:0-1".parse::<crate::config::inner::Config>()?,
240                "warboy(1)*1".parse::<crate::config::inner::Config>()?,
241                "warboy(2)*2".parse::<crate::config::inner::Config>()?,
242                "npu0pe0".parse::<crate::config::inner::Config>()?,
243            ]
244        );
245        Ok(())
246    }
247
248    #[test]
249    fn test_config_from_env() -> eyre::Result<()> {
250        let key = "ENV_KEY";
251        std::env::set_var(
252            key,
253            "npu:0:0,npu:0:1,npu:0:0-1,warboy(1)*1,warboy(2)*2,npu0pe0",
254        );
255        let config = DeviceConfig::from_env(key).build()?;
256
257        assert_eq!(
258            config.inner.cfgs,
259            vec![
260                "npu:0:0".parse::<crate::config::inner::Config>()?,
261                "npu:0:1".parse::<crate::config::inner::Config>()?,
262                "npu:0:0-1".parse::<crate::config::inner::Config>()?,
263                "warboy(1)*1".parse::<crate::config::inner::Config>()?,
264                "warboy(2)*2".parse::<crate::config::inner::Config>()?,
265                "npu0pe0".parse::<crate::config::inner::Config>()?,
266            ]
267        );
268        Ok(())
269    }
270
271    #[tokio::test]
272    async fn test_find_device_files_with_comma_separated() -> eyre::Result<()> {
273        // test directory contains 2 warboy NPUs
274        let devices =
275            list_devices_with("../test_data/test-0/dev", "../test_data/test-0/sys").await?;
276        let devices_with_statuses = expand_status(devices).await?;
277
278        // try lookup with various valid configs
279        let config = "npu:0:0,npu:0:1,npu:1:0,npu:1:1".parse::<DeviceConfig>()?;
280        let found_device_files = find_device_files_in(&config, &devices_with_statuses)?;
281        let mut found_device_file_names: Vec<&str> =
282            found_device_files.iter().map(|f| f.filename()).collect();
283        found_device_file_names.sort();
284        assert_eq!(
285            found_device_file_names,
286            &["npu0pe0", "npu0pe1", "npu1pe0", "npu1pe1"],
287        );
288
289        let config = "npu:0:0,npu0pe1,npu:1:0,npu1pe1".parse::<DeviceConfig>()?;
290        let found_device_files = find_device_files_in(&config, &devices_with_statuses)?;
291        let mut found_device_file_names: Vec<&str> =
292            found_device_files.iter().map(|f| f.filename()).collect();
293        found_device_file_names.sort();
294        assert_eq!(
295            found_device_file_names,
296            &["npu0pe0", "npu0pe1", "npu1pe0", "npu1pe1"],
297        );
298
299        let config = "warboy(1)*1,warboy(1)*1,warboy(1)*1,warboy(1)*1".parse::<DeviceConfig>()?;
300        let found_device_files = find_device_files_in(&config, &devices_with_statuses)?;
301        let mut found_device_file_names: Vec<&str> =
302            found_device_files.iter().map(|f| f.filename()).collect();
303        found_device_file_names.sort();
304        assert_eq!(
305            found_device_file_names,
306            &["npu0pe0", "npu0pe1", "npu1pe0", "npu1pe1"],
307        );
308
309        let config = "npu:0:0,npu:0:1,warboy(1)*2".parse::<DeviceConfig>()?;
310        let found_device_files = find_device_files_in(&config, &devices_with_statuses)?;
311        let mut found_device_file_names: Vec<&str> =
312            found_device_files.iter().map(|f| f.filename()).collect();
313        found_device_file_names.sort();
314        assert_eq!(
315            found_device_file_names,
316            &["npu0pe0", "npu0pe1", "npu1pe0", "npu1pe1"],
317        );
318
319        Ok(())
320    }
321
322    #[tokio::test]
323    async fn test_find_device_files_with_failing_cases() -> eyre::Result<()> {
324        // test directory contains 2 warboy NPUs
325        let devices =
326            list_devices_with("../test_data/test-0/dev", "../test_data/test-0/sys").await?;
327        let devices_with_statuses = expand_status(devices).await?;
328
329        // test duplicate configs
330        let config = "npu:0:0,npu:0:0".parse::<DeviceConfig>()?;
331        let found = find_device_files_in(&config, &devices_with_statuses);
332        match found {
333            Ok(_) => panic!("looking for duplicate devices should fail"),
334            Err(e) => assert!(matches!(e, DeviceError::DeviceNotFound { .. })),
335        }
336
337        let config = "npu:0:0-1,npu0pe0-1".parse::<DeviceConfig>()?;
338        let found = find_device_files_in(&config, &devices_with_statuses);
339        match found {
340            Ok(_) => panic!("looking for duplicate devices should fail"),
341            Err(e) => assert!(matches!(e, DeviceError::DeviceNotFound { .. })),
342        }
343
344        let config = "npu:2:0".parse::<DeviceConfig>()?;
345        let found = find_device_files_in(&config, &devices_with_statuses);
346        match found {
347            Ok(_) => panic!("looking for not exist device should fail"),
348            Err(e) => assert!(matches!(e, DeviceError::DeviceNotFound { .. })),
349        }
350
351        Ok(())
352    }
353}
furiosa_device/config/mod.rs

furiosa_device/config/
mod.rs
