Fix profiler state cleanup and bucket validation in profiling utilities.

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

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:

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()