fractal-server 2.5.1__py3-none-any.whl → 2.6.0a0__py3-none-any.whl

fractal_server/app/routes/aux/validate_user_settings.py

@@ -0,0 +1,74 @@
+from fastapi import HTTPException
+from fastapi import status
+from pydantic import BaseModel
+from pydantic import ValidationError
+
+from fractal_server.app.db import AsyncSession
+from fractal_server.app.models import UserOAuth
+from fractal_server.app.models import UserSettings
+from fractal_server.app.user_settings import SlurmSshUserSettings
+from fractal_server.app.user_settings import SlurmSudoUserSettings
+from fractal_server.logger import set_logger
+
+logger = set_logger(__name__)
+
+
+def verify_user_has_settings(user: UserOAuth) -> None:
+    """
+    Check that the `user.user_settings_id` foreign-key is set.
+
+    NOTE: This check will become useless when we make the foreign-key column
+    required, but for the moment (as of v2.6.0) we have to keep it in place.
+
+    Arguments:
+        user: The user to be checked.
+    """
+    if user.user_settings_id is None:
+        raise HTTPException(
+            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail=f"Error: user '{user.email}' has no settings.",
+        )
+
+
+async def validate_user_settings(
+    *, user: UserOAuth, backend: str, db: AsyncSession
+) -> UserSettings:
+    """
+    Get a UserSettings object and validate it based on a given Fractal backend.
+
+    Arguments:
+        user: The user whose settings we should validate.
+        backend: The value of `FRACTAL_RUNNER_BACKEND`
+        db: An async DB session
+
+    Returns:
+        `UserSetting` object associated to `user`, if valid.
+    """
+
+    verify_user_has_settings(user)
+
+    user_settings = await db.get(UserSettings, user.user_settings_id)
+
+    if backend == "slurm_ssh":
+        UserSettingsValidationModel = SlurmSshUserSettings
+    elif backend == "slurm":
+        UserSettingsValidationModel = SlurmSudoUserSettings
+    else:
+        # For other backends, we don't validate anything
+        UserSettingsValidationModel = BaseModel
+
+    try:
+        UserSettingsValidationModel(**user_settings.model_dump())
+    except ValidationError as e:
+        error_msg = (
+            "User settings are not valid for "
+            f"FRACTAL_RUNNER_BACKEND='{backend}'. "
+            f"Original error: {str(e)}"
+        )
+        logger.warning(error_msg)
+        raise HTTPException(
+            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail=error_msg,
+        )
+
+    return user_settings

fractal_server/app/runner/executors/slurm/ssh/executor.py

@@ -163,17 +163,34 @@ class FractalSlurmSSHExecutor(SlurmExecutor):
         settings = Inject(get_settings)
         self.python_remote = settings.FRACTAL_SLURM_WORKER_PYTHON
         if self.python_remote is None:
+            self._stop_and_join_wait_thread()
             raise ValueError("FRACTAL_SLURM_WORKER_PYTHON is not set. Exit.")
 
         # Initialize connection and perform handshake
         self.fractal_ssh = fractal_ssh
         logger.warning(self.fractal_ssh)
-        self.handshake()
+        try:
+            self.handshake()
+        except Exception as e:
+            logger.warning(
+                "Stop/join waiting thread and then "
+                f"re-raise original error {str(e)}"
+            )
+            self._stop_and_join_wait_thread()
+            raise e
 
         # Set/validate parameters for SLURM submission scripts
         self.slurm_account = slurm_account
         self.common_script_lines = common_script_lines or []
-        self._validate_common_script_lines()
+        try:
+            self._validate_common_script_lines()
+        except Exception as e:
+            logger.warning(
+                "Stop/join waiting thread and then "
+                f"re-raise original error {str(e)}"
+            )
+            self._stop_and_join_wait_thread()
+            raise e
 
         # Set/initialize some more options
         self.keep_pickle_files = keep_pickle_files
@@ -1385,6 +1402,10 @@ class FractalSlurmSSHExecutor(SlurmExecutor):
             self.fractal_ssh.run_command(cmd=scancel_command)
         logger.debug("Executor shutdown: end")
 
+    def _stop_and_join_wait_thread(self):
+        self.wait_thread.stop()
+        self.wait_thread.join()
+
     def __exit__(self, *args, **kwargs):
         """
         See
@@ -1393,8 +1414,7 @@ class FractalSlurmSSHExecutor(SlurmExecutor):
         logger.debug(
             "[FractalSlurmSSHExecutor.__exit__] Stop and join `wait_thread`"
         )
-        self.wait_thread.stop()
-        self.wait_thread.join()
+        self._stop_and_join_wait_thread()
         logger.debug("[FractalSlurmSSHExecutor.__exit__] End")
 
     def run_squeue(self, job_ids):

fractal_server/app/runner/executors/slurm/sudo/executor.py

@@ -259,6 +259,7 @@ class FractalSlurmExecutor(SlurmExecutor):
                 for line in self.common_script_lines
                 if line.startswith("#SBATCH --account=")
             )
+            self._stop_and_join_wait_thread()
             raise RuntimeError(
                 "Invalid line in `FractalSlurmExecutor.common_script_lines`: "
                 f"'{invalid_line}'.\n"
@@ -1287,6 +1288,10 @@ class FractalSlurmExecutor(SlurmExecutor):
 
         logger.debug("Executor shutdown: end")
 
+    def _stop_and_join_wait_thread(self):
+        self.wait_thread.stop()
+        self.wait_thread.join()
+
     def __exit__(self, *args, **kwargs):
         """
         See
@@ -1295,6 +1300,5 @@ class FractalSlurmExecutor(SlurmExecutor):
         logger.debug(
             "[FractalSlurmExecutor.__exit__] Stop and join `wait_thread`"
         )
-        self.wait_thread.stop()
-        self.wait_thread.join()
+        self._stop_and_join_wait_thread()
         logger.debug("[FractalSlurmExecutor.__exit__] End")

fractal_server/app/runner/v2/__init__.py

@@ -42,6 +42,7 @@ from .handle_failed_job import assemble_filters_failed_job
 from .handle_failed_job import assemble_history_failed_job
 from .handle_failed_job import assemble_images_failed_job
 from fractal_server import __VERSION__
+from fractal_server.app.models import UserSettings
 
 _backends = {}
 _backends["local"] = local_process_workflow
@@ -76,6 +77,7 @@ async def submit_workflow(
     workflow_id: int,
     dataset_id: int,
     job_id: int,
+    user_settings: UserSettings,
     worker_init: Optional[str] = None,
     slurm_user: Optional[str] = None,
     user_cache_dir: Optional[str] = None,
@@ -196,8 +198,7 @@ async def submit_workflow(
             elif FRACTAL_RUNNER_BACKEND == "slurm_ssh":
                 # Folder creation is deferred to _process_workflow
                 WORKFLOW_DIR_REMOTE = (
-                    Path(settings.FRACTAL_SLURM_SSH_WORKING_BASE_DIR)
-                    / WORKFLOW_DIR_LOCAL.name
+                    Path(user_settings.ssh_jobs_dir) / WORKFLOW_DIR_LOCAL.name
                 )
             else:
                 logger.error(
@@ -270,11 +271,8 @@ async def submit_workflow(
             logger.debug(f"slurm_account: {job.slurm_account}")
             logger.debug(f"worker_init: {worker_init}")
         elif FRACTAL_RUNNER_BACKEND == "slurm_ssh":
-            logger.debug(f"ssh_host: {settings.FRACTAL_SLURM_SSH_HOST}")
-            logger.debug(f"ssh_user: {settings.FRACTAL_SLURM_SSH_USER}")
-            logger.debug(
-                f"base dir: {settings.FRACTAL_SLURM_SSH_WORKING_BASE_DIR}"
-            )
+            logger.debug(f"ssh_user: {user_settings.ssh_username}")
+            logger.debug(f"base dir: {user_settings.ssh_tasks_dir}")
             logger.debug(f"worker_init: {worker_init}")
         logger.debug(f"job.id: {job.id}")
         logger.debug(f"job.working_dir: {job.working_dir}")

fractal_server/app/schemas/__init__.py

@@ -1 +1,3 @@
 from .user import *  # noqa: F401, F403
+from .user_group import *  # noqa: F401, F403
+from .user_settings import *  # noqa: F401, F403

fractal_server/app/schemas/user.py

@@ -3,14 +3,10 @@ from typing import Optional
 from fastapi_users import schemas
 from pydantic import BaseModel
 from pydantic import Extra
-from pydantic import Field
 from pydantic import validator
-from pydantic.types import StrictStr
 
-from ._validators import val_absolute_path
 from ._validators import val_unique_list
 from ._validators import valstr
-from fractal_server.string_tools import validate_cmd
 
 __all__ = (
     "UserRead",
@@ -41,16 +37,10 @@ class UserRead(schemas.BaseUser[int]):
     Schema for `User` read from database.
 
     Attributes:
-        slurm_user:
-        cache_dir:
         username:
-        slurm_accounts:
     """
 
-    slurm_user: Optional[str]
-    cache_dir: Optional[str]
     username: Optional[str]
-    slurm_accounts: list[str]
     group_names: Optional[list[str]] = None
     group_ids: Optional[list[int]] = None
     oauth_accounts: list[OAuthAccountRead]
@@ -61,32 +51,14 @@ class UserUpdate(schemas.BaseUserUpdate):
     Schema for `User` update.
 
     Attributes:
-        slurm_user:
-        cache_dir:
         username:
-        slurm_accounts:
     """
 
-    slurm_user: Optional[str]
-    cache_dir: Optional[str]
     username: Optional[str]
-    slurm_accounts: Optional[list[StrictStr]]
 
     # Validators
-    _slurm_user = validator("slurm_user", allow_reuse=True)(
-        valstr("slurm_user")
-    )
     _username = validator("username", allow_reuse=True)(valstr("username"))
 
-    _slurm_accounts = validator("slurm_accounts", allow_reuse=True)(
-        val_unique_list("slurm_accounts")
-    )
-
-    @validator("cache_dir")
-    def cache_dir_validator(cls, value):
-        validate_cmd(value)
-        return val_absolute_path("cache_dir")(value)
-
     @validator(
         "is_active",
         "is_verified",
@@ -106,21 +78,9 @@ class UserUpdateStrict(BaseModel, extra=Extra.forbid):
     Schema for `User` self-editing.
 
     Attributes:
-        cache_dir:
-        slurm_accounts:
     """
 
-    cache_dir: Optional[str]
-    slurm_accounts: Optional[list[StrictStr]]
-
-    _slurm_accounts = validator("slurm_accounts", allow_reuse=True)(
-        val_unique_list("slurm_accounts")
-    )
-
-    @validator("cache_dir")
-    def cache_dir_validator(cls, value):
-        validate_cmd(value)
-        return val_absolute_path("cache_dir")(value)
+    pass
 
 
 class UserUpdateWithNewGroupIds(UserUpdate):
@@ -136,32 +96,11 @@ class UserCreate(schemas.BaseUserCreate):
     Schema for `User` creation.
 
     Attributes:
-        slurm_user:
-        cache_dir:
         username:
-        slurm_accounts:
     """
 
-    slurm_user: Optional[str]
-    cache_dir: Optional[str]
     username: Optional[str]
-    slurm_accounts: list[StrictStr] = Field(default_factory=list)
 
     # Validators
 
-    @validator("slurm_accounts")
-    def slurm_accounts_validator(cls, value):
-        for i, element in enumerate(value):
-            value[i] = valstr(attribute=f"slurm_accounts[{i}]")(element)
-        val_unique_list("slurm_accounts")(value)
-        return value
-
-    _slurm_user = validator("slurm_user", allow_reuse=True)(
-        valstr("slurm_user")
-    )
     _username = validator("username", allow_reuse=True)(valstr("username"))
-
-    @validator("cache_dir")
-    def cache_dir_validator(cls, value):
-        validate_cmd(value)
-        return val_absolute_path("cache_dir")(value)

fractal_server/app/schemas/user_settings.py

@@ -0,0 +1,93 @@
+from typing import Optional
+
+from pydantic import BaseModel
+from pydantic import Extra
+from pydantic import validator
+from pydantic.types import StrictStr
+
+from ._validators import val_absolute_path
+from ._validators import val_unique_list
+from ._validators import valstr
+from fractal_server.string_tools import validate_cmd
+
+__all__ = (
+    "UserSettingsRead",
+    "UserSettingsReadStrict",
+    "UserSettingsUpdate",
+    "UserSettingsUpdateStrict",
+)
+
+
+class UserSettingsRead(BaseModel):
+    id: int
+    ssh_host: Optional[str] = None
+    ssh_username: Optional[str] = None
+    ssh_private_key_path: Optional[str] = None
+    ssh_tasks_dir: Optional[str] = None
+    ssh_jobs_dir: Optional[str] = None
+    slurm_user: Optional[str] = None
+    slurm_accounts: list[str]
+    cache_dir: Optional[str] = None
+
+
+class UserSettingsReadStrict(BaseModel):
+    slurm_user: Optional[str] = None
+    slurm_accounts: list[str]
+    cache_dir: Optional[str] = None
+
+
+class UserSettingsUpdate(BaseModel, extra=Extra.forbid):
+    ssh_host: Optional[str] = None
+    ssh_username: Optional[str] = None
+    ssh_private_key_path: Optional[str] = None
+    ssh_tasks_dir: Optional[str] = None
+    ssh_jobs_dir: Optional[str] = None
+    slurm_user: Optional[str] = None
+    slurm_accounts: Optional[list[StrictStr]] = None
+    cache_dir: Optional[str] = None
+
+    _ssh_host = validator("ssh_host", allow_reuse=True)(valstr("ssh_host"))
+    _ssh_username = validator("ssh_username", allow_reuse=True)(
+        valstr("ssh_username")
+    )
+    _ssh_private_key_path = validator(
+        "ssh_private_key_path", allow_reuse=True
+    )(val_absolute_path("ssh_private_key_path"))
+
+    _ssh_tasks_dir = validator("ssh_tasks_dir", allow_reuse=True)(
+        val_absolute_path("ssh_tasks_dir")
+    )
+    _ssh_jobs_dir = validator("ssh_jobs_dir", allow_reuse=True)(
+        val_absolute_path("ssh_jobs_dir")
+    )
+
+    _slurm_user = validator("slurm_user", allow_reuse=True)(
+        valstr("slurm_user")
+    )
+
+    @validator("slurm_accounts")
+    def slurm_accounts_validator(cls, value):
+        if value is None:
+            return value
+        for i, item in enumerate(value):
+            value[i] = valstr(f"slurm_accounts[{i}]")(item)
+        return val_unique_list("slurm_accounts")(value)
+
+    @validator("cache_dir")
+    def cache_dir_validator(cls, value):
+        validate_cmd(value)
+        return val_absolute_path("cache_dir")(value)
+
+
+class UserSettingsUpdateStrict(BaseModel, extra=Extra.forbid):
+    slurm_accounts: Optional[list[StrictStr]] = None
+    cache_dir: Optional[str] = None
+
+    _slurm_accounts = validator("slurm_accounts", allow_reuse=True)(
+        val_unique_list("slurm_accounts")
+    )
+
+    @validator("cache_dir")
+    def cache_dir_validator(cls, value):
+        validate_cmd(value)
+        return val_absolute_path("cache_dir")(value)

fractal_server/app/security/__init__.py

@@ -54,6 +54,7 @@ from fractal_server.app.models import LinkUserGroup
 from fractal_server.app.models import OAuthAccount
 from fractal_server.app.models import UserGroup
 from fractal_server.app.models import UserOAuth
+from fractal_server.app.models import UserSettings
 from fractal_server.app.schemas.user import UserCreate
 from fractal_server.logger import set_logger
 
@@ -193,6 +194,8 @@ class UserManager(IntegerIDMixin, BaseUserManager[UserOAuth, int]):
     async def on_after_register(
         self, user: UserOAuth, request: Optional[Request] = None
     ):
+        logger = set_logger("fractal_server.on_after_register")
+
         logger.info(
             f"New-user registration completed ({user.id=}, {user.email=})."
         )
@@ -204,24 +207,30 @@ class UserManager(IntegerIDMixin, BaseUserManager[UserOAuth, int]):
             res = await db.execute(stm)
             default_group = res.scalar_one_or_none()
             if default_group is None:
-                logger.error(
+                logger.warning(
                     f"No group found with name {FRACTAL_DEFAULT_GROUP_NAME}"
                 )
             else:
-                logger.warning(
-                    f"START adding {user.email} user to group "
-                    f"{default_group.id=}."
-                )
                 link = LinkUserGroup(
                     user_id=user.id, group_id=default_group.id
                 )
                 db.add(link)
                 await db.commit()
-                logger.warning(
-                    f"END   adding {user.email} user to group "
-                    f"{default_group.id=}."
+                logger.info(
+                    f"Added {user.email} user to group {default_group.id=}."
                 )
 
+            this_user = await db.get(UserOAuth, user.id)
+
+            this_user.settings = UserSettings()
+            await db.merge(this_user)
+            await db.commit()
+            await db.refresh(this_user)
+            logger.info(
+                f"Associated empty settings (id={this_user.user_settings_id}) "
+                f"to '{this_user.email}'."
+            )
+
 
 async def get_user_manager(
     user_db: SQLModelUserDatabaseAsync = Depends(get_user_db),
@@ -294,9 +303,13 @@ async def _create_first_user(
                         kwargs["username"] = username
                     user = await user_manager.create(UserCreate(**kwargs))
                     function_logger.info(f"User '{user.email}' created")
-
     except UserAlreadyExists:
         function_logger.warning(f"User '{email}' already exists")
+    except Exception as e:
+        function_logger.error(
+            f"ERROR in _create_first_user, original error {str(e)}"
+        )
+        raise e
     finally:
         function_logger.info(f"END   _create_first_user, with email '{email}'")
 

fractal_server/app/user_settings.py

@@ -0,0 +1,42 @@
+# TODO: move this file to the appropriate path
+from pydantic import BaseModel
+
+
+class SlurmSshUserSettings(BaseModel):
+    """
+    Subset of user settings which must be present for task collection and job
+    execution when using the Slurm-SSH runner.
+
+    Attributes:
+        ssh_host: SSH-reachable host where a SLURM client is available.
+        ssh_username: User on `ssh_host`.
+        ssh_private_key_path: Path of private SSH key for `ssh_username`.
+        ssh_tasks_dir: Task-venvs base folder on `ssh_host`.
+        ssh_jobs_dir: Jobs base folder on `ssh_host`.
+        slurm_accounts:
+            List of SLURM accounts, to be used upon Fractal job submission.
+    """
+
+    ssh_host: str
+    ssh_username: str
+    ssh_private_key_path: str
+    ssh_tasks_dir: str
+    ssh_jobs_dir: str
+    slurm_accounts: list[str]
+
+
+class SlurmSudoUserSettings(BaseModel):
+    """
+    Subset of user settings which must be present for task collection and job
+    execution when using the Slurm-sudo runner.
+
+    Attributes:
+        slurm_user: User to be impersonated via `sudo -u`.
+        cache_dir: Folder where `slurm_user` can write.
+        slurm_accounts:
+            List of SLURM accounts, to be used upon Fractal job submission.
+    """
+
+    slurm_user: str
+    cache_dir: str
+    slurm_accounts: list[str]

fractal_server/config.py

@@ -631,22 +631,6 @@ class Settings(BaseSettings):
                 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,

fractal_server/data_migrations/2_6_0.py

@@ -0,0 +1,49 @@
+import logging
+
+from sqlalchemy import select
+
+from fractal_server.app.db import get_sync_db
+from fractal_server.app.models import UserOAuth
+from fractal_server.app.models import UserSettings
+from fractal_server.config import get_settings
+from fractal_server.data_migrations.tools import _check_current_version
+from fractal_server.syringe import Inject
+
+
+def fix_db():
+    logger = logging.getLogger("fix_db")
+    logger.warning("START execution of fix_db function")
+    _check_current_version("2.6.0")
+
+    global_settings = Inject(get_settings)
+
+    with next(get_sync_db()) as db:
+        users = db.execute(select(UserOAuth)).scalars().unique().all()
+        for user in sorted(users, key=lambda x: x.id):
+            logger.warning(f"START handling user {user.id}: '{user.email}'")
+            user_settings = UserSettings(
+                # SSH
+                ssh_host=global_settings.FRACTAL_SLURM_SSH_HOST,
+                ssh_username=global_settings.FRACTAL_SLURM_SSH_USER,
+                ssh_private_key_path=(
+                    global_settings.FRACTAL_SLURM_SSH_PRIVATE_KEY_PATH
+                ),
+                ssh_tasks_dir=(
+                    global_settings.FRACTAL_SLURM_SSH_WORKING_BASE_DIR
+                ),
+                ssh_jobs_dir=(
+                    global_settings.FRACTAL_SLURM_SSH_WORKING_BASE_DIR
+                ),
+                # SUDO
+                slurm_user=user.slurm_user,
+                slurm_accounts=user.slurm_accounts,
+                cache_dir=user.cache_dir,
+            )
+            user.settings = user_settings
+            db.add(user)
+            db.commit()
+            db.refresh(user)
+            logger.warning(f"New user {user.id} settings:\n{user.settings}")
+            logger.warning(f"END handling user {user.id}: '{user.email}'")
+
+    logger.warning("END of execution of fix_db function")

fractal_server/data_migrations/tools.py

@@ -0,0 +1,17 @@
+from packaging.version import parse
+
+import fractal_server
+
+
+def _check_current_version(expected_version: str):
+    # Check that this module matches with the current version
+    module_version = parse(expected_version)
+    current_version = parse(fractal_server.__VERSION__)
+    if (
+        current_version.major != module_version.major
+        or current_version.minor != module_version.minor
+        or current_version.micro != module_version.micro
+    ):
+        raise RuntimeError(
+            f"{fractal_server.__VERSION__=} not matching with {__file__=}"
+        )

fractal_server/main.py

@@ -92,32 +92,34 @@ async def lifespan(app: FastAPI):
     settings = Inject(get_settings)
 
     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)
+        from fractal_server.ssh._fabric import FractalSSHList
+
+        app.state.fractal_ssh_list = FractalSSHList()
+
         logger.info(
-            f"Created SSH connection "
-            f"({app.state.fractal_ssh.is_connected=})."
+            "Added empty FractalSSHList to app.state "
+            f"(id={id(app.state.fractal_ssh_list)})."
         )
     else:
-        app.state.fractal_ssh = None
+        app.state.fractal_ssh_list = 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=})."
+            "Close FractalSSH connections "
+            f"(current size: {app.state.fractal_ssh_list.size})."
         )
 
-        app.state.fractal_ssh.close()
+        app.state.fractal_ssh_list.close_all()
 
     logger.info(
         f"Current worker with pid {os.getpid()} is shutting down. "