Dask computation confusion

[Jailbone]

Since I have not solved my problem (of understanding) till this day I want to come back to a question I posted here.

I will try to reformulate it to make my confusion clear.

import uproot
import numpy as np


# First create a root testfile:
with uproot.recreate("test.root") as file:
    file["test_tree"] = {"test_branch": np.random.random((100,10))}

# Open and load branch into numpy array
tree = uproot.open("./test.root:test_tree")
branch = tree["test_branch"].array(library = 'np')

#Calculate mean row-wise
mean_rows = np.mean(branch, axis = 1)

No problem till here, standard stuff. Now trying the same thing with the uproot functionality to load the branch into dask arrays:

tree = uproot.dask("./test.root:test_tree", library = 'np')
branch = tree["test_branch"]
mean_rows = np.mean(branch, axis = 1).compute()

#out: numpy.AxisError: Axis 1 is out of bounds for array of dimension 1

This does not work and I do not understand why.

Let's see what happens when trying to calculate the mean along axis = 0:

mean = np.mean(branch).compute()

#out: array([0.49147538, 0.49599721, 0.50983183, 0.53635986, 0.52194236,
       0.44889005, 0.53403688, 0.50172642, 0.49488129, 0.50677238])

Hm, that is also strange. That is the mean calculated column-wise (10 elements for the 10 columns). Even more confusing. So what am I failing to understand?

[Jailbone]

Anyone? :(

[jpivarski]

This is now Issue #1173. What was happening is that uproot.dask was reporting the wrong shape to Dask: it was saying that the array shape is (100,) but the array shape is (100, 10) because we didn't have any test examples with fixed-size arrays as the per-entry data type.

The result that you got for axis=0 is correct for a 2D array; Dask only appears to be checking and complaining about values of axis for which there is no answer. If Dask were more scrupulous, it would notice that the result it got disagrees with its expectations from the wrong metadata.

You can get everything to work by passing library="ak" and using ak.mean instead of np.mean.

I've opened an issue. I don't know this part of the codebase well, but there seems to be only one place where a Dask array (not dask_awkward) is created; that must be the place to pass the correct shape.

[Jailbone]

Hey @jpivarski, thanks for checking! I'm still wrapping my head around how dask can boost my analysis efficiency. Let's say I have a large root file (similar to the example below), and I'm wondering how dask could speed up both I/O and analysis on my local machine by leveraging multiple cores.

with uproot.recreate("test.root") as file:
    file["test_tree"] = {f"test_branch{i}": np.random.random((10000,1000)) for i in range(20)}

My understanding is that dask could parallelize the process in two ways:

1. Each dask worker loads a different chunk of the branch data into memory simultaneously.
2. Each worker computes the mean on its chunk in parallel.

I gave it a shot with the following code:

import dask
from dask.distributed import Client, LocalCluster

cluster = LocalCluster(n_workers = 4)
client = Client(cluster)
tree = uproot.dask(['./test.root:test_tree'], library = 'ak')
tasks = []
for branch in tree.fields:
    tasks.append(ak.mean(file[branch], axis = 1))
dask.compute(*tasks)

But it's much slower than the standard method:

tree = uproot.open('./test.root:test_tree', library = 'np')
[np.mean(file[branch], axis =1) for branch in tree.keys()]

I'm not sure what I'm missing. I've played around with parameters like step_size, but it usually makes things slower. Also, there's the num_workers parameter, but I've already set that in LocalCluster(n_workers = 4). What's the difference?

I'm a bit overwhelmed by all the options in the docs here. Does the handler option (uproot.source.chunk.Source class; None) play a role here?

Any help or clarification would be greatly appreciated! I'm excited about dask but feeling a bit stuck. Let me know if this isn't the right place to ask.

[jpivarski]

Most of those options are passed through to uproot.open, uproot.TTree.arrays, etc., and since this is reference documentation, it all has to be here as well.

The one thing that's important for tuning a parallel process (in general, not just Dask) is the task size. If a job is broken down into too many small tasks, then the new overhead of managing all of the tasks, not present in the single-threaded version, would overwhelm the parallel job and make it slower than the single-threaded version. If it is broken down into too few large tasks, those tasks might use too much memory to all be able to run at the same time, or they might underutilize the available CPUs—in the extreme case of "too few tasks" = "1 task," the job is not parallel: it's the same as the single-threaded version but also has some overhead for dealing with the parallel system.

In uproot.dask, the arguments that control data partition size (and therefore compute task size) are:

• step_size, which works the same way as its counterpart in uproot.iterate, but for parallel tasks, rather than sequential ones.
• steps_per_file, which divides up the partitions in a way that might not be equal-sized, but can be determined without opening all of the files to find out how many entries are in each.
• open_files is a boolean that determines whether you'll open all the files to find out how large each partition is before launching the job.

step_size and steps_per_file are mutually exclusive: you can only use one, and if open_files=False, you can only use steps_per_file.

You only have one file, which is not the usual case because Dask jobs usually run over more data than can reasonably fit in one file. But I guess that only means that open_files has no effect for you: it controls whether Uproot will open more than one file before sending the job to the workers.

The next-biggest issue for parallel processing is how the workers communicate. If the workers are on separate machines, they have to send intermediate results over the network, so if the intermediate results are large arrays, that will slow down the parallel job in a way that wouldn't slow down a single-machine job. This also applies to multiple processes on the same computer: they have to serialize their data because the operating system won't let them look at each other's RAM. This generally means opening a localhost socket, so it's all the software steps involved in sending data over the network, it just doesn't go through a hardware network (or even get to the network card).

I don't personally know whether distributed.LocalCluster is single-process, multiple-threads or if it's multiple processes, communicating through localhost sockets, but that's a question for the distributed documentation. (Note: Dask task graphs can be used with a variety of parallel processing engines; the package named distributed is just the default and not necessarily the best. WorkQueue/TaskVine and Ray are alternatives.) The Dask graph optimization tries to minimize how much data is sent between tasks, but it might have to send large arrays. (The URL that the distributed cluster prints out when it starts up, or that you can get as one of its attributes, can be viewed as a webpage with diagnostics for finding out things like how much data is passing between tasks.)

If there is a lot of data passing between tasks, then you might be tempted to opt for workers as threads in a single process, so that the data passing between tasks isn't serialized—but unfortunately, this would slow it down for another reason: multiple threads in a single Python interpreter compete with each other because Python has a Global Interpreter Lock. It's a catch-22: multi-processing has serialization issues (imposed by the operating system) and multi-threading has concurrency issues in Python (imposed by the language implementation). There's bleeding-edge development to fix this by running multiple Python interpreters in one process, but I don't think any libraries can take advantage of this feature yet. The best practice right now is to always do multiple processes and hope for low serialization costs.

Parallel processing, by its nature, can't be automatic in all situations, though embarrassingly parallel workflows can come close (up to constraints imposed by Amdahl's law; some of the overheads I talked about above are in the "serial part" that limits parallelization). Using a parallel processing system like Dask distributed (or WorkQueue/TaskVine, Ray, etc.) comes with some amount of tuning, at least looking at the scheduler diagnostics to see where the bottlenecks are. Your knobs to turn are:

• sizes of the partitions, and therefore the tasks (step_size or steps_per_file)
• choice of the end-state of the computation: if the job is dominated by serializing and sending a lot of data back to the main process, maybe you don't want to look at an intermediate result but instead want to finish the calculation by making summary statistics, such as a histogram (see dask-histogram) because a histogram has a lot fewer bytes than the data that fill a histogram
• parameters of the parallel processing system itself, which would be constructor arguments for distributed.LocalCluster. If your process has few GIL-bound Python steps (i.e. it's mostly compiled code that actually runs) but large datasets to transfer between tasks, then you might be in the unusual situation in which choosing multi-threaded, single-process is beneficial, and would have to choose a non-default setting to get that. The same is true of other knobs: sometimes, sockets can be replaced with pipes, or maybe the paralell processing system can use Unix interprocess shared memory (old school) or an Apache Arrow Plasma store (new school).

I hope this gives you some leads. As you can see from the complexity of the problem, I don't know what's slowing down your job in particular.

[Jailbone]

Wow @jpivarski thanks for this extensive answer, that helps tremendously! I am very grateful there are people like you taking their time to educate strangers on the internet! Comments like these should be much more visible/accessible to other beginners than now staying "hidden" here as I think this is valuable information! Cheers!

[jpivarski]

The sad thing is that although I feel compelled to help people when asked (good), I'm years behind in finishing the documentation (bad). I'm hoping to improve things by making all of these hidden discussions semantically searchable.