fixed daemon running system

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)

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):
+    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}", current_dir, daemon)
+        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 self.status():
-            if config.has(f"{self.name}-pid"):
+                choice = choices.already_running.ask()
 
-                # Check if PID still running
+                if choice is True:
-                if psutil.pid_exists(config.get(f"{self.name}-pid")):
+                    self.stop()
-                    choice = input(f"{self.name} is already running, do you want to restart it? (y/n): ")
+                    self._launch()
-
+                    return
-                    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
                 else:
-                    logger.warning(
+                    # TODO: attach to subprocess / redirect logs?
-                        f"Previous PID found for {self.name} but process is not running, continuing as stopped...")
+                    logger.info("Continuing without restarting...")
+                    return
             else:
-                logger.debug(f"No previous PID found for {self.name}, continuing as stopped...")
+                logger.debug(f"Running command as daemon: {cmd}")
-
+                cmd = f"{cmd} &"
-            logger.debug(f"Starting {self.name} as daemon with command: {cmd}")
+                process = subprocess.Popen(cmd, shell=True, preexec_fn=os.setpgrp,
-            cmd = f"{cmd} &"
+                                           stdout=config.open_file(f"{self.id}-stdout"),
-            process = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+                                           stderr=config.open_file(f"{self.id}-stderr"))
-            config.set(f"{self.name}-pid", process.pid + 1)
+                self.set_pid(find_correct_pid(process.pid))
-            return
+                return
         else:
             logger.debug(f"Running command: {cmd}")
 

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", current_dir="webui",
-        self.python(f"app.py {' '.join(args)}", current_dir="webui",
+                    env=["TORCH_BLAS_PREFER_HIPBLASLT=0", f"GRADIO_SERVER_PORT={self.port}"], daemon=True)
-                    env=["TORCH_BLAS_PREFER_HIPBLASLT=0"], daemon=True)

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)

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)

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)

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)

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",
+        self.python(f"server.py", current_dir="webui",
-                    env=["TORCH_BLAS_PREFER_HIPBLASLT=0"], daemon=True)
+                    env=["TORCH_BLAS_PREFER_HIPBLASLT=0"], args=args, daemon=True)

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")

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)