ModelHackathon/devenv.nix at main · ZantFoundation/ModelHackathon · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
{ pkgs, lib, ... }:
let
  buildInputs = with pkgs; [
    # cudaPackages.cuda_cudart
    # cudaPackages.cudatoolkit
    # cudaPackages.cudnn
    stdenv.cc.cc
    libuv
    zlib
  ];
in
{
  # packages = with pkgs; [
  #   cudaPackages.cuda_nvcc
  # ];

  env = {
    LD_LIBRARY_PATH = "${lib.makeLibraryPath buildInputs}:/run/opengl-driver/lib:/run/opengl-driver-32/lib";

    # XLA_FLAGS =
    #   "--xla_gpu_cuda_data_dir=${pkgs.cudaPackages.cudatoolkit} "
    #   + "--xla_gpu_triton_gemm_any=True "
    #   + "--xla_gpu_enable_custom_fusions=true "
    #   + "--xla_gpu_enable_cudnn_fmha=false "
    #   + "--xla_gpu_enable_latency_hiding_scheduler=true ";
    # XLA_PYTHON_CLIENT_MEM_FRACTION = 0.9;
    # CUDA_PATH = pkgs.cudaPackages.cudatoolkit;
  };

  languages.python = {
    enable = true;
    uv = {
      enable = true;
      sync.enable = true;
    };
  };

  enterShell = ''
    . .devenv/state/venv/bin/activate
  '';

}