fractal-server 2.15.0a0__py3-none-any.whl → 2.15.0a2__py3-none-any.whl

fractal_server/__init__.py

@@ -1 +1 @@
-__VERSION__ = "2.15.0a0"
+__VERSION__ = "2.15.0a2"

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

@@ -107,7 +107,7 @@ async def collect_task_pixi(
             raise HTTPException(
                 status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
                 detail=(
-                    f"Pixi version {pixi_version} is not available. Available"
+                    f"Pixi version {pixi_version} is not available. Available "
                     f"versions: {list(settings.pixi.versions.keys())}"
                 ),
             )

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

@@ -27,8 +27,11 @@ from fractal_server.syringe import Inject
 from fractal_server.tasks.v2.local import deactivate_local
 from fractal_server.tasks.v2.local import deactivate_local_pixi
 from fractal_server.tasks.v2.local import reactivate_local
+from fractal_server.tasks.v2.local import reactivate_local_pixi
 from fractal_server.tasks.v2.ssh import deactivate_ssh
+from fractal_server.tasks.v2.ssh import deactivate_ssh_pixi
 from fractal_server.tasks.v2.ssh import reactivate_ssh
+from fractal_server.tasks.v2.ssh import reactivate_ssh_pixi
 from fractal_server.utils import get_timestamp
 
 router = APIRouter()
@@ -126,9 +129,12 @@ async def deactivate_task_group(
             host=user_settings.ssh_host,
             key_path=user_settings.ssh_private_key_path,
         )
-
+        if task_group.origin == TaskGroupV2OriginEnum.PIXI:
+            deactivate_function = deactivate_ssh_pixi
+        else:
+            deactivate_function = deactivate_ssh
         background_tasks.add_task(
-            deactivate_ssh,
+            deactivate_function,
             task_group_id=task_group.id,
             task_group_activity_id=task_group_activity.id,
             ssh_config=ssh_config,
@@ -252,8 +258,12 @@ async def reactivate_task_group(
             key_path=user_settings.ssh_private_key_path,
         )
 
+        if task_group.origin == TaskGroupV2OriginEnum.PIXI:
+            reactivate_function = reactivate_ssh_pixi
+        else:
+            reactivate_function = reactivate_ssh
         background_tasks.add_task(
-            reactivate_ssh,
+            reactivate_function,
             task_group_id=task_group.id,
             task_group_activity_id=task_group_activity.id,
             ssh_config=ssh_config,
@@ -261,8 +271,12 @@ async def reactivate_task_group(
         )
 
     else:
+        if task_group.origin == TaskGroupV2OriginEnum.PIXI:
+            reactivate_function = reactivate_local_pixi
+        else:
+            reactivate_function = reactivate_local
         background_tasks.add_task(
-            reactivate_local,
+            reactivate_function,
             task_group_id=task_group.id,
             task_group_activity_id=task_group_activity.id,
         )

fractal_server/ssh/_fabric.py

@@ -204,6 +204,32 @@ class FractalSSH:
         self.logger.info(f"END reading remote JSON file {filepath}.")
         return data
 
+    def read_remote_text_file(self, filepath: str) -> dict[str, Any]:
+        """
+        Read a remote text file into a string.
+
+        Note from paramiko docs:
+        > The Python 'b' flag is ignored, since SSH treats all files as binary.
+        """
+        self.logger.info(f"START reading remote text file {filepath}.")
+        with _acquire_lock_with_timeout(
+            lock=self._lock,
+            label="read_remote_text_file",
+            timeout=self.default_lock_timeout,
+        ):
+            try:
+                with self._sftp_unsafe().open(filepath, "r") as f:
+                    data = f.read().decode()
+            except Exception as e:
+                self.log_and_raise(
+                    e=e,
+                    message=(
+                        f"Error in `read_remote_text_file`, for {filepath=}."
+                    ),
+                )
+        self.logger.info(f"END reading remote text file {filepath}.")
+        return data
+
     def check_connection(self) -> None:
         """
         Open the SSH connection and handle exceptions.

fractal_server/tasks/v2/local/__init__.py

@@ -3,3 +3,4 @@ from .collect_pixi import collect_local_pixi  # noqa
 from .deactivate import deactivate_local  # noqa
 from .deactivate_pixi import deactivate_local_pixi  # noqa
 from .reactivate import reactivate_local  # noqa
+from .reactivate_pixi import reactivate_local_pixi  # noqa

fractal_server/tasks/v2/local/collect.py

@@ -1,5 +1,4 @@
 import json
-import logging
 import shutil
 import time
 from pathlib import Path
@@ -8,7 +7,6 @@ from tempfile import TemporaryDirectory
 from ..utils_database import create_db_tasks_and_update_task_group_sync
 from ._utils import _customize_and_run_template
 from fractal_server.app.db import get_sync_db
-from fractal_server.app.models.v2 import TaskGroupActivityV2
 from fractal_server.app.models.v2 import TaskGroupV2
 from fractal_server.app.schemas.v2 import FractalUploadedFile
 from fractal_server.app.schemas.v2 import TaskGroupActivityActionV2
@@ -20,6 +18,9 @@ from fractal_server.tasks.utils import get_log_path
 from fractal_server.tasks.v2.local._utils import check_task_files_exist
 from fractal_server.tasks.v2.utils_background import add_commit_refresh
 from fractal_server.tasks.v2.utils_background import fail_and_cleanup
+from fractal_server.tasks.v2.utils_background import (
+    get_activity_and_task_group,
+)
 from fractal_server.tasks.v2.utils_background import get_current_log
 from fractal_server.tasks.v2.utils_background import prepare_tasks_metadata
 from fractal_server.tasks.v2.utils_package_names import compare_package_names
@@ -67,16 +68,12 @@ def collect_local(
         )
 
         with next(get_sync_db()) as db:
-            # Get main objects from db
-            activity = db.get(TaskGroupActivityV2, task_group_activity_id)
-            task_group = db.get(TaskGroupV2, task_group_id)
-            if activity is None or task_group is None:
-                # Use `logging` directly
-                logging.error(
-                    "Cannot find database rows with "
-                    f"{task_group_id=} and {task_group_activity_id=}:\n"
-                    f"{task_group=}\n{activity=}. Exit."
-                )
+            success, task_group, activity = get_activity_and_task_group(
+                task_group_activity_id=task_group_activity_id,
+                task_group_id=task_group_id,
+                db=db,
+            )
+            if not success:
                 return
 
             # Log some info

fractal_server/tasks/v2/local/collect_pixi.py

@@ -1,5 +1,4 @@
 import json
-import logging
 import shutil
 import time
 from pathlib import Path
@@ -9,8 +8,6 @@ from ..utils_database import create_db_tasks_and_update_task_group_sync
 from ..utils_pixi import parse_collect_stdout
 from ..utils_pixi import SOURCE_DIR_NAME
 from fractal_server.app.db import get_sync_db
-from fractal_server.app.models.v2 import TaskGroupActivityV2
-from fractal_server.app.models.v2 import TaskGroupV2
 from fractal_server.app.schemas.v2 import FractalUploadedFile
 from fractal_server.app.schemas.v2 import TaskGroupActivityActionV2
 from fractal_server.app.schemas.v2 import TaskGroupActivityStatusV2
@@ -24,6 +21,9 @@ from fractal_server.tasks.v2.local._utils import _customize_and_run_template
 from fractal_server.tasks.v2.local._utils import check_task_files_exist
 from fractal_server.tasks.v2.utils_background import add_commit_refresh
 from fractal_server.tasks.v2.utils_background import fail_and_cleanup
+from fractal_server.tasks.v2.utils_background import (
+    get_activity_and_task_group,
+)
 from fractal_server.tasks.v2.utils_background import get_current_log
 from fractal_server.tasks.v2.utils_background import prepare_tasks_metadata
 from fractal_server.tasks.v2.utils_templates import SCRIPTS_SUBFOLDER
@@ -48,14 +48,12 @@ def collect_local_pixi(
         )
 
         with next(get_sync_db()) as db:
-            activity = db.get(TaskGroupActivityV2, task_group_activity_id)
-            task_group = db.get(TaskGroupV2, task_group_id)
-            if activity is None or task_group is None:
-                logging.error(
-                    "Cannot find database rows with "
-                    f"{task_group_id=} and {task_group_activity_id=}:\n"
-                    f"{task_group=}\n{activity=}. Exit."
-                )
+            success, task_group, activity = get_activity_and_task_group(
+                task_group_activity_id=task_group_activity_id,
+                task_group_id=task_group_id,
+                db=db,
+            )
+            if not success:
                 return
 
             logger.info("START")
@@ -75,10 +73,6 @@ def collect_local_pixi(
                 )
                 return
 
-            # Set `pixi_bin` and check that it exists
-            pixi_home = settings.pixi.versions[task_group.pixi_version]
-            pixi_bin = Path(pixi_home, "bin/pixi").as_posix()
-
             try:
                 Path(task_group.path).mkdir(parents=True)
                 logger.info(f"Created {task_group.path}")
@@ -91,24 +85,20 @@ def collect_local_pixi(
                 task_group.archive_path = archive_path
                 task_group = add_commit_refresh(obj=task_group, db=db)
 
-                replacements = {
-                    ("__PIXI_HOME__", pixi_home),
-                    ("__PACKAGE_DIR__", task_group.path),
-                    ("__TAR_GZ_PATH__", archive_path),
-                    (
-                        "__IMPORT_PACKAGE_NAME__",
-                        task_group.pkg_name.replace("-", "_"),
-                    ),
-                    ("__SOURCE_DIR_NAME__", SOURCE_DIR_NAME),
-                }
-
-                activity.status = TaskGroupActivityStatusV2.ONGOING
-                activity.log = get_current_log(log_file_path)
-                activity = add_commit_refresh(obj=activity, db=db)
-
-                stdout = _customize_and_run_template(
-                    template_filename="pixi_1_collect.sh",
-                    replacements=replacements,
+                common_args = dict(
+                    replacements={
+                        (
+                            "__PIXI_HOME__",
+                            settings.pixi.versions[task_group.pixi_version],
+                        ),
+                        ("__PACKAGE_DIR__", task_group.path),
+                        ("__TAR_GZ_PATH__", archive_path),
+                        (
+                            "__IMPORT_PACKAGE_NAME__",
+                            task_group.pkg_name.replace("-", "_"),
+                        ),
+                        ("__SOURCE_DIR_NAME__", SOURCE_DIR_NAME),
+                    },
                     script_dir=Path(
                         task_group.path, SCRIPTS_SUBFOLDER
                     ).as_posix(),
@@ -118,6 +108,32 @@ def collect_local_pixi(
                     ),
                     logger_name=LOGGER_NAME,
                 )
+
+                activity.status = TaskGroupActivityStatusV2.ONGOING
+                activity.log = get_current_log(log_file_path)
+                activity = add_commit_refresh(obj=activity, db=db)
+
+                # Run script 1
+                _customize_and_run_template(
+                    template_filename="pixi_1_extract.sh",
+                    **common_args,
+                )
+                activity.log = get_current_log(log_file_path)
+                activity = add_commit_refresh(obj=activity, db=db)
+
+                # Run script 2
+                _customize_and_run_template(
+                    template_filename="pixi_2_install.sh",
+                    **common_args,
+                )
+                activity.log = get_current_log(log_file_path)
+                activity = add_commit_refresh(obj=activity, db=db)
+
+                # Run script 3
+                stdout = _customize_and_run_template(
+                    template_filename="pixi_3_post_install.sh",
+                    **common_args,
+                )
                 activity.log = get_current_log(log_file_path)
                 activity = add_commit_refresh(obj=activity, db=db)
 
@@ -126,6 +142,9 @@ def collect_local_pixi(
                 package_root = parsed_output["package_root"]
                 venv_size = parsed_output["venv_size"]
                 venv_file_number = parsed_output["venv_file_number"]
+                project_python_wrapper = parsed_output[
+                    "project_python_wrapper"
+                ]
 
                 # Read and validate manifest
                 # NOTE: we are only supporting the manifest path being relative
@@ -145,14 +164,7 @@ def collect_local_pixi(
                     package_manifest=pkg_manifest,
                     package_version=task_group.version,
                     package_root=Path(package_root),
-                    pixi_bin=pixi_bin,
-                    pixi_manifest_path=(
-                        Path(
-                            task_group.path,
-                            SOURCE_DIR_NAME,
-                            "pyproject.toml",
-                        ).as_posix()
-                    ),
+                    project_python_wrapper=Path(project_python_wrapper),
                 )
                 check_task_files_exist(task_list=task_list)
                 logger.info("_prepare_tasks_metadata - end")

fractal_server/tasks/v2/local/deactivate.py

@@ -1,4 +1,3 @@
-import logging
 import shutil
 import time
 from pathlib import Path
@@ -6,11 +5,10 @@ from tempfile import TemporaryDirectory
 
 from ..utils_background import add_commit_refresh
 from ..utils_background import fail_and_cleanup
+from ..utils_background import get_activity_and_task_group
 from ..utils_templates import get_collection_replacements
 from ._utils import _customize_and_run_template
 from fractal_server.app.db import get_sync_db
-from fractal_server.app.models.v2 import TaskGroupActivityV2
-from fractal_server.app.models.v2 import TaskGroupV2
 from fractal_server.app.schemas.v2 import TaskGroupActivityActionV2
 from fractal_server.app.schemas.v2 import TaskGroupV2OriginEnum
 from fractal_server.app.schemas.v2.task_group import TaskGroupActivityStatusV2
@@ -49,17 +47,12 @@ def deactivate_local(
         )
 
         with next(get_sync_db()) as db:
-
-            # Get main objects from db
-            activity = db.get(TaskGroupActivityV2, task_group_activity_id)
-            task_group = db.get(TaskGroupV2, task_group_id)
-            if activity is None or task_group is None:
-                # Use `logging` directly
-                logging.error(
-                    "Cannot find database rows with "
-                    f"{task_group_id=} and {task_group_activity_id=}:\n"
-                    f"{task_group=}\n{activity=}. Exit."
-                )
+            success, task_group, activity = get_activity_and_task_group(
+                task_group_activity_id=task_group_activity_id,
+                task_group_id=task_group_id,
+                db=db,
+            )
+            if not success:
                 return
 
             # Log some info

fractal_server/tasks/v2/local/deactivate_pixi.py

@@ -1,14 +1,12 @@
-import logging
 import shutil
 from pathlib import Path
 from tempfile import TemporaryDirectory
 
 from ..utils_background import add_commit_refresh
 from ..utils_background import fail_and_cleanup
+from ..utils_background import get_activity_and_task_group
 from ..utils_pixi import SOURCE_DIR_NAME
 from fractal_server.app.db import get_sync_db
-from fractal_server.app.models.v2 import TaskGroupActivityV2
-from fractal_server.app.models.v2 import TaskGroupV2
 from fractal_server.app.schemas.v2.task_group import TaskGroupActivityStatusV2
 from fractal_server.logger import reset_logger_handlers
 from fractal_server.logger import set_logger
@@ -23,7 +21,7 @@ def deactivate_local_pixi(
     task_group_id: int,
 ) -> None:
     """
-    Deactivate a task group venv.
+    Deactivate a pixi task group venv.
 
     This function is run as a background task, therefore exceptions must be
     handled.
@@ -43,17 +41,12 @@ def deactivate_local_pixi(
         )
 
         with next(get_sync_db()) as db:
-
-            # Get main objects from db
-            activity = db.get(TaskGroupActivityV2, task_group_activity_id)
-            task_group = db.get(TaskGroupV2, task_group_id)
-            if activity is None or task_group is None:
-                # Use `logging` directly
-                logging.error(
-                    "Cannot find database rows with "
-                    f"{task_group_id=} and {task_group_activity_id=}:\n"
-                    f"{task_group=}\n{activity=}. Exit."
-                )
+            success, task_group, activity = get_activity_and_task_group(
+                task_group_activity_id=task_group_activity_id,
+                task_group_id=task_group_id,
+                db=db,
+            )
+            if not success:
                 return
 
             # Log some info
@@ -63,7 +56,6 @@ def deactivate_local_pixi(
                 logger.debug(f"task_group.{key}: {value}")
 
             source_dir = Path(task_group.path, SOURCE_DIR_NAME)
-            # Check that the (local) task_group venv_path does exist
             if not source_dir.exists():
                 error_msg = f"'{source_dir.as_posix()}' does not exist."
                 logger.error(error_msg)

fractal_server/tasks/v2/local/reactivate.py

@@ -1,4 +1,3 @@
-import logging
 import shutil
 import time
 from pathlib import Path
@@ -6,11 +5,10 @@ from tempfile import TemporaryDirectory
 
 from ..utils_background import add_commit_refresh
 from ..utils_background import fail_and_cleanup
+from ..utils_background import get_activity_and_task_group
 from ..utils_templates import get_collection_replacements
 from ._utils import _customize_and_run_template
 from fractal_server.app.db import get_sync_db
-from fractal_server.app.models.v2 import TaskGroupActivityV2
-from fractal_server.app.models.v2 import TaskGroupV2
 from fractal_server.app.schemas.v2 import TaskGroupActivityActionV2
 from fractal_server.app.schemas.v2.task_group import TaskGroupActivityStatusV2
 from fractal_server.logger import reset_logger_handlers
@@ -50,17 +48,12 @@ def reactivate_local(
         )
 
         with next(get_sync_db()) as db:
-
-            # Get main objects from db
-            activity = db.get(TaskGroupActivityV2, task_group_activity_id)
-            task_group = db.get(TaskGroupV2, task_group_id)
-            if activity is None or task_group is None:
-                # Use `logging` directly
-                logging.error(
-                    "Cannot find database rows with "
-                    f"{task_group_id=} and {task_group_activity_id=}:\n"
-                    f"{task_group=}\n{activity=}. Exit."
-                )
+            success, task_group, activity = get_activity_and_task_group(
+                task_group_activity_id=task_group_activity_id,
+                task_group_id=task_group_id,
+                db=db,
+            )
+            if not success:
                 return
 
             # Log some info
@@ -99,7 +92,7 @@ def reactivate_local(
                 replacements.append(
                     ("__PIP_FREEZE_FILE__", f"{tmpdir}/pip_freeze.txt")
                 )
-                # Prepare common arguments for `_customize_and_run_template``
+                # Prepare common arguments for `_customize_and_run_template`
                 common_args = dict(
                     replacements=replacements,
                     script_dir=(

fractal_server/tasks/v2/local/reactivate_pixi.py

@@ -0,0 +1,166 @@
+import shutil
+import time
+from pathlib import Path
+from tempfile import TemporaryDirectory
+
+from ..utils_background import add_commit_refresh
+from ..utils_background import fail_and_cleanup
+from ..utils_background import get_activity_and_task_group
+from ..utils_pixi import SOURCE_DIR_NAME
+from fractal_server.app.db import get_sync_db
+from fractal_server.app.schemas.v2 import TaskGroupActivityActionV2
+from fractal_server.app.schemas.v2.task_group import TaskGroupActivityStatusV2
+from fractal_server.config import get_settings
+from fractal_server.logger import reset_logger_handlers
+from fractal_server.logger import set_logger
+from fractal_server.syringe import Inject
+from fractal_server.tasks.utils import get_log_path
+from fractal_server.tasks.v2.local._utils import _customize_and_run_template
+from fractal_server.tasks.v2.utils_background import get_current_log
+from fractal_server.tasks.v2.utils_templates import SCRIPTS_SUBFOLDER
+from fractal_server.utils import get_timestamp
+
+
+def reactivate_local_pixi(
+    *,
+    task_group_activity_id: int,
+    task_group_id: int,
+) -> None:
+    """
+    Reactivate a task group venv.
+
+    This function is run as a background task, therefore exceptions must be
+    handled.
+
+    Arguments:
+        task_group_id:
+        task_group_activity_id:
+    """
+
+    LOGGER_NAME = f"{__name__}.ID{task_group_activity_id}"
+
+    with TemporaryDirectory() as tmpdir:
+        log_file_path = get_log_path(Path(tmpdir))
+        logger = set_logger(
+            logger_name=LOGGER_NAME,
+            log_file_path=log_file_path,
+        )
+        with next(get_sync_db()) as db:
+            success, task_group, activity = get_activity_and_task_group(
+                task_group_activity_id=task_group_activity_id,
+                task_group_id=task_group_id,
+                db=db,
+            )
+            if not success:
+                return
+
+            # Log some info
+            logger.debug("START")
+
+            for key, value in task_group.model_dump().items():
+                logger.debug(f"task_group.{key}: {value}")
+
+            source_dir = Path(task_group.path, SOURCE_DIR_NAME).as_posix()
+            if Path(source_dir).exists():
+                error_msg = f"{source_dir} already exists."
+                logger.error(error_msg)
+                fail_and_cleanup(
+                    task_group=task_group,
+                    task_group_activity=activity,
+                    logger_name=LOGGER_NAME,
+                    log_file_path=log_file_path,
+                    exception=FileExistsError(error_msg),
+                    db=db,
+                )
+                return
+
+            try:
+                activity.status = TaskGroupActivityStatusV2.ONGOING
+                activity = add_commit_refresh(obj=activity, db=db)
+
+                logger.debug(f"start - writing {source_dir}/pixi.lock")
+                with Path(source_dir, "pixi.lock").open("w") as f:
+                    f.write(task_group.env_info)
+                logger.debug(f"end - writing {source_dir}/pixi.lock")
+
+                settings = Inject(get_settings)
+                common_args = dict(
+                    replacements={
+                        (
+                            "__PIXI_HOME__",
+                            settings.pixi.versions[task_group.pixi_version],
+                        ),
+                        ("__PACKAGE_DIR__", task_group.path),
+                        ("__TAR_GZ_PATH__", task_group.archive_path),
+                        (
+                            "__IMPORT_PACKAGE_NAME__",
+                            task_group.pkg_name.replace("-", "_"),
+                        ),
+                        ("__SOURCE_DIR_NAME__", SOURCE_DIR_NAME),
+                        ("__FROZEN_OPTION__", "true"),
+                    },
+                    script_dir=Path(
+                        task_group.path, SCRIPTS_SUBFOLDER
+                    ).as_posix(),
+                    prefix=(
+                        f"{int(time.time())}_"
+                        f"{TaskGroupActivityActionV2.REACTIVATE}"
+                    ),
+                    logger_name=LOGGER_NAME,
+                )
+
+                # Run script 1
+                _customize_and_run_template(
+                    template_filename="pixi_1_extract.sh",
+                    **common_args,
+                )
+                activity.log = get_current_log(log_file_path)
+                activity = add_commit_refresh(obj=activity, db=db)
+
+                # Run script 2
+                _customize_and_run_template(
+                    template_filename="pixi_2_install.sh",
+                    **common_args,
+                )
+                activity.log = get_current_log(log_file_path)
+                activity = add_commit_refresh(obj=activity, db=db)
+
+                # Run script 3
+                _customize_and_run_template(
+                    template_filename="pixi_3_post_install.sh",
+                    **common_args,
+                )
+                activity.log = get_current_log(log_file_path)
+                activity = add_commit_refresh(obj=activity, db=db)
+
+                activity.log = get_current_log(log_file_path)
+                activity.status = TaskGroupActivityStatusV2.OK
+                activity.timestamp_ended = get_timestamp()
+                activity = add_commit_refresh(obj=activity, db=db)
+                task_group.active = True
+                task_group = add_commit_refresh(obj=task_group, db=db)
+                logger.debug("END")
+
+                reset_logger_handlers(logger)
+
+            except Exception as reactivate_e:
+                # Delete corrupted source_dir
+                try:
+                    logger.info(f"Now delete folder {source_dir}")
+                    shutil.rmtree(source_dir)
+                    logger.info(f"Deleted folder {source_dir}")
+                except Exception as rm_e:
+                    logger.error(
+                        "Removing folder failed.\n"
+                        f"Original error:\n{str(rm_e)}"
+                    )
+
+                fail_and_cleanup(
+                    task_group=task_group,
+                    task_group_activity=activity,
+                    logger_name=LOGGER_NAME,
+                    log_file_path=log_file_path,
+                    exception=reactivate_e,
+                    db=db,
+                )
+        return

fractal_server/tasks/v2/ssh/__init__.py

@@ -1,4 +1,6 @@
 from .collect import collect_ssh  # noqa
 from .collect_pixi import collect_ssh_pixi  # noqa
 from .deactivate import deactivate_ssh  # noqa
+from .deactivate_pixi import deactivate_ssh_pixi  # noqa
 from .reactivate import reactivate_ssh  # noqa
+from .reactivate_pixi import reactivate_ssh_pixi  # noqa

fractal_server/tasks/v2/ssh/collect.py

@@ -1,15 +1,13 @@
-import logging
 import time
 from pathlib import Path
 from tempfile import TemporaryDirectory
 
 from ....ssh._fabric import SingleUseFractalSSH
 from ..utils_background import fail_and_cleanup
+from ..utils_background import get_activity_and_task_group
 from ..utils_background import prepare_tasks_metadata
 from ..utils_database import create_db_tasks_and_update_task_group_sync
 from fractal_server.app.db import get_sync_db
-from fractal_server.app.models.v2 import TaskGroupActivityV2
-from fractal_server.app.models.v2 import TaskGroupV2
 from fractal_server.app.schemas.v2 import FractalUploadedFile
 from fractal_server.app.schemas.v2 import TaskGroupActivityActionV2
 from fractal_server.app.schemas.v2 import TaskGroupActivityStatusV2
@@ -76,16 +74,12 @@ def collect_ssh(
         ) as fractal_ssh:
 
             with next(get_sync_db()) as db:
-                # Get main objects from db
-                activity = db.get(TaskGroupActivityV2, task_group_activity_id)
-                task_group = db.get(TaskGroupV2, task_group_id)
-                if activity is None or task_group is None:
-                    # Use `logging` directly
-                    logging.error(
-                        "Cannot find database rows with "
-                        f"{task_group_id=} and {task_group_activity_id=}:\n"
-                        f"{task_group=}\n{activity=}. Exit."
-                    )
+                success, task_group, activity = get_activity_and_task_group(
+                    task_group_activity_id=task_group_activity_id,
+                    task_group_id=task_group_id,
+                    db=db,
+                )
+                if not success:
                     return
 
                 # Log some info