Skip to content

Commit

Permalink
Make threads explicitly configurable. (#199)
Browse files Browse the repository at this point in the history
  • Loading branch information
maleadt authored Apr 24, 2023
1 parent 57f72ad commit 7df878d
Show file tree
Hide file tree
Showing 3 changed files with 33 additions and 7 deletions.
13 changes: 10 additions & 3 deletions src/sandbox.jl
Original file line number Diff line number Diff line change
Expand Up @@ -355,11 +355,18 @@ function setup_julia_sandbox(config::Configuration, args=``;

# restrict resource usage
if !isempty(config.cpus)
env["JULIA_CPU_THREADS"] = string(length(config.cpus)) # JuliaLang/julia#35787
env["OPENBLAS_NUM_THREADS"] = string(length(config.cpus)) # defaults to Sys.CPU_THREADS
env["JULIA_NUM_PRECOMPILE_TASKS"] = string(length(config.cpus)) # defaults to Sys.CPU_THREADS
# we might not always have CPU limiting capabilities (e.g. JuliaLang/julia#35787),
# so also instruct Julia to limit the number of threads it thinks are available
env["JULIA_CPU_THREADS"] = string(length(config.cpus))

# these should default to Sys.CPU_THREADS, but it can't hurt to be explicit
env["OPENBLAS_NUM_THREADS"] = string(length(config.cpus))
env["JULIA_NUM_PRECOMPILE_TASKS"] = string(length(config.cpus))
end

# configure threads
env["JULIA_NUM_THREADS"] = string(config.threads)

setup_generic_sandbox(config, `$cmd $(Cmd(config.julia_flags)) $args`;
env, mounts, kwargs...)
end
Expand Down
5 changes: 3 additions & 2 deletions src/types.jl
Original file line number Diff line number Diff line change
Expand Up @@ -69,8 +69,9 @@ Base.@kwdef struct Configuration
## a list of environment variables to set in the sandbox
env::Setting{Vector{String}} = Default(String[])
## a list of CPUs to restrict the Julia process to (or empty if unconstrained).
## if set, JULIA_CPU_THREADS will also be set to a number equaling the number of CPUs.
cpus::Setting{Vector{Int}} = Default(Int[])
## how many threads the Julia process should use
threads::Setting{Int} = Default(1)
## whether to spawn a virtual X server and expose that to the sandbox
xvfb::Setting{Bool} = Default(true)
## whether to run under record-replay (rr). traces will be uploaded to AWS S3, so you
Expand Down Expand Up @@ -118,7 +119,7 @@ function Base.show(io::IO, cfg::Configuration)
println(io)

println(io, " # Execution properties")
show_setting.(["env", "cpus", "xvfb", "rr", "precompile", "compiled", "process_limit"])
show_setting.(["env", "cpus", "threads", "xvfb", "rr", "precompile", "compiled", "process_limit"])
show_setting.(["log_limit", "memory_limit"], Base.format_bytes)
show_setting.(["time_limit", "compile_time_limit"], durationstring)

Expand Down
22 changes: 20 additions & 2 deletions test/runtests.jl
Original file line number Diff line number Diff line change
Expand Up @@ -170,15 +170,33 @@ end
@test results[1, :status] == :kill && results[1, :reason] == :log_limit
end

if "cpuset" in PkgEval.get_cgroup_controllers() && Sys.CPU_THREADS > 1
if "cpuset" in PkgEval.get_cgroup_controllers()
@testset "cpu" begin
let config = Configuration(config; cpus=[0])
cpu_threads = parse(Int, chomp(sprint(stdout->PkgEval.evaluate_script(config, "println(Sys.CPU_THREADS)"; stdout))))
cpu_threads = parse(Int, chomp(sprint(stdout->PkgEval.sandboxed_cmd(config, `/bin/sh -c "nproc"`; stdout))))
@test cpu_threads == 1
end
end
end

# not really a cgroup resource constraint, but it fits here nicely
@testset "threads" begin
let config = Configuration(config; cpus=[0])
cpu_threads = parse(Int, chomp(sprint(stdout->PkgEval.evaluate_script(config, "println(Sys.CPU_THREADS)"; stdout))))
@test cpu_threads == 1

julia_threads = parse(Int, chomp(sprint(stdout->PkgEval.evaluate_script(config, "println(Threads.nthreads())"; stdout))))
@test julia_threads == 1
end
let config = Configuration(config; cpus=[0], threads=8)
cpu_threads = parse(Int, chomp(sprint(stdout->PkgEval.evaluate_script(config, "println(Sys.CPU_THREADS)"; stdout))))
@test cpu_threads == 1

julia_threads = parse(Int, chomp(sprint(stdout->PkgEval.evaluate_script(config, "println(Threads.nthreads())"; stdout))))
@test julia_threads == 8
end
end

if "pids" in PkgEval.get_cgroup_controllers()
@testset "process" begin
let config = Configuration(config; process_limit=1)
Expand Down

0 comments on commit 7df878d

Please sign in to comment.