def data_impl()

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

• 54 lines of code
• 15 McCabe index (conditional complexity)

  def data_impl(self, request):
    """Retrieves and processes the tool data for a run and a host.

    Args:
      request: XMLHttpRequest

    Returns:
      A string that can be served to the frontend tool or None if tool,
        run or host is invalid.
    """
    run = request.args.get('run')
    tool = request.args.get('tag')
    host = request.args.get('host')
    tqx = request.args.get('tqx')
    run_dir = self._run_dir(run)
    # Profile plugin "run" is the last component of run dir.
    profile_run = os.path.basename(run_dir)

    if tool not in TOOLS and not use_xplane(tool):
      return None, None

    self.start_grpc_stub_if_necessary()
    if tool == 'trace_viewer@' and self.stub is not None:
      # Streaming trace viewer needs profiler_analysis service, which is only
      # supported in Cloud TPU. This code is unused when data was produced by
      # open-source TensorFlow. Only import the library when needed.
      # pylint: disable=g-import-not-at-top
      # pylint: disable=g-direct-tensorflow-import
      from tensorflow.core.profiler import profiler_analysis_pb2
      # pylint: enable=g-import-not-at-top
      # pylint: enable=g-direct-tensorflow-import
      grpc_request = profiler_analysis_pb2.ProfileSessionDataRequest()
      grpc_request.repository_root = os.path.dirname(run_dir)
      grpc_request.session_id = profile_run
      grpc_request.tool_name = 'trace_viewer'
      # Remove the trailing dot if present
      grpc_request.host_name = host.rstrip('.')

      grpc_request.parameters['resolution'] = request.args.get(
          'resolution', 8000)
      if request.args.get('start_time_ms') is not None:
        grpc_request.parameters['start_time_ms'] = request.args.get(
            'start_time_ms')
      if request.args.get('end_time_ms') is not None:
        grpc_request.parameters['end_time_ms'] = request.args.get('end_time_ms')
      grpc_response = self.stub.GetSessionToolData(grpc_request)
      return grpc_response.output, None

    asset_path = os.path.join(run_dir, make_filename(host, tool))

    data, content_encoding = None, None
    if use_xplane(tool):
      if host == ALL_HOSTS:
        file_pattern = make_filename('*', 'xplane')
        try:
          asset_paths = tf.io.gfile.glob(os.path.join(run_dir, file_pattern))
        except tf.errors.OpError as e:
          logger.warning('Cannot read asset directory: %s, OpError %s', run_dir,
                         e)
      else:
        asset_paths = [asset_path]

      try:
        data = convert.xspace_to_tool_data(asset_paths, tool, tqx)
      except AttributeError:
        logger.warning('XPlane converters are available after Tensorflow 2.4')
      return data, content_encoding

    raw_data = None
    try:
      with tf.io.gfile.GFile(asset_path, 'rb') as f:
        raw_data = f.read()
    except tf.errors.NotFoundError:
      logger.warning('Asset path %s not found', asset_path)
    except tf.errors.OpError as e:
      logger.warning("Couldn't read asset path: %s, OpError %s", asset_path, e)

    if raw_data is None:
      return None, None

    return get_data_content_encoding(raw_data, tool, tqx)