Skip to content

Instantly share code, notes, and snippets.

@erolm-a
Created January 20, 2025 04:21
Show Gist options
  • Save erolm-a/01586f294f3bae45020c85002652401f to your computer and use it in GitHub Desktop.
Save erolm-a/01586f294f3bae45020c85002652401f to your computer and use it in GitHub Desktop.
PoC of JAX, PyQt, numpy, tensorboard etc. in nix
{
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