flake.nix

{
  description = "raft";

  nixConfig = {
    bash-prompt-prefix = "(raft) ";
    bash-prompt = ''\[\033]0;\u@\h:\w\007\]\[\033[01;32m\]\u@\h\[\033[01;34m\] \w \$\[\033[00m\]'';
    bash-prompt-suffix = " ";
  };

  inputs = {
    nixpkgs = {
      url = "github:nixos/nixpkgs";

      # IMPORTANT: This needs to follow the `nixpkgs` from nativelink so that
      # the local LRE toolchains are in sync with the remote toolchains.
      follows = "nativelink/nixpkgs";
    };
    nativelink = {
      # Note: Keep this commit in sync with the LRE commit in `ll/init.bzl`.
      url = "github:TraceMachina/nativelink/8a632953b86395088e4ab8c1e160a650739549b7";

      # IMPORTANT: This repository provides the autogenerated LRE toolchains
      # which are dependent on the nixpkgs version in the nativelink repository.
      # To keep the local LRE toolchains aligned with remote LRE, we need to use
      # the nixpkgs used by nativelink as the the "global" nixpkgs. We do this
      # by setting `nixpkgs.follows = "nativelink/nixpkgs"` above.

      inputs = {
        flake-utils.follows = "flake-utils";
        flake-parts.follows = "flake-parts";
        pre-commit-hooks.follows = "git-hooks";
      };
    };
    flake-utils.url = "github:numtide/flake-utils";
    flake-parts = {
      url = "github:hercules-ci/flake-parts";
      inputs.nixpkgs-lib.follows = "nixpkgs";
    };
    git-hooks = {
      url = "github:cachix/git-hooks.nix";
      inputs.nixpkgs.follows = "nixpkgs";
    };
    rules_ll = {
      # Note: Keep this commit in sync with the one in MODULE.bazel.
      url = "github:eomii/rules_ll/3ee809512cfb605a00fe5eb938eab0e4f8705204";
      inputs = {
        nixpkgs.follows = "nixpkgs";
        flake-parts.follows = "flake-parts";
        nativelink.follows = "nativelink";
      };
    };
  };

  outputs =
    { self
    , flake-parts
    , git-hooks
    , rules_ll
    , nativelink
    , ...
    } @ inputs:
    flake-parts.lib.mkFlake { inherit inputs; }
      {
        systems = [
          "x86_64-linux"
        ];
        imports = [
          inputs.nativelink.flakeModule
          inputs.git-hooks.flakeModule
          inputs.rules_ll.flakeModule
        ];
        perSystem =
          { config
          , pkgs
          , system
          , lib
          , ...
          }:
          {
            _module.args.pkgs =
              let
                nixpkgs-patched = (import self.inputs.nixpkgs { inherit system; }).applyPatches {
                  name = "nixpkgs-patched";
                  src = self.inputs.nixpkgs;
                  # TODO(aaronmondal): Remove when this patch is merged upstream.
                  patches = [ ./patches/nixpkgs_cuda_12_4_1.diff ];
                };
              in
              import nixpkgs-patched {
                inherit system;

                # CUDA support
                config.allowUnfree = true;
                config.cudaSupport = true;
              };
            pre-commit.settings = {
              hooks = import ./pre-commit-hooks.nix { inherit pkgs; };
            };
            local-remote-execution.settings = {
              inherit (nativelink.packages.${system}.lre-cc.meta) Env;
            };
            rules_ll.settings.llEnv =
              let
                openssl = pkgs.openssl.override { static = true; };
                cudaPackages = { };
                merged-cuda = [
                  # This is a subset of the cudatoolkit. The full distribution
                  # contains duplicates of thrust, libcudacxx and a few others.
                  # Having those in the build graph risks mixing up headers of
                  # different versions, so we omit them entirely.
                  pkgs.cudaPackages_12_4.cuda_cudart
                  pkgs.cudaPackages_12_4.cuda_cupti
                  pkgs.cudaPackages_12_4.libcublas
                  pkgs.cudaPackages_12_4.libcurand
                  pkgs.cudaPackages_12_4.libcusolver
                  pkgs.cudaPackages_12_4.libcusparse
                ];
                customCudaPackages = {
                  inherit (pkgs.cudaPackages_12_4) cuda_nvcc;
                  cudatoolkit = pkgs.symlinkJoin rec {
                    name = "cuda-merged-lucene-${version}";
                    version = "12.4.1";
                    paths = merged-cuda;
                    passthru = {
                      lib = pkgs.symlinkJoin {
                        inherit name;
                        paths = map (p: lib.getLib p) merged-cuda;
                      };
                    };
                  };
                };
              in
              rules_ll.lib.defaultLlEnv {
                inherit pkgs;
                cudaPackages = customCudaPackages;
                LL_CFLAGS = "-I${openssl.dev}/include";
                LL_LDFLAGS = "-L${openssl.out}/lib";
              };
            devShells.default = pkgs.mkShell {
              nativeBuildInputs =
                let
                  inherit (inputs.rules_ll.packages.${system}) ll;
                  bazel = pkgs.writeShellScriptBin "bazel" ''
                    unset TMPDIR TMP
                    exec ${pkgs.bazelisk}/bin/bazelisk "$@"
                  '';
                in
                [ bazel ll pkgs.kubectl ];
              shellHook = ''
                # Generate the .pre-commit-config.yaml symlink when entering the
                # development shell.
                ${config.pre-commit.installationScript}

                # Generate .bazelrc.ll which containes action-env
                # configuration when rules_ll is run from a nix environment.
                ${config.rules_ll.installationScript}

                # Generate .bazelrc.lre which configures the LRE toolchains.
                ${config.local-remote-execution.installationScript}

                # Ensure that the ll command points to our ll binary.
                [[ $(type -t ll) == "alias" ]] && unalias ll

                # Ensure that the bazel command points to our custom wrapper.
                [[ $(type -t bazel) == "alias" ]] && unalias bazel
              '';
            };
          };
      };
}