xprof-nightly 2.22.3a20251208__cp311-none-manylinux2014_x86_64.whl

xprof/__init__.py

@@ -0,0 +1,22 @@
+# Copyright 2025 The XProf Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Entry point for the TensorBoard plugin package for XProf.
+
+Public submodules:
+  profile_plugin: The TensorBoard plugin integration.
+  profile_plugin_loader: TensorBoard's entrypoint for the plugin.
+  server: Standalone server entrypoint.
+  version: The version of the plugin.
+"""

xprof/convert/csv_writer.py

@@ -0,0 +1,87 @@
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Converts data between different formats."""
+
+import csv
+import io
+import json
+from typing import List, Optional
+
+
+def json_to_csv(
+    json_string: str,
+    columns_order: Optional[List[str]] = None,
+    separator: str = ",",
+) -> str:
+  """Converts a JSON string to a CSV string.
+
+  Args:
+    json_string: The JSON string to convert.
+    columns_order: Optional. Specifies the order of columns in the output table.
+      Specify a list of all column IDs in the order in which you want the table
+      created. Note that you must list all column IDs in this parameter, if you
+      use it.
+    separator: Optional. The separator to use between the values.
+
+  Returns:
+    A CSV string representing the data.
+    Example result:
+     'a','b','c'
+     1,'z',2
+     3,'w',''
+
+  Raises:
+    ValueError: The json_string is not a valid json or is not a single object.
+  """
+  try:
+    data = json.loads(json_string)
+  except json.JSONDecodeError as e:
+    raise ValueError(f"Invalid JSON string: {e}") from e
+
+  if not isinstance(data, dict):
+    raise ValueError("JSON data must be a single object")
+
+  headers = list(data.keys())
+  if columns_order is None:
+    columns_order = headers
+
+  if not all(col in headers for col in columns_order):
+    raise ValueError("columns_order must be a list of all column IDs")
+
+  if not columns_order:
+    return ""
+
+  csv_buffer = io.StringIO(newline="")
+  writer = csv.writer(csv_buffer, delimiter=separator, lineterminator="\n")
+
+  # Write header
+  writer.writerow([col for col in columns_order])
+
+  # Write row
+  cells_list = []
+  for col in columns_order:
+    value = ""
+    if col in data:
+      value = data.get(col)
+      if value is not None:
+        if isinstance(value, bool):
+          value = "true" if value else "false"
+        else:
+          value = str(value)
+    cells_list.append(value)
+  writer.writerow(cells_list)
+
+  return csv_buffer.getvalue()

xprof/convert/raw_to_tool_data.py

@@ -0,0 +1,232 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""For conversion of raw files to tool data.
+
+Usage:
+    data = xspace_to_tool_data(xplane, tool, params)
+    data = tool_proto_to_tool_data(tool_proto, tool, params)
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import logging
+
+from xprof.convert import csv_writer
+from xprof.convert import trace_events_json
+from xprof.protobuf import trace_events_old_pb2
+from xprof.convert import _pywrap_profiler_plugin
+
+
+logger = logging.getLogger('tensorboard')
+
+
+def process_raw_trace(raw_trace):
+  """Processes raw trace data and returns the UI data."""
+  trace = trace_events_old_pb2.Trace()
+  trace.ParseFromString(raw_trace)
+  return ''.join(trace_events_json.TraceEventsJsonStream(trace))
+
+
+def xspace_to_tools_data_from_byte_string(xspace_byte_list, filenames, tool,
+                                          params):
+  """Helper function for getting an XSpace tool from a bytes string.
+
+  Args:
+    xspace_byte_list: A list of byte strings read from a XSpace proto file.
+    filenames: Names of the read files.
+    tool: A string of tool name.
+    params: user input parameters.
+
+  Returns:
+    Returns a string of tool data.
+  """
+# pylint:disable=dangerous-default-value
+  def xspace_wrapper_func(xspace_arg, tool_arg, params={}):
+    return _pywrap_profiler_plugin.xspace_to_tools_data_from_byte_string(
+        xspace_arg, filenames, tool_arg, params)
+# pylint:enable=dangerous-default-value
+
+  return xspace_to_tool_data(xspace_byte_list, tool, params,
+                             xspace_wrapper_func)
+
+
+def xspace_to_tool_names(xspace_paths):
+  """Converts XSpace to all the available tool names.
+
+  Args:
+    xspace_paths: A list of XSpace paths.
+
+  Returns:
+    Returns a list of tool names.
+  """
+  raw_data, success = _pywrap_profiler_plugin.xspace_to_tools_data(
+      xspace_paths, 'tool_names')
+  if success:
+    return [tool for tool in raw_data.decode().split(',')]
+  return []
+
+
+def xspace_to_tool_data(
+    xspace_paths,
+    tool,
+    params,
+    xspace_wrapper_func=_pywrap_profiler_plugin.xspace_to_tools_data):
+  """Converts XSpace to tool data string.
+
+  Args:
+    xspace_paths: A list of XSpace paths.
+    tool: A string of tool name.
+    params: user input parameters.
+    xspace_wrapper_func: A callable that takes a list of strings and a tool and
+      returns the raw data. If failed, raw data contains the error message.
+
+  Returns:
+    Returns a string of tool data and the content type for the response.
+  """
+  if (tool[-1] == '^'):
+    old_tool = tool
+    tool = tool[:-1]  # Remove the trailing '^'
+    logger.warning(
+        'Received old tool format: %s; mapped to new format: %s', old_tool, tool
+    )
+  data = None
+  content_type = 'application/json'
+  # tqx: gViz output format
+  tqx = params.get('tqx', '')
+  options = {}
+  options['use_saved_result'] = params.get('use_saved_result', True)
+  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, options)
+    if success:
+      data = process_raw_trace(raw_data)
+  elif tool == 'trace_viewer@':
+    options = params.get('trace_viewer_options', {})
+    options['use_saved_result'] = params.get('use_saved_result', True)
+    options['hosts'] = params.get('hosts', [])
+    raw_data, success = xspace_wrapper_func(xspace_paths, tool, options)
+    if success:
+      data = raw_data
+  elif tool == 'overview_page':
+    json_data, success = xspace_wrapper_func(xspace_paths, tool, options)
+    if success:
+      data = json_data
+  elif tool == 'input_pipeline_analyzer':
+    json_data, success = xspace_wrapper_func(xspace_paths, tool, options)
+    if success:
+      data = json_data
+  elif tool == 'framework_op_stats':
+    json_data, success = xspace_wrapper_func(xspace_paths, tool, options)
+    if success:
+      if tqx == 'out:csv':
+        data = csv_writer.json_to_csv(json_data)
+      else:
+        data = json_data
+    # Try legacy tool name: Handle backward compatibility with lower TF version
+    else:
+      # TODO(b/419013992): Remove this tool completely as it has been deprecated
+      legacy_tool = 'tensorflow_stats'
+      json_data, success = xspace_wrapper_func(
+          xspace_paths, legacy_tool, options
+      )
+      if success:
+        if tqx == 'out:csv':
+          data = csv_writer.json_to_csv(json_data)
+        else:
+          data = json_data
+  elif tool == 'kernel_stats':
+    json_data, success = xspace_wrapper_func(xspace_paths, tool, options)
+    if success:
+      if tqx == 'out:csv':
+        data = csv_writer.json_to_csv(json_data)
+      else:
+        data = json_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, options)
+    if success:
+      data = raw_data
+  elif tool == 'pod_viewer':
+    raw_data, success = xspace_wrapper_func(xspace_paths, tool, options)
+    if success:
+      data = raw_data
+  elif tool == 'op_profile':
+    options['group_by'] = params.get('group_by', 'program')
+    raw_data, success = xspace_wrapper_func(xspace_paths, tool, options)
+    if success:
+      data = raw_data
+  elif tool == 'hlo_op_profile':
+    options['group_by'] = params.get('group_by', 'program')
+    raw_data, success = xspace_wrapper_func(xspace_paths, tool, options)
+    if success:
+      data = raw_data
+  elif tool == 'hlo_stats':
+    json_data, success = xspace_wrapper_func(xspace_paths, tool, options)
+    if success:
+      data = json_data
+  elif tool == 'roofline_model':
+    json_data, success = xspace_wrapper_func(xspace_paths, tool, options)
+    if success:
+      data = json_data
+  elif tool == 'graph_viewer':
+    download_hlo_types = ['pb', 'pbtxt', 'json', 'short_txt', 'long_txt']
+    graph_html_type = 'graph'
+    options = params.get('graph_viewer_options', {})
+    options['use_saved_result'] = params.get('use_saved_result', True)
+    raw_data, success = xspace_wrapper_func(xspace_paths, tool, options)
+    if success:
+      data = raw_data
+      content_type = 'text/plain'
+      data_type = options.get('type', '')
+      if (data_type in download_hlo_types):
+        content_type = 'application/octet-stream'
+      if data_type == graph_html_type:
+        content_type = 'text/html'
+    else:
+      # TODO(tf-profiler) Handle errors for other tools as well,
+      # to pass along the error message to client
+      if isinstance(raw_data, bytes):
+        raw_data = raw_data.decode('utf-8')
+      raise ValueError(raw_data)
+  elif tool == 'memory_viewer':
+    view_memory_allocation_timeline = params.get(
+        'view_memory_allocation_timeline', False
+    )
+    options = {
+        'module_name': params.get('module_name'),
+        'view_memory_allocation_timeline': view_memory_allocation_timeline,
+        'memory_space': params.get('memory_space', ''),
+    }
+    raw_data, success = xspace_wrapper_func(xspace_paths, tool, options)
+    if success:
+      data = raw_data
+      if view_memory_allocation_timeline:
+        content_type = 'text/html'
+  elif tool == 'megascale_stats':
+    options = {'host_name': params.get('host')}
+    json_data, success = xspace_wrapper_func(xspace_paths, tool, options)
+    if success:
+      data = json_data
+  elif tool == 'inference_profile':
+    json_data, success = xspace_wrapper_func(xspace_paths, tool, options)
+    if success:
+      data = json_data
+  else:
+    logger.warning('%s is not a known xplane tool', tool)
+  return data, content_type

xprof/convert/trace_events_json.py

@@ -0,0 +1,105 @@
+# -*- coding: utf-8 -*-
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Converts trace events to JSON format consumed by catapult trace viewer."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import json
+import six
+
+# Values for type (ph) and s (scope) parameters in catapult trace format.
+_TYPE_METADATA = 'M'
+_TYPE_COMPLETE = 'X'
+_TYPE_INSTANT = 'i'
+_SCOPE_THREAD = 't'
+
+
+class TraceEventsJsonStream(object):
+  """A streaming trace file in the format expected by catapult trace viewer.
+
+  Iterating over this yields a sequence of string chunks, so it is suitable for
+  returning in a werkzeug Response.
+  """
+
+  def __init__(self, proto):
+    """Create an iterable JSON stream over the supplied Trace.
+
+    Args:
+      proto: a tensorboard.profile.Trace protobuf
+    """
+    self._proto = proto
+
+  def _events(self):
+    """Iterator over all catapult trace events, as python values."""
+    for did, device in sorted(six.iteritems(self._proto.devices)):
+      if device.name:
+        yield dict(
+            ph=_TYPE_METADATA,
+            pid=did,
+            name='process_name',
+            args=dict(name=device.name))
+      yield dict(
+          ph=_TYPE_METADATA,
+          pid=did,
+          name='process_sort_index',
+          args=dict(sort_index=did))
+      for rid, resource in sorted(six.iteritems(device.resources)):
+        if resource.name:
+          yield dict(
+              ph=_TYPE_METADATA,
+              pid=did,
+              tid=rid,
+              name='thread_name',
+              args=dict(name=resource.name))
+        yield dict(
+            ph=_TYPE_METADATA,
+            pid=did,
+            tid=rid,
+            name='thread_sort_index',
+            args=dict(sort_index=rid))
+    for event in self._proto.trace_events:
+      yield self._event(event)
+
+  def _event(self, event):
+    """Converts a TraceEvent proto into a catapult trace event python value."""
+    result = dict(
+        pid=event.device_id,
+        tid=event.resource_id,
+        name=event.name,
+        ts=event.timestamp_ps / 1000000.0)
+    if event.duration_ps:
+      result['ph'] = _TYPE_COMPLETE
+      result['dur'] = event.duration_ps / 1000000.0
+    else:
+      result['ph'] = _TYPE_INSTANT
+      result['s'] = _SCOPE_THREAD
+    for key in dict(event.args):
+      if 'args' not in result:
+        result['args'] = {}
+      result['args'][key] = event.args[key]
+    return result
+
+  def __iter__(self):
+    """Returns an iterator of string chunks of a complete JSON document."""
+    yield '{"displayTimeUnit":"ns","metadata":{"highres-ticks":true},\n'
+    yield '"traceEvents":[\n'
+    for event in self._events():
+      yield json.dumps(event, sort_keys=True)
+      yield ',\n'
+    # Add one fake event to avoid dealing with no-trailing-comma rule.
+    yield '{}]}\n'

xprof/integration_tests/tf_mnist.py

@@ -0,0 +1,100 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tensorflow based mnist model and input datasets helper."""
+
+_NUM_CLASSES = 10
+
+
+def get_input_datasets(tf, use_bfloat16=False):
+  """Downloads the MNIST dataset and creates train and eval dataset objects.
+
+  Args:
+    tf: version specific tensorflow (either tf1 or tf2).
+    use_bfloat16: Boolean to determine if input should be cast to bfloat16
+
+  Returns:
+    Train dataset, eval dataset, test dataset and input shape.
+
+  """
+  # input image dimensions
+  img_rows, img_cols = 28, 28
+  cast_dtype = tf.bfloat16 if use_bfloat16 else tf.float32
+
+  # the data, split between train and test sets
+  (x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
+
+  if tf.keras.backend.image_data_format() == 'channels_first':
+    x_train = x_train.reshape(x_train.shape[0], 1, img_rows, img_cols)
+    x_test = x_test.reshape(x_test.shape[0], 1, img_rows, img_cols)
+    input_shape = (1, img_rows, img_cols)
+  else:
+    x_train = x_train.reshape(x_train.shape[0], img_rows, img_cols, 1)
+    x_test = x_test.reshape(x_test.shape[0], img_rows, img_cols, 1)
+    input_shape = (img_rows, img_cols, 1)
+
+  x_train = x_train.astype('float32')
+  x_test = x_test.astype('float32')
+  x_train /= 255
+  x_test /= 255
+
+  # convert class vectors to binary class matrices
+  y_train = tf.keras.utils.to_categorical(y_train, _NUM_CLASSES)
+  y_test = tf.keras.utils.to_categorical(y_test, _NUM_CLASSES)
+
+  # train dataset
+  train_ds = tf.data.Dataset.from_tensor_slices((x_train, y_train))
+  train_ds = train_ds.repeat()
+  train_ds = train_ds.map(lambda x, y: (tf.cast(x, cast_dtype), y))
+  train_ds = train_ds.batch(64, drop_remainder=True)
+
+  # eval dataset
+  eval_ds = tf.data.Dataset.from_tensor_slices((x_test, y_test))
+  eval_ds = eval_ds.repeat()
+  eval_ds = eval_ds.map(lambda x, y: (tf.cast(x, cast_dtype), y))
+  eval_ds = eval_ds.batch(64, drop_remainder=True)
+
+  # test dataset
+  test_ds = tf.data.Dataset.from_tensor_slices((x_test))
+  test_ds = test_ds.repeat()
+  test_ds = test_ds.map(lambda x: tf.cast(x, cast_dtype))
+  test_ds = test_ds.batch(64, drop_remainder=True)
+
+  return train_ds, eval_ds, test_ds, input_shape
+
+
+def get_model(tf, input_shape):
+  """Builds a Sequential CNN model to recognize MNIST digits.
+
+  Args:
+   tf: version specific tensorflow (either tf1 or tf2).
+   input_shape: Shape of the input depending on the `image_data_format`.
+
+  Returns:
+    a Keras model
+
+  """
+  # Define a CNN model to recognize MNIST digits.
+  model = tf.keras.models.Sequential()
+  model.add(
+      tf.keras.layers.Conv2D(
+          32, kernel_size=(3, 3), activation='relu', input_shape=input_shape))
+  model.add(tf.keras.layers.Conv2D(64, (3, 3), activation='relu'))
+  model.add(tf.keras.layers.MaxPooling2D(pool_size=(2, 2)))
+  model.add(tf.keras.layers.Dropout(0.25))
+  model.add(tf.keras.layers.Flatten())
+  model.add(tf.keras.layers.Dense(128, activation='relu'))
+  model.add(tf.keras.layers.Dropout(0.5))
+  model.add(tf.keras.layers.Dense(_NUM_CLASSES, activation='softmax'))
+  return model

xprof/integration_tests/tf_profiler_session.py

@@ -0,0 +1,40 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Programmatic functionality for tensorflow profile session."""
+
+import types
+
+
+class TensorflowProfilerSession:
+  """context manager for Tensorflow programmatic profile session."""
+
+  def __init__(
+      self, tf: types.ModuleType, log_dir: str, python_tracer_level: int = 1
+  ):
+    """tf object is version dependent object."""
+    self.profiler = tf.profiler.experimental
+    self.log_dir = log_dir
+    self.python_tracer_level = python_tracer_level
+
+  def __enter__(self):
+    """Starts profile session and serializes data to temp directory."""
+    options = self.profiler.ProfilerOptions(
+        host_tracer_level=2, python_tracer_level=self.python_tracer_level
+    )
+    self.profiler.start(self.log_dir, options=options)
+
+  def __exit__(self, exc_type, exc_value, traceback):
+    """Ends current profile session and verifies test expectations."""
+    self.profiler.stop()