fractal-server 2.2.0a0__py3-none-any.whl → 2.3.0a0__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,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/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,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)

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