i15-1: Automatically add background scans to queue

pyproject.toml

@@ -16,6 +16,7 @@ dependencies = [
     "blueapi>=1.18.0",
     "fastapi>=0.136.0",
     "pydantic>=2.13.2",
+    "tiled>=0.2.9",
 ]
 dynamic = ["version"]
 license.file = "LICENSE"

src/daq_queuing_service/api/errors.py

@@ -1,6 +1,10 @@
+from collections.abc import Awaitable, Callable
+from typing import TypeVar
+
 from fastapi import FastAPI, Request
 from fastapi.responses import JSONResponse
 
+from daq_queuing_service.log import LOGGER
 from daq_queuing_service.plugins.converter import ConverterError, ValidateError
 from daq_queuing_service.task_queue.queue_utils import (
     NegativePositionError,
@@ -12,59 +16,56 @@
 
 # pyright: reportUnusedFunction=false
 
+E = TypeVar("E", bound=Exception)
 
-def register_exception_handlers(app: FastAPI):
-    @app.exception_handler(TaskInProgressError)
-    async def task_in_progress_handler(
-        request: Request, exception: TaskInProgressError
-    ):
-        return JSONResponse(
-            status_code=409,
-            content={"error": "task_in_progress", "message": str(exception)},
-        )
+Handler = Callable[[Request, E], Awaitable[JSONResponse]]
 
-    @app.exception_handler(TaskNotFoundError)
-    async def task_not_found_handler(request: Request, exception: TaskNotFoundError):
-        return JSONResponse(
-            status_code=404,
-            content={"error": "task_not_found", "message": str(exception)},
-        )
 
-    @app.exception_handler(TaskNotInQueueError)
-    async def task_not_in_queue_handler(
-        request: Request, exception: TaskNotInQueueError
-    ):
+def make_exception_handler(
+    status_code: int, error_code: str
+) -> Callable[[Request, Exception], Awaitable[JSONResponse]]:
+    async def handler(request: Request, exception: Exception):
+        LOGGER.exception("Error while handling request: %s", request)
         return JSONResponse(
-            status_code=409,
-            content={"error": "task_not_in_queue", "message": str(exception)},
+            status_code=status_code,
+            content={"error": error_code, "message": str(exception)},
         )
 
-    @app.exception_handler(NegativePositionError)
-    async def negative_position_handler(
-        request: Request, exception: NegativePositionError
-    ):
-        return JSONResponse(
-            status_code=400,
-            content={"error": "negative_position", "message": str(exception)},
-        )
+    return handler
 
-    @app.exception_handler(QueueError)
-    async def queue_error_handler(request: Request, exception: QueueError):
-        return JSONResponse(
-            status_code=409,
-            content={"error": "queue_error", "message": str(exception)},
-        )
 
-    @app.exception_handler(ValidateError)
-    async def validation_error_handler(request: Request, exception: ValidateError):
-        return JSONResponse(
-            status_code=422,
-            content={"error": "validation_error", "message": str(exception)},
-        )
+def register_exception_handlers(app: FastAPI):
+    app.add_exception_handler(
+        TaskInProgressError,
+        make_exception_handler(409, "task_in_progress"),
+    )
 
-    @app.exception_handler(ConverterError)
-    async def converter_error_handler(request: Request, exception: ConverterError):
-        return JSONResponse(
-            status_code=422,
-            content={"error": "converter_error", "message": str(exception)},
-        )
+    app.add_exception_handler(
+        TaskNotFoundError,
+        make_exception_handler(404, "task_not_found"),
+    )
+
+    app.add_exception_handler(
+        TaskNotInQueueError,
+        make_exception_handler(409, "task_not_in_queue"),
+    )
+
+    app.add_exception_handler(
+        NegativePositionError,
+        make_exception_handler(400, "negative_position"),
+    )
+
+    app.add_exception_handler(
+        QueueError,
+        make_exception_handler(409, "queue_error"),
+    )
+
+    app.add_exception_handler(
+        ValidateError,
+        make_exception_handler(422, "validation_error"),
+    )
+
+    app.add_exception_handler(
+        ConverterError,
+        make_exception_handler(422, "converter_error"),
+    )

src/daq_queuing_service/plugins/i15_1/backgrounds.py

@@ -0,0 +1,24 @@
+from typing import Literal
+
+from pydantic import BaseModel, ConfigDict
+
+# This should be generated from the json schema
+# https://github.com/DiamondLightSource/daq-queuing-service/issues/78
+BACKGROUND_TYPES = Literal["air", "bs", "fq", "pi"]
+
+
+class BackgroundInfo(BaseModel):
+    # Currently only room temperatures scans are supported
+    # https://github.com/DiamondLightSource/daq-queuing-service/issues/84
+    model_config = ConfigDict(frozen=True)
+    bg_type: BACKGROUND_TYPES
+
+    def add_tiled_id(self, tiled_id: str) -> "TiledBackground":
+        return TiledBackground(
+            bg_type=self.bg_type,
+            tiled_id=tiled_id,
+        )
+
+
+class TiledBackground(BackgroundInfo):
+    tiled_id: str

src/daq_queuing_service/plugins/i15_1/i15_1_converter.py

@@ -0,0 +1,190 @@
+from typing import Any
+
+from blueapi.service.model import TaskRequest
+from tiled.client import from_uri  # type: ignore
+from tiled.client.container import Container
+
+from daq_queuing_service.blueapi_interaction.blueapi_call import BlueapiCall
+from daq_queuing_service.log import LOGGER
+from daq_queuing_service.plugins.converter import Converter
+from daq_queuing_service.plugins.i15_1.backgrounds import BackgroundInfo
+from daq_queuing_service.plugins.i15_1.tiled_interaction import get_background_tiled_id
+from daq_queuing_service.task_queue.task import (
+    Experiment,
+    ExperimentDefinition,
+    Sample,
+    Task,
+    TaskWithPosition,
+)
+
+BACKGROUND_SCAN = "Background"
+
+
+class I151Converter(Converter):
+    def __init__(self):
+        self.tiled_client: Container = from_uri("https://tiled.diamond.ac.uk/api/v1")
+
+    def pre_process(
+        self,
+        queue: list[Task],
+        history: list[TaskWithPosition],
+        call_history: list[BlueapiCall],
+    ) -> list[Task]:
+        return self._add_required_background_scans(queue)
+
+    def construct_blueapi_calls(
+        self,
+        queue: list[TaskWithPosition],
+        history: list[TaskWithPosition],
+        call_history: list[BlueapiCall],
+    ) -> list[BlueapiCall]:
+
+        call_list: list[BlueapiCall] = []
+
+        for task in queue:
+            match task.experiment:
+                case TaskRequest():
+                    call_list.append(
+                        BlueapiCall(
+                            task_request=task.experiment, parent_task_id=task.id
+                        )
+                    )
+                case Experiment():
+                    call_list.extend(
+                        [
+                            BlueapiCall(task_request=b_api_task, parent_task_id=task.id)
+                            for b_api_task in (
+                                self._construct_blueapi_tasks_from_experiment(
+                                    task.experiment
+                                )
+                            )
+                        ]
+                    )
+        return call_list
+
+    def _construct_blueapi_tasks_from_experiment(
+        self,
+        experiment: Experiment,
+    ) -> list[TaskRequest]:
+        sample_name = experiment.sample.name
+        # Assume sample name is of form test_8_1 to load from position 8 on puck 1
+        _, position, puck = sample_name.split("_")
+
+        # For air calibration scans, we need to not to robot load/unload.
+        # https://github.com/DiamondLightSource/daq-queuing-service/issues/83
+        return [
+            TaskRequest(
+                name="robot_load",
+                params={"puck": puck, "position": position},
+                instrument_session=experiment.instrument_session,
+            ),
+            TaskRequest(
+                name="centre_sample",
+                params={
+                    "start_z": -20,
+                    "end_z": 0,
+                    "steps": 20,
+                    "exposure_time": 0.01,
+                    "metadata": {
+                        "sample": experiment.sample,
+                        # This will include tiled background scan info
+                        "experiment_definition": experiment.experiment_definition,
+                    },
+                },
+                instrument_session=experiment.instrument_session,
+            ),
+            TaskRequest(
+                name="robot_unload",
+                params={},
+                instrument_session=experiment.instrument_session,
+            ),
+        ]
+
+    def _add_required_background_scans(self, tasks: list[Task]) -> list[Task]:
+        """Adds background scan tasks to the queue. Backgrounds will be added directly
+        in front of the first task in the queue that requires them.
+
+        Args:
+            tasks (list[Task]): Current list of tasks
+
+        Returns:
+            list[Task]: New list of tasks including backgrounds
+        """
+        LOGGER.info("Adding required background scans")
+
+        # This can be made more robust https://github.com/DiamondLightSource/daq-queuing-service/issues/80
+        new_tasks: list[Task] = []
+
+        for task in tasks:
+            experiment = task.experiment
+            if (
+                isinstance(experiment, Experiment)
+                and experiment.name != BACKGROUND_SCAN
+            ):
+                instrument_session = experiment.instrument_session
+                backgrounds = self._get_required_backgrounds(experiment)
+
+                for background in backgrounds:
+                    if tiled_id := get_background_tiled_id(
+                        self.tiled_client,
+                        background,
+                        instrument_session,
+                    ):
+                        self._add_tiled_background_to_md(
+                            experiment.experiment_definition.data, tiled_id, background
+                        )
+
+                    else:
+                        bg_experiment = self._construct_background_experiment(
+                            background, instrument_session
+                        )
+                        new_tasks.append(Task(experiment=bg_experiment))
+
+            new_tasks.append(task)
+        return self._remove_repeated_backgrounds(new_tasks)
+
+    def _remove_repeated_backgrounds(self, tasks: list[Task]) -> list[Task]:
+        LOGGER.info("Removing repeated background scans")
+        new_tasks: list[Task] = []
+        queued_background_experiments: list[Experiment] = []
+
+        for task in tasks:
+            if task.experiment.name != BACKGROUND_SCAN:
+                new_tasks.append(task)
+            elif task.experiment not in queued_background_experiments:
+                assert isinstance(task.experiment, Experiment)
+                queued_background_experiments.append(task.experiment)
+                new_tasks.append(task)
+            else:
+                LOGGER.debug(f"Removing repeated background scan: {task.experiment}")
+        return new_tasks
+
+    def _get_required_backgrounds(self, experiment: Experiment) -> list[BackgroundInfo]:
+        # This should be fleshed out https://github.com/DiamondLightSource/daq-queuing-service/issues/79
+        return [BackgroundInfo(bg_type="fq")]
+
+    def _add_tiled_background_to_md(
+        self, params: dict[str, Any], tiled_id: str, background: BackgroundInfo
+    ):
+        LOGGER.debug("Adding background scan tiled info to metadata")
+        if metadata := params.get("metadata"):
+            if tiled_backgrounds := metadata.get("tiled_backgrounds"):
+                tiled_backgrounds[tiled_id] = background
+            else:
+                metadata["tiled_backgrounds"] = {tiled_id: background}
+        else:
+            params["metadata"] = {"tiled_backgrounds": {tiled_id: background}}
+
+    def _construct_background_experiment(
+        self, background: BackgroundInfo, instrument_session: str
+    ) -> Experiment:
+        LOGGER.debug(f"Constructing experiment for background: {background}")
+        return Experiment(
+            name=BACKGROUND_SCAN,
+            instrument_session=instrument_session,
+            # Need to get sample info for test samples (air, empty capillary etc)
+            sample=Sample(name="fq_1_1", id="", data={}),
+            experiment_definition=ExperimentDefinition(
+                name="background_scan", id="", data={"background": background}
+            ),
+        )

src/daq_queuing_service/plugins/i15_1/tiled_interaction.py

@@ -0,0 +1,58 @@
+from cachetools import TTLCache, cached
+from tiled.client.container import Container
+from tiled.queries import Eq
+
+from daq_queuing_service.log import LOGGER
+from daq_queuing_service.plugins.i15_1.backgrounds import BackgroundInfo
+
+# Ignoring the following rules as the tiled client is poorly typed as scares the linter
+# pyright: reportUnknownMemberType=false
+# pyright: reportUnknownVariableType=false
+# pyright: reportUnknownArgumentType=false
+# pyright: reportUnknownLambdaType=false
+
+cache: TTLCache[tuple[BackgroundInfo, str], str | None] = TTLCache(maxsize=100, ttl=1)
+
+
+def get_background_tiled_id(
+    tiled_client: Container,
+    required_background: BackgroundInfo,
+    instrument_session: str,
+) -> str | None:
+
+    @cached(cache)
+    def _get_background_tiled_id(
+        required_background: BackgroundInfo, instrument_session: str
+    ) -> str | None:
+
+        result: Container = (
+            tiled_client.search(Eq("start.instrument_session", instrument_session))
+            .search(Eq("start.instrument", "i15-1"))
+            .search(
+                Eq(
+                    "start.experiment_definition.metadata.background",
+                    required_background.model_dump_json(),
+                )
+            )
+        )
+
+        if not len(result):
+            LOGGER.debug(
+                f"Found no scans in tiled matching background: {required_background}"
+            )
+            return
+
+        items = sorted(
+            ((key, value) for key, value in result.items()),
+            key=lambda item: item[1].metadata["start"]["time"],
+        )
+
+        # return the tiled ID
+        tiled_id = items[-1][0]
+        LOGGER.debug(
+            f"Found {len(items)} scans in tiled matching background: "
+            + f"{required_background}. Returning the first: {tiled_id}"
+        )
+        return tiled_id
+
+    return _get_background_tiled_id(required_background, instrument_session)