fractal-server 2.14.0a9__py3-none-any.whl → 2.14.0a11__py3-none-any.whl

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

@@ -6,19 +6,22 @@ import shlex
 import subprocess  # nosec
 import sys
 import time
+from copy import copy
 from pathlib import Path
 from typing import Any
+from typing import Literal
 from typing import Optional
 
 import cloudpickle
 from pydantic import BaseModel
 from pydantic import ConfigDict
 
-from ._check_jobs_status import get_finished_jobs
+from ..slurm_common._check_jobs_status import get_finished_jobs
+from ..slurm_common._check_jobs_status import run_squeue
 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.runner.components import _COMPONENT_KEY_
+from fractal_server.app.db import get_sync_db
 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
@@ -30,7 +33,8 @@ from fractal_server.app.runner.executors.slurm_common._slurm_config import (
 )
 from fractal_server.app.runner.filenames import SHUTDOWN_FILENAME
 from fractal_server.app.runner.task_files import TaskFiles
-from fractal_server.app.schemas.v2.task import TaskTypeType
+from fractal_server.app.runner.v2.db_tools import update_status_of_history_unit
+from fractal_server.app.schemas.v2 import HistoryUnitStatus
 from fractal_server.config import get_settings
 from fractal_server.logger import set_logger
 from fractal_server.syringe import Inject
@@ -97,40 +101,68 @@ class SlurmJob(BaseModel):
     tasks: list[SlurmTask]
 
     @property
-    def slurm_log_file_local(self) -> str:
+    def slurm_submission_script_local(self) -> str:
+        return (
+            self.workdir_local / f"slurm-{self.label}-submit.sh"
+        ).as_posix()
+
+    @property
+    def slurm_submission_script_remote(self) -> str:
+        return (
+            self.workdir_remote / f"slurm-{self.label}-submit.sh"
+        ).as_posix()
+
+    @property
+    def slurm_stdout_remote(self) -> str:
         if self.slurm_job_id:
             return (
-                self.workdir_local
-                / f"slurm-{self.label}-{self.slurm_job_id}.log"
+                self.workdir_remote
+                / f"slurm-{self.label}-{self.slurm_job_id}.out"
             ).as_posix()
+
         else:
             return (
-                self.workdir_local / f"slurm-{self.label}-%j.log"
+                self.workdir_remote / f"slurm-{self.label}-%j.out"
             ).as_posix()
 
     @property
-    def slurm_log_file_remote(self) -> str:
+    def slurm_stderr_remote(self) -> str:
         if self.slurm_job_id:
             return (
                 self.workdir_remote
-                / f"slurm-{self.label}-{self.slurm_job_id}.log"
+                / f"slurm-{self.label}-{self.slurm_job_id}.err"
             ).as_posix()
+
         else:
             return (
-                self.workdir_remote / f"slurm-{self.label}-%j.log"
+                self.workdir_remote / f"slurm-{self.label}-%j.err"
             ).as_posix()
 
     @property
-    def slurm_submission_script_local(self) -> str:
-        return (
-            self.workdir_local / f"slurm-{self.label}-submit.sh"
-        ).as_posix()
+    def slurm_stdout_local(self) -> str:
+        if self.slurm_job_id:
+            return (
+                self.workdir_local
+                / f"slurm-{self.label}-{self.slurm_job_id}.out"
+            ).as_posix()
+
+        else:
+            return (
+                self.workdir_local / f"slurm-{self.label}-%j.out"
+            ).as_posix()
 
     @property
-    def slurm_submission_script_remote(self) -> str:
-        return (
-            self.workdir_remote / f"slurm-{self.label}-submit.sh"
-        ).as_posix()
+    def slurm_stderr_local(self) -> str:
+        if self.slurm_job_id:
+            return (
+                self.workdir_local
+                / f"slurm-{self.label}-{self.slurm_job_id}.err"
+            ).as_posix()
+
+        else:
+            return (
+                self.workdir_local / f"slurm-{self.label}-%j.err"
+            ).as_posix()
 
     @property
     def log_files_local(self) -> list[str]:
@@ -276,9 +308,7 @@ class RunnerSlurmSudo(BaseRunner):
         slurm_job: SlurmJob,
         slurm_config: SlurmConfig,
     ) -> str:
-        # if len(slurm_job.tasks) > 1:
-        #     raise NotImplementedError()
-
+        logger.debug("[_submit_single_sbatch] START")
         # Prepare input pickle(s)
         versions = dict(
             python=sys.version_info[:3],
@@ -287,44 +317,87 @@ class RunnerSlurmSudo(BaseRunner):
         )
         for task in slurm_job.tasks:
             _args = []
-            _kwargs = dict(parameters=task.parameters)
+            _kwargs = dict(
+                parameters=task.parameters,
+                remote_files=task.task_files.remote_files_dict,
+            )
             funcser = cloudpickle.dumps((versions, func, _args, _kwargs))
             with open(task.input_pickle_file_local, "wb") as f:
                 f.write(funcser)
-
+            logger.debug(
+                "[_submit_single_sbatch] Written "
+                f"{task.input_pickle_file_local=}"
+            )
         # Prepare commands to be included in SLURM submission script
-
-        preamble_lines = [
-            "#!/bin/bash",
-            "#SBATCH --partition=main",
-            "#SBATCH --ntasks=1",
-            "#SBATCH --cpus-per-task=1",
-            "#SBATCH --mem=10M",
-            f"#SBATCH --err={slurm_job.slurm_log_file_remote}",
-            f"#SBATCH --out={slurm_job.slurm_log_file_remote}",
-            f"#SBATCH -D {slurm_job.workdir_remote}",
-            "#SBATCH --job-name=test",
-            "\n",
-        ]
-
+        settings = Inject(get_settings)
+        python_worker_interpreter = (
+            settings.FRACTAL_SLURM_WORKER_PYTHON or sys.executable
+        )
         cmdlines = []
         for task in slurm_job.tasks:
-            cmd = (
-                f"{self.python_worker_interpreter}"
-                " -m fractal_server.app.runner.executors.slurm_common.remote "
-                f"--input-file {task.input_pickle_file_local} "
-                f"--output-file {task.output_pickle_file_remote}"
-            )
-            cmdlines.append("whoami")
+            input_pickle_file = task.input_pickle_file_local
+            output_pickle_file = task.output_pickle_file_remote
             cmdlines.append(
-                f"srun --ntasks=1 --cpus-per-task=1 --mem=10MB {cmd} &"
+                (
+                    f"{python_worker_interpreter}"
+                    " -m fractal_server.app.runner."
+                    "executors.slurm_common.remote "
+                    f"--input-file {input_pickle_file} "
+                    f"--output-file {output_pickle_file}"
+                )
             )
-        cmdlines.append("wait\n")
+
+        # ...
+        num_tasks_max_running = slurm_config.parallel_tasks_per_job
+        mem_per_task_MB = slurm_config.mem_per_task_MB
+
+        # Set ntasks
+        ntasks = min(len(cmdlines), num_tasks_max_running)
+        slurm_config.parallel_tasks_per_job = ntasks
+
+        # Prepare SLURM preamble based on SlurmConfig object
+        script_lines = slurm_config.to_sbatch_preamble(
+            remote_export_dir=self.user_cache_dir
+        )
+
+        # Extend SLURM preamble with variable which are not in SlurmConfig, and
+        # fix their order
+        script_lines.extend(
+            [
+                f"#SBATCH --out={slurm_job.slurm_stdout_remote}",
+                f"#SBATCH --err={slurm_job.slurm_stderr_remote}",
+                f"#SBATCH -D {slurm_job.workdir_remote}",
+            ]
+        )
+        script_lines = slurm_config.sort_script_lines(script_lines)
+        logger.debug(script_lines)
+
+        # Always print output of `uname -n` and `pwd`
+        script_lines.append(
+            '"Hostname: `uname -n`; current directory: `pwd`"\n'
+        )
+
+        # Complete script preamble
+        script_lines.append("\n")
+
+        # Include command lines
+        tmp_list_commands = copy(cmdlines)
+        while tmp_list_commands:
+            if tmp_list_commands:
+                cmd = tmp_list_commands.pop(0)  # take first element
+                script_lines.append(
+                    "srun --ntasks=1 --cpus-per-task=$SLURM_CPUS_PER_TASK "
+                    f"--mem={mem_per_task_MB}MB "
+                    f"{cmd} &"
+                )
+        script_lines.append("wait\n")
+
+        script = "\n".join(script_lines)
 
         # Write submission script
-        submission_script_contents = "\n".join(preamble_lines + cmdlines)
+        # submission_script_contents = "\n".join(preamble_lines + cmdlines)
         with open(slurm_job.slurm_submission_script_local, "w") as f:
-            f.write(submission_script_contents)
+            f.write(script)
 
         # Run sbatch
         pre_command = f"sudo --set-home --non-interactive -u {self.slurm_user}"
@@ -350,8 +423,10 @@ class RunnerSlurmSudo(BaseRunner):
         """
         Note: this would differ for SSH
         """
+        logger.debug(f"[_copy_files_from_remote_to_local] {job.slurm_job_id=}")
         source_target_list = [
-            (job.slurm_log_file_remote, job.slurm_log_file_local)
+            (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(
@@ -419,21 +494,22 @@ class RunnerSlurmSudo(BaseRunner):
         self,
         func: callable,
         parameters: dict[str, Any],
-        history_item_id: int,
+        history_unit_id: int,
         task_files: TaskFiles,
-        slurm_config: SlurmConfig,
-        task_type: TaskTypeType,
+        task_type: Literal[
+            "non_parallel",
+            "converter_non_parallel",
+            "compound",
+            "converter_compound",
+        ],
+        config: SlurmConfig,
     ) -> tuple[Any, Exception]:
-        workdir_local = task_files.wftask_subfolder_local
-        workdir_remote = task_files.wftask_subfolder_remote
 
-        task_files = TaskFiles(
-            **task_files.model_dump(
-                exclude={"component"},
-            ),
-            component=parameters[_COMPONENT_KEY_],
-        )
+        if len(self.jobs) > 0:
+            raise RuntimeError(f"Cannot run .submit when {len(self.jobs)=}")
 
+        workdir_local = task_files.wftask_subfolder_local
+        workdir_remote = task_files.wftask_subfolder_remote
         if self.jobs != {}:
             raise JobExecutionError("Unexpected branch: jobs should be empty.")
 
@@ -441,7 +517,7 @@ class RunnerSlurmSudo(BaseRunner):
             raise JobExecutionError("Cannot continue after shutdown.")
 
         # Validation phase
-        self.validate_submit_parameters(parameters)
+        self.validate_submit_parameters(parameters, task_type=task_type)
 
         # Create task subfolder
         original_umask = os.umask(0)
@@ -460,7 +536,7 @@ class RunnerSlurmSudo(BaseRunner):
             tasks=[
                 SlurmTask(
                     index=0,
-                    component="0",
+                    component=task_files.component,
                     parameters=parameters,
                     workdir_remote=workdir_remote,
                     workdir_local=workdir_local,
@@ -468,25 +544,57 @@ class RunnerSlurmSudo(BaseRunner):
                 )
             ],
         )  # TODO: replace with actual values (BASED ON TASKFILES)
+
+        config.parallel_tasks_per_job = 1
         self._submit_single_sbatch(
             func,
             slurm_job=slurm_job,
-            slurm_config=slurm_config,
+            slurm_config=config,
         )
-
-        # LOGFILE = task_files.log_file_local
+        logger.debug("END SUBMISSION PHASE")
+        logger.debug(f"{self.jobs=}")
+        logger.debug(f"{self.job_ids=}")
+
+        # FIXME
+        jobs_that_started = set()
+        while len(jobs_that_started) != len(self.job_ids):
+            logger.debug("CALL SQUEUE")
+            res = run_squeue(self.job_ids)
+            new_jobs = set(out.split()[0] for out in res.stdout.splitlines())
+            jobs_that_started = jobs_that_started.union(new_jobs)
+            logger.debug(f"{new_jobs=}")
+            logger.debug(f"{len(jobs_that_started)=}")
+
+        logger.debug("START RETRIEVAL PHASE")
 
         # Retrieval phase
         while len(self.jobs) > 0:
             if self.is_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)
-                self._copy_files_from_remote_to_local(slurm_job)
-                result, exception = self._postprocess_single_task(
-                    task=slurm_job.tasks[0]
-                )
+            logger.debug(f"{finished_job_ids=}")
+            with next(get_sync_db()) as db:
+                for slurm_job_id in finished_job_ids:
+                    logger.debug(f"Now process {slurm_job_id=}")
+                    slurm_job = self.jobs.pop(slurm_job_id)
+                    self._copy_files_from_remote_to_local(slurm_job)
+                    result, exception = self._postprocess_single_task(
+                        task=slurm_job.tasks[0]
+                    )
+                    if result is not None:
+                        if task_type not in ["compound", "converter_compound"]:
+                            update_status_of_history_unit(
+                                history_unit_id=history_unit_id,
+                                status=HistoryUnitStatus.DONE,
+                                db_sync=db,
+                            )
+                    if exception is not None:
+                        update_status_of_history_unit(
+                            history_unit_id=history_unit_id,
+                            status=HistoryUnitStatus.FAILED,
+                            db_sync=db,
+                        )
+
             time.sleep(self.slurm_poll_interval)
 
         return result, exception
@@ -495,19 +603,38 @@ class RunnerSlurmSudo(BaseRunner):
         self,
         func: callable,
         list_parameters: list[dict],
-        history_item_id: int,
-        task_files: TaskFiles,
-        slurm_config: SlurmConfig,
-        task_type: TaskTypeType,
+        history_unit_ids: list[int],
+        list_task_files: list[TaskFiles],
+        task_type: Literal["parallel", "compound", "converter_compound"],
+        config: SlurmConfig,
     ):
-        # self.scancel_jobs()
+
+        if len(self.jobs) > 0:
+            raise RuntimeError(f"Cannot run .submit when {len(self.jobs)=}")
+
+        if task_type in ["compound", "converter_compound"]:
+            if len(history_unit_ids) != 1:
+                raise NotImplementedError(
+                    "We are breaking the assumption that compound/multisubmit "
+                    "is associated to a single HistoryUnit. This is not "
+                    "supported."
+                )
+        elif task_type == "parallel" and len(history_unit_ids) != len(
+            list_parameters
+        ):
+            raise ValueError(
+                f"{len(history_unit_ids)=} differs from "
+                f"{len(list_parameters)=}."
+            )
 
         self.validate_multisubmit_parameters(
-            list_parameters=list_parameters, task_type=task_type
+            list_parameters=list_parameters,
+            task_type=task_type,
+            list_task_files=list_task_files,
         )
 
-        workdir_local = task_files.wftask_subfolder_local
-        workdir_remote = task_files.wftask_subfolder_remote
+        workdir_local = list_task_files[0].wftask_subfolder_local
+        workdir_remote = list_task_files[0].wftask_subfolder_remote
 
         # Create local&remote task subfolders
         if task_type not in ["converter_compound", "compound"]:
@@ -525,7 +652,7 @@ class RunnerSlurmSudo(BaseRunner):
         results: dict[int, Any] = {}
         exceptions: dict[int, BaseException] = {}
 
-        original_task_files = task_files
+        original_task_files = list_task_files
         tot_tasks = len(list_parameters)
 
         # Set/validate parameters for task batching
@@ -533,21 +660,21 @@ class RunnerSlurmSudo(BaseRunner):
             # 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
+            tasks_per_job=config.tasks_per_job,
+            parallel_tasks_per_job=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,
+            cpus_per_task=config.cpus_per_task,
+            mem_per_task=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,
+            target_cpus_per_job=config.target_cpus_per_job,
+            target_mem_per_job=config.target_mem_per_job,
+            target_num_jobs=config.target_num_jobs,
+            max_cpus_per_job=config.max_cpus_per_job,
+            max_mem_per_job=config.max_mem_per_job,
+            max_num_jobs=config.max_num_jobs,
         )
-        slurm_config.parallel_tasks_per_job = parallel_tasks_per_job
-        slurm_config.tasks_per_job = tasks_per_job
+        config.parallel_tasks_per_job = parallel_tasks_per_job
+        config.tasks_per_job = tasks_per_job
 
         # Divide arguments in batches of `tasks_per_job` tasks each
         args_batches = []
@@ -561,24 +688,18 @@ class RunnerSlurmSudo(BaseRunner):
 
         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_]
+                index = (ind_batch * batch_size) + ind_chunk
                 tasks.append(
                     SlurmTask(
-                        index=(ind_batch * batch_size) + ind_chunk,
-                        component=component,
+                        index=index,
+                        component=original_task_files[index].component,
                         workdir_local=workdir_local,
                         workdir_remote=workdir_remote,
                         parameters=parameters,
                         zarr_url=parameters["zarr_url"],
-                        task_files=TaskFiles(
-                            **original_task_files.model_dump(
-                                exclude={"component"}
-                            ),
-                            component=component,
-                        ),
+                        task_files=original_task_files[index],
                     ),
                 )
 
@@ -591,26 +712,56 @@ class RunnerSlurmSudo(BaseRunner):
             self._submit_single_sbatch(
                 func,
                 slurm_job=slurm_job,
-                slurm_config=slurm_config,
+                slurm_config=config,
             )
         logger.info(f"END submission phase, {list(self.jobs.keys())=}")
 
+        # FIXME
+        jobs_that_started = set()
+        while len(jobs_that_started) != len(self.job_ids):
+            res = run_squeue(self.job_ids)
+            new_jobs = set(out.split()[0] for out in res.stdout.splitlines())
+            jobs_that_started = jobs_that_started.union(new_jobs)
+            logger.debug(f"{new_jobs=}")
+            logger.debug(f"{len(jobs_that_started)=}")
+
         # Retrieval phase
         while len(self.jobs) > 0:
             if self.is_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)
-                self._copy_files_from_remote_to_local(slurm_job)
-                for task in slurm_job.tasks:
-                    result, exception = self._postprocess_single_task(
-                        task=task
-                    )
-                    if exception is None:
-                        results[task.index] = result
-                    else:
-                        exceptions[task.index] = exception
+            logger.debug(f"{finished_job_ids=}")
+            with next(get_sync_db()) as db:
+                for slurm_job_id in finished_job_ids:
+                    logger.debug(f"Now processing {slurm_job_id=}")
+                    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 result is not None:
+                            results[task.index] = result
+                            if task_type == "parallel":
+                                update_status_of_history_unit(
+                                    history_unit_id=history_unit_ids[
+                                        task.index
+                                    ],
+                                    status=HistoryUnitStatus.DONE,
+                                    db_sync=db,
+                                )
+                        if exception is not None:
+                            exceptions[task.index] = exception
+                            if task_type == "parallel":
+                                update_status_of_history_unit(
+                                    history_unit_id=history_unit_ids[
+                                        task.index
+                                    ],
+                                    status=HistoryUnitStatus.FAILED,
+                                    db_sync=db,
+                                )
+
             time.sleep(self.slurm_poll_interval)
         return results, exceptions
 

fractal_server/app/runner/task_files.py

@@ -96,3 +96,11 @@ class TaskFiles(BaseModel):
         return (
             self.wftask_subfolder_remote / f"{self.component}-metadiff.json"
         ).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,
+        )