Replace time.sleep with asyncio.sleep in profiling server. Remove a log line that is emitted unnecessarily in unittests.

PiperOrigin-RevId: 878638182

Author: lukebaumann

pathwaysutils/__init__.py

@@ -20,4 +20,4 @@
 del _initialize
 
 # When changing this, also update the CHANGELOG.md.
-__version__ = "v0.1.6"
+__version__ = "v0.1.5"

pathwaysutils/profiling.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 """Profiling Utilities."""
 
+import asyncio
 import dataclasses
 import json
 import logging
@@ -93,7 +94,7 @@ def _start_pathways_trace_from_profile_request(
     try:
       _, result_future = _profile_state.executable.call()
       result_future.result()
-    except Exception as e:  # pylint: disable=broad-except
+    except Exception:  # pylint: disable=broad-except
       _logger.exception("Failed to start trace")
       _profile_state.reset()
       raise
@@ -194,9 +195,9 @@ class ProfilingConfig:
     async def profiling(pc: ProfilingConfig):  # pylint: disable=unused-variable
       _logger.debug("Capturing profiling data for %s ms", pc.duration_ms)
       _logger.debug("Writing profiling data to %s", pc.repository_path)
-      jax.profiler.start_trace(pc.repository_path)
-      time.sleep(pc.duration_ms / 1e3)
-      jax.profiler.stop_trace()
+      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")
@@ -210,7 +211,7 @@ async def profiling(pc: ProfilingConfig):  # pylint: disable=unused-variable
 
 
 def stop_server():
-  """Raises an error if there is not an active profiler server but otherwise does nothing.
+  """Raises an error if there is no active profiler server.
 
   Pathways profiling servers are not stoppable at this time.
   """

pathwaysutils/test/profiling_test.py

@@ -186,7 +186,11 @@ def test_start_trace_log_dir_error(self, log_dir):
       profiling.start_trace(log_dir)
 
   def test_lock_released_on_success(self):
-    """Tests that the lock is released after successful start_trace and stop_trace."""
+    """Tests lock release after successful start and stop trace.
+
+    Verifies that the profiling lock is released after both a successful
+    `start_trace` and `stop_trace` calls.
+    """
     profiling.start_trace("gs://test_bucket/test_dir")
     self.assertFalse(profiling._profile_state.lock.locked())
     profiling.stop_trace()
@@ -198,7 +202,9 @@ def test_lock_released_on_start_failure(self):
         self.mock_plugin_executable_cls.return_value.call.return_value[1]
     )
     mock_result.result.side_effect = RuntimeError("start failed")
-    with self.assertRaisesRegex(RuntimeError, "start failed"):
+    with self.assertRaisesRegex(
+        RuntimeError, "start failed"
+    ), mock.patch.object(profiling._logger, "exception"):
       profiling.start_trace("gs://test_bucket/test_dir2")
     self.assertFalse(profiling._profile_state.lock.locked())