Full Code of GuillaumeDesforges/fix-python for AI

Repository: GuillaumeDesforges/fix-python
Branch: master
Commit: 248e2ea9620f
Files: 5
Total size: 9.4 KB

Directory structure:
gitextract__09hmi36/

├── .github/
│   └── workflows/
│       └── build.yaml
├── .gitignore
├── README.md
├── fix-python
└── flake.nix

================================================
FILE CONTENTS
================================================

================================================
FILE: .github/workflows/build.yaml
================================================
name: Check

on:
  push:
    branches:
      - main
      - master
      - action
  pull_request:
    branches:
      - main
      - master

jobs:
  check-numpy-import:
    runs-on: ubuntu-latest

    steps:
      - name: Checkout repository
        uses: actions/checkout@v2

      - name: Set up Nix
        uses: cachix/install-nix-action@v22

      - name: Install Python 3.11 with Nix
        run: nix --extra-experimental-features nix-command --extra-experimental-features flakes run nixpkgs#python311 -- -m venv .venv --copies

      - name: check that using it on numpy works
        run: |
          echo "Install numpy"
          .venv/bin/pip install numpy
          echo "Run fix-python"
          ./fix-python --venv .venv
          echo "Check numpy works"
          .venv/bin/python -c 'import numpy'


================================================
FILE: .gitignore
================================================
result
.nix
.venv
libs.nix


================================================
FILE: README.md
================================================
# fix-python

Work with Python "normally" on NixOS in one command!

Tired of all these "*.so not found" errors?
Change the RPATH of all the binaries in your venv!

## Requirements

- Nix
- `nix-command` and `flakes` experimental features must be enabled

## Install

Use temporarily in a shell

```
nix shell github:GuillaumeDesforges/fix-python
```

Or add it to your profile

```
nix profile install github:GuillaumeDesforges/fix-python
```

## Usage

In your Python project, create a virtual environment `.venv` and use your preferred tool (pip, poetry, ...) to install your dependencies.

> [!IMPORTANT]
> `fix-python` will patch binary files in the virtual environment _without_ following symlinks.
> 
> For example, if you use the `venv` module from Python's standard library, please make sure you pass `--copies` when creating the environment.
> ```bash
> python -m venv .venv --copies
> ```

By default, `fix-python` patches the packages given in the following expression:
```nix
let pkgs = import (builtins.getFlake "nixpkgs") { };
in [
  pkgs.gcc.cc
  pkgs.glibc
  pkgs.zlib
]
```

> Note: these three packages are fundamental for most Python packages and should never be removed.

If you need to patch packages in addition to these, create a `.nix/libs.nix` file with a structure similar to the above that returns the array of packages that you want binaries to be linked with.

> Note: you may add this `.nix` folder to your project `.gitignore`.

Finally, call `fix-python`.

```console
fix-python --venv .venv [--libs .nix/libs.nix]
```

The list of options is:

```
$ ./fix-python --help
Usage: fix-python --venv .venv [--libs libs.nix] [--no-default-libs]
--help: show this help message
--venv: path to Python virtual environment
--libs: path to a Nix file which returns a list of derivations
--no-default-libs: don't patch C++ standard libraries, glibc, and zlib by default
--gpu: enable GPU support
--with-torch: fix pytorch dependencies issues
--deep: looks for anything executable to patch, very slow but needed sometimes (e.g. PyQt)
--verbose: increase verbosity
```


================================================
FILE: fix-python
================================================
#!/usr/bin/env bash
set -e

# This script fix issues with Python binaries on NixOS
# Usage:
# fix-python --venv .venv [--libs libs.nix] [--no-default-libs]

DEFAULT_LIBS_EXPRESSION="
(
  let pkgs = import (builtins.getFlake \"nixpkgs\") { };
  in [
    pkgs.gcc.cc
    pkgs.glibc
    pkgs.zlib
  ]
)
"

# Help
if [ "$1" = "--help" ]; then
  echo "Usage: fix-python --venv .venv [--libs libs.nix] [--no-default-libs]" >&2
  echo "--help: show this help message" >&2
  echo "--venv: path to Python virtual environment" >&2
  echo "--libs: path to a Nix file which returns a list of derivations" >&2
  echo "--no-default-libs: don't patch C++ standard libraries, glibc, and zlib by default" >&2
  echo "--gpu: enable GPU support" >&2
  echo "--with-torch: fix pytorch dependencies issues" >&2
  echo "--deep: looks for anything executable to patch, very slow but needed sometimes (e.g. PyQt)" >&2
  echo "--verbose: increase verbosity" >&2
  exit 0
fi

# arguments
while [ $# -gt 0 ]; do
  case "$1" in
    --venv)
      shift
      VENV_PATH="$1"
      ;;
    --libs)
      shift
      LIBS_PATH="$1"
      ;;
    --no-default-libs)
      shift
      DEFAULT_LIBS_EXPRESSION="[]"
      ;;
    --gpu)
      enable_gpu="1"
      ;;
    --with-torch)
      enable_torch="1"
      ;;
    --deep)
      deep="1"
      ;;
    --verbose)
      verbose="1"
      ;;
    *)
      echo "Unknown argument: $1" >&2
      echo "Usage: fix-python --venv .venv [--libs libs.nix] [--no-default-libs]" >&2
      exit 1
      ;;
  esac
  shift
done

# check arguments
if [ -z "$VENV_PATH" ]; then
  echo "Missing argument: --venv" >&2
  echo "Usage: fix-python --venv .venv [--libs libs.nix] [--no-default-libs]" >&2
  echo "or set VENV_PATH" >&2
  exit 1
fi

# check runtime dependencies are installed
if ! command -v file &> /dev/null
then
  echo "Automatically adding \"file\" to PATH." >&2
  dep_file_path="$(nix build --no-link --print-out-paths nixpkgs#file.out)/bin"
  export PATH="$dep_file_path:$PATH"
  if [ "$verbose" ]; then 
    echo "dep_file_path=$dep_file_path" >&2
  fi
fi
if ! command -v patchelf &> /dev/null
then
  echo "Automatically adding \"patchelf\" to PATH." >&2
  dep_patchelf_path="$(nix build --no-link --print-out-paths nixpkgs#patchelf)/bin"
  export PATH="$dep_patchelf_path:$PATH"
  if [ "$verbose" ]; then 
    echo "dep_patchelf_path=$dep_patchelf_path" >&2
  fi
fi

# load libs from Nix file
if [ "$LIBS_PATH" ];
then
  # if $LIBS_PATH is just a file in the current working directory,
  # specified without leading "./", we add "./" so that $LIBS_PATH
  # can be interpreted directly as a Nix path in the following
  # expression
  if [[ ! "$LIBS_PATH" == *"/"* ]];
  then
    LIBS_PATH="./$LIBS_PATH"
  fi
  custom_libs_expression="(import $LIBS_PATH)"
  mkdir -p .nix/fix-python
  nix_libs_build_status=$(
    nix build --impure --expr "$custom_libs_expression" -o .nix/fix-python/result
    echo $?
  )
  if [ "$nix_libs_build_status" -eq "1" ];
  then
    echo "Failed to load libs from Nix file $LIBS_PATH" >&2
    echo "" >&2
    echo "Try to debug this issue with the command:" >&2
    echo "    nix build --impure --expr \"import $LIBS_PATH\"" >&2
    exit 1
  fi
else
  custom_libs_expression="[]"
fi
all_nix_libs_expression="($custom_libs_expression ++ $DEFAULT_LIBS_EXPRESSION)"
nixos_python_nix_libs="$(nix eval --impure --expr "let pkgs = import (builtins.getFlake \"nixpkgs\") {}; in pkgs.lib.strings.makeLibraryPath $all_nix_libs_expression" | sed 's/^"\(.*\)"$/\1/')"
if [ "$verbose" ]; then 
  echo "nixos_python_nix_libs=$nixos_python_nix_libs" >&2
fi
libs="$nixos_python_nix_libs"

# load libs from virtual environment
python_venv_libs=$(echo "$(find "$(realpath "$VENV_PATH")" -name '*.libs'):$(find "$(realpath "$VENV_PATH")" -name 'lib')" | tr '\n' ':')
if [ "$verbose" ]; then 
  echo "nixos_python_venv_libs=$python_venv_libs" >&2
fi
libs="$libs:$python_venv_libs"

# load libs from NixOS for GPU support if requested
if [ "$enable_gpu" ]; then
  nixos_gpu_libs="$(readlink /run/opengl-driver)/lib"
  if [ "$verbose" ]; then 
    echo "nixos_gpu_libs=$nixos_gpu_libs" >&2
  fi
  libs="$libs:$nixos_gpu_libs"
fi

# put it all together
libs=$(echo "$libs" | sed 's/:\+/:/g' | sed 's/^://' | sed 's/:$//')
if [ "$verbose" ]; then 
  echo "libs=$libs" >&2
fi

# patch each binary file found in the virtual environment
# shellcheck disable=SC2156
echo "Searching for files to patch in $VENV_PATH" >&2
if [ "$deep" ]; then
  echo "Deep search for binary files" >&2
  # For context, see #19
  binary_files=$(find "$(realpath "$VENV_PATH")" -type f -exec sh -c "file -i '{}' | grep -qE 'application/x-(executable|sharedlib); charset=binary'" \; -print)
else
  echo "Fast search for binary files" >&2
  binary_files=$(find "$(realpath "$VENV_PATH")" -type f -executable -exec sh -c "file -i '{}' | grep -qE 'x-(.*); charset=binary'" \; -print)
fi
n_binary_files=$(wc -l <<< "$binary_files")
echo "Found $n_binary_files binary files" >&2

cat <<< "$binary_files" \
  | while read -r file
    do
      echo "Patching file: $file" >&2
      old_rpath="$(patchelf --print-rpath "$file" || true)"
      # prevent duplicates
      new_rpath="$(echo "$libs:$old_rpath"  | sed 's/:$//' | tr ':' '\n' | sort --unique | tr '\n' ':' | sed 's/^://' | sed 's/:$/\n/')"
      patchelf --set-rpath "$new_rpath" "$file" || true
      old_interpreter=$(patchelf --print-interpreter "$file" || true)
      if [ -n "$old_interpreter" ]; then
        interpreter_name="$(basename "$old_interpreter")"
        new_interpreter="$(echo "$new_rpath" | tr ':' '\n' | xargs -I {} find {} -name "$interpreter_name" | head -1)"
        patchelf --set-interpreter "$new_interpreter" "$file" || true
      fi
      echo >&2
    done

# `libtorch_global_deps.so` depends on libstdc++ but does not properly declare it, fix it manually see
# https://github.com/eth-sri/lmql/blob/main/scripts/flake.d/overrides.nix#L28-L40
if [ "$enable_torch" ]; then
    torch_files=$(find "$(realpath "$VENV_PATH")" -name libtorch_global_deps.so)
    cat <<< "$torch_files" \
      | while read -r file
        do
          echo "Patching torch file: $file" >&2
          patchelf $file --add-needed libstdc++.so
        done
fi


================================================
FILE: flake.nix
================================================
{
  outputs = { self, flake-utils, nixpkgs, ...}:
    flake-utils.lib.eachDefaultSystem (system:
      let
        pkgs = nixpkgs.legacyPackages.${system};
      in
        {
          packages.default =
            pkgs.stdenv.mkDerivation {
              name = "fix-python";
              src = self;
              phases = [ "unpackPhase" "installPhase" ];
              installPhase = ''
                mkdir -p $out/bin
                cp $src/fix-python $out/bin
              '';
            };
        }
    );
}