Remove Python bytecode after checks (#286)

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`.

Author: danieldk

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

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

pkgs/remove-bytecode-hook/default.nix

@@ -0,0 +1,5 @@
+{ makeSetupHook, python3 }:
+
+makeSetupHook {
+  name = "remove-bytecode-hook";
+} ./remove-bytecode-hook.sh

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