fractal-server 1.4.9__py3-none-any.whl → 2.0.0a0__py3-none-any.whl

fractal_server/app/runner/v2/_slurm/get_slurm_config.py

@@ -0,0 +1,179 @@
+from pathlib import Path
+from typing import Literal
+from typing import Optional
+
+from fractal_server.app.models.v2 import WorkflowTaskV2
+from fractal_server.app.runner.executors.slurm._slurm_config import (
+    _parse_mem_value,
+)
+from fractal_server.app.runner.executors.slurm._slurm_config import (
+    load_slurm_config_file,
+)
+from fractal_server.app.runner.executors.slurm._slurm_config import logger
+from fractal_server.app.runner.executors.slurm._slurm_config import SlurmConfig
+from fractal_server.app.runner.executors.slurm._slurm_config import (
+    SlurmConfigError,
+)
+
+
+def get_slurm_config(
+    wftask: WorkflowTaskV2,
+    workflow_dir: Path,
+    workflow_dir_user: Path,
+    which_type: Literal["non_parallel", "parallel"],
+    config_path: Optional[Path] = None,
+) -> SlurmConfig:
+    """
+    Prepare a `SlurmConfig` configuration object
+
+    The argument `which_type` determines whether we use `wftask.meta_parallel`
+    or `wftask.meta_non_parallel`. In the following descritpion, let us assume
+    that `which_type="parallel"`.
+
+    The sources for `SlurmConfig` attributes, in increasing priority order, are
+
+    1. The general content of the Fractal SLURM configuration file.
+    2. The GPU-specific content of the Fractal SLURM configuration file, if
+        appropriate.
+    3. Properties in `wftask.meta_parallel` (which typically include those in
+       `wftask.task.meta_parallel`). Note that `wftask.meta_parallel` may be
+       `None`.
+
+    Arguments:
+        wftask:
+            WorkflowTask for which the SLURM configuration is is to be
+            prepared.
+        workflow_dir:
+            Server-owned directory to store all task-execution-related relevant
+            files (inputs, outputs, errors, and all meta files related to the
+            job execution). Note: users cannot write directly to this folder.
+        workflow_dir_user:
+            User-side directory with the same scope as `workflow_dir`, and
+            where a user can write.
+        config_path:
+            Path of a Fractal SLURM configuration file; if `None`, use
+            `FRACTAL_SLURM_CONFIG_FILE` variable from settings.
+        which_type:
+            Determines whether to use `meta_parallel` or `meta_non_parallel`.
+
+    Returns:
+        slurm_config:
+            The SlurmConfig object
+    """
+
+    if which_type == "non_parallel":
+        wftask_meta = wftask.meta_non_parallel
+    elif which_type == "parallel":
+        wftask_meta = wftask.meta_parallel
+    else:
+        raise ValueError(
+            f"get_slurm_config received invalid argument {which_type=}."
+        )
+
+    logger.debug(
+        "[get_slurm_config] WorkflowTask meta attribute: {wftask_meta=}"
+    )
+
+    # Incorporate slurm_env.default_slurm_config
+    slurm_env = load_slurm_config_file(config_path=config_path)
+    slurm_dict = slurm_env.default_slurm_config.dict(
+        exclude_unset=True, exclude={"mem"}
+    )
+    if slurm_env.default_slurm_config.mem:
+        slurm_dict["mem_per_task_MB"] = slurm_env.default_slurm_config.mem
+
+    # Incorporate slurm_env.batching_config
+    for key, value in slurm_env.batching_config.dict().items():
+        slurm_dict[key] = value
+
+    # Incorporate slurm_env.user_local_exports
+    slurm_dict["user_local_exports"] = slurm_env.user_local_exports
+
+    logger.debug(
+        "[get_slurm_config] Fractal SLURM configuration file: "
+        f"{slurm_env.dict()=}"
+    )
+
+    # GPU-related options
+    # Notes about priority:
+    # 1. This block of definitions takes priority over other definitions from
+    #    slurm_env which are not under the `needs_gpu` subgroup
+    # 2. This block of definitions has lower priority than whatever comes next
+    #    (i.e. from WorkflowTask.meta).
+    if wftask_meta is not None:
+        needs_gpu = wftask_meta.get("needs_gpu", False)
+    else:
+        needs_gpu = False
+    logger.debug(f"[get_slurm_config] {needs_gpu=}")
+    if needs_gpu:
+        for key, value in slurm_env.gpu_slurm_config.dict(
+            exclude_unset=True, exclude={"mem"}
+        ).items():
+            slurm_dict[key] = value
+        if slurm_env.gpu_slurm_config.mem:
+            slurm_dict["mem_per_task_MB"] = slurm_env.gpu_slurm_config.mem
+
+    # Number of CPUs per task, for multithreading
+    if wftask_meta is not None and "cpus_per_task" in wftask_meta:
+        cpus_per_task = int(wftask_meta["cpus_per_task"])
+        slurm_dict["cpus_per_task"] = cpus_per_task
+
+    # Required memory per task, in MB
+    if wftask_meta is not None and "mem" in wftask_meta:
+        raw_mem = wftask_meta["mem"]
+        mem_per_task_MB = _parse_mem_value(raw_mem)
+        slurm_dict["mem_per_task_MB"] = mem_per_task_MB
+
+    # Job name
+    job_name = wftask.task.name.replace(" ", "_")
+    slurm_dict["job_name"] = job_name
+
+    # Optional SLURM arguments and extra lines
+    if wftask_meta is not None:
+        account = wftask_meta.get("account", None)
+        if account is not None:
+            error_msg = (
+                f"Invalid {account=} property in WorkflowTask `meta` "
+                "attribute.\n"
+                "SLURM account must be set in the request body of the "
+                "apply-workflow endpoint, or by modifying the user properties."
+            )
+            logger.error(error_msg)
+            raise SlurmConfigError(error_msg)
+        for key in ["time", "gres", "constraint"]:
+            value = wftask_meta.get(key, None)
+            if value:
+                slurm_dict[key] = value
+    if wftask_meta is not None:
+        extra_lines = wftask_meta.get("extra_lines", [])
+    else:
+        extra_lines = []
+    extra_lines = slurm_dict.get("extra_lines", []) + extra_lines
+    if len(set(extra_lines)) != len(extra_lines):
+        logger.debug(
+            "[get_slurm_config] Removing repeated elements "
+            f"from {extra_lines=}."
+        )
+        extra_lines = list(set(extra_lines))
+    slurm_dict["extra_lines"] = extra_lines
+
+    # Job-batching parameters (if None, they will be determined heuristically)
+    if wftask_meta is not None:
+        tasks_per_job = wftask_meta.get("tasks_per_job", None)
+        parallel_tasks_per_job = wftask_meta.get(
+            "parallel_tasks_per_job", None
+        )
+    else:
+        tasks_per_job = None
+        parallel_tasks_per_job = None
+    slurm_dict["tasks_per_job"] = tasks_per_job
+    slurm_dict["parallel_tasks_per_job"] = parallel_tasks_per_job
+
+    # Put everything together
+    logger.debug(
+        "[get_slurm_config] Now create a SlurmConfig object based "
+        f"on {slurm_dict=}"
+    )
+    slurm_config = SlurmConfig(**slurm_dict)
+
+    return slurm_config

fractal_server/app/runner/v2/components.py

@@ -0,0 +1,5 @@
+def _index_to_component(ind: int) -> str:
+    return f"{ind:07d}"
+
+
+_COMPONENT_KEY_ = "__FRACTAL_PARALLEL_COMPONENT__"

fractal_server/app/runner/v2/deduplicate_list.py

@@ -0,0 +1,24 @@
+from typing import TypeVar
+
+from pydantic.main import ModelMetaclass
+
+from ....images import SingleImage
+from .task_interface import InitArgsModel
+
+T = TypeVar("T", SingleImage, InitArgsModel)
+
+
+def deduplicate_list(
+    this_list: list[T], PydanticModel: ModelMetaclass
+) -> list[T]:
+    """
+    Custom replacement for `set(this_list)`, when items are Pydantic-model
+    instances and then non-hashable (e.g. SingleImage or InitArgsModel).
+    """
+    this_list_dict = [this_item.dict() for this_item in this_list]
+    new_list_dict = []
+    for this_dict in this_list_dict:
+        if this_dict not in new_list_dict:
+            new_list_dict.append(this_dict)
+    new_list = [PydanticModel(**this_dict) for this_dict in new_list_dict]
+    return new_list

fractal_server/app/runner/v2/handle_failed_job.py

@@ -0,0 +1,156 @@
+# Copyright 2022 (C) Friedrich Miescher Institute for Biomedical Research and
+# University of Zurich
+#
+# Original authors:
+# Tommaso Comparin <tommaso.comparin@exact-lab.it>
+# Marco Franzon <marco.franzon@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.
+"""
+Helper functions to handle Dataset history.
+"""
+import json
+import logging
+from pathlib import Path
+from typing import Any
+from typing import Optional
+
+from ...models.v2 import DatasetV2
+from ...models.v2 import JobV2
+from ...models.v2 import WorkflowTaskV2
+from ...models.v2 import WorkflowV2
+from ...schemas.v2 import WorkflowTaskStatusTypeV2
+from ..filenames import FILTERS_FILENAME
+from ..filenames import HISTORY_FILENAME
+from ..filenames import IMAGES_FILENAME
+
+
+def assemble_history_failed_job(
+    job: JobV2,
+    dataset: DatasetV2,
+    workflow: WorkflowV2,
+    logger: logging.Logger,
+    failed_wftask: Optional[WorkflowTaskV2] = None,
+) -> list[dict[str, Any]]:
+    """
+    Assemble `history` after a workflow-execution job fails.
+
+    Args:
+        job:
+            The failed `ApplyWorkflow` object.
+        output_dataset:
+            The `dataset` associated to `job`.
+        workflow:
+            The `workflow` associated to `job`.
+        logger: A logger instance.
+        failed_wftask:
+            If set, append it to `history` during step 3; if `None`, infer
+            it by comparing the job task list and the one in
+            `HISTORY_FILENAME`.
+
+    Returns:
+        The new value of `history`, to be merged into
+        `output_dataset.meta`.
+    """
+
+    # The final value of the history attribute should include up to three
+    # parts, coming from: the database, the temporary file, the failed-task
+    # information.
+
+    # Part 1: Read exising history from DB
+    new_history = dataset.history
+
+    # Part 2: Extend history based on temporary-file contents
+    tmp_history_file = Path(job.working_dir) / HISTORY_FILENAME
+    try:
+        with tmp_history_file.open("r") as f:
+            tmp_file_history = json.load(f)
+            new_history.extend(tmp_file_history)
+    except FileNotFoundError:
+        tmp_file_history = []
+
+    # Part 3/A: Identify failed task, if needed
+    if failed_wftask is None:
+        job_wftasks = workflow.task_list[
+            job.first_task_index : (job.last_task_index + 1)  # noqa
+        ]
+        tmp_file_wftasks = [
+            history_item["workflowtask"] for history_item in tmp_file_history
+        ]
+        if len(job_wftasks) <= len(tmp_file_wftasks):
+            n_tasks_job = len(job_wftasks)
+            n_tasks_tmp = len(tmp_file_wftasks)
+            logger.error(
+                "Cannot identify the failed task based on job task list "
+                f"(length {n_tasks_job}) and temporary-file task list "
+                f"(length {n_tasks_tmp})."
+            )
+            logger.error("Failed task not appended to history.")
+        else:
+            failed_wftask = job_wftasks[len(tmp_file_wftasks)]
+
+    # Part 3/B: Append failed task to history
+    if failed_wftask is not None:
+        failed_wftask_dump = failed_wftask.model_dump(exclude={"task"})
+        failed_wftask_dump["task"] = failed_wftask.task.model_dump()
+        new_history_item = dict(
+            workflowtask=failed_wftask_dump,
+            status=WorkflowTaskStatusTypeV2.FAILED,
+            parallelization=dict(),  # FIXME: re-include parallelization
+        )
+        new_history.append(new_history_item)
+
+    return new_history
+
+
+def assemble_images_failed_job(job: JobV2) -> Optional[dict[str, Any]]:
+    """
+    Assemble `DatasetV2.images` for a failed workflow-execution.
+
+    Assemble new value of `images` based on the last successful task, i.e.
+    based on the content of the temporary `IMAGES_FILENAME` file. If the file
+    is missing, return `None`.
+
+    Argumentss:
+        job:
+            The failed `JobV2` object.
+
+    Returns:
+        The new value of `dataset.images`, or `None` if `IMAGES_FILENAME`
+        is missing.
+    """
+    tmp_file = Path(job.working_dir) / IMAGES_FILENAME
+    try:
+        with tmp_file.open("r") as f:
+            new_images = json.load(f)
+        return new_images
+    except FileNotFoundError:
+        return None
+
+
+def assemble_filters_failed_job(job: JobV2) -> Optional[dict[str, Any]]:
+    """
+    Assemble `DatasetV2.filters` for a failed workflow-execution.
+
+    Assemble new value of `filters` based on the last successful task, i.e.
+    based on the content of the temporary `FILTERS_FILENAME` file. If the file
+    is missing, return `None`.
+
+    Argumentss:
+        job:
+            The failed `JobV2` object.
+
+    Returns:
+        The new value of `dataset.filters`, or `None` if `FILTERS_FILENAME`
+        is missing.
+    """
+    tmp_file = Path(job.working_dir) / FILTERS_FILENAME
+    try:
+        with tmp_file.open("r") as f:
+            new_filters = json.load(f)
+        return new_filters
+    except FileNotFoundError:
+        return None

fractal_server/app/runner/v2/merge_outputs.py

@@ -0,0 +1,41 @@
+from copy import copy
+
+from fractal_server.app.runner.v2.deduplicate_list import deduplicate_list
+from fractal_server.app.runner.v2.task_interface import TaskOutput
+from fractal_server.images import SingleImage
+
+
+def merge_outputs(task_outputs: list[TaskOutput]) -> TaskOutput:
+
+    final_image_list_updates = []
+    final_image_list_removals = []
+    last_new_filters = None
+
+    for ind, task_output in enumerate(task_outputs):
+
+        final_image_list_updates.extend(task_output.image_list_updates)
+        final_image_list_removals.extend(task_output.image_list_removals)
+
+        # Check that all filters are the same
+        current_new_filters = task_output.filters
+        if ind == 0:
+            last_new_filters = copy(current_new_filters)
+        if current_new_filters != last_new_filters:
+            raise ValueError(f"{current_new_filters=} but {last_new_filters=}")
+        last_new_filters = copy(current_new_filters)
+
+    final_image_list_updates = deduplicate_list(
+        final_image_list_updates, PydanticModel=SingleImage
+    )
+
+    additional_args = {}
+    if last_new_filters is not None:
+        additional_args["filters"] = last_new_filters
+
+    final_output = TaskOutput(
+        image_list_updates=final_image_list_updates,
+        image_list_removals=final_image_list_removals,
+        **additional_args,
+    )
+
+    return final_output

fractal_server/app/runner/v2/runner.py

@@ -0,0 +1,264 @@
+import json
+from concurrent.futures import ThreadPoolExecutor
+from copy import copy
+from copy import deepcopy
+from pathlib import Path
+from typing import Callable
+from typing import Optional
+
+from ....images import Filters
+from ....images import SingleImage
+from ....images.tools import _filter_image_list
+from ....images.tools import find_image_by_path
+from ....images.tools import match_filter
+from ..filenames import FILTERS_FILENAME
+from ..filenames import HISTORY_FILENAME
+from ..filenames import IMAGES_FILENAME
+from .runner_functions import no_op_submit_setup_call
+from .runner_functions import run_v1_task_parallel
+from .runner_functions import run_v2_task_compound
+from .runner_functions import run_v2_task_non_parallel
+from .runner_functions import run_v2_task_parallel
+from fractal_server.app.models.v2 import DatasetV2
+from fractal_server.app.models.v2 import WorkflowTaskV2
+from fractal_server.app.schemas.v2.dataset import _DatasetHistoryItemV2
+from fractal_server.app.schemas.v2.workflowtask import WorkflowTaskStatusTypeV2
+
+# FIXME: define RESERVED_ARGUMENTS = [", ...]
+
+
+def execute_tasks_v2(
+    wf_task_list: list[WorkflowTaskV2],
+    dataset: DatasetV2,
+    executor: ThreadPoolExecutor,
+    workflow_dir: Path,
+    workflow_dir_user: Optional[Path] = None,
+    logger_name: Optional[str] = None,
+    submit_setup_call: Callable = no_op_submit_setup_call,
+) -> DatasetV2:
+
+    if not workflow_dir.exists():  # FIXME: this should have already happened
+        workflow_dir.mkdir()
+
+    # Initialize local dataset attributes
+    zarr_dir = dataset.zarr_dir
+    tmp_images = deepcopy(dataset.images)
+    tmp_filters = deepcopy(dataset.filters)
+    tmp_history = []
+
+    for wftask in wf_task_list:
+        task = wftask.task
+
+        # PRE TASK EXECUTION
+
+        # Get filtered images
+        pre_type_filters = copy(tmp_filters["types"])
+        pre_type_filters.update(wftask.input_filters["types"])
+        pre_attribute_filters = copy(tmp_filters["attributes"])
+        pre_attribute_filters.update(wftask.input_filters["attributes"])
+        filtered_images = _filter_image_list(
+            images=tmp_images,
+            filters=Filters(
+                types=pre_type_filters,
+                attributes=pre_attribute_filters,
+            ),
+        )
+        # Verify that filtered images comply with task input_types
+        for image in filtered_images:
+            if not match_filter(image, Filters(types=task.input_types)):
+                raise ValueError(
+                    f"Filtered images include {image.dict()}, which does "
+                    f"not comply with {task.input_types=}."
+                )
+
+        # TASK EXECUTION (V2)
+        if not wftask.is_legacy_task:
+            if task.type == "non_parallel":
+                current_task_output = run_v2_task_non_parallel(
+                    images=filtered_images,
+                    zarr_dir=zarr_dir,
+                    wftask=wftask,
+                    task=wftask.task,
+                    workflow_dir=workflow_dir,
+                    workflow_dir_user=workflow_dir_user,
+                    executor=executor,
+                    logger_name=logger_name,
+                    submit_setup_call=submit_setup_call,
+                )
+            elif task.type == "parallel":
+                current_task_output = run_v2_task_parallel(
+                    images=filtered_images,
+                    wftask=wftask,
+                    task=wftask.task,
+                    workflow_dir=workflow_dir,
+                    workflow_dir_user=workflow_dir_user,
+                    executor=executor,
+                    logger_name=logger_name,
+                    submit_setup_call=submit_setup_call,
+                )
+            elif task.type == "compound":
+                current_task_output = run_v2_task_compound(
+                    images=filtered_images,
+                    zarr_dir=zarr_dir,
+                    wftask=wftask,
+                    task=wftask.task,
+                    workflow_dir=workflow_dir,
+                    workflow_dir_user=workflow_dir_user,
+                    executor=executor,
+                    logger_name=logger_name,
+                    submit_setup_call=submit_setup_call,
+                )
+            else:
+                raise ValueError(f"Invalid {task.type=}.")
+        # TASK EXECUTION (V1)
+        else:
+            current_task_output = run_v1_task_parallel(
+                images=filtered_images,
+                wftask=wftask,
+                task_legacy=wftask.task_legacy,
+                executor=executor,
+                logger_name=logger_name,
+                submit_setup_call=submit_setup_call,
+            )
+
+        # POST TASK EXECUTION
+
+        # Update image list
+        current_task_output.check_paths_are_unique()
+        for image_obj in current_task_output.image_list_updates:
+            image = image_obj.dict()
+            # Edit existing image
+            if image["path"] in [_image["path"] for _image in tmp_images]:
+                if (
+                    image["origin"] is not None
+                    and image["origin"] != image["path"]
+                ):
+                    raise ValueError(
+                        f"Trying to edit an image with {image['path']=} "
+                        f"and {image['origin']=}."
+                    )
+                image_search = find_image_by_path(
+                    images=tmp_images,
+                    path=image["path"],
+                )
+                if image_search is None:
+                    raise ValueError(
+                        f"Image with path {image['path']} not found, while "
+                        "updating image list."
+                    )
+                original_img = image_search["image"]
+                original_index = image_search["index"]
+                updated_attributes = copy(original_img["attributes"])
+                updated_types = copy(original_img["types"])
+
+                # Update image attributes/types with task output and manifest
+                updated_attributes.update(image["attributes"])
+                updated_types.update(image["types"])
+                updated_types.update(task.output_types)
+
+                # Update image in the dataset image list
+                tmp_images[original_index]["attributes"] = updated_attributes
+                tmp_images[original_index]["types"] = updated_types
+            # Add new image
+            else:
+                # Check that image['path'] is relative to zarr_dir
+                if not image["path"].startswith(zarr_dir):
+                    raise ValueError(
+                        f"{zarr_dir} is not a parent directory of "
+                        f"{image['path']}"
+                    )
+                # Propagate attributes and types from `origin` (if any)
+                updated_attributes = {}
+                updated_types = {}
+                if image["origin"] is not None:
+                    image_search = find_image_by_path(
+                        images=tmp_images,
+                        path=image["origin"],
+                    )
+                    if image_search is not None:
+                        original_img = image_search["image"]
+                        updated_attributes = copy(original_img["attributes"])
+                        updated_types = copy(original_img["types"])
+                # Update image attributes/types with task output and manifest
+                updated_attributes.update(image["attributes"])
+                updated_types.update(image["types"])
+                updated_types.update(task.output_types)
+                new_image = SingleImage(
+                    path=image["path"],
+                    origin=image["origin"],
+                    attributes=updated_attributes,
+                    types=updated_types,
+                )
+                # Add image into the dataset image list
+                tmp_images.append(new_image.dict())
+
+        # Remove images from tmp_images
+        for image in current_task_output.image_list_removals:
+            image_search = find_image_by_path(
+                images=tmp_images, path=image["path"]
+            )
+            if image_search["index"] is None:
+                raise
+            else:
+                tmp_images.pop(image_search["index"])
+
+        # Update filters.attributes:
+        # current + (task_output: not really, in current examples..)
+        if current_task_output.filters is not None:
+            tmp_filters["attributes"].update(
+                current_task_output.filters.attributes
+            )
+
+        # Update filters.types: current + (task_output + task_manifest)
+        if wftask.is_legacy_task:
+            types_from_manifest = {}
+        else:
+            types_from_manifest = task.output_types
+        if current_task_output.filters is not None:
+            types_from_task = current_task_output.filters.types
+        else:
+            types_from_task = {}
+        # Check that key sets are disjoint
+        set_types_from_manifest = set(types_from_manifest.keys())
+        set_types_from_task = set(types_from_task.keys())
+        if not set_types_from_manifest.isdisjoint(set_types_from_task):
+            overlap = set_types_from_manifest.intersection(set_types_from_task)
+            raise ValueError(
+                "Both task and task manifest did set the same"
+                f"output type. Overlapping keys: {overlap}."
+            )
+        # Update filters.types
+        tmp_filters["types"].update(types_from_manifest)
+        tmp_filters["types"].update(types_from_task)
+
+        # Update history (based on _DatasetHistoryItemV2)
+        history_item = _DatasetHistoryItemV2(
+            workflowtask=wftask,
+            status=WorkflowTaskStatusTypeV2.DONE,
+            parallelization=dict(
+                # task_type=wftask.task.type,  # FIXME: breaks for V1 tasks
+                # component_list=fil, #FIXME
+            ),
+        ).dict()
+        tmp_history.append(history_item)
+
+        # Write current dataset attributes (history, images, filters) into
+        # temporary files which can be used (1) to retrieve the latest state
+        # when the job fails, (2) from within endpoints that need up-to-date
+        # information
+        with open(workflow_dir / HISTORY_FILENAME, "w") as f:
+            json.dump(tmp_history, f, indent=2)
+        with open(workflow_dir / FILTERS_FILENAME, "w") as f:
+            json.dump(tmp_filters, f, indent=2)
+        with open(workflow_dir / IMAGES_FILENAME, "w") as f:
+            json.dump(tmp_images, f, indent=2)
+
+    # NOTE: tmp_history only contains the newly-added history items (to be
+    # appended to the original history), while tmp_filters and tmp_images
+    # represent the new attributes (to replace the original ones)
+    result = dict(
+        history=tmp_history,
+        filters=tmp_filters,
+        images=tmp_images,
+    )
+    return result