Passive profiling/observability without block_until_ready()?

[samuela]

How does one observe the runtime of various operations, eg. the time that it takes to process a batch, within a training loop? I don't want to decimate my performance with block_until_ready(), but I would like to get some timing stats. I just want to be able to peek into the execution of asynchronous dispatch a bit.

In a similar vein, is there a (non-blocking) method on jnp.arrays like is_ready()? Having this would be useful to work around issues with external logging libraries like wandb/wandb#3690.

[jakevdp]

block_until_ready() is useful for benchmarking of short snippets. If you're interested in profiling of full programs, you should refer to the documentation here: https://jax.readthedocs.io/en/latest/profiling.html

Those docs could definitely be improved, but hopefully that will get you on the right track!

[samuela]

Thanks @jakevdp! Is there anything in between benchmarking and profiling? I'm really looking for "observability"... info on large, running training runs with minimal overhead.

[YouJiacheng]

I think samuela need a light weight method to log inexact performance stats. Profiler will incur more overhead than block_until_ready() IIUC.
For that purpose, I think allow asynchronous dispatch is okay. Pytorch on gpu is asynchronous as well.
But, I guess https://jax.readthedocs.io/en/latest/jax.experimental.host_callback.html might meet your need, i.e.

from jax.experimental.host_callback import id_tap
def do_sth_after_result_ready(_, _):
    print("time")

def train_step(state, data):
    # do training
    new_state = id_tap(do_sth_after_result_ready, None, result=new_state)
    return new_state

https://jax.readthedocs.io/en/latest/_autosummary/jax.experimental.host_callback.id_tap.html

The order of execution is by data dependency: after all the arguments and the value of result if present, are computed and before the returned value is used. At least one of the returned values of id_tap must be used in the rest of the computation, or else this operation has no effect.

[samuela]

Thanks @YouJiacheng ! I think this is more or less what I'm looking for