Fix absl flag initialization in cloud_tpu_profiler (tensorflow#14426)

This fixes a regression caused by 2652704 where the command would not run. See also: tensorflow/tensorboard#716
Person11111 · Nov 29, 2017 · 4b7d79b · 4b7d79b
1 parent 6bbff83
commit 4b7d79b
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 9 deletions.
diff --git a/tensorflow/contrib/tpu/profiler/pip_package/cloud_tpu_profiler/main.py b/tensorflow/contrib/tpu/profiler/pip_package/cloud_tpu_profiler/main.py
@@ -24,22 +24,18 @@
 
 import tensorflow as tf
 
-
 tf.flags.DEFINE_string('service_addr', '',
                        'Address of TPU profiler service e.g. localhost:8466')
-
-
 tf.flags.DEFINE_string('logdir', '',
                        'Path of TensorBoard log directory e.g. /tmp/tb_log')
-
-
 tf.flags.DEFINE_integer('duration_ms', 2000, 'Duration of tracing in ms.')
 
-
 FLAGS = tf.flags.FLAGS
+EXECUTABLE = 'data/capture_tpu_profile'
 
 
-EXECUTABLE = 'data/capture_tpu_profile'
+def run_main():
+  tf.app.run(main)
 
 
 def main(unused_argv=None):
@@ -54,4 +50,4 @@ def main(unused_argv=None):
 
 
 if __name__ == '__main__':
-  tf.app.run(main)
+  run_main()
diff --git a/tensorflow/contrib/tpu/profiler/pip_package/setup.py b/tensorflow/contrib/tpu/profiler/pip_package/setup.py
@@ -23,7 +23,7 @@
 _VERSION = '1.3.0-a1'
 
 CONSOLE_SCRIPTS = [
-    'capture_tpu_profile=cloud_tpu_profiler.main:main',
+    'capture_tpu_profile=cloud_tpu_profiler.main:run_main',
 ]
 
 REQUIRED_PACKAGES = [