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

fractal_server/tasks/v2/background_operations.py

@@ -5,288 +5,230 @@ is used as a background task for the task-collection endpoint.
 import json
 from pathlib import Path
 from shutil import rmtree as shell_rmtree
+from typing import Optional
 
-from ..utils import _init_venv
-from ..utils import _normalize_package_name
+from sqlalchemy.orm import Session as DBSyncSession
+from sqlalchemy.orm.attributes import flag_modified
+
+from ..utils import get_absolute_venv_path
+from ..utils import get_collection_freeze
 from ..utils import get_collection_log
 from ..utils import get_collection_path
 from ..utils import get_log_path
 from ..utils import slugify_task_name
 from ._TaskCollectPip import _TaskCollectPip
-from fractal_server.app.db import DBSyncSession
 from fractal_server.app.db import get_sync_db
 from fractal_server.app.models.v2 import CollectionStateV2
 from fractal_server.app.models.v2 import TaskV2
-from fractal_server.app.schemas.v2 import TaskCollectStatusV2
+from fractal_server.app.schemas.v2 import CollectionStatusV2
 from fractal_server.app.schemas.v2 import TaskCreateV2
 from fractal_server.app.schemas.v2 import TaskReadV2
+from fractal_server.app.schemas.v2.manifest import ManifestV2
 from fractal_server.logger import get_logger
 from fractal_server.logger import reset_logger_handlers
 from fractal_server.logger import set_logger
-from fractal_server.utils import execute_command
+from fractal_server.tasks.v2._venv_pip import _create_venv_install_package_pip
 
 
-async def _pip_install(
-    venv_path: Path,
-    task_pkg: _TaskCollectPip,
-    logger_name: str,
-) -> Path:
-    """
-    Install package in venv
+def _get_task_type(task: TaskCreateV2) -> str:
+    if task.command_non_parallel is None:
+        return "parallel"
+    elif task.command_parallel is None:
+        return "non_parallel"
+    else:
+        return "compound"
 
-    Args:
-        venv_path:
-        task_pkg:
-        logger_name:
 
-    Returns:
-        The location of the package.
+def _insert_tasks(
+    task_list: list[TaskCreateV2],
+    db: DBSyncSession,
+    owner: Optional[str] = None,
+) -> list[TaskV2]:
+    """
+    Insert tasks into database
     """
 
-    logger = get_logger(logger_name)
-
-    pip = venv_path / "venv/bin/pip"
-
-    extras = f"[{task_pkg.package_extras}]" if task_pkg.package_extras else ""
+    owner_dict = dict(owner=owner) if owner is not None else dict()
 
-    if task_pkg.is_local_package:
-        pip_install_str = f"{task_pkg.package_path.as_posix()}{extras}"
-    else:
-        version_string = (
-            f"=={task_pkg.package_version}" if task_pkg.package_version else ""
-        )
-        pip_install_str = f"{task_pkg.package}{extras}{version_string}"
-
-    cmd_install = f"{pip} install {pip_install_str}"
-    cmd_inspect = f"{pip} show {task_pkg.package}"
+    task_db_list = [
+        TaskV2(**t.dict(), **owner_dict, type=_get_task_type(t))
+        for t in task_list
+    ]
+    db.add_all(task_db_list)
+    db.commit()
+    for t in task_db_list:
+        db.refresh(t)
+    db.close()
+    return task_db_list
 
-    await execute_command(
-        cwd=venv_path,
-        command=f"{pip} install --upgrade pip",
-        logger_name=logger_name,
-    )
-    await execute_command(
-        cwd=venv_path, command=cmd_install, logger_name=logger_name
-    )
-    if task_pkg.pinned_package_versions:
-        for (
-            pinned_pkg_name,
-            pinned_pkg_version,
-        ) in task_pkg.pinned_package_versions.items():
-
-            logger.debug(
-                "Specific version required: "
-                f"{pinned_pkg_name}=={pinned_pkg_version}"
-            )
-            logger.debug(
-                "Preliminary check: verify that "
-                f"{pinned_pkg_version} is already installed"
-            )
-            stdout_inspect = await execute_command(
-                cwd=venv_path,
-                command=f"{pip} show {pinned_pkg_name}",
-                logger_name=logger_name,
-            )
-            current_version = next(
-                line.split()[-1]
-                for line in stdout_inspect.split("\n")
-                if line.startswith("Version:")
-            )
-            if current_version != pinned_pkg_version:
-                logger.debug(
-                    f"Currently installed version of {pinned_pkg_name} "
-                    f"({current_version}) differs from pinned version "
-                    f"({pinned_pkg_version}); "
-                    f"install version {pinned_pkg_version}."
-                )
-                await execute_command(
-                    cwd=venv_path,
-                    command=(
-                        f"{pip} install "
-                        f"{pinned_pkg_name}=={pinned_pkg_version}"
-                    ),
-                    logger_name=logger_name,
-                )
-            else:
-                logger.debug(
-                    f"Currently installed version of {pinned_pkg_name} "
-                    f"({current_version}) already matches the pinned version."
-                )
 
-    # Extract package installation path from `pip show`
-    stdout_inspect = await execute_command(
-        cwd=venv_path, command=cmd_inspect, logger_name=logger_name
-    )
+def _set_collection_state_data_status(
+    *,
+    state_id: int,
+    new_status: CollectionStatusV2,
+    logger_name: str,
+    db: DBSyncSession,
+):
+    logger = get_logger(logger_name)
+    logger.debug(f"{state_id=} - set state.data['status'] to {new_status}")
+    collection_state = db.get(CollectionStateV2, state_id)
+    collection_state.data["status"] = CollectionStatusV2(new_status)
+    flag_modified(collection_state, "data")
+    db.commit()
 
-    location = Path(
-        next(
-            line.split()[-1]
-            for line in stdout_inspect.split("\n")
-            if line.startswith("Location:")
-        )
-    )
 
-    # NOTE
-    # https://packaging.python.org/en/latest/specifications/recording-installed-packages/
-    # This directory is named as {name}-{version}.dist-info, with name and
-    # version fields corresponding to Core metadata specifications. Both
-    # fields must be normalized (see the name normalization specification and
-    # the version normalization specification), and replace dash (-)
-    # characters with underscore (_) characters, so the .dist-info directory
-    # always has exactly one dash (-) character in its stem, separating the
-    # name and version fields.
-    package_root = location / (task_pkg.package.replace("-", "_"))
-    logger.debug(f"[_pip install] {location=}")
-    logger.debug(f"[_pip install] {task_pkg.package=}")
-    logger.debug(f"[_pip install] {package_root=}")
-    if not package_root.exists():
-        raise RuntimeError(
-            "Could not determine package installation location."
-        )
-    return package_root
+def _set_collection_state_data_log(
+    *,
+    state_id: int,
+    new_log: str,
+    logger_name: str,
+    db: DBSyncSession,
+):
+    logger = get_logger(logger_name)
+    logger.debug(f"{state_id=} - set state.data['log']")
+    collection_state = db.get(CollectionStateV2, state_id)
+    collection_state.data["log"] = new_log
+    flag_modified(collection_state, "data")
+    db.commit()
 
 
-async def _create_venv_install_package(
+def _set_collection_state_data_info(
     *,
-    task_pkg: _TaskCollectPip,
-    path: Path,
+    state_id: int,
+    new_info: str,
     logger_name: str,
-) -> tuple[Path, Path]:
-    """Create venv and install package
+    db: DBSyncSession,
+):
+    logger = get_logger(logger_name)
+    logger.debug(f"{state_id=} - set state.data['info']")
+    collection_state = db.get(CollectionStateV2, state_id)
+    collection_state.data["info"] = new_info
+    flag_modified(collection_state, "data")
+    db.commit()
 
-    Args:
-        path: the directory in which to create the environment
-        task_pkg: object containing the different metadata required to install
-            the package
 
-    Returns:
-        python_bin: path to venv's python interpreter
-        package_root: the location of the package manifest
+def _handle_failure(
+    state_id: int,
+    log_file_path: Path,
+    logger_name: str,
+    exception: Exception,
+    db: DBSyncSession,
+    venv_path: Optional[Path] = None,
+):
+    """
+    Note: `venv_path` is only required to trigger the folder deletion.
     """
 
-    # Normalize package name
-    task_pkg.package_name = _normalize_package_name(task_pkg.package_name)
-    task_pkg.package = _normalize_package_name(task_pkg.package)
+    logger = get_logger(logger_name)
+    logger.error(f"Task collection failed. Original error: {str(exception)}")
 
-    python_bin = await _init_venv(
-        path=path,
-        python_version=task_pkg.python_version,
+    _set_collection_state_data_status(
+        state_id=state_id,
+        new_status=CollectionStatusV2.FAIL,
         logger_name=logger_name,
+        db=db,
     )
-    package_root = await _pip_install(
-        venv_path=path, task_pkg=task_pkg, logger_name=logger_name
+
+    new_log = log_file_path.open().read()
+    _set_collection_state_data_log(
+        state_id=state_id,
+        new_log=new_log,
+        logger_name=logger_name,
+        db=db,
     )
-    return python_bin, package_root
+    # For backwards-compatibility, we also set state.data["info"]
+    _set_collection_state_data_info(
+        state_id=state_id,
+        new_info=f"Original error: {exception}",
+        logger_name=logger_name,
+        db=db,
+    )
+    # Delete corrupted package dir
+    if venv_path is not None:
+        logger.info(f"Now delete temporary folder {venv_path}")
+        shell_rmtree(venv_path)
+        logger.info("Temporary folder deleted")
+
+    reset_logger_handlers(logger)
+    return
 
 
-async def create_package_environment_pip(
+def _prepare_tasks_metadata(
     *,
-    task_pkg: _TaskCollectPip,
-    venv_path: Path,
-    logger_name: str,
+    package_manifest: ManifestV2,
+    package_source: str,
+    python_bin: Path,
+    package_root: Path,
+    package_version: Optional[str] = None,
 ) -> list[TaskCreateV2]:
     """
-    Create environment, install package, and prepare task list
-    """
-
-    logger = get_logger(logger_name)
+    Based on the package manifest and additional info, prepare the task list.
 
-    # Normalize package name
-    task_pkg.package_name = _normalize_package_name(task_pkg.package_name)
-    task_pkg.package = _normalize_package_name(task_pkg.package)
-
-    # Only proceed if package, version and manifest attributes are set
-    task_pkg.check()
-
-    try:
-
-        logger.debug("Creating venv and installing package")
-        python_bin, package_root = await _create_venv_install_package(
-            path=venv_path,
-            task_pkg=task_pkg,
-            logger_name=logger_name,
-        )
-        logger.debug("Venv creation and package installation ended correctly.")
-
-        # Prepare task_list with appropriate metadata
-        logger.debug("Creating task list from manifest")
-        task_list = []
-        for t in task_pkg.package_manifest.task_list:
-            # Fill in attributes for TaskCreate
-            task_attributes = {}
-            task_attributes["version"] = task_pkg.package_version
-            task_name_slug = slugify_task_name(t.name)
+    Args:
+        package_manifest:
+        package_source:
+        python_bin:
+        package_root:
+        package_version:
+    """
+    task_list = []
+    for _task in package_manifest.task_list:
+        # Set non-command attributes
+        task_attributes = {}
+        if package_version is not None:
+            task_attributes["version"] = package_version
+        task_name_slug = slugify_task_name(_task.name)
+        task_attributes["source"] = f"{package_source}:{task_name_slug}"
+        if package_manifest.has_args_schemas:
             task_attributes[
-                "source"
-            ] = f"{task_pkg.package_source}:{task_name_slug}"
-            # Executables
-            if t.executable_non_parallel is not None:
-                non_parallel_path = package_root / t.executable_non_parallel
-                if not non_parallel_path.exists():
-                    raise FileNotFoundError(
-                        f"Cannot find executable `{non_parallel_path}` "
-                        f"for task `{t.name}`"
-                    )
-                task_attributes[
-                    "command_non_parallel"
-                ] = f"{python_bin.as_posix()} {non_parallel_path.as_posix()}"
-            if t.executable_parallel is not None:
-                parallel_path = package_root / t.executable_parallel
-                if not parallel_path.exists():
-                    raise FileNotFoundError(
-                        f"Cannot find executable `{parallel_path}` "
-                        f"for task `{t.name}`"
-                    )
-                task_attributes[
-                    "command_parallel"
-                ] = f"{python_bin.as_posix()} {parallel_path.as_posix()}"
-
-            manifest = task_pkg.package_manifest
-            if manifest.has_args_schemas:
-                task_attributes[
-                    "args_schema_version"
-                ] = manifest.args_schema_version
-
-            this_task = TaskCreateV2(
-                **t.dict(
-                    exclude={"executable_non_parallel", "executable_parallel"}
-                ),
-                **task_attributes,
+                "args_schema_version"
+            ] = package_manifest.args_schema_version
+        # Set command attributes
+        if _task.executable_non_parallel is not None:
+            non_parallel_path = package_root / _task.executable_non_parallel
+            task_attributes["command_non_parallel"] = (
+                f"{python_bin.as_posix()} " f"{non_parallel_path.as_posix()}"
             )
-            task_list.append(this_task)
-        logger.debug("Task list created correctly")
-    except Exception as e:
-        logger.error("Task manifest loading failed")
-        raise e
+        if _task.executable_parallel is not None:
+            parallel_path = package_root / _task.executable_parallel
+            task_attributes[
+                "command_parallel"
+            ] = f"{python_bin.as_posix()} {parallel_path.as_posix()}"
+        # Create object
+        task_obj = TaskCreateV2(
+            **_task.dict(
+                exclude={
+                    "executable_non_parallel",
+                    "executable_parallel",
+                }
+            ),
+            **task_attributes,
+        )
+        task_list.append(task_obj)
     return task_list
 
 
-def _get_task_type(task: TaskCreateV2) -> str:
-    if task.command_non_parallel is None:
-        return "parallel"
-    elif task.command_parallel is None:
-        return "non_parallel"
-    else:
-        return "compound"
-
-
-async def _insert_tasks(
-    task_list: list[TaskCreateV2],
-    db: DBSyncSession,
-) -> list[TaskV2]:
-    """
-    Insert tasks into database
+def _check_task_files_exist(task_list: list[TaskCreateV2]) -> None:
     """
+    Check that the modules listed in task commands point to existing files.
 
-    task_db_list = [
-        TaskV2(**t.dict(), type=_get_task_type(t)) for t in task_list
-    ]
-    db.add_all(task_db_list)
-    db.commit()
-    for t in task_db_list:
-        db.refresh(t)
-    db.close()
-    return task_db_list
+    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}`."
+                )
 
 
 async def background_collect_pip(
@@ -295,89 +237,109 @@ async def background_collect_pip(
     task_pkg: _TaskCollectPip,
 ) -> None:
     """
-    Install package and collect tasks
-
-    Install a python package and collect the tasks it provides according to
-    the manifest.
-
-    In case of error, copy the log into the state and delete the package
-    directory.
+    Setup venv, install package, collect tasks.
+
+    This function (executed as background task), includes the several steps
+    associated to automated collection of a Python task package.
+    1. Preliminary checks
+    2. Create venv and run `pip install`
+    3. Collect tasks into db
+    4. Finalize things.
+    5. Handle failures by copying the log into the state and deleting the
+       package directory.
     """
     logger_name = task_pkg.package.replace("/", "_")
     logger = set_logger(
         logger_name=logger_name,
         log_file_path=get_log_path(venv_path),
     )
-    logger.debug("Start background task collection")
+
+    # Start
+    logger.debug("START")
     for key, value in task_pkg.dict(exclude={"package_manifest"}).items():
-        logger.debug(f"{key}: {value}")
+        logger.debug(f"task_pkg.{key}: {value}")
 
     with next(get_sync_db()) as db:
-        state: CollectionStateV2 = db.get(CollectionStateV2, state_id)
-        data = TaskCollectStatusV2(**state.data)
-        data.info = None
 
         try:
-            # install
-            logger.debug("Task-collection status: installing")
-            data.status = "installing"
-
-            state.data = data.sanitised_dict()
-            db.merge(state)
-            db.commit()
-            task_list = await create_package_environment_pip(
-                venv_path=venv_path,
+            # Block 1: preliminary checks (only proceed if version and
+            # manifest attributes are set).
+            # Required: task_pkg
+            task_pkg.check()
+
+            # Block 2: create venv and run pip install
+            # Required: state_id, venv_path, task_pkg
+            logger.debug("installing - START")
+            _set_collection_state_data_status(
+                state_id=state_id,
+                new_status=CollectionStatusV2.INSTALLING,
+                logger_name=logger_name,
+                db=db,
+            )
+            python_bin, package_root = await _create_venv_install_package_pip(
+                path=venv_path,
                 task_pkg=task_pkg,
                 logger_name=logger_name,
             )
+            logger.debug("installing - END")
+
+            # Block 3: create task metadata and create database entries
+            # Required: state_id, python_bin, package_root, task_pkg
+            logger.debug("collecting - START")
+            _set_collection_state_data_status(
+                state_id=state_id,
+                new_status=CollectionStatusV2.COLLECTING,
+                logger_name=logger_name,
+                db=db,
+            )
+            logger.debug("collecting - prepare tasks and update db " "- START")
+            task_list = _prepare_tasks_metadata(
+                package_manifest=task_pkg.package_manifest,
+                package_version=task_pkg.package_version,
+                package_source=task_pkg.package_source,
+                package_root=package_root,
+                python_bin=python_bin,
+            )
+            _check_task_files_exist(task_list=task_list)
+            tasks = _insert_tasks(task_list=task_list, db=db)
+            logger.debug("collecting -  prepare tasks and update db " "- END")
+            logger.debug("collecting - END")
 
-            # collect
-            logger.debug("Task-collection status: collecting")
-            data.status = "collecting"
-            state.data = data.sanitised_dict()
-            db.merge(state)
-            db.commit()
-            tasks = await _insert_tasks(task_list=task_list, db=db)
-
-            # finalise
-            logger.debug("Task-collection status: finalising")
+            # Block 4: finalize (write collection files, write metadata to DB)
+            logger.debug("finalising - START")
             collection_path = get_collection_path(venv_path)
-            data.task_list = [
-                TaskReadV2(**task.model_dump()) for task in tasks
+            collection_state = db.get(CollectionStateV2, state_id)
+            task_read_list = [
+                TaskReadV2(**task.model_dump()).dict() for task in tasks
             ]
+            collection_state.data["task_list"] = task_read_list
+            collection_state.data["log"] = get_collection_log(venv_path)
+            collection_state.data["freeze"] = get_collection_freeze(venv_path)
             with collection_path.open("w") as f:
-                json.dump(data.sanitised_dict(), f, indent=2)
+                json.dump(collection_state.data, f, indent=2)
 
-            # Update DB
-            data.status = "OK"
-            data.log = get_collection_log(venv_path)
-            state.data = data.sanitised_dict()
-            db.merge(state)
+            flag_modified(collection_state, "data")
             db.commit()
-
-            # Write last logs to file
-            logger.debug("Task-collection status: OK")
-            logger.info("Background task collection completed successfully")
-            reset_logger_handlers(logger)
-
-            db.close()
+            logger.debug("finalising - END")
 
         except Exception as e:
-            # Write last logs to file
-            logger.debug("Task-collection status: fail")
-            logger.info(f"Background collection failed. Original error: {e}")
-
-            # Update db
-            data.status = "fail"
-            data.info = f"Original error: {e}"
-            data.log = get_collection_log(venv_path)
-            state.data = data.sanitised_dict()
-            db.merge(state)
-            db.commit()
-            db.close()
+            logfile_path = get_log_path(get_absolute_venv_path(venv_path))
+            _handle_failure(
+                state_id=state_id,
+                log_file_path=logfile_path,
+                logger_name=logger_name,
+                exception=e,
+                db=db,
+                venv_path=venv_path,
+            )
+            return
 
-            # Delete corrupted package dir
-            logger.info(f"Now deleting temporary folder {venv_path}")
-            shell_rmtree(venv_path)
-            logger.info("Temporary folder deleted")
-            reset_logger_handlers(logger)
+        logger.debug("Task-collection status: OK")
+        logger.info("Background task collection completed successfully")
+        _set_collection_state_data_status(
+            state_id=state_id,
+            new_status=CollectionStatusV2.OK,
+            logger_name=logger_name,
+            db=db,
+        )
+        reset_logger_handlers(logger)