mathieu/ai-suite-rocm-local
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases 1 Wiki Activity

fixed daemon running system

Browse Source
This commit is contained in:
Mathieu Broillet 2024-08-29 12:03:37 +02:00
parent a6728d3d1c
commit f109e9016f
Signed by: mathieu
GPG Key ID: A08E484FE95074C1
10 changed files with 96 additions and 44 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

31
core/config.py
View File

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

85
core/stack.py
View File

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

5
services/background_removal_dis.py
View File

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

2
services/comfy_ui.py
View File

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

2
services/stable_diffusion_forge.py
View File

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

2
services/stable_diffusion_webui.py
View File

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

2
services/text_generation_webui.py
View File

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

4
services/xtts_webui.py
View File

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

5
ui/choices.py
View File

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

2
ui/interface.py
View File

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