def xspace_to_tool_data()

in plugin/tensorboard_plugin_profile/convert/raw_to_tool_data.py [0:0]

• 59 lines of code
• 19 McCabe index (conditional complexity)

def xspace_to_tool_data(
    xspace_paths,
    tool,
    tqx,
    xspace_wrapper_func=_pywrap_profiler.xspace_to_tools_data):
  """Converts XSpace to tool data string.

  Args:
    xspace_paths: A list of XSpace paths.
    tool: A string of tool name.
    tqx: Gviz output format.
    xspace_wrapper_func: A callable that takes a list of strings and a tool and
      returns the raw data.

  Returns:
    Returns a string of tool data.
  """
  assert tool[-1] == '^'
  tool = tool[:-1]  # xplane tool name ends with '^'
  data = None
  if tool == 'trace_viewer':
    # Trace viewer handles one host at a time.
    assert len(xspace_paths) == 1
    raw_data, success = xspace_wrapper_func(
        xspace_paths, tool)
    if success:
      data = process_raw_trace(raw_data)
  elif tool == 'overview_page':
    raw_data, success = xspace_wrapper_func(
        xspace_paths, tool)
    if success:
      data = overview_page_proto_to_gviz.to_json(raw_data)
  elif tool == 'input_pipeline_analyzer':
    raw_data, success = xspace_wrapper_func(
        xspace_paths, tool)
    if success:
      data = input_pipeline_proto_to_gviz.to_json(raw_data)
  elif tool == 'tensorflow_stats':
    raw_data, success = xspace_wrapper_func(
        xspace_paths, tool)
    if success:
      if tqx == 'out:csv':
        data = tf_stats_proto_to_gviz.to_csv(raw_data)
      else:
        data = tf_stats_proto_to_gviz.to_json(raw_data)
  elif tool == 'kernel_stats':
    raw_data, success = xspace_wrapper_func(
        xspace_paths, tool)
    if success:
      if tqx == 'out:csv;':
        data = kernel_stats_proto_to_gviz.to_csv(raw_data)
      else:
        data = kernel_stats_proto_to_gviz.to_json(raw_data)
  elif tool == 'memory_profile':
    # Memory profile handles one host at a time.
    assert len(xspace_paths) == 1
    raw_data, success = xspace_wrapper_func(
        xspace_paths, tool)
    if success:
      data = raw_data
  elif tool == 'pod_viewer':
    raw_data, success = xspace_wrapper_func(
        xspace_paths, tool)
    if success:
      data = raw_data
  elif tool == 'tf_data_bottleneck_analysis':
    raw_data, success = xspace_wrapper_func(
        xspace_paths, tool)
    if success:
      data = tf_data_stats_proto_to_gviz.to_json(raw_data)
  else:
    logger.warning('%s is not a known xplane tool', tool)
  return data