127 lines
3.8 KiB
Nix
127 lines
3.8 KiB
Nix
{
|
|
# Build Pyo3 package
|
|
inputs = {
|
|
nixpkgs.url = "github:nixos/nixpkgs/nixpkgs-unstable";
|
|
flake-utils.url = "github:numtide/flake-utils";
|
|
rust-overlay = {
|
|
url = "github:oxalica/rust-overlay";
|
|
inputs.nixpkgs.follows = "nixpkgs";
|
|
};
|
|
crane.url = "github:ipetkov/crane";
|
|
fenix = {
|
|
url = "github:nix-community/fenix";
|
|
inputs.nixpkgs.follows = "nixpkgs";
|
|
};
|
|
};
|
|
|
|
outputs =
|
|
inputs:
|
|
inputs.flake-utils.lib.eachDefaultSystem (
|
|
system:
|
|
let
|
|
pkgs = import inputs.nixpkgs {
|
|
inherit system;
|
|
overlays = [ inputs.rust-overlay.overlays.default ];
|
|
config = {
|
|
allowUnfree = true;
|
|
cudaSupport = true;
|
|
};
|
|
};
|
|
lib = pkgs.lib;
|
|
|
|
python_version = pkgs.python313;
|
|
wheel_tail = "cp313-cp313-linux_x86_64"; # Change if python_version changes
|
|
|
|
# Get a custom rust toolchain
|
|
craneLib = (inputs.crane.mkLib pkgs).overrideToolchain (
|
|
p: inputs.fenix.packages.${system}.complete.toolchain
|
|
);
|
|
|
|
project_name = (craneLib.crateNameFromCargoToml { cargoToml = ./game/Cargo.toml; }).pname;
|
|
project_version = (craneLib.crateNameFromCargoToml { cargoToml = ./game/Cargo.toml; }).version;
|
|
|
|
crate_cfg = {
|
|
src =
|
|
let
|
|
fs = lib.fileset;
|
|
in
|
|
fs.toSource {
|
|
root = ./game;
|
|
fileset = fs.unions [
|
|
./game/src/lib.rs
|
|
./game/Cargo.lock
|
|
./game/Cargo.toml
|
|
./game/pyproject.toml
|
|
];
|
|
};
|
|
nativeBuildInputs = [ python_version ];
|
|
# doCheck = true;
|
|
# buildInputs = [];
|
|
};
|
|
|
|
crate_artifacts = craneLib.buildDepsOnly (
|
|
crate_cfg
|
|
// {
|
|
pname = "${project_name}-artifacts";
|
|
version = project_version;
|
|
}
|
|
);
|
|
|
|
# Build the library, then re-use the target dir to generate the wheel file with maturin
|
|
crate_wheel =
|
|
(craneLib.buildPackage (
|
|
crate_cfg
|
|
// {
|
|
pname = project_name;
|
|
version = project_version;
|
|
cargoArtifacts = crate_artifacts;
|
|
}
|
|
)).overrideAttrs
|
|
(old: {
|
|
nativeBuildInputs = old.nativeBuildInputs ++ [ pkgs.maturin ];
|
|
buildPhase = old.buildPhase + ''
|
|
maturin build --offline --target-dir ./target
|
|
'';
|
|
installPhase = old.installPhase + ''
|
|
cp target/wheels/${project_name}-${project_version}-${wheel_tail}.whl $out/
|
|
'';
|
|
});
|
|
in
|
|
rec {
|
|
packages = {
|
|
default = crate_wheel; # The wheel itself
|
|
|
|
# A python version with the library installed
|
|
pythonpkg = python_version.withPackages (ps: [
|
|
(lib.python_package ps)
|
|
]);
|
|
};
|
|
devShells.default = craneLib.devShell {
|
|
packages = [
|
|
(pkgs.python3.withPackages (ppkgs: [
|
|
ppkgs.torch
|
|
ppkgs.tqdm
|
|
ppkgs.matplotlib
|
|
(lib.python_package ppkgs)
|
|
]))
|
|
];
|
|
shellHook = ''
|
|
export CUDA_PATH=${pkgs.cudatoolkit}
|
|
'';
|
|
};
|
|
lib = {
|
|
# To use in other builds with the "withPackages" call
|
|
python_package =
|
|
ps:
|
|
ps.buildPythonPackage {
|
|
pname = project_name;
|
|
format = "wheel";
|
|
version = project_version;
|
|
src = "${crate_wheel}/${project_name}-${project_version}-${wheel_tail}.whl";
|
|
doCheck = false;
|
|
pythonImportsCheck = [ project_name ];
|
|
};
|
|
};
|
|
}
|
|
);
|
|
}
|