diff --git a/docs/explanations/converters.md b/docs/explanations/converters.md index 1e3572ef..dfcea677 100644 --- a/docs/explanations/converters.md +++ b/docs/explanations/converters.md @@ -22,16 +22,16 @@ class MyConverter(Converter): def pre_process( self, queue: list[Task], - history: list[TaskWithPosition], - call_history: list[BlueapiCall], + history: list[QueueTask], + call_history: list[QueueBlueapiCall], ) -> list[Task]: ... def construct_blueapi_calls( self, - queue: list[TaskWithPosition], - history: list[TaskWithPosition], - call_history: list[BlueapiCall], - ) -> list[BlueapiCall]: ... + queue: list[QueueTask], + history: list[QueueTask], + call_history: list[QueueBlueapiCall], + ) -> list[QueueBlueapiCall]: ... ``` ### `validate()` @@ -44,7 +44,7 @@ This is called whenever anything changes in the queue (e.g a task is added, or a ### `construct_blueapi_calls()` -The queue contains ulims experiments and/or TaskRequests. While TaskRequests are already in the right format for blueapi to receive them, experiments need to be converted to blueapi calls for anything to run. This method is used for this conversion, constructing a list of BlueapiCall objects (which wrap TaskRequest objects) from the current queue contents. This is run whenever the queue changes, but after `pre_process`. The default converter assumes that the queue contains exclusively TaskRequests, and will raise a `NotImplementedError` if an `Experiment` is added. If you want to add ulims experiments to the queue, make sure you are using a custom converter that overrides this method! +The queue contains ulims experiments and/or TaskRequests. While TaskRequests are already in the right format for blueapi to receive them, experiments need to be converted to blueapi calls for anything to run. This method is used for this conversion, constructing a list of QueueBlueapiCall objects (which wrap TaskRequest objects) from the current queue contents. This is run whenever the queue changes, but after `pre_process`. The default converter assumes that the queue contains exclusively TaskRequests, and will raise a `NotImplementedError` if an `Experiment` is added. If you want to add ulims experiments to the queue, make sure you are using a custom converter that overrides this method! Pre-processing and conversion is run any time the queue changes and right before a call is made to blueapi, so the blueapi call list stays in sync. This means these methods can make use of current information on a beamline such as PV values to inform their logic. However, because they're run so frequently, it is important these methods do not take a large amount of time. diff --git a/docs/reference/rest_api.json b/docs/reference/rest_api.json index c93877bf..3aa33b58 100644 --- a/docs/reference/rest_api.json +++ b/docs/reference/rest_api.json @@ -139,7 +139,7 @@ "schema": { "type": "array", "items": { - "$ref": "#/components/schemas/TaskWithPosition" + "$ref": "#/components/schemas/QueueTask" }, "title": "Response Get Queued Tasks Queue Get" } @@ -238,7 +238,7 @@ "schema": { "type": "array", "items": { - "$ref": "#/components/schemas/TaskWithPosition" + "$ref": "#/components/schemas/QueueTask" }, "title": "Response Cancel All Tasks Queue Delete" } @@ -318,7 +318,7 @@ "application/json": { "schema": { "items": { - "$ref": "#/components/schemas/TaskWithPosition" + "$ref": "#/components/schemas/QueueTask" }, "type": "array", "title": "Response Cancel Tasks Queue Tasks Delete" @@ -362,7 +362,7 @@ "schema": { "anyOf": [ { - "$ref": "#/components/schemas/TaskWithPosition" + "$ref": "#/components/schemas/QueueTask" }, { "type": "null" @@ -416,7 +416,7 @@ "schema": { "type": "array", "items": { - "$ref": "#/components/schemas/TaskWithPosition" + "$ref": "#/components/schemas/QueueTask" }, "title": "Response Get All Tasks Tasks Get" } @@ -457,7 +457,7 @@ "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/TaskWithPosition" + "$ref": "#/components/schemas/QueueTask" } } } @@ -505,7 +505,7 @@ "schema": { "type": "array", "items": { - "$ref": "#/components/schemas/TaskWithPosition" + "$ref": "#/components/schemas/QueueTask" }, "title": "Response Get Completed Tasks History Get" } @@ -550,7 +550,7 @@ "application/json": { "schema": { "items": { - "$ref": "#/components/schemas/BlueapiCallResponse" + "$ref": "#/components/schemas/QueueBlueapiCall" }, "type": "array", "title": "Response Get Call Queue Call Queue Get" @@ -572,7 +572,7 @@ "application/json": { "schema": { "items": { - "$ref": "#/components/schemas/BlueapiCallResponse" + "$ref": "#/components/schemas/QueueBlueapiCall" }, "type": "array", "title": "Response Get Call History Call History Get" @@ -688,126 +688,6 @@ "title": "BasicAuthentication", "description": "User credentials for basic authentication" }, - "BlueapiCallResponse": { - "properties": { - "task_request": { - "$ref": "#/components/schemas/TaskRequest" - }, - "parent_task_id": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Parent Task Id" - }, - "status": { - "$ref": "#/components/schemas/CallStatus" - }, - "time_started": { - "anyOf": [ - { - "type": "string", - "format": "date-time" - }, - { - "type": "null" - } - ], - "title": "Time Started" - }, - "time_completed": { - "anyOf": [ - { - "type": "string", - "format": "date-time" - }, - { - "type": "null" - } - ], - "title": "Time Completed" - }, - "result": { - "anyOf": [ - { - "$ref": "#/components/schemas/TaskResult" - }, - { - "type": "null" - } - ] - }, - "errors": { - "items": { - "anyOf": [ - { - "type": "string" - }, - { - "$ref": "#/components/schemas/TaskError" - } - ] - }, - "type": "array", - "title": "Errors" - }, - "blueapi_id": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Blueapi Id" - }, - "tiled_ids": { - "items": { - "type": "string" - }, - "type": "array", - "title": "Tiled Ids" - }, - "scan_ids": { - "items": { - "type": "integer" - }, - "type": "array", - "title": "Scan Ids" - }, - "runtime": { - "anyOf": [ - { - "type": "number" - }, - { - "type": "null" - } - ], - "title": "Runtime" - } - }, - "type": "object", - "required": [ - "task_request", - "parent_task_id", - "status", - "time_started", - "time_completed", - "result", - "errors", - "blueapi_id", - "tiled_ids", - "scan_ids", - "runtime" - ], - "title": "BlueapiCallResponse" - }, "BlueapiConfig": { "properties": { "stomp": { @@ -1072,29 +952,9 @@ "OIDCConfig": { "properties": { "well_known_url": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], + "type": "string", "title": "Well Known Url", - "description": "URL to fetch OIDC config from the provider", - "deprecated": true - }, - "issuer": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Issuer", - "description": "URL of OIDC provider" + "description": "URL to fetch OIDC config from the provider" }, "client_id": { "type": "string", @@ -1117,6 +977,7 @@ "additionalProperties": false, "type": "object", "required": [ + "well_known_url", "client_id" ], "title": "OIDCConfig" @@ -1131,6 +992,114 @@ ], "title": "PauseReason" }, + "QueueBlueapiCall": { + "properties": { + "task_request": { + "$ref": "#/components/schemas/TaskRequest" + }, + "parent_task_id": { + "type": "string", + "title": "Parent Task Id" + }, + "status": { + "$ref": "#/components/schemas/CallStatus", + "default": "Waiting" + }, + "time_started": { + "anyOf": [ + { + "type": "string", + "format": "date-time" + }, + { + "type": "null" + } + ], + "title": "Time Started" + }, + "time_completed": { + "anyOf": [ + { + "type": "string", + "format": "date-time" + }, + { + "type": "null" + } + ], + "title": "Time Completed" + }, + "result": { + "anyOf": [ + { + "$ref": "#/components/schemas/TaskResult" + }, + { + "type": "null" + } + ] + }, + "errors": { + "items": { + "anyOf": [ + { + "type": "string" + }, + { + "$ref": "#/components/schemas/TaskError" + } + ] + }, + "type": "array", + "title": "Errors" + }, + "blueapi_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Blueapi Id" + }, + "tiled_ids": { + "items": { + "type": "string" + }, + "type": "array", + "title": "Tiled Ids" + }, + "scan_ids": { + "items": { + "type": "integer" + }, + "type": "array", + "title": "Scan Ids" + }, + "runtime": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Runtime", + "readOnly": true + } + }, + "type": "object", + "required": [ + "task_request", + "parent_task_id", + "runtime" + ], + "title": "QueueBlueapiCall", + "description": "A class representing a call to blueapi, containing the plan parameters and\nsome associated metadata. Constructed by a queue plugin." + }, "QueueState": { "properties": { "paused": { @@ -1172,6 +1141,93 @@ ], "title": "QueueStateUpdate" }, + "QueueTask": { + "properties": { + "experiment": { + "anyOf": [ + { + "$ref": "#/components/schemas/Experiment" + }, + { + "$ref": "#/components/schemas/TaskRequest" + } + ], + "title": "Experiment" + }, + "id": { + "type": "string", + "title": "Id" + }, + "blueapi_calls": { + "items": { + "$ref": "#/components/schemas/QueueBlueapiCall" + }, + "type": "array", + "title": "Blueapi Calls" + }, + "user": { + "anyOf": [ + { + "$ref": "#/components/schemas/User" + }, + { + "type": "null" + } + ] + }, + "estimated_runtime": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "title": "Estimated Runtime" + }, + "position": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "title": "Position" + }, + "status": { + "$ref": "#/components/schemas/Status" + }, + "kind": { + "$ref": "#/components/schemas/TaskKind", + "readOnly": true + }, + "runtime": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Runtime", + "readOnly": true + } + }, + "type": "object", + "required": [ + "experiment", + "position", + "status", + "kind", + "runtime" + ], + "title": "QueueTask", + "description": "Task object with a position field, for passing to a queue plugin or API layer." + }, "RestConfig": { "properties": { "url": { @@ -1392,94 +1448,6 @@ "title": "TaskResult", "description": "Serializable wrapper around the result of a plan\n\nIf the result is not serializable, the result will be None but the type\nwill be the name of the type. If the result is actually None, the type will\nbe 'NoneType'." }, - "TaskWithPosition": { - "properties": { - "experiment": { - "anyOf": [ - { - "$ref": "#/components/schemas/Experiment" - }, - { - "$ref": "#/components/schemas/TaskRequest" - } - ], - "title": "Experiment" - }, - "id": { - "type": "string", - "title": "Id" - }, - "status": { - "$ref": "#/components/schemas/Status" - }, - "blueapi_calls": { - "items": { - "$ref": "#/components/schemas/BlueapiCallResponse" - }, - "type": "array", - "title": "Blueapi Calls" - }, - "position": { - "anyOf": [ - { - "type": "integer" - }, - { - "type": "null" - } - ], - "title": "Position" - }, - "kind": { - "$ref": "#/components/schemas/TaskKind" - }, - "user": { - "anyOf": [ - { - "$ref": "#/components/schemas/User" - }, - { - "type": "null" - } - ] - }, - "estimated_runtime": { - "anyOf": [ - { - "type": "integer" - }, - { - "type": "null" - } - ], - "title": "Estimated Runtime" - }, - "runtime": { - "anyOf": [ - { - "type": "number" - }, - { - "type": "null" - } - ], - "title": "Runtime" - } - }, - "type": "object", - "required": [ - "experiment", - "id", - "status", - "blueapi_calls", - "position", - "kind", - "user", - "estimated_runtime", - "runtime" - ], - "title": "TaskWithPosition" - }, "User": { "properties": { "fedid": { diff --git a/src/daq_queuing_service/api/api.py b/src/daq_queuing_service/api/api.py index 1cbb849e..263b4774 100644 --- a/src/daq_queuing_service/api/api.py +++ b/src/daq_queuing_service/api/api.py @@ -13,7 +13,7 @@ from daq_queuing_service.app.config import AppConfig from daq_queuing_service.broadcaster import Broadcaster from daq_queuing_service.external_interaction.blueapi.blueapi_call import ( - BlueapiCallResponse, + QueueBlueapiCall, ) from daq_queuing_service.log import LOGGER from daq_queuing_service.plugins.converter import Converter, ValidateError @@ -21,8 +21,8 @@ QUEUE_EVENTS, PauseReason, QueueState, + QueueTask, TaskQueue, - TaskWithPosition, ) from daq_queuing_service.task_queue.task import Experiment, Status, Task @@ -45,9 +45,7 @@ class TaskCancelRequest(BaseModel): task_ids: list[str] -def _filter_by_status( - tasks: list[TaskWithPosition], status: Status | None -) -> list[TaskWithPosition]: +def _filter_by_status(tasks: list[QueueTask], status: Status | None) -> list[QueueTask]: if status is None: return tasks return [task for task in tasks if task.status == status] @@ -99,7 +97,7 @@ async def update_queue_state(payload: QueueStateUpdate) -> QueueState: return await queue.resume_queue() @router.get("/queue") - async def get_queued_tasks(status: Status | None = None) -> list[TaskWithPosition]: + async def get_queued_tasks(status: Status | None = None) -> list[QueueTask]: return _filter_by_status(await queue.get_queue(), status) @router.post("/queue") @@ -119,7 +117,7 @@ async def add_tasks_to_queue( return task_ids @router.delete("/queue") - async def cancel_all_tasks() -> list[TaskWithPosition]: + async def cancel_all_tasks() -> list[QueueTask]: return await queue.cancel_all_tasks() @router.post("/queue/move") @@ -127,25 +125,25 @@ async def move_task(task_id: str, new_position: int) -> int: return await queue.move_task(task_id, new_position) @router.delete("/queue/tasks") - async def cancel_tasks(payload: TaskCancelRequest) -> list[TaskWithPosition]: + async def cancel_tasks(payload: TaskCancelRequest) -> list[QueueTask]: return await queue.cancel_tasks(payload.task_ids) @router.get("/queue/{position}") - async def get_task_by_position(position: int) -> TaskWithPosition | None: + async def get_task_by_position(position: int) -> QueueTask | None: return await queue.get_task_by_position(position) @router.get("/tasks") - async def get_all_tasks(status: Status | None = None) -> list[TaskWithPosition]: + async def get_all_tasks(status: Status | None = None) -> list[QueueTask]: return _filter_by_status(await queue.get_tasks(), status) @router.get("/tasks/{task_id}") - async def get_task_by_id(task_id: str) -> TaskWithPosition: + async def get_task_by_id(task_id: str) -> QueueTask: return await queue.get_task_by_id(task_id) @router.get("/history") async def get_completed_tasks( status: Status | None = None, - ) -> list[TaskWithPosition]: + ) -> list[QueueTask]: return _filter_by_status(await queue.get_history(), status) @router.delete("/history") @@ -153,11 +151,11 @@ async def clear_history(): return await queue.clear_history() @router.get("/call_queue") - async def get_call_queue() -> list[BlueapiCallResponse]: + async def get_call_queue() -> list[QueueBlueapiCall]: return await queue.get_call_queue() @router.get("/call_history") - async def get_call_history() -> list[BlueapiCallResponse]: + async def get_call_history() -> list[QueueBlueapiCall]: return await queue.get_call_history() @router.get("/events") diff --git a/src/daq_queuing_service/external_interaction/blueapi/blueapi_call.py b/src/daq_queuing_service/external_interaction/blueapi/blueapi_call.py index ba0cf7aa..ca321c91 100644 --- a/src/daq_queuing_service/external_interaction/blueapi/blueapi_call.py +++ b/src/daq_queuing_service/external_interaction/blueapi/blueapi_call.py @@ -31,7 +31,11 @@ def allowed_transitions(self): return allowed_transitions[self] -class BlueapiCall(BaseModel): +class QueueBlueapiCall(BaseModel): + """A class representing a call to blueapi, containing the plan parameters and + some associated metadata. Constructed by a queue plugin. + """ + task_request: TaskRequest parent_task_id: str status: CallStatus = CallStatus.WAITING @@ -43,6 +47,18 @@ class BlueapiCall(BaseModel): tiled_ids: list[str] = Field(default_factory=list) scan_ids: list[int] = Field(default_factory=list[int]) + @computed_field + @property + def runtime(self) -> float | None: + if self.time_completed is not None and self.time_started is not None: + return self.time_completed.timestamp() - self.time_started.timestamp() + + +class BlueapiCall(QueueBlueapiCall): + """A QueueBlueapiCall for internal use inside the queue, with added methods for + state management. + """ + def _update_status(self, new_status: CallStatus): """Updates the status of the task, checking that the transition is valid @@ -103,29 +119,9 @@ def fail(self, errors: list[str | TaskError] | None = None): def skip(self): self._update_status(CallStatus.SKIPPED) - def to_response(self) -> "BlueapiCallResponse": - return BlueapiCallResponse.from_blueapi_call(self) - - @computed_field - @property - def runtime(self) -> float | None: - if self.time_completed is not None and self.time_started is not None: - return self.time_completed.timestamp() - self.time_started.timestamp() - - -class BlueapiCallResponse(BaseModel): - task_request: TaskRequest - parent_task_id: str | None - status: CallStatus - time_started: datetime | None - time_completed: datetime | None - result: TaskResult | None - errors: list[str | TaskError] - blueapi_id: str | None - tiled_ids: list[str] - scan_ids: list[int] - runtime: float | None + def to_parent(self) -> "QueueBlueapiCall": + return QueueBlueapiCall.model_validate(self.model_dump()) @classmethod - def from_blueapi_call(cls, blueapi_call: BlueapiCall): - return cls.model_validate({**blueapi_call.model_dump()}) + def from_parent(cls, queue_blueapi_call: QueueBlueapiCall) -> "BlueapiCall": + return cls.model_validate(queue_blueapi_call.model_dump()) diff --git a/src/daq_queuing_service/external_interaction/tiled/tiled.py b/src/daq_queuing_service/external_interaction/tiled/tiled.py index 3e3b7b96..b695dff1 100644 --- a/src/daq_queuing_service/external_interaction/tiled/tiled.py +++ b/src/daq_queuing_service/external_interaction/tiled/tiled.py @@ -9,7 +9,9 @@ from tiled.client.container import Container from tiled.queries import Eq -from daq_queuing_service.external_interaction.blueapi.blueapi_call import BlueapiCall +from daq_queuing_service.external_interaction.blueapi.blueapi_call import ( + BlueapiCall, +) from daq_queuing_service.log import LOGGER # Ignoring the following rule as the tiled client is poorly typed diff --git a/src/daq_queuing_service/plugins/converter.py b/src/daq_queuing_service/plugins/converter.py index 7142e31b..67090fb1 100644 --- a/src/daq_queuing_service/plugins/converter.py +++ b/src/daq_queuing_service/plugins/converter.py @@ -4,13 +4,12 @@ from blueapi.worker.event import TaskError, TaskResult from daq_queuing_service.external_interaction.blueapi.blueapi_call import ( - BlueapiCall, - BlueapiCallResponse, + QueueBlueapiCall, ) from daq_queuing_service.task_queue.task import ( Experiment, + QueueTask, Task, - TaskWithPosition, ) @@ -42,10 +41,10 @@ def validate(self, experiments: list[TaskRequest | Experiment]) -> None: def pre_process( self, - current_task: TaskWithPosition | None, + current_task: QueueTask | None, queue: list[Task], - history: list[TaskWithPosition], - call_history: list[BlueapiCallResponse], + history: list[QueueTask], + call_history: list[QueueBlueapiCall], ) -> list[Task]: """This gets run whenever something changes in the queue, before 'construct_blueapi_calls'. It is an opportunity to automatically modify the @@ -58,11 +57,11 @@ def pre_process( If not overridden, this method returns the queue as is with no modification. Args: - current_task (TaskWithPosition | None): The current running task, if one + current_task (QueueTask | None): The current running task, if one exists. queue (list[Task]): List of tasks currently in the queue - history (list[TaskWithPosition]): List of completed tasks - call_history (list[BlueapiCallResponse]): List of completed blueapi calls + history (list[QueueTask]): List of completed tasks + call_history (list[QueueBlueapiCall]): List of completed blueapi calls Returns: list[Task]: The new list of tasks in the queue @@ -71,27 +70,27 @@ def pre_process( def construct_blueapi_calls( self, - queue: list[TaskWithPosition], - history: list[TaskWithPosition], - call_history: list[BlueapiCallResponse], - ) -> list[BlueapiCall]: + queue: list[QueueTask], + history: list[QueueTask], + call_history: list[QueueBlueapiCall], + ) -> list[QueueBlueapiCall]: """Converts the list of queued tasks into a list of blueapi calls. This is needed if the queue contains ulims experiments, as these must be mapped onto bluesky plans. If the queue contains TaskRequests, they can just be - wrapped into BlueapiCall objects with no conversion, as this default + wrapped into QueueBlueapiCall objects with no conversion, as this default implementation does. Args: - queue (list[TaskWithPosition]): List of tasks in the queue - history (list[TaskWithPosition]): List of completed tasks - call_history (list[BlueapiCallResponse]): List of completed blueapi calls + queue (list[QueueTask]): List of tasks in the queue + history (list[QueueTask]): List of completed tasks + call_history (list[QueueBlueapiCall]): List of completed blueapi calls Returns: - list[BlueapiCall]: List of blueapi calls to execute. + list[QueueBlueapiCall]: List of blueapi calls to execute. """ call_list = [ - BlueapiCall( + QueueBlueapiCall( parent_task_id=task.id, task_request=self._construct_blueapi_task_request(task.experiment), ) @@ -100,12 +99,12 @@ def construct_blueapi_calls( return call_list def pause_on_complete_blueapi_call( - self, call: BlueapiCall, result: TaskResult + self, call: QueueBlueapiCall, result: TaskResult ) -> bool: """Decide whether to pause the queue after a blueapi call completes successfully Args: - call (BlueapiCall): The completed call + call (QueueBlueapiCall): The completed call result (TaskResult): The result of the call from blueapi Returns: @@ -114,12 +113,12 @@ def pause_on_complete_blueapi_call( return False def pause_on_fail_blueapi_call( - self, call: BlueapiCall, errors: list[str | TaskError] | None + self, call: QueueBlueapiCall, errors: list[str | TaskError] | None ) -> bool: """Decide whether to pause the queue after a blueapi call fails Args: - call (BlueapiCall): The failed call + call (QueueBlueapiCall): The failed call errors (list[str | TaskError] | None): The errors that occurred Returns: diff --git a/src/daq_queuing_service/plugins/i15_1/i15_1_converter.py b/src/daq_queuing_service/plugins/i15_1/i15_1_converter.py index 5819c899..0361e624 100644 --- a/src/daq_queuing_service/plugins/i15_1/i15_1_converter.py +++ b/src/daq_queuing_service/plugins/i15_1/i15_1_converter.py @@ -8,8 +8,7 @@ from tiled.client.container import Container as TiledContainer from daq_queuing_service.external_interaction.blueapi.blueapi_call import ( - BlueapiCall, - BlueapiCallResponse, + QueueBlueapiCall, ) from daq_queuing_service.external_interaction.tiled.tiled import get_tiled_client from daq_queuing_service.log import ( @@ -38,9 +37,9 @@ ContainerPosition, Experiment, ExperimentDefinition, + QueueTask, Sample, Task, - TaskWithPosition, ) CONFIG_SERVER_URL = "https://i15-1-daq-config.diamond.ac.uk" @@ -87,10 +86,10 @@ def _queue_params(self) -> QueueParameters: def pre_process( self, - current_task: TaskWithPosition | None, + current_task: QueueTask | None, queue: list[Task], - history: list[TaskWithPosition], - call_history: list[BlueapiCallResponse], + history: list[QueueTask], + call_history: list[QueueBlueapiCall], ) -> list[Task]: self._config_client.reset_cache() if self._queue_params.add_auxiliary_scans: @@ -99,25 +98,25 @@ def pre_process( def construct_blueapi_calls( self, - queue: list[TaskWithPosition], - history: list[TaskWithPosition], - call_history: list[BlueapiCallResponse], - ) -> list[BlueapiCall]: - call_list: list[BlueapiCall] = [] + queue: list[QueueTask], + history: list[QueueTask], + call_history: list[QueueBlueapiCall], + ) -> list[QueueBlueapiCall]: + call_list: list[QueueBlueapiCall] = [] for task in queue: with add_session_to_logger(task.experiment.instrument_session): match task.experiment: case TaskRequest(): call_list.append( - BlueapiCall( + QueueBlueapiCall( task_request=task.experiment, parent_task_id=task.id ) ) case Experiment(): call_list.extend( [ - BlueapiCall( + QueueBlueapiCall( task_request=b_api_task, parent_task_id=task.id ) for b_api_task in ( @@ -130,7 +129,7 @@ def construct_blueapi_calls( return call_list def pause_on_fail_blueapi_call( - self, call: BlueapiCall, errors: list[str | TaskError] | None + self, call: QueueBlueapiCall, errors: list[str | TaskError] | None ) -> bool: for error in errors or []: @@ -180,13 +179,13 @@ def _construct_blueapi_tasks_from_experiment( return plans def _add_required_auxiliary_scans( - self, current_task: TaskWithPosition | None, tasks: list[Task] + self, current_task: QueueTask | None, tasks: list[Task] ) -> list[Task]: """Adds auxiliary scan tasks to the queue. They will be added directly in front of the first task in the queue that requires them. Args: - current_task (TaskWithPosition | None): Current running task, if one exists. + current_task (QueueTask | None): Current running task, if one exists. tasks (list[Task]): Current list of tasks Returns: @@ -224,7 +223,7 @@ def _add_required_auxiliary_scans( def _ensure_auxiliary_in_queue_or_tiled( self, required_auxiliary: AuxiliaryScan, - current_task: TaskWithPosition | None, + current_task: QueueTask | None, new_tasks: list[Task], task_id: str, instrument_session: str, diff --git a/src/daq_queuing_service/task_queue/queue.py b/src/daq_queuing_service/task_queue/queue.py index 598bc4fd..8d307013 100644 --- a/src/daq_queuing_service/task_queue/queue.py +++ b/src/daq_queuing_service/task_queue/queue.py @@ -13,8 +13,8 @@ from daq_queuing_service.broadcaster import Broadcaster, Event from daq_queuing_service.external_interaction.blueapi.blueapi_call import ( BlueapiCall, - BlueapiCallResponse, CallStatus, + QueueBlueapiCall, ) from daq_queuing_service.external_interaction.runtime_estimator.client import ( RuntimeEstimatorClient, @@ -38,7 +38,7 @@ TaskNotFoundError, TaskNotInQueueError, ) -from daq_queuing_service.task_queue.task import Status, Task, TaskWithPosition +from daq_queuing_service.task_queue.task import QueueTask, Status, Task class TaskRegistry(dict[str, Task]): @@ -62,8 +62,8 @@ class QueueState(BaseModel): 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. + 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. @@ -221,21 +221,24 @@ def _sync(self): ] 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(), - ) + new_blueapi_calls = [ + BlueapiCall.from_parent(call) + for call in 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) + self._call_queue.extend(new_blueapi_calls) - for call in new_calls: + for call in new_blueapi_calls: # Add children to parent tasks if call.parent_task_id: self._tasks[call.parent_task_id].blueapi_calls.append(call) @@ -399,14 +402,14 @@ async def fail_call( call.fail(errors) self._call_history.append(call) - async def get_task_by_id(self, task_id: str) -> TaskWithPosition: + async def get_task_by_id(self, task_id: str) -> QueueTask: """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. + QueueTask: A copy of the task with a position field. Raises: TaskNotFoundError: Raised if the no task exists with the requested task ID. @@ -415,57 +418,57 @@ async def get_task_by_id(self, task_id: str) -> TaskWithPosition: async with self._lock: return self._get_task_by_id(task_id) - def _get_task_by_id(self, task_id: str) -> TaskWithPosition: + def _get_task_by_id(self, task_id: str) -> QueueTask: 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) + return QueueTask.from_task(task, position) - async def get_task_by_position(self, position: int) -> TaskWithPosition | None: + async def get_task_by_position(self, position: int) -> QueueTask | 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 + QueueTask | 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: + def _get_task_by_position(self, position: int) -> QueueTask | None: if position < -self.length or position >= self.length: return return self._get_task_by_id(self._queue[position]) - async def get_queue(self) -> list[TaskWithPosition]: + async def get_queue(self) -> list[QueueTask]: """Get the entire queue (not including history) Returns: - list[TaskWithPosition]: A list of the tasks in the queue, in the order they + list[QueueTask]: 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() - async def get_history(self) -> list[TaskWithPosition]: + async def get_history(self) -> list[QueueTask]: """Get the history list. Returns: - list[TaskWithPosition]: A list of the tasks in the history list, in + list[QueueTask]: 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() - async def get_tasks(self) -> list[TaskWithPosition]: + async def get_tasks(self) -> list[QueueTask]: """Get all the tasks in the queue and history list. Returns: - list[TaskWithPosition]: A list of tasks in chronological order, starting + list[QueueTask]: A list of tasks in chronological order, starting with the history. """ # Returns copies so don't have to be worried about caller modifying tasks. @@ -511,7 +514,7 @@ async def move_task(self, task_id: str, position: int) -> int: LOGGER.info(f"Successfully moved task {task_id} to position {new_position}") return new_position - async def cancel_tasks(self, task_ids: Sequence[str]) -> list[TaskWithPosition]: + async def cancel_tasks(self, task_ids: Sequence[str]) -> list[QueueTask]: """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. @@ -520,7 +523,7 @@ async def cancel_tasks(self, task_ids: Sequence[str]) -> list[TaskWithPosition]: task_ids (Sequence[str]): List of task IDs to cancel Returns: - list[TaskWithPosition]: List of the task objects that were removed from the + list[QueueTask]: List of the task objects that were removed from the queue. """ async with self._modifying: @@ -528,7 +531,7 @@ async def cancel_tasks(self, task_ids: Sequence[str]) -> list[TaskWithPosition]: LOGGER.info(f"Successfully cancelled tasks: {task_ids}") return cancelled_tasks - async def cancel_all_tasks(self) -> list[TaskWithPosition]: + async def cancel_all_tasks(self) -> list[QueueTask]: async with self._modifying: task_ids = [ id @@ -640,14 +643,14 @@ def _add_tasks(self, tasks: list[Task], position: int | None) -> None: for task in tasks: self._tasks[task.id] = task - def _cancel_tasks(self, task_ids: Sequence[str]) -> list[TaskWithPosition]: + def _cancel_tasks(self, task_ids: Sequence[str]) -> list[QueueTask]: 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] + return [QueueTask.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) @@ -695,33 +698,30 @@ def _validate_tasks_for_move_or_deletion(self, task_ids: list[str]): 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]: + def _get_queue_with_positions(self) -> list[QueueTask]: return [ - TaskWithPosition.from_task(self._tasks[task_id], i) + QueueTask.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 - ] + def _get_history(self) -> list[QueueTask]: + return [QueueTask.from_task(self._tasks[task_id]) for task_id in self._history] - async def get_call_queue(self) -> list[BlueapiCallResponse]: + async def get_call_queue(self) -> list[QueueBlueapiCall]: 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] + def _get_call_queue(self) -> list[QueueBlueapiCall]: + return [call.to_parent() for call in self._call_queue] - async def get_call_history(self) -> list[BlueapiCallResponse]: + async def get_call_history(self) -> list[QueueBlueapiCall]: 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_call_history(self) -> list[QueueBlueapiCall]: + return [call.to_parent() for call in self._call_history] - def _get_running_task(self) -> TaskWithPosition | None: + def _get_running_task(self) -> QueueTask | None: current_task = self._get_task_by_position(0) if current_task and current_task.status == Status.IN_PROGRESS: return current_task diff --git a/src/daq_queuing_service/task_queue/task.py b/src/daq_queuing_service/task_queue/task.py index 382cf43b..af496316 100644 --- a/src/daq_queuing_service/task_queue/task.py +++ b/src/daq_queuing_service/task_queue/task.py @@ -1,15 +1,15 @@ from enum import StrEnum -from typing import Any, Self +from typing import Any, Generic, Self, TypeVar from uuid import uuid4 from blueapi.service.model import TaskRequest -from pydantic import BaseModel, Field, computed_field +from pydantic import BaseModel, Field, PrivateAttr, computed_field from daq_queuing_service.app.authentication import User from daq_queuing_service.external_interaction.blueapi.blueapi_call import ( BlueapiCall, - BlueapiCallResponse, CallStatus, + QueueBlueapiCall, ) @@ -60,13 +60,52 @@ class TaskKind(StrEnum): PLAN = "Plan" -class Task(BaseModel): +T = TypeVar("T", bound=QueueBlueapiCall) + + +class BaseTask(BaseModel, Generic[T]): + """Base class for a queue task. Contains either a ulims experiment or blueapi task + request, and a list of blueapi calls to be made. + + Args: + Generic (QueueBlueapiCall | BlueapiCall): The type of BlueapiCall contained in + the blueapi_calls field. This is to differentiate between objects used by the + queue internally and ones copied to pass to queue plugins or the API layer. + """ + experiment: Experiment | TaskRequest - id: str = Field(default_factory=create_uuid_str) - blueapi_calls: list[BlueapiCall] = Field(default_factory=lambda: []) - user: User | None = None + id: str = Field(default_factory=create_uuid_str, frozen=True) + blueapi_calls: list[T] = Field(default_factory=lambda: []) + user: User | None = Field(default=None, frozen=True) estimated_runtime: int | None = None - _cancelled: bool = False + + @computed_field + @property + def kind(self) -> TaskKind: + match self.experiment: + case Experiment(): + return TaskKind.EXPERIMENT + case TaskRequest(): + return TaskKind.PLAN + + @computed_field + @property + def runtime(self) -> float | None: + # This should get nicer once we have separate types for queued, in progress and + # complete tasks. + if runtimes := [blueapi_call.runtime for blueapi_call in self.blueapi_calls]: + total = 0 + for runtime in runtimes: + if runtime is None: + return + total += runtime + return total + + +class Task(BaseTask[BlueapiCall]): + """Internal task object held by the queue.""" + + _cancelled: bool = PrivateAttr(False) def cancel(self): self._cancelled = True @@ -94,39 +133,12 @@ def status(self) -> Status: return Status.IN_PROGRESS return Status.QUEUED - @computed_field - @property - def kind(self) -> TaskKind: - match self.experiment: - case Experiment(): - return TaskKind.EXPERIMENT - case TaskRequest(): - return TaskKind.PLAN - - @computed_field - @property - def runtime(self) -> float | None: - # This should get nicer once we have separate types for queued, in progress and - # complete tasks. - if runtimes := [blueapi_call.runtime for blueapi_call in self.blueapi_calls]: - total = 0 - for runtime in runtimes: - if runtime is None: - return - total += runtime - return total +class QueueTask(BaseTask[QueueBlueapiCall]): + """Task object with a position field, for passing to a queue plugin or API layer.""" -class TaskWithPosition(BaseModel): - experiment: Experiment | TaskRequest - id: str - status: Status - blueapi_calls: list[BlueapiCallResponse] position: int | None - kind: TaskKind - user: User | None - estimated_runtime: int | None - runtime: float | None + status: Status @classmethod def from_task(cls, task: Task, position: int | None = None) -> Self: diff --git a/tests/unit_tests/conftest.py b/tests/unit_tests/conftest.py index bae0ca47..742f3e4d 100644 --- a/tests/unit_tests/conftest.py +++ b/tests/unit_tests/conftest.py @@ -21,8 +21,7 @@ ) from daq_queuing_service.broadcaster import Broadcaster from daq_queuing_service.external_interaction.blueapi.blueapi_call import ( - BlueapiCall, - BlueapiCallResponse, + QueueBlueapiCall, ) from daq_queuing_service.external_interaction.runtime_estimator.client import ( RuntimeEstimatorClient, @@ -36,9 +35,9 @@ ContainerPosition, Experiment, ExperimentDefinition, + QueueTask, Sample, Task, - TaskWithPosition, ) from daq_queuing_service.worker.worker import LOGGER as WORKER_LOGGER @@ -143,16 +142,16 @@ async def task_queue( @pytest.fixture async def task_queue_one_to_many(tasks: list[Task]): def construct_blueapi_calls( - queue: list[TaskWithPosition], - history: list[TaskWithPosition], - call_history: list[BlueapiCallResponse], - ) -> list[BlueapiCall]: - call_list: list[BlueapiCall] = [] + queue: list[QueueTask], + history: list[QueueTask], + call_history: list[QueueBlueapiCall], + ) -> list[QueueBlueapiCall]: + call_list: list[QueueBlueapiCall] = [] for task in queue: for _ in range(2): assert isinstance(task.experiment, Experiment) call_list.append( - BlueapiCall( + QueueBlueapiCall( parent_task_id=task.id, task_request=TaskRequest( name=task.experiment.experiment_definition.name, diff --git a/tests/unit_tests/plugins/i15_1/test_i15_1_converter.py b/tests/unit_tests/plugins/i15_1/test_i15_1_converter.py index 934eb708..bfbfc901 100644 --- a/tests/unit_tests/plugins/i15_1/test_i15_1_converter.py +++ b/tests/unit_tests/plugins/i15_1/test_i15_1_converter.py @@ -29,11 +29,11 @@ ContainerPosition, Experiment, ExperimentDefinition, + QueueTask, Sample, Status, Task, TaskKind, - TaskWithPosition, ) from ...conftest import make_sample @@ -193,12 +193,9 @@ def make_background_task(capillary: STANDARD_CAPILLARY, time_per_pdf: int) -> Ta ) -def assert_tasks_equal(task1: Task | TaskWithPosition, task2: Task | TaskWithPosition): +def assert_tasks_equal(task1: Task | QueueTask, task2: Task | QueueTask): # Check two tasks are equal other than the generated UUID - copy1 = type(task1).model_validate(task1) - copy2 = type(task2).model_validate(task2) - copy1.id = copy2.id = "" - assert task1 == task2 + assert task1.model_dump(exclude={"id"}) == task2.model_dump(exclude={"id"}) def create_task_dict(task_id: str, experiment: dict[str, Any]) -> dict[str, Any]: @@ -471,15 +468,13 @@ def test_experiment_with_correct_experiment_type_are_converted(): sample=make_sample("test_8_1", ""), instrument_session="cm12345-1", ) - task = TaskWithPosition( + task = QueueTask( experiment=experiment, id="1", status=Status.QUEUED, blueapi_calls=[], position=None, - kind=TaskKind.EXPERIMENT, user=None, - runtime=None, estimated_runtime=100, ) call_list = I151Converter().construct_blueapi_calls([task], [], []) @@ -645,15 +640,13 @@ def test_mix_of_experiments_with_correct_experiment_type_are_converted(): sample=make_sample("test_8_1", ""), instrument_session="cm12345-1", ) - good_task = TaskWithPosition( + good_task = QueueTask( experiment=good_experiment, id="1", status=Status.QUEUED, blueapi_calls=[], position=None, - kind=TaskKind.EXPERIMENT, user=None, - runtime=None, estimated_runtime=100, ) @@ -664,7 +657,6 @@ class BadExperiment: bad_task.experiment = BadExperiment() # type: ignore plan_task = deepcopy(good_task) plan_task.experiment = TaskRequest(name="", instrument_session="") - plan_task.kind = TaskKind.PLAN tasks = [good_task, bad_task, plan_task, good_task] call_list = I151Converter().construct_blueapi_calls(tasks, [], []) assert len(call_list) == 13 @@ -910,7 +902,7 @@ def test__ensure_background_in_queue_or_tiled_returns_if_current_task_is_suitabl new_tasks: list[Task] = [] result = i15_1_converter._ensure_auxiliary_in_queue_or_tiled( background, - TaskWithPosition.from_task(make_background_task("fq1.0", 25)), + QueueTask.from_task(make_background_task("fq1.0", 25)), new_tasks, "task_id", "cm12345-1", @@ -926,7 +918,7 @@ def test__ensure_background_in_queue_or_tiled_modifies_queued_background_if_poss pin=StandardsPin(capillary="fq1.0", contents=None), time_per_pdf=25, ) - current_task = TaskWithPosition.from_task(make_background_task("fq1.0", 25)) + current_task = QueueTask.from_task(make_background_task("fq1.0", 25)) # Current task is suitable but not a background assert isinstance(current_task.experiment, Experiment) current_task.experiment.experiment_definition.name = "Not a background" @@ -1071,7 +1063,7 @@ def test_no_robot_load_queued_if_turned_off_in_queue_params( queue_params.do_robot_load = False blueapi_calls = I151Converter().construct_blueapi_calls( - [TaskWithPosition.from_task(i15_1_tasks[0])], [], [] + [QueueTask.from_task(i15_1_tasks[0])], [], [] ) assert not any(call.task_request.name == "robot_load" for call in blueapi_calls) assert not any(call.task_request.name == "robot_unload" for call in blueapi_calls) @@ -1084,7 +1076,7 @@ def test_no_centring_queued_if_turned_off_in_queue_params( queue_params.do_centring = False blueapi_calls = I151Converter().construct_blueapi_calls( - [TaskWithPosition.from_task(i15_1_tasks[0])], [], [] + [QueueTask.from_task(i15_1_tasks[0])], [], [] ) assert not any(call.task_request.name == "centre_sample" for call in blueapi_calls) wait_for_beams = [ @@ -1100,6 +1092,6 @@ def test_no_wait_for_beam_queued_if_turned_off_in_queue_params( queue_params.do_wait_for_beam = False blueapi_calls = I151Converter().construct_blueapi_calls( - [TaskWithPosition.from_task(i15_1_tasks[0])], [], [] + [QueueTask.from_task(i15_1_tasks[0])], [], [] ) assert not any(call.task_request.name == "wait_for_beam" for call in blueapi_calls) diff --git a/tests/unit_tests/plugins/test_converter.py b/tests/unit_tests/plugins/test_converter.py index 24715322..29ff67ef 100644 --- a/tests/unit_tests/plugins/test_converter.py +++ b/tests/unit_tests/plugins/test_converter.py @@ -1,7 +1,9 @@ import pytest from blueapi.service.model import TaskRequest -from daq_queuing_service.external_interaction.blueapi.blueapi_call import BlueapiCall +from daq_queuing_service.external_interaction.blueapi.blueapi_call import ( + QueueBlueapiCall, +) from daq_queuing_service.plugins.converter import ( Converter, get_converter, @@ -9,8 +11,8 @@ from daq_queuing_service.task_queue.task import ( Experiment, ExperimentDefinition, + QueueTask, Task, - TaskWithPosition, ) from ..conftest import make_sample @@ -50,10 +52,10 @@ def test_default_converter_pre_process_does_not_modify_queue(tasks: list[Task]): def test_default_converter_construct_blueapi_calls_creates_one_blueapi_call_per_task( bluesky_tasks: list[Task], ): - task_copies = [TaskWithPosition.from_task(task) for task in bluesky_tasks] + task_copies = [QueueTask.from_task(task) for task in bluesky_tasks] blueapi_calls = Converter().construct_blueapi_calls(task_copies, [], []) for task, blueapi_call in zip(task_copies, blueapi_calls, strict=True): assert isinstance(task.experiment, TaskRequest) - assert blueapi_call == BlueapiCall( + assert blueapi_call == QueueBlueapiCall( task_request=task.experiment, parent_task_id=task.id ) diff --git a/tests/unit_tests/test_api.py b/tests/unit_tests/test_api.py index 1dd51775..1b659a8f 100644 --- a/tests/unit_tests/test_api.py +++ b/tests/unit_tests/test_api.py @@ -30,18 +30,17 @@ from daq_queuing_service.app.config import load_config from daq_queuing_service.broadcaster import Broadcaster, Event from daq_queuing_service.external_interaction.blueapi.blueapi_call import ( - BlueapiCallResponse, CallStatus, + QueueBlueapiCall, ) from daq_queuing_service.plugins.converter import Converter from daq_queuing_service.task_queue.queue import QUEUE_EVENTS, PauseReason, TaskQueue from daq_queuing_service.task_queue.queue_utils import QueueError from daq_queuing_service.task_queue.task import ( Experiment, + QueueTask, Status, Task, - TaskKind, - TaskWithPosition, ) @@ -410,7 +409,7 @@ async def test_add_tasks_to_queue_adds_to_queue_and_and_returns_task_ids( task_ids: list[str] = response.json() [uuid.UUID(task_id) for task_id in task_ids] - assert await task_queue_with_history.get_task_by_position(-1) == TaskWithPosition( + assert await task_queue_with_history.get_task_by_position(-1) == QueueTask( experiment=TaskRequest( name="add_tasks", params={"time": 10}, @@ -418,7 +417,7 @@ async def test_add_tasks_to_queue_adds_to_queue_and_and_returns_task_ids( ), id=task_ids[-1], blueapi_calls=[ - BlueapiCallResponse( + QueueBlueapiCall( task_request=TaskRequest( name="add_tasks", params={"time": 10}, instrument_session="abc" ), @@ -431,14 +430,11 @@ async def test_add_tasks_to_queue_adds_to_queue_and_and_returns_task_ids( blueapi_id=None, tiled_ids=[], scan_ids=[], - runtime=None, ) ], position=3, status=Status.QUEUED, - kind=TaskKind.PLAN, user=None, - runtime=None, estimated_runtime=100, ) @@ -495,10 +491,10 @@ async def test_if_sync_fails_after_tasks_added_then_contents_restored_and_error( class SomeError(Exception): ... def fail_conversion( - current_task: TaskWithPosition | None, + current_task: QueueTask | None, queue: list[Task], - history: list[TaskWithPosition], - call_history: list[BlueapiCallResponse], + history: list[QueueTask], + call_history: list[QueueBlueapiCall], ): raise SomeError("Conversion failed because xyz") @@ -630,7 +626,10 @@ async def test_move_task_moves_task_and_returns_position( "0", 5, 409, - {"error": "task_not_in_queue", "message": "Task 0 isn't present in queue"}, + { + "error": "task_not_in_queue", + "message": "Task 0 isn't present in queue", + }, ], [ "2", @@ -805,7 +804,10 @@ async def test_cancel_tasks_removes_task_from_queue_and_returns_tasks( [ TaskCancelRequest(task_ids=["1", "2", "3", "4"]).model_dump(), 409, - {"error": "task_not_in_queue", "message": "Task 1 isn't present in queue"}, + { + "error": "task_not_in_queue", + "message": "Task 1 isn't present in queue", + }, ], ], ) diff --git a/tests/unit_tests/test_queue.py b/tests/unit_tests/test_queue.py index 366d4736..af9c2229 100644 --- a/tests/unit_tests/test_queue.py +++ b/tests/unit_tests/test_queue.py @@ -13,17 +13,17 @@ from daq_queuing_service.broadcaster import Broadcaster from daq_queuing_service.external_interaction.blueapi.blueapi_call import ( BlueapiCall, - BlueapiCallResponse, CallStatus, + QueueBlueapiCall, ) from daq_queuing_service.plugins.converter import Converter, ConverterError from daq_queuing_service.task_queue.queue import ( PauseReason, QueueContents, QueueState, + QueueTask, TaskQueue, TaskRegistry, - TaskWithPosition, ) from daq_queuing_service.task_queue.queue_utils import ( NegativePositionError, @@ -37,7 +37,6 @@ ExperimentDefinition, Status, Task, - TaskKind, ) from .conftest import make_sample @@ -284,7 +283,7 @@ async def test_get_queue_only_returns_tasks_in_queue( assert task_queue_with_history._history == ["0", "1"] result = await task_queue_with_history.get_queue() assert result == [ - TaskWithPosition( + QueueTask( experiment=Experiment( name="test_experiment", instrument_session="", @@ -296,7 +295,7 @@ async def test_get_queue_only_returns_tasks_in_queue( id="2", status=Status.IN_PROGRESS, blueapi_calls=[ - BlueapiCallResponse( + QueueBlueapiCall( task_request=TaskRequest(name="test", instrument_session=""), parent_task_id="2", status=CallStatus.IN_PROGRESS, @@ -307,16 +306,13 @@ async def test_get_queue_only_returns_tasks_in_queue( blueapi_id=None, tiled_ids=[], scan_ids=[], - runtime=None, ) ], position=0, - kind=TaskKind.EXPERIMENT, user=None, - runtime=None, estimated_runtime=100, ), - TaskWithPosition( + QueueTask( experiment=Experiment( name="test_experiment", instrument_session="", @@ -328,7 +324,7 @@ async def test_get_queue_only_returns_tasks_in_queue( id="3", status=Status.QUEUED, blueapi_calls=[ - BlueapiCallResponse( + QueueBlueapiCall( task_request=TaskRequest(name="test", instrument_session=""), parent_task_id="3", status=CallStatus.WAITING, @@ -339,16 +335,13 @@ async def test_get_queue_only_returns_tasks_in_queue( blueapi_id=None, tiled_ids=[], scan_ids=[], - runtime=None, ) ], position=1, - kind=TaskKind.EXPERIMENT, user=None, - runtime=None, estimated_runtime=100, ), - TaskWithPosition( + QueueTask( experiment=Experiment( name="test_experiment", instrument_session="", @@ -360,7 +353,7 @@ async def test_get_queue_only_returns_tasks_in_queue( id="4", status=Status.QUEUED, blueapi_calls=[ - BlueapiCallResponse( + QueueBlueapiCall( task_request=TaskRequest(name="test", instrument_session=""), parent_task_id="4", status=CallStatus.WAITING, @@ -371,13 +364,10 @@ async def test_get_queue_only_returns_tasks_in_queue( blueapi_id=None, tiled_ids=[], scan_ids=[], - runtime=None, ) ], position=2, - kind=TaskKind.EXPERIMENT, user=None, - runtime=None, estimated_runtime=100, ), ] @@ -389,7 +379,7 @@ async def test_get_history_only_returns_tasks_in_history( assert task_queue_with_history._queue == ["2", "3", "4"] result = await task_queue_with_history.get_history() assert result == [ - TaskWithPosition( + QueueTask( experiment=Experiment( name="test_experiment", instrument_session="", @@ -401,7 +391,7 @@ async def test_get_history_only_returns_tasks_in_history( id="0", status=Status.ERROR, blueapi_calls=[ - BlueapiCallResponse( + QueueBlueapiCall( task_request=TaskRequest(name="test", instrument_session=""), parent_task_id="0", status=CallStatus.ERROR, @@ -418,16 +408,13 @@ async def test_get_history_only_returns_tasks_in_history( blueapi_id=None, tiled_ids=[], scan_ids=[], - runtime=59, ) ], position=None, - kind=TaskKind.EXPERIMENT, user=None, - runtime=59, estimated_runtime=100, ), - TaskWithPosition( + QueueTask( experiment=Experiment( name="test_experiment", instrument_session="", @@ -439,7 +426,7 @@ async def test_get_history_only_returns_tasks_in_history( id="1", status=Status.COMPLETE, blueapi_calls=[ - BlueapiCallResponse( + QueueBlueapiCall( task_request=TaskRequest(name="test", instrument_session=""), parent_task_id="1", status=CallStatus.SUCCESS, @@ -450,13 +437,10 @@ async def test_get_history_only_returns_tasks_in_history( blueapi_id=None, tiled_ids=[], scan_ids=[], - runtime=59, ) ], position=None, - kind=TaskKind.EXPERIMENT, user=None, - runtime=59, estimated_runtime=100, ), ] @@ -469,7 +453,7 @@ async def test_get_tasks_returns_tasks_in_queue_and_history( assert task_queue_with_history._history == ["0", "1"] result = await task_queue_with_history.get_tasks() assert result == [ - TaskWithPosition( + QueueTask( experiment=Experiment( name="test_experiment", instrument_session="", @@ -481,7 +465,7 @@ async def test_get_tasks_returns_tasks_in_queue_and_history( id="0", status=Status.ERROR, blueapi_calls=[ - BlueapiCallResponse( + QueueBlueapiCall( task_request=TaskRequest(name="test", instrument_session=""), parent_task_id="0", status=CallStatus.ERROR, @@ -498,16 +482,13 @@ async def test_get_tasks_returns_tasks_in_queue_and_history( blueapi_id=None, tiled_ids=[], scan_ids=[], - runtime=59, ) ], position=None, - kind=TaskKind.EXPERIMENT, user=None, - runtime=59, estimated_runtime=100, ), - TaskWithPosition( + QueueTask( experiment=Experiment( name="test_experiment", instrument_session="", @@ -519,7 +500,7 @@ async def test_get_tasks_returns_tasks_in_queue_and_history( id="1", status=Status.COMPLETE, blueapi_calls=[ - BlueapiCallResponse( + QueueBlueapiCall( task_request=TaskRequest(name="test", instrument_session=""), parent_task_id="1", status=CallStatus.SUCCESS, @@ -530,16 +511,13 @@ async def test_get_tasks_returns_tasks_in_queue_and_history( blueapi_id=None, tiled_ids=[], scan_ids=[], - runtime=59, ) ], position=None, - kind=TaskKind.EXPERIMENT, user=None, - runtime=59, estimated_runtime=100, ), - TaskWithPosition( + QueueTask( experiment=Experiment( name="test_experiment", instrument_session="", @@ -551,7 +529,7 @@ async def test_get_tasks_returns_tasks_in_queue_and_history( id="2", status=Status.IN_PROGRESS, blueapi_calls=[ - BlueapiCallResponse( + QueueBlueapiCall( task_request=TaskRequest(name="test", instrument_session=""), parent_task_id="2", status=CallStatus.IN_PROGRESS, @@ -562,16 +540,13 @@ async def test_get_tasks_returns_tasks_in_queue_and_history( blueapi_id=None, tiled_ids=[], scan_ids=[], - runtime=None, ) ], position=0, - kind=TaskKind.EXPERIMENT, user=None, - runtime=None, estimated_runtime=100, ), - TaskWithPosition( + QueueTask( experiment=Experiment( name="test_experiment", instrument_session="", @@ -583,7 +558,7 @@ async def test_get_tasks_returns_tasks_in_queue_and_history( id="3", status=Status.QUEUED, blueapi_calls=[ - BlueapiCallResponse( + QueueBlueapiCall( task_request=TaskRequest(name="test", instrument_session=""), parent_task_id="3", status=CallStatus.WAITING, @@ -594,16 +569,13 @@ async def test_get_tasks_returns_tasks_in_queue_and_history( blueapi_id=None, tiled_ids=[], scan_ids=[], - runtime=None, ) ], position=1, - kind=TaskKind.EXPERIMENT, user=None, - runtime=None, estimated_runtime=100, ), - TaskWithPosition( + QueueTask( experiment=Experiment( name="test_experiment", instrument_session="", @@ -615,7 +587,7 @@ async def test_get_tasks_returns_tasks_in_queue_and_history( id="4", status=Status.QUEUED, blueapi_calls=[ - BlueapiCallResponse( + QueueBlueapiCall( task_request=TaskRequest(name="test", instrument_session=""), parent_task_id="4", status=CallStatus.WAITING, @@ -626,13 +598,10 @@ async def test_get_tasks_returns_tasks_in_queue_and_history( blueapi_id=None, tiled_ids=[], scan_ids=[], - runtime=None, ) ], position=2, - kind=TaskKind.EXPERIMENT, user=None, - runtime=None, estimated_runtime=100, ), ] @@ -649,12 +618,8 @@ async def test_get_task_by_id_returns_task_in_queue_or_history( "4" in task_queue_with_history._queue and "4" not in task_queue_with_history._history ) - assert isinstance( - await task_queue_with_history.get_task_by_id("0"), TaskWithPosition - ) - assert isinstance( - await task_queue_with_history.get_task_by_id("4"), TaskWithPosition - ) + assert isinstance(await task_queue_with_history.get_task_by_id("0"), QueueTask) + assert isinstance(await task_queue_with_history.get_task_by_id("4"), QueueTask) async def test_get_task_by_id_raises_error_if_task_id_does_not_exist( @@ -668,7 +633,7 @@ async def test_get_task_by_pos_returns_task_in_queue( task_queue_with_history: TaskQueue, ): task = await task_queue_with_history.get_task_by_position(2) - assert isinstance(task, TaskWithPosition) + assert isinstance(task, QueueTask) assert task.position == 2 @@ -969,7 +934,7 @@ async def test_get_call_queue_returns_calls_in_call_queue( ): result = await task_queue_with_history.get_call_queue() assert result == [ - BlueapiCallResponse( + QueueBlueapiCall( task_request=TaskRequest(name="test", params={}, instrument_session=""), parent_task_id="2", status=CallStatus.IN_PROGRESS, @@ -980,9 +945,8 @@ async def test_get_call_queue_returns_calls_in_call_queue( blueapi_id=None, tiled_ids=[], scan_ids=[], - runtime=None, ), - BlueapiCallResponse( + QueueBlueapiCall( task_request=TaskRequest(name="test", params={}, instrument_session=""), parent_task_id="3", status=CallStatus.WAITING, @@ -993,9 +957,8 @@ async def test_get_call_queue_returns_calls_in_call_queue( blueapi_id=None, tiled_ids=[], scan_ids=[], - runtime=None, ), - BlueapiCallResponse( + QueueBlueapiCall( task_request=TaskRequest(name="test", params={}, instrument_session=""), parent_task_id="4", status=CallStatus.WAITING, @@ -1006,7 +969,6 @@ async def test_get_call_queue_returns_calls_in_call_queue( blueapi_id=None, tiled_ids=[], scan_ids=[], - runtime=None, ), ] @@ -1016,7 +978,7 @@ async def test_get_call_history_returns_calls_in_call_history( ): result = await task_queue_with_history.get_call_history() assert result == [ - BlueapiCallResponse( + QueueBlueapiCall( task_request=TaskRequest(name="test", params={}, instrument_session=""), parent_task_id="0", status=CallStatus.ERROR, @@ -1031,9 +993,8 @@ async def test_get_call_history_returns_calls_in_call_history( blueapi_id=None, tiled_ids=[], scan_ids=[], - runtime=59, ), - BlueapiCallResponse( + QueueBlueapiCall( task_request=TaskRequest(name="test", params={}, instrument_session=""), parent_task_id="1", status=CallStatus.SUCCESS, @@ -1044,7 +1005,6 @@ async def test_get_call_history_returns_calls_in_call_history( blueapi_id=None, tiled_ids=[], scan_ids=[], - runtime=59, ), ] @@ -1252,9 +1212,9 @@ async def test_if_error_during_conversion_then_error_handled_and_contents_restor task_queue: TaskQueue, ): def convert( - queue: list[TaskWithPosition], - history: list[TaskWithPosition], - call_history: list[BlueapiCallResponse], + queue: list[QueueTask], + history: list[QueueTask], + call_history: list[QueueBlueapiCall], ): for task_id in task_queue._queue: del task_queue._tasks[task_id] @@ -1275,9 +1235,9 @@ async def test_if_error_during_conversion_then__restore_latest_good_contents_cal task_queue: TaskQueue, ): def convert( - queue: list[TaskWithPosition], - history: list[TaskWithPosition], - call_history: list[BlueapiCallResponse], + queue: list[QueueTask], + history: list[QueueTask], + call_history: list[QueueBlueapiCall], ): raise ValueError("Conversion failed")