Port the collect_profile script from JAX to PathwaysUtils

PiperOrigin-RevId: 746578297

Author: lukebaumann

pathwaysutils/collect_profile.py

@@ -0,0 +1,66 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Module for collecting JAX profiles for Pathways on Cloud.
+
+This is a replacement for the `collect_profile` script in JAX that works with
+Pathways on Cloud.
+"""
+
+import argparse
+import logging
+
+from pathwaysutils import profiling
+
+_logger = logging.getLogger(__name__)
+
+
+_DESCRIPTION = """
+To profile running JAX programs, you first need to start the profiler server
+in the program of interest. You can do this via
+`jax.profiler.start_server(<port>)`. Once the program is running and the
+profiler server has started, you can run `collect_profile` to trace the execution
+for a provided duration. The trace file will be dumped into a directory
+(determined by `--log_dir`).
+(determined by `--log_dir`).
+"""
+parser = argparse.ArgumentParser(description=_DESCRIPTION)
+parser.add_argument(
+    "--log_dir",
+    default=None,
+    help=(
+        "Directory to store log files. "
+        "Uses a temporary directory if none provided."
+    ),
+    type=str,
+)
+parser.add_argument("port", help="Port to collect trace", type=int)
+parser.add_argument(
+    "duration_ms", help="Duration to collect trace in milliseconds", type=int
+)
+parser.add_argument(
+    "--host",
+    default="127.0.0.1",
+    help="Host to collect trace. Defaults to 127.0.0.1",
+    type=str,
+)
+
+
+def main(args):
+  profiling.collect_profile(
+      args.port, args.duration_ms, args.host, args.log_dir
+  )
+
+
+if __name__ == "__main__":
+  main(parser.parse_args())

pathwaysutils/profiling.py

@@ -15,19 +15,28 @@
 
 import dataclasses
 import logging
+import os
+import pathlib
+import tempfile
 import threading
 import time
+import urllib.parse
 
 import fastapi
 import jax
 from jax import numpy as jnp
 from pathwaysutils import plugin_executable
+import requests
 import uvicorn
 
-logger = logging.getLogger(__name__)
+
+_logger = logging.getLogger(__name__)
 
 
 class _ProfileState:
+  executable: plugin_executable.PluginExecutable | None = None
+  lock: threading.Lock
+
   def __init__(self):
     self.executable = None
     self.lock = threading.Lock()
@@ -88,7 +97,7 @@ def stop_trace():
   _original_stop_trace()
 
 
-_profiler_thread = None
+_profiler_thread: threading.Thread | None = None
 
 
 def start_server(port: int):
@@ -102,7 +111,7 @@ def start_server(port: int):
     port : The port to start the server on.
   """
   def server_loop(port: int):
-    logger.debug("Starting JAX profiler server on port %s", port)
+    _logger.debug("Starting JAX profiler server on port %s", port)
     app = fastapi.FastAPI()
 
     @dataclasses.dataclass
@@ -112,8 +121,8 @@ class ProfilingConfig:
 
     @app.post("/profiling")
     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)
+      _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()
@@ -138,6 +147,39 @@ def stop_server():
     raise ValueError("No active profiler server.")
 
 
+def collect_profile(
+    port: int,
+    duration_ms: int,
+    host: str,
+    log_dir: os.PathLike[str] | str | None,
+):
+  """Collects a JAX profile and saves it to the specified directory.
+
+  Args:
+    port: The port on which the JAX profiler server is running.
+    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 directory to save the profile data. If None, a temporary
+      directory will be used.
+  """
+  if log_dir is None:
+    log_dir = tempfile.mkdtemp()
+  prefix = ""
+  if isinstance(log_dir, str) and log_dir.startswith("gs://"):
+    prefix = "gs://"
+    log_dir = log_dir[5:]
+
+  log_dir = prefix + str(pathlib.Path(log_dir))
+
+  json = {
+      "duration_ms": duration_ms,
+      "repository_path": str(log_dir),
+  }
+  address = urllib.parse.urljoin(f"http://{host}:{port}", "profiling")
+  requests.post(address, json=json)
+  _logger.info("Dumped profiling information in: %s", log_dir)
+
+
 def monkey_patch_jax():
   """Monkey patches JAX with Pathways versions of functions.
 
@@ -158,25 +200,27 @@ def start_trace_patch(
       create_perfetto_link: bool = False,  # pylint: disable=unused-argument
       create_perfetto_trace: bool = False,  # pylint: disable=unused-argument
   ) -> None:
-    logger.debug("jax.profile.start_trace patched with pathways' start_trace")
+    _logger.debug("jax.profile.start_trace patched with pathways' start_trace")
     return start_trace(log_dir)
 
   jax.profiler.start_trace = start_trace_patch
 
   def stop_trace_patch() -> None:
-    logger.debug("jax.profile.stop_trace patched with pathways' stop_trace")
+    _logger.debug("jax.profile.stop_trace patched with pathways' stop_trace")
     return stop_trace()
 
   jax.profiler.stop_trace = stop_trace_patch
 
   def start_server_patch(port: int):
-    logger.debug("jax.profile.start_server patched with pathways' start_server")
+    _logger.debug(
+        "jax.profile.start_server patched with pathways' start_server"
+    )
     return start_server(port)
 
   jax.profiler.start_server = start_server_patch
 
   def stop_server_patch():
-    logger.debug("jax.profile.stop_server patched with pathways' stop_server")
+    _logger.debug("jax.profile.stop_server patched with pathways' stop_server")
     return stop_server()
 
   jax.profiler.stop_server = stop_server_patch

pathwaysutils/test/profiling_test.py

@@ -0,0 +1,125 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pathlib
+from unittest import mock
+
+from pathwaysutils import profiling
+import requests
+
+from absl.testing import absltest
+from absl.testing import parameterized
+
+
+class ProfilingTest(parameterized.TestCase):
+  """Tests for Pathways on Cloud profiling."""
+
+  def setUp(self):
+    super().setUp()
+    self.mock_post = self.enter_context(
+        mock.patch.object(requests, "post", autospec=True)
+    )
+
+  @parameterized.parameters(
+      (8000,),
+      (1234,),
+  )
+  def test_collect_profile_port(self, port):
+    profiling.collect_profile(
+        port=port,
+        duration_ms=1000,
+        host="127.0.0.1",
+        log_dir="/tmp/test_log_dir",
+    )
+
+    self.mock_post.assert_called_once_with(
+        f"http://127.0.0.1:{port}/profiling",
+        json={
+            "duration_ms": 1000,
+            "repository_path": "/tmp/test_log_dir",
+        },
+    )
+
+  @parameterized.parameters(
+      (1000,),
+      (1234,),
+  )
+  def test_collect_profile_duration_ms(self, duration_ms):
+    profiling.collect_profile(
+        port=8000,
+        duration_ms=duration_ms,
+        host="127.0.0.1",
+        log_dir="/tmp/test_log_dir",
+    )
+
+    self.mock_post.assert_called_once_with(
+        "http://127.0.0.1:8000/profiling",
+        json={
+            "duration_ms": duration_ms,
+            "repository_path": "/tmp/test_log_dir",
+        },
+    )
+
+  @parameterized.parameters(
+      ("127.0.0.1",),
+      ("localhost",),
+      ("192.168.1.1",),
+  )
+  def test_collect_profile_host(self, host):
+    profiling.collect_profile(
+        port=8000, duration_ms=1000, host=host, log_dir="/tmp/test_log_dir"
+    )
+
+    self.mock_post.assert_called_once_with(
+        f"http://{host}:8000/profiling",
+        json={
+            "duration_ms": 1000,
+            "repository_path": "/tmp/test_log_dir",
+        },
+    )
+
+  @parameterized.parameters(
+      ("gs://test_bucket/test_log_dir",),
+      ("/logs/test_log_dir",),
+      ("relative_path/my_log_dir",),
+      (pathlib.Path("/tmp/test_log_dir"),),
+  )
+  def test_collect_profile_log_dir(self, log_dir):
+    profiling.collect_profile(
+        port=8000, duration_ms=1000, host="127.0.0.1", log_dir=log_dir
+    )
+
+    self.mock_post.assert_called_once_with(
+        "http://127.0.0.1:8000/profiling",
+        json={
+            "duration_ms": 1000,
+            "repository_path": str(log_dir),
+        },
+    )
+
+  def test_collect_profile_no_log_dir(self):
+    port = 8000
+    duration_ms = 1000
+    host = "127.0.0.1"
+
+    profiling.collect_profile(port, duration_ms, host, None)
+
+    self.mock_post.assert_called_once()
+    _, mock_kwargs = self.mock_post.call_args
+    self.assertIn("repository_path", mock_kwargs["json"])
+    self.assertStartsWith(mock_kwargs["json"]["repository_path"], "/tmp/")
+
+
+if __name__ == "__main__":
+  absltest.main()