fractal-server 2.3.0a2__py3-none-any.whl → 2.3.1__py3-none-any.whl

fractal_server/ssh/_fabric.py

@@ -1,7 +1,16 @@
+import logging
 import time
+from contextlib import contextmanager
+from pathlib import Path
+from threading import Lock
+from typing import Any
+from typing import Generator
+from typing import Literal
 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
 
@@ -10,9 +19,293 @@ from ..logger import set_logger
 from fractal_server.config import get_settings
 from fractal_server.syringe import Inject
 
+
+class FractalSSHTimeoutError(RuntimeError):
+    pass
+
+
 logger = set_logger(__name__)
 
-MAX_ATTEMPTS = 5
+
+class FractalSSH(object):
+
+    """
+    FIXME SSH: Fix docstring
+
+    Attributes:
+        _lock:
+        connection:
+        default_lock_timeout:
+        logger_name:
+    """
+
+    _lock: Lock
+    _connection: Connection
+    default_lock_timeout: float
+    default_max_attempts: int
+    default_base_interval: float
+    logger_name: str
+
+    def __init__(
+        self,
+        connection: Connection,
+        default_timeout: float = 250,
+        default_max_attempts: int = 5,
+        default_base_interval: float = 3.0,
+        logger_name: str = __name__,
+    ):
+        self._lock = Lock()
+        self._connection = connection
+        self.default_lock_timeout = default_timeout
+        self.default_base_interval = default_base_interval
+        self.default_max_attempts = default_max_attempts
+        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
+
+    @property
+    def logger(self) -> logging.Logger:
+        return get_logger(self.logger_name)
+
+    def put(
+        self, *args, lock_timeout: Optional[float] = None, **kwargs
+    ) -> 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.put(*args, **kwargs)
+
+    def get(
+        self, *args, lock_timeout: Optional[float] = None, **kwargs
+    ) -> 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)
+
+    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
+        with self.acquire_timeout(timeout=actual_lock_timeout):
+            return self._connection.run(*args, **kwargs)
+
+    def sftp(self) -> paramiko.sftp_client.SFTPClient:
+        return self._connection.sftp()
+
+    def check_connection(self) -> None:
+        """
+        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.
+        """
+        if not self._connection.is_connected:
+            try:
+                self._connection.open()
+            except Exception as e:
+                raise RuntimeError(
+                    f"Cannot open SSH connection. Original error:\n{str(e)}"
+                )
+
+    def close(self) -> None:
+        return self._connection.close()
+
+    def run_command(
+        self,
+        *,
+        cmd: str,
+        max_attempts: Optional[int] = None,
+        base_interval: Optional[int] = None,
+        lock_timeout: Optional[int] = None,
+    ) -> str:
+        """
+        Run a command within an open SSH connection.
+
+        Args:
+            cmd: Command to be run
+            max_attempts:
+            base_interval:
+            lock_timeout:
+
+        Returns:
+            Standard output of the command, if successful.
+        """
+        actual_max_attempts = self.default_max_attempts
+        if max_attempts is not None:
+            actual_max_attempts = max_attempts
+
+        actual_base_interval = self.default_base_interval
+        if base_interval is not None:
+            actual_base_interval = base_interval
+
+        actual_lock_timeout = self.default_lock_timeout
+        if lock_timeout is not None:
+            actual_lock_timeout = lock_timeout
+
+        t_0 = time.perf_counter()
+        ind_attempt = 0
+        while ind_attempt <= actual_max_attempts:
+            ind_attempt += 1
+            prefix = f"[attempt {ind_attempt}/{actual_max_attempts}]"
+            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
+                )
+                t_1 = time.perf_counter()
+                self.logger.info(
+                    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}")
+                return res.stdout
+            except NoValidConnectionsError as e:
+                # Case 2: Command fails with a connection error
+                self.logger.warning(
+                    f"{prefix} Running command `{cmd}` over SSH failed.\n"
+                    f"Original NoValidConnectionError:\n{str(e)}.\n"
+                    f"{e.errors=}\n"
+                )
+                if ind_attempt < actual_max_attempts:
+                    sleeptime = actual_base_interval**ind_attempt
+                    self.logger.warning(
+                        f"{prefix} Now sleep {sleeptime:.3f} "
+                        "seconds and continue."
+                    )
+                    time.sleep(sleeptime)
+                else:
+                    self.logger.error(f"{prefix} Reached last attempt")
+                    break
+            except UnexpectedExit as e:
+                # Case 3: Command fails with an actual error
+                error_msg = (
+                    f"{prefix} Running command `{cmd}` over SSH failed.\n"
+                    f"Original error:\n{str(e)}."
+                )
+                self.logger.error(error_msg)
+                raise RuntimeError(error_msg)
+            except Exception as e:
+                self.logger.error(
+                    f"Running command `{cmd}` over SSH failed.\n"
+                    f"Original Error:\n{str(e)}."
+                )
+                raise e
+
+        raise RuntimeError(
+            f"Reached last attempt ({max_attempts=}) for running "
+            f"'{cmd}' over SSH"
+        )
+
+    def send_file(
+        self,
+        *,
+        local: str,
+        remote: str,
+        logger_name: Optional[str] = None,
+        lock_timeout: Optional[float] = None,
+    ) -> None:
+        """
+        Transfer a file via SSH
+
+        Args:
+            local: Local path to file
+            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)
+        except Exception as e:
+            logger = get_logger(logger_name=logger_name)
+            logger.error(
+                f"Transferring {local=} to {remote=} 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:
+            cmd = f"mkdir {folder}"
+        self.run_command(cmd=cmd)
+
+    def remove_folder(
+        self,
+        *,
+        folder: str,
+        safe_root: str,
+    ) -> None:
+        """
+        Removes a folder remotely via SSH.
+
+        This functions calls `rm -r`, after a few checks on `folder`.
+
+        Args:
+            folder: Absolute path to a folder that should be removed.
+            safe_root: If `folder` is not a subfolder of the absolute
+                `safe_root` path, raise an error.
+            fractal_ssh:
+        """
+        invalid_characters = {" ", "\n", ";", "$", "`"}
+
+        if (
+            not isinstance(folder, str)
+            or not isinstance(safe_root, str)
+            or len(invalid_characters.intersection(folder)) > 0
+            or len(invalid_characters.intersection(safe_root)) > 0
+            or not Path(folder).is_absolute()
+            or not Path(safe_root).is_absolute()
+            or not Path(folder).resolve().is_relative_to(safe_root)
+        ):
+            raise ValueError(
+                f"{folder=} argument is invalid or it is not "
+                f"relative to {safe_root=}."
+            )
+        else:
+            cmd = f"rm -r {folder}"
+            self.run_command(cmd=cmd)
 
 
 def get_ssh_connection(
@@ -44,147 +337,7 @@ def get_ssh_connection(
     connection = Connection(
         host=host,
         user=user,
+        forward_agent=False,
         connect_kwargs={"key_filename": key_filename},
     )
-    logger.debug(f"Now created {connection=}.")
     return connection
-
-
-def check_connection(connection: Connection) -> None:
-    """
-    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.
-
-    Args:
-        connection: Fabric connection object
-    """
-    if not connection.is_connected:
-        try:
-            connection.open()
-        except Exception as e:
-            raise RuntimeError(
-                f"Cannot open SSH connection (original error: '{str(e)}')."
-            )
-
-
-def run_command_over_ssh(
-    *,
-    cmd: str,
-    connection: Connection,
-    max_attempts: int = MAX_ATTEMPTS,
-    base_interval: float = 3.0,
-) -> str:
-    """
-    Run a command within an open SSH connection.
-
-    Args:
-        cmd: Command to be run
-        connection: Fabric connection object
-
-    Returns:
-        Standard output of the command, if successful.
-    """
-    t_0 = time.perf_counter()
-    ind_attempt = 0
-    while ind_attempt <= max_attempts:
-        ind_attempt += 1
-        prefix = f"[attempt {ind_attempt}/{max_attempts}]"
-        logger.info(f"{prefix} START running '{cmd}' over SSH.")
-        try:
-            # Case 1: Command runs successfully
-            res = connection.run(cmd, hide=True)
-            t_1 = time.perf_counter()
-            logger.info(
-                f"{prefix} END   running '{cmd}' over SSH, "
-                f"elapsed {t_1-t_0:.3f}"
-            )
-            logger.debug(f"STDOUT: {res.stdout}")
-            logger.debug(f"STDERR: {res.stderr}")
-            return res.stdout
-        except NoValidConnectionsError as e:
-            # Case 2: Command fails with a connection error
-            logger.warning(
-                f"{prefix} Running command `{cmd}` over SSH failed.\n"
-                f"Original NoValidConnectionError:\n{str(e)}.\n"
-                f"{e.errors=}\n"
-            )
-            if ind_attempt < max_attempts:
-                sleeptime = (
-                    base_interval**ind_attempt
-                )  # FIXME SSH: add jitter?
-                logger.warning(
-                    f"{prefix} Now sleep {sleeptime:.3f} seconds and continue."
-                )
-                time.sleep(sleeptime)
-                continue
-            else:
-                logger.error(f"{prefix} Reached last attempt")
-                break
-        except UnexpectedExit as e:
-            # Case 3: Command fails with an actual error
-            error_msg = (
-                f"{prefix} Running command `{cmd}` over SSH failed.\n"
-                f"Original error:\n{str(e)}."
-            )
-            logger.error(error_msg)
-            raise ValueError(error_msg)
-        except Exception as e:
-            logger.error(
-                f"Running command `{cmd}` over SSH failed.\n"
-                f"Original Error:\n{str(e)}."
-            )
-            raise e
-
-    raise ValueError(
-        f"Reached last attempt ({max_attempts=}) for running '{cmd}' over SSH"
-    )
-
-
-def put_over_ssh(
-    *,
-    local: str,
-    remote: str,
-    connection: Connection,
-    logger_name: Optional[str] = None,
-) -> None:
-    """
-    Transfer a file via SSH
-
-    Args:
-        local: Local path to file
-        remote: Target path on remote host
-        connection: Fabric connection object
-        logger_name: Name of the logger
-
-    """
-    try:
-        connection.put(local=local, remote=remote)
-    except Exception as e:
-        logger = get_logger(logger_name=logger_name)
-        logger.error(
-            f"Transferring {local=} to {remote=} over SSH failed.\n"
-            f"Original Error:\n{str(e)}."
-        )
-        raise e
-
-
-def _mkdir_over_ssh(
-    *, folder: str, connection: Connection, parents: bool = True
-) -> None:
-    """
-    Create a folder remotely via SSH.
-
-    Args:
-        folder:
-        connection:
-        parents:
-    """
-    # FIXME SSH: try using `mkdir` method of `paramiko.SFTPClient`
-    if parents:
-        cmd = f"mkdir -p {folder}"
-    else:
-        cmd = f"mkdir {folder}"
-    run_command_over_ssh(cmd=cmd, connection=connection)

fractal_server/string_tools.py

@@ -0,0 +1,39 @@
+import string
+
+__SPECIAL_CHARACTERS__ = f"{string.punctuation}{string.whitespace}"
+
+
+def sanitize_string(value: str) -> str:
+    """
+    Make string safe to be used in file/folder names and subprocess commands.
+
+    Make the string lower-case, and replace any special character with an
+    underscore, where special characters are:
+    ```python repl
+    >>> string.punctuation
+    '!"#$%&\'()*+,-./:;<=>?@[\\]^_`{|}~'
+    >>> string.whitespace
+    ' \t\n\r\x0b\x0c'
+    ```
+
+    Args:
+        value: Input string
+
+    Returns:
+        Sanitized value
+    """
+    new_value = value.lower()
+    for character in __SPECIAL_CHARACTERS__:
+        new_value = new_value.replace(character, "_")
+    return new_value
+
+
+def slugify_task_name_for_source(task_name: str) -> str:
+    """
+    NOTE: this function is used upon creation of tasks' sources, therefore
+    for the moment we cannot replace it with its more comprehensive version
+    from `fractal_server.string_tools.sanitize_string`, nor we can remove it.
+    As 2.3.1, we are renaming it to `slugify_task_name_for_source`, to make
+    it clear that it should not be used for other purposes.
+    """
+    return task_name.replace(" ", "_").lower()

fractal_server/tasks/utils.py

@@ -9,10 +9,6 @@ COLLECTION_LOG_FILENAME = "collection.log"
 COLLECTION_FREEZE_FILENAME = "collection_freeze.txt"
 
 
-def slugify_task_name(task_name: str) -> str:
-    return task_name.replace(" ", "_").lower()
-
-
 def get_absolute_venv_path(venv_path: Path) -> Path:
     """
     If a path is not absolute, make it a relative path of FRACTAL_TASKS_DIR.

fractal_server/tasks/v1/background_operations.py

@@ -6,11 +6,11 @@ import json
 from pathlib import Path
 from shutil import rmtree as shell_rmtree
 
+from ...string_tools import slugify_task_name_for_source
 from ..utils import _normalize_package_name
 from ..utils import get_collection_log
 from ..utils import get_collection_path
 from ..utils import get_log_path
-from ..utils import slugify_task_name
 from ._TaskCollectPip import _TaskCollectPip
 from .utils import _init_venv_v1
 from fractal_server.app.db import DBSyncSession
@@ -215,7 +215,7 @@ async def create_package_environment_pip(
             # Fill in attributes for TaskCreate
             task_executable = package_root / t.executable
             cmd = f"{python_bin.as_posix()} {task_executable.as_posix()}"
-            task_name_slug = slugify_task_name(t.name)
+            task_name_slug = slugify_task_name_for_source(t.name)
             task_source = f"{task_pkg.package_source}:{task_name_slug}"
             if not task_executable.exists():
                 raise FileNotFoundError(

fractal_server/tasks/v2/background_operations.py

@@ -10,12 +10,12 @@ from typing import Optional
 from sqlalchemy.orm import Session as DBSyncSession
 from sqlalchemy.orm.attributes import flag_modified
 
+from ...string_tools import slugify_task_name_for_source
 from ..utils import get_absolute_venv_path
 from ..utils import get_collection_freeze
 from ..utils import get_collection_log
 from ..utils import get_collection_path
 from ..utils import get_log_path
-from ..utils import slugify_task_name
 from ._TaskCollectPip import _TaskCollectPip
 from fractal_server.app.db import get_sync_db
 from fractal_server.app.models.v2 import CollectionStateV2
@@ -177,7 +177,7 @@ def _prepare_tasks_metadata(
         task_attributes = {}
         if package_version is not None:
             task_attributes["version"] = package_version
-        task_name_slug = slugify_task_name(_task.name)
+        task_name_slug = slugify_task_name_for_source(_task.name)
         task_attributes["source"] = f"{package_source}:{task_name_slug}"
         if package_manifest.has_args_schemas:
             task_attributes[

fractal_server/tasks/v2/background_operations_ssh.py

@@ -3,7 +3,6 @@ import os
 from pathlib import Path
 from tempfile import TemporaryDirectory
 
-from fabric import Connection
 from sqlalchemy.orm.attributes import flag_modified
 
 from ...app.models.v2 import CollectionStateV2
@@ -18,9 +17,7 @@ from fractal_server.app.schemas.v2.manifest import ManifestV2
 from fractal_server.config import get_settings
 from fractal_server.logger import get_logger
 from fractal_server.logger import set_logger
-from fractal_server.ssh._fabric import check_connection
-from fractal_server.ssh._fabric import put_over_ssh
-from fractal_server.ssh._fabric import run_command_over_ssh
+from fractal_server.ssh._fabric import FractalSSH
 from fractal_server.syringe import Inject
 from fractal_server.tasks.v2.utils import get_python_interpreter_v2
 
@@ -32,7 +29,7 @@ def _parse_script_5_stdout(stdout: str) -> dict[str, str]:
         ("Python interpreter:", "python_bin"),
         ("Package name:", "package_name"),
         ("Package version:", "package_version"),
-        ("Package parent folder:", "package_root_remote"),
+        ("Package parent folder:", "package_root_parent_remote"),
         ("Manifest absolute path:", "manifest_path_remote"),
     ]
     stdout_lines = stdout.splitlines()
@@ -59,7 +56,7 @@ def _customize_and_run_template(
     replacements: list[tuple[str, str]],
     tmpdir: str,
     logger_name: str,
-    connection: Connection,
+    fractal_ssh: FractalSSH,
 ) -> str:
     """
     Customize one of the template bash scripts, transfer it to the remote host
@@ -71,7 +68,7 @@ def _customize_and_run_template(
         replacements:
         tmpdir:
         logger_name:
-        connection:
+        fractal_ssh:
     """
     logger = get_logger(logger_name)
     logger.debug(f"_customize_and_run_template {script_filename} - START")
@@ -95,17 +92,15 @@ def _customize_and_run_template(
         f"script_{abs(hash(tmpdir))}{script_filename}",
     )
     logger.debug(f"Now transfer {script_path_local=} over SSH.")
-    put_over_ssh(
+    fractal_ssh.send_file(
         local=script_path_local,
         remote=script_path_remote,
-        connection=connection,
-        logger_name=logger_name,
     )
 
     # Execute script remotely
     cmd = f"bash {script_path_remote}"
     logger.debug(f"Now run '{cmd}' over SSH.")
-    stdout = run_command_over_ssh(cmd=cmd, connection=connection)
+    stdout = fractal_ssh.run_command(cmd=cmd)
     logger.debug(f"Standard output of '{cmd}':\n{stdout}")
 
     logger.debug(f"_customize_and_run_template {script_filename} - END")
@@ -115,7 +110,7 @@ def _customize_and_run_template(
 def background_collect_pip_ssh(
     state_id: int,
     task_pkg: _TaskCollectPip,
-    connection: Connection,
+    fractal_ssh: FractalSSH,
 ) -> None:
     """
     Collect a task package over SSH
@@ -127,6 +122,7 @@ def background_collect_pip_ssh(
     starlette/fastapi handling of background tasks (see
     https://github.com/encode/starlette/blob/master/starlette/background.py).
     """
+
     # Work within a temporary folder, where also logs will be placed
     with TemporaryDirectory() as tmpdir:
         LOGGER_NAME = "task_collection_ssh"
@@ -185,10 +181,10 @@ def background_collect_pip_ssh(
                     replacements=replacements,
                     tmpdir=tmpdir,
                     logger_name=LOGGER_NAME,
-                    connection=connection,
+                    fractal_ssh=fractal_ssh,
                 )
 
-                check_connection(connection)
+                fractal_ssh.check_connection()
 
                 logger.debug("installing - START")
                 _set_collection_state_data_status(
@@ -253,13 +249,23 @@ def background_collect_pip_ssh(
                     logger.error(error_msg)
                     raise ValueError(error_msg)
                 # Extract/drop parsed attributes
-                pkg_attrs.pop("package_name")
+                package_name = pkg_attrs.pop("package_name")
                 python_bin = pkg_attrs.pop("python_bin")
-                package_root_remote = pkg_attrs.pop("package_root_remote")
+                package_root_parent_remote = pkg_attrs.pop(
+                    "package_root_parent_remote"
+                )
                 manifest_path_remote = pkg_attrs.pop("manifest_path_remote")
 
+                # FIXME SSH: Use more robust logic to determine `package_root`,
+                # e.g. as in the custom task-collection endpoint (where we use
+                # `importlib.util.find_spec`)
+                package_name_underscore = package_name.replace("-", "_")
+                package_root_remote = (
+                    Path(package_root_parent_remote) / package_name_underscore
+                ).as_posix()
+
                 # Read and validate remote manifest file
-                with connection.sftp().open(manifest_path_remote, "r") as f:
+                with fractal_ssh.sftp().open(manifest_path_remote, "r") as f:
                     manifest = json.load(f)
                 logger.info(f"collecting - loaded {manifest_path_remote=}")
                 ManifestV2(**manifest)
@@ -298,6 +304,7 @@ def background_collect_pip_ssh(
                 logger.debug("END")
 
             except Exception as e:
+                # Delete corrupted package dir
                 _handle_failure(
                     state_id=state_id,
                     log_file_path=log_file_path,
@@ -305,4 +312,16 @@ def background_collect_pip_ssh(
                     exception=e,
                     db=db,
                 )
+                try:
+                    logger.info(f"Now delete remote folder {package_env_dir}")
+                    fractal_ssh.remove_folder(
+                        folder=package_env_dir,
+                        safe_root=settings.FRACTAL_SLURM_SSH_WORKING_BASE_DIR,
+                    )
+                    logger.info(f"Deleted remoted folder {package_env_dir}")
+                except Exception as e:
+                    logger.error(
+                        f"Deleting remote folder failed.\n"
+                        f"Original error:\n{str(e)}"
+                    )
                 return