Ensure background order

src/daq_queuing_service/plugins/i15_1/i15_1_converter.py

@@ -40,6 +40,8 @@
     ExperimentDefinition,
     Sample,
     Task,
+    TaskID,
+    TaskRules,
     TaskWithPosition,
 )
 
@@ -52,9 +54,11 @@
 )
 
 
-def _filter_auxiliary_scans(tasks: list[Task]) -> list[tuple[int, AuxiliaryScan]]:
+def _filter_auxiliary_scans(
+    tasks: list[Task],
+) -> list[tuple[int, TaskID, AuxiliaryScan]]:
     return [
-        (i, AuxiliaryScan.from_experiment(task.experiment))
+        (i, task.id, AuxiliaryScan.from_experiment(task.experiment))
         for i, task in enumerate(tasks)
         if isinstance(task.experiment, Experiment)
         and is_auxiliary_experiment(task.experiment)
@@ -63,9 +67,8 @@ def _filter_auxiliary_scans(tasks: list[Task]) -> list[tuple[int, AuxiliaryScan]
 
 class I151Converter(Converter):
     def __init__(self):
-        # First key is the ID of the task using the auxiliary scan
         self._tiled_auxiliary_scans: dict[
-            str, dict[AuxiliaryScanType, TiledAuxiliary]
+            TaskID, dict[AuxiliaryScanType, TiledAuxiliary]
         ] = {}
         self._config_client = ConfigClient.from_url(CONFIG_SERVER_URL)
 
@@ -144,7 +147,7 @@ def pause_on_fail_blueapi_call(
         return True
 
     def _construct_blueapi_tasks_from_experiment(
-        self, experiment: Experiment, task_id: str
+        self, experiment: Experiment, task_id: TaskID
     ) -> list[TaskRequest]:
         LOGGER.debug(f"Converting to blueapi calls, experiment = {experiment}")
 
@@ -214,7 +217,7 @@ def _add_required_auxiliary_scans(
                             required_auxiliary,
                             current_task,
                             new_tasks,
-                            task.id,
+                            task,
                             instrument_session,
                         )
 
@@ -226,7 +229,7 @@ def _ensure_auxiliary_in_queue_or_tiled(
         required_auxiliary: AuxiliaryScan,
         current_task: TaskWithPosition | None,
         new_tasks: list[Task],
-        task_id: str,
+        task: Task,
         instrument_session: str,
     ) -> list[Task]:
         if (
@@ -237,21 +240,27 @@ def _ensure_auxiliary_in_queue_or_tiled(
                 required_auxiliary
             )
         ):
+            task.rules.run_after.add(current_task.id)
             return new_tasks
 
         queued_auxiliary_scans = _filter_auxiliary_scans(new_tasks)
-        if any(
-            queued_auxiliary.is_suitable(required_auxiliary)
-            for _, queued_auxiliary in queued_auxiliary_scans
+        if suitable_task_id := next(
+            (
+                task_id
+                for _, task_id, auxiliary_scan in queued_auxiliary_scans
+                if auxiliary_scan.is_suitable(required_auxiliary)
+            ),
+            None,
         ):
+            task.rules.run_after.add(suitable_task_id)
             return new_tasks
 
         if tiled_scan := get_suitable_tiled_scan(
             self._tiled_client,
             required_auxiliary,
             self._queue_params.tiled_stale_time_s,
         ):
-            self._tiled_auxiliary_scans[task_id][tiled_scan.kind] = tiled_scan
+            self._tiled_auxiliary_scans[task.id][tiled_scan.kind] = tiled_scan
             return new_tasks
 
         LOGGER.info(
@@ -261,6 +270,7 @@ def _ensure_auxiliary_in_queue_or_tiled(
         return self._add_or_replace_auxiliary_scan(
             required_auxiliary,
             new_tasks,
+            task,
             queued_auxiliary_scans,
             instrument_session,
         )
@@ -269,26 +279,38 @@ def _add_or_replace_auxiliary_scan(
         self,
         required_auxiliary: AuxiliaryScan,
         new_tasks: list[Task],
-        queued_auxiliary_scans: list[tuple[int, AuxiliaryScan]],
+        task: Task,
+        queued_auxiliary_scans: list[tuple[int, TaskID, AuxiliaryScan]],
         instrument_session: str,
     ) -> list[Task]:
         combined_auxiliary_scan = None
         index = None
+        aux_task_id = None
 
-        for i, auxiliary_scan in queued_auxiliary_scans:
+        for i, task_id, auxiliary_scan in queued_auxiliary_scans:
             if combined_auxiliary_scan := auxiliary_scan.attempt_to_combine_with(
                 required_auxiliary
             ):
+                aux_task_id = task_id
                 index = i
                 break
 
         aux_experiment = self._construct_auxiliary_experiment(
             combined_auxiliary_scan or required_auxiliary, instrument_session
         )
-        if index is None:
-            new_tasks.append(Task(experiment=aux_experiment))
+        if index is None or aux_task_id is None:
+            new_tasks.append(
+                Task(experiment=aux_experiment, rules=TaskRules(movable=False))
+            )
+            aux_task_id = new_tasks[-1].id
         else:
-            new_tasks[index] = Task(experiment=aux_experiment)
+            new_tasks[index] = Task(
+                experiment=aux_experiment,
+                id=aux_task_id,
+                rules=TaskRules(movable=False),
+            )
+
+        task.rules.run_after.add(aux_task_id)
         return new_tasks
 
     def _get_required_auxiliary_scans(

src/daq_queuing_service/task_queue/queue.py

@@ -37,8 +37,10 @@
     TaskNotClaimedError,
     TaskNotFoundError,
     TaskNotInQueueError,
+    TaskNotMovableError,
+    validate_task_position,
 )
-from daq_queuing_service.task_queue.task import Status, Task, TaskWithPosition
+from daq_queuing_service.task_queue.task import Status, Task, TaskID, TaskWithPosition
 
 
 class TaskRegistry(dict[str, Task]):
@@ -65,10 +67,10 @@ class QueueContents(TypedDict):
     tasks (dict[str, Task]): A dict containing all the tasks contained in the queue,
         including history, with the task ID as the key.
 
-    queue (list[str]): A list of the task IDs of the tasks in the queue. The order of
+    queue (list[TaskID]): A list of the task IDs of the tasks in the queue. The order of
         the list represents the order of the queue.
 
-    history (list[str]): A list of the task IDs of any completed tasks.
+    history (list[TaskID]): A list of the task IDs of any completed tasks.
 
     call_queue: (list[BlueapiCall]): A list of blueapi calls that will be executed in
         order to complete the queued tasks.
@@ -77,8 +79,8 @@ class QueueContents(TypedDict):
     """
 
     tasks: dict[str, Task]
-    queue: list[str]
-    history: list[str]
+    queue: list[TaskID]
+    history: list[TaskID]
     call_queue: list[BlueapiCall]
     call_history: list[BlueapiCall]
 
@@ -135,8 +137,8 @@ def __init__(
     ):
         self.instrument = instrument
         self._tasks: TaskRegistry = TaskRegistry()
-        self._queue: list[str] = []
-        self._history: list[str] = []
+        self._queue: list[TaskID] = []
+        self._history: list[TaskID] = []
         self._call_queue: list[BlueapiCall] = []
         self._call_history: list[BlueapiCall] = []
         self._save_contents()
@@ -204,6 +206,9 @@ def _sync(self):
         ]
         self._queue.extend(task.id for task in new_tasks)
 
+        # Ensure new task positions don't break any rules
+        self._validate_queue_positions()
+
         # This is needed because pre_process could remove tasks
         self._remove_tasks_from_registry(list(self._tasks.keys()))
 
@@ -503,8 +508,8 @@ async def move_task(self, task_id: str, position: int) -> int:
             int: The new position of the task (may be different to what was requested)
         """
         async with self._modifying:
-            self._validate_tasks_for_move_or_deletion([task_id])
             position = self._get_valid_position(position)
+            self._validate_task_for_move(task_id, position)
             self._remove_tasks_from_queue([task_id])
             self._queue[position:position] = [task_id]
             new_position = self._queue.index(task_id)
@@ -649,7 +654,7 @@ def _cancel_tasks(self, task_ids: Sequence[str]) -> list[TaskWithPosition]:
             task.cancel()
         return [TaskWithPosition.from_task(task) for task in tasks]
 
-    def _remove_tasks_from_queue(self, task_ids: list[str]) -> list[str]:
+    def _remove_tasks_from_queue(self, task_ids: list[TaskID]) -> list[TaskID]:
         #  Only removes tasks in the queue (not history or registry)
         def should_be_removed(task_id: str):
             return (
@@ -662,7 +667,7 @@ def should_be_removed(task_id: str):
 
         return removed_ids
 
-    def _remove_tasks_from_registry(self, task_ids: list[str]) -> None:
+    def _remove_tasks_from_registry(self, task_ids: list[TaskID]) -> None:
         # Only removes tasks not present in the queue or history
         # So should remove tasks from queue/history before removing from registry
         def should_be_removed(task_id: str) -> bool:
@@ -682,7 +687,7 @@ def should_be_removed(task_id: str) -> bool:
             }
         )
 
-    def _validate_tasks_for_move_or_deletion(self, task_ids: list[str]):
+    def _validate_tasks_for_move_or_deletion(self, task_ids: list[TaskID]):
         for task_id in task_ids:
             task = self._tasks[task_id]
             if task_id not in self._queue:
@@ -692,6 +697,31 @@ def _validate_tasks_for_move_or_deletion(self, task_ids: list[str]):
                     f"Cannot move task '{task_id}', it is currently in progress!"
                 )
 
+    def _validate_task_for_move(self, task_id: TaskID, position: int):
+        task = self._tasks[task_id]
+        if not task.rules.movable:
+            raise TaskNotMovableError(
+                "Task isn't movable. See the following rules for task "
+                + f"{task_id}: {task.rules}"
+            )
+        self._validate_task_in_position(task_id, position)
+        self._validate_tasks_for_move_or_deletion([task_id])
+
+    def _validate_task_in_position(self, task_id: TaskID, position: int):
+        task = self._tasks[task_id]
+        queue_without_task = [_id for _id in self._queue if _id != task_id]
+        queue_after_move = (
+            queue_without_task[:position] + [task_id] + queue_without_task[position:]
+        )
+        positions = {_id: i for i, _id in enumerate(queue_after_move)}
+        validate_task_position(task, positions)
+
+    def _validate_queue_positions(self):
+        positions = {_id: i for i, _id in enumerate(self._queue)}
+        for task_id in self._queue:
+            task = self._tasks[task_id]
+            validate_task_position(task, positions)
+
     def _get_queue(self) -> list[Task]:
         return [self._tasks[task_id] for task_id in self._queue]
 
@@ -749,3 +779,5 @@ def _add_runtime_estimates_to_tasks(self):
     @cached_property
     def _tiled_client(self) -> Container:
         return get_tiled_client()
+
+    def _check_rules(self): ...

src/daq_queuing_service/task_queue/queue_utils.py

@@ -1,3 +1,6 @@
+from daq_queuing_service.task_queue.task import Task, TaskID
+
+
 class QueueError(Exception):
     pass
 
@@ -24,3 +27,33 @@ class NegativePositionError(QueueError):
 
 class TaskNotClaimedError(QueueError):
     pass
+
+
+class TaskNotMovableError(QueueError):
+    pass
+
+
+class TaskPositionError(QueueError):
+    pass
+
+
+def validate_task_position(task: Task, positions: dict[TaskID, int]):
+    task_position = positions[task.id]
+
+    for must_be_before_id in task.rules.run_before:
+        must_be_before = positions.get(must_be_before_id)
+        if must_be_before is not None and task_position > must_be_before:
+            raise TaskPositionError(
+                f"Task cannot be in position {task_position} as it would be after"
+                + f" {must_be_before_id} in position {must_be_before}. See the "
+                + f"following rules for task {task.id}: {task.rules}"
+            )
+
+    for must_be_after_id in task.rules.run_after:
+        must_be_after = positions.get(must_be_after_id)
+        if must_be_after is not None and task_position < must_be_after:
+            raise TaskPositionError(
+                f"Task cannot be in position {task_position} as it would be before"
+                + f" {must_be_after_id} in position {must_be_after}. See the following "
+                + f"rules for task {task.id}: {task.rules}"
+            )

src/daq_queuing_service/task_queue/task.py

@@ -1,9 +1,9 @@
 from enum import StrEnum
-from typing import Any, Self
+from typing import Any, Self, TypeAlias
 from uuid import uuid4
 
 from blueapi.service.model import TaskRequest
-from pydantic import BaseModel, Field, computed_field
+from pydantic import BaseModel, Field, computed_field, model_validator
 
 from daq_queuing_service.app.authentication import User
 from daq_queuing_service.external_interaction.blueapi.blueapi_call import (
@@ -47,6 +47,9 @@ def create_uuid_str() -> str:
     return str(uuid4())
 
 
+TaskID: TypeAlias = str
+
+
 class Status(StrEnum):
     QUEUED = "Queued"
     IN_PROGRESS = "In progress"
@@ -60,12 +63,28 @@ class TaskKind(StrEnum):
     PLAN = "Plan"
 
 
+class TaskRules(BaseModel):
+    movable: bool = True
+    run_before: set[TaskID] = Field(default_factory=set)
+    run_after: set[TaskID] = Field(default_factory=set)
+
+    @model_validator(mode="after")
+    def ensure_rules_do_not_contradict(self) -> Self:
+        if in_both := self.run_before.intersection(self.run_after):
+            raise ValueError(
+                f"{in_both} appear in both run_before and run_after. "
+                + "This is an impossible requirement."
+            )
+        return self
+
+
 class Task(BaseModel):
     experiment: Experiment | TaskRequest
-    id: str = Field(default_factory=create_uuid_str)
+    id: TaskID = Field(default_factory=create_uuid_str)
     blueapi_calls: list[BlueapiCall] = Field(default_factory=lambda: [])
     user: User | None = None
     estimated_runtime: int | None = None
+    rules: TaskRules = Field(default_factory=TaskRules)
     _cancelled: bool = False
 
     def cancel(self):
@@ -127,6 +146,7 @@ class TaskWithPosition(BaseModel):
     user: User | None
     estimated_runtime: int | None
     runtime: float | None
+    rules: TaskRules
 
     @classmethod
     def from_task(cls, task: Task, position: int | None = None) -> Self: