apply_ufunc parallelized atomic and logging

[gmoutso]

Hi all. I have two questions related to the following calculation. I have a function opt(vector, parameter) that is not a gufunc and takes time on each "atom". Let's say its signature is opt: [["dim1"], []] -> [["dim1"]], think of the first argument as the intial vector guess, the second is a parameter and the output is the solution vector to a scipy optimisation routine. xarray's apply_ufunc is extremely convenient to run this over a DataArray of parameter specs and possibly different initialisation vector values, mingled as appropriate with each parameter.

sol = xr.apply_ufunc(opt, vector_da, parameter_da, 
                                 input_core_dims=[["dim1"], []], output_core_dims=[["dim1"]], output_dtypes=[float],
                                 dask="parallelized", vectorize=True)

I obviously need the vectorize as opt can only work atomically. Then to make the slow optimisation faster overall, I chunk the inputs chunks={"dim2":1,...} before this line on all exterior dimensions, and finally I .compute. This is a sketch, I have wrapped all these in classes and it is working fine. The example is essentially: use vectorize and parallelize to the maximum, ie the minimum chunk size of exterior dims.

My first question is general. apply_ufunc is extremely convenient in taking care of reassembling exterior dimensions. I keep thinking though if there would be more mileage out of using dask tasks and not apply_ufunc. The reason for saying so is I feel I am abusing apply_ufunc however convenient it is.

My second question is specific. I have a lot of logging for each individual opt and the messages come from each chunk as they are made, and the logging stream is impossible to read. Are there ways to mitigate this problem? Ideally I could retrieve each individual opt log messages like the log messages from a celery task, somehow. Or be able to organize them and name them by worker id? Can they be written to the root logger delayed, flushed if an exception is raised? Any examples using any idea out there?

Thank you!

[gmoutso]

My calculation is essentially

import logging
import numpy as np
import xarray as xr
LOG = logging.getLogger()
logging.basicConfig()
import time
def opt(vector, parameter):
    LOG.info("starting %", parameter)
    slow_time = np.random.exponential(60) # don't run this!
    time.sleep(slow_time) # takes time
    solution = vector - vector.sum()
    LOG.info("found solution for %", parameter) # log messages won;t be aligned
    return solution
vector_da = xr.DataArray( np.random.normal(size=(5,100)), dims=["dim1", "dim2"]).chunk(chunks={"dim2":1})
parameter_da = xr.DataArray( np.linspace(0,1,100), dims=["dim2"]).chunk(chunks={"dim2":1})
sol = xr.apply_ufunc(opt, vector_da, parameter_da, 
                     input_core_dims=[["dim1"], []], output_core_dims=[["dim1"]],
                     output_dtypes=[float],
                     dask="parallelized", vectorize=True).compute()