From 77290679c49292f6eb1501404d6ad6c252251e22 Mon Sep 17 00:00:00 2001
From: ditsuke <ditsuke@protonmail.com>
Date: Tue, 20 Aug 2024 02:57:02 +0530
Subject: [PATCH] dev: Simplify devShells, restore the -extra devShell

---
 .devops/nix/devshells.nix      | 52 ++++++++++++++++++++++++++++------
 .devops/nix/package.nix        | 33 ++++++---------------
 .devops/nix/python-scripts.nix | 32 +++++++++++++--------
 3 files changed, 71 insertions(+), 46 deletions(-)

diff --git a/.devops/nix/devshells.nix b/.devops/nix/devshells.nix
index b7eed5c74e2026..9a1d1dd5525794 100644
--- a/.devops/nix/devshells.nix
+++ b/.devops/nix/devshells.nix
@@ -1,14 +1,48 @@
+{ inputs, ... }:
+
 {
   perSystem =
-    { config, lib, ... }:
     {
-      devShells = lib.pipe (config.packages) [
-        (lib.concatMapAttrs
-        (name: package: {
-          ${name} = package.passthru.shell or null;
-        }))
-        (lib.filterAttrs (name: value: value != null))
-      ];
+      config,
+      lib,
+      system,
+      ...
+    }:
+    {
+      devShells =
+        let
+          pkgs = import inputs.nixpkgs { inherit system; };
+          stdenv = pkgs.stdenv;
+          scripts = config.packages.python-scripts;
+        in
+        lib.pipe (config.packages) [
+          (lib.concatMapAttrs (
+            name: package: {
+              ${name} = pkgs.mkShell {
+                name = "${name}";
+                inputsFrom = [ package ];
+                shellHook = ''
+                  echo "Entering ${name} devShell"
+                '';
+              };
+              "${name}-extra" =
+                if (name == "python-scripts") then
+                  null
+                else
+                  pkgs.mkShell {
+                    name = "${name}-extra";
+                    inputsFrom = [
+                      package
+                      scripts
+                    ];
+                    shellHook = ''
+                    echo "Entering ${name} devShell"
+                    addToSearchPath "LD_LIBRARY_PATH" "${lib.getLib stdenv.cc.cc}/lib"
+                    '';
+                  };
+            }
+          ))
+          (lib.filterAttrs (name: value: value != null))
+        ];
     };
 }
-
diff --git a/.devops/nix/package.nix b/.devops/nix/package.nix
index d95c520fb7b848..d8c2a28d539cdb 100644
--- a/.devops/nix/package.nix
+++ b/.devops/nix/package.nix
@@ -20,12 +20,14 @@
   vulkan-loader,
   curl,
   shaderc,
-  useBlas ? builtins.all (x: !x) [
-    useCuda
-    useMetalKit
-    useRocm
-    useVulkan
-  ] && blas.meta.available,
+  useBlas ?
+    builtins.all (x: !x) [
+      useCuda
+      useMetalKit
+      useRocm
+      useVulkan
+    ]
+    && blas.meta.available,
   useCuda ? config.cudaSupport,
   useMetalKit ? stdenv.isAarch64 && stdenv.isDarwin,
   # Increases the runtime closure size by ~700M
@@ -214,25 +216,6 @@ effectiveStdenv.mkDerivation (finalAttrs: {
     cp $src/include/llama.h $out/include/
   '';
 
-  # Define the shells here, but don't add in the inputsFrom to avoid recursion.
-  passthru = {
-    inherit
-      useBlas
-      useCuda
-      useMetalKit
-      useMpi
-      useRocm
-      useVulkan
-      ;
-
-    shell = mkShell {
-      name = "shell-${finalAttrs.finalPackage.name}";
-      description = "contains numpy and sentencepiece";
-      nativeBuildInputs = [ cmake ];
-      inputsFrom = [ finalAttrs.finalPackage ];
-    };
-  };
-
   meta = {
     # Configurations we don't want even the CI to evaluate. Results in the
     # "unsupported platform" messages. This is mostly a no-op, because
diff --git a/.devops/nix/python-scripts.nix b/.devops/nix/python-scripts.nix
index 4b1a598f8601ca..392e9ffe41bf5b 100644
--- a/.devops/nix/python-scripts.nix
+++ b/.devops/nix/python-scripts.nix
@@ -3,7 +3,6 @@
   stdenv,
   buildPythonPackage,
   poetry-core,
-  breakpointHook,
   mkShell,
   python3Packages,
   gguf-py,
@@ -18,6 +17,25 @@ let
     torchWithoutCuda
     gguf-py
     tqdm
+
+    # for scripts/compare-llama-bench.py
+    gitpython
+    tabulate
+
+    # for examples/pydantic-models-to-grammar-examples.py
+    docstring-parser
+    pydantic
+
+  ];
+
+  llama-python-test-deps = with python3Packages; [
+    # Server bench
+    matplotlib
+
+    # server tests
+    openai
+    behave
+    prometheus-client
   ];
 in
 
@@ -43,16 +61,6 @@ buildPythonPackage ({
     src = lib.cleanSource ../../.;
   };
   nativeBuildInputs = [ poetry-core ];
+  nativeCheckInputs = llama-python-test-deps;
   dependencies = llama-python-deps;
-
-  passthru = {
-    shell = mkShell {
-      name = "shell-python-scripts";
-      description = "contains numpy and sentencepiece";
-      buildInputs = llama-python-deps;
-      shellHook = ''
-        addToSearchPath "LD_LIBRARY_PATH" "${lib.getLib stdenv.cc.cc}/lib"
-      '';
-    };
-  };
 })