Speeding up Julia with precompilation

Julia has a just-in-time compilation model, where it compiles on first execution, and just what it needs. For example:

function stuff(x)
    return x*x
end

stuff(2) # stuff(Int) is compiled, not Float, Complex, ...
stuff(3) # use compiled stuff(Int)

Usually this is fast enough, but when you need maximum performance, you can ask Julia to precompile, and you can tell it ahead of time what you will need. The performance difference on clusters can be quite large (x15). One of the reason is that Julia needs to acces $HOME/.julia to read and write precompiled files, $HOME is never a fast filesystem on HPC (I’ve never encountered one). You can setup a faster $HOME or $DEPOT_PATH in Julia, but there’s a more portable way.

To get this speedup, we’ll combine 3 things, your code has tests that execute the common paths, you have a Singularity image setup, and we add PreCompile.jl.

Setting up PackageCompiler.jl #

using Pkg; Pkg.add("PackagerCompiler")

Create a script that will run your tests.

# src/precompile.jl
import MyPackage
include(joinpath(pkgdir(Sperwer), "test", "runtests.jl"))

Next, we run it to trace execution:

julia --project=. --trace-compile=dc_precompile.jl src/precompile.jl

This will generate ‘dc_precompile.jl’, a Julia file with all the statements to precompile. Create ‘src/setupimage.jl’

# src/setupimage.jl
using MyPackage
using Pkg
using PackageCompiler
Pkg.activate(".")
create_sysimage(sysimage_path="sys_img.so", include_transitive_dependencies=false, cpu_target="generic", precompile_statements_file="dc_precompile.jl")

and execute

julia --project=. src/setupimage.jl

This will create a ‘system image’, a binary with all the needed code pre-compiled, and using a generic architecture we’re sure this runs everywhere.

Caution If you do not set cpu_target, it will default to your architecture, if you then run on an older generation CPU, you will crash.

Now we can use this to speedup execution

time julia --project=. --sysimage=sys_img.so -e 'using MyPackage; dostuff()'

Without --sysimage=... Julia will need to precompile, or load precompiled images, then execute, and as execution runs, precompile more on the fly. Because you knew what would run (or could), most of this overhead has been saved.

If you combine this with a Singularity image, you can store the precompiled image in the singularity image, and gain even more improvements.

The below is an example definition script, based on a skeleton definition file:

BootStrap: docker
From: fedora:35

%files
    MyPackage.jl.zip /opt/MyPackage.jl.zip

%post
    ## Get the tools we'll need
    dnf install -y wget unzip python3 g++
    dnf groupinstall -y 'Development Tools'

    ## Setup Julia
    export JLMJV=1.7
    export JLV=$JLMJV.1
    export JULIA_TGZ=julia-$JLV-linux-x86_64.tar.gz
    mkdir -p /opt/julia && cd /opt/julia
    wget https://julialang-s3.julialang.org/bin/linux/x64/$JLMJV/$JULIA_TGZ && tar -xf $JULIA_TGZ && rm $JULIA_TGZ
    export PATH=/opt/julia/julia-$JLV/bin:$PATH
    export JULIA_DEPOT_PATH=/opt/juliadepot
    mkdir -p $JULIA_DEPOT_PATH

    ## Setup local package
    export PKGNAME="MyPackage"
    cd /opt && unzip $PKGNAME.jl.zip
    export LOCALPKG=/opt/$PKGNAME.jl
    cd $LOCALPKG
    julia --project=$LOCALPKG -e 'using Pkg; Pkg.update(); ENV["PYTHON"]=""; Pkg.build(); Pkg.instantiate()'
    julia --project=$LOCALPKG -e 'using MyPackage'
    echo "Setting up precompile"
    julia --project=$LOCALPKG --trace-compile=dc_precompile.jl src/precompile.jl
    julia --project=$LOCALPKG src/setupimage.jl
    rm -rf /opt/juliadepot/logs
    ln -s /dev/shm/ /opt/juliadepot/logs

    ## Cleanup
    dnf remove -y wget unzip

%environment
    export LC_ALL=C
    export LOCALPKG=/opt/MyPackage.jl
    export JLMJV=1.7
    export JLV=$JLMJV.1
    export PATH=/opt/julia/julia-$JLV/bin:$PATH
    export JULIA_DEPOT_PATH=/opt/juliadepot

%runscript
    echo "Container was created $NOW"
    echo "Arguments received: $*"
    echo pwd
    exec echo "$@"

%labels
    Author Ben Cardoen, bcardoen@sfu.ca
    Version v0.0.1

First, build the image

git clone git@github.com:you/MyPackage.jl.git
singularity build image.sif recipe.def

You’d execute with

singularity exec image.sif julia --project=/opt/MyPackage.jl --sysimage=/opt/MyPackage.jl/sys_img.so <yourcode.jl>

On a cluster I use, this is 18x faster, which matters for compute loads executed 1000s of times.

Credits and more resources #

The documentation at PackageCompiler.jl is a superb guide, same goes for the Singularity docs.