fractal-server 2.6.4__py3-none-any.whl → 2.7.0__py3-none-any.whl

fractal_server/ssh/_fabric.py

@@ -1,3 +1,4 @@
+import json
 import logging
 import time
 from contextlib import contextmanager
@@ -23,21 +24,58 @@ class FractalSSHTimeoutError(RuntimeError):
     pass
 
 
-class FractalSSHListTimeoutError(RuntimeError):
-    pass
+logger = set_logger(__name__)
 
 
-logger = set_logger(__name__)
+@contextmanager
+def _acquire_lock_with_timeout(
+    lock: Lock,
+    label: str,
+    timeout: float,
+    logger_name: str = __name__,
+) -> Generator[Literal[True], Any, None]:
+    """
+    Given a `threading.Lock` object, try to acquire it within a given timeout.
+
+    Arguments:
+        lock:
+        label:
+        timeout:
+        logger_name:
+    """
+    logger = get_logger(logger_name)
+    logger.info(f"Trying to acquire lock for '{label}', with {timeout=}")
+    result = lock.acquire(timeout=timeout)
+    try:
+        if not result:
+            logger.error(f"Lock for '{label}' was *not* acquired.")
+            raise FractalSSHTimeoutError(
+                f"Failed to acquire lock for '{label}' within "
+                f"{timeout} seconds"
+            )
+        logger.info(f"Lock for '{label}' was acquired.")
+        yield result
+    finally:
+        if result:
+            lock.release()
+            logger.info(f"Lock for '{label}' was released.")
 
 
 class FractalSSH(object):
     """
-    FIXME SSH: Fix docstring
+    Wrapper of `fabric.Connection` object, enriched with locks.
+
+    Note: methods marked as `_unsafe` should not be used directly,
+    since they do not enforce locking.
 
     Attributes:
         _lock:
-        connection:
+        _connection:
         default_lock_timeout:
+        default_max_attempts:
+        default_base_interval:
+        sftp_get_prefetch:
+        sftp_get_max_requests:
         logger_name:
     """
 
@@ -46,6 +84,8 @@ class FractalSSH(object):
     default_lock_timeout: float
     default_max_attempts: int
     default_base_interval: float
+    sftp_get_prefetch: bool
+    sftp_get_max_requests: int
     logger_name: str
 
     def __init__(
@@ -54,6 +94,8 @@ class FractalSSH(object):
         default_timeout: float = 250,
         default_max_attempts: int = 5,
         default_base_interval: float = 3.0,
+        sftp_get_prefetch: bool = False,
+        sftp_get_max_requests: int = 64,
         logger_name: str = __name__,
     ):
         self._lock = Lock()
@@ -61,28 +103,11 @@ class FractalSSH(object):
         self.default_lock_timeout = default_timeout
         self.default_base_interval = default_base_interval
         self.default_max_attempts = default_max_attempts
+        self.sftp_get_prefetch = sftp_get_prefetch
+        self.sftp_get_max_requests = sftp_get_max_requests
         self.logger_name = logger_name
         set_logger(self.logger_name)
 
-    @contextmanager
-    def acquire_timeout(
-        self, timeout: float
-    ) -> Generator[Literal[True], Any, None]:
-        self.logger.debug(f"Trying to acquire lock, with {timeout=}")
-        result = self._lock.acquire(timeout=timeout)
-        try:
-            if not result:
-                self.logger.error("Lock was *NOT* acquired.")
-                raise FractalSSHTimeoutError(
-                    f"Failed to acquire lock within {timeout} seconds"
-                )
-            self.logger.debug("Lock was acquired.")
-            yield result
-        finally:
-            if result:
-                self._lock.release()
-                self.logger.debug("Lock was released")
-
     @property
     def is_connected(self) -> bool:
         return self._connection.is_connected
@@ -91,36 +116,82 @@ class FractalSSH(object):
     def logger(self) -> logging.Logger:
         return get_logger(self.logger_name)
 
-    def put(
-        self, *args, lock_timeout: Optional[float] = None, **kwargs
+    def _put(
+        self,
+        *,
+        local: str,
+        remote: str,
+        label: str,
+        lock_timeout: Optional[float] = None,
     ) -> Result:
+        """
+        Transfer a local file to a remote path, via SFTP.
+        """
         actual_lock_timeout = self.default_lock_timeout
         if lock_timeout is not None:
             actual_lock_timeout = lock_timeout
-        with self.acquire_timeout(timeout=actual_lock_timeout):
-            return self._connection.put(*args, **kwargs)
+        with _acquire_lock_with_timeout(
+            lock=self._lock,
+            label=label,
+            timeout=actual_lock_timeout,
+        ):
+            return self._sftp_unsafe().put(local, remote)
 
-    def get(
-        self, *args, lock_timeout: Optional[float] = None, **kwargs
+    def _get(
+        self,
+        *,
+        local: str,
+        remote: str,
+        label: str,
+        lock_timeout: Optional[float] = None,
     ) -> Result:
         actual_lock_timeout = self.default_lock_timeout
         if lock_timeout is not None:
             actual_lock_timeout = lock_timeout
-        with self.acquire_timeout(timeout=actual_lock_timeout):
-            return self._connection.get(*args, **kwargs)
+        with _acquire_lock_with_timeout(
+            lock=self._lock,
+            label=label,
+            timeout=actual_lock_timeout,
+        ):
+            return self._sftp_unsafe().get(
+                remote,
+                local,
+                prefetch=self.sftp_get_prefetch,
+                max_concurrent_prefetch_requests=self.sftp_get_max_requests,
+            )
 
-    def run(
-        self, *args, lock_timeout: Optional[float] = None, **kwargs
+    def _run(
+        self, *args, label: str, 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
-        with self.acquire_timeout(timeout=actual_lock_timeout):
+        with _acquire_lock_with_timeout(
+            lock=self._lock,
+            label=label,
+            timeout=actual_lock_timeout,
+        ):
             return self._connection.run(*args, **kwargs)
 
-    def sftp(self) -> paramiko.sftp_client.SFTPClient:
+    def _sftp_unsafe(self) -> paramiko.sftp_client.SFTPClient:
+        """
+        This is marked as unsafe because you should only use its methods
+        after acquiring a lock.
+        """
         return self._connection.sftp()
 
+    def read_remote_json_file(self, filepath: str) -> dict[str, Any]:
+        self.logger.info(f"START reading remote JSON file {filepath}.")
+        with _acquire_lock_with_timeout(
+            lock=self._lock,
+            label="read_remote_json_file",
+            timeout=self.default_lock_timeout,
+        ):
+            with self._sftp_unsafe().open(filepath, "r") as f:
+                data = json.load(f)
+        self.logger.info(f"END reading remote JSON file {filepath}.")
+        return data
+
     def check_connection(self) -> None:
         """
         Open the SSH connection and handle exceptions.
@@ -131,7 +202,12 @@ class FractalSSH(object):
         """
         if not self._connection.is_connected:
             try:
-                self._connection.open()
+                with _acquire_lock_with_timeout(
+                    lock=self._lock,
+                    label="_connection.open",
+                    timeout=self.default_lock_timeout,
+                ):
+                    self._connection.open()
             except Exception as e:
                 raise RuntimeError(
                     f"Cannot open SSH connection. Original error:\n{str(e)}"
@@ -146,8 +222,12 @@ class FractalSSH(object):
         because we observed cases where `is_connected=False` but the underlying
         `Transport` object was not closed.
         """
-
-        self._connection.close()
+        with _acquire_lock_with_timeout(
+            lock=self._lock,
+            label="_connection.close",
+            timeout=self.default_lock_timeout,
+        ):
+            self._connection.close()
 
         if self._connection.client is not None:
             self._connection.client.close()
@@ -197,8 +277,11 @@ class FractalSSH(object):
             self.logger.info(f"{prefix} START running '{cmd}' over SSH.")
             try:
                 # Case 1: Command runs successfully
-                res = self.run(
-                    cmd, lock_timeout=actual_lock_timeout, hide=True
+                res = self._run(
+                    cmd,
+                    label=f"run {cmd}",
+                    lock_timeout=actual_lock_timeout,
+                    hide=True,
                 )
                 t_1 = time.perf_counter()
                 self.logger.info(
@@ -250,7 +333,6 @@ class FractalSSH(object):
         *,
         local: str,
         remote: str,
-        logger_name: Optional[str] = None,
         lock_timeout: Optional[float] = None,
     ) -> None:
         """
@@ -261,28 +343,65 @@ class FractalSSH(object):
             remote: Target path on remote host
             fractal_ssh: FractalSSH connection object with custom lock
             logger_name: Name of the logger
-
         """
         try:
-            self.put(local=local, remote=remote, lock_timeout=lock_timeout)
+            prefix = "[send_file]"
+            self.logger.info(f"{prefix} START transfer of '{local}' over SSH.")
+            self._put(
+                local=local,
+                remote=remote,
+                lock_timeout=lock_timeout,
+                label=f"send_file {local=} {remote=}",
+            )
+            self.logger.info(f"{prefix} END transfer of '{local}' over SSH.")
         except Exception as e:
-            logger = get_logger(logger_name=logger_name)
-            logger.error(
+            self.logger.error(
                 f"Transferring {local=} to {remote=} over SSH failed.\n"
                 f"Original Error:\n{str(e)}."
             )
             raise e
 
+    def fetch_file(
+        self,
+        *,
+        local: str,
+        remote: str,
+        lock_timeout: Optional[float] = None,
+    ) -> None:
+        """
+        Transfer a file via SSH
+
+        Args:
+            local: Local path to file
+            remote: Target path on remote host
+            logger_name: Name of the logger
+            lock_timeout:
+        """
+        try:
+            prefix = "[fetch_file] "
+            self.logger.info(f"{prefix} START fetching '{remote}' over SSH.")
+            self._get(
+                local=local,
+                remote=remote,
+                lock_timeout=lock_timeout,
+                label=f"fetch_file {local=} {remote=}",
+            )
+            self.logger.info(f"{prefix} END fetching '{remote}' over SSH.")
+        except Exception as e:
+            self.logger.error(
+                f"Transferring {remote=} to {local=} over SSH failed.\n"
+                f"Original Error:\n{str(e)}."
+            )
+            raise e
+
     def mkdir(self, *, folder: str, 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:
@@ -339,12 +458,18 @@ class FractalSSH(object):
             contents: File contents
             lock_timeout:
         """
+        self.logger.info(f"START writing to remote file {path}.")
         actual_lock_timeout = self.default_lock_timeout
         if lock_timeout is not None:
             actual_lock_timeout = lock_timeout
-        with self.acquire_timeout(timeout=actual_lock_timeout):
-            with self.sftp().open(filename=path, mode="w") as f:
+        with _acquire_lock_with_timeout(
+            lock=self._lock,
+            label=f"write_remote_file {path=}",
+            timeout=actual_lock_timeout,
+        ):
+            with self._sftp_unsafe().open(filename=path, mode="w") as f:
                 f.write(content)
+        self.logger.info(f"END writing to remote file {path}.")
 
 
 class FractalSSHList(object):
@@ -425,7 +550,11 @@ class FractalSSHList(object):
                     "look_for_keys": False,
                 },
             )
-            with self.acquire_lock_with_timeout():
+            with _acquire_lock_with_timeout(
+                lock=self._lock,
+                label="FractalSSHList.get",
+                timeout=self._timeout,
+            ):
                 self._data[key] = FractalSSH(connection=connection)
                 return self._data[key]
 
@@ -465,7 +594,11 @@ class FractalSSHList(object):
             key_path:
         """
         key = (host, user, key_path)
-        with self.acquire_lock_with_timeout():
+        with _acquire_lock_with_timeout(
+            lock=self._lock,
+            timeout=self._timeout,
+            label="FractalSSHList.remove",
+        ):
             self.logger.info(
                 f"Removing FractalSSH object for {user}@{host} "
                 "from collection."
@@ -492,24 +625,4 @@ class FractalSSHList(object):
                 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_ssh_obj.close()

fractal_server/string_tools.py

@@ -33,14 +33,18 @@ def sanitize_string(value: str) -> str:
     return new_value
 
 
-def slugify_task_name_for_source(task_name: str) -> str:
+def slugify_task_name_for_source_v1(task_name: str) -> str:
     """
     NOTE: this function is used upon creation of tasks' sources, therefore
     for the moment we cannot replace it with its more comprehensive version
     from `fractal_server.string_tools.sanitize_string`, nor we can remove it.
-    As 2.3.1, we are renaming it to `slugify_task_name_for_source`, to make
+
+    As of 2.3.1, we are renaming it to `slugify_task_name_for_source`, to make
     it clear that it should not be used for other purposes.
 
+    As of 2.7.0, we are renaming it to `slugify_task_name_for_source_v1`, to
+    make it clear that it is not used for v2.
+
     Args:
         task_name:
 

fractal_server/tasks/utils.py

@@ -9,9 +9,11 @@ COLLECTION_LOG_FILENAME = "collection.log"
 COLLECTION_FREEZE_FILENAME = "collection_freeze.txt"
 
 
-def get_absolute_venv_path(venv_path: Path) -> Path:
+def get_absolute_venv_path_v1(venv_path: Path) -> Path:
     """
     If a path is not absolute, make it a relative path of FRACTAL_TASKS_DIR.
+
+    As of v2.7.0, we rename this to v1 since it is only to be used in v1.
     """
     if venv_path.is_absolute():
         package_path = venv_path
@@ -33,20 +35,32 @@ 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)
+def get_collection_log_v1(path: Path) -> str:
+    package_path = get_absolute_venv_path_v1(path)
     log_path = get_log_path(package_path)
     log = log_path.open().read()
     return log
 
 
-def get_collection_freeze(venv_path: Path) -> str:
-    package_path = get_absolute_venv_path(venv_path)
+def get_collection_log_v2(path: Path) -> str:
+    log_path = get_log_path(path)
+    log = log_path.open().read()
+    return log
+
+
+def get_collection_freeze_v1(venv_path: Path) -> str:
+    package_path = get_absolute_venv_path_v1(venv_path)
     freeze_path = get_freeze_path(package_path)
     freeze = freeze_path.open().read()
     return freeze
 
 
+def get_collection_freeze_v2(path: Path) -> str:
+    freeze_path = get_freeze_path(path)
+    freeze = freeze_path.open().read()
+    return freeze
+
+
 def _normalize_package_name(name: str) -> str:
     """
     Implement PyPa specifications for package-name normalization

fractal_server/tasks/v1/_TaskCollectPip.py

@@ -77,7 +77,7 @@ class _TaskCollectPip(TaskCollectPipV1):
         if self.python_version:
             python_version = f"py{self.python_version}"
         else:
-            python_version = ""  # FIXME: can we allow this?
+            python_version = ""
 
         source = ":".join(
             (

fractal_server/tasks/v1/background_operations.py

@@ -6,9 +6,9 @@ import json
 from pathlib import Path
 from shutil import rmtree as shell_rmtree
 
-from ...string_tools import slugify_task_name_for_source
+from ...string_tools import slugify_task_name_for_source_v1
 from ..utils import _normalize_package_name
-from ..utils import get_collection_log
+from ..utils import get_collection_log_v1
 from ..utils import get_collection_path
 from ..utils import get_log_path
 from ._TaskCollectPip import _TaskCollectPip
@@ -215,7 +215,7 @@ async def create_package_environment_pip(
             # Fill in attributes for TaskCreate
             task_executable = package_root / t.executable
             cmd = f"{python_bin.as_posix()} {task_executable.as_posix()}"
-            task_name_slug = slugify_task_name_for_source(t.name)
+            task_name_slug = slugify_task_name_for_source_v1(t.name)
             task_source = f"{task_pkg.package_source}:{task_name_slug}"
             if not task_executable.exists():
                 raise FileNotFoundError(
@@ -321,7 +321,7 @@ async def background_collect_pip(
 
             # Update DB
             data.status = "OK"
-            data.log = get_collection_log(venv_path)
+            data.log = get_collection_log_v1(venv_path)
             state.data = data.sanitised_dict()
             db.add(state)
             db.merge(state)
@@ -342,7 +342,7 @@ async def background_collect_pip(
             # Update db
             data.status = "fail"
             data.info = f"Original error: {e}"
-            data.log = get_collection_log(venv_path)
+            data.log = get_collection_log_v1(venv_path)
             state.data = data.sanitised_dict()
             db.merge(state)
             db.commit()

fractal_server/tasks/v1/get_collection_data.py

@@ -2,12 +2,12 @@ import json
 from pathlib import Path
 
 from fractal_server.app.schemas.v1 import TaskCollectStatusV1
-from fractal_server.tasks.utils import get_absolute_venv_path
+from fractal_server.tasks.utils import get_absolute_venv_path_v1
 from fractal_server.tasks.utils import get_collection_path
 
 
 def get_collection_data(venv_path: Path) -> TaskCollectStatusV1:
-    package_path = get_absolute_venv_path(venv_path)
+    package_path = get_absolute_venv_path_v1(venv_path)
     collection_path = get_collection_path(package_path)
     with collection_path.open() as f:
         data = json.load(f)