fix: prefer nix based library inclusion

build.toml

@@ -35,18 +35,4 @@ src = [
     "csrc/new_replicate.h",
     "csrc/new_sort.h",
     "csrc/new_sort.cu",
-]
-
-[test]
-python-git-packages = [
-    { url = "https://github.com/stanford-futuredata/stk.git", rev = "7363137", sha256 = "0m6g5l9nlwaiwybg5j8dhnz159wdpabdnkzapnn3dsifxrsb59vz" },
-    { url = "https://github.com/mosaicml/composer.git", rev = "v0.9.0", sha256 = "ekJ5nE6JwYY6Ld9kIk72R/a3iI943Gd5yvAkBHQs5aI=" },
-    # { url = "https://github.com/tgale96/grouped_gemm.git", rev = "v0.3.0", sha256 = "sha256-fS6MuDj6yQ00CSzFrmAmM20/ccvtLJ1MFjfeqdwuPl8=" }
-]
-python-packages = [
-    "tqdm",
-    "py-cpuinfo",
-    "importlib-metadata",
-    "torchmetrics",
-    # "yahp"
 ]

flake.lock

@@ -1,5 +1,21 @@
 {
   "nodes": {
+    "composer": {
+      "flake": false,
+      "locked": {
+        "lastModified": 1749592532,
+        "narHash": "sha256-VKfSWtf+Z20nP1cHiBNwFzYCuMGL0xelvd6HMyDnIhc=",
+        "owner": "mosaicml",
+        "repo": "composer",
+        "rev": "0eec49da42e7f617329f035853800211f0a54ca3",
+        "type": "github"
+      },
+      "original": {
+        "owner": "mosaicml",
+        "repo": "composer",
+        "type": "github"
+      }
+    },
     "flake-compat": {
       "locked": {
         "lastModified": 1747046372,
@@ -97,16 +113,16 @@
         ]
       },
       "locked": {
-        "lastModified": 1749577327,
-        "narHash": "sha256-Muo9lpVIlgasP56avMht54WlCO++ILy1O9Fxto9xl0Q=",
+        "lastModified": 1749765513,
+        "narHash": "sha256-/Tyhxb1v4ks5G7eewZiK9BHSn2606wH/KdCMmdeAw2c=",
         "owner": "huggingface",
         "repo": "kernel-builder",
-        "rev": "73d15b5e6f1797e3a9cfef82ed619f3cb9de87f6",
+        "rev": "a9c26e450a81296525128e3860c21ab028aa8d07",
         "type": "github"
       },
       "original": {
         "owner": "huggingface",
-        "ref": "support-custom-python-libraries-in-dev-shell",
+        "ref": "support-custom-python-libraries-in-dev-shell-nixland",
         "repo": "kernel-builder",
         "type": "github"
       }
@@ -129,7 +145,25 @@
     },
     "root": {
       "inputs": {
-        "kernel-builder": "kernel-builder"
+        "composer": "composer",
+        "kernel-builder": "kernel-builder",
+        "stk": "stk"
+      }
+    },
+    "stk": {
+      "flake": false,
+      "locked": {
+        "lastModified": 1724272107,
+        "narHash": "sha256-f6eydO4u6jasvepP25a6jacSvoUNyfKW51FxahMtz1Q=",
+        "owner": "stanford-futuredata",
+        "repo": "stk",
+        "rev": "736313768ef697ce13a0594a41b2512a0fbc9884",
+        "type": "github"
+      },
+      "original": {
+        "owner": "stanford-futuredata",
+        "repo": "stk",
+        "type": "github"
       }
     },
     "systems": {

flake.nix

@@ -1,17 +1,52 @@
 {
   description = "Flake for megablocks_moe kernel";
-
+  
   inputs = {
-    kernel-builder.url = "github:huggingface/kernel-builder/support-custom-python-libraries-in-dev-shell";
-  };
+    kernel-builder.url = "github:huggingface/kernel-builder/support-custom-python-libraries-in-dev-shell-nixland";
+    # Add libraries as inputs
+    composer = {
+      url = "github:mosaicml/composer";
+      flake = false;
+    };
+    stk = {
+      url = "github:stanford-futuredata/stk";
+      flake = false;
+    };
 
-  outputs =
-    {
-      self,
-      kernel-builder,
-    }:
+    # TODO: update to build with the correct torch version
+    # grouped_gemm = {
+    #   url = "github:tgale96/grouped_gemm";
+    #   flake = false;
+    # };
+  };
+  
+  outputs = {
+    self,
+    kernel-builder,
+    composer,
+    stk,
+    # grouped_gemm,
+  }:
     kernel-builder.lib.genFlakeOutputs {
       path = ./.;
       rev = self.shortRev or self.dirtyShortRev or self.lastModifiedDate;
+      
+      # Map custom packages to their sources
+      customPythonPackages = {
+        composer = composer;
+        stk = stk;
+        # grouped_gemm = grouped_gemm;
+      };
+      
+      pythonTestDeps = [
+        "tqdm"
+        "py-cpuinfo"
+        "importlib-metadata"
+        "torchmetrics"
+        "composer"
+        "stk" 
+        # "grouped_gemm"
+        # "yahp" # may be needed for some testing plugin
+      ];
     };
 }