Updated some docstrings in pathwaysutils/__init__.py, removed a warning, updated tests, and changed warnings to debug logs.

PiperOrigin-RevId: 746481798

Author: lukebaumann

pathwaysutils/__init__.py

@@ -16,7 +16,6 @@
 import datetime
 import logging
 import os
-import warnings
 
 import jax
 from pathwaysutils import cloud_logging
@@ -33,11 +32,29 @@
 
 #  This is a brittle implementation since the platforms value is not necessarily
 #  which backend is ultimately selected
-def _is_pathways_used():
+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():
+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
@@ -51,31 +68,29 @@ def _is_persistence_enabled():
   return False
 
 
-def initialize():
-  """Initializes pathwaysutils."""
+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
 
-  if _initialization_count == 1:
-    warnings.warn(
-        "pathwaysutils: Legacy initialization. Ensure you also call"
-        " pathwaysutils.initialize(). This warning will be removed in a future"
-        " release."
-    )
-
-  # Ignoring the second call to initialize() is a temporary measure so that this warning 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.
+  # 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 > 2:
-    warnings.warn(
-        "pathwaysutils: Already initialized. Ignoring duplicate call."
-    )
+    _logger.debug("Already initialized. Ignoring duplicate call.")
 
   if _initialization_count > 1:
     return
 
-  if _is_pathways_used():
-    _logger.debug(
-        "pathwaysutils: Detected Pathways-on-Cloud backend. Applying changes."
-    )
+  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
@@ -90,14 +105,13 @@ def initialize():
       cloud_logging.setup()
     except Exception as error:  # pylint: disable=broad-except
       _logger.debug(
-          "pathwaysutils: Failed to set up cloud logging due to the following"
-          " error: %s",
+          "Failed to set up cloud logging due to the following error: %s",
           error,
       )
   else:
     _logger.debug(
-        "pathwaysutils: Did not detect Pathways-on-Cloud backend. No changes"
-        " applied."
+        "Did not detect Pathways-on-Cloud backend. No changes applied."
     )
 
+
 initialize()

pathwaysutils/test/pathwaysutils_test.py

@@ -14,8 +14,8 @@
 
 import os
 from unittest import mock
-import warnings
 
+import google.cloud.logging
 import jax
 import pathwaysutils
 from pathwaysutils import cloud_logging
@@ -32,37 +32,65 @@ def setUp(self):
         cloud_logging, "setup", autospec=True
     )
 
-  def test_legacy_initialize(self):
+  def test_first_initialize(self):
+    jax.config.update("jax_platforms", "proxy")
     pathwaysutils._initialization_count = 0
 
-    with self.assertWarns(UserWarning, msg="Legacy initialization"):
+    self.enter_context(
+        mock.patch.object(google.cloud.logging, "Client", autospec=True)
+    )
+
+    with self.assertLogs(pathwaysutils._logger, level="DEBUG") as logs:
       pathwaysutils.initialize()
 
-  def test_legacy_and_new_initialize(self):
+    self.assertLen(logs.output, 1)
+    self.assertIn(
+        "Detected Pathways-on-Cloud backend. Applying changes.", logs.output[0]
+    )
+
+  def test_second_initialize(self):
+    jax.config.update("jax_platforms", "proxy")
     pathwaysutils._initialization_count = 1
 
-    with warnings.catch_warnings(record=True) as caught_warnings:
+    with self.assertNoLogs(pathwaysutils._logger, level="DEBUG"):
       pathwaysutils.initialize()
 
-    self.assertEmpty(caught_warnings)
-
   @parameterized.named_parameters(
       ("initialization_count 2", 2),
       ("initialization_count 5", 5),
+      ("initialization_count 1000", 1000),
   )
-  def test_initialize_more_than_once(self, initialization_count):
+  def test_initialize_more_than_twice(self, initialization_count):
     pathwaysutils._initialization_count = initialization_count
 
-    with self.assertWarns(UserWarning, msg="Already initialized"):
+    with self.assertLogs(pathwaysutils._logger, level="DEBUG") as logs:
       pathwaysutils.initialize()
 
-  def test_is_pathways_used(self):
-    for platform in ["", "cpu", "tpu", "gpu", "cpu,tpu,gpu"]:
-      jax.config.update("jax_platforms", platform)
-      self.assertFalse(pathwaysutils._is_pathways_used())
-    for platform in ["proxy", "proxy,cpu", "cpu,proxy", "tpu,cpu,proxy,gpu"]:
-      jax.config.update("jax_platforms", platform)
-      self.assertTrue(pathwaysutils._is_pathways_used())
+    self.assertLen(logs.output, 1)
+    self.assertIn(
+        "Already initialized. Ignoring duplicate call.", logs.output[0]
+    )
+
+  @parameterized.named_parameters(
+      ("empty", ""),
+      ("cpu", "cpu"),
+      ("tpu", "tpu"),
+      ("gpu", "gpu"),
+      ("cpu,tpu,gpu", "cpu,tpu,gpu"),
+  )
+  def test_not_is_pathways_backend_used(self, platform: str):
+    jax.config.update("jax_platforms", platform)
+    self.assertFalse(pathwaysutils.is_pathways_backend_used())
+
+  @parameterized.named_parameters(
+      ("proxy", "proxy"),
+      ("proxy,cpu", "proxy,cpu"),
+      ("cpu,proxy", "cpu,proxy"),
+      ("tpu,cpu,proxy,gpu", "tpu,cpu,proxy,gpu"),
+  )
+  def test_is_pathways_backend_used(self, platform: str):
+    jax.config.update("jax_platforms", platform)
+    self.assertTrue(pathwaysutils.is_pathways_backend_used())
 
   def test_persistence_enabled(self):
     os.environ["ENABLE_PATHWAYS_PERSISTENCE"] = "1"