code: def collect_profile(
      port: int,
      duration_in_ms: int,
      host: str,
      log_dir: str | None,
      host_tracer_level: int,
      device_tracer_level: int,
      python_tracer_level: int,
      no_perfetto_link: bool,
  ):
      options = profiler.ProfilerOptions(
          host_tracer_level=host_tracer_level,
          device_tracer_level=device_tracer_level,
          python_tracer_level=python_tracer_level,
      )
      log_dir_ = pathlib.Path(
          log_dir if log_dir is not None else tempfile.mkdtemp()
      )
      profiler_client.trace(
          f"{host}:{port}",
          str(log_dir_),
          duration_in_ms,
          options=options,
      )
      print(f"Dumped profiling information in: {log_dir_}")
      # The profiler dumps `xplane.pb` to the logging directory. To upload it to
      # the Perfetto trace viewer, we need to convert it to a `trace.json` file.
      # We do this by first finding the `xplane.pb` file, then passing it into
      # tensorflow_profile_plugin's `xplane` conversion function.
      curr_path = log_dir_.resolve()
      root_trace_folder = curr_path / "plugins" / "profile"
      trace_folders = [
          root_trace_folder / trace_folder
          for trace_folder in root_trace_folder.iterdir()
      ]
      latest_folder = max(trace_folders, key=os.path.getmtime)
      xplane = next(latest_folder.glob("*.xplane.pb"))
      result, _ = convert.xspace_to_tool_data(
          [xplane], "trace_viewer^", {}
      )
  
      with gzip.open(
          str(latest_folder / "remote.trace.json.gz"), "wb"
      ) as fp:
          fp.write(result.encode("utf-8"))
  
      if not no_perfetto_link:
          path = jax_profiler._write_perfetto_trace_file(str(log_dir_))
          jax_profiler._host_perfetto_trace_file(path)
  
file_path: jax/jax/collect_profile.py
imports: from __future__ import annotations
  import argparse
  import gzip
  import os
  import pathlib
  import tempfile
  from jax._src import profiler
label_desc: Use image node_id 1 for calling a global function with example usage: collect_profile(port, duration_in_ms, host, log_dir, host_tracer_level, device_tracer_level, python_tracer_level, no_perfetto_link) without return types