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

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

fractal_server/app/runner/v2/runner_functions.py

@@ -0,0 +1,339 @@
+import functools
+import traceback
+from concurrent.futures import Executor
+from pathlib import Path
+from typing import Callable
+from typing import Literal
+from typing import Optional
+
+from ....images import SingleImage
+from .deduplicate_list import deduplicate_list
+from .merge_outputs import merge_outputs
+from .runner_functions_low_level import run_single_task
+from .task_interface import InitArgsModel
+from .task_interface import InitTaskOutput
+from .task_interface import TaskOutput
+from .v1_compat import convert_v2_args_into_v1
+from fractal_server.app.models.v1 import Task as TaskV1
+from fractal_server.app.models.v2 import TaskV2
+from fractal_server.app.models.v2 import WorkflowTaskV2
+from fractal_server.app.runner.v2.components import _COMPONENT_KEY_
+from fractal_server.app.runner.v2.components import _index_to_component
+
+
+__all__ = [
+    "run_v2_task_non_parallel",
+    "run_v2_task_parallel",
+    "run_v2_task_compound",
+    "run_v1_task_parallel",
+]
+
+MAX_PARALLELIZATION_LIST_SIZE = 20_000
+
+
+def no_op_submit_setup_call(
+    *,
+    wftask: WorkflowTaskV2,
+    workflow_dir: Path,
+    workflow_dir_user: Path,
+    which_type: Literal["non_parallel", "parallel"],
+) -> dict:
+    """
+    Default (no-operation) interface of submit_setup_call in V2.
+    """
+    return {}
+
+
+# Backend-specific configuration
+def _get_executor_options(
+    *,
+    wftask: WorkflowTaskV2,
+    workflow_dir: Path,
+    workflow_dir_user: Path,
+    submit_setup_call: Callable,
+    which_type: Literal["non_parallel", "parallel"],
+) -> dict:
+    try:
+        options = submit_setup_call(
+            wftask=wftask,
+            workflow_dir=workflow_dir,
+            workflow_dir_user=workflow_dir_user,
+            which_type=which_type,
+        )
+    except Exception as e:
+        tb = "".join(traceback.format_tb(e.__traceback__))
+        raise RuntimeError(
+            f"{type(e)} error in {submit_setup_call=}\n"
+            f"Original traceback:\n{tb}"
+        )
+    return options
+
+
+def _check_parallelization_list_size(my_list):
+    if len(my_list) > MAX_PARALLELIZATION_LIST_SIZE:
+        raise ValueError(
+            "Too many parallelization items.\n"
+            f"   {len(my_list)}\n"
+            f"   {MAX_PARALLELIZATION_LIST_SIZE=}\n"
+        )
+
+
+def run_v2_task_non_parallel(
+    *,
+    images: list[SingleImage],
+    zarr_dir: str,
+    task: TaskV2,
+    wftask: WorkflowTaskV2,
+    workflow_dir: Path,
+    workflow_dir_user: Optional[Path] = None,
+    executor: Executor,
+    logger_name: Optional[str] = None,
+    submit_setup_call: Callable = no_op_submit_setup_call,
+) -> TaskOutput:
+    """
+    This runs server-side (see `executor` argument)
+    """
+
+    if not workflow_dir_user:
+        workflow_dir_user = workflow_dir
+
+    executor_options = _get_executor_options(
+        wftask=wftask,
+        workflow_dir=workflow_dir,
+        workflow_dir_user=workflow_dir_user,
+        submit_setup_call=submit_setup_call,
+        which_type="non_parallel",
+    )
+
+    function_kwargs = dict(
+        paths=[image["path"] for image in images],
+        zarr_dir=zarr_dir,
+        **(wftask.args_non_parallel or {}),
+    )
+    future = executor.submit(
+        functools.partial(
+            run_single_task,
+            wftask=wftask,
+            command=task.command_non_parallel,
+            workflow_dir=workflow_dir,
+            workflow_dir_user=workflow_dir_user,
+        ),
+        function_kwargs,
+        **executor_options,
+    )
+    output = future.result()
+    # FIXME V2: handle validation errors
+    if output is None:
+        return TaskOutput()
+    else:
+        validated_output = TaskOutput(**output)
+        return validated_output
+
+
+def run_v2_task_parallel(
+    *,
+    images: list[SingleImage],
+    task: TaskV2,
+    wftask: WorkflowTaskV2,
+    executor: Executor,
+    workflow_dir: Path,
+    workflow_dir_user: Optional[Path] = None,
+    logger_name: Optional[str] = None,
+    submit_setup_call: Callable = no_op_submit_setup_call,
+) -> TaskOutput:
+
+    _check_parallelization_list_size(images)
+
+    executor_options = _get_executor_options(
+        wftask=wftask,
+        workflow_dir=workflow_dir,
+        workflow_dir_user=workflow_dir_user,
+        submit_setup_call=submit_setup_call,
+        which_type="parallel",
+    )
+
+    list_function_kwargs = []
+    for ind, image in enumerate(images):
+        list_function_kwargs.append(
+            dict(
+                path=image["path"],
+                **(wftask.args_parallel or {}),
+            ),
+        )
+        list_function_kwargs[-1][_COMPONENT_KEY_] = _index_to_component(ind)
+
+    results_iterator = executor.map(
+        functools.partial(
+            run_single_task,
+            wftask=wftask,
+            command=task.command_parallel,
+            workflow_dir=workflow_dir,
+            workflow_dir_user=workflow_dir_user,
+        ),
+        list_function_kwargs,
+        **executor_options,
+    )
+    # Explicitly iterate over the whole list, so that all futures are waited
+    outputs = list(results_iterator)
+
+    # Validate all non-None outputs
+    for ind, output in enumerate(outputs):
+        if output is None:
+            outputs[ind] = TaskOutput()
+        else:
+            # FIXME: improve handling of validation errors
+            validated_output = TaskOutput(**output)
+            outputs[ind] = validated_output
+
+    merged_output = merge_outputs(outputs)
+    return merged_output
+
+
+def run_v2_task_compound(
+    *,
+    images: list[SingleImage],
+    zarr_dir: str,
+    task: TaskV2,
+    wftask: WorkflowTaskV2,
+    executor: Executor,
+    workflow_dir: Path,
+    workflow_dir_user: Optional[Path] = None,
+    logger_name: Optional[str] = None,
+    submit_setup_call: Callable = no_op_submit_setup_call,
+) -> TaskOutput:
+
+    executor_options_init = _get_executor_options(
+        wftask=wftask,
+        workflow_dir=workflow_dir,
+        workflow_dir_user=workflow_dir_user,
+        submit_setup_call=submit_setup_call,
+        which_type="non_parallel",
+    )
+    executor_options_compute = _get_executor_options(
+        wftask=wftask,
+        workflow_dir=workflow_dir,
+        workflow_dir_user=workflow_dir_user,
+        submit_setup_call=submit_setup_call,
+        which_type="parallel",
+    )
+
+    # 3/A: non-parallel init task
+    function_kwargs = dict(
+        paths=[image["path"] for image in images],
+        zarr_dir=zarr_dir,
+        **(wftask.args_non_parallel or {}),
+    )
+    future = executor.submit(
+        functools.partial(
+            run_single_task,
+            wftask=wftask,
+            command=task.command_non_parallel,
+            workflow_dir=workflow_dir,
+            workflow_dir_user=workflow_dir_user,
+        ),
+        function_kwargs,
+        **executor_options_init,
+    )
+    output = future.result()
+    if output is None:
+        init_task_output = InitTaskOutput()
+    else:
+        init_task_output = InitTaskOutput(**output)
+    parallelization_list = init_task_output.parallelization_list
+    parallelization_list = deduplicate_list(
+        parallelization_list, PydanticModel=InitArgsModel
+    )
+
+    # 3/B: parallel part of a compound task
+    _check_parallelization_list_size(parallelization_list)
+
+    list_function_kwargs = []
+    for ind, parallelization_item in enumerate(parallelization_list):
+        list_function_kwargs.append(
+            dict(
+                path=parallelization_item.path,
+                init_args=parallelization_item.init_args,
+                **(wftask.args_parallel or {}),
+            ),
+        )
+        list_function_kwargs[-1][_COMPONENT_KEY_] = _index_to_component(ind)
+
+    results_iterator = executor.map(
+        functools.partial(
+            run_single_task,
+            wftask=wftask,
+            command=task.command_parallel,
+            workflow_dir=workflow_dir,
+            workflow_dir_user=workflow_dir_user,
+        ),
+        list_function_kwargs,
+        **executor_options_compute,
+    )
+    # Explicitly iterate over the whole list, so that all futures are waited
+    outputs = list(results_iterator)
+
+    # Validate all non-None outputs
+    for ind, output in enumerate(outputs):
+        if output is None:
+            outputs[ind] = TaskOutput()
+        else:
+            # FIXME: improve handling of validation errors
+            validated_output = TaskOutput(**output)
+            outputs[ind] = validated_output
+
+    merged_output = merge_outputs(outputs)
+    return merged_output
+
+
+def run_v1_task_parallel(
+    *,
+    images: list[SingleImage],
+    task_legacy: TaskV1,
+    wftask: WorkflowTaskV2,
+    executor: Executor,
+    workflow_dir: Path,
+    workflow_dir_user: Optional[Path] = None,
+    logger_name: Optional[str] = None,
+    submit_setup_call: Callable = no_op_submit_setup_call,
+) -> TaskOutput:
+
+    _check_parallelization_list_size(images)
+
+    executor_options = _get_executor_options(
+        wftask=wftask,
+        workflow_dir=workflow_dir,
+        workflow_dir_user=workflow_dir_user,
+        submit_setup_call=submit_setup_call,
+        which_type="parallel",
+    )
+
+    list_function_kwargs = []
+    for ind, image in enumerate(images):
+        list_function_kwargs.append(
+            convert_v2_args_into_v1(
+                dict(
+                    path=image["path"],
+                    **(wftask.args_parallel or {}),
+                )
+            ),
+        )
+        list_function_kwargs[-1][_COMPONENT_KEY_] = _index_to_component(ind)
+
+    results_iterator = executor.map(
+        functools.partial(
+            run_single_task,
+            wftask=wftask,
+            command=task_legacy.command,
+            workflow_dir=workflow_dir,
+            workflow_dir_user=workflow_dir_user,
+            is_task_v1=True,
+        ),
+        list_function_kwargs,
+        **executor_options,
+    )
+    # Explicitly iterate over the whole list, so that all futures are waited
+    list(results_iterator)
+
+    # Ignore any output metadata for V1 tasks, and return an empty object
+    out = TaskOutput()
+    return out