pkgs/development/cuda-modules/generic-builders/multiplex.nix

   1 {
   2   # callPackage-provided arguments
   3   lib,
   4   cudaVersion,
   5   flags,
   6   stdenv,
   7   # Expected to be passed by the caller
   8   mkVersionedPackageName,
   9   # pname :: String
  10   pname,
  11   # releasesModule :: Path
  12   # A path to a module which provides a `releases` attribute
  13   releasesModule,
  14   # shims :: Path
  15   # A path to a module which provides a `shims` attribute
  16   # The redistribRelease is only used in ./manifest.nix for the package version
  17   # and the package description (which NVIDIA's manifest calls the "name").
  18   # It's also used for fetching the source, but we override that since we can't
  19   # re-use that portion of the functionality (different URLs, etc.).
  20   # The featureRelease is used to populate meta.platforms (by way of looking at the attribute names)
  21   # and to determine the outputs of the package.
  22   # shimFn :: {package, redistArch} -> AttrSet
  23   shimsFn ? (throw "shimsFn must be provided"),
  24   # fixupFn :: Path
  25   # A path (or nix expression) to be evaluated with callPackage and then
  26   # provided to the package's overrideAttrs function.
  27   # It must accept at least the following arguments:
  28   # - final
  29   # - cudaVersion
  30   # - mkVersionedPackageName
  31   # - package
  32   # - ...
  33   fixupFn ? (throw "fixupFn must be provided"),
  34 }:
  35 let
  36   inherit (lib)
  37     attrsets
  38     lists
  39     modules
  40     strings
  41     ;
  42
  43   inherit (stdenv) hostPlatform;
  44
  45   evaluatedModules = modules.evalModules {
  46     modules = [
  47       ../modules
  48       releasesModule
  49     ];
  50   };
  51
  52   # NOTE: Important types:
  53   # - Releases: ../modules/${pname}/releases/releases.nix
  54   # - Package: ../modules/${pname}/releases/package.nix
  55
  56   # FIXME: do this at the module system level
  57   propagatePlatforms = lib.mapAttrs (
  58     redistArch: packages: map (p: { inherit redistArch; } // p) packages
  59   );
  60
  61   # All releases across all platforms
  62   # See ../modules/${pname}/releases/releases.nix
  63   releaseSets = propagatePlatforms evaluatedModules.config.${pname}.releases;
  64
  65   # Compute versioned attribute name to be used in this package set
  66   # Patch version changes should not break the build, so we only use major and minor
  67   # computeName :: Package -> String
  68   computeName = { version, ... }: mkVersionedPackageName pname version;
  69
  70   # Check whether a package supports our CUDA version and platform.
  71   # isSupported :: Package -> Bool
  72   isSupported =
  73     package:
  74     redistArch == package.redistArch
  75     && strings.versionAtLeast cudaVersion package.minCudaVersion
  76     && strings.versionAtLeast package.maxCudaVersion cudaVersion;
  77
  78   # Get all of the packages for our given platform.
  79   # redistArch :: String
  80   # Value is `"unsupported"` if the platform is not supported.
  81   redistArch = flags.getRedistArch hostPlatform.system;
  82
  83   preferable =
  84     p1: p2: (isSupported p2 -> isSupported p1) && (strings.versionAtLeast p1.version p2.version);
  85
  86   # All the supported packages we can build for our platform.
  87   # perSystemReleases :: List Package
  88   allReleases = lib.pipe releaseSets [
  89     (lib.attrValues)
  90     (lists.flatten)
  91     (lib.groupBy (p: lib.versions.majorMinor p.version))
  92     (lib.mapAttrs (_: builtins.sort preferable))
  93     (lib.mapAttrs (_: lib.take 1))
  94     (lib.attrValues)
  95     (lib.concatMap lib.trivial.id)
  96   ];
  97
  98   newest = builtins.head (builtins.sort preferable allReleases);
  99
 100   # A function which takes the `final` overlay and the `package` being built and returns
 101   # a function to be consumed via `overrideAttrs`.
 102   overrideAttrsFixupFn =
 103     final: package:
 104     final.callPackage fixupFn {
 105       inherit
 106         final
 107         cudaVersion
 108         mkVersionedPackageName
 109         package
 110         ;
 111     };
 112
 113   extension =
 114     final: _:
 115     let
 116       # Builds our package into derivation and wraps it in a nameValuePair, where the name is the versioned name
 117       # of the package.
 118       buildPackage =
 119         package:
 120         let
 121           shims = final.callPackage shimsFn {
 122             inherit package;
 123             inherit (package) redistArch;
 124           };
 125           name = computeName package;
 126           drv = final.callPackage ./manifest.nix {
 127             inherit pname;
 128             redistName = pname;
 129             inherit (shims) redistribRelease featureRelease;
 130           };
 131           fixedDrv = drv.overrideAttrs (overrideAttrsFixupFn final package);
 132         in
 133         attrsets.nameValuePair name fixedDrv;
 134
 135       # versionedDerivations :: AttrSet Derivation
 136       versionedDerivations = builtins.listToAttrs (lists.map buildPackage allReleases);
 137
 138       defaultDerivation = {
 139         ${pname} = (buildPackage newest).value;
 140       };
 141     in
 142     versionedDerivations // defaultDerivation;
 143 in
 144 extension