fractal-server 2.5.1__tar.gz → 2.5.2__tar.gz

{fractal_server-2.5.1 → fractal_server-2.5.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: fractal-server
-Version: 2.5.1
+Version: 2.5.2
 Summary: Server component of the Fractal analytics platform
 Home-page: https://github.com/fractal-analytics-platform/fractal-server
 License: BSD-3-Clause

fractal_server-2.5.2/fractal_server/__init__.py

@@ -0,0 +1 @@
+__VERSION__ = "2.5.2"

{fractal_server-2.5.1 → fractal_server-2.5.2}/fractal_server/app/routes/api/v2/submit.py

@@ -238,6 +238,18 @@ async def apply_workflow(
     await db.merge(job)
     await db.commit()
 
+    # User appropriate FractalSSH object
+    if settings.FRACTAL_RUNNER_BACKEND == "slurm_ssh":
+        ssh_credentials = dict(
+            user=settings.FRACTAL_SLURM_SSH_USER,
+            host=settings.FRACTAL_SLURM_SSH_HOST,
+            key_path=settings.FRACTAL_SLURM_SSH_PRIVATE_KEY_PATH,
+        )
+        fractal_ssh_list = request.app.state.fractal_ssh_list
+        fractal_ssh = fractal_ssh_list.get(**ssh_credentials)
+    else:
+        fractal_ssh = None
+
     background_tasks.add_task(
         submit_workflow,
         workflow_id=workflow.id,
@@ -246,7 +258,7 @@ async def apply_workflow(
         worker_init=job.worker_init,
         slurm_user=user.slurm_user,
         user_cache_dir=user.cache_dir,
-        fractal_ssh=request.app.state.fractal_ssh,
+        fractal_ssh=fractal_ssh,
     )
     request.app.state.jobsV2.append(job.id)
     logger.info(

{fractal_server-2.5.1 → fractal_server-2.5.2}/fractal_server/app/routes/api/v2/task_collection.py

@@ -124,11 +124,20 @@ async def collect_tasks_pip(
         db.add(state)
         await db.commit()
 
+        # User appropriate FractalSSH object
+        ssh_credentials = dict(
+            user=settings.FRACTAL_SLURM_SSH_USER,
+            host=settings.FRACTAL_SLURM_SSH_HOST,
+            key_path=settings.FRACTAL_SLURM_SSH_PRIVATE_KEY_PATH,
+        )
+        fractal_ssh_list = request.app.state.fractal_ssh_list
+        fractal_ssh = fractal_ssh_list.get(**ssh_credentials)
+
         background_tasks.add_task(
             background_collect_pip_ssh,
             state.id,
             task_pkg,
-            request.app.state.fractal_ssh,
+            fractal_ssh,
         )
 
         response.status_code = status.HTTP_201_CREATED

{fractal_server-2.5.1 → fractal_server-2.5.2}/fractal_server/app/runner/executors/slurm/ssh/executor.py

@@ -163,17 +163,34 @@ class FractalSlurmSSHExecutor(SlurmExecutor):
         settings = Inject(get_settings)
         self.python_remote = settings.FRACTAL_SLURM_WORKER_PYTHON
         if self.python_remote is None:
+            self._stop_and_join_wait_thread()
             raise ValueError("FRACTAL_SLURM_WORKER_PYTHON is not set. Exit.")
 
         # Initialize connection and perform handshake
         self.fractal_ssh = fractal_ssh
         logger.warning(self.fractal_ssh)
-        self.handshake()
+        try:
+            self.handshake()
+        except Exception as e:
+            logger.warning(
+                "Stop/join waiting thread and then "
+                f"re-raise original error {str(e)}"
+            )
+            self._stop_and_join_wait_thread()
+            raise e
 
         # Set/validate parameters for SLURM submission scripts
         self.slurm_account = slurm_account
         self.common_script_lines = common_script_lines or []
-        self._validate_common_script_lines()
+        try:
+            self._validate_common_script_lines()
+        except Exception as e:
+            logger.warning(
+                "Stop/join waiting thread and then "
+                f"re-raise original error {str(e)}"
+            )
+            self._stop_and_join_wait_thread()
+            raise e
 
         # Set/initialize some more options
         self.keep_pickle_files = keep_pickle_files
@@ -1385,6 +1402,10 @@ class FractalSlurmSSHExecutor(SlurmExecutor):
             self.fractal_ssh.run_command(cmd=scancel_command)
         logger.debug("Executor shutdown: end")
 
+    def _stop_and_join_wait_thread(self):
+        self.wait_thread.stop()
+        self.wait_thread.join()
+
     def __exit__(self, *args, **kwargs):
         """
         See
@@ -1393,8 +1414,7 @@ class FractalSlurmSSHExecutor(SlurmExecutor):
         logger.debug(
             "[FractalSlurmSSHExecutor.__exit__] Stop and join `wait_thread`"
         )
-        self.wait_thread.stop()
-        self.wait_thread.join()
+        self._stop_and_join_wait_thread()
         logger.debug("[FractalSlurmSSHExecutor.__exit__] End")
 
     def run_squeue(self, job_ids):

{fractal_server-2.5.1 → fractal_server-2.5.2}/fractal_server/app/runner/executors/slurm/sudo/executor.py

@@ -259,6 +259,7 @@ class FractalSlurmExecutor(SlurmExecutor):
                 for line in self.common_script_lines
                 if line.startswith("#SBATCH --account=")
             )
+            self._stop_and_join_wait_thread()
             raise RuntimeError(
                 "Invalid line in `FractalSlurmExecutor.common_script_lines`: "
                 f"'{invalid_line}'.\n"
@@ -1287,6 +1288,10 @@ class FractalSlurmExecutor(SlurmExecutor):
 
         logger.debug("Executor shutdown: end")
 
+    def _stop_and_join_wait_thread(self):
+        self.wait_thread.stop()
+        self.wait_thread.join()
+
     def __exit__(self, *args, **kwargs):
         """
         See
@@ -1295,6 +1300,5 @@ class FractalSlurmExecutor(SlurmExecutor):
         logger.debug(
             "[FractalSlurmExecutor.__exit__] Stop and join `wait_thread`"
         )
-        self.wait_thread.stop()
-        self.wait_thread.join()
+        self._stop_and_join_wait_thread()
         logger.debug("[FractalSlurmExecutor.__exit__] End")

{fractal_server-2.5.1 → fractal_server-2.5.2}/fractal_server/main.py

@@ -92,32 +92,34 @@ async def lifespan(app: FastAPI):
     settings = Inject(get_settings)
 
     if settings.FRACTAL_RUNNER_BACKEND == "slurm_ssh":
-        from fractal_server.ssh._fabric import get_ssh_connection
-        from fractal_server.ssh._fabric import FractalSSH
 
-        connection = get_ssh_connection()
-        app.state.fractal_ssh = FractalSSH(connection=connection)
+        from fractal_server.ssh._fabric import FractalSSHList
+
+        app.state.fractal_ssh_list = FractalSSHList()
+
         logger.info(
-            f"Created SSH connection "
-            f"({app.state.fractal_ssh.is_connected=})."
+            "Added empty FractalSSHList to app.state "
+            f"(id={id(app.state.fractal_ssh_list)})."
         )
     else:
-        app.state.fractal_ssh = None
+        app.state.fractal_ssh_list = None
 
     config_uvicorn_loggers()
     logger.info("End application startup")
     reset_logger_handlers(logger)
+
     yield
+
     logger = get_logger("fractal_server.lifespan")
     logger.info("Start application shutdown")
 
     if settings.FRACTAL_RUNNER_BACKEND == "slurm_ssh":
         logger.info(
-            f"Closing SSH connection "
-            f"(current: {app.state.fractal_ssh.is_connected=})."
+            "Close FractalSSH connections "
+            f"(current size: {app.state.fractal_ssh_list.size})."
         )
 
-        app.state.fractal_ssh.close()
+        app.state.fractal_ssh_list.close_all()
 
     logger.info(
         f"Current worker with pid {os.getpid()} is shutting down. "

{fractal_server-2.5.1 → fractal_server-2.5.2}/fractal_server/ssh/_fabric.py

@@ -16,20 +16,21 @@ from paramiko.ssh_exception import NoValidConnectionsError
 
 from ..logger import get_logger
 from ..logger import set_logger
-from fractal_server.config import get_settings
 from fractal_server.string_tools import validate_cmd
-from fractal_server.syringe import Inject
 
 
 class FractalSSHTimeoutError(RuntimeError):
     pass
 
 
+class FractalSSHListTimeoutError(RuntimeError):
+    pass
+
+
 logger = set_logger(__name__)
 
 
 class FractalSSH(object):
-
     """
     FIXME SSH: Fix docstring
 
@@ -111,7 +112,6 @@ class FractalSSH(object):
     def run(
         self, *args, lock_timeout: Optional[float] = None, **kwargs
     ) -> Any:
-
         actual_lock_timeout = self.default_lock_timeout
         if lock_timeout is not None:
             actual_lock_timeout = lock_timeout
@@ -138,7 +138,19 @@ class FractalSSH(object):
                 )
 
     def close(self) -> None:
-        return self._connection.close()
+        """
+        Aggressively close `self._connection`.
+
+        When `Connection.is_connected` is `False`, `Connection.close()` does
+        not call `Connection.client.close()`. Thus we do this explicitly here,
+        because we observed cases where `is_connected=False` but the underlying
+        `Transport` object was not closed.
+        """
+
+        self._connection.close()
+
+        if self._connection.client is not None:
+            self._connection.client.close()
 
     def run_command(
         self,
@@ -335,36 +347,169 @@ class FractalSSH(object):
                 f.write(content)
 
 
-def get_ssh_connection(
-    *,
-    host: Optional[str] = None,
-    user: Optional[str] = None,
-    key_filename: Optional[str] = None,
-) -> Connection:
+class FractalSSHList(object):
     """
-    Create a `fabric.Connection` object based on fractal-server settings
-    or explicit arguments.
+    Collection of `FractalSSH` objects
 
-    Args:
-        host:
-        user:
-        key_filename:
+    Attributes are all private, and access to this collection must be
+    through methods (mostly the `get` one).
 
-    Returns:
-        Fabric connection object
+    Attributes:
+        _data:
+            Mapping of unique keys (the SSH-credentials tuples) to
+            `FractalSSH` objects.
+        _lock:
+            A `threading.Lock object`, to be acquired when changing `_data`.
+        _timeout: Timeout for `_lock` acquisition.
+        _logger_name: Logger name.
     """
-    settings = Inject(get_settings)
-    if host is None:
-        host = settings.FRACTAL_SLURM_SSH_HOST
-    if user is None:
-        user = settings.FRACTAL_SLURM_SSH_USER
-    if key_filename is None:
-        key_filename = settings.FRACTAL_SLURM_SSH_PRIVATE_KEY_PATH
-
-    connection = Connection(
-        host=host,
-        user=user,
-        forward_agent=False,
-        connect_kwargs={"key_filename": key_filename},
-    )
-    return connection
+
+    _data: dict[tuple[str, str, str], FractalSSH]
+    _lock: Lock
+    _timeout: float
+    _logger_name: str
+
+    def __init__(
+        self,
+        *,
+        timeout: float = 5.0,
+        logger_name: str = "fractal_server.FractalSSHList",
+    ):
+        self._lock = Lock()
+        self._data = {}
+        self._timeout = timeout
+        self._logger_name = logger_name
+        set_logger(self._logger_name)
+
+    @property
+    def logger(self) -> logging.Logger:
+        """
+        This property exists so that we never have to propagate the
+        `Logger` object.
+        """
+        return get_logger(self._logger_name)
+
+    @property
+    def size(self) -> int:
+        """
+        Number of current key-value pairs in `self._data`.
+        """
+        return len(self._data.values())
+
+    def get(self, *, host: str, user: str, key_path: str) -> FractalSSH:
+        """
+        Get the `FractalSSH` for the current credentials, or create one.
+
+        Note: Changing `_data` requires acquiring `_lock`.
+
+        Arguments:
+            host:
+            user:
+            key_path:
+        """
+        key = (host, user, key_path)
+        fractal_ssh = self._data.get(key, None)
+        if fractal_ssh is not None:
+            self.logger.info(
+                f"Return existing FractalSSH object for {user}@{host}"
+            )
+            return fractal_ssh
+        else:
+            self.logger.info(f"Add new FractalSSH object for {user}@{host}")
+            connection = Connection(
+                host=host,
+                user=user,
+                forward_agent=False,
+                connect_kwargs={
+                    "key_filename": key_path,
+                    "look_for_keys": False,
+                },
+            )
+            with self.acquire_lock_with_timeout():
+                self._data[key] = FractalSSH(connection=connection)
+                return self._data[key]
+
+    def contains(
+        self,
+        *,
+        host: str,
+        user: str,
+        key_path: str,
+    ) -> bool:
+        """
+        Return whether a given key is present in the collection.
+
+        Arguments:
+            host:
+            user:
+            key_path:
+        """
+        key = (host, user, key_path)
+        return key in self._data.keys()
+
+    def remove(
+        self,
+        *,
+        host: str,
+        user: str,
+        key_path: str,
+    ) -> None:
+        """
+        Remove a key from `_data` and close the corresponding connection.
+
+        Note: Changing `_data` requires acquiring `_lock`.
+
+        Arguments:
+            host:
+            user:
+            key_path:
+        """
+        key = (host, user, key_path)
+        with self.acquire_lock_with_timeout():
+            self.logger.info(
+                f"Removing FractalSSH object for {user}@{host} "
+                "from collection."
+            )
+            fractal_ssh_obj = self._data.pop(key)
+            self.logger.info(
+                f"Closing FractalSSH object for {user}@{host} "
+                f"({fractal_ssh_obj.is_connected=})."
+            )
+            fractal_ssh_obj.close()
+
+    def close_all(self, *, timeout: float = 5.0):
+        """
+        Close all `FractalSSH` objects in the collection.
+
+        Arguments:
+            timeout:
+                Timeout for `FractalSSH._lock` acquisition, to be obtained
+                before closing.
+        """
+        for key, fractal_ssh_obj in self._data.items():
+            host, user, _ = key[:]
+            self.logger.info(
+                f"Closing FractalSSH object for {user}@{host} "
+                f"({fractal_ssh_obj.is_connected=})."
+            )
+            with fractal_ssh_obj.acquire_timeout(timeout=timeout):
+                fractal_ssh_obj.close()
+
+    @contextmanager
+    def acquire_lock_with_timeout(self) -> Generator[Literal[True], Any, None]:
+        self.logger.debug(
+            f"Trying to acquire lock, with timeout {self._timeout} s"
+        )
+        result = self._lock.acquire(timeout=self._timeout)
+        try:
+            if not result:
+                self.logger.error("Lock was *NOT* acquired.")
+                raise FractalSSHListTimeoutError(
+                    f"Failed to acquire lock within {self._timeout} ss"
+                )
+            self.logger.debug("Lock was acquired.")
+            yield result
+        finally:
+            if result:
+                self._lock.release()
+                self.logger.debug("Lock was released")

{fractal_server-2.5.1 → fractal_server-2.5.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "fractal-server"
-version = "2.5.1"
+version = "2.5.2"
 description = "Server component of the Fractal analytics platform"
 authors = [
     "Tommaso Comparin <tommaso.comparin@exact-lab.it>",
@@ -92,7 +92,7 @@ filterwarnings = [
 ]
 
 [tool.bumpver]
-current_version = "2.5.1"
+current_version = "2.5.2"
 version_pattern = "MAJOR.MINOR.PATCH[PYTAGNUM]"
 commit_message = "bump version {old_version} -> {new_version}"
 commit = true

fractal_server-2.5.1/fractal_server/__init__.py

@@ -1 +0,0 @@
-__VERSION__ = "2.5.1"