{ 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 } ); mkPythonScript = name: path: pkgs.writeShellApplication { inherit name; runtimeInputs = [ exoVenv ]; runtimeEnv = { EXO_DASHBOARD_DIR = self'.packages.dashboard; EXO_RESOURCES_DIR = inputs.self + /resources; }; text = ''exec python ${path} "$@"''; }; benchVenv = pythonSet.mkVirtualEnv "exo-bench-env" { exo-bench = [ ]; }; mkBenchScript = name: path: pkgs.writeShellApplication { inherit name; runtimeInputs = [ benchVenv ]; text = ''exec python ${path} "$@"''; }; mkSimplePythonScript = name: path: pkgs.writeShellApplication { inherit name; runtimeInputs = [ pkgs.python313 ]; text = ''exec python ${path} "$@"''; }; exoPackage = pkgs.runCommand "exo" { nativeBuildInputs = [ pkgs.makeWrapper ]; } '' mkdir -p $out/bin # Create wrapper script makeWrapper ${exoVenv}/bin/exo $out/bin/exo \ --set EXO_DASHBOARD_DIR ${self'.packages.dashboard} \ --set EXO_RESOURCES_DIR ${inputs.self + /resources} \ ${lib.optionalString pkgs.stdenv.hostPlatform.isDarwin "--prefix PATH : ${pkgs.macmon}/bin"} ''; 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; } // { exo-bench = mkBenchScript "exo-bench" (inputs.self + /bench/exo_bench.py); exo-get-all-models-on-cluster = mkSimplePythonScript "exo-get-all-models-on-cluster" (inputs.self + /tests/get_all_models_on_cluster.py); }; 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 ''; }; }; }