Doc: change suggested way of starting the profiler

[wrzadkow]

When using TensorBoard profiler on Google Compute Platform VM, I noticed that

import jax.profiler
jax.profiler.start_server(9999)	   

results in TensorBoard not "seeing" the profiling server, i.e. displaying "Failed to capture profile: empty trace result." message. This seems to be fixed after simply using:

import jax.profiler
server = jax.profiler.start_server(9999)	   

Background is that the server is alive as long as the object returned by jax.profiler.start_server(9999) is not destroyed and I think without assigning it to a variable that might happen immediately.

This PR aims at changing the suggested way of starting the server as well as making users aware of server's lifetime.

[froystig]

Maybe we should keep the resulting object alive within the profiler module? Then this requirement of caller behavior (and documentation) won't be needed, if I understand correctly.

[skye]

Thanks @wrzadkow for tracking this down! I'm not sure how it ever worked for me...

I'm gonna commit the documentation change. I think we should expose the underlying server object in case someone does wanna destroy it for some reason (to restart it?), and while we could also expose a separate function that keeps it alive, I think it's only slightly inconvenient to keep it alive oneself. Also, in the long run, we'd like to expose functionality to collect a trace without manually starting a server at all.

[skye]

The above is a lightly-held position though, if anyone feels strongly that start_server should keep the server alive let's do that! I'm gonna commit this for now but we can revert if someone wants to change the behavior.