Add step time to step time graph in overview page.

PiperOrigin-RevId: 616919828
tensorflow · Mar 18, 2024 · e0af893 · e0af893
1 parent 7b160ff
commit e0af893
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/frontend/app/components/overview/step_time_graph/step_time_graph.ts b/frontend/app/components/overview/step_time_graph/step_time_graph.ts
@@ -5,6 +5,7 @@ import {type InputPipelineAnalysis} from 'org_xprof/frontend/app/common/interfac
 const MAX_CHART_WIDTH = 800;
 const COLUMN_IDS_FOR_TPU = [
   'stepnum',
+  'stepTimeMs',
   'computeTimeMs',
   'inputTimeMs',
   'idleTimeMs',

diff --git a/plugin/tensorboard_plugin_profile/convert/input_pipeline_proto_to_gviz.py b/plugin/tensorboard_plugin_profile/convert/input_pipeline_proto_to_gviz.py
@@ -40,6 +40,7 @@ def get_step_breakdown_table_args(ipa):
 
   table_description = [
       ("stepnum", "string", "Step number"),
+      ("stepTimeMs", "number", "Step time"),
       ("deviceComputeTimeMs", "number", "Device compute"),
       ("deviceToDeviceTimeMs", "number", "Device to device"),
       ("deviceCollectivesTimeMs", "number", "Device collectives"),
@@ -89,7 +90,7 @@ def get_step_breakdown_table_args(ipa):
                    details.device_compute_ms)
 
     row = [
-        details.step_name, details.device_compute_ms,
+        details.step_name, details.step_time_ms, details.device_compute_ms,
         details.device_to_device_ms, details.device_collectives_ms,
         details.host_compute_ms, details.host_prepare_ms,
         details.host_wait_input_ms + details.host_to_device_ms,