flake.nix 3.43 KB
Newer Older
1
2
{
  inputs = {
3
4
5
6
    crate2nix = {
      url = "github:nix-community/crate2nix";
      inputs.nixpkgs.follows = "tgi-nix/nixpkgs";
    };
7
8
9
    tgi-nix.url = "github:danieldk/tgi-nix";
    nixpkgs.follows = "tgi-nix/nixpkgs";
    flake-utils.url = "github:numtide/flake-utils";
10
11
12
13
    naersk = {
      url = "github:nix-community/naersk";
      inputs.nixpkgs.follows = "tgi-nix/nixpkgs";
    };
Nicolas Patry's avatar
Nicolas Patry committed
14
    poetry2nix.url = "github:nix-community/poetry2nix";
15
16
17
18
    rust-overlay = {
      url = "github:oxalica/rust-overlay";
      inputs.nixpkgs.follows = "tgi-nix/nixpkgs";
    };
19
20
21
22
  };
  outputs =
    {
      self,
23
      crate2nix,
24
      naersk,
25
26
      nixpkgs,
      flake-utils,
27
      rust-overlay,
28
      tgi-nix,
Nicolas Patry's avatar
Nicolas Patry committed
29
      poetry2nix,
30
31
32
33
    }:
    flake-utils.lib.eachDefaultSystem (
      system:
      let
34
35
36
37
        cargoNix = crate2nix.tools.${system}.appliedCargoNix {
          name = "tgi";
          src = ./.;
        };
38
39
40
41
42
43
        config = {
          allowUnfree = true;
          cudaSupport = true;
        };
        pkgs = import nixpkgs {
          inherit config system;
44
45
46
47
          overlays = [
            rust-overlay.overlays.default
            tgi-nix.overlay
          ];
48
        };
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
        naersk' = pkgs.callPackage naersk { };
        router =
          with pkgs;
          naersk'.buildPackage {
            name = "router";
            src = ./.;
            cargoBuildOptions =
              x:
              x
              ++ [
                "-p"
                "text-generation-router-v3"
              ];
            nativeBuildInputs = [ pkg-config ];
            buildInputs = [
              openssl.dev
              protobuf
            ];
            doCheck = false;
          };
Nicolas Patry's avatar
Nicolas Patry committed
69
70

        inherit (poetry2nix.lib.mkPoetry2Nix { inherit pkgs; }) mkPoetryEditablePackage;
71
        text-generation-server = mkPoetryEditablePackage { editablePackageSources = ./server; };
72
73
      in
      {
74
        defaultPackage = router;
75
76
77
78
79
80
81
        devShells.default =
          with pkgs;
          mkShell {
            buildInputs =
              [
                openssl.dev
                pkg-config
82
83
84
85
86
87
                (rust-bin.stable.latest.default.override {
                  extensions = [
                    "rust-analyzer"
                    "rust-src"
                  ];
                })
88
89
90
91
92
              ]
              ++ (with python3.pkgs; [
                venvShellHook
                pip

93
                causal-conv1d
Nicolas Patry's avatar
Nicolas Patry committed
94
                click
95
96
                einops
                fbgemm-gpu
Nicolas Patry's avatar
Nicolas Patry committed
97
                flashinfer
98
99
100
101
102
103
                flash-attn
                flash-attn-layer-norm
                flash-attn-rotary
                grpc-interceptor
                grpcio-reflection
                grpcio-status
Nicolas Patry's avatar
Nicolas Patry committed
104
                grpcio-tools
105
106
                hf-transfer
                loguru
107
                mamba-ssm
108
109
110
111
112
113
114
115
116
117
                marlin-kernels
                opentelemetry-api
                opentelemetry-exporter-otlp
                opentelemetry-instrumentation-grpc
                opentelemetry-semantic-conventions
                peft
                tokenizers
                torch
                transformers
                vllm
118

119
120
                cargoNix.workspaceMembers.text-generation-launcher.build
                router
121
122
123
124
125
126
127
128
129
130
131
132
133
134
              ]);

            venvDir = "./.venv";

            postVenv = ''
              unset SOURCE_DATE_EPOCH
            '';
            postShellHook = ''
              unset SOURCE_DATE_EPOCH
            '';
          };
      }
    );
}