Allow user to specify where they want to capture telemetry

[renan-souza]

Especially for the EA@Dask use case, user wants to capture telemetry in the middle of the task, not at the end. Currently, it's hacky to do so. Make this customizable.