Add setcpuaffinity(cmd, cpus) for setting CPU affinity of subprocesses

NEWS.md

@@ -58,6 +58,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

base/cmd.jl

@@ -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)
@@ -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
@@ -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, ")")
@@ -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)

base/exports.jl

@@ -944,6 +944,7 @@ export
     run,
     setenv,
     addenv,
+    setcpuaffinity,
     success,
     withenv,
 

base/process.jl

@@ -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)
@@ -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

doc/src/base/base.md

@@ -314,6 +314,7 @@ Base.Cmd
 Base.setenv
 Base.addenv
 Base.withenv
+Base.setcpuaffinity
 Base.pipeline(::Any, ::Any, ::Any, ::Any...)
 Base.pipeline(::Base.AbstractCmd)
 Base.Libc.gethostname

src/jl_uv.c

@@ -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;
@@ -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;

test/print_process_affinity.jl

@@ -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

test/show.jl

@@ -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

test/threads.jl

@@ -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