fractal-server 2.14.0a2__py3-none-any.whl → 2.14.0a4__py3-none-any.whl

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

@@ -1,5 +1,7 @@
 import json
 import logging
+import math
+import os
 import shlex
 import subprocess  # nosec
 import sys
@@ -16,13 +18,13 @@ from ._check_jobs_status import get_finished_jobs
 from ._subprocess_run_as_user import _mkdir_as_user
 from ._subprocess_run_as_user import _run_command_as_user
 from fractal_server import __VERSION__
-from fractal_server.app.history import HistoryItemImageStatus
-from fractal_server.app.history import update_all_images
-from fractal_server.app.history import update_single_image
 from fractal_server.app.runner.components import _COMPONENT_KEY_
 from fractal_server.app.runner.exceptions import JobExecutionError
 from fractal_server.app.runner.exceptions import TaskExecutionError
 from fractal_server.app.runner.executors.base_runner import BaseRunner
+from fractal_server.app.runner.executors.slurm_common._batching import (
+    heuristics,
+)
 from fractal_server.app.runner.executors.slurm_common._slurm_config import (
     SlurmConfig,
 )
@@ -32,6 +34,11 @@ from fractal_server.config import get_settings
 from fractal_server.logger import set_logger
 from fractal_server.syringe import Inject
 
+# from fractal_server.app.history import ImageStatus
+# from fractal_server.app.history import update_all_images
+# from fractal_server.app.history import update_single_image
+# from fractal_server.app.history import update_single_image_logfile
+
 
 logger = set_logger(__name__)
 
@@ -56,8 +63,10 @@ class SlurmTask(BaseModel):
     component: 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(self) -> str:
@@ -89,7 +98,7 @@ class SlurmJob(BaseModel):
     label: str
     workdir_local: Path
     workdir_remote: Path
-    tasks: tuple[SlurmTask]
+    tasks: list[SlurmTask]
 
     @property
     def slurm_log_file_local(self) -> str:
@@ -213,7 +222,9 @@ class RunnerSlurmSudo(BaseRunner):
         self.root_dir_remote = root_dir_remote
 
         # Create folders
-        self.root_dir_local.mkdir(parents=True, exist_ok=True)
+        original_umask = os.umask(0)
+        self.root_dir_local.mkdir(parents=True, exist_ok=True, mode=0o755)
+        os.umask(original_umask)
         _mkdir_as_user(
             folder=self.root_dir_remote.as_posix(),
             user=self.slurm_user,
@@ -242,9 +253,8 @@ class RunnerSlurmSudo(BaseRunner):
     def is_shutdown(self) -> bool:
         return self.shutdown_file.exists()
 
-    def scancel_if_shutdown(self) -> None:
-
-        logger.debug("[exit_if_shutdown] START")
+    def scancel_jobs(self) -> None:
+        logger.debug("[scancel_jobs] START")
 
         if self.jobs:
             scancel_string = " ".join(self.job_ids)
@@ -258,22 +268,20 @@ class RunnerSlurmSudo(BaseRunner):
                 )
             except RuntimeError as e:
                 logger.warning(
-                    "[exit_if_shutdown] `scancel` command failed. "
+                    "[scancel_jobs] `scancel` command failed. "
                     f"Original error:\n{str(e)}"
                 )
 
-        logger.debug("[exit_if_shutdown] END")
+        logger.debug("[scancel_jobs] END")
 
     def _submit_single_sbatch(
         self,
         func,
-        parameters,  # FIXME this should be per-task
         slurm_job: SlurmJob,
         slurm_config: SlurmConfig,
     ) -> str:
-
-        if len(slurm_job.tasks) > 1:
-            raise NotImplementedError()
+        # if len(slurm_job.tasks) > 1:
+        #     raise NotImplementedError()
 
         # Prepare input pickle(s)
         versions = dict(
@@ -283,10 +291,7 @@ class RunnerSlurmSudo(BaseRunner):
         )
         for task in slurm_job.tasks:
             _args = []
-            # TODO: make parameters task-dependent
-            _kwargs = dict(
-                parameters=parameters
-            )  # FIXME: this should be per-tas
+            _kwargs = dict(parameters=task.parameters)
             funcser = cloudpickle.dumps((versions, func, _args, _kwargs))
             with open(task.input_pickle_file_local, "wb") as f:
                 f.write(funcser)
@@ -339,6 +344,7 @@ class RunnerSlurmSudo(BaseRunner):
 
         # Add job to self.jobs
         self.jobs[slurm_job.slurm_job_id] = slurm_job
+        logger.debug(f"Added {slurm_job.slurm_job_id} to self.jobs.")
 
     @property
     def job_ids(self) -> list[str]:
@@ -352,22 +358,25 @@ class RunnerSlurmSudo(BaseRunner):
             (job.slurm_log_file_remote, job.slurm_log_file_local)
         ]
         for task in job.tasks:
-            source_target_list.append(
-                (task.output_pickle_file_remote, task.output_pickle_file_local)
-            )
-            source_target_list.append(
-                (
-                    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,
-                ),
+            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:
@@ -383,12 +392,12 @@ class RunnerSlurmSudo(BaseRunner):
                 # Write local file
                 with open(target, "wb") as f:
                     f.write(res.stdout)
+                logger.critical(f"Copied {source} into {target}")
             except RuntimeError as e:
                 logger.warning(
                     f"SKIP copy {source} into {target}. "
                     f"Original error: {str(e)}"
                 )
-            logger.debug(f"Copied {source} into {target}")
 
     def _postprocess_single_task(
         self, *, task: SlurmTask
@@ -416,11 +425,9 @@ class RunnerSlurmSudo(BaseRunner):
         parameters: dict[str, Any],
         history_item_id: int,
         task_files: TaskFiles,
+        slurm_config: SlurmConfig,
         in_compound_task: bool = False,
-        slurm_config: Optional[SlurmConfig] = None,
-        **kwargs,
     ) -> tuple[Any, Exception]:
-
         workdir_local = task_files.wftask_subfolder_local
         workdir_remote = task_files.wftask_subfolder_remote
 
@@ -433,25 +440,29 @@ class RunnerSlurmSudo(BaseRunner):
 
         if self.jobs != {}:
             if not in_compound_task:
-                update_all_images(
-                    history_item_id=history_item_id,
-                    status=HistoryItemImageStatus.FAILED,
-                )
+                pass
+                # update_all_images(
+                #     history_item_id=history_item_id,
+                #     status=ImageStatus.FAILED,
+                # )
             raise JobExecutionError("Unexpected branch: jobs should be empty.")
 
         if self.is_shutdown():
             if not in_compound_task:
-                update_all_images(
-                    history_item_id=history_item_id,
-                    status=HistoryItemImageStatus.FAILED,
-                )
+                pass
+                # update_all_images(
+                #     history_item_id=history_item_id,
+                #     status=ImageStatus.FAILED,
+                # )
             raise JobExecutionError("Cannot continue after shutdown.")
 
         # Validation phase
         self.validate_submit_parameters(parameters)
 
         # Create task subfolder
-        workdir_local.mkdir(parents=True, exist_ok=True)
+        original_umask = os.umask(0)
+        workdir_local.mkdir(parents=True, mode=0o755)
+        os.umask(original_umask)
         _mkdir_as_user(
             folder=workdir_remote.as_posix(),
             user=self.slurm_user,
@@ -464,7 +475,9 @@ class RunnerSlurmSudo(BaseRunner):
             workdir_remote=workdir_remote,
             tasks=[
                 SlurmTask(
+                    index=0,
                     component="0",
+                    parameters=parameters,
                     workdir_remote=workdir_remote,
                     workdir_local=workdir_local,
                     task_files=task_files,
@@ -473,16 +486,16 @@ class RunnerSlurmSudo(BaseRunner):
         )  # TODO: replace with actual values (BASED ON TASKFILES)
         self._submit_single_sbatch(
             func,
-            parameters=parameters,
             slurm_job=slurm_job,
+            slurm_config=slurm_config,
         )
 
-        LOGFILE = task_files.log_file_local
+        # LOGFILE = task_files.log_file_local
 
         # Retrieval phase
         while len(self.jobs) > 0:
             if self.is_shutdown():
-                self.scancel_if_shutdown()
+                self.scancel_jobs()
             finished_job_ids = get_finished_jobs(job_ids=self.job_ids)
             for slurm_job_id in finished_job_ids:
                 slurm_job = self.jobs.pop(slurm_job_id)
@@ -494,17 +507,19 @@ class RunnerSlurmSudo(BaseRunner):
 
         if not in_compound_task:
             if exception is None:
-                update_all_images(
-                    history_item_id=history_item_id,
-                    status=HistoryItemImageStatus.DONE,
-                    logfile=LOGFILE,
-                )
+                pass
+                # update_all_images(
+                #     history_item_id=history_item_id,
+                #     status=ImageStatus.DONE,
+                #     logfile=LOGFILE,
+                # )
             else:
-                update_all_images(
-                    history_item_id=history_item_id,
-                    status=HistoryItemImageStatus.FAILED,
-                    logfile=LOGFILE,
-                )
+                pass
+                # update_all_images(
+                #     history_item_id=history_item_id,
+                #     status=ImageStatus.FAILED,
+                #     logfile=LOGFILE,
+                # )
 
         return result, exception
 
@@ -514,10 +529,10 @@ class RunnerSlurmSudo(BaseRunner):
         list_parameters: list[dict],
         history_item_id: int,
         task_files: TaskFiles,
+        slurm_config: SlurmConfig,
         in_compound_task: bool = False,
-        **kwargs,
     ):
-        self.scancel_if_shutdown(active_slurm_jobs=[])
+        # self.scancel_jobs()
 
         self.validate_multisubmit_parameters(
             list_parameters=list_parameters,
@@ -527,82 +542,130 @@ class RunnerSlurmSudo(BaseRunner):
         workdir_local = task_files.wftask_subfolder_local
         workdir_remote = task_files.wftask_subfolder_remote
 
-        # Create folders
-        workdir_local.mkdir(parents=True, exist_ok=True)
-        _mkdir_as_user(
-            folder=workdir_remote.as_posix(),
-            user=self.slurm_user,
-        )
+        # Create local&remote task subfolders
+        if not in_compound_task:
+            original_umask = os.umask(0)
+            workdir_local.mkdir(parents=True, mode=0o755)
+            os.umask(original_umask)
+            _mkdir_as_user(
+                folder=workdir_remote.as_posix(),
+                user=self.slurm_user,
+            )
 
         # Execute tasks, in chunks of size `parallel_tasks_per_job`
         # TODO Pick a data structure for results and exceptions, or review the
         # interface
-        results = []
-        exceptions = []
-        jobs: dict[str, SlurmJob] = {}
+        results: dict[int, Any] = {}
+        exceptions: dict[int, BaseException] = {}
 
         original_task_files = task_files
-        # TODO: Add batching
-        for ind, parameters in enumerate(list_parameters):
-            # TODO: replace with actual values
+        tot_tasks = len(list_parameters)
+
+        # 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=slurm_config.tasks_per_job,
+            parallel_tasks_per_job=slurm_config.parallel_tasks_per_job,  # noqa
+            # Task requirements (multiple possible sources):
+            cpus_per_task=slurm_config.cpus_per_task,
+            mem_per_task=slurm_config.mem_per_task_MB,
+            # Fractal configuration variables (soft/hard limits):
+            target_cpus_per_job=slurm_config.target_cpus_per_job,
+            target_mem_per_job=slurm_config.target_mem_per_job,
+            target_num_jobs=slurm_config.target_num_jobs,
+            max_cpus_per_job=slurm_config.max_cpus_per_job,
+            max_mem_per_job=slurm_config.max_mem_per_job,
+            max_num_jobs=slurm_config.max_num_jobs,
+        )
+        slurm_config.parallel_tasks_per_job = parallel_tasks_per_job
+        slurm_config.tasks_per_job = tasks_per_job
+
+        # Divide arguments in batches of `tasks_per_job` tasks each
+        args_batches = []
+        batch_size = tasks_per_job
+        for ind_chunk in range(0, tot_tasks, batch_size):
+            args_batches.append(
+                list_parameters[ind_chunk : ind_chunk + batch_size]  # noqa
+            )
+        if len(args_batches) != math.ceil(tot_tasks / tasks_per_job):
+            raise RuntimeError("Something wrong here while batching tasks")
 
-            component = parameters[_COMPONENT_KEY_]
-            slurm_job = SlurmJob(
-                label=f"{ind:06d}",
-                workdir_local=workdir_local,
-                workdir_remote=workdir_remote,
-                tasks=[
+        logger.info(f"START submission phase, {list(self.jobs.keys())=}")
+        for ind_batch, chunk in enumerate(args_batches):
+            # TODO: replace with actual values
+            tasks = []
+            for ind_chunk, parameters in enumerate(chunk):
+                component = parameters[_COMPONENT_KEY_]
+                tasks.append(
                     SlurmTask(
+                        index=(ind_batch * batch_size) + ind_chunk,
                         component=component,
                         workdir_local=workdir_local,
                         workdir_remote=workdir_remote,
+                        parameters=parameters,
                         zarr_url=parameters["zarr_url"],
                         task_files=TaskFiles(
-                            **original_task_files,
+                            **original_task_files.model_dump(
+                                exclude={"component"}
+                            ),
                             component=component,
                         ),
-                    )
-                ],
+                    ),
+                )
+
+            slurm_job = SlurmJob(
+                label=f"{ind_batch:06d}",
+                workdir_local=workdir_local,
+                workdir_remote=workdir_remote,
+                tasks=tasks,
             )
-            slurm_job_id = self._submit_single_sbatch(
+            self._submit_single_sbatch(
                 func,
-                parameters=parameters,
                 slurm_job=slurm_job,
+                slurm_config=slurm_config,
             )
-            slurm_job.slurm_job_id = slurm_job_id
-            jobs[slurm_job_id] = slurm_job
+        logger.info(f"END submission phase, {list(self.jobs.keys())=}")
 
         # Retrieval phase
-        while len(jobs) > 0:
+        while len(self.jobs) > 0:
             if self.is_shutdown():
-                self.scancel_if_shutdown(active_slurm_jobs=jobs)
-            remaining_jobs = list(self.job_ids)
-            finished_jobs = get_finished_jobs(job_ids=remaining_jobs)
-            for slurm_job_id in finished_jobs:
-                slurm_job = jobs.pop(slurm_job_id)
+                self.scancel_jobs()
+            finished_job_ids = get_finished_jobs(job_ids=self.job_ids)
+            for slurm_job_id in finished_job_ids:
+                slurm_job = self.jobs.pop(slurm_job_id)
                 self._copy_files_from_remote_to_local(slurm_job)
                 for task in slurm_job.tasks:
                     result, exception = self._postprocess_single_task(
                         task=task
                     )
+                    if not in_compound_task:
+                        pass
+                        # update_single_image_logfile(
+                        #     history_item_id=history_item_id,
+                        #     zarr_url=task.zarr_url,
+                        #     logfile=task.task_files.log_file_local,
+                        # )
                     if not in_compound_task:
                         if exception is None:
-                            update_single_image(
-                                zarr_url=task.zarr_url,
-                                history_item_id=history_item_id,
-                                status=HistoryItemImageStatus.DONE,
-                                logfile=task.task_files.log_file_local,
-                            )
+                            pass
+                            # update_single_image(
+                            #     zarr_url=task.zarr_url,
+                            #     history_item_id=history_item_id,
+                            #     status=ImageStatus.DONE,
+                            # )
                         else:
-                            update_single_image(
-                                zarr_url=task.zarr_url,
-                                history_item_id=history_item_id,
-                                status=HistoryItemImageStatus.FAILED,
-                                logfile=task.task_files.log_file_local,
-                            )
-                    # TODO: Now just appending, but this should be done better
-                    results.append(result)
-                    exceptions.append(exception)
+                            pass
+                            # update_single_image(
+                            #     zarr_url=task.zarr_url,
+                            #     history_item_id=history_item_id,
+                            #     status=ImageStatus.FAILED,
+                            # )
+                    if exception is None:
+                        results[task.index] = result
+                    else:
+                        exceptions[task.index] = exception
             time.sleep(self.slurm_poll_interval)
         return results, exceptions
 

fractal_server/app/runner/v2/__init__.py

@@ -2,7 +2,7 @@
 Runner backend subsystem root V2
 
 This module is the single entry point to the runner backend subsystem V2.
-Other subystems should only import this module and not its submodules or
+Other subsystems should only import this module and not its submodules or
 the individual backends.
 """
 import os
@@ -29,7 +29,6 @@ from ..exceptions import JobExecutionError
 from ..exceptions import TaskExecutionError
 from ..executors.slurm_sudo._subprocess_run_as_user import _mkdir_as_user
 from ..filenames import WORKFLOW_LOG_FILENAME
-from ..task_files import task_subfolder_name
 from ._local import process_workflow as local_process_workflow
 from ._slurm_ssh import process_workflow as slurm_ssh_process_workflow
 from ._slurm_sudo import process_workflow as slurm_sudo_process_workflow
@@ -118,7 +117,7 @@ def submit_workflow(
             )
         except Exception as e:
             logger.error(
-                f"Error conneting to the database. Original error: {str(e)}"
+                f"Error connecting to the database. Original error: {str(e)}"
             )
             reset_logger_handlers(logger)
             return
@@ -200,25 +199,6 @@ def submit_workflow(
                     f"{settings.FRACTAL_RUNNER_BACKEND}."
                 )
 
-            # Create all tasks subfolders # FIXME: do this with Runner
-            for order in range(job.first_task_index, job.last_task_index + 1):
-                this_wftask = workflow.task_list[order]
-                task_name = this_wftask.task.name
-                subfolder_name = task_subfolder_name(
-                    order=order,
-                    task_name=task_name,
-                )
-                if FRACTAL_RUNNER_BACKEND == "slurm":
-                    # Create local subfolder (with 755) and remote one
-                    # (via `sudo -u`)
-                    original_umask = os.umask(0)
-                    (WORKFLOW_DIR_LOCAL / subfolder_name).mkdir(mode=0o755)
-                    os.umask(original_umask)
-                    _mkdir_as_user(
-                        folder=str(WORKFLOW_DIR_REMOTE / subfolder_name),
-                        user=slurm_user,
-                    )
-
         except Exception as e:
             error_type = type(e).__name__
             fail_job(

fractal_server/app/runner/v2/_slurm_ssh.py

@@ -11,7 +11,7 @@
 # Institute for Biomedical Research and Pelkmans Lab from the University of
 # Zurich.
 """
-Slurm Bakend
+Slurm Backend
 
 This backend runs fractal workflows in a SLURM cluster using Clusterfutures
 Executor objects.

fractal_server/app/runner/v2/_slurm_sudo.py

@@ -11,7 +11,7 @@
 # Institute for Biomedical Research and Pelkmans Lab from the University of
 # Zurich.
 """
-Slurm Bakend
+Slurm Backend
 
 This backend runs fractal workflows in a SLURM cluster using Clusterfutures
 Executor objects.