fractal-server 2.14.16__py3-none-any.whl → 2.15.0__py3-none-any.whl

fractal_server/tasks/v2/ssh/collect_pixi.py

@@ -0,0 +1,315 @@
+import time
+from pathlib import Path
+from tempfile import TemporaryDirectory
+
+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 ..utils_pixi import parse_collect_stdout
+from ..utils_pixi import SOURCE_DIR_NAME
+from ._utils import check_ssh_or_fail_and_cleanup
+from fractal_server.app.db import get_sync_db
+from fractal_server.app.schemas.v2 import FractalUploadedFile
+from fractal_server.app.schemas.v2 import TaskGroupActivityActionV2
+from fractal_server.app.schemas.v2 import TaskGroupActivityStatusV2
+from fractal_server.app.schemas.v2.manifest import ManifestV2
+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.ssh._fabric import SingleUseFractalSSH
+from fractal_server.ssh._fabric import SSHConfig
+from fractal_server.syringe import Inject
+from fractal_server.tasks.v2.ssh._utils import _customize_and_run_template
+from fractal_server.tasks.v2.utils_background import add_commit_refresh
+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 collect_ssh_pixi(
+    *,
+    task_group_id: int,
+    task_group_activity_id: int,
+    ssh_config: SSHConfig,
+    tasks_base_dir: str,
+    tar_gz_file: FractalUploadedFile,
+) -> None:
+    """
+    Collect a task package over SSH
+
+    This function runs as a background task, therefore exceptions must be
+    handled.
+
+    NOTE: since this function is sync, it runs within a thread - due to
+    starlette/fastapi handling of background tasks (see
+    https://github.com/encode/starlette/blob/master/starlette/background.py).
+
+
+    Arguments:
+        task_group_id:
+        task_group_activity_id:
+        ssh_config:
+        tasks_base_dir:
+            Only used as a `safe_root` in `remove_dir`, and typically set to
+            `user_settings.ssh_tasks_dir`.
+        tar_gz_file:
+    """
+
+    LOGGER_NAME = f"{__name__}.ID{task_group_activity_id}"
+
+    # Work within a temporary folder, where also logs will be placed
+    with TemporaryDirectory() as tmpdir:
+        log_file_path = Path(tmpdir) / "log"
+        logger = set_logger(
+            logger_name=LOGGER_NAME,
+            log_file_path=log_file_path,
+        )
+        logger.info("START")
+        with next(get_sync_db()) as db:
+            db_objects_ok, task_group, activity = get_activity_and_task_group(
+                task_group_activity_id=task_group_activity_id,
+                task_group_id=task_group_id,
+                db=db,
+                logger_name=LOGGER_NAME,
+            )
+            if not db_objects_ok:
+                return
+
+            with SingleUseFractalSSH(
+                ssh_config=ssh_config,
+                logger_name=LOGGER_NAME,
+            ) as fractal_ssh:
+
+                try:
+                    # Check SSH connection
+                    ssh_ok = check_ssh_or_fail_and_cleanup(
+                        fractal_ssh=fractal_ssh,
+                        task_group=task_group,
+                        task_group_activity=activity,
+                        logger_name=LOGGER_NAME,
+                        log_file_path=log_file_path,
+                        db=db,
+                    )
+                    if not ssh_ok:
+                        return
+
+                    # Check that the (remote) task_group path does not exist
+                    if fractal_ssh.remote_exists(task_group.path):
+                        error_msg = f"{task_group.path} 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
+
+                    # Create remote `task_group.path` and `script_dir_remote`
+                    # folders
+                    script_dir_remote = Path(
+                        task_group.path, SCRIPTS_SUBFOLDER
+                    ).as_posix()
+                    fractal_ssh.mkdir(folder=task_group.path, parents=True)
+                    fractal_ssh.mkdir(folder=script_dir_remote, parents=True)
+
+                    # Write tar.gz file locally and send it to remote path,
+                    # and set task_group.archive_path
+                    tar_gz_filename = tar_gz_file.filename
+                    archive_path = (
+                        Path(task_group.path) / tar_gz_filename
+                    ).as_posix()
+                    tmp_archive_path = Path(tmpdir, tar_gz_filename).as_posix()
+                    logger.info(f"Write tar.gz file into {tmp_archive_path}")
+                    with open(tmp_archive_path, "wb") as f:
+                        f.write(tar_gz_file.contents)
+                    fractal_ssh.send_file(
+                        local=tmp_archive_path,
+                        remote=archive_path,
+                    )
+                    task_group.archive_path = archive_path
+                    task_group = add_commit_refresh(obj=task_group, db=db)
+
+                    settings = Inject(get_settings)
+                    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__", ""),
+                        (
+                            "__TOKIO_WORKER_THREADS__",
+                            str(settings.pixi.TOKIO_WORKER_THREADS),
+                        ),
+                        (
+                            "__PIXI_CONCURRENT_SOLVES__",
+                            str(settings.pixi.PIXI_CONCURRENT_SOLVES),
+                        ),
+                        (
+                            "__PIXI_CONCURRENT_DOWNLOADS__",
+                            str(settings.pixi.PIXI_CONCURRENT_DOWNLOADS),
+                        ),
+                    }
+
+                    logger.info("installing - START")
+
+                    # Set status to ONGOING and refresh logs
+                    activity.status = TaskGroupActivityStatusV2.ONGOING
+                    activity.log = get_current_log(log_file_path)
+                    activity = add_commit_refresh(obj=activity, db=db)
+
+                    common_args = dict(
+                        script_dir_local=(
+                            Path(tmpdir, SCRIPTS_SUBFOLDER)
+                        ).as_posix(),
+                        script_dir_remote=script_dir_remote,
+                        prefix=(
+                            f"{int(time.time())}_"
+                            f"{TaskGroupActivityActionV2.COLLECT}"
+                        ),
+                        logger_name=LOGGER_NAME,
+                        fractal_ssh=fractal_ssh,
+                    )
+
+                    # Run the three pixi-related scripts
+                    stdout = _customize_and_run_template(
+                        template_filename="pixi_1_extract.sh",
+                        replacements=replacements,
+                        **common_args,
+                    )
+                    logger.debug(f"STDOUT: {stdout}")
+                    activity.log = get_current_log(log_file_path)
+                    activity = add_commit_refresh(obj=activity, db=db)
+
+                    stdout = _customize_and_run_template(
+                        template_filename="pixi_2_install.sh",
+                        replacements=replacements,
+                        **common_args,
+                    )
+                    logger.debug(f"STDOUT: {stdout}")
+                    activity.log = get_current_log(log_file_path)
+                    activity = add_commit_refresh(obj=activity, db=db)
+
+                    stdout = _customize_and_run_template(
+                        template_filename="pixi_3_post_install.sh",
+                        replacements=replacements,
+                        **common_args,
+                    )
+                    logger.debug(f"STDOUT: {stdout}")
+                    activity.log = get_current_log(log_file_path)
+                    activity = add_commit_refresh(obj=activity, db=db)
+
+                    # Parse stdout
+                    parsed_output = parse_collect_stdout(stdout)
+                    package_root_remote = 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"
+                    ]
+
+                    source_dir = Path(
+                        task_group.path, SOURCE_DIR_NAME
+                    ).as_posix()
+                    fractal_ssh.run_command(cmd=f"chmod 755 {source_dir} -R")
+
+                    # Read and validate remote manifest file
+                    manifest_path_remote = (
+                        f"{package_root_remote}/__FRACTAL_MANIFEST__.json"
+                    )
+                    pkg_manifest_dict = fractal_ssh.read_remote_json_file(
+                        manifest_path_remote
+                    )
+                    logger.info(f"Loaded {manifest_path_remote=}")
+                    pkg_manifest = ManifestV2(**pkg_manifest_dict)
+                    logger.info("Manifest is a valid ManifestV2")
+
+                    logger.info("_prepare_tasks_metadata - start")
+                    task_list = prepare_tasks_metadata(
+                        package_manifest=pkg_manifest,
+                        package_version=task_group.version,
+                        package_root=Path(package_root_remote),
+                        project_python_wrapper=Path(project_python_wrapper),
+                    )
+                    logger.info("_prepare_tasks_metadata - end")
+
+                    logger.info(
+                        "create_db_tasks_and_update_task_group - " "start"
+                    )
+                    create_db_tasks_and_update_task_group_sync(
+                        task_list=task_list,
+                        task_group_id=task_group.id,
+                        db=db,
+                    )
+                    logger.info("create_db_tasks_and_update_task_group - end")
+
+                    # NOTE: see issue 2626 about whether to keep `pixi.lock`
+                    # files in the database
+                    remote_pixi_lock_file = Path(
+                        task_group.path,
+                        SOURCE_DIR_NAME,
+                        "pixi.lock",
+                    ).as_posix()
+                    pixi_lock_contents = fractal_ssh.read_remote_text_file(
+                        remote_pixi_lock_file
+                    )
+
+                    # Update task_group data
+                    logger.info(
+                        "Add env_info, venv_size and venv_file_number "
+                        "to TaskGroupV2 - start"
+                    )
+                    task_group.env_info = pixi_lock_contents
+                    task_group.venv_size_in_kB = int(venv_size)
+                    task_group.venv_file_number = int(venv_file_number)
+                    task_group = add_commit_refresh(obj=task_group, db=db)
+                    logger.info(
+                        "Add env_info, venv_size and venv_file_number "
+                        "to TaskGroupV2 - end"
+                    )
+
+                    # Finalize (write metadata to DB)
+                    logger.info("finalising - START")
+                    activity.status = TaskGroupActivityStatusV2.OK
+                    activity.timestamp_ended = get_timestamp()
+                    activity = add_commit_refresh(obj=activity, db=db)
+                    logger.info("finalising - END")
+                    logger.info("END")
+                    reset_logger_handlers(logger)
+
+                except Exception as collection_e:
+                    # Delete corrupted package dir
+                    try:
+                        logger.info(
+                            f"Now delete remote folder {task_group.path}"
+                        )
+                        fractal_ssh.remove_folder(
+                            folder=task_group.path,
+                            safe_root=tasks_base_dir,
+                        )
+                        logger.info(
+                            f"Deleted remoted folder {task_group.path}"
+                        )
+                    except Exception as e_rm:
+                        logger.error(
+                            "Removing folder failed. "
+                            f"Original error: {str(e_rm)}"
+                        )
+                    fail_and_cleanup(
+                        task_group=task_group,
+                        task_group_activity=activity,
+                        log_file_path=log_file_path,
+                        logger_name=LOGGER_NAME,
+                        exception=collection_e,
+                        db=db,
+                    )

fractal_server/tasks/v2/ssh/deactivate.py

@@ -1,16 +1,15 @@
-import logging
 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_templates import get_collection_replacements
 from ._utils import _copy_wheel_file_ssh
 from ._utils import _customize_and_run_template
+from ._utils import check_ssh_or_fail_and_cleanup
 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
@@ -55,46 +54,33 @@ def deactivate_ssh(
             logger_name=LOGGER_NAME,
             log_file_path=log_file_path,
         )
-        with SingleUseFractalSSH(
-            ssh_config=ssh_config,
-            logger_name=LOGGER_NAME,
-        ) 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."
-                    )
-                    return
-
-                # Log some info
-                logger.debug("START")
-                for key, value in task_group.model_dump().items():
-                    logger.debug(f"task_group.{key}: {value}")
-
-                # Check that SSH connection works
+        logger.debug("START")
+        with next(get_sync_db()) as db:
+            db_objects_ok, task_group, activity = get_activity_and_task_group(
+                task_group_activity_id=task_group_activity_id,
+                task_group_id=task_group_id,
+                db=db,
+                logger_name=LOGGER_NAME,
+            )
+            if not db_objects_ok:
+                return
+
+            with SingleUseFractalSSH(
+                ssh_config=ssh_config,
+                logger_name=LOGGER_NAME,
+            ) as fractal_ssh:
                 try:
-                    fractal_ssh.check_connection()
-                except Exception as e:
-                    logger.error("Cannot establish SSH connection.")
-                    fail_and_cleanup(
+                    # Check SSH connection
+                    ssh_ok = check_ssh_or_fail_and_cleanup(
+                        fractal_ssh=fractal_ssh,
                         task_group=task_group,
                         task_group_activity=activity,
                         logger_name=LOGGER_NAME,
                         log_file_path=log_file_path,
-                        exception=e,
                         db=db,
                     )
-                    return
-
-                try:
+                    if not ssh_ok:
+                        return
 
                     # Check that the (local) task_group venv_path does exist
                     if not fractal_ssh.remote_exists(task_group.venv_path):
@@ -113,10 +99,10 @@ def deactivate_ssh(
                     activity.status = TaskGroupActivityStatusV2.ONGOING
                     activity = add_commit_refresh(obj=activity, db=db)
 
-                    if task_group.pip_freeze is None:
+                    if task_group.env_info is None:
                         logger.warning(
                             "Recreate pip-freeze information, since "
-                            f"{task_group.pip_freeze=}. NOTE: this should "
+                            f"{task_group.env_info=}. NOTE: this should "
                             "only happen for task groups created before 2.9.0."
                         )
 
@@ -161,7 +147,7 @@ def deactivate_ssh(
                         )
                         activity.log = get_current_log(log_file_path)
                         activity = add_commit_refresh(obj=activity, db=db)
-                        task_group.pip_freeze = pip_freeze_stdout
+                        task_group.env_info = pip_freeze_stdout
                         task_group = add_commit_refresh(obj=task_group, db=db)
                         logger.info(
                             "Add pip freeze stdout to TaskGroupV2 - end"
@@ -174,18 +160,19 @@ def deactivate_ssh(
                             f"Handle specific cases for {task_group.origin=}."
                         )
 
-                        # Blocking situation: `wheel_path` is not set or points
-                        # to a missing path
+                        # Blocking situation: `archive_path` is not set or
+                        # points to a missing path
                         if (
-                            task_group.wheel_path is None
+                            task_group.archive_path is None
                             or not fractal_ssh.remote_exists(
-                                task_group.wheel_path
+                                task_group.archive_path
                             )
                         ):
                             error_msg = (
                                 "Invalid wheel path for task group with "
-                                f"{task_group_id=}. {task_group.wheel_path=} "
-                                "is unset or does not exist."
+                                f"{task_group_id=}. "
+                                f"{task_group.archive_path=}  is unset or "
+                                "does not exist."
                             )
                             logger.error(error_msg)
                             fail_and_cleanup(
@@ -198,58 +185,58 @@ def deactivate_ssh(
                             )
                             return
 
-                        # Recoverable situation: `wheel_path` was not yet
+                        # Recoverable situation: `archive_path` was not yet
                         # copied over to the correct server-side folder
-                        wheel_path_parent_dir = Path(
-                            task_group.wheel_path
+                        archive_path_parent_dir = Path(
+                            task_group.archive_path
                         ).parent
-                        if wheel_path_parent_dir != Path(task_group.path):
+                        if archive_path_parent_dir != Path(task_group.path):
                             logger.warning(
-                                f"{wheel_path_parent_dir.as_posix()} differs "
-                                f"from {task_group.path}. NOTE: this should "
-                                "only happen for task groups created before "
-                                "2.9.0."
+                                f"{archive_path_parent_dir.as_posix()} "
+                                f"differs from {task_group.path}. "
+                                "NOTE: this should only happen for task "
+                                "groups created before 2.9.0."
                             )
 
                             if (
-                                task_group.wheel_path
-                                not in task_group.pip_freeze
+                                task_group.archive_path
+                                not in task_group.env_info
                             ):
                                 raise ValueError(
-                                    f"Cannot find {task_group.wheel_path=} in "
-                                    "pip-freeze data. Exit."
+                                    f"Cannot find {task_group.archive_path=} "
+                                    "in pip-freeze data. Exit."
                                 )
 
                             logger.info(
                                 f"Now copy wheel file into {task_group.path}."
                             )
-                            new_wheel_path = _copy_wheel_file_ssh(
+                            new_archive_path = _copy_wheel_file_ssh(
                                 task_group=task_group,
                                 fractal_ssh=fractal_ssh,
                                 logger_name=LOGGER_NAME,
                             )
                             logger.info(
-                                f"Copied wheel file to {new_wheel_path}."
+                                f"Copied wheel file to {new_archive_path}."
                             )
 
-                            task_group.wheel_path = new_wheel_path
-                            new_pip_freeze = task_group.pip_freeze.replace(
-                                task_group.wheel_path,
-                                new_wheel_path,
+                            task_group.archive_path = new_archive_path
+                            new_pip_freeze = task_group.env_info.replace(
+                                task_group.archive_path,
+                                new_archive_path,
                             )
-                            task_group.pip_freeze = new_pip_freeze
+                            task_group.env_info = new_pip_freeze
                             task_group = add_commit_refresh(
                                 obj=task_group, db=db
                             )
                             logger.info(
-                                "Updated `wheel_path` and `pip_freeze` "
+                                "Updated `archive_path` and `env_info` "
                                 "task-group attributes."
                             )
 
-                    # Fail if `pip_freeze` includes "github", see
+                    # Fail if `env_info` includes "github", see
                     # https://github.com/fractal-analytics-platform/fractal-server/issues/2142
                     for forbidden_string in FORBIDDEN_DEPENDENCY_STRINGS:
-                        if forbidden_string in task_group.pip_freeze:
+                        if forbidden_string in task_group.env_info:
                             raise ValueError(
                                 "Deactivation and reactivation of task "
                                 f"packages with direct {forbidden_string} "

fractal_server/tasks/v2/ssh/deactivate_pixi.py

@@ -0,0 +1,122 @@
+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 ._utils import check_ssh_or_fail_and_cleanup
+from fractal_server.app.db import get_sync_db
+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
+from fractal_server.ssh._fabric import SingleUseFractalSSH
+from fractal_server.ssh._fabric import SSHConfig
+from fractal_server.tasks.utils import get_log_path
+from fractal_server.tasks.v2.utils_background import get_current_log
+from fractal_server.utils import get_timestamp
+
+
+def deactivate_ssh_pixi(
+    *,
+    task_group_activity_id: int,
+    task_group_id: int,
+    ssh_config: SSHConfig,
+    tasks_base_dir: str,
+) -> None:
+    """
+    Deactivate a pixi task group venv.
+
+    This function is run as a background task, therefore exceptions must be
+    handled.
+
+    Arguments:
+        task_group_id:
+        task_group_activity_id:
+        ssh_config:
+        tasks_base_dir:
+            Only used as a `safe_root` in `remove_dir`, and typically set to
+            `user_settings.ssh_tasks_dir`.
+    """
+
+    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,
+        )
+        logger.debug("START")
+        with next(get_sync_db()) as db:
+            db_objects_ok, task_group, activity = get_activity_and_task_group(
+                task_group_activity_id=task_group_activity_id,
+                task_group_id=task_group_id,
+                db=db,
+                logger_name=LOGGER_NAME,
+            )
+            if not db_objects_ok:
+                return
+
+            with SingleUseFractalSSH(
+                ssh_config=ssh_config,
+                logger_name=LOGGER_NAME,
+            ) as fractal_ssh:
+                try:
+                    # Check SSH connection
+                    ssh_ok = check_ssh_or_fail_and_cleanup(
+                        fractal_ssh=fractal_ssh,
+                        task_group=task_group,
+                        task_group_activity=activity,
+                        logger_name=LOGGER_NAME,
+                        log_file_path=log_file_path,
+                        db=db,
+                    )
+                    if not ssh_ok:
+                        return
+
+                    # Check that the (remote) task_group venv_path does exist
+                    source_dir = Path(
+                        task_group.path, SOURCE_DIR_NAME
+                    ).as_posix()
+                    if not fractal_ssh.remote_exists(source_dir):
+                        error_msg = f"{source_dir} does not exist."
+                        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=FileNotFoundError(error_msg),
+                            db=db,
+                        )
+                        return
+
+                    # Actually mark the task group as non-active
+                    logger.info("Now setting `active=False`.")
+                    task_group.active = False
+                    task_group = add_commit_refresh(obj=task_group, db=db)
+
+                    # Proceed with deactivation
+                    logger.info(f"Now removing {source_dir}.")
+                    fractal_ssh.remove_folder(
+                        folder=source_dir,
+                        safe_root=tasks_base_dir,
+                    )
+                    logger.info(f"All good, {source_dir} removed.")
+                    activity.status = TaskGroupActivityStatusV2.OK
+                    activity.log = get_current_log(log_file_path)
+                    activity.timestamp_ended = get_timestamp()
+                    activity = add_commit_refresh(obj=activity, db=db)
+
+                    reset_logger_handlers(logger)
+
+                except Exception as e:
+                    fail_and_cleanup(
+                        task_group=task_group,
+                        task_group_activity=activity,
+                        logger_name=LOGGER_NAME,
+                        log_file_path=log_file_path,
+                        exception=e,
+                        db=db,
+                    )