fractal-server 2.2.0a1__py3-none-any.whl → 2.3.0a0__py3-none-any.whl

fractal_server/app/runner/compress_folder.py

@@ -0,0 +1,120 @@
+import shlex
+import subprocess  # nosec
+import sys
+import tarfile
+import time
+from pathlib import Path
+from typing import Optional
+
+
+# COMPRESS_FOLDER_MODALITY = "python"
+COMPRESS_FOLDER_MODALITY = "cp-tar-rmtree"
+
+
+def _filter(info: tarfile.TarInfo) -> Optional[tarfile.TarInfo]:
+    if info.name.endswith(".pickle"):
+        filename = info.name.split("/")[-1]
+        parts = filename.split("_")
+        if len(parts) == 3 and parts[1] == "in":
+            return None
+        elif len(parts) == 5 and parts[3] == "in":
+            return None
+    elif info.name.endswith("slurm_submit.sbatch"):
+        return None
+    return info
+
+
+if __name__ == "__main__":
+    help_msg = (
+        "Expected use:\n"
+        "python -m fractal_server.app.runner.compress_folder "
+        "path/to/folder"
+    )
+
+    if len(sys.argv[1:]) != 1:
+        raise ValueError(
+            "Invalid argument(s).\n" f"{help_msg}\n" f"Provided: {sys.argv=}"
+        )
+
+    subfolder_path = Path(sys.argv[1])
+    t_0 = time.perf_counter()
+    print("[compress_folder.py] START")
+    print(f"[compress_folder.py] {COMPRESS_FOLDER_MODALITY=}")
+    print(f"[compress_folder.py] {subfolder_path=}")
+
+    job_folder = subfolder_path.parent
+    subfolder_name = subfolder_path.name
+    tarfile_path = (job_folder / f"{subfolder_name}.tar.gz").as_posix()
+    print(f"[compress_folder.py] {tarfile_path=}")
+
+    if COMPRESS_FOLDER_MODALITY == "python":
+        raise NotImplementedError()
+        with tarfile.open(tarfile_path, "w:gz") as tar:
+            tar.add(
+                subfolder_path,
+                arcname=".",  # ????
+                recursive=True,
+                filter=_filter,
+            )
+    elif COMPRESS_FOLDER_MODALITY == "cp-tar-rmtree":
+        import shutil
+        import time
+
+        subfolder_path_tmp_copy = (
+            subfolder_path.parent / f"{subfolder_path.name}_copy"
+        )
+
+        t0 = time.perf_counter()
+        # shutil.copytree(subfolder_path, subfolder_path_tmp_copy)
+        cmd_cp = (
+            "cp -r "
+            f"{subfolder_path.as_posix()} "
+            f"{subfolder_path_tmp_copy.as_posix()}"
+        )
+        res = subprocess.run(  # nosec
+            shlex.split(cmd_cp),
+            check=True,
+            capture_output=True,
+            encoding="utf-8",
+        )
+        t1 = time.perf_counter()
+        print("[compress_folder.py] `cp -r` END - " f"elapsed: {t1-t0:.3f} s")
+
+        cmd_tar = (
+            "tar czf "
+            f"{tarfile_path} "
+            "--exclude *sbatch --exclude *_in_*.pickle "
+            f"--directory={subfolder_path_tmp_copy.as_posix()} "
+            "."
+        )
+
+        print(f"[compress_folder.py] cmd tar:\n{cmd_tar}")
+        t0 = time.perf_counter()
+        res = subprocess.run(  # nosec
+            shlex.split(cmd_tar),
+            capture_output=True,
+            encoding="utf-8",
+        )
+        t1 = time.perf_counter()
+        t_1 = time.perf_counter()
+        print(f"[compress_folder.py] tar END - elapsed: {t1-t0:.3f} s")
+
+        print(f"[compress_folder] END - elapsed {t_1 - t_0:.3f} seconds")
+
+        if res.returncode != 0:
+            print("[compress_folder.py] ERROR in tar")
+            print(f"[compress_folder.py] tar stdout:\n{res.stdout}")
+            print(f"[compress_folder.py] tar stderr:\n{res.stderr}")
+
+            shutil.rmtree(subfolder_path_tmp_copy)
+            sys.exit(1)
+
+        t0 = time.perf_counter()
+        shutil.rmtree(subfolder_path_tmp_copy)
+        t1 = time.perf_counter()
+        print(
+            f"[compress_folder.py] shutil.rmtree END - elapsed: {t1-t0:.3f} s"
+        )
+
+    t_1 = time.perf_counter()
+    print(f"[compress_folder] END - elapsed {t_1 - t_0:.3f} seconds")

fractal_server/app/runner/executors/slurm/__init__.py

@@ -1,3 +0,0 @@
-from .executor import SlurmExecutor
-
-__all__ = ["SlurmExecutor"]

fractal_server/app/runner/executors/slurm/_batching.py

@@ -33,7 +33,6 @@ def _estimate_parallel_tasks_per_job(
     """
     Compute how many parallel tasks can fit in a given SLURM job
 
-
     Note: If more resources than available are requested, return 1. This
     assumes that further checks will be performed on the output of the current
     function, as is the case in the `heuristics` function below.

fractal_server/app/runner/executors/slurm/_slurm_config.py

@@ -312,7 +312,7 @@ class SlurmConfig(BaseModel, extra=Extra.forbid):
 
     def to_sbatch_preamble(
         self,
-        user_cache_dir: Optional[str] = None,
+        remote_export_dir: Optional[str] = None,
     ) -> list[str]:
         """
         Compile `SlurmConfig` object into the preamble of a SLURM submission
@@ -345,14 +345,14 @@ class SlurmConfig(BaseModel, extra=Extra.forbid):
             if value is not None:
                 # Handle the `time` parameter
                 if key == "time" and self.parallel_tasks_per_job > 1:
+                    # FIXME SSH: time setting must be handled better. Right now
+                    # we simply propagate `time`, but this is not enough when
+                    # several `srun` are combined in a single script.
                     logger.warning(
-                        "Ignore `#SBATCH --time=...` line (given: "
-                        f"{self.time=}) for parallel_tasks_per_job>1"
-                        f" (given: {self.parallel_tasks_per_job}), "
-                        "since scaling of time with number of tasks is "
-                        "not implemented."
+                        f"`time` SLURM parameter is set to {self.time}, "
+                        "but this does not take into account the number of "
+                        f"SLURM tasks ({self.parallel_tasks_per_job})."
                     )
-                    continue
                 option = key.replace("_", "-")
                 lines.append(f"{self.prefix} --{option}={value}")
 
@@ -361,12 +361,12 @@ class SlurmConfig(BaseModel, extra=Extra.forbid):
                 lines.append(line)
 
         if self.user_local_exports:
-            if user_cache_dir is None:
+            if remote_export_dir is None:
                 raise ValueError(
                     f"user_cache_dir=None but {self.user_local_exports=}"
                 )
             for key, value in self.user_local_exports.items():
-                tmp_value = str(Path(user_cache_dir) / value)
+                tmp_value = str(Path(remote_export_dir) / value)
                 lines.append(f"export {key}={tmp_value}")
 
         """

fractal_server/app/runner/executors/slurm/ssh/__init__.py

@@ -0,0 +1,3 @@
+from .executor import SlurmExecutor
+
+__all__ = ["SlurmExecutor"]

fractal_server/app/runner/executors/slurm/ssh/_executor_wait_thread.py

@@ -0,0 +1,112 @@
+import os
+import time
+import traceback
+from itertools import count
+from typing import Callable
+
+from cfut import FileWaitThread
+
+from ......logger import set_logger
+
+logger = set_logger(__name__)
+
+
+class FractalSlurmWaitThread(FileWaitThread):
+    """
+    Overrides the original clusterfutures.FileWaitThread, so that:
+
+    1. Each jobid in the waiting list is associated to a tuple of filenames,
+       rather than a single one.
+    2. In the `check` method, we avoid output-file existence checks (which
+       would require `sudo -u user ls` calls), and we rather check for the
+       existence of the shutdown file. All the logic to check whether a job is
+       complete is deferred to the `cfut.slurm.jobs_finished` function.
+    3. There are additional attributes (...).
+
+    This class is based on clusterfutures 0.5. Original Copyright: 2022
+    Adrian Sampson, released under the MIT licence
+    """
+
+    shutdown_file: str
+    shutdown_callback: Callable
+    jobs_finished_callback: Callable
+    slurm_poll_interval = 30
+    active_job_ids: list[str]
+
+    def __init__(self, *args, **kwargs):
+        """
+        Init method
+
+        This method is executed on the main thread.
+        """
+        super().__init__(*args, **kwargs)
+        self.active_job_ids = []
+
+    def wait(self, *, job_id: str):
+        """
+        Add a a new job to the set of jobs being waited for.
+
+        This method is executed on the main thread.
+        """
+        with self.lock:
+            self.active_job_ids.append(job_id)
+
+    def check_shutdown(self):
+        """
+        Check whether the shutdown file exists
+
+        This method is executed on the waiting thread.
+        """
+        if os.path.exists(self.shutdown_file):
+            logger.info(
+                f"Detected executor-shutdown file {self.shutdown_file}"
+            )
+            self.shutdown = True
+
+    def check_jobs(self):
+        """
+        Check whether some jobs are over, and call callback.
+
+        This method is executed on the waiting thread.
+        """
+        try:
+            if self.active_job_ids == []:
+                return
+            finished_jobs = self.jobs_finished_callback(self.active_job_ids)
+            if finished_jobs == set(self.active_job_ids):
+                self.callback(self.active_job_ids)
+                self.active_job_ids = []
+
+        except Exception:
+            # If anything goes wrong, print an exception without re-raising
+            traceback.print_exc()
+
+    def run(self):
+        """
+        Run forever (until a shutdown takes place) and trigger callback
+
+        This method is executed on the waiting thread.
+
+        Note that `shutdown_callback` only takes care of cleaning up the
+        FractalSlurmExecutor variables, and then the `return` here is enough
+        to fully clean up the `FractalFileWaitThread` object.
+        """
+
+        # FIXME SSH: are those try/except below needed?
+
+        skip = max(self.slurm_poll_interval // self.interval, 1)
+        for ind in count():
+            self.check_shutdown()
+            if self.shutdown:
+                try:
+                    self.shutdown_callback()
+                except Exception:  # nosec
+                    pass
+                return
+            if ind % skip == 0:
+                with self.lock:
+                    try:
+                        self.check_jobs()
+                    except Exception:  # nosec
+                        pass
+            time.sleep(self.interval)

fractal_server/app/runner/executors/slurm/ssh/_slurm_job.py

@@ -0,0 +1,120 @@
+from pathlib import Path
+from typing import Optional
+
+from cfut.util import random_string
+
+from fractal_server.app.runner.executors.slurm._slurm_config import (
+    SlurmConfig,
+)
+
+
+class SlurmJob:
+    """
+    Collect information related to a FractalSlurmSSHExecutor job
+
+    This includes three groups of attributes:
+
+    1. Attributes related to the (possibly multi-task) SLURM job, e.g.
+       submission-file path.
+    2. Attributes related to single tasks, e.g. the paths of their input/output
+       pickle files.
+    3. SLURM configuration options, encoded in a SlurmConfig object.
+
+    Note: A SlurmJob object is generally defined as a multi-task job. Jobs
+    coming from the `map` method must have `single_task_submission=False` (even
+    if `num_tasks_tot=1`), while jobs coming from `submit` must have it set to
+    `True`.
+
+    Attributes:
+        num_tasks_tot:
+            Total number of tasks to be executed as part of this SLURM job.
+        single_task_submission:
+            This must be `True` for jobs submitted as part of the `submit`
+            method, and `False` for jobs coming from the `map` method.
+        slurm_file_prefix:
+            Prefix for SLURM-job related files (submission script and SLURM
+            stdout/stderr); this is also needed in the
+            `_copy_files_from_remote_to_local` method.
+        wftask_file_prefixes:
+            Prefix for files that are created as part of the functions
+            submitted for execution on the `FractalSlurmSSHExecutor`; this is
+            needed in the `_copy_files_from_remote_to_local` method, and also
+            to construct the names of per-task input/output pickle files.
+        wftask_subfolder_name:
+            Name of the per-task subfolder (e.g. `7_task_name`).
+        slurm_script:
+            Path of SLURM submission script.
+        slurm_stdout:
+            Path of SLURM stdout file; if this includes `"%j"`, then this
+            string will be replaced by the SLURM job ID upon `sbatch`
+            submission.
+        slurm_stderr:
+            Path of SLURM stderr file; see `slurm_stdout` concerning `"%j"`.
+        workerids:
+            IDs that enter in the per-task input/output pickle files (one per
+            task).
+        input_pickle_files:
+            Input pickle files (one per task).
+        output_pickle_files:
+            Output pickle files (one per task).
+        slurm_config:
+            `SlurmConfig` object.
+    """
+
+    # Job-related attributes
+    num_tasks_tot: int
+    single_task_submission: bool
+    slurm_file_prefix: str
+    slurm_script_local: Path
+    slurm_script_remote: Path
+    slurm_stdout_local: Path
+    slurm_stdout_remote: Path
+    slurm_stderr_local: Path
+    slurm_stderr_remote: Path
+
+    # Per-task attributes
+    wftask_subfolder_name: str
+    workerids: tuple[str, ...]
+    wftask_file_prefixes: tuple[str, ...]
+    input_pickle_files_local: tuple[Path, ...]
+    input_pickle_files_remote: tuple[Path, ...]
+    output_pickle_files_local: tuple[Path, ...]
+    output_pickle_files_remote: tuple[Path, ...]
+
+    # Slurm configuration
+    slurm_config: SlurmConfig
+
+    def __init__(
+        self,
+        num_tasks_tot: int,
+        slurm_config: SlurmConfig,
+        workflow_task_file_prefix: Optional[str] = None,
+        slurm_file_prefix: Optional[str] = None,
+        wftask_file_prefixes: Optional[tuple[str, ...]] = None,
+        single_task_submission: bool = False,
+    ):
+        if single_task_submission and num_tasks_tot > 1:
+            raise ValueError(
+                "Trying to initialize SlurmJob with"
+                f"{single_task_submission=} and {num_tasks_tot=}."
+            )
+        self.num_tasks_tot = num_tasks_tot
+        self.single_task_submission = single_task_submission
+        self.slurm_file_prefix = slurm_file_prefix or "default_slurm_prefix"
+        if wftask_file_prefixes is None:
+            self.wftask_file_prefixes = tuple(
+                "default_wftask_prefix" for i in range(self.num_tasks_tot)
+            )
+        else:
+            self.wftask_file_prefixes = wftask_file_prefixes
+        self.workerids = tuple(
+            random_string() for i in range(self.num_tasks_tot)
+        )
+        self.slurm_config = slurm_config
+
+    def get_clean_output_pickle_files(self) -> tuple[str, ...]:
+        """
+        Transform all pathlib.Path objects in self.output_pickle_files to
+        strings
+        """
+        return tuple(str(f.as_posix()) for f in self.output_pickle_files_local)