Upstream some overrides from cog-triton

datakami · Apr 4, 2024 · be26c1e · be26c1e
1 parent af421fd
commit be26c1e
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 1 deletion.
diff --git a/modules/cog.nix b/modules/cog.nix
@@ -1,4 +1,4 @@
-{ config, lib, dream2nix, pkgs, ... }:
+{ config, lib, dream2nix, pkgs, extendModules, ... }:
 let
   cognixcfg = config.cognix;
   cfg = config.cog.build;
@@ -102,6 +102,8 @@ in {
       nixpkgsVer = pkgs.lib.trivial.revisionWithDefault "nixos-unstable";
       in lib.mkIf config.cognix.includeNix
       "nixpkgs=https://github.com/nixos/nixpkgs/archive/refs/heads/${nixpkgsVer}.tar.xz";
+    # allow overriding via extendModules
+    public.extendModules = extendModules;
 
     dockerTools.streamLayeredImage = {
       passthru.entirePackage = entirePackage;

diff --git a/overrides.nix b/overrides.nix
@@ -12,4 +12,12 @@
   nvidia-cublas-cu12.env.appendRunpaths = [ "/usr/lib64" "$ORIGIN" ];
   nvidia-cudnn-cu12.env.appendRunpaths = [ "/usr/lib64" "$ORIGIN" ];
   nvidia-curand-cu12.env.appendRunpaths = [ "/usr/lib64" "$ORIGIN" ];
+
+  mpi4py = { config, lib, ... }: {
+    mkDerivation.buildInputs = [ config.deps.openmpi ];
+    mkDerivation.nativeBuildInputs = [ config.deps.openmpi ];
+    deps = { nixpkgs, ... }: {
+      openmpi = lib.mkDefault nixpkgs.openmpi;
+    };
+  };
 }