fractal-server 2.16.6__py3-none-any.whl → 2.17.0a0__py3-none-any.whl

fractal_server/runner/config/_slurm.py

@@ -0,0 +1,128 @@
+from typing import Annotated
+
+from pydantic import AfterValidator
+from pydantic import BaseModel
+from pydantic import ConfigDict
+from pydantic.types import PositiveInt
+
+from fractal_server.runner.config.slurm_mem_to_MB import slurm_mem_to_MB
+from fractal_server.types import DictStrStr
+from fractal_server.types import NonEmptyStr
+
+
+MemMBType = Annotated[
+    PositiveInt | NonEmptyStr, AfterValidator(slurm_mem_to_MB)
+]
+
+
+class _SlurmConfigSet(BaseModel):
+    """
+    Options for the default or gpu SLURM config.
+
+    Attributes:
+        partition:
+        cpus_per_task:
+        mem:
+            See `_parse_mem_value` for details on allowed values.
+        constraint:
+        gres:
+        time:
+        exclude:
+        nodelist:
+        account:
+        extra_lines:
+    """
+
+    model_config = ConfigDict(extra="forbid")
+
+    partition: NonEmptyStr | None = None
+    cpus_per_task: PositiveInt | None = None
+    mem: MemMBType | None = None
+    constraint: NonEmptyStr | None = None
+    gres: NonEmptyStr | None = None
+    exclude: NonEmptyStr | None = None
+    nodelist: NonEmptyStr | None = None
+    time: NonEmptyStr | None = None
+    account: NonEmptyStr | None = None
+    extra_lines: list[NonEmptyStr] | None = None
+    gpus: NonEmptyStr | None = None
+
+
+class _BatchingConfigSet(BaseModel):
+    """
+    Options to configure the batching strategy (that is, how to combine
+    several tasks in a single SLURM job).
+
+    Attributes:
+        target_cpus_per_job:
+        max_cpus_per_job:
+        target_mem_per_job:
+            (see `_parse_mem_value` for details on allowed values)
+        max_mem_per_job:
+            (see `_parse_mem_value` for details on allowed values)
+        target_num_jobs:
+        max_num_jobs:
+    """
+
+    model_config = ConfigDict(extra="forbid")
+
+    target_num_jobs: PositiveInt
+    max_num_jobs: PositiveInt
+    target_cpus_per_job: PositiveInt
+    max_cpus_per_job: PositiveInt
+    target_mem_per_job: MemMBType
+    max_mem_per_job: MemMBType
+
+
+class JobRunnerConfigSLURM(BaseModel):
+    """
+    Common SLURM configuration.
+
+    Note: this is a common and abstract class, which gets transformed into
+    more specific configuration objects during job execution.
+
+    Valid JSON example
+    ```JSON
+    {
+      "default_slurm_config": {
+          "partition": "main",
+          "cpus_per_task": 1
+      },
+      "gpu_slurm_config": {
+          "partition": "gpu",
+          "extra_lines": ["#SBATCH --gres=gpu:v100:1"]
+      },
+      "batching_config": {
+          "target_cpus_per_job": 1,
+          "max_cpus_per_job": 1,
+          "target_mem_per_job": 200,
+          "max_mem_per_job": 500,
+          "target_num_jobs": 2,
+          "max_num_jobs": 4
+      },
+      "user_local_exports": {
+          "CELLPOSE_LOCAL_MODELS_PATH": "CELLPOSE_LOCAL_MODELS_PATH",
+          "NAPARI_CONFIG": "napari_config.json"
+      }
+    }
+    ```
+
+    Attributes:
+        default_slurm_config:
+            Common default options for all tasks.
+        gpu_slurm_config:
+            Default configuration for all GPU tasks.
+        batching_config:
+            Configuration of the batching strategy.
+        user_local_exports:
+            Key-value pairs to be included as `export`-ed variables in SLURM
+            submission script, after prepending values with the user's cache
+            directory.
+    """
+
+    model_config = ConfigDict(extra="forbid")
+
+    default_slurm_config: _SlurmConfigSet
+    gpu_slurm_config: _SlurmConfigSet | None = None
+    batching_config: _BatchingConfigSet
+    user_local_exports: DictStrStr | None = None

fractal_server/runner/config/slurm_mem_to_MB.py

@@ -0,0 +1,63 @@
+from fractal_server.logger import set_logger
+from fractal_server.runner.exceptions import SlurmConfigError
+
+
+logger = set_logger(__name__)
+
+
+def slurm_mem_to_MB(raw_mem: str | int) -> int:
+    """
+    Convert a memory-specification string into an integer (in MB units), or
+    simply return the input if it is already an integer.
+
+    Supported units are `"M", "G", "T"`, with `"M"` being the default; some
+    parsing examples are: `"10M" -> 10000`, `"3G" -> 3000000`.
+
+    Args:
+        raw_mem:
+            A string (e.g. `"100M"`) or an integer (in MB).
+
+    Returns:
+        Integer value of memory in MB units.
+    """
+
+    info = f"[_parse_mem_value] {raw_mem=}"
+    error_msg = (
+        f"{info}, invalid specification of memory requirements "
+        "(valid examples: 93, 71M, 93G, 71T)."
+    )
+
+    # Handle integer argument
+    if type(raw_mem) is int:
+        return raw_mem
+
+    # Handle string argument
+    if not raw_mem[0].isdigit():  # fail e.g. for raw_mem="M100"
+        logger.error(error_msg)
+        raise SlurmConfigError(error_msg)
+    if raw_mem.isdigit():
+        mem_MB = int(raw_mem)
+    elif raw_mem.endswith("M"):
+        stripped_raw_mem = raw_mem.strip("M")
+        if not stripped_raw_mem.isdigit():
+            logger.error(error_msg)
+            raise SlurmConfigError(error_msg)
+        mem_MB = int(stripped_raw_mem)
+    elif raw_mem.endswith("G"):
+        stripped_raw_mem = raw_mem.strip("G")
+        if not stripped_raw_mem.isdigit():
+            logger.error(error_msg)
+            raise SlurmConfigError(error_msg)
+        mem_MB = int(stripped_raw_mem) * 10**3
+    elif raw_mem.endswith("T"):
+        stripped_raw_mem = raw_mem.strip("T")
+        if not stripped_raw_mem.isdigit():
+            logger.error(error_msg)
+            raise SlurmConfigError(error_msg)
+        mem_MB = int(stripped_raw_mem) * 10**6
+    else:
+        logger.error(error_msg)
+        raise SlurmConfigError(error_msg)
+
+    logger.debug(f"{info}, return {mem_MB}")
+    return mem_MB

fractal_server/runner/exceptions.py

@@ -63,3 +63,7 @@ class JobExecutionError(RuntimeError):
             content = str(self)
         message = f"JobExecutionError\n{content}"
         return message
+
+
+class SlurmConfigError(ValueError):
+    pass

fractal_server/runner/executors/base_runner.py

@@ -3,6 +3,8 @@ from typing import Any
 
 from fractal_server.app.schemas.v2.task import TaskType
 from fractal_server.logger import set_logger
+from fractal_server.runner.config import JobRunnerConfigLocal
+from fractal_server.runner.config import JobRunnerConfigSLURM
 from fractal_server.runner.task_files import TaskFiles
 
 
@@ -39,21 +41,24 @@ class BaseRunner:
     Base class for Fractal runners.
     """
 
+    shared_config: JobRunnerConfigLocal | JobRunnerConfigSLURM
+
     executor_error_log: str | None = None
 
     def submit(
         self,
+        *,
         base_command: str,
         workflow_task_order: int,
         workflow_task_id: int,
         task_name: str,
         parameters: dict[str, Any],
         history_unit_id: int,
-        task_type: TaskType,
+        task_type: SubmitTaskType,
         task_files: TaskFiles,
-        config: Any,
         user_id: int,
-    ) -> tuple[Any, BaseException]:
+        config: Any,
+    ) -> tuple[Any, BaseException | None]:
         """
         Run a single fractal task.
 
@@ -74,20 +79,25 @@ class BaseRunner:
 
     def multisubmit(
         self,
+        *,
         base_command: str,
         workflow_task_order: int,
         workflow_task_id: int,
         task_name: str,
-        list_parameters: list[dict[str, Any]],
+        list_parameters: list[dict],
         history_unit_ids: list[int],
         list_task_files: list[TaskFiles],
-        task_type: TaskType,
+        task_type: MultisubmitTaskType,
         config: Any,
         user_id: int,
     ) -> tuple[dict[int, Any], dict[int, BaseException]]:
         """
         Run a parallel fractal task.
 
+        Note: `list_parameters`, `list_task_files` and `history_unit_ids`
+        have the same size. For parallel tasks, this is also the number of
+        input images, while for compound tasks these can differ.
+
         Args:
             base_command:
             workflow_task_order:
@@ -108,7 +118,7 @@ class BaseRunner:
     def validate_submit_parameters(
         self,
         parameters: dict[str, Any],
-        task_type: TaskType,
+        task_type: SubmitTaskType | MultisubmitTaskType,
     ) -> None:
         """
         Validate parameters for `submit` method
@@ -143,7 +153,7 @@ class BaseRunner:
     def validate_multisubmit_parameters(
         self,
         *,
-        task_type: TaskType,
+        task_type: MultisubmitTaskType,
         list_parameters: list[dict[str, Any]],
         list_task_files: list[TaskFiles],
         history_unit_ids: list[int],

fractal_server/runner/executors/local/get_local_config.py

@@ -1,116 +1,51 @@
-# Copyright 2022 (C) Friedrich Miescher Institute for Biomedical Research and
-# University of Zurich
-#
-# Original authors:
-# Tommaso Comparin <tommaso.comparin@exact-lab.it>
-#
-# This file is part of Fractal and was originally developed by eXact lab S.r.l.
-# <exact-lab.it> under contract with Liberali Lab from the Friedrich Miescher
-# Institute for Biomedical Research and Pelkmans Lab from the University of
-# Zurich.
 """
 Submodule to handle the local-backend configuration for a WorkflowTask
 """
-import json
-from pathlib import Path
 from typing import Literal
 
-from pydantic import BaseModel
-from pydantic import ConfigDict
-from pydantic import ValidationError
-
 from fractal_server.app.models.v2 import WorkflowTaskV2
-from fractal_server.config import get_settings
-from fractal_server.syringe import Inject
-
-
-class LocalBackendConfigError(ValueError):
-    """
-    Local-backend configuration error
-    """
-
-    pass
-
-
-class LocalBackendConfig(BaseModel):
-    """
-    Specifications of the local-backend configuration
-
-    Attributes:
-        parallel_tasks_per_job:
-            Maximum number of tasks to be run in parallel as part of a call to
-            `FractalThreadPoolExecutor.map`; if `None`, then all tasks will
-            start at the same time.
-    """
-
-    model_config = ConfigDict(extra="forbid")
-    parallel_tasks_per_job: int | None = None
-
-    @property
-    def batch_size(self) -> int:
-        return self.parallel_tasks_per_job or 1
+from fractal_server.runner.config import JobRunnerConfigLocal
 
 
 def get_local_backend_config(
+    shared_config: JobRunnerConfigLocal,
     wftask: WorkflowTaskV2,
     which_type: Literal["non_parallel", "parallel"],
-    config_path: Path | None = None,
     tot_tasks: int = 1,
-) -> LocalBackendConfig:
+) -> JobRunnerConfigLocal:
     """
-    Prepare a `LocalBackendConfig` configuration object
+    Prepare a specific `LocalBackendConfig` configuration.
 
-    The sources for `parallel_tasks_per_job` attributes, starting from the
-    highest-priority one, are
+    The base configuration is the runner-level `shared_config` object, based
+    on `resource.jobs_runner_config`. We then incorporate attributes from
+    `wftask.meta_{non_parallel,parallel}` - with higher priority.
 
-    1. Properties in `wftask.meta_parallel` or `wftask.meta_non_parallel`
-       (depending on `which_type`);
-    2. The general content of the local-backend configuration file;
-    3. The default value (`None`).
-
-    Arguments:
+    Args:
+        shared_config:
+            Configuration object based on `resource.jobs_runner_config`.
         wftask:
             WorkflowTaskV2 for which the backend configuration should
             be prepared.
-        config_path:
-            Path of local-backend configuration file; if `None`, use
-            `FRACTAL_LOCAL_CONFIG_FILE` variable from settings.
+        which_type:
+            Whether we should look at the non-parallel or parallel part
+            of `wftask`.
+        tot_tasks: Not used here, only present as a common interface.
 
     Returns:
-        A local-backend configuration object
+        A ready-to-use local-backend configuration object.
     """
 
-    key = "parallel_tasks_per_job"
-    default_value = None
-
     if which_type == "non_parallel":
         wftask_meta = wftask.meta_non_parallel
     elif which_type == "parallel":
         wftask_meta = wftask.meta_parallel
     else:
         raise ValueError(
-            "`get_local_backend_config` received an invalid argument"
-            f" {which_type=}."
+            f"Invalid {which_type=} in `get_local_backend_config`."
         )
 
-    if wftask_meta and key in wftask_meta:
-        parallel_tasks_per_job = wftask_meta[key]
-    else:
-        if not config_path:
-            settings = Inject(get_settings)
-            config_path = settings.FRACTAL_LOCAL_CONFIG_FILE
-        if config_path is None:
-            parallel_tasks_per_job = default_value
-        else:
-            with config_path.open("r") as f:
-                env = json.load(f)
-            try:
-                _ = LocalBackendConfig(**env)
-            except ValidationError as e:
-                raise LocalBackendConfigError(
-                    f"Error while loading {config_path=}. "
-                    f"Original error:\n{str(e)}"
-                )
-
-            parallel_tasks_per_job = env.get(key, default_value)
-    return LocalBackendConfig(parallel_tasks_per_job=parallel_tasks_per_job)
+    __KEY__ = "parallel_tasks_per_job"
+    output = shared_config.model_copy(deep=True)
+    if wftask_meta and __KEY__ in wftask_meta:
+        output.parallel_tasks_per_job = wftask_meta[__KEY__]
+    return output

fractal_server/runner/executors/local/runner.py

@@ -5,11 +5,13 @@ from pathlib import Path
 from typing import Any
 
 from ..call_command_wrapper import call_command_wrapper
-from .get_local_config import LocalBackendConfig
 from fractal_server.app.db import get_sync_db
+from fractal_server.app.models import Profile
+from fractal_server.app.models import Resource
 from fractal_server.app.schemas.v2 import HistoryUnitStatus
 from fractal_server.app.schemas.v2 import TaskType
 from fractal_server.logger import set_logger
+from fractal_server.runner.config import JobRunnerConfigLocal
 from fractal_server.runner.exceptions import TaskExecutionError
 from fractal_server.runner.executors.base_runner import BaseRunner
 from fractal_server.runner.executors.base_runner import MultisubmitTaskType
@@ -56,15 +58,21 @@ def run_single_task(
 class LocalRunner(BaseRunner):
     executor: ThreadPoolExecutor
     root_dir_local: Path
+    shared_config: JobRunnerConfigLocal
 
     def __init__(
         self,
         root_dir_local: Path,
+        resource: Resource,
+        profile: Profile,
     ):
         self.root_dir_local = root_dir_local
         self.root_dir_local.mkdir(parents=True, exist_ok=True)
         self.executor = ThreadPoolExecutor()
         logger.debug("Create LocalRunner")
+        self.shared_config = JobRunnerConfigLocal(
+            **resource.jobs_runner_config
+        )
 
     def __enter__(self):
         logger.debug("Enter LocalRunner")
@@ -80,6 +88,7 @@ class LocalRunner(BaseRunner):
 
     def submit(
         self,
+        *,
         base_command: str,
         workflow_task_order: int,
         workflow_task_id: int,
@@ -87,10 +96,26 @@ class LocalRunner(BaseRunner):
         parameters: dict[str, Any],
         history_unit_id: int,
         task_files: TaskFiles,
-        config: LocalBackendConfig,
+        config: JobRunnerConfigLocal,
         task_type: SubmitTaskType,
         user_id: int,
-    ) -> tuple[Any, Exception]:
+    ) -> tuple[Any, Exception | None]:
+        """
+        Run a single fractal task.
+
+        Args:
+            base_command:
+            workflow_task_order:
+            workflow_task_id:
+            task_name:
+            parameters: Dictionary of parameters.
+            history_unit_id:
+                Database ID of the corresponding `HistoryUnit` entry.
+            task_type: Task type.
+            task_files: `TaskFiles` object.
+            config: Runner-specific parameters.
+            user_id:
+        """
         logger.debug("[submit] START")
 
         try:
@@ -146,21 +171,39 @@ class LocalRunner(BaseRunner):
 
     def multisubmit(
         self,
+        *,
         base_command: str,
         workflow_task_order: int,
         workflow_task_id: int,
         task_name: str,
-        list_parameters: list[dict],
+        list_parameters: list[dict[str, Any]],
         history_unit_ids: list[int],
         list_task_files: list[TaskFiles],
         task_type: MultisubmitTaskType,
-        config: LocalBackendConfig,
+        config: JobRunnerConfigLocal,
         user_id: int,
     ) -> tuple[dict[int, Any], dict[int, BaseException]]:
         """
+        Run a parallel fractal task.
+
         Note: `list_parameters`, `list_task_files` and `history_unit_ids`
         have the same size. For parallel tasks, this is also the number of
         input images, while for compound tasks these can differ.
+
+        Args:
+            base_command:
+            workflow_task_order:
+            workflow_task_id:
+            task_name:
+            list_parameters:
+                List of dictionaries of parameters (each one must include
+                `zarr_urls` key).
+            history_unit_ids:
+                Database IDs of the corresponding `HistoryUnit` entries.
+            list_task_files: `TaskFiles` objects.
+            task_type: Task type.
+            config: Runner-specific parameters.
+            user_id:
         """
 
         logger.debug(f"[multisubmit] START, {len(list_parameters)=}")

fractal_server/runner/executors/slurm_common/_batching.py

@@ -36,7 +36,7 @@ def _estimate_parallel_tasks_per_job(
     assumes that further checks will be performed on the output of the current
     function, as is the case in the `heuristics` function below.
 
-    Arguments:
+    Args:
         cpus_per_task: Number of CPUs needed for one task.
         mem_per_task: Memory (in MB) needed for one task.
         max_cpus_per_job: Maximum number of CPUs available for one job.
@@ -91,7 +91,7 @@ def heuristics(
        introduce in-job queues to satisfy the hard constraint on the maximum
        number of jobs.
 
-    Arguments:
+    Args:
         tot_tasks:
             Total number of elements to be processed (e.g. number of images in
             a OME-NGFF array).