|
import os
|
|
import gradio as gr
|
|
import subprocess
|
|
import time
|
|
import webbrowser
|
|
|
|
try:
|
|
os.environ["TF_ENABLE_ONEDNN_OPTS"] = "0"
|
|
import tensorflow
|
|
|
|
visibility = True
|
|
except ImportError:
|
|
visibility = False
|
|
|
|
from easygui import msgbox
|
|
from threading import Thread, Event
|
|
from .custom_logging import setup_logging
|
|
from .common_gui import setup_environment
|
|
|
|
|
|
class TensorboardManager:
|
|
DEFAULT_TENSORBOARD_PORT = 6006
|
|
|
|
def __init__(self, logging_dir, headless: bool = False, wait_time=5):
|
|
self.logging_dir = logging_dir
|
|
self.headless = headless
|
|
self.wait_time = wait_time
|
|
self.tensorboard_proc = None
|
|
self.tensorboard_port = os.environ.get(
|
|
"TENSORBOARD_PORT", self.DEFAULT_TENSORBOARD_PORT
|
|
)
|
|
self.log = setup_logging()
|
|
self.thread = None
|
|
self.stop_event = Event()
|
|
|
|
self.gradio_interface()
|
|
|
|
def get_button_states(self, started=False):
|
|
return gr.Button(
|
|
visible=visibility and (not started or self.headless)
|
|
), gr.Button(visible=visibility and (started or self.headless))
|
|
|
|
def open_tensorboard_url(self):
|
|
tensorboard_url = f"http://localhost:{self.tensorboard_port}"
|
|
self.log.info(f"Opening TensorBoard URL in browser: {tensorboard_url}")
|
|
webbrowser.open(tensorboard_url)
|
|
|
|
def start_tensorboard(self, logging_dir=None):
|
|
if self.tensorboard_proc is not None:
|
|
self.log.info(
|
|
"Tensorboard is already running. Terminating existing process before starting new one..."
|
|
)
|
|
self.stop_tensorboard()
|
|
|
|
if not os.path.exists(logging_dir) or not os.listdir(logging_dir):
|
|
self.log.error(
|
|
"Error: logging folder does not exist or does not contain logs."
|
|
)
|
|
msgbox(msg="Error: logging folder does not exist or does not contain logs.")
|
|
return self.get_button_states(started=False)
|
|
|
|
run_cmd = [
|
|
"tensorboard",
|
|
"--logdir",
|
|
logging_dir,
|
|
"--host",
|
|
"0.0.0.0",
|
|
"--port",
|
|
str(self.tensorboard_port),
|
|
]
|
|
|
|
self.log.info(run_cmd)
|
|
|
|
self.log.info("Starting TensorBoard on port {}".format(self.tensorboard_port))
|
|
try:
|
|
env = setup_environment()
|
|
self.tensorboard_proc = subprocess.Popen(run_cmd, env=env)
|
|
except Exception as e:
|
|
self.log.error("Failed to start Tensorboard:", e)
|
|
return self.get_button_states(started=False)
|
|
|
|
if not self.headless:
|
|
self.stop_event.clear()
|
|
|
|
time.sleep(self.wait_time)
|
|
if not self.stop_event.is_set():
|
|
self.thread = Thread(target=self.open_tensorboard_url)
|
|
self.thread.start()
|
|
|
|
return self.get_button_states(started=True)
|
|
|
|
def stop_tensorboard(self):
|
|
if self.tensorboard_proc is not None:
|
|
self.log.info("Stopping tensorboard process...")
|
|
try:
|
|
self.tensorboard_proc.terminate()
|
|
self.tensorboard_proc = None
|
|
self.log.info("...process stopped")
|
|
except Exception as e:
|
|
self.log.error("Failed to stop Tensorboard:", e)
|
|
|
|
if self.thread is not None:
|
|
self.stop_event.set()
|
|
self.thread.join()
|
|
self.thread = None
|
|
self.log.info("Thread terminated successfully.")
|
|
|
|
return self.get_button_states(started=False)
|
|
|
|
def gradio_interface(self):
|
|
|
|
with gr.Row():
|
|
button_start_tensorboard = gr.Button(
|
|
value="Start tensorboard",
|
|
elem_id="myTensorButton",
|
|
visible=visibility,
|
|
)
|
|
button_stop_tensorboard = gr.Button(
|
|
value="Stop tensorboard",
|
|
visible=visibility and self.headless,
|
|
elem_id="myTensorButtonStop",
|
|
)
|
|
button_open_tensorboard = gr.Button(
|
|
value="Open tensorboard",
|
|
elem_id="myTensorButton",
|
|
visible=not visibility,
|
|
link=f"http://localhost:{self.tensorboard_port}",
|
|
)
|
|
button_start_tensorboard.click(
|
|
self.start_tensorboard,
|
|
inputs=[self.logging_dir],
|
|
outputs=[button_start_tensorboard, button_stop_tensorboard],
|
|
show_progress=False,
|
|
)
|
|
button_stop_tensorboard.click(
|
|
self.stop_tensorboard,
|
|
outputs=[button_start_tensorboard, button_stop_tensorboard],
|
|
show_progress=False,
|
|
)
|
|
|
