From 3e62e800273068652a0937da249347abafeb69c0 Mon Sep 17 00:00:00 2001 From: Profiler Team Date: Thu, 9 Jul 2026 01:28:23 -0700 Subject: [PATCH] Add counter extraction to get counter names from files. PiperOrigin-RevId: 944966991 --- plugin/build_pip_package.sh | 4 + plugin/setup.py | 1 + plugin/xprof/convert/BUILD | 18 ++ plugin/xprof/convert/counter_extractor.py | 49 +++++ .../xprof/convert/counter_extractor_test.py | 167 ++++++++++++++++++ .../xprof/convert/scratch/test_extractor.py | 8 + .../xprof/convert/scratch/test_resources.py | 24 +++ plugin/xprof/profile_plugin.py | 17 ++ plugin/xprof/profile_plugin_test.py | 41 +++++ plugin/xprof/profile_plugin_test_utils.py | 14 +- xprof/utils/BUILD | 9 + 11 files changed, 349 insertions(+), 3 deletions(-) create mode 100644 plugin/xprof/convert/counter_extractor.py create mode 100644 plugin/xprof/convert/counter_extractor_test.py create mode 100644 plugin/xprof/convert/scratch/test_extractor.py create mode 100644 plugin/xprof/convert/scratch/test_resources.py diff --git a/plugin/build_pip_package.sh b/plugin/build_pip_package.sh index ca46623c7..ee50c259a 100755 --- a/plugin/build_pip_package.sh +++ b/plugin/build_pip_package.sh @@ -93,6 +93,10 @@ fi # Copy static files. cd xprof +mkdir -p utils +cp "${ROOT_RUNFILE_DIR}/xprof/utils/tpu_counter_ids_v7x.h" utils/ +cp "${ROOT_RUNFILE_DIR}/xprof/utils/tpu_counter_ids_v6e.h" utils/ + mkdir -p static cd static cp "$PLUGIN_RUNFILE_DIR/xprof/static/index.html" . diff --git a/plugin/setup.py b/plugin/setup.py index f6dfe275b..48d4a75d0 100644 --- a/plugin/setup.py +++ b/plugin/setup.py @@ -83,6 +83,7 @@ def get_readme(): package_data={ 'xprof': [ 'static/**', + 'utils/*.h', 'convert/profiler_plugin_c_api.so', 'convert/profiler_plugin_c_api.pyd', 'convert/profiler_plugin_c_api.dylib', diff --git a/plugin/xprof/convert/BUILD b/plugin/xprof/convert/BUILD index 96b6f59c3..f5c2ac32c 100644 --- a/plugin/xprof/convert/BUILD +++ b/plugin/xprof/convert/BUILD @@ -14,6 +14,7 @@ py_library( name = "all_libs", # strict_deps = False, deps = [ + ":counter_extractor", ":raw_to_tool_data", ":trace_events_json", ], @@ -59,6 +60,23 @@ py_library( ], ) +py_library( + name = "counter_extractor", + srcs = ["counter_extractor.py"], + data = ["@org_xprof//xprof/utils:tpu_counter_ids_headers"], + # strict_deps = False, +) + +py_test( + name = "counter_extractor_test", + srcs = ["counter_extractor_test.py"], + # strict_deps = False, + deps = [ + ":counter_extractor", + "@absl_py//absl/testing:absltest", + ], +) + py_test( name = "raw_to_tool_data_test", srcs = ["raw_to_tool_data_test.py"], diff --git a/plugin/xprof/convert/counter_extractor.py b/plugin/xprof/convert/counter_extractor.py new file mode 100644 index 000000000..b86d81811 --- /dev/null +++ b/plugin/xprof/convert/counter_extractor.py @@ -0,0 +1,49 @@ +"""Defines the counter extractor routine for TPU counters.""" + +import functools +import importlib.resources +import re + +VALID_DEVICE_TYPES = frozenset({'v6e', 'v7x'}) + + +@functools.lru_cache() +def get_all_counters(device_type: str) -> list[dict[str, int]]: + """Parses tpu_counter_ids_{device_type}.h and returns a list of counters with names and indices.""" + if device_type not in VALID_DEVICE_TYPES: + raise ValueError(f'Unsupported device_type: {device_type}') + filename = f'tpu_counter_ids_{device_type}.h' + + try: + ref = importlib.resources.files('xprof') / 'utils' / filename + except ImportError: + pass + content = ref.read_text(encoding='utf-8') + # Remove C++ style line comments from content. + content = re.sub(r'//.*', '', content) + + # The counter file is a C++ header file with an enum. + # Example: + # enum TpuCounterIdsTpu{version} : uint64_t { + # TPU_COUNTER_ID_FOO = 12345, + # TPU_COUNTER_ID_BAR = 67890, + # ... + # }; + match = re.search(r'enum\s+\w+\s*:\s*uint64_t\s*\{([^}]+)\}', content) + if not match: + return [] + + enum_body = match.group(1) + counters = [] + for line in enum_body.split(','): + line = line.strip() + if '=' in line: + name, val_str = line.split('=', 1) + try: + name = name.strip() + val = int(val_str.strip(), 0) + counters.append({'name': name.lower(), 'val': val}) + except ValueError: + continue + + return counters diff --git a/plugin/xprof/convert/counter_extractor_test.py b/plugin/xprof/convert/counter_extractor_test.py new file mode 100644 index 000000000..a70c9e400 --- /dev/null +++ b/plugin/xprof/convert/counter_extractor_test.py @@ -0,0 +1,167 @@ +# Copyright 2026 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. +# ============================================================================== + +"""Tests for the counter_extractor module.""" + +from unittest import mock + +from absl.testing import absltest + +from xprof.convert import counter_extractor + + +class CounterExtractorTest(absltest.TestCase): + + def setUp(self): + super().setUp() + counter_extractor.get_all_counters.cache_clear() + + @mock.patch("importlib.resources.files") + def test_get_all_counters_success(self, mock_files): + mock_content = """ + #ifndef THIRD_PARTY_XPROF_UTILS_TPU_COUNTER_IDS_H_ + #define THIRD_PARTY_XPROF_UTILS_TPU_COUNTER_IDS_H_ + namespace xprof { + enum TpuCounterIdsTpu7x : uint64_t { + // NOLINTBEGIN + TPU_COUNTER_ID_FOO = 12345, + // Some comment + TPU_COUNTER_ID_BAR = + 67890, + TPU_COUNTER_ID_HEX = 0x1234, + TPU_COUNTER_ID_MULTI_EQ = 1 = 2, + // NOLINTEND + }; + } + #endif + """ + + mock_file = mock.MagicMock() + mock_file.read_text.return_value = mock_content + # Chain: files() -> / 'utils' -> / filename + mock_pkg = mock.MagicMock() + mock_utils = mock.MagicMock() + mock_pkg.__truediv__.return_value = mock_utils + mock_utils.__truediv__.return_value = mock_file + mock_files.return_value = mock_pkg + + counters = counter_extractor.get_all_counters("v7x") + self.assertEqual( + counters, + [ + {"name": "tpu_counter_id_foo", "val": 12345}, + {"name": "tpu_counter_id_bar", "val": 67890}, + {"name": "tpu_counter_id_hex", "val": 4660}, + ], + ) + + @mock.patch("importlib.resources.files") + def test_get_all_counters_fallback_success(self, mock_files): + mock_content = """ + #ifndef THIRD_PARTY_XPROF_UTILS_TPU_COUNTER_IDS_H_ + #define THIRD_PARTY_XPROF_UTILS_TPU_COUNTER_IDS_H_ + namespace xprof { + enum TpuCounterIdsTpu7x : uint64_t { + TPU_COUNTER_ID_FOO = 12345, + }; + } + #endif + """ + + mock_file = mock.MagicMock() + mock_file.read_text.return_value = mock_content + mock_pkg = mock.MagicMock() + mock_utils = mock.MagicMock() + mock_pkg.__truediv__.return_value = mock_utils + mock_utils.__truediv__.return_value = mock_file + + def side_effect(package_name): + if package_name == "google3.third_party.xprof": + raise ImportError("No module named google3") + elif package_name == "xprof": + return mock_pkg + raise ValueError(f"Unexpected package: {package_name}") + + mock_files.side_effect = side_effect + + counters = counter_extractor.get_all_counters("v7x") + self.assertEqual( + counters, + [ + {"name": "tpu_counter_id_foo", "val": 12345}, + ], + ) + + @mock.patch("importlib.resources.files") + def test_get_all_counters_missing_file(self, mock_files): + mock_utils = mock.MagicMock() + # Simulate FileNotFound when reading text + mock_file = mock.MagicMock() + mock_file.read_text.side_effect = FileNotFoundError("File not found") + mock_pkg = mock.MagicMock() + mock_pkg.__truediv__.return_value = mock_utils + mock_utils.__truediv__.return_value = mock_file + mock_files.return_value = mock_pkg + + with self.assertRaises(FileNotFoundError): + counter_extractor.get_all_counters("v7x") + + @mock.patch("importlib.resources.files") + def test_get_all_counters_no_match(self, mock_files): + mock_content = "some file without enum" + mock_file = mock.MagicMock() + mock_file.read_text.return_value = mock_content + mock_pkg = mock.MagicMock() + mock_utils = mock.MagicMock() + mock_pkg.__truediv__.return_value = mock_utils + mock_utils.__truediv__.return_value = mock_file + mock_files.return_value = mock_pkg + + counters = counter_extractor.get_all_counters("v7x") + self.assertEqual(counters, []) + + @mock.patch("importlib.resources.files") + def test_get_all_counters_empty_enum(self, mock_files): + mock_content_empty = "enum TpuCounterIdsTpu7x : uint64_t {};" + mock_content_whitespace = "enum TpuCounterIdsTpu7x : uint64_t { \n };" + + mock_file = mock.MagicMock() + mock_pkg = mock.MagicMock() + mock_utils = mock.MagicMock() + mock_pkg.__truediv__.return_value = mock_utils + mock_utils.__truediv__.return_value = mock_file + mock_files.return_value = mock_pkg + + # Test completely empty + mock_file.read_text.return_value = mock_content_empty + counters = counter_extractor.get_all_counters("v7x") + self.assertEqual(counters, []) + + # Test whitespace only + mock_file.read_text.return_value = mock_content_whitespace + counters = counter_extractor.get_all_counters("v7x") + self.assertEqual(counters, []) + + def test_get_all_counters_invalid_device_type(self): + with self.assertRaises(ValueError): + counter_extractor.get_all_counters("../invalid") + + def test_get_all_counters_security_path_separators(self): + with self.assertRaises(ValueError): + counter_extractor.get_all_counters("v7x/../etc/passwd") + + +if __name__ == "__main__": + absltest.main() diff --git a/plugin/xprof/convert/scratch/test_extractor.py b/plugin/xprof/convert/scratch/test_extractor.py new file mode 100644 index 000000000..fbaf70372 --- /dev/null +++ b/plugin/xprof/convert/scratch/test_extractor.py @@ -0,0 +1,8 @@ +"""Test reading TPU counter IDs resources.""" +from xprof.convert import counter_extractor + +try: + counters = counter_extractor.get_all_counters("v7x") + print(f"Success! Found {len(counters)} counters.") +except (ImportError, ValueError) as e: + print(f"Failed: {e}") diff --git a/plugin/xprof/convert/scratch/test_resources.py b/plugin/xprof/convert/scratch/test_resources.py new file mode 100644 index 000000000..912bd1af5 --- /dev/null +++ b/plugin/xprof/convert/scratch/test_resources.py @@ -0,0 +1,24 @@ +"""Test reading TPU counter IDs resources.""" +import importlib.resources + +try: + # Try to read using importlib.resources + content = importlib.resources.read_text( + "google3.third_party.xprof.utils", "tpu_counter_ids_v7x.h" + ) + print("Success with importlib.resources!") + print(content[:100]) +except ImportError as e: + print(f"Failed with importlib.resources: {e}") + +try: + # Try with files() + ref = ( + importlib.resources.files("google3.third_party.xprof.utils") + / "tpu_counter_ids_v7x.h" + ) + content = ref.read_text() + print("Success with importlib.resources.files!") + print(content[:100]) +except ImportError as e: + print(f"Failed with importlib.resources.files: {e}") diff --git a/plugin/xprof/profile_plugin.py b/plugin/xprof/profile_plugin.py index 94fc7a439..55fe91a9b 100644 --- a/plugin/xprof/profile_plugin.py +++ b/plugin/xprof/profile_plugin.py @@ -37,6 +37,7 @@ from xprof import profile_io from xprof import version +from xprof.convert import counter_extractor from xprof.convert import raw_to_tool_data as convert from xprof.standalone.tensorboard_shim import base_plugin from xprof.standalone.tensorboard_shim import plugin_asset_util @@ -1211,6 +1212,22 @@ def data_impl( """ run = request.args.get('run') tool = request.args.get('tag') + + if tool == 'perf_counters' and request.args.get('names_only') == '1': + device_type = request.args.get('device_type') + if not device_type: + raise ValueError( + 'device_type is required for perf_counters with names_only' + ) + try: + names = counter_extractor.get_all_counters(device_type) + except FileNotFoundError: + logger.warning( + 'Failed to get counter names for device type: %s', device_type + ) + names = [] + return json.dumps(names), 'application/json', None + hosts_param = request.args.get('hosts') host = request.args.get('host') module_name = request.args.get('module_name') diff --git a/plugin/xprof/profile_plugin_test.py b/plugin/xprof/profile_plugin_test.py index f5a578087..2e6f89461 100644 --- a/plugin/xprof/profile_plugin_test.py +++ b/plugin/xprof/profile_plugin_test.py @@ -377,6 +377,47 @@ def testData(self): ) ) + def testDataNamesOnly_HappyPath(self): + data, content_type, _ = self.plugin.data_impl( + utils.make_data_request( + utils.DataRequestOptions( + tool='perf_counters', names_only='1', device_type='v7x' + ) + ) + ) + self.assertEqual(content_type, 'application/json') + names = json.loads(data) + self.assertIsInstance(names, list) + self.assertNotEmpty(names) + self.assertIn('name', names[0]) + self.assertIn('val', names[0]) + + def testDataNamesOnly_MissingDeviceType(self): + with self.assertRaisesRegex(ValueError, 'device_type is required'): + self.plugin.data_impl( + utils.make_data_request( + utils.DataRequestOptions(tool='perf_counters', names_only='1') + ) + ) + + def testDataNamesOnly_FileNotFoundFallback(self): + with mock.patch.object( + profile_plugin.counter_extractor, 'get_all_counters' + ) as mock_get_counters: + mock_get_counters.side_effect = FileNotFoundError( + 'Simulated missing file' + ) + data, content_type, _ = self.plugin.data_impl( + utils.make_data_request( + utils.DataRequestOptions( + tool='perf_counters', names_only='1', device_type='v7x' + ) + ) + ) + self.assertEqual(content_type, 'application/json') + names = json.loads(data) + self.assertEqual(names, []) + def testDataWithCache(self): generate_testdata(self.logdir) self.multiplexer.AddRunsFromDirectory(self.logdir) diff --git a/plugin/xprof/profile_plugin_test_utils.py b/plugin/xprof/profile_plugin_test_utils.py index 639827570..322d16fc6 100644 --- a/plugin/xprof/profile_plugin_test_utils.py +++ b/plugin/xprof/profile_plugin_test_utils.py @@ -23,7 +23,7 @@ from typing import Any, Callable, Sequence from etils import epath -from werkzeug import Request +from werkzeug import wrappers from xprof import profile_plugin from xprof import version @@ -104,6 +104,8 @@ class DataRequestOptions: search_metadata: Whether to search event metadata. format: Data format, e.g., 'pb'. event_name: Name of the event to select. + names_only: Whether to only return counter names instead of full data. + device_type: Device type for perf counter names, e.g., 'v7x'. """ run: str | None = None @@ -119,9 +121,11 @@ class DataRequestOptions: search_metadata: str | None = None format: str | None = None event_name: str | None = None + names_only: str | None = None + device_type: str | None = None -def make_data_request(options: DataRequestOptions) -> Request: +def make_data_request(options: DataRequestOptions) -> wrappers.Request: """Creates a werkzeug.Request to pass as argument to ProfilePlugin.data_impl. Args: @@ -130,7 +134,7 @@ def make_data_request(options: DataRequestOptions) -> Request: Returns: A werkzeug.Request to pass to ProfilePlugin.data_impl. """ - req = Request({}) + req = wrappers.Request({}) req.args = {} if options.run: req.args['run'] = options.run @@ -158,4 +162,8 @@ def make_data_request(options: DataRequestOptions) -> Request: req.args['format'] = options.format if options.event_name is not None: req.args['event_name'] = options.event_name + if options.names_only is not None: + req.args['names_only'] = options.names_only + if options.device_type: + req.args['device_type'] = options.device_type return req diff --git a/xprof/utils/BUILD b/xprof/utils/BUILD index e8163f0f4..0d6d32916 100644 --- a/xprof/utils/BUILD +++ b/xprof/utils/BUILD @@ -750,6 +750,15 @@ xprof_proto_library( srcs = ["backend_configs.proto"], ) +filegroup( + name = "tpu_counter_ids_headers", + srcs = [ + "tpu_counter_ids_v6e.h", + "tpu_counter_ids_v7x.h", + ], + visibility = ["//plugin:internal"], +) + cc_library( name = "tpu_counter_ids_v7x", hdrs = ["tpu_counter_ids_v7x.h"],