fractal-server 2.9.0a1__py3-none-any.whl → 2.9.0a3__py3-none-any.whl

fractal_server/tasks/v2/ssh/_utils.py

@@ -0,0 +1,87 @@
+import os
+from pathlib import Path
+
+from fractal_server.app.models.v2 import TaskGroupV2
+from fractal_server.logger import get_logger
+from fractal_server.ssh._fabric import FractalSSH
+from fractal_server.tasks.v2.utils_templates import customize_template
+
+
+def _customize_and_run_template(
+    *,
+    template_filename: str,
+    replacements: list[tuple[str, str]],
+    script_dir_local: str,
+    prefix: str,
+    fractal_ssh: FractalSSH,
+    script_dir_remote: str,
+    logger_name: str,
+) -> str:
+    """
+    Customize one of the template bash scripts, transfer it to the remote host
+    via SFTP and then run it via SSH.
+
+    Args:
+        template_filename: Filename of the template file (ends with ".sh").
+        replacements: Dictionary of replacements.
+        script_dir: Local folder where the script will be placed.
+        prefix: Prefix for the script filename.
+        fractal_ssh: FractalSSH object
+        script_dir_remote: Remote scripts directory
+    """
+    logger = get_logger(logger_name=logger_name)
+    logger.debug(f"_customize_and_run_template {template_filename} - START")
+    # Prepare name and path of script
+    if not template_filename.endswith(".sh"):
+        raise ValueError(
+            f"Invalid {template_filename=} (it must end with '.sh')."
+        )
+    script_filename = f"{prefix}_{template_filename}"
+    script_path_local = (Path(script_dir_local) / script_filename).as_posix()
+
+    customize_template(
+        template_name=template_filename,
+        replacements=replacements,
+        script_path=script_path_local,
+    )
+
+    # Transfer script to remote host
+    script_path_remote = os.path.join(
+        script_dir_remote,
+        script_filename,
+    )
+    logger.debug(f"Now transfer {script_path_local=} over SSH.")
+    fractal_ssh.send_file(
+        local=script_path_local,
+        remote=script_path_remote,
+    )
+
+    # Execute script remotely
+    cmd = f"bash {script_path_remote}"
+    logger.debug(f"Now run '{cmd}' over SSH.")
+    stdout = fractal_ssh.run_command(cmd=cmd)
+
+    logger.debug(f"_customize_and_run_template {template_filename} - END")
+    return stdout
+
+
+def _copy_wheel_file_ssh(
+    *, task_group: TaskGroupV2, fractal_ssh: FractalSSH, logger_name: str
+) -> str:
+    """
+    Handle the situation where `task_group.wheel_path` is not part of
+    `task_group.path`, by copying `wheel_path` into `path`.
+
+    Returns:
+        The new `wheel_path`.
+    """
+    logger = get_logger(logger_name=logger_name)
+    source = task_group.wheel_path
+    dest = (
+        Path(task_group.path) / Path(task_group.wheel_path).name
+    ).as_posix()
+    cmd = f"cp {source} {dest}"
+    logger.debug(f"[_copy_wheel_file] START {source=} {dest=}")
+    fractal_ssh.run_command(cmd=cmd)
+    logger.debug(f"[_copy_wheel_file] END {source=} {dest=}")
+    return dest

fractal_server/tasks/v2/ssh/collect.py

@@ -1,5 +1,4 @@
 import logging
-import os
 import time
 from pathlib import Path
 from tempfile import TemporaryDirectory
@@ -13,16 +12,16 @@ from fractal_server.app.models.v2 import TaskGroupV2
 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.logger import get_logger
 from fractal_server.logger import set_logger
 from fractal_server.ssh._fabric import FractalSSH
+from fractal_server.tasks.v2.ssh._utils import _copy_wheel_file_ssh
+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_package_names import compare_package_names
 from fractal_server.tasks.v2.utils_python_interpreter import (
     get_python_interpreter_v2,
 )
-from fractal_server.tasks.v2.utils_templates import customize_template
 from fractal_server.tasks.v2.utils_templates import get_collection_replacements
 from fractal_server.tasks.v2.utils_templates import (
     parse_script_pip_show_stdout,
@@ -33,81 +32,6 @@ from fractal_server.utils import get_timestamp
 LOGGER_NAME = __name__
 
 
-def _customize_and_run_template(
-    *,
-    template_filename: str,
-    replacements: list[tuple[str, str]],
-    script_dir_local: str,
-    prefix: str,
-    fractal_ssh: FractalSSH,
-    script_dir_remote: str,
-    logger_name: str,
-) -> str:
-    """
-    Customize one of the template bash scripts, transfer it to the remote host
-    via SFTP and then run it via SSH.
-
-    Args:
-        template_filename: Filename of the template file (ends with ".sh").
-        replacements: Dictionary of replacements.
-        script_dir: Local folder where the script will be placed.
-        prefix: Prefix for the script filename.
-        fractal_ssh: FractalSSH object
-        script_dir_remote: Remote scripts directory
-    """
-    logger = get_logger(logger_name=logger_name)
-    logger.debug(f"_customize_and_run_template {template_filename} - START")
-
-    # Prepare name and path of script
-    if not template_filename.endswith(".sh"):
-        raise ValueError(
-            f"Invalid {template_filename=} (it must end with '.sh')."
-        )
-    script_filename = f"{prefix}_{template_filename}"
-    script_path_local = Path(script_dir_local) / script_filename
-
-    customize_template(
-        template_name=template_filename,
-        replacements=replacements,
-        script_path=script_path_local,
-    )
-
-    # Transfer script to remote host
-    script_path_remote = os.path.join(
-        script_dir_remote,
-        script_filename,
-    )
-    logger.debug(f"Now transfer {script_path_local=} over SSH.")
-    fractal_ssh.send_file(
-        local=script_path_local,
-        remote=script_path_remote,
-    )
-
-    # Execute script remotely
-    cmd = f"bash {script_path_remote}"
-    logger.debug(f"Now run '{cmd}' over SSH.")
-    stdout = fractal_ssh.run_command(cmd=cmd)
-    logger.debug(f"Standard output of '{cmd}':\n{stdout}")
-
-    logger.debug(f"_customize_and_run_template {template_filename} - END")
-    return stdout
-
-
-def _copy_wheel_file_ssh(
-    task_group: TaskGroupV2, fractal_ssh: FractalSSH
-) -> str:
-    logger = get_logger(LOGGER_NAME)
-    source = task_group.wheel_path
-    dest = (
-        Path(task_group.path) / Path(task_group.wheel_path).name
-    ).as_posix()
-    cmd = f"cp {source} {dest}"
-    logger.debug(f"[_copy_wheel_file] START {source=} {dest=}")
-    fractal_ssh.run_command(cmd=cmd)
-    logger.debug(f"[_copy_wheel_file] END {source=} {dest=}")
-    return dest
-
-
 def collect_ssh(
     *,
     task_group_id: int,
@@ -232,6 +156,7 @@ def collect_ssh(
                     new_wheel_path = _copy_wheel_file_ssh(
                         task_group=task_group,
                         fractal_ssh=fractal_ssh,
+                        logger_name=LOGGER_NAME,
                     )
                     task_group.wheel_path = new_wheel_path
                     task_group = add_commit_refresh(obj=task_group, db=db)

fractal_server/tasks/v2/ssh/deactivate.py

@@ -1,2 +1,243 @@
-def deactivate_ssh():
-    pass
+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_templates import get_collection_replacements
+from ._utils import _copy_wheel_file_ssh
+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
+from fractal_server.logger import set_logger
+from fractal_server.ssh._fabric import FractalSSH
+from fractal_server.tasks.utils import get_log_path
+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
+
+LOGGER_NAME = __name__
+
+
+def deactivate_ssh(
+    *,
+    task_group_activity_id: int,
+    task_group_id: int,
+    fractal_ssh: FractalSSH,
+    tasks_base_dir: str,
+) -> None:
+    """
+    Deactivate 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:
+        fractal_ssh:
+        tasks_base_dir:
+            Only used as a `safe_root` in `remove_dir`, and typically set to
+            `user_settings.ssh_tasks_dir`.
+    """
+
+    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:
+
+            # 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
+            try:
+                fractal_ssh.check_connection()
+            except Exception as e:
+                logger.error("Cannot establish SSH connection.")
+                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,
+                )
+                return
+
+            # Check that the (local) task_group venv_path does exist
+            if not fractal_ssh.remote_exists(task_group.venv_path):
+                error_msg = f"{task_group.venv_path} 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
+
+            try:
+
+                activity.status = TaskGroupActivityStatusV2.ONGOING
+                activity = add_commit_refresh(obj=activity, db=db)
+
+                if task_group.pip_freeze is None:
+                    logger.warning(
+                        "Recreate pip-freeze information, since "
+                        f"{task_group.pip_freeze=}. NOTE: this should only "
+                        "happen for task groups created before 2.9.0."
+                    )
+
+                    # Prepare replacements for templates
+                    replacements = get_collection_replacements(
+                        task_group=task_group,
+                        python_bin="/not/applicable",
+                    )
+
+                    # Prepare arguments for `_customize_and_run_template`
+                    script_dir_remote = (
+                        Path(task_group.path) / SCRIPTS_SUBFOLDER
+                    ).as_posix()
+                    common_args = dict(
+                        replacements=replacements,
+                        script_dir_local=(
+                            Path(tmpdir) / SCRIPTS_SUBFOLDER
+                        ).as_posix(),
+                        script_dir_remote=script_dir_remote,
+                        prefix=(
+                            f"{int(time.time())}_"
+                            f"{TaskGroupActivityActionV2.DEACTIVATE}"
+                        ),
+                        fractal_ssh=fractal_ssh,
+                        logger_name=LOGGER_NAME,
+                    )
+
+                    # Run `pip freeze`
+                    pip_freeze_stdout = _customize_and_run_template(
+                        template_filename="3_pip_freeze.sh",
+                        **common_args,
+                    )
+
+                    # Update pip-freeze data
+                    logger.info("Add pip freeze stdout to TaskGroupV2 - start")
+                    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 = add_commit_refresh(obj=task_group, db=db)
+                    logger.info("Add pip freeze stdout to TaskGroupV2 - end")
+
+                # Handle some specific cases for wheel-file case
+                if task_group.origin == TaskGroupV2OriginEnum.WHEELFILE:
+
+                    logger.info(
+                        f"Handle specific cases for {task_group.origin=}."
+                    )
+
+                    # Blocking situation: `wheel_path` is not set or points
+                    # to a missing path
+                    if (
+                        task_group.wheel_path is None
+                        or not fractal_ssh.remote_exists(task_group.wheel_path)
+                    ):
+                        error_msg = (
+                            "Invalid wheel path for task group with "
+                            f"{task_group_id=}. {task_group.wheel_path=} is "
+                            "unset or 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
+
+                    # Recoverable situation: `wheel_path` was not yet copied
+                    # over to the correct server-side folder
+                    wheel_path_parent_dir = Path(task_group.wheel_path).parent
+                    if wheel_path_parent_dir != Path(task_group.path):
+                        logger.warning(
+                            f"{wheel_path_parent_dir.as_posix()} differs from "
+                            f"{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:
+                            raise ValueError(
+                                f"Cannot find {task_group.wheel_path=} in "
+                                "pip-freeze data. Exit."
+                            )
+
+                        logger.info(
+                            f"Now copy wheel file into {task_group.path}."
+                        )
+                        new_wheel_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}.")
+
+                        task_group.wheel_path = new_wheel_path
+                        new_pip_freeze = task_group.pip_freeze.replace(
+                            task_group.wheel_path,
+                            new_wheel_path,
+                        )
+                        task_group.pip_freeze = new_pip_freeze
+                        task_group = add_commit_refresh(obj=task_group, db=db)
+                        logger.info(
+                            "Updated `wheel_path` and `pip_freeze` "
+                            "task-group attributes."
+                        )
+
+                # We now have all required information for reactivating the
+                # virtual environment at a later point
+                logger.info(f"Now removing {task_group.venv_path}.")
+                fractal_ssh.remove_folder(
+                    folder=task_group.venv_path,
+                    safe_root=tasks_base_dir,
+                )
+                logger.info(f"All good, {task_group.venv_path} 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)
+
+            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,
+                )
+        return

fractal_server/tasks/v2/ssh/reactivate.py

@@ -1,2 +1,199 @@
-def reactivate_ssh():
-    pass
+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_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 set_logger
+from fractal_server.ssh._fabric import FractalSSH
+from fractal_server.tasks.utils import get_log_path
+from fractal_server.tasks.v2.utils_background import get_current_log
+from fractal_server.tasks.v2.utils_python_interpreter import (
+    get_python_interpreter_v2,
+)
+from fractal_server.tasks.v2.utils_templates import SCRIPTS_SUBFOLDER
+from fractal_server.utils import get_timestamp
+
+LOGGER_NAME = __name__
+
+
+def reactivate_ssh(
+    *,
+    task_group_activity_id: int,
+    task_group_id: int,
+    fractal_ssh: FractalSSH,
+    tasks_base_dir: str,
+) -> 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:
+        fractal_ssh:
+        tasks_base_dir:
+            Only used as a `safe_root` in `remove_dir`, and typically set to
+            `user_settings.ssh_tasks_dir`.
+    """
+
+    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:
+
+            # 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
+            try:
+                fractal_ssh.check_connection()
+            except Exception as e:
+                logger.error("Cannot establish SSH connection.")
+                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,
+                )
+                return
+
+            # Check that the (remote) task_group venv_path does not exist
+            if fractal_ssh.remote_exists(task_group.venv_path):
+                error_msg = f"{task_group.venv_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
+
+            try:
+                activity.status = TaskGroupActivityStatusV2.ONGOING
+                activity = add_commit_refresh(obj=activity, db=db)
+
+                # Prepare replacements for templates
+                replacements = get_collection_replacements(
+                    task_group=task_group,
+                    python_bin=get_python_interpreter_v2(
+                        python_version=task_group.python_version
+                    ),
+                )
+
+                # Prepare replacements for templates
+                pip_freeze_file_local = f"{tmpdir}/pip_freeze.txt"
+                pip_freeze_file_remote = (
+                    Path(task_group.path) / "_tmp_pip_freeze.txt"
+                ).as_posix()
+                with open(pip_freeze_file_local, "w") as f:
+                    f.write(task_group.pip_freeze)
+                fractal_ssh.send_file(
+                    local=pip_freeze_file_local, remote=pip_freeze_file_remote
+                )
+                replacements.append(
+                    ("__PIP_FREEZE_FILE__", pip_freeze_file_remote)
+                )
+
+                # Prepare common arguments for `_customize_and_run_template``
+                script_dir_remote = (
+                    Path(task_group.path) / SCRIPTS_SUBFOLDER
+                ).as_posix()
+                common_args = dict(
+                    replacements=replacements,
+                    script_dir_local=(
+                        Path(tmpdir) / SCRIPTS_SUBFOLDER
+                    ).as_posix(),
+                    script_dir_remote=script_dir_remote,
+                    prefix=(
+                        f"{int(time.time())}_"
+                        f"{TaskGroupActivityActionV2.REACTIVATE}"
+                    ),
+                    fractal_ssh=fractal_ssh,
+                    logger_name=LOGGER_NAME,
+                )
+
+                # Create remote directory for scripts
+                fractal_ssh.mkdir(folder=script_dir_remote)
+
+                logger.debug("start - create venv")
+                _customize_and_run_template(
+                    template_filename="1_create_venv.sh",
+                    **common_args,
+                )
+                logger.debug("end - create venv")
+                activity.log = get_current_log(log_file_path)
+                activity.timestamp_ended = get_timestamp()
+                activity = add_commit_refresh(obj=activity, db=db)
+
+                logger.debug("start - install from pip freeze")
+                _customize_and_run_template(
+                    template_filename="6_pip_install_from_freeze.sh",
+                    **common_args,
+                )
+                logger.debug("end - install from pip freeze")
+                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")
+
+            except Exception as reactivate_e:
+                # Delete corrupted venv_path
+                try:
+                    logger.info(f"Now delete folder {task_group.venv_path}")
+                    fractal_ssh.remove_folder(
+                        folder=task_group.venv_path,
+                        safe_root=tasks_base_dir,
+                    )
+                    logger.info(f"Deleted folder {task_group.venv_path}")
+                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/utils_background.py

@@ -119,30 +119,6 @@ def _prepare_tasks_metadata(
     return task_list
 
 
-def check_task_files_exist(task_list: list[TaskCreateV2]) -> None:
-    """
-    Check that the modules listed in task commands point to existing files.
-
-    Args:
-        task_list:
-    """
-    for _task in task_list:
-        if _task.command_non_parallel is not None:
-            _task_path = _task.command_non_parallel.split()[1]
-            if not Path(_task_path).exists():
-                raise FileNotFoundError(
-                    f"Task `{_task.name}` has `command_non_parallel` "
-                    f"pointing to missing file `{_task_path}`."
-                )
-        if _task.command_parallel is not None:
-            _task_path = _task.command_parallel.split()[1]
-            if not Path(_task_path).exists():
-                raise FileNotFoundError(
-                    f"Task `{_task.name}` has `command_parallel` "
-                    f"pointing to missing file `{_task_path}`."
-                )
-
-
 def get_current_log(logger_file_path: str) -> str:
     with open(logger_file_path, "r") as f:
         return f.read()