fractal-server 2.2.0a1__py3-none-any.whl → 2.3.0__py3-none-any.whl

fractal_server/ssh/_fabric.py

@@ -0,0 +1,245 @@
+import time
+from contextlib import contextmanager
+from threading import Lock
+from typing import Any
+from typing import Optional
+
+from fabric import Connection
+from fabric import Result
+from invoke import UnexpectedExit
+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.syringe import Inject
+
+logger = set_logger(__name__)
+
+MAX_ATTEMPTS = 5
+
+
+class TimeoutException(Exception):
+    pass
+
+
+@contextmanager
+def acquire_timeout(lock: Lock, timeout: int) -> Any:
+    logger.debug(f"Trying to acquire lock, with {timeout=}")
+    result = lock.acquire(timeout=timeout)
+    try:
+        if not result:
+            raise TimeoutException(
+                f"Failed to acquire lock within {timeout} seconds"
+            )
+        logger.debug("Lock was acquired.")
+        yield result
+    finally:
+        if result:
+            lock.release()
+            logger.debug("Lock was released")
+
+
+class FractalSSH(object):
+    lock: Lock
+    connection: Connection
+    default_timeout: int
+
+    # FIXME SSH: maybe extend the actual_timeout logic to other methods
+
+    def __init__(self, connection: Connection, default_timeout: int = 250):
+        self.lock = Lock()
+        self.conn = connection
+        self.default_timeout = default_timeout
+
+    @property
+    def is_connected(self) -> bool:
+        return self.conn.is_connected
+
+    def put(self, *args, timeout: Optional[int] = None, **kwargs) -> Result:
+        actual_timeout = timeout or self.default_timeout
+        with acquire_timeout(self.lock, timeout=actual_timeout):
+            return self.conn.put(*args, **kwargs)
+
+    def get(self, *args, **kwargs) -> Result:
+        with acquire_timeout(self.lock, timeout=self.default_timeout):
+            return self.conn.get(*args, **kwargs)
+
+    def run(self, *args, **kwargs) -> Any:
+        with acquire_timeout(self.lock, timeout=self.default_timeout):
+            return self.conn.run(*args, **kwargs)
+
+    def close(self):
+        return self.conn.close()
+
+    def sftp(self):
+        return self.conn.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.conn.is_connected:
+            try:
+                self.conn.open()
+            except Exception as e:
+                raise RuntimeError(
+                    f"Cannot open SSH connection (original error: '{str(e)}')."
+                )
+
+
+def get_ssh_connection(
+    *,
+    host: Optional[str] = None,
+    user: Optional[str] = None,
+    key_filename: Optional[str] = None,
+) -> Connection:
+    """
+    Create a `fabric.Connection` object based on fractal-server settings
+    or explicit arguments.
+
+    Args:
+        host:
+        user:
+        key_filename:
+
+    Returns:
+        Fabric connection object
+    """
+    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,
+        connect_kwargs={"key_filename": key_filename},
+    )
+    logger.debug(f"Now created {connection=}.")
+    return connection
+
+
+def run_command_over_ssh(
+    *,
+    cmd: str,
+    fractal_ssh: FractalSSH,
+    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
+        fractal_ssh: FractalSSH connection object with custom lock
+
+    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 = fractal_ssh.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
+                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,
+    fractal_ssh: FractalSSH,
+    logger_name: Optional[str] = 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:
+        fractal_ssh.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, fractal_ssh: FractalSSH, 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}"
+    run_command_over_ssh(cmd=cmd, fractal_ssh=fractal_ssh)

fractal_server/tasks/utils.py

@@ -1,42 +1,12 @@
 import re
-import shutil
-import sys
 from pathlib import Path
-from typing import Optional
 
 from fractal_server.config import get_settings
-from fractal_server.logger import get_logger
 from fractal_server.syringe import Inject
-from fractal_server.utils import execute_command
 
 COLLECTION_FILENAME = "collection.json"
 COLLECTION_LOG_FILENAME = "collection.log"
-
-
-def get_python_interpreter(version: Optional[str] = None) -> str:
-    """
-    Return the path to the python interpreter
-
-    Args:
-        version: Python version
-
-    Raises:
-        ValueError: If the python version requested is not available on the
-                    host.
-
-    Returns:
-        interpreter: string representing the python executable or its path
-    """
-    if version:
-        interpreter = shutil.which(f"python{version}")
-        if not interpreter:
-            raise ValueError(
-                f"Python version {version} not available on host."
-            )
-    else:
-        interpreter = sys.executable
-
-    return interpreter
+COLLECTION_FREEZE_FILENAME = "collection_freeze.txt"
 
 
 def slugify_task_name(task_name: str) -> str:
@@ -63,6 +33,10 @@ 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(venv_path: Path) -> str:
     package_path = get_absolute_venv_path(venv_path)
     log_path = get_log_path(package_path)
@@ -70,6 +44,13 @@ def get_collection_log(venv_path: Path) -> str:
     return log
 
 
+def get_collection_freeze(venv_path: Path) -> str:
+    package_path = get_absolute_venv_path(venv_path)
+    freeze_path = get_freeze_path(package_path)
+    freeze = freeze_path.open().read()
+    return freeze
+
+
 def _normalize_package_name(name: str) -> str:
     """
     Implement PyPa specifications for package-name normalization
@@ -86,36 +67,3 @@ def _normalize_package_name(name: str) -> str:
         The normalized package name.
     """
     return re.sub(r"[-_.]+", "-", name).lower()
-
-
-async def _init_venv(
-    *,
-    path: Path,
-    python_version: Optional[str] = None,
-    logger_name: str,
-) -> Path:
-    """
-    Set a virtual environment at `path/venv`
-
-    Args:
-        path : Path
-            path to directory in which to set up the virtual environment
-        python_version : default=None
-            Python version the virtual environment will be based upon
-
-    Returns:
-        python_bin : Path
-            path to python interpreter
-    """
-    logger = get_logger(logger_name)
-    logger.debug(f"[_init_venv] {path=}")
-    interpreter = get_python_interpreter(version=python_version)
-    logger.debug(f"[_init_venv] {interpreter=}")
-    await execute_command(
-        cwd=path,
-        command=f"{interpreter} -m venv venv",
-        logger_name=logger_name,
-    )
-    python_bin = path / "venv/bin/python"
-    logger.debug(f"[_init_venv] {python_bin=}")
-    return python_bin

fractal_server/tasks/v1/background_operations.py

@@ -6,13 +6,13 @@ import json
 from pathlib import Path
 from shutil import rmtree as shell_rmtree
 
-from ..utils import _init_venv
 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
 from fractal_server.app.db import get_sync_db
 from fractal_server.app.models.v1 import State
@@ -168,7 +168,7 @@ async def _create_venv_install_package(
     task_pkg.package_name = _normalize_package_name(task_pkg.package_name)
     task_pkg.package = _normalize_package_name(task_pkg.package)
 
-    python_bin = await _init_venv(
+    python_bin = await _init_venv_v1(
         path=path,
         python_version=task_pkg.python_version,
         logger_name=logger_name,

fractal_server/tasks/{endpoint_operations.py → v1/endpoint_operations.py}

@@ -4,12 +4,10 @@ from typing import Optional
 from typing import Union
 from zipfile import ZipFile
 
-from .utils import _normalize_package_name
-from .utils import get_python_interpreter
-from .v1._TaskCollectPip import _TaskCollectPip as _TaskCollectPipV1
-from .v2._TaskCollectPip import _TaskCollectPip as _TaskCollectPipV2
+from ..utils import _normalize_package_name
+from ._TaskCollectPip import _TaskCollectPip as _TaskCollectPipV1
+from .utils import get_python_interpreter_v1
 from fractal_server.app.schemas.v1 import ManifestV1
-from fractal_server.app.schemas.v2 import ManifestV2
 from fractal_server.config import get_settings
 from fractal_server.logger import get_logger
 from fractal_server.syringe import Inject
@@ -21,13 +19,13 @@ FRACTAL_PUBLIC_TASK_SUBDIR = ".fractal"
 
 async def download_package(
     *,
-    task_pkg: Union[_TaskCollectPipV1, _TaskCollectPipV2],
+    task_pkg: _TaskCollectPipV1,
     dest: Union[str, Path],
 ) -> Path:
     """
     Download package to destination
     """
-    interpreter = get_python_interpreter(version=task_pkg.python_version)
+    interpreter = get_python_interpreter_v1(version=task_pkg.python_version)
     pip = f"{interpreter} -m pip"
     version = (
         f"=={task_pkg.package_version}" if task_pkg.package_version else ""
@@ -43,7 +41,7 @@ async def download_package(
 
 def _load_manifest_from_wheel(
     path: Path, wheel: ZipFile, logger_name: Optional[str] = None
-) -> Union[ManifestV1, ManifestV2]:
+) -> ManifestV1:
     logger = get_logger(logger_name)
     namelist = wheel.namelist()
     try:
@@ -60,9 +58,6 @@ def _load_manifest_from_wheel(
     if manifest_version == "1":
         pkg_manifest = ManifestV1(**manifest_dict)
         return pkg_manifest
-    elif manifest_version == "2":
-        pkg_manifest = ManifestV2(**manifest_dict)
-        return pkg_manifest
     else:
         msg = f"Manifest version {manifest_version=} not supported"
         logger.error(msg)
@@ -140,7 +135,7 @@ def inspect_package(path: Path, logger_name: Optional[str] = None) -> dict:
 
 def create_package_dir_pip(
     *,
-    task_pkg: Union[_TaskCollectPipV1, _TaskCollectPipV2],
+    task_pkg: _TaskCollectPipV1,
     create: bool = True,
 ) -> Path:
     """

fractal_server/tasks/v1/utils.py

@@ -0,0 +1,67 @@
+from pathlib import Path
+from typing import Optional
+
+from fractal_server.logger import get_logger
+from fractal_server.utils import execute_command
+
+
+def get_python_interpreter_v1(version: Optional[str] = None) -> str:
+    """
+    Return the path to the python interpreter
+
+    Args:
+        version: Python version
+
+    Raises:
+        ValueError: If the python version requested is not available on the
+                    host.
+
+    Returns:
+        interpreter: string representing the python executable or its path
+    """
+    import shutil
+    import sys
+
+    if version:
+        interpreter = shutil.which(f"python{version}")
+        if not interpreter:
+            raise ValueError(
+                f"Python version {version} not available on host."
+            )
+    else:
+        interpreter = sys.executable
+
+    return interpreter
+
+
+async def _init_venv_v1(
+    *,
+    path: Path,
+    python_version: Optional[str] = None,
+    logger_name: str,
+) -> Path:
+    """
+    Set a virtual environment at `path/venv`
+
+    Args:
+        path : Path
+            path to directory in which to set up the virtual environment
+        python_version : default=None
+            Python version the virtual environment will be based upon
+
+    Returns:
+        python_bin : Path
+            path to python interpreter
+    """
+    logger = get_logger(logger_name)
+    logger.debug(f"[_init_venv] {path=}")
+    interpreter = get_python_interpreter_v1(version=python_version)
+    logger.debug(f"[_init_venv] {interpreter=}")
+    await execute_command(
+        cwd=path,
+        command=f"{interpreter} -m venv venv",
+        logger_name=logger_name,
+    )
+    python_bin = path / "venv/bin/python"
+    logger.debug(f"[_init_venv] {python_bin=}")
+    return python_bin

fractal_server/tasks/v2/_TaskCollectPip.py

@@ -1,56 +1,90 @@
+import logging
 from pathlib import Path
 from typing import Optional
 
+from pydantic import BaseModel
+from pydantic import Extra
+from pydantic import Field
 from pydantic import root_validator
+from pydantic import validator
 
+from fractal_server.app.schemas._validators import valdictkeys
+from fractal_server.app.schemas._validators import valstr
 from fractal_server.app.schemas.v2 import ManifestV2
-from fractal_server.app.schemas.v2 import TaskCollectPipV2
+from fractal_server.tasks.utils import _normalize_package_name
+from fractal_server.tasks.v2.utils import _parse_wheel_filename
 
 
-class _TaskCollectPip(TaskCollectPipV2):
+class _TaskCollectPip(BaseModel, extra=Extra.forbid):
     """
-    Internal TaskCollectPip schema
+    Internal task-collection model.
 
-    Differences with its parent class (`TaskCollectPip`):
+    This model is similar to the API request-body model (`TaskCollectPip`), but
+    with enough differences that we keep them separated (and they do not have a
+    common base).
 
-        1. We check if the package corresponds to a path in the filesystem, and
-           whether it exists (via new validator `check_local_package`, new
-           method `is_local_package` and new attribute `package_path`).
-        2. We include an additional `package_manifest` attribute.
-        3. We expose an additional attribute `package_name`, which is filled
-           during task collection.
+    Attributes:
+        package: Either a PyPI package name or the absolute path to a wheel
+            file.
+        package_name: The actual normalized name of the package, which is set
+            internally through a validator.
+        package_version: Package version. For local packages, it is set
+            internally through a validator.
     """
 
-    package_name: Optional[str] = None
+    package: str
+    package_name: str
+    python_version: str
+    package_extras: Optional[str] = None
+    pinned_package_versions: dict[str, str] = Field(default_factory=dict)
+    package_version: Optional[str] = None
     package_path: Optional[Path] = None
     package_manifest: Optional[ManifestV2] = None
 
+    _pinned_package_versions = validator(
+        "pinned_package_versions", allow_reuse=True
+    )(valdictkeys("pinned_package_versions"))
+    _package_extras = validator("package_extras", allow_reuse=True)(
+        valstr("package_extras")
+    )
+    _python_version = validator("python_version", allow_reuse=True)(
+        valstr("python_version")
+    )
+
     @property
     def is_local_package(self) -> bool:
         return bool(self.package_path)
 
     @root_validator(pre=True)
-    def check_local_package(cls, values):
+    def set_package_info(cls, values):
         """
-        Checks if package corresponds to an existing path on the filesystem
-
-        In this case, the user is providing directly a package file, rather
-        than a remote one from PyPI. We set the `package_path` attribute and
-        get the actual package name and version from the package file name.
+        Depending on whether the package is a local wheel file or a PyPI
+        package, set some of its metadata.
         """
         if "/" in values["package"]:
+            # Local package: parse wheel filename
             package_path = Path(values["package"])
             if not package_path.is_absolute():
                 raise ValueError("Package path must be absolute")
-            if package_path.exists():
-                values["package_path"] = package_path
-                (
-                    values["package"],
-                    values["version"],
-                    *_,
-                ) = package_path.name.split("-")
-            else:
-                raise ValueError(f"Package {package_path} does not exist.")
+            if not package_path.exists():
+                logging.warning(
+                    f"Package {package_path} does not exist locally."
+                )
+            values["package_path"] = package_path
+            wheel_metadata = _parse_wheel_filename(package_path.name)
+            values["package_name"] = _normalize_package_name(
+                wheel_metadata["distribution"]
+            )
+            values["package_version"] = wheel_metadata["version"]
+        else:
+            # Remote package: use `package` as `package_name`
+            _package = values["package"]
+            if _package.endswith(".whl"):
+                raise ValueError(
+                    f"ERROR: package={_package} ends with '.whl' "
+                    "but it is not the absolute path to a wheel file."
+                )
+            values["package_name"] = _normalize_package_name(values["package"])
         return values
 
     @property
@@ -74,10 +108,7 @@ class _TaskCollectPip(TaskCollectPipV2):
             collection_type = "pip_remote"
 
         package_extras = self.package_extras or ""
-        if self.python_version:
-            python_version = f"py{self.python_version}"
-        else:
-            python_version = ""  # FIXME: can we allow this?
+        python_version = f"py{self.python_version}"
 
         source = ":".join(
             (
@@ -95,8 +126,6 @@ class _TaskCollectPip(TaskCollectPipV2):
         Verify that the package has all attributes that are needed to continue
         with task collection
         """
-        if not self.package_name:
-            raise ValueError("`package_name` attribute is not set")
         if not self.package_version:
             raise ValueError("`package_version` attribute is not set")
         if not self.package_manifest: