k3s.nix

{
  pkgs,
  config,
  lib,
  ...
}:
with lib; let
  cfg = config.myModules.railbird-k3s;
  mount-path = "/var/lib/railbird/bucket";
  bucket-name = "railbird-dev-videos";
  plugins-path = pkgs.buildEnv {
    name = "combined-cni-plugins";
    paths = [
      pkgs.cni-plugins
      pkgs.calico-cni-plugin
      pkgs.calico-kube-controllers
      pkgs.cni-plugin-flannel
    ];
  };
  nvidia-device-plugin-version = "v0.19.1";
  nvidia-device-plugin-manifest = pkgs.writeText "nvidia-device-plugin.yaml" ''
    apiVersion: node.k8s.io/v1
    handler: nvidia
    kind: RuntimeClass
    metadata:
      name: nvidia
      labels:
        app.kubernetes.io/component: gpu-operator
    ---
    apiVersion: apps/v1
    kind: DaemonSet
    metadata:
      name: nvidia-device-plugin-daemonset
      namespace: kube-system
      labels:
        app.kubernetes.io/name: nvidia-device-plugin
    spec:
      selector:
        matchLabels:
          app.kubernetes.io/name: nvidia-device-plugin
      updateStrategy:
        type: RollingUpdate
      template:
        metadata:
          labels:
            app.kubernetes.io/name: nvidia-device-plugin
        spec:
          runtimeClassName: nvidia
          priorityClassName: system-node-critical
          nodeSelector:
            nvidia.com/gpu.present: "true"
          tolerations:
            - key: nvidia.com/gpu
              operator: Exists
              effect: NoSchedule
          containers:
            - name: nvidia-device-plugin-ctr
              image: nvcr.io/nvidia/k8s-device-plugin:${nvidia-device-plugin-version}
              imagePullPolicy: IfNotPresent
              command: ["nvidia-device-plugin"]
              env:
                - name: DEVICE_ID_STRATEGY
                  value: uuid
                - name: NVIDIA_VISIBLE_DEVICES
                  value: all
                - name: NVIDIA_DRIVER_CAPABILITIES
                  value: compute,utility
              securityContext:
                allowPrivilegeEscalation: false
                capabilities:
                  drop: ["ALL"]
              volumeMounts:
                - name: kubelet-device-plugins-dir
                  mountPath: /var/lib/kubelet/device-plugins
                - name: cdi-specs
                  mountPath: /var/run/cdi
                  readOnly: true
          volumes:
            - name: kubelet-device-plugins-dir
              hostPath:
                path: /var/lib/kubelet/device-plugins
                type: Directory
            - name: cdi-specs
              hostPath:
                path: /var/run/cdi
                type: DirectoryOrCreate
  '';
  gpu-test-pod = pkgs.writeText "gpu-test-pod.yaml" ''
    apiVersion: v1
    kind: Pod
    metadata:
      name: gpu-test
      namespace: default
    spec:
      restartPolicy: Never
      runtimeClassName: nvidia
      containers:
        - name: cuda-test
          image: nvcr.io/nvidia/cuda:12.6.3-base-ubuntu24.04
          command: ["nvidia-smi"]
          resources:
            limits:
              nvidia.com/gpu: 1
  '';
in {
  options = {
    myModules.railbird-k3s = {
      enable = mkEnableOption "railbird k3s";
      serverAddr = mkOption {
        type = lib.types.str;
        default = "";
        description = ''
          Existing k3s server URL to join. Leave empty for the first or only
          server; that enables cluster initialization with no peers.
        '';
      };
      extraFlags = mkOption {
        type = lib.types.listOf lib.types.str;
        default = [];
      };
    };
  };
  config = mkIf cfg.enable {
    users.users.railbird = {
      isSystemUser = true;
      group = "users";
      home = "/var/lib/railbird";
      createHome = true;
    };

    age.secrets."1896Folsom-k3s-token.age".file = ./secrets/1896Folsom-k3s-token.age;
    age.secrets."k3s-registry.yaml.age".file = ./secrets/k3s-registry.yaml.age;
    age.secrets.api-service-key = {
      file = ./secrets/api_service_account_key.json.age;
      owner = "railbird";
      group = "users";
    };
    environment.etc."rancher/k3s/registries.yaml".source = config.age.secrets."k3s-registry.yaml.age".path;
    services.dockerRegistry = {
      enable = true;
      listenAddress = "0.0.0.0";
      port = 5279;
      enableDelete = true;
      enableGarbageCollect = true;
    };

    # The upstream NixOS unit runs registry GC without --delete-untagged, so
    # every manifest ever pushed to the mutable `dev` tag remains reachable.
    # Stop writes during GC, discard historical untagged manifests, and always
    # bring the registry back even if collection fails.
    systemd.services.docker-registry-garbage-collect.script = mkForce ''
      ${pkgs.systemd}/bin/systemctl stop docker-registry.service
      trap '${pkgs.systemd}/bin/systemctl start docker-registry.service' EXIT
      ${config.services.dockerRegistry.package}/bin/registry garbage-collect \
        --delete-untagged ${config.services.dockerRegistry.configFile}
    '';

    virtualisation.containerd = {
      enable = true;
      settings = {
        plugins."io.containerd.cri.v1.runtime" = {
          enable_cdi = true;
          cdi_spec_dirs = ["/var/run/cdi"];
        };
        plugins."io.containerd.grpc.v1.cri" = {
          enable_cdi = true;
          cdi_spec_dirs = ["/var/run/cdi"];
          cni.bin_dir = "${plugins-path}/bin";
        };
      };
    };

    hardware.nvidia-container-toolkit = {
      enable = true;
      device-name-strategy = "uuid";
      mount-nvidia-executables = true;
    };
    virtualisation.containers = {
      containersConf.cniPlugins = [
        pkgs.cni-plugins
        pkgs.calico-cni-plugin
        pkgs.calico-kube-controllers
        pkgs.cni-plugin-flannel
      ];
    };

    systemd.services = {
      # k3s can sit in sd_notify startup indefinitely while waiting for remote
      # etcd peers. Treat it as a long-running service so nixos-rebuild switch
      # does not block on cluster readiness.
      k3s.serviceConfig = {
        Type = mkForce "simple";
      };

      nvidia-container-toolkit-cdi-generator = {
        # Even with `--library-search-path`, `nvidia-ctk` won't find the libs
        # unless I bodge their path into the environment.
        environment.LD_LIBRARY_PATH = "${config.hardware.nvidia.package}/lib";
      };
    };

    systemd.services.mount-railbird-bucket = {
      after = ["agenix.service"];
      wantedBy = ["multi-user.target"];
      description = "Mount railbird bucket";
      serviceConfig = {
        Type = "simple";
        RemainAfterExit = true;
        Restart = "on-failure"; # Restart the service on failure
        RestartSec = 5; # Wait 5 seconds before restarti
        TimeoutStopSec = 2;
        ExecStartPre = [
          "-${pkgs.util-linux}/bin/umount -f ${mount-path}"
          "${pkgs.coreutils}/bin/mkdir -p ${mount-path}"
          "${pkgs.coreutils}/bin/chown railbird:users ${mount-path}"
          "${pkgs.coreutils}/bin/chmod 0775 ${mount-path}"
        ];
        ExecStart = let
          key-file = config.age.secrets.api-service-key.path;
        in
          pkgs.writeShellScript "mount-railbird-bucket" ''
            while true; do
            if ${pkgs.util-linux}/bin/mount | grep -q "${mount-path}" && [ -d "${mount-path}/dev" ]; then
            echo "Mount path ${mount-path} is mounted and valid (contains directory 'dev')."
            else
            echo "Mount path is not valid or not mounted, attempting remount."
            ${pkgs.util-linux}/bin/umount -f "${mount-path}" || true
            ${pkgs.gcsfuse}/bin/gcsfuse --implicit-dirs --key-file "${key-file}" "${bucket-name}" "${mount-path}"
            fi
            echo "Sleeping"
            sleep 30
            done
          '';
        User = "root";
      };
    };

    services.k3s = {
      enable = true;
      clusterInit = cfg.serverAddr == "";
      serverAddr = cfg.serverAddr;
      configPath = pkgs.writeTextFile {
        name = "k3s-config.yaml";
        text = ''
          kubelet-arg:
          - "eviction-hard=nodefs.available<2Gi,imagefs.available<2Gi"
          - "eviction-soft=nodefs.available<5Gi,imagefs.available<5Gi"
          - "eviction-soft-grace-period=nodefs.available=5m,imagefs.available=5m"
        '';
      };
      tokenFile = config.age.secrets."1896Folsom-k3s-token.age".path;
      extraFlags =
        [
          "--tls-san ryzen-shine.local"
          "--tls-san biskcomp.local"
          "--tls-san jimi-hendnix.local"
          "--tls-san dev.railbird.ai"
          "--disable=traefik"
          "--disable=servicelb"
          "--node-label nixos-nvidia-cdi=enabled"
          "--node-label nvidia.com/gpu.present=true"
          "--resolv-conf=/run/systemd/resolve/resolv.conf"
          "--etcd-arg=quota-backend-bytes=8589934592"
        ]
        ++ cfg.extraFlags;
      containerdConfigTemplate = ''
        {{ template "base" . }}

        plugins."io.containerd.grpc.v1.cri".cdi_spec_dirs = [ "/var/run/cdi" ]
        plugins."io.containerd.grpc.v1.cri".enable_cdi = true

        [plugins."io.containerd.grpc.v1.cri".containerd.runtimes.nvidia]
        privileged_without_host_devices = false
        runtime_engine = ""
        runtime_root = ""
        runtime_type = "io.containerd.runc.v2"

        [plugins."io.containerd.grpc.v1.cri".containerd.runtimes.nvidia.options]
        BinaryName = "${lib.getOutput "tools" config.hardware.nvidia-container-toolkit.package}/bin/nvidia-container-runtime.cdi"

        [debug]
        level = "trace"
      '';
      gracefulNodeShutdown = {
        enable = true;
      };
    };

    environment.systemPackages = with pkgs; [
      kubectl
      kubernetes-helm
      nvidia-container-toolkit
      nvidia-container-toolkit.tools
    ];

    environment.etc."k3s/gpu-test-pod.yaml".source = gpu-test-pod;
    environment.etc."k3s/nvidia-device-plugin.yaml".source = nvidia-device-plugin-manifest;

    systemd.services.k3s-gpu-plugin-deploy = {
      description = "Deploy NVIDIA device plugin to k3s";
      after = ["k3s.service"];
      wants = ["k3s.service"];
      wantedBy = ["multi-user.target"];
      path = [pkgs.kubectl pkgs.coreutils];
      serviceConfig = {
        Type = "oneshot";
        RemainAfterExit = true;
        ExecStart = pkgs.writeShellScript "deploy-nvidia-device-plugin" ''
          export KUBECONFIG=/etc/rancher/k3s/k3s.yaml

          echo "Waiting for k3s API server..."
          for i in $(seq 1 60); do
            if kubectl get nodes &>/dev/null; then
              echo "k3s API server is ready"
              break
            fi
            sleep 5
          done

          kubectl delete daemonset -n kube-system generic-cdi-plugin --ignore-not-found=true
          kubectl apply -f ${nvidia-device-plugin-manifest}
          kubectl rollout status daemonset/nvidia-device-plugin-daemonset -n kube-system --timeout=120s || true
        '';
      };
    };
  };
}