Created
January 20, 2025 04:21
-
-
Save erolm-a/01586f294f3bae45020c85002652401f to your computer and use it in GitHub Desktop.
PoC of JAX, PyQt, numpy, tensorboard etc. in nix
This file contains 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 = "Dingo project"; | |
inputs = { | |
nixpkgs.url = "github:NixOS/nixpkgs/nixos-24.11"; | |
pyproject-nix = { | |
url = "github:pyproject-nix/pyproject.nix"; | |
inputs.nixpkgs.follows = "nixpkgs"; | |
}; | |
uv2nix = { | |
url = "github:pyproject-nix/uv2nix"; | |
inputs.pyproject-nix.follows = "pyproject-nix"; | |
inputs.nixpkgs.follows = "nixpkgs"; | |
}; | |
pyproject-build-systems = { | |
url = "github:pyproject-nix/build-system-pkgs"; | |
inputs.pyproject-nix.follows = "pyproject-nix"; | |
inputs.uv2nix.follows = "uv2nix"; | |
inputs.nixpkgs.follows = "nixpkgs"; | |
}; | |
}; | |
outputs = { nixpkgs, uv2nix, pyproject-nix, pyproject-build-systems, ... }: | |
let | |
inherit (nixpkgs) lib; | |
# Load a uv workspace from a workspace root. | |
# Uv2nix treats all uv projects as workspace projects. | |
workspace = uv2nix.lib.workspace.loadWorkspace { workspaceRoot = ./.; }; | |
# Create package overlay from workspace. | |
overlay = workspace.mkPyprojectOverlay { | |
# Prefer prebuilt binary wheels as a package source. | |
# Sdists are less likely to "just work" because of the metadata missing from uv.lock. | |
# Binary wheels are more likely to, but may still require overrides for library dependencies. | |
sourcePreference = "wheel"; # or sourcePreference = "sdist"; | |
# Optionally customise PEP 508 environment | |
# environ = { | |
# platform_release = "5.10.65"; | |
# }; | |
}; | |
pkgs = import nixpkgs { | |
system = "x86_64-linux"; | |
config = { | |
allowUnfree = true; | |
cudaSupport = true; | |
}; | |
}; | |
python = pkgs.python312; | |
ppacks = pkgs.python312Packages; | |
hacks = pkgs.callPackage pyproject-nix.build.hacks { }; | |
# Uv2nix can only work with what it has, and uv.lock is missing essential metadata to perform some builds. | |
# This is an additional overlay implementing build fixups. | |
# See: | |
# - https://pyproject-nix.github.io/uv2nix/FAQ.html | |
pyprojectOverrides = (_final: _prev: { | |
# pyqt6 ships its own qt library, which we definitely do not need and is also fairly heavy. | |
# It is also a major PITA to set up due to its many required libraries that require patchelf-ing. | |
pyqt6 = hacks.nixpkgsPrebuilt { | |
from = ppacks.pyqt6; | |
prev = _prev.pyqt6.overrideAttrs (old: { | |
passthru = old.passthru // { | |
dependencies = | |
lib.filterAttrs (name: _: !lib.hasSuffix "-qt6" name) | |
old.passthru.dependencies; | |
}; | |
}); | |
}; | |
numpy = hacks.nixpkgsPrebuilt { | |
from = ppacks.numpy; | |
prev = _prev.numpy; | |
}; | |
# https://github.com/nix-community/poetry2nix/blob/master/docs/edgecases.md#modulenotfounderror-no-module-named-packagename | |
neovim = _prev.neovim.overrideAttrs (attrs: ({ | |
buildInputs = attrs.buildInputs or [ ] ++ [ _final.setuptools ]; | |
})); | |
rockpool = _prev.rockpool.overrideAttrs (attrs: ({ | |
buildInputs = attrs.buildInputs or [ ] ++ [ _final.setuptools ]; | |
})); | |
}); | |
# Construct package set | |
pythonSet = | |
# Use base package set from pyproject.nix builders | |
(pkgs.callPackage pyproject-nix.build.packages { | |
inherit python; | |
}).overrideScope (lib.composeManyExtensions [ | |
pyproject-build-systems.overlays.default | |
overlay | |
pyprojectOverrides | |
]); | |
in { | |
# Package a virtual environment as our main application. | |
# | |
# Enable no optional dependencies for production build. | |
packages.x86_64-linux.default = | |
pythonSet.mkVirtualEnv "dingo-env" workspace.deps.default; | |
# This example provides two different modes of development: | |
# - Impurely using uv to manage virtual environments | |
# - Pure development using uv2nix to manage virtual environments | |
devShells.x86_64-linux = { | |
# It is of course perfectly OK to keep using an impure virtualenv workflow and only use uv2nix to build packages. | |
# This devShell simply adds Python and undoes the dependency leakage done by Nixpkgs Python infrastructure. | |
impure = pkgs.mkShell { | |
packages = [ python pkgs.uv ]; | |
env = { | |
# Prevent uv from managing Python downloads | |
UV_PYTHON_DOWNLOADS = "never"; | |
# Force uv to use nixpkgs Python interpreter | |
UV_PYTHON = python.interpreter; | |
} // lib.optionalAttrs pkgs.stdenv.isLinux { | |
# Python libraries often load native shared objects using dlopen(3). | |
# Setting LD_LIBRARY_PATH makes the dynamic library loader aware of libraries without using RPATH for lookup. | |
LD_LIBRARY_PATH = | |
lib.makeLibraryPath pkgs.pythonManylinuxPackages.manylinux1; | |
}; | |
shellHook = '' | |
unset PYTHONPATH | |
''; | |
}; | |
# This devShell uses uv2nix to construct a virtual environment purely from Nix, using the same dependency specification as the application. | |
# The notable difference is that we also apply another overlay here enabling editable mode ( https://setuptools.pypa.io/en/latest/userguide/development_mode.html ). | |
# | |
# This means that any changes done to your local files do not require a rebuild. | |
uv2nix = let | |
# Create an overlay enabling editable mode for all local dependencies. | |
editableOverlay = workspace.mkEditablePyprojectOverlay { | |
# Use environment variable | |
root = "$REPO_ROOT"; | |
# Optional: Only enable editable for these packages | |
# members = [ "hello-world" ]; | |
}; | |
# Override previous set with our overrideable overlay. | |
editablePythonSet = pythonSet.overrideScope editableOverlay; | |
# Build virtual environment, with local packages being editable. | |
# | |
# Enable all optional dependencies for development. | |
virtualenv = | |
editablePythonSet.mkVirtualEnv "dingo-dev-env" workspace.deps.all; | |
in pkgs.mkShell { | |
packages = with pkgs.cudaPackages; [ | |
pkgs.linuxPackages.nvidia_x11 | |
virtualenv | |
pkgs.uv | |
# https://jax.readthedocs.io/en/latest/installation.html#pip-installation-nvidia-gpu-cuda-installed-locally-harder | |
cudatoolkit | |
cudnn | |
nccl | |
# TODO profiling (e.g. via perfetto or tensorboard) require some extra work | |
pkgs.pprof | |
pkgs.graphviz | |
]; | |
env = { | |
# Don't create venv using uv | |
UV_NO_SYNC = "1"; | |
# Force uv to use Python interpreter from venv | |
UV_PYTHON = "${virtualenv}/bin/python"; | |
# Prevent uv from downloading managed Python's | |
UV_PYTHON_DOWNLOADS = "never"; | |
LD_LIBRARY_PATH = with pkgs.cudaPackages; | |
lib.makeLibraryPath ([ | |
virtualenv | |
cudatoolkit | |
cudnn | |
nccl | |
pkgs.linuxPackages.nvidia_x11 | |
] | |
++ pkgs.pythonManylinuxPackages.manylinux1); # the latter one is hardly useful unless we use uv directly and patchelf has not run yet | |
XLA_FLAGS = | |
"--xla_gpu_cuda_data_dir=${pkgs.cudaPackages.cudatoolkit}"; | |
IPYTHONDIR = ".ipython"; | |
SHELL = "${pkgs.zsh}/bin/zsh"; | |
}; | |
shellHook = '' | |
# Undo dependency propagation by nixpkgs. | |
unset PYTHONPATH | |
# Get repository root using git. This is expanded at runtime by the editable `.pth` machinery. | |
export REPO_ROOT=$(git rev-parse --show-toplevel) | |
# Make all exported dev bins accessible | |
source ${virtualenv}/bin/activate | |
#alias ipy=ipython --profile=dingo | |
#exec zsh | |
''; | |
}; | |
}; | |
}; | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment