From 1a336c714ca4c95b961e1e7bd66189ea5a5f26d5 Mon Sep 17 00:00:00 2001 From: Ari Lotter Date: Thu, 22 Jan 2026 11:28:26 -0500 Subject: [PATCH 01/11] VLLM nix packaging second attempt! - add vllm package and update the world fixes some cuda lib path problems in vllm updates mdbook fixes some clippy lints later, we should consider splitting the vllm code into a separate set of deps, so the regular psyche binaries don't get vllm included Original PR was https://github.com/PsycheFoundation/psyche/pull/464 --- Cargo.lock | 51 +++--- Cargo.toml | 6 +- flake.lock | 66 ++++---- nix/devShell.nix | 10 +- nix/nixpkgs.nix | 16 +- nix/python.nix | 1 + psyche-book/book.toml | 1 - psyche-book/default.nix | 49 +++--- psyche-book/src/development/python.md | 6 +- psyche-book/theme/index.hbs | 210 ++++++++++++++---------- shared/client/src/state/steps.rs | 2 +- shared/inference/src/protocol.rs | 4 +- shared/modeling/src/python_causal_lm.rs | 1 - website/common.nix | 7 +- website/wasm/Cargo.toml | 2 +- 15 files changed, 236 insertions(+), 196 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 24ee94226..6a6b5d6ab 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4841,9 +4841,9 @@ checksum = "00810f1d8b74be64b13dbf3db89ac67740615d6c891f0e7b6179326533011a07" [[package]] name = "js-sys" -version = "0.3.81" +version = "0.3.85" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ec48937a97411dcb524a265206ccd4c90bb711fca92b2792c407f268825b9305" +checksum = "8c942ebf8e95485ca0d52d97da7c5a2c387d0e7f0ba4c35e93bfcaee045955b3" dependencies = [ "once_cell", "wasm-bindgen", @@ -7480,7 +7480,7 @@ dependencies = [ [[package]] name = "pyo3-tch" version = "0.22.0" -source = "git+https://github.com/jquesnelle/tch-rs.git?rev=11d1ca2ef6dbd3f1e5b0986fab0a90fbb6734496#11d1ca2ef6dbd3f1e5b0986fab0a90fbb6734496" +source = "git+https://github.com/jquesnelle/tch-rs.git?rev=dda507e05a776547a112b6854d1e611684f8c729#dda507e05a776547a112b6854d1e611684f8c729" dependencies = [ "pyo3", "tch", @@ -11810,7 +11810,7 @@ dependencies = [ [[package]] name = "tch" version = "0.22.0" -source = "git+https://github.com/jquesnelle/tch-rs.git?rev=11d1ca2ef6dbd3f1e5b0986fab0a90fbb6734496#11d1ca2ef6dbd3f1e5b0986fab0a90fbb6734496" +source = "git+https://github.com/jquesnelle/tch-rs.git?rev=dda507e05a776547a112b6854d1e611684f8c729#dda507e05a776547a112b6854d1e611684f8c729" dependencies = [ "half", "lazy_static", @@ -12308,7 +12308,7 @@ dependencies = [ [[package]] name = "torch-sys" version = "0.22.0" -source = "git+https://github.com/jquesnelle/tch-rs.git?rev=11d1ca2ef6dbd3f1e5b0986fab0a90fbb6734496#11d1ca2ef6dbd3f1e5b0986fab0a90fbb6734496" +source = "git+https://github.com/jquesnelle/tch-rs.git?rev=dda507e05a776547a112b6854d1e611684f8c729#dda507e05a776547a112b6854d1e611684f8c729" dependencies = [ "anyhow", "cc", @@ -12870,9 +12870,9 @@ dependencies = [ [[package]] name = "wasm-bindgen" -version = "0.2.104" +version = "0.2.108" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c1da10c01ae9f1ae40cbfac0bac3b1e724b320abfcf52229f80b547c0d250e2d" +checksum = "64024a30ec1e37399cf85a7ffefebdb72205ca1c972291c51512360d90bd8566" dependencies = [ "cfg-if", "once_cell", @@ -12881,27 +12881,14 @@ dependencies = [ "wasm-bindgen-shared", ] -[[package]] -name = "wasm-bindgen-backend" -version = "0.2.104" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "671c9a5a66f49d8a47345ab942e2cb93c7d1d0339065d4f8139c486121b43b19" -dependencies = [ - "bumpalo", - "log", - "proc-macro2", - "quote", - "syn 2.0.106", - "wasm-bindgen-shared", -] - [[package]] name = "wasm-bindgen-futures" -version = "0.4.54" +version = "0.4.58" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7e038d41e478cc73bae0ff9b36c60cff1c98b8f38f8d7e8061e79ee63608ac5c" +checksum = "70a6e77fd0ae8029c9ea0063f87c46fde723e7d887703d74ad2616d792e51e6f" dependencies = [ "cfg-if", + "futures-util", "js-sys", "once_cell", "wasm-bindgen", @@ -12910,9 +12897,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro" -version = "0.2.104" +version = "0.2.108" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7ca60477e4c59f5f2986c50191cd972e3a50d8a95603bc9434501cf156a9a119" +checksum = "008b239d9c740232e71bd39e8ef6429d27097518b6b30bdf9086833bd5b6d608" dependencies = [ "quote", "wasm-bindgen-macro-support", @@ -12920,22 +12907,22 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro-support" -version = "0.2.104" +version = "0.2.108" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9f07d2f20d4da7b26400c9f4a0511e6e0345b040694e8a75bd41d578fa4421d7" +checksum = "5256bae2d58f54820e6490f9839c49780dff84c65aeab9e772f15d5f0e913a55" dependencies = [ + "bumpalo", "proc-macro2", "quote", "syn 2.0.106", - "wasm-bindgen-backend", "wasm-bindgen-shared", ] [[package]] name = "wasm-bindgen-shared" -version = "0.2.104" +version = "0.2.108" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bad67dc8b2a1a6e5448428adec4c3e84c43e561d8c9ee8a9e5aabeb193ec41d1" +checksum = "1f01b580c9ac74c8d8f0c0e4afb04eeef2acf145458e52c03845ee9cd23e3d12" dependencies = [ "unicode-ident", ] @@ -12955,9 +12942,9 @@ dependencies = [ [[package]] name = "web-sys" -version = "0.3.81" +version = "0.3.85" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9367c417a924a74cae129e6a2ae3b47fabb1f8995595ab474029da749a8be120" +checksum = "312e32e551d92129218ea9a2452120f4aabc03529ef03e4d0d82fb2780608598" dependencies = [ "js-sys", "wasm-bindgen", diff --git a/Cargo.toml b/Cargo.toml index 7b3968ade..36f1ba565 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -79,9 +79,9 @@ indicatif = "0.17.5" tokenizers = { version = "0.20.0", default-features = false, features = [ "onig", ] } -tch = { git = "https://github.com/jquesnelle/tch-rs.git", rev = "11d1ca2ef6dbd3f1e5b0986fab0a90fbb6734496" } -torch-sys = { git = "https://github.com/jquesnelle/tch-rs.git", rev = "11d1ca2ef6dbd3f1e5b0986fab0a90fbb6734496" } -pyo3-tch = { git = "https://github.com/jquesnelle/tch-rs.git", rev = "11d1ca2ef6dbd3f1e5b0986fab0a90fbb6734496" } +tch = { git = "https://github.com/jquesnelle/tch-rs.git", rev = "dda507e05a776547a112b6854d1e611684f8c729" } +torch-sys = { git = "https://github.com/jquesnelle/tch-rs.git", rev = "dda507e05a776547a112b6854d1e611684f8c729" } +pyo3-tch = { git = "https://github.com/jquesnelle/tch-rs.git", rev = "dda507e05a776547a112b6854d1e611684f8c729" } #tch = { path = "../tch-rs" } #torch-sys = { path = "../tch-rs/torch-sys" } #pyo3-tch = { path = "../tch-rs/pyo3-tch" } diff --git a/flake.lock b/flake.lock index 67a262751..e690f0f15 100644 --- a/flake.lock +++ b/flake.lock @@ -10,11 +10,11 @@ "systems": "systems" }, "locked": { - "lastModified": 1761656077, - "narHash": "sha256-lsNWuj4Z+pE7s0bd2OKicOFq9bK86JE0ZGeKJbNqb94=", + "lastModified": 1762618334, + "narHash": "sha256-wyT7Pl6tMFbFrs8Lk/TlEs81N6L+VSybPfiIgzU8lbQ=", "owner": "ryantm", "repo": "agenix", - "rev": "9ba0d85de3eaa7afeab493fed622008b6e4924f5", + "rev": "fcdea223397448d35d9b31f798479227e80183f6", "type": "github" }, "original": { @@ -54,11 +54,11 @@ }, "crane": { "locked": { - "lastModified": 1762189950, - "narHash": "sha256-aotggLUXjlDGqKWibGPQcMZJGgdr79S21ISrv1Wz6RI=", + "lastModified": 1768873933, + "narHash": "sha256-CfyzdaeLNGkyAHp3kT5vjvXhA1pVVK7nyDziYxCPsNk=", "owner": "ipetkov", "repo": "crane", - "rev": "50700219af884287ad7c85507e2f163b23a027a9", + "rev": "0bda7e7d005ccb5522a76d11ccfbf562b71953ca", "type": "github" }, "original": { @@ -90,11 +90,11 @@ ] }, "locked": { - "lastModified": 1762040540, - "narHash": "sha256-z5PlZ47j50VNF3R+IMS9LmzI5fYRGY/Z5O5tol1c9I4=", + "lastModified": 1768135262, + "narHash": "sha256-PVvu7OqHBGWN16zSi6tEmPwwHQ4rLPU9Plvs8/1TUBY=", "owner": "hercules-ci", "repo": "flake-parts", - "rev": "0010412d62a25d959151790968765a70c436598b", + "rev": "80daad04eddbbf5a4d883996a73f3f542fa437ac", "type": "github" }, "original": { @@ -213,11 +213,11 @@ ] }, "locked": { - "lastModified": 1760663237, - "narHash": "sha256-BflA6U4AM1bzuRMR8QqzPXqh8sWVCNDzOdsxXEguJIc=", + "lastModified": 1769069492, + "narHash": "sha256-Efs3VUPelRduf3PpfPP2ovEB4CXT7vHf8W+xc49RL/U=", "owner": "cachix", "repo": "git-hooks.nix", - "rev": "ca5b894d3e3e151ffc1db040b6ce4dcc75d31c37", + "rev": "a1ef738813b15cf8ec759bdff5761b027e3e1d23", "type": "github" }, "original": { @@ -300,11 +300,11 @@ }, "nixpkgs": { "locked": { - "lastModified": 1762111121, - "narHash": "sha256-4vhDuZ7OZaZmKKrnDpxLZZpGIJvAeMtK6FKLJYUtAdw=", + "lastModified": 1769018530, + "narHash": "sha256-MJ27Cy2NtBEV5tsK+YraYr2g851f3Fl1LpNHDzDX15c=", "owner": "NixOS", "repo": "nixpkgs", - "rev": "b3d51a0365f6695e7dd5cdf3e180604530ed33b4", + "rev": "88d3861acdd3d2f0e361767018218e51810df8a1", "type": "github" }, "original": { @@ -339,11 +339,11 @@ "uv2nix": "uv2nix" }, "locked": { - "lastModified": 1761781027, - "narHash": "sha256-YDvxPAm2WnxrznRqWwHLjryBGG5Ey1ATEJXrON+TWt8=", + "lastModified": 1763662255, + "narHash": "sha256-4bocaOyLa3AfiS8KrWjZQYu+IAta05u3gYZzZ6zXbT0=", "owner": "pyproject-nix", "repo": "build-system-pkgs", - "rev": "795a980d25301e5133eca37adae37283ec3c8e66", + "rev": "042904167604c681a090c07eb6967b4dd4dae88c", "type": "github" }, "original": { @@ -380,11 +380,11 @@ ] }, "locked": { - "lastModified": 1760402624, - "narHash": "sha256-jF6UKLs2uGc2rtved8Vrt58oTWjTQoAssuYs/0578Z4=", + "lastModified": 1768984832, + "narHash": "sha256-Q/5ICDALWlG+MbfLoghvXb3K9S0jqGt0mj8Y8RN+60k=", "owner": "pyproject-nix", "repo": "pyproject.nix", - "rev": "84c4ea102127c77058ea1ed7be7300261fafc7d2", + "rev": "c372512cad744dbdeb65dba864376a0a5ab339a7", "type": "github" }, "original": { @@ -401,11 +401,11 @@ ] }, "locked": { - "lastModified": 1760402624, - "narHash": "sha256-jF6UKLs2uGc2rtved8Vrt58oTWjTQoAssuYs/0578Z4=", + "lastModified": 1763716960, + "narHash": "sha256-PUlomle4klGbnZr0wOn8z61Mbt7tXh6Yp3hZ9/CQkq0=", "owner": "pyproject-nix", "repo": "pyproject.nix", - "rev": "84c4ea102127c77058ea1ed7be7300261fafc7d2", + "rev": "d6c61dbe0be75e2f4cf0efcdc62428175be4cfb5", "type": "github" }, "original": { @@ -461,11 +461,11 @@ ] }, "locked": { - "lastModified": 1762137611, - "narHash": "sha256-sTqb10FR/YQCuGbw16qxliX0NFlYg6evSEjN8w+9IYE=", + "lastModified": 1769091129, + "narHash": "sha256-Jj/vIHjiu4OdDIrDXZ3xOPCJrMZZKzhE2UIVXV/NYzY=", "owner": "oxalica", "repo": "rust-overlay", - "rev": "3a0ebe5d2965692f990cb27e62f501ad35e3deeb", + "rev": "131e22d6a6d54ab72aeef6a5a661ab7005b4c596", "type": "github" }, "original": { @@ -529,11 +529,11 @@ ] }, "locked": { - "lastModified": 1761311587, - "narHash": "sha256-Msq86cR5SjozQGCnC6H8C+0cD4rnx91BPltZ9KK613Y=", + "lastModified": 1768158989, + "narHash": "sha256-67vyT1+xClLldnumAzCTBvU0jLZ1YBcf4vANRWP3+Ak=", "owner": "numtide", "repo": "treefmt-nix", - "rev": "2eddae033e4e74bf581c2d1dfa101f9033dbd2dc", + "rev": "e96d59dff5c0d7fddb9d113ba108f03c3ef99eca", "type": "github" }, "original": { @@ -575,11 +575,11 @@ "pyproject-nix": "pyproject-nix_3" }, "locked": { - "lastModified": 1761872265, - "narHash": "sha256-i25GRgp2vUOebY70L3NTAgkd+Pr1hnn5xM3qHxH0ONU=", + "lastModified": 1768526251, + "narHash": "sha256-uT0ICn7tknFtdOP7zhkgOkusq38EcttqUCS5Ho7nIvc=", "owner": "pyproject-nix", "repo": "uv2nix", - "rev": "74dfb62871be152ad3673b143b0cc56105a4f3c5", + "rev": "1b8f4d4c874f909de24639142c1141c84119b0f0", "type": "github" }, "original": { diff --git a/nix/devShell.nix b/nix/devShell.nix index 1fe818760..45a670c45 100644 --- a/nix/devShell.nix +++ b/nix/devShell.nix @@ -101,10 +101,7 @@ echo "Welcome to the Psyche development shell."; ''; }; - in - { - default = craneLib.devShell defaultShell; - dev-python = craneLib.devShell ( + pythonShell = craneLib.devShell ( defaultShell // { packages = defaultShell.packages ++ [ @@ -119,6 +116,11 @@ ''; } ); + in + { + default = craneLib.devShell defaultShell; + python = pythonShell; + dev-python = pythonShell; # old name, kept for backwards compatibility }; }; } diff --git a/nix/nixpkgs.nix b/nix/nixpkgs.nix index 61d3c5b6b..c1f9a070d 100644 --- a/nix/nixpkgs.nix +++ b/nix/nixpkgs.nix @@ -15,13 +15,27 @@ in overlays = lib.optionals cudaSupported [ inputs.nix-gl-host.overlays.default ] ++ [ inputs.rust-overlay.overlays.default (final: prev: { + # temporary fix until https://github.com/NixOS/nixpkgs/pull/471394 is merged + # that lets us use `torch` instead of `torch-bin` + cudaPackages = prev.cudaPackages // { + cudnn = prev.cudaPackages.cudnn.overrideAttrs (old: { + appendRunpaths = (old.appendRunpaths or [ ]) ++ [ + "${prev.lib.getLib prev.cudaPackages.cuda_nvrtc}/lib" + ]; + }); + }; + + # provide packages for uv2pip to include python312Packages = prev.python312Packages.override { overrides = pyfinal: pyprev: { - torch = pyfinal.torch-bin; flash-attn = pyfinal.callPackage ../python/flash-attn.nix { }; liger-kernel = pyfinal.callPackage ../python/liger-kernel.nix { }; torchtitan = pyfinal.callPackage ../python/torchtitan.nix { }; torchdata = pyfinal.callPackage ../python/torchdata.nix { }; + keras = pyprev.keras.overrideAttrs (oldAttrs: { + dontUsePytestCheck = "please dont"; + + }); }; }; }) diff --git a/nix/python.nix b/nix/python.nix index 16f09970c..9d96434cb 100644 --- a/nix/python.nix +++ b/nix/python.nix @@ -31,6 +31,7 @@ let # packages that we provide to the venv via nix derivations topLevelNixPkgs = [ "torch" + "vllm" ] ++ lib.optionals stdenvNoCC.hostPlatform.isLinux [ "flash-attn" diff --git a/psyche-book/book.toml b/psyche-book/book.toml index b15598f64..697126b5b 100644 --- a/psyche-book/book.toml +++ b/psyche-book/book.toml @@ -1,7 +1,6 @@ [book] authors = ["Nous Research"] language = "en" -multilingual = false src = "src" title = "Psyche" diff --git a/psyche-book/default.nix b/psyche-book/default.nix index 0bf0e4a34..712bcf857 100644 --- a/psyche-book/default.nix +++ b/psyche-book/default.nix @@ -1,38 +1,14 @@ { lib, - mdbook, - fetchFromGitHub, - rustPlatform, stdenvNoCC, + mdbook, mdbook-mermaid, mdbook-linkcheck, + fetchFromGitHub, # custom args rustPackages, }: -let - mdbook-0-4-47 = mdbook.overrideAttrs ( - oldAttrs: - let - version = "0.4.47"; - src = fetchFromGitHub { - owner = "rust-lang"; - repo = "mdBook"; - tag = "v${version}"; - hash = "sha256-XTvC2pGRVat0kOybNb9TziG32wDVexnFx2ahmpUFmaA="; - }; - in - { - inherit version src; - cargoDeps = rustPlatform.fetchCargoVendor { - inherit (oldAttrs) pname; - inherit version src; - allowGitDependencies = false; - hash = "sha256-ASPRBAB+elJuyXpPQBm3WI97wD3mjoO1hw0fNHc+KAw="; - }; - } - ); -in let # Pull crate binary package names from rustPackages # prefer -nopython suffix if available, otherwise use normal version @@ -52,9 +28,26 @@ stdenvNoCC.mkDerivation { src = ./.; nativeBuildInputs = [ - mdbook-0-4-47 + mdbook mdbook-mermaid - mdbook-linkcheck + (mdbook-linkcheck.overrideAttrs ( + final: prev: { + version = "unstable-2025-12-04"; + src = fetchFromGitHub { + owner = "schilkp"; + repo = "mdbook-linkcheck"; + rev = "ed981be6ded11562e604fff290ae4c08f1c419c5"; + sha256 = "sha256-GTVWc/vkqY9Hml2fmm3iCHOzd/HPP1i/8NIIjFqGGbQ="; + }; + + cargoDeps = prev.cargoDeps.overrideAttrs (previousAttrs: { + vendorStaging = previousAttrs.vendorStaging.overrideAttrs { + inherit (final) src; + outputHash = "sha256-+73aI/jt5mu6dR6PR9Q08hPdOsWukb/z9crIdMMeF7U="; + }; + }); + } + )) ]; postPatch = '' diff --git a/psyche-book/src/development/python.md b/psyche-book/src/development/python.md index 62339776d..aaad835a4 100644 --- a/psyche-book/src/development/python.md +++ b/psyche-book/src/development/python.md @@ -24,7 +24,7 @@ This shell provides: ### Development Workflow -You can use `uv pip` to install arbitrary packages. Dependencies are tracked via `uv.lock`, so if you don't have `direnv` set up, you must exit and re-enter the development shell with `nix develop .#dev-python`. +You can use `uv pip` to install arbitrary packages. Dependencies are tracked via `uv.lock`, so if you don't have `direnv` set up, you must exit and re-enter the development shell with `nix develop .#python`. When you enter the dev shell, it compiles the Rust extension that provides the `psyche` Python module. **If you modify any Rust code in the Python extension or its dependencies, you must exit and re-enter the dev shell** to recompile the extension. @@ -33,7 +33,7 @@ We recommend running commands directly through the dev shell without entering it For example, to run the `train` program using python: ```bash -nix develop .#dev-python --command just train-model-python \ +nix develop .#python --command just train-model-python \ --model emozilla/llama2-20m-init \ --data-path ./data/fineweb-10bt/ \ --total-batch 2 \ @@ -44,7 +44,7 @@ nix develop .#dev-python --command just train-model-python \ Alternatively, you _could_ enter the shell and run the commands with: ```bash -nix develop .#dev-python +nix develop .#python ``` but **this is likely to be a footgun** as it's easy to forget to exit and re-enter the shell. diff --git a/psyche-book/theme/index.hbs b/psyche-book/theme/index.hbs index 373287457..c8d98fe91 100644 --- a/psyche-book/theme/index.hbs +++ b/psyche-book/theme/index.hbs @@ -33,15 +33,12 @@ {{/if}} - - {{#if copy_fonts}} - {{/if}} - - - + + + {{#each additional_css}} @@ -53,21 +50,38 @@ {{/if}} - + -
+
+
+

Keyboard shortcuts

+
+

Press ← or → to navigate between chapters

+ {{#if search_enabled}} +

Press S or / to search in the book

+ {{/if}} +

Press ? to show this help

+

Press Esc to hide this help

+
+
+
+
- + -