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

fractal_server/app/routes/api/v2/_aux_functions.py

@@ -0,0 +1,497 @@
+"""
+Auxiliary functions to get object from the database or perform simple checks
+"""
+from typing import Any
+from typing import Literal
+from typing import Optional
+from typing import Union
+
+from fastapi import HTTPException
+from fastapi import status
+from sqlmodel import select
+from sqlmodel.sql.expression import SelectOfScalar
+
+from ....db import AsyncSession
+from ....models.v1 import Task
+from ....models.v2 import DatasetV2
+from ....models.v2 import JobV2
+from ....models.v2 import LinkUserProjectV2
+from ....models.v2 import ProjectV2
+from ....models.v2 import TaskV2
+from ....models.v2 import WorkflowTaskV2
+from ....models.v2 import WorkflowV2
+from ....schemas.v1 import JobStatusTypeV1
+from ....security import User
+from fractal_server.images import Filters
+
+
+async def _get_project_check_owner(
+    *,
+    project_id: int,
+    user_id: int,
+    db: AsyncSession,
+) -> ProjectV2:
+    """
+    Check that user is a member of project and return the project.
+
+    Args:
+        project_id:
+        user_id:
+        db:
+        version:
+
+    Returns:
+        The project object
+
+    Raises:
+        HTTPException(status_code=403_FORBIDDEN):
+            If the user is not a member of the project
+        HTTPException(status_code=404_NOT_FOUND):
+            If the project does not exist
+    """
+    project = await db.get(ProjectV2, project_id)
+
+    link_user_project = await db.get(LinkUserProjectV2, (project_id, user_id))
+    if not project:
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND, detail="Project not found"
+        )
+    if not link_user_project:
+        raise HTTPException(
+            status_code=status.HTTP_403_FORBIDDEN,
+            detail=f"Not allowed on project {project_id}",
+        )
+
+    return project
+
+
+async def _get_workflow_check_owner(
+    *,
+    workflow_id: int,
+    project_id: int,
+    user_id: int,
+    db: AsyncSession,
+) -> WorkflowV2:
+    """
+    Get a workflow and a project, after access control on the project.
+
+    Args:
+        workflow_id:
+        project_id:
+        user_id:
+        db:
+
+    Returns:
+        The workflow object.
+
+    Raises:
+        HTTPException(status_code=404_NOT_FOUND):
+            If the workflow does not exist
+        HTTPException(status_code=422_UNPROCESSABLE_ENTITY):
+            If the workflow is not associated to the project
+    """
+
+    # Access control for project
+    project = await _get_project_check_owner(
+        project_id=project_id, user_id=user_id, db=db
+    )
+    # Get workflow
+    workflow = await db.get(WorkflowV2, workflow_id)
+    if not workflow:
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND, detail="Workflow not found"
+        )
+    if workflow.project_id != project.id:
+        raise HTTPException(
+            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail=(f"Invalid {project_id=} for {workflow_id=}."),
+        )
+
+    # Refresh so that workflow.project relationship is loaded (see discussion
+    # in issue #1063)
+    await db.refresh(workflow)
+
+    return workflow
+
+
+async def _get_workflow_task_check_owner(
+    *,
+    project_id: int,
+    workflow_id: int,
+    workflow_task_id: int,
+    user_id: int,
+    db: AsyncSession,
+) -> tuple[WorkflowTaskV2, WorkflowV2]:
+    """
+    Check that user has access to Workflow and WorkflowTask.
+
+    Args:
+        project_id:
+        workflow_id:
+        workflow_task_id:
+        user_id:
+        db:
+
+    Returns:
+        Tuple of WorkflowTask and Workflow objects.
+
+    Raises:
+        HTTPException(status_code=404_NOT_FOUND):
+            If the WorkflowTask does not exist
+        HTTPException(status_code=422_UNPROCESSABLE_ENTITY):
+            If the WorkflowTask is not associated to the Workflow
+    """
+
+    # Access control for workflow
+    workflow = await _get_workflow_check_owner(
+        workflow_id=workflow_id,
+        project_id=project_id,
+        user_id=user_id,
+        db=db,
+    )
+
+    # If WorkflowTask is not in the db, exit
+    workflow_task = await db.get(WorkflowTaskV2, workflow_task_id)
+    if not workflow_task:
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND,
+            detail="WorkflowTask not found",
+        )
+
+    # If WorkflowTask is not part of the expected Workflow, exit
+    if workflow_id != workflow_task.workflow_id:
+        raise HTTPException(
+            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail=f"Invalid {workflow_id=} for {workflow_task_id=}",
+        )
+
+    return workflow_task, workflow
+
+
+async def _check_workflow_exists(
+    *,
+    name: str,
+    project_id: int,
+    db: AsyncSession,
+) -> None:
+    """
+    Check that no other workflow of this project has the same name.
+
+    Args:
+        name: Workflow name
+        project_id: Project ID
+        db:
+
+    Raises:
+        HTTPException(status_code=422_UNPROCESSABLE_ENTITY):
+            If such a workflow already exists
+    """
+    stm = (
+        select(WorkflowV2)
+        .where(WorkflowV2.name == name)
+        .where(WorkflowV2.project_id == project_id)
+    )
+    res = await db.execute(stm)
+    if res.scalars().all():
+        raise HTTPException(
+            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail=f"Workflow with {name=} and {project_id=} already exists.",
+        )
+
+
+async def _check_project_exists(
+    *,
+    project_name: str,
+    user_id: int,
+    db: AsyncSession,
+) -> None:
+    """
+    Check that no other project with this name exists for this user.
+
+    Args:
+        project_name: Project name
+        user_id: User ID
+        db:
+
+    Raises:
+        HTTPException(status_code=422_UNPROCESSABLE_ENTITY):
+            If such a project already exists
+    """
+    stm = (
+        select(ProjectV2)
+        .join(LinkUserProjectV2)
+        .where(ProjectV2.name == project_name)
+        .where(LinkUserProjectV2.user_id == user_id)
+    )
+    res = await db.execute(stm)
+    if res.scalars().all():
+        raise HTTPException(
+            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail=f"Project name ({project_name}) already in use",
+        )
+
+
+async def _get_dataset_check_owner(
+    *,
+    project_id: int,
+    dataset_id: int,
+    user_id: int,
+    db: AsyncSession,
+) -> dict[Literal["dataset", "project"], Union[DatasetV2, ProjectV2]]:
+    """
+    Get a dataset and a project, after access control on the project
+
+    Args:
+        project_id:
+        dataset_id:
+        user_id:
+        db:
+
+    Returns:
+        Dictionary with the dataset and project objects (keys: `dataset`,
+            `project`).
+
+    Raises:
+        HTTPException(status_code=422_UNPROCESSABLE_ENTITY):
+            If the dataset is not associated to the project
+    """
+
+    # Access control for project
+    project = await _get_project_check_owner(
+        project_id=project_id, user_id=user_id, db=db
+    )
+    # Get dataset
+    dataset = await db.get(DatasetV2, dataset_id)
+    if not dataset:
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND, detail="Dataset not found"
+        )
+    if dataset.project_id != project_id:
+        raise HTTPException(
+            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail=f"Invalid {project_id=} for {dataset_id=}",
+        )
+
+    # Refresh so that dataset.project relationship is loaded (see discussion
+    # in issue #1063)
+    await db.refresh(dataset)
+
+    return dict(dataset=dataset, project=project)
+
+
+async def _get_job_check_owner(
+    *,
+    project_id: int,
+    job_id: int,
+    user_id: int,
+    db: AsyncSession,
+) -> dict[Literal["job", "project"], Union[JobV2, ProjectV2]]:
+    """
+    Get a job and a project, after access control on the project
+
+    Args:
+        project_id:
+        job_id:
+        user_id:
+        db:
+
+    Returns:
+        Dictionary with the job and project objects (keys: `job`,
+            `project`).
+
+    Raises:
+        HTTPException(status_code=422_UNPROCESSABLE_ENTITY):
+            If the job is not associated to the project
+    """
+    # Access control for project
+    project = await _get_project_check_owner(
+        project_id=project_id,
+        user_id=user_id,
+        db=db,
+    )
+    # Get dataset
+    job = await db.get(JobV2, job_id)
+    if not job:
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND, detail="Job not found"
+        )
+    if job.project_id != project_id:
+        raise HTTPException(
+            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail=f"Invalid {project_id=} for {job_id=}",
+        )
+    return dict(job=job, project=project)
+
+
+async def _get_task_check_owner(
+    *,
+    task_id: int,
+    user: User,
+    db: AsyncSession,
+) -> TaskV2:
+    """
+    Get a task, after access control.
+
+    This check constitutes a preliminary version of access control:
+    if the current user is not a superuser and differs from the task owner
+    (including when `owner is None`), we raise an 403 HTTP Exception.
+
+    Args:
+        task_id:
+        user:
+        db:
+
+    Returns:
+        The task object.
+
+    Raises:
+        HTTPException(status_code=404_NOT_FOUND):
+            If the task does not exist
+        HTTPException(status_code=403_FORBIDDEN):
+            If the user does not have rights to edit this task.
+    """
+    task = await db.get(TaskV2, task_id)
+    if not task:
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND,
+            detail=f"TaskV2 {task_id} not found.",
+        )
+
+    if not user.is_superuser:
+        if task.owner is None:
+            raise HTTPException(
+                status_code=status.HTTP_403_FORBIDDEN,
+                detail=(
+                    "Only a superuser can modify a TaskV2 with `owner=None`."
+                ),
+            )
+        else:
+            owner = user.username or user.slurm_user
+            if owner != task.owner:
+                raise HTTPException(
+                    status_code=status.HTTP_403_FORBIDDEN,
+                    detail=(
+                        f"Current user ({owner}) cannot modify TaskV2 "
+                        f"{task.id} with different owner ({task.owner})."
+                    ),
+                )
+    return task
+
+
+def _get_submitted_jobs_statement() -> SelectOfScalar:
+    """
+    Returns:
+        A sqlmodel statement that selects all `ApplyWorkflow`s with
+        `ApplyWorkflow.status` equal to `submitted`.
+    """
+    stm = select(JobV2).where(JobV2.status == JobStatusTypeV1.SUBMITTED)
+    return stm
+
+
+async def _workflow_insert_task(
+    *,
+    workflow_id: int,
+    task_id: int,
+    is_legacy_task: bool = False,
+    order: Optional[int] = None,
+    meta_parallel: Optional[dict[str, Any]] = None,
+    meta_non_parallel: Optional[dict[str, Any]] = None,
+    args_non_parallel: Optional[dict[str, Any]] = None,
+    args_parallel: Optional[dict[str, Any]] = None,
+    input_filters: Optional[Filters] = None,
+    db: AsyncSession,
+) -> WorkflowTaskV2:
+    """
+    Insert a new WorkflowTask into Workflow.task_list
+
+    Args:
+        task_id: TBD
+        args: TBD
+        meta: TBD
+        order: TBD
+        db: TBD
+    """
+    db_workflow = await db.get(WorkflowV2, workflow_id)
+    if db_workflow is None:
+        raise ValueError(f"Workflow {workflow_id} does not exist")
+
+    if order is None:
+        order = len(db_workflow.task_list)
+
+    # Get task from db, and extract default arguments via a Task property
+    # method
+    if is_legacy_task is True:
+        db_task = await db.get(Task, task_id)
+        if db_task is None:
+            raise ValueError(f"Task {task_id} not found.")
+        task_type = "parallel"
+
+        final_args_parallel = db_task.default_args_from_args_schema.copy()
+        final_args_non_parallel = {}
+        final_meta_parallel = (db_task.meta or {}).copy()
+        final_meta_non_parallel = {}
+
+    else:
+        db_task = await db.get(TaskV2, task_id)
+        if db_task is None:
+            raise ValueError(f"TaskV2 {task_id} not found.")
+        task_type = db_task.type
+
+        final_args_non_parallel = (
+            db_task.default_args_non_parallel_from_args_schema.copy()
+        )
+        final_args_parallel = (
+            db_task.default_args_parallel_from_args_schema.copy()
+        )
+        final_meta_parallel = (db_task.meta_parallel or {}).copy()
+        final_meta_non_parallel = (db_task.meta_non_parallel or {}).copy()
+
+    # Combine arg_parallel
+    if args_parallel is not None:
+        for k, v in args_parallel.items():
+            final_args_parallel[k] = v
+    if final_args_parallel == {}:
+        final_args_parallel = None
+    # Combine arg_non_parallel
+    if args_non_parallel is not None:
+        for k, v in args_non_parallel.items():
+            final_args_non_parallel[k] = v
+    if final_args_non_parallel == {}:
+        final_args_non_parallel = None
+
+    # Combine meta_parallel (higher priority)
+    # and db_task.meta_parallel (lower priority)
+    final_meta_parallel.update(meta_parallel or {})
+    if final_meta_parallel == {}:
+        final_meta_parallel = None
+    # Combine meta_non_parallel (higher priority)
+    # and db_task.meta_non_parallel (lower priority)
+    final_meta_non_parallel.update(meta_non_parallel or {})
+    if final_meta_non_parallel == {}:
+        final_meta_non_parallel = None
+
+    # Prepare input_filters attribute
+    if input_filters is None:
+        input_filters_kwarg = {}
+    else:
+        input_filters_kwarg = dict(input_filters=input_filters)
+
+    # Create DB entry
+    wf_task = WorkflowTaskV2(
+        task_type=task_type,
+        is_legacy_task=is_legacy_task,
+        task_id=(task_id if not is_legacy_task else None),
+        task_legacy_id=(task_id if is_legacy_task else None),
+        args_non_parallel=final_args_non_parallel,
+        args_parallel=final_args_parallel,
+        meta_parallel=final_meta_parallel,
+        meta_non_parallel=final_meta_non_parallel,
+        **input_filters_kwarg,
+    )
+    db.add(wf_task)
+    db_workflow.task_list.insert(order, wf_task)
+    db_workflow.task_list.reorder()  # type: ignore
+    await db.commit()
+    await db.refresh(wf_task)
+
+    return wf_task

fractal_server/app/routes/api/v2/apply.py

@@ -0,0 +1,220 @@
+from datetime import datetime
+from datetime import timedelta
+from datetime import timezone
+from typing import Optional
+
+from fastapi import APIRouter
+from fastapi import BackgroundTasks
+from fastapi import Depends
+from fastapi import HTTPException
+from fastapi import status
+from sqlmodel import select
+
+from .....config import get_settings
+from .....syringe import Inject
+from ....db import AsyncSession
+from ....db import get_async_db
+from ....models.v2 import JobV2
+from ....runner.set_start_and_last_task_index import (
+    set_start_and_last_task_index,
+)  # FIXME V2
+from ....runner.v2 import submit_workflow
+from ....schemas.v2 import JobCreateV2
+from ....schemas.v2 import JobReadV2
+from ....schemas.v2 import JobStatusTypeV2
+from ....security import current_active_verified_user
+from ....security import User
+from ._aux_functions import _get_dataset_check_owner
+from ._aux_functions import _get_workflow_check_owner
+
+
+def _encode_as_utc(dt: datetime):
+    return dt.replace(tzinfo=timezone.utc).isoformat()
+
+
+router = APIRouter()
+
+
+@router.post(
+    "/project/{project_id}/workflow/{workflow_id}/apply/",
+    status_code=status.HTTP_202_ACCEPTED,
+    response_model=JobReadV2,
+)
+async def apply_workflow(
+    project_id: int,
+    workflow_id: int,
+    dataset_id: int,
+    job_create: JobCreateV2,
+    background_tasks: BackgroundTasks,
+    user: User = Depends(current_active_verified_user),
+    db: AsyncSession = Depends(get_async_db),
+) -> Optional[JobReadV2]:
+
+    output = await _get_dataset_check_owner(
+        project_id=project_id,
+        dataset_id=dataset_id,
+        user_id=user.id,
+        db=db,
+    )
+    project = output["project"]
+    dataset = output["dataset"]
+
+    if dataset.read_only:
+        raise HTTPException(
+            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail=(
+                "Cannot apply workflow because dataset "
+                f"({dataset_id=}) is read_only."
+            ),
+        )
+
+    workflow = await _get_workflow_check_owner(
+        project_id=project_id, workflow_id=workflow_id, user_id=user.id, db=db
+    )
+
+    if not workflow.task_list:
+        raise HTTPException(
+            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail=f"Workflow {workflow_id} has empty task list",
+        )
+
+    # Set values of first_task_index and last_task_index
+    num_tasks = len(workflow.task_list)
+    try:
+        first_task_index, last_task_index = set_start_and_last_task_index(
+            num_tasks,
+            first_task_index=job_create.first_task_index,
+            last_task_index=job_create.last_task_index,
+        )
+        job_create.first_task_index = first_task_index
+        job_create.last_task_index = last_task_index
+    except ValueError as e:
+        raise HTTPException(
+            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail=(
+                "Invalid values for first_task_index or last_task_index "
+                f"(with {num_tasks=}).\n"
+                f"Original error: {str(e)}"
+            ),
+        )
+
+    # If backend is SLURM, check that the user has required attributes
+    settings = Inject(get_settings)
+    backend = settings.FRACTAL_RUNNER_BACKEND
+    if backend == "slurm":
+        if not user.slurm_user:
+            raise HTTPException(
+                status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+                detail=(
+                    f"FRACTAL_RUNNER_BACKEND={backend}, "
+                    f"but {user.slurm_user=}."
+                ),
+            )
+        if not user.cache_dir:
+            raise HTTPException(
+                status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+                detail=(
+                    f"FRACTAL_RUNNER_BACKEND={backend}, "
+                    f"but {user.cache_dir=}."
+                ),
+            )
+
+    # Check that no other job with the same dataset_id is SUBMITTED
+    stm = (
+        select(JobV2)
+        .where(JobV2.dataset_id == dataset_id)
+        .where(JobV2.status == JobStatusTypeV2.SUBMITTED)
+    )
+    res = await db.execute(stm)
+    if res.scalars().all():
+        raise HTTPException(
+            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail=(
+                f"Dataset {dataset_id} is already in use "
+                "in submitted job(s)."
+            ),
+        )
+
+    if job_create.slurm_account is not None:
+        if job_create.slurm_account not in user.slurm_accounts:
+            raise HTTPException(
+                status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+                detail=(
+                    f"SLURM account '{job_create.slurm_account}' is not "
+                    "among those available to the current user"
+                ),
+            )
+    else:
+        if len(user.slurm_accounts) > 0:
+            job_create.slurm_account = user.slurm_accounts[0]
+
+    # Add new ApplyWorkflow object to DB
+    job = JobV2(
+        project_id=project_id,
+        dataset_id=dataset_id,
+        workflow_id=workflow_id,
+        user_email=user.email,
+        dataset_dump=dict(
+            **dataset.model_dump(exclude={"timestamp_created"}),
+            timestamp_created=_encode_as_utc(dataset.timestamp_created),
+        ),
+        workflow_dump=dict(
+            **workflow.model_dump(exclude={"task_list", "timestamp_created"}),
+            timestamp_created=_encode_as_utc(workflow.timestamp_created),
+        ),
+        project_dump=dict(
+            **project.model_dump(exclude={"user_list", "timestamp_created"}),
+            timestamp_created=_encode_as_utc(project.timestamp_created),
+        ),
+        **job_create.dict(),
+    )
+
+    # Rate Limiting:
+    # raise `429 TOO MANY REQUESTS` if this endpoint has been called with the
+    # same database keys (Project, Workflow and Datasets) during the last
+    # `settings.FRACTAL_API_SUBMIT_RATE_LIMIT` seconds.
+    stm = (
+        select(JobV2)
+        .where(JobV2.project_id == project_id)
+        .where(JobV2.workflow_id == workflow_id)
+        .where(JobV2.dataset_id == dataset_id)
+    )
+    res = await db.execute(stm)
+    db_jobs = res.scalars().all()
+    if db_jobs and any(
+        abs(
+            job.start_timestamp
+            - db_job.start_timestamp.replace(tzinfo=timezone.utc)
+        )
+        < timedelta(seconds=settings.FRACTAL_API_SUBMIT_RATE_LIMIT)
+        for db_job in db_jobs
+    ):
+        raise HTTPException(
+            status_code=status.HTTP_429_TOO_MANY_REQUESTS,
+            detail=(
+                f"The endpoint 'POST /api/v2/project/{project_id}/workflow/"
+                f"{workflow_id}/apply/' "
+                "was called several times within an interval of less "
+                f"than {settings.FRACTAL_API_SUBMIT_RATE_LIMIT} seconds, using"
+                " the same foreign keys. If it was intentional, please wait "
+                "and try again."
+            ),
+        )
+
+    db.add(job)
+    await db.commit()
+    await db.refresh(job)
+
+    background_tasks.add_task(
+        submit_workflow,
+        workflow_id=workflow.id,
+        dataset_id=dataset.id,
+        job_id=job.id,
+        worker_init=job.worker_init,
+        slurm_user=user.slurm_user,
+        user_cache_dir=user.cache_dir,
+    )
+
+    await db.close()
+
+    return job