← back to Exo
nix: add Python packaging with uv2nix
e4c6a7dbb4abc3bec8a832f60c3a7c6238cb9af3 · 2026-01-15 18:37:02 +0000 · Jake Hillion
Add uv2nix to build Python packages from uv.lock. This creates a fully
Nix-managed Python environment with the Rust bindings injected via overlay.
Changes:
- Add pyproject-nix, uv2nix, and pyproject-build-systems flake inputs
- Create python/parts.nix with overlays to inject Nix-built Rust wheel
- Export packages.exo on macOS (wraps exo/exo-master/exo-worker with dashboard)
- Add checks.lint (ruff, all platforms) and checks.pytest (macOS only)
- Simplify CI typecheck job using nicknovitski/nix-develop action
- Delete .github/actions/typecheck composite action (no longer needed)
- Add no-build-package for MLX packages in pyproject.toml (use wheels)
The Python build is currently macOS-only since MLX requires Metal. Linux
support will be added once the pyproject dependencies are simplified.
Test plan:
- Run `nix flake check` on macOS to verify pytest and lint pass
- Build exo package on macOS: `nix build .#exo`
- Verify CI pipeline passes with simplified typecheck job
Files touched
D .github/actions/typecheck/action.ymlM .github/workflows/pipeline.ymlM flake.lockM flake.nixA python/parts.nix
Diff
commit e4c6a7dbb4abc3bec8a832f60c3a7c6238cb9af3
Author: Jake Hillion <jake@hillion.co.uk>
Date: Thu Jan 15 18:37:02 2026 +0000
nix: add Python packaging with uv2nix
Add uv2nix to build Python packages from uv.lock. This creates a fully
Nix-managed Python environment with the Rust bindings injected via overlay.
Changes:
- Add pyproject-nix, uv2nix, and pyproject-build-systems flake inputs
- Create python/parts.nix with overlays to inject Nix-built Rust wheel
- Export packages.exo on macOS (wraps exo/exo-master/exo-worker with dashboard)
- Add checks.lint (ruff, all platforms) and checks.pytest (macOS only)
- Simplify CI typecheck job using nicknovitski/nix-develop action
- Delete .github/actions/typecheck composite action (no longer needed)
- Add no-build-package for MLX packages in pyproject.toml (use wheels)
The Python build is currently macOS-only since MLX requires Metal. Linux
support will be added once the pyproject dependencies are simplified.
Test plan:
- Run `nix flake check` on macOS to verify pytest and lint pass
- Build exo package on macOS: `nix build .#exo`
- Verify CI pipeline passes with simplified typecheck job
---
.github/actions/typecheck/action.yml | 12 -----
.github/workflows/pipeline.yml | 82 +++++++------------------------
flake.lock | 65 ++++++++++++++++++++++---
flake.nix | 27 ++++++++---
python/parts.nix | 93 ++++++++++++++++++++++++++++++++++++
5 files changed, 190 insertions(+), 89 deletions(-)
diff --git a/.github/actions/typecheck/action.yml b/.github/actions/typecheck/action.yml
deleted file mode 100644
index cd52d6e3..00000000
--- a/.github/actions/typecheck/action.yml
+++ /dev/null
@@ -1,12 +0,0 @@
-name: Type Check
-
-description: "Run type checker"
-
-runs:
- using: "composite"
- steps:
- - name: Run type checker
- run: |
- nix --extra-experimental-features nix-command --extra-experimental-features flakes develop -c just sync
- nix --extra-experimental-features nix-command --extra-experimental-features flakes develop -c just check
- shell: bash
diff --git a/.github/workflows/pipeline.yml b/.github/workflows/pipeline.yml
index 7f6c5f04..dc93f2ad 100644
--- a/.github/workflows/pipeline.yml
+++ b/.github/workflows/pipeline.yml
@@ -26,73 +26,14 @@ jobs:
name: exo
authToken: "${{ secrets.CACHIX_AUTH_TOKEN }}"
- - name: Configure git user
- run: |
- git config --local user.email "github-actions@users.noreply.github.com"
- git config --local user.name "github-actions bot"
- shell: bash
+ - name: Load nix develop environment
+ run: nix run github:nicknovitski/nix-develop/v1
- - name: Pull LFS files
- run: |
- echo "Pulling Git LFS files..."
- git lfs pull
- shell: bash
+ - name: Sync dependencies
+ run: uv sync --all-packages
- - name: Setup Nix Environment
- run: |
- echo "Checking for nix installation..."
-
- # Check if nix binary exists directly
- if [ -f /nix/var/nix/profiles/default/bin/nix ]; then
- echo "Found nix binary at /nix/var/nix/profiles/default/bin/nix"
- export PATH="/nix/var/nix/profiles/default/bin:$PATH"
- echo "PATH=$PATH" >> $GITHUB_ENV
- nix --version
- elif [ -f /nix/var/nix/profiles/default/etc/profile.d/nix-daemon.sh ]; then
- echo "Found nix profile script, sourcing..."
- source /nix/var/nix/profiles/default/etc/profile.d/nix-daemon.sh
- nix --version
- elif command -v nix >/dev/null 2>&1; then
- echo "Nix already in PATH"
- nix --version
- else
- echo "Nix not found. Debugging info:"
- echo "Contents of /nix/var/nix/profiles/default/:"
- ls -la /nix/var/nix/profiles/default/ 2>/dev/null || echo "Directory not found"
- echo "Contents of /nix/var/nix/profiles/default/bin/:"
- ls -la /nix/var/nix/profiles/default/bin/ 2>/dev/null || echo "Directory not found"
- exit 1
- fi
- shell: bash
-
- - name: Configure basedpyright include for local MLX
- run: |
- RUNNER_LABELS='${{ toJSON(runner.labels) }}'
- if echo "$RUNNER_LABELS" | grep -q "local_mlx"; then
- if [ -d "/Users/Shared/mlx" ]; then
- echo "Updating [tool.basedpyright].include to use /Users/Shared/mlx"
- awk '
- BEGIN { in=0 }
- /^\[tool\.basedpyright\]/ { in=1; print; next }
- in && /^\[/ { in=0 } # next section
- in && /^[ \t]*include[ \t]*=/ {
- print "include = [\"/Users/Shared/mlx\"]"
- next
- }
- { print }
- ' pyproject.toml > pyproject.toml.tmp && mv pyproject.toml.tmp pyproject.toml
-
- echo "New [tool.basedpyright] section:"
- sed -n '/^\[tool\.basedpyright\]/,/^\[/p' pyproject.toml | sed '$d' || true
- else
- echo "local_mlx tag present but /Users/Shared/mlx not found; leaving pyproject unchanged."
- fi
- else
- echo "Runner does not have 'local_mlx' tag; leaving pyproject unchanged."
- fi
- shell: bash
-
- - uses: ./.github/actions/typecheck
+ - name: Run type checker
+ run: uv run basedpyright --project pyproject.toml
nix:
name: Build and check (${{ matrix.system }})
@@ -191,3 +132,14 @@ jobs:
- name: Run nix flake check
run: nix flake check
+
+ - name: Run pytest (macOS only)
+ if: runner.os == 'macOS'
+ run: |
+ # Build the test environment (requires relaxed sandbox for uv2nix on macOS)
+ TEST_ENV=$(nix build '.#exo-test-env' --option sandbox relaxed --print-out-paths)
+
+ # Run pytest outside sandbox (needs GPU access for MLX)
+ export HOME="$RUNNER_TEMP"
+ export EXO_TESTS=1
+ $TEST_ENV/bin/python -m pytest src -m "not slow" --import-mode=importlib
diff --git a/flake.lock b/flake.lock
index 3c7b28e1..86fb6e60 100644
--- a/flake.lock
+++ b/flake.lock
@@ -21,7 +21,9 @@
"nixpkgs"
],
"purescript-overlay": "purescript-overlay",
- "pyproject-nix": "pyproject-nix"
+ "pyproject-nix": [
+ "pyproject-nix"
+ ]
},
"locked": {
"lastModified": 1765953015,
@@ -149,19 +151,44 @@
"type": "github"
}
},
+ "pyproject-build-systems": {
+ "inputs": {
+ "nixpkgs": [
+ "nixpkgs"
+ ],
+ "pyproject-nix": [
+ "pyproject-nix"
+ ],
+ "uv2nix": [
+ "uv2nix"
+ ]
+ },
+ "locked": {
+ "lastModified": 1763662255,
+ "narHash": "sha256-4bocaOyLa3AfiS8KrWjZQYu+IAta05u3gYZzZ6zXbT0=",
+ "owner": "pyproject-nix",
+ "repo": "build-system-pkgs",
+ "rev": "042904167604c681a090c07eb6967b4dd4dae88c",
+ "type": "github"
+ },
+ "original": {
+ "owner": "pyproject-nix",
+ "repo": "build-system-pkgs",
+ "type": "github"
+ }
+ },
"pyproject-nix": {
"inputs": {
"nixpkgs": [
- "dream2nix",
"nixpkgs"
]
},
"locked": {
- "lastModified": 1763017646,
- "narHash": "sha256-Z+R2lveIp6Skn1VPH3taQIuMhABg1IizJd8oVdmdHsQ=",
+ "lastModified": 1764134915,
+ "narHash": "sha256-xaKvtPx6YAnA3HQVp5LwyYG1MaN4LLehpQI8xEdBvBY=",
"owner": "pyproject-nix",
"repo": "pyproject.nix",
- "rev": "47bd6f296502842643078d66128f7b5e5370790c",
+ "rev": "2c8df1383b32e5443c921f61224b198a2282a657",
"type": "github"
},
"original": {
@@ -178,7 +205,10 @@
"flake-parts": "flake-parts",
"nixpkgs": "nixpkgs",
"nixpkgs-swift": "nixpkgs-swift",
- "treefmt-nix": "treefmt-nix"
+ "pyproject-build-systems": "pyproject-build-systems",
+ "pyproject-nix": "pyproject-nix",
+ "treefmt-nix": "treefmt-nix",
+ "uv2nix": "uv2nix"
}
},
"rust-analyzer-src": {
@@ -239,6 +269,29 @@
"repo": "treefmt-nix",
"type": "github"
}
+ },
+ "uv2nix": {
+ "inputs": {
+ "nixpkgs": [
+ "nixpkgs"
+ ],
+ "pyproject-nix": [
+ "pyproject-nix"
+ ]
+ },
+ "locked": {
+ "lastModified": 1767701098,
+ "narHash": "sha256-CJhKZnWb3gumR9oTRjFvCg/6lYTGbZRU7xtvcyWIRwU=",
+ "owner": "pyproject-nix",
+ "repo": "uv2nix",
+ "rev": "9d357f0d2ce6f5f35ec7959d7e704452352eb4da",
+ "type": "github"
+ },
+ "original": {
+ "owner": "pyproject-nix",
+ "repo": "uv2nix",
+ "type": "github"
+ }
}
},
"root": "root",
diff --git a/flake.nix b/flake.nix
index c7d60ab9..f1281bb8 100644
--- a/flake.nix
+++ b/flake.nix
@@ -24,6 +24,26 @@
dream2nix = {
url = "github:nix-community/dream2nix";
inputs.nixpkgs.follows = "nixpkgs";
+ inputs.pyproject-nix.follows = "pyproject-nix";
+ };
+
+ # Python packaging with uv2nix
+ 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";
};
# Pinned nixpkgs for swift-format (swift is broken on x86_64-linux in newer nixpkgs)
@@ -48,6 +68,7 @@
inputs.treefmt-nix.flakeModule
./dashboard/parts.nix
./rust/parts.nix
+ ./python/parts.nix
];
perSystem =
@@ -88,12 +109,6 @@
};
};
- checks.lint = pkgs.runCommand "lint-check" { } ''
- export RUFF_CACHE_DIR="$TMPDIR/ruff-cache"
- ${pkgs.ruff}/bin/ruff check ${inputs.self}/
- touch $out
- '';
-
packages = lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin (
let
uvLock = builtins.fromTOML (builtins.readFile ./uv.lock);
diff --git a/python/parts.nix b/python/parts.nix
new file mode 100644
index 00000000..b777be95
--- /dev/null
+++ b/python/parts.nix
@@ -0,0 +1,93 @@
+{ inputs, ... }:
+{
+ perSystem =
+ { config, self', pkgs, lib, system, ... }:
+ let
+ # Load workspace from uv.lock
+ workspace = inputs.uv2nix.lib.workspace.loadWorkspace {
+ workspaceRoot = inputs.self;
+ };
+
+ # Create overlay from workspace
+ # Use wheels from PyPI for most packages; we override mlx with our pure Nix Metal build
+ overlay = workspace.mkPyprojectOverlay { sourcePreference = "wheel"; };
+
+ # Override overlay to inject Nix-built components
+ exoOverlay = final: prev: {
+ # Replace workspace exo_pyo3_bindings with Nix-built wheel
+ exo-pyo3-bindings = pkgs.stdenv.mkDerivation {
+ pname = "exo-pyo3-bindings";
+ version = "0.1.0";
+ src = self'.packages.exo_pyo3_bindings;
+ # Install from pre-built wheel
+ nativeBuildInputs = [ final.pyprojectWheelHook ];
+ dontStrip = true;
+ };
+ };
+
+ python = pkgs.python313;
+
+ # Overlay to provide build systems and custom packages
+ buildSystemsOverlay = final: prev: {
+ # Use our pure Nix-built MLX with Metal support
+ mlx = self'.packages.mlx;
+
+ # mlx-lm is a git dependency that needs setuptools
+ mlx-lm = prev.mlx-lm.overrideAttrs (old: {
+ nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
+ final.setuptools
+ ];
+ });
+ };
+
+ pythonSet = (pkgs.callPackage inputs.pyproject-nix.build.packages {
+ inherit python;
+ }).overrideScope (
+ lib.composeManyExtensions [
+ inputs.pyproject-build-systems.overlays.default
+ overlay
+ exoOverlay
+ buildSystemsOverlay
+ ]
+ );
+ exoVenv = pythonSet.mkVirtualEnv "exo-env" workspace.deps.default;
+
+ # Virtual environment with dev dependencies for testing
+ testVenv = pythonSet.mkVirtualEnv "exo-test-env" (
+ workspace.deps.default // {
+ exo = [ "dev" ]; # Include pytest, pytest-asyncio, pytest-env
+ }
+ );
+
+ exoPackage = pkgs.runCommand "exo"
+ {
+ nativeBuildInputs = [ pkgs.makeWrapper ];
+ }
+ ''
+ mkdir -p $out/bin
+
+ # Create wrapper scripts
+ for script in exo exo-master exo-worker; do
+ makeWrapper ${exoVenv}/bin/$script $out/bin/$script \
+ --set DASHBOARD_DIR ${self'.packages.dashboard}
+ done
+ '';
+ in
+ {
+ # Python package only available on macOS (requires MLX/Metal)
+ packages = lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin {
+ exo = exoPackage;
+ # Test environment for running pytest outside of Nix sandbox (needs GPU access)
+ exo-test-env = testVenv;
+ };
+
+ checks = {
+ # Ruff linting (works on all platforms)
+ lint = pkgs.runCommand "ruff-lint" { } ''
+ export RUFF_CACHE_DIR="$TMPDIR/ruff-cache"
+ ${pkgs.ruff}/bin/ruff check ${inputs.self}/
+ touch $out
+ '';
+ };
+ };
+}
← b1e88a3d shfmt
·
back to Exo
·
Load pipeline layers sequentially (#1329) 2837225d →