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

fractal_server/app/runner/executors/{slurm/sudo → slurm_sudo}/_subprocess_run_as_user.py

@@ -19,7 +19,7 @@ import shlex
 import subprocess  # nosec
 from typing import Optional
 
-from ......logger import set_logger
+from fractal_server.logger import set_logger
 from fractal_server.string_tools import validate_cmd
 
 logger = set_logger(__name__)
@@ -65,10 +65,7 @@ def _run_command_as_user(
 
     if check and not res.returncode == 0:
         raise RuntimeError(
-            f"{cmd=}\n\n"
-            f"{res.returncode=}\n\n"
-            f"{res.stdout=}\n\n"
-            f"{res.stderr=}\n"
+            f"{cmd=}\n\n{res.returncode=}\n\n{res.stdout=}\n\n{res.stderr=}\n"
         )
 
     return res
@@ -91,81 +88,3 @@ def _mkdir_as_user(*, folder: str, user: str) -> None:
 
     cmd = f"mkdir -p {folder}"
     _run_command_as_user(cmd=cmd, user=user, check=True)
-
-
-def _glob_as_user(
-    *, folder: str, user: str, startswith: Optional[str] = None
-) -> list[str]:
-    """
-    Run `ls` in a folder (as a user) and filter results
-
-    Execute `ls` on a folder (impersonating a user, if `user` is not `None`)
-    and select results that start with `startswith` (if not `None`).
-
-    Arguments:
-        folder: Absolute path to the folder
-        user: If not `None`, the user to be impersonated via `sudo -u`
-        startswith: If not `None`, this is used to filter output of `ls`.
-    """
-
-    res = _run_command_as_user(cmd=f"ls {folder}", user=user, check=True)
-    output = res.stdout.split()
-    if startswith:
-        output = [f for f in output if f.startswith(startswith)]
-    return output
-
-
-def _glob_as_user_strict(
-    *,
-    folder: str,
-    user: str,
-    startswith: str,
-) -> list[str]:
-    """
-    Run `ls` in a folder (as a user) and filter results
-
-    Execute `ls` on a folder (impersonating a user, if `user` is not `None`)
-    and select results that comply with a set of rules. They all start with
-    `startswith` (if not `None`), and they match one of the known filename
-    patterns. See details in
-    https://github.com/fractal-analytics-platform/fractal-server/issues/1240
-
-
-    Arguments:
-        folder: Absolute path to the folder
-        user: If not `None`, the user to be impersonated via `sudo -u`
-        startswith: If not `None`, this is used to filter output of `ls`.
-    """
-
-    res = _run_command_as_user(cmd=f"ls {folder}", user=user, check=True)
-    output = res.stdout.split()
-
-    new_output = []
-    known_filenames = [
-        f"{startswith}{suffix}"
-        for suffix in [".args.json", ".metadiff.json", ".err", ".out", ".log"]
-    ]
-    for filename in output:
-        if filename in known_filenames:
-            new_output.append(filename)
-        elif filename.startswith(f"{startswith}_out_") and filename.endswith(
-            ".pickle"
-        ):
-            new_output.append(filename)
-
-    return new_output
-
-
-def _path_exists_as_user(*, path: str, user: Optional[str] = None) -> bool:
-    """
-    Impersonate a user and check if `path` exists via `ls`
-
-    Arguments:
-        path: Absolute file/folder path
-        user: If not `None`, user to be impersonated
-    """
-    res = _run_command_as_user(cmd=f"ls {path}", user=user)
-    if res.returncode == 0:
-        return True
-    else:
-        return False

fractal_server/app/runner/executors/slurm_sudo/runner.py

@@ -0,0 +1,193 @@
+import logging
+import os
+import shlex
+import subprocess  # nosec
+import sys
+from concurrent.futures import ThreadPoolExecutor
+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 ._subprocess_run_as_user import _mkdir_as_user
+from ._subprocess_run_as_user import _run_command_as_user
+from fractal_server.app.runner.exceptions import JobExecutionError
+from fractal_server.config import get_settings
+from fractal_server.logger import set_logger
+from fractal_server.syringe import Inject
+
+logger = set_logger(__name__)
+
+
+def _subprocess_run_or_raise(
+    full_command: str,
+) -> Optional[subprocess.CompletedProcess]:
+    try:
+        output = subprocess.run(  # nosec
+            shlex.split(full_command),
+            capture_output=True,
+            check=True,
+            encoding="utf-8",
+        )
+        return output
+    except subprocess.CalledProcessError as e:
+        error_msg = (
+            f"Submit command `{full_command}` failed. "
+            f"Original error:\n{str(e)}\n"
+            f"Original stdout:\n{e.stdout}\n"
+            f"Original stderr:\n{e.stderr}\n"
+        )
+        logging.error(error_msg)
+        raise JobExecutionError(info=error_msg)
+
+
+class SudoSlurmRunner(BaseSlurmRunner):
+    slurm_user: str
+    slurm_account: Optional[str] = None
+
+    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
+        slurm_account: Optional[str] = None,
+        slurm_user: str,
+    ) -> None:
+        """
+        Set parameters that are the same for different Fractal tasks and for
+        different SLURM jobs/tasks.
+        """
+
+        self.slurm_user = slurm_user
+        self.slurm_account = slurm_account
+        settings = Inject(get_settings)
+
+        super().__init__(
+            slurm_runner_type="sudo",
+            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 or sys.executable
+            ),
+        )
+
+    def _mkdir_local_folder(self, folder: str) -> None:
+        original_umask = os.umask(0)
+        Path(folder).mkdir(parents=True, mode=0o755)
+        os.umask(original_umask)
+
+    def _mkdir_remote_folder(self, folder: str) -> None:
+        _mkdir_as_user(folder=folder, user=self.slurm_user)
+
+    def _fetch_artifacts_single_job(self, job: SlurmJob) -> None:
+        """
+        Fetch artifacts for a single SLURM jobs.
+        """
+        logger.debug(
+            f"[_fetch_artifacts_single_job] {job.slurm_job_id=} START"
+        )
+        source_target_list = [
+            (job.slurm_stdout_remote, job.slurm_stdout_local),
+            (job.slurm_stderr_remote, job.slurm_stderr_local),
+        ]
+        for task in job.tasks:
+            source_target_list.extend(
+                [
+                    (
+                        task.output_pickle_file_remote,
+                        task.output_pickle_file_local,
+                    ),
+                    (
+                        task.task_files.log_file_remote,
+                        task.task_files.log_file_local,
+                    ),
+                    (
+                        task.task_files.args_file_remote,
+                        task.task_files.args_file_local,
+                    ),
+                    (
+                        task.task_files.metadiff_file_remote,
+                        task.task_files.metadiff_file_local,
+                    ),
+                ]
+            )
+
+        for source, target in source_target_list:
+            # NOTE: By setting encoding=None, we read/write bytes instead
+            # of strings; this is needed to also handle pickle files.
+            try:
+                res = _run_command_as_user(
+                    cmd=f"cat {source}",
+                    user=self.slurm_user,
+                    encoding=None,
+                    check=True,
+                )
+                # Write local file
+                with open(target, "wb") as f:
+                    f.write(res.stdout)
+                logger.debug(
+                    f"[_fetch_artifacts_single_job] Copied {source} into "
+                    f"{target}"
+                )
+            except RuntimeError as e:
+                logger.warning(
+                    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(
+            "[_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(
+            cmd=cmd,
+            user=self.slurm_user,
+            encoding="utf-8",
+            check=True,
+        )
+        return res.stdout
+
+    def run_squeue(self, job_ids: list[str]) -> str:
+        """
+        Run `squeue` for a set of SLURM job IDs.
+        """
+
+        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}"
+        )
+        res = _subprocess_run_or_raise(cmd)
+        return res.stdout

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)

fractal_server/app/runner/task_files.py

@@ -2,10 +2,19 @@ from pathlib import Path
 from typing import Optional
 from typing import Union
 
+from pydantic import BaseModel
+
+from fractal_server.app.runner.components import _index_to_component
 from fractal_server.string_tools import sanitize_string
 
+SUBMIT_PREFIX = "non_par"
+MULTISUBMIT_PREFIX = "par"
+
 
-def task_subfolder_name(order: Union[int, str], task_name: str) -> str:
+def task_subfolder_name(
+    order: Union[int, str],
+    task_name: str,
+) -> str:
     """
     Get name of task-specific subfolder.
 
@@ -17,108 +26,146 @@ def task_subfolder_name(order: Union[int, str], task_name: str) -> str:
     return f"{order}_{task_name_slug}"
 
 
-class TaskFiles:
+class TaskFiles(BaseModel):
     """
-    Group all file paths pertaining to a task
+    Files related to a task.
 
     Attributes:
-        workflow_dir_local:
-            Server-owned directory to store all task-execution-related relevant
-            files. Note: users cannot write directly to this folder.
-        workflow_dir_remote:
-            User-side directory with the same scope as `workflow_dir_local`,
-            and where a user can write.
-        subfolder_name:
-            Name of task-specific subfolder
-        remote_subfolder:
-            Path to user-side task-specific subfolder
+        root_dir_local:
+        root_dir_remote:
         task_name:
-            Name of the task
         task_order:
-            Positional order of the task within a workflow.
         component:
-            Specific component to run the task for (relevant for tasks to be
-            executed in parallel over many components).
-        file_prefix:
-            Prefix for all task-related files.
-        args:
-            Path for input json file.
-        metadiff:
-            Path for output json file with metadata update.
-        out:
-            Path for task-execution stdout.
-        err:
-            Path for task-execution stderr.
+        prefix:
     """
 
-    workflow_dir_local: Path
-    workflow_dir_remote: Path
-    remote_subfolder: Path
-    subfolder_name: str
+    # Parent directory
+    root_dir_local: Path
+    root_dir_remote: Path
+
+    # Per-wftask
     task_name: str
-    task_order: Optional[int] = None
+    task_order: int
+
+    # Per-single-component
     component: Optional[str] = None
+    prefix: Optional[str] = None
 
-    file_prefix: str
-    file_prefix_with_subfolder: str
-    args: Path
-    out: Path
-    err: Path
-    log: Path
-    metadiff: Path
-
-    def __init__(
-        self,
-        workflow_dir_local: Path,
-        workflow_dir_remote: Path,
-        task_name: str,
-        task_order: Optional[int] = None,
-        component: Optional[str] = None,
-    ):
-        self.workflow_dir_local = workflow_dir_local
-        self.workflow_dir_remote = workflow_dir_remote
-        self.task_order = task_order
-        self.task_name = task_name
-        self.component = component
-
-        if self.component is not None:
-            component_safe = sanitize_string(str(self.component))
-            component_safe = f"_par_{component_safe}"
-        else:
-            component_safe = ""
+    def _check_component(self):
+        if self.component is None:
+            raise ValueError("`component` cannot be None")
 
-        if self.task_order is not None:
-            order = str(self.task_order)
+    @property
+    def subfolder_name(self) -> str:
+        order = str(self.task_order or 0)
+        return task_subfolder_name(
+            order=order,
+            task_name=self.task_name,
+        )
+
+    @property
+    def wftask_subfolder_remote(self) -> Path:
+        return self.root_dir_remote / self.subfolder_name
+
+    @property
+    def wftask_subfolder_local(self) -> Path:
+        return self.root_dir_local / self.subfolder_name
+
+    @property
+    def prefix_component(self):
+        if self.prefix is None:
+            return self.component
         else:
-            order = "0"
-        self.file_prefix = f"{order}{component_safe}"
-        self.subfolder_name = task_subfolder_name(
-            order=order, task_name=self.task_name
+            return f"{self.prefix}-{self.component}"
+
+    @property
+    def log_file_local(self) -> str:
+        self._check_component()
+        return (
+            self.wftask_subfolder_local / f"{self.prefix_component}-log.txt"
+        ).as_posix()
+
+    @property
+    def log_file_remote_path(self) -> Path:
+        self._check_component()
+        return (
+            self.wftask_subfolder_remote / f"{self.prefix_component}-log.txt"
         )
-        self.remote_subfolder = self.workflow_dir_remote / self.subfolder_name
-        self.args = self.remote_subfolder / f"{self.file_prefix}.args.json"
-        self.out = self.remote_subfolder / f"{self.file_prefix}.out"
-        self.err = self.remote_subfolder / f"{self.file_prefix}.err"
-        self.log = self.remote_subfolder / f"{self.file_prefix}.log"
-        self.metadiff = (
-            self.remote_subfolder / f"{self.file_prefix}.metadiff.json"
+
+    @property
+    def log_file_remote(self) -> str:
+        return self.log_file_remote_path.as_posix()
+
+    @property
+    def args_file_local(self) -> str:
+        self._check_component()
+        return (
+            self.wftask_subfolder_local / f"{self.prefix_component}-args.json"
+        ).as_posix()
+
+    @property
+    def args_file_remote_path(self) -> Path:
+        self._check_component()
+        return (
+            self.wftask_subfolder_remote / f"{self.prefix_component}-args.json"
         )
 
+    @property
+    def args_file_remote(self) -> str:
+        return self.args_file_remote_path.as_posix()
 
-def get_task_file_paths(
-    workflow_dir_local: Path,
-    workflow_dir_remote: Path,
-    task_name: str,
-    task_order: Optional[int] = None,
-    component: Optional[str] = None,
-) -> TaskFiles:
+    @property
+    def metadiff_file_local(self) -> str:
+        self._check_component()
+        return (
+            self.wftask_subfolder_local
+            / f"{self.prefix_component}-metadiff.json"
+        ).as_posix()
+
+    @property
+    def metadiff_file_remote_path(self) -> Path:
+        self._check_component()
+        return (
+            self.wftask_subfolder_remote
+            / f"{self.prefix_component}-metadiff.json"
+        )
+
+    @property
+    def metadiff_file_remote(self) -> str:
+        return self.metadiff_file_remote_path.as_posix()
+
+    @property
+    def remote_files_dict(self) -> dict[str, str]:
+        return dict(
+            args_file_remote=self.args_file_remote,
+            metadiff_file_remote=self.metadiff_file_remote,
+            log_file_remote=self.log_file_remote,
+        )
+
+
+def enrich_task_files_multisubmit(
+    *,
+    tot_tasks: int,
+    batch_size: int,
+    base_task_files: TaskFiles,
+) -> list[TaskFiles]:
     """
-    Return the corrisponding TaskFiles object
+    Expand `TaskFiles` objects with `component` and `prefix`.
     """
-    return TaskFiles(
-        workflow_dir_local=workflow_dir_local,
-        workflow_dir_remote=workflow_dir_remote,
-        task_name=task_name,
-        task_order=task_order,
-        component=component,
-    )
+
+    new_list_task_files: list[TaskFiles] = []
+    for absolute_index in range(tot_tasks):
+        ind_batch = absolute_index // batch_size
+        new_list_task_files.append(
+            TaskFiles(
+                **base_task_files.model_dump(
+                    exclude={
+                        "component",
+                        "prefix",
+                    }
+                ),
+                prefix=f"{MULTISUBMIT_PREFIX}-{ind_batch:06d}",
+                component=_index_to_component(absolute_index),
+            )
+        )
+    return new_list_task_files