diff --git a/pathwaysutils/experimental/profiling.py b/pathwaysutils/experimental/profiling.py index 0816302..ec77d61 100644 --- a/pathwaysutils/experimental/profiling.py +++ b/pathwaysutils/experimental/profiling.py @@ -48,16 +48,17 @@ def start_trace( 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._validate_gcs_bucket(str(log_dir)) # pylint: disable=protected-access 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, - ) + try: + profiling._original_start_trace( # pylint: disable=protected-access + log_dir=log_dir, + create_perfetto_link=create_perfetto_link, + create_perfetto_trace=create_perfetto_trace, + ) + except Exception: + with profiling._profile_state.lock: # pylint: disable=protected-access + profiling._profile_state.reset() # pylint: disable=protected-access + raise diff --git a/pathwaysutils/profiling.py b/pathwaysutils/profiling.py index 1f773e3..c6f83b8 100644 --- a/pathwaysutils/profiling.py +++ b/pathwaysutils/profiling.py @@ -76,11 +76,7 @@ def call_profile_executable(self) -> None: ): out_avals = [jax.core.ShapedArray((1,), jnp.object_)] out_shardings = [ - getattr( - jax.sharding, - "make_single_device_sharding", - jax.sharding.SingleDeviceSharding, - )(backend.get_default_device()) + jax.sharding.SingleDeviceSharding(backend.get_default_device()) ] else: out_avals = () @@ -502,14 +498,31 @@ async def profiling(pc: ProfilingConfig) -> Mapping[str, str]: _logger.debug("Capturing profiling data for %s ms", pc.duration_ms) log_dir = pc.repository_path.strip() _logger.debug("Writing profiling data to %s", log_dir) - _validate_gcs_bucket(log_dir) + try: + _validate_gcs_bucket(log_dir) + except ValueError as e: + raise fastapi.HTTPException( + status_code=fastapi.status.HTTP_400_BAD_REQUEST, + detail=str(e), + ) from e options = None if jax.version.__version_info__ >= (0, 9, 2): options = jax.profiler.ProfileOptions() options.duration_ms = pc.duration_ms - await asyncio.to_thread(start_trace, log_dir, profiler_options=options) + try: + await asyncio.to_thread(start_trace, log_dir, profiler_options=options) + except ValueError as e: + raise fastapi.HTTPException( + status_code=fastapi.status.HTTP_400_BAD_REQUEST, + detail=str(e), + ) from e + except RuntimeError as e: + raise fastapi.HTTPException( + status_code=fastapi.status.HTTP_409_CONFLICT, + detail=str(e), + ) from e try: await asyncio.sleep(pc.duration_ms / 1e3) finally: diff --git a/pathwaysutils/test/profiling_test.py b/pathwaysutils/test/profiling_test.py index dd6a10f..3c75505 100644 --- a/pathwaysutils/test/profiling_test.py +++ b/pathwaysutils/test/profiling_test.py @@ -25,6 +25,7 @@ import jax from jax import numpy as jnp from pathwaysutils import profiling +from pathwaysutils.experimental import profiling as exp_profiling import requests @@ -1159,6 +1160,65 @@ def test_profiling_server_passes_duration_ms_to_start_trace(self): self.assertIsNotNone(call_kwargs.get("profiler_options")) self.assertEqual(call_kwargs["profiler_options"].duration_ms, 2500) + @unittest.skipIf( + os.environ.get("GITHUB_ACTIONS") == "true", + "Skipping FastAPI server test in GitHub CI", + ) + def test_profiling_server_error_status_codes(self): + from fastapi import testclient + + with ( + mock.patch.dict( + profiling.os.environ, + {"PATHWAYS_PROFILING_ALLOWED_GCS_BUCKETS": "gs://allowed"}, + clear=True, + ), + mock.patch.object(profiling, "start_trace") as mock_start_trace, + mock.patch.object(profiling, "stop_trace"), + mock.patch.object(profiling.asyncio, "sleep"), + ): + app = self._get_server_app() + client = testclient.TestClient(app) + + # Disallowed bucket -> 400 Bad Request + resp_bad_bucket = client.post( + "/profiling", + json={"duration_ms": 100, "repository_path": "gs://disallowed/dir"}, + ) + self.assertEqual(resp_bad_bucket.status_code, 400) + self.assertIn("not in allowed buckets list", resp_bad_bucket.json()["detail"]) + + # Concurrent trace -> 409 Conflict + mock_start_trace.side_effect = RuntimeError( + "start_trace called while a trace is already being taken!" + ) + resp_conflict = client.post( + "/profiling", + json={"duration_ms": 100, "repository_path": "gs://allowed/dir"}, + ) + self.assertEqual(resp_conflict.status_code, 409) + self.assertIn( + "trace is already being taken", resp_conflict.json()["detail"] + ) + + def test_experimental_start_trace_validates_allowed_bucket_and_rolls_back( + self, + ): + with mock.patch.dict( + profiling.os.environ, + {"PATHWAYS_PROFILING_ALLOWED_GCS_BUCKETS": "gs://allowed_bucket"}, + ): + with self.assertRaisesRegex(ValueError, "is not in allowed buckets list"): + exp_profiling.start_trace({"traceLocation": "gs://other_bucket/dir"}) + + self.mock_original_start_trace.side_effect = RuntimeError( + "original start trace error" + ) + with self.assertRaisesRegex(RuntimeError, "original start trace error"): + exp_profiling.start_trace({"traceLocation": "gs://test_bucket/dir"}) + self.assertIsNone(profiling._profile_state.executable) + self.assertFalse(profiling._profile_state.lock.locked()) + if __name__ == "__main__": absltest.main()