245 lines
8.9 KiB
Python
245 lines
8.9 KiB
Python
from __future__ import annotations
|
|
|
|
import time
|
|
from typing import Any
|
|
|
|
import redis
|
|
|
|
from aare.common.automation_models import (
|
|
QueueItem,
|
|
WorkflowEvent,
|
|
ControlState,
|
|
RuntimeState,
|
|
QueueItemStatus,
|
|
WorkflowStepRecord,
|
|
WorkflowStateKind,
|
|
)
|
|
|
|
CURRENT_QUEUE_LIMIT = 576
|
|
|
|
def build_default_steps() -> list[WorkflowStepRecord]:
|
|
return [
|
|
WorkflowStepRecord(kind=WorkflowStateKind.MOUNT.value),
|
|
WorkflowStepRecord(kind=WorkflowStateKind.LOOP_CENTRE.value),
|
|
WorkflowStepRecord(kind=WorkflowStateKind.RASTER.value),
|
|
WorkflowStepRecord(kind=WorkflowStateKind.DATA_COLLECTION.value),
|
|
]
|
|
|
|
|
|
class WorkflowRedisManager:
|
|
def __init__(self, client: redis.Redis, beamline: str):
|
|
self._client = client
|
|
self._bl = beamline.lower()
|
|
|
|
def _key(self, suffix: str) -> str:
|
|
return f"{self._bl}:workflow:{suffix}"
|
|
|
|
def _item_key(self, item_id: str) -> str:
|
|
return self._key(f"item:{item_id}")
|
|
|
|
def _queue_key(self) -> str:
|
|
return self._key("queue")
|
|
|
|
def _runtime_key(self) -> str:
|
|
return self._key("runtime")
|
|
|
|
def _control_key(self) -> str:
|
|
return self._key("control")
|
|
|
|
def _events_key(self) -> str:
|
|
return self._key("events")
|
|
|
|
def _next_item_id(self) -> str:
|
|
n = int(self._client.incr(self._key("item_seq")))
|
|
return f"wf_{int(time.time())}_{n:06d}"
|
|
|
|
# ─────────────────────────────────────────────
|
|
# Queue item operations
|
|
# ─────────────────────────────────────────────
|
|
|
|
def create_item(self, item: QueueItem) -> QueueItem:
|
|
if not item.item_id:
|
|
item.item_id = self._next_item_id()
|
|
|
|
current_count = self._client.zcard(self._queue_key())
|
|
if current_count >= CURRENT_QUEUE_LIMIT:
|
|
raise ValueError(f"Queue size limit reached ({CURRENT_QUEUE_LIMIT} items). Please clear the queue.")
|
|
|
|
pipe = self._client.pipeline(transaction=True)
|
|
pipe.set(self._item_key(item.item_id), item.model_dump_json())
|
|
pipe.zadd(self._queue_key(), {item.item_id: float(item.order_index)})
|
|
pipe.execute()
|
|
|
|
self.append_event(WorkflowEvent(
|
|
beamline=self._bl,
|
|
item_id=item.item_id,
|
|
event_type="item_created",
|
|
message="Queue item created",
|
|
payload={"status": item.status.value},
|
|
))
|
|
|
|
return item
|
|
|
|
def get_item(self, item_id: str) -> QueueItem | None:
|
|
raw = self._client.get(self._item_key(item_id))
|
|
if raw is None:
|
|
return None
|
|
return QueueItem.model_validate_json(raw)
|
|
|
|
def update_item(self, item_id: str, patch: dict[str, Any]) -> QueueItem:
|
|
item = self.get_item(item_id)
|
|
if item is None:
|
|
raise KeyError(f"Queue item not found: {item_id}")
|
|
|
|
updated = item.model_copy(update=patch)
|
|
self._client.set(self._item_key(item_id), updated.model_dump_json())
|
|
return updated
|
|
|
|
def delete_item(self, item_id: str) -> None:
|
|
pipe = self._client.pipeline(transaction=True)
|
|
pipe.delete(self._item_key(item_id))
|
|
pipe.zrem(self._queue_key(), item_id)
|
|
pipe.execute()
|
|
|
|
def list_queue_order(self) -> list[str]:
|
|
return [str(x) for x in self._client.zrange(self._queue_key(), 0, -1)]
|
|
|
|
def list_items(self, *, include_finished: bool = True) -> list[QueueItem]:
|
|
out: list[QueueItem] = []
|
|
for item_id in self.list_queue_order():
|
|
item = self.get_item(item_id)
|
|
if item is None:
|
|
continue
|
|
if not include_finished and item.status in {
|
|
QueueItemStatus.COMPLETED,
|
|
QueueItemStatus.FAILED,
|
|
QueueItemStatus.SKIPPED,
|
|
QueueItemStatus.ABORTED,
|
|
}:
|
|
continue
|
|
out.append(item)
|
|
return out
|
|
|
|
def get_next_pending_item(self) -> QueueItem | None:
|
|
for item_id in self.list_queue_order():
|
|
item = self.get_item(item_id)
|
|
if item is not None and item.status == QueueItemStatus.PENDING:
|
|
return item
|
|
return None
|
|
|
|
def move_item(self, item_id: str, new_order_index: int) -> None:
|
|
if self.get_item(item_id) is None:
|
|
raise KeyError(f"Queue item not found: {item_id}")
|
|
self._client.zadd(self._queue_key(), {item_id: float(new_order_index)})
|
|
self.update_item(item_id, {"order_index": new_order_index})
|
|
|
|
# ─────────────────────────────────────────────
|
|
# Runtime state
|
|
# ─────────────────────────────────────────────
|
|
|
|
def get_runtime(self) -> RuntimeState:
|
|
raw = self._client.get(self._runtime_key())
|
|
if raw is None:
|
|
return RuntimeState()
|
|
return RuntimeState.model_validate_json(raw)
|
|
|
|
def set_runtime(self, runtime: RuntimeState) -> RuntimeState:
|
|
runtime.last_update = time.time()
|
|
self._client.set(self._runtime_key(), runtime.model_dump_json())
|
|
return runtime
|
|
|
|
def patch_runtime(self, patch: dict[str, Any]) -> RuntimeState:
|
|
runtime = self.get_runtime()
|
|
updated = runtime.model_copy(update=patch)
|
|
return self.set_runtime(updated)
|
|
|
|
# ─────────────────────────────────────────────
|
|
# Control state
|
|
# ─────────────────────────────────────────────
|
|
|
|
def get_control(self) -> ControlState:
|
|
raw = self._client.get(self._control_key())
|
|
if raw is None:
|
|
return ControlState()
|
|
return ControlState.model_validate_json(raw)
|
|
|
|
def request_control(self, patch: dict[str, Any], *, requested_by: str) -> ControlState:
|
|
control = self.get_control()
|
|
updated = control.model_copy(update={
|
|
**patch,
|
|
"requested_by": requested_by,
|
|
"requested_at": time.time(),
|
|
})
|
|
self._client.set(self._control_key(), updated.model_dump_json())
|
|
|
|
self.append_event(WorkflowEvent(
|
|
beamline=self._bl,
|
|
item_id=self.get_runtime().current_item_id,
|
|
event_type="control_requested",
|
|
actor=requested_by,
|
|
payload=patch,
|
|
))
|
|
|
|
return updated
|
|
|
|
def clear_control(self) -> ControlState:
|
|
control = ControlState()
|
|
self._client.set(self._control_key(), control.model_dump_json())
|
|
return control
|
|
|
|
# ─────────────────────────────────────────────
|
|
# Event log
|
|
# ─────────────────────────────────────────────
|
|
|
|
def append_event(self, event: WorkflowEvent) -> None:
|
|
payload = event.model_dump_json()
|
|
self._client.xadd(self._events_key(), {"json": payload}, maxlen=5000, approximate=True)
|
|
|
|
def read_events(self, *, limit: int = 200) -> list[WorkflowEvent]:
|
|
rows = self._client.xrevrange(self._events_key(), count=limit)
|
|
out: list[WorkflowEvent] = []
|
|
for _, fields in reversed(rows):
|
|
raw = fields.get("json")
|
|
if raw:
|
|
out.append(WorkflowEvent.model_validate_json(raw))
|
|
return out
|
|
|
|
# ─────────────────────────────────────────────
|
|
# Step status updates
|
|
# ─────────────────────────────────────────────
|
|
|
|
def update_step(
|
|
self,
|
|
item_id: str,
|
|
step_index: int,
|
|
*,
|
|
status: str | None = None,
|
|
message: str | None = None,
|
|
error_detail: str | None = None,
|
|
) -> QueueItem:
|
|
item = self.get_item(item_id)
|
|
if item is None:
|
|
raise KeyError(f"Queue item not found: {item_id}")
|
|
|
|
if not (0 <= step_index < len(item.steps)):
|
|
raise IndexError(f"Step index out of range: {step_index}")
|
|
|
|
step = item.steps[step_index]
|
|
|
|
if status is not None:
|
|
step.status = status
|
|
if status == "running" and step.started_at is None:
|
|
step.started_at = time.time()
|
|
elif status in ("completed", "failed", "skipped"):
|
|
step.completed_at = time.time()
|
|
|
|
if message is not None:
|
|
step.message = message
|
|
|
|
if error_detail is not None:
|
|
step.error_detail = error_detail
|
|
|
|
item.steps[step_index] = step
|
|
self._client.set(self._item_key(item_id), item.model_dump_json())
|
|
|
|
return item |