Add all_hosts information to the session_snapshot and move the device collision logic for trace_viewer to CreateTraceEventsContainer.

PiperOrigin-RevId: 822654197

Author: muditgokhale2

plugin/xprof/convert/raw_to_tool_data.py

@@ -41,12 +41,13 @@ def process_raw_trace(raw_trace):
   return ''.join(trace_events_json.TraceEventsJsonStream(trace))
 
 
-def xspace_to_tools_data_from_byte_string(xspace_byte_list, filenames, tool,
-                                          params):
+def xspace_to_tools_data_from_byte_string(xspace_byte_list, all_hosts,
+                                          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.
+    all_hosts: A list of all hosts in the session.
     filenames: Names of the read files.
     tool: A string of tool name.
     params: user input parameters.
@@ -57,7 +58,7 @@ def xspace_to_tools_data_from_byte_string(xspace_byte_list, filenames, tool,
 # 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)
+        xspace_arg, all_hosts, filenames, tool_arg, params)
 # pylint:enable=dangerous-default-value
 
   return xspace_to_tool_data(xspace_byte_list, tool, params,
@@ -73,22 +74,26 @@ def xspace_to_tool_names(xspace_paths):
   Returns:
     Returns a list of tool names.
   """
+  # xspace_to_tools_data expects all_hosts as the second argument, passing an
+  # empty list.
   raw_data, success = _pywrap_profiler_plugin.xspace_to_tools_data(
-      xspace_paths, 'tool_names')
+      xspace_paths, [], 'tool_names', {})
   if success:
     return [tool for tool in raw_data.decode().split(',')]
   return []
 
 
 def xspace_to_tool_data(
     xspace_paths,
+    all_hosts,
     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.
+    all_hosts: A list of all hosts in the session.
     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
@@ -112,26 +117,31 @@ def xspace_to_tool_data(
   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)
+    raw_data, success = xspace_wrapper_func(
+        xspace_paths, all_hosts, 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)
+    options['hosts'] = all_hosts
+    raw_data, success = xspace_wrapper_func(
+        xspace_paths, all_hosts, tool, options)
     if success:
       data = raw_data
   elif tool == 'overview_page':
-    json_data, success = xspace_wrapper_func(xspace_paths, tool, options)
+    json_data, success = xspace_wrapper_func(
+        xspace_paths, all_hosts, tool, options)
     if success:
       data = json_data
   elif tool == 'input_pipeline_analyzer':
-    json_data, success = xspace_wrapper_func(xspace_paths, tool, options)
+    json_data, success = xspace_wrapper_func(
+        xspace_paths, all_hosts, tool, options)
     if success:
       data = json_data
   elif tool == 'framework_op_stats':
-    json_data, success = xspace_wrapper_func(xspace_paths, tool, options)
+    json_data, success = xspace_wrapper_func(
+        xspace_paths, all_hosts, tool, options)
     if success:
       if tqx == 'out:csv':
         data = csv_writer.json_to_csv(json_data)
@@ -142,15 +152,16 @@ def xspace_to_tool_data(
       # 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
+          xspace_paths, all_hosts, 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)
+    json_data, success = xspace_wrapper_func(
+        xspace_paths, all_hosts, tool, options)
     if success:
       if tqx == 'out:csv':
         data = csv_writer.json_to_csv(json_data)
@@ -159,37 +170,44 @@ def xspace_to_tool_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)
+    raw_data, success = xspace_wrapper_func(
+        xspace_paths, all_hosts, tool, options)
     if success:
       data = raw_data
   elif tool == 'pod_viewer':
-    raw_data, success = xspace_wrapper_func(xspace_paths, tool, options)
+    raw_data, success = xspace_wrapper_func(
+        xspace_paths, all_hosts, 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)
+    raw_data, success = xspace_wrapper_func(
+        xspace_paths, all_hosts, 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)
+    raw_data, success = xspace_wrapper_func(
+        xspace_paths, all_hosts, tool, options)
     if success:
       data = raw_data
   elif tool == 'hlo_stats':
-    json_data, success = xspace_wrapper_func(xspace_paths, tool, options)
+    json_data, success = xspace_wrapper_func(
+        xspace_paths, all_hosts, tool, options)
     if success:
       data = json_data
   elif tool == 'roofline_model':
-    json_data, success = xspace_wrapper_func(xspace_paths, tool, options)
+    json_data, success = xspace_wrapper_func(
+        xspace_paths, all_hosts, 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)
+    raw_data, success = xspace_wrapper_func(
+        xspace_paths, all_hosts, tool, options)
     if success:
       data = raw_data
       content_type = 'text/plain'
@@ -213,18 +231,21 @@ def xspace_to_tool_data(
         'view_memory_allocation_timeline': view_memory_allocation_timeline,
         'memory_space': params.get('memory_space', ''),
     }
-    raw_data, success = xspace_wrapper_func(xspace_paths, tool, options)
+    raw_data, success = xspace_wrapper_func(
+        xspace_paths, all_hosts, 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)
+    json_data, success = xspace_wrapper_func(
+        xspace_paths, all_hosts, tool, options)
     if success:
       data = json_data
   elif tool == 'inference_profile':
-    json_data, success = xspace_wrapper_func(xspace_paths, tool, options)
+    json_data, success = xspace_wrapper_func(
+        xspace_paths, all_hosts, tool, options)
     if success:
       data = json_data
   else:

plugin/xprof/convert/raw_to_tool_data_test.py

@@ -27,7 +27,11 @@ def test_using_old_tool_format_maps_to_new_format(self):
         xspace_paths=["/path/to/xspace"],
         tool="trace_viewer@^",
         params={},
-        xspace_wrapper_func=lambda paths, tool, options: (tool.encode(), True),
+        all_hosts=[],
+        xspace_wrapper_func=lambda paths, hosts, tool, options: (
+            tool.encode(),
+            True,
+        ),
     )
 
     self.assertEqual(data, b"trace_viewer@")
@@ -38,7 +42,11 @@ def test_using_new_tool_format_does_not_map_to_old_format(self):
         xspace_paths=["/path/to/xspace"],
         tool="trace_viewer@",
         params={},
-        xspace_wrapper_func=lambda paths, tool, options: (tool.encode(), True),
+        all_hosts=[],
+        xspace_wrapper_func=lambda paths, hosts, tool, options: (
+            tool.encode(),
+            True,
+        ),
     )
 
     self.assertEqual(data, b"trace_viewer@")

plugin/xprof/integration_tests/tpu/tensorflow/tpu_tf2_keras_test.py

@@ -114,7 +114,7 @@ def test_tools_are_in_list(self):
 
   def test_overview_page(self):
     xspace_filenames = self._get_session_snapshot()
-    result, _ = raw_to_tool_data.xspace_to_tool_data(xspace_filenames,
+    result, _ = raw_to_tool_data.xspace_to_tool_data(xspace_filenames, [],
                                                      'overview_page', {})
     result = json.loads(result)
     run_environment = result[2]
@@ -123,7 +123,9 @@ def test_overview_page(self):
 
   def test_overview_page_creates_cache(self):
     xspace_filenames = self._get_session_snapshot()
-    raw_to_tool_data.xspace_to_tool_data(xspace_filenames, 'overview_page', {})
+    raw_to_tool_data.xspace_to_tool_data(
+        xspace_filenames, [], 'overview_page', {}
+    )
     profile_plugin_root = os.path.join(log_dir, 'plugins/profile')
     # The session exists under a director whose name is time-dependent.
     cache_glob = os.path.join(profile_plugin_root, '*', '*.op_stats.pb')
@@ -132,7 +134,7 @@ def test_overview_page_creates_cache(self):
   def test_op_profile(self):
     xspace_filenames = self._get_session_snapshot()
     result, _ = raw_to_tool_data.xspace_to_tool_data(
-        xspace_filenames, 'op_profile', {'group_by': 'category'}
+        xspace_filenames, [], 'op_profile', {'group_by': 'category'}
     )
     result = json.loads(result)
     logging.info(result)
@@ -151,7 +153,7 @@ def test_op_profile(self):
   def test_device_trace_contains_threads(self):
     xspace_filenames = self._get_session_snapshot()
     result, _ = raw_to_tool_data.xspace_to_tool_data(
-        xspace_filenames, 'trace_viewer', {}
+        xspace_filenames, [], 'trace_viewer', {}
     )
     result = json.loads(result)
     thread_names = []

plugin/xprof/profile_plugin.py

@@ -707,7 +707,7 @@ def hlo_module_list_route(
 
   def _get_valid_hosts(
       self, run_dir: str, run: str, tool: str, hosts_param: str, host: str
-  ) -> tuple[List[str], List[epath.Path]]:
+  ) -> tuple[List[str], List[epath.Path], List[str]]:
     """Retrieves and validates the hosts and asset paths for a run and tool.
 
     Args:
@@ -718,7 +718,7 @@ def _get_valid_hosts(
       host: The single host parameter.
 
     Returns:
-      A tuple containing (selected_hosts, asset_paths).
+      A tuple containing (selected_hosts, asset_paths, all_hosts).
 
     Raises:
       FileNotFoundError: If a required xplane file for the specified host(s)
@@ -781,7 +781,9 @@ def _get_valid_hosts(
             'Host must be specified for tool %s in run %s' % (tool, run)
         )
 
-    return selected_hosts, asset_paths
+    all_hosts = list(all_xplane_files.keys())
+
+    return selected_hosts, asset_paths, all_hosts
 
   def data_impl(
       self, request: wrappers.Request
@@ -870,7 +872,7 @@ def data_impl(
 
     _, content_encoding = None, None
     if use_xplane(tool):
-      selected_hosts, asset_paths = self._get_valid_hosts(
+      selected_hosts, asset_paths, all_hosts = self._get_valid_hosts(
           run_dir, run, tool, hosts_param, host
       )
       if not asset_paths:
@@ -879,7 +881,7 @@ def data_impl(
       params['hosts'] = selected_hosts
       try:
         data, content_type = convert.xspace_to_tool_data(
-            asset_paths, tool, params)
+            asset_paths, all_hosts, tool, params)
       except AttributeError as e:
         logger.warning('Error generating analysis results due to %s', e)
         raise AttributeError(

plugin/xprof/profile_plugin_test.py

@@ -465,7 +465,7 @@ def testDataImplTraceViewerOptions(self, mock_xspace_to_tool_data):
     )
 
     mock_xspace_to_tool_data.assert_called_once_with(
-        [mock.ANY], 'trace_viewer@', expected_params
+        [mock.ANY], ['host0', 'host1'], 'trace_viewer@', expected_params
     )
     args, _ = mock_xspace_to_tool_data.call_args
     actual_path_list = args[0]

xprof/convert/BUILD

@@ -191,6 +191,7 @@ cc_library(
         ":repository",
         ":tool_options",
         ":xplane_to_trace_container",
+        "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",

xprof/convert/repository.cc

@@ -58,7 +58,8 @@ static auto* kHostDataSuffixes =
 
 absl::StatusOr<SessionSnapshot> SessionSnapshot::Create(
     std::vector<std::string> xspace_paths,
-    std::optional<std::vector<std::unique_ptr<XSpace>>> xspaces) {
+    std::optional<std::vector<std::unique_ptr<XSpace>>> xspaces,
+    std::optional<std::vector<std::string>> all_hosts) {
   if (xspace_paths.empty()) {
     return absl::InvalidArgumentError("Can not find XSpace path.");
   }
@@ -85,7 +86,26 @@ absl::StatusOr<SessionSnapshot> SessionSnapshot::Create(
     }
   }
 
-  return SessionSnapshot(std::move(xspace_paths), std::move(xspaces));
+  return SessionSnapshot(std::move(xspace_paths), std::move(xspaces),
+                         std::move(all_hosts));
+}
+
+SessionSnapshot::SessionSnapshot(
+    std::vector<std::string> xspace_paths,
+    std::optional<std::vector<std::unique_ptr<XSpace>>> xspaces,
+    std::optional<std::vector<std::string>> all_hosts)
+    : xspace_paths_(std::move(xspace_paths)),
+      all_hosts_(std::move(all_hosts)),
+      // If the snapshot was initialized by xspaces, the file path and run dir
+      // is a path tensorflow can't read from or write to so any file IO
+      // encapsulated in this class will be disabled in this mode.
+      has_accessible_run_dir_(!xspaces.has_value()),
+      xspaces_(std::move(xspaces)) {
+  session_run_dir_ = tsl::io::Dirname(xspace_paths_.at(0));
+  for (size_t i = 0; i < xspace_paths_.size(); ++i) {
+    std::string host_name = GetHostname(i);
+    hostname_map_[host_name] = i;
+  }
 }
 
 absl::StatusOr<XSpace*> SessionSnapshot::GetXSpace(size_t index,
@@ -126,6 +146,10 @@ std::string SessionSnapshot::GetHostname(size_t index) const {
   return GetHostnameByPath(xspace_paths_.at(index));
 }
 
+std::optional<std::vector<std::string>> SessionSnapshot::GetAllHosts() const {
+  return all_hosts_;
+}
+
 std::optional<std::string> SessionSnapshot::GetFilePath(
     absl::string_view toolname, absl::string_view hostname) const {
   if (!has_accessible_run_dir_) return std::nullopt;