Merge pull request #24 from Never-Over/bridge-stop

restart celery/flower
gauge-sh · Apr 16, 2024 · 4f1b0da · 4f1b0da
2 parents 2ff66ca + 10ba542
commit 4f1b0da
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 27 deletions.
diff --git a/bridge/cli/stop.py b/bridge/cli/stop.py
@@ -29,6 +29,7 @@ def stop():
                         pass
             os.remove(cid_path)
         # Processes - celery, flower
+        # TODO make this a helper method
         for proc in psutil.process_iter(["pid", "name", "cmdline"]):
             try:
                 # Check if the name fragment is in the command line; this field is a list

diff --git a/bridge/framework/django.py b/bridge/framework/django.py
@@ -1,11 +1,12 @@
-import contextlib
 import os
+import signal
 import socket
 import subprocess
 import sys
 from time import sleep
 from typing import Any
 
+import psutil
 from rich.console import Console
 
 from bridge.config import get_config
@@ -146,18 +147,28 @@ def start_local_worker(self) -> None:
             with log_task("Starting local worker", "Local worker started"):
                 from bridge.service.django_celery import app
 
-                # Check if celery is already running
-                if not app.control.inspect().ping():
-                    subprocess.Popen(
-                        "nohup "
-                        "celery -A bridge.service.django_celery worker -c 1 -l INFO"
-                        " > /dev/null 2>&1 &",
-                        shell=True,
-                        stdin=subprocess.DEVNULL,
-                        stdout=subprocess.DEVNULL,
-                        stderr=subprocess.STDOUT,
-                        start_new_session=True,
-                    )
+                # Kill any active celery processes
+                for proc in psutil.process_iter(["pid", "name", "cmdline"]):
+                    try:
+                        proc_name = proc.info["cmdline"]
+                        if (
+                            proc_name
+                            and "bridge.service.django_celery" in proc_name
+                            and "worker" in proc_name
+                        ):
+                            proc.send_signal(signal.SIGTERM)
+                    except (psutil.NoSuchProcess, psutil.AccessDenied):
+                        pass
+                subprocess.Popen(
+                    "nohup "
+                    "celery -A bridge.service.django_celery worker -c 1 -l INFO"
+                    " > /dev/null 2>&1 &",
+                    shell=True,
+                    stdin=subprocess.DEVNULL,
+                    stdout=subprocess.DEVNULL,
+                    stderr=subprocess.STDOUT,
+                    start_new_session=True,
+                )
                 while not app.control.inspect().ping():
                     # Wait for celery to start
                     sleep(0.1)
@@ -175,20 +186,30 @@ def start_local_flower(self) -> None:
                 "[white]bridge_flower[/white]..."
             )
             with log_task("Starting flower", "Flower started"):
-                # Account for flower already running
-                with contextlib.suppress(OSError):
-                    dot_bridge_path = resolve_dot_bridge() / "flower_db"
-                    subprocess.Popen(
-                        "nohup "
-                        "celery -A bridge.service.django_celery flower "
-                        f"--persistent=True --db='{dot_bridge_path}'"
-                        " > /dev/null 2>&1 &",
-                        shell=True,
-                        stdin=subprocess.DEVNULL,
-                        stdout=subprocess.DEVNULL,
-                        stderr=subprocess.STDOUT,
-                        start_new_session=True,
-                    )
+                # Kill any active flower process
+                for proc in psutil.process_iter(["pid", "name", "cmdline"]):
+                    try:
+                        proc_name = proc.info["cmdline"]
+                        if (
+                            proc_name
+                            and "bridge.service.django_celery" in proc_name
+                            and "flower" in proc_name
+                        ):
+                            proc.send_signal(signal.SIGTERM)
+                    except (psutil.NoSuchProcess, psutil.AccessDenied):
+                        pass
+                dot_bridge_path = resolve_dot_bridge() / "flower_db"
+                subprocess.Popen(
+                    "nohup "
+                    "celery -A bridge.service.django_celery flower "
+                    f"--persistent=True --db='{dot_bridge_path}'"
+                    " > /dev/null 2>&1 &",
+                    shell=True,
+                    stdin=subprocess.DEVNULL,
+                    stdout=subprocess.DEVNULL,
+                    stderr=subprocess.STDOUT,
+                    start_new_session=True,
+                )
                 port_bound = False
                 while not port_bound:
                     with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: