Merge pull request #1664 from lcolok/master

Improved Environment Variable Handling for Enhanced Flexibility in TensorBoard Launch
bmaltais · Nov 11, 2023 · 1fc8b54 · 1fc8b54
2 parents 504c602 + 43bd029
commit 1fc8b54
Showing 1 changed file with 10 additions and 4 deletions.
diff --git a/library/tensorboard_gui.py b/library/tensorboard_gui.py
@@ -12,12 +12,17 @@
 tensorboard_proc = None
 TENSORBOARD = 'tensorboard' if os.name == 'posix' else 'tensorboard.exe'
 
+# Set the default tensorboard port
+DEFAULT_TENSORBOARD_PORT = 6006
 
 def start_tensorboard(headless, logging_dir, wait_time=5):
     global tensorboard_proc
 
     headless_bool = True if headless.get('label') == 'True' else False
 
+    # Read the TENSORBOARD_PORT from the environment, or use the default
+    tensorboard_port = os.environ.get('TENSORBOARD_PORT', DEFAULT_TENSORBOARD_PORT)
+
     if not os.listdir(logging_dir):
         log.info('Error: log folder is empty')
         msgbox(msg='Error: log folder is empty')
@@ -30,7 +35,7 @@ def start_tensorboard(headless, logging_dir, wait_time=5):
         '--host',
         '0.0.0.0',
         '--port',
-        '6006',
+        str(tensorboard_port),
     ]
 
     log.info(run_cmd)
@@ -41,7 +46,7 @@ def start_tensorboard(headless, logging_dir, wait_time=5):
         stop_tensorboard()
 
     # Start background process
-    log.info('Starting tensorboard...')
+    log.info('Starting TensorBoard on port {}'.format(tensorboard_port))
     try:
         tensorboard_proc = subprocess.Popen(run_cmd)
     except Exception as e:
@@ -53,8 +58,9 @@ def start_tensorboard(headless, logging_dir, wait_time=5):
         time.sleep(wait_time)
 
         # Open the TensorBoard URL in the default browser
-        log.info('Opening tensorboard url in browser...')
-        webbrowser.open('http://localhost:6006')
+        tensorboard_url = f'http://localhost:{tensorboard_port}'
+        log.info(f'Opening TensorBoard URL in browser: {tensorboard_url}')
+        webbrowser.open(tensorboard_url)
 
 
 def stop_tensorboard():