fractal-server 2.6.4__py3-none-any.whl → 2.7.0__py3-none-any.whl

fractal_server/tasks/v2/_venv_pip.py

@@ -2,17 +2,48 @@ from pathlib import Path
 from typing import Optional
 
 from ..utils import COLLECTION_FREEZE_FILENAME
+from fractal_server.app.models.v2 import TaskGroupV2
 from fractal_server.config import get_settings
 from fractal_server.logger import get_logger
 from fractal_server.syringe import Inject
-from fractal_server.tasks.v2._TaskCollectPip import _TaskCollectPip
 from fractal_server.tasks.v2.utils import get_python_interpreter_v2
 from fractal_server.utils import execute_command
 
 
-async def _pip_install(
+async def _init_venv_v2(
+    *,
     venv_path: Path,
-    task_pkg: _TaskCollectPip,
+    python_version: Optional[str] = None,
+    logger_name: str,
+) -> Path:
+    """
+    Set a virtual environment at `path/venv`
+
+    Args:
+        path : Path
+            path to the venv actual directory (not its parent).
+        python_version : default=None
+            Python version the virtual environment will be based upon
+
+    Returns:
+        python_bin : Path
+            path to python interpreter
+    """
+    logger = get_logger(logger_name)
+    logger.debug(f"[_init_venv_v2] {venv_path=}")
+    interpreter = get_python_interpreter_v2(python_version=python_version)
+    logger.debug(f"[_init_venv_v2] {interpreter=}")
+    await execute_command(
+        command=f"{interpreter} -m venv {venv_path}",
+        logger_name=logger_name,
+    )
+    python_bin = venv_path / "bin/python"
+    logger.debug(f"[_init_venv_v2] {python_bin=}")
+    return python_bin
+
+
+async def _pip_install(
+    task_group: TaskGroupV2,
     logger_name: str,
 ) -> Path:
     """
@@ -30,48 +61,45 @@ async def _pip_install(
 
     logger = get_logger(logger_name)
 
-    pip = venv_path / "venv/bin/pip"
-
-    extras = f"[{task_pkg.package_extras}]" if task_pkg.package_extras else ""
-
-    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_name}{extras}{version_string}"
+    python_bin = Path(task_group.venv_path) / "bin/python"
+    pip_install_str = task_group.pip_install_string
+    logger.info(f"{pip_install_str=}")
 
     await execute_command(
-        cwd=venv_path,
+        cwd=Path(task_group.venv_path),
         command=(
-            f"{pip} install --upgrade "
+            f"{python_bin} -m pip install --upgrade "
             f"'pip<={settings.FRACTAL_MAX_PIP_VERSION}'"
         ),
         logger_name=logger_name,
     )
     await execute_command(
-        cwd=venv_path,
-        command=f"{pip} install {pip_install_str}",
+        cwd=Path(task_group.venv_path),
+        command=f"{python_bin} -m pip install setuptools",
+        logger_name=logger_name,
+    )
+    await execute_command(
+        cwd=Path(task_group.venv_path),
+        command=f"{python_bin} -m pip install {pip_install_str}",
         logger_name=logger_name,
     )
-    if task_pkg.pinned_package_versions:
+
+    if task_group.pinned_package_versions:
         for (
             pinned_pkg_name,
             pinned_pkg_version,
-        ) in task_pkg.pinned_package_versions.items():
-
+        ) in task_group.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"
+                f"{pinned_pkg_name} is already installed"
             )
             stdout_show = await execute_command(
-                cwd=venv_path,
-                command=f"{pip} show {pinned_pkg_name}",
+                cwd=Path(task_group.venv_path),
+                command=f"{python_bin} -m pip show {pinned_pkg_name}",
                 logger_name=logger_name,
             )
             current_version = next(
@@ -87,9 +115,9 @@ async def _pip_install(
                     f"install version {pinned_pkg_version}."
                 )
                 await execute_command(
-                    cwd=venv_path,
+                    cwd=Path(task_group.venv_path),
                     command=(
-                        f"{pip} install "
+                        f"{python_bin} -m pip install "
                         f"{pinned_pkg_name}=={pinned_pkg_version}"
                     ),
                     logger_name=logger_name,
@@ -102,8 +130,8 @@ async def _pip_install(
 
     # Extract package installation path from `pip show`
     stdout_show = await execute_command(
-        cwd=venv_path,
-        command=f"{pip} show {task_pkg.package_name}",
+        cwd=Path(task_group.venv_path),
+        command=f"{python_bin} -m pip show {task_group.pkg_name}",
         logger_name=logger_name,
     )
 
@@ -124,58 +152,26 @@ async def _pip_install(
     # 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_name.replace("-", "_"))
+    package_root = location / (task_group.pkg_name.replace("-", "_"))
     logger.debug(f"[_pip install] {location=}")
-    logger.debug(f"[_pip install] {task_pkg.package_name=}")
+    logger.debug(f"[_pip install] {task_group.pkg_name=}")
     logger.debug(f"[_pip install] {package_root=}")
 
     # Run `pip freeze --all` and store its output
     stdout_freeze = await execute_command(
-        cwd=venv_path, command=f"{pip} freeze --all", logger_name=logger_name
+        cwd=Path(task_group.venv_path),
+        command=f"{python_bin} -m pip freeze --all",
+        logger_name=logger_name,
     )
-    with (venv_path / COLLECTION_FREEZE_FILENAME).open("w") as f:
+    with (Path(task_group.path) / COLLECTION_FREEZE_FILENAME).open("w") as f:
         f.write(stdout_freeze)
 
     return package_root
 
 
-async def _init_venv_v2(
-    *,
-    path: Path,
-    python_version: Optional[str] = None,
-    logger_name: str,
-) -> Path:
-    """
-    Set a virtual environment at `path/venv`
-
-    Args:
-        path : Path
-            path to directory in which to set up the virtual environment
-        python_version : default=None
-            Python version the virtual environment will be based upon
-
-    Returns:
-        python_bin : Path
-            path to python interpreter
-    """
-    logger = get_logger(logger_name)
-    logger.debug(f"[_init_venv] {path=}")
-    interpreter = get_python_interpreter_v2(python_version=python_version)
-    logger.debug(f"[_init_venv] {interpreter=}")
-    await execute_command(
-        cwd=path,
-        command=f"{interpreter} -m venv venv",
-        logger_name=logger_name,
-    )
-    python_bin = path / "venv/bin/python"
-    logger.debug(f"[_init_venv] {python_bin=}")
-    return python_bin
-
-
 async def _create_venv_install_package_pip(
     *,
-    task_pkg: _TaskCollectPip,
-    path: Path,
+    task_group: TaskGroupV2,
     logger_name: str,
 ) -> tuple[Path, Path]:
     """
@@ -191,11 +187,12 @@ async def _create_venv_install_package_pip(
         package_root: the location of the package manifest
     """
     python_bin = await _init_venv_v2(
-        path=path,
-        python_version=task_pkg.python_version,
+        venv_path=Path(task_group.venv_path),
+        python_version=task_group.python_version,
         logger_name=logger_name,
     )
     package_root = await _pip_install(
-        venv_path=path, task_pkg=task_pkg, logger_name=logger_name
+        task_group=task_group,
+        logger_name=logger_name,
     )
     return python_bin, package_root

fractal_server/tasks/v2/background_operations.py

@@ -5,21 +5,23 @@ 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 tempfile import TemporaryDirectory
 from typing import Optional
+from typing import Union
+from zipfile import ZipFile
 
 from sqlalchemy.orm import Session as DBSyncSession
 from sqlalchemy.orm.attributes import flag_modified
+from sqlmodel import select
 
-from ...string_tools import slugify_task_name_for_source
-from ..utils import get_absolute_venv_path
-from ..utils import get_collection_freeze
-from ..utils import get_collection_log
+from ..utils import get_collection_freeze_v2
+from ..utils import get_collection_log_v2
 from ..utils import get_collection_path
 from ..utils import get_log_path
-from ._TaskCollectPip import _TaskCollectPip
+from .database_operations import create_db_tasks_and_update_task_group
 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.models.v2 import TaskGroupV2
 from fractal_server.app.schemas.v2 import CollectionStatusV2
 from fractal_server.app.schemas.v2 import TaskCreateV2
 from fractal_server.app.schemas.v2 import TaskReadV2
@@ -28,38 +30,8 @@ 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.tasks.v2._venv_pip import _create_venv_install_package_pip
-
-
-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"
-
-
-def _insert_tasks(
-    task_list: list[TaskCreateV2],
-    db: DBSyncSession,
-    owner: Optional[str] = None,
-) -> list[TaskV2]:
-    """
-    Insert tasks into database
-    """
-
-    owner_dict = dict(owner=owner) if owner is not None else dict()
-
-    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
+from fractal_server.tasks.v2.utils import get_python_interpreter_v2
+from fractal_server.utils import execute_command
 
 
 def _set_collection_state_data_status(
@@ -113,7 +85,8 @@ def _handle_failure(
     logger_name: str,
     exception: Exception,
     db: DBSyncSession,
-    venv_path: Optional[Path] = None,
+    task_group_id: int,
+    path: Optional[Path] = None,
 ):
     """
     Note: `venv_path` is only required to trigger the folder deletion.
@@ -144,11 +117,32 @@ def _handle_failure(
         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)
+    if path is not None and Path(path).exists():
+        logger.info(f"Now delete temporary folder {path}")
+        shell_rmtree(path)
         logger.info("Temporary folder deleted")
 
+    # Delete TaskGroupV2 object / and apply cascade operation to FKs
+    logger.info(f"Now delete TaskGroupV2 with {task_group_id=}")
+    logger.info("Start of CollectionStateV2 cascade operations.")
+    stm = select(CollectionStateV2).where(
+        CollectionStateV2.taskgroupv2_id == task_group_id
+    )
+    res = db.execute(stm)
+    collection_states = res.scalars().all()
+    for collection_state in collection_states:
+        logger.info(
+            f"Setting CollectionStateV2[{collection_state.id}].taskgroupv2_id "
+            "to None."
+        )
+        collection_state.taskgroupv2_id = None
+        db.add(collection_state)
+    logger.info("End of CollectionStateV2 cascade operations.")
+    task_group = db.get(TaskGroupV2, task_group_id)
+    db.delete(task_group)
+    db.commit()
+    logger.info(f"TaskGroupV2 with {task_group_id=} deleted")
+
     reset_logger_handlers(logger)
     return
 
@@ -156,7 +150,6 @@ def _handle_failure(
 def _prepare_tasks_metadata(
     *,
     package_manifest: ManifestV2,
-    package_source: str,
     python_bin: Path,
     package_root: Path,
     package_version: Optional[str] = None,
@@ -166,7 +159,6 @@ def _prepare_tasks_metadata(
 
     Args:
         package_manifest:
-        package_source:
         python_bin:
         package_root:
         package_version:
@@ -177,8 +169,6 @@ def _prepare_tasks_metadata(
         task_attributes = {}
         if package_version is not None:
             task_attributes["version"] = package_version
-        task_name_slug = slugify_task_name_for_source(_task.name)
-        task_attributes["source"] = f"{package_source}:{task_name_slug}"
         if package_manifest.has_args_schemas:
             task_attributes[
                 "args_schema_version"
@@ -203,6 +193,7 @@ def _prepare_tasks_metadata(
                 }
             ),
             **task_attributes,
+            authors=package_manifest.authors,
         )
         task_list.append(task_obj)
     return task_list
@@ -231,10 +222,95 @@ def _check_task_files_exist(task_list: list[TaskCreateV2]) -> None:
                 )
 
 
+async def _download_package(
+    *,
+    python_version: str,
+    pkg_name: str,
+    version: str,
+    dest: Union[str, Path],
+) -> Path:
+    """
+    Download package to destination and return wheel-file path.
+    """
+    python_bin = get_python_interpreter_v2(python_version=python_version)
+    pip = f"{python_bin} -m pip"
+    package_and_version = f"{pkg_name}=={version}"
+    cmd = f"{pip} download --no-deps {package_and_version} -d {dest}"
+    stdout = await execute_command(command=cmd)
+    pkg_file = next(
+        line.split()[-1] for line in stdout.split("\n") if "Saved" in line
+    )
+    return Path(pkg_file)
+
+
+def _load_manifest_from_wheel(
+    wheel_file_path: str,
+    logger_name: str,
+) -> ManifestV2:
+    """
+    Given a wheel file on-disk, extract the Fractal manifest.
+    """
+    logger = get_logger(logger_name)
+
+    with ZipFile(wheel_file_path) as wheel:
+        namelist = wheel.namelist()
+        try:
+            manifest = next(
+                name
+                for name in namelist
+                if "__FRACTAL_MANIFEST__.json" in name
+            )
+        except StopIteration:
+            msg = (
+                f"{wheel_file_path} does not include __FRACTAL_MANIFEST__.json"
+            )
+            logger.error(msg)
+            raise ValueError(msg)
+        with wheel.open(manifest) as manifest_fd:
+            manifest_dict = json.load(manifest_fd)
+    manifest_version = str(manifest_dict["manifest_version"])
+    if manifest_version != "2":
+        msg = f"Manifest version {manifest_version=} not supported"
+        logger.error(msg)
+        raise ValueError(msg)
+    pkg_manifest = ManifestV2(**manifest_dict)
+    return pkg_manifest
+
+
+async def _get_package_manifest(
+    *,
+    task_group: TaskGroupV2,
+    logger_name: str,
+) -> ManifestV2:
+    wheel_file_path = task_group.wheel_path
+    if wheel_file_path is None:
+        with TemporaryDirectory() as tmpdir:
+            # Copy or download the package wheel file to tmpdir
+            wheel_file_path = await _download_package(
+                python_version=task_group.python_version,
+                pkg_name=task_group.pkg_name,
+                version=task_group.version,
+                dest=tmpdir,
+            )
+            wheel_file_path = wheel_file_path.as_posix()
+            # Read package manifest from temporary wheel file
+            manifest = _load_manifest_from_wheel(
+                wheel_file_path=wheel_file_path,
+                logger_name=logger_name,
+            )
+    else:
+        # Read package manifest from wheel file
+        manifest = _load_manifest_from_wheel(
+            wheel_file_path=wheel_file_path,
+            logger_name=logger_name,
+        )
+    return manifest
+
+
 async def background_collect_pip(
+    *,
     state_id: int,
-    venv_path: Path,
-    task_pkg: _TaskCollectPip,
+    task_group: TaskGroupV2,
 ) -> None:
     """
     Setup venv, install package, collect tasks.
@@ -248,24 +324,48 @@ async def background_collect_pip(
     5. Handle failures by copying the log into the state and deleting the
        package directory.
     """
-    logger_name = task_pkg.package.replace("/", "_")
+    logger_name = (
+        f"{task_group.user_id}-{task_group.pkg_name}-{task_group.version}"
+    )
+
+    try:
+        Path(task_group.path).mkdir(parents=True, exist_ok=False)
+    except FileExistsError as e:
+        logger = set_logger(
+            logger_name=logger_name,
+            log_file_path=get_log_path(Path(task_group.path)),
+        )
+
+        logfile_path = get_log_path(Path(task_group.path))
+        with next(get_sync_db()) as db:
+            _handle_failure(
+                state_id=state_id,
+                log_file_path=logfile_path,
+                logger_name=logger_name,
+                exception=e,
+                db=db,
+                path=None,  # Do not remove an existing path
+                task_group_id=task_group.id,
+            )
+            return
+
     logger = set_logger(
         logger_name=logger_name,
-        log_file_path=get_log_path(venv_path),
+        log_file_path=get_log_path(Path(task_group.path)),
     )
 
     # Start
     logger.debug("START")
-    for key, value in task_pkg.dict(exclude={"package_manifest"}).items():
-        logger.debug(f"task_pkg.{key}: {value}")
+    for key, value in task_group.model_dump().items():
+        logger.debug(f"task_group.{key}: {value}")
 
     with next(get_sync_db()) as db:
-
         try:
-            # Block 1: preliminary checks (only proceed if version and
-            # manifest attributes are set).
-            # Required: task_pkg
-            task_pkg.check()
+            # Block 1: get and validate manfifest
+            pkg_manifest = await _get_package_manifest(
+                task_group=task_group,
+                logger_name=logger_name,
+            )
 
             # Block 2: create venv and run pip install
             # Required: state_id, venv_path, task_pkg
@@ -277,8 +377,7 @@ async def background_collect_pip(
                 db=db,
             )
             python_bin, package_root = await _create_venv_install_package_pip(
-                path=venv_path,
-                task_pkg=task_pkg,
+                task_group=task_group,
                 logger_name=logger_name,
             )
             logger.debug("installing - END")
@@ -294,27 +393,38 @@ async def background_collect_pip(
             )
             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_manifest=pkg_manifest,
+                package_version=task_group.version,
                 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)
+
+            # Prepare some task-group attributes
+            task_group = create_db_tasks_and_update_task_group(
+                task_list=task_list,
+                task_group_id=task_group.id,
+                db=db,
+            )
+
             logger.debug("collecting -  prepare tasks and update db " "- END")
             logger.debug("collecting - END")
 
             # Block 4: finalize (write collection files, write metadata to DB)
             logger.debug("finalising - START")
-            collection_path = get_collection_path(venv_path)
+            collection_path = get_collection_path(Path(task_group.path))
             collection_state = db.get(CollectionStateV2, state_id)
             task_read_list = [
-                TaskReadV2(**task.model_dump()).dict() for task in tasks
+                TaskReadV2(**task.model_dump()).dict()
+                for task in task_group.task_list
             ]
             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)
+            collection_state.data["log"] = get_collection_log_v2(
+                Path(task_group.path)
+            )
+            collection_state.data["freeze"] = get_collection_freeze_v2(
+                Path(task_group.path)
+            )
             with collection_path.open("w") as f:
                 json.dump(collection_state.data, f, indent=2)
 
@@ -323,14 +433,15 @@ async def background_collect_pip(
             logger.debug("finalising - END")
 
         except Exception as e:
-            logfile_path = get_log_path(get_absolute_venv_path(venv_path))
+            logfile_path = get_log_path(Path(task_group.path))
             _handle_failure(
                 state_id=state_id,
                 log_file_path=logfile_path,
                 logger_name=logger_name,
                 exception=e,
                 db=db,
-                venv_path=venv_path,
+                path=task_group.path,
+                task_group_id=task_group.id,
             )
             return