fractal-server 2.8.0__py3-none-any.whl → 2.9.0__py3-none-any.whl

fractal_server/ssh/_fabric.py

@@ -11,7 +11,6 @@ from typing import Optional
 
 import paramiko.sftp_client
 from fabric import Connection
-from fabric import Result
 from invoke import UnexpectedExit
 from paramiko.ssh_exception import NoValidConnectionsError
 
@@ -116,50 +115,34 @@ class FractalSSH(object):
     def logger(self) -> logging.Logger:
         return get_logger(self.logger_name)
 
-    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.
+    def log_and_raise(self, *, e: Exception, message: str) -> None:
         """
-        actual_lock_timeout = self.default_lock_timeout
-        if lock_timeout is not None:
-            actual_lock_timeout = lock_timeout
-        with _acquire_lock_with_timeout(
-            lock=self._lock,
-            label=label,
-            timeout=actual_lock_timeout,
-        ):
-            return self._sftp_unsafe().put(local, remote)
+        Log and re-raise an exception from a FractalSSH method.
 
-    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 _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,
+        Arguments:
+            message: Additional message to be logged.
+            e: Original exception
+        """
+        try:
+            self.logger.error(message)
+            self.logger.error(f"Original Error {type(e)} : \n{str(e)}")
+            # Handle the specific case of `NoValidConnectionsError`s from
+            # paramiko, which store relevant information in the `errors`
+            # attribute
+            if hasattr(e, "errors"):
+                self.logger.error(f"{type(e)=}")
+                for err in e.errors:
+                    self.logger.error(f"{err}")
+        except Exception as exception:
+            # Handle unexpected cases, e.g. (1) `e` has no `type`, or
+            # (2) `errors` is not iterable.
+            self.logger.error(
+                "Unexpected Error while handling exception above: "
+                f"{str(exception)}"
             )
 
+        raise e
+
     def _run(
         self, *args, label: str, lock_timeout: Optional[float] = None, **kwargs
     ) -> Any:
@@ -187,8 +170,17 @@ class FractalSSH(object):
             label="read_remote_json_file",
             timeout=self.default_lock_timeout,
         ):
-            with self._sftp_unsafe().open(filepath, "r") as f:
-                data = json.load(f)
+
+            try:
+                with self._sftp_unsafe().open(filepath, "r") as f:
+                    data = json.load(f)
+            except Exception as e:
+                self.log_and_raise(
+                    e=e,
+                    message=(
+                        f"Error in `read_remote_json_file`, for {filepath=}."
+                    ),
+                )
         self.logger.info(f"END reading remote JSON file {filepath}.")
         return data
 
@@ -196,22 +188,56 @@ class FractalSSH(object):
         """
         Open the SSH connection and handle exceptions.
 
-        This function can be called from within other functions that use
-        `connection`, so that we can provide a meaningful error in case the
-        SSH connection cannot be opened.
+        This method should always be called at the beginning of background
+        operations that use FractalSSH, so that:
+
+        1. We try to restore unusable connections (e.g. due to closed socket).
+        2. We provide an informative error if connection cannot be established.
         """
-        if not self._connection.is_connected:
+        self.logger.debug(
+            f"[check_connection] {self._connection.is_connected=}"
+        )
+        if self._connection.is_connected:
+            # Even if the connection appears open, it could be broken for
+            # external reasons (e.g. the socket is closed because the SSH
+            # server was restarted). In these cases, we catch the error and
+            # try to re-open the connection.
             try:
-                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)}"
+                self.logger.info(
+                    "[check_connection] Run dummy command to check connection."
                 )
+                # Run both an SFTP and an SSH command, as they correspond to
+                # different sockets
+                self.remote_exists("/dummy/path/")
+                self.run_command(cmd="whoami")
+                self.logger.info(
+                    "[check_connection] SSH connection is already OK, exit."
+                )
+                return
+            except (OSError, EOFError) as e:
+                self.logger.warning(
+                    f"[check_connection] Detected error {str(e)}, re-open."
+                )
+        # Try opening the connection (if it was closed) or to re-open it (if
+        # an error happened).
+        try:
+            self.close()
+            with _acquire_lock_with_timeout(
+                lock=self._lock,
+                label="_connection.open",
+                timeout=self.default_lock_timeout,
+                logger_name=self.logger_name,
+            ):
+                self._connection.open()
+                self._connection.client.open_sftp()
+                self.logger.info(
+                    "[check_connection] SSH connection opened, exit."
+                )
+
+        except Exception as e:
+            raise RuntimeError(
+                f"Cannot open SSH connection. Original error:\n{str(e)}"
+            )
 
     def close(self) -> None:
         """
@@ -228,9 +254,8 @@ class FractalSSH(object):
             timeout=self.default_lock_timeout,
         ):
             self._connection.close()
-
-        if self._connection.client is not None:
-            self._connection.client.close()
+            if self._connection.client is not None:
+                self._connection.client.close()
 
     def run_command(
         self,
@@ -288,8 +313,10 @@ class FractalSSH(object):
                     f"{prefix} END   running '{cmd}' over SSH, "
                     f"elapsed {t_1-t_0:.3f}"
                 )
-                self.logger.debug(f"STDOUT: {res.stdout}")
-                self.logger.debug(f"STDERR: {res.stderr}")
+                self.logger.debug("STDOUT:")
+                self.logger.debug(res.stdout)
+                self.logger.debug("STDERR:")
+                self.logger.debug(res.stderr)
                 return res.stdout
             except NoValidConnectionsError as e:
                 # Case 2: Command fails with a connection error
@@ -345,21 +372,29 @@ class FractalSSH(object):
             logger_name: Name of the logger
         """
         try:
-            prefix = "[send_file]"
-            self.logger.info(f"{prefix} START transfer of '{local}' over SSH.")
-            self._put(
-                local=local,
-                remote=remote,
-                lock_timeout=lock_timeout,
+            self.logger.info(
+                f"[send_file] START transfer of '{local}' over SSH."
+            )
+            actual_lock_timeout = self.default_lock_timeout
+            if lock_timeout is not None:
+                actual_lock_timeout = lock_timeout
+            with _acquire_lock_with_timeout(
+                lock=self._lock,
                 label=f"send_file {local=} {remote=}",
+                timeout=actual_lock_timeout,
+            ):
+                self._sftp_unsafe().put(local, remote)
+            self.logger.info(
+                f"[send_file] END transfer of '{local}' over SSH."
             )
-            self.logger.info(f"{prefix} END transfer of '{local}' over SSH.")
         except Exception as e:
-            self.logger.error(
-                f"Transferring {local=} to {remote=} over SSH failed.\n"
-                f"Original Error:\n{str(e)}."
+            self.log_and_raise(
+                e=e,
+                message=(
+                    "Error in `send_file`, while "
+                    f"transferring {local=} to {remote=}."
+                ),
             )
-            raise e
 
     def fetch_file(
         self,
@@ -380,19 +415,29 @@ class FractalSSH(object):
         try:
             prefix = "[fetch_file] "
             self.logger.info(f"{prefix} START fetching '{remote}' over SSH.")
-            self._get(
-                local=local,
-                remote=remote,
-                lock_timeout=lock_timeout,
+            actual_lock_timeout = self.default_lock_timeout
+            if lock_timeout is not None:
+                actual_lock_timeout = lock_timeout
+            with _acquire_lock_with_timeout(
+                lock=self._lock,
                 label=f"fetch_file {local=} {remote=}",
-            )
+                timeout=actual_lock_timeout,
+            ):
+                self._sftp_unsafe().get(
+                    remote,
+                    local,
+                    prefetch=self.sftp_get_prefetch,
+                    max_concurrent_prefetch_requests=self.sftp_get_max_requests,  # noqa E501
+                )
             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)}."
+            self.log_and_raise(
+                e=e,
+                message=(
+                    "Error in `fetch_file`, while "
+                    f"Transferring {remote=} to {local=}."
+                ),
             )
-            raise e
 
     def mkdir(self, *, folder: str, parents: bool = True) -> None:
         """
@@ -467,10 +512,38 @@ class FractalSSH(object):
             label=f"write_remote_file {path=}",
             timeout=actual_lock_timeout,
         ):
-            with self._sftp_unsafe().open(filename=path, mode="w") as f:
-                f.write(content)
+            try:
+                with self._sftp_unsafe().open(filename=path, mode="w") as f:
+                    f.write(content)
+            except Exception as e:
+                self.log_and_raise(
+                    e=e, message=f"Error in `write_remote_file`, for {path=}."
+                )
+
         self.logger.info(f"END writing to remote file {path}.")
 
+    def remote_exists(self, path: str) -> bool:
+        """
+        Return whether a remote file/folder exists
+        """
+        self.logger.info(f"START remote_file_exists {path}")
+        with _acquire_lock_with_timeout(
+            lock=self._lock,
+            label=f"remote_file_exists {path=}",
+            timeout=self.default_lock_timeout,
+        ):
+            try:
+                self._sftp_unsafe().stat(path)
+                self.logger.info(f"END   remote_file_exists {path} / True")
+                return True
+            except FileNotFoundError:
+                self.logger.info(f"END   remote_file_exists {path} / False")
+                return False
+            except Exception as e:
+                self.log_and_raise(
+                    e=e, message=f"Error in `remote_exists`, for {path=}."
+                )
+
 
 class FractalSSHList(object):
     """

fractal_server/string_tools.py

@@ -54,7 +54,12 @@ def slugify_task_name_for_source_v1(task_name: str) -> str:
     return task_name.replace(" ", "_").lower()
 
 
-def validate_cmd(command: str, allow_char: Optional[str] = None):
+def validate_cmd(
+    command: str,
+    *,
+    allow_char: Optional[str] = None,
+    attribute_name: str = "Command",
+):
     """
     Assert that the provided `command` does not contain any of the forbidden
     characters for commands
@@ -63,6 +68,7 @@ def validate_cmd(command: str, allow_char: Optional[str] = None):
     Args:
         command: command to validate.
         allow_char: chars to accept among the forbidden ones
+        attribute_name: Name of the attribute, to be used in error message.
     """
     if not isinstance(command, str):
         raise ValueError(f"{command=} is not a string.")
@@ -71,6 +77,7 @@ def validate_cmd(command: str, allow_char: Optional[str] = None):
         forbidden = forbidden - set(allow_char)
     if not forbidden.isdisjoint(set(command)):
         raise ValueError(
-            f"Command must not contain any of this characters: '{forbidden}'\n"
-            f"Provided command: '{command}'."
+            f"{attribute_name} must not contain any of this characters: "
+            f"'{forbidden}'\n"
+            f"Provided {attribute_name.lower()}: '{command}'."
         )

fractal_server/tasks/utils.py

@@ -30,19 +30,9 @@ def get_log_path(base: Path) -> Path:
     return base / COLLECTION_LOG_FILENAME
 
 
-def get_freeze_path(base: Path) -> Path:
-    return base / COLLECTION_FREEZE_FILENAME
-
-
 def get_collection_log_v1(path: Path) -> str:
     package_path = get_absolute_venv_path_v1(path)
     log_path = get_log_path(package_path)
-    log = log_path.open().read()
+    with log_path.open("r") as f:
+        log = f.read()
     return log
-
-
-def get_collection_freeze_v1(venv_path: Path) -> str:
-    package_path = get_absolute_venv_path_v1(venv_path)
-    freeze_path = get_freeze_path(package_path)
-    freeze = freeze_path.open().read()
-    return freeze

fractal_server/tasks/v2/local/__init__.py

@@ -0,0 +1,3 @@
+from .collect import collect_local  # noqa
+from .deactivate import deactivate_local  # noqa
+from .reactivate import reactivate_local  # noqa

fractal_server/tasks/v2/local/_utils.py

@@ -0,0 +1,70 @@
+from pathlib import Path
+
+from fractal_server.app.schemas.v2 import TaskCreateV2
+from fractal_server.logger import get_logger
+from fractal_server.tasks.v2.utils_templates import customize_template
+from fractal_server.utils import execute_command_sync
+
+
+def _customize_and_run_template(
+    template_filename: str,
+    replacements: list[tuple[str, str]],
+    script_dir: str,
+    logger_name: str,
+    prefix: int,
+) -> str:
+    """
+    Customize one of the template bash scripts.
+
+    Args:
+        template_filename: Filename of the template file (ends with ".sh").
+        replacements: Dictionary of replacements.
+        script_dir: Local folder where the script will be placed.
+        prefix: Prefix for the script filename.
+    """
+    logger = get_logger(logger_name=logger_name)
+    logger.debug(f"_customize_and_run_template {template_filename} - START")
+
+    # Prepare name and path of script
+    if not template_filename.endswith(".sh"):
+        raise ValueError(
+            f"Invalid {template_filename=} (it must end with '.sh')."
+        )
+
+    script_filename = f"{prefix}{template_filename}"
+    script_path_local = Path(script_dir) / script_filename
+    # Read template
+    customize_template(
+        template_name=template_filename,
+        replacements=replacements,
+        script_path=script_path_local,
+    )
+    cmd = f"bash {script_path_local}"
+    logger.debug(f"Now run '{cmd}' ")
+    stdout = execute_command_sync(command=cmd, logger_name=logger_name)
+    logger.debug(f"_customize_and_run_template {template_filename} - END")
+    return stdout
+
+
+def check_task_files_exist(task_list: list[TaskCreateV2]) -> None:
+    """
+    Check that the modules listed in task commands point to existing files.
+
+    Args:
+        task_list:
+    """
+    for _task in task_list:
+        if _task.command_non_parallel is not None:
+            _task_path = _task.command_non_parallel.split()[1]
+            if not Path(_task_path).exists():
+                raise FileNotFoundError(
+                    f"Task `{_task.name}` has `command_non_parallel` "
+                    f"pointing to missing file `{_task_path}`."
+                )
+        if _task.command_parallel is not None:
+            _task_path = _task.command_parallel.split()[1]
+            if not Path(_task_path).exists():
+                raise FileNotFoundError(
+                    f"Task `{_task.name}` has `command_parallel` "
+                    f"pointing to missing file `{_task_path}`."
+                )