idun-apptainer-nix/flake.nix

{
  inputs.nixpkgs.url = github:NixOS/nixpkgs/nixpkgs-unstable;
  inputs.nixpkgs-unfree.url = github:SomeoneSerge/nixpkgs-unfree;
  inputs.nixpkgs-unfree.inputs.nixpkgs.follows = "nixpkgs";
  #inputs.nix2container.url = "github:nlewo/nix2container";
  #inputs.nix2container.inputs.nixpkgs.follows = "nixpkgs";

  nixConfig.extra-substituters = [
    "https://cuda-maintainers.cachix.org"
  ];
  nixConfig.extra-trusted-public-keys = [
    "cuda-maintainers.cachix.org-1:0dq3bujKpuEPMCX6U4WylrUDZ9JyUG0VpVZa7CNfq5E="
  ];

  # TODO: https://www.canva.dev/blog/engineering/supporting-gpu-accelerated-machine-learning-with-kubernetes-and-nix/

  outputs = {
    self,
    nixpkgs,
    nixpkgs-unfree,
    ... } @ inputs:
  let
    #flake = inputs: system: nixpkgs.lib.mapAttrs (name: flake: {
    #  nixos = flake.nixosModules
    #    or null;
    #  pkgs = flake.packages.${system}
    #    or flake.legacyPackages.${system}
    #    or null;
    #  lib = flake.lib.${system}
    #    or flake.lib
    #    or null;
    #}) inputs;
    forSystems = systems: f: nixpkgs.lib.genAttrs systems (system: f {
      inherit system;
      pkgs = nixpkgs-unfree.legacyPackages.${system};
      #pkgs = nixpkgs.legacyPackages.${system};
      #pkgs = nixpkgs { config.allowUnfree = true; config.cudaSupport = true; };
      lib  = nixpkgs-unfree.legacyPackages.${system}.lib;
    #  flakes = flake inputs system;
    });
    forAllSystems = forSystems [
      "x86_64-linux"
      "aarch64-linux"
    ];
  in {
    packages = forAllSystems ({ system, pkgs, lib, ...}: rec {

      # to try this, inside the default devShell do:
      #   apptainer shell $(nix build .#apptainer --print-out-paths --no-link)
      apptainer = let
        # https://github.com/NixOS/nixpkgs/issues/177908#issuecomment-1495625986
        mk-singularity = name: {
          contents,
          runscript ? "#!/bin/sh\nexec ${pkgs.hello}/bin/hello",
          startscript ? "#!/bin/sh\nexec ${pkgs.hello}/bin/hello",
          env ? {},
          shellHook ? "",
        }:
        pkgs.runCommand "${name}.sqfs" {
          outputs = [ "out" "tree" ];
          nativeBuildInputs = [ pkgs.squashfsTools ];
          env.shellHookData = shellHook;
          env.closureInfo = pkgs.closureInfo {
            rootPaths = contents ++ [ pkgs.bashInteractive ];
          };
          env.environVars = pkgs.writeText "env" (lib.pipe env [
            (lib.mapAttrsToList (key: val: "${key}=${lib.escapeShellArg val}"))
            (lib.concatStringsSep "\n")
          ]);
        } ''
          set -o pipefail
          set -x
          mkdir -p $tree/{bin,etc/ssl/certs,dev,proc,sys,usr/bin,.singularity.d/{actions,env,libs}}
          cd $tree
          cp -na --parents $(cat $closureInfo/store-paths) .
          touch etc/{passwd,group}
          #ln -s /bin usr/
          #ln -s ${pkgs.bashInteractive}/bin/bash bin/sh
          cp -a ${pkgs.pkgsStatic.bashInteractive}/bin/bash bin/sh
          cp -a ${pkgs.pkgsStatic.nix}/bin/* bin/
          cp -a ${pkgs.pkgsStatic.nix}/etc/profile.d/nix.sh .singularity.d/env/
          cp -a ${pkgs.cacert}/etc/ssl/certs/ca-bundle.crt etc/ssl/certs/ca-certificates.crt
          for p in ${lib.concatStringsSep " " contents}; do
            ln -sn $p/bin/* bin/ || true
          done
          echo "${runscript}" >.singularity.d/runscript
          echo "${startscript}" >.singularity.d/startscript
          chmod +x .singularity.d/{runscript,startscript}
          cat "$environVars" >.singularity.d/env/99-nix-env.sh
          echo "$shellHookData" >.singularity.d/env/99-nix-shell-hook.sh
          mksquashfs $tree $out -no-hardlinks -all-root
        '';
      in mk-singularity "testing123" {
        env.PS1 = "\\033[01;32m\\u@\\h\\033[33m(nix) \\033[01;34m\\W\\033[01;32m\\$\\033[00m ";
        #shellHook = ''
        #  export PS1="\033[33m(nix)\033[00m $PS1"
        #'';
        contents = with pkgs; [
          #pkgsStatic.nix
        ];
      };
      #} ''
      #    mkdir unpack
      #    tar xzvf ${docker-img}/image.tgz -C unpack
      #    # Singularity can't handle .gz
      #    tar -C unpack/ -cvf layer.tar .
      #    # TODO: Allow for module of user defined nightly, opposed to using src
      #    singularity build $out Singularity.nightly
      #  '';
      #};

      # https://nixos.org/manual/nixpkgs/stable/#ssec-pkgs-dockerTools-buildImage
      hpc-oci = pkgs.dockerTools.buildLayeredImage {
        name = "hpc-oci";
        #config.Cmd = [ "${pkgs.mysql}/bin/mysqld" ];
        config.Cmd = [ "/bin/bash" ];
        config.WorkingDir = "/data";
        config.Volumes."/data" = { };
        #copyToRoot = pkgs.buildEnv {
        #  name = "image-root";
        #  pathsToLink = [ "/bin" ];
        #  paths = with pkgs; [
        #    redis
        #  ];
        #};
      };
      #hpc-oci2 = flakes.nix2container.pkgs.nix2container.buildImage {
      #  name = "hello";
      #  config.entrypoint = ["${pkgs.hello}/bin/hello" ];
      #};
    });
    devShells = forAllSystems ({ pkgs, ...}: {


      default = pkgs.mkShellNoCC {
        #env.APPTAINER_BINDPATH = ".direnv/nix:/nix";
        #env.SINGULARITY_BINDPATH = ".direnv/nix:/nix";
        env.APPTAINER_BINDPATH = "/usr,/lib,/lib64,.direnv/nix:/nix";
        env.SINGULARITY_BINDPATH = "/usr,/lib,/lib64,.direnv/nix:/nix";
        #env.APPTAINER_BINDPATH = "/usr,/lib,/lib64";
        #env.SINGULARITY_BINDPATH = "/usr,/lib,/lib64";
        packages = with pkgs; [
          remote-exec
          (python3.withPackages (ps: with ps; [
            typer
          ]))
        ];
      };

    });
  };
}