Execute Function for each Julia Thread

General Usage
, , ,
1

I have an external library which requires thread local resources to function. For every thread Julia creates, it must notify this library via an initialize_thread_local() function. Is there a way to run this code for each OS thread that Julia spawns? @everywhere only runs on each process and not thread. Is the best option Threads.@threads :static?

2

The best option is probably a TaskLocalValue (OncePerTask on Julia 1.12+) from TaskLocalValues.jl. (That assumes this is a C-style “threading context”
that you pass into the ccalls you’re making).

Julia tasks can migrate threads, so it’s not safe to simply assume that a task will remain on the thread it started on.

1 Like
3

In this case I do not need to track any thread local context. The external library sets up thread local contextes when I call initialize_thread_local()

4

Hmm, that might be tricky with Julia’s mutlithreading. Thankfully StableTasks.jl provides a @spawnat macro that you can use to specify that your task should stick to the indicated thread.

If you just need to run a single initializer function per thread, here’s what you could do:

using StableTasks: @spawnat
tasks = [@spawnat i initialize_thread_local() for i in 1:Threads.nthreads(:default)]
foreach(fetch, tasks) # wait for all tasks to finish

and then you could make sure that all your other library calls also use @spawnat - you don’t want the thread-local state getting corrupted by a Julia task suddenly switching threads.

But many C libraries like this have “reentrant” APIs so that is worth exploring if it’s an option.

1 Like
5

I think what you want is OncePerThread

3 Likes
6

OncePerThread seems to be the right tool. I tried to run it in a module’s __init__ function:

    Threads.@threads :static for _ = 1:Threads.nthreads()
        @info "Static initialize: $(Threads.threadid())"
    end

    thread_local_initializer = Base.OncePerThread() do
        @info "Perthread initialize: $(Threads.threadid())"
    end
    thread_local_initializer()

but it printed out

[ Info: Static initialize: 2
[ Info: Perthread initialize: 1

indicating that the second initializer was not run for thread #2. Is there a reason why it would not run on #2?

7

I think you just didn’t call it in Thread 2. The way I think OncePerThread is essentially by lazy caching. So calling it multiple times on the same thread does not invoke the wrapped function multiple times:

    thread_local_initializer = Base.OncePerThread() do
        @info "Perthread initialize: $(Threads.threadid())"
    end
    thread_local_initializer()
    thread_local_initializer()
    thread_local_initializer()

Should only print once.

If you wish to initialize all threads then you’ll have ti call this on all threads.

This sounds to me like this library would not like Julia migrating Tasks between threads, so you’ll probably have to use static scheduling? In that case you could just put a call to initialize_thread_local() at the top of the threaded loop’s body to ensure that the thread the Task runs in is initialized. That would probably be the safest option.

1 Like