fractal-server 2.14.0a21__tar.gz → 2.14.0a23__tar.gz

{fractal_server-2.14.0a21 → fractal_server-2.14.0a23}/PKG-INFO

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

fractal_server-2.14.0a23/fractal_server/__init__.py

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

{fractal_server-2.14.0a21 → fractal_server-2.14.0a23}/fractal_server/app/routes/api/v2/history.py

@@ -1,3 +1,4 @@
+from copy import deepcopy
 from typing import Any
 from typing import Optional
 
@@ -92,7 +93,7 @@ async def get_workflow_tasks_statuses(
         db=db,
     )
 
-    response = {}
+    response: dict[int, dict[str, int | str] | None] = {}
     for wftask in workflow.task_list:
         res = await db.execute(
             select(HistoryRun)
@@ -130,7 +131,18 @@ async def get_workflow_tasks_statuses(
                 f"num_{target_status.value}_images"
             ] = num_images
 
-    return JSONResponse(content=response, status_code=200)
+    new_response = deepcopy(response)
+    for key, value in response.items():
+        if value is not None:
+            num_total_images = sum(
+                value[f"num_{target_status.value}_images"]
+                for target_status in HistoryUnitStatus
+            )
+            if num_total_images > value["num_available_images"]:
+                value["num_available_images"] = None
+        new_response[key] = value
+
+    return JSONResponse(content=new_response, status_code=200)
 
 
 @router.get("/project/{project_id}/status/run/")

{fractal_server-2.14.0a21 → fractal_server-2.14.0a23}/fractal_server/app/runner/compress_folder.py

@@ -11,8 +11,8 @@ built-in `tarfile` library has to do with performance issues we observed
 when handling files which were just created within a SLURM job, and in the
 context of a CephFS filesystem.
 """
-import shutil
 import sys
+import time
 from pathlib import Path
 
 from fractal_server.app.runner.run_subprocess import run_subprocess
@@ -20,48 +20,66 @@ from fractal_server.logger import get_logger
 from fractal_server.logger import set_logger
 
 
-def copy_subfolder(src: Path, dest: Path, logger_name: str):
+def _copy_subfolder(src: Path, dest: Path, logger_name: str):
+    t_start = time.perf_counter()
     cmd_cp = f"cp -r {src.as_posix()} {dest.as_posix()}"
     logger = get_logger(logger_name=logger_name)
     logger.debug(f"{cmd_cp=}")
     res = run_subprocess(cmd=cmd_cp, logger_name=logger_name)
+    elapsed = time.perf_counter() - t_start
+    logger.debug(f"[_copy_subfolder] END {elapsed=} s ({dest.as_posix()})")
     return res
 
 
-def create_tar_archive(
-    tarfile_path: Path,
+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)
-
-    if remote_to_local:
-        exclude_options = "--exclude *sbatch --exclude *_in_*.pickle "
+    logger.debug(f"[_create_tar_archive] START ({tarfile_path})")
+    t_start = time.perf_counter()
+
+    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"
+        )
+
+    logger.critical(f"cmd tar:\n{cmd_tar}")
 
-    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})")
 
 
-def remove_temp_subfolder(subfolder_path_tmp_copy: Path, logger_name: str):
+def _remove_temp_subfolder(subfolder_path_tmp_copy: Path, logger_name: str):
     logger = get_logger(logger_name)
+    t_start = time.perf_counter()
     try:
-        logger.debug(f"Now remove {subfolder_path_tmp_copy}")
-        shutil.rmtree(subfolder_path_tmp_copy)
+        cmd_rm = f"rm -rf {subfolder_path_tmp_copy}"
+        logger.debug(f"cmd rm:\n{cmd_rm}")
+        run_subprocess(cmd=cmd_rm, logger_name=logger_name, allow_char="*")
     except Exception as e:
-        logger.debug(f"ERROR during shutil.rmtree: {e}")
+        logger.debug(f"ERROR during {cmd_rm}: {e}")
+    elapsed = time.perf_counter() - t_start
+    logger.debug(
+        f"[_remove_temp_subfolder] END {elapsed=} s "
+        f"({subfolder_path_tmp_copy=})"
+    )
 
 
 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`
@@ -91,14 +109,16 @@ def compress_folder(
         subfolder_path.parent / f"{subfolder_path.name}_copy"
     )
     try:
-        copy_subfolder(
-            subfolder_path, subfolder_path_tmp_copy, logger_name=logger_name
+        _copy_subfolder(
+            subfolder_path,
+            subfolder_path_tmp_copy,
+            logger_name=logger_name,
         )
-        create_tar_archive(
+        _create_tar_archive(
             tarfile_path,
             subfolder_path_tmp_copy,
             logger_name=logger_name,
-            remote_to_local=remote_to_local,
+            filelist_path=filelist_path,
         )
         return tarfile_path
 
@@ -107,7 +127,9 @@ def compress_folder(
         sys.exit(1)
 
     finally:
-        remove_temp_subfolder(subfolder_path_tmp_copy, logger_name=logger_name)
+        _remove_temp_subfolder(
+            subfolder_path_tmp_copy, logger_name=logger_name
+        )
 
 
 def main(sys_argv: list[str]):
@@ -115,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.0a21 → fractal_server-2.14.0a23}/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,19 @@ 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.
+        """
+        # Extract `workdir_remote` and `workdir_local`
+        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.0a21 → fractal_server-2.14.0a23}/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.0a23/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 iterations?
+        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