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

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

@@ -0,0 +1,210 @@
+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,
+)
+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}/job/submit/",
+    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"]
+
+    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 Job 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}/job/submit/'"
+                " 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

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

@@ -0,0 +1,222 @@
+from copy import deepcopy  # noqa
+from typing import Optional
+
+from fastapi import APIRouter
+from fastapi import Depends
+from fastapi import HTTPException
+from fastapi import Response
+from fastapi import status
+from sqlmodel import select
+
+from .....logger import set_logger
+from ....db import AsyncSession
+from ....db import get_async_db
+from ....models.v1 import Task as TaskV1
+from ....models.v2 import TaskV2
+from ....models.v2 import WorkflowTaskV2
+from ....schemas.v2 import TaskCreateV2
+from ....schemas.v2 import TaskReadV2
+from ....schemas.v2 import TaskUpdateV2
+from ....security import current_active_user
+from ....security import current_active_verified_user
+from ....security import User
+from ._aux_functions import _get_task_check_owner
+
+router = APIRouter()
+
+logger = set_logger(__name__)
+
+
+@router.get("/", response_model=list[TaskReadV2])
+async def get_list_task(
+    args_schema_parallel: bool = True,
+    args_schema_non_parallel: bool = True,
+    user: User = Depends(current_active_user),
+    db: AsyncSession = Depends(get_async_db),
+) -> list[TaskReadV2]:
+    """
+    Get list of available tasks
+    """
+    stm = select(TaskV2)
+    res = await db.execute(stm)
+    task_list = res.scalars().all()
+    await db.close()
+    if args_schema_parallel is False:
+        for task in task_list:
+            setattr(task, "args_schema_parallel", None)
+    if args_schema_non_parallel is False:
+        for task in task_list:
+            setattr(task, "args_schema_non_parallel", None)
+
+    return task_list
+
+
+@router.get("/{task_id}/", response_model=TaskReadV2)
+async def get_task(
+    task_id: int,
+    user: User = Depends(current_active_user),
+    db: AsyncSession = Depends(get_async_db),
+) -> TaskReadV2:
+    """
+    Get info on a specific task
+    """
+    task = await db.get(TaskV2, task_id)
+    await db.close()
+    if not task:
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND, detail="TaskV2 not found"
+        )
+    return task
+
+
+@router.patch("/{task_id}/", response_model=TaskReadV2)
+async def patch_task(
+    task_id: int,
+    task_update: TaskUpdateV2,
+    user: User = Depends(current_active_verified_user),
+    db: AsyncSession = Depends(get_async_db),
+) -> Optional[TaskReadV2]:
+    """
+    Edit a specific task (restricted to superusers and task owner)
+    """
+
+    # Retrieve task from database
+    db_task = await _get_task_check_owner(task_id=task_id, user=user, db=db)
+    update = task_update.dict(exclude_unset=True)
+
+    # Forbid changes that set a previously unset command
+    if db_task.type == "non_parallel" and "command_parallel" in update:
+        raise HTTPException(
+            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail="Cannot set an unset `command_parallel`.",
+        )
+    if db_task.type == "parallel" and "command_non_parallel" in update:
+        raise HTTPException(
+            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail="Cannot set an unset `command_non_parallel`.",
+        )
+
+    for key, value in update.items():
+        setattr(db_task, key, value)
+
+    await db.commit()
+    await db.refresh(db_task)
+    await db.close()
+    return db_task
+
+
+@router.post(
+    "/", response_model=TaskReadV2, status_code=status.HTTP_201_CREATED
+)
+async def create_task(
+    task: TaskCreateV2,
+    user: User = Depends(current_active_verified_user),
+    db: AsyncSession = Depends(get_async_db),
+) -> Optional[TaskReadV2]:
+    """
+    Create a new task
+    """
+
+    if task.command_non_parallel is None:
+        task_type = "parallel"
+    elif task.command_parallel is None:
+        task_type = "non_parallel"
+    else:
+        task_type = "compound"
+
+    if task_type == "parallel" and (
+        task.args_schema_non_parallel is not None
+        or task.meta_non_parallel is not None
+    ):
+        raise HTTPException(
+            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail=(
+                "Cannot set `TaskV2.args_schema_non_parallel` or "
+                "`TaskV2.args_schema_non_parallel` if TaskV2 is parallel"
+            ),
+        )
+    elif task_type == "non_parallel" and (
+        task.args_schema_parallel is not None or task.meta_parallel is not None
+    ):
+        raise HTTPException(
+            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail=(
+                "Cannot set `TaskV2.args_schema_parallel` or "
+                "`TaskV2.args_schema_parallel` if TaskV2 is non_parallel"
+            ),
+        )
+
+    # Set task.owner attribute
+    if user.username:
+        owner = user.username
+    elif user.slurm_user:
+        owner = user.slurm_user
+    else:
+        raise HTTPException(
+            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail=(
+                "Cannot add a new task because current user does not "
+                "have `username` or `slurm_user` attributes."
+            ),
+        )
+
+    # Prepend owner to task.source
+    task.source = f"{owner}:{task.source}"
+
+    # Verify that source is not already in use (note: this check is only useful
+    # to provide a user-friendly error message, but `task.source` uniqueness is
+    # already guaranteed by a constraint in the table definition).
+    stm = select(TaskV2).where(TaskV2.source == task.source)
+    res = await db.execute(stm)
+    if res.scalars().all():
+        raise HTTPException(
+            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail=f"Source '{task.source}' already used by some TaskV2",
+        )
+    stm = select(TaskV1).where(TaskV1.source == task.source)
+    res = await db.execute(stm)
+    if res.scalars().all():
+        raise HTTPException(
+            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail=f"Source '{task.source}' already used by some TaskV1",
+        )
+    # Add task
+    db_task = TaskV2(**task.dict(), owner=owner, type=task_type)
+    db.add(db_task)
+    await db.commit()
+    await db.refresh(db_task)
+    await db.close()
+    return db_task
+
+
+@router.delete("/{task_id}/", status_code=204)
+async def delete_task(
+    task_id: int,
+    user: User = Depends(current_active_user),
+    db: AsyncSession = Depends(get_async_db),
+) -> Response:
+    """
+    Delete a task
+    """
+
+    db_task = await _get_task_check_owner(task_id=task_id, user=user, db=db)
+
+    # Check that the TaskV2 is not in relationship with some WorkflowTaskV2
+    stm = select(WorkflowTaskV2).filter(WorkflowTaskV2.task_id == task_id)
+    res = await db.execute(stm)
+    workflowtask_list = res.scalars().all()
+    if workflowtask_list:
+        raise HTTPException(
+            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail=(
+                f"Cannot remove TaskV2 {task_id} because it is currently "
+                "imported in WorkflowsV2 "
+                f"{[x.workflow_id for x in workflowtask_list]}. "
+                "If you want to remove this task, then you should first remove"
+                " the workflows.",
+            ),
+        )
+
+    await db.delete(db_task)
+    await db.commit()
+    return Response(status_code=status.HTTP_204_NO_CONTENT)

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

@@ -0,0 +1,239 @@
+from pathlib import Path
+from shutil import copy as shell_copy
+from tempfile import TemporaryDirectory
+
+from fastapi import APIRouter
+from fastapi import BackgroundTasks
+from fastapi import Depends
+from fastapi import HTTPException
+from fastapi import Response
+from fastapi import status
+from pydantic.error_wrappers import ValidationError
+from sqlmodel import select
+
+from .....config import get_settings
+from .....logger import close_logger
+from .....logger import set_logger
+from .....syringe import Inject
+from ....db import AsyncSession
+from ....db import get_async_db
+from ....models.v2 import CollectionStateV2
+from ....models.v2 import TaskV2
+from ....schemas.state import StateRead
+from ....schemas.v2 import TaskCollectPipV2
+from ....schemas.v2 import TaskCollectStatusV2
+from ....security import current_active_user
+from ....security import current_active_verified_user
+from ....security import User
+from fractal_server.tasks.endpoint_operations import create_package_dir_pip
+from fractal_server.tasks.endpoint_operations import download_package
+from fractal_server.tasks.endpoint_operations import inspect_package
+from fractal_server.tasks.utils import get_collection_log
+from fractal_server.tasks.utils import slugify_task_name
+from fractal_server.tasks.v2._TaskCollectPip import _TaskCollectPip
+from fractal_server.tasks.v2.background_operations import (
+    background_collect_pip,
+)
+from fractal_server.tasks.v2.get_collection_data import get_collection_data
+
+router = APIRouter()
+
+logger = set_logger(__name__)
+
+
+@router.post(
+    "/collect/pip/",
+    response_model=StateRead,
+    responses={
+        201: dict(
+            description=(
+                "Task collection successfully started in the background"
+            )
+        ),
+        200: dict(
+            description=(
+                "Package already collected. Returning info on already "
+                "available tasks"
+            )
+        ),
+    },
+)
+async def collect_tasks_pip(
+    task_collect: TaskCollectPipV2,
+    background_tasks: BackgroundTasks,
+    response: Response,
+    user: User = Depends(current_active_verified_user),
+    db: AsyncSession = Depends(get_async_db),
+) -> StateRead:  # State[TaskCollectStatus]
+    """
+    Task collection endpoint
+
+    Trigger the creation of a dedicated virtual environment, the installation
+    of a package and the collection of tasks as advertised in the manifest.
+    """
+
+    logger = set_logger(logger_name="collect_tasks_pip")
+
+    # Validate payload as _TaskCollectPip, which has more strict checks than
+    # TaskCollectPip
+    try:
+        task_pkg = _TaskCollectPip(**task_collect.dict(exclude_unset=True))
+    except ValidationError as e:
+        raise HTTPException(
+            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail=f"Invalid task-collection object. Original error: {e}",
+        )
+
+    with TemporaryDirectory() as tmpdir:
+        try:
+            # Copy or download the package wheel file to tmpdir
+            if task_pkg.is_local_package:
+                shell_copy(task_pkg.package_path.as_posix(), tmpdir)
+                pkg_path = Path(tmpdir) / task_pkg.package_path.name
+            else:
+                pkg_path = await download_package(
+                    task_pkg=task_pkg, dest=tmpdir
+                )
+            # Read package info from wheel file, and override the ones coming
+            # from the request body
+            pkg_info = inspect_package(pkg_path)
+            task_pkg.package_name = pkg_info["pkg_name"]
+            task_pkg.package_version = pkg_info["pkg_version"]
+            task_pkg.package_manifest = pkg_info["pkg_manifest"]
+            task_pkg.check()
+        except Exception as e:
+            raise HTTPException(
+                status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+                detail=f"Invalid package or manifest. Original error: {e}",
+            )
+
+    try:
+        venv_path = create_package_dir_pip(task_pkg=task_pkg)
+    except FileExistsError:
+        venv_path = create_package_dir_pip(task_pkg=task_pkg, create=False)
+        try:
+            task_collect_status = get_collection_data(venv_path)
+            for task in task_collect_status.task_list:
+                db_task = await db.get(TaskV2, task.id)
+                if (
+                    (not db_task)
+                    or db_task.source != task.source
+                    or db_task.name != task.name
+                ):
+                    await db.close()
+                    raise HTTPException(
+                        status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+                        detail=(
+                            "Cannot collect package. Folder already exists, "
+                            f"but task {task.id} does not exists or it does "
+                            f"not have the expected source ({task.source}) or "
+                            f"name ({task.name})."
+                        ),
+                    )
+        except FileNotFoundError as e:
+            await db.close()
+            raise HTTPException(
+                status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+                detail=(
+                    "Cannot collect package. Possible reason: another "
+                    "collection of the same package is in progress. "
+                    f"Original FileNotFoundError: {e}"
+                ),
+            )
+        except ValidationError as e:
+            await db.close()
+            raise HTTPException(
+                status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+                detail=(
+                    "Cannot collect package. Possible reason: an old version "
+                    "of the same package has already been collected. "
+                    f"Original ValidationError: {e}"
+                ),
+            )
+        task_collect_status.info = "Already installed"
+        state = CollectionStateV2(data=task_collect_status.sanitised_dict())
+        response.status_code == status.HTTP_200_OK
+        await db.close()
+        return state
+    settings = Inject(get_settings)
+
+    # Check that tasks are not already in the DB
+    for new_task in task_pkg.package_manifest.task_list:
+        new_task_name_slug = slugify_task_name(new_task.name)
+        new_task_source = f"{task_pkg.package_source}:{new_task_name_slug}"
+        stm = select(TaskV2).where(TaskV2.source == new_task_source)
+        res = await db.execute(stm)
+        if res.scalars().all():
+            raise HTTPException(
+                status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+                detail=(
+                    "Cannot collect package. Task with source "
+                    f'"{new_task_source}" already exists in the database.'
+                ),
+            )
+
+    # All checks are OK, proceed with task collection
+    full_venv_path = venv_path.relative_to(settings.FRACTAL_TASKS_DIR)
+    collection_status = TaskCollectStatusV2(
+        status="pending", venv_path=full_venv_path, package=task_pkg.package
+    )
+
+    # Create State object (after casting venv_path to string)
+    collection_status_dict = collection_status.dict()
+    collection_status_dict["venv_path"] = str(collection_status.venv_path)
+    state = CollectionStateV2(data=collection_status_dict)
+    db.add(state)
+    await db.commit()
+    await db.refresh(state)
+
+    background_tasks.add_task(
+        background_collect_pip,
+        state_id=state.id,
+        venv_path=venv_path,
+        task_pkg=task_pkg,
+    )
+    logger.debug(
+        "Task-collection endpoint: start background collection "
+        "and return state"
+    )
+    close_logger(logger)
+    info = (
+        "Collecting tasks in the background. "
+        f"GET /task/collect/{state.id} to query collection status"
+    )
+    state.data["info"] = info
+    response.status_code = status.HTTP_201_CREATED
+    await db.close()
+
+    return state
+
+
+@router.get("/collect/{state_id}/", response_model=StateRead)
+async def check_collection_status(
+    state_id: int,
+    user: User = Depends(current_active_user),
+    verbose: bool = False,
+    db: AsyncSession = Depends(get_async_db),
+) -> StateRead:  # State[TaskCollectStatus]
+    """
+    Check status of background task collection
+    """
+    logger = set_logger(logger_name="check_collection_status")
+    logger.debug(f"Querying state for state.id={state_id}")
+    state = await db.get(CollectionStateV2, state_id)
+    if not state:
+        await db.close()
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND,
+            detail=f"No task collection info with id={state_id}",
+        )
+    data = TaskCollectStatusV2(**state.data)
+
+    # In some cases (i.e. a successful or ongoing task collection), data.log is
+    # not set; if so, we collect the current logs
+    if verbose and not data.log:
+        data.log = get_collection_log(data.venv_path)
+        state.data = data.sanitised_dict()
+    close_logger(logger)
+    await db.close()
+    return state