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

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

@@ -0,0 +1,309 @@
+import json
+from pathlib import Path
+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 ....db import AsyncSession
+from ....db import get_async_db
+from ....models.v2 import DatasetV2
+from ....models.v2 import JobV2
+from ....models.v2 import ProjectV2
+from ....schemas.v2 import DatasetCreateV2
+from ....schemas.v2 import DatasetReadV2
+from ....schemas.v2 import DatasetUpdateV2
+from ....schemas.v2.dataset import DatasetStatusReadV2
+from ....schemas.v2.dataset import WorkflowTaskStatusTypeV2
+from ....security import current_active_user
+from ....security import User
+from ._aux_functions import _get_dataset_check_owner
+from ._aux_functions import _get_project_check_owner
+from ._aux_functions import _get_submitted_jobs_statement
+from ._aux_functions import _get_workflow_check_owner
+from fractal_server.app.runner.filenames import HISTORY_FILENAME
+
+router = APIRouter()
+
+
+@router.post(
+    "/project/{project_id}/dataset/",
+    response_model=DatasetReadV2,
+    status_code=status.HTTP_201_CREATED,
+)
+async def create_dataset(
+    project_id: int,
+    dataset: DatasetCreateV2,
+    user: User = Depends(current_active_user),
+    db: AsyncSession = Depends(get_async_db),
+) -> Optional[DatasetReadV2]:
+    """
+    Add new dataset to current project
+    """
+    await _get_project_check_owner(
+        project_id=project_id, user_id=user.id, db=db
+    )
+    db_dataset = DatasetV2(project_id=project_id, **dataset.dict())
+    db.add(db_dataset)
+    await db.commit()
+    await db.refresh(db_dataset)
+    await db.close()
+
+    return db_dataset
+
+
+@router.get(
+    "/project/{project_id}/dataset/",
+    response_model=list[DatasetReadV2],
+)
+async def read_dataset_list(
+    project_id: int,
+    history: bool = True,
+    user: User = Depends(current_active_user),
+    db: AsyncSession = Depends(get_async_db),
+) -> Optional[list[DatasetReadV2]]:
+    """
+    Get dataset list for given project
+    """
+    # Access control
+    project = await _get_project_check_owner(
+        project_id=project_id, user_id=user.id, db=db
+    )
+    # Find datasets of the current project. Note: this select/where approach
+    # has much better scaling than refreshing all elements of
+    # `project.dataset_list` - ref
+    # https://github.com/fractal-analytics-platform/fractal-server/pull/1082#issuecomment-1856676097.
+    stm = select(DatasetV2).where(DatasetV2.project_id == project.id)
+    res = await db.execute(stm)
+    dataset_list = res.scalars().all()
+    await db.close()
+    if not history:
+        for ds in dataset_list:
+            setattr(ds, "history", [])
+    return dataset_list
+
+
+@router.get(
+    "/project/{project_id}/dataset/{dataset_id}/",
+    response_model=DatasetReadV2,
+)
+async def read_dataset(
+    project_id: int,
+    dataset_id: int,
+    user: User = Depends(current_active_user),
+    db: AsyncSession = Depends(get_async_db),
+) -> Optional[DatasetReadV2]:
+    """
+    Get info on a dataset associated to the current project
+    """
+    output = await _get_dataset_check_owner(
+        project_id=project_id,
+        dataset_id=dataset_id,
+        user_id=user.id,
+        db=db,
+    )
+    dataset = output["dataset"]
+    await db.close()
+    return dataset
+
+
+@router.patch(
+    "/project/{project_id}/dataset/{dataset_id}/",
+    response_model=DatasetReadV2,
+)
+async def update_dataset(
+    project_id: int,
+    dataset_id: int,
+    dataset_update: DatasetUpdateV2,
+    user: User = Depends(current_active_user),
+    db: AsyncSession = Depends(get_async_db),
+) -> Optional[DatasetReadV2]:
+    """
+    Edit a dataset associated to the current project
+    """
+
+    output = await _get_dataset_check_owner(
+        project_id=project_id,
+        dataset_id=dataset_id,
+        user_id=user.id,
+        db=db,
+    )
+    db_dataset = output["dataset"]
+
+    for key, value in dataset_update.dict(exclude_unset=True).items():
+        setattr(db_dataset, key, value)
+
+    await db.commit()
+    await db.refresh(db_dataset)
+    await db.close()
+    return db_dataset
+
+
+@router.delete(
+    "/project/{project_id}/dataset/{dataset_id}/",
+    status_code=204,
+)
+async def delete_dataset(
+    project_id: int,
+    dataset_id: int,
+    user: User = Depends(current_active_user),
+    db: AsyncSession = Depends(get_async_db),
+) -> Response:
+    """
+    Delete a dataset associated to the current project
+    """
+    output = await _get_dataset_check_owner(
+        project_id=project_id,
+        dataset_id=dataset_id,
+        user_id=user.id,
+        db=db,
+    )
+    dataset = output["dataset"]
+
+    # Fail if there exist jobs that are submitted and in relation with the
+    # current dataset.
+    stm = _get_submitted_jobs_statement().where(JobV2.dataset_id == dataset_id)
+    res = await db.execute(stm)
+    jobs = res.scalars().all()
+    if jobs:
+        string_ids = str([job.id for job in jobs])[1:-1]
+        raise HTTPException(
+            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail=(
+                f"Cannot delete dataset {dataset.id} because it "
+                f"is linked to active job(s) {string_ids}."
+            ),
+        )
+
+    # Cascade operations: set foreign-keys to null for jobs which are in
+    # relationship with the current dataset
+    stm = select(JobV2).where(JobV2.dataset_id == dataset_id)
+    res = await db.execute(stm)
+    jobs = res.scalars().all()
+    for job in jobs:
+        job.dataset_id = None
+    await db.commit()
+
+    # Delete dataset
+    await db.delete(dataset)
+    await db.commit()
+
+    return Response(status_code=status.HTTP_204_NO_CONTENT)
+
+
+@router.get("/dataset/", response_model=list[DatasetReadV2])
+async def get_user_datasets(
+    history: bool = True,
+    user: User = Depends(current_active_user),
+    db: AsyncSession = Depends(get_async_db),
+) -> list[DatasetReadV2]:
+    """
+    Returns all the datasets of the current user
+    """
+    stm = select(DatasetV2)
+    stm = stm.join(ProjectV2).where(
+        ProjectV2.user_list.any(User.id == user.id)
+    )
+
+    res = await db.execute(stm)
+    dataset_list = res.scalars().all()
+    await db.close()
+    if not history:
+        for ds in dataset_list:
+            setattr(ds, "history", [])
+    return dataset_list
+
+
+@router.get(
+    "/project/{project_id}/dataset/{dataset_id}/status/",
+    response_model=DatasetStatusReadV2,
+)
+async def get_workflowtask_status(
+    project_id: int,
+    dataset_id: int,
+    user: User = Depends(current_active_user),
+    db: AsyncSession = Depends(get_async_db),
+) -> Optional[DatasetStatusReadV2]:
+    """
+    Extract the status of all `WorkflowTask`s that ran on a given `DatasetV2`.
+    """
+    # Get the dataset DB entry
+    output = await _get_dataset_check_owner(
+        project_id=project_id,
+        dataset_id=dataset_id,
+        user_id=user.id,
+        db=db,
+    )
+    dataset = output["dataset"]
+
+    # Check whether there exists a job such that
+    # 1. `job.dataset_id == dataset_id`, and
+    # 2. `job.status` is submitted
+    # If one such job exists, it will be used later. If there are multiple
+    # jobs, raise an error.
+    stm = _get_submitted_jobs_statement().where(JobV2.dataset_id == dataset_id)
+    res = await db.execute(stm)
+    running_jobs = res.scalars().all()
+    if len(running_jobs) == 0:
+        running_job = None
+    elif len(running_jobs) == 1:
+        running_job = running_jobs[0]
+    else:
+        string_ids = str([job.id for job in running_jobs])[1:-1]
+        raise HTTPException(
+            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail=(
+                f"Cannot get WorkflowTaskV2 statuses as DatasetV2 {dataset.id}"
+                f" is linked to multiple active jobs: {string_ids}."
+            ),
+        )
+
+    # Initialize empty dictionary for WorkflowTaskV2 status
+    workflow_tasks_status_dict: dict = {}
+
+    # Lowest priority: read status from DB, which corresponds to jobs that are
+    # not running
+    history = dataset.history
+    for history_item in history:
+        wftask_id = history_item["workflowtask"]["id"]
+        wftask_status = history_item["status"]
+        workflow_tasks_status_dict[wftask_id] = wftask_status
+
+    # If a job is running, then gather more up-to-date information
+    if running_job is not None:
+        # Get the workflow DB entry
+        running_workflow = await _get_workflow_check_owner(
+            project_id=project_id,
+            workflow_id=running_job.workflow_id,
+            user_id=user.id,
+            db=db,
+        )
+        # Mid priority: Set all WorkflowTask's that are part of the running job
+        # as "submitted"
+        start = running_job.first_task_index
+        end = running_job.last_task_index + 1
+        for wftask in running_workflow.task_list[start:end]:
+            workflow_tasks_status_dict[
+                wftask.id
+            ] = WorkflowTaskStatusTypeV2.SUBMITTED
+
+        # Highest priority: Read status updates coming from the running-job
+        # temporary file. Note: this file only contains information on
+        # # WorkflowTask's that ran through successfully.
+        tmp_file = Path(running_job.working_dir) / HISTORY_FILENAME
+        try:
+            with tmp_file.open("r") as f:
+                history = json.load(f)
+        except FileNotFoundError:
+            history = []
+        for history_item in history:
+            wftask_id = history_item["workflowtask"]["id"]
+            wftask_status = history_item["status"]
+            workflow_tasks_status_dict[wftask_id] = wftask_status
+
+    response_body = DatasetStatusReadV2(status=workflow_tasks_status_dict)
+    return response_body

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

@@ -0,0 +1,207 @@
+from typing import Any
+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 pydantic import BaseModel
+from pydantic import Field
+from sqlalchemy.orm.attributes import flag_modified
+
+from ._aux_functions import _get_dataset_check_owner
+from fractal_server.app.db import AsyncSession
+from fractal_server.app.db import get_async_db
+from fractal_server.app.security import current_active_user
+from fractal_server.app.security import User
+from fractal_server.images import Filters
+from fractal_server.images import SingleImage
+from fractal_server.images.tools import match_filter
+
+router = APIRouter()
+
+
+class ImagePage(BaseModel):
+
+    total_count: int
+    page_size: int
+    current_page: int
+
+    attributes: dict[str, list[Any]]
+    types: list[str]
+
+    images: list[SingleImage]
+
+
+class ImageQuery(BaseModel):
+    path: Optional[str]
+    filters: Filters = Field(default_factory=Filters)
+
+
+@router.post(
+    "/project/{project_id}/dataset/{dataset_id}/images/",
+    status_code=status.HTTP_201_CREATED,
+)
+async def post_new_image(
+    project_id: int,
+    dataset_id: int,
+    new_image: SingleImage,
+    user: User = Depends(current_active_user),
+    db: AsyncSession = Depends(get_async_db),
+) -> Response:
+
+    output = await _get_dataset_check_owner(
+        project_id=project_id, dataset_id=dataset_id, user_id=user.id, db=db
+    )
+    dataset = output["dataset"]
+
+    if new_image.path in dataset.image_paths:
+        raise HTTPException(
+            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail=(
+                f"Image with path '{new_image.path}' "
+                f"already in DatasetV2 {dataset_id}",
+            ),
+        )
+
+    dataset.images.append(new_image.dict())
+    flag_modified(dataset, "images")
+
+    await db.commit()
+
+    return Response(status_code=status.HTTP_201_CREATED)
+
+
+@router.post(
+    "/project/{project_id}/dataset/{dataset_id}/images/query/",
+    response_model=ImagePage,
+    status_code=status.HTTP_200_OK,
+)
+async def query_dataset_images(
+    project_id: int,
+    dataset_id: int,
+    use_dataset_filters: bool = False,  # query param
+    page: int = 1,  # query param
+    page_size: Optional[int] = None,  # query param
+    query: Optional[ImageQuery] = None,  # body
+    user: User = Depends(current_active_user),
+    db: AsyncSession = Depends(get_async_db),
+) -> ImagePage:
+
+    if page < 1:
+        raise HTTPException(
+            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail=f"Invalid pagination parameter: page={page} < 1",
+        )
+
+    output = await _get_dataset_check_owner(
+        project_id=project_id, dataset_id=dataset_id, user_id=user.id, db=db
+    )
+    dataset = output["dataset"]
+    images = dataset.images
+
+    if use_dataset_filters is True:
+        images = [
+            image
+            for image in images
+            if match_filter(image, Filters(**dataset.filters))
+        ]
+
+    attributes = {}
+    for image in images:
+        for k, v in image["attributes"].items():
+            attributes.setdefault(k, []).append(v)
+        for k, v in attributes.items():
+            attributes[k] = list(set(v))
+
+    types = list(
+        set(type for image in images for type in image["types"].keys())
+    )
+
+    if query is not None:
+
+        if query.path is not None:
+            image = next(
+                (image for image in images if image["path"] == query.path),
+                None,
+            )
+            if image is None:
+                images = []
+            else:
+                images = [image]
+
+        if query.filters.attributes or query.filters.types:
+            images = [
+                image
+                for image in images
+                if match_filter(
+                    image,
+                    Filters(**query.filters.dict()),
+                )
+            ]
+
+    total_count = len(images)
+
+    if page_size is not None:
+        if page_size <= 0:
+            raise HTTPException(
+                status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+                detail=(
+                    f"Invalid pagination parameter: page_size={page_size} <= 0"
+                ),
+            )
+    else:
+        page_size = total_count
+
+    if total_count == 0:
+        page = 1
+    else:
+        last_page = (total_count // page_size) + (total_count % page_size > 0)
+        if page > last_page:
+            page = last_page
+        offset = (page - 1) * page_size
+        images = images[offset : offset + page_size]  # noqa E203
+
+    return ImagePage(
+        total_count=total_count,
+        current_page=page,
+        page_size=page_size,
+        attributes=attributes,
+        types=types,
+        images=images,
+    )
+
+
+@router.delete(
+    "/project/{project_id}/dataset/{dataset_id}/images/",
+    status_code=status.HTTP_204_NO_CONTENT,
+)
+async def delete_dataset_images(
+    project_id: int,
+    dataset_id: int,
+    path: str,
+    user: User = Depends(current_active_user),
+    db: AsyncSession = Depends(get_async_db),
+) -> Response:
+
+    output = await _get_dataset_check_owner(
+        project_id=project_id, dataset_id=dataset_id, user_id=user.id, db=db
+    )
+    dataset = output["dataset"]
+
+    image_to_remove = next(
+        (image for image in dataset.images if image["path"] == path), None
+    )
+    if image_to_remove is None:
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND,
+            detail=f"No image with path '{path}' in DatasetV2 {dataset_id}.",
+        )
+
+    dataset.images.remove(image_to_remove)
+    flag_modified(dataset, "images")
+
+    await db.commit()
+
+    return Response(status_code=status.HTTP_204_NO_CONTENT)

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

@@ -0,0 +1,200 @@
+from pathlib import Path
+from typing import Optional
+
+from fastapi import APIRouter
+from fastapi import Depends
+from fastapi import Response
+from fastapi import status
+from fastapi.responses import StreamingResponse
+from sqlmodel import select
+
+from ....db import AsyncSession
+from ....db import get_async_db
+from ....models.v2 import JobV2
+from ....models.v2 import ProjectV2
+from ....runner.filenames import WORKFLOW_LOG_FILENAME  # FIXME
+from ....schemas.v2 import JobReadV2
+from ....schemas.v2 import JobStatusTypeV2
+from ....security import current_active_user
+from ....security import User
+from ...aux._job import _write_shutdown_file
+from ...aux._job import _zip_folder_to_byte_stream
+from ...aux._runner import _check_backend_is_slurm
+from ._aux_functions import _get_job_check_owner
+from ._aux_functions import _get_project_check_owner
+from ._aux_functions import _get_workflow_check_owner
+
+router = APIRouter()
+
+
+@router.get("/job/", response_model=list[JobReadV2])
+async def get_user_jobs(
+    user: User = Depends(current_active_user),
+    log: bool = True,
+    db: AsyncSession = Depends(get_async_db),
+) -> list[JobReadV2]:
+    """
+    Returns all the jobs of the current user
+    """
+    stm = (
+        select(JobV2)
+        .join(ProjectV2)
+        .where(ProjectV2.user_list.any(User.id == user.id))
+    )
+    res = await db.execute(stm)
+    job_list = res.scalars().all()
+    await db.close()
+    if not log:
+        for job in job_list:
+            setattr(job, "log", None)
+
+    return job_list
+
+
+@router.get(
+    "/project/{project_id}/workflow/{workflow_id}/job/",
+    response_model=list[JobReadV2],
+)
+async def get_workflow_jobs(
+    project_id: int,
+    workflow_id: int,
+    user: User = Depends(current_active_user),
+    db: AsyncSession = Depends(get_async_db),
+) -> Optional[list[JobReadV2]]:
+    """
+    Returns all the jobs related to a specific workflow
+    """
+    await _get_workflow_check_owner(
+        project_id=project_id, workflow_id=workflow_id, user_id=user.id, db=db
+    )
+    stm = select(JobV2).where(JobV2.workflow_id == workflow_id)
+    res = await db.execute(stm)
+    job_list = res.scalars().all()
+    return job_list
+
+
+@router.get(
+    "/project/{project_id}/job/{job_id}/",
+    response_model=JobReadV2,
+)
+async def read_job(
+    project_id: int,
+    job_id: int,
+    show_tmp_logs: bool = False,
+    user: User = Depends(current_active_user),
+    db: AsyncSession = Depends(get_async_db),
+) -> Optional[JobReadV2]:
+    """
+    Return info on an existing job
+    """
+
+    output = await _get_job_check_owner(
+        project_id=project_id,
+        job_id=job_id,
+        user_id=user.id,
+        db=db,
+    )
+    job = output["job"]
+    await db.close()
+
+    if show_tmp_logs and (job.status == JobStatusTypeV2.SUBMITTED):
+        try:
+            with open(f"{job.working_dir}/{WORKFLOW_LOG_FILENAME}", "r") as f:
+                job.log = f.read()
+        except FileNotFoundError:
+            pass
+
+    return job
+
+
+@router.get(
+    "/project/{project_id}/job/{job_id}/download/",
+    response_class=StreamingResponse,
+)
+async def download_job_logs(
+    project_id: int,
+    job_id: int,
+    user: User = Depends(current_active_user),
+    db: AsyncSession = Depends(get_async_db),
+) -> StreamingResponse:
+    """
+    Download job folder
+    """
+    output = await _get_job_check_owner(
+        project_id=project_id,
+        job_id=job_id,
+        user_id=user.id,
+        db=db,
+    )
+    job = output["job"]
+
+    # Create and return byte stream for zipped log folder
+    PREFIX_ZIP = Path(job.working_dir).name
+    zip_filename = f"{PREFIX_ZIP}_archive.zip"
+    byte_stream = _zip_folder_to_byte_stream(
+        folder=job.working_dir, zip_filename=zip_filename
+    )
+    return StreamingResponse(
+        iter([byte_stream.getvalue()]),
+        media_type="application/x-zip-compressed",
+        headers={"Content-Disposition": f"attachment;filename={zip_filename}"},
+    )
+
+
+@router.get(
+    "/project/{project_id}/job/",
+    response_model=list[JobReadV2],
+)
+async def get_job_list(
+    project_id: int,
+    user: User = Depends(current_active_user),
+    log: bool = True,
+    db: AsyncSession = Depends(get_async_db),
+) -> Optional[list[JobReadV2]]:
+    """
+    Get job list for given project
+    """
+    project = await _get_project_check_owner(
+        project_id=project_id, user_id=user.id, db=db
+    )
+
+    stm = select(JobV2).where(JobV2.project_id == project.id)
+    res = await db.execute(stm)
+    job_list = res.scalars().all()
+    await db.close()
+    if not log:
+        for job in job_list:
+            setattr(job, "log", None)
+
+    return job_list
+
+
+@router.get(
+    "/project/{project_id}/job/{job_id}/stop/",
+    status_code=202,
+)
+async def stop_job(
+    project_id: int,
+    job_id: int,
+    user: User = Depends(current_active_user),
+    db: AsyncSession = Depends(get_async_db),
+) -> Response:
+    """
+    Stop execution of a workflow job (only available for slurm backend)
+    """
+
+    # This endpoint is only implemented for SLURM backend
+    _check_backend_is_slurm()
+
+    # Get job from DB
+    output = await _get_job_check_owner(
+        project_id=project_id,
+        job_id=job_id,
+        user_id=user.id,
+        db=db,
+    )
+    job = output["job"]
+
+    _write_shutdown_file(job=job)
+
+    return Response(status_code=status.HTTP_202_ACCEPTED)