import asyncio
from collections.abc import Callable, Sequence
from copy import deepcopy
from enum import StrEnum
from functools import cached_property
from types import TracebackType
from typing import Literal, TypedDict
from blueapi.worker.event import TaskError, TaskResult
from pydantic import BaseModel
from tiled.client.container import Container
from daq_queuing_service.broadcaster import Broadcaster, Event
from daq_queuing_service.external_interaction.blueapi.blueapi_call import (
BlueapiCall,
BlueapiCallResponse,
CallStatus,
)
from daq_queuing_service.external_interaction.runtime_estimator.client import (
RuntimeEstimatorClient,
)
from daq_queuing_service.external_interaction.tiled.tiled import (
get_tiled_and_scan_ids,
get_tiled_client,
)
from daq_queuing_service.log import (
LOGGER,
add_session_to_logger,
get_session_filter,
notify,
)
from daq_queuing_service.plugins.converter import Converter, ConverterError
from daq_queuing_service.task_queue.queue_utils import (
NegativePositionError,
TaskIdInUseError,
TaskInProgressError,
TaskNotClaimedError,
TaskNotFoundError,
TaskNotInQueueError,
)
from daq_queuing_service.task_queue.task import Status, Task, TaskWithPosition
class TaskRegistry(dict[str, Task]):
def __missing__(self, task_id: str) -> Task:
raise TaskNotFoundError(f"No task found matching id: {task_id}")
class PauseReason(StrEnum):
USER_REQUESTED = "Pause requested by user"
EMPTY_QUEUE = "Paused as queue completed"
ERROR = "Paused as last task errored"
PLUGIN_REQUESTED = "Pause requested by plugin"
class QueueState(BaseModel):
paused: bool
last_pause_reason: PauseReason
estimated_runtime: int | None = None
class QueueContents(TypedDict):
"""A class to represent the contents of the queue
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
the list represents the order of the queue.
history (list[str]): 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.
call_history: (list[BlueapiCall]): A list of completed blueapi calls.
"""
tasks: dict[str, Task]
queue: list[str]
history: list[str]
call_queue: list[BlueapiCall]
call_history: list[BlueapiCall]
class Modifying(asyncio.Condition):
def __init__(
self,
lock: asyncio.Lock,
on_exit: Callable[[], None],
on_error: Callable[[], None],
):
super().__init__(lock=lock)
self._on_exit = on_exit
self._on_error = on_error
async def __aexit__(
self,
exc_type: type[BaseException] | None,
exc: BaseException | None,
tb: TracebackType | None,
):
try:
if exc is not None:
self._on_error()
return
self._on_exit()
except Exception:
self._on_error()
raise
finally:
await super().__aexit__(exc_type, exc, tb)
QUEUE_EVENTS = Literal[
"state_update",
"queue_update",
"history_update",
"tasks_update",
"call_queue_update",
"call_history_update",
]
[docs]
class TaskQueue:
[docs]
def __init__(
self,
converter: Converter,
broadcaster: Broadcaster[QUEUE_EVENTS],
runtime_estimator_client: RuntimeEstimatorClient | None = None,
instrument: str | None = None,
):
self.instrument = instrument
self._tasks: TaskRegistry = TaskRegistry()
self._queue: list[str] = []
self._history: list[str] = []
self._call_queue: list[BlueapiCall] = []
self._call_history: list[BlueapiCall] = []
self._save_contents()
self._state: QueueState = QueueState(
paused=True, last_pause_reason=PauseReason.EMPTY_QUEUE
)
self._converter = converter
self._broadcaster = broadcaster
self._lock = asyncio.Lock()
self._modifying = Modifying(
lock=self._lock,
on_exit=self._sync,
on_error=self._restore_latest_good_contents,
)
self._runtime_client = runtime_estimator_client
def _sync(self):
"""Syncs the task queue with the call queue, applying a conversion from queue of
tasks to a queue of blueapi calls. This is called every time the queue is
modified, and also right before a call is popped off the front of the queue.
"""
assert get_session_filter(LOGGER).instrument_session_var.get() is None, (
"Logger instrument session should not be set at this point"
)
LOGGER.debug(f"Syncing. Queue before sync: {self._queue}")
for task_id in list(self._queue):
task = self._tasks[task_id]
if task.status in (Status.COMPLETE, Status.ERROR):
# Move task from queue to history if all blueapi calls complete
# Or an error occurred during one of them
self._queue.remove(task_id)
self._history.append(task_id)
for call in filter(
lambda call: call.status == CallStatus.WAITING, task.blueapi_calls
):
call.skip()
elif task.status == Status.QUEUED:
# If task is not in progress calls will be re-calculated
task.blueapi_calls = []
try:
new_tasks = self._converter.pre_process(
self._get_running_task(),
[
self._tasks[task_id]
for task_id in self._queue
if self._tasks[task_id].status == Status.QUEUED
],
self._get_history(),
self._get_call_history(),
)
except Exception as e:
raise ConverterError(e) from e
# Update task_registry to match new tasks
# Not needed as long as pre_process modifies in place
for task in new_tasks:
self._tasks[task.id] = task
# Update queue to match new tasks
self._queue = [
task_id
for task_id in self._queue
if not self._tasks[task_id].status == Status.QUEUED
]
self._queue.extend(task.id for task in new_tasks)
# This is needed because pre_process could remove tasks
self._remove_tasks_from_registry(list(self._tasks.keys()))
self._call_queue = [
# Persist calls which aren't complete but who's parent task is in progress
# Once a task is in progress it is not provided to the converter
# More work needed to allow for interleaved calls from different tasks,
# and for in progress tasks to inform conversion
call
for call in self._call_queue
if call.parent_task_id
and call.parent_task_id in self._queue
and self._tasks[call.parent_task_id].status == Status.IN_PROGRESS
and call.status not in (CallStatus.SUCCESS, CallStatus.ERROR)
]
try:
new_calls = self._converter.construct_blueapi_calls(
[
task
for task in self._get_queue_with_positions()
if task.status == Status.QUEUED
],
self._get_history(),
self._get_call_history(),
)
except Exception as e:
raise ConverterError(e) from e
self._call_queue.extend(new_calls)
for call in new_calls:
# Add children to parent tasks
if call.parent_task_id:
self._tasks[call.parent_task_id].blueapi_calls.append(call)
if not self._call_queue:
self._pause_queue(PauseReason.EMPTY_QUEUE)
self._add_runtime_estimates_to_tasks()
self._save_contents()
self._broadcast_changes()
self._modifying.notify_all()
LOGGER.debug(f"Queue after sync: {self._queue}")
def _copy_contents(self) -> QueueContents:
return deepcopy(
{
"tasks": self._tasks,
"queue": self._queue,
"history": self._history,
"call_queue": self._call_queue,
"call_history": self._call_history,
}
)
def _save_contents(self):
self._last_good_contents = self._copy_contents()
def _restore_from_contents(self, contents: QueueContents):
LOGGER.info(f"Restoring to contents: {contents}")
restored = deepcopy(contents)
self._tasks = TaskRegistry(restored["tasks"])
self._queue = restored["queue"]
self._history = restored["history"]
self._call_queue = restored["call_queue"]
self._call_history = restored["call_history"]
def _restore_latest_good_contents(self):
LOGGER.info("Restoring to last good contents")
self._restore_from_contents(self._last_good_contents)
def _broadcast_changes(self):
queue = self._get_queue_with_positions()
history = self._get_history()
# broadcast only does anything if there have been changes
self._broadcaster.broadcast(Event(type="queue_update", data=queue))
self._broadcaster.broadcast(Event(type="history_update", data=history))
self._broadcaster.broadcast(Event(type="tasks_update", data=history + queue))
self._broadcaster.broadcast(
Event(type="call_queue_update", data=self._get_call_queue())
)
self._broadcaster.broadcast(
Event(type="call_history_update", data=self._get_call_history())
)
self._broadcaster.broadcast(Event(type="state_update", data=self._state))
[docs]
async def get_next_call_once_available(self) -> BlueapiCall:
"""Waits until a call is available before returning the call. A call is
available if it's at the top of the call queue, is not already in progress or
claimed, and the queue is not paused.
Returns:
BlueapiCall: The call at the top of the queue
"""
async with self._modifying:
while not self._task_available():
await self._modifying.wait()
self._sync() # Do conversion here to ensure conditions are up to date
call = self._call_queue[0]
with add_session_to_logger(call.task_request.instrument_session):
call.claim()
LOGGER.info(f"Plan {call} has been claimed")
return call
[docs]
async def wait_until_call_available(self):
"""Waits until a task is available before returning. A task is
available if it's at the top of the queue, is not already in progress or
claimed, and the queue is not paused.
"""
async with self._lock:
while not self._task_available():
await self._modifying.wait()
[docs]
async def return_call_to_queue(self, call: BlueapiCall):
"""Returns a task to the queue that had previously been claimed
Args:
task (Task): The task to return
Raises:
TaskNotClaimedError: Raised if the task's status is not 'Claimed'
"""
self._check_call_valid_to_be_returned(call)
async with self._modifying:
match call.status:
case CallStatus.CLAIMED:
assert call == self._call_queue[0]
call.wait()
case _:
raise TaskNotClaimedError(
f"Cannot return call {call}, "
+ f"it's status is {call.status}."
)
with add_session_to_logger(call.task_request.instrument_session):
LOGGER.info(f"Call {call} has been returned to the queue")
[docs]
async def complete_call(self, call: BlueapiCall, result: TaskResult):
"""Sets a task to complete, removes it from the queue and adds it to history
Args:
task (Task): Task to be completed
result (TaskResult): The result of the task from blueapi
"""
with add_session_to_logger(call.task_request.instrument_session):
tiled_ids, scan_ids = get_tiled_and_scan_ids(self._tiled_client, call)
should_pause = self._converter.pause_on_complete_blueapi_call(
call.model_copy(), result
)
LOGGER.info(f"Plan {call} completed successfully. Result: {result}")
async with self._modifying:
call.tiled_ids = tiled_ids
call.scan_ids = scan_ids
self._check_call_valid_to_be_returned(call)
call.succeed(result)
self._call_history.append(call)
if should_pause:
self._pause_queue(reason=PauseReason.PLUGIN_REQUESTED)
[docs]
async def fail_call(
self, call: BlueapiCall, errors: list[str | TaskError] | None = None
):
"""Sets a task to failed, removes it from the call queue and adds it to history.
Args:
task (Task): The task to fail
errors (list[str | TaskError] | None, optional): A list of errors that
occurred when trying to run the task. Defaults to None.
"""
with add_session_to_logger(call.task_request.instrument_session):
should_pause = self._converter.pause_on_fail_blueapi_call(
call.model_copy(), errors
)
tiled_ids, scan_ids = get_tiled_and_scan_ids(self._tiled_client, call)
LOGGER.error(f"Plan {call} has failed with the following errors: {errors}")
async with self._modifying:
call.tiled_ids = tiled_ids
call.scan_ids = scan_ids
if should_pause:
self._pause_queue(PauseReason.ERROR)
else:
LOGGER.warning(
"Call failed but queue not paused due to plugin decision"
)
self._check_call_valid_to_be_returned(call)
call.fail(errors)
self._call_history.append(call)
[docs]
async def get_task_by_id(self, task_id: str) -> TaskWithPosition:
"""Returns a task based on it's task ID
Args:
task_id (str): Task ID of the task
Returns:
TaskWithPosition: A copy of the task with a position field.
Raises:
TaskNotFoundError: Raised if the no task exists with the requested task ID.
"""
# Returns copy so don't have to be worried about caller modifying task.
async with self._lock:
return self._get_task_by_id(task_id)
def _get_task_by_id(self, task_id: str) -> TaskWithPosition:
task = self._tasks[task_id]
position = self._queue.index(task.id) if task.id in self._queue else None
return TaskWithPosition.from_task(task, position)
[docs]
async def get_task_by_position(self, position: int) -> TaskWithPosition | None:
"""Return a task based on it's position in the queue
Args:
position (int): The position of the task to be returned.
Returns:
TaskWithPosition | None: A copy of the task with a position field, or None
if no task exists at the requested position.
"""
# Returns copy so don't have to be worried about caller modifying task.
async with self._lock:
return self._get_task_by_position(position)
def _get_task_by_position(self, position: int) -> TaskWithPosition | None:
if position < -self.length or position >= self.length:
return
return self._get_task_by_id(self._queue[position])
[docs]
async def get_queue(self) -> list[TaskWithPosition]:
"""Get the entire queue (not including history)
Returns:
list[TaskWithPosition]: A list of the tasks in the queue, in the order they
will be run in.
"""
# Returns copies so don't have to be worried about caller modifying tasks.
async with self._lock:
return self._get_queue_with_positions()
[docs]
async def get_history(self) -> list[TaskWithPosition]:
"""Get the history list.
Returns:
list[TaskWithPosition]: A list of the tasks in the history list, in
chronological order.
"""
# Returns copies so don't have to be worried about caller modifying tasks.
async with self._lock:
return self._get_history()
[docs]
async def get_tasks(self) -> list[TaskWithPosition]:
"""Get all the tasks in the queue and history list.
Returns:
list[TaskWithPosition]: A list of tasks in chronological order, starting
with the history.
"""
# Returns copies so don't have to be worried about caller modifying tasks.
async with self._lock:
return self._get_history() + self._get_queue_with_positions()
[docs]
async def add_tasks(self, tasks: list[Task], position: int | None = None) -> None:
"""Add tasks to the queue. Can specify a position to add the tasks in at.
This position will apply to the first task in the list, with each subsequent
task having a position of 1 more than the previous. By default adds tasks to the
end of the queue.
Args:
tasks (list[Task]): List of tasks to add
position (int | None, optional): Position of the tasks. Defaults to None.
"""
async with self._modifying:
self._validate_new_tasks(tasks)
if position is not None:
position = self._get_valid_position(position)
self._add_tasks(tasks, position)
LOGGER.info(f"Successfully added tasks to queue: {[task.id for task in tasks]}")
[docs]
async def move_task(self, task_id: str, position: int) -> int:
"""Move a task into a different position. If the requested position is
unavailable, will move to the closest available position. If the requested task
cannot be moved because it is complete or in progress, an error will be raised.
Args:
task_id (str): ID of the task to be moved
new_position (int): New position of the task
Returns:
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._remove_tasks_from_queue([task_id])
self._queue[position:position] = [task_id]
new_position = self._queue.index(task_id)
LOGGER.info(f"Successfully moved task {task_id} to position {new_position}")
return new_position
[docs]
async def cancel_tasks(self, task_ids: Sequence[str]) -> list[TaskWithPosition]:
"""Remove tasks from the queue. If one or more of the requested tasks cannot be
cancelled as they are complete or in progress, an error will be raised and none
of the requested tasks will be cancelled.
Args:
task_ids (Sequence[str]): List of task IDs to cancel
Returns:
list[TaskWithPosition]: List of the task objects that were removed from the
queue.
"""
async with self._modifying:
cancelled_tasks = self._cancel_tasks(task_ids)
LOGGER.info(f"Successfully cancelled tasks: {task_ids}")
return cancelled_tasks
[docs]
async def cancel_all_tasks(self) -> list[TaskWithPosition]:
async with self._modifying:
task_ids = [
id
for id in self._queue
if self._get_task_by_id(id).status == Status.QUEUED
]
cancelled_tasks = self._cancel_tasks(task_ids)
LOGGER.info("Successfully cleared all queued tasks")
return cancelled_tasks
[docs]
async def clear_history(self):
"""Clears the history list. Any task in the history list at the time will be
deleted permanently and inaccessible.
"""
async with self._modifying:
self._history.clear()
LOGGER.info("Successfully cleared history")
def _pause_queue(self, reason: PauseReason):
reason = (
reason
if reason == PauseReason.ERROR or not self._state.paused
else self._state.last_pause_reason
)
self._state.paused = True
self._state.last_pause_reason = reason
notify(f"Queue paused: {reason}", "INFO")
[docs]
async def pause_queue(self, reason: PauseReason) -> QueueState:
"""Pause the queue if running.
Args:
paused (PauseReason): The reason for pausing the queue.
Returns:
QueueState: The new state of the queue.
"""
async with self._modifying:
self._pause_queue(reason)
return self._state
[docs]
async def resume_queue(self) -> QueueState:
"""Resume the queue if paused.
Returns:
QueueState: The new state of the queue.
"""
LOGGER.info("Resuming queue")
async with self._modifying:
self._state.paused = False
return self._state
@property
def state(self):
return self._state
@property
def length(self):
return len(self._queue)
def _task_available(self) -> bool:
"""Predicate that determines whether the queue has a task available. This
returns True if the first task in the queue has a status of WAITING, and the
queue is not paused.
Returns:
bool: Whether or not the queue has a task available.
"""
if self._state.paused or not self._call_queue:
return False
return self._call_queue[0].status == CallStatus.WAITING
def _check_call_valid_to_be_returned(self, call: BlueapiCall):
# Check caller has actual task object not copy
# This ensures the caller has claimed the task, reducing the chance a task is
# returned that is actually still being run/modified by a different process.
# However if the worker crashes we then lose the Task object and can't return
# the task? Needs discussion with others.
assert call is self._call_queue[0]
assert call.parent_task_id in self._queue, (
f"This call has no parent task: {call}"
)
def _get_valid_position(self, position: int) -> int:
if position < 0:
raise NegativePositionError(f"Position must be >= 0, got {position}")
if ( # if position 0 requested but a task is in progress, return position 1
position == 0
and self.length
and self._tasks[self._queue[0]].status != Status.QUEUED
):
return 1
return position
def _validate_new_tasks(self, tasks: list[Task]):
for task in tasks:
if task.id in self._tasks:
raise TaskIdInUseError(f"Task ID '{task.id}' already in use!")
def _add_tasks(self, tasks: list[Task], position: int | None) -> None:
task_ids = [task.id for task in tasks]
if position is None:
self._queue.extend(task_ids)
else:
self._queue[position:position] = task_ids
for task in tasks:
self._tasks[task.id] = task
def _cancel_tasks(self, task_ids: Sequence[str]) -> list[TaskWithPosition]:
task_ids = list(task_ids)
self._validate_tasks_for_move_or_deletion(task_ids)
self._remove_tasks_from_queue(task_ids)
tasks = [self._tasks[task_id] for task_id in task_ids]
for task in tasks:
task.cancel()
return [TaskWithPosition.from_task(task) for task in tasks]
def _remove_tasks_from_queue(self, task_ids: list[str]) -> list[str]:
# Only removes tasks in the queue (not history or registry)
def should_be_removed(task_id: str):
return (
task_id in self._queue
and self._tasks[task_id].status != Status.IN_PROGRESS
)
removed_ids = [task_id for task_id in task_ids if should_be_removed(task_id)]
self._queue = [task_id for task_id in self._queue if task_id not in removed_ids]
return removed_ids
def _remove_tasks_from_registry(self, task_ids: list[str]) -> 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:
return (
task_id in self._tasks
and self._tasks[task_id].status != Status.IN_PROGRESS
and task_id not in self._queue
and task_id not in self._history
)
removed_ids = [task_id for task_id in task_ids if should_be_removed(task_id)]
self._tasks = TaskRegistry(
{
task_id: task
for task_id, task in self._tasks.items()
if task.id not in removed_ids
}
)
def _validate_tasks_for_move_or_deletion(self, task_ids: list[str]):
for task_id in task_ids:
task = self._tasks[task_id]
if task_id not in self._queue:
raise TaskNotInQueueError(f"Task {task_id} isn't present in queue")
if task.status != Status.QUEUED:
raise TaskInProgressError(
f"Cannot move task '{task_id}', it is currently in progress!"
)
def _get_queue(self) -> list[Task]:
return [self._tasks[task_id] for task_id in self._queue]
def _get_queue_with_positions(self) -> list[TaskWithPosition]:
return [
TaskWithPosition.from_task(self._tasks[task_id], i)
for i, task_id in enumerate(self._queue)
]
def _get_history(self) -> list[TaskWithPosition]:
return [
TaskWithPosition.from_task(self._tasks[task_id])
for task_id in self._history
]
[docs]
async def get_call_queue(self) -> list[BlueapiCallResponse]:
async with self._lock:
return self._get_call_queue()
def _get_call_queue(self) -> list[BlueapiCallResponse]:
return [call.to_response() for call in self._call_queue]
[docs]
async def get_call_history(self) -> list[BlueapiCallResponse]:
async with self._lock:
return self._get_call_history()
def _get_call_history(self) -> list[BlueapiCallResponse]:
return [call.to_response() for call in self._call_history]
def _get_running_task(self) -> TaskWithPosition | None:
current_task = self._get_task_by_position(0)
if current_task and current_task.status == Status.IN_PROGRESS:
return current_task
def _add_runtime_estimates_to_tasks(self):
if self._runtime_client and self.instrument:
tasks = self._get_queue()
experiments = [task.experiment for task in tasks]
try:
total_runtime = self._runtime_client.estimate_runtime(
experiments, self.instrument
)
runtimes = self._runtime_client.estimate_runtime_per_task(
experiments, self.instrument
)
except Exception as e:
# This shouldn't prevent syncing so catch and log
LOGGER.exception(e)
return
self._state.estimated_runtime = total_runtime
for task, runtime in zip(tasks, runtimes, strict=True):
if task.status == Status.QUEUED:
task.estimated_runtime = runtime
@cached_property
def _tiled_client(self) -> Container:
return get_tiled_client()
