Updates for ThreadPinning v1

.github/workflows/ci.yml

@@ -27,9 +27,9 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - '1.6'
-          - '1.9'
+          - '1.10'
           - '1'
+          - 'pre'
           - 'nightly'
         os:
           - ubuntu-latest

Project.toml

@@ -1,6 +1,6 @@
 name = "ParallelProcessingTools"
 uuid = "8e8a01fc-6193-5ca1-a2f1-20776dae4199"
-version = "0.4.5"
+version = "0.4.6"
 
 [deps]
 ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197"
@@ -10,8 +10,8 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 Parameters = "d96e819e-fc66-5662-9728-84c9c7592b0a"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Sockets = "6462fe0b-24de-5631-8697-dd941f90decc"
-ThreadPinning = "811555cd-349b-4f26-b7bc-1f208b848042"
 
 [weakdeps]
 ThreadPinning = "811555cd-349b-4f26-b7bc-1f208b848042"
@@ -27,6 +27,7 @@ LinearAlgebra = "1"
 Logging = "1"
 Parameters = "0.12, 0.13"
 Pkg = "1"
+Random = "1"
 Sockets = "1"
-ThreadPinning = "0.7.22, 1"
-julia = "1.6"
+ThreadPinning = "1"
+julia = "1.10"

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

ext/ParallelProcessingToolsThreadPinningExt.jl

@@ -5,34 +5,199 @@
 import Distributed
 import ThreadPinning
 
-# ThreadPinning.jl does not support all operating systems, currently:
-const _threadpinning_supported = isdefined(ThreadPinning, :affinitymask2cpuids)
+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
+    @assert convert(Int, ThreadPinning.getcpuid()) isa Int
+    true
+catch err
+    false
+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.affinitymask2cpuids(ThreadPinning.get_affinity_mask())
-            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

src/ParallelProcessingTools.jl

@@ -9,6 +9,7 @@ using Distributed
 
 import LinearAlgebra
 import Pkg
+import Random # Required by ThreadPinning extention
 import Sockets
 
 import ClusterManagers
@@ -44,8 +45,4 @@ include("slurm.jl")
 include("htcondor.jl")
 include("deprecated.jl")
 
-@static if !isdefined(Base, :get_extension)
-    include("../ext/ParallelProcessingToolsThreadPinningExt.jl")
-end
-
 end # module

src/deprecated.jl

@@ -47,3 +47,12 @@
 
 
 @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

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 @@
 
 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()
@@ -88,7 +109,7 @@
 Example:
 
 ```julia
-task, n = runworkers(OnLocalhost(nprocs = 4))
+task, n = runworkers(OnLocalhost(n = 4))
 ```
 
 See also [`worker_resources()`](@ref).

test/runtests.jl

@@ -3,7 +3,10 @@
 import Test
 
 import ThreadPinning
-ThreadPinning.Prefs.set_os_warning(false)
+
+if isdefined(ThreadPinning, :Prefs)
+    ThreadPinning.Prefs.set_os_warning(false)
+end
 
 Test.@testset "Package ParallelProcessingTools" begin
     @info "Testing with $(Base.Threads.nthreads()) Julia threads."
@@ -21,6 +24,7 @@ Test.@testset "Package ParallelProcessingTools" begin
     include("test_procinit.jl")
     include("test_workerpool.jl")
     include("test_onworkers.jl")
+    include("test_ext_threadpinning.jl")
     include("test_deprecated.jl")
     include("test_docs.jl")
 end # testset

test/test_deprecated.jl

@@ -52,4 +52,6 @@ include("testtools.jl")
         end
     end
     rmprocs(pids)
+
+    @test_deprecated pinthreads_auto() isa Nothing
 end