From 5722b3d682aa876c39662c5e0d1a85c8fe6ac9cd Mon Sep 17 00:00:00 2001 From: Ari Lotter Date: Sun, 1 Feb 2026 19:58:09 -0500 Subject: [PATCH 1/3] nix: only wrap dev env with vllm when building a package that needs it --- .../inference-node/packages.nix | 4 ++ nix/lib.nix | 55 ++++++++++++++----- nix/python.nix | 3 +- 3 files changed, 46 insertions(+), 16 deletions(-) diff --git a/architectures/inference-only/inference-node/packages.nix b/architectures/inference-only/inference-node/packages.nix index 0bd98dbb9..91c37284f 100644 --- a/architectures/inference-only/inference-node/packages.nix +++ b/architectures/inference-only/inference-node/packages.nix @@ -4,9 +4,13 @@ psycheLib.buildRustPackage { needsPython = true; needsGpu = true; cratePath = ./.; + # vllm doesn't build on macos supportedSystems = [ "x86_64-linux" "aarch64-linux" ]; + pythonBuildInputs.main = [ + "vllm" + ]; } diff --git a/nix/lib.nix b/nix/lib.nix index b0f8107f4..dd78f7573 100644 --- a/nix/lib.nix +++ b/nix/lib.nix @@ -88,6 +88,20 @@ let # python venv without the psyche extension (vllm, etc) pythonDeps = { inherit (inputs) uv2nix pyproject-nix pyproject-build-systems; }; + + # function to create a python venv with specific python packages + mkPythonVenvForPackages = + pythonPkgs: + pkgs.callPackage ./python.nix ( + { + extraPackages = { + psyche = psychePythonExtension; + }; + additionalNixPackages = pythonPkgs; + } + // pythonDeps + ); + psychePythonVenv = pkgs.callPackage ./python.nix ( { extraPackages = { }; @@ -95,15 +109,8 @@ let // pythonDeps ); - # python venv with the psyche extension - psychePythonVenvWithExtension = pkgs.callPackage ./python.nix ( - { - extraPackages = { - psyche = psychePythonExtension; - }; - } - // pythonDeps - ); + # python venv with the psyche extension (legacy, kept for compatibility) + psychePythonVenvWithExtension = mkPythonVenvForPackages [ ]; # builds a rust package # Returns an attrset of packages: { packageName = ...; packageName-nopython = ...; } # Automatically discovers and builds examples from the crate's examples/ directory @@ -116,6 +123,7 @@ let # - buildInputs.main = [ deps ] applies to src/main.rs # - buildInputs. = [ deps ] applies to all binaries of type (bin/test/example) # - buildInputs.. = [ deps ] applies to specific binary + # pythonBuildInputs: attrset of runtime python deps for each built binary, same format as buildInputs buildRustPackage = { needsPython ? false, @@ -123,14 +131,15 @@ let cratePath, # path to the crate dir supportedSystems ? null, buildInputs ? { }, + pythonBuildInputs ? { }, }: let # type: "main" | "bin" | "test" | "example" # name: the binary / test name getRuntimeDepsForArtifact = - type: name: + inputsAttrset: inputsAttrsetName: type: name: let - typeConfig = buildInputs.${type} or null; + typeConfig = inputsAttrset.${type} or null; # there's only one "main" mainDeps = @@ -139,7 +148,7 @@ let if lib.isList typeConfig then typeConfig else if typeConfig != null then - throw "buildInputs.main must be a list, got ${builtins.typeOf typeConfig}" + throw "${inputsAttrsetName}.main must be a list, got ${builtins.typeOf typeConfig}" else [ ] ) @@ -168,7 +177,23 @@ let "example" ]) "type must be 'bin', 'test', or 'example', got: ${type}"; let - workspaceArgs = if withPython then rustWorkspaceArgsWithPython else rustWorkspaceArgsNoPython; + # custom venv for this package with its specific python dependencies + pythonRuntimeDeps = + getRuntimeDepsForArtifact pythonBuildInputs "pythonBuildInputs" type + originalName; + customPythonVenv = if withPython then mkPythonVenvForPackages pythonRuntimeDeps else null; + + workspaceArgs = + if withPython then + ( + rustWorkspaceArgs + // { + buildInputs = rustWorkspaceArgs.buildInputs ++ [ customPythonVenv ]; + NIX_LDFLAGS = "-L${customPythonVenv}/lib -lpython3.12"; + } + ) + else + rustWorkspaceArgsNoPython; artifacts = if withPython then cargoArtifacts else cargoArtifactsNoPython; # delete conflicting bins from other crates to prevent ambiguous --bin/--example/--test @@ -277,8 +302,8 @@ let } ); - runtimeDeps = getRuntimeDepsForArtifact type originalName; - allRuntimeDeps = (lib.optionals withPython [ psychePythonVenvWithExtension ]) ++ runtimeDeps; + runtimeDeps = getRuntimeDepsForArtifact buildInputs "buildInputs" type originalName; + allRuntimeDeps = (lib.optionals withPython [ customPythonVenv ]) ++ runtimeDeps; wrappedRustPackage = pkgs.runCommand "${name}" diff --git a/nix/python.nix b/nix/python.nix index b24335d7a..9fbffa1a3 100644 --- a/nix/python.nix +++ b/nix/python.nix @@ -8,6 +8,7 @@ pyproject-nix, pyproject-build-systems, extraPackages ? { }, # attrset of package names to derivations to include in the venv + additionalNixPackages ? [ ], # list of additional nix-provided python package names (e.g., ["vllm"]) }: let getAllTransitiveDeps = @@ -32,8 +33,8 @@ let topLevelNixPkgs = [ "torch" ] + ++ additionalNixPackages ++ lib.optionals stdenvNoCC.hostPlatform.isLinux [ - "vllm" # for inference package "flash-attn" "liger-kernel" # i'm really not a fan of providing torchtitan like this. i'd much rather have it be built as a git dep via uv2nix. From 28975f61ff1520f17b80f1c9a07e6463e28403c1 Mon Sep 17 00:00:00 2001 From: Ari Lotter Date: Tue, 3 Feb 2026 11:36:36 -0500 Subject: [PATCH 2/3] nix: refactor to make all python deps an available output --- architectures/centralized/client/packages.nix | 4 +- .../centralized/local-testnet/packages.nix | 4 +- architectures/centralized/server/packages.nix | 4 +- .../decentralized/solana-client/packages.nix | 4 +- .../decentralized/testing/packages.nix | 4 +- .../inference-node/packages.nix | 4 +- nix/lib.nix | 114 ++++++++++++++++-- nix/pkgs.nix | 62 +--------- shared/modeling/packages.nix | 4 +- shared/network/packages.nix | 4 +- tools/rust-tools/expand-distro/packages.nix | 4 +- tools/rust-tools/preview-lr/packages.nix | 4 +- tools/rust-tools/psyche-sidecar/packages.nix | 4 +- tools/rust-tools/run-manager/packages.nix | 4 +- 14 files changed, 128 insertions(+), 96 deletions(-) diff --git a/architectures/centralized/client/packages.nix b/architectures/centralized/client/packages.nix index 1edaf4220..d9874cc6c 100644 --- a/architectures/centralized/client/packages.nix +++ b/architectures/centralized/client/packages.nix @@ -1,6 +1,6 @@ -{ psycheLib, ... }: +{ buildRustPackage, ... }: -psycheLib.buildRustPackage { +buildRustPackage { needsPython = "optional"; needsGpu = true; cratePath = ./.; diff --git a/architectures/centralized/local-testnet/packages.nix b/architectures/centralized/local-testnet/packages.nix index ca0354241..2638bd280 100644 --- a/architectures/centralized/local-testnet/packages.nix +++ b/architectures/centralized/local-testnet/packages.nix @@ -1,5 +1,5 @@ -{ psycheLib, ... }: +{ buildRustPackage, ... }: -psycheLib.buildRustPackage { +buildRustPackage { cratePath = ./.; } diff --git a/architectures/centralized/server/packages.nix b/architectures/centralized/server/packages.nix index ca0354241..2638bd280 100644 --- a/architectures/centralized/server/packages.nix +++ b/architectures/centralized/server/packages.nix @@ -1,5 +1,5 @@ -{ psycheLib, ... }: +{ buildRustPackage, ... }: -psycheLib.buildRustPackage { +buildRustPackage { cratePath = ./.; } diff --git a/architectures/decentralized/solana-client/packages.nix b/architectures/decentralized/solana-client/packages.nix index 1edaf4220..d9874cc6c 100644 --- a/architectures/decentralized/solana-client/packages.nix +++ b/architectures/decentralized/solana-client/packages.nix @@ -1,6 +1,6 @@ -{ psycheLib, ... }: +{ buildRustPackage, ... }: -psycheLib.buildRustPackage { +buildRustPackage { needsPython = "optional"; needsGpu = true; cratePath = ./.; diff --git a/architectures/decentralized/testing/packages.nix b/architectures/decentralized/testing/packages.nix index ccc2531a6..866e3aa20 100644 --- a/architectures/decentralized/testing/packages.nix +++ b/architectures/decentralized/testing/packages.nix @@ -1,5 +1,5 @@ { - psycheLib, + buildRustPackage, pkgs, inputs, ... @@ -8,7 +8,7 @@ let system = pkgs.stdenv.hostPlatform.system; in -psycheLib.buildRustPackage { +buildRustPackage { cratePath = ./.; # all tests need solana CLI and just buildInputs.test = [ diff --git a/architectures/inference-only/inference-node/packages.nix b/architectures/inference-only/inference-node/packages.nix index 91c37284f..a887ca378 100644 --- a/architectures/inference-only/inference-node/packages.nix +++ b/architectures/inference-only/inference-node/packages.nix @@ -1,6 +1,6 @@ -{ psycheLib, ... }: +{ buildRustPackage, ... }: -psycheLib.buildRustPackage { +buildRustPackage { needsPython = true; needsGpu = true; cratePath = ./.; diff --git a/nix/lib.nix b/nix/lib.nix index dd78f7573..02431cf04 100644 --- a/nix/lib.nix +++ b/nix/lib.nix @@ -483,18 +483,37 @@ let else assert (builtins.all (c: c == null) typeChecks); x; + flatPythonDeps = + let + getDepsFromType = + typeValue: + if lib.isList typeValue then + typeValue + else if lib.isAttrs typeValue then + lib.attrValues typeValue + else + [ ]; + in + lib.pipe pythonBuildInputs [ + lib.attrValues + (map getDepsFromType) + lib.flatten + lib.unique + ]; in - validateBuildInputs ( - lib.optionalAttrs shouldBuildForThisSystem ( - util.mergeAttrsetsNoConflicts "can't merge binary package sets" [ - mainRsPackage - binDirPackages - examplePackages - testPackages - ] - ) - ); - + { + inherit flatPythonDeps; + packages = validateBuildInputs ( + lib.optionalAttrs shouldBuildForThisSystem ( + util.mergeAttrsetsNoConflicts "can't merge binary package sets" [ + mainRsPackage + binDirPackages + examplePackages + testPackages + ] + ) + ); + }; # TODO: i can't set the rust build target to WASM for the build deps for wasm-pack, since *some* of them don't build. # really, i want like a wasm-only set of deps to build... can I do that? # like do the buildDepsOnly for not the workspace, but my specific package that *happens* to be in a workspace. @@ -637,6 +656,77 @@ let doInstallCargoArtifacts = false; } ); + + workspaceCargoToml = builtins.fromTOML (builtins.readFile ../Cargo.toml); + + # expand globs in workspace members from cargo.toml + expandWorkspaceMembers = + members: + lib.flatten ( + lib.map ( + memberPattern: + if lib.hasSuffix "/*" memberPattern then + let + dir = lib.removeSuffix "/*" memberPattern; + dirPath = ../${dir}; + entries = builtins.readDir dirPath; + subdirs = lib.filterAttrs (n: v: v == "directory") entries; + in + lib.mapAttrsToList (name: _: "${dir}/${name}") subdirs + else + [ memberPattern ] + ) members + ); + + expandedMembers = expandWorkspaceMembers workspaceCargoToml.workspace.members; + + # find all crates with packages.nix + discoverCratesWithPackagesNix = + members: + lib.filter (pkg: pkg != null) ( + lib.map ( + memberPath: + let + fullPath = ../${memberPath}; + packagesNixPath = fullPath + "/packages.nix"; + cargoTomlPath = fullPath + "/Cargo.toml"; + + isExcluded = builtins.elem memberPath [ + "python/" # python venv with special dependencies + ]; + + hasCargoToml = builtins.pathExists cargoTomlPath; + hasPackagesNix = builtins.pathExists packagesNixPath; + in + if hasCargoToml && hasPackagesNix && !isExcluded then + let + cargoToml = builtins.fromTOML (builtins.readFile cargoTomlPath); + packageName = cargoToml.package.name or (baseNameOf memberPath); + in + { + name = packageName; + path = fullPath; + } + else + null + ) members + ); + + rustPackageSets = lib.map ( + pkg: import (pkg.path + "/packages.nix") { inherit buildRustPackage pkgs inputs; } + ) (discoverCratesWithPackagesNix expandedMembers); + # a packages.nix returns an attrset of packages (including examples) + rustPackages = util.mergeAttrsetsNoConflicts "can't merge rust package sets." ( + lib.map (pkg: pkg.packages) rustPackageSets + ); + + allPythonDeps = lib.pipe rustPackageSets [ + lib.attrValues + (map (pkg: pkg.flatPythonDeps or [ ])) + lib.flatten + lib.unique + ]; + in { inherit @@ -654,6 +744,8 @@ in gitcommit psychePythonVenv psychePythonVenvWithExtension + allPythonDeps + rustPackages ; mkWebsitePackage = pkgs.callPackage ../website/common.nix { }; diff --git a/nix/pkgs.nix b/nix/pkgs.nix index 29bfb1b16..6baafcf58 100644 --- a/nix/pkgs.nix +++ b/nix/pkgs.nix @@ -11,60 +11,7 @@ lib.makeScope pkgs.newScope ( }; util = import ./util.nix; - workspaceCargoToml = builtins.fromTOML (builtins.readFile ../Cargo.toml); - - # expand globs in workspace members from cargo.toml - expandWorkspaceMembers = - members: - lib.flatten ( - lib.map ( - memberPattern: - if lib.hasSuffix "/*" memberPattern then - let - dir = lib.removeSuffix "/*" memberPattern; - dirPath = ../${dir}; - entries = builtins.readDir dirPath; - subdirs = lib.filterAttrs (n: v: v == "directory") entries; - in - lib.mapAttrsToList (name: _: "${dir}/${name}") subdirs - else - [ memberPattern ] - ) members - ); - - expandedMembers = expandWorkspaceMembers workspaceCargoToml.workspace.members; - - # find all crates with packages.nix - discoverCratesWithPackagesNix = - members: - lib.filter (pkg: pkg != null) ( - lib.map ( - memberPath: - let - fullPath = ../${memberPath}; - packagesNixPath = fullPath + "/packages.nix"; - cargoTomlPath = fullPath + "/Cargo.toml"; - - isExcluded = builtins.elem memberPath [ - "python/" # python venv with special dependencies - ]; - - hasCargoToml = builtins.pathExists cargoTomlPath; - hasPackagesNix = builtins.pathExists packagesNixPath; - in - if hasCargoToml && hasPackagesNix && !isExcluded then - let - cargoToml = builtins.fromTOML (builtins.readFile cargoTomlPath); - packageName = cargoToml.package.name or (baseNameOf memberPath); - in - { - name = packageName; - path = fullPath; - } - else - null - ) members - ); + inherit (psycheLib) rustPackages; externalRustPackages = { solana-toolbox-cli = pkgs.rustPlatform.buildRustPackage rec { @@ -86,13 +33,6 @@ lib.makeScope pkgs.newScope ( }; }; - # a packages.nix returns an attrset of packages (including examples) - rustPackages = util.mergeAttrsetsNoConflicts "can't merge rust package sets." ( - lib.map (pkg: import (pkg.path + "/packages.nix") { inherit psycheLib pkgs inputs; }) ( - discoverCratesWithPackagesNix expandedMembers - ) - ); - dockerPackages = import ./docker.nix { inherit pkgs diff --git a/shared/modeling/packages.nix b/shared/modeling/packages.nix index 0add9d0fb..2c7c724e0 100644 --- a/shared/modeling/packages.nix +++ b/shared/modeling/packages.nix @@ -1,6 +1,6 @@ -{ psycheLib, ... }: +{ buildRustPackage, ... }: -psycheLib.buildRustPackage { +buildRustPackage { needsGpu = true; needsPython = "optional"; cratePath = ./.; diff --git a/shared/network/packages.nix b/shared/network/packages.nix index 856bdb919..31a0327aa 100644 --- a/shared/network/packages.nix +++ b/shared/network/packages.nix @@ -1,6 +1,6 @@ -{ psycheLib, ... }: +{ buildRustPackage, ... }: -psycheLib.buildRustPackage { +buildRustPackage { needsPython = "optional"; cratePath = ./.; } diff --git a/tools/rust-tools/expand-distro/packages.nix b/tools/rust-tools/expand-distro/packages.nix index 1edaf4220..d9874cc6c 100644 --- a/tools/rust-tools/expand-distro/packages.nix +++ b/tools/rust-tools/expand-distro/packages.nix @@ -1,6 +1,6 @@ -{ psycheLib, ... }: +{ buildRustPackage, ... }: -psycheLib.buildRustPackage { +buildRustPackage { needsPython = "optional"; needsGpu = true; cratePath = ./.; diff --git a/tools/rust-tools/preview-lr/packages.nix b/tools/rust-tools/preview-lr/packages.nix index ca0354241..2638bd280 100644 --- a/tools/rust-tools/preview-lr/packages.nix +++ b/tools/rust-tools/preview-lr/packages.nix @@ -1,5 +1,5 @@ -{ psycheLib, ... }: +{ buildRustPackage, ... }: -psycheLib.buildRustPackage { +buildRustPackage { cratePath = ./.; } diff --git a/tools/rust-tools/psyche-sidecar/packages.nix b/tools/rust-tools/psyche-sidecar/packages.nix index 1edaf4220..d9874cc6c 100644 --- a/tools/rust-tools/psyche-sidecar/packages.nix +++ b/tools/rust-tools/psyche-sidecar/packages.nix @@ -1,6 +1,6 @@ -{ psycheLib, ... }: +{ buildRustPackage, ... }: -psycheLib.buildRustPackage { +buildRustPackage { needsPython = "optional"; needsGpu = true; cratePath = ./.; diff --git a/tools/rust-tools/run-manager/packages.nix b/tools/rust-tools/run-manager/packages.nix index ca0354241..2638bd280 100644 --- a/tools/rust-tools/run-manager/packages.nix +++ b/tools/rust-tools/run-manager/packages.nix @@ -1,5 +1,5 @@ -{ psycheLib, ... }: +{ buildRustPackage, ... }: -psycheLib.buildRustPackage { +buildRustPackage { cratePath = ./.; } From b92240e662f0433b8a1830242a12b7bd7b106dee Mon Sep 17 00:00:00 2001 From: Ari Lotter Date: Tue, 3 Feb 2026 12:24:45 -0500 Subject: [PATCH 3/3] nix: refactor to build venvs dynamically for devShell / checks --- nix/checks.nix | 12 +++++++- nix/devShell.nix | 13 +++++++-- nix/lib.nix | 72 +++++++++++++++++++++--------------------------- 3 files changed, 54 insertions(+), 43 deletions(-) diff --git a/nix/checks.nix b/nix/checks.nix index c1981dac9..d71ae1790 100644 --- a/nix/checks.nix +++ b/nix/checks.nix @@ -10,9 +10,19 @@ inherit (pkgs.psycheLib) craneLib rustWorkspaceArgs - rustWorkspaceArgsWithPython + mkRustWorkspaceArgsWithPython + mkPythonVenv + psychePythonExtension + allPythonDeps cargoArtifacts ; + + rustWorkspaceArgsWithPython = mkRustWorkspaceArgsWithPython (mkPythonVenv { + extraPackages = { + psyche = psychePythonExtension; + }; + additionalNixPackages = allPythonDeps; + }); in { checks = diff --git a/nix/devShell.nix b/nix/devShell.nix index 45a670c45..9af87c6f1 100644 --- a/nix/devShell.nix +++ b/nix/devShell.nix @@ -13,9 +13,18 @@ rustWorkspaceArgs craneLib env - psychePythonVenv - psychePythonVenvWithExtension + mkPythonVenv + psychePythonExtension + allPythonDeps ; + + psychePythonVenv = mkPythonVenv { }; + psychePythonVenvWithExtension = mkPythonVenv { + extraPackages = { + psyche = psychePythonExtension; + }; + additionalNixPackages = allPythonDeps; + }; in { # fmt as precommit hook diff --git a/nix/lib.nix b/nix/lib.nix index 02431cf04..271e59c8c 100644 --- a/nix/lib.nix +++ b/nix/lib.nix @@ -66,12 +66,25 @@ let cargoExtraArgs = "--features python" + lib.optionalString (pkgs.config.cudaSupport) ",parallelism"; }; - rustWorkspaceArgsWithPython = rustWorkspaceArgs // { - buildInputs = rustWorkspaceArgs.buildInputs ++ [ - psychePythonVenvWithExtension - ]; - NIX_LDFLAGS = "-L${psychePythonVenvWithExtension}/lib -lpython3.12"; - }; + psychePythonExtension = pkgs.callPackage ../python { }; + + mkPythonVenv = + { + extraPackages ? { }, # actual packages from pkgs.callPackage + additionalNixPackages ? [ ], # like "vllm", etc + }: + pkgs.callPackage ./python.nix { + inherit (inputs) uv2nix pyproject-nix pyproject-build-systems; + inherit extraPackages additionalNixPackages; + }; + + mkRustWorkspaceArgsWithPython = + venv: + rustWorkspaceArgs + // { + buildInputs = rustWorkspaceArgs.buildInputs ++ [ venv ]; + NIX_LDFLAGS = "-L${venv}/lib -lpython3.12"; + }; rustWorkspaceArgsNoPython = rustWorkspaceDeps // { inherit env src; @@ -82,35 +95,6 @@ let cargoArtifacts = craneLib.buildDepsOnly rustWorkspaceArgs; cargoArtifactsNoPython = craneLib.buildDepsOnly rustWorkspaceArgsNoPython; - - psychePythonExtension = pkgs.callPackage ../python { }; - - # python venv without the psyche extension (vllm, etc) - - pythonDeps = { inherit (inputs) uv2nix pyproject-nix pyproject-build-systems; }; - - # function to create a python venv with specific python packages - mkPythonVenvForPackages = - pythonPkgs: - pkgs.callPackage ./python.nix ( - { - extraPackages = { - psyche = psychePythonExtension; - }; - additionalNixPackages = pythonPkgs; - } - // pythonDeps - ); - - psychePythonVenv = pkgs.callPackage ./python.nix ( - { - extraPackages = { }; - } - // pythonDeps - ); - - # python venv with the psyche extension (legacy, kept for compatibility) - psychePythonVenvWithExtension = mkPythonVenvForPackages [ ]; # builds a rust package # Returns an attrset of packages: { packageName = ...; packageName-nopython = ...; } # Automatically discovers and builds examples from the crate's examples/ directory @@ -181,7 +165,16 @@ let pythonRuntimeDeps = getRuntimeDepsForArtifact pythonBuildInputs "pythonBuildInputs" type originalName; - customPythonVenv = if withPython then mkPythonVenvForPackages pythonRuntimeDeps else null; + customPythonVenv = + if withPython then + mkPythonVenv { + extraPackages = { + psyche = psychePythonExtension; + }; + additionalNixPackages = pythonRuntimeDeps; + } + else + null; workspaceArgs = if withPython then @@ -721,7 +714,6 @@ let ); allPythonDeps = lib.pipe rustPackageSets [ - lib.attrValues (map (pkg: pkg.flatPythonDeps or [ ])) lib.flatten lib.unique @@ -734,7 +726,7 @@ in craneLib buildSolanaIdl rustWorkspaceArgs - rustWorkspaceArgsWithPython + mkRustWorkspaceArgsWithPython cargoArtifacts buildRustPackage buildRustWasmTsPackage @@ -742,8 +734,8 @@ in env src gitcommit - psychePythonVenv - psychePythonVenvWithExtension + mkPythonVenv + psychePythonExtension allPythonDeps rustPackages ;