fractal-server 2.14.0a22__tar.gz → 2.14.0a24__tar.gz

{fractal_server-2.14.0a22 → fractal_server-2.14.0a24}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: fractal-server
-Version: 2.14.0a22
+Version: 2.14.0a24
 Summary: Backend component of the Fractal analytics platform
 License: BSD-3-Clause
 Author: Tommaso Comparin

fractal_server-2.14.0a24/fractal_server/__init__.py

@@ -0,0 +1 @@
+__VERSION__ = "2.14.0a24"

{fractal_server-2.14.0a22 → fractal_server-2.14.0a24}/fractal_server/app/runner/compress_folder.py

@@ -35,24 +35,27 @@ def _create_tar_archive(
     tarfile_path: str,
     subfolder_path_tmp_copy: Path,
     logger_name: str,
-    remote_to_local: bool,
+    filelist_path: str | None,
 ):
     logger = get_logger(logger_name)
     logger.debug(f"[_create_tar_archive] START ({tarfile_path})")
     t_start = time.perf_counter()
 
-    if remote_to_local:
-        exclude_options = "--exclude *sbatch --exclude *_in_*.pickle "
+    if filelist_path is None:
+        cmd_tar = (
+            f"tar -c -z -f {tarfile_path} "
+            f"--directory={subfolder_path_tmp_copy.as_posix()} "
+            "."
+        )
     else:
-        exclude_options = ""
+        cmd_tar = (
+            f"tar -c -z -f {tarfile_path} "
+            f"--directory={subfolder_path_tmp_copy.as_posix()} "
+            f"--files-from={filelist_path} --ignore-failed-read"
+        )
 
-    cmd_tar = (
-        f"tar czf {tarfile_path} "
-        f"{exclude_options} "
-        f"--directory={subfolder_path_tmp_copy.as_posix()} "
-        "."
-    )
     logger.debug(f"cmd tar:\n{cmd_tar}")
+
     run_subprocess(cmd=cmd_tar, logger_name=logger_name, allow_char="*")
     elapsed = time.perf_counter() - t_start
     logger.debug(f"[_create_tar_archive] END {elapsed=} s ({tarfile_path})")
@@ -75,7 +78,8 @@ def _remove_temp_subfolder(subfolder_path_tmp_copy: Path, logger_name: str):
 
 
 def compress_folder(
-    subfolder_path: Path, remote_to_local: bool = False
+    subfolder_path: Path,
+    filelist_path: str | None,
 ) -> str:
     """
     Compress e.g. `/path/archive` into `/path/archive.tar.gz`
@@ -114,7 +118,7 @@ def compress_folder(
             tarfile_path,
             subfolder_path_tmp_copy,
             logger_name=logger_name,
-            remote_to_local=remote_to_local,
+            filelist_path=filelist_path,
         )
         return tarfile_path
 
@@ -133,15 +137,21 @@ def main(sys_argv: list[str]):
     help_msg = (
         "Expected use:\n"
         "python -m fractal_server.app.runner.compress_folder "
-        "path/to/folder [--remote-to-local]\n"
+        "path/to/folder [--filelist /path/to/filelist]\n"
     )
     num_args = len(sys_argv[1:])
     if num_args == 0:
         sys.exit(f"Invalid argument.\n{help_msg}\nProvided: {sys_argv[1:]=}")
     elif num_args == 1:
-        compress_folder(subfolder_path=Path(sys_argv[1]))
-    elif num_args == 2 and sys_argv[2] == "--remote-to-local":
-        compress_folder(subfolder_path=Path(sys_argv[1]), remote_to_local=True)
+        compress_folder(
+            subfolder_path=Path(sys_argv[1]),
+            filelist_path=None,
+        )
+    elif num_args == 3 and sys_argv[2] == "--filelist":
+        compress_folder(
+            subfolder_path=Path(sys_argv[1]),
+            filelist_path=sys_argv[3],
+        )
     else:
         sys.exit(f"Invalid argument.\n{help_msg}\nProvided: {sys_argv[1:]=}")
 

{fractal_server-2.14.0a22 → fractal_server-2.14.0a24}/fractal_server/app/runner/executors/slurm_common/base_slurm_runner.py

@@ -60,6 +60,7 @@ class BaseSlurmRunner(BaseRunner):
         root_dir_local: Path,
         root_dir_remote: Path,
         slurm_runner_type: Literal["ssh", "sudo"],
+        python_worker_interpreter: str,
         common_script_lines: Optional[list[str]] = None,
         user_cache_dir: Optional[str] = None,
         poll_interval: Optional[int] = None,
@@ -70,6 +71,7 @@ class BaseSlurmRunner(BaseRunner):
         self.common_script_lines = common_script_lines or []
         self._check_slurm_account()
         self.user_cache_dir = user_cache_dir
+        self.python_worker_interpreter = python_worker_interpreter
 
         settings = Inject(get_settings)
 
@@ -327,9 +329,9 @@ class BaseSlurmRunner(BaseRunner):
         )
         logger.info("[_submit_single_sbatch] END")
 
-    def _copy_files_from_remote_to_local(
+    def _fetch_artifacts(
         self,
-        slurm_job: SlurmJob,
+        finished_slurm_jobs: list[SlurmJob],
     ) -> None:
         raise NotImplementedError("Implement in child class.")
 
@@ -530,14 +532,14 @@ class BaseSlurmRunner(BaseRunner):
             # Look for finished jobs
             finished_job_ids = self._get_finished_jobs(job_ids=self.job_ids)
             logger.debug(f"[submit] {finished_job_ids=}")
-
+            finished_jobs = [
+                self.jobs[_slurm_job_id] for _slurm_job_id in finished_job_ids
+            ]
+            self._fetch_artifacts(finished_jobs)
             with next(get_sync_db()) as db:
                 for slurm_job_id in finished_job_ids:
                     logger.debug(f"[submit] Now process {slurm_job_id=}")
                     slurm_job = self.jobs.pop(slurm_job_id)
-                    self._copy_files_from_remote_to_local(
-                        slurm_job
-                    )  # FIXME: add prefix  # noqa
                     was_job_scancelled = slurm_job_id in scancelled_job_ids
                     result, exception = self._postprocess_single_task(
                         task=slurm_job.tasks[0],
@@ -653,7 +655,9 @@ class BaseSlurmRunner(BaseRunner):
         if len(args_batches) != math.ceil(tot_tasks / tasks_per_job):
             raise RuntimeError("Something wrong here while batching tasks")
 
-        logger.info(f"START submission phase, {list(self.jobs.keys())=}")
+        # Part 1/3: Iterate over chunks, prepare SlurmJob objects
+        logger.info("[multisubmit] Prepare `SlurmJob`s.")
+        jobs_to_submit = []
         for ind_batch, chunk in enumerate(args_batches):
             prefix = f"{MULTISUBMIT_PREFIX}-{ind_batch:06d}"
             tasks = []
@@ -673,17 +677,26 @@ class BaseSlurmRunner(BaseRunner):
                     ),
                 )
 
-            slurm_job = SlurmJob(
-                prefix=prefix,
-                workdir_local=workdir_local,
-                workdir_remote=workdir_remote,
-                tasks=tasks,
+            jobs_to_submit.append(
+                SlurmJob(
+                    prefix=prefix,
+                    workdir_local=workdir_local,
+                    workdir_remote=workdir_remote,
+                    tasks=tasks,
+                )
             )
+
+        # FIXME: split parts 2 and 3
+        # Part 2/3. Transfer all relevant input files (for SSH)
+        # Part 3/3. Run all `sbatch`es and update `self.jobs`
+        logger.info("[multisubmit] Transfer files and submit jobs.")
+        for slurm_job in jobs_to_submit:
             self._submit_single_sbatch(
                 func,
                 slurm_job=slurm_job,
                 slurm_config=config,
             )
+
         if task_type == "parallel":
             # FIXME: replace loop with a `bulk_update_history_unit` function
             for ind, task_files in enumerate(list_task_files):
@@ -711,20 +724,21 @@ class BaseSlurmRunner(BaseRunner):
 
         # Retrieval phase
         logger.info("[multisubmit] START retrieval phase")
+        scancelled_job_ids = []
         while len(self.jobs) > 0:
 
             # Look for finished jobs
             finished_job_ids = self._get_finished_jobs(job_ids=self.job_ids)
             logger.debug(f"[multisubmit] {finished_job_ids=}")
+            finished_jobs = [
+                self.jobs[_slurm_job_id] for _slurm_job_id in finished_job_ids
+            ]
+            self._fetch_artifacts(finished_jobs)
 
-            scancelled_job_ids = []
             with next(get_sync_db()) as db:
                 for slurm_job_id in finished_job_ids:
                     logger.info(f"[multisubmit] Now process {slurm_job_id=}")
                     slurm_job = self.jobs.pop(slurm_job_id)
-                    self._copy_files_from_remote_to_local(
-                        slurm_job
-                    )  # FIXME: add prefix  # noqa
                     for task in slurm_job.tasks:
                         logger.info(f"[multisubmit] Now process {task.index=}")
                         was_job_scancelled = slurm_job_id in scancelled_job_ids
@@ -810,3 +824,17 @@ class BaseSlurmRunner(BaseRunner):
                 )
         logger.info("[scancel_jobs] END")
         return scancelled_job_ids
+
+    def validate_slurm_jobs_workdirs(
+        self,
+        slurm_jobs: list[SlurmJob],
+    ) -> None:
+        """
+        Check that a list of `SlurmJob`s have homogeneous working folders.
+        """
+        set_workdir_local = set(_job.workdir_local for _job in slurm_jobs)
+        set_workdir_remote = set(_job.workdir_remote for _job in slurm_jobs)
+        if len(set_workdir_local) > 1:
+            raise ValueError(f"Non-unique values in {set_workdir_local=}.")
+        if len(set_workdir_remote) > 1:
+            raise ValueError(f"Non-unique values in {set_workdir_remote=}.")

{fractal_server-2.14.0a22 → fractal_server-2.14.0a24}/fractal_server/app/runner/executors/slurm_common/slurm_job_task_models.py

@@ -20,31 +20,47 @@ class SlurmTask(BaseModel):
     index: int
 
     @property
-    def input_pickle_file_local(self) -> str:
+    def input_pickle_file_local_path(self) -> Path:
         return (
             self.workdir_local / f"{self.prefix}-{self.component}-input.pickle"
-        ).as_posix()
+        )
 
     @property
-    def input_pickle_file_remote(self) -> str:
+    def input_pickle_file_remote_path(self) -> Path:
         return (
             self.workdir_remote
             / f"{self.prefix}-{self.component}-input.pickle"
-        ).as_posix()
+        )
 
     @property
-    def output_pickle_file_local(self) -> str:
+    def output_pickle_file_local_path(self) -> Path:
         return (
             self.workdir_local
             / f"{self.prefix}-{self.component}-output.pickle"
-        ).as_posix()
+        )
 
     @property
-    def output_pickle_file_remote(self) -> str:
+    def output_pickle_file_remote_path(self) -> Path:
         return (
             self.workdir_remote
             / f"{self.prefix}-{self.component}-output.pickle"
-        ).as_posix()
+        )
+
+    @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):
@@ -74,29 +90,45 @@ class SlurmJob(BaseModel):
             return "%j"
 
     @property
-    def slurm_stdout_remote(self) -> str:
+    def slurm_stdout_remote_path(self) -> Path:
         return (
             self.workdir_remote
             / f"{self.prefix}-slurm-{self.slurm_job_id_placeholder}.out"
-        ).as_posix()
+        )
 
     @property
-    def slurm_stderr_remote(self) -> str:
+    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"
-        ).as_posix()
+        )
 
     @property
-    def slurm_stdout_local(self) -> str:
+    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"
-        ).as_posix()
+        )
 
     @property
-    def slurm_stderr_local(self) -> str:
+    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"
-        ).as_posix()
+        )
+
+    @property
+    def slurm_stderr_local(self) -> str:
+        return self.slurm_stderr_local_path.as_posix()

fractal_server-2.14.0a24/fractal_server/app/runner/executors/slurm_ssh/runner.py

@@ -0,0 +1,208 @@
+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.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
+        # # FIXME can we make this more efficient with iterators?
+        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"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(
+            f"Subfolder archive transferred back 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"[_get_subfolder_sftp] 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

{fractal_server-2.14.0a22 → fractal_server-2.14.0a24}/fractal_server/app/runner/executors/slurm_sudo/runner.py

@@ -3,6 +3,7 @@ import os
 import shlex
 import subprocess  # nosec
 import sys
+from concurrent.futures import ThreadPoolExecutor
 from pathlib import Path
 from typing import Optional
 
@@ -15,7 +16,6 @@ from fractal_server.config import get_settings
 from fractal_server.logger import set_logger
 from fractal_server.syringe import Inject
 
-
 logger = set_logger(__name__)
 
 
@@ -67,10 +67,6 @@ class SudoSlurmRunner(BaseSlurmRunner):
         self.slurm_account = slurm_account
         settings = Inject(get_settings)
 
-        self.python_worker_interpreter = (
-            settings.FRACTAL_SLURM_WORKER_PYTHON or sys.executable
-        )
-
         super().__init__(
             slurm_runner_type="sudo",
             root_dir_local=root_dir_local,
@@ -78,6 +74,9 @@ class SudoSlurmRunner(BaseSlurmRunner):
             common_script_lines=common_script_lines,
             user_cache_dir=user_cache_dir,
             poll_interval=poll_interval,
+            python_worker_interpreter=(
+                settings.FRACTAL_SLURM_WORKER_PYTHON or sys.executable
+            ),
         )
 
     def _mkdir_local_folder(self, folder: str) -> None:
@@ -88,12 +87,12 @@ class SudoSlurmRunner(BaseSlurmRunner):
     def _mkdir_remote_folder(self, folder: str) -> None:
         _mkdir_as_user(folder=folder, user=self.slurm_user)
 
-    def _copy_files_from_remote_to_local(self, job: SlurmJob) -> None:
+    def _fetch_artifacts_single_job(self, job: SlurmJob) -> None:
         """
-        Note: this would differ for SSH
+        Fetch artifacts for a single SLURM jobs.
         """
         logger.debug(
-            f"[_copy_files_from_remote_to_local] {job.slurm_job_id=} START"
+            f"[_fetch_artifacts_single_job] {job.slurm_job_id=} START"
         )
         source_target_list = [
             (job.slurm_stdout_remote, job.slurm_stdout_local),
@@ -140,9 +139,30 @@ class SudoSlurmRunner(BaseSlurmRunner):
                     f"SKIP copy {source} into {target}. "
                     f"Original error: {str(e)}"
                 )
+        logger.debug(f"[_fetch_artifacts_single_job] {job.slurm_job_id=} END")
+
+    def _fetch_artifacts(
+        self,
+        finished_slurm_jobs: list[SlurmJob],
+    ) -> None:
+        """
+        Fetch artifacts for a list of SLURM jobs.
+        """
+        MAX_NUM_THREADS = 4
+        THREAD_NAME_PREFIX = "fetch_artifacts"
         logger.debug(
-            f"[_copy_files_from_remote_to_local] {job.slurm_job_id=} END"
+            "[_fetch_artifacts] START "
+            f"({MAX_NUM_THREADS=}, {len(finished_slurm_jobs)=})."
         )
+        with ThreadPoolExecutor(
+            max_workers=MAX_NUM_THREADS,
+            thread_name_prefix=THREAD_NAME_PREFIX,
+        ) as executor:
+            executor.map(
+                self._fetch_artifacts_single_job,
+                finished_slurm_jobs,
+            )
+        logger.debug("[_fetch_artifacts] END.")
 
     def _run_remote_cmd(self, cmd: str) -> str:
         res = _run_command_as_user(

{fractal_server-2.14.0a22 → fractal_server-2.14.0a24}/fractal_server/app/runner/extract_archive.py

@@ -57,9 +57,7 @@ def extract_archive(archive_path: Path):
 
     # Run tar command
     cmd_tar = (
-        f"tar -xzvf {archive_path} "
-        f"--directory={subfolder_path.as_posix()} "
-        "."
+        f"tar -xzvf {archive_path} --directory={subfolder_path.as_posix()}"
     )
     logger.debug(f"{cmd_tar=}")
     run_subprocess(cmd=cmd_tar, logger_name=logger_name)