def capture_route_impl()

in plugin/tensorboard_plugin_profile/profile_plugin.py [0:0]

• 68 lines of code
• 12 McCabe index (conditional complexity)

  def capture_route_impl(self, request):
    """Runs the client trace for capturing profiling information."""
    service_addr = request.args.get('service_addr')
    duration = int(request.args.get('duration', '1000'))
    is_tpu_name = request.args.get('is_tpu_name') == 'true'
    worker_list = request.args.get('worker_list')
    num_tracing_attempts = int(request.args.get('num_retry', '0')) + 1
    options = None
    try:
      options = profiler.ProfilerOptions(
          host_tracer_level=int(request.args.get('host_tracer_level', '2')),
          device_tracer_level=int(request.args.get('device_tracer_level', '1')),
          python_tracer_level=int(request.args.get('python_tracer_level', '0')),
      )
      # For preserving backwards compatibility with TensorFlow 2.3 and older.
      if 'delay_ms' in options._fields:
        options.delay_ms = int(request.args.get('delay', '0'))
    except AttributeError:
      logger.warning('ProfilerOptions are available after tensorflow 2.3')

    if is_tpu_name:
      try:
        tpu_cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver(
            service_addr)
        master_grpc_addr = tpu_cluster_resolver.get_master()
      except (ImportError, RuntimeError) as err:
        return respond({'error': err.message}, 'application/json', code=200)
      except (ValueError, TypeError):
        return respond(
            {'error': 'no TPUs with the specified names exist.'},
            'application/json',
            code=200,
        )
      if not worker_list:
        worker_list = get_worker_list(tpu_cluster_resolver)
      # TPU cluster resolver always returns port 8470. Replace it with 8466
      # on which profiler service is running.
      master_ip = master_grpc_addr.replace('grpc://', '').replace(':8470', '')
      service_addr = master_ip + ':8466'
      # Set the master TPU for streaming trace viewer.
      self.master_tpu_unsecure_channel = master_ip
    try:
      if options:
        profiler_client.trace(
            service_addr,
            self.logdir,
            duration,
            worker_list,
            num_tracing_attempts,
            options=options)
      else:
        profiler_client.trace(
            service_addr,
            self.logdir,
            duration,
            worker_list,
            num_tracing_attempts,
        )
      return respond(
          {'result': 'Capture profile successfully. Please refresh.'},
          'application/json',
      )
    except tf.errors.UnavailableError:
      return respond(
          {'error': 'empty trace result.'},
          'application/json',
          code=200,
      )
    except Exception as e:  # pylint: disable=broad-except
      return respond(
          {'error': str(e)},
          'application/json',
          code=200,
      )