Rework after doing the same split as poc

Author: lars-petter-hauge

.github/workflows/python.yaml

@@ -18,8 +18,6 @@ jobs:
 
     steps:
     - uses: actions/checkout@v4
-      with:
-        sparse-checkout: backend
 
     - uses: actions/setup-python@v5
       with:
@@ -28,17 +26,62 @@ jobs:
     - uses: astral-sh/setup-uv@v5
 
     - name: Install backend
-      run: uv sync
+      run: uv sync --package acidwatch-api
 
     - name: Run ruff
       run: |
-        uv run ruff format --check
-        uv run ruff check
+        uv run --package acidwatch-api ruff format --check src tests packages ../workers
+        uv run --package acidwatch-api ruff check src tests packages ../workers
 
     - name: Run mypy
-      run: uv run mypy --strict src
+      run: uv run --package acidwatch-api mypy --strict src packages/acidwatch-models/src packages/acidwatch-messaging/src
 
     - name: Run pytest
       run: |
         cp .env.example .env
-        uv run pytest tests
+        uv run --package acidwatch-api pytest tests
+
+  worker-packages:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        include:
+          - package: acidwatch-worker-example
+            test_path: workers/example/tests
+          - package: acidwatch-worker-tocomo
+            test_path: workers/tocomo/tests
+          - package: acidwatch-worker-arcs
+            test_path: ""
+          - package: acidwatch-worker-arcs-exp
+            test_path: workers/arcs-exp/tests
+          - package: acidwatch-worker-solubilityccs
+            test_path: ""
+          - package: acidwatch-worker-gibbs-minimization
+            test_path: workers/gibbs-minimization/tests
+          - package: acidwatch-worker-phpitz-reactive
+            test_path: ""
+          - package: acidwatch-worker-phpitz-solubility
+            test_path: ""
+
+    steps:
+    - uses: actions/checkout@v4
+
+    - uses: actions/setup-python@v5
+      with:
+        python-version: '3.12'
+
+    - uses: astral-sh/setup-uv@v5
+
+    - name: Install worker
+      run: uv sync --package ${{ matrix.package }}
+
+    - name: Build worker package
+      run: uv build --package ${{ matrix.package }}
+
+    - name: Install Tocomo calculation package
+      if: matrix.package == 'acidwatch-worker-tocomo'
+      run: uv pip install --no-deps "git+https://github.com/equinor/tocomo@74e23131bd27d63dcf71b4a62eb81cdb76d981c9#subdirectory=backend"
+
+    - name: Test worker
+      if: matrix.test_path != ''
+      run: uv run --package ${{ matrix.package }} pytest ${{ matrix.test_path }}

backend/packages/acidwatch-messaging/pyproject.toml

@@ -0,0 +1,14 @@
+[project]
+name = "acidwatch-messaging"
+version = "0.1.0"
+description = "Broker transports and worker runtime for AcidWatch models"
+requires-python = ">=3.12, <3.14"
+dependencies = [
+  "acidwatch-models",
+  "aio-pika>=10.0.1",
+  "azure-servicebus>=7.14.2,<8.0.0",
+]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"

backend/packages/acidwatch-messaging/src/acidwatch_messaging/__init__.py

@@ -0,0 +1,39 @@
+from .contracts import AdapterJob, AdapterResult, Heartbeat
+from .queues import (
+    DEAD_LETTER_QUEUE,
+    HEARTBEATS_QUEUE,
+    RESULTS_QUEUE,
+    job_queue_name,
+)
+from .transport import (
+    AzureServiceBusTransport,
+    Message,
+    RabbitMQTransport,
+    Transport,
+    create_transport,
+)
+from .worker import (
+    AdapterWorker,
+    run_adapter_job,
+    run_worker,
+    run_worker_from_environment,
+)
+
+__all__ = [
+    "AdapterJob",
+    "AdapterResult",
+    "AdapterWorker",
+    "AzureServiceBusTransport",
+    "DEAD_LETTER_QUEUE",
+    "HEARTBEATS_QUEUE",
+    "Heartbeat",
+    "Message",
+    "RESULTS_QUEUE",
+    "RabbitMQTransport",
+    "Transport",
+    "create_transport",
+    "job_queue_name",
+    "run_adapter_job",
+    "run_worker",
+    "run_worker_from_environment",
+]

backend/packages/acidwatch-messaging/src/acidwatch_messaging/contracts.py

@@ -0,0 +1,32 @@
+from datetime import datetime
+from uuid import UUID
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from acidwatch_models.datamodel import AnyPanel, Conditions, Phase
+
+
+class _Message(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+
+class AdapterJob(_Message):
+    model_input_id: UUID
+    model_id: str
+    concentrations: dict[str, int | float]
+    parameters: dict[str, bool | float | int | str]
+    conditions: Conditions
+
+
+class AdapterResult(_Message):
+    model_input_id: UUID
+    phases: list[Phase] = Field(default_factory=list)
+    panels: list[AnyPanel] = Field(default_factory=list)
+    error: str | None = None
+
+
+class Heartbeat(_Message):
+    model_id: str
+    instance_id: str
+    timestamp: datetime
+    job_id: str | None = None

…ckend/src/acidwatch_api/broker/queues.py …saging/src/acidwatch_messaging/queues.pybackend/src/acidwatch_api/broker/queues.py renamed to backend/packages/acidwatch-messaging/src/acidwatch_messaging/queues.py backend/src/acidwatch_api/broker/queues.py renamed to backend/packages/acidwatch-messaging/src/acidwatch_messaging/queues.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import json
+import asyncio
 from collections.abc import AsyncIterator
 from types import TracebackType
 from typing import Any, Protocol, Self
@@ -9,8 +10,18 @@
 from azure.servicebus import ServiceBusMessage
 from azure.servicebus.aio import AutoLockRenewer, ServiceBusClient
 from pamqp.common import FieldTable
+from pydantic import BaseModel
 
-from acidwatch_api.broker.queues import DEAD_LETTER_QUEUE, RESULTS_QUEUE
+from acidwatch_messaging.queues import DEAD_LETTER_QUEUE, RESULTS_QUEUE
+
+MESSAGE_LOCK_RENEWAL_SECONDS = 7200
+Payload = BaseModel | dict[str, Any]
+
+
+def _payload_json(payload: Payload) -> str:
+    if isinstance(payload, BaseModel):
+        return payload.model_dump_json()
+    return json.dumps(payload)
 
 
 class Message(Protocol):
@@ -22,7 +33,7 @@ async def reject(self) -> None: ...
 
 
 class Transport(Protocol):
-    async def publish(self, queue_name: str, payload: dict[str, Any]) -> None: ...
+    async def publish(self, queue_name: str, payload: Payload) -> None: ...
 
     def subscribe(self, queue_name: str) -> AsyncIterator[Message]: ...
 
@@ -85,13 +96,13 @@ async def _declare_queues(
             )
 
     async def publish(
-        self, queue_name: str, payload: dict[str, Any]
+        self, queue_name: str, payload: Payload
     ) -> None:
         if self._publish_channel is None:
             raise RuntimeError("Transport is not open")
         await self._publish_channel.default_exchange.publish(
             aio_pika.Message(
-                body=json.dumps(payload).encode(),
+                body=_payload_json(payload).encode(),
                 delivery_mode=aio_pika.DeliveryMode.PERSISTENT,
                 content_type="application/json",
             ),
@@ -131,6 +142,8 @@ class AzureServiceBusTransport:
     def __init__(self, connection_string: str, queue_names: list[str]):
         self._connection_string = connection_string
         self._client: ServiceBusClient | None = None
+        self._senders: dict[str, Any] = {}
+        self._sender_locks: dict[str, asyncio.Lock] = {}
 
     async def __aenter__(self) -> Self:
         self._client = ServiceBusClient.from_connection_string(
@@ -145,18 +158,26 @@ async def __aexit__(
         traceback: TracebackType | None,
     ) -> None:
         if self._client is not None:
+            for sender in self._senders.values():
+                await sender.close()
+            self._senders.clear()
+            self._sender_locks.clear()
             await self._client.close()
 
     async def publish(
-        self, queue_name: str, payload: dict[str, Any]
+        self, queue_name: str, payload: Payload
     ) -> None:
         if self._client is None:
             raise RuntimeError("Transport is not open")
-        sender = self._client.get_queue_sender(queue_name)
-        async with sender:
+        sender = self._senders.get(queue_name)
+        if sender is None:
+            sender = self._client.get_queue_sender(queue_name)
+            self._senders[queue_name] = sender
+            self._sender_locks[queue_name] = asyncio.Lock()
+        async with self._sender_locks[queue_name]:
             await sender.send_messages(
                 ServiceBusMessage(
-                    json.dumps(payload),
+                    _payload_json(payload),
                     content_type="application/json",
                 )
             )
@@ -169,8 +190,21 @@ async def subscribe(self, queue_name: str) -> AsyncIterator[Message]:
             prefetch_count=1,
         )
         async with receiver, AutoLockRenewer(
-            max_lock_renewal_duration=7200
+            max_lock_renewal_duration=MESSAGE_LOCK_RENEWAL_SECONDS
         ) as renewer:
             async for message in receiver:
                 renewer.register(receiver, message)
                 yield AzureServiceBusMessage(message, receiver)
+
+
+def create_transport(
+    broker_url: str,
+    queue_names: list[str],
+    backend: str = "",
+) -> RabbitMQTransport | AzureServiceBusTransport:
+    selected = backend.strip().lower()
+    if selected in {"azure_service_bus", "servicebus"} or (
+        not selected and broker_url.startswith("Endpoint=")
+    ):
+        return AzureServiceBusTransport(broker_url, queue_names)
+    return RabbitMQTransport(broker_url, queue_names)