From 448e080c752a1b5f54c422401b41e824b0274a91 Mon Sep 17 00:00:00 2001 From: Ran Chen Date: Tue, 17 Nov 2020 14:06:20 -0800 Subject: [PATCH] Guess test binary path from TEST_TARGET env var PiperOrigin-RevId: 342941997 Change-Id: Ibb7d8da784b869927c26f9eb4fe7564aee05a21d --- tensorflow/python/distribute/BUILD | 1 + .../python/distribute/multi_process_lib.py | 43 +++++++++---------- 2 files changed, 22 insertions(+), 22 deletions(-) diff --git a/tensorflow/python/distribute/BUILD b/tensorflow/python/distribute/BUILD index ea0dafac69a..20d77bd40ed 100644 --- a/tensorflow/python/distribute/BUILD +++ b/tensorflow/python/distribute/BUILD @@ -1668,6 +1668,7 @@ py_library( deps = [ "//tensorflow/python:client_testlib", "//tensorflow/python/eager:test", + "@absl_py//absl/logging", ], ) diff --git a/tensorflow/python/distribute/multi_process_lib.py b/tensorflow/python/distribute/multi_process_lib.py index 81ee53a285f..175dea7d184 100644 --- a/tensorflow/python/distribute/multi_process_lib.py +++ b/tensorflow/python/distribute/multi_process_lib.py @@ -23,6 +23,7 @@ import platform import sys import unittest from absl import app +from absl import logging from tensorflow.python.eager import test @@ -97,31 +98,29 @@ def _set_spawn_exe_path(): """ # TODO(b/150264776): This does not work with Windows. Find a solution. if sys.argv[0].endswith('.py'): + path = None # If all we have is a python module path, we'll need to make a guess for - # the actual executable path. Since the binary path may correspond to the - # parent's path of the python module, we are making guesses by reducing - # directories one at a time. E.g., - # tensorflow/python/some/path/my_test.py - # -> tensorflow/python/some/path/my_test - # -> tensorflow/python/some/my_test - # -> tensorflow/python/my_test - path_to_use = None - guess_path = sys.argv[0][:-3] - guess_path = guess_path.split(os.sep) - for path_reduction in range(-1, -len(guess_path), -1): - possible_path = os.sep.join(guess_path[:path_reduction] + - [guess_path[-1]]) + # the actual executable path. + if 'bazel-out' in sys.argv[0]: + # Guess the binary path under bazel. For target + # //tensorflow/python/distribute:input_lib_test_multiworker_gpu, the + # argv[0] is in the form of + # /.../org_tensorflow/tensorflow/python/distribute/input_lib_test.py + # and the binary is + # /.../org_tensorflow/tensorflow/python/distribute/input_lib_test_multiworker_gpu + org_tensorflow_path = sys.argv[0][:sys.argv[0].rfind('/tensorflow')] + if org_tensorflow_path.endswith('/org_tensorflow'): + binary = os.environ['TEST_TARGET'][2:].replace(':', '/', 1) + possible_path = os.path.join(org_tensorflow_path, binary) + logging.info('Guessed test binary path: %s', possible_path) if os.access(possible_path, os.X_OK): - path_to_use = possible_path - break - # The binary can possibly have _gpu suffix. - possible_path += '_gpu' - if os.access(possible_path, os.X_OK): - path_to_use = possible_path - break - if path_to_use is None: + path = possible_path + if path is None: + logging.error( + 'Cannot determine binary path. sys.argv[0]=%s os.environ=%s', + sys.argv[0], os.environ) raise RuntimeError('Cannot determine binary path') - sys.argv[0] = path_to_use + sys.argv[0] = path # Note that this sets the executable for *all* contexts. multiprocessing.get_context().set_executable(sys.argv[0])