fractal-server 2.7.0a2__py3-none-any.whl → 2.7.0a4__py3-none-any.whl

fractal_server/tasks/v2/background_operations.py

@@ -5,30 +5,33 @@ 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_path
 from ..utils import get_log_path
-from ._TaskCollectPip import _TaskCollectPip
-from .database_operations import create_db_task_group_and_tasks
+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 TaskGroupV2
 from fractal_server.app.schemas.v2 import CollectionStatusV2
 from fractal_server.app.schemas.v2 import TaskCreateV2
-from fractal_server.app.schemas.v2 import TaskGroupCreateV2
 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.tasks.v2._venv_pip import _create_venv_install_package_pip
+from fractal_server.tasks.v2.utils import get_python_interpreter_v2
+from fractal_server.utils import execute_command
 
 
 def _set_collection_state_data_status(
@@ -82,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.
@@ -113,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
 
@@ -125,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,
@@ -135,7 +159,6 @@ def _prepare_tasks_metadata(
 
     Args:
         package_manifest:
-        package_source:
         python_bin:
         package_root:
         package_version:
@@ -146,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"
@@ -201,13 +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,
-    user_id: int,
-    user_group_id: Optional[int],
+    task_group: TaskGroupV2,
 ) -> None:
     """
     Setup venv, install package, collect tasks.
@@ -221,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
@@ -250,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")
@@ -267,29 +393,17 @@ 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)
 
             # Prepare some task-group attributes
-            task_group_attrs = dict(
-                pkg_name=task_pkg.package_name,
-                version=task_pkg.package_version,
-            )
-            if task_pkg.is_local_package:
-                task_group_attrs["origin"] = "wheel-file"
-            else:
-                task_group_attrs["origin"] = "pypi"
-
-            task_group = create_db_task_group_and_tasks(
+            task_group = create_db_tasks_and_update_task_group(
                 task_list=task_list,
-                task_group_obj=TaskGroupCreateV2(**task_group_attrs),
-                user_id=user_id,
-                user_group_id=user_group_id,
+                task_group_id=task_group.id,
                 db=db,
             )
 
@@ -298,15 +412,19 @@ async def background_collect_pip(
 
             # 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 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(
+                Path(task_group.path)
+            )
+            collection_state.data["freeze"] = get_collection_freeze(
+                Path(task_group.path)
+            )
             with collection_path.open("w") as f:
                 json.dump(collection_state.data, f, indent=2)
 
@@ -315,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
 

fractal_server/tasks/v2/background_operations_ssh.py

@@ -2,19 +2,17 @@ import json
 import os
 from pathlib import Path
 from tempfile import TemporaryDirectory
-from typing import Optional
 
 from sqlalchemy.orm.attributes import flag_modified
 
-from ...app.models.v2 import CollectionStateV2
-from ._TaskCollectPip import _TaskCollectPip
 from .background_operations import _handle_failure
 from .background_operations import _prepare_tasks_metadata
 from .background_operations import _set_collection_state_data_status
-from .database_operations import create_db_task_group_and_tasks
+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 TaskGroupV2
 from fractal_server.app.schemas.v2 import CollectionStatusV2
-from fractal_server.app.schemas.v2 import TaskGroupCreateV2
 from fractal_server.app.schemas.v2.manifest import ManifestV2
 from fractal_server.config import get_settings
 from fractal_server.logger import get_logger
@@ -112,11 +110,9 @@ def _customize_and_run_template(
 def background_collect_pip_ssh(
     *,
     state_id: int,
-    task_pkg: _TaskCollectPip,
+    task_group: TaskGroupV2,
     fractal_ssh: FractalSSH,
     tasks_base_dir: str,
-    user_id: int,
-    user_group_id: Optional[int],
 ) -> None:
     """
     Collect a task package over SSH
@@ -131,7 +127,7 @@ def background_collect_pip_ssh(
 
     Arguments:
         state_id:
-        task_pkg:
+        task_group:
         fractal_ssh:
         tasks_base_dir:
     """
@@ -144,47 +140,23 @@ def background_collect_pip_ssh(
             logger_name=LOGGER_NAME,
             log_file_path=log_file_path,
         )
-
         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}")
 
         # Open a DB session soon, since it is needed for updating `state`
         with next(get_sync_db()) as db:
             try:
                 # Prepare replacements for task-collection scripts
                 python_bin = get_python_interpreter_v2(
-                    python_version=task_pkg.python_version
-                )
-                package_version = (
-                    ""
-                    if task_pkg.package_version is None
-                    else task_pkg.package_version
+                    python_version=task_group.python_version
                 )
-
-                install_string = task_pkg.package
-                if task_pkg.package_extras is not None:
-                    install_string = (
-                        f"{install_string}[{task_pkg.package_extras}]"
-                    )
-                if (
-                    task_pkg.package_version is not None
-                    and not task_pkg.is_local_package
-                ):
-                    install_string = (
-                        f"{install_string}=={task_pkg.package_version}"
-                    )
-                package_env_dir = (
-                    Path(tasks_base_dir)
-                    / ".fractal"
-                    / f"{task_pkg.package_name}{package_version}"
-                ).as_posix()
-                logger.debug(f"{package_env_dir=}")
+                install_string = task_group.pip_install_string
                 settings = Inject(get_settings)
                 replacements = [
-                    ("__PACKAGE_NAME__", task_pkg.package_name),
-                    ("__PACKAGE_ENV_DIR__", package_env_dir),
-                    ("__PACKAGE__", task_pkg.package),
+                    ("__PACKAGE_NAME__", task_group.pkg_name),
+                    ("__TASK_GROUP_DIR__", task_group.path),
+                    ("__PACKAGE_ENV_DIR__", task_group.venv_path),
                     ("__PYTHON__", python_bin),
                     ("__INSTALL_STRING__", install_string),
                     (
@@ -263,13 +235,13 @@ def background_collect_pip_ssh(
                         f"collecting - parsed from pip-show: {key}={value}"
                     )
                 # Check package_name match
-                # FIXME SSH: Does this work for non-canonical `package_name`?
+                # FIXME SSH: Does this work well for non-canonical names?
                 package_name_pip_show = pkg_attrs.get("package_name")
-                package_name_task_pkg = task_pkg.package_name
-                if package_name_pip_show != package_name_task_pkg:
+                package_name_task_group = task_group.pkg_name
+                if package_name_pip_show != package_name_task_group:
                     error_msg = (
                         f"`package_name` mismatch: "
-                        f"{package_name_task_pkg=} but "
+                        f"{package_name_task_group=} but "
                         f"{package_name_pip_show=}"
                     )
                     logger.error(error_msg)
@@ -292,47 +264,32 @@ def background_collect_pip_ssh(
 
                 # Read and validate remote manifest file
                 with fractal_ssh.sftp().open(manifest_path_remote, "r") as f:
-                    manifest = json.load(f)
+                    pkg_manifest_dict = json.load(f)
                 logger.info(f"collecting - loaded {manifest_path_remote=}")
-                ManifestV2(**manifest)
+                pkg_manifest = ManifestV2(**pkg_manifest_dict)
                 logger.info("collecting - manifest is a valid ManifestV2")
 
-                # Create new _TaskCollectPip object
-                new_pkg = _TaskCollectPip(
-                    **task_pkg.dict(
-                        exclude={"package_version", "package_name"},
-                        exclude_unset=True,
-                        exclude_none=True,
-                    ),
-                    package_manifest=manifest,
-                    **pkg_attrs,
-                )
-
+                logger.info("collecting - _prepare_tasks_metadata - start")
                 task_list = _prepare_tasks_metadata(
-                    package_manifest=new_pkg.package_manifest,
-                    package_version=new_pkg.package_version,
-                    package_source=new_pkg.package_source,
+                    package_manifest=pkg_manifest,
+                    package_version=task_group.version,
                     package_root=Path(package_root_remote),
                     python_bin=Path(python_bin),
                 )
+                logger.info("collecting - _prepare_tasks_metadata - end")
 
-                # Prepare some task-group attributes
-                task_group_attrs = dict(
-                    pkg_name=task_pkg.package_name,
-                    version=task_pkg.package_version,
+                logger.info(
+                    "collecting - create_db_tasks_and_update_task_group - "
+                    "start"
                 )
-                if task_pkg.is_local_package:
-                    task_group_attrs["origin"] = "wheel-file"
-                else:
-                    task_group_attrs["origin"] = "pypi"
-
-                create_db_task_group_and_tasks(
+                create_db_tasks_and_update_task_group(
                     task_list=task_list,
-                    task_group_obj=TaskGroupCreateV2(**task_group_attrs),
-                    user_id=user_id,
-                    user_group_id=user_group_id,
+                    task_group_id=task_group.id,
                     db=db,
                 )
+                logger.info(
+                    "collecting - create_db_tasks_and_update_task_group - end"
+                )
 
                 logger.debug("collecting - END")
 
@@ -356,18 +313,19 @@ def background_collect_pip_ssh(
                     logger_name=LOGGER_NAME,
                     exception=e,
                     db=db,
+                    task_group_id=task_group.id,
                 )
                 if remove_venv_folder_upon_failure:
                     try:
                         logger.info(
-                            f"Now delete remote folder {package_env_dir}"
+                            f"Now delete remote folder {task_group.path}"
                         )
                         fractal_ssh.remove_folder(
-                            folder=package_env_dir,
+                            folder=task_group.path,
                             safe_root=tasks_base_dir,
                         )
                         logger.info(
-                            f"Deleted remoted folder {package_env_dir}"
+                            f"Deleted remoted folder {task_group.path}"
                         )
                     except Exception as e:
                         logger.error(
@@ -377,6 +335,6 @@ def background_collect_pip_ssh(
                     else:
                         logger.info(
                             "Not trying to remove remote folder "
-                            f"{package_env_dir}."
+                            f"{task_group.path}."
                         )
                 return

fractal_server/tasks/v2/database_operations.py

@@ -1,11 +1,8 @@
-from typing import Optional
-
 from sqlalchemy.orm import Session as DBSyncSession
 
 from fractal_server.app.models.v2 import TaskGroupV2
 from fractal_server.app.models.v2 import TaskV2
 from fractal_server.app.schemas.v2 import TaskCreateV2
-from fractal_server.app.schemas.v2 import TaskGroupCreateV2
 
 
 def _get_task_type(task: TaskCreateV2) -> str:
@@ -17,22 +14,18 @@ def _get_task_type(task: TaskCreateV2) -> str:
         return "compound"
 
 
-def create_db_task_group_and_tasks(
+def create_db_tasks_and_update_task_group(
     *,
+    task_group_id: int,
     task_list: list[TaskCreateV2],
-    task_group_obj: TaskGroupCreateV2,
-    user_id: int,
     db: DBSyncSession,
-    user_group_id: Optional[int] = None,
 ) -> TaskGroupV2:
     """
     Create a `TaskGroupV2` with N `TaskV2`s, and insert them into the database.
 
     Arguments:
-        task_group:
-        task_list:
-        user_id:
-        user_group_id: Can be `None`
+        task_group: ID of an existing TaskGroupV2 object.
+        task_list: A list of TaskCreateV2 objects to be inserted into the db.
         db: A synchronous database session
     """
     actual_task_list = [
@@ -42,13 +35,10 @@ def create_db_task_group_and_tasks(
         )
         for task in task_list
     ]
-    task_group = TaskGroupV2(
-        user_id=user_id,
-        user_group_id=user_group_id,
-        task_list=actual_task_list,
-        **task_group_obj.dict(),
-    )
+    task_group = db.get(TaskGroupV2, task_group_id)
+    task_group.task_list = actual_task_list
     db.add(task_group)
     db.commit()
     db.refresh(task_group)
+
     return task_group