From 0c89c768ca9a3bf5fc8b2cc82fb7565c12476fdc Mon Sep 17 00:00:00 2001 From: Luke Baumann Date: Wed, 6 Aug 2025 10:15:37 -0700 Subject: [PATCH] Internal change to copybara PiperOrigin-RevId: 791741957 --- pathwaysutils/__init__.py | 91 +-------------- pathwaysutils/_initialize.py | 104 ++++++++++++++++++ ...thwaysutils_test.py => initialize_test.py} | 33 +++--- pathwaysutils/test/proxy_backend_test.py | 1 + 4 files changed, 123 insertions(+), 106 deletions(-) create mode 100644 pathwaysutils/_initialize.py rename pathwaysutils/test/{pathwaysutils_test.py => initialize_test.py} (71%) diff --git a/pathwaysutils/__init__.py b/pathwaysutils/__init__.py index 0be743f..e45a156 100644 --- a/pathwaysutils/__init__.py +++ b/pathwaysutils/__init__.py @@ -12,95 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. """Package of Pathways-on-Cloud utilities.""" +from pathwaysutils import _initialize -import datetime -import logging -import os +initialize = _initialize.initialize +del _initialize -import jax -from pathwaysutils import profiling -from pathwaysutils import proxy_backend -from pathwaysutils.persistence import orbax_handler - - -_logger = logging.getLogger(__name__) -_initialization_count = 0 # When changing this, also update the CHANGELOG.md. __version__ = "v0.1.1" - - -# This is a brittle implementation since the platforms value is not necessarily -# which backend is ultimately selected -def is_pathways_backend_used() -> bool: - """Returns whether Pathways backend is used. - - This function checks the JAX platforms configuration to determine whether - Pathways is used. If the platforms configuration contains the string "proxy", - Pathways is used. This is a brittle implementation since the platforms value - is not necessarily which backend is ultimately selected or there may be more - than one platform specified and another may have higher priority. - """ - return jax.config.jax_platforms and "proxy" in jax.config.jax_platforms - - -def _is_persistence_enabled() -> bool: - """Returns whether persistence is enabled. - - This function checks the environment variable ENABLE_PATHWAYS_PERSISTENCE to - determine whether persistence is enabled. If the variable is set to "1", - persistence is enabled. If the variable is set to "0" or unset, persistence is - disabled. - - Returns: - True if persistence is enabled, False otherwise. - """ - if "ENABLE_PATHWAYS_PERSISTENCE" in os.environ: - if os.environ["ENABLE_PATHWAYS_PERSISTENCE"] == "1": - return True - if os.environ["ENABLE_PATHWAYS_PERSISTENCE"] == "0": - return False - else: - raise ValueError( - "ENABLE_PATHWAYS_PERSISTENCE must be set to 1/0 or unset, got: " - + os.environ["ENABLE_PATHWAYS_PERSISTENCE"] - ) - return False - - -def initialize() -> None: - """Initializes pathwaysutils. - - This function is called by the user to initialize pathwaysutils. It is - responsible for setting up the logging, profiling, and persistence handlers - through various monkey patching functions. It is also responsible for - registering the proxy backend factory. - """ - global _initialization_count - _initialization_count += 1 - - # Ignoring the second call to initialize() is a temporary measure so that this - # debug log is not triggered for customers who are following our instructions - # and using the new initialize() function only once but have already had the - # legacy initialization triggered. - if _initialization_count > 1: - _logger.debug("Already initialized. Ignoring duplicate call.") - return - - _logger.debug("Starting initialize.") - - if is_pathways_backend_used(): - _logger.debug("Detected Pathways-on-Cloud backend. Applying changes.") - proxy_backend.register_backend_factory() - profiling.monkey_patch_jax() - # TODO: b/365549911 - Remove when OCDBT-compatible - if _is_persistence_enabled(): - orbax_handler.register_pathways_handlers(datetime.timedelta(hours=1)) - - # Turn off JAX compilation cache because Pathways handles its own - # compilation cache. - jax.config.update("jax_enable_compilation_cache", False) - - else: - _logger.debug( - "Did not detect Pathways-on-Cloud backend. No changes applied." - ) diff --git a/pathwaysutils/_initialize.py b/pathwaysutils/_initialize.py new file mode 100644 index 0000000..0e42e59 --- /dev/null +++ b/pathwaysutils/_initialize.py @@ -0,0 +1,104 @@ +# Copyright 2025 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. +"""Initialization functions for Pathways-on-Cloud utilities.""" + +import datetime +import logging +import os + +import jax +from pathwaysutils import profiling +from pathwaysutils import proxy_backend +from pathwaysutils.persistence import orbax_handler + + +_logger = logging.getLogger(__name__) +_initialization_count = 0 + + +# This is a brittle implementation since the platforms value is not necessarily +# which backend is ultimately selected +def is_pathways_backend_used() -> bool: + """Returns whether Pathways backend is used. + + This function checks the JAX platforms configuration to determine whether + Pathways is used. If the platforms configuration contains the string "proxy", + Pathways is used. This is a brittle implementation since the platforms value + is not necessarily which backend is ultimately selected or there may be more + than one platform specified and another may have higher priority. + """ + return jax.config.jax_platforms and "proxy" in jax.config.jax_platforms + + +def _is_persistence_enabled() -> bool: + """Returns whether persistence is enabled. + + This function checks the environment variable ENABLE_PATHWAYS_PERSISTENCE to + determine whether persistence is enabled. If the variable is set to "1", + persistence is enabled. If the variable is set to "0" or unset, persistence is + disabled. + + Returns: + True if persistence is enabled, False otherwise. + """ + if "ENABLE_PATHWAYS_PERSISTENCE" in os.environ: + if os.environ["ENABLE_PATHWAYS_PERSISTENCE"] == "1": + return True + if os.environ["ENABLE_PATHWAYS_PERSISTENCE"] == "0": + return False + else: + raise ValueError( + "ENABLE_PATHWAYS_PERSISTENCE must be set to 1/0 or unset, got: " + + os.environ["ENABLE_PATHWAYS_PERSISTENCE"] + ) + return False + + +def initialize() -> None: + """Initializes pathwaysutils. + + This function is called by the user to initialize pathwaysutils. It is + responsible for setting up the logging, profiling, and persistence handlers + through various monkey patching functions. It is also responsible for + registering the proxy backend factory. + """ + global _initialization_count + _initialization_count += 1 + + # Ignoring the second call to initialize() is a temporary measure so that this + # debug log is not triggered for customers who are following our instructions + # and using the new initialize() function only once but have already had the + # legacy initialization triggered. + if _initialization_count > 1: + _logger.debug("Already initialized. Ignoring duplicate call.") + return + + _logger.debug("Starting initialize.") + + if is_pathways_backend_used(): + _logger.debug("Detected Pathways-on-Cloud backend. Applying changes.") + proxy_backend.register_backend_factory() + profiling.monkey_patch_jax() + # TODO: b/365549911 - Remove when OCDBT-compatible + if _is_persistence_enabled(): + orbax_handler.register_pathways_handlers(datetime.timedelta(hours=1)) + + # Turn off JAX compilation cache because Pathways handles its own + # compilation cache. + jax.config.update("jax_enable_compilation_cache", False) + + else: + _logger.debug( + "Did not detect Pathways-on-Cloud backend. No changes applied." + ) diff --git a/pathwaysutils/test/pathwaysutils_test.py b/pathwaysutils/test/initialize_test.py similarity index 71% rename from pathwaysutils/test/pathwaysutils_test.py rename to pathwaysutils/test/initialize_test.py index a2fedd9..8b5c087 100644 --- a/pathwaysutils/test/pathwaysutils_test.py +++ b/pathwaysutils/test/initialize_test.py @@ -13,28 +13,25 @@ # limitations under the License. import os -from unittest import mock import jax -import pathwaysutils +from pathwaysutils import _initialize from absl.testing import absltest from absl.testing import parameterized -class PathwaysutilsTest(parameterized.TestCase): +class InitializeTest(parameterized.TestCase): def test_first_initialize(self): jax.config.update("jax_platforms", "proxy") - pathwaysutils._initialization_count = 0 + _initialize._initialization_count = 0 - with self.assertLogs(pathwaysutils._logger, level="DEBUG") as logs: - pathwaysutils.initialize() + with self.assertLogs(_initialize._logger, level="DEBUG") as logs: + _initialize.initialize() self.assertLen(logs.output, 2) - self.assertIn( - "Starting initialize.", logs.output[0] - ) + self.assertIn("Starting initialize.", logs.output[0]) self.assertIn( "Detected Pathways-on-Cloud backend. Applying changes.", logs.output[1] ) @@ -46,10 +43,10 @@ def test_first_initialize(self): ("initialization_count 1000", 1000), ) def test_initialize_more_than_once(self, initialization_count): - pathwaysutils._initialization_count = initialization_count + _initialize._initialization_count = initialization_count - with self.assertLogs(pathwaysutils._logger, level="DEBUG") as logs: - pathwaysutils.initialize() + with self.assertLogs(_initialize._logger, level="DEBUG") as logs: + _initialize.initialize() self.assertLen(logs.output, 1) self.assertIn( @@ -65,7 +62,7 @@ def test_initialize_more_than_once(self, initialization_count): ) def test_not_is_pathways_backend_used(self, platform: str): jax.config.update("jax_platforms", platform) - self.assertFalse(pathwaysutils.is_pathways_backend_used()) + self.assertFalse(_initialize.is_pathways_backend_used()) @parameterized.named_parameters( ("proxy", "proxy"), @@ -75,20 +72,20 @@ def test_not_is_pathways_backend_used(self, platform: str): ) def test_is_pathways_backend_used(self, platform: str): jax.config.update("jax_platforms", platform) - self.assertTrue(pathwaysutils.is_pathways_backend_used()) + self.assertTrue(_initialize.is_pathways_backend_used()) def test_persistence_enabled(self): os.environ["ENABLE_PATHWAYS_PERSISTENCE"] = "1" - self.assertTrue(pathwaysutils._is_persistence_enabled()) + self.assertTrue(_initialize._is_persistence_enabled()) os.environ["ENABLE_PATHWAYS_PERSISTENCE"] = "0" - self.assertFalse(pathwaysutils._is_persistence_enabled()) + self.assertFalse(_initialize._is_persistence_enabled()) os.environ["ENABLE_PATHWAYS_PERSISTENCE"] = "" - self.assertRaises(ValueError, pathwaysutils._is_persistence_enabled) + self.assertRaises(ValueError, _initialize._is_persistence_enabled) del os.environ["ENABLE_PATHWAYS_PERSISTENCE"] - self.assertFalse(pathwaysutils._is_persistence_enabled()) + self.assertFalse(_initialize._is_persistence_enabled()) if __name__ == "__main__": diff --git a/pathwaysutils/test/proxy_backend_test.py b/pathwaysutils/test/proxy_backend_test.py index 08988de..d102fbd 100644 --- a/pathwaysutils/test/proxy_backend_test.py +++ b/pathwaysutils/test/proxy_backend_test.py @@ -20,6 +20,7 @@ from jax.lib.xla_extension import ifrt_proxy from pathwaysutils import proxy_backend + from absl.testing import absltest