Came across this when working on dask/distributed#4347 . The new annotations machinery is attached to high level graph Layers. However, in the optimization step, layers can be rewritten, and new layers can lose their annotations. This can especially cause unintended effects when these annotation affect scheduling/execution (as is the case when specifying workers, priority, retries, etc).

A MWE that reproduces the issue on dask master:

import dask
import dask.array as da

with dask.annotate(foo="bar"):
    a = da.ones((1,1))
    
print(a.__dask_graph__().layers[a.name].annotations)
aa, = dask.optimize(a)
print(aa.__dask_graph__().layers[aa.name].annotations)

produces

{'foo': 'bar'}
None

Perhaps @sjperkins or @rjzamora may have some insight into the best fix here (it looks like @sjperkins may have come across this before?)