code: def testSingleWorkerSamplingMode(self, delay_ms=None):
      # TODO(jieying): remove after 01/10/2023.
      if not jtu.pjrt_c_api_version_at_least(0, 34):
          raise unittest.SkipTest(
              "Profiler is not supported on PJRT C API version < 0.34."
          )
  
      def on_worker(port, worker_start):
          jax.profiler.start_server(port)
          worker_start.set()
          x = jnp.ones((1000, 1000))
          while True:
              with jax.profiler.TraceAnnotation("atraceannotation"):
                  jnp.dot(x, x.T).block_until_ready()
                  if self.profile_done:
                      jax.profiler.stop_server()
                      break
  
      def on_profile(port, logdir, worker_start):
          worker_start.wait()
          options = tf_profiler.ProfilerOptions(
              host_tracer_level=2,
              python_tracer_level=2,
              device_tracer_level=1,
              delay_ms=delay_ms,
          )
  
          # Request for 1000 milliseconds of profile.
          duration_ms = 1000
          profiler_client.trace(
              f"localhost:{port}",
              logdir,
              duration_ms,
              "",
              1000,
              options,
          )
          self.profile_done = True
  
      logdir = absltest.get_default_test_tmpdir()
      # Remove any existing log files.
      shutil.rmtree(logdir, ignore_errors=True)
      port = portpicker.pick_unused_port()
      thread_profiler = threading.Thread(
          target=on_profile, args=(port, logdir, self.worker_start)
      )
      thread_worker = threading.Thread(
          target=on_worker, args=(port, self.worker_start)
      )
      thread_worker.start()
      thread_profiler.start()
      thread_profiler.join()
      thread_worker.join(120)
      self._check_xspace_pb_exist(logdir)
  
imports: from functools import partial
  import glob
  import os
  import shutil
  import sys
  import tempfile
  import threading
  import time
  import unittest
  from absl.testing import absltest
  import jax
  import jax.numpy
  import jax.profiler
  from jax import config
  import jax._src.test_util
label_desc: Use image node_id 13 for calling the ProfilerTest obj's underlying member method code with example usage: obj.testSingleWorkerSamplingMode(delay_ms) without return types