manager: launch TENSORBOARD_BINARY if available (#2386)

Summary:
This enables specifying the TensorBoard binary with an environment
variable. Callers previously needed to set up a `PATH` with a custom
binary; this is simpler and less intrusive.

(Googlers, see comments on <http://cl/236333099> for context.)

Test Plan:
The newly added test fails before this change and passes after it.

wchargin-branch: tensorboard-binary

Author: wchargin

tensorboard/manager.py

@@ -397,9 +397,10 @@ def start(arguments, timeout=datetime.timedelta(seconds=60)):
   (stdout_fd, stdout_path) = tempfile.mkstemp(prefix=".tensorboard-stdout-")
   (stderr_fd, stderr_path) = tempfile.mkstemp(prefix=".tensorboard-stderr-")
   start_time_seconds = time.time()
+  tensorboard_binary = os.environ.get("TENSORBOARD_BINARY", "tensorboard")
   try:
     p = subprocess.Popen(
-        ["tensorboard"] + arguments,
+        [tensorboard_binary] + arguments,
         stdout=stdout_fd,
         stderr=stderr_fd,
     )

tensorboard/manager_e2e_test.py

@@ -319,6 +319,40 @@ def test_timeout(self):
     self.assertEqual(start_result, manager.StartTimedOut(pid=expected_pid))
     self.assertEqual(manager.get_all(), [])
 
+  def test_tensorboard_binary_environment_variable(self):
+    if os.name == "nt":
+      # TODO(@wchargin): This could in principle work on Windows.
+      self.skipTest("Requires a POSIX shell for the stub script.")
+    tempdir = tempfile.mkdtemp()
+    filepath = os.path.join(tempdir, "tensorbad")
+    program = textwrap.dedent(
+        r"""
+        #!/bin/sh
+        printf >&2 'tensorbad: fatal: something bad happened\n'
+        printf 'tensorbad: also some stdout\n'
+        exit 77
+        """.lstrip()
+    )
+    with open(filepath, "w") as outfile:
+      outfile.write(program)
+    os.chmod(filepath, 0o777)
+    environ = {"TENSORBOARD_BINARY": filepath}
+    environ_patcher = mock.patch.dict(os.environ, environ)
+    environ_patcher.start()
+    self.addCleanup(environ_patcher.stop)
+
+    start_result = manager.start(["--logdir=./logs", "--port=0"])
+    self.assertIsInstance(start_result, manager.StartFailed)
+    self.assertEqual(
+        start_result,
+        manager.StartFailed(
+            exit_code=77,
+            stderr="tensorbad: fatal: something bad happened\n",
+            stdout="tensorbad: also some stdout\n",
+        ),
+    )
+    self.assertEqual(manager.get_all(), [])
+
 
 if __name__ == "__main__":
   tf.test.main()