fractal-server 2.16.2a0__py3-none-any.whl → 2.16.4__py3-none-any.whl

fractal_server/{app/runner → runner}/v2/submit_workflow.py

@@ -11,26 +11,26 @@ from pathlib import Path
 
 from sqlalchemy.orm import Session as DBSyncSession
 
-from ....config import get_settings
-from ....logger import get_logger
-from ....logger import reset_logger_handlers
-from ....logger import set_logger
-from ....ssh._fabric import FractalSSH
-from ....syringe import Inject
-from ....utils import get_timestamp
-from ....zip_tools import _zip_folder_to_file_and_remove
-from ...db import DB
-from ...models.v2 import DatasetV2
-from ...models.v2 import JobV2
-from ...models.v2 import WorkflowV2
-from ...schemas.v2 import JobStatusTypeV2
+from ...config import get_settings
+from ...logger import get_logger
+from ...logger import reset_logger_handlers
+from ...logger import set_logger
+from ...ssh._fabric import FractalSSH
+from ...syringe import Inject
+from ...utils import get_timestamp
+from ...zip_tools import _zip_folder_to_file_and_remove
 from ..exceptions import JobExecutionError
 from ..filenames import WORKFLOW_LOG_FILENAME
 from ._local import process_workflow as local_process_workflow
 from ._slurm_ssh import process_workflow as slurm_ssh_process_workflow
 from ._slurm_sudo import process_workflow as slurm_sudo_process_workflow
 from fractal_server import __VERSION__
+from fractal_server.app.db import DB
 from fractal_server.app.models import UserSettings
+from fractal_server.app.models.v2 import DatasetV2
+from fractal_server.app.models.v2 import JobV2
+from fractal_server.app.models.v2 import WorkflowV2
+from fractal_server.app.schemas.v2 import JobStatusTypeV2
 
 
 _backends = {}

fractal_server/{app/runner → runner}/v2/task_interface.py

@@ -5,8 +5,8 @@ from pydantic import ConfigDict
 from pydantic import Field
 from pydantic import ValidationError
 
-from ....images import SingleImageTaskOutput
-from fractal_server.app.runner.exceptions import TaskOutputValidationError
+from fractal_server.images import SingleImageTaskOutput
+from fractal_server.runner.exceptions import TaskOutputValidationError
 from fractal_server.types import ZarrUrlStr
 
 

fractal_server/ssh/_fabric.py

@@ -1,5 +1,6 @@
 import json
 import logging
+import os
 import time
 from collections.abc import Generator
 from contextlib import contextmanager
@@ -15,11 +16,15 @@ from invoke import UnexpectedExit
 from paramiko.ssh_exception import NoValidConnectionsError
 from pydantic import BaseModel
 
+from ..logger import close_logger
 from ..logger import get_logger
 from ..logger import set_logger
 from fractal_server.string_tools import validate_cmd
 
 
+SSH_MONITORING_LOGGER_NAME = "ssh-log"
+
+
 class FractalSSHTimeoutError(RuntimeError):
     pass
 
@@ -42,9 +47,6 @@ class SSHConfig(BaseModel):
     key_path: str
 
 
-logger = set_logger(__name__)
-
-
 def retry_if_socket_error(func):
     @wraps(func)
     def func_with_retry(*args, **kwargs):
@@ -76,7 +78,8 @@ def _acquire_lock_with_timeout(
     lock: Lock,
     label: str,
     timeout: float,
-    logger_name: str = __name__,
+    pid: int,
+    logger_name: str,
 ) -> Generator[Literal[True], Any, None]:
     """
     Given a `threading.Lock` object, try to acquire it within a given timeout.
@@ -88,8 +91,9 @@ def _acquire_lock_with_timeout(
         logger_name:
     """
     logger = get_logger(logger_name)
+    ssh_logger = get_logger(SSH_MONITORING_LOGGER_NAME)
     logger.info(f"Trying to acquire lock for '{label}', with {timeout=}")
-    t_start_lock_acquire = time.perf_counter()
+    t_lock_request = time.perf_counter()
     result = lock.acquire(timeout=timeout)
     try:
         if not result:
@@ -98,14 +102,25 @@ def _acquire_lock_with_timeout(
                 f"Failed to acquire lock for '{label}' within "
                 f"{timeout} seconds"
             )
-        t_end_lock_acquire = time.perf_counter()
-        elapsed = t_end_lock_acquire - t_start_lock_acquire
+        t_lock_acquisition = time.perf_counter()
+        elapsed = t_lock_acquisition - t_lock_request
         logger.info(f"Lock for '{label}' was acquired - {elapsed=:.4f} s")
         yield result
     finally:
         if result:
             lock.release()
             logger.info(f"Lock for '{label}' was released.")
+            t_lock_release = time.perf_counter()
+            lock_was_acquired = 1
+        else:
+            t_lock_release = time.perf_counter()
+            t_lock_acquisition = t_lock_release
+            lock_was_acquired = 0
+        lock_waiting_time = t_lock_acquisition - t_lock_request
+        lock_holding_time = t_lock_release - t_lock_acquisition
+        ssh_logger.info(
+            f"{pid} {lock_waiting_time:.6e} {lock_holding_time:.6e} {lock_was_acquired} {label.replace(' ', '_')}"  # noqa
+        )
 
 
 class FractalSSH:
@@ -130,11 +145,12 @@ class FractalSSH:
     sftp_get_prefetch: bool
     sftp_get_max_requests: int
     logger_name: str
+    _pid: int
 
     def __init__(
         self,
         connection: Connection,
-        default_timeout: float = 250,
+        default_timeout: float = 500.0,
         sftp_get_prefetch: bool = False,
         sftp_get_max_requests: int = 64,
         logger_name: str = __name__,
@@ -146,6 +162,8 @@ class FractalSSH:
         self.sftp_get_max_requests = sftp_get_max_requests
         self.logger_name = logger_name
         set_logger(self.logger_name)
+        set_logger(SSH_MONITORING_LOGGER_NAME)
+        self._pid = os.getpid()
 
     @property
     def is_connected(self) -> bool:
@@ -197,6 +215,8 @@ class FractalSSH:
             lock=self._lock,
             label=label,
             timeout=actual_lock_timeout,
+            pid=self._pid,
+            logger_name=self.logger_name,
         ):
             return self._connection.run(*args, **kwargs)
 
@@ -212,8 +232,10 @@ class FractalSSH:
         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,
+            logger_name=self.logger_name,
+            pid=self._pid,
+            label=f"read_remote_json_file({filepath})",
         ):
             try:
                 with self._sftp_unsafe().open(filepath, "r") as f:
@@ -239,8 +261,10 @@ class FractalSSH:
         self.logger.info(f"START reading remote text file {filepath}.")
         with _acquire_lock_with_timeout(
             lock=self._lock,
-            label="read_remote_text_file",
+            label=f"read_remote_text_file({filepath})",
             timeout=self.default_lock_timeout,
+            pid=self._pid,
+            logger_name=self.logger_name,
         ):
             try:
                 with self._sftp_unsafe().open(filepath, "r") as f:
@@ -298,9 +322,10 @@ class FractalSSH:
             self.close()
             with _acquire_lock_with_timeout(
                 lock=self._lock,
-                label="_connection.open",
+                label="FractalSSH._connection.{open,open_sftp}()",
                 timeout=self.default_lock_timeout,
                 logger_name=self.logger_name,
+                pid=self._pid,
             ):
                 self._connection.open()
                 self._connection.client.open_sftp()
@@ -324,12 +349,16 @@ class FractalSSH:
         """
         with _acquire_lock_with_timeout(
             lock=self._lock,
-            label="_connection.close",
+            label="FractalSSH._connection.close()",
             timeout=self.default_lock_timeout,
+            pid=self._pid,
+            logger_name=self.logger_name,
         ):
             self._connection.close()
             if self._connection.client is not None:
                 self._connection.client.close()
+        close_logger(get_logger(self.logger_name))
+        close_logger(get_logger(SSH_MONITORING_LOGGER_NAME))
 
     @retry_if_socket_error
     def run_command(
@@ -364,14 +393,14 @@ class FractalSSH:
             # Case 1: Command runs successfully
             res = self._run(
                 cmd,
-                label=f"run {cmd}",
+                label=cmd,
                 lock_timeout=actual_lock_timeout,
                 hide=True,
                 in_stream=False,
             )
             t_1 = time.perf_counter()
             self.logger.info(
-                f"END   running '{cmd}' over SSH, " f"elapsed={t_1 - t_0:.3f}"
+                f"END   running '{cmd}' over SSH, elapsed={t_1 - t_0:.3f}"
             )
             self.logger.debug("STDOUT:")
             self.logger.debug(res.stdout)
@@ -423,8 +452,10 @@ class FractalSSH:
                 actual_lock_timeout = lock_timeout
             with _acquire_lock_with_timeout(
                 lock=self._lock,
-                label=f"send_file {local=} {remote=}",
+                label=f"send_file({local},{remote})",
                 timeout=actual_lock_timeout,
+                pid=self._pid,
+                logger_name=self.logger_name,
             ):
                 self._sftp_unsafe().put(local, remote)
             self.logger.info(
@@ -463,8 +494,10 @@ class FractalSSH:
                 actual_lock_timeout = lock_timeout
             with _acquire_lock_with_timeout(
                 lock=self._lock,
-                label=f"fetch_file {local=} {remote=}",
+                label=f"fetch_file({local},{remote})",
                 timeout=actual_lock_timeout,
+                pid=self._pid,
+                logger_name=self.logger_name,
             ):
                 self._sftp_unsafe().get(
                     remote,
@@ -554,8 +587,10 @@ class FractalSSH:
             actual_lock_timeout = lock_timeout
         with _acquire_lock_with_timeout(
             lock=self._lock,
-            label=f"write_remote_file {path=}",
+            label=f"write_remote_file({path})",
             timeout=actual_lock_timeout,
+            pid=self._pid,
+            logger_name=self.logger_name,
         ):
             try:
                 with self._sftp_unsafe().open(filename=path, mode="w") as f:
@@ -576,8 +611,10 @@ class FractalSSH:
         self.logger.info(f"START remote_file_exists {path}")
         with _acquire_lock_with_timeout(
             lock=self._lock,
-            label=f"remote_file_exists {path=}",
+            label=f"remote_file_exists({path})",
             timeout=self.default_lock_timeout,
+            pid=self._pid,
+            logger_name=self.logger_name,
         ):
             try:
                 self._sftp_unsafe().stat(path)
@@ -613,6 +650,7 @@ class FractalSSHList:
     _lock: Lock
     _timeout: float
     _logger_name: str
+    _pid: int
 
     def __init__(
         self,
@@ -625,6 +663,7 @@ class FractalSSHList:
         self._timeout = timeout
         self._logger_name = logger_name
         set_logger(self._logger_name)
+        self._pid = os.getpid()
 
     @property
     def logger(self) -> logging.Logger:
@@ -677,6 +716,8 @@ class FractalSSHList:
                 lock=self._lock,
                 label="FractalSSHList.get",
                 timeout=self._timeout,
+                pid=self._pid,
+                logger_name=self._logger_name,
             ):
                 self._data[key] = FractalSSH(connection=connection)
                 return self._data[key]
@@ -721,6 +762,8 @@ class FractalSSHList:
             lock=self._lock,
             timeout=self._timeout,
             label="FractalSSHList.remove",
+            pid=self._pid,
+            logger_name=self._logger_name,
         ):
             self.logger.info(
                 f"Removing FractalSSH object for {user}@{host} "

fractal_server/tasks/v2/ssh/_pixi_slurm_ssh.py

@@ -2,12 +2,18 @@ import os
 import time
 from pathlib import Path
 
+from sqlalchemy.orm import Session
+
+from fractal_server.app.models.v2 import TaskGroupActivityV2
 from fractal_server.config import get_settings
 from fractal_server.config import PixiSLURMConfig
 from fractal_server.logger import get_logger
 from fractal_server.ssh._fabric import FractalSSH
 from fractal_server.syringe import Inject
+from fractal_server.tasks.v2.utils_background import add_commit_refresh
+from fractal_server.tasks.v2.utils_background import get_current_log
 
+FRACTAL_SQUEUE_ERROR_STATE = "__FRACTAL_SQUEUE_ERROR__"
 
 # https://slurm.schedmd.com/squeue.html#lbAG
 STATES_FINISHED = {
@@ -21,20 +27,130 @@ STATES_FINISHED = {
     "PREEMPTED",
     "SPECIAL_EXIT",
     "TIMEOUT",
+    FRACTAL_SQUEUE_ERROR_STATE,
 }
 
 
+def _get_workdir_remote(script_paths: list[str]) -> str:
+    """
+    Check that there is one and only one `workdir`, and return it.
+
+    Note: The `is_absolute` check is to filter out a `chmod` command.
+    """
+    workdirs = [
+        Path(script_path).parent.as_posix()
+        for script_path in script_paths
+        if Path(script_path).is_absolute()
+    ]
+    if not len(set(workdirs)) == 1:
+        raise ValueError(f"Invalid {script_paths=}.")
+    return workdirs[0]
+
+
+def _read_file_if_exists(
+    *,
+    fractal_ssh: FractalSSH,
+    path: str,
+) -> str:
+    """
+    Read a remote file if it exists, or return an empty string.
+    """
+    if fractal_ssh.remote_exists(path=path):
+        return fractal_ssh.read_remote_text_file(path)
+    else:
+        return ""
+
+
+def _log_change_of_job_state(
+    *,
+    old_state: str | None,
+    new_state: str,
+    logger_name: str,
+) -> None:
+    """
+    Emit a log for state changes.
+
+    Args:
+        old_state:
+        new_state:
+        logger_name:
+    """
+    if new_state != old_state:
+        logger = get_logger(logger_name=logger_name)
+        logger.debug(
+            f"SLURM-job state changed from {old_state=} to {new_state=}."
+        )
+
+
+def _run_squeue(
+    *,
+    fractal_ssh: FractalSSH,
+    squeue_cmd: str,
+    logger_name: str,
+) -> str:
+    """
+    Run a `squeue` command and handle exceptions.
+
+    Args:
+        fractal_ssh:
+        logger_name:
+        squeue_cmd:
+
+    Return:
+        state: The SLURM-job state.
+    """
+    try:
+        cmd_stdout = fractal_ssh.run_command(cmd=squeue_cmd)
+        state = cmd_stdout.strip().split()[1]
+        return state
+    except Exception as e:
+        logger = get_logger(logger_name=logger_name)
+        logger.info(f"`squeue` command failed (original error: {e})")
+        return FRACTAL_SQUEUE_ERROR_STATE
+
+
+def _verify_success_file_exists(
+    *,
+    fractal_ssh: FractalSSH,
+    success_file_remote: str,
+    logger_name: str,
+    stderr_remote: str,
+) -> None:
+    """
+    Fail if the success sentinel file does not exist remotely.
+
+    Note: the `FractalSSH` methods in this function may fail, and such failures
+    are not handled in this function. Any such failure, however, will lead to
+    a "failed" task-group lifecycle activity (because it will raise an
+    exception from within `run_script_on_remote_slurm`, which will then be
+    handled at the calling-function level.
+    """
+    if not fractal_ssh.remote_exists(path=success_file_remote):
+        logger = get_logger(logger_name=logger_name)
+        error_msg = f"{success_file_remote=} missing."
+        logger.info(error_msg)
+
+        stderr = _read_file_if_exists(
+            fractal_ssh=fractal_ssh, path=stderr_remote
+        )
+        if stderr:
+            logger.info(f"SLURM-job stderr:\n{stderr}")
+        raise RuntimeError(error_msg)
+
+
 def run_script_on_remote_slurm(
     *,
-    script_path: str,
+    script_paths: list[str],
     slurm_config: PixiSLURMConfig,
     fractal_ssh: FractalSSH,
     logger_name: str,
+    log_file_path: Path,
     prefix: str,
+    db: Session,
+    activity: TaskGroupActivityV2,
 ):
     """
-    FIXME
-
+    Run a `pixi install` script as a SLURM job.
 
     NOTE: This is called from within a try/except, thus we can use exceptions
     as a mechanism to propagate failure/errors.
@@ -44,7 +160,7 @@ def run_script_on_remote_slurm(
     settings = Inject(get_settings)
 
     # (1) Prepare remote submission script
-    workdir_remote = Path(script_path).parent.as_posix()
+    workdir_remote = _get_workdir_remote(script_paths)
     submission_script_remote = os.path.join(
         workdir_remote, f"{prefix}-submit.sh"
     )
@@ -61,55 +177,79 @@ def run_script_on_remote_slurm(
         f"#SBATCH --out={stdout_remote}",
         f"#SBATCH -D {workdir_remote}",
         "",
-        f"bash {script_path}",
-        f"touch {success_file_remote}",
-        "",
     ]
+    for script_path in script_paths:
+        script_lines.append(f"bash {script_path}")
+    script_lines.append(f"touch {success_file_remote}")
+
     script_contents = "\n".join(script_lines)
     fractal_ssh.write_remote_file(
         path=submission_script_remote,
         content=script_contents,
     )
+    logger.debug(f"Written {submission_script_remote=}.")
+
+    activity.log = get_current_log(log_file_path)
+    activity = add_commit_refresh(obj=activity, db=db)
 
     # (2) Submit SLURM job
-    sbatch_cmd = f"sbatch --parsable {submission_script_remote} "
+    logger.debug("Now submit SLURM job.")
+    sbatch_cmd = f"sbatch --parsable {submission_script_remote}"
     try:
         stdout = fractal_ssh.run_command(cmd=sbatch_cmd)
+        job_id = int(stdout)
+        logger.debug(f"SLURM-job submission successful ({job_id=}).")
     except Exception as e:
         logger.error(
-            f"Submission of {submission_script_remote} failed. "
-            f"Original error: {str(e)}"
+            (
+                f"Submission of {submission_script_remote} failed. "
+                f"Original error: {str(e)}"
+            )
         )
         raise e
-    logger.debug(f"Now submit job {submission_script_remote} to SLURM.")
-    job_id = int(stdout)
-    logger.debug(f"SLURM-job submission successful ({job_id=}).")
+    finally:
+        activity.log = get_current_log(log_file_path)
+        activity = add_commit_refresh(obj=activity, db=db)
 
     # (3) Monitor job
     squeue_cmd = (
         f"squeue --noheader --format='%i %T' --states=all --jobs={job_id}"
     )
+    logger.debug(f"Start monitoring job with {squeue_cmd=}.")
+    old_state = None
     while True:
-        try:
-            stdout = fractal_ssh.run_command(cmd=squeue_cmd)
-        except Exception as e:
-            # FIXME: review this logic
-            logger.info(
-                f"`squeue` command failed (original error: {e}), "
-                "consider the job as complete."
-            )
-            break
-        state = stdout.strip().split()[1]
-        logger.debug(f"Status of SLURM job {job_id}: {state}")
-        if state in STATES_FINISHED:
-            logger.debug(f"Exit retrieval loop ({state=}).")
+        new_state = _run_squeue(
+            fractal_ssh=fractal_ssh,
+            squeue_cmd=squeue_cmd,
+            logger_name=logger_name,
+        )
+        _log_change_of_job_state(
+            old_state=old_state,
+            new_state=new_state,
+            logger_name=logger_name,
+        )
+        activity.log = get_current_log(log_file_path)
+        activity = add_commit_refresh(obj=activity, db=db)
+        if new_state in STATES_FINISHED:
+            logger.debug(f"Exit retrieval loop (state={new_state}).")
             break
+        old_state = new_state
         time.sleep(settings.FRACTAL_SLURM_POLL_INTERVAL)
 
-    if fractal_ssh.remote_exists(path=success_file_remote):
-        logger.info(f"{success_file_remote=} exists.")
-    else:
-        raise RuntimeError(
-            "SLURM job did not complete correctly "
-            f"({success_file_remote=} missing)."
-        )
+    _verify_success_file_exists(
+        fractal_ssh=fractal_ssh,
+        logger_name=logger_name,
+        success_file_remote=success_file_remote,
+        stderr_remote=stderr_remote,
+    )
+
+    stdout = _read_file_if_exists(
+        fractal_ssh=fractal_ssh,
+        path=stdout_remote,
+    )
+
+    logger.info("SLURM-job execution completed successfully, continue.")
+    activity.log = get_current_log(log_file_path)
+    activity = add_commit_refresh(obj=activity, db=db)
+
+    return stdout

fractal_server/tasks/v2/ssh/collect_pixi.py

@@ -203,12 +203,21 @@ def collect_ssh_pixi(
                         pyproject_toml_path=pyproject_toml_path,
                     )
 
-                    # Run script 2 - run pixi-install command
+                    # Prepare scripts 2 and 3
                     remote_script2_path = _customize_and_send_template(
                         template_filename="pixi_2_install.sh",
                         replacements=replacements,
                         **common_args,
                     )
+                    remote_script3_path = _customize_and_send_template(
+                        template_filename="pixi_3_post_install.sh",
+                        replacements=replacements,
+                        **common_args,
+                    )
+                    logger.debug(
+                        "Post-installation script written to "
+                        f"{remote_script3_path=}."
+                    )
                     logger.debug(
                         "Installation script written to "
                         f"{remote_script2_path=}."
@@ -216,26 +225,24 @@ def collect_ssh_pixi(
                     activity.log = get_current_log(log_file_path)
                     activity = add_commit_refresh(obj=activity, db=db)
 
-                    run_script_on_remote_slurm(
-                        script_path=remote_script2_path,
+                    # Run scripts 2 and 3
+                    stdout = run_script_on_remote_slurm(
+                        script_paths=[
+                            remote_script2_path,
+                            remote_script3_path,
+                            f"chmod -R 755 {source_dir}",
+                        ],
                         slurm_config=settings.pixi.SLURM_CONFIG,
                         fractal_ssh=fractal_ssh,
                         logger_name=LOGGER_NAME,
                         prefix=common_args["prefix"],
+                        db=db,
+                        activity=activity,
+                        log_file_path=log_file_path,
                     )
                     activity.log = get_current_log(log_file_path)
                     activity = add_commit_refresh(obj=activity, db=db)
 
-                    # Run script 3 - post-install
-                    stdout = _customize_and_run_template(
-                        template_filename="pixi_3_post_install.sh",
-                        replacements=replacements,
-                        **common_args,
-                    )
-                    logger.debug(f"STDOUT: {stdout}")
-                    activity.log = get_current_log(log_file_path)
-                    activity = add_commit_refresh(obj=activity, db=db)
-
                     # Parse stdout
                     parsed_output = parse_collect_stdout(stdout)
                     package_root_remote = parsed_output["package_root"]
@@ -245,8 +252,6 @@ def collect_ssh_pixi(
                         "project_python_wrapper"
                     ]
 
-                    fractal_ssh.run_command(cmd=f"chmod -R 755 {source_dir}")
-
                     # Read and validate remote manifest file
                     manifest_path_remote = (
                         f"{package_root_remote}/__FRACTAL_MANIFEST__.json"