UNDER CONSTRUCTION

Nix, Rust, Python

I have lately been working on coupled cluster Monte Carlo.1 Our pilot code is pure Python: easy to write, but the performance is not optimal. We are thus in the process of rewriting the most computationally intensive parts of the code in a compiled language.

Originally, that compiled language was going to be C++. pybind11 indeed works perfectly for writing Python extensions. Then I learned about Rust and decided to try that instead:

  1. I thought it would be a good idea to learn a new language;
  2. The memory safety guarantees and their potential benefits for parallel and concurrent programming are very enthusing;
  3. There is not runtime library apart from GLIBC;
  4. The tooling is excellent: unobtrusive and shipped with the language.

Both Python and Rust can talk to C, so that’s ultimately what we’ll have to do to get our extension to work properly. Fortunately, we don’t have to go all the way to C. There are 2 tools that can help in writing Python extensions in Rust:

  • Using CFFI. Slightly lower level and it can involve a lot of boilerplate.
  • Using PyO3. Aims at doing what pybind11 (and Boost.Python before it) did for C++ extensions. It is more bleeding edge and what I chose to work with.

This presentation has more information on writing Python extensions in Rust.

In this post I’ll explain:

  1. How to set up a mixed-language development environment.
  2. How I set up my mixed-language development environment using Nix, Pipenv, and direnv.
  3. How to deploy a manylinux1-compatible package using Travis CI for tagged releases.

The public example repository in on GitHub: [robertodr/rustafarian**

Development environment

Rust dependencies

These are for developers only. First and foremost, we need the Rust toolchain. Easy enough with rustup:

Other dependencies will be specified in the Cargo.toml file and installed when we build the extension.

Python dependencies

These are both for final users (e.g., Click for the command-line interface) and for developers (e.g., pytest for testing). Preferably one should only install Python dependencies in virtual environments, never globally. I am partial towards Pipenv, but [Conda] could work as well.

Running pipenv install --dev with the following Pipfile:

[[source]]
name = "pypi"
url = "https://pypi.org/simple"
verify_ssl = true

[dev-packages]
maturin = ">=0.7.6"
pytest = ">=4.0"

[packages]
click = ">=7.0"

will create the virtual environment and install all needed packages. With pipenv shell, you can jump in the virtual environment. The maturin package is essential: this is the tool that will orchestrate building the extension and packaging for PyPI.

Power-up using Nix and direnv

Nix and direnv help keep your development environment tidy and reproducible, but special care needs to be taken when interacting with Python. In addition to Pipfile and Cargo.toml, we also have:

  • .envrc, used by direnv to set up the local environment:

I am using this version of the use_nix() function of direnv.

  • shell.nix, used by nix-shell and setting up the Rust toolchain and the Pipenv shell.

It is particularly convenient to use a Nix shell to not propagate per-developer dependencies. I use Emacs with language server protocol (LSP). The LSP for Python is provided by some extra packages that are thus my “private” development dependencies. These can be specified in the shell.nix file.

A closer look at shell.nix

This is a bit involved, so I’ll describe it bit by bit.

Pinning nixpkgs and local overlays

Here I pin a specific version of the nixpkgs repository and declare an override for the python-language-server package in an overlay.

with import (builtins.fetchGit {
  name = "nixos-19.03";
  url = "https://github.com/NixOS/nixpkgs-channels";
  ref = "nixos-19.03";
  # Commit hash for nixos-19.03 as of 2019-08-18
  # `git ls-remote https://github.com/nixos/nixpkgs-channels nixos-19.03`
  rev = "67135fbcc5d5d28390c127ef519b09a362ef2466";
}) {
  overlays = [(self: super:
    {
      python3 = super.python3.override {
        packageOverrides = py-self: py-super: {
          python-language-server = py-super.python-language-server.override {
            providers = [
              "rope"
              "pyflakes"
              "mccabe"
              "pycodestyle"
              "pydocstyle"
            ];
          };
        };
      };
    }
  )];
};
Rust from the Mozilla overlay

We need a nightly build of Rust, which we can get from the Mozilla Nix overlay.

let src = fetchFromGitHub {
  owner = "mozilla";
  repo = "nixpkgs-mozilla";
  # commit from: 2019-09-04
  rev = "b52a8b7de89b1fac49302cbaffd4caed4551515f";
  sha256 = "1np4fmcrg6kwlmairyacvhprqixrk7x9h89k813safnlgbgqwrqb";
};
in
with import "${src.out}/rust-overlay.nix" pkgs pkgs;
Declare the shell

We separate the packages into nativeBuildInputs and buildInputs. The latter are used to specify the per-developer private dependencies, which might include any libraries needed to compile extensions of Python dependencies (e.g. freetype for matplotlib). The shellHook sets up the Pipenv shell to cooperate with the Nix shell.

mkShell {
  name = "rustafarian";
  nativeBuildInputs = [
    # Note: to use stable, just replace nightly with stable
    latest.rustChannels.nightly.rust

    # Build-time Additional Dependencies
    #pkgconfig
  ];

  # Run-time Additional Dependencies
  buildInputs = [
    pipenv

    python3

    # System libraries needed for Python packages
    #freetype # Demanded by matplotlib

    # Development tools
    lldb
    python3.pkgs.black
    python3.pkgs.epc
    python3.pkgs.importmagic
    python3.pkgs.isort
    python3.pkgs.jedi
    python3.pkgs.mypy
    python3.pkgs.pyls-black
    python3.pkgs.pyls-isort
    python3.pkgs.pyls-mypy
    python3.pkgs.python-language-server
    travis
  ];

  # Set Environment Variables
  RUST_BACKTRACE = 1;
  shellHook = ''
    SOURCE_DATE_EPOCH=$(date +%s) # required for python wheels

    local venv=$(pipenv --bare --venv &>> /dev/null)

    if [[ -z $venv || ! -d $venv ]]; then
      pipenv install --dev &>> /dev/null
    fi

    export VIRTUAL_ENV=$(pipenv --venv)
    export PIPENV_ACTIVE=1
    export PYTHONPATH="$VIRTUAL_ENV/${python3.sitePackages}:$PYTHONPATH"
    export PATH="$VIRTUAL_ENV/bin:$PATH"
  '';
}

Workflow

OK, now we have a working development environment! To build the extension:

Yes, that’s it. maturin will compile the module and link it to a dynamic shared object (DSO) alongside the Python code. We can thus import it and use it in an interactive Python session spawned in our development environment. The Cargo.toml lists the Rust dependencies of our project and these will be automatically downloaded and compiled if not already available.

Running ldd on the DSO shows that there is no Rust runtime to link to:

Packaging and deployment

It is now time to distribute this little Rust+Python package by deploying it to PyPI. PEP 517 introduced the pyproject.toml file to specify build systems for Python packages. Ours will specify maturin as our build tool and thus look like:

[build-system]
requires = ["maturin"]
build-backend = "maturin"

As you can see, there is no mention of user dependencies (Click, in our case) in this file. Nor of console scripts to be installed alongside the module. This information is contained in the Cargo.toml:

[package.metadata.maturin]
requires-python = ">=3.6"
requires-dist = ["click>=7.0"]
scripts = {rustafarian = "rustafarian.cli:cli"}
classifier = ["Programming Language :: Python"]

Automated deploy to PyPI with Travis CI

maturin has a publish CLI switch to upload a package to PyPI. The process can be automatised: whenever a new tag is pushed to GitHub, Travis will run a CI job. If successful, it will deploy the artifact to PyPI:

The deployment script contains the following:

where the MATURIN_PASSWORD is the password to the PyPI account and is encrypted via the Travis CI web UI. We are building in a Docker container to ensure compliance with the manylinux1 policy.

References

(1) Scott, C. J. C.; Di Remigio, R.; Crawford, T. D.; Thom, A. J. W. Diagrammatic Coupled Cluster Monte Carlo. J. Phys. Chem. Lett. 2019, 10 (5), 925–935. https://doi.org/10.1021/acs.jpclett.9b00067.