From 9e34c03be82204ff96e207e57661c048adc3fcf5 Mon Sep 17 00:00:00 2001 From: Pathways-on-Cloud Team Date: Sun, 28 Jun 2026 23:04:59 -0700 Subject: [PATCH] Secure Pathways Profiling API using gRPC PiperOrigin-RevId: 939630828 --- pathwaysutils/profiling.py | 197 ++++++++++----- pathwaysutils/proto/pathways_profiler.proto | 46 ++++ pathwaysutils/test/profiling_test.py | 255 +++++++++++++++----- 3 files changed, 371 insertions(+), 127 deletions(-) create mode 100644 pathwaysutils/proto/pathways_profiler.proto diff --git a/pathwaysutils/profiling.py b/pathwaysutils/profiling.py index 5a54586..15101b6 100644 --- a/pathwaysutils/profiling.py +++ b/pathwaysutils/profiling.py @@ -13,24 +13,23 @@ # limitations under the License. """Profiling Utilities.""" -import asyncio from collections.abc import Mapping -import dataclasses +import concurrent.futures import datetime import json import logging import os import threading +import time from typing import Any -import urllib.parse -import fastapi +import grpc import jax from jax import numpy as jnp from jax.extend import backend from pathwaysutils import plugin_executable -import requests -import uvicorn +from pathwaysutils.proto import pathways_profiler_pb2 +from pathwaysutils.proto import pathways_profiler_pb2_grpc _logger = logging.getLogger(__name__) @@ -297,19 +296,27 @@ def start_trace( _start_pathways_trace_from_profile_request(profile_request) - 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, + try: + 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, + ) + except Exception: + _logger.exception( + "Failed to start JAX local trace, resetting pathways trace state" ) + with _profile_state.lock: + _profile_state.reset() + raise def stop_trace() -> None: @@ -326,54 +333,83 @@ def stop_trace() -> None: _original_stop_trace() -_profiler_thread: threading.Thread | None = None +_profiler_server: grpc.Server | None = None +_profiler_server_lock = threading.Lock() -def start_server(port: int) -> None: - """Starts the profiling server on port `port`. +class PathwaysProfilerServicer( + pathways_profiler_pb2_grpc.PathwaysProfilerServiceServicer +): + """gRPC servicer for Pathways Profiler Service.""" - The signature is slightly different from `jax.profiler.start_server` - because no handle to the server is returned because there is no - `xla_client.profiler.ProfilerServer` to return. + def Profile( + self, + request: pathways_profiler_pb2.ProfileRequest, + context: grpc.ServicerContext, + ) -> pathways_profiler_pb2.ProfileResponse: + _logger.info("Received gRPC profile request for %s ms", request.duration_ms) + _logger.info("Writing profiling data to %s", request.repository_path) - Args: - port : The port to start the server on. - """ - def server_loop(port: int): - _logger.debug("Starting JAX profiler server on port %s", port) - app = fastapi.FastAPI() - - @dataclasses.dataclass - class ProfilingConfig: - duration_ms: int - repository_path: str + try: + # jax.profiler.start_trace is monkey-patched to start pathways trace + jax.profiler.start_trace(request.repository_path) + + elapsed = 0.0 + duration_secs = request.duration_ms / 1000.0 + while elapsed < duration_secs: + if not context.is_active(): + _logger.warning("Client disconnected, aborting profile.") + raise RuntimeError("Client disconnected") + time.sleep(0.1) + elapsed += 0.1 + + except Exception as e: # pylint: disable=broad-exception-caught + _logger.exception("Error during profiling") + context.set_code(grpc.StatusCode.INTERNAL) + context.set_details(str(e)) + return pathways_profiler_pb2.ProfileResponse(status=f"Failed: {e}") + finally: + _logger.info("Stopping trace") + try: + jax.profiler.stop_trace() + except Exception: # pylint: disable=broad-exception-caught + _logger.exception("Failed to stop trace") - @app.post("/profiling") - async def profiling(pc: ProfilingConfig) -> Mapping[str, str]: - _logger.debug("Capturing profiling data for %s ms", pc.duration_ms) - _logger.debug("Writing profiling data to %s", pc.repository_path) - await asyncio.to_thread(jax.profiler.start_trace, pc.repository_path) - await asyncio.sleep(pc.duration_ms / 1e3) - await asyncio.to_thread(jax.profiler.stop_trace) - return {"response": "profiling completed"} + return pathways_profiler_pb2.ProfileResponse(status="Completed") - uvicorn.run(app, host="0.0.0.0", port=port, log_level="debug") - global _profiler_thread - if _profiler_thread is not None: - raise RuntimeError("Only one profiler server can be active at a time.") +def start_server(port: int) -> None: + """Starts the profiling server on port `port`. - _profiler_thread = threading.Thread(target=server_loop, args=(port,)) - _profiler_thread.start() + Args: + port: The port to start the server on. + """ + global _profiler_server + with _profiler_server_lock: + if _profiler_server is not None: + raise RuntimeError("Only one profiler server can be active at a time.") + + _logger.info("Starting JAX pathways profiler gRPC server on port %s", port) + server = grpc.server(concurrent.futures.ThreadPoolExecutor(max_workers=2)) + pathways_profiler_pb2_grpc.add_PathwaysProfilerServiceServicer_to_server( + PathwaysProfilerServicer(), server + ) + # Use ALTS credentials for secure communication inside Google + server_creds = grpc.alts_server_credentials() + server.add_secure_port(f"[::]:{port}", server_creds) + server.start() + _profiler_server = server def stop_server() -> None: - """Raises an error if there is no active profiler server. - - Pathways profiling servers are not stoppable at this time. - """ - if _profiler_thread is None: - raise RuntimeError("No active profiler server.") + """Stops the active profiler server.""" + global _profiler_server + with _profiler_server_lock: + if _profiler_server is None: + raise RuntimeError("No active profiler server.") + _logger.info("Stopping JAX pathways profiler gRPC server") + _profiler_server.stop(grace=5.0) + _profiler_server = None def collect_profile( @@ -399,16 +435,49 @@ def collect_profile( if not str(log_dir).startswith("gs://"): raise ValueError(f"log_dir must be a GCS bucket path, got {log_dir}") - request_json = { - "duration_ms": duration_ms, - "repository_path": log_dir, - } - address = urllib.parse.urljoin(f"http://{host}:{port}", "profiling") + # Use ALTS credentials for secure client connection + creds = grpc.alts_channel_credentials() + target = f"{host}:{port}" + _logger.info("Connecting to profiling server at %s using ALTS", target) try: - response = requests.post(address, json=request_json) - response.raise_for_status() - except requests.exceptions.RequestException: - _logger.exception("Failed to collect profiling data") + with grpc.secure_channel(target, creds) as channel: + stub = pathways_profiler_pb2_grpc.PathwaysProfilerServiceStub(channel) + request = pathways_profiler_pb2.ProfileRequest( + duration_ms=duration_ms, + repository_path=str(log_dir), + ) + timeout = (duration_ms / 1000.0) + 10.0 + _logger.info("Triggering profile for %s ms", duration_ms) + response = stub.Profile(request, timeout=timeout) + _logger.info("Profiling response: %s", response.status) + if "Failed" in response.status: + return False + except grpc.RpcError as e: + e_call: Any = e + if e_call.code() == grpc.StatusCode.UNAVAILABLE: + _logger.error( + "Failed to connect to the profiling server at %s. " + "Please verify that the server is running on this port. " + "Note: If the server is running an older version of pathwaysutils, " + "it may be expecting HTTP (FastAPI) connections instead of gRPC, " + "which is incompatible with this client. " + "Error details: %s", + target, + e_call, + ) + elif e_call.code() == grpc.StatusCode.DEADLINE_EXCEEDED: + _logger.error( + "Profiling request timed out. The server might be unresponsive. " + "Error details: %s", + e_call, + ) + else: + _logger.error( + "gRPC error occurred while collecting profile. " + "Error code: %s, details: %s", + e_call.code(), + e_call.details(), + ) return False return True diff --git a/pathwaysutils/proto/pathways_profiler.proto b/pathwaysutils/proto/pathways_profiler.proto new file mode 100644 index 0000000..a1a6b8e --- /dev/null +++ b/pathwaysutils/proto/pathways_profiler.proto @@ -0,0 +1,46 @@ +// Copyright 2026 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. + +syntax = "proto3"; + +package pathwaysutils.profiler; + +option java_multiple_files = true; +option java_package = "com.google.pathwaysutils.profiler"; +option java_outer_classname = "PathwaysProfilerProto"; + +// PathwaysProfilerService provides an API to trigger and stop distributed +// profiling traces across Pathways workers. +service PathwaysProfilerService { + // Profiles the Pathways execution for the requested duration. + // The profiling data will be dumped to the specified repository path on GCS. + rpc Profile(ProfileRequest) returns (ProfileResponse) {} +} + +message ProfileRequest { + // Duration in milliseconds to collect the profile. + int64 duration_ms = 1; + + // The GCS repository path (e.g., gs://my-bucket/profiles) to save the + // profiling data. + string repository_path = 2; + + // Optional session ID to group traces. + string session_id = 3; +} + +message ProfileResponse { + // Status message indicating the result of the profiling operation. + string status = 1; +} diff --git a/pathwaysutils/test/profiling_test.py b/pathwaysutils/test/profiling_test.py index d6d926b..76f120c 100644 --- a/pathwaysutils/test/profiling_test.py +++ b/pathwaysutils/test/profiling_test.py @@ -19,10 +19,25 @@ from absl.testing import absltest from absl.testing import parameterized +import grpc import jax from jax import numpy as jnp from pathwaysutils import profiling -import requests +from pathwaysutils.proto import pathways_profiler_pb2 +from pathwaysutils.proto import pathways_profiler_pb2_grpc + + +class _MockRpcError(grpc.RpcError): + + def __init__(self, code, details=""): + self._code = code + self._details = details + + def code(self): + return self._code + + def details(self): + return self._details class ProfilingTest(parameterized.TestCase): @@ -30,12 +45,33 @@ class ProfilingTest(parameterized.TestCase): def setUp(self): super().setUp() - self.mock_post = self.enter_context( - mock.patch.object(requests, "post", autospec=True) + # Mock grpc channel and server + self.mock_secure_channel = self.enter_context( + mock.patch.object(grpc, "secure_channel", autospec=True) + ) + self.mock_grpc_server = self.enter_context( + mock.patch.object(grpc, "server", autospec=True) + ) + self.mock_alts_server_creds = self.enter_context( + mock.patch.object(grpc, "alts_server_credentials", autospec=True) + ) + self.mock_alts_channel_creds = self.enter_context( + mock.patch.object(grpc, "alts_channel_credentials", autospec=True) + ) + + # Mock the gRPC stub + self.mock_stub_cls = self.enter_context( + mock.patch.object( + pathways_profiler_pb2_grpc, + "PathwaysProfilerServiceStub", + ) ) + self.mock_stub = self.mock_stub_cls.return_value + profiling._profile_state.reset() profiling._first_profile_start = True - profiling._profiler_thread = None + profiling._profiler_server = None + self.mock_plugin_executable_cls = self.enter_context( mock.patch.object( profiling.plugin_executable, "PluginExecutable", autospec=True @@ -91,6 +127,9 @@ def _get_expected_profile_request( @parameterized.parameters(8000, 1234) def test_collect_profile_port(self, port): + self.mock_stub.Profile.return_value = pathways_profiler_pb2.ProfileResponse( + status="Completed" + ) result = profiling.collect_profile( port=port, duration_ms=1000, @@ -99,16 +138,22 @@ def test_collect_profile_port(self, port): ) self.assertTrue(result) - self.mock_post.assert_called_once_with( - f"http://127.0.0.1:{port}/profiling", - json={ - "duration_ms": 1000, - "repository_path": "gs://test_bucket/test_dir", - }, + self.mock_secure_channel.assert_called_once_with( + f"127.0.0.1:{port}", self.mock_alts_channel_creds.return_value + ) + self.mock_stub.Profile.assert_called_once_with( + pathways_profiler_pb2.ProfileRequest( + duration_ms=1000, + repository_path="gs://test_bucket/test_dir", + ), + timeout=11.0, ) @parameterized.parameters(1000, 1234) def test_collect_profile_duration_ms(self, duration_ms): + self.mock_stub.Profile.return_value = pathways_profiler_pb2.ProfileResponse( + status="Completed" + ) result = profiling.collect_profile( port=8000, duration_ms=duration_ms, @@ -117,16 +162,22 @@ def test_collect_profile_duration_ms(self, duration_ms): ) self.assertTrue(result) - self.mock_post.assert_called_once_with( - "http://127.0.0.1:8000/profiling", - json={ - "duration_ms": duration_ms, - "repository_path": "gs://test_bucket/test_dir", - }, + self.mock_secure_channel.assert_called_once_with( + "127.0.0.1:8000", self.mock_alts_channel_creds.return_value + ) + self.mock_stub.Profile.assert_called_once_with( + pathways_profiler_pb2.ProfileRequest( + duration_ms=duration_ms, + repository_path="gs://test_bucket/test_dir", + ), + timeout=(duration_ms / 1000.0) + 10.0, ) @parameterized.parameters("127.0.0.1", "localhost", "192.168.1.1") def test_collect_profile_host(self, host): + self.mock_stub.Profile.return_value = pathways_profiler_pb2.ProfileResponse( + status="Completed" + ) result = profiling.collect_profile( port=8000, duration_ms=1000, @@ -135,12 +186,15 @@ def test_collect_profile_host(self, host): ) self.assertTrue(result) - self.mock_post.assert_called_once_with( - f"http://{host}:8000/profiling", - json={ - "duration_ms": 1000, - "repository_path": "gs://test_bucket/test_dir", - }, + self.mock_secure_channel.assert_called_once_with( + f"{host}:8000", self.mock_alts_channel_creds.return_value + ) + self.mock_stub.Profile.assert_called_once_with( + pathways_profiler_pb2.ProfileRequest( + duration_ms=1000, + repository_path="gs://test_bucket/test_dir", + ), + timeout=11.0, ) @parameterized.parameters( @@ -149,17 +203,23 @@ def test_collect_profile_host(self, host): "gs://test_bucket3/test/log/dir", ) def test_collect_profile_log_dir(self, log_dir): + self.mock_stub.Profile.return_value = pathways_profiler_pb2.ProfileResponse( + status="Completed" + ) result = profiling.collect_profile( port=8000, duration_ms=1000, host="127.0.0.1", log_dir=log_dir ) self.assertTrue(result) - self.mock_post.assert_called_once_with( - "http://127.0.0.1:8000/profiling", - json={ - "duration_ms": 1000, - "repository_path": log_dir, - }, + self.mock_secure_channel.assert_called_once_with( + "127.0.0.1:8000", self.mock_alts_channel_creds.return_value + ) + self.mock_stub.Profile.assert_called_once_with( + pathways_profiler_pb2.ProfileRequest( + duration_ms=1000, + repository_path=log_dir, + ), + timeout=11.0, ) @parameterized.parameters("/logs/test_log_dir", "relative_path/my_log_dir") @@ -169,15 +229,38 @@ def test_collect_profile_log_dir_error(self, log_dir): port=8000, duration_ms=1000, host="127.0.0.1", log_dir=log_dir ) - @parameterized.parameters( - requests.exceptions.ConnectionError("Connection error"), - requests.exceptions.Timeout("Timeout"), - requests.exceptions.TooManyRedirects("Too many redirects"), - requests.exceptions.RequestException("Request exception"), - requests.exceptions.HTTPError("HTTP error"), + @parameterized.named_parameters( + dict( + testcase_name="unavailable", + status_code=grpc.StatusCode.UNAVAILABLE, + expected_log=( + "Failed to connect to the profiling server at 127.0.0.1:8000." + " Please verify that the server is running on this port. Note: If" + " the server is running an older version of pathwaysutils, it may" + " be expecting HTTP (FastAPI) connections instead of gRPC, which" + " is incompatible with this client. Error details:" + ), + ), + dict( + testcase_name="timeout", + status_code=grpc.StatusCode.DEADLINE_EXCEEDED, + expected_log=( + "Profiling request timed out. The server might be unresponsive." + ), + ), + dict( + testcase_name="other_error", + status_code=grpc.StatusCode.INTERNAL, + expected_log=( + "gRPC error occurred while collecting profile. " + "Error code: StatusCode.INTERNAL" + ), + ), ) - def test_collect_profile_request_error(self, exception): - self.mock_post.side_effect = exception + def test_collect_profile_request_error(self, status_code, expected_log): + self.mock_stub.Profile.side_effect = _MockRpcError( + status_code, "Error details" + ) with self.assertLogs(profiling._logger, level=logging.ERROR) as logs: result = profiling.collect_profile( @@ -188,26 +271,22 @@ def test_collect_profile_request_error(self, exception): ) self.assertLen(logs.output, 1) - self.assertIn("Failed to collect profiling data", logs.output[0]) - self.assertIn(str(exception), logs.output[0]) + self.assertIn(expected_log, logs.output[0]) self.assertFalse(result) - self.mock_post.assert_called_once() - - def test_collect_profile_success(self): - mock_response = mock.Mock() - mock_response.raise_for_status.return_value = None - self.mock_post.return_value = mock_response + self.mock_stub.Profile.assert_called_once() + def test_collect_profile_response_failed(self): + self.mock_stub.Profile.return_value = pathways_profiler_pb2.ProfileResponse( + status="Failed: some error" + ) result = profiling.collect_profile( port=8000, duration_ms=1000, host="127.0.0.1", log_dir="gs://test_bucket/test_dir", ) - - self.assertTrue(result) - self.mock_post.assert_called_once() - mock_response.raise_for_status.assert_called_once() + self.assertFalse(result) + self.mock_stub.Profile.assert_called_once() @parameterized.parameters( "/logs/test_log_dir", @@ -400,19 +479,18 @@ def test_stop_trace_before_start_error(self): ): profiling.stop_trace() - def test_start_server_starts_thread(self): - mock_thread = self.enter_context( - mock.patch.object(profiling.threading, "Thread", autospec=True) - ) + def test_start_server_starts_grpc_server(self): + mock_server = self.mock_grpc_server.return_value profiling.start_server(9000) - mock_thread.assert_called_once_with(target=mock.ANY, args=(9000,)) - mock_thread.return_value.start.assert_called_once() - self.assertIsNotNone(profiling._profiler_thread) - def test_start_server_twice_raises_error(self): - self.enter_context( - mock.patch.object(profiling.threading, "Thread", autospec=True) + self.mock_grpc_server.assert_called_once() + mock_server.add_secure_port.assert_called_once_with( + "[::]:9000", self.mock_alts_server_creds.return_value ) + mock_server.start.assert_called_once() + self.assertEqual(profiling._profiler_server, mock_server) + + def test_start_server_twice_raises_error(self): profiling.start_server(9000) with self.assertRaisesRegex( RuntimeError, "Only one profiler server can be active" @@ -423,12 +501,12 @@ def test_stop_server_no_server_raises_error(self): with self.assertRaisesRegex(RuntimeError, "No active profiler server"): profiling.stop_server() - def test_stop_server_does_nothing_if_server_exists(self): - self.enter_context( - mock.patch.object(profiling.threading, "Thread", autospec=True) - ) + def test_stop_server_stops_server(self): + mock_server = self.mock_grpc_server.return_value profiling.start_server(9000) - profiling.stop_server() # Should not raise + profiling.stop_server() + mock_server.stop.assert_called_once_with(grace=5.0) + self.assertIsNone(profiling._profiler_server) def _setup_monkey_patch(self): """Saves originals, applies monkey patch, and sets up mocks.""" @@ -705,6 +783,57 @@ def test_start_trace_compatibility_error(self): "gs://test_bucket/test_dir", profiler_options=options ) + @mock.patch.object(jax.profiler, "start_trace", autospec=True) + @mock.patch.object(jax.profiler, "stop_trace", autospec=True) + def test_servicer_profile_success(self, mock_stop_trace, mock_start_trace): + servicer = profiling.PathwaysProfilerServicer() + request = pathways_profiler_pb2.ProfileRequest( + duration_ms=100, repository_path="gs://test_bucket/test_dir" + ) + mock_context = mock.create_autospec(grpc.ServicerContext, instance=True) + mock_context.is_active.return_value = True + + response = servicer.Profile(request, mock_context) + + self.assertEqual(response.status, "Completed") + mock_start_trace.assert_called_once_with("gs://test_bucket/test_dir") + mock_stop_trace.assert_called_once() + + @mock.patch.object(jax.profiler, "start_trace", autospec=True) + @mock.patch.object(jax.profiler, "stop_trace", autospec=True) + def test_servicer_profile_failure(self, mock_stop_trace, mock_start_trace): + servicer = profiling.PathwaysProfilerServicer() + request = pathways_profiler_pb2.ProfileRequest( + duration_ms=100, repository_path="gs://test_bucket/test_dir" + ) + mock_context = mock.create_autospec(grpc.ServicerContext, instance=True) + mock_context.is_active.return_value = True + mock_start_trace.side_effect = RuntimeError("start failed") + + with mock.patch.object(profiling._logger, "exception"): + response = servicer.Profile(request, mock_context) + + self.assertIn("Failed: start failed", response.status) + mock_context.set_code.assert_called_once_with(grpc.StatusCode.INTERNAL) + mock_context.set_details.assert_called_once_with("start failed") + mock_stop_trace.assert_called_once() + + @mock.patch.object(jax.profiler, "start_trace", autospec=True) + @mock.patch.object(jax.profiler, "stop_trace", autospec=True) + def test_servicer_profile_cancelled(self, mock_stop_trace, mock_start_trace): + servicer = profiling.PathwaysProfilerServicer() + request = pathways_profiler_pb2.ProfileRequest( + duration_ms=500, repository_path="gs://test_bucket/test_dir" + ) + mock_context = mock.create_autospec(grpc.ServicerContext, instance=True) + mock_context.is_active.side_effect = [True, False] + + with mock.patch.object(profiling._logger, "exception"): + response = servicer.Profile(request, mock_context) + + self.assertIn("Failed: Client disconnected", response.status) + mock_start_trace.assert_called_once_with("gs://test_bucket/test_dir") + mock_stop_trace.assert_called_once() if __name__ == "__main__": absltest.main()