fractal-server 2.0.6__py3-none-any.whl → 2.2.0__py3-none-any.whl

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

@@ -110,12 +110,14 @@ class SlurmJob:
         slurm_file_prefix:
             Prefix for SLURM-job related files (submission script and SLURM
             stdout/stderr); this is also needed in the
-            `_copy_files_from_user_to_server` method.
+            `_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 `FractalSlurmExecutor`; this is
-            needed in the `_copy_files_from_user_to_server` method, and also to
-            construct the names of per-task input/output pickle files.
+            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:
@@ -145,6 +147,7 @@ class SlurmJob:
     # Per-task attributes
     workerids: tuple[str, ...]
     wftask_file_prefixes: tuple[str, ...]
+    wftask_subfolder_name: str
     input_pickle_files: tuple[Path, ...]
     output_pickle_files: tuple[Path, ...]
     # Slurm configuration
@@ -196,9 +199,9 @@ class FractalSlurmExecutor(SlurmExecutor):
         common_script_lines:
             Arbitrary script lines that will always be included in the
             sbatch script
-        working_dir:
+        workflow_dir_local:
             Directory for both the cfut/SLURM and fractal-server files and logs
-        working_dir_user:
+        workflow_dir_remote:
             Directory for both the cfut/SLURM and fractal-server files and logs
         map_jobid_to_slurm_files:
             Dictionary with paths of slurm-related files for active jobs
@@ -209,17 +212,18 @@ class FractalSlurmExecutor(SlurmExecutor):
     shutdown_file: str
     common_script_lines: list[str]
     user_cache_dir: str
-    working_dir: Path
-    working_dir_user: Path
+    workflow_dir_local: Path
+    workflow_dir_remote: Path
     map_jobid_to_slurm_files: dict[str, tuple[str, str, str]]
     keep_pickle_files: bool
     slurm_account: Optional[str]
+    jobs: dict[str, tuple[Future, SlurmJob]]
 
     def __init__(
         self,
         slurm_user: str,
-        working_dir: Path,
-        working_dir_user: Path,
+        workflow_dir_local: Path,
+        workflow_dir_remote: Path,
         shutdown_file: Optional[str] = None,
         user_cache_dir: Optional[str] = None,
         common_script_lines: Optional[list[str]] = None,
@@ -262,14 +266,14 @@ class FractalSlurmExecutor(SlurmExecutor):
         except StopIteration:
             pass
 
-        self.working_dir = working_dir
+        self.workflow_dir_local = workflow_dir_local
         if not _path_exists_as_user(
-            path=str(working_dir_user), user=self.slurm_user
+            path=str(workflow_dir_remote), user=self.slurm_user
         ):
-            logger.info(f"Missing folder {working_dir_user=}")
+            logger.info(f"Missing folder {workflow_dir_remote=}")
         self.user_cache_dir = user_cache_dir
 
-        self.working_dir_user = working_dir_user
+        self.workflow_dir_remote = workflow_dir_remote
         self.map_jobid_to_slurm_files = {}
 
         # Set the attribute slurm_poll_interval for self.wait_thread (see
@@ -281,7 +285,8 @@ class FractalSlurmExecutor(SlurmExecutor):
         self.wait_thread.slurm_user = self.slurm_user
 
         self.wait_thread.shutdown_file = (
-            shutdown_file or (self.working_dir / SHUTDOWN_FILENAME).as_posix()
+            shutdown_file
+            or (self.workflow_dir_local / SHUTDOWN_FILENAME).as_posix()
         )
         self.wait_thread.shutdown_callback = self.shutdown
 
@@ -294,32 +299,64 @@ class FractalSlurmExecutor(SlurmExecutor):
             self.map_jobid_to_slurm_files.pop(jobid)
 
     def get_input_pickle_file_path(
-        self, arg: str, prefix: Optional[str] = None
+        self, *, arg: str, subfolder_name: str, prefix: Optional[str] = None
     ) -> Path:
+
         prefix = prefix or "cfut"
-        return self.working_dir / f"{prefix}_in_{arg}.pickle"
+        output = (
+            self.workflow_dir_local
+            / subfolder_name
+            / f"{prefix}_in_{arg}.pickle"
+        )
+        return output
 
     def get_output_pickle_file_path(
-        self, arg: str, prefix: Optional[str] = None
+        self, *, arg: str, subfolder_name: str, prefix: Optional[str] = None
     ) -> Path:
         prefix = prefix or "cfut"
-        return self.working_dir_user / f"{prefix}_out_{arg}.pickle"
+        return (
+            self.workflow_dir_remote
+            / subfolder_name
+            / f"{prefix}_out_{arg}.pickle"
+        )
 
-    def get_slurm_script_file_path(self, prefix: Optional[str] = None) -> Path:
+    def get_slurm_script_file_path(
+        self, *, subfolder_name: str, prefix: Optional[str] = None
+    ) -> Path:
         prefix = prefix or "_temp"
-        return self.working_dir / f"{prefix}_slurm_submit.sbatch"
+        return (
+            self.workflow_dir_local
+            / subfolder_name
+            / f"{prefix}_slurm_submit.sbatch"
+        )
 
     def get_slurm_stdout_file_path(
-        self, arg: str = "%j", prefix: Optional[str] = None
+        self,
+        *,
+        subfolder_name: str,
+        arg: str = "%j",
+        prefix: Optional[str] = None,
     ) -> Path:
         prefix = prefix or "slurmpy.stdout"
-        return self.working_dir_user / f"{prefix}_slurm_{arg}.out"
+        return (
+            self.workflow_dir_remote
+            / subfolder_name
+            / f"{prefix}_slurm_{arg}.out"
+        )
 
     def get_slurm_stderr_file_path(
-        self, arg: str = "%j", prefix: Optional[str] = None
+        self,
+        *,
+        subfolder_name: str,
+        arg: str = "%j",
+        prefix: Optional[str] = None,
     ) -> Path:
         prefix = prefix or "slurmpy.stderr"
-        return self.working_dir_user / f"{prefix}_slurm_{arg}.err"
+        return (
+            self.workflow_dir_remote
+            / subfolder_name
+            / f"{prefix}_slurm_{arg}.err"
+        )
 
     def submit(
         self,
@@ -599,6 +636,8 @@ class FractalSlurmExecutor(SlurmExecutor):
                 )
             job.single_task_submission = True
             job.wftask_file_prefixes = (task_files.file_prefix,)
+            job.wftask_subfolder_name = task_files.subfolder_name
+
         else:
             if not components or len(components) < 1:
                 raise ValueError(
@@ -613,33 +652,60 @@ class FractalSlurmExecutor(SlurmExecutor):
             )
 
             _prefixes = []
+            _subfolder_names = []
             for component in components:
                 if isinstance(component, dict):
                     # This is needed for V2
                     actual_component = component.get(_COMPONENT_KEY_, None)
                 else:
                     actual_component = component
-                _prefixes.append(
-                    get_task_file_paths(
-                        workflow_dir=task_files.workflow_dir,
-                        workflow_dir_user=task_files.workflow_dir_user,
-                        task_order=task_files.task_order,
-                        component=actual_component,
-                    ).file_prefix
+                _task_file_paths = get_task_file_paths(
+                    workflow_dir_local=task_files.workflow_dir_local,
+                    workflow_dir_remote=task_files.workflow_dir_remote,
+                    task_name=task_files.task_name,
+                    task_order=task_files.task_order,
+                    component=actual_component,
                 )
+                _prefixes.append(_task_file_paths.file_prefix)
+                _subfolder_names.append(_task_file_paths.subfolder_name)
             job.wftask_file_prefixes = tuple(_prefixes)
 
+            num_subfolders = len(set(_subfolder_names))
+            if num_subfolders != 1:
+                error_msg_short = (
+                    f"[_submit_job] Subfolder list has {num_subfolders} "
+                    "different values, but it must have only one (since "
+                    "workflow tasks are executed one by one)."
+                )
+                error_msg_detail = (
+                    "[_submit_job] Current unique subfolder names: "
+                    f"{set(_subfolder_names)}"
+                )
+                logger.error(error_msg_short)
+                logger.error(error_msg_detail)
+                raise ValueError(error_msg_short)
+            job.wftask_subfolder_name = _subfolder_names[0]
+
+        # Check that server-side subfolder exists
+        subfolder_path = self.workflow_dir_local / job.wftask_subfolder_name
+        if not subfolder_path.exists():
+            raise FileNotFoundError(
+                f"Missing folder {subfolder_path.as_posix()}."
+            )
+
         # Define I/O pickle file names/paths
         job.input_pickle_files = tuple(
             self.get_input_pickle_file_path(
-                job.workerids[ind],
+                arg=job.workerids[ind],
+                subfolder_name=job.wftask_subfolder_name,
                 prefix=job.wftask_file_prefixes[ind],
             )
             for ind in range(job.num_tasks_tot)
         )
         job.output_pickle_files = tuple(
             self.get_output_pickle_file_path(
-                job.workerids[ind],
+                arg=job.workerids[ind],
+                subfolder_name=job.wftask_subfolder_name,
                 prefix=job.wftask_file_prefixes[ind],
             )
             for ind in range(job.num_tasks_tot)
@@ -647,13 +713,16 @@ class FractalSlurmExecutor(SlurmExecutor):
 
         # Define SLURM-job file names/paths
         job.slurm_script = self.get_slurm_script_file_path(
-            prefix=job.slurm_file_prefix
+            subfolder_name=job.wftask_subfolder_name,
+            prefix=job.slurm_file_prefix,
         )
         job.slurm_stdout = self.get_slurm_stdout_file_path(
-            prefix=job.slurm_file_prefix
+            subfolder_name=job.wftask_subfolder_name,
+            prefix=job.slurm_file_prefix,
         )
         job.slurm_stderr = self.get_slurm_stderr_file_path(
-            prefix=job.slurm_file_prefix
+            subfolder_name=job.wftask_subfolder_name,
+            prefix=job.slurm_file_prefix,
         )
 
         # Dump serialized versions+function+args+kwargs to pickle file
@@ -706,10 +775,10 @@ class FractalSlurmExecutor(SlurmExecutor):
         Prepare the `JobExecutionError` for a given job
 
         This method creates a `JobExecutionError` object and sets its attribute
-        to the appropriate SLURM-related file names. Note that the method shoul
-        always be called after values in `self.map_jobid_to_slurm_files` have
-        been updated, so that they point to `self.working_dir` files which are
-        readable from `fractal-server`.
+        to the appropriate SLURM-related file names. Note that the method
+        should always be called after values in `self.map_jobid_to_slurm_files`
+        have been updated, so that they point to `self.workflow_dir_local`
+        files which are readable from `fractal-server`.
 
         Arguments:
             jobid:
@@ -758,13 +827,13 @@ class FractalSlurmExecutor(SlurmExecutor):
                 if not self.jobs:
                     self.jobs_empty_cond.notify_all()
 
-            # Copy all relevant files from self.working_dir_user to
-            # self.working_dir
+            # Copy all relevant files from self.workflow_dir_remote to
+            # self.workflow_dir_local
 
-            self._copy_files_from_user_to_server(job)
+            self._copy_files_from_remote_to_local(job)
 
-            # Update the paths to use the files in self.working_dir (rather
-            # than the user's ones in self.working_dir_user)
+            # Update the paths to use the files in self.workflow_dir_local
+            # (rather than the user's ones in self.workflow_dir_remote)
             with self.jobs_lock:
                 self.map_jobid_to_slurm_files[jobid]
                 (
@@ -773,10 +842,14 @@ class FractalSlurmExecutor(SlurmExecutor):
                     slurm_stderr_file,
                 ) = self.map_jobid_to_slurm_files[jobid]
             new_slurm_stdout_file = str(
-                self.working_dir / Path(slurm_stdout_file).name
+                self.workflow_dir_local
+                / job.wftask_subfolder_name
+                / Path(slurm_stdout_file).name
             )
             new_slurm_stderr_file = str(
-                self.working_dir / Path(slurm_stderr_file).name
+                self.workflow_dir_local
+                / job.wftask_subfolder_name
+                / Path(slurm_stderr_file).name
             )
             with self.jobs_lock:
                 self.map_jobid_to_slurm_files[jobid] = (
@@ -787,7 +860,8 @@ class FractalSlurmExecutor(SlurmExecutor):
 
             in_paths = job.input_pickle_files
             out_paths = tuple(
-                self.working_dir / f.name for f in job.output_pickle_files
+                (self.workflow_dir_local / job.wftask_subfolder_name / f.name)
+                for f in job.output_pickle_files
             )
 
             outputs = []
@@ -908,16 +982,16 @@ class FractalSlurmExecutor(SlurmExecutor):
                     " FractalSlurmExecutor._completion."
                 )
 
-    def _copy_files_from_user_to_server(
+    def _copy_files_from_remote_to_local(
         self,
         job: SlurmJob,
     ):
         """
         Impersonate the user and copy task-related files
 
-        For all files in `self.working_dir_user` that start with
+        For all files in `self.workflow_dir_remote` that start with
         `job.file_prefix`, read them (with `sudo -u` impersonation) and write
-        them to `self.working_dir`.
+        them to `self.workflow_dir_local`.
 
         Files to copy:
         * Job-related files (SLURM stderr/stdout files); with prefix
@@ -932,36 +1006,48 @@ class FractalSlurmExecutor(SlurmExecutor):
         Raises:
             JobExecutionError: If a `cat` command fails.
         """
-        logger.debug("Enter _copy_files_from_user_to_server")
-        if self.working_dir_user == self.working_dir:
+        logger.debug("[_copy_files_from_remote_to_local] Start")
+
+        if self.workflow_dir_remote == self.workflow_dir_local:
+            logger.debug(
+                "[_copy_files_from_remote_to_local] "
+                "workflow_dir_local corresponds to workflow_dir_remote, "
+                "return."
+            )
             return
 
+        subfolder_name = job.wftask_subfolder_name
         prefixes = set(
             [job.slurm_file_prefix] + list(job.wftask_file_prefixes)
         )
 
-        logger.debug(f"[_copy_files_from_user_to_server] {prefixes=}")
         logger.debug(
-            f"[_copy_files_from_user_to_server] {str(self.working_dir_user)=}"
+            "[_copy_files_from_remote_to_local] "
+            f"WorkflowTask subfolder_name: {subfolder_name}"
+        )
+        logger.debug(f"[_copy_files_from_remote_to_local] {prefixes=}")
+        logger.debug(
+            "[_copy_files_from_remote_to_local] "
+            f"{str(self.workflow_dir_remote)=}"
         )
 
         for prefix in prefixes:
 
             if prefix == job.slurm_file_prefix:
                 files_to_copy = _glob_as_user(
-                    folder=str(self.working_dir_user),
+                    folder=str(self.workflow_dir_remote / subfolder_name),
                     user=self.slurm_user,
                     startswith=prefix,
                 )
             else:
                 files_to_copy = _glob_as_user_strict(
-                    folder=str(self.working_dir_user),
+                    folder=str(self.workflow_dir_remote / subfolder_name),
                     user=self.slurm_user,
                     startswith=prefix,
                 )
 
             logger.debug(
-                "[_copy_files_from_user_to_server] "
+                "[_copy_files_from_remote_to_local] "
                 f"{prefix=}, {len(files_to_copy)=}"
             )
 
@@ -972,7 +1058,9 @@ class FractalSlurmExecutor(SlurmExecutor):
                         "contains whitespaces"
                     )
                 source_file_path = str(
-                    self.working_dir_user / source_file_name
+                    self.workflow_dir_remote
+                    / subfolder_name
+                    / source_file_name
                 )
 
                 # Read source_file_path (requires sudo)
@@ -991,10 +1079,12 @@ class FractalSlurmExecutor(SlurmExecutor):
                     logger.error(info)
                     raise JobExecutionError(info)
                 # Write to dest_file_path (including empty files)
-                dest_file_path = str(self.working_dir / source_file_name)
+                dest_file_path = str(
+                    self.workflow_dir_local / subfolder_name / source_file_name
+                )
                 with open(dest_file_path, "wb") as f:
                     f.write(res.stdout)
-        logger.debug("[_copy_files_from_user_to_server] End")
+        logger.debug("[_copy_files_from_remote_to_local] End")
 
     def _start(
         self,
@@ -1099,7 +1189,7 @@ class FractalSlurmExecutor(SlurmExecutor):
             [
                 f"#SBATCH --err={slurm_err_path}",
                 f"#SBATCH --out={slurm_out_path}",
-                f"#SBATCH -D {self.working_dir_user}",
+                f"#SBATCH -D {self.workflow_dir_remote}",
             ]
         )
         script_lines = slurm_config.sort_script_lines(script_lines)
@@ -1131,12 +1221,11 @@ class FractalSlurmExecutor(SlurmExecutor):
         This will be called when self.submit or self.map are called from
         outside fractal-server, and then lack some optional arguments.
         """
-        import random
-
         task_files = TaskFiles(
-            workflow_dir=self.working_dir,
-            workflow_dir_user=self.working_dir_user,
-            task_order=random.randint(10000, 99999),  # nosec
+            workflow_dir_local=self.workflow_dir_local,
+            workflow_dir_remote=self.workflow_dir_remote,
+            task_order=None,
+            task_name="name",
         )
         return task_files
 
@@ -1154,7 +1243,7 @@ class FractalSlurmExecutor(SlurmExecutor):
             while self.jobs:
                 jobid, fut_and_job = self.jobs.popitem()
                 slurm_jobs_to_scancel.append(jobid)
-                fut, job = fut_and_job[:]
+                fut = fut_and_job[0]
                 self.map_jobid_to_slurm_files.pop(jobid)
                 if not fut.cancelled():
                     fut.set_exception(
@@ -1188,3 +1277,15 @@ class FractalSlurmExecutor(SlurmExecutor):
                 raise JobExecutionError(info=error_msg)
 
         logger.debug("Executor shutdown: end")
+
+    def __exit__(self, *args, **kwargs):
+        """
+        See
+        https://github.com/fractal-analytics-platform/fractal-server/issues/1508
+        """
+        logger.debug(
+            "[FractalSlurmExecutor.__exit__] Stop and join `wait_thread`"
+        )
+        self.wait_thread.stop()
+        self.wait_thread.join()
+        logger.debug("[FractalSlurmExecutor.__exit__] End")

fractal_server/app/runner/shutdown.py

@@ -0,0 +1,88 @@
+import time
+
+from sqlmodel import select
+
+from fractal_server.app.db import get_async_db
+from fractal_server.app.models.v1 import ApplyWorkflow
+from fractal_server.app.models.v1.job import JobStatusTypeV1
+from fractal_server.app.models.v2 import JobV2
+from fractal_server.app.models.v2.job import JobStatusTypeV2
+from fractal_server.app.routes.aux._job import _write_shutdown_file
+from fractal_server.config import get_settings
+from fractal_server.logger import get_logger
+from fractal_server.syringe import Inject
+
+
+async def cleanup_after_shutdown(
+    *, jobsV1: list[int], jobsV2: list[int], logger_name: str
+):
+    logger = get_logger(logger_name)
+    logger.info("Cleanup function after shutdown")
+    stm_v2 = (
+        select(JobV2)
+        .where(JobV2.id.in_(jobsV2))
+        .where(JobV2.status == JobStatusTypeV2.SUBMITTED)
+    )
+
+    stm_v1 = (
+        select(ApplyWorkflow)
+        .where(ApplyWorkflow.id.in_(jobsV1))
+        .where(ApplyWorkflow.status == JobStatusTypeV1.SUBMITTED)
+    )
+
+    async for session in get_async_db():
+        jobsV2_db = (await session.execute(stm_v2)).scalars().all()
+        jobsV1_db = (await session.execute(stm_v1)).scalars().all()
+
+        for job in jobsV2_db:
+            _write_shutdown_file(job=job)
+
+        for job in jobsV1_db:
+            _write_shutdown_file(job=job)
+
+        settings = Inject(get_settings)
+
+        t_start = time.perf_counter()
+        while (
+            time.perf_counter() - t_start
+        ) < settings.FRACTAL_GRACEFUL_SHUTDOWN_TIME:  # 30 seconds
+            logger.info("Waiting 3 seconds before checking")
+            time.sleep(3)
+            jobsV2_db = (await session.execute(stm_v2)).scalars().all()
+            jobsV1_db = (await session.execute(stm_v1)).scalars().all()
+
+            if len(jobsV2_db) == 0 and len(jobsV1_db) == 0:
+                logger.info(
+                    (
+                        "All jobs associated to this app are "
+                        "either done or failed. Exit."
+                    )
+                )
+                return
+            else:
+                logger.info(
+                    (
+                        f"Some jobs are still 'submitted' "
+                        f"{jobsV1_db=}, {jobsV2_db=}"
+                    )
+                )
+        logger.info(
+            (
+                "Graceful shutdown reached its maximum time, "
+                "but some jobs are still submitted"
+            )
+        )
+
+        for job in jobsV2_db:
+            job.status = "failed"
+            job.log = (job.log or "") + "\nJob stopped due to app shutdown\n"
+            session.add(job)
+        await session.commit()
+
+        for job in jobsV1_db:
+            job.status = "failed"
+            job.log = (job.log or "") + "\nJob stopped due to app shutdown\n"
+            session.add(job)
+        await session.commit()
+
+        logger.info("Exit from shutdown logic")