Skip to content

Commit

Permalink
Add setcpuaffinity(cmd, cpus) for setting CPU affinity of subproces…
Browse files Browse the repository at this point in the history
…ses (JuliaLang#42469)

Support running commands with cpumask.
  • Loading branch information
tkf authored and LilithHafner committed Feb 22, 2022
1 parent d811fcb commit 9f19050
Show file tree
Hide file tree
Showing 9 changed files with 122 additions and 13 deletions.
1 change: 1 addition & 0 deletions NEWS.md
Original file line number Diff line number Diff line change
Expand Up @@ -72,6 +72,7 @@ New library functions
---------------------

* `hardlink(src, dst)` can be used to create hard links. ([#41639])
* `setcpuaffinity(cmd, cpus)` can be used to set CPU affinity of sub-processes. ([#42469])
* `diskstat(path=pwd())` can be used to return statistics about the disk. ([#42248])

New library features
Expand Down
53 changes: 48 additions & 5 deletions base/cmd.jl
Original file line number Diff line number Diff line change
Expand Up @@ -13,28 +13,31 @@ struct Cmd <: AbstractCmd
flags::UInt32 # libuv process flags
env::Union{Vector{String},Nothing}
dir::String
cpus::Union{Nothing,Vector{UInt16}}
Cmd(exec::Vector{String}) =
new(exec, false, 0x00, nothing, "")
Cmd(cmd::Cmd, ignorestatus, flags, env, dir) =
new(exec, false, 0x00, nothing, "", nothing)
Cmd(cmd::Cmd, ignorestatus, flags, env, dir, cpus = nothing) =
new(cmd.exec, ignorestatus, flags, env,
dir === cmd.dir ? dir : cstr(dir))
dir === cmd.dir ? dir : cstr(dir), cpus)
function Cmd(cmd::Cmd; ignorestatus::Bool=cmd.ignorestatus, env=cmd.env, dir::AbstractString=cmd.dir,
cpus::Union{Nothing,Vector{UInt16}} = cmd.cpus,
detach::Bool = 0 != cmd.flags & UV_PROCESS_DETACHED,
windows_verbatim::Bool = 0 != cmd.flags & UV_PROCESS_WINDOWS_VERBATIM_ARGUMENTS,
windows_hide::Bool = 0 != cmd.flags & UV_PROCESS_WINDOWS_HIDE)
flags = detach * UV_PROCESS_DETACHED |
windows_verbatim * UV_PROCESS_WINDOWS_VERBATIM_ARGUMENTS |
windows_hide * UV_PROCESS_WINDOWS_HIDE
new(cmd.exec, ignorestatus, flags, byteenv(env),
dir === cmd.dir ? dir : cstr(dir))
dir === cmd.dir ? dir : cstr(dir), cpus)
end
end

has_nondefault_cmd_flags(c::Cmd) =
c.ignorestatus ||
c.flags != 0x00 ||
c.env !== nothing ||
c.dir !== ""
c.dir !== "" ||
c.cpus !== nothing

"""
Cmd(cmd::Cmd; ignorestatus, detach, windows_verbatim, windows_hide, env, dir)
Expand Down Expand Up @@ -114,6 +117,8 @@ function show(io::IO, cmd::Cmd)
print_env = cmd.env !== nothing
print_dir = !isempty(cmd.dir)
(print_env || print_dir) && print(io, "setenv(")
print_cpus = cmd.cpus !== nothing
print_cpus && print(io, "setcpuaffinity(")
print(io, '`')
join(io, map(cmd.exec) do arg
replace(sprint(context=io) do io
Expand All @@ -123,6 +128,11 @@ function show(io::IO, cmd::Cmd)
end, '`' => "\\`")
end, ' ')
print(io, '`')
if print_cpus
print(io, ", ")
show(io, collect(Int, cmd.cpus))
print(io, ")")
end
print_env && (print(io, ","); show(io, cmd.env))
print_dir && (print(io, "; dir="); show(io, cmd.dir))
(print_dir || print_env) && print(io, ")")
Expand Down Expand Up @@ -294,6 +304,39 @@ function addenv(cmd::Cmd, env::Vector{<:AbstractString}; inherit::Bool = true)
return addenv(cmd, Dict(k => v for (k, v) in eachsplit.(env, "=")); inherit)
end

"""
setcpuaffinity(original_command::Cmd, cpus) -> command::Cmd
Set the CPU affinity of the `command` by a list of CPU IDs (1-based) `cpus`. Passing
`cpus = nothing` means to unset the CPU affinity if the `original_command` has any.
This function is supported only in Linux and Windows. It is not supported in macOS because
libuv does not support affinity setting.
!!! compat "Julia 1.8"
This function requires at least Julia 1.8.
# Examples
In Linux, the `taskset` command line program can be used to see how `setcpuaffinity` works.
```julia
julia> run(setcpuaffinity(`sh -c 'taskset -p \$\$'`, [1, 2, 5]));
pid 2273's current affinity mask: 13
```
Note that the mask value `13` reflects that the first, second, and the fifth bits (counting
from the least significant position) are turned on:
```julia
julia> 0b010011
0x13
```
"""
function setcpuaffinity end
setcpuaffinity(cmd::Cmd, ::Nothing) = Cmd(cmd; cpus = nothing)
setcpuaffinity(cmd::Cmd, cpus) = Cmd(cmd; cpus = collect(UInt16, cpus))

(&)(left::AbstractCmd, right::AbstractCmd) = AndCmds(left, right)
redir_out(src::AbstractCmd, dest::AbstractCmd) = OrCmds(src, dest)
redir_err(src::AbstractCmd, dest::AbstractCmd) = ErrOrCmds(src, dest)
Expand Down
1 change: 1 addition & 0 deletions base/exports.jl
Original file line number Diff line number Diff line change
Expand Up @@ -944,6 +944,7 @@ export
run,
setenv,
addenv,
setcpuaffinity,
success,
withenv,

Expand Down
15 changes: 14 additions & 1 deletion base/process.jl
Original file line number Diff line number Diff line change
Expand Up @@ -71,9 +71,20 @@ end

const SpawnIOs = Vector{Any} # convenience name for readability

function as_cpumask(cpus::Vector{UInt16})
n = max(Int(maximum(cpus)), Int(ccall(:uv_cpumask_size, Cint, ())))
cpumask = zeros(Bool, n)
for i in cpus
cpumask[i] = true
end
return cpumask
end

# handle marshalling of `Cmd` arguments from Julia to C
@noinline function _spawn_primitive(file, cmd::Cmd, stdio::SpawnIOs)
loop = eventloop()
cpumask = cmd.cpus
cpumask === nothing || (cpumask = as_cpumask(cmd.cpus))
GC.@preserve stdio begin
iohandles = Tuple{Cint, UInt}[ # assuming little-endian layout
let h = rawhandle(io)
Expand All @@ -89,12 +100,14 @@ const SpawnIOs = Vector{Any} # convenience name for readability
err = ccall(:jl_spawn, Int32,
(Cstring, Ptr{Cstring}, Ptr{Cvoid}, Ptr{Cvoid},
Ptr{Tuple{Cint, UInt}}, Int,
UInt32, Ptr{Cstring}, Cstring, Ptr{Cvoid}),
UInt32, Ptr{Cstring}, Cstring, Ptr{Bool}, Csize_t, Ptr{Cvoid}),
file, exec, loop, handle,
iohandles, length(iohandles),
flags,
env === nothing ? C_NULL : env,
isempty(dir) ? C_NULL : dir,
cpumask === nothing ? C_NULL : cpumask,
cpumask === nothing ? 0 : length(cpumask),
@cfunction(uv_return_spawn, Cvoid, (Ptr{Cvoid}, Int64, Int32)))
end
if err != 0
Expand Down
1 change: 1 addition & 0 deletions doc/src/base/base.md
Original file line number Diff line number Diff line change
Expand Up @@ -316,6 +316,7 @@ Base.Cmd
Base.setenv
Base.addenv
Base.withenv
Base.setcpuaffinity
Base.pipeline(::Any, ::Any, ::Any, ::Any...)
Base.pipeline(::Base.AbstractCmd)
Base.Libc.gethostname
Expand Down
7 changes: 4 additions & 3 deletions src/jl_uv.c
Original file line number Diff line number Diff line change
Expand Up @@ -285,7 +285,8 @@ JL_DLLEXPORT void jl_uv_disassociate_julia_struct(uv_handle_t *handle)
JL_DLLEXPORT int jl_spawn(char *name, char **argv,
uv_loop_t *loop, uv_process_t *proc,
uv_stdio_container_t *stdio, int nstdio,
uint32_t flags, char **env, char *cwd, uv_exit_cb cb)
uint32_t flags, char **env, char *cwd, char* cpumask,
size_t cpumask_size, uv_exit_cb cb)
{
uv_process_options_t opts = {0};
opts.stdio = stdio;
Expand All @@ -295,8 +296,8 @@ JL_DLLEXPORT int jl_spawn(char *name, char **argv,
// unused fields:
//opts.uid = 0;
//opts.gid = 0;
//opts.cpumask = NULL;
//opts.cpumask_size = 0;
opts.cpumask = cpumask;
opts.cpumask_size = cpumask_size;
opts.cwd = cwd;
opts.args = argv;
opts.stdio_count = nstdio;
Expand Down
31 changes: 31 additions & 0 deletions test/print_process_affinity.jl
Original file line number Diff line number Diff line change
@@ -0,0 +1,31 @@
# This file is a part of Julia. License is MIT: https://julialang.org/license

const uv_thread_t = UInt # TODO: this is usually correct (or tolerated by the API), but not guaranteed

function uv_thread_getaffinity()
masksize = ccall(:uv_cpumask_size, Cint, ())
self = ccall(:uv_thread_self, uv_thread_t, ())
ref = Ref(self)
cpumask = zeros(Bool, masksize)
err = ccall(
:uv_thread_getaffinity,
Cint,
(Ref{uv_thread_t}, Ptr{Bool}, Cssize_t),
ref,
cpumask,
masksize,
)
Base.uv_error("getaffinity", err)
n = something(findlast(cpumask)) # we must have at least one active core
resize!(cpumask, n)
return cpumask
end

function print_process_affinity()
join(stdout, findall(uv_thread_getaffinity()), ",")
println()
end

if Base.Filesystem.samefile(PROGRAM_FILE, @__FILE__)
print_process_affinity()
end
8 changes: 8 additions & 0 deletions test/show.jl
Original file line number Diff line number Diff line change
Expand Up @@ -2355,3 +2355,11 @@ end
@test_repr "T[1;;; 2 3;;; 4]"
@test_repr "T[1;;; 2;;;; 3;;; 4]"
end

@testset "Cmd" begin
@test sprint(show, `true`) == "`true`"
@test sprint(show, setenv(`true`, "A" => "B")) == """setenv(`true`,["A=B"])"""
@test sprint(show, setcpuaffinity(`true`, [1, 2])) == "setcpuaffinity(`true`, [1, 2])"
@test sprint(show, setenv(setcpuaffinity(`true`, [1, 2]), "A" => "B")) ==
"""setenv(setcpuaffinity(`true`, [1, 2]),["A=B"])"""
end
18 changes: 14 additions & 4 deletions test/threads.jl
Original file line number Diff line number Diff line change
Expand Up @@ -31,16 +31,26 @@ let cmd = `$(Base.julia_cmd()) --depwarn=error --rr-detach --startup-file=no thr
end
end

function run_with_affinity(cpus)
script = joinpath(@__DIR__, "print_process_affinity.jl")
return readchomp(setcpuaffinity(`$(Base.julia_cmd()) $script`, cpus))
end

# issue #34415 - make sure external affinity settings work
if Sys.islinux()
const SYS_rrcall_check_presence = 1008
global running_under_rr() = 0 == ccall(:syscall, Int,
(Int, Int, Int, Int, Int, Int, Int),
SYS_rrcall_check_presence, 0, 0, 0, 0, 0, 0)
if Sys.CPU_THREADS > 1 && Sys.which("taskset") !== nothing && !running_under_rr()
run_with_affinity(spec) = readchomp(`taskset -c $spec $(Base.julia_cmd()) -e "run(\`taskset -p \$(getpid())\`)"`)
@test endswith(run_with_affinity("1"), "2")
@test endswith(run_with_affinity("0,1"), "3")
else
global running_under_rr() = false
end
# Note also that libuv does not support affinity in macOS and it is known to
# hang in FreeBSD. So, it's tested only in Linux and Windows:
if Sys.islinux() || Sys.iswindows()
if Sys.CPU_THREADS > 1 && !running_under_rr()
@test run_with_affinity([2]) == "2"
@test run_with_affinity([1, 2]) == "1,2"
end
end

Expand Down

0 comments on commit 9f19050

Please sign in to comment.