PR 3119 introduced overtouching #3185
Comments
ikrommyd
added
the
bug (unverified)
ikrommyd
changed the title
|
a ping on this one @jpivarski @agoose77 ... this is rather badly noticeable. |
|
I'm trying to reproduce this using only Awkward, both to narrow the scope of search for the bug and to add a test to our test suite. Since you've narrowed in on #3119 as a cause, I think it must have something to do with RecordArrays in IndexedArrays, and But so far, I haven't been able to reproduce it. I'm starting from a dataset that's already grouped into "electron" and "muon" records, for convenience: dataset = ak.from_parquet("https://github.com/jpivarski-talks/2024-07-24-codas-hep-ml/raw/main/data/SMHiggsToZZTo4L.parquet")I make leptons = ak.concatenate([dataset.electron, dataset.muon], axis=1)then make a TypeTracerArray layout, report = ak.typetracer.typetracer_with_report(leptons.layout.form_with_key())and it's initially empty: print(report.data_touched, report.shape_touched)
# [] []Slicing through both the electrons and muons touches only the expected buffers: ak.Array(layout).ptprint(report.data_touched, report.shape_touched)
# ['node1', 'node3', 'node14'] ['node0', 'node1', 'node3', 'node14']The None of the following touch any more buffers than the above:
Printing arrays to the screen would touch more buffers (hence the precaution of assigning to Okay, ak.local_index(ak.Array(layout))and I can't think of why that would be. I'm looking into it now... |
|
Actually, the above only touches the shape of everything. Is that consistent with what you were seeing? That the shape of everything was being touched (as opposed to the data)? |
|
But ak.local_index(ak.Array(layout))touching the shape of many buffers also happens if I revert #3119, so it might not be your issue. |
|
awkward/src/awkward/_layout.py Line 387 in bf1c874 Perhaps l3 = ak.local_index(events.leptons)line? (If you stop computing just before this line, do you get no error? If you stop computing just after this line, do you get the error?) Also, the l3 = ak.local_index(events.leptons, axis=1)(instead of the default |
|
No, this does not fix the issue. :-/ I'll see if I can synthesize an awkward only reproducer. |
|
Got it a little more minimal for now. This is still overtouching import awkward as ak
import numpy as np
from coffea.nanoevents import NanoEventsFactory
import dask_awkward as dak
events = NanoEventsFactory.from_root(
{
"https://github.com/CoffeaTeam/coffea/raw/master/tests/samples/nano_dy.root": "Events"
},
).events()
events["leptons"] = ak.concatenate(
[events.Electron, events.Muon],
axis=1,
)
pair = ak.argcombinations(events.leptons, 2, fields=["l1", "l2"])
pair = pair[
ak.argmin(
(events.leptons[pair.l1] + events.leptons[pair.l2]).pt,
axis=1,
)
]
events = events[ak.num(pair) > 0]
l3 = events.leptons
print(dak.necessary_columns(l3.pt))
print(l3.pt.compute())Everything here seems improtant. If I try do remove anything, like the |
|
This is also happening. There's overtouching on import awkward as ak
import numpy as np
from coffea.nanoevents import NanoEventsFactory
import dask_awkward as dak
events = NanoEventsFactory.from_root(
{
"https://github.com/CoffeaTeam/coffea/raw/master/tests/samples/nano_dy.root": "Events"
},
).events()
events["leptons"] = ak.concatenate(
[events.Electron, events.Muon],
axis=1,
)
pair = ak.argcombinations(events.leptons, 2, fields=["l1", "l2"])
pair = pair[
ak.argmin(
(events.leptons[pair.l1] + events.leptons[pair.l2]).pt,
axis=1,
)
]
events = events[ak.num(pair) > 0]
print(dak.necessary_columns(events.Muon.pt)) |
|
Could dask-contrib/dask-awkward#526 be somehow part of the problem as well since |
|
Here is a minimal reproducer using awkward and one coffea behavior: import awkward as ak
from coffea.nanoevents.methods import nanoaod
def test_necessary_columns():
with open("nanoevents_form.json", "r") as fin:
form = ak.forms.from_json(fin.read())
events_layout, report = ak.typetracer.typetracer_with_report(form)
events = ak.Array(events_layout, behavior=nanoaod.behavior)
events["leptons"] = ak.concatenate(
[events.Electron, events.Muon],
axis=1,
)
pair = ak.argcombinations(events.leptons, 2, fields=["l1", "l2"])
# looks like all that is needed to trigger it is the sum with behaviors?
ptsum = (events.leptons[pair.l1] + events.leptons[pair.l2]).pt
if len(str(report.data_touched)) > 1000:
print("bad!")
return 1
else:
print("good!")
return 0
exit(test_necessary_columns())I have attached the json of the input form to this post. |
|
@pfackeldey, @maxgalli, and I looked into this. The over-touching problem comes from broadcasting through a union (the leptons) to apply a ufunc (the awkward/src/awkward/_broadcasting.py Line 946 in bf1c874 Before PR #3119, After PR #3119, PR #3119 was necessary because assignments like cat["another", "w"] = three.xthe size of an internal buffer doubled. (So imagine assigning 20 fields into such a record.) The exponential growth of that internal buffer had something to do with the interplay between broadcasting and lazy projection, and so I fixed it by removing lazy projection. @pfackeldey is looking into the problem, to see if there's some other way to avoid exponential memory growth of that buffer, while keeping this one project operation lazy so that it doesn't touch all fields. Both of these problems are tightly coupled to the fact that these are UnionArrays. They wouldn't happen with any other array type. |
|
Thanks for investigating! |
jpivarski
added
bug
* [Indexed*Array] adding '_trim()' method to fix exponential memory growth * [test] enhance test_3118 with a third array * [Indexed*Array] _trim: safer index treating * [test] add two tests, one with 2 and one with 3 arrays * style: pre-commit fixes * [Indexed*Array] _trim: properly handle typetracers (and all negative indices for IndexOptionArrays) * [test] revert commit 698b015 for 'tests/test_2713_from_buffers_allow_noncanonical.py' --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
|
Fixed by #3193. |
Version of Awkward Array
>=2.6.5Description and code to reproduce
Reproducer:
The bug is introduced by #3119
My
git bisectreport:The text was updated successfully, but these errors were encountered: