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

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. "

fractal_server/migrations/versions/9c5ae74c9b98_add_user_settings_table.py

@@ -0,0 +1,74 @@
+"""Add user_settings table
+
+Revision ID: 9c5ae74c9b98
+Revises: d9a140db5d42
+Create Date: 2024-09-24 12:01:13.393326
+
+"""
+import sqlalchemy as sa
+import sqlmodel
+from alembic import op
+
+
+# revision identifiers, used by Alembic.
+revision = "9c5ae74c9b98"
+down_revision = "d9a140db5d42"
+branch_labels = None
+depends_on = None
+
+# Manually define constraint name, see issue #1777
+CONSTRAINT_NAME = "fk_user_oauth_user_settings_id_user_settings"
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table(
+        "user_settings",
+        sa.Column("id", sa.Integer(), nullable=False),
+        sa.Column(
+            "slurm_accounts", sa.JSON(), server_default="[]", nullable=False
+        ),
+        sa.Column(
+            "ssh_host", sqlmodel.sql.sqltypes.AutoString(), nullable=True
+        ),
+        sa.Column(
+            "ssh_username", sqlmodel.sql.sqltypes.AutoString(), nullable=True
+        ),
+        sa.Column(
+            "ssh_private_key_path",
+            sqlmodel.sql.sqltypes.AutoString(),
+            nullable=True,
+        ),
+        sa.Column(
+            "ssh_tasks_dir", sqlmodel.sql.sqltypes.AutoString(), nullable=True
+        ),
+        sa.Column(
+            "ssh_jobs_dir", sqlmodel.sql.sqltypes.AutoString(), nullable=True
+        ),
+        sa.Column(
+            "slurm_user", sqlmodel.sql.sqltypes.AutoString(), nullable=True
+        ),
+        sa.Column(
+            "cache_dir", sqlmodel.sql.sqltypes.AutoString(), nullable=True
+        ),
+        sa.PrimaryKeyConstraint("id"),
+    )
+    with op.batch_alter_table("user_oauth", schema=None) as batch_op:
+        batch_op.add_column(
+            sa.Column("user_settings_id", sa.Integer(), nullable=True)
+        )
+        batch_op.create_foreign_key(
+            CONSTRAINT_NAME, "user_settings", ["user_settings_id"], ["id"]
+        )
+
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("user_oauth", schema=None) as batch_op:
+        batch_op.drop_constraint(CONSTRAINT_NAME, type_="foreignkey")
+        batch_op.drop_column("user_settings_id")
+
+    op.drop_table("user_settings")
+    # ### end Alembic commands ###

fractal_server/ssh/_fabric.py

@@ -16,20 +16,21 @@ 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.string_tools import validate_cmd
-from fractal_server.syringe import Inject
 
 
 class FractalSSHTimeoutError(RuntimeError):
     pass
 
 
+class FractalSSHListTimeoutError(RuntimeError):
+    pass
+
+
 logger = set_logger(__name__)
 
 
 class FractalSSH(object):
-
     """
     FIXME SSH: Fix docstring
 
@@ -111,7 +112,6 @@ class FractalSSH(object):
     def run(
         self, *args, lock_timeout: Optional[float] = None, **kwargs
     ) -> Any:
-
         actual_lock_timeout = self.default_lock_timeout
         if lock_timeout is not None:
             actual_lock_timeout = lock_timeout
@@ -138,7 +138,19 @@ class FractalSSH(object):
                 )
 
     def close(self) -> None:
-        return self._connection.close()
+        """
+        Aggressively close `self._connection`.
+
+        When `Connection.is_connected` is `False`, `Connection.close()` does
+        not call `Connection.client.close()`. Thus we do this explicitly here,
+        because we observed cases where `is_connected=False` but the underlying
+        `Transport` object was not closed.
+        """
+
+        self._connection.close()
+
+        if self._connection.client is not None:
+            self._connection.client.close()
 
     def run_command(
         self,
@@ -293,21 +305,21 @@ class FractalSSH(object):
             safe_root: If `folder` is not a subfolder of the absolute
                 `safe_root` path, raise an error.
         """
-        invalid_characters = {" ", "\n", ";", "$", "`"}
-
-        if (
-            not isinstance(folder, str)
-            or not isinstance(safe_root, str)
-            or len(invalid_characters.intersection(folder)) > 0
-            or len(invalid_characters.intersection(safe_root)) > 0
-            or not Path(folder).is_absolute()
-            or not Path(safe_root).is_absolute()
-            or not Path(folder).resolve().is_relative_to(safe_root)
+        validate_cmd(folder)
+        validate_cmd(safe_root)
+
+        if " " in folder:
+            raise ValueError(f"folder='{folder}' includes whitespace.")
+        elif " " in safe_root:
+            raise ValueError(f"safe_root='{safe_root}' includes whitespace.")
+        elif not Path(folder).is_absolute():
+            raise ValueError(f"{folder=} is not an absolute path.")
+        elif not Path(safe_root).is_absolute():
+            raise ValueError(f"{safe_root=} is not an absolute path.")
+        elif not (
+            Path(folder).resolve().is_relative_to(Path(safe_root).resolve())
         ):
-            raise ValueError(
-                f"{folder=} argument is invalid or it is not "
-                f"relative to {safe_root=}."
-            )
+            raise ValueError(f"{folder=} is not a subfolder of {safe_root=}.")
         else:
             cmd = f"rm -r {folder}"
             self.run_command(cmd=cmd)
@@ -335,36 +347,169 @@ class FractalSSH(object):
                 f.write(content)
 
 
-def get_ssh_connection(
-    *,
-    host: Optional[str] = None,
-    user: Optional[str] = None,
-    key_filename: Optional[str] = None,
-) -> Connection:
+class FractalSSHList(object):
     """
-    Create a `fabric.Connection` object based on fractal-server settings
-    or explicit arguments.
+    Collection of `FractalSSH` objects
 
-    Args:
-        host:
-        user:
-        key_filename:
+    Attributes are all private, and access to this collection must be
+    through methods (mostly the `get` one).
 
-    Returns:
-        Fabric connection object
+    Attributes:
+        _data:
+            Mapping of unique keys (the SSH-credentials tuples) to
+            `FractalSSH` objects.
+        _lock:
+            A `threading.Lock object`, to be acquired when changing `_data`.
+        _timeout: Timeout for `_lock` acquisition.
+        _logger_name: Logger name.
     """
-    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,
-        forward_agent=False,
-        connect_kwargs={"key_filename": key_filename},
-    )
-    return connection
+
+    _data: dict[tuple[str, str, str], FractalSSH]
+    _lock: Lock
+    _timeout: float
+    _logger_name: str
+
+    def __init__(
+        self,
+        *,
+        timeout: float = 5.0,
+        logger_name: str = "fractal_server.FractalSSHList",
+    ):
+        self._lock = Lock()
+        self._data = {}
+        self._timeout = timeout
+        self._logger_name = logger_name
+        set_logger(self._logger_name)
+
+    @property
+    def logger(self) -> logging.Logger:
+        """
+        This property exists so that we never have to propagate the
+        `Logger` object.
+        """
+        return get_logger(self._logger_name)
+
+    @property
+    def size(self) -> int:
+        """
+        Number of current key-value pairs in `self._data`.
+        """
+        return len(self._data.values())
+
+    def get(self, *, host: str, user: str, key_path: str) -> FractalSSH:
+        """
+        Get the `FractalSSH` for the current credentials, or create one.
+
+        Note: Changing `_data` requires acquiring `_lock`.
+
+        Arguments:
+            host:
+            user:
+            key_path:
+        """
+        key = (host, user, key_path)
+        fractal_ssh = self._data.get(key, None)
+        if fractal_ssh is not None:
+            self.logger.info(
+                f"Return existing FractalSSH object for {user}@{host}"
+            )
+            return fractal_ssh
+        else:
+            self.logger.info(f"Add new FractalSSH object for {user}@{host}")
+            connection = Connection(
+                host=host,
+                user=user,
+                forward_agent=False,
+                connect_kwargs={
+                    "key_filename": key_path,
+                    "look_for_keys": False,
+                },
+            )
+            with self.acquire_lock_with_timeout():
+                self._data[key] = FractalSSH(connection=connection)
+                return self._data[key]
+
+    def contains(
+        self,
+        *,
+        host: str,
+        user: str,
+        key_path: str,
+    ) -> bool:
+        """
+        Return whether a given key is present in the collection.
+
+        Arguments:
+            host:
+            user:
+            key_path:
+        """
+        key = (host, user, key_path)
+        return key in self._data.keys()
+
+    def remove(
+        self,
+        *,
+        host: str,
+        user: str,
+        key_path: str,
+    ) -> None:
+        """
+        Remove a key from `_data` and close the corresponding connection.
+
+        Note: Changing `_data` requires acquiring `_lock`.
+
+        Arguments:
+            host:
+            user:
+            key_path:
+        """
+        key = (host, user, key_path)
+        with self.acquire_lock_with_timeout():
+            self.logger.info(
+                f"Removing FractalSSH object for {user}@{host} "
+                "from collection."
+            )
+            fractal_ssh_obj = self._data.pop(key)
+            self.logger.info(
+                f"Closing FractalSSH object for {user}@{host} "
+                f"({fractal_ssh_obj.is_connected=})."
+            )
+            fractal_ssh_obj.close()
+
+    def close_all(self, *, timeout: float = 5.0):
+        """
+        Close all `FractalSSH` objects in the collection.
+
+        Arguments:
+            timeout:
+                Timeout for `FractalSSH._lock` acquisition, to be obtained
+                before closing.
+        """
+        for key, fractal_ssh_obj in self._data.items():
+            host, user, _ = key[:]
+            self.logger.info(
+                f"Closing FractalSSH object for {user}@{host} "
+                f"({fractal_ssh_obj.is_connected=})."
+            )
+            with fractal_ssh_obj.acquire_timeout(timeout=timeout):
+                fractal_ssh_obj.close()
+
+    @contextmanager
+    def acquire_lock_with_timeout(self) -> Generator[Literal[True], Any, None]:
+        self.logger.debug(
+            f"Trying to acquire lock, with timeout {self._timeout} s"
+        )
+        result = self._lock.acquire(timeout=self._timeout)
+        try:
+            if not result:
+                self.logger.error("Lock was *NOT* acquired.")
+                raise FractalSSHListTimeoutError(
+                    f"Failed to acquire lock within {self._timeout} ss"
+                )
+            self.logger.debug("Lock was acquired.")
+            yield result
+        finally:
+            if result:
+                self._lock.release()
+                self.logger.debug("Lock was released")

fractal_server/string_tools.py

@@ -60,6 +60,8 @@ def validate_cmd(command: str, allow_char: Optional[str] = None):
         command: command to validate.
         allow_char: chars to accept among the forbidden ones
     """
+    if not isinstance(command, str):
+        raise ValueError(f"{command=} is not a string.")
     forbidden = set(__NOT_ALLOWED_FOR_COMMANDS__)
     if allow_char is not None:
         forbidden = forbidden - set(allow_char)

fractal_server/tasks/v2/background_operations_ssh.py

@@ -57,6 +57,7 @@ def _customize_and_run_template(
     tmpdir: str,
     logger_name: str,
     fractal_ssh: FractalSSH,
+    tasks_base_dir: str,
 ) -> str:
     """
     Customize one of the template bash scripts, transfer it to the remote host
@@ -72,7 +73,6 @@ def _customize_and_run_template(
     """
     logger = get_logger(logger_name)
     logger.debug(f"_customize_and_run_template {script_filename} - START")
-    settings = Inject(get_settings)
 
     # Read template
     template_path = templates_folder / script_filename
@@ -88,7 +88,7 @@ def _customize_and_run_template(
 
     # Transfer script to remote host
     script_path_remote = os.path.join(
-        settings.FRACTAL_SLURM_SSH_WORKING_BASE_DIR,
+        tasks_base_dir,
         f"script_{abs(hash(tmpdir))}{script_filename}",
     )
     logger.debug(f"Now transfer {script_path_local=} over SSH.")
@@ -111,6 +111,7 @@ def background_collect_pip_ssh(
     state_id: int,
     task_pkg: _TaskCollectPip,
     fractal_ssh: FractalSSH,
+    tasks_base_dir: str,
 ) -> None:
     """
     Collect a task package over SSH
@@ -121,6 +122,13 @@ def background_collect_pip_ssh(
     NOTE: by making this function sync, it will run within a thread - due to
     starlette/fastapi handling of background tasks (see
     https://github.com/encode/starlette/blob/master/starlette/background.py).
+
+
+    Arguments:
+        state_id:
+        task_pkg:
+        fractal_ssh:
+        tasks_base_dir:
     """
 
     # Work within a temporary folder, where also logs will be placed
@@ -140,7 +148,6 @@ def background_collect_pip_ssh(
         with next(get_sync_db()) as db:
             try:
                 # Prepare replacements for task-collection scripts
-                settings = Inject(get_settings)
                 python_bin = get_python_interpreter_v2(
                     python_version=task_pkg.python_version
                 )
@@ -163,11 +170,12 @@ def background_collect_pip_ssh(
                         f"{install_string}=={task_pkg.package_version}"
                     )
                 package_env_dir = (
-                    Path(settings.FRACTAL_SLURM_SSH_WORKING_BASE_DIR)
+                    Path(tasks_base_dir)
                     / ".fractal"
                     / f"{task_pkg.package_name}{package_version}"
                 ).as_posix()
                 logger.debug(f"{package_env_dir=}")
+                settings = Inject(get_settings)
                 replacements = [
                     ("__PACKAGE_NAME__", task_pkg.package_name),
                     ("__PACKAGE_ENV_DIR__", package_env_dir),
@@ -186,6 +194,7 @@ def background_collect_pip_ssh(
                     tmpdir=tmpdir,
                     logger_name=LOGGER_NAME,
                     fractal_ssh=fractal_ssh,
+                    tasks_base_dir=tasks_base_dir,
                 )
 
                 fractal_ssh.check_connection()
@@ -332,7 +341,7 @@ def background_collect_pip_ssh(
                         )
                         fractal_ssh.remove_folder(
                             folder=package_env_dir,
-                            safe_root=settings.FRACTAL_SLURM_SSH_WORKING_BASE_DIR,  # noqa: E501
+                            safe_root=tasks_base_dir,
                         )
                         logger.info(
                             f"Deleted remoted folder {package_env_dir}"

{fractal_server-2.5.1.dist-info → fractal_server-2.6.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: fractal-server
-Version: 2.5.1
+Version: 2.6.0
 Summary: Server component of the Fractal analytics platform
 Home-page: https://github.com/fractal-analytics-platform/fractal-server
 License: BSD-3-Clause