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

fractal_server/app/runner/v2/runner_functions.py

@@ -0,0 +1,374 @@
+import functools
+import logging
+import traceback
+from concurrent.futures import Executor
+from pathlib import Path
+from typing import Any
+from typing import Callable
+from typing import Literal
+from typing import Optional
+
+from pydantic import ValidationError
+
+from ..exceptions import JobExecutionError
+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 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.components import _COMPONENT_KEY_
+from fractal_server.app.runner.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 _cast_and_validate_TaskOutput(
+    task_output: dict[str, Any]
+) -> Optional[TaskOutput]:
+    try:
+        validated_task_output = TaskOutput(**task_output)
+        return validated_task_output
+    except ValidationError as e:
+        raise JobExecutionError(
+            "Validation of task output failed.\n"
+            f"Original error: {str(e)}\n"
+            f"Original data: {task_output}."
+        )
+
+
+def _cast_and_validate_InitTaskOutput(
+    init_task_output: dict[str, Any],
+) -> Optional[InitTaskOutput]:
+    try:
+        validated_init_task_output = InitTaskOutput(**init_task_output)
+        return validated_init_task_output
+    except ValidationError as e:
+        raise JobExecutionError(
+            "Validation of init-task output failed.\n"
+            f"Original error: {str(e)}\n"
+            f"Original data: {init_task_output}."
+        )
+
+
+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 JobExecutionError(
+            "Too many parallelization items.\n"
+            f"   {len(my_list)}\n"
+            f"   {MAX_PARALLELIZATION_LIST_SIZE=}\n"
+        )
+
+
+def run_v2_task_non_parallel(
+    *,
+    images: list[dict[str, Any]],
+    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 workflow_dir_user is None:
+        workflow_dir_user = workflow_dir
+        logging.warning(
+            "In `run_single_task`, workflow_dir_user=None. Is this right?"
+        )
+        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(
+        zarr_urls=[image["zarr_url"] 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()
+    if output is None:
+        return TaskOutput()
+    else:
+        return _cast_and_validate_TaskOutput(output)
+
+
+def run_v2_task_parallel(
+    *,
+    images: list[dict[str, Any]],
+    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:
+
+    if len(images) == 0:
+        return 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(
+                zarr_url=image["zarr_url"],
+                **(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:
+            outputs[ind] = _cast_and_validate_TaskOutput(output)
+
+    merged_output = merge_outputs(outputs)
+    return merged_output
+
+
+def run_v2_task_compound(
+    *,
+    images: list[dict[str, Any]],
+    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(
+        zarr_urls=[image["zarr_url"] 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 = _cast_and_validate_InitTaskOutput(output)
+    parallelization_list = init_task_output.parallelization_list
+    parallelization_list = deduplicate_list(parallelization_list)
+
+    # 3/B: parallel part of a compound task
+    _check_parallelization_list_size(parallelization_list)
+
+    if len(parallelization_list) == 0:
+        return TaskOutput()
+
+    list_function_kwargs = []
+    for ind, parallelization_item in enumerate(parallelization_list):
+        list_function_kwargs.append(
+            dict(
+                zarr_url=parallelization_item.zarr_url,
+                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:
+            validated_output = _cast_and_validate_TaskOutput(output)
+            outputs[ind] = validated_output
+
+    merged_output = merge_outputs(outputs)
+    return merged_output
+
+
+def run_v1_task_parallel(
+    *,
+    images: list[dict[str, Any]],
+    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(
+                kwargs_v2=dict(
+                    zarr_url=image["zarr_url"],
+                    **(wftask.args_parallel or {}),
+                ),
+                parallelization_level=task_legacy.parallelization_level,
+            ),
+        )
+        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

fractal_server/app/runner/v2/runner_functions_low_level.py

@@ -0,0 +1,130 @@
+import json
+import logging
+import shutil
+import subprocess  # nosec
+from pathlib import Path
+from shlex import split as shlex_split
+from typing import Any
+from typing import Optional
+
+from ..components import _COMPONENT_KEY_
+from ..exceptions import JobExecutionError
+from ..exceptions import TaskExecutionError
+from fractal_server.app.models.v2 import WorkflowTaskV2
+from fractal_server.app.runner.task_files import get_task_file_paths
+
+
+def _call_command_wrapper(cmd: str, log_path: Path) -> None:
+    """
+    Call a command and write its stdout and stderr to files
+
+    Raises:
+        TaskExecutionError: If the `subprocess.run` call returns a positive
+                            exit code
+        JobExecutionError:  If the `subprocess.run` call returns a negative
+                            exit code (e.g. due to the subprocess receiving a
+                            TERM or KILL signal)
+    """
+
+    # Verify that task command is executable
+    if shutil.which(shlex_split(cmd)[0]) is None:
+        msg = (
+            f'Command "{shlex_split(cmd)[0]}" is not valid. '
+            "Hint: make sure that it is executable."
+        )
+        raise TaskExecutionError(msg)
+
+    fp_log = open(log_path, "w")
+    try:
+        result = subprocess.run(  # nosec
+            shlex_split(cmd),
+            stderr=fp_log,
+            stdout=fp_log,
+        )
+    except Exception as e:
+        raise e
+    finally:
+        fp_log.close()
+
+    if result.returncode > 0:
+        with log_path.open("r") as fp_stderr:
+            err = fp_stderr.read()
+        raise TaskExecutionError(err)
+    elif result.returncode < 0:
+        raise JobExecutionError(
+            info=f"Task failed with returncode={result.returncode}"
+        )
+
+
+def run_single_task(
+    args: dict[str, Any],
+    command: str,
+    wftask: WorkflowTaskV2,
+    workflow_dir: Path,
+    workflow_dir_user: Optional[Path] = None,
+    logger_name: Optional[str] = None,
+    is_task_v1: bool = False,
+) -> dict[str, Any]:
+    """
+    Runs within an executor.
+    """
+
+    logger = logging.getLogger(logger_name)
+    logger.debug(f"Now start running {command=}")
+
+    if not workflow_dir_user:
+        workflow_dir_user = workflow_dir
+
+    component = args.pop(_COMPONENT_KEY_, None)
+    task_files = get_task_file_paths(
+        workflow_dir=workflow_dir,
+        workflow_dir_user=workflow_dir_user,
+        task_order=wftask.order,
+        component=component,
+    )
+
+    # Write arguments to args.json file
+    with task_files.args.open("w") as f:
+        json.dump(args, f, indent=2)
+
+    # Assemble full command
+    if is_task_v1:
+        full_command = (
+            f"{command} "
+            f"--json {task_files.args.as_posix()} "
+            f"--metadata-out {task_files.metadiff.as_posix()}"
+        )
+    else:
+        full_command = (
+            f"{command} "
+            f"--args-json {task_files.args.as_posix()} "
+            f"--out-json {task_files.metadiff.as_posix()}"
+        )
+
+    try:
+        _call_command_wrapper(
+            full_command,
+            log_path=task_files.log,
+        )
+    except TaskExecutionError as e:
+        e.workflow_task_order = wftask.order
+        e.workflow_task_id = wftask.id
+        if wftask.is_legacy_task:
+            e.task_name = wftask.task_legacy.name
+        else:
+            e.task_name = wftask.task.name
+        raise e
+
+    try:
+        with task_files.metadiff.open("r") as f:
+            out_meta = json.load(f)
+    except FileNotFoundError as e:
+        logger.debug(
+            "Task did not produce output metadata. "
+            f"Original FileNotFoundError: {str(e)}"
+        )
+        out_meta = None
+
+    if out_meta == {}:
+        return None
+    return out_meta

fractal_server/app/runner/v2/task_interface.py

@@ -0,0 +1,62 @@
+from typing import Any
+
+from pydantic import BaseModel
+from pydantic import Field
+from pydantic import validator
+
+from ....images import SingleImageTaskOutput
+from fractal_server.images import Filters
+from fractal_server.urls import normalize_url
+
+
+class TaskOutput(BaseModel):
+    class Config:
+        extra = "forbid"
+
+    image_list_updates: list[SingleImageTaskOutput] = Field(
+        default_factory=list
+    )
+    image_list_removals: list[str] = Field(default_factory=list)
+    filters: Filters = Field(default_factory=Filters)
+
+    def check_zarr_urls_are_unique(self) -> None:
+        zarr_urls = [img.zarr_url for img in self.image_list_updates]
+        zarr_urls.extend(self.image_list_removals)
+        if len(zarr_urls) != len(set(zarr_urls)):
+            duplicates = [
+                zarr_url
+                for zarr_url in set(zarr_urls)
+                if zarr_urls.count(zarr_url) > 1
+            ]
+            msg = (
+                "TaskOutput "
+                f"({len(self.image_list_updates)} image_list_updates and "
+                f"{len(self.image_list_removals)} image_list_removals) "
+                "has non-unique zarr_urls:"
+            )
+            for duplicate in duplicates:
+                msg = f"{msg}\n{duplicate}"
+            raise ValueError(msg)
+
+    @validator("image_list_removals")
+    def normalize_paths(cls, v: list[str]) -> list[str]:
+        return [normalize_url(zarr_url) for zarr_url in v]
+
+
+class InitArgsModel(BaseModel):
+    class Config:
+        extra = "forbid"
+
+    zarr_url: str
+    init_args: dict[str, Any] = Field(default_factory=dict)
+
+    @validator("zarr_url")
+    def normalize_path(cls, v: str) -> str:
+        return normalize_url(v)
+
+
+class InitTaskOutput(BaseModel):
+    class Config:
+        extra = "forbid"
+
+    parallelization_list: list[InitArgsModel] = Field(default_factory=list)

fractal_server/app/runner/v2/v1_compat.py

@@ -0,0 +1,31 @@
+from copy import deepcopy
+from pathlib import Path
+from typing import Any
+
+
+def convert_v2_args_into_v1(
+    kwargs_v2: dict[str, Any],
+    parallelization_level: str = "image",
+) -> dict[str, Any]:
+    kwargs_v1 = deepcopy(kwargs_v2)
+
+    zarr_url = kwargs_v1.pop("zarr_url")
+    input_path = Path(zarr_url).parents[3].as_posix()
+    image_component = zarr_url.replace(input_path, "").lstrip("/")
+    if parallelization_level == "image":
+        component = image_component
+    elif parallelization_level == "well":
+        component = str(Path(image_component).parent)
+    elif parallelization_level == "plate":
+        component = str(Path(image_component).parents[2])
+    else:
+        raise ValueError(f"Invalid {parallelization_level=}.")
+
+    kwargs_v1.update(
+        input_paths=[input_path],
+        output_path=input_path,
+        metadata={},
+        component=component,
+    )
+
+    return kwargs_v1

fractal_server/app/schemas/__init__.py

@@ -1,42 +1 @@
-"""
-Schemas for API request/response bodies
-"""
-from .applyworkflow import ApplyWorkflowCreate  # noqa: F401
-from .applyworkflow import ApplyWorkflowRead  # noqa: F401
-from .applyworkflow import ApplyWorkflowUpdate  # noqa: F401
-from .applyworkflow import JobStatusType  # noqa: F401
-from .dataset import DatasetCreate  # noqa: F401
-from .dataset import DatasetRead  # noqa: F401
-from .dataset import DatasetStatusRead  # noqa: F401
-from .dataset import DatasetUpdate  # noqa: F401
-from .dataset import ResourceCreate  # noqa: F401
-from .dataset import ResourceRead  # noqa: F401
-from .dataset import ResourceUpdate  # noqa: F401
-from .manifest import ManifestV1  # noqa: F401
-from .manifest import TaskManifestV1  # noqa: F401
-from .project import ProjectCreate  # noqa: F401
-from .project import ProjectRead  # noqa: F401
-from .project import ProjectUpdate  # noqa: F401
-from .state import _StateBase  # noqa: F401
-from .state import StateRead  # noqa: F401
-from .task import TaskCreate  # noqa: F401
-from .task import TaskImport  # noqa: F401
-from .task import TaskRead  # noqa: F401
-from .task import TaskUpdate  # noqa: F401
-from .task_collection import TaskCollectPip  # noqa: F401
-from .task_collection import TaskCollectStatus  # noqa: F401
-from .user import UserCreate  # noqa: F401
-from .user import UserRead  # noqa: F401
-from .user import UserUpdate  # noqa: F401
-from .user import UserUpdateStrict  # noqa: F401
-from .workflow import WorkflowCreate  # noqa: F401
-from .workflow import WorkflowExport  # noqa: F401
-from .workflow import WorkflowImport  # noqa: F401
-from .workflow import WorkflowRead  # noqa: F401
-from .workflow import WorkflowTaskCreate  # noqa: F401
-from .workflow import WorkflowTaskExport  # noqa: F401
-from .workflow import WorkflowTaskImport  # noqa: F401
-from .workflow import WorkflowTaskRead  # noqa: F401
-from .workflow import WorkflowTaskStatusType  # noqa: F401
-from .workflow import WorkflowTaskUpdate  # noqa: F401
-from .workflow import WorkflowUpdate  # noqa: F401
+from .user import *  # noqa: F401, F403