fractal-server 2.7.0a9__tar.gz → 2.7.0a10__tar.gz

{fractal_server-2.7.0a9 → fractal_server-2.7.0a10}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: fractal-server
-Version: 2.7.0a9
+Version: 2.7.0a10
 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.7.0a10/fractal_server/__init__.py

@@ -0,0 +1 @@
+__VERSION__ = "2.7.0a10"

{fractal_server-2.7.0a9 → fractal_server-2.7.0a10}/fractal_server/app/runner/executors/slurm/ssh/executor.py

@@ -861,7 +861,7 @@ class FractalSlurmSSHExecutor(SlurmExecutor):
 
         # Transfer archive
         t_0_put = time.perf_counter()
-        self.fractal_ssh.put(
+        self.fractal_ssh.send_file(
             local=tarfile_path_local,
             remote=tarfile_path_remote,
         )
@@ -1255,16 +1255,9 @@ class FractalSlurmSSHExecutor(SlurmExecutor):
             self.workflow_dir_remote / f"{subfolder_name}.tar.gz"
         ).as_posix()
 
-        # Remove local tarfile - FIXME SSH: is this needed?
-        logger.warning(f"In principle I just removed {tarfile_path_local}")
-        logger.warning(f"{Path(tarfile_path_local).exists()=}")
-
-        # Remove remote tarfile - FIXME SSH: is this needed?
-        # rm_command = f"rm {tarfile_path_remote}"
-        # _run_command_over_ssh(cmd=rm_command, fractal_ssh=self.fractal_ssh)
-        logger.warning(f"Unlink {tarfile_path_remote=} - START")
-        self.fractal_ssh.sftp().unlink(tarfile_path_remote)
-        logger.warning(f"Unlink {tarfile_path_remote=} - STOP")
+        # Remove remote tarfile
+        rm_command = f"rm {tarfile_path_remote}"
+        self.fractal_ssh.run_command(cmd=rm_command)
 
         # Create remote tarfile
         tar_command = (
@@ -1278,7 +1271,7 @@ class FractalSlurmSSHExecutor(SlurmExecutor):
 
         # Fetch tarfile
         t_0_get = time.perf_counter()
-        self.fractal_ssh.get(
+        self.fractal_ssh.fetch_file(
             remote=tarfile_path_remote,
             local=tarfile_path_local,
         )
@@ -1291,6 +1284,11 @@ class FractalSlurmSSHExecutor(SlurmExecutor):
         # Extract tarfile locally
         extract_archive(Path(tarfile_path_local))
 
+        # Remove local tarfile
+        if Path(tarfile_path_local).exists():
+            logger.warning(f"Remove existing file {tarfile_path_local}.")
+            Path(tarfile_path_local).unlink()
+
         t_1 = time.perf_counter()
         logger.info("[_get_subfolder_sftp] End - " f"elapsed: {t_1-t_0:.3f} s")
 

{fractal_server-2.7.0a9 → fractal_server-2.7.0a10}/fractal_server/ssh/_fabric.py

@@ -1,3 +1,4 @@
+import json
 import logging
 import time
 from contextlib import contextmanager
@@ -23,21 +24,58 @@ class FractalSSHTimeoutError(RuntimeError):
     pass
 
 
-class FractalSSHListTimeoutError(RuntimeError):
-    pass
+logger = set_logger(__name__)
 
 
-logger = set_logger(__name__)
+@contextmanager
+def _acquire_lock_with_timeout(
+    lock: Lock,
+    label: str,
+    timeout: float,
+    logger_name: str = __name__,
+) -> Generator[Literal[True], Any, None]:
+    """
+    Given a `threading.Lock` object, try to acquire it within a given timeout.
+
+    Arguments:
+        lock:
+        label:
+        timeout:
+        logger_name:
+    """
+    logger = get_logger(logger_name)
+    logger.info(f"Trying to acquire lock for '{label}', with {timeout=}")
+    result = lock.acquire(timeout=timeout)
+    try:
+        if not result:
+            logger.error(f"Lock for '{label}' was *not* acquired.")
+            raise FractalSSHTimeoutError(
+                f"Failed to acquire lock for '{label}' within "
+                f"{timeout} seconds"
+            )
+        logger.info(f"Lock for '{label}' was acquired.")
+        yield result
+    finally:
+        if result:
+            lock.release()
+            logger.info(f"Lock for '{label}' was released.")
 
 
 class FractalSSH(object):
     """
-    FIXME SSH: Fix docstring
+    Wrapper of `fabric.Connection` object, enriched with locks.
+
+    Note: methods marked as `_unsafe` should not be used directly,
+    since they do not enforce locking.
 
     Attributes:
         _lock:
-        connection:
+        _connection:
         default_lock_timeout:
+        default_max_attempts:
+        default_base_interval:
+        sftp_get_prefetch:
+        sftp_get_max_requests:
         logger_name:
     """
 
@@ -46,6 +84,8 @@ class FractalSSH(object):
     default_lock_timeout: float
     default_max_attempts: int
     default_base_interval: float
+    sftp_get_prefetch: bool
+    sftp_get_max_requests: int
     logger_name: str
 
     def __init__(
@@ -54,6 +94,8 @@ class FractalSSH(object):
         default_timeout: float = 250,
         default_max_attempts: int = 5,
         default_base_interval: float = 3.0,
+        sftp_get_prefetch: bool = False,
+        sftp_get_max_requests: int = 64,
         logger_name: str = __name__,
     ):
         self._lock = Lock()
@@ -61,28 +103,11 @@ class FractalSSH(object):
         self.default_lock_timeout = default_timeout
         self.default_base_interval = default_base_interval
         self.default_max_attempts = default_max_attempts
+        self.sftp_get_prefetch = sftp_get_prefetch
+        self.sftp_get_max_requests = sftp_get_max_requests
         self.logger_name = logger_name
         set_logger(self.logger_name)
 
-    @contextmanager
-    def acquire_timeout(
-        self, timeout: float
-    ) -> Generator[Literal[True], Any, None]:
-        self.logger.debug(f"Trying to acquire lock, with {timeout=}")
-        result = self._lock.acquire(timeout=timeout)
-        try:
-            if not result:
-                self.logger.error("Lock was *NOT* acquired.")
-                raise FractalSSHTimeoutError(
-                    f"Failed to acquire lock within {timeout} seconds"
-                )
-            self.logger.debug("Lock was acquired.")
-            yield result
-        finally:
-            if result:
-                self._lock.release()
-                self.logger.debug("Lock was released")
-
     @property
     def is_connected(self) -> bool:
         return self._connection.is_connected
@@ -91,36 +116,82 @@ class FractalSSH(object):
     def logger(self) -> logging.Logger:
         return get_logger(self.logger_name)
 
-    def put(
-        self, *args, lock_timeout: Optional[float] = None, **kwargs
+    def _put(
+        self,
+        *,
+        local: str,
+        remote: str,
+        label: str,
+        lock_timeout: Optional[float] = None,
     ) -> Result:
+        """
+        Transfer a local file to a remote path, via SFTP.
+        """
         actual_lock_timeout = self.default_lock_timeout
         if lock_timeout is not None:
             actual_lock_timeout = lock_timeout
-        with self.acquire_timeout(timeout=actual_lock_timeout):
-            return self._connection.put(*args, **kwargs)
+        with _acquire_lock_with_timeout(
+            lock=self._lock,
+            label=label,
+            timeout=actual_lock_timeout,
+        ):
+            return self._sftp_unsafe().put(local, remote)
 
-    def get(
-        self, *args, lock_timeout: Optional[float] = None, **kwargs
+    def _get(
+        self,
+        *,
+        local: str,
+        remote: str,
+        label: str,
+        lock_timeout: Optional[float] = None,
     ) -> Result:
         actual_lock_timeout = self.default_lock_timeout
         if lock_timeout is not None:
             actual_lock_timeout = lock_timeout
-        with self.acquire_timeout(timeout=actual_lock_timeout):
-            return self._connection.get(*args, **kwargs)
+        with _acquire_lock_with_timeout(
+            lock=self._lock,
+            label=label,
+            timeout=actual_lock_timeout,
+        ):
+            return self._sftp_unsafe().get(
+                remote,
+                local,
+                prefetch=self.sftp_get_prefetch,
+                max_concurrent_prefetch_requests=self.sftp_get_max_requests,
+            )
 
-    def run(
-        self, *args, lock_timeout: Optional[float] = None, **kwargs
+    def _run(
+        self, *args, label: str, 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
-        with self.acquire_timeout(timeout=actual_lock_timeout):
+        with _acquire_lock_with_timeout(
+            lock=self._lock,
+            label=label,
+            timeout=actual_lock_timeout,
+        ):
             return self._connection.run(*args, **kwargs)
 
-    def sftp(self) -> paramiko.sftp_client.SFTPClient:
+    def _sftp_unsafe(self) -> paramiko.sftp_client.SFTPClient:
+        """
+        This is marked as unsafe because you should only use its methods
+        after acquiring a lock.
+        """
         return self._connection.sftp()
 
+    def read_remote_json_file(self, filepath: str) -> dict[str, Any]:
+        self.logger.info(f"START reading remote JSON file {filepath}.")
+        with _acquire_lock_with_timeout(
+            lock=self._lock,
+            label="read_remote_json_file",
+            timeout=self.default_lock_timeout,
+        ):
+            with self._sftp_unsafe().open(filepath, "r") as f:
+                data = json.load(f)
+        self.logger.info(f"END reading remote JSON file {filepath}.")
+        return data
+
     def check_connection(self) -> None:
         """
         Open the SSH connection and handle exceptions.
@@ -131,7 +202,12 @@ class FractalSSH(object):
         """
         if not self._connection.is_connected:
             try:
-                self._connection.open()
+                with _acquire_lock_with_timeout(
+                    lock=self._lock,
+                    label="_connection.open",
+                    timeout=self.default_lock_timeout,
+                ):
+                    self._connection.open()
             except Exception as e:
                 raise RuntimeError(
                     f"Cannot open SSH connection. Original error:\n{str(e)}"
@@ -146,8 +222,12 @@ class FractalSSH(object):
         because we observed cases where `is_connected=False` but the underlying
         `Transport` object was not closed.
         """
-
-        self._connection.close()
+        with _acquire_lock_with_timeout(
+            lock=self._lock,
+            label="_connection.close",
+            timeout=self.default_lock_timeout,
+        ):
+            self._connection.close()
 
         if self._connection.client is not None:
             self._connection.client.close()
@@ -197,8 +277,11 @@ class FractalSSH(object):
             self.logger.info(f"{prefix} START running '{cmd}' over SSH.")
             try:
                 # Case 1: Command runs successfully
-                res = self.run(
-                    cmd, lock_timeout=actual_lock_timeout, hide=True
+                res = self._run(
+                    cmd,
+                    label=f"run {cmd}",
+                    lock_timeout=actual_lock_timeout,
+                    hide=True,
                 )
                 t_1 = time.perf_counter()
                 self.logger.info(
@@ -250,7 +333,6 @@ class FractalSSH(object):
         *,
         local: str,
         remote: str,
-        logger_name: Optional[str] = None,
         lock_timeout: Optional[float] = None,
     ) -> None:
         """
@@ -261,28 +343,65 @@ class FractalSSH(object):
             remote: Target path on remote host
             fractal_ssh: FractalSSH connection object with custom lock
             logger_name: Name of the logger
-
         """
         try:
-            self.put(local=local, remote=remote, lock_timeout=lock_timeout)
+            prefix = "[send_file]"
+            self.logger.info(f"{prefix} START transfer of '{local}' over SSH.")
+            self._put(
+                local=local,
+                remote=remote,
+                lock_timeout=lock_timeout,
+                label=f"send_file {local=} {remote=}",
+            )
+            self.logger.info(f"{prefix} END transfer of '{local}' over SSH.")
         except Exception as e:
-            logger = get_logger(logger_name=logger_name)
-            logger.error(
+            self.logger.error(
                 f"Transferring {local=} to {remote=} over SSH failed.\n"
                 f"Original Error:\n{str(e)}."
             )
             raise e
 
+    def fetch_file(
+        self,
+        *,
+        local: str,
+        remote: str,
+        lock_timeout: Optional[float] = None,
+    ) -> None:
+        """
+        Transfer a file via SSH
+
+        Args:
+            local: Local path to file
+            remote: Target path on remote host
+            logger_name: Name of the logger
+            lock_timeout:
+        """
+        try:
+            prefix = "[fetch_file] "
+            self.logger.info(f"{prefix} START fetching '{remote}' over SSH.")
+            self._get(
+                local=local,
+                remote=remote,
+                lock_timeout=lock_timeout,
+                label=f"fetch_file {local=} {remote=}",
+            )
+            self.logger.info(f"{prefix} END fetching '{remote}' over SSH.")
+        except Exception as e:
+            self.logger.error(
+                f"Transferring {remote=} to {local=} over SSH failed.\n"
+                f"Original Error:\n{str(e)}."
+            )
+            raise e
+
     def mkdir(self, *, folder: str, parents: bool = True) -> None:
         """
         Create a folder remotely via SSH.
 
         Args:
             folder:
-            fractal_ssh:
             parents:
         """
-        # FIXME SSH: try using `mkdir` method of `paramiko.SFTPClient`
         if parents:
             cmd = f"mkdir -p {folder}"
         else:
@@ -339,12 +458,18 @@ class FractalSSH(object):
             contents: File contents
             lock_timeout:
         """
+        self.logger.info(f"START writing to remote file {path}.")
         actual_lock_timeout = self.default_lock_timeout
         if lock_timeout is not None:
             actual_lock_timeout = lock_timeout
-        with self.acquire_timeout(timeout=actual_lock_timeout):
-            with self.sftp().open(filename=path, mode="w") as f:
+        with _acquire_lock_with_timeout(
+            lock=self._lock,
+            label=f"write_remote_file {path=}",
+            timeout=actual_lock_timeout,
+        ):
+            with self._sftp_unsafe().open(filename=path, mode="w") as f:
                 f.write(content)
+        self.logger.info(f"END writing to remote file {path}.")
 
 
 class FractalSSHList(object):
@@ -425,7 +550,11 @@ class FractalSSHList(object):
                     "look_for_keys": False,
                 },
             )
-            with self.acquire_lock_with_timeout():
+            with _acquire_lock_with_timeout(
+                lock=self._lock,
+                label="FractalSSHList.get",
+                timeout=self._timeout,
+            ):
                 self._data[key] = FractalSSH(connection=connection)
                 return self._data[key]
 
@@ -465,7 +594,11 @@ class FractalSSHList(object):
             key_path:
         """
         key = (host, user, key_path)
-        with self.acquire_lock_with_timeout():
+        with _acquire_lock_with_timeout(
+            lock=self._lock,
+            timeout=self._timeout,
+            label="FractalSSHList.remove",
+        ):
             self.logger.info(
                 f"Removing FractalSSH object for {user}@{host} "
                 "from collection."
@@ -492,24 +625,4 @@ class FractalSSHList(object):
                 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_ssh_obj.close()

{fractal_server-2.7.0a9 → fractal_server-2.7.0a10}/fractal_server/tasks/v2/background_operations_ssh.py

@@ -1,4 +1,3 @@
-import json
 import os
 from pathlib import Path
 from tempfile import TemporaryDirectory
@@ -144,6 +143,13 @@ def background_collect_pip_ssh(
         for key, value in task_group.model_dump().items():
             logger.debug(f"task_group.{key}: {value}")
 
+        # `remove_venv_folder_upon_failure` is set to True only if
+        # script 1 goes through, which means that the remote folder
+        # `package_env_dir` did not already exist. If this remote
+        # folder already existed, then script 1 fails and the boolean
+        # flag `remove_venv_folder_upon_failure` remains false.
+        remove_venv_folder_upon_failure = False
+
         # Open a DB session soon, since it is needed for updating `state`
         with next(get_sync_db()) as db:
             try:
@@ -187,12 +193,11 @@ def background_collect_pip_ssh(
                 # long operations that do not use the db
                 db.close()
 
-                # `remove_venv_folder_upon_failure` is set to True only if
-                # script 1 goes through, which means that the remote folder
-                # `package_env_dir` did not already exist. If this remote
-                # folder already existed, then script 1 fails and the boolean
-                # flag `remove_venv_folder_upon_failure` remains false.
-                remove_venv_folder_upon_failure = False
+                # Create remote folder (note that because of `parents=True` we
+                # are in the `no error if existing, make parent directories as
+                # needed` scenario)
+                fractal_ssh.mkdir(folder=tasks_base_dir, parents=True)
+
                 stdout = _customize_and_run_template(
                     script_filename="_1_create_venv.sh",
                     **common_args,
@@ -263,8 +268,9 @@ def background_collect_pip_ssh(
                 ).as_posix()
 
                 # Read and validate remote manifest file
-                with fractal_ssh.sftp().open(manifest_path_remote, "r") as f:
-                    pkg_manifest_dict = json.load(f)
+                pkg_manifest_dict = fractal_ssh.read_remote_json_file(
+                    manifest_path_remote
+                )
                 logger.info(f"collecting - loaded {manifest_path_remote=}")
                 pkg_manifest = ManifestV2(**pkg_manifest_dict)
                 logger.info("collecting - manifest is a valid ManifestV2")

{fractal_server-2.7.0a9 → fractal_server-2.7.0a10}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "fractal-server"
-version = "2.7.0a9"
+version = "2.7.0a10"
 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.7.0a9"
+current_version = "2.7.0a10"
 version_pattern = "MAJOR.MINOR.PATCH[PYTAGNUM]"
 commit_message = "bump version {old_version} -> {new_version}"
 commit = true

fractal_server-2.7.0a9/fractal_server/__init__.py

@@ -1 +0,0 @@
-__VERSION__ = "2.7.0a9"