From a4efbda6c0ad9df1ef4fadbbd9eb1e02d13eccde Mon Sep 17 00:00:00 2001 From: Etienne Pot Date: Tue, 5 Dec 2023 05:14:44 -0800 Subject: [PATCH] Attempt to fix XManager cli PiperOrigin-RevId: 588033997 --- etils/epath/gpath.py | 2 +- etils/epy/binary_import.py | 15 ++++----------- 2 files changed, 5 insertions(+), 12 deletions(-) diff --git a/etils/epath/gpath.py b/etils/epath/gpath.py index 990e81c..8ee11d2 100644 --- a/etils/epath/gpath.py +++ b/etils/epath/gpath.py @@ -63,7 +63,7 @@ @functools.cache def _is_tf_installed() -> bool: """Checks whether TensorFlow is installed.""" - if os.environ.get('EPATH_USE_TF', '').lower() in ['false', 'no', 'f', '0']: + if os.environ['EPATH_USE_TF'].lower() in ['false', 'no', 'f', '0']: return False return importlib.util.find_spec('tensorflow') is not None diff --git a/etils/epy/binary_import.py b/etils/epy/binary_import.py index 2c07d18..15c1fbd 100644 --- a/etils/epy/binary_import.py +++ b/etils/epy/binary_import.py @@ -20,7 +20,6 @@ import sys from typing import Any -import __main__ # pylint: disable=g-bad-import-order from etils.epy import py_utils @@ -28,16 +27,10 @@ def _is_ipython_terminal() -> bool: """Returns True if running in a IPython terminal/XManager CLI environment.""" # XManager CLI trigger binary imports - - # `epy` is imported before the `runpy.run_module(`, so main is still the - # XManager binary - if __main__.__file__.endswith('xmanager2/client/cli/xm_cli.py'): - return True - - # In case `epy` is imported after the XManager CLI, detecting we're in - # `xmanager launch` is non-trivial because the script is launched with - # `runpy.run_module(`, hiding some XManager internals and overwriting - # `__main__`. + # Detecting we're in `xmanager launch` is non-trivial because the script + # is launched with `runpy.run_module(`, hiding some XManager internals. + # Otherwise we could have checked if `__main__.__file__.endswith('xm_cli')`, + # but `__main__` get overwritten here. if any(flag.startswith('--xm_launch_script=') for flag in sys.argv): return True