diff --git a/pathwaysutils/collect_profile.py b/pathwaysutils/collect_profile.py index 57c971e..a3da3e2 100644 --- a/pathwaysutils/collect_profile.py +++ b/pathwaysutils/collect_profile.py @@ -58,6 +58,12 @@ def _get_parser() -> argparse.ArgumentParser: ), type=str, ) + parser.add_argument( + "--auth_token", + default=None, + help="Authentication token for accessing the profiling server.", + type=str, + ) return parser @@ -67,7 +73,11 @@ def main() -> None: args = parser.parse_args() if profiling.collect_profile( - args.port, args.duration_ms, args.host, args.log_dir + args.port, + args.duration_ms, + args.host, + args.log_dir, + auth_token=args.auth_token, ): _logger.info("Dumped profiling information in: %s", args.log_dir) else: diff --git a/pathwaysutils/profiling.py b/pathwaysutils/profiling.py index 2a39b85..a42cd83 100644 --- a/pathwaysutils/profiling.py +++ b/pathwaysutils/profiling.py @@ -225,6 +225,25 @@ def _start_pathways_trace_from_profile_request( raise +def _validate_gcs_bucket(log_dir: os.PathLike[str] | str) -> None: + """Validates that log_dir is a valid GCS path and optionally in allowed buckets.""" + log_dir_str = str(log_dir) + if not log_dir_str.startswith("gs://"): + raise ValueError(f"log_dir must be a GCS bucket path, got {log_dir}") + + env_buckets = os.environ.get("PATHWAYS_ALLOWED_GCS_BUCKETS") + if env_buckets: + allowed_buckets = [b.strip() for b in env_buckets.split(",") if b.strip()] + allowed_set = set(allowed_buckets) + path_without_scheme = log_dir_str[5:] + bucket_name = path_without_scheme.split("/")[0] + if not bucket_name or bucket_name not in allowed_set: + raise ValueError( + f"GCS bucket '{bucket_name}' is not in allowed buckets list: " + f"{allowed_set}" + ) + + def start_trace( log_dir: os.PathLike[str] | str, *, @@ -263,8 +282,7 @@ def start_trace( max_num_hosts: An optional integer to limit the number of hosts profiled (defaults to 1). """ - if not str(log_dir).startswith("gs://"): - raise ValueError(f"log_dir must be a GCS bucket path, got {log_dir}") + _validate_gcs_bucket(log_dir) if create_perfetto_link or create_perfetto_trace: _logger.warning( @@ -297,19 +315,24 @@ 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: + with _profile_state.lock: + _profile_state.reset() + raise def stop_trace() -> None: @@ -341,8 +364,11 @@ def start_server(port: int, requires_backend: bool = True) -> None: requires_backend: Unused in Pathways; accepted for parameter parity. """ del requires_backend - def server_loop(port: int): - _logger.debug("Starting JAX profiler server on port %s", port) + host = os.environ.get("PATHWAYS_PROFILING_SERVER_HOST", "127.0.0.1") + token_to_verify = os.environ.get("PATHWAYS_PROFILING_AUTH_TOKEN") + + def server_loop(port: int, host: str, token_to_verify: str | None): + _logger.debug("Starting JAX profiler server on host %s port %s", host, port) app = fastapi.FastAPI() @dataclasses.dataclass @@ -351,21 +377,43 @@ class ProfilingConfig: repository_path: str @app.post("/profiling") - async def profiling(pc: ProfilingConfig) -> Mapping[str, str]: + async def profiling( + pc: ProfilingConfig, + req: fastapi.Request, + ) -> Mapping[str, str]: + if token_to_verify is not None: + token = req.headers.get("X-Auth-Token") or req.headers.get( + "Authorization" + ) + if token and token.startswith("Bearer "): + token = token[7:] + if token != token_to_verify: + raise fastapi.HTTPException( + status_code=fastapi.status.HTTP_401_UNAUTHORIZED, + detail="Unauthorized: invalid or missing authentication token", + ) _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) + _validate_gcs_bucket(pc.repository_path) + + await asyncio.to_thread(start_trace, pc.repository_path) + try: + await asyncio.sleep(pc.duration_ms / 1e3) + finally: + await asyncio.to_thread(stop_trace) + return {"response": "profiling completed"} - uvicorn.run(app, host="0.0.0.0", port=port, log_level="debug") + uvicorn.run(app, host=host, 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 = threading.Thread( + target=server_loop, + args=(port, host, token_to_verify), + ) _profiler_thread.start() @@ -383,6 +431,7 @@ def collect_profile( duration_ms: int, host: str, log_dir: os.PathLike[str] | str, + auth_token: str | None = None, ) -> bool: """Collects a JAX profile and saves it to the specified directory. @@ -391,6 +440,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. + auth_token: Optional authentication token to send to the profiling server. Returns: True if the profile was collected successfully, False otherwise. @@ -398,16 +448,24 @@ def collect_profile( Raises: ValueError: If the log_dir is not a GCS path. """ - if not str(log_dir).startswith("gs://"): - raise ValueError(f"log_dir must be a GCS bucket path, got {log_dir}") + _validate_gcs_bucket(log_dir) request_json = { "duration_ms": duration_ms, "repository_path": log_dir, } + headers = {} + effective_token = ( + auth_token + if auth_token is not None + else os.environ.get("PATHWAYS_PROFILING_AUTH_TOKEN") + ) + if effective_token: + headers["X-Auth-Token"] = effective_token + address = urllib.parse.urljoin(f"http://{host}:{port}", "profiling") try: - response = requests.post(address, json=request_json) + response = requests.post(address, json=request_json, headers=headers) response.raise_for_status() except requests.exceptions.RequestException: _logger.exception("Failed to collect profiling data") diff --git a/pathwaysutils/proxy_backend.py b/pathwaysutils/proxy_backend.py index cf1f806..2507a8a 100644 --- a/pathwaysutils/proxy_backend.py +++ b/pathwaysutils/proxy_backend.py @@ -13,17 +13,27 @@ # limitations under the License. """Register the IFRT Proxy as a backend for JAX.""" +import os import jax from jax.extend import backend from jax.extend.backend import ifrt_proxy def register_backend_factory() -> None: + """Registers the IFRT Proxy backend factory with JAX.""" + + def make_client(): + options = ifrt_proxy.ClientConnectionOptions() + timeout_secs = os.environ.get("PATHWAYS_PROXY_CONNECTION_TIMEOUT_SECS") + if timeout_secs: + options.connection_timeout_in_seconds = int(timeout_secs) + return ifrt_proxy.get_client( + jax.config.read("jax_backend_target"), + options, + ) + backend.register_backend_factory( "proxy", - lambda: ifrt_proxy.get_client( - jax.config.read("jax_backend_target"), - ifrt_proxy.ClientConnectionOptions(), - ), + make_client, priority=-1, ) diff --git a/pathwaysutils/test/profiling_test.py b/pathwaysutils/test/profiling_test.py index 0909919..21ee0b1 100644 --- a/pathwaysutils/test/profiling_test.py +++ b/pathwaysutils/test/profiling_test.py @@ -14,8 +14,8 @@ 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 @@ -105,6 +105,7 @@ def test_collect_profile_port(self, port): "duration_ms": 1000, "repository_path": "gs://test_bucket/test_dir", }, + headers={}, ) @parameterized.parameters(1000, 1234) @@ -123,6 +124,7 @@ def test_collect_profile_duration_ms(self, duration_ms): "duration_ms": duration_ms, "repository_path": "gs://test_bucket/test_dir", }, + headers={}, ) @parameterized.parameters("127.0.0.1", "localhost", "192.168.1.1") @@ -141,6 +143,7 @@ def test_collect_profile_host(self, host): "duration_ms": 1000, "repository_path": "gs://test_bucket/test_dir", }, + headers={}, ) @parameterized.parameters( @@ -160,6 +163,7 @@ def test_collect_profile_log_dir(self, log_dir): "duration_ms": 1000, "repository_path": log_dir, }, + headers={}, ) @parameterized.parameters("/logs/test_log_dir", "relative_path/my_log_dir") @@ -405,7 +409,9 @@ def test_start_server_starts_thread(self): mock.patch.object(profiling.threading, "Thread", autospec=True) ) profiling.start_server(9000) - mock_thread.assert_called_once_with(target=mock.ANY, args=(9000,)) + mock_thread.assert_called_once_with( + target=mock.ANY, args=(9000, "127.0.0.1", None) + ) mock_thread.return_value.start.assert_called_once() self.assertIsNotNone(profiling._profiler_thread) @@ -526,7 +532,10 @@ def test_monkey_patched_start_server(self, profiler_module): profiler_module.start_server(1234, requires_backend=False) - mocks["start_server"].assert_called_once_with(1234, requires_backend=False) + mocks["start_server"].assert_called_once_with( + 1234, + requires_backend=False, + ) @parameterized.named_parameters( dict(testcase_name="jax_profiler", profiler_module=jax.profiler), @@ -715,6 +724,47 @@ def test_start_trace_compatibility_error(self): "gs://test_bucket/test_dir", profiler_options=options ) + def test_validate_gcs_bucket_env_var(self): + with mock.patch.dict( + profiling.os.environ, + {"PATHWAYS_ALLOWED_GCS_BUCKETS": "bucket1,bucket2"}, + ): + profiling._validate_gcs_bucket("gs://bucket1/dir") + with self.assertRaisesRegex(ValueError, "is not in allowed buckets list"): + profiling._validate_gcs_bucket("gs://bucket3/dir") + + def test_start_trace_rollback_on_original_failure(self): + self.mock_original_start_trace.side_effect = RuntimeError( + "original start trace error" + ) + with self.assertRaisesRegex(RuntimeError, "original start trace error"): + profiling.start_trace("gs://test_bucket/test_dir") + + self.assertIsNone(profiling._profile_state.executable) + self.assertFalse(profiling._profile_state.lock.locked()) + + def test_collect_profile_with_auth_token(self): + with mock.patch.dict( + profiling.os.environ, + {"PATHWAYS_PROFILING_AUTH_TOKEN": "secret_token"}, + ): + 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_with( + "http://127.0.0.1:8000/profiling", + json={ + "duration_ms": 1000, + "repository_path": "gs://test_bucket/test_dir", + }, + headers={"X-Auth-Token": "secret_token"}, + ) + if __name__ == "__main__": absltest.main() diff --git a/pathwaysutils/test/proxy_backend_test.py b/pathwaysutils/test/proxy_backend_test.py index fb8ad8c..13fce79 100644 --- a/pathwaysutils/test/proxy_backend_test.py +++ b/pathwaysutils/test/proxy_backend_test.py @@ -13,6 +13,7 @@ # limitations under the License. """Tests for the proxy backend module.""" +import os from unittest import mock from absl.testing import absltest @@ -54,6 +55,46 @@ def test_proxy_backend_registration(self): proxy_backend.register_backend_factory() self.assertIn("proxy", backend.backends()) + def test_proxy_backend_registration_with_timeout(self): + mock_get_client = self.enter_context( + mock.patch.object( + ifrt_proxy, + "get_client", + return_value=mock.MagicMock(), + ) + ) + self.enter_context( + mock.patch.dict( + os.environ, {"PATHWAYS_PROXY_CONNECTION_TIMEOUT_SECS": "42"} + ) + ) + proxy_backend.register_backend_factory() + self.assertIn("proxy", backend.backends()) + mock_get_client.assert_called_once() + args, _ = mock_get_client.call_args + self.assertEqual(args[0], "grpc://localhost:12345") + options = args[1] + self.assertEqual(options.connection_timeout_in_seconds, 42) + + def test_proxy_backend_registration_without_timeout(self): + mock_get_client = self.enter_context( + mock.patch.object( + ifrt_proxy, + "get_client", + return_value=mock.MagicMock(), + ) + ) + self.enter_context(mock.patch.dict(os.environ)) + os.environ.pop("PATHWAYS_PROXY_CONNECTION_TIMEOUT_SECS", None) + + proxy_backend.register_backend_factory() + self.assertIn("proxy", backend.backends()) + mock_get_client.assert_called_once() + args, _ = mock_get_client.call_args + self.assertEqual(args[0], "grpc://localhost:12345") + options = args[1] + self.assertNotEqual(options.connection_timeout_in_seconds, 42) + if __name__ == "__main__": absltest.main()