diff --git a/modules/cog.nix b/modules/cog.nix index ce1fcc9..4969178 100644 --- 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 index 1e04615..89317b2 100644 --- 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; + }; + }; }