fractal-server 2.8.0__py3-none-any.whl → 2.9.0a0__py3-none-any.whl

fractal_server/app/schemas/v2/task_collection.py

@@ -1,7 +1,4 @@
-from datetime import datetime
-from enum import Enum
 from pathlib import Path
-from typing import Any
 from typing import Literal
 from typing import Optional
 
@@ -11,19 +8,10 @@ from pydantic import root_validator
 from pydantic import validator
 
 from .._validators import valstr
-from fractal_server.app.schemas._validators import valutc
 from fractal_server.app.schemas.v2 import ManifestV2
 from fractal_server.string_tools import validate_cmd
 
 
-class CollectionStatusV2(str, Enum):
-    PENDING = "pending"
-    INSTALLING = "installing"
-    COLLECTING = "collecting"
-    FAIL = "fail"
-    OK = "OK"
-
-
 class TaskCollectPipV2(BaseModel, extra=Extra.forbid):
     """
     TaskCollectPipV2 class
@@ -57,16 +45,8 @@ class TaskCollectPipV2(BaseModel, extra=Extra.forbid):
     python_version: Optional[Literal["3.9", "3.10", "3.11", "3.12"]] = None
     pinned_package_versions: Optional[dict[str, str]] = None
 
-    _package = validator("package", allow_reuse=True)(valstr("package"))
-    _package_version = validator("package_version", allow_reuse=True)(
-        valstr("package_version")
-    )
-    _package_extras = validator("package_extras", allow_reuse=True)(
-        valstr("package_extras")
-    )
-
     @validator("pinned_package_versions")
-    def pinned_package_validator(cls, value):
+    def pinned_package_versions_validator(cls, value):
         if value is None:
             return value
         old_keys = list(value.keys())
@@ -87,6 +67,7 @@ class TaskCollectPipV2(BaseModel, extra=Extra.forbid):
 
     @validator("package")
     def package_validator(cls, value):
+        value = valstr("package")(value)
         if "/" in value or value.endswith(".whl"):
             if not value.endswith(".whl"):
                 raise ValueError(
@@ -97,17 +78,27 @@ class TaskCollectPipV2(BaseModel, extra=Extra.forbid):
                 raise ValueError(
                     f"Local-package path must be absolute: (given {value})."
                 )
+        validate_cmd(value, attribute_name="package")
         return value
 
     @validator("package_version")
-    def package_version_validator(cls, v, values):
+    def package_version_validator(
+        cls, v: Optional[str], values
+    ) -> Optional[str]:
         v = valstr("package_version")(v)
         if values["package"].endswith(".whl"):
             raise ValueError(
                 "Cannot provide package version when package is a wheel file."
             )
+        validate_cmd(v, attribute_name="package_version")
         return v
 
+    @validator("package_extras")
+    def package_extras_validator(cls, value: Optional[str]) -> Optional[str]:
+        value = valstr("package_extras")(value)
+        validate_cmd(value, attribute_name="package_extras")
+        return value
+
 
 class TaskCollectCustomV2(BaseModel, extra=Extra.forbid):
     """
@@ -188,12 +179,3 @@ class TaskCollectCustomV2(BaseModel, extra=Extra.forbid):
                 f"Python interpreter path must be absolute: (given {value})."
             )
         return value
-
-
-class CollectionStateReadV2(BaseModel):
-
-    id: Optional[int]
-    data: dict[str, Any]
-    timestamp: datetime
-
-    _timestamp = validator("timestamp", allow_reuse=True)(valutc("timestamp"))

fractal_server/app/schemas/v2/task_group.py

@@ -20,6 +20,19 @@ class TaskGroupV2OriginEnum(str, Enum):
     OTHER = "other"
 
 
+class TaskGroupActivityStatusV2(str, Enum):
+    PENDING = "pending"
+    ONGOING = "ongoing"
+    FAILED = "failed"
+    OK = "OK"
+
+
+class TaskGroupActivityActionV2(str, Enum):
+    COLLECT = "collect"
+    DEACTIVATE = "deactivate"
+    REACTIVATE = "reactivate"
+
+
 class TaskGroupCreateV2(BaseModel, extra=Extra.forbid):
     user_id: int
     user_group_id: Optional[int] = None
@@ -32,6 +45,7 @@ class TaskGroupCreateV2(BaseModel, extra=Extra.forbid):
     venv_path: Optional[str] = None
     wheel_path: Optional[str] = None
     pip_extras: Optional[str] = None
+    pip_freeze: Optional[str] = None
     pinned_package_versions: dict[str, str] = Field(default_factory=dict)
 
     # Validators
@@ -67,19 +81,29 @@ class TaskGroupReadV2(BaseModel):
     path: Optional[str] = None
     venv_path: Optional[str] = None
     wheel_path: Optional[str] = None
+    pip_freeze: Optional[str] = None
     pip_extras: Optional[str] = None
     pinned_package_versions: dict[str, str] = Field(default_factory=dict)
 
+    venv_size_in_kB: Optional[int] = None
+    venv_file_number: Optional[int] = None
+
     active: bool
     timestamp_created: datetime
 
 
 class TaskGroupUpdateV2(BaseModel, extra=Extra.forbid):
     user_group_id: Optional[int] = None
-    active: Optional[bool] = None
 
-    @validator("active")
-    def active_cannot_be_None(cls, value):
-        if value is None:
-            raise ValueError("`active` cannot be set to None")
-        return value
+
+class TaskGroupActivityV2Read(BaseModel):
+    id: int
+    user_id: int
+    taskgroupv2_id: Optional[int] = None
+    timestamp_started: datetime
+    timestamp_ended: Optional[datetime] = None
+    pkg_name: str
+    version: str
+    status: TaskGroupActivityStatusV2
+    action: TaskGroupActivityActionV2
+    log: Optional[str] = None

fractal_server/migrations/versions/3082479ac4ea_taskgroup_activity_and_venv_info_to_.py

@@ -0,0 +1,105 @@
+"""TaskGroup Activity and venv-info to TaskGroup
+
+Revision ID: 3082479ac4ea
+Revises: 19eca0dd47a9
+Create Date: 2024-11-12 14:39:34.035859
+
+"""
+import sqlalchemy as sa
+import sqlmodel
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "3082479ac4ea"
+down_revision = "19eca0dd47a9"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table(
+        "taskgroupactivityv2",
+        sa.Column("id", sa.Integer(), nullable=False),
+        sa.Column("user_id", sa.Integer(), nullable=False),
+        sa.Column("taskgroupv2_id", sa.Integer(), nullable=True),
+        sa.Column(
+            "timestamp_started", sa.DateTime(timezone=True), nullable=False
+        ),
+        sa.Column(
+            "pkg_name", sqlmodel.sql.sqltypes.AutoString(), nullable=False
+        ),
+        sa.Column(
+            "version", sqlmodel.sql.sqltypes.AutoString(), nullable=False
+        ),
+        sa.Column(
+            "status", sqlmodel.sql.sqltypes.AutoString(), nullable=False
+        ),
+        sa.Column(
+            "action", sqlmodel.sql.sqltypes.AutoString(), nullable=False
+        ),
+        sa.Column("log", sqlmodel.sql.sqltypes.AutoString(), nullable=True),
+        sa.Column(
+            "timestamp_ended", sa.DateTime(timezone=True), nullable=True
+        ),
+        sa.ForeignKeyConstraint(
+            ["taskgroupv2_id"],
+            ["taskgroupv2.id"],
+            name=op.f("fk_taskgroupactivityv2_taskgroupv2_id_taskgroupv2"),
+        ),
+        sa.ForeignKeyConstraint(
+            ["user_id"],
+            ["user_oauth.id"],
+            name=op.f("fk_taskgroupactivityv2_user_id_user_oauth"),
+        ),
+        sa.PrimaryKeyConstraint("id", name=op.f("pk_taskgroupactivityv2")),
+    )
+    op.drop_table("collectionstatev2")
+    with op.batch_alter_table("taskgroupv2", schema=None) as batch_op:
+        batch_op.add_column(
+            sa.Column(
+                "pip_freeze", sqlmodel.sql.sqltypes.AutoString(), nullable=True
+            )
+        )
+        batch_op.add_column(
+            sa.Column("venv_size_in_kB", sa.Integer(), nullable=True)
+        )
+        batch_op.add_column(
+            sa.Column("venv_file_number", sa.Integer(), nullable=True)
+        )
+
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("taskgroupv2", schema=None) as batch_op:
+        batch_op.drop_column("venv_file_number")
+        batch_op.drop_column("venv_size_in_kB")
+        batch_op.drop_column("pip_freeze")
+
+    op.create_table(
+        "collectionstatev2",
+        sa.Column("id", sa.INTEGER(), autoincrement=True, nullable=False),
+        sa.Column(
+            "data",
+            sa.JSON(),
+            nullable=True,
+        ),
+        sa.Column(
+            "timestamp",
+            sa.DateTime(timezone=True),
+            nullable=True,
+        ),
+        sa.Column(
+            "taskgroupv2_id", sa.INTEGER(), autoincrement=False, nullable=True
+        ),
+        sa.ForeignKeyConstraint(
+            ["taskgroupv2_id"],
+            ["taskgroupv2.id"],
+            name="fk_collectionstatev2_taskgroupv2_id_taskgroupv2",
+        ),
+        sa.PrimaryKeyConstraint("id", name="pk_collectionstatev2"),
+    )
+    op.drop_table("taskgroupactivityv2")
+    # ### end Alembic commands ###

fractal_server/ssh/_fabric.py

@@ -471,6 +471,24 @@ class FractalSSH(object):
                 f.write(content)
         self.logger.info(f"END writing to remote file {path}.")
 
+    def remote_exists(self, path: str) -> bool:
+        """
+        Return whether a remote file/folder exists
+        """
+        self.logger.info(f"START remote_file_exists {path}")
+        with _acquire_lock_with_timeout(
+            lock=self._lock,
+            label=f"remote_file_exists {path=}",
+            timeout=self.default_lock_timeout,
+        ):
+            try:
+                self._sftp_unsafe().stat(path)
+                self.logger.info(f"END   remote_file_exists {path} / True")
+                return True
+            except FileNotFoundError:
+                self.logger.info(f"END   remote_file_exists {path} / False")
+                return False
+
 
 class FractalSSHList(object):
     """

fractal_server/string_tools.py

@@ -54,7 +54,12 @@ def slugify_task_name_for_source_v1(task_name: str) -> str:
     return task_name.replace(" ", "_").lower()
 
 
-def validate_cmd(command: str, allow_char: Optional[str] = None):
+def validate_cmd(
+    command: str,
+    *,
+    allow_char: Optional[str] = None,
+    attribute_name: str = "Command",
+):
     """
     Assert that the provided `command` does not contain any of the forbidden
     characters for commands
@@ -63,6 +68,7 @@ def validate_cmd(command: str, allow_char: Optional[str] = None):
     Args:
         command: command to validate.
         allow_char: chars to accept among the forbidden ones
+        attribute_name: Name of the attribute, to be used in error message.
     """
     if not isinstance(command, str):
         raise ValueError(f"{command=} is not a string.")
@@ -71,6 +77,7 @@ def validate_cmd(command: str, allow_char: Optional[str] = None):
         forbidden = forbidden - set(allow_char)
     if not forbidden.isdisjoint(set(command)):
         raise ValueError(
-            f"Command must not contain any of this characters: '{forbidden}'\n"
-            f"Provided command: '{command}'."
+            f"{attribute_name} must not contain any of this characters: "
+            f"'{forbidden}'\n"
+            f"Provided {attribute_name.lower()}: '{command}'."
         )

fractal_server/tasks/utils.py

@@ -30,19 +30,9 @@ def get_log_path(base: Path) -> Path:
     return base / COLLECTION_LOG_FILENAME
 
 
-def get_freeze_path(base: Path) -> Path:
-    return base / COLLECTION_FREEZE_FILENAME
-
-
 def get_collection_log_v1(path: Path) -> str:
     package_path = get_absolute_venv_path_v1(path)
     log_path = get_log_path(package_path)
-    log = log_path.open().read()
+    with log_path.open("r") as f:
+        log = f.read()
     return log
-
-
-def get_collection_freeze_v1(venv_path: Path) -> str:
-    package_path = get_absolute_venv_path_v1(venv_path)
-    freeze_path = get_freeze_path(package_path)
-    freeze = freeze_path.open().read()
-    return freeze

fractal_server/tasks/v2/local/__init__.py

@@ -0,0 +1,3 @@
+from .collect import collect_package_local  # noqa
+from .deactivate import deactivate_local  # noqa
+from .reactivate import reactivate_local  # noqa

fractal_server/tasks/v2/local/collect.py

@@ -0,0 +1,291 @@
+import json
+import logging
+import shutil
+import time
+from pathlib import Path
+from tempfile import TemporaryDirectory
+
+from ..utils_database import create_db_tasks_and_update_task_group
+from .utils_local 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 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.tasks.utils import get_log_path
+from fractal_server.tasks.v2.utils_background import _prepare_tasks_metadata
+from fractal_server.tasks.v2.utils_background import add_commit_refresh
+from fractal_server.tasks.v2.utils_background import check_task_files_exist
+from fractal_server.tasks.v2.utils_background import fail_and_cleanup
+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 get_collection_replacements
+from fractal_server.tasks.v2.utils_templates import (
+    parse_script_pip_show_stdout,
+)
+from fractal_server.tasks.v2.utils_templates import SCRIPTS_SUBFOLDER
+from fractal_server.utils import get_timestamp
+
+LOGGER_NAME = __name__
+
+
+def _copy_wheel_file_local(task_group: TaskGroupV2) -> str:
+    logger = get_logger(LOGGER_NAME)
+    source = task_group.wheel_path
+    dest = (
+        Path(task_group.path) / Path(task_group.wheel_path).name
+    ).as_posix()
+    logger.debug(f"[_copy_wheel_file] START {source=} {dest=}")
+    shutil.copy(task_group.wheel_path, task_group.path)
+    logger.debug(f"[_copy_wheel_file] END {source=} {dest=}")
+    return dest
+
+
+def collect_package_local(
+    *,
+    task_group_activity_id: int,
+    task_group_id: int,
+) -> None:
+    """
+    Collect a task package.
+
+    This function is run as a background task, therefore exceptions must be
+    handled.
+
+    NOTE: by making this function sync, it runs within a thread - due to
+    starlette/fastapi handling of background tasks (see
+    https://github.com/encode/starlette/blob/master/starlette/background.py).
+
+
+    Arguments:
+        task_group_id:
+        task_group_activity_id:
+    """
+
+    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 the (local) task_group path does exist
+            if Path(task_group.path).exists():
+                error_msg = f"{task_group.path} already exists."
+                logger.error(error_msg)
+                fail_and_cleanup(
+                    task_group=task_group,
+                    task_group_activity=activity,
+                    logger_name=LOGGER_NAME,
+                    log_file_path=log_file_path,
+                    exception=FileExistsError(error_msg),
+                    db=db,
+                )
+                return
+
+            try:
+
+                # Create task_group.path folder
+                Path(task_group.path).mkdir(parents=True)
+                logger.debug(f"Created {task_group.path}")
+
+                # Copy wheel file into task group path
+                if task_group.wheel_path:
+                    new_wheel_path = _copy_wheel_file_local(
+                        task_group=task_group
+                    )
+                    task_group.wheel_path = new_wheel_path
+                    task_group = add_commit_refresh(obj=task_group, 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 common arguments for `_customize_and_run_template``
+                common_args = dict(
+                    replacements=replacements,
+                    script_dir=(
+                        Path(task_group.path) / SCRIPTS_SUBFOLDER
+                    ).as_posix(),
+                    prefix=(
+                        f"{int(time.time())}_"
+                        f"{TaskGroupActivityActionV2.COLLECT}_"
+                    ),
+                    logger_name=LOGGER_NAME,
+                )
+
+                # Set status to ONGOING and refresh logs
+                activity.status = TaskGroupActivityStatusV2.ONGOING
+                activity.log = get_current_log(log_file_path)
+                activity = add_commit_refresh(obj=activity, db=db)
+
+                # Run script 1
+                stdout = _customize_and_run_template(
+                    template_filename="1_create_venv.sh",
+                    **common_args,
+                )
+                activity.log = get_current_log(log_file_path)
+                activity = add_commit_refresh(obj=activity, db=db)
+
+                # Run script 2
+                stdout = _customize_and_run_template(
+                    template_filename="2_pip_install.sh",
+                    **common_args,
+                )
+                activity.log = get_current_log(log_file_path)
+                activity = add_commit_refresh(obj=activity, db=db)
+
+                # Run script 3
+                pip_freeze_stdout = _customize_and_run_template(
+                    template_filename="3_pip_freeze.sh",
+                    **common_args,
+                )
+                activity.log = get_current_log(log_file_path)
+                activity = add_commit_refresh(obj=activity, db=db)
+
+                # Run script 4
+                stdout = _customize_and_run_template(
+                    template_filename="4_pip_show.sh",
+                    **common_args,
+                )
+                activity.log = get_current_log(log_file_path)
+                activity = add_commit_refresh(obj=activity, db=db)
+
+                # Run script 5
+                venv_info = _customize_and_run_template(
+                    template_filename="5_get_venv_size_and_file_number.sh",
+                    **common_args,
+                )
+                venv_size, venv_file_number = venv_info.split()
+                activity.log = get_current_log(log_file_path)
+                activity = add_commit_refresh(obj=activity, db=db)
+
+                pkg_attrs = parse_script_pip_show_stdout(stdout)
+                for key, value in pkg_attrs.items():
+                    logger.debug(f"Parsed from pip-show: {key}={value}")
+                # Check package_name match between pip show and task-group
+                task_group = db.get(TaskGroupV2, task_group_id)
+                package_name_pip_show = pkg_attrs.get("package_name")
+                package_name_task_group = task_group.pkg_name
+                compare_package_names(
+                    pkg_name_pip_show=package_name_pip_show,
+                    pkg_name_task_group=package_name_task_group,
+                    logger_name=LOGGER_NAME,
+                )
+                # Extract/drop parsed attributes
+                package_name = package_name_task_group
+                python_bin = pkg_attrs.pop("python_bin")
+                package_root_parent = pkg_attrs.pop("package_root_parent")
+
+                # TODO : Use more robust logic to determine `package_root`.
+                # Examples: use `importlib.util.find_spec`, or parse the
+                # output of `pip show --files {package_name}`.
+                package_name_underscore = package_name.replace("-", "_")
+                package_root = (
+                    Path(package_root_parent) / package_name_underscore
+                ).as_posix()
+
+                # Read and validate manifest file
+                manifest_path = pkg_attrs.pop("manifest_path")
+                logger.info(f"now loading {manifest_path=}")
+                with open(manifest_path) as json_data:
+                    pkg_manifest_dict = json.load(json_data)
+                logger.info(f"loaded {manifest_path=}")
+                logger.info("now validating manifest content")
+                pkg_manifest = ManifestV2(**pkg_manifest_dict)
+                logger.info("validated manifest content")
+                activity.log = get_current_log(log_file_path)
+                activity = add_commit_refresh(obj=activity, db=db)
+
+                logger.info("_prepare_tasks_metadata - start")
+                task_list = _prepare_tasks_metadata(
+                    package_manifest=pkg_manifest,
+                    package_version=task_group.version,
+                    package_root=Path(package_root),
+                    python_bin=Path(python_bin),
+                )
+                check_task_files_exist(task_list=task_list)
+                logger.info("_prepare_tasks_metadata - end")
+                activity.log = get_current_log(log_file_path)
+                activity = add_commit_refresh(obj=activity, db=db)
+
+                logger.info("create_db_tasks_and_update_task_group - " "start")
+                create_db_tasks_and_update_task_group(
+                    task_list=task_list,
+                    task_group_id=task_group.id,
+                    db=db,
+                )
+                logger.info("create_db_tasks_and_update_task_group - end")
+
+                # Update task_group data
+                logger.info(
+                    "Add pip_freeze, venv_size and venv_file_number "
+                    "to TaskGroupV2 - start"
+                )
+                task_group.pip_freeze = pip_freeze_stdout
+                task_group.venv_size_in_kB = int(venv_size)
+                task_group.venv_file_number = int(venv_file_number)
+                task_group = add_commit_refresh(obj=task_group, db=db)
+                logger.info(
+                    "Add pip_freeze, venv_size and venv_file_number "
+                    "to TaskGroupV2 - end"
+                )
+
+                # Finalize (write metadata to DB)
+                logger.debug("finalising - START")
+                activity.status = TaskGroupActivityStatusV2.OK
+                activity.timestamp_ended = get_timestamp()
+                activity = add_commit_refresh(obj=activity, db=db)
+                logger.debug("finalising - END")
+                logger.debug("END")
+
+            except Exception as collection_e:
+                # Delete corrupted package dir
+                try:
+                    logger.info(f"Now delete folder {task_group.path}")
+                    shutil.rmtree(task_group.path)
+                    logger.info(f"Deleted folder {task_group.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=collection_e,
+                    db=db,
+                )
+        return