xarray: applying ufunc to month+day groups of dask based Dataset

I have a function that I'd like to apply along the time axis for each combination of month & day of a dataset, eg. to all time slices that have month=1 and day=1 over all years (and so on).

My current solution is to loop over the combinations, subset the dataset, apply the function and persist the result on the dask cluster:

from distributed import Client
import numpy as np
import xarray as xr

client = Client()


# fake function
def fun_along_time(x):
    return x*2

# testdata
x = xr.tutorial.load_dataset("air_temperature")

# time dimension needs to be single chunk
x = x.chunk({"time":-1, "lat":10, "lon":10})

timeindex = x.time.to_index()

for month in timeindex.month.unique():
    for day in timeindex.day.unique():
        xsel = x.sel(
            time=np.logical_and(x.time.dt.month == month, x.time.dt.day == day)
        )
        
        xres = xr.apply_ufunc(fun_along_time,
                              xsel,
                              input_core_dims=[["time"]],
                              output_core_dims=[["time"]],
                              dask="parallelized")

        xres = xres.persist()
        
        # some downstream tasks following

While this approach works, it does not seem very elegant and it calls persist repeatedly which is hard on the scheduler (and a bad practice?).

To make it work for the scheduler I've tried persisting the arrays together using client.persist, which worked OK but doesn't seem to be efficient either, and waiting for each array to finish persisting then computing the downstream tasks before moving on top the next iteration - this works too with obvious downsides.

So I'm wondering what would be a better way to achieve this, keeping in mind that this needs to scale well with dask.

I've looked into groupby and map but I don't think it works well for me as it generates a ton of dask tasks and messes with the chunksize which generates all kinds of issues itself.

I though about adding a new dimension / coordinate layer with month/day which I can then use for map/apply_ufunc but I cound't wrap my head around it.