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

NEWS.md

@@ -40,6 +40,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])
 
 New library features
 --------------------

base/cmd.jl

@@ -7,26 +7,30 @@ const UV_PROCESS_WINDOWS_VERBATIM_ARGUMENTS = UInt8(1 << 2)
 const UV_PROCESS_DETACHED = UInt8(1 << 3)
 const UV_PROCESS_WINDOWS_HIDE = UInt8(1 << 4)
 
+const RawCPUMask = Union{Nothing,Vector{Cchar}}
+
 struct Cmd <: AbstractCmd
     exec::Vector{String}
     ignorestatus::Bool
     flags::UInt32 # libuv process flags
     env::Union{Vector{String},Nothing}
     dir::String
+    cpumask::RawCPUMask
     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, cpumask = nothing) =
         new(cmd.exec, ignorestatus, flags, env,
-            dir === cmd.dir ? dir : cstr(dir))
+            dir === cmd.dir ? dir : cstr(dir), cpumask)
     function Cmd(cmd::Cmd; ignorestatus::Bool=cmd.ignorestatus, env=cmd.env, dir::AbstractString=cmd.dir,
+                 cpumask::RawCPUMask = cmd.cpumask,
                  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), cpumask)
     end
 end
 
@@ -287,6 +291,34 @@ 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 is supported on Unix and Windows but not in macOS.
+
+# Examples
+```julia
+julia> run(setcpuaffinity(`sh -c 'taskset -p \$\$'`, [1, 2, 5]));
+pid 2273's current affinity mask: 13
+
+julia> 0b010011
+0x13
+```
+"""
+function setcpuaffinity end
+setcpuaffinity(cmd::Cmd, ::Nothing) = Cmd(cmd; cpumask = nothing)
+function setcpuaffinity(cmd::Cmd, cpus::AbstractVector{<:Integer})
+    n = max(maximum(cpus), ccall(:uv_cpumask_size, Cint, ()))
+    cpumask = zeros(Cchar, n)
+    for i in cpus
+        cpumask[i] = true
+    end
+    return Cmd(cmd; cpumask)
+end
+
 (&)(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

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

base/process.jl

@@ -89,12 +89,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{Cchar}, Csize_t, Ptr{Cvoid}),
             file, exec, loop, handle,
             iohandles, length(iohandles),
             flags,
             env === nothing ? C_NULL : env,
             isempty(dir) ? C_NULL : dir,
+            cmd.cpumask === nothing ? C_NULL : cmd.cpumask,
+            cmd.cpumask === nothing ? 0 : length(cmd.cpumask),
             @cfunction(uv_return_spawn, Cvoid, (Ptr{Cvoid}, Int64, Int32)))
     end
     if err != 0

doc/src/base/base.md

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

@@ -290,7 +290,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;
@@ -300,8 +301,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,42 @@
+# This file is a part of Julia. License is MIT: https://julialang.org/license
+
+const pthread_t = Culong
+const uv_thread_t = pthread_t
+
+function uv_thread_getaffinity()
+    masksize = ccall(:uv_cpumask_size, Cint, ())
+    self = ccall(:uv_thread_self, uv_thread_t, ())
+    selfref = Ref(self)
+    cpumask = zeros(Cchar, masksize)
+    err = ccall(
+        :uv_thread_getaffinity,
+        Cint,
+        (Ptr{uv_thread_t}, Ptr{Cchar}, Cssize_t),
+        selfref,
+        cpumask,
+        masksize,
+    )
+    @assert err == 0
+    n = findlast(isone, cpumask)
+    resize!(cpumask, n)
+    return cpumask
+end
+
+function print_process_affinity()
+    isfirst = true
+    for (i, m) in enumerate(uv_thread_getaffinity())
+        if m != 0
+            if isfirst
+                isfirst = false
+            else
+                print(",")
+            end
+            print(i)
+        end
+    end
+    println()
+end
+
+if Base.Filesystem.samefile(PROGRAM_FILE, @__FILE__)
+    print_process_affinity()
+end

test/threads.jl

@@ -10,16 +10,24 @@ 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
+if Sys.islinux() || Sys.iswindows() || Sys.isfreebsd()
+    if Sys.CPU_THREADS > 1 && !running_under_rr()
+        @test run_with_affinity([2]) == "2"
+        @test run_with_affinity([1, 2]) == "1,2"
     end
 end