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

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
@@ -166,7 +167,7 @@ class Settings(BaseSettings):
     ###########################################################################
     # DATABASE
     ###########################################################################
-    DB_ENGINE: Literal["sqlite", "postgres"] = "sqlite"
+    DB_ENGINE: Literal["sqlite", "postgres", "postgres-psycopg"] = "sqlite"
     """
     Select which database engine to use (supported: `sqlite` and `postgres`).
     """
@@ -201,39 +202,51 @@ class Settings(BaseSettings):
     """
 
     @property
-    def DATABASE_URL(self) -> URL:
-        if self.DB_ENGINE == "sqlite":
-            if not self.SQLITE_PATH:
-                raise FractalConfigurationError(
-                    "SQLITE_PATH path cannot be None"
-                )
-            sqlite_path = abspath(self.SQLITE_PATH)
+    def DATABASE_ASYNC_URL(self) -> URL:
+        if self.DB_ENGINE == "postgres":
             url = URL.create(
-                drivername="sqlite+aiosqlite",
-                database=sqlite_path,
+                drivername="postgresql+asyncpg",
+                username=self.POSTGRES_USER,
+                password=self.POSTGRES_PASSWORD,
+                host=self.POSTGRES_HOST,
+                port=self.POSTGRES_PORT,
+                database=self.POSTGRES_DB,
             )
-            return url
-        elif "postgres":
+        elif self.DB_ENGINE == "postgres-psycopg":
             url = URL.create(
-                drivername="postgresql+asyncpg",
+                drivername="postgresql+psycopg",
                 username=self.POSTGRES_USER,
                 password=self.POSTGRES_PASSWORD,
                 host=self.POSTGRES_HOST,
                 port=self.POSTGRES_PORT,
                 database=self.POSTGRES_DB,
             )
-            return url
+        else:
+            if not self.SQLITE_PATH:
+                raise FractalConfigurationError(
+                    "SQLITE_PATH path cannot be None"
+                )
+            sqlite_path = abspath(self.SQLITE_PATH)
+            url = URL.create(
+                drivername="sqlite+aiosqlite",
+                database=sqlite_path,
+            )
+        return url
 
     @property
     def DATABASE_SYNC_URL(self):
-        if self.DB_ENGINE == "sqlite":
+        if self.DB_ENGINE == "postgres":
+            return self.DATABASE_ASYNC_URL.set(
+                drivername="postgresql+psycopg2"
+            )
+        elif self.DB_ENGINE == "postgres-psycopg":
+            return self.DATABASE_ASYNC_URL.set(drivername="postgresql+psycopg")
+        else:
             if not self.SQLITE_PATH:
                 raise FractalConfigurationError(
                     "SQLITE_PATH path cannot be None"
                 )
-            return self.DATABASE_URL.set(drivername="sqlite")
-        elif self.DB_ENGINE == "postgres":
-            return self.DATABASE_URL.set(drivername="postgresql+psycopg2")
+            return self.DATABASE_ASYNC_URL.set(drivername="sqlite")
 
     ###########################################################################
     # FRACTAL SPECIFIC
@@ -311,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.
@@ -354,9 +370,125 @@ 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
     """
@@ -380,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
@@ -420,6 +571,14 @@ class Settings(BaseSettings):
                     "DB engine is `postgres` but `psycopg2` or `asyncpg` "
                     "are not available"
                 )
+        elif self.DB_ENGINE == "postgres-psycopg":
+            try:
+                import psycopg  # noqa: F401
+            except ModuleNotFoundError:
+                raise FractalConfigurationError(
+                    "DB engine is `postgres-psycopg` but `psycopg` is not "
+                    "available"
+                )
         else:
             if not self.SQLITE_PATH:
                 raise FractalConfigurationError(
@@ -460,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,40 @@ 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
+        from fractal_server.ssh._fabric import FractalSSH
+
+        connection = get_ssh_connection()
+        app.state.fractal_ssh = FractalSSH(connection=connection)
+        logger.info(
+            f"Created SSH connection "
+            f"({app.state.fractal_ssh.is_connected=})."
+        )
+    else:
+        app.state.fractal_ssh = 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.fractal_ssh.is_connected=})."
+        )
+
+        app.state.fractal_ssh.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 +144,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/migrations/env.py

@@ -54,7 +54,7 @@ def run_migrations_offline() -> None:
     settings = Inject(get_settings)
     settings.check_db()
     context.configure(
-        url=settings.DATABASE_URL,
+        url=settings.DATABASE_ASYNC_URL,
         target_metadata=target_metadata,
         literal_binds=True,
         dialect_opts={"paramstyle": "named"},

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,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)