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

fractal_server/app/schemas/v2/task_collection.py

@@ -1,14 +1,26 @@
+from datetime import datetime
+from enum import Enum
 from pathlib import Path
+from typing import Any
 from typing import Literal
 from typing import Optional
 
 from pydantic import BaseModel
-from pydantic import Field
+from pydantic import root_validator
 from pydantic import validator
 
 from .._validators import valdictkeys
 from .._validators import valstr
-from .task import TaskReadV2
+from fractal_server.app.schemas._validators import valutc
+from fractal_server.app.schemas.v2 import ManifestV2
+
+
+class CollectionStatusV2(str, Enum):
+    PENDING = "pending"
+    INSTALLING = "installing"
+    COLLECTING = "collecting"
+    FAIL = "fail"
+    OK = "OK"
 
 
 class TaskCollectPipV2(BaseModel):
@@ -41,7 +53,7 @@ class TaskCollectPipV2(BaseModel):
     package: str
     package_version: Optional[str] = None
     package_extras: Optional[str] = None
-    python_version: Optional[str] = None
+    python_version: Optional[Literal["3.9", "3.10", "3.11", "3.12"]] = None
     pinned_package_versions: Optional[dict[str, str]] = None
 
     _pinned_package_versions = validator(
@@ -70,40 +82,98 @@ class TaskCollectPipV2(BaseModel):
 
     @validator("package_version")
     def package_version_validator(cls, v, values):
-
-        valstr("package_version")(v)
-
+        v = valstr("package_version")(v)
         if values["package"].endswith(".whl"):
             raise ValueError(
-                "Cannot provide version when package is a Wheel file."
+                "Cannot provide package version when package is a wheel file."
             )
         return v
 
 
-class TaskCollectStatusV2(BaseModel):
+class TaskCollectCustomV2(BaseModel):
     """
-    TaskCollectStatus class
-
     Attributes:
-        status:
-        package:
-        venv_path:
-        task_list:
-        log:
-        info:
+        manifest: Manifest of a Fractal task package (this is typically the
+            content of `__FRACTAL_MANIFEST__.json`).
+        python_interpreter: Absolute path to the Python interpreter to be used
+            for running tasks.
+        source: A common label identifying this package.
+        package_root: The folder where the package is installed.
+            If not provided, it will be extracted via `pip show`
+            (requires `package_name` to be set).
+        package_name: Name of the package, as used for `import <package_name>`;
+            this is then used to extract the package directory (`package_root`)
+            via `pip show <package_name>`.
+        version: Optional version of tasks to be collected.
     """
 
-    status: Literal["pending", "installing", "collecting", "fail", "OK"]
-    package: str
-    venv_path: Path
-    task_list: Optional[list[TaskReadV2]] = Field(default=[])
-    log: Optional[str]
-    info: Optional[str]
+    manifest: ManifestV2
+    python_interpreter: str
+    source: str
+    package_root: Optional[str]
+    package_name: Optional[str]
+    version: Optional[str]
+
+    @root_validator(pre=True)
+    def one_of_package_root_or_name(cls, values):
+        package_root = values["package_root"]
+        package_name = values["package_name"]
+        if (package_root is None and package_name is None) or (
+            package_root is not None and package_name is not None
+        ):
+            raise ValueError(
+                "One and only one must be set between "
+                "'package_root' and 'package_name'"
+            )
+        return values
 
-    def sanitised_dict(self):
+    @validator("package_name")
+    def package_name_prevent_injection(cls, value: str):
         """
-        Return `self.dict()` after casting `self.venv_path` to a string
+        Remove all whitespace characters, and reject values containing `;`.
         """
-        d = self.dict()
-        d["venv_path"] = str(self.venv_path)
-        return d
+        if value is not None:
+            if ";" in value:
+                raise ValueError(f"Invalid package_name: {value}")
+            value = value.replace(" ", "")
+        return value
+
+    @validator("package_root")
+    def package_root_validator(cls, value):
+        if (value is not None) and (not Path(value).is_absolute()):
+            raise ValueError(
+                f"'package_root' must be an absolute path: (given {value})."
+            )
+        return value
+
+    @validator("python_interpreter")
+    def python_interpreter_validator(cls, value):
+        if not Path(value).is_absolute():
+            raise ValueError(
+                f"Python interpreter path must be absolute: (given {value})."
+            )
+        return value
+
+    # Valstr
+    _python_interpreter = validator("python_interpreter", allow_reuse=True)(
+        valstr("python_interpreter")
+    )
+    _source = validator("source", allow_reuse=True)(valstr("source"))
+    _package_root = validator("package_root", allow_reuse=True)(
+        valstr("package_root", accept_none=True)
+    )
+    _package_name = validator("package_name", allow_reuse=True)(
+        valstr("package_name", accept_none=True)
+    )
+    _version = validator("version", allow_reuse=True)(
+        valstr("version", accept_none=True)
+    )
+
+
+class CollectionStateReadV2(BaseModel):
+
+    id: Optional[int]
+    data: dict[str, Any]
+    timestamp: datetime
+
+    _timestamp = validator("timestamp", allow_reuse=True)(valutc("timestamp"))

fractal_server/config.py

@@ -13,6 +13,7 @@
 # Zurich.
 import logging
 import shutil
+import sys
 from os import environ
 from os import getenv
 from os.path import abspath
@@ -323,7 +324,10 @@ class Settings(BaseSettings):
         return FRACTAL_RUNNER_WORKING_BASE_DIR_path
 
     FRACTAL_RUNNER_BACKEND: Literal[
-        "local", "local_experimental", "slurm"
+        "local",
+        "local_experimental",
+        "slurm",
+        "slurm_ssh",
     ] = "local"
     """
     Select which runner backend to use.
@@ -366,10 +370,126 @@ class Settings(BaseSettings):
 
     FRACTAL_SLURM_WORKER_PYTHON: Optional[str] = None
     """
-    Path to Python interpreter that will run the jobs on the SLURM nodes. If
-    not specified, the same interpreter that runs the server is used.
+    Absolute path to Python interpreter that will run the jobs on the SLURM
+    nodes. If not specified, the same interpreter that runs the server is used.
     """
 
+    @validator("FRACTAL_SLURM_WORKER_PYTHON", always=True)
+    def absolute_FRACTAL_SLURM_WORKER_PYTHON(cls, v):
+        """
+        If `FRACTAL_SLURM_WORKER_PYTHON` is a relative path, fail.
+        """
+        if v is None:
+            return None
+        elif not Path(v).is_absolute():
+            raise FractalConfigurationError(
+                f"Non-absolute value for FRACTAL_SLURM_WORKER_PYTHON={v}"
+            )
+        else:
+            return v
+
+    FRACTAL_TASKS_PYTHON_DEFAULT_VERSION: Optional[
+        Literal["3.9", "3.10", "3.11", "3.12"]
+    ] = None
+    """
+    Default Python version to be used for task collection. Defaults to the
+    current version. Requires the corresponding variable (e.g
+    `FRACTAL_TASKS_PYTHON_3_10`) to be set.
+    """
+
+    FRACTAL_TASKS_PYTHON_3_9: Optional[str] = None
+    """
+    Absolute path to the Python 3.9 interpreter that serves as base for virtual
+    environments tasks. Note that this interpreter must have the `venv` module
+    installed. If set, this must be an absolute path. If the version specified
+    in `FRACTAL_TASKS_PYTHON_DEFAULT_VERSION` is `"3.9"` and this attribute is
+    unset, `sys.executable` is used as a default.
+    """
+
+    FRACTAL_TASKS_PYTHON_3_10: Optional[str] = None
+    """
+    Same as `FRACTAL_TASKS_PYTHON_3_9`, for Python 3.10.
+    """
+
+    FRACTAL_TASKS_PYTHON_3_11: Optional[str] = None
+    """
+    Same as `FRACTAL_TASKS_PYTHON_3_9`, for Python 3.11.
+    """
+
+    FRACTAL_TASKS_PYTHON_3_12: Optional[str] = None
+    """
+    Same as `FRACTAL_TASKS_PYTHON_3_9`, for Python 3.12.
+    """
+
+    @root_validator(pre=True)
+    def check_tasks_python(cls, values) -> None:
+        """
+        Perform multiple checks of the Python-intepreter variables.
+
+        1. Each `FRACTAL_TASKS_PYTHON_X_Y` variable must be an absolute path,
+            if set.
+        2. If `FRACTAL_TASKS_PYTHON_DEFAULT_VERSION` is unset, use
+            `sys.executable` and set the corresponding
+            `FRACTAL_TASKS_PYTHON_X_Y` (and unset all others).
+        """
+
+        # `FRACTAL_TASKS_PYTHON_X_Y` variables can only be absolute paths
+        for version in ["3_9", "3_10", "3_11", "3_12"]:
+            key = f"FRACTAL_TASKS_PYTHON_{version}"
+            value = values.get(key)
+            if value is not None and not Path(value).is_absolute():
+                raise FractalConfigurationError(
+                    f"Non-absolute value {key}={value}"
+                )
+
+        default_version = values.get("FRACTAL_TASKS_PYTHON_DEFAULT_VERSION")
+
+        if default_version is not None:
+            # "production/slurm" branch
+            # If a default version is set, then the corresponding interpreter
+            # must also be set
+            default_version_undescore = default_version.replace(".", "_")
+            key = f"FRACTAL_TASKS_PYTHON_{default_version_undescore}"
+            value = values.get(key)
+            if value is None:
+                msg = (
+                    f"FRACTAL_TASKS_PYTHON_DEFAULT_VERSION={default_version} "
+                    f"but {key}={value}."
+                )
+                logging.error(msg)
+                raise FractalConfigurationError(msg)
+
+        else:
+            # If no default version is set, then only `sys.executable` is made
+            # available
+            _info = sys.version_info
+            current_version = f"{_info.major}_{_info.minor}"
+            current_version_dot = f"{_info.major}.{_info.minor}"
+            values[
+                "FRACTAL_TASKS_PYTHON_DEFAULT_VERSION"
+            ] = current_version_dot
+            logging.info(
+                "Setting FRACTAL_TASKS_PYTHON_DEFAULT_VERSION to "
+                f"{current_version_dot}"
+            )
+
+            # Unset all existing intepreters variable
+            for _version in ["3_9", "3_10", "3_11", "3_12"]:
+                key = f"FRACTAL_TASKS_PYTHON_{_version}"
+                if _version == current_version:
+                    values[key] = sys.executable
+                    logging.info(f"Setting {key} to {sys.executable}.")
+                else:
+                    value = values.get(key)
+                    if value is not None:
+                        logging.info(
+                            f"Setting {key} to None (given: {value}), "
+                            "because FRACTAL_TASKS_PYTHON_DEFAULT_VERSION was "
+                            "not set."
+                        )
+                    values[key] = None
+        return values
+
     FRACTAL_SLURM_POLL_INTERVAL: int = 5
     """
     Interval to wait (in seconds) before checking whether unfinished job are
@@ -392,6 +512,25 @@ class Settings(BaseSettings):
     `JobExecutionError`.
     """
 
+    FRACTAL_SLURM_SSH_HOST: Optional[str] = None
+    """
+    SSH-reachable host where a SLURM client is available.
+    """
+    FRACTAL_SLURM_SSH_USER: Optional[str] = None
+    """
+    User on `FRACTAL_SLURM_SSH_HOST`.
+    """
+    FRACTAL_SLURM_SSH_PRIVATE_KEY_PATH: Optional[str] = None
+    """
+    Private key for connecting to `FRACTAL_SLURM_SSH_HOST` as
+    `FRACTAL_SLURM_SSH_USER`.
+    """
+    # FIXME SSH: Split this into two folders (for tasks and for jobs)
+    FRACTAL_SLURM_SSH_WORKING_BASE_DIR: Optional[str] = None
+    """
+    Remote folder on `FRACTAL_SLURM_SSH_HOST`.
+    """
+
     FRACTAL_API_SUBMIT_RATE_LIMIT: int = 2
     """
     Interval to wait (in seconds) to be allowed to call again
@@ -480,6 +619,48 @@ class Settings(BaseSettings):
                     raise FractalConfigurationError(
                         f"{info} but `squeue` command not found."
                     )
+        elif self.FRACTAL_RUNNER_BACKEND == "slurm_ssh":
+            if self.FRACTAL_SLURM_WORKER_PYTHON is None:
+                raise FractalConfigurationError(
+                    f"Must set FRACTAL_SLURM_WORKER_PYTHON when {info}"
+                )
+            if self.FRACTAL_SLURM_SSH_USER is None:
+                raise FractalConfigurationError(
+                    f"Must set FRACTAL_SLURM_SSH_USER when {info}"
+                )
+            if self.FRACTAL_SLURM_SSH_HOST is None:
+                raise FractalConfigurationError(
+                    f"Must set FRACTAL_SLURM_SSH_HOST when {info}"
+                )
+            if self.FRACTAL_SLURM_SSH_PRIVATE_KEY_PATH is None:
+                raise FractalConfigurationError(
+                    f"Must set FRACTAL_SLURM_SSH_PRIVATE_KEY_PATH when {info}"
+                )
+            if self.FRACTAL_SLURM_SSH_WORKING_BASE_DIR is None:
+                raise FractalConfigurationError(
+                    f"Must set FRACTAL_SLURM_SSH_WORKING_BASE_DIR when {info}"
+                )
+
+            from fractal_server.app.runner.executors.slurm._slurm_config import (  # noqa: E501
+                load_slurm_config_file,
+            )
+
+            if not self.FRACTAL_SLURM_CONFIG_FILE:
+                raise FractalConfigurationError(
+                    f"Must set FRACTAL_SLURM_CONFIG_FILE when {info}"
+                )
+            else:
+                if not self.FRACTAL_SLURM_CONFIG_FILE.exists():
+                    raise FractalConfigurationError(
+                        f"{info} but FRACTAL_SLURM_CONFIG_FILE="
+                        f"{self.FRACTAL_SLURM_CONFIG_FILE} not found."
+                    )
+
+                load_slurm_config_file(self.FRACTAL_SLURM_CONFIG_FILE)
+                if not shutil.which("ssh"):
+                    raise FractalConfigurationError(
+                        f"{info} but `ssh` command not found."
+                    )
         else:  # i.e. self.FRACTAL_RUNNER_BACKEND == "local"
             if self.FRACTAL_LOCAL_CONFIG_FILE:
                 if not self.FRACTAL_LOCAL_CONFIG_FILE.exists():

fractal_server/main.py

@@ -20,6 +20,7 @@ from contextlib import asynccontextmanager
 
 from fastapi import FastAPI
 
+from .app.routes.aux._runner import _backend_supports_shutdown  # FIXME: change
 from .app.runner.shutdown import cleanup_after_shutdown
 from .app.security import _create_first_user
 from .config import get_settings
@@ -97,17 +98,38 @@ async def lifespan(app: FastAPI):
         is_superuser=True,
         is_verified=True,
     )
+
+    if settings.FRACTAL_RUNNER_BACKEND == "slurm_ssh":
+        from fractal_server.ssh._fabric import get_ssh_connection
+
+        app.state.connection = get_ssh_connection()
+        logger.info(
+            f"Created SSH connection "
+            f"({app.state.connection.is_connected=})."
+        )
+    else:
+        app.state.connection = 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.connection.is_connected=})."
+        )
+
+        app.state.connection.close()
+
     logger.info(
         f"Current worker with pid {os.getpid()} is shutting down. "
         f"Current jobs: {app.state.jobsV1=}, {app.state.jobsV2=}"
     )
-    if settings.FRACTAL_RUNNER_BACKEND == "slurm":
+    if _backend_supports_shutdown(settings.FRACTAL_RUNNER_BACKEND):
         try:
             await cleanup_after_shutdown(
                 jobsV1=app.state.jobsV1,
@@ -120,6 +142,8 @@ async def lifespan(app: FastAPI):
                 "some of running jobs are not shutdown properly. "
                 f"Original error: {e}"
             )
+    else:
+        logger.info("Shutdown not available for this backend runner.")
 
     logger.info("End application shutdown")
     reset_logger_handlers(logger)

fractal_server/ssh/__init__.py

@@ -0,0 +1,4 @@
+"""
+The `fractal_server.ssh` subpackage is meant as a layer in front of some SSH
+library (e.g. `fabric` or `asyncssh`).
+"""

fractal_server/ssh/_fabric.py

@@ -0,0 +1,190 @@
+import time
+from typing import Optional
+
+from fabric import Connection
+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
+
+
+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 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)