deformops/flake.nix at main · tue-p8n/deformops · 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
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
# Support for various CUDA versions.
# See https://github.com/NixOS/nixpkgs/blob/nixos-unstable/doc/languages-frameworks/cuda.section.md

{
  description = "Deformable grid sampling operations.";

  inputs = {
    nixpkgs.url = "github:NixOS/nixpkgs/nixos-unstable";
    nix-gl-host.url = "github:numtide/nix-gl-host";
    flake-parts.url = "github:hercules-ci/flake-parts";
  };

  outputs =
    inputs@{ self, flake-parts, ... }:
    flake-parts.lib.mkFlake { inherit inputs; } {
      imports = [
      ];
      systems = [
        "x86_64-linux"
      ];
      perSystem =
        {
          config,
          system,
          ...
        }:
        let
          # Nixpkgs package set with CUDA support enabled.
          pkgs = import inputs.nixpkgs {
            inherit system;
            config = {
              #cudaCapabilities = [ <target-architectures> ];
              cudaForwardCompat = true;
              cudaSupport = true;
              allowUnfree = true;
            };
          };

          # OCI images
          ociImages = import ./environments/oci.nix { inherit inputs pkgs; };
          ociPackages = pkgs.lib.mapAttrs' (name: value: {
            name = "oci-${name}";
            value = value;
          }) ociImages;

          # UV
          uvBuild = import ./environments/uv.nix;
          uvFHS = [
            (uvBuild {
              inherit inputs;
              pkgs = pkgs.cudaPackages_12_6.pkgs;
              name = "cu126";
            })
            (uvBuild {
              inherit inputs;
              pkgs = pkgs.cudaPackages_12_8.pkgs;
              name = "cu128";
            })
            (uvBuild {
              inherit inputs;
              pkgs = pkgs.cudaPackages_12_9.pkgs;
              name = "cu129";
            })
          ];
          uvShells = builtins.listToAttrs (
            builtins.map (fhs: {
              name = fhs.name;
              value = fhs.env;
            }) uvFHS
          );

          # Micromamba
          mmBuild = import ./environments/micromamba.nix pkgs;
          mmFHS = builtins.map mmBuild [
            "py313cu128"
            "py313cu129"
          ];
          mmShells = builtins.listToAttrs (
            builtins.map (fhs: {
              name = fhs.name;
              value = fhs.env;
            }) mmFHS
          );
        in
        {
          packages = ociPackages;
          devShells = {
            default = self.devShells.${system}.uv-cu129;
          }
          // uvShells
          // mmShells;
          formatter = pkgs.alejandra;
        };
      flake = { };
    };
}