-
Notifications
You must be signed in to change notification settings - Fork 19
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Refactor column optimization based on typetracer updates #148
Refactor column optimization based on typetracer updates #148
Conversation
6c6e956
to
0e5c760
Compare
Codecov Report
📣 This organization is not using Codecov’s GitHub App Integration. We recommend you install it so Codecov can continue to function properly for your repositories. Learn more @@ Coverage Diff @@
## main #148 +/- ##
==========================================
- Coverage 94.94% 94.45% -0.49%
==========================================
Files 19 19
Lines 1919 2002 +83
==========================================
+ Hits 1822 1891 +69
- Misses 97 111 +14
📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more |
Good to go, right? |
I think it would be reasonable to bring this in as is and start addressing the warnings in follow up PRs @lgray would you be able to run a quick workflow of yours that you may be experimenting with? I think it would be useful to see how this PR affects things you've been working on.
There are actually a few more now that @lgray's #146 PR is in (which is good, just more information we can report upstream as holes in the typetracer) |
@douglasdavis Sure, I can give this a try later today. |
I don't have any introspection as to what part of my code is causing it though. Interestingly the results come back just fine? |
The dict I'm passing into
|
Also it would be great if |
We do know about the UnknownSclalar error (this one:
So the route we've gone down here for now is to catch the exceptions that are raised during the optimization and then throw that warning that you are seeing (the warning that says the optimization failed). We do still expect the computation to be correct, we just bailed on the column projection optimization (so I'm glad you see that the results come back just fine). Martin and I were just talking about a way to avoid the exceptions from awkward and still run the column projection optimization. This would be a workaround we can use until the upstream typetracer development is complete and the exceptions no longer get raised by awkward. I'm going to start experimenting with the workaround today.
a |
a necessary columns example: >>> import dask_awkward as dak
>>> ds = dak.from_parquet("s3://ddavistemp/hpq", storage_options={"anon": True})
>>> ds.muons.pt - ds.gen.pt
dask.awkward<subtract, npartitions=50>
>>> dak.necessary_columns(ds.muons.pt - ds.gen.pt)
{'read-parquet-7ec21ca9ab2b70cee1e98308ed59125d': ['muons.pt', 'gen.pt']} |
Ah it's already public. ok! Let me try again in a bit here. |
I guess my question is then: can I reliably use |
tl;dr: Since the optimization is based on typetracer arrays, right now the operations in the workflow need to be typetracer array aware. More explanation taking dask-histogram as an example: If we create a graph from the following code: import dask_histogram.boost as dhb
import dask_awkward as dak
ds = dak.from_parquet("s3://ddavistemp/hpq", storage_options={"anon": True})
pt = dak.flatten(ds.muons.pt)
h = dhb.Histogram(dhb.axis.Regular(25, 0, 250))
h.fill(pt)
h.compute() This compute will work, but only the dask_histogram optimizations will run because we called compute on the dask_histogram object; that is, this: function is called. Which does vanilla dask optimizations (not specific to dask-awkward). you can see dask-histogram's optimize is just calling functions from upstream dask However, if we call dak.necessary_columns(h) We see /Users/ddavis/software/repos/dask-awkward/src/dask_awkward/lib/optimize.py:150: UserWarning: Column projection optimization failed: <class 'AssertionError'>, bug in Awkward Array: attempt to convert TypeTracerArray into a concrete array
See if this has been reported at https://github.com/scikit-hep/awkward-1.0/issues
warnings.warn(f"Column projection optimization failed: {type(err)}, {err}") In this case the graph is getting analyzed by dask-awkward, which knows about
It should be possible to add to dask-histogram a feature to make it run dask-awkward optimizations if dask-awkward is installed. But we still need to solve the issue of exceptions getting raised due to external libraries not knowing about typetracers. This is do-able too, but we need to figure out the best way to mock these layers in the type-tracer-only graph. |
Ah, right this would be taken care of by the |
Ultimately, any code executed by Dask with a typetracer can lead to operations on unknown values. We're planning for There are two facets to this problem of probing functions with typetracer:
In both cases, the current idea that we're pursuing is to export an Note that |
@agoose77 There are a few places in coffea that need to call a |
Though I guess that can also be solved by checking if it's a typetracer and then manually calling |
def fill_user_hist(data):
p = np.sqrt(data.px**2 + data.py**2 + data.pz**2)
p_is_small = p < 10
# Maybe this will only be a single-arg function, who knows.
# Ensure bogus data are available for histogramming
p, p_is_small = ak.empty_if_typetracer(p, p_is_small)
p_hist = Hist.new.Bool().Reg(128, 0, 1000).Int64()
p_hist.fill(p_is_small, p)
dak.map_partitions(fill_user_hist, array) |
Ah - I see. This kind of stuff never goes anywhere near the common user in my perception of things, it should be there for them as a pressure relief valve or for power users, but the libraries should do a good portion of common functionality themselves without much complaining. |
Yes, this is an escape hatch. Really, it's more of a convenience L2 function (public API for library developer uses) who need to know about typetracers. For developers of libraries, you'd already use So in this case, |
It is, however, considerably more concise and tidy in library code! :D |
I agree agree with @lgray's comment on it being a pressure relief valve. We can benefit (on the dask-awkward side as library authors) from usng |
And second, we can include dask-awkward optimizations in dask-histogram for graphs where dask-awkward's PS it can be already be done explicitly by optimizing the dask-awkward graph before building more task graph with a dask-histogram step: ds = dask_awkward_dataset_factory()
pt_ratio = ds.pt_reco / ds.pt_true
(pt_ratio,) = dask.optimize(pt_ratio) # <-- modifies the read step in the graph to only grab ['pt_true', 'pt_reco']
h = dhb.Histogram(...)
h.fill(pt_ratio)
h.compute() |
Just to be clear, the empty_if_typetracer proposal is something users would need to add themselves to the functions that get passed in to map_partitions. Otherwise, all you can do is materialise/touch everything in the array, which really should work most of the time (unless the function actually cares about array lengths). Then of course there's the existing 'run one partition, take the meta' strategy as a third option. |
Yea this makes sense, we already mock output layers in our typetracer-only-graph (which is one of the main components of this PR!) and touch the data without writing anything to disk. def _mock_output_func(*args: Any, **kwargs: Any) -> Any:
import awkward as ak
for arg in args + tuple(kwargs.values()):
if isinstance(arg, ak.Array):
arg.layout._touch_data(recursive=True)
def _mock_output(layer: Any) -> Any:
assert len(layer.dsk) == 1
new_layer = copy.deepcopy(layer)
mp = new_layer.mapping.copy()
for k in iter(mp.keys()):
# this layer tuple is originally is of the form (ak.to_parquet, ...), but we
# replace the to_parquet or to_json, etc. with our mocked_output_func
# which just touches the data
mp[k] = (_mock_output_func,) + mp[k][1:]
new_layer.mapping = mp
return What @martindurant and I have in mind is to mock (really wrap) all other layers in the typetracer-only-graph and try to call whatever function |
Code in dask-histogram is like: def _blocked_dak(data: Any, *, histref: bh.Histogram | None = None) -> bh.Histogram:
import awkward as ak
thedata = data
if isinstance(thedata, ak.Array) and ak.backend(thedata) == "typetracer":
thedata.layout._touch_data(recursive=True)
thedata = data.layout.form.length_zero_array()
return clone(histref).fill(thedata)
def _blocked_dak_ma(*data: Any, histref: bh.Histogram | None = None) -> bh.Histogram:
import awkward as ak
thedata = list(data)
for idata, adatum in enumerate(thedata):
if isinstance(adatum, ak.Array) and ak.backend(adatum) == "typetracer":
adatum.layout._touch_data(recursive=True)
thedata[idata] = adatum.layout.form.length_zero_array()
return clone(histref).fill(*tuple(thedata)) |
So sorry in advance - this needs a few PRs to go in (and I have to wrap up some work in coffea0 before it's easy to repeat. However, I did find some interesting outcomes in the column optimization for things that aren't zips, but are more complex things we tend to find in analysis code. For instance with coffea nanoevents it's not so uncommon you want to match objects together get the nearest one, and maybe ask for the pt or some other single kinematic quantity. We have some convenience functions that wrap that up as follows, which works fine if daskified:
However when looking at the necessary columns. It's interesting that the column opt picks up the needed columns for Muons correctly (calculating delta_r only needs Muon.eta and Muon.phi), but the entire Jet object is requested from disk which is not necessary since we likewise only need the Jet's eta, phi, and pt in this case.
I'll follow up with instructions to repeat once everything is in and the recipe is easily shareable. |
84a6f6d
to
9665a62
Compare
Add behavior propagation in the iolayer
…ies with out-of-class state access
Push map_partitions in awkward mixins back to the user
for more information, see https://pre-commit.ci
Found a bug that will be run into regularly, here's a repro: import uproot
import dask_awkward as dak
x = uproot.dask({"https://github.com/CoffeaTeam/coffea/blob/master/tests/samples/nano_dy.root?raw=true": "Events"})
print(dak.necessary_columns(dak.broadcast_arrays(x.fixedGridRhoFastjetAll, x.Jet_pt))) prints:
As opposed to needing the data from both columns! @agoose77 @jpivarski since this could possibly include an issue with typetracer reports as well, but I can't find where it's failing. FYI - I can fix it by manually touching all the data that goes into |
Before this PR we recently created two different optimization strategies for column projection:
getitem
layers that use strings or lists of stringsWith this PR we now use the new
typetracer_with_report
API from upstreamawkward
to create a mutable report object, compute the task graph on purely the metadata (typetracer arrays), and return a list of necessary columns for each IO layer in the graph.There are still some outstanding upstream typetracer array issues that need to be ironed out before this optimization works for the current dask-awkward test suite.