Propagating context information to child-tasks and remote calls?

[oschulz]

@tkf and me have been discussing ways to propagate information about available workers (or resources in general) in distributed hierarchical computations:

https://discourse.julialang.org/t/propagation-of-available-assigned-worker-ids-in-hierarchical-computations

Adding resource arguments to every function call would be impractical, and using Cassette & friends to add add context by rewriting the whole computation would be very heavy-handed, since it might tough large code stacks (and may also be not be a complete solution when remote calls are involved?).

Could we add something like a context field to Task, in addition to the storage field - with the difference that context is automatically propagated to spawned tasks and remote calls? Adding the possibility to propagate a context through a computation in this fashion could also be useful for other use cases too, I think.

[vtjnash]

C.f. #34543

[tkf]

Thanks for opening the issue.

I'm actually against copying task_local_storage API. I think something similar to contextvars in Python (see PEP 567 for discussion) would be better:

const var1 = ContextVar(:var1, 42)  # with default
@contextvar var1 = 42 # possible sugar

const var2 = ContextVar{Lockable{Vector{Int}}}(:var2)  # without default, with type
@contextvar var2::Lockable{Vector{Int}} # possible sugar
# Lockable from https://github.com/JuliaLang/julia/pull/34400

const var3 = ContextVar(:var3)  # equivalent to ContextVar{Any}(:var3)
@contextvar var3 # possible sugar


function f()
    @sync @async @show var1[] # var1[] => 42
    var1[] = 0
    @sync @async @show var1[] # var1[] => 0

    setting(var1 => 1) do
        @show var1[] # var1[] => 1
    end
    @show var1[] # var1[] => 0
end

In contrast to task_local_storage(:var):

1. var[] can be inferred
2. var is forced to be namespaced (i.e., it has to exist in some module name space)
3. var can be backed up by an efficient concrete key type (e.g., UUID)
4. var allows small-size optimization when the value type can be inlined into the context storage

We may want to use something like HAMT (as in PEP 567) for the context storage.

---

A possible/reference implementation of ContextVar may be something like

struct ContextVar{T}
    name::Symbol  # only for human readability
    key::UUID     # (for example)
    default::T    # or `Union{Some{T},Nothing}`?
    has_default::Bool
end

getindex(var::ContextVar{T}) where {T} =
    if var.has_default
        get(task_local_storage(:CONTEXT), var.key, var.default)
    else
        task_local_storage(:CONTEXT)[var.key]
    end :: T

[tkf]

Here is a proof-of-concept and its usage: https://github.com/tkf/ContextVariables.jl/blob/ebc2b550e2177c4ebc059b6974f1ea36309d7f4f/test/runtests.jl#L14-L29

[oschulz]

Oh, nice! Will this work across remote-call boundaries?

Maybe use with_context instead of with_variables? Variable is such a generic term. :-)

[tkf]

Yeah, with_context sounds like a much better name. :)

Will this work across remote-call boundaries?

I think it'd work with local context variables (i.e., captured in closures) but not with global const context variables. But I think it's possible in principle. We need to tweak Distributed (or you need to propagate it manually) though.

Notes: Global const doesn't work because I'm using uuid4 (which is based on the global RNG) to generate the key. I can use uuid5 to make the key deterministic (i.e., hashing the package's UUID, module namespaces, and the variable name). We need a macro to make this easier.

[c42f]

Yes we need something better here, thanks again for opening this issue! I feel an alternative name for this issue could be "taking dynamic scope seriously" ;-) Though perhaps "context variable" is a more straightforward name for this.

If we somehow had efficient codegen for context variables, there's a lot of interesting language facilities which can be improved based on this. For example

• It could replace the special cases for the logger state in the existing logging system
• Error handling with an "algebraic effects" or "condition system" - see some discussion at julep: "chain of custody" error handling #7026
• Deterministic task local random number generators WIP: Tasklocal xoshiro #34852 (see also https://discourse.julialang.org/t/parallel-mersenne-twister/27567). Interesting and very difficult wrinkle: serving the RNG use case requires running code at task creation time, not just efficient access.

At #35690 (comment) I wondered about one possible way to specialize code based on context variables. As noted in the OP, specializing a whole call stack based on the presence/type of a context variable is just really heavy handed. But on the other hand, I feel like specializing a leaf function (or a few innermost inlined frames) on a context var could be very powerful, and possibly of acceptable cost.

Getting this idea to be sound and practical within inference seems quite tricky. I guess it would require a context calling convention to allow inference to reason about the innermost frames in a systematic way. If it actually panned out, I'm imagining the compiler could hoist context variable access out of innermost loops and across inlined function calls.

Kotlin seems to have some interesting APIs around context and coroutines (and they have taken structured concurrency seriously!) so we might learn something from examining that closer https://kotlinlang.org/docs/reference/coroutines/coroutine-context-and-dispatchers.html

[oschulz]

Getting this idea to be sound and practical within inference seems quite tricky.

I guess it would depend on how deep we want to take it. Context attached to tasks could be done with just a few small, lightweight changes. I think that would cover many use cases already.

Context at a deeper level, propagating through all function calls will come with a higher cost, I guess - but there are quite a few uses cases that would need that, of course, though I think less concerned with scheduling and resource allocation questions.

[tkf]

@c42f What do you think about the design of ContextVar{T} I posted above? I think it solves the problem of inferrability. There is still a dynamic code inside getindex(::ContextVar{T}) but LICM'ing it (or maybe rather get(::ContextVar{T}) :: Union{Some{T}, Nothing} which won't throw) sounds much easier with the current infrastructure than implementing the new effect specialization facility. I think ContextVar{T} is a neat design that completely eliminates the type-instability on the user side of the code.

Also, I don't think it's possible to eliminate some dynamic typing. Since we need to allow arbitrary value types, and we can't specialize for the whole set of the context variables currently set (i.e., equivalent to using NamedTuple as the context), the valtype would be Any. So I think type-stabilizing at use-site is a good strategy.

(I wonder if it can be generalized to the strategy for implementing dynamically scoped effect handler. That is to say, type information is statically "hoisted out" to the global/lexical scope and then there is a tiny bits of run-time information invoking those pre-defined entrypoints.) (Edit: actually, I guess this happens automatically anyway in effect handlers like chain-of-custody error handling.)

I also thought about supporting the deterministic parallel RNG. It probably is possible via implementing some kind of hook system via AbstractContextVar. But I fear there will be too much dynamic dispatches on the task creation. Though it looks like Kotlin's context system allows user-defined restoreThreadContext/updateThreadContext? I wonder how they make it efficient. Is it doable because Kotlin is statically typed or something?

[c42f]

What do you think about the design of ContextVar{T} I posted above? I

Well, I like it! It's simple and can be implemented right away with only some small changes to the runtime. My comments above about effect systems, etc, are all pretty much pie-in-the-sky speculation at this stage, though it would be nice to have a rough feeling for where things might go.

So anyway, if we provided ContextVar, that at least gives us strongly typed task local storage which would be reasonably efficient, convenient and solve some problems with context propagation which are currently unsolvable for users. It also doesn't seem too far-fetched to teach the Julia optimizer about the scoping semantics so that some form of LICM will work in the medium term. Specialization based on type of context vars could eventually be introduced as an optimization (if it's even feasible).

[oschulz]

I also thought about supporting the deterministic parallel RNG

That's actually something I already have in BAT.jl

https://github.com/bat/BAT.jl/blob/master/src/rngs/rng_init.jl

This provides reproducible random numbers for hierarchical parallel applications, via hierarchical partitioning of counter-based RNGs .

I'll make that more widely available via ParallelProcessingTools.jl. It's just in BAT.jl currently because I needed it fast and didn't have time to do a nice API at the time. I'll get on it.

[oschulz]

Context variables would be super-useful to propagate those RNG partitions!

[JeffBezanson]

I think the core problem is the cost of the key lookup. I believe it would be much too slow for something like RNG.

[oschulz]

I think the core problem is the cost of the key lookup. I believe it would be much too slow for something like RNG.

Well, I guess in many use cases, the task would get it's RNG once, and then pass it on to the functions it calls as an explicit parameter. And then pass partitions of that RNG to sub-tasks, when they spawn. I really wasn't trying to go for a full resource-injection solution with this issue, just something on the task-level for low/medium-frequency access.

The RNG-story is a bit different from the worker-availability story, though. Which RNG is used for which part of the computation shouldn't depend on number or workers/threads/tasks, for reproducibility, so that often may need to be handled explicitly.

[tkf]

My comments above about effect systems, etc, are all pretty much pie-in-the-sky speculation at this stage, though it would be nice to have a rough feeling for where things might go.

@c42f Yeah, I do like pie-in-the-sky speculations and am very interested in effect systems! Also, I think context variable API helps implementing POC effect handlers, even though it won't be as efficient as we'd want. Some effect handlers don't require too much optimizations when the overhead of the handler itself is large (e.g., a process+thread pool abstracting Distributed and Threads). I think it'd be nice to have a building block for playing with effect handler interfaces and assessing programming experience with it, before start thinking about optimizing the hell out of it (though it's also important to think about compiler-friendly interface at the same time).