def tool_proto_to_tool_data(tool_proto, tool, tqx):
"""Converts the serialized tool proto to tool data string.
Args:
tool_proto: A serialized XSpace proto string.
tool: A string of tool name.
tqx: Gviz output format.
Returns:
Returns a string of tool data.
"""
data = ''
if tool == 'trace_viewer':
data = process_raw_trace(tool_proto)
elif tool == 'tensorflow_stats':
if tqx == 'out:csv;':
data = tf_stats_proto_to_gviz.to_csv(tool_proto)
else:
data = tf_stats_proto_to_gviz.to_json(tool_proto)
elif tool == 'overview_page@':
data = overview_page_proto_to_gviz.to_json(tool_proto)
elif tool == 'input_pipeline_analyzer@':
data = input_pipeline_proto_to_gviz.to_json(tool_proto)
elif tool == 'kernel_stats':
if tqx == 'out:csv;':
data = kernel_stats_proto_to_gviz.to_csv(tool_proto)
else:
data = kernel_stats_proto_to_gviz.to_json(tool_proto)
else:
logger.warning('%s is not a known tool', tool)
return data
def xspace_to_tool_data(xspace, tool, tqx):
"""Converts the serialized XSpace proto to tool data string.
Args:
xspace: A serialized XSpace proto string.
tool: A string of tool name.
tqx: Gviz output format.
Returns:
Returns a string of tool data.
"""
assert tool[-1] == '^'
tool = tool[:-1] # xplane tool name ends with '^'
data = ''
if tool == 'trace_viewer':
data = process_raw_trace(
_pywrap_profiler.xspace_to_trace_events(xspace))
elif tool == 'overview_page':
data = overview_page_proto_to_gviz.to_json(
_pywrap_profiler.xspace_to_overview_page(xspace))
elif tool == 'input_pipeline_analyzer':
data = input_pipeline_proto_to_gviz.to_json(
_pywrap_profiler.xspace_to_input_pipeline(xspace))
elif tool == 'tensorflow_stats':
data = _pywrap_profiler.xspace_to_tf_stats(xspace)
if tqx == 'out:csv;':
data = tf_stats_proto_to_gviz.to_csv(data)
else:
data = tf_stats_proto_to_gviz.to_json(data)
elif tool == 'kernel_stats':
data = _pywrap_profiler.xspace_to_kernel_stats(xspace)
if tqx == 'out:csv;':
data = kernel_stats_proto_to_gviz.to_csv(data)
else:
data = kernel_stats_proto_to_gviz.to_json(data)
elif tool == 'memory_profile':
data = _pywrap_profiler.xspace_to_memory_profile(xspace)
else:
logger.warning('%s is not a known xplane tool', tool)
return data
def xspace_to_tool_data(xspace_paths, tool, tqx):
"""Converts XSpace to tool data string.
Args:
xspace_paths: A list of XSpace paths.
tool: A string of tool name.
tqx: Gviz output format.
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 = _pywrap_profiler.xspace_to_tools_data(
xspace_paths, tool)
if success:
data = process_raw_trace(raw_data)
elif tool == 'overview_page':
raw_data, success = _pywrap_profiler.xspace_to_tools_data(
xspace_paths, tool)
if success:
data = overview_page_proto_to_gviz.to_json(raw_data)
elif tool == 'input_pipeline_analyzer':
raw_data, success = _pywrap_profiler.xspace_to_tools_data(
xspace_paths, tool)
if success:
data = input_pipeline_proto_to_gviz.to_json(raw_data)
elif tool == 'tensorflow_stats':
raw_data, success = _pywrap_profiler.xspace_to_tools_data(
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 = _pywrap_profiler.xspace_to_tools_data(
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 = _pywrap_profiler.xspace_to_tools_data(
xspace_paths, tool)
if success:
data = raw_data
elif tool == 'pod_viewer':
raw_data, success = _pywrap_profiler.xspace_to_tools_data(
xspace_paths, tool)
if success:
data = raw_data
else:
logger.warning('%s is not a known xplane tool', tool)
return data
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:
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')
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))
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
data, content_encoding = None, None
if tool == 'trace_viewer':
data = process_raw_trace(raw_data)
elif tool == 'tensorflow_stats':
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 == 'overview_page@':
data = overview_page_proto_to_gviz.to_json(raw_data)
elif tool == 'input_pipeline_analyzer@':
data = input_pipeline_proto_to_gviz.to_json(raw_data)
elif tool == 'kernel_stats':
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':
logger.info('Accessing memory_profile data')
data = raw_data
content_encoding = 'gzip'
elif tool in _RAW_DATA_TOOLS:
data = raw_data
if tool[-1] == '#':
content_encoding = 'gzip'
return data, content_encoding