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

fractal_server/data_migrations/2_14_10.py

@@ -1,48 +0,0 @@
-import logging
-
-from sqlmodel import select
-
-from fractal_server.app.db import get_sync_db
-from fractal_server.app.models import HistoryRun
-from fractal_server.app.models import TaskV2
-from fractal_server.app.models import WorkflowTaskV2
-
-logger = logging.getLogger("fix_db")
-logger.setLevel(logging.INFO)
-
-
-def fix_db():
-    logger.info("START execution of fix_db function")
-
-    with next(get_sync_db()) as db:
-        stm = select(HistoryRun).order_by(HistoryRun.id)
-        history_runs = db.execute(stm).scalars().all()
-
-        for hr in history_runs:
-            logger.info(f"HistoryRun[{hr.id}] START")
-            if hr.workflowtask_id is None:
-                continue
-            wft = db.get(WorkflowTaskV2, hr.workflowtask_id)
-            if wft is None:
-                logger.warning(
-                    f"WorkflowTaskV2[{hr.workflowtask_id}] not found. "
-                    "Trying to use HistoryRun.workflowtask_dump"
-                )
-                task_id = hr.workflowtask_dump.get("task_id")
-                if task_id is not None and db.get(TaskV2, task_id) is not None:
-                    hr.task_id = task_id
-                else:
-                    logger.warning(f"TaskV2[{task_id}] not found")
-            else:
-                hr.task_id = wft.task_id
-                logger.info(
-                    f"HistoryRun[{hr.id}].task_id set to {wft.task_id}"
-                )
-
-            db.add(hr)
-            logger.info(f"HistoryRun[{hr.id}] END")
-
-        db.commit()
-        logger.info("Changes committed.")
-
-    logger.info("END execution of fix_db function")

fractal_server/runner/executors/slurm_common/_slurm_config.py

@@ -1,471 +0,0 @@
-# Copyright 2022 (C) Friedrich Miescher Institute for Biomedical Research and
-# University of Zurich
-#
-# Original authors:
-# Jacopo Nespolo <jacopo.nespolo@exact-lab.it>
-# 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 SLURM configuration for a WorkflowTask
-"""
-import json
-from pathlib import Path
-
-from pydantic import BaseModel
-from pydantic import ConfigDict
-from pydantic import Field
-from pydantic import ValidationError
-
-from fractal_server.config import get_settings
-from fractal_server.logger import set_logger
-from fractal_server.syringe import Inject
-
-logger = set_logger(__name__)
-
-
-class SlurmConfigError(ValueError):
-    """
-    Slurm configuration error
-    """
-
-    pass
-
-
-class _SlurmConfigSet(BaseModel):
-    """
-    Options that can be set in `FRACTAL_SLURM_CONFIG_FILE` for the default/gpu
-    SLURM config. Only used as part of `SlurmConfigFile`.
-
-    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: str | None = None
-    cpus_per_task: int | None = None
-    mem: int | str | None = None
-    constraint: str | None = None
-    gres: str | None = None
-    exclude: str | None = None
-    nodelist: str | None = None
-    time: str | None = None
-    account: str | None = None
-    extra_lines: list[str] | None = None
-    gpus: str | None = None
-
-
-class _BatchingConfigSet(BaseModel):
-    """
-    Options that can be set in `FRACTAL_SLURM_CONFIG_FILE` to configure the
-    batching strategy (that is, how to combine several tasks in a single SLURM
-    job). Only used as part of `SlurmConfigFile`.
-
-    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_cpus_per_job: int
-    max_cpus_per_job: int
-    target_mem_per_job: int | str
-    max_mem_per_job: int | str
-    target_num_jobs: int
-    max_num_jobs: int
-
-
-class SlurmConfigFile(BaseModel):
-    """
-    Specifications for the content of `FRACTAL_SLURM_CONFIG_FILE`
-
-    This must be a JSON file, and a valid example is
-    ```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"
-      }
-    }
-    ```
-
-    See `_SlurmConfigSet` and `_BatchingConfigSet` for more details.
-
-    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: dict[str, str] | None = None
-
-
-def load_slurm_config_file(
-    config_path: Path | None = None,
-) -> SlurmConfigFile:
-    """
-    Load a SLURM configuration file and validate its content with
-    `SlurmConfigFile`.
-
-    Arguments:
-        config_path:
-    """
-
-    if not config_path:
-        settings = Inject(get_settings)
-        config_path = settings.FRACTAL_SLURM_CONFIG_FILE
-
-    # Load file
-    logger.debug(f"[get_slurm_config] Now loading {config_path=}")
-    try:
-        with config_path.open("r") as f:
-            slurm_env = json.load(f)
-    except Exception as e:
-        raise SlurmConfigError(
-            f"Error while loading {config_path=}. "
-            f"Original error:\n{str(e)}"
-        )
-
-    # Validate file content
-    logger.debug(f"[load_slurm_config_file] Now validating {config_path=}")
-    logger.debug(f"[load_slurm_config_file] {slurm_env=}")
-    try:
-        obj = SlurmConfigFile(**slurm_env)
-    except ValidationError as e:
-        raise SlurmConfigError(
-            f"Error while loading {config_path=}. "
-            f"Original error:\n{str(e)}"
-        )
-
-    # Convert memory to MB units, in all relevant attributes
-    if obj.default_slurm_config.mem:
-        obj.default_slurm_config.mem = _parse_mem_value(
-            obj.default_slurm_config.mem
-        )
-    if obj.gpu_slurm_config and obj.gpu_slurm_config.mem:
-        obj.gpu_slurm_config.mem = _parse_mem_value(obj.gpu_slurm_config.mem)
-    obj.batching_config.target_mem_per_job = _parse_mem_value(
-        obj.batching_config.target_mem_per_job
-    )
-    obj.batching_config.max_mem_per_job = _parse_mem_value(
-        obj.batching_config.max_mem_per_job
-    )
-
-    return obj
-
-
-class SlurmConfig(BaseModel):
-    """
-    Abstraction for SLURM parameters
-
-    **NOTE**: `SlurmConfig` objects are created internally in `fractal-server`,
-    and they are not meant to be initialized by the user; the same holds for
-    `SlurmConfig` attributes (e.g. `mem_per_task_MB`), which are not meant to
-    be part of the `FRACTAL_SLURM_CONFIG_FILE` JSON file (details on the
-    expected file content are defined in
-    [`SlurmConfigFile`](#fractal_server.runner._slurm._slurm_config.SlurmConfigFile)).
-
-    Part of the attributes map directly to some of the SLURM attributes (see
-    https://slurm.schedmd.com/sbatch.html), e.g. `partition`. Other attributes
-    are metaparameters which are needed in fractal-server to combine multiple
-    tasks in the same SLURM job (e.g. `parallel_tasks_per_job` or
-    `max_num_jobs`).
-
-    Attributes:
-        partition: Corresponds to SLURM option.
-        cpus_per_task: Corresponds to SLURM option.
-        mem_per_task_MB: Corresponds to `mem` SLURM option.
-        job_name: Corresponds to `name` SLURM option.
-        constraint: Corresponds to SLURM option.
-        gres: Corresponds to SLURM option.
-        account: Corresponds to SLURM option.
-        gpus: Corresponds to SLURM option.
-        time: Corresponds to SLURM option (WARNING: not fully supported).
-        nodelist: Corresponds to SLURM option.
-        exclude: Corresponds to SLURM option.
-        prefix: Prefix of configuration lines in SLURM submission scripts.
-        shebang_line: Shebang line for SLURM submission scripts.
-        extra_lines: Additional lines to include in SLURM submission scripts.
-        tasks_per_job: Number of tasks for each SLURM job.
-        parallel_tasks_per_job: Number of tasks to run in parallel for
-                                each SLURM job.
-        target_cpus_per_job: Optimal number of CPUs to be requested in each
-                             SLURM job.
-        max_cpus_per_job: Maximum number of CPUs that can be requested in each
-                          SLURM job.
-        target_mem_per_job: Optimal amount of memory (in MB) to be requested in
-                            each SLURM job.
-        max_mem_per_job: Maximum amount of memory (in MB) that can be requested
-                         in each SLURM job.
-        target_num_jobs: Optimal number of SLURM jobs for a given WorkflowTask.
-        max_num_jobs: Maximum number of SLURM jobs for a given WorkflowTask.
-        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")
-
-    # Required SLURM parameters (note that the integer attributes are those
-    # that will need to scale up with the number of parallel tasks per job)
-    partition: str
-    cpus_per_task: int
-    mem_per_task_MB: int
-    prefix: str = "#SBATCH"
-    shebang_line: str = "#!/bin/sh"
-
-    # Optional SLURM parameters
-    job_name: str | None = None
-    constraint: str | None = None
-    gres: str | None = None
-    gpus: str | None = None
-    time: str | None = None
-    account: str | None = None
-    nodelist: str | None = None
-    exclude: str | None = None
-
-    # Free-field attribute for extra lines to be added to the SLURM job
-    # preamble
-    extra_lines: list[str] | None = Field(default_factory=list)
-
-    # Variables that will be `export`ed in the SLURM submission script
-    user_local_exports: dict[str, str] | None = None
-
-    # Metaparameters needed to combine multiple tasks in each SLURM job
-    tasks_per_job: int | None = None
-    parallel_tasks_per_job: int | None = None
-    target_cpus_per_job: int
-    max_cpus_per_job: int
-    target_mem_per_job: int
-    max_mem_per_job: int
-    target_num_jobs: int
-    max_num_jobs: int
-
-    def _sorted_extra_lines(self) -> list[str]:
-        """
-        Return a copy of `self.extra_lines`, where lines starting with
-        `self.prefix` are listed first.
-        """
-
-        def _no_prefix(_line):
-            if _line.startswith(self.prefix):
-                return 0
-            else:
-                return 1
-
-        return sorted(self.extra_lines, key=_no_prefix)
-
-    def sort_script_lines(self, script_lines: list[str]) -> list[str]:
-        """
-        Return a copy of `script_lines`, where lines are sorted as in:
-
-        1. `self.shebang_line` (if present);
-        2. Lines starting with `self.prefix`;
-        3. Other lines.
-
-        Arguments:
-            script_lines:
-        """
-
-        def _sorting_function(_line):
-            if _line == self.shebang_line:
-                return 0
-            elif _line.startswith(self.prefix):
-                return 1
-            else:
-                return 2
-
-        return sorted(script_lines, key=_sorting_function)
-
-    def to_sbatch_preamble(
-        self,
-        remote_export_dir: str | None = None,
-    ) -> list[str]:
-        """
-        Compile `SlurmConfig` object into the preamble of a SLURM submission
-        script.
-
-        Arguments:
-            remote_export_dir:
-                Base directory for exports defined in
-                `self.user_local_exports`.
-        """
-        if self.parallel_tasks_per_job is None:
-            raise ValueError(
-                "SlurmConfig.sbatch_preamble requires that "
-                f"{self.parallel_tasks_per_job=} is not None."
-            )
-        if self.extra_lines:
-            if len(self.extra_lines) != len(set(self.extra_lines)):
-                raise ValueError(f"{self.extra_lines=} contains repetitions")
-
-        mem_per_job_MB = self.parallel_tasks_per_job * self.mem_per_task_MB
-        lines = [
-            self.shebang_line,
-            f"{self.prefix} --partition={self.partition}",
-            f"{self.prefix} --ntasks={self.parallel_tasks_per_job}",
-            f"{self.prefix} --cpus-per-task={self.cpus_per_task}",
-            f"{self.prefix} --mem={mem_per_job_MB}M",
-        ]
-        for key in [
-            "job_name",
-            "constraint",
-            "gres",
-            "gpus",
-            "time",
-            "account",
-            "exclude",
-            "nodelist",
-        ]:
-            value = getattr(self, key)
-            if value is not None:
-                # Handle the `time` parameter
-                if key == "time" and self.parallel_tasks_per_job > 1:
-                    # NOTE: see issue #1632
-                    logger.warning(
-                        f"`time` SLURM parameter is set to {self.time}, "
-                        "but this does not take into account the number of "
-                        f"SLURM tasks ({self.parallel_tasks_per_job})."
-                    )
-                option = key.replace("_", "-")
-                lines.append(f"{self.prefix} --{option}={value}")
-
-        if self.extra_lines:
-            for line in self._sorted_extra_lines():
-                lines.append(line)
-
-        if self.user_local_exports:
-            if remote_export_dir is None:
-                raise ValueError(
-                    f"remote_export_dir=None but {self.user_local_exports=}"
-                )
-            for key, value in self.user_local_exports.items():
-                tmp_value = str(Path(remote_export_dir) / value)
-                lines.append(f"export {key}={tmp_value}")
-
-        """
-        FIXME export SRUN_CPUS_PER_TASK
-        # From https://slurm.schedmd.com/sbatch.html: Beginning with 22.05,
-        # srun will not inherit the --cpus-per-task value requested by salloc
-        # or sbatch.  It must be requested again with the call to srun or set
-        # with the SRUN_CPUS_PER_TASK environment variable if desired for the
-        # task(s).
-        if config.cpus_per_task:
-            #additional_setup_lines.append(
-                f"export SRUN_CPUS_PER_TASK={config.cpus_per_task}"
-            )
-        """
-
-        return lines
-
-    @property
-    def batch_size(self) -> int:
-        return self.tasks_per_job
-
-
-def _parse_mem_value(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`.
-
-    Arguments:
-        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