Spaces:
Build error
Build error
| import os | |
| import gradio as gr | |
| import subprocess | |
| import time | |
| import webbrowser | |
| try: | |
| os.environ["TF_ENABLE_ONEDNN_OPTS"] = "0" | |
| import tensorflow # Attempt to import tensorflow to check if it is installed | |
| 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() # Wait for the thread to finish | |
| 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, | |
| ) | |