Add AutoThreadPinning

oschulz · oschulz · commit b0deee8d8af5 · 2024-08-31T02:00:55.000+02:00
diff --git a/docs/src/index.md b/docs/src/index.md
@@ -14,7 +14,7 @@ The elastic cluster manager automatically adds new workers to an automatically c
 Since workers can appear and disappear dynamically, initializing them (loading packages, etc.) via the standard `Distributed.@everywhere` macro is problematic, as workers added afterwards won't be initialized. Parallel processing tools provides the macro [`@always_everywhere`](@ref) to run code globally on all current processes, but also store the code so it can be run again on future new worker processes. Workers that are part of a [`FlexWorkerPool`](@ref) will be updated automatically on `take!` and `onworker`. You can also use [`ensure_procinit`](@ref) to manually update all workers
 to all `@always_everywhere` used so far.
 
-The function [`pinthreads_auto`](@ref)  (used inside of `@always_everywhere`) provides a convenient way to perform some automatic thread pinning on all processes. Note that it needs to follow an [`import ThreadPinning`](https://github.com/carstenbauer/ThreadPinning.jl/), and that more complex use cases may require customized thread pinning for best performance.
+[`AutoThreadPinning`](@ref), in conjunction with the package [`ThreadPinning`](https://github.com/carstenbauer/ThreadPinning.jl/), provides a convenient way to perform automatic thread pinning (e.g. inside of `@always_everywhere`, to apply thead pinning to all processes). Note that `ThreadPinning.pinthreads(AutoThreadPinning())` works on a best-effort basis and that advanced applications may require customized thread pinning for best performance.
 
 Some batch system configurations can result in whole Julia processes, or even a whole batch job, being terminated if a process exceeds its memory limit. In such cases, you can try to gain a softer failure mode by setting a custom (slightly smaller) memory limit using [`memory_limit!`](@ref).
 
diff --git a/ext/ParallelProcessingToolsThreadPinningExt.jl b/ext/ParallelProcessingToolsThreadPinningExt.jl
@@ -5,6 +5,17 @@ import LinearAlgebra
 import Distributed
 import ThreadPinning
 
+using ThreadPinning: ispinned, getaffinity,
+    cpuids, node, core, ncores, socket, nsockets, nnuma, numa, ishyperthread, isefficiencycore,
+    getcpuids, pinthreads, openblas_getcpuids, openblas_pinthreads
+
+using ThreadPinning.Utility: affinitymask2cpuids
+
+using LinearAlgebra.BLAS: get_num_threads as blas_nthreads
+using LinearAlgebra.BLAS: set_num_threads as set_blas_nthreads
+using Random
+
+
 
 # ThreadPinning.jl does not fully support all operating systems, currently:
 const _threadpinning_supported = try
@@ -17,27 +28,176 @@ end
 
 @static if _threadpinning_supported
 
-function ParallelProcessingTools._pinthreads_auto_impl(::Val{true})
-    pid = Distributed.myid()
-    if Distributed.myid() == 1
-        @debug "On process $pid, leaving Julia threads unpinned"
-        let n_juliathreads = Threads.nthreads()
-            if n_juliathreads > 1
-                LinearAlgebra.BLAS.set_num_threads(n_juliathreads)
-            end
+
+function _get_core_map()
+    core_map = IdDict{Int, Int}()
+    for core_id in 1:ncores()
+        for cpu_id in core(core_id)
+            core_map[cpu_id] = core_id
         end
+    end
+    return core_map
+end
+
+
+function _maybe_set_blas_nthreads(avail_ncores::Integer = 0)
+    if haskey(ENV, "OPENBLAS_NUM_THREADS")
+        @info("OPENBLAS_NUM_THREADS set, not changing number of BLAS threads.")
+    elseif haskey(ENV, "OMP_NUM_THREADS")
+        @info("OMP_NUM_THREADS set, not changing number of BLAS threads.")
+    elseif avail_ncores > 0
+        @info "Setting number of BLAS threads to $avail_ncores (number of physical CPU cores in affinity mask)"
+        set_blas_nthreads(avail_ncores)
     else
-        @debug "On process $pid, pinning threads according to affinity mask"
-        let available_cpus = _ThreadPinning.Utility.affinitymask2cpuids(ThreadPinning.getaffinity())
-            ThreadPinning.pinthreads(:affinitymask)
-            LinearAlgebra.BLAS.set_num_threads(length(available_cpus))
+        n = Threads.nthreads()
+        if blas_nthreads() != n
+            @info "Setting number of BLAS threads to $n (same as number of Julia threads)"
+            set_blas_nthreads(n)
         end
     end
 end
 
 
+function _pin_threads_to(unsorted_sel_cpus_julia::AbstractVector{Int}, unsorted_sel_cpus_blas::AbstractVector{Int}, pin_blas::Bool)
+    sel_cpus_julia = sort(unsorted_sel_cpus_julia)
+    sel_cpus_blas = sort(unsorted_sel_cpus_blas)
+    if length(sel_cpus_julia) >= Threads.nthreads()
+        @info "Pinning Julia threads to CPU IDs $sel_cpus_julia"
+        pinthreads(sel_cpus_julia)
+
+        if pin_blas
+            if isempty(intersect(sel_cpus_julia, sel_cpus_blas))
+                if length(sel_cpus_blas) >= blas_nthreads()
+                    @info "Pinning OpenBLAS threads to CPU IDs $sel_cpus_blas"
+                    openblas_pinthreads(sel_cpus_blas)
+                    # Partial mitigation for ThreadPinning issue #105, ensure
+                    # Julia threads are pinned correcty, at least:
+                    pinthreads(sel_cpus_julia)
+                else
+                    @warn "Can't pin $(blas_nthreads()) BLAS threads, found only $(length(sel_cpus_blas)) suitable CPU IDs."
+                end
+            else
+                @warn "Won't pin BLAS threads on same CPU IDs as the Julia threads"
+            end
+        end    
+    else
+        @warn "Can't pin $(Threads.nthreads()) Julia threads, found only $(length(sel_cpus_julia)) suitable CPU IDs."
+    end
+end
+
+
+function _log_thread_pinning()
+    if ispinned()
+        julia_cpuids = sort(getcpuids())
+        @info "Julia threads pinned to CPU IDs $julia_cpuids"
+    else
+        @info "Julia threads not pinned."
+    end
+
+    try
+        blas_cpuids = sort(openblas_getcpuids())
+        @info "OpenBLAS threads pinned to CPU IDs $blas_cpuids"
+    catch err
+        if err isa ErrorException
+            if contains(err.msg, "could not load library")
+                @warn "Could not get OpenBLAS thread pinning information"
+            else
+                @info "OpenBLAS threads don't seem to be pinned."
+            end
+        else
+            rethrow()
+        end
+    end
+end
+
+
+function ThreadPinning.pinthreads(mode::ParallelProcessingTools.AutoThreadPinning)
+    pin_blas = mode.blas
+
+    if ispinned()
+        @info "Thread pinning already in effect, not changing it."
+    elseif any(iszero, getaffinity())
+        @info "Thread affinity mask available, using it."
+
+        core_map = _get_core_map()
+        # Order Index that has hyperthread CPU IDs last:
+        cpuid_order_idx = IdDict(cpu_id => order_idx for (order_idx, cpu_id) in pairs(node()))
+
+        avail_cpuids = sort(affinitymask2cpuids(getaffinity()), by = i -> cpuid_order_idx[i])
+        avail_ncores = length(unique([core_map[i] for i in avail_cpuids]))
+
+        _maybe_set_blas_nthreads(avail_ncores)
+
+        sel_cpus_julia = avail_cpuids[begin:begin+Threads.nthreads()-1]
+        sel_cpus_blas = avail_cpuids[end-blas_nthreads()+1:end]
+        _pin_threads_to(sel_cpus_julia, sel_cpus_blas, pin_blas)
+    elseif Threads.nthreads() < 2
+        @info "Julia running single-threaded with no thread affinity mask, not pinning threads."
+    elseif mode.random
+        _maybe_set_blas_nthreads()
+        n_julia_rest::Int = Threads.nthreads()
+        n_blas_rest::Int = blas_nthreads()
+        sel_cpus_julia = Int[]
+        sel_cpus_blas = Int[]
+        for sid in shuffle(1:nsockets())
+            socket_cpus = socket(sid)
+            for nid in shuffle(1:nnuma())
+                cids = filter(!isefficiencycore, intersect(socket_cpus, numa(nid)))
+                if !isempty(cids)
+                    mainthreads_here = filter(!ishyperthread, cids)
+                    perm = shuffle(eachindex(mainthreads_here))
+                    mainthreads_here = mainthreads_here[perm]
+
+                    hyperthreads_here = filter(ishyperthread, cids)
+                    if axes(hyperthreads_here) == axes(mainthreads_here)
+                        hyperthreads_here = hyperthreads_here[perm]
+                    else
+                        hyperthreads_here = shuffle(hyperthreads_here)
+                    end
+
+                    julia_threadsource = mainthreads_here
+                    blas_threadsource = !isempty(hyperthreads_here) ? hyperthreads_here : mainthreads_here
+                    if n_julia_rest > 0
+                        n_julia_here = min(n_julia_rest, length(julia_threadsource))
+                        append!(sel_cpus_julia, julia_threadsource[1:n_julia_here])
+                        n_julia_rest -= n_julia_here
+                    end
+                    if n_blas_rest > 0
+                        n_blas_here = min(n_blas_rest, length(blas_threadsource))
+                        append!(sel_cpus_blas, blas_threadsource[1:n_blas_here])
+                        n_blas_rest -= n_blas_here
+                    end
+                    !(n_julia_rest > 0) && !(n_blas_rest > 0) && break
+                end
+            end
+            !(n_julia_rest > 0) && !(n_blas_rest > 0) && break
+        end
+
+        _pin_threads_to(sel_cpus_julia, sel_cpus_blas, pin_blas)
+    else
+        _maybe_set_blas_nthreads()
+        @info "No thread affinity set and random pinning not enabled, not pinning threads."
+    end
+
+    _log_thread_pinning()
+
+    return nothing
+end
+
+function ParallelProcessingTools._pinthreads_auto_impl(::Val{true})
+    pinthreads(ParallelProcessingTools.AutoThreadPinning())
+end
+
 ParallelProcessingTools._getcpuids_impl(::Val{true}) = ThreadPinning.getcpuids()
 
+
+else #! _threadpinning_supported
+
+
+ThreadPinning.pinthreads(ParallelProcessingTools.AutoThreadPinning) = nothing
+
+
 end # if _threadpinning_supported
+    
 
 end # module ChangesOfVariablesInverseFunctionsExt
diff --git a/src/deprecated.jl b/src/deprecated.jl
@@ -47,3 +47,12 @@ export @mp_async
 
 
 @deprecate isdefined_local(x) isassigned(x)
+
+
+@noinline function pinthreads_auto()
+    Base.depwarn("`pinthreads_auto()` is deprecated, use `ThreadPinning.pinthreads(AutoThreadPinning())` instead.",:pinthreads_auto)
+    _pinthreads_auto_impl(Val(true))
+end
+export pinthreads_auto
+
+_pinthreads_auto_impl(::Val) = nothing
diff --git a/src/runworkers.jl b/src/runworkers.jl
@@ -1,24 +1,41 @@
 # This file is a part of ParallelProcessingTools.jl, licensed under the MIT License (MIT).
 
 
-
 """
-    pinthreads_auto()
+    struct AutoThreadPinning
 
-!!! note
-    Only has an effect if
-    [`ThreadPinning`](https://github.com/carstenbauer/ThreadPinning.jl/) is
-    loaded, and only on operating systems supported by `ThreadPinning`.
-"""
-function pinthreads_auto end
-export pinthreads_auto
+ParallelProcessingTools default thread pinning mode.
+
+Constructor:
+
+```julia
+AutoThreadPinning(; random::Bool = false, pin_blas::Bool = false)
+```
+
+Arguments:
 
-pinthreads_auto() = _pinthreads_auto_impl(Val(true))
-_pinthreads_auto_impl(::Val) = nothing
+- `random`: Use system topology based random thread pinning if no thread
+          affinity mask is set (e.g. via SLURM, `taskset`).
 
+- `blas`: Try to pin BLAS threads. Not fully functional due to bugs in
+  BLAS thread pinning (see ThreadPinning issue
+  [#105](https://github.com/carstenbauer/ThreadPinning.jl/issues/105)).
+
+Use with `ThreadPinning.pinthreads`:
+
+```julia
+using ParallelProcessingTools, ThreadPinning
+pinthreads(AutoThreadPinning())
+```
+"""
+@kwdef struct AutoThreadPinning
+    random::Bool = false
+    blas::Bool = false
+end
+export AutoThreadPinning
 
 _getcpuids() = _getcpuids_impl(Val(true))
-_getcpuids_impl(::Val) = nothing
+_getcpuids_impl(::Val) = missing
 
 
 """
@@ -28,6 +45,10 @@ Get the distributed Julia worker process resources currently available.
 
 This may take some time as some code needs to be loaded on all processes.
 Automatically runs `ensure_procinit()` before querying worker resources.
+
+Note: CPU ID information will only be available if
+[`ThreadPinning`](https://github.com/carstenbauer/ThreadPinning.jl) is
+loaded.
 """
 function worker_resources()
     ensure_procinit()
diff --git a/test/test_deprecated.jl b/test/test_deprecated.jl
@@ -52,4 +52,6 @@ include("testtools.jl")
         end
     end
     rmprocs(pids)
+
+    @test_deprecated pinthreads_auto() isa Nothing
 end
diff --git a/test/test_ext_threadpinning.jl b/test/test_ext_threadpinning.jl
@@ -6,8 +6,5 @@ using Test
 import ThreadPinning
 
 @testset "ext_threadpinning" begin
-    ParallelProcessingToolsThreadPinningExt = Base.get_extension(ParallelProcessingTools, :ParallelProcessingToolsThreadPinningExt)
-    @test ParallelProcessingToolsThreadPinningExt isa Module
-
-    @test pinthreads_auto() isa Nothing
+    @test ThreadPinning.pinthreads(AutoThreadPinning()) isa Nothing
 end
