fractal-server 2.14.4a0__py3-none-any.whl → 2.14.6__py3-none-any.whl

fractal_server/app/runner/executors/slurm_common/base_slurm_runner.py

@@ -5,9 +5,9 @@ import time
 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 ..slurm_common._slurm_config import SlurmConfig
 from ..slurm_common.slurm_job_task_models import SlurmJob
@@ -36,6 +36,17 @@ SHUTDOWN_EXCEPTION = JobExecutionError(SHUTDOWN_ERROR_MESSAGE)
 logger = set_logger(__name__)
 
 
+class RemoteInputData(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    python_version: tuple[int, int, int]
+    fractal_server_version: str
+    full_command: str
+
+    metadiff_file_remote: str
+    log_file_remote: str
+
+
 def create_accounting_record_slurm(
     *,
     user_id: int,
@@ -69,9 +80,9 @@ class BaseSlurmRunner(BaseRunner):
         root_dir_remote: Path,
         slurm_runner_type: Literal["ssh", "sudo"],
         python_worker_interpreter: str,
-        common_script_lines: Optional[list[str]] = None,
-        user_cache_dir: Optional[str] = None,
-        poll_interval: Optional[int] = None,
+        common_script_lines: list[str] | None = None,
+        user_cache_dir: str | None = None,
+        poll_interval: int | None = None,
     ):
         self.slurm_runner_type = slurm_runner_type
         self.root_dir_local = root_dir_local
@@ -121,7 +132,6 @@ class BaseSlurmRunner(BaseRunner):
         raise NotImplementedError("Implement in child class.")
 
     def _get_finished_jobs(self, job_ids: list[str]) -> set[str]:
-
         #  If there is no Slurm job to check, return right away
         if not job_ids:
             return set()
@@ -168,59 +178,73 @@ class BaseSlurmRunner(BaseRunner):
 
     def _submit_single_sbatch(
         self,
-        func,
+        *,
+        base_command: str,
         slurm_job: SlurmJob,
         slurm_config: SlurmConfig,
     ) -> str:
         logger.debug("[_submit_single_sbatch] START")
-        # Prepare input pickle(s)
-        versions = dict(
-            python=sys.version_info[:3],
-            cloudpickle=cloudpickle.__version__,
-            fractal_server=__VERSION__,
-        )
+
         for task in slurm_job.tasks:
-            # Write input pickle
-            _args = []
-            _kwargs = dict(
-                parameters=task.parameters,
-                remote_files=task.task_files.remote_files_dict,
+            # Write input file
+            if self.slurm_runner_type == "ssh":
+                args_file_remote = task.task_files.args_file_remote
+            else:
+                args_file_remote = task.task_files.args_file_local
+            metadiff_file_remote = task.task_files.metadiff_file_remote
+            full_command = (
+                f"{base_command} "
+                f"--args-json {args_file_remote} "
+                f"--out-json {metadiff_file_remote}"
             )
-            funcser = cloudpickle.dumps((versions, func, _args, _kwargs))
-            with open(task.input_pickle_file_local, "wb") as f:
-                f.write(funcser)
+
+            input_data = RemoteInputData(
+                full_command=full_command,
+                python_version=sys.version_info[:3],
+                fractal_server_version=__VERSION__,
+                metadiff_file_remote=task.task_files.metadiff_file_remote,
+                log_file_remote=task.task_files.log_file_remote,
+            )
+
+            with open(task.input_file_local, "w") as f:
+                json.dump(input_data.model_dump(), f, indent=2)
+
+            with open(task.task_files.args_file_local, "w") as f:
+                json.dump(task.parameters, f, indent=2)
+
             logger.debug(
-                "[_submit_single_sbatch] Written "
-                f"{task.input_pickle_file_local=}"
+                "[_submit_single_sbatch] Written " f"{task.input_file_local=}"
             )
 
             if self.slurm_runner_type == "ssh":
-                # Send input pickle (only relevant for SSH)
+                # Send input file (only relevant for SSH)
+                self.fractal_ssh.send_file(
+                    local=task.input_file_local,
+                    remote=task.input_file_remote,
+                )
                 self.fractal_ssh.send_file(
-                    local=task.input_pickle_file_local,
-                    remote=task.input_pickle_file_remote,
+                    local=task.task_files.args_file_local,
+                    remote=task.task_files.args_file_remote,
                 )
                 logger.debug(
                     "[_submit_single_sbatch] Transferred "
-                    f"{task.input_pickle_file_local=}"
+                    f"{task.input_file_local=}"
                 )
 
         # Prepare commands to be included in SLURM submission script
         cmdlines = []
         for task in slurm_job.tasks:
             if self.slurm_runner_type == "ssh":
-                input_pickle_file = task.input_pickle_file_remote
+                input_file = task.input_file_remote
             else:
-                input_pickle_file = task.input_pickle_file_local
-            output_pickle_file = task.output_pickle_file_remote
+                input_file = task.input_file_local
+            output_file = task.output_file_remote
             cmdlines.append(
-                (
-                    f"{self.python_worker_interpreter}"
-                    " -m fractal_server.app.runner."
-                    "executors.slurm_common.remote "
-                    f"--input-file {input_pickle_file} "
-                    f"--output-file {output_pickle_file}"
-                )
+                f"{self.python_worker_interpreter}"
+                " -m fractal_server.app.runner."
+                "executors.slurm_common.remote "
+                f"--input-file {input_file} "
+                f"--output-file {output_file}"
             )
 
         # Set ntasks
@@ -363,12 +387,12 @@ class BaseSlurmRunner(BaseRunner):
         was_job_scancelled: bool = False,
     ) -> tuple[Any, Exception]:
         try:
-            with open(task.output_pickle_file_local, "rb") as f:
-                outdata = f.read()
-            success, output = cloudpickle.loads(outdata)
+            with open(task.output_file_local) as f:
+                output = json.load(f)
+            success = output[0]
             if success:
                 # Task succeeded
-                result = output
+                result = output[1]
                 return (result, None)
             else:
                 # Task failed in a controlled way, and produced an `output`
@@ -376,21 +400,18 @@ class BaseSlurmRunner(BaseRunner):
                 # `exc_type_name` and `traceback_string` and with optional
                 # keys `workflow_task_order`, `workflow_task_id` and
                 # `task_name`.
-                exc_type_name = output.get("exc_type_name")
+                exc_proxy = output[1]
+                exc_type_name = exc_proxy.get("exc_type_name")
                 logger.debug(
-                    f"Output pickle contains a '{exc_type_name}' exception."
+                    f"Output file contains a '{exc_type_name}' exception."
+                )
+                traceback_string = output[1].get("traceback_string")
+                exception = TaskExecutionError(
+                    traceback_string,
+                    workflow_task_id=task.workflow_task_id,
+                    workflow_task_order=task.workflow_task_order,
+                    task_name=task.task_name,
                 )
-                traceback_string = output.get("traceback_string")
-                kwargs = {
-                    key: output[key]
-                    for key in [
-                        "workflow_task_order",
-                        "workflow_task_id",
-                        "task_name",
-                    ]
-                    if key in output.keys()
-                }
-                exception = TaskExecutionError(traceback_string, **kwargs)
                 return (None, exception)
 
         except Exception as e:
@@ -405,8 +426,8 @@ class BaseSlurmRunner(BaseRunner):
                 exception = SHUTDOWN_EXCEPTION
             return (None, exception)
         finally:
-            Path(task.input_pickle_file_local).unlink(missing_ok=True)
-            Path(task.output_pickle_file_local).unlink(missing_ok=True)
+            Path(task.input_file_local).unlink(missing_ok=True)
+            Path(task.output_file_local).unlink(missing_ok=True)
 
     def is_shutdown(self) -> bool:
         return self.shutdown_file.exists()
@@ -451,7 +472,10 @@ class BaseSlurmRunner(BaseRunner):
 
     def submit(
         self,
-        func: callable,
+        base_command: str,
+        workflow_task_order: int,
+        workflow_task_id: int,
+        task_name: str,
         parameters: dict[str, Any],
         history_unit_id: int,
         task_files: TaskFiles,
@@ -507,13 +531,16 @@ class BaseSlurmRunner(BaseRunner):
                         workdir_remote=workdir_remote,
                         workdir_local=workdir_local,
                         task_files=task_files,
+                        workflow_task_order=workflow_task_order,
+                        workflow_task_id=workflow_task_id,
+                        task_name=task_name,
                     )
                 ],
             )
 
             config.parallel_tasks_per_job = 1
             self._submit_single_sbatch(
-                func,
+                base_command=base_command,
                 slurm_job=slurm_job,
                 slurm_config=config,
             )
@@ -586,7 +613,10 @@ class BaseSlurmRunner(BaseRunner):
 
     def multisubmit(
         self,
-        func: callable,
+        base_command: str,
+        workflow_task_order: int,
+        workflow_task_id: int,
+        task_name: str,
         list_parameters: list[dict],
         history_unit_ids: list[int],
         list_task_files: list[TaskFiles],
@@ -602,7 +632,6 @@ class BaseSlurmRunner(BaseRunner):
 
         logger.debug(f"[multisubmit] START, {len(list_parameters)=}")
         try:
-
             if self.is_shutdown():
                 if task_type == "parallel":
                     with next(get_sync_db()) as db:
@@ -672,6 +701,9 @@ class BaseSlurmRunner(BaseRunner):
                             parameters=parameters,
                             zarr_url=parameters["zarr_url"],
                             task_files=list_task_files[index],
+                            workflow_task_order=workflow_task_order,
+                            workflow_task_id=workflow_task_id,
+                            task_name=task_name,
                         ),
                     )
                 jobs_to_submit.append(
@@ -687,7 +719,7 @@ class BaseSlurmRunner(BaseRunner):
             logger.debug("[multisubmit] Transfer files and submit jobs.")
             for slurm_job in jobs_to_submit:
                 self._submit_single_sbatch(
-                    func,
+                    base_command=base_command,
                     slurm_job=slurm_job,
                     slurm_config=config,
                 )
@@ -850,8 +882,8 @@ class BaseSlurmRunner(BaseRunner):
         """
         Check that a list of `SlurmJob`s have homogeneous working folders.
         """
-        set_workdir_local = set(_job.workdir_local for _job in slurm_jobs)
-        set_workdir_remote = set(_job.workdir_remote for _job in slurm_jobs)
+        set_workdir_local = {_job.workdir_local for _job in slurm_jobs}
+        set_workdir_remote = {_job.workdir_remote for _job in slurm_jobs}
         if len(set_workdir_local) > 1:
             raise ValueError(f"Non-unique values in {set_workdir_local=}.")
         if len(set_workdir_remote) > 1:

fractal_server/app/runner/executors/slurm_common/get_slurm_config.py

@@ -1,6 +1,5 @@
 from pathlib import Path
 from typing import Literal
-from typing import Optional
 
 from ._batching import heuristics
 from ._slurm_config import _parse_mem_value
@@ -14,7 +13,7 @@ from fractal_server.app.models.v2 import WorkflowTaskV2
 def get_slurm_config_internal(
     wftask: WorkflowTaskV2,
     which_type: Literal["non_parallel", "parallel"],
-    config_path: Optional[Path] = None,
+    config_path: Path | None = None,
 ) -> SlurmConfig:
     """
     Prepare a `SlurmConfig` configuration object
@@ -168,7 +167,7 @@ def get_slurm_config_internal(
 def get_slurm_config(
     wftask: WorkflowTaskV2,
     which_type: Literal["non_parallel", "parallel"],
-    config_path: Optional[Path] = None,
+    config_path: Path | None = None,
     tot_tasks: int = 1,
 ) -> SlurmConfig:
     config = get_slurm_config_internal(

fractal_server/app/runner/executors/slurm_common/remote.py

@@ -1,27 +1,10 @@
-# This adapts clusterfutures <https://github.com/sampsyo/clusterfutures>
-# Original Copyright
-# Copyright 2021 Adrian Sampson <asampson@cs.washington.edu>
-# License: MIT
-#
-# Modified by:
-# Jacopo Nespolo <jacopo.nespolo@exact-lab.it>
-# Tommaso Comparin <tommaso.comparin@exact-lab.it>
-#
-# Copyright 2022 (C) Friedrich Miescher Institute for Biomedical Research and
-# University of Zurich
-"""
-This module provides a simple self-standing script that executes arbitrary
-python code received via pickled files on a cluster node.
-"""
 import argparse
+import json
 import logging
 import os
 import sys
-from typing import Literal
-from typing import Union
-
-import cloudpickle
 
+from ..call_command_wrapper import call_command_wrapper
 from fractal_server import __VERSION__
 
 
@@ -33,59 +16,6 @@ class FractalVersionMismatch(RuntimeError):
     pass
 
 
-def _check_versions_mismatch(
-    server_versions: dict[
-        Literal["python", "fractal_server", "cloudpickle"],
-        Union[str, tuple[int]],
-    ]
-):
-    """
-    Compare the server {python,cloudpickle,fractal_server} versions with the
-    ones available to the current worker
-
-    Arguments:
-        server_versions:
-            The version used in the fractal-server instance that created the
-            cloudpickle file
-
-    Raises:
-        FractalVersionMismatch: If the cloudpickle or fractal_server versions
-                                do not match with the ones on the server
-    """
-
-    server_python_version = list(server_versions["python"])
-    worker_python_version = list(sys.version_info[:3])
-    if worker_python_version != server_python_version:
-        if worker_python_version[:2] != server_python_version[:2]:
-            # FIXME: Turn this into an error, in some version post 2.14.
-            logging.error(
-                f"{server_python_version=} but {worker_python_version=}. "
-                "This configuration will be deprecated in a future version, "
-                "please contact the admin of this Fractal instance."
-            )
-        else:
-            # Major.minor versions match, patch versions differ
-            logging.warning(
-                f"{server_python_version=} but {worker_python_version=}."
-            )
-
-    server_cloudpickle_version = server_versions["cloudpickle"]
-    worker_cloudpickle_version = cloudpickle.__version__
-    if worker_cloudpickle_version != server_cloudpickle_version:
-        raise FractalVersionMismatch(
-            f"{server_cloudpickle_version=} but "
-            f"{worker_cloudpickle_version=}"
-        )
-
-    server_fractal_server_version = server_versions["fractal_server"]
-    worker_fractal_server_version = __VERSION__
-    if worker_fractal_server_version != server_fractal_server_version:
-        raise FractalVersionMismatch(
-            f"{server_fractal_server_version=} but "
-            f"{worker_fractal_server_version=}"
-        )
-
-
 def worker(
     *,
     in_fname: str,
@@ -95,8 +25,8 @@ def worker(
     Execute a job, possibly on a remote node.
 
     Arguments:
-        in_fname: Absolute path to the input pickle file (must be readable).
-        out_fname: Absolute path of the output pickle file (must be writeable).
+        in_fname: Absolute path to the input file (must be readable).
+        out_fname: Absolute path of the output file (must be writeable).
     """
 
     # Create output folder, if missing
@@ -107,19 +37,49 @@ def worker(
 
     # Execute the job and capture exceptions
     try:
-        with open(in_fname, "rb") as f:
-            indata = f.read()
-        server_versions, fun, args, kwargs = cloudpickle.loads(indata)
-        _check_versions_mismatch(server_versions)
+        with open(in_fname) as f:
+            input_data = json.load(f)
+
+        server_python_version = input_data["python_version"]
+        server_fractal_server_version = input_data["fractal_server_version"]
+
+        # Fractal-server version must be identical
+        worker_fractal_server_version = __VERSION__
+        if worker_fractal_server_version != server_fractal_server_version:
+            raise FractalVersionMismatch(
+                f"{server_fractal_server_version=} but "
+                f"{worker_fractal_server_version=}"
+            )
+
+        # Python version mismatch only raises a warning
+        worker_python_version = tuple(sys.version_info[:3])
+        if worker_python_version != server_python_version:
+            if worker_python_version[:2] != server_python_version[:2]:
+                logging.warning(
+                    f"{server_python_version=} but {worker_python_version=}."
+                )
+
+        # Extract some useful paths
+        metadiff_file_remote = input_data["metadiff_file_remote"]
+        log_path = input_data["log_file_remote"]
+
+        # Execute command
+        full_command = input_data["full_command"]
+        call_command_wrapper(cmd=full_command, log_path=log_path)
+
+        try:
+            with open(metadiff_file_remote) as f:
+                out_meta = json.load(f)
+            result = (True, out_meta)
+        except FileNotFoundError:
+            # Command completed, but it produced no metadiff file
+            result = (True, None)
 
-        result = (True, fun(*args, **kwargs))
-        out = cloudpickle.dumps(result)
     except Exception as e:
         # Exception objects are not serialisable. Here we save the relevant
         # exception contents in a serializable dictionary. Note that whenever
         # the task failed "properly", the exception is a `TaskExecutionError`
         # and it has additional attributes.
-
         import traceback
 
         exc_type, exc_value, traceback_obj = sys.exc_info()
@@ -131,33 +91,28 @@ def worker(
         )
         traceback_string = "".join(traceback_list)
         exc_proxy = dict(
-            exc_type_name=exc_type.__name__,
+            exc_type_name=type(e).__name__,
             traceback_string=traceback_string,
-            workflow_task_order=getattr(e, "workflow_task_order", None),
-            workflow_task_id=getattr(e, "workflow_task_id", None),
-            task_name=getattr(e, "task_name", None),
         )
         result = (False, exc_proxy)
-        out = cloudpickle.dumps(result)
 
-    # Write the output pickle file
-    with open(out_fname, "wb") as f:
-        f.write(out)
+    # Write output file
+    with open(out_fname, "w") as f:
+        json.dump(result, f, indent=2)
 
 
 if __name__ == "__main__":
-
     parser = argparse.ArgumentParser()
     parser.add_argument(
         "--input-file",
         type=str,
-        help="Path of input pickle file",
+        help="Path of input JSON file",
         required=True,
     )
     parser.add_argument(
         "--output-file",
         type=str,
-        help="Path of output pickle file",
+        help="Path of output JSON file",
         required=True,
     )
     parsed_args = parser.parse_args()

fractal_server/app/runner/executors/slurm_common/slurm_job_task_models.py

@@ -1,6 +1,5 @@
 from pathlib import Path
 from typing import Any
-from typing import Optional
 
 from pydantic import BaseModel
 from pydantic import ConfigDict
@@ -15,56 +14,57 @@ class SlurmTask(BaseModel):
     workdir_local: Path
     workdir_remote: Path
     parameters: dict[str, Any]
-    zarr_url: Optional[str] = None
+    zarr_url: str | None = None
     task_files: TaskFiles
     index: int
 
+    workflow_task_order: int
+    workflow_task_id: int
+    task_name: str
+
     @property
-    def input_pickle_file_local_path(self) -> Path:
+    def input_file_local_path(self) -> Path:
         return (
-            self.workdir_local / f"{self.prefix}-{self.component}-input.pickle"
+            self.workdir_local / f"{self.prefix}-{self.component}-input.json"
         )
 
     @property
-    def input_pickle_file_remote_path(self) -> Path:
+    def input_file_remote_path(self) -> Path:
         return (
-            self.workdir_remote
-            / f"{self.prefix}-{self.component}-input.pickle"
+            self.workdir_remote / f"{self.prefix}-{self.component}-input.json"
         )
 
     @property
-    def output_pickle_file_local_path(self) -> Path:
+    def output_file_local_path(self) -> Path:
         return (
-            self.workdir_local
-            / f"{self.prefix}-{self.component}-output.pickle"
+            self.workdir_local / f"{self.prefix}-{self.component}-output.json"
         )
 
     @property
-    def output_pickle_file_remote_path(self) -> Path:
+    def output_file_remote_path(self) -> Path:
         return (
-            self.workdir_remote
-            / f"{self.prefix}-{self.component}-output.pickle"
+            self.workdir_remote / f"{self.prefix}-{self.component}-output.json"
         )
 
     @property
-    def input_pickle_file_local(self) -> str:
-        return self.input_pickle_file_local_path.as_posix()
+    def input_file_local(self) -> str:
+        return self.input_file_local_path.as_posix()
 
     @property
-    def input_pickle_file_remote(self) -> str:
-        return self.input_pickle_file_remote_path.as_posix()
+    def input_file_remote(self) -> str:
+        return self.input_file_remote_path.as_posix()
 
     @property
-    def output_pickle_file_local(self) -> str:
-        return self.output_pickle_file_local_path.as_posix()
+    def output_file_local(self) -> str:
+        return self.output_file_local_path.as_posix()
 
     @property
-    def output_pickle_file_remote(self) -> str:
-        return self.output_pickle_file_remote_path.as_posix()
+    def output_file_remote(self) -> str:
+        return self.output_file_remote_path.as_posix()
 
 
 class SlurmJob(BaseModel):
-    slurm_job_id: Optional[str] = None
+    slurm_job_id: str | None = None
     prefix: str
     workdir_local: Path
     workdir_remote: Path

fractal_server/app/runner/executors/slurm_ssh/run_subprocess.py

@@ -1,6 +1,5 @@
 import shlex
 import subprocess  # nosec
-from typing import Optional
 
 from fractal_server.logger import get_logger
 from fractal_server.string_tools import validate_cmd
@@ -8,8 +7,8 @@ from fractal_server.string_tools import validate_cmd
 
 def run_subprocess(
     cmd: str,
-    allow_char: Optional[str] = None,
-    logger_name: Optional[str] = None,
+    allow_char: str | None = None,
+    logger_name: str | None = None,
 ) -> subprocess.CompletedProcess:
     validate_cmd(cmd, allow_char=allow_char)
     logger = get_logger(logger_name)

fractal_server/app/runner/executors/slurm_ssh/runner.py

@@ -1,6 +1,5 @@
 import time
 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
@@ -27,9 +26,9 @@ class SlurmSSHRunner(BaseSlurmRunner):
         # 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,
+        common_script_lines: list[str] | None = None,
+        user_cache_dir: str | None = None,
+        poll_interval: int | None = None,
         # Specific
         fractal_ssh: FractalSSH,
     ) -> None:
@@ -99,9 +98,8 @@ class SlurmSSHRunner(BaseSlurmRunner):
             for task in _slurm_job.tasks:
                 _single_job_filelist.extend(
                     [
-                        task.output_pickle_file_remote_path.name,
+                        task.output_file_remote_path.name,
                         task.task_files.log_file_remote_path.name,
-                        task.task_files.args_file_remote_path.name,
                         task.task_files.metadiff_file_remote_path.name,
                     ]
                 )

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

@@ -17,7 +17,6 @@ another user. Note that this requires appropriate sudo permissions.
 """
 import shlex
 import subprocess  # nosec
-from typing import Optional
 
 from fractal_server.logger import set_logger
 from fractal_server.string_tools import validate_cmd
@@ -28,8 +27,7 @@ logger = set_logger(__name__)
 def _run_command_as_user(
     *,
     cmd: str,
-    user: Optional[str] = None,
-    encoding: Optional[str] = "utf-8",
+    user: str | None = None,
     check: bool = False,
 ) -> subprocess.CompletedProcess:
     """
@@ -38,8 +36,6 @@ def _run_command_as_user(
     Arguments:
         cmd: Command to be run
         user: User to be impersonated
-        encoding: Argument for `subprocess.run`. Note that this must be `None`
-                  to have stdout/stderr as bytes.
         check: If `True`, check that `returncode=0` and fail otherwise.
 
     Raises:
@@ -57,7 +53,7 @@ def _run_command_as_user(
     res = subprocess.run(  # nosec
         shlex.split(new_cmd),
         capture_output=True,
-        encoding=encoding,
+        encoding="utf-8",
     )
     logger.debug(f"[_run_command_as_user] {res.returncode=}")
     logger.debug(f"[_run_command_as_user] {res.stdout=}")