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

fractal_server/images/status_tools.py

@@ -1,5 +1,4 @@
 import time
-from copy import deepcopy
 from typing import Any
 
 from sqlalchemy import Select
@@ -11,7 +10,6 @@ from fractal_server.app.models.v2 import HistoryImageCache
 from fractal_server.app.models.v2 import HistoryUnit
 from fractal_server.app.schemas.v2 import HistoryUnitStatusWithUnset
 from fractal_server.logger import set_logger
-from fractal_server.types import ImageAttributeValue
 
 logger = set_logger(__name__)
 
@@ -19,36 +17,84 @@ logger = set_logger(__name__)
 IMAGE_STATUS_KEY = "__wftask_dataset_image_status__"
 
 
-def _enriched_image(*, img: dict[str, Any], status: str) -> dict[str, Any]:
-    img["attributes"][IMAGE_STATUS_KEY] = status
-    return img
+def _enriched_image(
+    *,
+    img: dict[str, Any],
+    status: str,
+) -> dict[str, Any]:
+    return img | {
+        "attributes": (img["attributes"] | {IMAGE_STATUS_KEY: status})
+    }
 
 
 def _prepare_query(
     *,
     dataset_id: int,
     workflowtask_id: int,
-    zarr_urls: list[str],
 ) -> Select:
+    """
+    Note: the query does not include `.order_by`.
+    """
     stm = (
         select(HistoryImageCache.zarr_url, HistoryUnit.status)
         .join(HistoryUnit)
         .where(HistoryImageCache.dataset_id == dataset_id)
         .where(HistoryImageCache.workflowtask_id == workflowtask_id)
         .where(HistoryImageCache.latest_history_unit_id == HistoryUnit.id)
-        .where(HistoryImageCache.zarr_url.in_(zarr_urls))
-        .order_by(HistoryImageCache.zarr_url)
     )
     return stm
 
 
-async def enrich_images_async(
+def _postprocess_image_lists(
+    target_images: list[dict[str, Any]],
+    list_query_url_status: list[tuple[str, str]],
+) -> list[dict[str, Any]]:
+    """ """
+    t_1 = time.perf_counter()
+
+    # Select only processed images that are part of the target image set
+    zarr_url_to_image = {img["zarr_url"]: img for img in target_images}
+    target_zarr_urls = zarr_url_to_image.keys()
+    list_processed_url_status = [
+        url_status
+        for url_status in list_query_url_status
+        if url_status[0] in target_zarr_urls
+    ]
+
+    set_processed_urls = set(
+        url_status[0] for url_status in list_processed_url_status
+    )
+    processed_images_with_status = [
+        _enriched_image(
+            img=zarr_url_to_image[item[0]],
+            status=item[1],
+        )
+        for item in list_processed_url_status
+    ]
+
+    non_processed_urls = target_zarr_urls - set_processed_urls
+    non_processed_images_with_status = [
+        _enriched_image(
+            img=zarr_url_to_image[zarr_url],
+            status=HistoryUnitStatusWithUnset.UNSET,
+        )
+        for zarr_url in non_processed_urls
+    ]
+    t_2 = time.perf_counter()
+    logger.debug(
+        f"[enrich_images_async] post-processing, elapsed={t_2 - t_1:.5f} s"
+    )
+
+    return processed_images_with_status + non_processed_images_with_status
+
+
+async def enrich_images_unsorted_async(
     *,
     images: list[dict[str, Any]],
     dataset_id: int,
     workflowtask_id: int,
     db: AsyncSession,
-) -> list[dict[str, ImageAttributeValue]]:
+) -> list[dict[str, Any]]:
     """
     Enrich images with a status-related attribute.
 
@@ -59,116 +105,75 @@ async def enrich_images_async(
         db: An async db session
 
     Returns:
-        The list of enriched images
+        The list of enriched images, not necessarily in the same order as
+        the input.
     """
     t_0 = time.perf_counter()
     logger.info(
         f"[enrich_images_async] START, {dataset_id=}, {workflowtask_id=}"
     )
 
-    zarr_url_to_image = {img["zarr_url"]: deepcopy(img) for img in images}
-
+    # Get `(zarr_url, status)` for _all_ processed images (including those that
+    # are not part of the target image set)
     res = await db.execute(
         _prepare_query(
             dataset_id=dataset_id,
             workflowtask_id=workflowtask_id,
-            zarr_urls=zarr_url_to_image.keys(),
         )
     )
-    list_processed_url_status = res.all()
+    list_query_url_status = res.all()
     t_1 = time.perf_counter()
-    logger.debug(f"[enrich_images_async] db-query, elapsed={t_1 - t_0:.3f} s")
+    logger.debug(f"[enrich_images_async] query, elapsed={t_1 - t_0:.5f} s")
 
-    set_processed_urls = set(item[0] for item in list_processed_url_status)
-    processed_images_with_status = [
-        _enriched_image(
-            img=zarr_url_to_image[item[0]],
-            status=item[1],
-        )
-        for item in list_processed_url_status
-    ]
-    t_2 = time.perf_counter()
-    logger.debug(
-        "[enrich_images_async] processed-images, " f"elapsed={t_2 - t_1:.3f} s"
-    )
-
-    non_processed_urls = zarr_url_to_image.keys() - set_processed_urls
-    non_processed_images_with_status = [
-        _enriched_image(
-            img=zarr_url_to_image[zarr_url],
-            status=HistoryUnitStatusWithUnset.UNSET,
-        )
-        for zarr_url in non_processed_urls
-    ]
-    t_3 = time.perf_counter()
-    logger.debug(
-        "[enrich_images_async] non-processed-images, "
-        f"elapsed={t_3 - t_2:.3f} s"
+    output = _postprocess_image_lists(
+        target_images=images,
+        list_query_url_status=list_query_url_status,
     )
 
-    return processed_images_with_status + non_processed_images_with_status
+    return output
 
 
-def enrich_images_sync(
+def enrich_images_unsorted_sync(
     *,
     images: list[dict[str, Any]],
     dataset_id: int,
     workflowtask_id: int,
-) -> list[dict[str, ImageAttributeValue]]:
+) -> list[dict[str, Any]]:
     """
     Enrich images with a status-related attribute.
 
+
     Args:
         images: The input image list
         dataset_id: The dataset ID
         workflowtask_id: The workflow-task ID
 
     Returns:
-        The list of enriched images
+        The list of enriched images, not necessarily in the same order as
+        the input.
     """
+
     t_0 = time.perf_counter()
     logger.info(
         f"[enrich_images_async] START, {dataset_id=}, {workflowtask_id=}"
     )
 
-    zarr_url_to_image = {img["zarr_url"]: deepcopy(img) for img in images}
+    # Get `(zarr_url, status)` for _all_ processed images (including those that
+    # are not part of the target image set)
     with next(get_sync_db()) as db:
         res = db.execute(
             _prepare_query(
                 dataset_id=dataset_id,
                 workflowtask_id=workflowtask_id,
-                zarr_urls=zarr_url_to_image.keys(),
             )
         )
-    list_processed_url_status = res.all()
+        list_query_url_status = res.all()
     t_1 = time.perf_counter()
-    logger.debug(f"[enrich_images_async] db-query, elapsed={t_1 - t_0:.3f} s")
-
-    set_processed_urls = set(item[0] for item in list_processed_url_status)
-    processed_images_with_status = [
-        _enriched_image(
-            img=zarr_url_to_image[item[0]],
-            status=item[1],
-        )
-        for item in list_processed_url_status
-    ]
-    t_2 = time.perf_counter()
-    logger.debug(
-        "[enrich_images_async] processed-images, " f"elapsed={t_2 - t_1:.3f} s"
-    )
+    logger.debug(f"[enrich_images_async] query, elapsed={t_1 - t_0:.5f} s")
 
-    non_processed_urls = zarr_url_to_image.keys() - set_processed_urls
-    non_processed_images_with_status = [
-        _enriched_image(
-            img=zarr_url_to_image[zarr_url],
-            status=HistoryUnitStatusWithUnset.UNSET,
-        )
-        for zarr_url in non_processed_urls
-    ]
-    t_3 = time.perf_counter()
-    logger.debug(
-        "[enrich_images_async] non-processed-images, "
-        f"elapsed={t_3 - t_2:.3f} s"
+    output = _postprocess_image_lists(
+        target_images=images,
+        list_query_url_status=list_query_url_status,
     )
 
-    return processed_images_with_status + non_processed_images_with_status
+    return output

fractal_server/migrations/versions/791ce783d3d8_add_indices.py

@@ -0,0 +1,41 @@
+"""Add indices
+
+Revision ID: 791ce783d3d8
+Revises: 969d84257cac
+Create Date: 2025-06-03 09:32:30.757651
+
+"""
+from alembic import op
+
+
+# revision identifiers, used by Alembic.
+revision = "791ce783d3d8"
+down_revision = "969d84257cac"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("historyimagecache", schema=None) as batch_op:
+        batch_op.create_index(
+            batch_op.f("ix_historyimagecache_dataset_id"),
+            ["dataset_id"],
+            unique=False,
+        )
+        batch_op.create_index(
+            batch_op.f("ix_historyimagecache_workflowtask_id"),
+            ["workflowtask_id"],
+            unique=False,
+        )
+
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("historyimagecache", schema=None) as batch_op:
+        batch_op.drop_index(batch_op.f("ix_historyimagecache_workflowtask_id"))
+        batch_op.drop_index(batch_op.f("ix_historyimagecache_dataset_id"))
+
+    # ### end Alembic commands ###

fractal_server/migrations/versions/b1e7f7a1ff71_task_group_for_pixi.py

@@ -0,0 +1,53 @@
+"""Task group for pixi
+
+Revision ID: b1e7f7a1ff71
+Revises: 791ce783d3d8
+Create Date: 2025-05-29 16:31:17.565973
+
+"""
+import sqlalchemy as sa
+import sqlmodel
+from alembic import op
+
+
+# revision identifiers, used by Alembic.
+revision = "b1e7f7a1ff71"
+down_revision = "791ce783d3d8"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    with op.batch_alter_table("taskgroupv2", schema=None) as batch_op:
+        batch_op.add_column(
+            sa.Column(
+                "pixi_version",
+                sqlmodel.sql.sqltypes.AutoString(),
+                nullable=True,
+            )
+        )
+        batch_op.alter_column(
+            "wheel_path",
+            nullable=True,
+            new_column_name="archive_path",
+        )
+        batch_op.alter_column(
+            "pip_freeze",
+            nullable=True,
+            new_column_name="env_info",
+        )
+
+
+def downgrade() -> None:
+    with op.batch_alter_table("taskgroupv2", schema=None) as batch_op:
+        batch_op.alter_column(
+            "archive_path",
+            nullable=True,
+            new_column_name="wheel_path",
+        )
+        batch_op.alter_column(
+            "env_info",
+            nullable=True,
+            new_column_name="pip_freeze",
+        )
+        batch_op.drop_column("pixi_version")

fractal_server/ssh/_fabric.py

@@ -642,6 +642,9 @@ class FractalSSHList:
                 connect_kwargs={
                     "key_filename": key_path,
                     "look_for_keys": False,
+                    "banner_timeout": 30,
+                    "auth_timeout": 30,  # default value
+                    "channel_timeout": 60 * 60,  # default value
                 },
             )
             with _acquire_lock_with_timeout(

fractal_server/tasks/v2/local/__init__.py

@@ -1,3 +1,5 @@
 from .collect import collect_local  # noqa
+from .collect_pixi import collect_local_pixi  # noqa
 from .deactivate import deactivate_local  # noqa
+from .deactivate_pixi import deactivate_local_pixi  # noqa
 from .reactivate import reactivate_local  # noqa

fractal_server/tasks/v2/local/_utils.py

@@ -50,19 +50,24 @@ def check_task_files_exist(task_list: list[TaskCreateV2]) -> None:
     """
     Check that the modules listed in task commands point to existing files.
 
+    Note: commands may be like `/one/python /another/task.py` or
+    `/one/pixi [...] /another/task.py`, and in both cases `split()[-1]`
+    returns `/another/task.py`.
+
     Args:
         task_list:
     """
+
     for _task in task_list:
         if _task.command_non_parallel is not None:
-            _task_path = _task.command_non_parallel.split()[1]
+            _task_path = _task.command_non_parallel.split()[-1]
             if not Path(_task_path).exists():
                 raise FileNotFoundError(
                     f"Task `{_task.name}` has `command_non_parallel` "
                     f"pointing to missing file `{_task_path}`."
                 )
         if _task.command_parallel is not None:
-            _task_path = _task.command_parallel.split()[1]
+            _task_path = _task.command_parallel.split()[-1]
             if not Path(_task_path).exists():
                 raise FileNotFoundError(
                     f"Task `{_task.name}` has `command_parallel` "

fractal_server/tasks/v2/local/collect.py

@@ -10,18 +10,18 @@ from ._utils 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 FractalUploadedFile
 from fractal_server.app.schemas.v2 import TaskGroupActivityActionV2
 from fractal_server.app.schemas.v2 import TaskGroupActivityStatusV2
-from fractal_server.app.schemas.v2 import WheelFile
 from fractal_server.app.schemas.v2.manifest import ManifestV2
 from fractal_server.logger import reset_logger_handlers
 from fractal_server.logger import set_logger
 from fractal_server.tasks.utils import get_log_path
 from fractal_server.tasks.v2.local._utils import check_task_files_exist
-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 fail_and_cleanup
 from fractal_server.tasks.v2.utils_background import get_current_log
+from fractal_server.tasks.v2.utils_background import prepare_tasks_metadata
 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,
@@ -38,7 +38,7 @@ def collect_local(
     *,
     task_group_activity_id: int,
     task_group_id: int,
-    wheel_file: WheelFile | None = None,
+    wheel_file: FractalUploadedFile | None = None,
 ) -> None:
     """
     Collect a task package.
@@ -103,16 +103,18 @@ def collect_local(
                 Path(task_group.path).mkdir(parents=True)
                 logger.info(f"Created {task_group.path}")
 
-                # Write wheel file and set task_group.wheel_path
+                # Write wheel file and set task_group.archive_path
                 if wheel_file is not None:
 
-                    wheel_path = (
+                    archive_path = (
                         Path(task_group.path) / wheel_file.filename
                     ).as_posix()
-                    logger.info(f"Write wheel-file contents into {wheel_path}")
-                    with open(wheel_path, "wb") as f:
+                    logger.info(
+                        f"Write wheel-file contents into {archive_path}"
+                    )
+                    with open(archive_path, "wb") as f:
                         f.write(wheel_file.contents)
-                    task_group.wheel_path = wheel_path
+                    task_group.archive_path = archive_path
                     task_group = add_commit_refresh(obj=task_group, db=db)
 
                 # Prepare replacements for templates
@@ -220,7 +222,7 @@ def collect_local(
                 activity = add_commit_refresh(obj=activity, db=db)
 
                 logger.info("_prepare_tasks_metadata - start")
-                task_list = _prepare_tasks_metadata(
+                task_list = prepare_tasks_metadata(
                     package_manifest=pkg_manifest,
                     package_version=task_group.version,
                     package_root=Path(package_root),
@@ -241,15 +243,15 @@ def collect_local(
 
                 # Update task_group data
                 logger.info(
-                    "Add pip_freeze, venv_size and venv_file_number "
+                    "Add env_info, venv_size and venv_file_number "
                     "to TaskGroupV2 - start"
                 )
-                task_group.pip_freeze = pip_freeze_stdout
+                task_group.env_info = 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 "
+                    "Add env_info, venv_size and venv_file_number "
                     "to TaskGroupV2 - end"
                 )