fractal-server 2.2.0a1__py3-none-any.whl → 2.3.0__py3-none-any.whl

fractal_server/__init__.py

@@ -1 +1 @@
-__VERSION__ = "2.2.0a1"
+__VERSION__ = "2.3.0"

fractal_server/app/models/v1/state.py

@@ -9,10 +9,9 @@ from sqlmodel import Field
 from sqlmodel import SQLModel
 
 from ....utils import get_timestamp
-from ...schemas.state import _StateBase
 
 
-class State(_StateBase, SQLModel, table=True):
+class State(SQLModel, table=True):
     """
     Store arbitrary data in the database
 

fractal_server/app/routes/admin/v1.py

@@ -35,7 +35,7 @@ from ...schemas.v1 import WorkflowReadV1
 from ...security import current_active_superuser
 from ..aux._job import _write_shutdown_file
 from ..aux._job import _zip_folder_to_byte_stream
-from ..aux._runner import _is_shutdown_available
+from ..aux._runner import _check_shutdown_is_supported
 
 router_admin_v1 = APIRouter()
 
@@ -351,7 +351,7 @@ async def stop_job(
     Stop execution of a workflow job.
     """
 
-    _is_shutdown_available()
+    _check_shutdown_is_supported()
 
     job = await db.get(ApplyWorkflow, job_id)
     if job is None:

fractal_server/app/routes/admin/v2.py

@@ -38,7 +38,7 @@ from ...schemas.v2 import ProjectReadV2
 from ...security import current_active_superuser
 from ..aux._job import _write_shutdown_file
 from ..aux._job import _zip_folder_to_byte_stream
-from ..aux._runner import _is_shutdown_available
+from ..aux._runner import _check_shutdown_is_supported
 
 router_admin_v2 = APIRouter()
 
@@ -238,7 +238,7 @@ async def stop_job(
     Stop execution of a workflow job.
     """
 
-    _is_shutdown_available()
+    _check_shutdown_is_supported()
 
     job = await db.get(JobV2, job_id)
     if job is None:

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

@@ -19,7 +19,7 @@ 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 _is_shutdown_available
+from ...aux._runner import _check_shutdown_is_supported
 from ._aux_functions import _get_job_check_owner
 from ._aux_functions import _get_project_check_owner
 from ._aux_functions import _get_workflow_check_owner
@@ -180,7 +180,7 @@ async def stop_job(
     Stop execution of a workflow job.
     """
 
-    _is_shutdown_available()
+    _check_shutdown_is_supported()
 
     # Get job from DB
     output = await _get_job_check_owner(

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

@@ -19,21 +19,21 @@ from ....db import AsyncSession
 from ....db import get_async_db
 from ....models.v1 import State
 from ....models.v1 import Task
-from ....schemas.state import StateRead
+from ....schemas.v1 import StateRead
 from ....schemas.v1 import TaskCollectPipV1
 from ....schemas.v1 import TaskCollectStatusV1
 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.v1._TaskCollectPip import _TaskCollectPip
 from fractal_server.tasks.v1.background_operations import (
     background_collect_pip,
 )
+from fractal_server.tasks.v1.endpoint_operations import create_package_dir_pip
+from fractal_server.tasks.v1.endpoint_operations import download_package
+from fractal_server.tasks.v1.endpoint_operations import inspect_package
 from fractal_server.tasks.v1.get_collection_data import get_collection_data
 
 router = APIRouter()

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

@@ -11,9 +11,14 @@ from .status import router as status_router_v2
 from .submit import router as submit_job_router_v2
 from .task import router as task_router_v2
 from .task_collection import router as task_collection_router_v2
+from .task_collection_custom import router as task_collection_router_v2_custom
+from .task_collection_ssh import router as task_collection_router_v2_ssh
 from .task_legacy import router as task_legacy_router_v2
 from .workflow import router as workflow_router_v2
 from .workflowtask import router as workflowtask_router_v2
+from fractal_server.config import get_settings
+from fractal_server.syringe import Inject
+
 
 router_api_v2 = APIRouter()
 
@@ -22,9 +27,24 @@ router_api_v2.include_router(job_router_v2, tags=["V2 Job"])
 router_api_v2.include_router(images_routes_v2, tags=["V2 Images"])
 router_api_v2.include_router(project_router_v2, tags=["V2 Project"])
 router_api_v2.include_router(submit_job_router_v2, tags=["V2 Job"])
-router_api_v2.include_router(
-    task_collection_router_v2, prefix="/task", tags=["V2 Task Collection"]
-)
+
+
+settings = Inject(get_settings)
+if settings.FRACTAL_RUNNER_BACKEND == "slurm_ssh":
+    router_api_v2.include_router(
+        task_collection_router_v2_ssh,
+        prefix="/task",
+        tags=["V2 Task Collection"],
+    )
+else:
+    router_api_v2.include_router(
+        task_collection_router_v2, prefix="/task", tags=["V2 Task Collection"]
+    )
+    router_api_v2.include_router(
+        task_collection_router_v2_custom,
+        prefix="/task",
+        tags=["V2 Task Collection"],
+    )
 router_api_v2.include_router(task_router_v2, prefix="/task", tags=["V2 Task"])
 router_api_v2.include_router(
     task_legacy_router_v2, prefix="/task-legacy", tags=["V2 Task Legacy"]

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

@@ -19,7 +19,7 @@ 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 _is_shutdown_available
+from ...aux._runner import _check_shutdown_is_supported
 from ._aux_functions import _get_job_check_owner
 from ._aux_functions import _get_project_check_owner
 from ._aux_functions import _get_workflow_check_owner
@@ -183,7 +183,7 @@ async def stop_job(
     Stop execution of a workflow job.
     """
 
-    _is_shutdown_available()
+    _check_shutdown_is_supported()
 
     # Get job from DB
     output = await _get_job_check_owner(

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

@@ -226,6 +226,11 @@ async def apply_workflow(
         WORKFLOW_DIR_REMOTE = (
             Path(user.cache_dir) / f"{WORKFLOW_DIR_LOCAL.name}"
         )
+    elif FRACTAL_RUNNER_BACKEND == "slurm_ssh":
+        WORKFLOW_DIR_REMOTE = (
+            Path(settings.FRACTAL_SLURM_SSH_WORKING_BASE_DIR)
+            / f"{WORKFLOW_DIR_LOCAL.name}"
+        )
 
     # Update job folders in the db
     job.working_dir = WORKFLOW_DIR_LOCAL.as_posix()
@@ -241,6 +246,7 @@ async def apply_workflow(
         worker_init=job.worker_init,
         slurm_user=user.slurm_user,
         user_cache_dir=user.cache_dir,
+        fractal_ssh=request.app.state.fractal_ssh,
     )
     request.app.state.jobsV2.append(job.id)
     logger.info(

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

@@ -1,3 +1,4 @@
+import json
 from pathlib import Path
 from shutil import copy as shell_copy
 from tempfile import TemporaryDirectory
@@ -19,22 +20,25 @@ 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 CollectionStateReadV2
+from ....schemas.v2 import CollectionStatusV2
 from ....schemas.v2 import TaskCollectPipV2
-from ....schemas.v2 import TaskCollectStatusV2
+from ....schemas.v2 import TaskReadV2
 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_absolute_venv_path
 from fractal_server.tasks.utils import get_collection_log
+from fractal_server.tasks.utils import get_collection_path
 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
+from fractal_server.tasks.v2.endpoint_operations import create_package_dir_pip
+from fractal_server.tasks.v2.endpoint_operations import download_package
+from fractal_server.tasks.v2.endpoint_operations import inspect_package
+
 
 router = APIRouter()
 
@@ -43,7 +47,7 @@ logger = set_logger(__name__)
 
 @router.post(
     "/collect/pip/",
-    response_model=StateRead,
+    response_model=CollectionStateReadV2,
     responses={
         201: dict(
             description=(
@@ -64,7 +68,7 @@ async def collect_tasks_pip(
     response: Response,
     user: User = Depends(current_active_verified_user),
     db: AsyncSession = Depends(get_async_db),
-) -> StateRead:  # State[TaskCollectStatus]
+) -> CollectionStateReadV2:
     """
     Task collection endpoint
 
@@ -74,6 +78,13 @@ async def collect_tasks_pip(
 
     logger = set_logger(logger_name="collect_tasks_pip")
 
+    # Set default python version
+    if task_collect.python_version is None:
+        settings = Inject(get_settings)
+        task_collect.python_version = (
+            settings.FRACTAL_TASKS_PYTHON_DEFAULT_VERSION
+        )
+
     # Validate payload as _TaskCollectPip, which has more strict checks than
     # TaskCollectPip
     try:
@@ -89,18 +100,18 @@ async def collect_tasks_pip(
             # 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
+                wheel_path = Path(tmpdir) / task_pkg.package_path.name
             else:
-                pkg_path = await download_package(
+                logger.info(f"Now download {task_pkg}")
+                wheel_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"]
+            # from the request body. Note that `package_name` was already set
+            # (and normalized) as part of `_TaskCollectPip` initialization.
+            pkg_info = inspect_package(wheel_path)
             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,
@@ -112,8 +123,35 @@ async def collect_tasks_pip(
     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:
+            package_path = get_absolute_venv_path(venv_path)
+            collection_path = get_collection_path(package_path)
+            with collection_path.open("r") as f:
+                task_collect_data = json.load(f)
+
+            err_msg = (
+                "Cannot collect package, possible reason: an old version of "
+                "the same package has already been collected.\n"
+                f"{str(collection_path)} has invalid content: "
+            )
+            if not isinstance(task_collect_data, dict):
+                raise HTTPException(
+                    status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+                    detail=f"{err_msg} it's not a Python dictionary.",
+                )
+            if "task_list" not in task_collect_data.keys():
+                raise HTTPException(
+                    status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+                    detail=f"{err_msg} it has no key 'task_list'.",
+                )
+            if not isinstance(task_collect_data["task_list"], list):
+                raise HTTPException(
+                    status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+                    detail=f"{err_msg} 'task_list' is not a Python list.",
+                )
+
+            for task_dict in task_collect_data["task_list"]:
+
+                task = TaskReadV2(**task_dict)
                 db_task = await db.get(TaskV2, task.id)
                 if (
                     (not db_task)
@@ -150,8 +188,8 @@ async def collect_tasks_pip(
                     f"Original ValidationError: {e}"
                 ),
             )
-        task_collect_status.info = "Already installed"
-        state = CollectionStateV2(data=task_collect_status.sanitised_dict())
+        task_collect_data["info"] = "Already installed"
+        state = CollectionStateV2(data=task_collect_data)
         response.status_code == status.HTTP_200_OK
         await db.close()
         return state
@@ -173,15 +211,12 @@ async def collect_tasks_pip(
             )
 
     # 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
+    collection_status = dict(
+        status=CollectionStatusV2.PENDING,
+        venv_path=venv_path.relative_to(settings.FRACTAL_TASKS_DIR).as_posix(),
+        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)
+    state = CollectionStateV2(data=collection_status)
     db.add(state)
     await db.commit()
     await db.refresh(state)
@@ -208,13 +243,13 @@ async def collect_tasks_pip(
     return state
 
 
-@router.get("/collect/{state_id}/", response_model=StateRead)
+@router.get("/collect/{state_id}/", response_model=CollectionStateReadV2)
 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]
+) -> CollectionStateReadV2:  # State[TaskCollectStatus]
     """
     Check status of background task collection
     """
@@ -227,13 +262,18 @@ async def check_collection_status(
             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()
+    # In some cases (i.e. a successful or ongoing task collection),
+    # state.data.log is not set; if so, we collect the current logs.
+    if verbose and not state.data.get("log"):
+        if "venv_path" not in state.data.keys():
+            await db.close()
+            raise HTTPException(
+                status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+                detail=f"No 'venv_path' in CollectionStateV2[{state_id}].data",
+            )
+        state.data["log"] = get_collection_log(Path(state.data["venv_path"]))
+        state.data["venv_path"] = str(state.data["venv_path"])
     reset_logger_handlers(logger)
     await db.close()
     return state

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

@@ -0,0 +1,170 @@
+import shlex
+import subprocess  # nosec
+from pathlib import Path
+
+from fastapi import APIRouter
+from fastapi import Depends
+from fastapi import HTTPException
+from fastapi import status
+from sqlmodel import select
+
+from .....config import get_settings
+from .....logger import set_logger
+from .....syringe import Inject
+from ....db import DBSyncSession
+from ....db import get_sync_db
+from ....models.v1 import Task as TaskV1
+from ....models.v2 import TaskV2
+from ....schemas.v2 import TaskCollectCustomV2
+from ....schemas.v2 import TaskCreateV2
+from ....schemas.v2 import TaskReadV2
+from ....security import current_active_verified_user
+from ....security import User
+from fractal_server.tasks.v2.background_operations import _insert_tasks
+from fractal_server.tasks.v2.background_operations import (
+    _prepare_tasks_metadata,
+)
+
+
+router = APIRouter()
+
+logger = set_logger(__name__)
+
+
+@router.post(
+    "/collect/custom/", status_code=201, response_model=list[TaskReadV2]
+)
+async def collect_task_custom(
+    task_collect: TaskCollectCustomV2,
+    user: User = Depends(current_active_verified_user),
+    db: DBSyncSession = Depends(get_sync_db),
+) -> list[TaskReadV2]:
+
+    settings = Inject(get_settings)
+
+    if settings.FRACTAL_RUNNER_BACKEND == "slurm_ssh":
+        if task_collect.package_root is None:
+            raise HTTPException(
+                status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+                detail="Cannot infer 'package_root' with 'slurm_ssh' backend.",
+            )
+    else:
+        if not Path(task_collect.python_interpreter).is_file():
+            raise HTTPException(
+                status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+                detail=(
+                    f"{task_collect.python_interpreter=} "
+                    "doesn't exist or is not a file."
+                ),
+            )
+        if (
+            task_collect.package_root is not None
+            and not Path(task_collect.package_root).is_dir()
+        ):
+            raise HTTPException(
+                status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+                detail=(
+                    f"{task_collect.package_root=} "
+                    "doesn't exist or is not a directory."
+                ),
+            )
+
+    if task_collect.package_root is None:
+
+        package_name_underscore = task_collect.package_name.replace("-", "_")
+        # Note that python_command is then used as part of a subprocess.run
+        # statement: be careful with mixing `'` and `"`.
+        python_command = (
+            "import importlib.util; "
+            "from pathlib import Path; "
+            "init_path=importlib.util.find_spec"
+            f'("{package_name_underscore}").origin; '
+            "print(Path(init_path).parent.as_posix())"
+        )
+        logger.debug(
+            f"Now running {python_command=} through "
+            f"{task_collect.python_interpreter}."
+        )
+        res = subprocess.run(  # nosec
+            shlex.split(
+                f"{task_collect.python_interpreter} -c '{python_command}'"
+            ),
+            capture_output=True,
+            encoding="utf8",
+        )
+
+        if (
+            res.returncode != 0
+            or res.stdout is None
+            or ("\n" in res.stdout.strip("\n"))
+        ):
+            raise HTTPException(
+                status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+                detail=(
+                    "Cannot determine 'package_root'.\n"
+                    f"Original output: {res.stdout}\n"
+                    f"Original error: {res.stderr}"
+                ),
+            )
+        package_root = Path(res.stdout.strip("\n"))
+    else:
+        package_root = Path(task_collect.package_root)
+
+    # Set task.owner attribute
+    owner = user.username or user.slurm_user
+    if owner is None:
+        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."
+            ),
+        )
+    source = f"{owner}:{task_collect.source}"
+
+    task_list: list[TaskCreateV2] = _prepare_tasks_metadata(
+        package_manifest=task_collect.manifest,
+        package_source=source,
+        python_bin=Path(task_collect.python_interpreter),
+        package_root=package_root,
+        package_version=task_collect.version,
+    )
+    # 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).
+    sources = [task.source for task in task_list]
+    stm = select(TaskV2).where(TaskV2.source.in_(sources))
+    res = db.execute(stm)
+    overlapping_sources_v2 = res.scalars().all()
+    if overlapping_sources_v2:
+        overlapping_tasks_v2_source_and_id = [
+            f"TaskV2 with ID {task.id} already has source='{task.source}'"
+            for task in overlapping_sources_v2
+        ]
+        raise HTTPException(
+            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail="\n".join(overlapping_tasks_v2_source_and_id),
+        )
+    stm = select(TaskV1).where(TaskV1.source.in_(sources))
+    res = db.execute(stm)
+    overlapping_sources_v1 = res.scalars().all()
+    if overlapping_sources_v1:
+        overlapping_tasks_v1_source_and_id = [
+            f"TaskV1 with ID {task.id} already has source='{task.source}'\n"
+            for task in overlapping_sources_v1
+        ]
+        raise HTTPException(
+            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail="\n".join(overlapping_tasks_v1_source_and_id),
+        )
+
+    task_list_db: list[TaskV2] = _insert_tasks(
+        task_list=task_list, owner=owner, db=db
+    )
+
+    logger.debug(
+        f"Custom-environment task collection by user {user.email} completed, "
+        f"for package with {source=}"
+    )
+
+    return task_list_db

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

@@ -0,0 +1,125 @@
+from fastapi import APIRouter
+from fastapi import BackgroundTasks
+from fastapi import Depends
+from fastapi import HTTPException
+from fastapi import Request
+from fastapi import Response
+from fastapi import status
+from pydantic.error_wrappers import ValidationError
+
+from .....config import get_settings
+from .....logger import reset_logger_handlers
+from .....logger import set_logger
+from .....syringe import Inject
+from .....tasks.v2._TaskCollectPip import _TaskCollectPip
+from ....db import AsyncSession
+from ....db import get_async_db
+from ....models.v2 import CollectionStateV2
+from ....schemas.v2 import CollectionStateReadV2
+from ....schemas.v2 import CollectionStatusV2
+from ....schemas.v2 import TaskCollectPipV2
+from ....security import current_active_user
+from ....security import current_active_verified_user
+from ....security import User
+from fractal_server.tasks.v2.background_operations_ssh import (
+    background_collect_pip_ssh,
+)
+
+router = APIRouter()
+
+logger = set_logger(__name__)
+
+
+@router.post(
+    "/collect/pip/",
+    response_model=CollectionStateReadV2,
+    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,
+    request: Request,
+    user: User = Depends(current_active_verified_user),
+    db: AsyncSession = Depends(get_async_db),
+) -> CollectionStateReadV2:
+    """
+    Task collection endpoint (SSH version)
+    """
+
+    # Set default python version
+    if task_collect.python_version is None:
+        settings = Inject(get_settings)
+        task_collect.python_version = (
+            settings.FRACTAL_TASKS_PYTHON_DEFAULT_VERSION
+        )
+
+    # 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}",
+        )
+
+    # Note: we don't use TaskCollectStatusV2 here for the JSON column `data`
+    state = CollectionStateV2(
+        data=dict(
+            status=CollectionStatusV2.PENDING, package=task_collect.package
+        )
+    )
+    db.add(state)
+    await db.commit()
+
+    background_tasks.add_task(
+        background_collect_pip_ssh,
+        state.id,
+        task_pkg,
+        request.app.state.fractal_ssh,
+    )
+
+    response.status_code = status.HTTP_201_CREATED
+    return state
+
+
+# FIXME SSH: check_collection_status code is almost identical to the
+# one in task_collection.py
+@router.get("/collect/{state_id}/", response_model=CollectionStateReadV2)
+async def check_collection_status(
+    state_id: int,
+    verbose: bool = False,
+    user: User = Depends(current_active_user),
+    db: AsyncSession = Depends(get_async_db),
+) -> CollectionStateReadV2:
+    """
+    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}",
+        )
+
+    # FIXME SSH: add logic for when data.state["log"] is empty
+
+    reset_logger_handlers(logger)
+    await db.close()
+    return state