From 516392d7909a9dfc306b2fe617be404cd59a9104 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20de=20Kok?= Date: Tue, 20 Aug 2024 22:07:33 +0200 Subject: [PATCH] nix: add pure server to flake, add both pure and impure devshells (#2430) * nix: pure server and support both pure and impure devShells * nix: remove unused poetry2nix input It is not wired up and we now have a pure server. * nix: add ipdb to impure devshell --- flake.lock | 115 ------------------------------------------------- flake.nix | 64 +++++++++++---------------- nix/server.nix | 105 ++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 129 insertions(+), 155 deletions(-) create mode 100644 nix/server.nix diff --git a/flake.lock b/flake.lock index cd5d6d2a..c20dd98a 100644 --- a/flake.lock +++ b/flake.lock @@ -492,24 +492,6 @@ "type": "github" } }, - "flake-utils_7": { - "inputs": { - "systems": "systems_7" - }, - "locked": { - "lastModified": 1710146030, - "narHash": "sha256-SZ5L6eA7HJ/nmkzGG7/ISclqe6oZdOZTNoesiInkXPQ=", - "owner": "numtide", - "repo": "flake-utils", - "rev": "b1d9ab70662946ef0850d488da1c9019f3a9752a", - "type": "github" - }, - "original": { - "owner": "numtide", - "repo": "flake-utils", - "type": "github" - } - }, "gitignore": { "inputs": { "nixpkgs": [ @@ -594,27 +576,6 @@ "type": "github" } }, - "nix-github-actions": { - "inputs": { - "nixpkgs": [ - "poetry2nix", - "nixpkgs" - ] - }, - "locked": { - "lastModified": 1703863825, - "narHash": "sha256-rXwqjtwiGKJheXB43ybM8NwWB8rO2dSRrEqes0S7F5Y=", - "owner": "nix-community", - "repo": "nix-github-actions", - "rev": "5163432afc817cf8bd1f031418d1869e4c9d5547", - "type": "github" - }, - "original": { - "owner": "nix-community", - "repo": "nix-github-actions", - "type": "github" - } - }, "nix-test-runner": { "flake": false, "locked": { @@ -753,31 +714,6 @@ "type": "github" } }, - "poetry2nix": { - "inputs": { - "flake-utils": "flake-utils_7", - "nix-github-actions": "nix-github-actions", - "nixpkgs": [ - "tgi-nix", - "nixpkgs" - ], - "systems": "systems_8", - "treefmt-nix": "treefmt-nix" - }, - "locked": { - "lastModified": 1723854676, - "narHash": "sha256-+BrHfNuXrqeE7PoV6xDaoh0joYiJkvTTCIV0fFR3THw=", - "owner": "nix-community", - "repo": "poetry2nix", - "rev": "d650118bce34c0238b9b54f23f7f173f9e4db867", - "type": "github" - }, - "original": { - "owner": "nix-community", - "repo": "poetry2nix", - "type": "github" - } - }, "pre-commit-hooks": { "inputs": { "flake-compat": [ @@ -887,7 +823,6 @@ "tgi-nix", "nixpkgs" ], - "poetry2nix": "poetry2nix", "rust-overlay": "rust-overlay", "tgi-nix": "tgi-nix" } @@ -1003,35 +938,6 @@ "type": "github" } }, - "systems_7": { - "locked": { - "lastModified": 1681028828, - "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=", - "owner": "nix-systems", - "repo": "default", - "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e", - "type": "github" - }, - "original": { - "owner": "nix-systems", - "repo": "default", - "type": "github" - } - }, - "systems_8": { - "locked": { - "lastModified": 1681028828, - "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=", - "owner": "nix-systems", - "repo": "default", - "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e", - "type": "github" - }, - "original": { - "id": "systems", - "type": "indirect" - } - }, "tgi-nix": { "inputs": { "flake-compat": "flake-compat_4", @@ -1050,27 +956,6 @@ "repo": "tgi-nix", "type": "github" } - }, - "treefmt-nix": { - "inputs": { - "nixpkgs": [ - "poetry2nix", - "nixpkgs" - ] - }, - "locked": { - "lastModified": 1719749022, - "narHash": "sha256-ddPKHcqaKCIFSFc/cvxS14goUhCOAwsM1PbMr0ZtHMg=", - "owner": "numtide", - "repo": "treefmt-nix", - "rev": "8df5ff62195d4e67e2264df0b7f5e8c9995fd0bd", - "type": "github" - }, - "original": { - "owner": "numtide", - "repo": "treefmt-nix", - "type": "github" - } } }, "root": "root", diff --git a/flake.nix b/flake.nix index 299e6b3d..adc70fd1 100644 --- a/flake.nix +++ b/flake.nix @@ -8,10 +8,6 @@ tgi-nix.url = "github:danieldk/tgi-nix"; nixpkgs.follows = "tgi-nix/nixpkgs"; flake-utils.url = "github:numtide/flake-utils"; - poetry2nix = { - url = "github:nix-community/poetry2nix"; - inputs.nixpkgs.follows = "tgi-nix/nixpkgs"; - }; rust-overlay = { url = "github:oxalica/rust-overlay"; inputs.nixpkgs.follows = "tgi-nix/nixpkgs"; @@ -26,7 +22,6 @@ flake-utils, rust-overlay, tgi-nix, - poetry2nix, }: flake-utils.lib.eachDefaultSystem ( system: @@ -47,14 +42,28 @@ tgi-nix.overlay ]; }; - inherit (poetry2nix.lib.mkPoetry2Nix { inherit pkgs; }) mkPoetryEditablePackage; - text-generation-server = mkPoetryEditablePackage { editablePackageSources = ./server; }; crateOverrides = import ./nix/crate-overrides.nix { inherit pkgs nix-filter; }; + launcher = cargoNix.workspaceMembers.text-generation-launcher.build.override { + inherit crateOverrides; + }; + router = cargoNix.workspaceMembers.text-generation-router-v3.build.override { + inherit crateOverrides; + }; + server = pkgs.python3.pkgs.callPackage ./nix/server.nix { inherit nix-filter; }; in { - devShells.default = - with pkgs; - mkShell { + devShells = with pkgs; rec { + default = pure; + + pure = mkShell { + buildInputs = [ + launcher + router + server + ]; + }; + + impure = mkShell { buildInputs = [ openssl.dev @@ -65,43 +74,16 @@ "rust-src" ]; }) + protobuf ] ++ (with python3.pkgs; [ venvShellHook pip - - causal-conv1d - click - einops - exllamav2 - fbgemm-gpu - flashinfer - flash-attn - flash-attn-layer-norm - flash-attn-rotary - grpc-interceptor - grpcio-reflection - grpcio-status - grpcio-tools - hf-transfer ipdb - loguru - mamba-ssm - marlin-kernels - opentelemetry-api - opentelemetry-exporter-otlp - opentelemetry-instrumentation-grpc - opentelemetry-semantic-conventions - peft - tokenizers - torch - transformers - vllm - - (cargoNix.workspaceMembers.text-generation-launcher.build.override { inherit crateOverrides; }) - (cargoNix.workspaceMembers.text-generation-router-v3.build.override { inherit crateOverrides; }) ]); + inputsFrom = [ server ]; + venvDir = "./.venv"; postVenv = '' @@ -109,8 +91,10 @@ ''; postShellHook = '' unset SOURCE_DATE_EPOCH + export PATH=$PATH:~/.cargo/bin ''; }; + }; } ); } diff --git a/nix/server.nix b/nix/server.nix new file mode 100644 index 00000000..ff40757a --- /dev/null +++ b/nix/server.nix @@ -0,0 +1,105 @@ +{ + nix-filter, + buildPythonPackage, + poetry-core, + mypy-protobuf, + causal-conv1d, + einops, + exllamav2, + fbgemm-gpu, + flashinfer, + flash-attn, + flash-attn-layer-norm, + flash-attn-rotary, + grpc-interceptor, + grpcio-reflection, + grpcio-status, + grpcio-tools, + hf-transfer, + loguru, + mamba-ssm, + marlin-kernels, + opentelemetry-api, + opentelemetry-exporter-otlp, + opentelemetry-instrumentation-grpc, + opentelemetry-semantic-conventions, + peft, + safetensors, + tokenizers, + sentencepiece, + transformers, + typer, + vllm, +}: + +let + filter = nix-filter.lib; +in +buildPythonPackage { + name = "text-generation-server"; + + src = filter { + root = ../.; + include = with filter; [ + isDirectory + (and (inDirectory "server") (or_ (matchExt "py") (matchExt "pyi"))) + "server/pyproject.toml" + (and (inDirectory "proto/v3") (matchExt "proto")) + ]; + }; + + pyproject = true; + + build-system = [ poetry-core ]; + + nativeBuildInputs = [ mypy-protobuf ]; + + pythonRelaxDeps = [ + "einops" + "huggingface-hub" + "loguru" + "opentelemetry-instrumentation-grpc" + "sentencepiece" + "typer" + ]; + + pythonRemoveDeps = [ "scipy" ]; + + dependencies = [ + causal-conv1d + einops + exllamav2 + fbgemm-gpu + flashinfer + flash-attn + flash-attn-layer-norm + flash-attn-rotary + grpc-interceptor + grpcio-reflection + grpcio-status + grpcio-tools + hf-transfer + loguru + mamba-ssm + marlin-kernels + opentelemetry-api + opentelemetry-exporter-otlp + opentelemetry-instrumentation-grpc + opentelemetry-semantic-conventions + peft + safetensors + sentencepiece + tokenizers + transformers + typer + vllm + ]; + + prePatch = '' + python -m grpc_tools.protoc -Iproto/v3 --python_out=server/text_generation_server/pb \ + --grpc_python_out=server/text_generation_server/pb --mypy_out=server/text_generation_server/pb proto/v3/generate.proto + find server/text_generation_server/pb/ -type f -name "*.py" -print0 -exec sed -i -e 's/^\(import.*pb2\)/from . \1/g' {} \; + touch server/text_generation_server/pb/__init__.py + cd server + ''; +}