Last active
February 24, 2026 19:38
-
-
Save lunalunaa/c5b9b7f56f3799ca2e7144ef8198f5a6 to your computer and use it in GitHub Desktop.
Burn setup on nixos
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| { | |
| description = "Rust machine learning environment using Burn"; | |
| inputs = { | |
| nixpkgs.url = "github:NixOS/nixpkgs/nixos-unstable"; | |
| fenix = { | |
| url = "github:nix-community/fenix"; | |
| inputs.nixpkgs.follows = "nixpkgs"; | |
| }; | |
| crane.url = "github:ipetkov/crane"; | |
| flake-parts.url = "github:hercules-ci/flake-parts"; | |
| }; | |
| outputs = | |
| inputs@{ flake-parts, ... }: | |
| flake-parts.lib.mkFlake { inherit inputs; } { | |
| systems = [ | |
| "x86_64-linux" | |
| "aarch64-linux" | |
| "x86_64-darwin" | |
| "aarch64-darwin" | |
| ]; | |
| perSystem = | |
| { | |
| system, | |
| pkgs, | |
| inputs', | |
| ... | |
| }: | |
| let | |
| cudaPkgs = import inputs.nixpkgs { | |
| inherit system; | |
| config = { | |
| allowUnfree = true; | |
| cudaSupport = true; | |
| }; | |
| }; | |
| rustToolchain = inputs'.fenix.packages.stable.withComponents [ | |
| "cargo" | |
| "clippy" | |
| "rust-src" | |
| "rustc" | |
| "rustfmt" | |
| "rust-analyzer" | |
| ]; | |
| craneLib = (inputs.crane.mkLib pkgs).overrideToolchain rustToolchain; | |
| nativeBuildInputs = [ | |
| pkgs.pkg-config | |
| rustToolchain | |
| ]; | |
| buildInputs = [ | |
| pkgs.openssl | |
| # --- WGPU / Vulkan backend (GPU) --- | |
| # pkgs.vulkan-loader | |
| # pkgs.vulkan-validation-layers | |
| # pkgs.libGL | |
| # pkgs.wayland | |
| # pkgs.libxkbcommon | |
| # pkgs.xorg.libX11 | |
| # pkgs.xorg.libXcursor | |
| # pkgs.xorg.libXrandr | |
| # pkgs.xorg.libXi | |
| # --- CUDA backend --- | |
| cudaPkgs.cudaPackages.cudatoolkit | |
| cudaPkgs.cudaPackages.cudnn | |
| cudaPkgs.cudaPackages.cuda_nvcc | |
| ]; | |
| # Build-essential env vars — merged into crane derivations. | |
| # Only include vars that affect compilation; omit dev-convenience | |
| # vars (RUST_BACKTRACE, RUST_LOG) to avoid polluting the drv hash. | |
| buildEnv = { | |
| PKG_CONFIG_PATH = "${pkgs.openssl.dev}/lib/pkgconfig"; | |
| CUDA_PATH = "${cudaPkgs.cudaPackages.cudatoolkit}"; | |
| CUDA_ROOT = "${cudaPkgs.cudaPackages.cudatoolkit}"; | |
| LD_LIBRARY_PATH = | |
| pkgs.lib.makeLibraryPath [ | |
| cudaPkgs.cudaPackages.cudatoolkit | |
| cudaPkgs.cudaPackages.cudnn | |
| ] | |
| + ":/run/opengl-driver/lib"; | |
| # VK_ICD_FILENAMES = "${pkgs.mesa}/share/vulkan/icd.d/radeon_icd.x86_64.json"; | |
| }; | |
| # Dev-convenience env vars — devShell only, not baked into drv hash. | |
| shellEnv = buildEnv // { | |
| RUST_BACKTRACE = "1"; | |
| RUST_LOG = "info"; | |
| }; | |
| commonArgs = { | |
| src = craneLib.cleanCargoSource ./.; | |
| strictDeps = true; | |
| inherit nativeBuildInputs buildInputs; | |
| } | |
| // buildEnv; | |
| cargoArtifacts = craneLib.buildDepsOnly commonArgs; | |
| defaultPackage = craneLib.buildPackage ( | |
| commonArgs | |
| // { | |
| inherit cargoArtifacts; | |
| } | |
| ); | |
| in | |
| { | |
| packages.default = defaultPackage; | |
| checks = { | |
| build = defaultPackage; | |
| clippy = craneLib.cargoClippy ( | |
| commonArgs | |
| // { | |
| inherit cargoArtifacts; | |
| cargoClippyExtraArgs = "-- --deny warnings"; | |
| } | |
| ); | |
| test = craneLib.cargoTest ( | |
| commonArgs | |
| // { | |
| inherit cargoArtifacts; | |
| } | |
| ); | |
| }; | |
| devShells.default = pkgs.mkShell ( | |
| { | |
| inputsFrom = [ defaultPackage ]; | |
| packages = [ | |
| pkgs.cargo-watch | |
| pkgs.cargo-expand | |
| pkgs.cargo-flamegraph | |
| pkgs.hyperfine | |
| ]; | |
| shellHook = '' | |
| echo "" | |
| echo " 🔥 Burn ML dev environment (CUDA)" | |
| echo " Rust $(rustc --version)" | |
| echo " CUDA $(nvcc --version 2>/dev/null | grep release | awk '{print $6}' || echo 'not found')" | |
| echo "" | |
| echo " Commands:" | |
| echo " cargo run — run the project" | |
| echo " cargo test — run tests" | |
| echo " cargo watch -x run — auto-rebuild on changes" | |
| echo "" | |
| ''; | |
| } | |
| // shellEnv | |
| ); | |
| }; | |
| }; | |
| } |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment