diff --git a/core/config.py b/core/config.py index c56e8e8..d4c9221 100644 --- a/core/config.py +++ b/core/config.py @@ -28,12 +28,15 @@ def write(): json.dump(data, f) -def get(key: str): +def get(key: str, default=None): global data + if key not in data: + return default + return data.get(key) -def set(key: str, value): +def put(key: str, value): global data data[key] = value write() @@ -46,7 +49,8 @@ def has(key: str): def remove(key: str): global data - data.pop(key) + if key in data: + data.pop(key) write() @@ -54,3 +58,24 @@ def clear(): global data data = {} write() + + +def create_file(filename: str, content: str): + with open(os.path.join(os.path.dirname(file), filename), "w") as f: + f.write(content) + + +def remove_file(filename: str): + os.remove(os.path.join(os.path.dirname(file), filename)) + + +def file_exists(filename: str): + return os.path.exists(os.path.join(os.path.dirname(file), filename)) + + +def get_file_path(filename: str): + return os.path.join(os.path.dirname(file), filename) + + +def open_file(filename: str, mode: str = 'w'): + return open(os.path.join(os.path.dirname(file), filename), mode) diff --git a/core/stack.py b/core/stack.py index 3d9cb89..452a12c 100644 --- a/core/stack.py +++ b/core/stack.py @@ -2,11 +2,30 @@ import logging import os import shutil import subprocess +import time import psutil from core import utils, config from core.vars import logger, PYTHON_EXEC +from ui import choices + + +def find_correct_pid(parent_pid: int) -> int: + processes: list[psutil.Process] = [psutil.Process(parent_pid)] + create_time = processes[0].create_time() + + time.sleep(0.5) # Wait for child processes to spawn + + for i in range(1, 10): + if psutil.pid_exists(processes[0].pid + i): + child_process = psutil.Process(processes[0].pid + i) + if child_process.create_time() - create_time < 1: + processes.append(psutil.Process(processes[0].pid + i)) + else: + time.sleep(0.5 / i) + + return processes[-1].pid class Stack: @@ -18,7 +37,7 @@ class Stack: self.url = url self.port = port - self.process = None + self.pid = config.get(f"{self.name}-pid") def install(self): self.create_file('.installed', 'true') @@ -61,14 +80,28 @@ class Stack: pass def stop(self): - pass + if self.status(): + logger.debug(f"Killing {self.name} with PID: {self.pid}") + psutil.Process(self.pid).kill() + + self.set_pid(None) + + def set_pid(self, pid): + self.pid = pid + if pid is not None: + config.put(f"{self.name}-pid", pid) + else: + config.remove(f"{self.name}-pid") def restart(self): self.stop() self.start() def status(self) -> bool: - pass + if self.pid is None: + return False + + return psutil.pid_exists(self.pid) # Python/Bash utils def create_venv(self): @@ -100,42 +133,34 @@ class Stack: self.pip(f"install -r {filename}", env=env) def pip(self, cmd: str, env=[], args=[], current_dir: str = None): - self.bash(f"{' '.join(env)} {self.path}/venv/bin/pip {cmd} {' '.join(args)}", current_dir) + self.python(f"-m pip {cmd}", env=env, args=args, current_dir=current_dir) - def python(self, cmd: str, env=[], current_dir: str = None, daemon: bool = False): - self.bash(f"{' '.join(env)} {self.path}/venv/bin/python {cmd}", current_dir, daemon) + def python(self, cmd: str, env=[], args=[], current_dir: str = None, daemon: bool = False): + self.bash(f"{' '.join(env)} {self.path}/venv/bin/python {cmd} {' '.join(args)}", current_dir, daemon) def bash(self, cmd: str, current_dir: str = None, daemon: bool = False): cmd = f"cd {self.path if current_dir is None else os.path.join(self.path, current_dir)} && {cmd}" if daemon: - # Check if previous run process is saved - if config.has(f"{self.name}-pid"): + if self.status(): + choice = choices.already_running.ask() - # Check if PID still running - if psutil.pid_exists(config.get(f"{self.name}-pid")): - choice = input(f"{self.name} is already running, do you want to restart it? (y/n): ") - - if choice.lower() == 'y': - pid = config.get(f"{self.name}-pid") - logger.debug(f"Killing previous daemon with PID: {pid}") - psutil.Process(pid).kill() - else: - # TODO: attach to subprocess? - logger.info("Continuing without restarting...") - - return + if choice is True: + self.stop() + self._launch() + return else: - logger.warning( - f"Previous PID found for {self.name} but process is not running, continuing as stopped...") + # TODO: attach to subprocess / redirect logs? + logger.info("Continuing without restarting...") + return else: - logger.debug(f"No previous PID found for {self.name}, continuing as stopped...") - - logger.debug(f"Starting {self.name} as daemon with command: {cmd}") - cmd = f"{cmd} &" - process = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE) - config.set(f"{self.name}-pid", process.pid + 1) - return + logger.debug(f"Running command as daemon: {cmd}") + cmd = f"{cmd} &" + process = subprocess.Popen(cmd, shell=True, preexec_fn=os.setpgrp, + stdout=config.open_file(f"{self.id}-stdout"), + stderr=config.open_file(f"{self.id}-stderr")) + self.set_pid(find_correct_pid(process.pid)) + return else: logger.debug(f"Running command: {cmd}") diff --git a/services/background_removal_dis.py b/services/background_removal_dis.py index c45636e..45ff627 100644 --- a/services/background_removal_dis.py +++ b/services/background_removal_dis.py @@ -29,6 +29,5 @@ class BackgroundRemovalDis(Stack): super().install() def _launch(self): - args = ["--port", str(self.port)] - self.python(f"app.py {' '.join(args)}", current_dir="webui", - env=["TORCH_BLAS_PREFER_HIPBLASLT=0"], daemon=True) + self.python(f"app.py", current_dir="webui", + env=["TORCH_BLAS_PREFER_HIPBLASLT=0", f"GRADIO_SERVER_PORT={self.port}"], daemon=True) diff --git a/services/comfy_ui.py b/services/comfy_ui.py index 18fb988..11535c5 100644 --- a/services/comfy_ui.py +++ b/services/comfy_ui.py @@ -30,5 +30,5 @@ class ComfyUi(Stack): def _launch(self): args = ["--port", str(self.port)] - self.python(f"main.py {' '.join(args)}", current_dir="webui", + self.python(f"main.py", args=args, current_dir="webui", env=["TORCH_BLAS_PREFER_HIPBLASLT=0"], daemon=True) diff --git a/services/stable_diffusion_forge.py b/services/stable_diffusion_forge.py index b8e2d96..950d843 100644 --- a/services/stable_diffusion_forge.py +++ b/services/stable_diffusion_forge.py @@ -23,5 +23,5 @@ class StableDiffusionForge(Stack): def _launch(self): args = ["--listen", "--enable-insecure-extension-access", "--port", str(self.port)] - self.python(f"launch.py {' '.join(args)}", current_dir="webui", + self.python(f"launch.py", args=args, current_dir="webui", env=["TORCH_BLAS_PREFER_HIPBLASLT=0"], daemon=True) diff --git a/services/stable_diffusion_webui.py b/services/stable_diffusion_webui.py index da84989..5066a8b 100644 --- a/services/stable_diffusion_webui.py +++ b/services/stable_diffusion_webui.py @@ -23,5 +23,5 @@ class StableDiffusionWebui(Stack): def _launch(self): args = ["--listen", "--enable-insecure-extension-access", "--port", str(self.port)] - self.python(f"launch.py {' '.join(args)}", current_dir="webui", + self.python(f"launch.py", args=args, current_dir="webui", env=["TORCH_BLAS_PREFER_HIPBLASLT=0"], daemon=True) diff --git a/services/text_generation_webui.py b/services/text_generation_webui.py index b0c7d12..2c8dd8c 100644 --- a/services/text_generation_webui.py +++ b/services/text_generation_webui.py @@ -52,5 +52,5 @@ class TextGenerationWebui(Stack): def _launch(self): args = ["--listen", "--listen-port", str(self.port)] - self.python(f"server.py {' '.join(args)}", current_dir="webui", + self.python(f"server.py", args=args, current_dir="webui", env=["TORCH_BLAS_PREFER_HIPBLASLT=0"], daemon=True) diff --git a/services/xtts_webui.py b/services/xtts_webui.py index 8ef2480..ffd7ce3 100644 --- a/services/xtts_webui.py +++ b/services/xtts_webui.py @@ -34,5 +34,5 @@ class XttsWebui(Stack): def _launch(self): args = ["--host", "0.0.0.0", "--port", str(self.port)] - self.python(f"server.py {' '.join(args)}", current_dir="webui", - env=["TORCH_BLAS_PREFER_HIPBLASLT=0"], daemon=True) + self.python(f"server.py", current_dir="webui", + env=["TORCH_BLAS_PREFER_HIPBLASLT=0"], args=args, daemon=True) diff --git a/ui/choices.py b/ui/choices.py index 97b5263..3c63e97 100644 --- a/ui/choices.py +++ b/ui/choices.py @@ -10,10 +10,11 @@ install_service = None uninstall_service = None are_you_sure = None any_key = None +already_running = None def update_choices(): - global start, start_service, stop_service, install_service, uninstall_service, are_you_sure, any_key + global start, start_service, stop_service, install_service, uninstall_service, are_you_sure, any_key, already_running start = questionary.select( "Choose an option:", @@ -50,4 +51,6 @@ def update_choices(): are_you_sure = questionary.confirm("Are you sure?") + already_running = questionary.confirm("Service is already running, do you want to restart it?") + any_key = questionary.text("Press any key to continue") diff --git a/ui/interface.py b/ui/interface.py index c7c753d..3f7bb0b 100644 --- a/ui/interface.py +++ b/ui/interface.py @@ -60,6 +60,6 @@ def run_interactive_cmd_ui(): service = choices.uninstall_service.ask() handle_services("uninstall", service) - elif choice == "Exit": + elif choice == "exit": print("Exiting...") exit(0)