From 6107386c13244867215727130e0a4467424c1f59 Mon Sep 17 00:00:00 2001 From: Pathways-on-Cloud Team Date: Mon, 29 Jun 2026 06:04:59 +0000 Subject: [PATCH] Secure Pathways Profiling API using gRPC PiperOrigin-RevId: 939630828 --- pathwaysutils/profiling.py | 159 ++++++++++++------ pathwaysutils/test/profiling_test.py | 236 +++++++++++++++++++-------- 2 files changed, 280 insertions(+), 115 deletions(-) diff --git a/pathwaysutils/profiling.py b/pathwaysutils/profiling.py index 2a39b85..bc739d2 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 tensorflow.core.profiler.protobuf import profiler_service_pb2 # pylint: disable=g-direct-tensorflow-import +from tensorflow.core.profiler.protobuf import profiler_service_pb2_grpc # pylint: disable=g-direct-tensorflow-import _logger = logging.getLogger(__name__) @@ -326,10 +325,39 @@ 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, requires_backend: bool = True) -> None: +class PathwaysProfilerServicer( + profiler_service_pb2_grpc.ProfilerServiceServicer +): + """gRPC servicer for Pathways Profiler Service implementing tensorflow.ProfilerService.""" + + def Profile( + self, + request: profiler_service_pb2.ProfileRequest, + context: grpc.ServicerContext, + ) -> profiler_service_pb2.ProfileResponse: + del context + duration_ms = request.duration_ms + log_dir = request.repository_root + _logger.info("Received gRPC profile request for %s ms", duration_ms) + _logger.info("Writing profiling data to %s", log_dir) + + try: + jax.profiler.start_trace(log_dir) + time.sleep(duration_ms / 1000.0) + finally: + _logger.info("Stopping trace") + jax.profiler.stop_trace() + + return profiler_service_pb2.ProfileResponse() + + +def start_server( + port: int, requires_backend: bool = True, use_alts: bool | None = None +) -> None: """Starts the profiling server on port `port`. The signature matches `jax.profiler.start_server`, though no handle @@ -339,43 +367,41 @@ def start_server(port: int, requires_backend: bool = True) -> None: Args: port: The port to start the server on. requires_backend: Unused in Pathways; accepted for parameter parity. + use_alts: Whether to use ALTS credentials. If None, defaults to checking + the PATHWAYS_PROFILER_USE_ALTS environment variable (defaulting to True). """ del requires_backend - 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 - - @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"} - - 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.") - - _profiler_thread = threading.Thread(target=server_loop, args=(port,)) - _profiler_thread.start() + 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.") + + if use_alts is None: + use_alts = os.environ.get("PATHWAYS_PROFILER_USE_ALTS", "1") == "1" + + _logger.info("Starting JAX pathways profiler gRPC server on port %s", port) + server = grpc.server(concurrent.futures.ThreadPoolExecutor(max_workers=2)) + profiler_service_pb2_grpc.add_ProfilerServiceServicer_to_server( + PathwaysProfilerServicer(), server + ) + if use_alts: + server_creds = grpc.alts_server_credentials() + server.add_secure_port(f"[::]:{port}", server_creds) + else: + server.add_insecure_port(f"[::]:{port}") + 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( @@ -383,6 +409,7 @@ def collect_profile( duration_ms: int, host: str, log_dir: os.PathLike[str] | str, + use_alts: bool = True, ) -> bool: """Collects a JAX profile and saves it to the specified directory. @@ -391,6 +418,7 @@ def collect_profile( duration_ms: The duration in milliseconds for which to collect the profile. host: The host on which the JAX profiler server is running. log_dir: The GCS path to save the profile data. + use_alts: Whether to connect using ALTS credentials. Returns: True if the profile was collected successfully, False otherwise. @@ -401,16 +429,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") + target = f"{host}:{port}" + _logger.info("Connecting to profiling server at %s", target) try: - response = requests.post(address, json=request_json) - response.raise_for_status() - except requests.exceptions.RequestException: - _logger.exception("Failed to collect profiling data") + if use_alts: + creds = grpc.alts_channel_credentials() + channel = grpc.secure_channel(target, creds) + else: + channel = grpc.insecure_channel(target) + + with channel: + stub = profiler_service_pb2_grpc.ProfilerServiceStub(channel) + request = profiler_service_pb2.ProfileRequest( + duration_ms=duration_ms, + repository_root=str(log_dir), + ) + timeout = (duration_ms / 1000.0) + 10.0 + _logger.info("Triggering profile for %s ms", duration_ms) + stub.Profile(request, timeout=timeout) + _logger.info("Profiling response completed successfully") + return True + 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. " + "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/test/profiling_test.py b/pathwaysutils/test/profiling_test.py index 0909919..d5a41ef 100644 --- a/pathwaysutils/test/profiling_test.py +++ b/pathwaysutils/test/profiling_test.py @@ -14,15 +14,30 @@ import json import logging -from unittest import mock from typing import Any +from unittest import mock 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 tensorflow.core.profiler.protobuf import profiler_service_pb2 # pylint: disable=g-direct-tensorflow-import +from tensorflow.core.profiler.protobuf import profiler_service_pb2_grpc # pylint: disable=g-direct-tensorflow-import + + +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( + profiler_service_pb2_grpc, + "ProfilerServiceStub", + ) ) + 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,7 @@ def _get_expected_profile_request( @parameterized.parameters(8000, 1234) def test_collect_profile_port(self, port): + self.mock_stub.Profile.return_value = profiler_service_pb2.ProfileResponse() result = profiling.collect_profile( port=port, duration_ms=1000, @@ -99,16 +136,17 @@ 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_stub.Profile.assert_called_once_with( + profiler_service_pb2.ProfileRequest( + duration_ms=1000, + repository_root="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 = profiler_service_pb2.ProfileResponse() result = profiling.collect_profile( port=8000, duration_ms=duration_ms, @@ -117,16 +155,17 @@ 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_stub.Profile.assert_called_once_with( + profiler_service_pb2.ProfileRequest( + duration_ms=duration_ms, + repository_root="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 = profiler_service_pb2.ProfileResponse() result = profiling.collect_profile( port=8000, duration_ms=1000, @@ -135,12 +174,12 @@ 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_stub.Profile.assert_called_once_with( + profiler_service_pb2.ProfileRequest( + duration_ms=1000, + repository_root="gs://test_bucket/test_dir", + ), + timeout=11.0, ) @parameterized.parameters( @@ -149,17 +188,18 @@ 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 = profiler_service_pb2.ProfileResponse() 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_stub.Profile.assert_called_once_with( + profiler_service_pb2.ProfileRequest( + duration_ms=1000, + repository_root=log_dir, + ), + timeout=11.0, ) @parameterized.parameters("/logs/test_log_dir", "relative_path/my_log_dir") @@ -169,15 +209,34 @@ 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." + ), + ), + 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 +247,24 @@ 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_use_alts(self): + self.mock_stub.Profile.return_value = profiler_service_pb2.ProfileResponse() result = profiling.collect_profile( port=8000, duration_ms=1000, host="127.0.0.1", log_dir="gs://test_bucket/test_dir", + use_alts=True, ) - self.assertTrue(result) - self.mock_post.assert_called_once() - mock_response.raise_for_status.assert_called_once() + 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() @parameterized.parameters( "/logs/test_log_dir", @@ -309,10 +366,14 @@ def test_start_trace_with_max_num_hosts(self): 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) + profiling.start_trace( + "gs://test_bucket/test_dir", profiler_options=options + ) expected_request = self._get_expected_profile_request( - "gs://test_bucket/test_dir", max_num_hosts=1, session_id="options_session" + "gs://test_bucket/test_dir", + max_num_hosts=1, + session_id="options_session", ) self.mock_plugin_executable_cls.assert_called_once_with( json.dumps(expected_request) @@ -323,7 +384,9 @@ def test_start_trace_with_session_id_in_options(self): 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") + 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") @@ -400,19 +463,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 +485,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.""" @@ -715,6 +777,48 @@ def test_start_trace_compatibility_error(self): "gs://test_bucket/test_dir", profiler_options=options ) + def test_servicer_profile_success(self): + with ( + mock.patch.object( + jax.profiler, "start_trace", autospec=True + ) as mock_start_trace, + mock.patch.object( + jax.profiler, "stop_trace", autospec=True + ) as mock_stop_trace, + ): + servicer = profiling.PathwaysProfilerServicer() + request = profiler_service_pb2.ProfileRequest( + duration_ms=100, repository_root="gs://test_bucket/test_dir" + ) + mock_context = mock.create_autospec(grpc.ServicerContext, instance=True) + + response = servicer.Profile(request, mock_context) + + self.assertIsInstance(response, profiler_service_pb2.ProfileResponse) + mock_start_trace.assert_called_once_with("gs://test_bucket/test_dir") + mock_stop_trace.assert_called_once() + + def test_servicer_profile_failure(self): + with ( + mock.patch.object( + jax.profiler, "start_trace", autospec=True + ) as mock_start_trace, + mock.patch.object( + jax.profiler, "stop_trace", autospec=True + ) as mock_stop_trace, + ): + servicer = profiling.PathwaysProfilerServicer() + request = profiler_service_pb2.ProfileRequest( + duration_ms=100, repository_root="gs://test_bucket/test_dir" + ) + mock_context = mock.create_autospec(grpc.ServicerContext, instance=True) + mock_start_trace.side_effect = RuntimeError("start failed") + + with self.assertRaisesRegex(RuntimeError, "start failed"): + servicer.Profile(request, mock_context) + + mock_stop_trace.assert_called_once() + if __name__ == "__main__": absltest.main()