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

fractal_server/config.py

@@ -11,6 +11,7 @@
 # <exact-lab.it> under contract with Liberali Lab from the Friedrich Miescher
 # Institute for Biomedical Research and Pelkmans Lab from the University of
 # Zurich.
+import json
 import logging
 import shutil
 import sys
@@ -34,6 +35,7 @@ from sqlalchemy.engine import URL
 
 import fractal_server
 from fractal_server.types import AbsolutePathStr
+from fractal_server.types import DictStrStr
 
 
 class MailSettings(BaseModel):
@@ -62,6 +64,63 @@ class MailSettings(BaseModel):
     use_login: bool
 
 
+class PixiSettings(BaseModel):
+    """
+    Configuration for Pixi task collection
+
+    See https://pixi.sh/latest/reference/cli/pixi/install/#config-options for
+    `pixi install` concurrency options.
+    See https://docs.rs/tokio/latest/tokio/#cpu-bound-tasks-and-blocking-code
+    for `tokio` configuration.
+
+    versions:
+        Available `pixi` versions and their `PIXI_HOME` folders.
+    default_version:
+        Default `pixi` version to use for task collection - must be one
+        of `versions` keys.
+    PIXI_CONCURRENT_SOLVES:
+        Value of `--concurrent-solves` for `pixi install`.
+    PIXI_CONCURRENT_DOWNLOADS:
+        Value of `--concurrent-downloads for `pixi install`.
+    TOKIO_WORKER_THREADS:
+        From tokio docs, "The core threads are where all asynchronous code
+        runs, and Tokio will by default spawn one for each CPU core. You can
+        use the environment variable TOKIO_WORKER_THREADS to override the
+        default value."
+    """
+
+    versions: DictStrStr
+    default_version: str
+
+    PIXI_CONCURRENT_SOLVES: int = 4
+    PIXI_CONCURRENT_DOWNLOADS: int = 4
+    TOKIO_WORKER_THREADS: int = 2
+
+    @model_validator(mode="after")
+    def check_pixi_settings(self):
+
+        if self.default_version not in self.versions:
+            raise ValueError(
+                f"Default version '{self.default_version}' not in "
+                f"available version {list(self.versions.keys())}."
+            )
+
+        pixi_base_dir = Path(self.versions[self.default_version]).parent
+
+        for key, value in self.versions.items():
+
+            pixi_path = Path(value)
+
+            if pixi_path.parent != pixi_base_dir:
+                raise ValueError(
+                    f"{pixi_path=} is not located within the {pixi_base_dir=}."
+                )
+            if pixi_path.name != key:
+                raise ValueError(f"{pixi_path.name=} is not equal to {key=}")
+
+        return self
+
+
 class FractalConfigurationError(RuntimeError):
     pass
 
@@ -513,6 +572,17 @@ class Settings(BaseSettings):
     FRACTAL_VIEWER_AUTHORIZATION_SCHEME is set to "users-folders".
     """
 
+    FRACTAL_PIXI_CONFIG_FILE: Path | None = None
+
+    pixi: PixiSettings | None = None
+
+    @model_validator(mode="after")
+    def populate_pixi_settings(self):
+        if self.FRACTAL_PIXI_CONFIG_FILE is not None:
+            with self.FRACTAL_PIXI_CONFIG_FILE.open("r") as f:
+                self.pixi = PixiSettings(**json.load(f))
+        return self
+
     ###########################################################################
     # SMTP SERVICE
     ###########################################################################

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")