Nested distributed loops in Julia with RemoteChannels

The Serial Approach: Simple but Slow

First, let’s establish a baseline with a serial implementation. Here’s how we might solve this problem sequentially:

data_sets = randn(10)          # 10 datasets
params_grid = [-1, 0, 1]       # Parameter grid for search

function crunch_heavy(data, param)
    sleep(1)                   # Simulate heavy computation
    return data * param
end

function summarize(runs)
    findmax(runs)              # Returns (max_value, index)
end

results = map(data_sets) do data
    all_runs = map(params_grid) do param
        crunch_heavy(data, param)
    end
    summary = summarize(all_runs)
    Dict(
        :all_runs => all_runs,  # Vector of results for each parameter
        :summary => summary     # Best result and its index
    )
end

This code produces a Vector of Dicts, one per dataset. Each dictionary contains :all_runs, a vector of results corresponding to params_grid, and :summary, which holds the “best fit” (in this case, just the maximum) value and the index in params_grid that gave rise to it. It’s straightforward and works fine for small problems. However, if data_sets grows to hundreds or thousands of elements, or if params_grid expands, the serial approach becomes painfully slow. Since crunch_heavy is CPU-bound and each computation is independent, this is a prime candidate for parallelization using Julia’s multiprocessing capabilities. Let’s try that next.

Naive Parallelization: A Deadlock Trap

Given we have multiple cores—say, 8 workers—our first instinct might be to parallelize both loops using pmap, Julia’s distrituted version of map. We will spawn 8 additional Julia processes on the current machine and give them tasks to run. Here’s the attempt:

using Distributed
addprocs(8)                    # Add 8 worker processes

data_sets = randn(10)
@everywhere params_grid = [-1, 0, 1]  # Define on all processes

@everywhere function crunch_heavy(data, param)
    sleep(1)
    return data * param
end

@everywhere function summarize(runs)
    findmax(runs)
end

results = pmap(data_sets) do data
    all_runs = pmap(params_grid) do param
        crunch_heavy(data, param)
    end
    summary = summarize(all_runs)
    Dict(
        :all_runs => all_runs,
        :summary => summary
    )
end

This looks promising, but if you run it, you’ll notice something alarming: the program hangs, CPU usage drops to 0%, and no progress is made. What’s happening? The outer pmap distributes the 10 datasets across the 8 workers, occupying all available processes. When each worker reaches the inner pmap, it tries to distribute tasks for the parameters, but there are no free workers left to handle them. The workers wait indefinitely for resources that never become available—a classic deadlock. Nesting pmap calls like this exhausts the worker pool, rendering the approach unusable.

Alternative Approaches: Partial Solutions

Could we parallelize just one loop instead? Let’s consider the options.

With a parallel outer loop and serial inner loop, each worker processes one dataset, handling all parameters serially. This works without deadlocking, but it’s inefficient if we have fewer datasets than workers. For example, with 2 datasets and 4 parameters across 4 workers, only 2 workers are active, each taking 4 seconds (due to sleep(1)), totaling 4 seconds. Ideally, all 4 workers could process the 8 tasks (2 datasets × 4 parameters) in parallel, finishing in 2 seconds. Resources sit idle, which isn’t optimal.

Processing one dataset at a time, parallelizing the parameter grid search, is equally bad when there are more parameters than datasets. For 2 datasets and 4 parameters with 4 workers, each dataset’s 4 parameters are computed in 1 second, totaling 2 seconds. This is better, but if params_grid has fewer elements than workers (e.g., 2 parameters), only 2 workers are used per dataset, leaving others idle. Neither approach fully utilizes resources consistently.

What if we split the workers and allocate, say, 2 of them to the outer loop and 6 to the inner loop? The outer workers delegate tasks to the inner pool, but here’s the catch: after delegating, the outer workers sleep until the inner tasks complete, doing no useful work. For 2 datasets and 4 parameters, the 2 outer workers could finish their datasets in parallel, but they’d wait 4 seconds each while the 6 inner workers handle only 2 tasks at a time. This underutilizes the outer workers, wasting potential compute power. We need a way to keep all workers busy with the heavy lifting—crunch_heavy—without idling.

To illustrate this inefficiency, consider a small example where we keep track of which worker IDs executed each crunch:

results = pmap(CachingPool([2, 3, 4, 5]), data_sets) do data
    all_runs = pmap(CachingPool([7, 8, 9]), params_grid) do param
        crunch_heavy(data, param)
    end
    summary = summarize(all_runs)
    Dict(
        :all_runs => all_runs,
        :summary => summary
    )
end

If we check results, we’ll see that only 3 workers (7, 8, and 9) ever run the heavy lifting. During that time, the outer workers are idling (they do come back alive when it’s time to summarize, but that step is very fast compared to the heavy crunching). Moreover, the summarize step requires all all_runs to be present, so if the computational load is heterogeneous across parameters, and for one particular parameter it takes much longer than the others, even the inner-loop workers will sit around idle, waiting until the last worker finishes its task. This happens because the loop structure enforces a synchronization point: the outer loop can only advance once all the iterations of the inner loop are done. So, splitting workers is possibly the worst solution for this problem, because workers from both pools are potentially left idling.

The Efficient Solution: Linearize with RemoteChannel

The core issue is that nesting parallel loops ties up workers unnecessarily. Instead, we should linearize the nested structure—submit all crunch_heavy calls as independent tasks to available workers—while preserving the results’ structure. A RemoteChannel offers a powerful way to achieve this in a multiprocessing context.

In Julia, an AbstractChannel is a thread-safe, first-in, first-out (FIFO) queue for communication between tasks. Its interface consists of four methods: put! (add an item), take! (remove and return an item), fetch (view without removing), and wait (block until ready). Channels shine in multithreaded programs, but we’re using multiprocessing. Enter RemoteChannel: a wrapper around an AbstractChannel that enables inter-process communication. The actual channel resides in one process (typically the master), and workers interact with it via references. When a worker calls put! on a RemoteChannel, the data is sent to the master process, which updates the underlying channel atomically. Because only the master process is in charge of modifying the actual channel, no concurrency issues arise.

To maintain our nested structure, we’ll use a custom DictChannel (adapted from Julia’s manual) that stores results in a dictionary, keyed by dataset index, with values accumulating as lists. The key modification is in put!: instead of overwriting values, it appends to a list, allowing multiple results (one per parameter) per dataset. Here’s the implementation:

@everywhere begin

    struct DictChannel{T} <: AbstractChannel{T}
        d::Dict                    # Stores key-value pairs
        cond_take::Threads.Condition  # Signals data availability
        DictChannel{T}() where {T} = new(Dict(), Threads.Condition())
        DictChannel() = DictChannel{Any}()
    end

    function Base.put!(D::DictChannel, k, v)
        @lock D.cond_take begin
            D.d[k] = push!(get(D.d, k, []), v)  # Append v to list at key k
            notify(D.cond_take)
        end
        return D
    end

    function Base.take!(D::DictChannel, k)
        @lock D.cond_take begin
            v = fetch(D, k)
            delete!(D.d, k)         # Remove key after fetching
            return v
        end
    end

    Base.isready(D::DictChannel) = @lock D.cond_take !isempty(D.d)
    Base.isready(D::DictChannel, k) = @lock D.cond_take haskey(D.d, k)

    function Base.fetch(D::DictChannel, k)
        @lock D.cond_take begin
            wait(D, k)
            return D.d[k]
        end
    end

    function Base.wait(D::DictChannel, k)
        @lock D.cond_take begin
            while !isready(D, k)
                wait(D.cond_take)
            end
        end
    end

end

Now let’s put it all together. To ensure that all workers are kept busy, te’ll linearize the loops using Iterators.product, which generates all the combination of data and params, while using a DictChannel within a RemoteChannel to preserve the nested structure of the original loop.

# Create RemoteChannel with DictChannel
jobs = RemoteChannel(() -> DictChannel())

# Linearize nested loops into task iterator
data_params_iter = zip(
    Iterators.repeated(jobs),
    Iterators.product(pairs(data_sets), pairs(params_grid))
)

# Distribute heavy computations
pmap(data_params_iter) do (jobs, ((data_idx, data), (param_idx, param)))
    run = crunch_heavy(data, param)
    put!(jobs, data_idx, (param_idx, run))
end

# Process results in parallel
results_iter = zip(Iterators.repeated(jobs), eachindex(data))
results = pmap(results_iter) do (jobs, data_idx)
    all_runs_tuples = take!(jobs, data_idx)  # List of (param_idx, run)
    # Sort to align with params_grid (optional)
    # sort!(all_runs_tuples, by = x -> x[1])
    all_runs = [run for (param_idx, run) in all_runs_tuples]
    summary = summarize(all_runs)
    Dict(
        :all_runs => all_runs,
        :summary => summary
    )
end

Conclusion

Nested parallel loops in Julia can lead to deadlocks or inefficient resource usage if not handled carefully. The naive approach of nesting pmap calls exhausts workers, while partial parallelization leaves resources idle. By linearizing tasks and using a RemoteChannel with a custom DictChannel, we achieve efficient parallelization: all workers tackle the heavy computations as they become available, and the nested structure is reconstructed seamlessly. This pattern scales well and can be extended to deeper nesting levels—perhaps an exercise for the curious reader! For production use, consider timing comparisons (serial vs. parallel) or error handling, but I hope this post exposed some common pitfalls of distributed computing.

References