Speed up execution of `@everywhere` in julia

README.md

As described in detail here, julia can take really excessive amounts of time to execute the first @everywhere statement on many processes — around 1 hour for thousands of processes — even if the actual code being executed everywhere is trivial. Basically, the Distributed functions need to be precompiled to make this happen quickly.

This gist provides a simple way to do so — at least on Slurm clusters (though the same principles should apply elsewhere). Just submit precompile.jl as a batch job (adjusting the SBATCH directives as needed), and it should create a sysimage that you can use to run future batch jobs. Check end of the log of the Slurm job to see exactly how to use the sysimage.

Note that both the original julia process and all processes created with addprocs should use the --sysimage=/path/to/sys_everywhere.so argument. Doing so reduces the time taken to execute the first @everywhere statement by a factor of ~20 for ~100 processes, and possibly more for a greater number of processes.

precompile.jl

#!/bin/bash -l
# -*- mode: julia -*-
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=2
#SBATCH --time=00:05:00
#SBATCH --partition=development
#SBATCH --output=%x_%j.log
#SBATCH --job-name=precompile_everywhere

# Edit the SBATCH directives above as needed, then submit this script as a
# batch job with something like `sbatch precompile.jl`

#=
# This block will execute in bash.  The lines with "${BASH_SOURCE[0]}" will run
# the remainder of this file as a julia script with the given arguments,
# but should return and proceed with the other commands, eventually generating
# `sys_everywhere.so`.

# Echo each command in bash
set -x

# Exit on any error
set -e

# Prepare the project; I don't know why these are all used, but they are
julia -e 'using Pkg; Pkg.add(["Distributed", "ClusterManagers", "Sockets", "Serialization", "Logging", "LinearAlgebra", "REPL"])'

# Run @everywhere 1+2 on 2 processes and trace
julia --trace-compile=precompile01.jl "${BASH_SOURCE[0]}" everywhere

# Combine the precompilation files into one
echo "using Distributed, ClusterManagers, Sockets, Serialization, Logging, LinearAlgebra, REPL" > precompile_everywhere.jl
cat precompile01.jl precompile02.jl >> precompile_everywhere.jl

# Create the sysimage
julia "${BASH_SOURCE[0]}" precompile

# Finally, just print out instructions
exec cat << EOF

The sysimage `sys_everywhere.so` has now been built.  To use it, run julia as

    julia --sysimage $PWD/sys_everywhere.so [other switches] [program file] [args...]

You should be able to move the .so file around on this system, but you probably
won't be able to move it to different systems.

EOF
=#

using PackageCompiler
using Distributed
using ClusterManagers
using Sockets
using Serialization
using Logging
using LinearAlgebra
using REPL

if "everywhere" ∈ ARGS

    addprocs_slurm(
        1;
        partition=get(ENV, "SBATCH_PARTITION", "development"),
        time="00:04:50",
        exeflags="--trace-compile=precompile02.jl"
    )
    while nprocs() < 2
        sleep(1)
    end
    @everywhere 1+2
    map(rmprocs, workers())

elseif "precompile" ∈ ARGS

    create_sysimage(
        [:Distributed, :ClusterManagers, :Sockets, :Serialization, :Logging, :LinearAlgebra, :REPL],
        sysimage_path="sys_everywhere.so",
        precompile_execution_file="precompile_everywhere.jl"
    )

else

    ArgumentError("""Accepted args are "everywhere" or "precompile"; got $ARGS""")

end