Merge branch 'main' into feat/extract_archive

Author: dotsdl

alchemiscale/cli.py

@@ -7,7 +7,6 @@
 import click
 import yaml
 import json
-import signal
 
 from .security.models import (
     CredentialedEntity,
@@ -382,31 +381,23 @@ def synchronous(config_file, name, compute_manager_id):
     from alchemiscale.models import Scope
     from alchemiscale.compute.service import SynchronousComputeService
     from alchemiscale.compute.settings import ComputeServiceSettings
+    from alchemiscale.compute.signals import install_stop_handlers
 
     params = yaml.safe_load(config_file)
 
-    params_init = params.get("init", {})
-    params_start = params.get("start", {})
-
     if name is not None:
-        params_init["name"] = name
+        params["name"] = name
 
     if compute_manager_id is not None:
-        params_init["compute_manager_id"] = compute_manager_id
-
-    service = SynchronousComputeService(ComputeServiceSettings(**params_init))
-
-    # add signal handling
-    for signame in {"SIGHUP", "SIGINT", "SIGTERM"}:
+        params["compute_manager_id"] = compute_manager_id
 
-        def stop(*args, **kwargs):
-            service.stop()
-            raise KeyboardInterrupt()
+    service = SynchronousComputeService(ComputeServiceSettings(**params))
 
-        signal.signal(getattr(signal, signame), stop)
+    # install handlers so SIGHUP/SIGINT/SIGTERM stop the service cleanly
+    install_stop_handlers(service)
 
     try:
-        service.start(**params_start)
+        service.start()
     except KeyboardInterrupt:
         pass
 
@@ -551,6 +542,7 @@ def strategist(config_file):
     from alchemiscale.models import Scope
     from alchemiscale.strategist.service import StrategistService
     from alchemiscale.strategist.settings import StrategistSettings
+    from alchemiscale.compute.signals import install_stop_handlers
 
     params = yaml.safe_load(config_file)
 
@@ -559,18 +551,12 @@ def strategist(config_file):
 
     service = StrategistService(StrategistSettings(**params))
 
-    # add signal handling
-    for signame in {"SIGHUP", "SIGINT", "SIGTERM"}:
-
-        def stop(*args, **kwargs):
-            service.stop()
+    # install handlers so SIGHUP/SIGINT/SIGTERM stop the service cleanly.
+    # do *not* raise KeyboardInterrupt --- the strategist's stop() triggers a
+    # ProcessPoolExecutor shutdown that must not be interrupted partway through.
+    install_stop_handlers(service, raise_keyboard_interrupt=False)
 
-        signal.signal(getattr(signal, signame), stop)
-
-    try:
-        service.start()
-    except KeyboardInterrupt:
-        pass
+    service.start()
 
 
 @cli.group(help="Subcommands for managing identities")

alchemiscale/compute/manager.py

@@ -5,6 +5,7 @@
 """
 
 from abc import abstractmethod
+from contextlib import contextmanager
 import logging
 import time
 
@@ -18,6 +19,7 @@
     AlchemiscaleComputeManagerClientError,
 )
 from .settings import ComputeManagerSettings, ComputeServiceSettings
+from ..sleep import InterruptableSleep, SleepInterrupted
 
 
 class ComputeManager:
@@ -44,6 +46,7 @@ def __init__(
         )
 
         self._stop = False
+        self.int_sleep = InterruptableSleep()
 
         logger = logging.getLogger("AlchemiscaleComputeManager")
         logger.setLevel(self.settings.loglevel)
@@ -75,34 +78,73 @@ def _register(self, steal=False):
     def _deregister(self):
         self.client.deregister(self.compute_manager_id)
 
-    def start(self, max_cycles: int | None = None, steal=False):
-        self.logger.info(f"Starting up compute manager '{self.settings.name}'")
-        self._register(steal=steal)
-        self.logger.info(f"Registered compute manager '{self.compute_manager_id}'")
-        self._stop = False
+    @contextmanager
+    def _running(self, steal=False):
+        """Register this compute manager for the lifetime of the context.
+
+        This guarantees that if anything interrupts startup after registration,
+        including int_sleep.clear(), the manager is still deregistered.
+        """
+        registered = False
+
         try:
-            count = 0
-            self.logger.info("Starting main loop")
-            while self.cycle():
-                count += 1
-                if max_cycles and count >= max_cycles:
-                    self.logger.info("Reached maximum number of cycles")
-                    break
-                self.logger.info(f"Sleeping for {self.settings.sleep_interval} seconds")
-                time.sleep(self.settings.sleep_interval)
-        except Exception as e:
-            self.logger.error(f"Unknown exception raised: '{str(e)}'")
-            self.logger.info(f"Updating manager status to 'ERROR'")
-            self.client.update_status(
-                self.compute_manager_id, ComputeManagerStatus.ERROR, detail=repr(e)
-            )
-            raise e
-        except KeyboardInterrupt:
-            self.logger.info("Caught SIGINT/Keyboard interrupt.")
+            self._register(steal=steal)
+            registered = True
+
+            self.logger.info(f"Registered compute manager '{self.compute_manager_id}'")
+
+            self._stop = False
+            self.int_sleep.clear()
+
+            yield
+
         finally:
-            self.logger.info(f"Deregistering '{self.compute_manager_id}'")
-            self._deregister()
-            self.logger.info(f"Deregistration successful")
+            if registered:
+                self.logger.info(f"Deregistering '{self.compute_manager_id}'")
+
+                # kept here in case we add additional cleanup to stop later, such as other threads
+                self.stop()
+
+                self._deregister()
+                self.logger.info("Deregistration successful")
+
+    def start(self, max_cycles: int | None = None, steal=False):
+        self.logger.info(f"Starting up compute manager '{self.settings.name}'")
+
+        with self._running(steal=steal):
+            try:
+                count = 0
+                self.logger.info("Starting main loop")
+
+                while self.cycle():
+                    count += 1
+
+                    if max_cycles and count >= max_cycles:
+                        self.logger.info("Reached maximum number of cycles")
+                        break
+
+                    self.logger.info(
+                        f"Sleeping for {self.settings.sleep_interval} seconds"
+                    )
+                    self.int_sleep(self.settings.sleep_interval)
+
+            except SleepInterrupted:
+                self.logger.info("Compute manager stopping.")
+
+            except KeyboardInterrupt:
+                self.logger.info("Caught SIGINT/Keyboard interrupt.")
+
+            except Exception as e:
+                self.logger.error(f"Unknown exception raised: '{str(e)}'")
+                self.logger.info("Updating manager status to 'ERROR'")
+
+                self.client.update_status(
+                    self.compute_manager_id,
+                    ComputeManagerStatus.ERROR,
+                    detail=repr(e),
+                )
+
+                raise
 
     @abstractmethod
     def create_compute_services(self, data: dict, target: int) -> int:
@@ -189,6 +231,7 @@ def _compute_jobs_to_create(self, num_tasks: int, num_active_services: int) -> i
         return jobs or 1
 
     def stop(self):
+        self.int_sleep.interrupt()
         self._stop = True
 
     def cycle(self) -> bool: