From 2f794ccf82dbe27d63c1486b11f510d48d3bbbbf Mon Sep 17 00:00:00 2001 From: Marco Mascheroni Date: Mon, 19 Oct 2026 16:34:59 +0200 Subject: [PATCH 1/8] feat(core): add transformation packer contract models (DX-ADR-006) Add the proposal models a packer plugin yields (ComputeParcel, DataParcel, BadInput, DelayInput), the PackerInput it receives, the diracx.transformations.packers entry-point group, and the TransformationNotFoundError and InputStatusConflictError exceptions. Refs #1083 Co-Authored-By: Claude Opus 5.5 --- diracx-core/src/diracx/core/exceptions.py | 32 +++++ diracx-core/src/diracx/core/extensions.py | 1 + .../src/diracx/core/models/__init__.py | 14 +++ .../src/diracx/core/models/transformation.py | 114 +++++++++++++++++- .../tests/test_transformation_models.py | 20 +++ 5 files changed, 180 insertions(+), 1 deletion(-) diff --git a/diracx-core/src/diracx/core/exceptions.py b/diracx-core/src/diracx/core/exceptions.py index ad9dd460f..8487b412b 100644 --- a/diracx-core/src/diracx/core/exceptions.py +++ b/diracx-core/src/diracx/core/exceptions.py @@ -8,6 +8,7 @@ "DocumentUpsertError", "IAMClientError", "IAMServerError", + "InputStatusConflictError", "InvalidCredentialsError", "InvalidQueryError", "NotReadyError", @@ -19,6 +20,7 @@ "SandboxAlreadyInsertedError", "SandboxNotFoundError", "TokenNotFoundError", + "TransformationNotFoundError", ] @@ -185,3 +187,33 @@ class PilotAlreadyExistsError(DiracError): class PilotAlreadyAssociatedWithJobError(DiracError): """We can't associate a pilot with the same job twice.""" + + +class TransformationNotFoundError(DiracError): + """Raised when a transformation cannot be found. + + Attributes: + transformation_id: Identifier of the missing transformation. + """ + + def __init__(self, transformation_id, detail: str = ""): + self.transformation_id = transformation_id + super().__init__( + f"Transformation {transformation_id} not found" + + (f" ({detail})" if detail else "") + ) + + +class InputStatusConflictError(DiracError): + """Raised when transformation inputs are not in the status a change expects. + + Attributes: + transformation_id: Transformation the inputs belong to. + """ + + def __init__(self, transformation_id, detail: str = ""): + self.transformation_id = transformation_id + super().__init__( + f"Inputs of transformation {transformation_id} changed status" + + (f" ({detail})" if detail else "") + ) diff --git a/diracx-core/src/diracx/core/extensions.py b/diracx-core/src/diracx/core/extensions.py index eba98c436..1c728a276 100644 --- a/diracx-core/src/diracx/core/extensions.py +++ b/diracx-core/src/diracx/core/extensions.py @@ -34,6 +34,7 @@ class DiracEntryPoint(StrEnum): RESOURCES = "diracx.resources" SERVICES = "diracx.services" LOCK_OBJECT_TYPES = "diracx.lock_object_types" + TRANSFORMATION_PACKERS = "diracx.transformations.packers" @cached(cache=LRUCache(maxsize=1)) diff --git a/diracx-core/src/diracx/core/models/__init__.py b/diracx-core/src/diracx/core/models/__init__.py index 51ebbab54..7a840af8c 100644 --- a/diracx-core/src/diracx/core/models/__init__.py +++ b/diracx-core/src/diracx/core/models/__init__.py @@ -14,8 +14,12 @@ "WORKGRAPH_TRANSITIONS", "AccessTokenPayload", "ActionResult", + "BadInput", "ChecksumAlgorithm", + "ComputeParcel", "CounterEntity", + "DataParcel", + "DelayInput", "GrantType", "GroupInfo", "HeartbeatData", @@ -33,6 +37,8 @@ "JobStatusUpdate", "Metadata", "OpenIDConfiguration", + "PackerInput", + "PackerProposal", "ParcelStatus", "RefreshTokenPayload", "ReplicaMap", @@ -48,6 +54,7 @@ "SetJobStatusReturn", "SortDirection", "SortSpec", + "StorageDelta", "SummaryParams", "TokenPayload", "TokenResponse", @@ -114,9 +121,16 @@ WORKGRAPH_MEMBER_STATUSES, WORKGRAPH_TRANSITIONS, ActionResult, + BadInput, + ComputeParcel, CounterEntity, + DataParcel, + DelayInput, InputStatus, + PackerInput, + PackerProposal, ParcelStatus, + StorageDelta, TransformationKind, TransformationStatus, WorkgraphStatus, diff --git a/diracx-core/src/diracx/core/models/transformation.py b/diracx-core/src/diracx/core/models/transformation.py index a14e55624..97a747eaf 100644 --- a/diracx-core/src/diracx/core/models/transformation.py +++ b/diracx-core/src/diracx/core/models/transformation.py @@ -1,4 +1,7 @@ -"""Enumerations of the Transformation System (DX-ADR-004, DX-ADR-005). +"""Enumerations and plugin contracts of the Transformation System. + +The enumerations and state machines follow DX-ADR-004 and DX-ADR-005; the +packer contract at the end of the module follows DX-ADR-006. The members of each status enum are written in state-machine order and new states must only ever be appended: the database stores them as native @@ -17,6 +20,12 @@ from __future__ import annotations from enum import StrEnum +from typing import Any +from uuid import UUID + +from pydantic import BaseModel, ConfigDict, Field + +from .types import UTCDatetime class WorkgraphStatus(StrEnum): @@ -290,3 +299,106 @@ class CounterEntity(StrEnum): ), WorkgraphStatus.CANCELLED: frozenset({TransformationStatus.CANCELLED}), } + + +# -------------------------------------------------------------------------- +# Packer contract (DX-ADR-006) +# -------------------------------------------------------------------------- + + +class PackerInput(BaseModel): + """An ``Unassigned`` input claimed by the core and offered to a packer. + + Attributes: + input_id: Identifier of the input. + parent_input_id: Input this one was split from, if any. + lfn: Logical file name, if the input is a file. + lfn_size: Size of the whole file, whatever the descriptor selects. + descriptor: Everything else describing the input. + error_count: Failed parcels the input was part of since the last reset. + """ + + model_config = ConfigDict(frozen=True) + + input_id: UUID + parent_input_id: UUID | None = None + lfn: str | None = None + lfn_size: int | None = None + descriptor: dict[str, Any] = {} + error_count: int = 0 + + +class ComputeParcel(BaseModel): + """A packer's proposal for a parcel that runs as a job. + + Attributes: + inputs: Inputs the parcel is made of. + metadata: Free-form annotations stored on the parcel. + requirements: Additions merged over the transformation's + requirements template. + parameters: The CWL input object passed to the runner. + """ + + inputs: list[UUID] = Field(min_length=1) + metadata: dict[str, Any] = {} + requirements: dict[str, Any] | None = None + parameters: dict[str, Any] = {} + + +class StorageDelta(BaseModel): + """What a data parcel's request costs one storage element. + + Attributes: + se: Storage element affected. + files: Files added there if positive, freed if negative. + bytes: Bytes added there if positive, freed if negative. + """ + + se: str + files: int + bytes: int + + +class DataParcel(BaseModel): + """A packer's proposal for a parcel that runs as a request. + + Attributes: + inputs: Inputs the parcel is made of. + metadata: Free-form annotations stored on the parcel. + request: DIRAC RMS request body (DX-ADR-003). + deltas: What the request costs each storage element. + """ + + inputs: list[UUID] = Field(min_length=1) + metadata: dict[str, Any] = {} + request: dict[str, Any] + deltas: list[StorageDelta] = [] + + +class BadInput(BaseModel): + """A packer's verdict that an input cannot be used at all. + + The core moves the input to ``Problematic`` with the reason. + + Attributes: + input_id: The input rejected. + reason: Why the packer rejected it. + """ + + input_id: UUID + reason: str = "Unknown reason" + + +class DelayInput(BaseModel): + """A packer's request to withhold inputs from packing for a while. + + Attributes: + input_ids: The inputs to withhold. + delay_until: Time before which they are not offered again. + """ + + input_ids: list[UUID] = Field(min_length=1) + delay_until: UTCDatetime + + +PackerProposal = ComputeParcel | DataParcel | BadInput | DelayInput diff --git a/diracx-core/tests/test_transformation_models.py b/diracx-core/tests/test_transformation_models.py index ac7a368b6..99899aba6 100644 --- a/diracx-core/tests/test_transformation_models.py +++ b/diracx-core/tests/test_transformation_models.py @@ -1,6 +1,9 @@ from __future__ import annotations +from uuid import uuid4 + import pytest +from pydantic import ValidationError from diracx.core.models.transformation import ( INPUT_TRANSITIONS, @@ -9,8 +12,11 @@ WORKGRAPH_MEMBER_STATUSES, WORKGRAPH_TRANSITIONS, ActionResult, + ComputeParcel, CounterEntity, + DelayInput, InputStatus, + PackerInput, ParcelStatus, TransformationKind, TransformationStatus, @@ -66,3 +72,17 @@ def test_enum_values_are_append_only(enum_type): ) def test_state_tables_cover_every_state(transitions, enum_type): assert set(transitions) == set(enum_type) + + +def test_packer_proposals_need_inputs(): + with pytest.raises(ValidationError): + ComputeParcel(inputs=[]) + with pytest.raises(ValidationError): + DelayInput(input_ids=[], delay_until="2026-01-01T00:00:00Z") + + +def test_packer_input_is_frozen(): + packer_input = PackerInput(input_id=uuid4(), lfn="/vo/file") + assert packer_input.descriptor == {} + with pytest.raises(ValidationError): + packer_input.lfn = "/vo/other" # type: ignore[misc] From 33bae8bdaa0268b50084a0c125c93ee1c1491d54 Mon Sep 17 00:00:00 2001 From: Marco Mascheroni Date: Mon, 19 Oct 2026 16:38:52 +0200 Subject: [PATCH 2/8] feat(db): add TransformationDB methods for packing inputs Add the bulk operations the packer needs: reading a transformation and its compute facet, counting inputs by status, claiming Unassigned inputs with FOR UPDATE SKIP LOCKED, inserting compute parcels and assigning their inputs, marking inputs Problematic and delaying them. Counter journalling is left to #1085. Refs #1083 Co-Authored-By: Claude Opus 5.5 --- .../src/diracx/db/sql/transformation/db.py | 392 +++++++++++++++++- .../transformation/test_transformation_db.py | 331 +++++++++++++++ 2 files changed, 722 insertions(+), 1 deletion(-) diff --git a/diracx-db/src/diracx/db/sql/transformation/db.py b/diracx-db/src/diracx/db/sql/transformation/db.py index 4a89a88b2..06cc69e7e 100644 --- a/diracx-db/src/diracx/db/sql/transformation/db.py +++ b/diracx-db/src/diracx/db/sql/transformation/db.py @@ -2,13 +2,60 @@ from __future__ import annotations -from sqlalchemy import insert +import uuid +from collections.abc import Iterable +from datetime import datetime +from typing import Any, NamedTuple, TypeAlias + +from sqlalchemy import bindparam, func, insert, select, update from sqlalchemy.ext.asyncio import AsyncConnection +from uuid_utils import UUID +from diracx.core.exceptions import ( + InputStatusConflictError, + TransformationNotFoundError, +) +from diracx.core.models.transformation import ( + InputStatus, + PackerInput, + TransformationKind, + TransformationStatus, +) from diracx.db.sql.utils import BaseSQLDB from diracx.db.sql.utils.counters import counter_locks_table, journalled_counters from .schema import Base as TransformationDBBase +from .schema import ( + ComputeParcels, + ComputeTransformations, + Inputs, + ParcelInputs, + Parcels, + Requirements, + TransformationLog, + Transformations, +) + +# BinaryUUID binds either kind; rows come back as uuid_utils.UUID +AnyUUID: TypeAlias = UUID | uuid.UUID + + +class NewComputeParcel(NamedTuple): + """A compute parcel ready to be inserted. + + Attributes: + parcel_id: UUIDv7 identifier minted for the parcel. + requirements_id: Content hash of its stored requirements, if any. + parameters: The CWL input object passed to the runner. + annotations: Free-form annotations stored on the parcel. + input_ids: Inputs the parcel is made of. + """ + + parcel_id: AnyUUID + requirements_id: bytes | None + parameters: dict[str, Any] + annotations: dict[str, Any] + input_ids: list[AnyUUID] class TransformationDB(BaseSQLDB): @@ -44,3 +91,346 @@ async def post_create(cls, conn: AsyncConnection) -> None: .prefix_with("OR IGNORE", dialect="sqlite"), rows, ) + + async def get_transformation(self, transformation_id: AnyUUID) -> dict[str, Any]: + """Get what the packer needs to know about a transformation. + + Args: + transformation_id: The transformation. + + Returns: + Its ``status``, ``kind``, ``feeder``, ``feeder_enabled``, + ``packer`` and ``packer_args``. + + Raises: + TransformationNotFoundError: If the transformation does not exist. + """ + stmt = select( + Transformations.status.label("status"), + Transformations.kind.label("kind"), + Transformations.feeder.label("feeder"), + Transformations.feeder_enabled.label("feeder_enabled"), + Transformations.packer.label("packer"), + Transformations.packer_args.label("packer_args"), + ).where(Transformations.transformation_id == transformation_id) + row = (await self.conn.execute(stmt)).one_or_none() + if row is None: + raise TransformationNotFoundError(transformation_id) + return dict(row._mapping) + + async def list_transformation_ids(self, status: TransformationStatus) -> list[UUID]: + """List the transformations in a status. + + Args: + status: Status to filter on. + + Returns: + The identifiers of the matching transformations. + """ + stmt = select(Transformations.transformation_id).where( + Transformations.status == status + ) + return list((await self.conn.scalars(stmt)).all()) + + async def get_input_overview( + self, transformation_id: AnyUUID + ) -> dict[InputStatus, int]: + """Count a transformation's inputs by status. + + Args: + transformation_id: The transformation. + + Returns: + The number of inputs in each status that has any. + """ + # TODO(#1085): read the Transformation counter instead + stmt = ( + select(Inputs.status, func.count()) + .where(Inputs.transformation_id == transformation_id) + .group_by(Inputs.status) + ) + return { + InputStatus(status): count + for status, count in (await self.conn.execute(stmt)).all() + } + + async def claim_unassigned_inputs( + self, + transformation_id: AnyUUID, + *, + now: datetime, + after_input_id: AnyUUID | None = None, + limit: int, + ) -> list[PackerInput]: + """Lock a page of a transformation's packable inputs. + + Inputs are ``Unassigned`` ones whose ``DelayedUntil`` has passed, in + ``InputID`` order. Rows locked by another transaction are skipped, and + the locks are held until this transaction ends. + + Args: + transformation_id: The transformation. + now: Time against which ``DelayedUntil`` is compared. + after_input_id: Last input of the previous page, if any. + limit: Maximum number of inputs to claim. + + Returns: + The claimed inputs. + """ + stmt = ( + select( + Inputs.input_id, + Inputs.parent_input_id, + Inputs.lfn, + Inputs.lfn_size, + Inputs.descriptor, + Inputs.error_count, + ) + .where( + Inputs.transformation_id == transformation_id, + Inputs.status == InputStatus.UNASSIGNED, + Inputs.delayed_until.is_(None) | (Inputs.delayed_until <= now), + ) + .order_by(Inputs.input_id) + .limit(limit) + .with_for_update(skip_locked=True) + ) + if after_input_id is not None: + stmt = stmt.where(Inputs.input_id > after_input_id) + return [ + PackerInput( + input_id=uuid.UUID(bytes=row.InputID.bytes), + parent_input_id=( + uuid.UUID(bytes=row.ParentInputID.bytes) + if row.ParentInputID is not None + else None + ), + lfn=row.LFN, + lfn_size=row.LFNSize, + descriptor=row.Descriptor or {}, + error_count=row.ErrorCount, + ) + for row in await self.conn.execute(stmt) + ] + + async def get_compute_facet( + self, transformation_id: AnyUUID + ) -> tuple[bytes, dict[str, Any]]: + """Get what a compute transformation's parcels run and require. + + Args: + transformation_id: The transformation. + + Returns: + The process hash and the requirements template, empty if the + transformation has none. + + Raises: + TransformationNotFoundError: If the transformation has no compute + facet. + """ + stmt = ( + select(ComputeTransformations.process_id, Requirements.requirements) + .outerjoin( + Requirements, + Requirements.requirements_id + == ComputeTransformations.requirements_template_id, + ) + .where(ComputeTransformations.transformation_id == transformation_id) + ) + row = (await self.conn.execute(stmt)).one_or_none() + if row is None: + raise TransformationNotFoundError(transformation_id, "no compute facet") + return row.ProcessID, row.Requirements or {} + + async def insert_requirements( + self, requirements_id: bytes, requirements: dict[str, Any] + ) -> None: + """Store a requirement set unless one with its hash already exists. + + Args: + requirements_id: Content hash of the canonicalised set. + requirements: The requirement set. + """ + await self.conn.execute( + insert(Requirements) + .prefix_with("IGNORE", dialect="mysql") + .prefix_with("OR IGNORE", dialect="sqlite") + .values(requirements_id=requirements_id, requirements=requirements) + ) + + async def insert_compute_parcels( + self, + transformation_id: AnyUUID, + process_id: bytes, + parcels: list[NewComputeParcel], + ) -> None: + """Insert compute parcels and assign their inputs to them. + + Args: + transformation_id: Transformation the parcels belong to. + process_id: Content hash of what the parcels run. + parcels: The parcels to insert. + + Raises: + InputStatusConflictError: If any of the inputs is not + ``Unassigned``. + """ + if not parcels: + return + await self.conn.execute( + insert(Parcels), + [ + { + "ParcelID": parcel.parcel_id, + "TransformationID": transformation_id, + "Kind": TransformationKind.COMPUTE, + "Annotations": parcel.annotations or None, + } + for parcel in parcels + ], + ) + await self.conn.execute( + insert(ComputeParcels), + [ + { + "ParcelID": parcel.parcel_id, + "ProcessID": process_id, + "RequirementsID": parcel.requirements_id, + "Parameters": parcel.parameters or None, + } + for parcel in parcels + ], + ) + links = [ + { + "TransformationID": transformation_id, + "ParcelID": parcel.parcel_id, + "InputID": input_id, + } + for parcel in parcels + for input_id in parcel.input_ids + ] + await self.conn.execute(insert(ParcelInputs), links) + + await self._update_unassigned_inputs( + transformation_id, + (link["InputID"] for link in links), + status=InputStatus.ASSIGNED, + ) + # Each input points at its own parcel, so this one is an executemany + await self.conn.execute( + update(Inputs) + .where( + Inputs.transformation_id == bindparam("b_transformation_id"), + Inputs.input_id == bindparam("b_input_id"), + ) + .values(last_parcel_id=bindparam("b_parcel_id")), + [ + { + "b_transformation_id": transformation_id, + "b_input_id": link["InputID"], + "b_parcel_id": link["ParcelID"], + } + for link in links + ], + ) + # TODO(#1085): journal the inputs' Unassigned -> Assigned transition + # and the new parcels as Unassigned + + async def set_inputs_problematic( + self, + transformation_id: AnyUUID, + reasons: dict[AnyUUID, str], + author: str, + ) -> None: + """Move ``Unassigned`` inputs to ``Problematic``, logging the reasons. + + ``Inputs`` has no column for the reason, so each one is recorded as an + ``InputProblematic`` event in ``TransformationLog``. + + Args: + transformation_id: Transformation the inputs belong to. + reasons: Why each input is problematic. + author: Component responsible for the change. + + Raises: + InputStatusConflictError: If any of the inputs is not + ``Unassigned``. + """ + if not reasons: + return + await self._update_unassigned_inputs( + transformation_id, reasons, status=InputStatus.PROBLEMATIC + ) + await self.conn.execute( + insert(TransformationLog), + [ + { + "TransformationID": transformation_id, + "EventType": "InputProblematic", + "Author": author, + "Message": reason[:512], + "Data": {"InputID": str(input_id), "Reason": reason}, + } + for input_id, reason in reasons.items() + ], + ) + # TODO(#1085): journal the inputs' Unassigned -> Problematic transition + + async def delay_inputs( + self, + transformation_id: AnyUUID, + input_ids: Iterable[AnyUUID], + until: datetime, + ) -> None: + """Withhold ``Unassigned`` inputs from the packer until a given time. + + Args: + transformation_id: Transformation the inputs belong to. + input_ids: The inputs to withhold. + until: Time before which they are not claimed again. + + Raises: + InputStatusConflictError: If any of the inputs is not + ``Unassigned``. + """ + await self._update_unassigned_inputs( + transformation_id, input_ids, delayed_until=until + ) + + async def _update_unassigned_inputs( + self, + transformation_id: AnyUUID, + input_ids: Iterable[AnyUUID], + **values: Any, + ) -> None: + """Update ``Unassigned`` inputs, failing if any is in another status. + + Args: + transformation_id: Transformation the inputs belong to. + input_ids: The inputs to update. + **values: Columns to set. + + Raises: + InputStatusConflictError: If any of the inputs is not + ``Unassigned``. + """ + input_ids = list(input_ids) + if not input_ids: + return + result = await self.conn.execute( + update(Inputs) + .where( + Inputs.transformation_id == transformation_id, + Inputs.input_id.in_(input_ids), + Inputs.status == InputStatus.UNASSIGNED, + ) + .values(**values) + ) + if result.rowcount != len(input_ids): + raise InputStatusConflictError( + transformation_id, + f"{len(input_ids) - result.rowcount} of {len(input_ids)} inputs " + "were not Unassigned", + ) diff --git a/diracx-db/tests/transformation/test_transformation_db.py b/diracx-db/tests/transformation/test_transformation_db.py index 1c23eb954..7af3aa10c 100644 --- a/diracx-db/tests/transformation/test_transformation_db.py +++ b/diracx-db/tests/transformation/test_transformation_db.py @@ -1,10 +1,17 @@ from __future__ import annotations +import uuid +from datetime import datetime, timedelta, timezone + import pytest from sqlalchemy import delete, event, func, insert, select, update from sqlalchemy.exc import IntegrityError from uuid_utils import UUID, uuid7 +from diracx.core.exceptions import ( + InputStatusConflictError, + TransformationNotFoundError, +) from diracx.core.models.transformation import ( InputStatus, ParcelStatus, @@ -12,13 +19,20 @@ TransformationStatus, ) from diracx.db.sql import TransformationDB +from diracx.db.sql.transformation.db import NewComputeParcel from diracx.db.sql.transformation.schema import ( COUNTER_STATUSES, + ComputeParcels, + ComputeTransformations, DataParcelDeltas, DataParcels, Inputs, Ownership, + ParcelInputs, Parcels, + Processes, + Requirements, + TransformationLog, Transformations, Workgraphs, ) @@ -229,3 +243,320 @@ def test_counter_statuses_cover_inputs_and_parcels(): assert set(COUNTER_STATUSES) == {s.value for s in InputStatus} | { s.value for s in ParcelStatus } + + +# -------------------------------------------------------------------------- +# Packing (#1083) +# -------------------------------------------------------------------------- + +PROCESS_ID = b"\x0a" * 32 +TEMPLATE_ID = b"\x0b" * 32 + + +@pytest.fixture +async def compute_facet(transformation_db: TransformationDB, transformation_id: UUID): + async with transformation_db as db: + await db.conn.execute( + insert(Processes).values(process_id=PROCESS_ID, process={"class": "Tool"}) + ) + await db.conn.execute( + insert(Requirements).values( + requirements_id=TEMPLATE_ID, requirements={"cpu": 1} + ) + ) + await db.conn.execute( + insert(ComputeTransformations).values( + transformation_id=transformation_id, + process_id=PROCESS_ID, + requirements_template_id=TEMPLATE_ID, + ) + ) + + +async def add_inputs( + db: TransformationDB, transformation_id: UUID, n: int, **values +) -> list[UUID]: + input_ids = [uuid7() for _ in range(n)] + await db.conn.execute( + insert(Inputs), + [ + { + "TransformationID": transformation_id, + "InputID": input_id, + "InputHash": input_id.bytes, + "LFN": f"/lhcb/{input_id}", + "LFNSize": 10, + "Descriptor": {"Run": 1}, + **values, + } + for input_id in input_ids + ], + ) + return input_ids + + +async def input_statuses(db: TransformationDB) -> dict[UUID, InputStatus]: + return dict((await db.conn.execute(select(Inputs.input_id, Inputs.status))).all()) + + +async def test_get_transformation( + transformation_db: TransformationDB, transformation_id: UUID +): + async with transformation_db as db: + row = await db.get_transformation(transformation_id) + assert row == { + "status": TransformationStatus.NEW, + "kind": TransformationKind.COMPUTE, + "feeder": None, + "feeder_enabled": True, + "packer": "Standard", + "packer_args": None, + } + with pytest.raises(TransformationNotFoundError): + await db.get_transformation(uuid7()) + + +async def test_list_transformation_ids( + transformation_db: TransformationDB, transformation_id: UUID +): + async with transformation_db as db: + assert await db.list_transformation_ids(TransformationStatus.NEW) == [ + transformation_id + ] + assert await db.list_transformation_ids(TransformationStatus.ACTIVE) == [] + + +async def test_input_overview( + transformation_db: TransformationDB, transformation_id: UUID +): + async with transformation_db as db: + assert await db.get_input_overview(transformation_id) == {} + await add_inputs(db, transformation_id, 3) + await add_inputs(db, transformation_id, 2, Status=InputStatus.PROCESSED) + async with transformation_db as db: + assert await db.get_input_overview(transformation_id) == { + InputStatus.UNASSIGNED: 3, + InputStatus.PROCESSED: 2, + } + + +async def test_claim_unassigned_inputs( + transformation_db: TransformationDB, transformation_id: UUID +): + now = datetime.now(tz=timezone.utc) + async with transformation_db as db: + packable = await add_inputs(db, transformation_id, 3) + packable += await add_inputs( + db, transformation_id, 1, DelayedUntil=now - timedelta(hours=1) + ) + await add_inputs( + db, transformation_id, 1, DelayedUntil=now + timedelta(hours=1) + ) + await add_inputs(db, transformation_id, 1, Status=InputStatus.ASSIGNED) + + async with transformation_db as db: + claimed = await db.claim_unassigned_inputs( + transformation_id, now=now, limit=100 + ) + # Nothing of another transformation is claimed + assert await db.claim_unassigned_inputs(uuid7(), now=now, limit=100) == [] + assert [i.input_id for i in claimed] == [uuid.UUID(bytes=i.bytes) for i in packable] + assert all(type(i.input_id) is uuid.UUID for i in claimed) + assert claimed[0].lfn == f"/lhcb/{packable[0]}" + assert claimed[0].lfn_size == 10 + assert claimed[0].descriptor == {"Run": 1} + + # Paging by keyset + async with transformation_db as db: + first = await db.claim_unassigned_inputs(transformation_id, now=now, limit=2) + second = await db.claim_unassigned_inputs( + transformation_id, now=now, after_input_id=first[-1].input_id, limit=2 + ) + assert first + second == claimed + + +async def test_get_compute_facet( + transformation_db: TransformationDB, transformation_id: UUID, compute_facet +): + async with transformation_db as db: + assert await db.get_compute_facet(transformation_id) == ( + PROCESS_ID, + {"cpu": 1}, + ) + await db.conn.execute( + update(ComputeTransformations).values(requirements_template_id=None) + ) + assert await db.get_compute_facet(transformation_id) == (PROCESS_ID, {}) + with pytest.raises(TransformationNotFoundError): + await db.get_compute_facet(uuid7()) + + +async def test_insert_requirements_deduplicates(transformation_db: TransformationDB): + async with transformation_db as db: + await db.insert_requirements(b"\x01" * 32, {"cpu": 4}) + async with transformation_db as db: + await db.insert_requirements(b"\x01" * 32, {"cpu": 4}) + count = await db.conn.scalar(select(func.count()).select_from(Requirements)) + assert count == 1 + + +async def test_insert_compute_parcels( + transformation_db: TransformationDB, transformation_id: UUID, compute_facet +): + async with transformation_db as db: + input_ids = await add_inputs(db, transformation_id, 3) + leftover = await add_inputs(db, transformation_id, 1) + + parcels = [ + NewComputeParcel(uuid7(), TEMPLATE_ID, {"x": 1}, {"note": "a"}, input_ids[:2]), + NewComputeParcel(uuid7(), None, {}, {}, input_ids[2:]), + ] + async with transformation_db as db: + await db.insert_compute_parcels(transformation_id, PROCESS_ID, parcels) + + async with transformation_db as db: + parcel_rows = ( + await db.conn.execute( + select( + Parcels.parcel_id, + Parcels.kind, + Parcels.status, + Parcels.annotations, + ).order_by(Parcels.parcel_id) + ) + ).all() + compute_rows = ( + await db.conn.execute( + select( + ComputeParcels.parcel_id, + ComputeParcels.process_id, + ComputeParcels.requirements_id, + ComputeParcels.parameters, + ).order_by(ComputeParcels.parcel_id) + ) + ).all() + links = set( + ( + await db.conn.execute( + select(ParcelInputs.parcel_id, ParcelInputs.input_id) + ) + ).all() + ) + inputs = { + row.InputID: row + for row in await db.conn.execute( + select(Inputs.input_id, Inputs.status, Inputs.last_parcel_id) + ) + } + + (p1, p2) = parcels + assert parcel_rows == [ + ( + p1.parcel_id, + TransformationKind.COMPUTE, + ParcelStatus.UNASSIGNED, + {"note": "a"}, + ), + (p2.parcel_id, TransformationKind.COMPUTE, ParcelStatus.UNASSIGNED, None), + ] + assert compute_rows == [ + (p1.parcel_id, PROCESS_ID, TEMPLATE_ID, {"x": 1}), + (p2.parcel_id, PROCESS_ID, None, None), + ] + assert links == { + (p1.parcel_id, input_ids[0]), + (p1.parcel_id, input_ids[1]), + (p2.parcel_id, input_ids[2]), + } + for parcel in parcels: + for input_id in parcel.input_ids: + assert inputs[input_id].Status == InputStatus.ASSIGNED + assert inputs[input_id].LastParcelID == parcel.parcel_id + assert inputs[leftover[0]].Status == InputStatus.UNASSIGNED + assert inputs[leftover[0]].LastParcelID is None + + +async def test_insert_compute_parcels_refuses_assigned_inputs( + transformation_db: TransformationDB, transformation_id: UUID, compute_facet +): + async with transformation_db as db: + input_ids = await add_inputs( + db, transformation_id, 1, Status=InputStatus.ASSIGNED + ) + input_ids += await add_inputs(db, transformation_id, 1) + + with pytest.raises(InputStatusConflictError): + async with transformation_db as db: + await db.insert_compute_parcels( + transformation_id, + PROCESS_ID, + [NewComputeParcel(uuid7(), None, {}, {}, input_ids)], + ) + + # The whole transaction was rolled back + async with transformation_db as db: + assert await db.conn.scalar(select(func.count()).select_from(Parcels)) == 0 + statuses = await input_statuses(db) + assert statuses[input_ids[1]] == InputStatus.UNASSIGNED + + +async def test_set_inputs_problematic( + transformation_db: TransformationDB, transformation_id: UUID +): + async with transformation_db as db: + input_ids = await add_inputs(db, transformation_id, 2) + + async with transformation_db as db: + await db.set_inputs_problematic( + transformation_id, {input_ids[0]: "corrupt file"}, author="packer" + ) + + async with transformation_db as db: + statuses = await input_statuses(db) + logs = ( + await db.conn.execute( + select( + TransformationLog.event_type, + TransformationLog.author, + TransformationLog.message, + TransformationLog.data, + ) + ) + ).all() + assert statuses == { + input_ids[0]: InputStatus.PROBLEMATIC, + input_ids[1]: InputStatus.UNASSIGNED, + } + assert logs == [ + ( + "InputProblematic", + "packer", + "corrupt file", + {"InputID": str(input_ids[0]), "Reason": "corrupt file"}, + ) + ] + + with pytest.raises(InputStatusConflictError): + async with transformation_db as db: + await db.set_inputs_problematic( + transformation_id, {input_ids[0]: "again"}, author="packer" + ) + + +async def test_delay_inputs( + transformation_db: TransformationDB, transformation_id: UUID +): + now = datetime.now(tz=timezone.utc) + async with transformation_db as db: + input_ids = await add_inputs(db, transformation_id, 2) + + async with transformation_db as db: + await db.delay_inputs( + transformation_id, input_ids[:1], now + timedelta(hours=1) + ) + + async with transformation_db as db: + claimed = await db.claim_unassigned_inputs(transformation_id, now=now, limit=10) + statuses = await input_statuses(db) + assert [i.input_id for i in claimed] == [uuid.UUID(bytes=input_ids[1].bytes)] + assert set(statuses.values()) == {InputStatus.UNASSIGNED} From a02efafe1d022ce0ad5876cd378736d8222f8d09 Mon Sep 17 00:00:00 2001 From: Marco Mascheroni Date: Mon, 19 Oct 2026 16:44:13 +0200 Subject: [PATCH 3/8] feat(logic): add the transformation packing core and FixedSize packer pack_transformation claims a transformation's Unassigned inputs, streams them to the packer plugin named by the transformation, checks every proposal (offered, not reused) and applies it: compute parcels with their requirements merged over the template, Problematic inputs and delays. Data parcels and dependency injection into plugins are not supported yet. FixedSize is a simple packer making parcels of a fixed number of inputs, holding back the remainder until the feeder is off or on a flush. Refs #1083 Co-Authored-By: Claude Opus 5.5 --- diracx-logic/pyproject.toml | 3 + .../diracx/logic/transformations/__init__.py | 13 + .../diracx/logic/transformations/packers.py | 56 +++ .../diracx/logic/transformations/packing.py | 281 ++++++++++++ .../tests/transformations/test_packers.py | 53 +++ .../tests/transformations/test_packing.py | 401 ++++++++++++++++++ 6 files changed, 807 insertions(+) create mode 100644 diracx-logic/src/diracx/logic/transformations/__init__.py create mode 100644 diracx-logic/src/diracx/logic/transformations/packers.py create mode 100644 diracx-logic/src/diracx/logic/transformations/packing.py create mode 100644 diracx-logic/tests/transformations/test_packers.py create mode 100644 diracx-logic/tests/transformations/test_packing.py diff --git a/diracx-logic/pyproject.toml b/diracx-logic/pyproject.toml index 0ddfb5949..cf39e6991 100644 --- a/diracx-logic/pyproject.toml +++ b/diracx-logic/pyproject.toml @@ -29,6 +29,9 @@ compute_status = "diracx.logic.rss.source:ComputeElementStatusSource" fts_status = "diracx.logic.rss.source:FTSStatusSource" site_status = "diracx.logic.rss.source:SiteStatusSource" +[project.entry-points."diracx.transformations.packers"] +FixedSize = "diracx.logic.transformations.packers:fixed_size_packer" + [project.optional-dependencies] testing = ["diracx-testing", "freezegun"] types = [ diff --git a/diracx-logic/src/diracx/logic/transformations/__init__.py b/diracx-logic/src/diracx/logic/transformations/__init__.py new file mode 100644 index 000000000..587779262 --- /dev/null +++ b/diracx-logic/src/diracx/logic/transformations/__init__.py @@ -0,0 +1,13 @@ +from __future__ import annotations + +__all__ = [ + "FixedSizePackerArgs", + "PackingSummary", + "content_hash", + "fixed_size_packer", + "load_packer", + "pack_transformation", +] + +from .packers import FixedSizePackerArgs, fixed_size_packer +from .packing import PackingSummary, content_hash, load_packer, pack_transformation diff --git a/diracx-logic/src/diracx/logic/transformations/packers.py b/diracx-logic/src/diracx/logic/transformations/packers.py new file mode 100644 index 000000000..a55052f01 --- /dev/null +++ b/diracx-logic/src/diracx/logic/transformations/packers.py @@ -0,0 +1,56 @@ +"""Packer plugins shipped with DiracX (DX-ADR-006).""" + +from __future__ import annotations + +from collections.abc import AsyncGenerator, AsyncIterator + +from pydantic import BaseModel, PositiveInt + +from diracx.core.models.transformation import ( + ComputeParcel, + InputStatus, + PackerInput, +) + + +class FixedSizePackerArgs(BaseModel): + """Arguments of the fixed-size packer. + + Attributes: + group_size: Number of inputs in each parcel. + """ + + group_size: PositiveInt = 1 + + +async def fixed_size_packer( + *, + packer_args: FixedSizePackerArgs, + inputs: AsyncIterator[PackerInput], + flush: bool = False, + feeder_active: bool, + input_overview: dict[InputStatus, int], +) -> AsyncGenerator[ComputeParcel, None]: + """Pack inputs into parcels of ``group_size``, in the order offered. + + A last, smaller group is held back while more input may arrive, and packed + once the feeder is off or on a flush. + + Args: + packer_args: The packer's arguments. + inputs: The claimed ``Unassigned`` inputs. + flush: Whether to pack the last group even if it is small. + feeder_active: Whether more inputs may still arrive. + input_overview: Number of the transformation's inputs in each status. + + Yields: + One compute parcel per group. + """ + group = [] + async for packer_input in inputs: + group.append(packer_input.input_id) + if len(group) == packer_args.group_size: + yield ComputeParcel(inputs=group) + group = [] + if group and (flush or not feeder_active): + yield ComputeParcel(inputs=group) diff --git a/diracx-logic/src/diracx/logic/transformations/packing.py b/diracx-logic/src/diracx/logic/transformations/packing.py new file mode 100644 index 000000000..73cb450ba --- /dev/null +++ b/diracx-logic/src/diracx/logic/transformations/packing.py @@ -0,0 +1,281 @@ +"""The core side of packing a transformation's inputs into parcels (DX-ADR-006). + +The packer plugin only proposes; this module checks every proposal and is the +only thing that writes. Everything happens inside the caller's transaction, so +a proposal that fails a check rolls back the whole invocation. +""" + +from __future__ import annotations + +import hashlib +import inspect +import json +import logging +from collections.abc import AsyncGenerator, AsyncIterator, Callable, Iterator +from datetime import datetime, timezone +from typing import Any, get_type_hints + +from pydantic import BaseModel +from uuid_utils import uuid7 + +from diracx.core.exceptions import ConfigurationError +from diracx.core.extensions import DiracEntryPoint, select_from_extension +from diracx.core.models.transformation import ( + INPUT_TRANSITIONS, + BadInput, + ComputeParcel, + DataParcel, + DelayInput, + InputStatus, + PackerInput, + PackerProposal, + TransformationKind, + TransformationStatus, +) +from diracx.db.sql.transformation.db import ( + AnyUUID, + NewComputeParcel, + TransformationDB, +) + +logger = logging.getLogger(__name__) + +# Inputs claimed from the database per query +CLAIM_BATCH_SIZE = 1000 +# Parcels buffered before they are written +PARCEL_BATCH_SIZE = 500 +# Author recorded on the log rows the packer writes +PACKER_AUTHOR = "TransformationPacker" + + +class PackingSummary(BaseModel): + """What one packer invocation did. + + Attributes: + skipped: Whether the transformation was not ``Active``, so nothing ran. + parcels_created: Parcels inserted. + inputs_assigned: Inputs moved to ``Assigned``. + inputs_problematic: Inputs moved to ``Problematic``. + inputs_delayed: Inputs given a ``DelayedUntil``. + inputs_left: Inputs claimed but left ``Unassigned`` by the packer. + """ + + skipped: bool = False + parcels_created: int = 0 + inputs_assigned: int = 0 + inputs_problematic: int = 0 + inputs_delayed: int = 0 + inputs_left: int = 0 + + +def load_packer(name: str) -> tuple[Callable[..., Any], type[BaseModel]]: + """Find a packer plugin and the model of its arguments. + + The highest priority extension registering ``name`` wins. The arguments + model is the annotation of the plugin's ``packer_args`` parameter. + + Args: + name: Entry point name of the packer. + + Returns: + The plugin function and its arguments model. + + Raises: + ConfigurationError: If no packer is registered under ``name``, or its + ``packer_args`` is not annotated with a pydantic model. + """ + entry_points = select_from_extension( + group=DiracEntryPoint.TRANSFORMATION_PACKERS, name=name + ) + if not entry_points: + raise ConfigurationError( + f"No packer named {name!r} is registered under " + f"{DiracEntryPoint.TRANSFORMATION_PACKERS.value!r}" + ) + plugin = entry_points[0].load() + args_model = get_type_hints(plugin).get("packer_args") + if not (isinstance(args_model, type) and issubclass(args_model, BaseModel)): + raise ConfigurationError( + f"Packer {name!r} must annotate packer_args with a pydantic model" + ) + return plugin, args_model + + +def content_hash(document: dict[str, Any]) -> bytes: + """Hash a JSON document by its canonical form. + + Args: + document: The document to hash. + + Returns: + The SHA-256 digest of the document with sorted keys and no whitespace. + """ + canonical = json.dumps(document, sort_keys=True, separators=(",", ":")) + return hashlib.sha256(canonical.encode()).digest() + + +def _require_transition(old: InputStatus, new: InputStatus) -> None: + if new not in INPUT_TRANSITIONS[old]: + raise ValueError(f"Inputs cannot move from {old} to {new}") + + +async def pack_transformation( + transformation_db: TransformationDB, + transformation_id: AnyUUID, + *, + flush: bool = False, + max_inputs: int, +) -> PackingSummary: + """Run a transformation's packer once and apply what it proposes. + + Args: + transformation_db: Database, inside the caller's transaction. + transformation_id: The transformation to pack. + flush: Ask the packer to also pack groups it would hold back as too + small. + max_inputs: Maximum number of inputs offered to the packer. + + Returns: + What the invocation did. + + Raises: + NotImplementedError: For data transformations and data parcels, which + the core does not support yet. + ValueError: If the packer proposes an input it was not offered, or the + same input twice. + ConfigurationError: If the packer is unknown or its arguments are + invalid. + """ + transformation = await transformation_db.get_transformation(transformation_id) + if transformation["status"] != TransformationStatus.ACTIVE: + return PackingSummary(skipped=True) + if transformation["kind"] != TransformationKind.COMPUTE: + raise NotImplementedError("Packing data transformations is not supported yet") + + plugin, args_model = load_packer(transformation["packer"]) + try: + packer_args = args_model.model_validate(transformation["packer_args"] or {}) + except ValueError as e: + raise ConfigurationError( + f"Invalid arguments for packer {transformation['packer']!r}: {e}" + ) from e + feeder_active = ( + transformation["feeder"] is not None and transformation["feeder_enabled"] + ) + input_overview = await transformation_db.get_input_overview(transformation_id) + process_id, template = await transformation_db.get_compute_facet(transformation_id) + + now = datetime.now(tz=timezone.utc) + claimed: set[AnyUUID] = set() + + async def claim() -> AsyncGenerator[PackerInput, None]: + after = None + while len(claimed) < max_inputs: + page = await transformation_db.claim_unassigned_inputs( + transformation_id, + now=now, + after_input_id=after, + limit=min(CLAIM_BATCH_SIZE, max_inputs - len(claimed)), + ) + for packer_input in page: + claimed.add(packer_input.input_id) + yield packer_input + if len(page) < CLAIM_BATCH_SIZE: + return + after = page[-1].input_id + + kwargs = { + "packer_args": packer_args, + "flush": flush, + "feeder_active": feeder_active, + "input_overview": input_overview, + } + proposals: AsyncIterator[PackerProposal] + if inspect.isasyncgenfunction(plugin): + proposals = plugin(inputs=claim(), **kwargs) + else: + # A synchronous packer cannot await the claim, so it is offered every + # input up front, which max_inputs keeps bounded + offered = [packer_input async for packer_input in claim()] + proposals = _aiter(plugin(inputs=iter(offered), **kwargs)) + + summary = PackingSummary() + used: set[AnyUUID] = set() + parcels: list[NewComputeParcel] = [] + problematic: dict[AnyUUID, str] = {} + # Requirement sets awaiting insertion, and those already inserted + new_requirements: dict[bytes, dict[str, Any]] = {} + stored_requirements: set[bytes] = set() + + def consume(input_ids: list[AnyUUID]) -> None: + for input_id in input_ids: + if input_id not in claimed: + raise ValueError(f"Packer proposed input {input_id} it was not offered") + if input_id in used: + raise ValueError(f"Packer proposed input {input_id} twice") + used.add(input_id) + + async def write() -> None: + for requirements_id, requirements in new_requirements.items(): + await transformation_db.insert_requirements(requirements_id, requirements) + stored_requirements.update(new_requirements) + new_requirements.clear() + if parcels: + _require_transition(InputStatus.UNASSIGNED, InputStatus.ASSIGNED) + await transformation_db.insert_compute_parcels( + transformation_id, process_id, parcels + ) + summary.parcels_created += len(parcels) + summary.inputs_assigned += sum(len(p.input_ids) for p in parcels) + parcels.clear() + if problematic: + _require_transition(InputStatus.UNASSIGNED, InputStatus.PROBLEMATIC) + await transformation_db.set_inputs_problematic( + transformation_id, problematic, author=PACKER_AUTHOR + ) + summary.inputs_problematic += len(problematic) + problematic.clear() + + async for proposal in proposals: + match proposal: + case ComputeParcel(): + consume(proposal.inputs) + requirements = {**template, **(proposal.requirements or {})} + requirements_id = content_hash(requirements) if requirements else None + if requirements_id and requirements_id not in stored_requirements: + new_requirements[requirements_id] = requirements + parcels.append( + NewComputeParcel( + parcel_id=uuid7(), + requirements_id=requirements_id, + parameters=proposal.parameters, + annotations=proposal.metadata, + input_ids=list(proposal.inputs), + ) + ) + case BadInput(): + consume([proposal.input_id]) + problematic[proposal.input_id] = proposal.reason + case DelayInput(): + consume(proposal.input_ids) + await transformation_db.delay_inputs( + transformation_id, proposal.input_ids, proposal.delay_until + ) + summary.inputs_delayed += len(proposal.input_ids) + case DataParcel(): + raise NotImplementedError("Data parcels are not supported yet") + case _: + raise TypeError(f"Packer yielded {proposal!r}, not a proposal") + if len(parcels) + len(problematic) >= PARCEL_BATCH_SIZE: + await write() + await write() + + # Inputs the packer held back stay Unassigned; their locks go at commit + summary.inputs_left = len(claimed) - len(used) + logger.info("Packed transformation %s: %s", transformation_id, summary) + return summary + + +async def _aiter(iterator: Iterator[PackerProposal]) -> AsyncIterator[PackerProposal]: + for item in iterator: + yield item diff --git a/diracx-logic/tests/transformations/test_packers.py b/diracx-logic/tests/transformations/test_packers.py new file mode 100644 index 000000000..57caa9d21 --- /dev/null +++ b/diracx-logic/tests/transformations/test_packers.py @@ -0,0 +1,53 @@ +from __future__ import annotations + +from uuid import uuid4 + +import pytest + +from diracx.core.models.transformation import PackerInput +from diracx.logic.transformations.packers import ( + FixedSizePackerArgs, + fixed_size_packer, +) + + +async def run(n: int, group_size: int, *, flush=False, feeder_active=True): + inputs = [PackerInput(input_id=uuid4()) for _ in range(n)] + + async def stream(): + for packer_input in inputs: + yield packer_input + + parcels = [ + parcel.inputs + async for parcel in fixed_size_packer( + packer_args=FixedSizePackerArgs(group_size=group_size), + inputs=stream(), + flush=flush, + feeder_active=feeder_active, + input_overview={}, + ) + ] + return [i.input_id for i in inputs], parcels + + +async def test_full_groups_only_while_feeder_active(): + ids, parcels = await run(7, 3) + assert parcels == [ids[0:3], ids[3:6]] + + +@pytest.mark.parametrize( + "kwargs", [{"flush": True}, {"feeder_active": False}], ids=["flush", "drained"] +) +async def test_remainder_packed_when_draining(kwargs): + ids, parcels = await run(7, 3, **kwargs) + assert parcels == [ids[0:3], ids[3:6], ids[6:]] + + +async def test_no_inputs_no_parcels(): + assert await run(0, 3, flush=True) == ([], []) + + +def test_group_size_must_be_positive(): + with pytest.raises(ValueError): + FixedSizePackerArgs(group_size=0) diff --git a/diracx-logic/tests/transformations/test_packing.py b/diracx-logic/tests/transformations/test_packing.py new file mode 100644 index 000000000..47ef3e2ee --- /dev/null +++ b/diracx-logic/tests/transformations/test_packing.py @@ -0,0 +1,401 @@ +from __future__ import annotations + +from datetime import datetime, timedelta, timezone +from uuid import uuid4 + +import pytest +from sqlalchemy import event, insert, select, update +from uuid_utils import UUID, uuid7 + +from diracx.core.exceptions import ConfigurationError +from diracx.core.models.transformation import ( + BadInput, + ComputeParcel, + DataParcel, + DelayInput, + InputStatus, + TransformationKind, + TransformationStatus, +) +from diracx.db.sql import TransformationDB +from diracx.db.sql.transformation.schema import ( + ComputeParcels, + ComputeTransformations, + Inputs, + Ownership, + ParcelInputs, + Parcels, + Processes, + Requirements, + TransformationLog, + Transformations, + Workgraphs, +) +from diracx.logic.transformations import packing +from diracx.logic.transformations.packers import ( + FixedSizePackerArgs, + fixed_size_packer, +) +from diracx.logic.transformations.packing import ( + content_hash, + load_packer, + pack_transformation, +) + +PROCESS_ID = b"\x0a" * 32 +TEMPLATE = {"cpu": 1, "platform": "x86_64"} + + +@pytest.fixture +async def transformation_db(): + db = TransformationDB("sqlite+aiosqlite:///:memory:") + async with db.engine_context(): + event.listen( + db.engine.sync_engine, + "connect", + lambda dbapi_conn, _: dbapi_conn.cursor().execute("PRAGMA foreign_keys=ON"), + ) + async with db.engine.begin() as conn: + await conn.run_sync(db.metadata.create_all) + await db.post_create(conn) + yield db + + +@pytest.fixture +async def transformation_id(transformation_db: TransformationDB) -> UUID: + workgraph_id, ownership_id, transformation_id = uuid7(), uuid7(), uuid7() + template_id = content_hash(TEMPLATE) + async with transformation_db as db: + await db.conn.execute( + insert(Workgraphs).values( + workgraph_id=workgraph_id, type="MCSimulation", stripped_spec={} + ) + ) + await db.conn.execute( + insert(Ownership).values( + ownership_id=ownership_id, + vo="lhcb", + sub="some-subject", + workgraph_id=workgraph_id, + ) + ) + await db.conn.execute( + insert(Transformations).values( + transformation_id=transformation_id, + ownership_id=ownership_id, + type="MCSimulation", + kind=TransformationKind.COMPUTE, + status=TransformationStatus.ACTIVE, + feeder="SomeFeeder", + packer="Test", + packer_args={"group_size": 2}, + ) + ) + await db.conn.execute( + insert(Processes).values(process_id=PROCESS_ID, process={"class": "Tool"}) + ) + await db.conn.execute( + insert(Requirements).values( + requirements_id=template_id, requirements=TEMPLATE + ) + ) + await db.conn.execute( + insert(ComputeTransformations).values( + transformation_id=transformation_id, + process_id=PROCESS_ID, + requirements_template_id=template_id, + ) + ) + return transformation_id + + +@pytest.fixture +def use_packer(monkeypatch): + def use(plugin, args_model=FixedSizePackerArgs): + monkeypatch.setattr(packing, "load_packer", lambda name: (plugin, args_model)) + + use(fixed_size_packer) + return use + + +async def add_inputs(transformation_db, transformation_id, n: int) -> list[UUID]: + input_ids = [uuid7() for _ in range(n)] + async with transformation_db as db: + await db.conn.execute( + insert(Inputs), + [ + { + "TransformationID": transformation_id, + "InputID": input_id, + "InputHash": input_id.bytes, + "LFN": f"/lhcb/{input_id}", + } + for input_id in input_ids + ], + ) + return input_ids + + +async def pack(transformation_db, transformation_id, *, flush=False, max_inputs=100): + async with transformation_db as db: + return await pack_transformation( + db, transformation_id, flush=flush, max_inputs=max_inputs + ) + + +async def input_statuses(transformation_db) -> dict[UUID, InputStatus]: + async with transformation_db as db: + return dict( + (await db.conn.execute(select(Inputs.input_id, Inputs.status))).all() + ) + + +async def parcel_inputs(transformation_db) -> list[set[UUID]]: + async with transformation_db as db: + rows = ( + await db.conn.execute( + select(ParcelInputs.parcel_id, ParcelInputs.input_id).order_by( + ParcelInputs.parcel_id + ) + ) + ).all() + groups: dict[UUID, set[UUID]] = {} + for parcel_id, input_id in rows: + groups.setdefault(parcel_id, set()).add(input_id) + return list(groups.values()) + + +async def set_transformation(transformation_db, **values): + async with transformation_db as db: + await db.conn.execute(update(Transformations).values(**values)) + + +async def test_fixed_size_holds_remainder_while_feeding( + transformation_db, transformation_id, use_packer +): + input_ids = await add_inputs(transformation_db, transformation_id, 5) + + summary = await pack(transformation_db, transformation_id) + + assert summary.model_dump() == { + "skipped": False, + "parcels_created": 2, + "inputs_assigned": 4, + "inputs_problematic": 0, + "inputs_delayed": 0, + "inputs_left": 1, + } + assert await parcel_inputs(transformation_db) == [ + set(input_ids[0:2]), + set(input_ids[2:4]), + ] + statuses = await input_statuses(transformation_db) + assert statuses[input_ids[4]] == InputStatus.UNASSIGNED + + # Running again packs nothing new while the feeder is still active + summary = await pack(transformation_db, transformation_id) + assert (summary.parcels_created, summary.inputs_left) == (0, 1) + + +@pytest.mark.parametrize("drain", ["flush", "feeder_disabled", "no_feeder"]) +async def test_fixed_size_packs_remainder_when_draining( + transformation_db, transformation_id, use_packer, drain +): + input_ids = await add_inputs(transformation_db, transformation_id, 3) + if drain == "feeder_disabled": + await set_transformation(transformation_db, feeder_enabled=False) + elif drain == "no_feeder": + await set_transformation(transformation_db, feeder=None) + + summary = await pack(transformation_db, transformation_id, flush=drain == "flush") + + assert (summary.parcels_created, summary.inputs_left) == (2, 0) + assert await parcel_inputs(transformation_db) == [ + set(input_ids[0:2]), + set(input_ids[2:]), + ] + + +async def test_max_inputs_bounds_the_claim( + transformation_db, transformation_id, use_packer +): + await add_inputs(transformation_db, transformation_id, 5) + summary = await pack(transformation_db, transformation_id, max_inputs=3) + assert (summary.parcels_created, summary.inputs_left) == (1, 1) + + +async def test_paging_through_many_inputs( + transformation_db, transformation_id, use_packer, monkeypatch +): + monkeypatch.setattr(packing, "CLAIM_BATCH_SIZE", 2) + monkeypatch.setattr(packing, "PARCEL_BATCH_SIZE", 2) + input_ids = await add_inputs(transformation_db, transformation_id, 7) + await set_transformation(transformation_db, feeder_enabled=False) + + summary = await pack(transformation_db, transformation_id) + + assert summary.parcels_created == 4 + assert summary.inputs_assigned == 7 + statuses = await input_statuses(transformation_db) + assert set(statuses.values()) == {InputStatus.ASSIGNED} + assert len(statuses) == len(input_ids) + + +async def test_requirements_merged_over_template( + transformation_db, transformation_id, use_packer +): + await add_inputs(transformation_db, transformation_id, 3) + + async def plugin(*, packer_args, inputs, **_): + group = [i.input_id async for i in inputs] + yield ComputeParcel( + inputs=group[:1], + requirements={"cpu": 8}, + parameters={"seed": 1}, + metadata={"note": "big"}, + ) + yield ComputeParcel(inputs=group[1:2], requirements={"cpu": 8}) + yield ComputeParcel(inputs=group[2:]) + + use_packer(plugin) + await pack(transformation_db, transformation_id) + + async with transformation_db as db: + rows = ( + await db.conn.execute( + select( + Requirements.requirements, + ComputeParcels.process_id, + ComputeParcels.parameters, + Parcels.annotations, + ) + .join(Requirements) + .join(Parcels) + .order_by(Parcels.parcel_id) + ) + ).all() + n_requirements = len((await db.conn.execute(select(Requirements))).all()) + merged = {"cpu": 8, "platform": "x86_64"} + assert rows == [ + (merged, PROCESS_ID, {"seed": 1}, {"note": "big"}), + (merged, PROCESS_ID, None, None), + (TEMPLATE, PROCESS_ID, None, None), + ] + # The template and one merged set, stored once each + assert n_requirements == 2 + + +async def test_bad_and_delayed_inputs(transformation_db, transformation_id, use_packer): + input_ids = await add_inputs(transformation_db, transformation_id, 3) + later = datetime.now(tz=timezone.utc) + timedelta(hours=1) + + async def plugin(*, packer_args, inputs, **_): + offered = [i.input_id async for i in inputs] + yield BadInput(input_id=offered[0], reason="corrupt") + yield DelayInput(input_ids=offered[1:], delay_until=later) + + use_packer(plugin) + summary = await pack(transformation_db, transformation_id) + + assert (summary.inputs_problematic, summary.inputs_delayed) == (1, 2) + assert summary.inputs_left == 0 + statuses = await input_statuses(transformation_db) + assert statuses[input_ids[0]] == InputStatus.PROBLEMATIC + async with transformation_db as db: + data = await db.conn.scalar(select(TransformationLog.data)) + assert data == {"InputID": str(input_ids[0]), "Reason": "corrupt"} + + # Delayed inputs are not offered again before their time + use_packer(fixed_size_packer) + summary = await pack(transformation_db, transformation_id, flush=True) + assert summary.parcels_created == 0 + + +async def test_sync_packer(transformation_db, transformation_id, use_packer): + await add_inputs(transformation_db, transformation_id, 2) + + def plugin(*, packer_args, inputs, **_): + yield ComputeParcel(inputs=[i.input_id for i in inputs]) + + use_packer(plugin) + summary = await pack(transformation_db, transformation_id) + assert summary.inputs_assigned == 2 + + +@pytest.mark.parametrize("mistake", ["unoffered", "twice"]) +async def test_invalid_proposals_roll_back( + transformation_db, transformation_id, use_packer, mistake +): + await add_inputs(transformation_db, transformation_id, 2) + + async def plugin(*, packer_args, inputs, **_): + offered = [i.input_id async for i in inputs] + yield ComputeParcel(inputs=offered[:1]) + if mistake == "unoffered": + yield ComputeParcel(inputs=[uuid4()]) + else: + yield BadInput(input_id=offered[0]) + + use_packer(plugin) + with pytest.raises(ValueError, match=mistake.replace("unoffered", "not offered")): + await pack(transformation_db, transformation_id) + statuses = await input_statuses(transformation_db) + assert set(statuses.values()) == {InputStatus.UNASSIGNED} + + +async def test_inactive_transformation_is_skipped( + transformation_db, transformation_id, use_packer +): + await add_inputs(transformation_db, transformation_id, 2) + await set_transformation(transformation_db, status=TransformationStatus.PAUSED) + summary = await pack(transformation_db, transformation_id, flush=True) + assert summary == packing.PackingSummary(skipped=True) + + +async def test_invalid_packer_args(transformation_db, transformation_id, use_packer): + await set_transformation(transformation_db, packer_args={"group_size": 0}) + with pytest.raises(ConfigurationError, match="Invalid arguments"): + await pack(transformation_db, transformation_id) + + +async def test_data_is_not_supported(transformation_db, transformation_id, use_packer): + await add_inputs(transformation_db, transformation_id, 1) + + async def plugin(*, packer_args, inputs, **_): + offered = [i.input_id async for i in inputs] + yield DataParcel(inputs=offered, request={}) + + use_packer(plugin) + with pytest.raises(NotImplementedError): + await pack(transformation_db, transformation_id) + + await set_transformation(transformation_db, kind=TransformationKind.DATA) + with pytest.raises(NotImplementedError): + await pack(transformation_db, transformation_id) + + +def test_load_packer_finds_fixed_size(): + assert load_packer("FixedSize") == (fixed_size_packer, FixedSizePackerArgs) + + +def test_load_packer_unknown(): + with pytest.raises(ConfigurationError, match="No packer named"): + load_packer("DoesNotExist") + + +def test_load_packer_needs_an_args_model(monkeypatch): + class EntryPoint: + def load(self): + def plugin(*, packer_args: dict, inputs): ... + + return plugin + + monkeypatch.setattr(packing, "select_from_extension", lambda **_: [EntryPoint()]) + with pytest.raises(ConfigurationError, match="pydantic model"): + load_packer("Untyped") + + +def test_content_hash_is_canonical(): + assert content_hash({"a": 1, "b": [1, 2]}) == content_hash({"b": [1, 2], "a": 1}) + assert len(content_hash({})) == 32 + assert content_hash({"a": 1}) != content_hash({"a": 2}) From 4bb41929f948e3ecbe38d16dc053f13b78bb966d Mon Sep 17 00:00:00 2001 From: Marco Mascheroni Date: Mon, 19 Oct 2026 16:49:15 +0200 Subject: [PATCH 4/8] feat(tasks): add the transformation packer sweep and per-transformation task PackerSweepTask periodically schedules a PackTransformationTask for every Active transformation; each one holds a mutex on its transformation and calls pack_transformation. The sweep is disabled by default and configured through DIRACX_TASKS_PACKER_*. Refs #1083 Co-Authored-By: Claude Opus 5.5 --- diracx-tasks/pyproject.toml | 4 + .../diracx/tasks/transformations/__init__.py | 8 ++ .../diracx/tasks/transformations/packer.py | 98 +++++++++++++++++++ diracx-tasks/tests/test_packer_tasks.py | 97 ++++++++++++++++++ 4 files changed, 207 insertions(+) create mode 100644 diracx-tasks/src/diracx/tasks/transformations/__init__.py create mode 100644 diracx-tasks/src/diracx/tasks/transformations/packer.py create mode 100644 diracx-tasks/tests/test_packer_tasks.py diff --git a/diracx-tasks/pyproject.toml b/diracx-tasks/pyproject.toml index fb8442162..730d77814 100644 --- a/diracx-tasks/pyproject.toml +++ b/diracx-tasks/pyproject.toml @@ -40,6 +40,10 @@ CleanSandboxStoreTask = "diracx.tasks.jobs:CleanSandboxStoreTask" DummyJobExecutorMonitorTask = "diracx.tasks.jobs:DummyJobExecutorMonitorTask" DummyJobExecutorTask = "diracx.tasks.jobs:DummyJobExecutorTask" +[project.entry-points."diracx.tasks.transformation"] +PackerSweepTask = "diracx.tasks.transformations:PackerSweepTask" +PackTransformationTask = "diracx.tasks.transformations:PackTransformationTask" + [build-system] requires = ["hatchling", "hatch-vcs"] build-backend = "hatchling.build" diff --git a/diracx-tasks/src/diracx/tasks/transformations/__init__.py b/diracx-tasks/src/diracx/tasks/transformations/__init__.py new file mode 100644 index 000000000..d70f76ec3 --- /dev/null +++ b/diracx-tasks/src/diracx/tasks/transformations/__init__.py @@ -0,0 +1,8 @@ +from __future__ import annotations + +__all__ = [ + "PackTransformationTask", + "PackerSweepTask", +] + +from .packer import PackerSweepTask, PackTransformationTask diff --git a/diracx-tasks/src/diracx/tasks/transformations/packer.py b/diracx-tasks/src/diracx/tasks/transformations/packer.py new file mode 100644 index 000000000..0c863d2f5 --- /dev/null +++ b/diracx-tasks/src/diracx/tasks/transformations/packer.py @@ -0,0 +1,98 @@ +"""Tasks that run transformation packers (DX-ADR-006).""" + +from __future__ import annotations + +import dataclasses +import logging +from typing import Any +from uuid import UUID + +from pydantic import PositiveInt + +from diracx.core.models.transformation import TransformationStatus +from diracx.core.settings import ServiceSettingsBase +from diracx.db.sql import TransformationDB +from diracx.logic.transformations import pack_transformation +from diracx.tasks.plumbing.base_task import BaseTask, PeriodicBaseTask +from diracx.tasks.plumbing.enums import Priority, Size +from diracx.tasks.plumbing.lock_registry import TRANSFORMATION +from diracx.tasks.plumbing.locks import BaseLock, MutexLock +from diracx.tasks.plumbing.retry_policies import ExponentialBackoff +from diracx.tasks.plumbing.schedules import IntervalSeconds + +logger = logging.getLogger(__name__) + + +class PackerSettings(ServiceSettingsBase): + """Settings controlling the transformation packer.""" + + model_config = ServiceSettingsBase.model_config | { + "env_prefix": "DIRACX_TASKS_PACKER_", + "use_attribute_docstrings": True, + } + + enabled: bool = False + """Whether the packer sweep is scheduled automatically.""" + + interval_seconds: PositiveInt = 60 + """How often the sweep schedules a packer run for each active transformation.""" + + max_inputs_per_invocation: PositiveInt = 10_000 + """Maximum number of inputs offered to a packer in one run.""" + + +_settings = PackerSettings() + + +@dataclasses.dataclass +class PackerSweepTask(PeriodicBaseTask): + """Schedule one packer run per ``Active`` transformation. + + Each run holds a mutex on its transformation, so thousands of + transformations pack in parallel while no two runs for one overlap. + """ + + priority = Priority.BACKGROUND + size = Size.SMALL + _enabled = _settings.enabled + default_schedule = IntervalSeconds(_settings.interval_seconds) + + async def execute(self, transformation_db: TransformationDB) -> int: + transformation_ids = await transformation_db.list_transformation_ids( + TransformationStatus.ACTIVE + ) + for transformation_id in transformation_ids: + await PackTransformationTask( + transformation_id=str(transformation_id) + ).schedule() + return len(transformation_ids) + + +@dataclasses.dataclass +class PackTransformationTask(BaseTask): + """Run one transformation's packer and apply its proposals. + + Attributes: + transformation_id: The transformation to pack. + flush: Ask the packer to also pack groups it would hold back. + """ + + priority = Priority.NORMAL + size = Size.MEDIUM + retry_policy = ExponentialBackoff(base_delay_seconds=30, max_retries=3) + + transformation_id: str + flush: bool = False + + @property + def execution_locks(self) -> list[BaseLock]: + return [MutexLock(TRANSFORMATION, self.transformation_id)] + + async def execute(self, transformation_db: TransformationDB) -> dict[str, Any]: + summary = await pack_transformation( + transformation_db, + UUID(self.transformation_id), + flush=self.flush, + max_inputs=_settings.max_inputs_per_invocation, + ) + return summary.model_dump() diff --git a/diracx-tasks/tests/test_packer_tasks.py b/diracx-tasks/tests/test_packer_tasks.py new file mode 100644 index 000000000..dd585b5b6 --- /dev/null +++ b/diracx-tasks/tests/test_packer_tasks.py @@ -0,0 +1,97 @@ +"""Tests for the transformation packer tasks.""" + +from __future__ import annotations + +from unittest.mock import AsyncMock +from uuid import UUID, uuid4 + +import pytest +from pydantic import ValidationError + +from diracx.core.models.transformation import TransformationStatus +from diracx.logic.transformations import PackingSummary +from diracx.tasks.plumbing.enums import Priority, Size +from diracx.tasks.plumbing.factory import load_task_registry +from diracx.tasks.plumbing.locks import MutexLock +from diracx.tasks.transformations import PackerSweepTask, PackTransformationTask +from diracx.tasks.transformations import packer as packer_module + + +def test_tasks_are_registered(): + registry = load_task_registry() + assert registry["transformation:PackerSweepTask"] is PackerSweepTask + assert registry["transformation:PackTransformationTask"] is PackTransformationTask + + +def test_task_metadata(): + assert PackerSweepTask.priority == Priority.BACKGROUND + assert PackerSweepTask.size == Size.SMALL + assert PackerSweepTask._enabled is False + assert PackTransformationTask.retry_policy.max_retries == 3 + assert PackTransformationTask.dlq_eligible is False + + +def test_serialize_round_trip(): + task = PackTransformationTask(transformation_id=str(uuid4()), flush=True) + assert PackTransformationTask(*task.serialize()) == task + + +def test_pack_task_takes_a_per_transformation_mutex(): + first, second = str(uuid4()), str(uuid4()) + (lock,) = PackTransformationTask(transformation_id=first).execution_locks + assert isinstance(lock, MutexLock) + assert lock.redis_key == f"lock:mutex:transformation:{first}" + (other,) = PackTransformationTask(transformation_id=second).execution_locks + assert other.redis_key != lock.redis_key + + +def test_settings_must_be_positive(): + with pytest.raises(ValidationError): + packer_module.PackerSettings(interval_seconds=0) + with pytest.raises(ValidationError): + packer_module.PackerSettings(max_inputs_per_invocation=0) + + +async def test_sweep_schedules_one_task_per_active_transformation(monkeypatch): + active = [uuid4(), uuid4()] + scheduled = [] + + async def fake_schedule(self, **kwargs): + scheduled.append(self) + return "task-id" + + monkeypatch.setattr(PackTransformationTask, "schedule", fake_schedule) + transformation_db = AsyncMock() + transformation_db.list_transformation_ids.return_value = active + + assert await PackerSweepTask().execute(transformation_db) == 2 + + transformation_db.list_transformation_ids.assert_awaited_once_with( + TransformationStatus.ACTIVE + ) + assert scheduled == [ + PackTransformationTask(transformation_id=str(transformation_id)) + for transformation_id in active + ] + + +async def test_pack_task_delegates_to_logic(monkeypatch): + pack_transformation = AsyncMock( + return_value=PackingSummary(parcels_created=1, inputs_assigned=2) + ) + monkeypatch.setattr(packer_module, "pack_transformation", pack_transformation) + transformation_db = AsyncMock() + transformation_id = uuid4() + + result = await PackTransformationTask( + transformation_id=str(transformation_id), flush=True + ).execute(transformation_db) + + assert result["parcels_created"] == 1 + assert result["inputs_assigned"] == 2 + pack_transformation.assert_awaited_once_with( + transformation_db, + UUID(str(transformation_id)), + flush=True, + max_inputs=packer_module._settings.max_inputs_per_invocation, + ) From 60ba5ed28b1846c510bd76b1553ec9e007b6b1b2 Mon Sep 17 00:00:00 2001 From: Marco Mascheroni Date: Mon, 19 Oct 2026 16:51:00 +0200 Subject: [PATCH 5/8] docs: explain the transformation packer and how to write one Add a developer page on the packer contract, what the core checks and how to register a packer through the diracx.transformations.packers entry point, and document the DIRACX_TASKS_PACKER_* settings. Refs #1083 Co-Authored-By: Claude Opus 5.5 --- docs/admin/reference/env-variables.md | 22 +++ docs/admin/reference/env-variables.md.j2 | 2 + .../transformation-system/packer.md | 125 ++++++++++++++++++ mkdocs.yml | 2 + 4 files changed, 151 insertions(+) create mode 100644 docs/dev/explanations/transformation-system/packer.md diff --git a/docs/admin/reference/env-variables.md b/docs/admin/reference/env-variables.md index 1d1d3302d..1c35b52f3 100644 --- a/docs/admin/reference/env-variables.md +++ b/docs/admin/reference/env-variables.md @@ -484,3 +484,25 @@ Whether the monitor is scheduled automatically. *Optional*, default value: `10` How often the enabled monitor searches for received jobs. + +## PackerSettings + +Settings controlling the transformation packer. + +### `DIRACX_TASKS_PACKER_ENABLED` + +*Optional*, default value: `False` + +Whether the packer sweep is scheduled automatically. + +### `DIRACX_TASKS_PACKER_INTERVAL_SECONDS` + +*Optional*, default value: `60` + +How often the sweep schedules a packer run for each active transformation. + +### `DIRACX_TASKS_PACKER_MAX_INPUTS_PER_INVOCATION` + +*Optional*, default value: `10000` + +Maximum number of inputs offered to a packer in one run. diff --git a/docs/admin/reference/env-variables.md.j2 b/docs/admin/reference/env-variables.md.j2 index 0c8cda89a..5bf5e3b2a 100644 --- a/docs/admin/reference/env-variables.md.j2 +++ b/docs/admin/reference/env-variables.md.j2 @@ -24,3 +24,5 @@ _X, where X is a number. The files will be loaded in order. ## Tasks {{ render_class('DummyJobExecutorSettings') }} + +{{ render_class('PackerSettings') }} diff --git a/docs/dev/explanations/transformation-system/packer.md b/docs/dev/explanations/transformation-system/packer.md new file mode 100644 index 000000000..a72c69e06 --- /dev/null +++ b/docs/dev/explanations/transformation-system/packer.md @@ -0,0 +1,125 @@ +# The packer + +The packer is the part of the Transformation System that turns a transformation's `Unassigned` inputs into parcels. It is split in two, following [DX-ADR-006](../../../adr/DX-ADR-006_extensions.md): + +- a **packer plugin**, chosen per transformation by its `Packer` column, decides how inputs are grouped; +- the **core**, `pack_transformation` in `diracx.logic.transformations`, claims the inputs, checks what the plugin proposes and is the only thing that writes to the database. + +One rule applies throughout: **plugins propose, the core writes.** + +## How a packer run works + +`PackerSweepTask` runs periodically and schedules one `PackTransformationTask` for each `Active` transformation. Each of those tasks holds a mutex on its transformation id, so two runs for one transformation never overlap while different transformations pack in parallel. A run is one database transaction: + +1. The core loads the transformation. A transformation that is not `Active` is skipped. +2. It resolves the plugin from the `diracx.transformations.packers` entry point named by `Packer`, and validates `PackerArgs` against the plugin's arguments model. +3. It claims `Unassigned` inputs whose `DelayedUntil` has passed, with `SELECT ... FOR UPDATE SKIP LOCKED`, in pages of 1000 and up to `DIRACX_TASKS_PACKER_MAX_INPUTS_PER_INVOCATION`, and streams them to the plugin. +4. It checks and applies each proposal the plugin yields, writing parcels in batches. +5. Inputs the plugin did not use stay `Unassigned`; their row locks are released when the transaction commits. + +If any check fails, the task fails and the whole transaction rolls back, so a run either applies everything the plugin proposed or nothing. + +## The plugin contract + +A packer is a generator function, preferably asynchronous, called with keyword arguments only: + +| Argument | Type | Meaning | +| ---------------- | ------------------------ | ------------------------------------------------------------------------------------- | +| `packer_args` | a pydantic model | The transformation's `PackerArgs`, validated against the annotation of this parameter | +| `inputs` | stream of `PackerInput` | The claimed inputs, in `InputID` order | +| `flush` | `bool` | Pack groups that would otherwise be held back as too small | +| `feeder_active` | `bool` | Whether more inputs may still arrive: the transformation has an enabled feeder | +| `input_overview` | `dict[InputStatus, int]` | How many of the transformation's inputs are in each status | + +`PackerInput` carries `input_id`, `parent_input_id`, `lfn`, `lfn_size`, `descriptor` and `error_count`. + +The plugin yields proposals, all defined in `diracx.core.models.transformation`: + +| Proposal | Effect | +| --------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `ComputeParcel` | A new `Unassigned` parcel made of the given inputs, which move to `Assigned`. Its `requirements` are merged over the transformation's requirements template; `parameters` is the CWL input object; `metadata` is stored as the parcel's annotations. | +| `BadInput` | The input moves to `Problematic`, and the reason is recorded as an `InputProblematic` event in `TransformationLog`. | +| `DelayInput` | The inputs get `DelayedUntil` and are not offered again before that time. | +| `DataParcel` | Not supported yet. | + +### What the core checks + +- Every input in a proposal must have been offered to the plugin in this run. +- No input may appear in more than one proposal in a run. +- Every input must still be `Unassigned` when it is written, and the move must be allowed by the input state machine of [DX-ADR-005](../../../adr/DX-ADR-005_state_machines.md). +- `PackerArgs` must validate against the plugin's arguments model. + +### Using `flush` and `feeder_active` + +Deciding when to pack a small remainder is the plugin's job, not the core's. While `feeder_active` is true, more inputs may arrive, so a packer that groups inputs can hold back an incomplete group. Once the feeder is off, or when `flush` is set, it should pack what is left, otherwise the transformation never finishes. `input_overview` is there for packers whose thresholds depend on how much of the pool is left. + +## Writing and registering a packer + +A packer that groups inputs by run number: + +```python +from collections import defaultdict +from collections.abc import AsyncGenerator, AsyncIterator + +from pydantic import BaseModel, PositiveInt + +from diracx.core.models.transformation import ( + BadInput, + ComputeParcel, + InputStatus, + PackerInput, +) + + +class ByRunPackerArgs(BaseModel): + group_size: PositiveInt = 10 + + +async def by_run_packer( + *, + packer_args: ByRunPackerArgs, + inputs: AsyncIterator[PackerInput], + flush: bool = False, + feeder_active: bool, + input_overview: dict[InputStatus, int], +) -> AsyncGenerator[ComputeParcel | BadInput, None]: + groups = defaultdict(list) + async for packer_input in inputs: + run = packer_input.descriptor.get("Run") + if run is None: + yield BadInput(input_id=packer_input.input_id, reason="No run number") + continue + groups[run].append(packer_input.input_id) + if len(groups[run]) == packer_args.group_size: + yield ComputeParcel(inputs=groups.pop(run), metadata={"Run": run}) + if flush or not feeder_active: + for run, group in groups.items(): + yield ComputeParcel(inputs=group, metadata={"Run": run}) +``` + +Register it in your extension's `pyproject.toml`: + +```toml +[project.entry-points."diracx.transformations.packers"] +ByRun = "myext.logic.transformations.packers:by_run_packer" +``` + +A transformation then uses it with `Packer = "ByRun"` and, for example, `PackerArgs = {"group_size": 20}`. If several installed extensions register the same name, the one with the highest extension priority wins, so an extension can replace a packer shipped by DiracX. + +A synchronous generator function also works. Since it cannot wait for the database between pages, it is offered all the claimed inputs at once, which `DIRACX_TASKS_PACKER_MAX_INPUTS_PER_INVOCATION` keeps bounded. + +DiracX ships one packer, `FixedSize`, which makes parcels of `group_size` inputs in the order they are offered and holds back a smaller last group until the feeder is off or a flush is requested. + +## Configuration + +The tasks are configured with environment variables, listed in the [environment variables reference](../../../admin/reference/env-variables.md#packersettings): + +- `DIRACX_TASKS_PACKER_ENABLED` turns the periodic sweep on. It is off by default. +- `DIRACX_TASKS_PACKER_INTERVAL_SECONDS` sets how often the sweep runs. +- `DIRACX_TASKS_PACKER_MAX_INPUTS_PER_INVOCATION` caps the inputs offered to a packer in one run. + +## Current limitations + +- `DataParcel` proposals and transformations of the Data kind raise `NotImplementedError`. +- The journalled counters of DX-ADR-009 are not written yet, and `input_overview` is computed with a `GROUP BY` over `Inputs` ([#1085](https://github.com/DIRACGrid/diracx/issues/1085)). +- Plugins receive only the arguments above. Injecting databases and settings into them, as DX-ADR-006 describes, is not implemented yet. diff --git a/mkdocs.yml b/mkdocs.yml index 971d3a15e..c189a7af9 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -254,6 +254,8 @@ nav: - dev/explanations/tasks/index.md - Class details: dev/explanations/tasks/class-details.md - Broker lifecycle: dev/explanations/tasks/broker-lifecycle.md + - Transformation System: + - Packer: dev/explanations/transformation-system/packer.md - Frontend: - CLI: - API: dev/explanations/components/api.md From d8d6ad08731d71373038a6a08bc7ce14b71cdfbb Mon Sep 17 00:00:00 2001 From: Marco Mascheroni Date: Tue, 20 Oct 2026 16:45:27 +0200 Subject: [PATCH 6/8] docs: note that failed packer runs commit until #1101 is fixed The task worker commits an injected DB transaction even when the task raises, so a packer run that fails part-way keeps its earlier writes. Point to the issue where the packer task and pack_transformation claim a rollback. Refs #1083, #1101 Co-Authored-By: Claude Opus 5.5 --- diracx-logic/src/diracx/logic/transformations/packing.py | 4 +++- diracx-tasks/src/diracx/tasks/transformations/packer.py | 3 +++ 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/diracx-logic/src/diracx/logic/transformations/packing.py b/diracx-logic/src/diracx/logic/transformations/packing.py index 73cb450ba..587455ab3 100644 --- a/diracx-logic/src/diracx/logic/transformations/packing.py +++ b/diracx-logic/src/diracx/logic/transformations/packing.py @@ -2,7 +2,9 @@ The packer plugin only proposes; this module checks every proposal and is the only thing that writes. Everything happens inside the caller's transaction, so -a proposal that fails a check rolls back the whole invocation. +a proposal that fails a check rolls back the whole invocation, provided the +caller rolls back on an exception. The task worker does not yet +(DIRACGrid/diracx#1101). """ from __future__ import annotations diff --git a/diracx-tasks/src/diracx/tasks/transformations/packer.py b/diracx-tasks/src/diracx/tasks/transformations/packer.py index 0c863d2f5..135a18e8d 100644 --- a/diracx-tasks/src/diracx/tasks/transformations/packer.py +++ b/diracx-tasks/src/diracx/tasks/transformations/packer.py @@ -89,6 +89,9 @@ def execution_locks(self) -> list[BaseLock]: return [MutexLock(TRANSFORMATION, self.transformation_id)] async def execute(self, transformation_db: TransformationDB) -> dict[str, Any]: + # The worker commits an injected transaction even when the task raises + # (DIRACGrid/diracx#1101), so until that is fixed a failed run keeps the + # writes it made before the failure instead of rolling them back. summary = await pack_transformation( transformation_db, UUID(self.transformation_id), From 2950ea31fbe573d9ecf2c5721f06de74fc5f384d Mon Sep 17 00:00:00 2001 From: Marco Mascheroni Date: Wed, 21 Oct 2026 16:07:16 +0200 Subject: [PATCH 7/8] Rename RequirementsID to RequirementsTemplateID Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- diracx-db/src/diracx/db/sql/transformation/db.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/diracx-db/src/diracx/db/sql/transformation/db.py b/diracx-db/src/diracx/db/sql/transformation/db.py index 06cc69e7e..17e0c86ee 100644 --- a/diracx-db/src/diracx/db/sql/transformation/db.py +++ b/diracx-db/src/diracx/db/sql/transformation/db.py @@ -296,7 +296,7 @@ async def insert_compute_parcels( { "ParcelID": parcel.parcel_id, "ProcessID": process_id, - "RequirementsID": parcel.requirements_id, + "RequirementsTemplateID": parcel.requirements_id, "Parameters": parcel.parameters or None, } for parcel in parcels From b57e364888708831b5ecfa3a9ded08110e1a6526 Mon Sep 17 00:00:00 2001 From: Marco Mascheroni Date: Thu, 22 Oct 2026 13:08:21 +0200 Subject: [PATCH 8/8] fix(db): use the RequirementsTemplates table renamed in #1099 The merge with main brought in the rename of Requirements to RequirementsTemplates, which TransformationDB and the packer tests still imported under the old name, so importing diracx.db.sql failed. Refs #1083 --- .../src/diracx/db/sql/transformation/db.py | 20 ++++++++++++------- .../transformation/test_transformation_db.py | 12 ++++++----- .../tests/transformations/test_packing.py | 14 +++++++------ 3 files changed, 28 insertions(+), 18 deletions(-) diff --git a/diracx-db/src/diracx/db/sql/transformation/db.py b/diracx-db/src/diracx/db/sql/transformation/db.py index 17e0c86ee..6cfda50c3 100644 --- a/diracx-db/src/diracx/db/sql/transformation/db.py +++ b/diracx-db/src/diracx/db/sql/transformation/db.py @@ -31,7 +31,7 @@ Inputs, ParcelInputs, Parcels, - Requirements, + RequirementsTemplates, TransformationLog, Transformations, ) @@ -230,10 +230,13 @@ async def get_compute_facet( facet. """ stmt = ( - select(ComputeTransformations.process_id, Requirements.requirements) + select( + ComputeTransformations.process_id, + RequirementsTemplates.requirements_template, + ) .outerjoin( - Requirements, - Requirements.requirements_id + RequirementsTemplates, + RequirementsTemplates.requirements_template_id == ComputeTransformations.requirements_template_id, ) .where(ComputeTransformations.transformation_id == transformation_id) @@ -241,7 +244,7 @@ async def get_compute_facet( row = (await self.conn.execute(stmt)).one_or_none() if row is None: raise TransformationNotFoundError(transformation_id, "no compute facet") - return row.ProcessID, row.Requirements or {} + return row.ProcessID, row.RequirementsTemplate or {} async def insert_requirements( self, requirements_id: bytes, requirements: dict[str, Any] @@ -253,10 +256,13 @@ async def insert_requirements( requirements: The requirement set. """ await self.conn.execute( - insert(Requirements) + insert(RequirementsTemplates) .prefix_with("IGNORE", dialect="mysql") .prefix_with("OR IGNORE", dialect="sqlite") - .values(requirements_id=requirements_id, requirements=requirements) + .values( + requirements_template_id=requirements_id, + requirements_template=requirements, + ) ) async def insert_compute_parcels( diff --git a/diracx-db/tests/transformation/test_transformation_db.py b/diracx-db/tests/transformation/test_transformation_db.py index 7af3aa10c..bd1aec344 100644 --- a/diracx-db/tests/transformation/test_transformation_db.py +++ b/diracx-db/tests/transformation/test_transformation_db.py @@ -31,7 +31,7 @@ ParcelInputs, Parcels, Processes, - Requirements, + RequirementsTemplates, TransformationLog, Transformations, Workgraphs, @@ -260,8 +260,8 @@ async def compute_facet(transformation_db: TransformationDB, transformation_id: insert(Processes).values(process_id=PROCESS_ID, process={"class": "Tool"}) ) await db.conn.execute( - insert(Requirements).values( - requirements_id=TEMPLATE_ID, requirements={"cpu": 1} + insert(RequirementsTemplates).values( + requirements_template_id=TEMPLATE_ID, requirements_template={"cpu": 1} ) ) await db.conn.execute( @@ -396,7 +396,9 @@ async def test_insert_requirements_deduplicates(transformation_db: Transformatio await db.insert_requirements(b"\x01" * 32, {"cpu": 4}) async with transformation_db as db: await db.insert_requirements(b"\x01" * 32, {"cpu": 4}) - count = await db.conn.scalar(select(func.count()).select_from(Requirements)) + count = await db.conn.scalar( + select(func.count()).select_from(RequirementsTemplates) + ) assert count == 1 @@ -430,7 +432,7 @@ async def test_insert_compute_parcels( select( ComputeParcels.parcel_id, ComputeParcels.process_id, - ComputeParcels.requirements_id, + ComputeParcels.requirements_template_id, ComputeParcels.parameters, ).order_by(ComputeParcels.parcel_id) ) diff --git a/diracx-logic/tests/transformations/test_packing.py b/diracx-logic/tests/transformations/test_packing.py index 47ef3e2ee..5b4450262 100644 --- a/diracx-logic/tests/transformations/test_packing.py +++ b/diracx-logic/tests/transformations/test_packing.py @@ -26,7 +26,7 @@ ParcelInputs, Parcels, Processes, - Requirements, + RequirementsTemplates, TransformationLog, Transformations, Workgraphs, @@ -95,8 +95,8 @@ async def transformation_id(transformation_db: TransformationDB) -> UUID: insert(Processes).values(process_id=PROCESS_ID, process={"class": "Tool"}) ) await db.conn.execute( - insert(Requirements).values( - requirements_id=template_id, requirements=TEMPLATE + insert(RequirementsTemplates).values( + requirements_template_id=template_id, requirements_template=TEMPLATE ) ) await db.conn.execute( @@ -264,17 +264,19 @@ async def plugin(*, packer_args, inputs, **_): rows = ( await db.conn.execute( select( - Requirements.requirements, + RequirementsTemplates.requirements_template, ComputeParcels.process_id, ComputeParcels.parameters, Parcels.annotations, ) - .join(Requirements) + .join(RequirementsTemplates) .join(Parcels) .order_by(Parcels.parcel_id) ) ).all() - n_requirements = len((await db.conn.execute(select(Requirements))).all()) + n_requirements = len( + (await db.conn.execute(select(RequirementsTemplates))).all() + ) merged = {"cpu": 8, "platform": "x86_64"} assert rows == [ (merged, PROCESS_ID, {"seed": 1}, {"note": "big"}),