fractal-server 2.13.0__py3-none-any.whl → 2.14.0__py3-none-any.whl

fractal_server/app/runner/{v2/_slurm_common → executors/slurm_common}/get_slurm_config.py

@@ -2,21 +2,16 @@ from pathlib import Path
 from typing import Literal
 from typing import Optional
 
+from ._batching import heuristics
+from ._slurm_config import _parse_mem_value
+from ._slurm_config import load_slurm_config_file
+from ._slurm_config import logger
+from ._slurm_config import SlurmConfig
+from ._slurm_config import SlurmConfigError
 from fractal_server.app.models.v2 import WorkflowTaskV2
-from fractal_server.app.runner.executors.slurm._slurm_config import (
-    _parse_mem_value,
-)
-from fractal_server.app.runner.executors.slurm._slurm_config import (
-    load_slurm_config_file,
-)
-from fractal_server.app.runner.executors.slurm._slurm_config import logger
-from fractal_server.app.runner.executors.slurm._slurm_config import SlurmConfig
-from fractal_server.app.runner.executors.slurm._slurm_config import (
-    SlurmConfigError,
-)
 
 
-def get_slurm_config(
+def get_slurm_config_internal(
     wftask: WorkflowTaskV2,
     which_type: Literal["non_parallel", "parallel"],
     config_path: Optional[Path] = None,
@@ -25,7 +20,7 @@ def get_slurm_config(
     Prepare a `SlurmConfig` configuration object
 
     The argument `which_type` determines whether we use `wftask.meta_parallel`
-    or `wftask.meta_non_parallel`. In the following descritpion, let us assume
+    or `wftask.meta_non_parallel`. In the following description, let us assume
     that `which_type="parallel"`.
 
     The sources for `SlurmConfig` attributes, in increasing priority order, are
@@ -142,8 +137,8 @@ def get_slurm_config(
     extra_lines = slurm_dict.get("extra_lines", []) + extra_lines
     if len(set(extra_lines)) != len(extra_lines):
         logger.debug(
-            "[get_slurm_config] Removing repeated elements "
-            f"from {extra_lines=}."
+            "[get_slurm_config] Removing repeated elements from "
+            f"{extra_lines=}."
         )
         extra_lines = list(set(extra_lines))
     slurm_dict["extra_lines"] = extra_lines
@@ -162,9 +157,45 @@ def get_slurm_config(
 
     # Put everything together
     logger.debug(
-        "[get_slurm_config] Now create a SlurmConfig object based "
-        f"on {slurm_dict=}"
+        "[get_slurm_config] Now create a SlurmConfig object based on "
+        f"{slurm_dict=}"
     )
     slurm_config = SlurmConfig(**slurm_dict)
 
     return slurm_config
+
+
+def get_slurm_config(
+    wftask: WorkflowTaskV2,
+    which_type: Literal["non_parallel", "parallel"],
+    config_path: Optional[Path] = None,
+    tot_tasks: int = 1,
+) -> SlurmConfig:
+    config = get_slurm_config_internal(
+        wftask,
+        which_type,
+        config_path,
+    )
+
+    # Set/validate parameters for task batching
+    tasks_per_job, parallel_tasks_per_job = heuristics(
+        # Number of parallel components (always known)
+        tot_tasks=tot_tasks,
+        # Optional WorkflowTask attributes:
+        tasks_per_job=config.tasks_per_job,
+        parallel_tasks_per_job=config.parallel_tasks_per_job,  # noqa
+        # Task requirements (multiple possible sources):
+        cpus_per_task=config.cpus_per_task,
+        mem_per_task=config.mem_per_task_MB,
+        # Fractal configuration variables (soft/hard limits):
+        target_cpus_per_job=config.target_cpus_per_job,
+        target_mem_per_job=config.target_mem_per_job,
+        target_num_jobs=config.target_num_jobs,
+        max_cpus_per_job=config.max_cpus_per_job,
+        max_mem_per_job=config.max_mem_per_job,
+        max_num_jobs=config.max_num_jobs,
+    )
+    config.parallel_tasks_per_job = parallel_tasks_per_job
+    config.tasks_per_job = tasks_per_job
+
+    return config

fractal_server/app/runner/executors/{slurm → slurm_common}/remote.py

@@ -19,7 +19,6 @@ import os
 import sys
 from typing import Literal
 from typing import Optional
-from typing import Type
 from typing import Union
 
 import cloudpickle
@@ -27,30 +26,6 @@ import cloudpickle
 from fractal_server import __VERSION__
 
 
-class ExceptionProxy:
-    """
-    Proxy class to serialise exceptions
-
-    In general exceptions are not serialisable. This proxy class saves the
-    serialisable content of an exception. On the receiving end, it can be used
-    to reconstruct a TaskExecutionError.
-
-    Attributes:
-        exc_type_name: Name of the exception type
-        tb: TBD
-        args: TBD
-        kwargs: TBD
-    """
-
-    def __init__(
-        self, exc_type: Type[BaseException], tb: str, *args, **kwargs
-    ):
-        self.exc_type_name: str = exc_type.__name__
-        self.tb: str = tb
-        self.args = args
-        self.kwargs: dict = kwargs
-
-
 class FractalVersionMismatch(RuntimeError):
     """
     Custom exception for version mismatch
@@ -79,18 +54,21 @@ def _check_versions_mismatch(
                                 do not match with the ones on the server
     """
 
-    server_python_version = server_versions["python"]
-    worker_python_version = sys.version_info[:3]
+    server_python_version = list(server_versions["python"])
+    worker_python_version = list(sys.version_info[:3])
     if worker_python_version != server_python_version:
-        # FIXME: turn this into an error, after fixing a broader CI issue, see
-        # https://github.com/fractal-analytics-platform/fractal-server/issues/375
-        logging.critical(
-            f"{server_python_version=} but {worker_python_version=}. "
-            "cloudpickle is not guaranteed to correctly load "
-            "pickle files created with different python versions. "
-            "Note, however, that if you reached this line it means that "
-            "the pickle file was likely loaded correctly."
-        )
+        if worker_python_version[:2] != server_python_version[:2]:
+            # FIXME: Turn this into an error, in some version post 2.14.
+            logging.error(
+                f"{server_python_version=} but {worker_python_version=}. "
+                "This configuration will be deprecated in a future version, "
+                "please contact the admin of this Fractal instance."
+            )
+        else:
+            # Major.minor versions match, patch versions differ
+            logging.warning(
+                f"{server_python_version=} but {worker_python_version=}."
+            )
 
     server_cloudpickle_version = server_versions["cloudpickle"]
     worker_cloudpickle_version = cloudpickle.__version__
@@ -134,28 +112,39 @@ def worker(
         _extra_import_paths = extra_import_paths.split(":")
         sys.path[:0] = _extra_import_paths
 
-    # Execute the job and catpure exceptions
+    # Execute the job and capture exceptions
     try:
         with open(in_fname, "rb") as f:
             indata = f.read()
         server_versions, fun, args, kwargs = cloudpickle.loads(indata)
         _check_versions_mismatch(server_versions)
 
-        result = True, fun(*args, **kwargs)
+        result = (True, fun(*args, **kwargs))
         out = cloudpickle.dumps(result)
     except Exception as e:
+        # Exception objects are not serialisable. Here we save the relevant
+        # exception contents in a serializable dictionary. Note that whenever
+        # the task failed "properly", the exception is a `TaskExecutionError`
+        # and it has additional attributes.
+
         import traceback
 
-        typ, value, tb = sys.exc_info()
-        tb = tb.tb_next
-        exc_proxy = ExceptionProxy(
-            typ,
-            "".join(traceback.format_exception(typ, value, tb)),
-            *e.args,
-            **e.__dict__,
+        exc_type, exc_value, traceback_obj = sys.exc_info()
+        traceback_obj = traceback_obj.tb_next
+        traceback_list = traceback.format_exception(
+            exc_type,
+            exc_value,
+            traceback_obj,
         )
-
-        result = False, exc_proxy
+        traceback_string = "".join(traceback_list)
+        exc_proxy = dict(
+            exc_type_name=exc_type.__name__,
+            traceback_string=traceback_string,
+            workflow_task_order=getattr(e, "workflow_task_order", None),
+            workflow_task_id=getattr(e, "workflow_task_id", None),
+            task_name=getattr(e, "task_name", None),
+        )
+        result = (False, exc_proxy)
         out = cloudpickle.dumps(result)
 
     # Write the output pickle file

fractal_server/app/runner/executors/slurm_common/slurm_job_task_models.py

@@ -0,0 +1,134 @@
+from pathlib import Path
+from typing import Any
+from typing import Optional
+
+from pydantic import BaseModel
+from pydantic import ConfigDict
+
+from fractal_server.app.runner.task_files import TaskFiles
+
+
+class SlurmTask(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+    component: str
+    prefix: str
+    workdir_local: Path
+    workdir_remote: Path
+    parameters: dict[str, Any]
+    zarr_url: Optional[str] = None
+    task_files: TaskFiles
+    index: int
+
+    @property
+    def input_pickle_file_local_path(self) -> Path:
+        return (
+            self.workdir_local / f"{self.prefix}-{self.component}-input.pickle"
+        )
+
+    @property
+    def input_pickle_file_remote_path(self) -> Path:
+        return (
+            self.workdir_remote
+            / f"{self.prefix}-{self.component}-input.pickle"
+        )
+
+    @property
+    def output_pickle_file_local_path(self) -> Path:
+        return (
+            self.workdir_local
+            / f"{self.prefix}-{self.component}-output.pickle"
+        )
+
+    @property
+    def output_pickle_file_remote_path(self) -> Path:
+        return (
+            self.workdir_remote
+            / f"{self.prefix}-{self.component}-output.pickle"
+        )
+
+    @property
+    def input_pickle_file_local(self) -> str:
+        return self.input_pickle_file_local_path.as_posix()
+
+    @property
+    def input_pickle_file_remote(self) -> str:
+        return self.input_pickle_file_remote_path.as_posix()
+
+    @property
+    def output_pickle_file_local(self) -> str:
+        return self.output_pickle_file_local_path.as_posix()
+
+    @property
+    def output_pickle_file_remote(self) -> str:
+        return self.output_pickle_file_remote_path.as_posix()
+
+
+class SlurmJob(BaseModel):
+    slurm_job_id: Optional[str] = None
+    prefix: str
+    workdir_local: Path
+    workdir_remote: Path
+    tasks: list[SlurmTask]
+
+    @property
+    def slurm_submission_script_local(self) -> str:
+        return (
+            self.workdir_local / f"{self.prefix}-slurm-submit.sh"
+        ).as_posix()
+
+    @property
+    def slurm_submission_script_remote(self) -> str:
+        return (
+            self.workdir_remote / f"{self.prefix}-slurm-submit.sh"
+        ).as_posix()
+
+    @property
+    def slurm_job_id_placeholder(self) -> str:
+        if self.slurm_job_id:
+            return self.slurm_job_id
+        else:
+            return "%j"
+
+    @property
+    def slurm_stdout_remote_path(self) -> Path:
+        return (
+            self.workdir_remote
+            / f"{self.prefix}-slurm-{self.slurm_job_id_placeholder}.out"
+        )
+
+    @property
+    def slurm_stdout_remote(self) -> str:
+        return self.slurm_stdout_remote_path.as_posix()
+
+    @property
+    def slurm_stderr_remote_path(self) -> Path:
+        return (
+            self.workdir_remote
+            / f"{self.prefix}-slurm-{self.slurm_job_id_placeholder}.err"
+        )
+
+    @property
+    def slurm_stderr_remote(self) -> str:
+        return self.slurm_stderr_remote_path.as_posix()
+
+    @property
+    def slurm_stdout_local_path(self) -> str:
+        return (
+            self.workdir_local
+            / f"{self.prefix}-slurm-{self.slurm_job_id_placeholder}.out"
+        )
+
+    @property
+    def slurm_stdout_local(self) -> str:
+        return self.slurm_stdout_local_path.as_posix()
+
+    @property
+    def slurm_stderr_local_path(self) -> Path:
+        return (
+            self.workdir_local
+            / f"{self.prefix}-slurm-{self.slurm_job_id_placeholder}.err"
+        )
+
+    @property
+    def slurm_stderr_local(self) -> str:
+        return self.slurm_stderr_local_path.as_posix()

fractal_server/app/runner/executors/slurm_ssh/runner.py

@@ -0,0 +1,268 @@
+import time
+from pathlib import Path
+from typing import Optional
+
+from ..slurm_common.base_slurm_runner import BaseSlurmRunner
+from ..slurm_common.slurm_job_task_models import SlurmJob
+from fractal_server.app.runner.compress_folder import compress_folder
+from fractal_server.app.runner.extract_archive import extract_archive
+from fractal_server.config import get_settings
+from fractal_server.logger import set_logger
+from fractal_server.ssh._fabric import FractalSSH
+from fractal_server.ssh._fabric import FractalSSHCommandError
+from fractal_server.ssh._fabric import FractalSSHTimeoutError
+from fractal_server.syringe import Inject
+
+
+logger = set_logger(__name__)
+
+
+class SlurmSSHRunner(BaseSlurmRunner):
+    fractal_ssh: FractalSSH
+
+    def __init__(
+        self,
+        *,
+        # Common
+        root_dir_local: Path,
+        root_dir_remote: Path,
+        common_script_lines: Optional[list[str]] = None,
+        user_cache_dir: Optional[str] = None,
+        poll_interval: Optional[int] = None,
+        # Specific
+        fractal_ssh: FractalSSH,
+    ) -> None:
+        """
+        Set parameters that are the same for different Fractal tasks and for
+        different SLURM jobs/tasks.
+        """
+        self.fractal_ssh = fractal_ssh
+        logger.warning(self.fractal_ssh)
+
+        settings = Inject(get_settings)
+
+        super().__init__(
+            slurm_runner_type="ssh",
+            root_dir_local=root_dir_local,
+            root_dir_remote=root_dir_remote,
+            common_script_lines=common_script_lines,
+            user_cache_dir=user_cache_dir,
+            poll_interval=poll_interval,
+            python_worker_interpreter=settings.FRACTAL_SLURM_WORKER_PYTHON,
+        )
+
+    def _mkdir_local_folder(self, folder: str) -> None:
+        Path(folder).mkdir(parents=True)
+
+    def _mkdir_remote_folder(self, folder: str):
+        self.fractal_ssh.mkdir(
+            folder=folder,
+            parents=True,
+        )
+
+    def _fetch_artifacts(
+        self,
+        finished_slurm_jobs: list[SlurmJob],
+    ) -> None:
+        """
+        Fetch artifacts for a list of SLURM jobs.
+        """
+
+        # Check length
+        if len(finished_slurm_jobs) == 0:
+            logger.debug(f"[_fetch_artifacts] EXIT ({finished_slurm_jobs=}).")
+            return None
+
+        t_0 = time.perf_counter()
+        logger.debug(
+            f"[_fetch_artifacts] START ({len(finished_slurm_jobs)=})."
+        )
+
+        # Extract `workdir_remote` and `workdir_local`
+        self.validate_slurm_jobs_workdirs(finished_slurm_jobs)
+        workdir_local = finished_slurm_jobs[0].workdir_local
+        workdir_remote = finished_slurm_jobs[0].workdir_remote
+
+        # Define local/remote tarfile paths
+        tarfile_path_local = (
+            workdir_local.parent / f"{workdir_local.name}.tar.gz"
+        ).as_posix()
+        tarfile_path_remote = (
+            workdir_remote.parent / f"{workdir_remote.name}.tar.gz"
+        ).as_posix()
+
+        # Create file list
+        # NOTE: see issue 2483
+        filelist = []
+        for _slurm_job in finished_slurm_jobs:
+            _single_job_filelist = [
+                _slurm_job.slurm_stdout_remote_path.name,
+                _slurm_job.slurm_stderr_remote_path.name,
+            ]
+            for task in _slurm_job.tasks:
+                _single_job_filelist.extend(
+                    [
+                        task.output_pickle_file_remote_path.name,
+                        task.task_files.log_file_remote_path.name,
+                        task.task_files.args_file_remote_path.name,
+                        task.task_files.metadiff_file_remote_path.name,
+                    ]
+                )
+            filelist.extend(_single_job_filelist)
+        filelist_string = "\n".join(filelist)
+        elapsed = time.perf_counter() - t_0
+        logger.debug(
+            "[_fetch_artifacts] Created filelist "
+            f"({len(filelist)=}, from start: {elapsed=:.3f} s)."
+        )
+
+        # Write filelist to file remotely
+        tmp_filelist_path = workdir_remote / f"filelist_{time.time()}.txt"
+        self.fractal_ssh.write_remote_file(
+            path=tmp_filelist_path.as_posix(),
+            content=f"{filelist_string}\n",
+        )
+        elapsed = time.perf_counter() - t_0
+        logger.debug(
+            f"[_fetch_artifacts] File list written to {tmp_filelist_path} "
+            f"(from start: {elapsed=:.3f} s)."
+        )
+
+        # Create remote tarfile
+        t_0_tar = time.perf_counter()
+        tar_command = (
+            f"{self.python_worker_interpreter} "
+            "-m fractal_server.app.runner.compress_folder "
+            f"{workdir_remote.as_posix()} "
+            f"--filelist {tmp_filelist_path}"
+        )
+        self.fractal_ssh.run_command(cmd=tar_command)
+        t_1_tar = time.perf_counter()
+        logger.info(
+            f"[_fetch_artifacts] Remote archive {tarfile_path_remote} created"
+            f" - elapsed={t_1_tar - t_0_tar:.3f} s"
+        )
+
+        # Fetch tarfile
+        t_0_get = time.perf_counter()
+        self.fractal_ssh.fetch_file(
+            remote=tarfile_path_remote,
+            local=tarfile_path_local,
+        )
+        t_1_get = time.perf_counter()
+        logger.info(
+            "[_fetch_artifacts] Subfolder archive transferred back "
+            f"to {tarfile_path_local}"
+            f" - elapsed={t_1_get - t_0_get:.3f} s"
+        )
+
+        # Extract tarfile locally
+        extract_archive(Path(tarfile_path_local))
+
+        # Remove local tarfile
+        Path(tarfile_path_local).unlink(missing_ok=True)
+
+        t_1 = time.perf_counter()
+        logger.info(f"[_fetch_artifacts] End - elapsed={t_1 - t_0:.3f} s")
+
+    def _send_inputs(self, jobs: list[SlurmJob]) -> None:
+        """
+        Transfer the jobs subfolder to the remote host.
+        """
+        for job in jobs:
+            # Create local archive
+            tarfile_path_local = compress_folder(
+                job.workdir_local,
+                filelist_path=None,
+            )
+            tarfile_name = Path(tarfile_path_local).name
+            logger.info(f"Subfolder archive created at {tarfile_path_local}")
+
+            # Transfer archive
+            tarfile_path_remote = (
+                job.workdir_remote.parent / tarfile_name
+            ).as_posix()
+            t_0_put = time.perf_counter()
+            self.fractal_ssh.send_file(
+                local=tarfile_path_local,
+                remote=tarfile_path_remote,
+            )
+            t_1_put = time.perf_counter()
+            logger.info(
+                f"Subfolder archive transferred to {tarfile_path_remote}"
+                f" - elapsed={t_1_put - t_0_put:.3f} s"
+            )
+
+            # Remove local archive
+            Path(tarfile_path_local).unlink()
+            logger.debug(f"Local archive {tarfile_path_local} removed")
+
+            # Uncompress remote archive
+            tar_command = (
+                f"{self.python_worker_interpreter} -m "
+                "fractal_server.app.runner.extract_archive "
+                f"{tarfile_path_remote}"
+            )
+            self.fractal_ssh.run_command(cmd=tar_command)
+
+    def _run_remote_cmd(self, cmd: str) -> str:
+        stdout = self.fractal_ssh.run_command(cmd=cmd)
+        return stdout
+
+    def run_squeue(
+        self,
+        *,
+        job_ids: list[str],
+        base_interval: float = 2.0,
+        max_attempts: int = 7,
+    ) -> str:
+        """
+        Run `squeue` for a set of SLURM job IDs.
+
+        Different scenarios:
+
+        1. When `squeue -j` succeeds (with exit code 0), return its stdout.
+        2. When `squeue -j` fails (typical example:
+           `squeue -j {invalid_job_id}` fails with exit code 1), re-raise.
+           The error will be handled upstream.
+        3. When the SSH command fails because another thread is keeping the
+           lock of the `FractalSSH` object for a long time, mock the standard
+           output of the `squeue` command so that it looks like jobs are not
+           completed yet.
+        4. When the SSH command fails for other reasons, despite a forgiving
+           setup (7 connection attempts with base waiting interval of 2
+           seconds, with a cumulative timeout of 126 seconds), return an empty
+           string. This will be treated upstream as an empty `squeu` output,
+           indirectly resulting in marking the job as completed.
+        """
+
+        if len(job_ids) == 0:
+            return ""
+
+        job_id_single_str = ",".join([str(j) for j in job_ids])
+        cmd = (
+            "squeue --noheader --format='%i %T' --states=all "
+            f"--jobs={job_id_single_str}"
+        )
+
+        try:
+            stdout = self.fractal_ssh.run_command(
+                cmd=cmd,
+                base_interval=base_interval,
+                max_attempts=max_attempts,
+            )
+            return stdout
+        except FractalSSHCommandError as e:
+            raise e
+        except FractalSSHTimeoutError:
+            logger.warning(
+                "[run_squeue] Could not acquire lock, use stdout placeholder."
+            )
+            FAKE_STATUS = "FRACTAL_STATUS_PLACEHOLDER"
+            placeholder_stdout = "\n".join(
+                [f"{job_id} {FAKE_STATUS}" for job_id in job_ids]
+            )
+            return placeholder_stdout
+        except Exception as e:
+            logger.error(f"Ignoring `squeue` command failure {e}")
+            return ""