From e0d1e75191f9753486394f36f4ccabea5693960d Mon Sep 17 00:00:00 2001
From: mulatta <67085791+mulatta@users.noreply.github.com>
Date: Sat, 26 Sep 2026 23:14:24 +0900
Subject: [PATCH] nix: make python shells reliable across hosts

Binary CUDA wheels depend on libraries outside their individual derivations, while source distributions can omit build metadata. Apply targeted overrides so pure environments evaluate consistently without imposing Linux behavior on Darwin.\n\nReuse host driver paths and avoid unnecessary environment synchronization so development remains practical on GPU machines and slow networks.
---
 nix/devshell.nix | 19 ++++++++++++----
 nix/python.nix   | 59 ++++++++++++++++++++++++++++++++----------------
 2 files changed, 54 insertions(+), 24 deletions(-)

diff --git a/nix/devshell.nix b/nix/devshell.nix
index 9288558..f389f2b 100644
--- a/nix/devshell.nix
+++ b/nix/devshell.nix
@@ -19,6 +19,8 @@ let
     env = {
       UV_PYTHON = pythonSet.python.interpreter;
       UV_PYTHON_DOWNLOADS = "never";
+      UV_CONCURRENT_DOWNLOADS = "4";
+      UV_HTTP_TIMEOUT = "600";
     };
     shellHook = ''
       unset PYTHONPATH
@@ -38,23 +40,32 @@ let
       }
     );
 
+  driverLibs = lib.optionalString pkgs.stdenv.hostPlatform.isLinux "${pkgs.addDriverRunpath.driverLink}/lib";
+
   uv2nix = mkDevShell {
     packages = [ virtualenv ];
-    env.UV_NO_SYNC = "1";
+    env = {
+      UV_NO_SYNC = "1";
+    }
+    // lib.optionalAttrs pkgs.stdenv.hostPlatform.isLinux {
+      LD_LIBRARY_PATH = driverLibs;
+    };
   };
 
   impure = mkDevShell {
     packages = [ pythonSet.python ];
     env = lib.optionalAttrs pkgs.stdenv.hostPlatform.isLinux {
-      LD_LIBRARY_PATH = lib.makeLibraryPath [ pkgs.stdenv.cc.cc.lib ];
+      LD_LIBRARY_PATH = lib.makeLibraryPath [ pkgs.stdenv.cc.cc.lib ] + ":" + driverLibs;
     };
     shellHook = ''
-      uv sync
+      if [ ! -x .venv/bin/python ] || [ uv.lock -nt .venv/pyvenv.cfg ]; then
+        uv sync --frozen
+      fi
       . .venv/bin/activate
     '';
   };
 in
 {
   inherit impure uv2nix;
-  default = impure;
+  default = uv2nix;
 }
diff --git a/nix/python.nix b/nix/python.nix
index 1851a86..87aa3df 100644
--- a/nix/python.nix
+++ b/nix/python.nix
@@ -9,34 +9,53 @@ let
 
   pythonOverlay = workspace.mkPyprojectOverlay { sourcePreference = "wheel"; };
 
-  # sdist-only deps that build with setuptools but omit it from
-  # build-system.requires, so uv2nix has nothing to bootstrap them with.
-  needsSetuptools = [
-    "watchdog"
-  ];
+  packageOverrides = final: prev: {
+    # uv.lock omits build-system metadata required for this sdist build.
+    watchdog = prev.watchdog.overrideAttrs (old: {
+      nativeBuildInputs =
+        old.nativeBuildInputs
+        ++ final.resolveBuildSystem {
+          setuptools = [ ];
+        };
+    });
+  };
 
-  pythonOverrides =
-    final: prev:
-    lib.genAttrs needsSetuptools (
-      name:
-      prev.${name}.overrideAttrs (old: {
-        nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [ final.setuptools ];
-      })
-    );
+  isCudaWheel =
+    name:
+    lib.hasPrefix "nvidia-" name
+    || builtins.elem name [
+      "torch"
+      "triton"
+    ];
+
+  cudaWheelOverrides =
+    _final: prev:
+    lib.mapAttrs (
+      _name: package:
+      package.overrideAttrs {
+        # Linux torch depends on CUDA libraries shipped by sibling wheels.
+        # autoPatchelf cannot resolve them across separate derivations.
+        autoPatchelfIgnoreMissingDeps = true;
+      }
+    ) (lib.filterAttrs (name: _package: isCudaWheel name) prev);
 in
 {
   inherit workspace;
 
   mkPythonSet =
     pkgs:
+    let
+      overlays = [
+        inputs.pyproject-build-systems.overlays.wheel
+        pythonOverlay
+        packageOverrides
+      ]
+      ++ lib.optionals pkgs.stdenv.hostPlatform.isLinux [
+        cudaWheelOverrides
+      ];
+    in
     (pkgs.callPackage inputs.pyproject-nix.build.packages {
       python = pkgs.python314;
     }).overrideScope
-      (
-        lib.composeManyExtensions [
-          inputs.pyproject-build-systems.overlays.wheel
-          pythonOverlay
-          pythonOverrides
-        ]
-      );
+      (lib.composeManyExtensions overlays);
 }