Add start_trace and stop_trace API in profiler

[lsy323]

Add start_trace and stop_trace APIs to programmatically start and stop profiling session. Before this PR, we can only start the profiling with a time duration, or within a context manager. This support allows better control over the profiling session.

The implementation is based on the profiler implementation in JAX.

Example usage:

server = xp.start_server(8001)
xp.start_trace(profilng_dir)
# Run some computation
...
xp.stop_trace()

[tengyifei]

That's amazing

[yaochengji]

Thanks, Siyuan!

[miladm]

@lsy323

• plz add this form of profiling to our documentations?
• I wonder how large of a profile file we can create. Do we know the user experience impact if the profiling duration is super long?

cc @mikegre-google

[miladm]

do we need a long-duration profile test?

[lsy323]

No, I don't think we need a long-duration profile test in torchxla. The goal of this PR is to provide better usability to users, the capability of profiler is out of scope of this PR (Should be in the underlying tsl library)

[lsy323]

@lsy323

• plz add this form of profiling to our documentations?
• I wonder how large of a profile file we can create. Do we know the user experience impact if the profiling duration is super long?

cc @mikegre-google

plz add this form of profiling to our documentations?

Let me add in a follow up PR I just realized we don't have a user guide on how to use the profiler.

I wonder how large of a profile file we can create. Do we know the user experience impact if the profiling duration is super long?

If the profiling time is super long, the traced content will be omitted in the tensorboard.