From a5165a531cd37d4a7549d4119bde9ac099b3df68 Mon Sep 17 00:00:00 2001 From: Luke Baumann Date: Thu, 4 Jun 2026 00:07:26 -0700 Subject: [PATCH] Delete experimental profiling utilities in pathwaysutils. PiperOrigin-RevId: 926509277 --- pathwaysutils/experimental/profiling.py | 63 -------------------- pathwaysutils/profiling.py | 53 +++++++++++------ pathwaysutils/test/profiling_test.py | 79 ++++++++++++++++++++++--- 3 files changed, 106 insertions(+), 89 deletions(-) delete mode 100644 pathwaysutils/experimental/profiling.py diff --git a/pathwaysutils/experimental/profiling.py b/pathwaysutils/experimental/profiling.py deleted file mode 100644 index 0816302..0000000 --- a/pathwaysutils/experimental/profiling.py +++ /dev/null @@ -1,63 +0,0 @@ -# Copyright 2025 Google LLC -# -# 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 -# -# https://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. -"""Experimental profiling utilites.""" - -from collections.abc import Mapping -from typing import Any - -from pathwaysutils import profiling - - -def start_trace( - profile_request: Mapping[str, Any], - *, - create_perfetto_link: bool = False, - create_perfetto_trace: bool = False, -) -> None: - """Starts a profiler trace. - - This is primarily for internal use where we can experiment with profile - requests with additional fields without needing to change - `pathwaysutils.profiling.start_trace` for each experiment. - - Use `jax.profiler.stop_trace` to end profiling. - - Args: - profile_request: A mapping containing the profile request options. - create_perfetto_link: A boolean which, if true, creates and prints link to - the Perfetto trace viewer UI (https://ui.perfetto.dev). The program will - block until the link is opened and Perfetto loads the trace. This feature - is experimental for Pathways on Cloud and may not be fully supported. - create_perfetto_trace: A boolean which, if true, additionally dumps a - ``perfetto_trace.json.gz`` file that is compatible for upload with the - Perfetto trace viewer UI (https://ui.perfetto.dev). The file will also be - generated if ``create_perfetto_link`` is true. This could be useful if you - want to generate a Perfetto-compatible trace without blocking the process. - This feature is experimental for Pathways on Cloud and may not be fully - supported. - """ - log_dir = profile_request["traceLocation"] - if not str(log_dir).startswith("gs://"): - raise ValueError( - "profile_request['traceLocation'] must be a GCS bucket path, got" - f" {log_dir}" - ) - - profiling._start_pathways_trace_from_profile_request(profile_request) # pylint: disable=protected-access - - profiling._original_start_trace( # pylint: disable=protected-access - log_dir=log_dir, - create_perfetto_link=create_perfetto_link, - create_perfetto_trace=create_perfetto_trace, - ) diff --git a/pathwaysutils/profiling.py b/pathwaysutils/profiling.py index 7fcffd1..7e067bc 100644 --- a/pathwaysutils/profiling.py +++ b/pathwaysutils/profiling.py @@ -16,6 +16,7 @@ import asyncio from collections.abc import Mapping import dataclasses +import datetime import json import logging import os @@ -112,8 +113,8 @@ def _is_default_profile_options( and profiler_options.python_tracer_level == default_options.python_tracer_level and profiler_options.duration_ms == default_options.duration_ms - and not getattr(profiler_options, "advanced_configuration", None) - and not getattr(profiler_options, "session_id", None) + and not profiler_options.advanced_configuration + and not profiler_options.session_id ) @@ -132,9 +133,9 @@ def _create_profile_request( return profile_request advanced_config = None - if getattr(profiler_options, "advanced_configuration", None): + if profiler_options.advanced_configuration: advanced_config = {} - for k, v in getattr(profiler_options, "advanced_configuration").items(): + for k, v in profiler_options.advanced_configuration.items(): # Convert python dict to tensorflow.ProfileOptions.AdvancedConfigValue # json-compatible dict if isinstance(v, bool): @@ -168,7 +169,7 @@ def _create_profile_request( if pw_trace_opts: xprof_options["pwTraceOptions"] = pw_trace_opts - if getattr(profiler_options, "session_id", None): + if profiler_options.session_id: xprof_options["traceSessionName"] = profiler_options.session_id profile_request["xprofTraceOptions"] = xprof_options @@ -270,26 +271,44 @@ def start_trace( "features for Pathways on Cloud and may not be fully supported." ) - if jax.version.__version_info__ < (0, 9, 2) and profiler_options is not None: - _logger.warning( - "ProfileOptions are not supported until JAX 0.9.2 and will be omitted. " - "Some options can be specified via command line flags." - ) - profiler_options = None + if jax.version.__version_info__ < (0, 9, 2): + if profiler_options is not None: + _logger.warning( + "ProfileOptions are not supported until JAX 0.9.2 and will be omitted. " + "Some options can be specified via command line flags." + ) + profiler_options = None + else: + if profiler_options is None: + profiler_options = jax.profiler.ProfileOptions() + if not profiler_options.session_id: + profiler_options.session_id = datetime.datetime.now().strftime( + "%Y_%m_%d_%H_%M_%S" + ) profile_request = _create_profile_request( - log_dir, profiler_options, max_num_hosts=max_num_hosts + log_dir, + profiler_options, + max_num_hosts=max_num_hosts, ) _logger.debug("Profile request: %s", profile_request) _start_pathways_trace_from_profile_request(profile_request) - _original_start_trace( - log_dir=log_dir, - create_perfetto_link=create_perfetto_link, - create_perfetto_trace=create_perfetto_trace, - ) + if jax.version.__version_info__ >= (0, 9, 2): + _original_start_trace( + log_dir=log_dir, + create_perfetto_link=create_perfetto_link, + create_perfetto_trace=create_perfetto_trace, + profiler_options=profiler_options, + ) + else: + _original_start_trace( + log_dir=log_dir, + create_perfetto_link=create_perfetto_link, + create_perfetto_trace=create_perfetto_trace, + ) def stop_trace() -> None: diff --git a/pathwaysutils/test/profiling_test.py b/pathwaysutils/test/profiling_test.py index 5cb2fa8..bf12857 100644 --- a/pathwaysutils/test/profiling_test.py +++ b/pathwaysutils/test/profiling_test.py @@ -53,6 +53,13 @@ def setUp(self): self.mock_original_stop_trace = self.enter_context( mock.patch.object(profiling, "_original_stop_trace", autospec=True) ) + self.mock_datetime = self.enter_context( + mock.patch.object(profiling.datetime, "datetime", autospec=True) + ) + self.mock_datetime.now.return_value.strftime.return_value = ( + "2026_06_04_05_29_33" + ) + @parameterized.parameters(8000, 1234) def test_collect_profile_port(self, port): @@ -233,15 +240,26 @@ def test_start_trace_success(self): "profileRequest": { "traceLocation": "gs://test_bucket/test_dir", "maxNumHosts": 1, + "xprofTraceOptions": { + "traceDirectory": "gs://test_bucket/test_dir", + "pwTraceOptions": { + "enablePythonTracer": True, + }, + "traceSessionName": "2026_06_04_05_29_33", + }, } }) ) self.mock_plugin_executable_cls.return_value.call.assert_called_once() - self.mock_original_start_trace.assert_called_once_with( - log_dir="gs://test_bucket/test_dir", - create_perfetto_link=False, - create_perfetto_trace=False, - ) + self.mock_original_start_trace.assert_called_once() + call_args = self.mock_original_start_trace.call_args[1] + self.assertEqual(call_args["log_dir"], "gs://test_bucket/test_dir") + self.assertFalse(call_args["create_perfetto_link"]) + self.assertFalse(call_args["create_perfetto_trace"]) + if jax.version.__version_info__ >= (0, 9, 2): + self.assertEqual( + call_args["profiler_options"].session_id, "2026_06_04_05_29_33" + ) self.assertIsNotNone(profiling._profile_state.executable) def test_start_trace_with_max_num_hosts(self): @@ -253,15 +271,58 @@ def test_start_trace_with_max_num_hosts(self): "profileRequest": { "traceLocation": "gs://test_bucket/test_dir", "maxNumHosts": 10, + "xprofTraceOptions": { + "traceDirectory": "gs://test_bucket/test_dir", + "pwTraceOptions": { + "enablePythonTracer": True, + }, + "traceSessionName": "2026_06_04_05_29_33", + }, } }) ) self.mock_plugin_executable_cls.return_value.call.assert_called_once() - self.mock_original_start_trace.assert_called_once_with( - log_dir="gs://test_bucket/test_dir", - create_perfetto_link=False, - create_perfetto_trace=False, + self.mock_original_start_trace.assert_called_once() + call_args = self.mock_original_start_trace.call_args[1] + self.assertEqual(call_args["log_dir"], "gs://test_bucket/test_dir") + self.assertFalse(call_args["create_perfetto_link"]) + self.assertFalse(call_args["create_perfetto_trace"]) + if jax.version.__version_info__ >= (0, 9, 2): + self.assertEqual( + call_args["profiler_options"].session_id, "2026_06_04_05_29_33" + ) + + @absltest.skipIf( + jax.version.__version_info__ < (0, 9, 2), + "ProfileOptions requires JAX 0.9.2 or newer", + ) + def test_start_trace_with_session_id_in_options(self): + options = jax.profiler.ProfileOptions() + options.session_id = "options_session" + profiling.start_trace("gs://test_bucket/test_dir", profiler_options=options) + + self.mock_plugin_executable_cls.assert_called_once_with( + json.dumps({ + "profileRequest": { + "traceLocation": "gs://test_bucket/test_dir", + "maxNumHosts": 1, + "xprofTraceOptions": { + "traceDirectory": "gs://test_bucket/test_dir", + "pwTraceOptions": { + "enablePythonTracer": True, + }, + "traceSessionName": "options_session", + }, + } + }) ) + self.assertEqual(options.session_id, "options_session") + self.mock_original_start_trace.assert_called_once() + call_args = self.mock_original_start_trace.call_args[1] + self.assertEqual(call_args["log_dir"], "gs://test_bucket/test_dir") + self.assertFalse(call_args["create_perfetto_link"]) + self.assertFalse(call_args["create_perfetto_trace"]) + self.assertEqual(call_args["profiler_options"].session_id, "options_session") def test_start_trace_no_toy_computation_second_time(self): profiling.start_trace("gs://test_bucket/test_dir")