flake.nix 4.91 KB
Newer Older
1
2
{
  inputs = {
3
4
5
6
    crate2nix = {
      url = "github:nix-community/crate2nix";
      inputs.nixpkgs.follows = "tgi-nix/nixpkgs";
    };
7
    nix-filter.url = "github:numtide/nix-filter";
8
9
10
    tgi-nix.url = "github:danieldk/tgi-nix";
    nixpkgs.follows = "tgi-nix/nixpkgs";
    flake-utils.url = "github:numtide/flake-utils";
11
12
13
14
    rust-overlay = {
      url = "github:oxalica/rust-overlay";
      inputs.nixpkgs.follows = "tgi-nix/nixpkgs";
    };
15
16
17
18
  };
  outputs =
    {
      self,
19
      crate2nix,
20
      nix-filter,
21
22
      nixpkgs,
      flake-utils,
23
      rust-overlay,
24
25
26
27
28
      tgi-nix,
    }:
    flake-utils.lib.eachDefaultSystem (
      system:
      let
29
30
31
        cargoNix = crate2nix.tools.${system}.appliedCargoNix {
          name = "tgi";
          src = ./.;
32
          additionalCargoNixArgs = [ "--all-features" ];
33
        };
34
        pkgs = import nixpkgs {
35
36
          inherit system;
          inherit (tgi-nix.lib) config;
37
38
          overlays = [
            rust-overlay.overlays.default
39
            tgi-nix.overlays.default
40
          ];
41
        };
42
        crateOverrides = import ./nix/crate-overrides.nix { inherit pkgs nix-filter; };
43
44
45
        benchmark = cargoNix.workspaceMembers.text-generation-benchmark.build.override {
          inherit crateOverrides;
        };
46
47
48
        launcher = cargoNix.workspaceMembers.text-generation-launcher.build.override {
          inherit crateOverrides;
        };
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
        router =
          let
            routerUnwrapped = cargoNix.workspaceMembers.text-generation-router-v3.build.override {
              inherit crateOverrides;
            };
            packagePath =
              with pkgs.python3.pkgs;
              makePythonPath [
                protobuf
                sentencepiece
                torch
                transformers
              ];
          in
          pkgs.writeShellApplication {
            name = "text-generation-router";
            text = ''
              PYTHONPATH="${packagePath}" ${routerUnwrapped}/bin/text-generation-router "$@"
            '';
          };
69
        server = pkgs.python3.pkgs.callPackage ./nix/server.nix { inherit nix-filter; };
70
71
      in
      {
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
        checks = {
          rust = with pkgs; rustPlatform.buildRustPackage {
            name = "rust-checks";
            src = ./.;
            cargoLock = {
              lockFile = ./Cargo.lock;
            };
            buildInputs = [ openssl.dev ];
            nativeBuildInputs = [ clippy pkg-config protobuf python3 rustfmt ];
            buildPhase = ''
              cargo check
            '';
            checkPhase = ''
              cargo fmt -- --check
              cargo test -j $NIX_BUILD_CORES
              cargo clippy
            '';
            installPhase = "touch $out";
          } ;
        };

Nicolas Patry's avatar
Nicolas Patry committed
93
        formatter = pkgs.nixfmt-rfc-style;
94

95
96
97
98
99
        devShells = with pkgs; rec {
          default = pure;

          pure = mkShell {
            buildInputs = [
100
              benchmark
101
102
103
104
105
              launcher
              router
              server
            ];
          };
Nicolas Patry's avatar
Nicolas Patry committed
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
          test = mkShell {
            buildInputs =
              [
                # benchmark
                # launcher
                # router
                server
                openssl.dev
                pkg-config
                cargo
                rustfmt
                clippy
              ]
              ++ (with python3.pkgs; [
                docker
                pytest
                pytest-asyncio
                syrupy
                pre-commit
                ruff
              ]);

          };
129
130

          impure = mkShell {
131
132
133
134
            buildInputs =
              [
                openssl.dev
                pkg-config
135
136
137
138
139
140
                (rust-bin.stable.latest.default.override {
                  extensions = [
                    "rust-analyzer"
                    "rust-src"
                  ];
                })
141
                protobuf
142
143
144
              ]
              ++ (with python3.pkgs; [
                venvShellHook
145
                docker
146
                pip
Nicolas Patry's avatar
Nicolas Patry committed
147
                ipdb
Nicolas Patry's avatar
Nicolas Patry committed
148
                click
149
                pyright
150
151
                pytest
                pytest-asyncio
152
                ruff
153
                syrupy
154
155
              ]);

156
157
            inputsFrom = [ server ];

158
159
            venvDir = "./.venv";

160
            postVenvCreation = ''
161
              unset SOURCE_DATE_EPOCH
162
163
              ( cd server ; python -m pip install --no-dependencies -e . )
              ( cd clients/python ; python -m pip install --no-dependencies -e . )
164
165
166
            '';
            postShellHook = ''
              unset SOURCE_DATE_EPOCH
167
              export PATH=$PATH:~/.cargo/bin
168
169
            '';
          };
170
        };
171
172
173
174
175
176
177
178
179
180
181

        packages.default = pkgs.writeShellApplication {
          name = "text-generation-inference";
          runtimeInputs = [
            server
            router
          ];
          text = ''
            ${launcher}/bin/text-generation-launcher "$@"
          '';
        };
182
183
184
      }
    );
}