Refactor/stop-workers #597

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged

bgunnar5 merged 100 commits into llnl:develop-2.0 from bgunnar5:refactor/stop-workers

Jan 14, 2026

CHANGELOG.md

-Original file line number
+Diff line change
@@ Expand Up @@
     ### Changed
     - Removed old monitor code from v1.0
+    - Moved `stop-workers` functionality to the `CeleryWorkerHandlers` class
     ## [2.0.0b4]
@@ Expand Down @@

merlin/cli/commands/stop_workers.py

-Original file line number
+Diff line change
@@ Expand Up / @@ -19,9 +19,9 @@ @@
     from merlin.ascii_art import banner_small
     from merlin.cli.commands.command_entry_point import CommandEntryPoint
-    from merlin.router import stop_workers
     from merlin.spec.specification import MerlinSpec
     from merlin.utils import verify_filepath
+    from merlin.workers.handlers.handler_factory import worker_handler_factory
     LOG = logging.getLogger("merlin")
@@ Expand Down Expand Up / @@ -67,6 +67,12 @@ def add_parser(self, subparsers: ArgumentParser): @@
                 default=None,
                 help="regex match for specific workers to stop",
             )
+            stop.add_argument(
+                "-d",
+                "--dry-run",
+                action="store_true",
+                help="Display which workers would be stopped without actually stopping them",
+            )
         def process_command(self, args: Namespace):
             """
@@ Expand All / @@ -87,13 +93,27 @@ def process_command(self, args: Namespace): @@
             worker_names = []
             # Load in the spec if one was provided via the CLI
+            spec = None
             if args.spec:
                 spec_path = verify_filepath(args.spec)
                 spec = MerlinSpec.load_specification(spec_path)
                 worker_names = spec.get_worker_names()
                 for worker_name in worker_names:
                     if "$" in worker_name:
                         LOG.warning(f"Worker '{worker_name}' is unexpanded. Target provenance spec instead?")
+                LOG.debug(f"Searching for the following workers to stop based on the spec {args.spec}: {worker_names}")
+            # If we have workers from --workers flag, add them to the list
+            if args.workers:
+                worker_names.extend(args.workers)
-            # Send stop command to router
-            stop_workers(args.task_server, worker_names, args.queues, args.workers)
+            # Get the task server from spec or CLI argument
+            task_server = spec.merlin["resources"]["task_server"] if spec else args.task_server
+            # Create the handler and send stop command
+            worker_handler = worker_handler_factory.create(task_server)
+            worker_handler.stop_workers(
+                queues=args.queues,
+                workers=worker_names if worker_names else None,
+                dry_run=args.dry_run,
+            )

merlin/common/tasks.py

-Original file line number
+Diff line change
@@ Expand Up / @@ -41,13 +41,13 @@ @@
         RestartException,
         RetryException,
     )
-    from merlin.router import stop_workers
     from merlin.spec.expansion import parameter_substitutions_for_cmd, parameter_substitutions_for_sample
     from merlin.study.dag import DAG
     from merlin.study.status import read_status, status_conflict_handler
     from merlin.study.step import Step
     from merlin.study.study import MerlinStudy
     from merlin.utils import dict_deep_merge
+    from merlin.workers.handlers.celery_handler import CeleryWorkerHandler
     retry_exceptions = (
@@ Expand Down Expand Up @@
         name="merlin:shutdown_workers",
         priority=get_priority(Priority.HIGH),
     )
-    def shutdown_workers(self: Task, shutdown_queues: List[str]):  # pylint: disable=W0613
+    def shutdown_workers(self: Task, shutdown_queues: List[str] = None):  # pylint: disable=W0613
         """
         Initiates the shutdown of Celery workers.
-        This task wraps the [`stop_celery_workers`][study.celeryadapter.stop_celery_workers]
-        function, allowing for the graceful shutdown of specified Celery worker queues. It is
+        This task wraps the [`stop_workers`][workers.handlers.celery_handler.CeleryWorkerHandler.stop_workers]
+        method of the [`CeleryWorkerHandler`][workers.handlers.celery_handler.CeleryWorkerHandler]
+        class, allowing for the graceful shutdown of specified Celery worker queues. It is
         acknowledged immediately upon execution, ensuring that it will not be requeued, even
         if executed by a worker.
@@ Expand All @@
             shutdown_queues: A list of specific queues to shut down. If None, all queues will
                 be shut down.
         """
-        if shutdown_queues is not None:
-            LOG.warning(f"Shutting down workers in queues {shutdown_queues}!")
-        else:
-            LOG.warning("Shutting down workers in all queues!")
-        return stop_workers("celery", None, shutdown_queues, None)
+        worker_handler = CeleryWorkerHandler()
+        worker_handler.stop_workers(queues=shutdown_queues)
     # Pylint complains that these args are unused but celery passes args
@@ Expand Down @@

merlin/router.py

-Original file line number
+Diff line change
@@ Expand Up / @@ -21,7 +21,6 @@ @@
         purge_celery_tasks,
         query_celery_queues,
         run_celery,
-        stop_celery_workers,
     )
     from merlin.study.study import MerlinStudy
@@ Expand Down Expand Up / @@ -150,24 +149,3 @@ def query_queues( @@
         else:
             LOG.error("Celery is not specified as the task server!")
             return {}
-    def stop_workers(task_server: str, spec_worker_names: List[str], queues: List[str], workers_regex: str):
-        """
-        This function sends a command to stop workers that match the specified
-        criteria from the designated task server.
-        Args:
-            task_server: The task server from which to stop workers.
-            spec_worker_names: A list of worker names to stop, as defined
-                in a specification.
-            queues: A list of queues from which to stop associated workers.
-            workers_regex: A regex pattern used to filter the workers to stop.
-        """
-        LOG.info("Stopping workers...")
-        if task_server == "celery":  # pylint: disable=R1705
-            # Stop workers
-            stop_celery_workers(queues, spec_worker_names, workers_regex)
-        else:
-            LOG.error("Celery is not specified as the task server!")

merlin/study/celeryadapter.py

-Original file line number
+Diff line change
@@ Expand Up / @@ -20,7 +20,7 @@ @@
     from merlin.config import Config
     from merlin.spec.specification import MerlinSpec
     from merlin.study.study import MerlinStudy
-    from merlin.utils import apply_list_of_regex, get_procs, is_running
+    from merlin.utils import get_procs, is_running
     LOG = logging.getLogger(__name__)
@@ Expand Down Expand Up / @@ -458,82 +458,3 @@ def purge_celery_tasks(queues: str, force: bool) -> int: @@
         purge_command = " ".join(["celery -A merlin purge", force_com, "-Q", queues])
         LOG.debug(purge_command)
         return subprocess.run(purge_command, shell=True).returncode
-    def stop_celery_workers(
-        queues: List[str] = None, spec_worker_names: List[str] = None, worker_regex: List[str] = None
-    ):  # pylint: disable=R0912
-        """
-        Send a stop command to Celery workers.
-        This function sends a shutdown command to Celery workers associated with
-        specified queues. By default, it stops all connected workers, but it can
-        be configured to target specific workers based on queue names or regular
-        expression patterns.
-        Args:
-            queues: A list of queue names to which the stop command will be sent.
-                If None, all connected workers across all queues will be stopped.
-            spec_worker_names: A list of specific worker names to stop, in addition
-                to those matching the `worker_regex`.
-            worker_regex: A regular expression string used to match worker names.
-                If None, no regex filtering will be applied.
-        Side Effects:
-            - Broadcasts a shutdown signal to Celery workers
-        Example:
-            ```python
-            stop_celery_workers(queues=['hello'], worker_regex='celery@*my_machine*')
-            stop_celery_workers()
-            ```
-        """
-        from merlin.celery import app  # pylint: disable=C0415
-        LOG.debug(f"Sending stop to queues: {queues}, worker_regex: {worker_regex}, spec_worker_names: {spec_worker_names}")
-        active_queues, _ = get_active_celery_queues(app)
-        # If not specified, get all the queues
-        if queues is None:
-            queues = [*active_queues]
-        # Celery adds the queue tag in front of each queue so we add that here
-        else:
-            celerize_queues(queues)
-        # Find the set of all workers attached to all of those queues
-        all_workers = set()
-        for queue in queues:
-            try:
-                all_workers.update(active_queues[queue])
-                LOG.debug(f"Workers attached to queue {queue}: {active_queues[queue]}")
-            except KeyError:
-                LOG.warning(f"No workers are connected to queue {queue}")
-        all_workers = list(all_workers)
-        LOG.debug(f"Pre-filter worker stop list: {all_workers}")
-        # Stop workers with no flags
-        if (spec_worker_names is None or len(spec_worker_names) == 0) and worker_regex is None:
-            workers_to_stop = list(all_workers)
-        # Flag handling
-        else:
-            workers_to_stop = []
-            # --spec flag
-            if (spec_worker_names is not None) and len(spec_worker_names) > 0:
-                apply_list_of_regex(spec_worker_names, all_workers, workers_to_stop)
-            # --workers flag
-            if worker_regex is not None:
-                LOG.debug(f"Searching for workers to stop based on the following regex's: {worker_regex}")
-                apply_list_of_regex(worker_regex, all_workers, workers_to_stop)
-        # Remove duplicates
-        workers_to_stop = list(set(workers_to_stop))
-        LOG.debug(f"Post-filter worker stop list: {workers_to_stop}")
-        if workers_to_stop:
-            LOG.info(f"Sending stop to these workers: {workers_to_stop}")
-            # Send the shutdown signal
-            app.control.broadcast("shutdown", destination=workers_to_stop)
-        else:
-            LOG.warning("No workers found to stop")

merlin/study/manager.py

-Original file line number
+Diff line change
@@ Expand Up / @@ -21,7 +21,8 @@ @@
     from merlin.db_scripts.merlin_db import MerlinDatabase
     from merlin.exceptions import RunNotFoundError, StudyNotFoundError
     from merlin.spec.specification import MerlinSpec
-    from merlin.study.celeryadapter import purge_celery_tasks, stop_celery_workers
+    from merlin.study.celeryadapter import purge_celery_tasks
+    from merlin.workers.handlers.celery_handler import CeleryWorkerHandler
     LOG = logging.getLogger(__name__)
@@ Expand Down Expand Up / @@ -97,12 +98,13 @@ def cancel( @@
             # Step 1: Stop the workers
             if stop_workers:
-                # TODO when we refactor `stop-workers`, update this
                 worker_names = spec.get_worker_names()
                 for worker_name in worker_names:
                     if "$" in worker_name:
                         LOG.warning(f"Worker '{worker_name}' is unexpanded. Target provenance spec instead?")
-                stop_celery_workers(spec_worker_names=worker_names)
+                worker_handler = CeleryWorkerHandler()
+                worker_handler.stop_workers(workers=worker_names)
                 # TODO when we refactor `stop-workers`, may want to do some extra validation here to ensure
                 # all of these workers have actually been stopped
@@ Expand Down @@

merlin/workers/celery_worker.py

-Original file line number
+Diff line change
@@ Expand Up / @@ -85,6 +85,7 @@ def __init__( @@
             self.batch = self.config.get("batch", {})
             self.machines = self.config.get("machines", [])
             self.overlap = overlap
+            self.pid = None  # Set when the worker is launched
             # Add this worker to the database
             merlin_db = MerlinDatabase()
@@ Expand Down Expand Up @@
             if self.should_launch():
                 launch_cmd = self.get_launch_command(override_args=override_args, disable_logs=disable_logs)
                 try:
-                    subprocess.Popen(launch_cmd, env=self.env, shell=True, universal_newlines=True)  # pylint: disable=R1732
+                    worker_proc = subprocess.Popen(
+                        launch_cmd, env=self.env, shell=True, universal_newlines=True
+                    )  # pylint: disable=R1732
+                    self.pid = worker_proc.pid
                     LOG.debug(f"Launched worker '{self.name}' with command: {launch_cmd}.")
                 except Exception as e:  # pylint: disable=C0103
                     LOG.error(f"Cannot start celery workers, {e}")
                     raise MerlinWorkerLaunchError from e
+        def stop(self):
+            """
+            Stop the worker process.
+            If the worker has a known PID, sends a SIGTERM to terminate it.
+            Otherwise, logs a warning that the worker cannot be stopped.
+            """
+            if self.pid:
+                try:
+                    os.kill(self.pid, 15)  # Send SIGTERM
+                    LOG.debug(f"Stopped worker '{self.name}' with PID {self.pid}.")
+                    self.pid = None  # Reset PID after stopping
+                except Exception as e:  # pylint: disable=C0103
+                    LOG.error(f"Cannot stop celery worker '{self.name}', {e}")
+            else:
+                LOG.warning(f"Worker '{self.name}' is not running or PID is unknown; cannot stop.")
         def get_metadata(self) -> Dict:
             """
             Return metadata about this worker instance.
@@ Expand Down @@

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Refactor/stop-workers #597

Uh oh!

Diff view

Diff view

There are no files selected for viewing

Uh oh!