fractal-server 2.13.0__py3-none-any.whl → 2.14.0a0__py3-none-any.whl

fractal_server/app/runner/v2/runner.py

@@ -1,5 +1,5 @@
+import json
 import logging
-from concurrent.futures import ThreadPoolExecutor
 from copy import copy
 from copy import deepcopy
 from pathlib import Path
@@ -18,10 +18,14 @@ from .runner_functions import run_v2_task_non_parallel
 from .runner_functions import run_v2_task_parallel
 from .task_interface import TaskOutput
 from fractal_server.app.db import get_sync_db
+from fractal_server.app.history.status_enum import HistoryItemImageStatus
+from fractal_server.app.models.v2 import AccountingRecord
 from fractal_server.app.models.v2 import DatasetV2
+from fractal_server.app.models.v2 import HistoryItemV2
+from fractal_server.app.models.v2 import ImageStatus
+from fractal_server.app.models.v2 import TaskGroupV2
 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
+from fractal_server.app.runner.executors.base_runner import BaseRunner
 from fractal_server.images.models import AttributeFiltersType
 from fractal_server.images.tools import merge_type_filters
 
@@ -30,7 +34,8 @@ def execute_tasks_v2(
     *,
     wf_task_list: list[WorkflowTaskV2],
     dataset: DatasetV2,
-    executor: ThreadPoolExecutor,
+    runner: BaseRunner,
+    user_id: int,
     workflow_dir_local: Path,
     workflow_dir_remote: Optional[Path] = None,
     logger_name: Optional[str] = None,
@@ -41,8 +46,8 @@ def execute_tasks_v2(
 
     if not workflow_dir_local.exists():
         logger.warning(
-            f"Now creating {workflow_dir_local}, "
-            "but it should have already happened."
+            f"Now creating {workflow_dir_local}, but it "
+            "should have already happened."
         )
         workflow_dir_local.mkdir()
 
@@ -58,69 +63,116 @@ def execute_tasks_v2(
 
         # PRE TASK EXECUTION
 
-        # Get filtered images
+        # Filter images by types and attributes (in two steps)
         type_filters = copy(current_dataset_type_filters)
         type_filters_patch = merge_type_filters(
             task_input_types=task.input_types,
             wftask_type_filters=wftask.type_filters,
         )
         type_filters.update(type_filters_patch)
-        filtered_images = filter_image_list(
+        type_filtered_images = filter_image_list(
             images=tmp_images,
             type_filters=type_filters,
+            attribute_filters=None,
+        )
+        filtered_images = filter_image_list(
+            images=type_filtered_images,
+            type_filters=None,
             attribute_filters=job_attribute_filters,
         )
 
-        # First, set status SUBMITTED in dataset.history for each wftask
+        # Create history item
         with next(get_sync_db()) as db:
-            db_dataset = db.get(DatasetV2, dataset.id)
-            new_history_item = _DatasetHistoryItemV2(
-                workflowtask=dict(
-                    **wftask.model_dump(exclude={"task"}),
-                    task=wftask.task.model_dump(),
-                ),
-                status=WorkflowTaskStatusTypeV2.SUBMITTED,
-                parallelization=dict(),  # FIXME: re-include parallelization
-            ).model_dump()
-            db_dataset.history.append(new_history_item)
-            flag_modified(db_dataset, "history")
-            db.merge(db_dataset)
+            workflowtask_dump = dict(
+                **wftask.model_dump(exclude={"task"}),
+                task=wftask.task.model_dump(),
+            )
+            # Exclude timestamps since they'd need to be serialized properly
+            task_group = db.get(TaskGroupV2, wftask.task.taskgroupv2_id)
+            task_group_dump = task_group.model_dump(
+                exclude={
+                    "timestamp_created",
+                    "timestamp_last_used",
+                }
+            )
+            parameters_hash = str(
+                hash(
+                    json.dumps(
+                        [workflowtask_dump, task_group_dump],
+                        sort_keys=True,
+                        indent=None,
+                    ).encode("utf-8")
+                )
+            )
+            images = {
+                image["zarr_url"]: HistoryItemImageStatus.SUBMITTED
+                for image in filtered_images
+            }
+            history_item = HistoryItemV2(
+                dataset_id=dataset.id,
+                workflowtask_id=wftask.id,
+                workflowtask_dump=workflowtask_dump,
+                task_group_dump=task_group_dump,
+                parameters_hash=parameters_hash,
+                num_available_images=len(type_filtered_images),
+                num_current_images=len(filtered_images),
+                images=images,
+            )
+            db.add(history_item)
+            for image in filtered_images:
+                db.merge(
+                    ImageStatus(
+                        zarr_url=image["zarr_url"],
+                        workflowtask_id=wftask.id,
+                        dataset_id=dataset.id,
+                        parameters_hash=parameters_hash,
+                        status=HistoryItemImageStatus.SUBMITTED,
+                        logfile="/placeholder",
+                    )
+                )
             db.commit()
+            db.refresh(history_item)
+            history_item_id = history_item.id
+
         # TASK EXECUTION (V2)
         if task.type == "non_parallel":
-            current_task_output = run_v2_task_non_parallel(
+            (
+                current_task_output,
+                num_tasks,
+                exceptions,
+            ) = run_v2_task_non_parallel(
                 images=filtered_images,
                 zarr_dir=zarr_dir,
                 wftask=wftask,
                 task=task,
                 workflow_dir_local=workflow_dir_local,
                 workflow_dir_remote=workflow_dir_remote,
-                executor=executor,
-                logger_name=logger_name,
+                executor=runner,
                 submit_setup_call=submit_setup_call,
+                history_item_id=history_item_id,
             )
         elif task.type == "parallel":
-            current_task_output = run_v2_task_parallel(
+            current_task_output, num_tasks, exceptions = run_v2_task_parallel(
                 images=filtered_images,
                 wftask=wftask,
                 task=task,
                 workflow_dir_local=workflow_dir_local,
                 workflow_dir_remote=workflow_dir_remote,
-                executor=executor,
-                logger_name=logger_name,
+                executor=runner,
                 submit_setup_call=submit_setup_call,
+                history_item_id=history_item_id,
             )
         elif task.type == "compound":
-            current_task_output = run_v2_task_compound(
+            current_task_output, num_tasks, exceptions = run_v2_task_compound(
                 images=filtered_images,
                 zarr_dir=zarr_dir,
                 wftask=wftask,
                 task=task,
                 workflow_dir_local=workflow_dir_local,
                 workflow_dir_remote=workflow_dir_remote,
-                executor=executor,
-                logger_name=logger_name,
+                executor=runner,
                 submit_setup_call=submit_setup_call,
+                history_item_id=history_item_id,
             )
         else:
             raise ValueError(f"Unexpected error: Invalid {task.type=}.")
@@ -144,7 +196,10 @@ def execute_tasks_v2(
             )
 
         # Update image list
+        num_new_images = 0
         current_task_output.check_zarr_urls_are_unique()
+        # FIXME: Introduce for loop over task outputs, and processe them sequentially
+        # each failure should lead to an update of the specific image status
         for image_obj in current_task_output.image_list_updates:
             image = image_obj.model_dump()
             # Edit existing image
@@ -246,6 +301,7 @@ def execute_tasks_v2(
                 SingleImage(**new_image)
                 # Add image into the dataset image list
                 tmp_images.append(new_image)
+                num_new_images += 1
 
         # Remove images from tmp_images
         for img_zarr_url in current_task_output.image_list_removals:
@@ -269,7 +325,6 @@ def execute_tasks_v2(
         # information
         with next(get_sync_db()) as db:
             db_dataset = db.get(DatasetV2, dataset.id)
-            db_dataset.history[-1]["status"] = WorkflowTaskStatusTypeV2.DONE
             db_dataset.type_filters = current_dataset_type_filters
             db_dataset.images = tmp_images
             for attribute_name in [
@@ -281,4 +336,24 @@ def execute_tasks_v2(
             db.merge(db_dataset)
             db.commit()
 
+            # Create accounting record
+            record = AccountingRecord(
+                user_id=user_id,
+                num_tasks=num_tasks,
+                num_new_images=num_new_images,
+            )
+            db.add(record)
+            db.commit()
+
+        if exceptions != {}:
+            logger.error(
+                f'END    {wftask.order}-th task (name="{task_name}") '
+                "- ERROR."
+            )
+            # Raise first error
+            for key, value in exceptions.items():
+                raise JobExecutionError(
+                    info=(f"An error occurred.\nOriginal error:\n{value}")
+                )
+
         logger.debug(f'END    {wftask.order}-th task (name="{task_name}")')

fractal_server/app/runner/v2/runner_functions.py

@@ -1,7 +1,6 @@
 import functools
 import logging
 import traceback
-from concurrent.futures import Executor
 from pathlib import Path
 from typing import Any
 from typing import Callable
@@ -20,6 +19,7 @@ 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
+from fractal_server.app.runner.executors.base_runner import BaseRunner
 
 
 __all__ = [
@@ -59,13 +59,7 @@ def _cast_and_validate_InitTaskOutput(
         )
 
 
-def no_op_submit_setup_call(
-    *,
-    wftask: WorkflowTaskV2,
-    workflow_dir_local: Path,
-    workflow_dir_remote: Path,
-    which_type: Literal["non_parallel", "parallel"],
-) -> dict:
+def no_op_submit_setup_call(*args, **kwargs) -> dict:
     """
     Default (no-operation) interface of submit_setup_call in V2.
     """
@@ -84,8 +78,8 @@ def _get_executor_options(
     try:
         options = submit_setup_call(
             wftask=wftask,
-            workflow_dir_local=workflow_dir_local,
-            workflow_dir_remote=workflow_dir_remote,
+            root_dir_local=workflow_dir_local,
+            root_dir_remote=workflow_dir_remote,
             which_type=which_type,
         )
     except Exception as e:
@@ -114,10 +108,10 @@ def run_v2_task_non_parallel(
     wftask: WorkflowTaskV2,
     workflow_dir_local: Path,
     workflow_dir_remote: Optional[Path] = None,
-    executor: Executor,
-    logger_name: Optional[str] = None,
+    executor: BaseRunner,
     submit_setup_call: Callable = no_op_submit_setup_call,
-) -> TaskOutput:
+    history_item_id: int,
+) -> tuple[TaskOutput, int, dict[int, BaseException]]:
     """
     This runs server-side (see `executor` argument)
     """
@@ -142,22 +136,29 @@ def run_v2_task_non_parallel(
         zarr_dir=zarr_dir,
         **(wftask.args_non_parallel or {}),
     )
-    future = executor.submit(
+    function_kwargs[_COMPONENT_KEY_] = _index_to_component(0)
+
+    result, exception = executor.submit(
         functools.partial(
             run_single_task,
             wftask=wftask,
             command=task.command_non_parallel,
-            workflow_dir_local=workflow_dir_local,
-            workflow_dir_remote=workflow_dir_remote,
+            root_dir_local=workflow_dir_local,
+            root_dir_remote=workflow_dir_remote,
         ),
-        function_kwargs,
+        parameters=function_kwargs,
+        history_item_id=history_item_id,
         **executor_options,
     )
-    output = future.result()
-    if output is None:
-        return TaskOutput()
+
+    num_tasks = 1
+    if exception is None:
+        if result is None:
+            return (TaskOutput(), num_tasks, {})
+        else:
+            return (_cast_and_validate_TaskOutput(result), num_tasks, {})
     else:
-        return _cast_and_validate_TaskOutput(output)
+        return (TaskOutput(), num_tasks, {0: exception})
 
 
 def run_v2_task_parallel(
@@ -165,15 +166,15 @@ def run_v2_task_parallel(
     images: list[dict[str, Any]],
     task: TaskV2,
     wftask: WorkflowTaskV2,
-    executor: Executor,
+    executor: BaseRunner,
     workflow_dir_local: Path,
     workflow_dir_remote: Optional[Path] = None,
-    logger_name: Optional[str] = None,
     submit_setup_call: Callable = no_op_submit_setup_call,
-) -> TaskOutput:
+    history_item_id: int,
+) -> tuple[TaskOutput, int, dict[int, BaseException]]:
 
     if len(images) == 0:
-        return TaskOutput()
+        return (TaskOutput(), 0, {})
 
     _check_parallelization_list_size(images)
 
@@ -195,29 +196,36 @@ def run_v2_task_parallel(
         )
         list_function_kwargs[-1][_COMPONENT_KEY_] = _index_to_component(ind)
 
-    results_iterator = executor.map(
+    results, exceptions = executor.multisubmit(
         functools.partial(
             run_single_task,
             wftask=wftask,
             command=task.command_parallel,
-            workflow_dir_local=workflow_dir_local,
-            workflow_dir_remote=workflow_dir_remote,
+            root_dir_local=workflow_dir_local,
+            root_dir_remote=workflow_dir_remote,
         ),
-        list_function_kwargs,
+        list_parameters=list_function_kwargs,
+        history_item_id=history_item_id,
         **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()
+    outputs = []
+    for ind in range(len(list_function_kwargs)):
+        if ind in results.keys():
+            result = results[ind]
+            if result is None:
+                output = TaskOutput()
+            else:
+                output = _cast_and_validate_TaskOutput(result)
+            outputs.append(output)
+        elif ind in exceptions.keys():
+            print(f"Bad: {exceptions[ind]}")
         else:
-            outputs[ind] = _cast_and_validate_TaskOutput(output)
+            print("VERY BAD - should have not reached this point")
 
+    num_tasks = len(images)
     merged_output = merge_outputs(outputs)
-    return merged_output
+    return (merged_output, num_tasks, exceptions)
 
 
 def run_v2_task_compound(
@@ -226,12 +234,12 @@ def run_v2_task_compound(
     zarr_dir: str,
     task: TaskV2,
     wftask: WorkflowTaskV2,
-    executor: Executor,
+    executor: BaseRunner,
     workflow_dir_local: Path,
     workflow_dir_remote: Optional[Path] = None,
-    logger_name: Optional[str] = None,
     submit_setup_call: Callable = no_op_submit_setup_call,
-) -> TaskOutput:
+    history_item_id: int,
+) -> tuple[TaskOutput, int, dict[int, BaseException]]:
 
     executor_options_init = _get_executor_options(
         wftask=wftask,
@@ -254,30 +262,40 @@ def run_v2_task_compound(
         zarr_dir=zarr_dir,
         **(wftask.args_non_parallel or {}),
     )
-    future = executor.submit(
+    function_kwargs[_COMPONENT_KEY_] = f"init_{_index_to_component(0)}"
+    result, exception = executor.submit(
         functools.partial(
             run_single_task,
             wftask=wftask,
             command=task.command_non_parallel,
-            workflow_dir_local=workflow_dir_local,
-            workflow_dir_remote=workflow_dir_remote,
+            root_dir_local=workflow_dir_local,
+            root_dir_remote=workflow_dir_remote,
         ),
-        function_kwargs,
+        parameters=function_kwargs,
+        history_item_id=history_item_id,
+        in_compound_task=True,
         **executor_options_init,
     )
-    output = future.result()
-    if output is None:
-        init_task_output = InitTaskOutput()
+
+    num_tasks = 1
+    if exception is None:
+        if result is None:
+            init_task_output = InitTaskOutput()
+        else:
+            init_task_output = _cast_and_validate_InitTaskOutput(result)
     else:
-        init_task_output = _cast_and_validate_InitTaskOutput(output)
+        return (TaskOutput(), num_tasks, {0: exception})
+
     parallelization_list = init_task_output.parallelization_list
     parallelization_list = deduplicate_list(parallelization_list)
 
+    num_tasks = 1 + len(parallelization_list)
+
     # 3/B: parallel part of a compound task
     _check_parallelization_list_size(parallelization_list)
 
     if len(parallelization_list) == 0:
-        return TaskOutput()
+        return (TaskOutput(), 0, {})
 
     list_function_kwargs = []
     for ind, parallelization_item in enumerate(parallelization_list):
@@ -288,29 +306,35 @@ def run_v2_task_compound(
                 **(wftask.args_parallel or {}),
             ),
         )
-        list_function_kwargs[-1][_COMPONENT_KEY_] = _index_to_component(ind)
+        list_function_kwargs[-1][
+            _COMPONENT_KEY_
+        ] = f"compute_{_index_to_component(ind)}"
 
-    results_iterator = executor.map(
+    results, exceptions = executor.multisubmit(
         functools.partial(
             run_single_task,
             wftask=wftask,
             command=task.command_parallel,
-            workflow_dir_local=workflow_dir_local,
-            workflow_dir_remote=workflow_dir_remote,
+            root_dir_local=workflow_dir_local,
+            root_dir_remote=workflow_dir_remote,
         ),
-        list_function_kwargs,
+        list_parameters=list_function_kwargs,
+        history_item_id=history_item_id,
+        in_compound_task=True,
         **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
+    outputs = []
+    for ind in range(len(list_function_kwargs)):
+        if ind in results.keys():
+            result = results[ind]
+            if result is None:
+                output = TaskOutput()
+            else:
+                output = _cast_and_validate_TaskOutput(result)
+            outputs.append(output)
+        elif ind in exceptions.keys():
+            print(f"Bad: {exceptions[ind]}")
 
     merged_output = merge_outputs(outputs)
-    return merged_output
+    return (merged_output, num_tasks, exceptions)

fractal_server/app/runner/v2/runner_functions_low_level.py

@@ -11,11 +11,11 @@ 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
+from fractal_server.app.runner.task_files import TaskFiles
 from fractal_server.string_tools import validate_cmd
 
 
-def _call_command_wrapper(cmd: str, log_path: Path) -> None:
+def _call_command_wrapper(cmd: str, log_path: str) -> None:
     """
     Call a command and write its stdout and stderr to files
 
@@ -50,7 +50,7 @@ def _call_command_wrapper(cmd: str, log_path: Path) -> None:
             raise e
 
     if result.returncode > 0:
-        with log_path.open("r") as fp_stderr:
+        with open(log_path, "r") as fp_stderr:
             err = fp_stderr.read()
         raise TaskExecutionError(err)
     elif result.returncode < 0:
@@ -60,49 +60,49 @@ def _call_command_wrapper(cmd: str, log_path: Path) -> None:
 
 
 def run_single_task(
-    args: dict[str, Any],
+    parameters: dict[str, Any],
     command: str,
     wftask: WorkflowTaskV2,
-    workflow_dir_local: Path,
-    workflow_dir_remote: Optional[Path] = None,
+    root_dir_local: Path,
+    root_dir_remote: Optional[Path] = None,
     logger_name: Optional[str] = None,
 ) -> dict[str, Any]:
     """
-    Runs within an executor.
+    Runs within an executor (AKA on the SLURM cluster).
     """
 
     logger = logging.getLogger(logger_name)
     logger.debug(f"Now start running {command=}")
 
-    if not workflow_dir_remote:
-        workflow_dir_remote = workflow_dir_local
+    if not root_dir_remote:
+        root_dir_remote = root_dir_local
 
     task_name = wftask.task.name
 
-    component = args.pop(_COMPONENT_KEY_, None)
-    task_files = get_task_file_paths(
-        workflow_dir_local=workflow_dir_local,
-        workflow_dir_remote=workflow_dir_remote,
-        task_order=wftask.order,
+    component = parameters.pop(_COMPONENT_KEY_)
+    task_files = TaskFiles(
+        root_dir_local=root_dir_local,
+        root_dir_remote=root_dir_remote,
         task_name=task_name,
+        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)
+    with open(task_files.args_file_remote, "w") as f:
+        json.dump(parameters, f, indent=2)
 
     # Assemble full command
     full_command = (
         f"{command} "
-        f"--args-json {task_files.args.as_posix()} "
-        f"--out-json {task_files.metadiff.as_posix()}"
+        f"--args-json {task_files.args_file_remote} "
+        f"--out-json {task_files.metadiff_file_remote}"
     )
 
     try:
         _call_command_wrapper(
             full_command,
-            log_path=task_files.log,
+            log_path=task_files.log_file_remote,
         )
     except TaskExecutionError as e:
         e.workflow_task_order = wftask.order
@@ -111,7 +111,7 @@ def run_single_task(
         raise e
 
     try:
-        with task_files.metadiff.open("r") as f:
+        with open(task_files.metadiff_file_remote, "r") as f:
             out_meta = json.load(f)
     except FileNotFoundError as e:
         logger.debug(

fractal_server/app/schemas/v2/__init__.py

@@ -1,3 +1,4 @@
+from .accounting import AccountingRecordRead  # noqa F401
 from .dataset import DatasetCreateV2  # noqa F401
 from .dataset import DatasetExportV2  # noqa F401
 from .dataset import DatasetImportV2  # noqa F401

fractal_server/app/schemas/v2/accounting.py

@@ -0,0 +1,18 @@
+from datetime import datetime
+
+from pydantic import BaseModel
+from pydantic import field_serializer
+from pydantic.types import AwareDatetime
+
+
+class AccountingRecordRead(BaseModel):
+
+    id: int
+    user_id: int
+    timestamp: AwareDatetime
+    num_tasks: int
+    num_new_images: int
+
+    @field_serializer("timestamp")
+    def serialize_datetime(v: datetime) -> str:
+        return v.isoformat()

fractal_server/app/schemas/v2/dataset.py

@@ -14,27 +14,12 @@ from .._filter_validators import validate_attribute_filters
 from .._filter_validators import validate_type_filters
 from .._validators import root_validate_dict_keys
 from .._validators import valstr
-from .dumps import WorkflowTaskDumpV2
 from .project import ProjectReadV2
-from .workflowtask import WorkflowTaskStatusTypeV2
 from fractal_server.images import SingleImage
 from fractal_server.images.models import AttributeFiltersType
 from fractal_server.urls import normalize_url
 
 
-class _DatasetHistoryItemV2(BaseModel):
-    """
-    Class for an item of `Dataset.history`.
-    """
-
-    workflowtask: WorkflowTaskDumpV2
-    status: WorkflowTaskStatusTypeV2
-    parallelization: Optional[dict] = None
-
-
-# CRUD
-
-
 class DatasetCreateV2(BaseModel):
     model_config = ConfigDict(extra="forbid")
 
@@ -74,8 +59,6 @@ class DatasetReadV2(BaseModel):
     project_id: int
     project: ProjectReadV2
 
-    history: list[_DatasetHistoryItemV2]
-
     timestamp_created: AwareDatetime
 
     zarr_dir: str