[Profiler] Correct xprof trace call and update tests
When migrating away from using Tensorflow.profiler, a translation layer was missed. This PR copies over the logic from Tensorflow (https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/profiler/profiler_client.py#L25) into collect_profile.
PiperOrigin-RevId: 782139578
