Allow to run main locally

PiperOrigin-RevId: 702292043
google-research · Dec 3, 2024 · 9b5b71b · 9b5b71b
1 parent a27f847
commit 9b5b71b
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 4 deletions.
diff --git a/kauldron/main.py b/kauldron/main.py
@@ -26,8 +26,12 @@
 import jax
 import tensorflow as tf
 
-from kauldron import kd
-from kauldron.utils import sweep_utils
+# Imports adhoc-imported when running with `ml_python`
+with epy.binary_adhoc():
+  # pylint: disable=g-import-not-at-top
+  from kauldron import kd
+  from kauldron.utils.status_utils import status  # pylint: disable=g-importing-member
+  # pylint: enable=g-import-not-at-top
 
 _CONFIG = kd.konfig.DEFINE_config_file(
     "cfg",
@@ -63,7 +67,10 @@ def main(_):
 
 @contextlib.contextmanager
 def _wu_error_handling(post_mortem: bool = False):
-  yield  # not yet supported externally
+  """Catch and log error."""
+  if not status.on_xmanager:
+    yield
+    return
 
 
 def _flags_parser(args: list[str]) -> None:

diff --git a/kauldron/utils/status_utils.py b/kauldron/utils/status_utils.py
@@ -51,7 +51,7 @@ class _Status:
 
   @functools.cached_property
   def on_xmanager(self) -> bool:
-    return xmanager_api.is_running_under_xmanager()
+    return False
 
   @functools.cached_property
   def is_lead_host(self) -> bool: