From ab44922a56e16d80997ad728a3151717830d962e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Dani=C3=ABl=20de=20Kok?= <me@danieldk.eu>
Date: Fri, 31 Oct 2025 15:29:44 +0000
Subject: [PATCH] Remove Python bytecode after checks

The included byte code will be symlinked in snapshots in the Hub cache.
However, Python replaces them by regular files, causing havoc for e.g.
hash computation in `kernels`.
---
 lib/torch-extension/arch.nix                      | 15 +++++++++++----
 overlay.nix                                       |  2 ++
 pkgs/remove-bytecode-hook/default.nix             |  5 +++++
 pkgs/remove-bytecode-hook/remove-bytecode-hook.sh | 12 ++++++++++++
 4 files changed, 30 insertions(+), 4 deletions(-)
 create mode 100644 pkgs/remove-bytecode-hook/default.nix
 create mode 100755 pkgs/remove-bytecode-hook/remove-bytecode-hook.sh

diff --git a/lib/torch-extension/arch.nix b/lib/torch-extension/arch.nix
index 20e72be7..0d5468c6 100644
--- a/lib/torch-extension/arch.nix
+++ b/lib/torch-extension/arch.nix
@@ -5,19 +5,25 @@
 
   lib,
   stdenv,
-  cudaPackages,
+
+  # Native build inputs
+  build2cmake,
   cmake,
   cmakeNvccThreadsHook,
-  ninja,
-  build2cmake,
   get-kernel-check,
   kernel-abi-check,
+  ninja,
   python3,
+  remove-bytecode-hook,
   rewrite-nix-paths-macho,
-  rocmPackages,
   writeScriptBin,
+
+  # Framework packages
+  cudaPackages,
+  rocmPackages,
   xpuPackages,
 
+  # Build inputs
   apple-sdk_15,
   clr,
   oneapi-torch-dev,
@@ -102,6 +108,7 @@ stdenv.mkDerivation (prevAttrs: {
     cmake
     ninja
     build2cmake
+    remove-bytecode-hook
   ]
   ++ lib.optionals doGetKernelCheck [
     get-kernel-check
diff --git a/overlay.nix b/overlay.nix
index 8907f6e7..97976395 100644
--- a/overlay.nix
+++ b/overlay.nix
@@ -11,6 +11,8 @@ final: prev: {
 
   rewrite-nix-paths-macho = prev.callPackage ./pkgs/rewrite-nix-paths-macho { };
 
+  remove-bytecode-hook = prev.callPackage ./pkgs/remove-bytecode-hook { };
+
   stdenvGlibc_2_27 = prev.callPackage ./pkgs/stdenv-glibc-2_27 { };
 
   # Python packages
diff --git a/pkgs/remove-bytecode-hook/default.nix b/pkgs/remove-bytecode-hook/default.nix
new file mode 100644
index 00000000..12c7f577
--- /dev/null
+++ b/pkgs/remove-bytecode-hook/default.nix
@@ -0,0 +1,5 @@
+{ makeSetupHook, python3 }:
+
+makeSetupHook {
+  name = "remove-bytecode-hook";
+} ./remove-bytecode-hook.sh
diff --git a/pkgs/remove-bytecode-hook/remove-bytecode-hook.sh b/pkgs/remove-bytecode-hook/remove-bytecode-hook.sh
new file mode 100755
index 00000000..47b6a180
--- /dev/null
+++ b/pkgs/remove-bytecode-hook/remove-bytecode-hook.sh
@@ -0,0 +1,12 @@
+#!/bin/sh
+
+echo "Sourcing remove-bytecode-hook.sh"
+
+removeBytecodeHook() {
+  echo "Removing Python bytecode"
+  find $out -type d -name '__pycache__' -exec rm -rf {} +
+}
+
+if [ -z "${dontRemoveBytecode-}" ]; then
+  appendToVar preDistPhases removeBytecodeHook
+fi