Merge branch 'main' into elastic-readme

Author: lukebaumann

CHANGELOG.md

@@ -23,6 +23,13 @@ To release a new version (e.g. from `1.0.0` -> `2.0.0`):
 
 ## [Unreleased]
 
+## [0.1.1] - 2025-04-25
+
+* Port the `collect_profile` script from JAX to PathwaysUtils
+* Remove support for legacy initialize
+* Add collect_profile as a script of pathwaysutils
+* Make CloudPathwaysArrayHandler compatible with async directory creation feature in orbax
+
 ## [0.1.0] - 2025-04-07
 * Bump the JAX requirement to 0.5.1
 * Introduce `pathwaysutils.initialize()` to remove relying on side-effects from `import pathwaysutils`. by @copybara-service in https://github.com/AI-Hypercomputer/pathways-utils/pull/47
@@ -47,5 +54,5 @@ To release a new version (e.g. from `1.0.0` -> `2.0.0`):
 * Persistence enabled
 * General argument type fixes
 
-[Unreleased]: https://github.com/AI-Hypercomputer/pathways-utils/compare/v0.1.0...HEAD
+[Unreleased]: https://github.com/AI-Hypercomputer/pathways-utils/compare/v0.1.1...HEAD
 

pathwaysutils/sidecar/README.md

@@ -218,7 +218,7 @@ spec:
           - -c
           - |
             pip install --upgrade pip
-            pip install -U --pre jax jaxlib -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html
+            pip install -U --pre jax jaxlib --extra-index-url https://us-python.pkg.dev/ml-oss-artifacts-published/jax-public-nightly-artifacts-registry/simple/
             pip install pathwaysutils
             python -c "import jax; import pathwaysutils; print(\"Number of JAX devices is\", len(jax.devices()))"
 ```

pathwaysutils/test/profiling_test.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import logging
 from unittest import mock
 
 from pathwaysutils import profiling
@@ -32,13 +33,14 @@ def setUp(self):
 
   @parameterized.parameters(8000, 1234)
   def test_collect_profile_port(self, port):
-    profiling.collect_profile(
+    result = profiling.collect_profile(
         port=port,
         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(
         f"http://127.0.0.1:{port}/profiling",
         json={
@@ -49,13 +51,14 @@ def test_collect_profile_port(self, port):
 
   @parameterized.parameters(1000, 1234)
   def test_collect_profile_duration_ms(self, duration_ms):
-    profiling.collect_profile(
+    result = profiling.collect_profile(
         port=8000,
         duration_ms=duration_ms,
         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={
@@ -66,13 +69,14 @@ def test_collect_profile_duration_ms(self, duration_ms):
 
   @parameterized.parameters("127.0.0.1", "localhost", "192.168.1.1")
   def test_collect_profile_host(self, host):
-    profiling.collect_profile(
+    result = profiling.collect_profile(
         port=8000,
         duration_ms=1000,
         host=host,
         log_dir="gs://test_bucket/test_dir",
     )
 
+    self.assertTrue(result)
     self.mock_post.assert_called_once_with(
         f"http://{host}:8000/profiling",
         json={
@@ -87,10 +91,11 @@ def test_collect_profile_host(self, host):
       "gs://test_bucket3/test/log/dir",
   )
   def test_collect_profile_log_dir(self, log_dir):
-    profiling.collect_profile(
+    result = profiling.collect_profile(
         port=8000, duration_ms=1000, host="127.0.0.1", log_dir=log_dir
     )
 
+    self.assertTrue(result)
     self.mock_post.assert_called_once_with(
         "http://127.0.0.1:8000/profiling",
         json={
@@ -107,22 +112,27 @@ def test_collect_profile_log_dir_error(self, log_dir):
       )
 
   @parameterized.parameters(
-      requests.exceptions.ConnectionError,
-      requests.exceptions.Timeout,
-      requests.exceptions.TooManyRedirects,
-      requests.exceptions.RequestException,
-      requests.exceptions.HTTPError,
+      requests.exceptions.ConnectionError("Connection error"),
+      requests.exceptions.Timeout("Timeout"),
+      requests.exceptions.TooManyRedirects("Too many redirects"),
+      requests.exceptions.RequestException("Request exception"),
+      requests.exceptions.HTTPError("HTTP error"),
   )
-  def test_collect_profile_request_error(self, exception_type):
-    self.mock_post.side_effect = exception_type
+  def test_collect_profile_request_error(self, exception):
+    self.mock_post.side_effect = exception
+
+    with self.assertLogs(profiling._logger, level=logging.ERROR) as logs:
+      result = profiling.collect_profile(
+          port=8000,
+          duration_ms=1000,
+          host="127.0.0.1",
+          log_dir="gs://test_bucket/test_dir",
+      )
 
-    result = profiling.collect_profile(
-        port=8000,
-        duration_ms=1000,
-        host="127.0.0.1",
-        log_dir="gs://test_bucket/test_dir",
+    self.assertLen(logs.output, 1)
+    self.assertIn(
+        f"Failed to collect profiling data: {exception}", logs.output[0]
     )
-
     self.assertFalse(result)
     self.mock_post.assert_called_once()
 

pyproject.toml

@@ -14,10 +14,10 @@ keywords = []
 dependencies = [
     "absl-py",
     "fastapi",
-    "google-cloud-logging",
     "jax>=0.4.26",
     "orbax-checkpoint",
     "uvicorn",
+    "requests",
 ]
 
 # `version` is automatically set by flit to use `my_project.__version__`