fractal-server 2.14.15__py3-none-any.whl → 2.15.0__py3-none-any.whl

fractal_server/app/routes/api/v2/task_collection_pixi.py

@@ -0,0 +1,219 @@
+import os
+from pathlib import Path
+
+from fastapi import APIRouter
+from fastapi import BackgroundTasks
+from fastapi import Depends
+from fastapi import Form
+from fastapi import HTTPException
+from fastapi import Response
+from fastapi import status
+from fastapi import UploadFile
+
+from fractal_server.app.db import AsyncSession
+from fractal_server.app.db import get_async_db
+from fractal_server.app.models import UserOAuth
+from fractal_server.app.models.v2 import TaskGroupActivityV2
+from fractal_server.app.models.v2 import TaskGroupV2
+from fractal_server.app.routes.api.v2._aux_functions_tasks import (
+    _get_valid_user_group_id,
+)
+from fractal_server.app.routes.api.v2._aux_functions_tasks import (
+    _verify_non_duplication_group_constraint,
+)
+from fractal_server.app.routes.api.v2._aux_functions_tasks import (
+    _verify_non_duplication_group_path,
+)
+from fractal_server.app.routes.api.v2._aux_functions_tasks import (
+    _verify_non_duplication_user_constraint,
+)
+from fractal_server.app.routes.auth import current_active_verified_user
+from fractal_server.app.routes.aux.validate_user_settings import (
+    validate_user_settings,
+)
+from fractal_server.app.schemas.v2 import FractalUploadedFile
+from fractal_server.app.schemas.v2 import TaskGroupActivityActionV2
+from fractal_server.app.schemas.v2 import TaskGroupActivityStatusV2
+from fractal_server.app.schemas.v2 import TaskGroupActivityV2Read
+from fractal_server.app.schemas.v2.task_group import TaskGroupV2OriginEnum
+from fractal_server.config import get_settings
+from fractal_server.logger import set_logger
+from fractal_server.ssh._fabric import SSHConfig
+from fractal_server.syringe import Inject
+from fractal_server.tasks.v2.local import collect_local_pixi
+from fractal_server.tasks.v2.ssh import collect_ssh_pixi
+from fractal_server.tasks.v2.utils_package_names import normalize_package_name
+from fractal_server.types import NonEmptyStr
+
+
+router = APIRouter()
+
+logger = set_logger(__name__)
+
+
+def validate_pkgname_and_version(filename: str) -> tuple[str, str]:
+    if not filename.endswith(".tar.gz"):
+        raise HTTPException(
+            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail=f"{filename=} does not end with '.tar.gz'.",
+        )
+    filename_splitted = filename.split("-")
+    if len(filename_splitted) != 2:
+        raise HTTPException(
+            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail=(
+                f"Invalid filename: '{filename}' must contain a single `-` "
+                "character, separating the package name from the version "
+                "(expected format: 'pkg_name-version')."
+            ),
+        )
+
+    pkg_name = filename_splitted[0]
+    version = filename.removeprefix(f"{pkg_name}-").removesuffix(".tar.gz")
+
+    return normalize_package_name(pkg_name), version
+
+
+@router.post(
+    "/collect/pixi/",
+    status_code=202,
+    response_model=TaskGroupActivityV2Read,
+)
+async def collect_task_pixi(
+    response: Response,
+    background_tasks: BackgroundTasks,
+    file: UploadFile,
+    pixi_version: NonEmptyStr | None = Form(None),
+    private: bool = False,
+    user_group_id: int | None = None,
+    user: UserOAuth = Depends(current_active_verified_user),
+    db: AsyncSession = Depends(get_async_db),
+) -> TaskGroupActivityV2Read:
+
+    settings = Inject(get_settings)
+    # Check if Pixi is available
+    if settings.pixi is None:
+        raise HTTPException(
+            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail="Pixi task collection is not available.",
+        )
+    # Check if provided Pixi version is available. Use default if not provided
+    if pixi_version is None:
+        pixi_version = settings.pixi.default_version
+    else:
+        if pixi_version not in settings.pixi.versions:
+            raise HTTPException(
+                status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+                detail=(
+                    f"Pixi version {pixi_version} is not available. Available "
+                    f"versions: {list(settings.pixi.versions.keys())}"
+                ),
+            )
+
+    pkg_name, version = validate_pkgname_and_version(file.filename)
+    tar_gz_content = await file.read()
+    tar_gz_file = FractalUploadedFile(
+        filename=file.filename,
+        contents=tar_gz_content,
+    )
+
+    user_group_id = await _get_valid_user_group_id(
+        user_group_id=user_group_id,
+        private=private,
+        user_id=user.id,
+        db=db,
+    )
+
+    user_settings = await validate_user_settings(
+        user=user, backend=settings.FRACTAL_RUNNER_BACKEND, db=db
+    )
+
+    if settings.FRACTAL_RUNNER_BACKEND == "slurm_ssh":
+        base_tasks_path = user_settings.ssh_tasks_dir
+    else:
+        base_tasks_path = settings.FRACTAL_TASKS_DIR.as_posix()
+    task_group_path = (
+        Path(base_tasks_path) / str(user.id) / pkg_name / version
+    ).as_posix()
+
+    task_group_attrs = dict(
+        user_id=user.id,
+        user_group_id=user_group_id,
+        origin=TaskGroupV2OriginEnum.PIXI,
+        pixi_version=pixi_version,
+        pkg_name=pkg_name,
+        version=version,
+        path=task_group_path,
+    )
+
+    await _verify_non_duplication_user_constraint(
+        user_id=user.id,
+        pkg_name=task_group_attrs["pkg_name"],
+        version=task_group_attrs["version"],
+        db=db,
+    )
+    await _verify_non_duplication_group_constraint(
+        user_group_id=task_group_attrs["user_group_id"],
+        pkg_name=task_group_attrs["pkg_name"],
+        version=task_group_attrs["version"],
+        db=db,
+    )
+    await _verify_non_duplication_group_path(
+        path=task_group_attrs["path"],
+        db=db,
+    )
+
+    if settings.FRACTAL_RUNNER_BACKEND != "slurm_ssh":
+        if Path(task_group_path).exists():
+            raise HTTPException(
+                status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+                detail=f"{task_group_path} already exists.",
+            )
+
+    task_group = TaskGroupV2(**task_group_attrs)
+    db.add(task_group)
+    await db.commit()
+    await db.refresh(task_group)
+    db.expunge(task_group)
+
+    task_group_activity = TaskGroupActivityV2(
+        user_id=task_group.user_id,
+        taskgroupv2_id=task_group.id,
+        status=TaskGroupActivityStatusV2.PENDING,
+        action=TaskGroupActivityActionV2.COLLECT,
+        pkg_name=task_group.pkg_name,
+        version=task_group.version,
+    )
+    db.add(task_group_activity)
+    await db.commit()
+    await db.refresh(task_group_activity)
+
+    if settings.FRACTAL_RUNNER_BACKEND == "slurm_ssh":
+        ssh_config = SSHConfig(
+            user=user_settings.ssh_username,
+            host=user_settings.ssh_host,
+            key_path=user_settings.ssh_private_key_path,
+        )
+
+        background_tasks.add_task(
+            collect_ssh_pixi,
+            task_group_id=task_group.id,
+            task_group_activity_id=task_group_activity.id,
+            ssh_config=ssh_config,
+            tasks_base_dir=user_settings.ssh_tasks_dir,
+            tar_gz_file=tar_gz_file,
+        )
+    else:
+        background_tasks.add_task(
+            collect_local_pixi,
+            task_group_id=task_group.id,
+            task_group_activity_id=task_group_activity.id,
+            tar_gz_file=tar_gz_file,
+        )
+    logger.info(
+        "Task-collection endpoint: start background collection "
+        "and return task_group_activity. "
+        f"Current pid is {os.getpid()}. "
+    )
+    response.status_code = status.HTTP_202_ACCEPTED
+    return task_group_activity

fractal_server/app/routes/api/v2/task_group.py

@@ -12,6 +12,7 @@ from pydantic.types import AwareDatetime
 from sqlmodel import or_
 from sqlmodel import select
 
+from ._aux_functions_task_lifecycle import check_no_ongoing_activity
 from ._aux_functions_tasks import _get_task_group_full_access
 from ._aux_functions_tasks import _get_task_group_read_access
 from ._aux_functions_tasks import _verify_non_duplication_group_constraint
@@ -216,6 +217,8 @@ async def delete_task_group(
         db=db,
     )
 
+    await check_no_ongoing_activity(task_group_id=task_group_id, db=db)
+
     stm = select(WorkflowTaskV2).where(
         WorkflowTaskV2.task_id.in_({task.id for task in task_group.task_list})
     )

fractal_server/app/routes/api/v2/task_group_lifecycle.py

@@ -2,7 +2,6 @@ from fastapi import APIRouter
 from fastapi import BackgroundTasks
 from fastapi import Depends
 from fastapi import HTTPException
-from fastapi import Request
 from fastapi import Response
 from fastapi import status
 
@@ -25,9 +24,13 @@ from fractal_server.logger import set_logger
 from fractal_server.ssh._fabric import SSHConfig
 from fractal_server.syringe import Inject
 from fractal_server.tasks.v2.local import deactivate_local
+from fractal_server.tasks.v2.local import deactivate_local_pixi
 from fractal_server.tasks.v2.local import reactivate_local
+from fractal_server.tasks.v2.local import reactivate_local_pixi
 from fractal_server.tasks.v2.ssh import deactivate_ssh
+from fractal_server.tasks.v2.ssh import deactivate_ssh_pixi
 from fractal_server.tasks.v2.ssh import reactivate_ssh
+from fractal_server.tasks.v2.ssh import reactivate_ssh_pixi
 from fractal_server.utils import get_timestamp
 
 router = APIRouter()
@@ -44,7 +47,6 @@ async def deactivate_task_group(
     task_group_id: int,
     background_tasks: BackgroundTasks,
     response: Response,
-    request: Request,
     user: UserOAuth = Depends(current_active_user),
     db: AsyncSession = Depends(get_async_db),
 ) -> TaskGroupReadV2:
@@ -125,9 +127,12 @@ async def deactivate_task_group(
             host=user_settings.ssh_host,
             key_path=user_settings.ssh_private_key_path,
         )
-
+        if task_group.origin == TaskGroupV2OriginEnum.PIXI:
+            deactivate_function = deactivate_ssh_pixi
+        else:
+            deactivate_function = deactivate_ssh
         background_tasks.add_task(
-            deactivate_ssh,
+            deactivate_function,
             task_group_id=task_group.id,
             task_group_activity_id=task_group_activity.id,
             ssh_config=ssh_config,
@@ -135,8 +140,12 @@ async def deactivate_task_group(
         )
 
     else:
+        if task_group.origin == TaskGroupV2OriginEnum.PIXI:
+            deactivate_function = deactivate_local_pixi
+        else:
+            deactivate_function = deactivate_local
         background_tasks.add_task(
-            deactivate_local,
+            deactivate_function,
             task_group_id=task_group.id,
             task_group_activity_id=task_group_activity.id,
         )
@@ -157,7 +166,6 @@ async def reactivate_task_group(
     task_group_id: int,
     background_tasks: BackgroundTasks,
     response: Response,
-    request: Request,
     user: UserOAuth = Depends(current_active_user),
     db: AsyncSession = Depends(get_async_db),
 ) -> TaskGroupReadV2:
@@ -210,12 +218,12 @@ async def reactivate_task_group(
         response.status_code = status.HTTP_202_ACCEPTED
         return task_group_activity
 
-    if task_group.pip_freeze is None:
+    if task_group.env_info is None:
         raise HTTPException(
             status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
             detail=(
                 "Cannot reactivate a task group with "
-                f"{task_group.pip_freeze=}."
+                f"{task_group.env_info=}."
             ),
         )
 
@@ -247,8 +255,12 @@ async def reactivate_task_group(
             key_path=user_settings.ssh_private_key_path,
         )
 
+        if task_group.origin == TaskGroupV2OriginEnum.PIXI:
+            reactivate_function = reactivate_ssh_pixi
+        else:
+            reactivate_function = reactivate_ssh
         background_tasks.add_task(
-            reactivate_ssh,
+            reactivate_function,
             task_group_id=task_group.id,
             task_group_activity_id=task_group_activity.id,
             ssh_config=ssh_config,
@@ -256,8 +268,12 @@ async def reactivate_task_group(
         )
 
     else:
+        if task_group.origin == TaskGroupV2OriginEnum.PIXI:
+            reactivate_function = reactivate_local_pixi
+        else:
+            reactivate_function = reactivate_local
         background_tasks.add_task(
-            reactivate_local,
+            reactivate_function,
             task_group_id=task_group.id,
             task_group_activity_id=task_group_activity.id,
         )

fractal_server/app/runner/executors/slurm_common/_slurm_config.py

@@ -48,6 +48,8 @@ class _SlurmConfigSet(BaseModel):
         constraint:
         gres:
         time:
+        exclude:
+        nodelist:
         account:
         extra_lines:
     """
@@ -59,6 +61,8 @@ class _SlurmConfigSet(BaseModel):
     mem: int | str | None = None
     constraint: str | None = None
     gres: str | None = None
+    exclude: str | None = None
+    nodelist: str | None = None
     time: str | None = None
     account: str | None = None
     extra_lines: list[str] | None = None
@@ -227,6 +231,8 @@ class SlurmConfig(BaseModel):
         account: Corresponds to SLURM option.
         gpus: Corresponds to SLURM option.
         time: Corresponds to SLURM option (WARNING: not fully supported).
+        nodelist: Corresponds to SLURM option.
+        exclude: Corresponds to SLURM option.
         prefix: Prefix of configuration lines in SLURM submission scripts.
         shebang_line: Shebang line for SLURM submission scripts.
         extra_lines: Additional lines to include in SLURM submission scripts.
@@ -268,6 +274,8 @@ class SlurmConfig(BaseModel):
     gpus: str | None = None
     time: str | None = None
     account: str | None = None
+    nodelist: str | None = None
+    exclude: str | None = None
 
     # Free-field attribute for extra lines to be added to the SLURM job
     # preamble
@@ -361,6 +369,8 @@ class SlurmConfig(BaseModel):
             "gpus",
             "time",
             "account",
+            "exclude",
+            "nodelist",
         ]:
             value = getattr(self, key)
             if value is not None:

fractal_server/app/runner/executors/slurm_common/base_slurm_runner.py

@@ -182,33 +182,53 @@ class BaseSlurmRunner(BaseRunner):
     def _mkdir_remote_folder(self, folder: str) -> None:
         raise NotImplementedError("Implement in child class.")
 
-    def _submit_single_sbatch(
+    def _enrich_slurm_config(
         self,
-        *,
-        base_command: str,
-        slurm_job: SlurmJob,
         slurm_config: SlurmConfig,
-    ) -> str:
-        logger.debug("[_submit_single_sbatch] START")
+    ) -> SlurmConfig:
+        """
+        Return an enriched `SlurmConfig` object
 
-        # Include SLURM account in `slurm_config`. Note: we make this change
-        # here, rather than exposing a new argument of `get_slurm_config`,
-        # because it's a backend-specific argument while `get_slurm_config` has
-        # a generic interface.
+        Include `self.account` and `self.common_script_lines` into a
+        `SlurmConfig` object. Extracting this logic into an independent
+        class method is useful to fix issue #2659 (which was due to
+        performing this same operation multiple times rather than once).
+
+        Args:
+            slurm_config: The original `SlurmConfig` object.
+
+        Returns:
+            A new, up-to-date, `SlurmConfig` object.
+        """
+
+        new_slurm_config = slurm_config.model_copy()
+
+        # Include SLURM account in `slurm_config`.
         if self.slurm_account is not None:
-            slurm_config.account = self.slurm_account
+            new_slurm_config.account = self.slurm_account
 
         # Include common_script_lines in extra_lines
         if len(self.common_script_lines) > 0:
             logger.debug(
                 f"Add {self.common_script_lines} to "
-                f"{slurm_config.extra_lines=}."
+                f"{new_slurm_config.extra_lines=}."
             )
-            current_extra_lines = slurm_config.extra_lines or []
-            slurm_config.extra_lines = (
+            current_extra_lines = new_slurm_config.extra_lines or []
+            new_slurm_config.extra_lines = (
                 current_extra_lines + self.common_script_lines
             )
 
+        return new_slurm_config
+
+    def _submit_single_sbatch(
+        self,
+        *,
+        base_command: str,
+        slurm_job: SlurmJob,
+        slurm_config: SlurmConfig,
+    ) -> str:
+        logger.debug("[_submit_single_sbatch] START")
+
         for task in slurm_job.tasks:
             # Write input file
             if self.slurm_runner_type == "ssh":
@@ -508,6 +528,9 @@ class BaseSlurmRunner(BaseRunner):
         user_id: int,
     ) -> tuple[Any, Exception]:
         logger.debug("[submit] START")
+
+        config = self._enrich_slurm_config(config)
+
         try:
             workdir_local = task_files.wftask_subfolder_local
             workdir_remote = task_files.wftask_subfolder_remote
@@ -649,6 +672,8 @@ class BaseSlurmRunner(BaseRunner):
         input images, while for compound tasks these can differ.
         """
 
+        config = self._enrich_slurm_config(config)
+
         logger.debug(f"[multisubmit] START, {len(list_parameters)=}")
         try:
             if self.is_shutdown():

fractal_server/app/runner/executors/slurm_common/get_slurm_config.py

@@ -125,7 +125,14 @@ def get_slurm_config_internal(
             )
             logger.error(error_msg)
             raise SlurmConfigError(error_msg)
-        for key in ["time", "gres", "gpus", "constraint"]:
+        for key in [
+            "time",
+            "gres",
+            "gpus",
+            "constraint",
+            "nodelist",
+            "exclude",
+        ]:
             value = wftask_meta.get(key, None)
             if value is not None:
                 slurm_dict[key] = value

fractal_server/app/runner/executors/slurm_ssh/runner.py

@@ -40,8 +40,10 @@ class SlurmSSHRunner(BaseSlurmRunner):
         self.fractal_ssh = fractal_ssh
         logger.warning(self.fractal_ssh)
 
-        settings = Inject(get_settings)
+        # Check SSH connection and try to recover from a closed-socket error
+        self.fractal_ssh.check_connection()
 
+        settings = Inject(get_settings)
         super().__init__(
             slurm_runner_type="ssh",
             root_dir_local=root_dir_local,

fractal_server/app/runner/v2/runner.py

@@ -32,7 +32,7 @@ from fractal_server.app.schemas.v2 import TaskDumpV2
 from fractal_server.app.schemas.v2 import TaskGroupDumpV2
 from fractal_server.app.schemas.v2 import TaskType
 from fractal_server.images import SingleImage
-from fractal_server.images.status_tools import enrich_images_sync
+from fractal_server.images.status_tools import enrich_images_unsorted_sync
 from fractal_server.images.status_tools import IMAGE_STATUS_KEY
 from fractal_server.images.tools import filter_image_list
 from fractal_server.images.tools import find_image_by_zarr_url
@@ -147,7 +147,7 @@ def execute_tasks_v2(
 
             if ind_wftask == 0 and ENRICH_IMAGES_WITH_STATUS:
                 # FIXME: Could this be done on `type_filtered_images`?
-                tmp_images = enrich_images_sync(
+                tmp_images = enrich_images_unsorted_sync(
                     images=tmp_images,
                     dataset_id=dataset.id,
                     workflowtask_id=wftask.id,

fractal_server/app/schemas/v2/__init__.py

@@ -33,9 +33,9 @@ from .task import TaskImportV2Legacy  # noqa F401
 from .task import TaskReadV2  # noqa F401
 from .task import TaskType  # noqa F401
 from .task import TaskUpdateV2  # noqa F401
+from .task_collection import FractalUploadedFile  # noqa F401
 from .task_collection import TaskCollectCustomV2  # noqa F401
 from .task_collection import TaskCollectPipV2  # noqa F401
-from .task_collection import WheelFile  # noqa F401
 from .task_group import TaskGroupActivityActionV2  # noqa F401
 from .task_group import TaskGroupActivityStatusV2  # noqa F401
 from .task_group import TaskGroupActivityV2Read  # noqa F401

fractal_server/app/schemas/v2/dumps.py

@@ -86,4 +86,4 @@ class TaskGroupDumpV2(BaseModel):
 
     path: str | None = None
     venv_path: str | None = None
-    wheel_path: str | None = None
+    archive_path: str | None = None

fractal_server/app/schemas/v2/task_collection.py

@@ -12,7 +12,7 @@ from fractal_server.types import DictStrStr
 from fractal_server.types import NonEmptyStr
 
 
-class WheelFile(BaseModel):
+class FractalUploadedFile(BaseModel):
     """
     Model for data sent from the endpoint to the background task.
     """

fractal_server/app/schemas/v2/task_group.py

@@ -16,6 +16,7 @@ from fractal_server.types import NonEmptyStr
 class TaskGroupV2OriginEnum(StrEnum):
     PYPI = "pypi"
     WHEELFILE = "wheel-file"
+    PIXI = "pixi"
     OTHER = "other"
 
 
@@ -41,11 +42,12 @@ class TaskGroupCreateV2(BaseModel):
     pkg_name: str
     version: str | None = None
     python_version: NonEmptyStr = None
+    pixi_version: NonEmptyStr = None
     path: AbsolutePathStr = None
     venv_path: AbsolutePathStr = None
-    wheel_path: AbsolutePathStr = None
+    archive_path: AbsolutePathStr = None
     pip_extras: NonEmptyStr = None
-    pip_freeze: str | None = None
+    env_info: str | None = None
     pinned_package_versions: DictStrStr = Field(default_factory=dict)
 
 
@@ -55,8 +57,8 @@ class TaskGroupCreateV2Strict(TaskGroupCreateV2):
     """
 
     path: AbsolutePathStr
-    venv_path: AbsolutePathStr
     version: NonEmptyStr
+    venv_path: AbsolutePathStr
     python_version: NonEmptyStr
 
 
@@ -71,10 +73,10 @@ class TaskGroupReadV2(BaseModel):
     pkg_name: str
     version: str | None = None
     python_version: str | None = None
+    pixi_version: str | None = None
     path: str | None = None
     venv_path: str | None = None
-    wheel_path: str | None = None
-    pip_freeze: str | None = None
+    archive_path: str | None = None
     pip_extras: str | None = None
     pinned_package_versions: dict[str, str] = Field(default_factory=dict)