Learning Rust -- Correcting PyTorch dependency hell in Poetry with a CLI

[green sail]

Hey all! I am brand new to Rust, and wanted to build a project to solve (poorly) a pain point for managing PyTorch source repositories as an excuse to learn the language.

Any advice to do things in a more Rust-ic manner (Like Pythonic?) would be greatly appreciated. If it sucks, I have to know lol

https://crates.io/crates/torch_poetry_bootstrap

https://github.com/drossotto/torch_poetry_bootstrap

[opal rune]

Good job with you first crate! 🎉
Here are some things you might consider.

You can remove serde and serde_json crates, if you use values directly in the code, like const TORCH_SOURCES: &[TorchSource] = ..., but you cannot make Strings in the const context so you will ether need to change the TorchSource to use &'static str or use the LazyCell<Vec<TourchSource>> (i think that would be easier)

// the fact that rust-fmt will splay them like that is lame.
// `String` implements `From<&str>` so "thatever".into() turns to Strings
static TORCH_SOURCES: LazyCell<Vec<TorchSource>> = LazyCell::new(|| {
    vec![
        TorchSource {
            cuda: "12.8".into(),
            source: "cu128".into(),
            url: "https://download.pytorch.org/whl/cu128".into(),
        },
        TorchSource {
            cuda: "12.1".into(),
            source: "cu121".into(),
            url: "https://download.pytorch.org/whl/cu121".into(),
        },
        TorchSource {
            cuda: "11.8".into(),
            source: "cu118".into(),
            url: "https://download.pytorch.org/whl/cu118".into(),
        },
        TorchSource {
            cuda: "11.7".into(),
            source: "cu117".into(),
            url: "https://download.pytorch.org/whl/cu117".into(),
        },
        TorchSource {
            cuda: "11.6".into(),
            source: "cu116".into(),
            url: "https://download.pytorch.org/whl/cu116".into(),
        },
        TorchSource {
            cuda: "11.3".into(),
            source: "cu113".into(),
            url: "https://download.pytorch.org/whl/cu113".into(),
        },
        TorchSource {
            cuda: "10.2".into(),
            source: "cu102".into(),
            url: "https://download.pytorch.org/whl/cu102".into(),
        },
        TorchSource {
            cuda: "cpu".into(),
            source: "cpu".into(),
            url: "https://download.pytorch.org/whl/cpu".into(),
        },
    ]
});

Instead of shelling-out to run nvidia-smi you can use nvml-wraper (https://crates.io/crates/nvml-wrapper) crate to get the cuda version directly from nvml library. This is what nvidia-smi uses internally. This way you won't need to parse strings, so regex can go away too.

The crate includes the library dynamically and seems to work for both Linux and Windows (it should be available as a part of driver)

You seem to have included once_cell library, but do not use it. In rust 1.70+ we have various cells in std::cell module (OnceCell/Lock, LazyCell/Lock), so check them out.

panic = "abort" is fine but unless you really need it, you should use unwind instead (default).

make_loger function should be inside the logs module. You will probably appreciate using log crate. It has nice syntax, but that's another crate to add.

You should not use String as Error, use the BootstrapError::Other, or make additional enum variants if you want to track the errors.

You can annotate your error as #[non_exhaustive] to indicate that you intend to add other errors in the future but it's ok if you don't.

In rust you can shadow the vars even in if-let expressions for example:
if let Some(target) = target { ... }.

let-else can be helpful to get values or return early like:
let Some(target) = parse_version(detected_version) else { return ...; }.

Iterators can be easier to follow than loops, you should probably see the docs for the Iterator trait (https://doc.rust-lang.org/std/iter/trait.Iterator.html). Also the Default trait is nice to use.

use std::default::Default;

impl Default for TorchSource {
    fn default() -> Self {
        TorchSource {
            cuda: "cpu".to_string(),
            source: "pypi".to_string(),
            url: "https://pypi.org/simple".to_string(),
        }
    }
}

pub fn resolve_best_source(detected_version: &str) -> TorchSource {
    let Some(target) = parse_version(detected_version) else {
        return TorchSource::default();
    };

    // this probably has the same behaviour as the for loop but is more consise 
    let best = TORCH_SOURCES
        .iter()
        .filter_map(|src| parse_version(&src.cuda).map(|v| (src, v))) // yes, you can map the `Option`, also notice no src.clone() since iterators use refs 
        .filter(|&(_src, v)| v <= target)
        .max_by_key(|&(_src, v)| v)
        .map(|(src, _v)| src.clone());

    best.unwrap_or_default()
}

[green sail]

Man, this is exactly the type of criticism I was looking for, thank you! idk Rust just seems to have amazing syntax for docgen and DSL stuff. Hopefully I start writing some idiomatic Rust eventually