diff --git a/README.md b/README.md
index a0d096f..9a4290e 100644
--- a/README.md
+++ b/README.md
@@ -63,7 +63,7 @@ This is my personal nix config.
 
 #### Initial deploy
 
-- Create a iso disk image for the system by using `nix build --print-out-paths --no-link .#installer-image-<host>`
+- Create a iso disk image for the system with `nix build --print-out-paths --no-link .#installer-image-<host>`
 - dd the resulting image to a stick and boot from it on the target
 - (Optional) ssh into the target (keys are already set up)
 - Run `install-system` and reboot
diff --git a/hosts/common/zfs.nix b/hosts/common/zfs.nix
index c1dda8a..7f3df45 100644
--- a/hosts/common/zfs.nix
+++ b/hosts/common/zfs.nix
@@ -6,6 +6,10 @@
   boot.supportedFilesystems = ["zfs"];
   boot.kernelPackages = config.boot.zfs.package.latestCompatibleLinuxPackages;
 
+  # The root pool should never be imported forcefully.
+  # Failure to import is important to notice!
+  boot.zfs.forceImportRoot = false;
+
   environment.systemPackages = with pkgs; [zfs];
 
   services.zfs = {
diff --git a/hosts/nom/net.nix b/hosts/nom/net.nix
index 0723e58..32b38fa 100644
--- a/hosts/nom/net.nix
+++ b/hosts/nom/net.nix
@@ -1,4 +1,8 @@
-{nodeSecrets, ...}: {
+{
+  config,
+  nodeSecrets,
+  ...
+}: {
   networking = {
     inherit (nodeSecrets.networking) hostId;
     wireless.iwd.enable = true;
diff --git a/hosts/ward/fs.nix b/hosts/ward/fs.nix
index 0ea96fd..6d305dc 100644
--- a/hosts/ward/fs.nix
+++ b/hosts/ward/fs.nix
@@ -1,5 +1,6 @@
 {
   extraLib,
+  nodeSecrets,
   pkgs,
   ...
 }: {
@@ -7,49 +8,38 @@
     disk = {
       m2-ssd = {
         type = "disk";
-        device = "/dev/disk/by-id/nvme-Samsung_SSD_980_1TB_S649NL0TC36758M";
-        content = {
+        device = "/dev/disk/by-id/${nodeSecrets.disk.m2-ssd}";
+        content = with extraLib.disko.gpt; {
           type = "table";
           format = "gpt";
           partitions = [
-            {
-              name = "efi";
-              start = "0%";
-              end = "1GiB";
-              fs-type = "fat32";
-              bootable = true;
-              content = {
-                type = "filesystem";
-                format = "vfat";
-                mountpoint = "/boot";
-              };
-            }
-            {
-              name = "swap";
-              start = "1GiB";
-              end = "17GiB";
-              fs-type = "linux-swap";
-              content = {
-                type = "swap";
-                randomEncryption = true;
-              };
-            }
-            {
-              name = "rpool";
-              start = "17GiB";
-              end = "100%";
-              content = {
-                type = "zfs";
-                pool = "rpool";
-              };
-            }
+            (partEfi "efi" "0%" "1GiB")
+            (partSwap "swap" "1GiB" "17GiB")
+            (partZfs "rpool" "17GiB" "100%")
           ];
         };
       };
     };
-    zpool = extraLib.disko.defineEncryptedZpool "rpool" {};
+    zpool = with extraLib.disko.zfs; {
+      rpool =
+        encryptedZpool
+        // {
+          datasets = {
+            "local" = unmountable;
+            "local/root" =
+              filesystem "/"
+              // {
+                postCreateHook = "zfs snapshot rpool/local/root@blank";
+              };
+            "local/nix" = filesystem "/nix";
+            "safe" = unmountable;
+            "safe/persist" = filesystem "/persist";
+          };
+        };
+    };
   };
 
+  # After importing the rpool, rollback the root system to be empty.
   boot.initrd.systemd.services = {
     impermanence-root = {
       wantedBy = ["initrd.target"];
diff --git a/hosts/ward/secrets/secrets.nix.age b/hosts/ward/secrets/secrets.nix.age
index 7990cb5..b264915 100644
Binary files a/hosts/ward/secrets/secrets.nix.age and b/hosts/ward/secrets/secrets.nix.age differ
diff --git a/nix/lib.nix b/nix/lib.nix
index e1b3a2c..284a341 100644
--- a/nix/lib.nix
+++ b/nix/lib.nix
@@ -19,6 +19,7 @@
     mapAttrs'
     mergeAttrs
     nameValuePair
+    optionalAttrs
     partition
     recursiveUpdate
     removeSuffix
@@ -48,11 +49,36 @@ in rec {
   # True if the path or string starts with /
   isAbsolutePath = x: substring 0 1 x == "/";
 
-  # Defines a simple encrypted and compressed pool
-  # with datasets necessary datasets for use with impermanence
-  disko.defineEncryptedZpool = name:
-    recursiveUpdate {
-      ${name} = {
+  disko = {
+    gpt = {
+      partEfi = name: start: end: {
+        inherit name start end;
+        fs-type = "fat32";
+        bootable = true;
+        content = {
+          type = "filesystem";
+          format = "vfat";
+          mountpoint = "/boot";
+        };
+      };
+      partSwap = name: start: end: {
+        inherit name start end;
+        fs-type = "linux-swap";
+        content = {
+          type = "swap";
+          randomEncryption = true;
+        };
+      };
+      partZfs = name: start: end: {
+        inherit name start end;
+        content = {
+          type = "zfs";
+          pool = name;
+        };
+      };
+    };
+    zfs = {
+      encryptedZpool = {
         type = "zpool";
         mountRoot = "/mnt";
         rootFsOptions = {
@@ -69,37 +95,23 @@ in rec {
           keylocation = "prompt";
         };
         options.ashift = "12";
-        datasets = {
-          "local".type = "zfs_fs";
-          "local/root" = {
-            type = "zfs_fs";
-            postCreateHook = "zfs snapshot ${name}/local/root@blank";
-            options = {
-              canmount = "on";
-              mountpoint = "/";
-            };
-            mountpoint = "/";
-          };
-          "local/nix" = {
-            type = "zfs_fs";
-            options = {
-              canmount = "on";
-              mountpoint = "/nix";
-            };
-            mountpoint = "/nix";
-          };
-          "safe".type = "zfs_fs";
-          "safe/persist" = {
-            type = "zfs_fs";
-            options = {
-              canmount = "on";
-              mountpoint = "/persist";
-            };
-            mountpoint = "/persist";
-          };
-        };
       };
+
+      unmountable = {type = "zfs_fs";};
+      filesystem = mountpoint:
+        {
+          type = "zfs_fs";
+          options = {
+            canmount = "on";
+            inherit mountpoint;
+          };
+        }
+        // optionalAttrs (mountpoint == "/") {
+          # Required to add dependencies for initrd
+          inherit mountpoint;
+        };
     };
+  };
 
   rageMasterIdentityArgs = concatMapStrings (x: ''-i ${escapeShellArg x} '') self.secrets.masterIdentities;
   rageExtraEncryptionPubkeys =