"""
Each task process will have a ``LabView`` instance, through which it can request
the lab resources (devices and sample positions).
It can also update the position of a sample in the lab.
"""
import sys
import threading
import time
from contextlib import contextmanager, suppress
from traceback import format_exc
from typing import Any
from bson import ObjectId
from alab_management.device_manager import DevicesClient
from alab_management.device_view.device import BaseDevice
from alab_management.experiment_view.experiment_view import ExperimentView
from alab_management.logger import DBLogger
from alab_management.resource_manager.resource_requester import ResourceRequester
from alab_management.sample_view.sample import Sample
from alab_management.sample_view.sample_view import SampleView
from alab_management.task_view.task import BaseTask
from alab_management.task_view.task_enums import TaskPriority, TaskStatus
from alab_management.task_view.task_view import TaskView
from alab_management.user_input import request_user_input, request_user_input_with_note
[docs]
class DeviceRunningException(Exception):
"""Raise when a task try to release a device that is still running."""
[docs]
class LabView:
"""
LabView is a wrapper over device view and sample view.
A task can get access to that to request resources, query sample and
update sample positions.
"""
def __init__(self, task_id: ObjectId):
self._task_view = TaskView()
self.__task_entry = self._task_view.get_task(
task_id=task_id
) # will throw error if task_id does not exist
self._experiment_view = ExperimentView()
self._task_id = task_id
self._sample_view = SampleView()
self._resource_requester = ResourceRequester(task_id=task_id)
self._device_client = DevicesClient(task_id=task_id, timeout=None)
self.logger = DBLogger(task_id=task_id)
self._priority = TaskPriority.NORMAL.value
self.__cancellation_event: threading.Event | None = None
self.__cancellation_lock = threading.Lock()
@property
def task_id(self) -> ObjectId:
"""Get the task id of the current task."""
return self._task_id
[docs]
@contextmanager
def request_resources(
self,
resource_request: dict[type[BaseDevice] | str | None, dict[str, str | int]],
priority: int | None = None,
timeout: float | None = None,
exact_positions: set[str] | None = None,
):
"""
Request devices and sample positions. This function is a context manager, which should be used in
a with statement to ensure all the devices are released when the task is done.
resource_request format is: {device: {position: number, ...}, ...} device can be a name of a specific device
(str), a type of device, or None. If device is a type, the resource request will look for any available
device of that type. If device is None, the resource request will look for sample positions that do not
belong to a device. position is the name of a sample position that should be reserved on the device,
and number is the number of such positions that should be reserved. If the device is required but no
positions are required, this can be left as an empty dictionary.
Examples -------- {TubeFurnace: {"tray": 4}, "arm1": {}} will find the first available TubeFurnace device,
then reserve 4 sample positions of "{tubefurnacename}/tray/{tray_index}" on that device. It will also find
the device named "arm1".
By default, position names are matched by prefix (e.g., "input_rack/slot/1" would match both
"input_rack/slot/1" and "input_rack/slot/10"). To match exact positions, use the ``exact_positions``
parameter to specify which position names should be matched exactly.
Args:
resource_request: Dictionary mapping devices to position requests
priority: Optional priority for the request (0-40, default 20). Higher number = higher priority.
Numbers >= 100 are reserved for urgent/error correcting requests.
timeout: Optional timeout for the request in seconds
exact_positions: Set of position names that should be matched exactly (not by prefix).
Position names should be specified as they appear in the resource_request dictionary,
before device prefixes are added. For example:
- For {None: {"input_rack/slot/1": 1}}, use exact_positions={"input_rack/slot/1"}
- For {Furnace: {"slot/1": 1}}, use exact_positions={"slot/1"} (device prefix added automatically)
**Important**: Exact matching only works when number=1. If you request multiple positions
(e.g., {"slot": 11}), you cannot use exact matching for that position. Use exact_positions
only for positions where number=1.
If None (default), all positions use prefix matching for backward compatibility.
Example:
.. code-block:: python
# Prefix matching (default) - "slot/1" could match "slot/1" or "slot/10"
with self.lab_view.request_resources({Furnace: {"slot/1": 1}}) as (devices, positions):
...
# Exact matching - only matches "slot/1"
with self.lab_view.request_resources(
{Furnace: {"slot/1": 1}},
exact_positions={"slot/1"}
) as (devices, positions):
...
# Mix of prefix and exact matching
with self.lab_view.request_resources(
{None: {"input_rack/slot/1": 1, "input_rack/slot/2": 1}},
exact_positions={"input_rack/slot/1"} # Only slot/1 is exact
) as (devices, positions):
# slot/1 is exact, slot/2 uses prefix matching
...
# Invalid: Cannot use exact matching with number > 1
# This will raise ValueError:
# with self.lab_view.request_resources(
# {"BFT_input_rack": {"slot": 11}},
# exact_positions={"slot"} # ERROR: exact matching requires number=1
# ) as (devices, positions):
# ...
"""
priority = priority or self.priority
self._task_view.update_status(
task_id=self.task_id, status=TaskStatus.REQUESTING_RESOURCES
)
result = self._resource_requester.request_resources(
resource_request=resource_request,
timeout=timeout,
priority=priority,
exact_positions=exact_positions,
)
request_id = result["request_id"]
devices = result["devices"]
sample_positions = result["sample_positions"]
devices = {
device_type: self._device_client.create_device_wrapper(device_name)
for device_type, device_name in devices.items()
} # type: ignore
self._task_view.update_status(task_id=self.task_id, status=TaskStatus.RUNNING)
yield devices, sample_positions
self._resource_requester.release_resources(request_id=request_id)
def _sample_name_to_id(self, sample_name: str) -> ObjectId:
"""
Get a sample id by name.
Looks up sample name->id mapping for the experiment `self.task_id` belongs to.
"""
for sample in self.__task_entry["samples"]:
if sample["name"] == sample_name:
return sample["sample_id"]
raise ValueError(
f"No sample with name \"{sample_name}\" found for task \"{self.__task_entry['type']}\""
)
[docs]
def get_sample(self, sample: ObjectId | str) -> Sample:
"""
Get a sample by either an ObjectId corresponding to sample_id, or as a string corresponding to the sample's
name within the experiment., see also :py:meth:`get_sample
<alab_management.sample_view.sample_view.SampleView.get_sample>`.
"""
if isinstance(sample, str):
sample_id = self._sample_name_to_id(sample)
elif isinstance(sample, ObjectId):
sample_id = sample
else:
raise TypeError("sample must be a sample name (str) or id (ObjectId)")
return self._sample_view.get_sample(sample_id=sample_id)
[docs]
def move_sample(self, sample: ObjectId | str, position: str | None):
"""
Move a sample to a new position. `sample` can be given as either an ObjectId corresponding to sample_id,
or as a string corresponding to the sample's name within the experiment.
See Also
--------
:py:meth:`move_sample <alab_management.sample_view.sample_view.SampleView.move_sample>`
"""
# check if this sample position is locked by current task
if (
position is not None
and self._sample_view.get_sample_position_status(position)[1]
!= self._task_id
):
# Wait a few seconds for the sample position to be locked in case it is not
# locked by the current task yet.
for _ in range(5):
time.sleep(1)
if (
self._sample_view.get_sample_position_status(position)[1]
== self._task_id
):
break
if (
self._sample_view.get_sample_position_status(position)[1]
!= self._task_id
):
raise ValueError(
f"Cannot move sample to the new sample position ({position}) without locking it."
)
# check if this sample is owned by current task
sample_entry = self.get_sample(sample=sample)
if sample_entry.task_id != self._task_id:
raise ValueError("Cannot move a sample that does not belong to this task.")
self._sample_view.move_sample(
sample_id=sample_entry.sample_id, position=position
)
[docs]
def set_sample_in_transit(
self, sample: ObjectId | str, source: str | None, destination: str | None
):
"""
Record that a sample is being physically moved from ``source`` to ``destination``.
Call this right before a robot move begins. It does not change the recorded position, so if
the move crashes mid-transfer the last known position plus intended destination remain
visible. A subsequent successful :py:meth:`move_sample` clears the record automatically.
See Also
--------
:py:meth:`set_sample_in_transit <alab_management.sample_view.sample_view.SampleView.set_sample_in_transit>`
"""
# check if this sample is owned by current task
sample_entry = self.get_sample(sample=sample)
if sample_entry.task_id != self._task_id:
raise ValueError(
"Cannot set in-transit for a sample that does not belong to this task."
)
self._sample_view.set_sample_in_transit(
sample_id=sample_entry.sample_id, source=source, destination=destination
)
[docs]
def get_locked_sample_positions(self) -> list[str]:
"""Get a list of sample positions that are occupied by this task."""
return self._sample_view.get_sample_positions_by_task(task_id=self._task_id)
[docs]
def lock_sample_position(self, position: str):
"""
Lock an exact sample position for this task.
Unlike ``request_resources``, this method locks an exact position name rather than
using prefix matching. This is useful when you need to lock a specific position like
"input_rack/slot/1" without it potentially matching "input_rack/slot/10".
Args:
position: The exact name of the sample position to lock (e.g., "input_rack/slot/1")
Raises
------
ValueError: If the position is currently occupied or locked by another task.
Example:
.. code-block:: python
# Lock an exact position
self.lab_view.lock_sample_position("input_rack/slot/1")
# Use the position
self.lab_view.move_sample(sample=self.sample, position="input_rack/slot/1")
# Release when done
self.lab_view.release_sample_position("input_rack/slot/1")
"""
self._sample_view.lock_sample_position(task_id=self._task_id, position=position)
[docs]
def release_sample_position(self, position: str):
"""
Release a locked sample position.
Args:
position: The exact name of the sample position to release
Raises
------
ValueError: If the position is invalid
Example:
.. code-block:: python
self.lab_view.lock_sample_position("input_rack/slot/1")
try:
# Use the position
self.lab_view.move_sample(sample=self.sample, position="input_rack/slot/1")
finally:
self.lab_view.release_sample_position("input_rack/slot/1")
"""
# Verify the position is locked by this task before releasing
_status, locked_by_task_id = self._sample_view.get_sample_position_status(
position
)
if locked_by_task_id != self._task_id:
raise ValueError(
f"Cannot release position {position} - it is not locked by this task "
f"(locked by: {locked_by_task_id})"
)
self._sample_view.release_sample_position(position=position)
[docs]
@contextmanager
def lock_exact_sample_positions(self, positions: list[str]):
"""
Lock exact sample positions as a context manager.
This method locks exact position names (not prefixes) and automatically releases
them when exiting the context. This is useful when you need to lock specific
positions like "input_rack/slot/1" without prefix matching potentially selecting
"input_rack/slot/10".
Args:
positions: List of exact sample position names to lock
Yields
------
list[str]: The list of locked positions (same as input)
Raises
------
ValueError: If any position is currently occupied or locked by another task.
Example:
.. code-block:: python
# Lock exact positions
with self.lab_view.lock_exact_sample_positions([
"input_rack/slot/1",
"input_rack/slot/2"
]) as locked_positions:
# Use the positions
self.lab_view.move_sample(
sample=self.sample,
position=locked_positions[0]
)
# Positions are automatically released here
"""
locked_positions = []
try:
for position in positions:
self.lock_sample_position(position)
locked_positions.append(position)
yield locked_positions
finally:
# Release all positions that were successfully locked
for position in locked_positions:
with suppress(ValueError):
# Position might have been released already or doesn't exist
# Continue releasing other positions
self.release_sample_position(position)
[docs]
def get_sample_position_parent_device(self, position: str) -> str | None:
"""Get the name of the device that owns the sample position."""
return self._sample_view.get_sample_position_parent_device(position=position)
[docs]
def run_subtask(
self, task: type[BaseTask], samples: list[ObjectId | str], **kwargs
):
"""
Run a task as a subtask within the task. basically fills in task_id and lab_view for you.
this command blocks until the subtask is completed.
Args:
task (Type[BaseTask]): The type/class of the Task to run.
samples (List[Union[ObjectId, str]]): List of sample IDs or names.
**kwargs: will be passed to the Task method via the parameters entry in the task collection.
"""
if not issubclass(task, BaseTask):
raise TypeError("task must be a subclass of BaseTask!")
# TODO maybe check if task is in task_registry? for future if tasks are somehow checked when adding to registry
# task_id and lab_view kwargs forced to match that of current LabView instance
kwargs.pop("task_id", None)
kwargs.pop("lab_view", None)
task_id = self._task_id
lab_view = self
subtask_id = self._task_view.create_subtask(
task_id=task_id,
subtask_type=task.__name__,
samples=samples,
parameters=kwargs,
)
try:
subtask: BaseTask = task(
_offline_mode=False,
task_id=task_id,
lab_view=lab_view,
samples=samples,
**kwargs,
)
except Exception as exc:
self._task_view.update_subtask_status(
task_id=task_id, subtask_id=subtask_id, status=TaskStatus.ERROR
)
self._task_view.update_subtask_result(
task_id=task_id, subtask_id=subtask_id, result=str(exc)
)
raise Exception(
"Failed to create subtask of type {} within task {} of type {}".format(
task,
task_id,
self._task_view.get_task(task_id=task_id, encode=True)["type"],
)
) from exc
self.logger.system_log(
level="INFO",
log_data={
"logged_by": "TaskActor",
"type": "SubTaskStart",
"task_id": task_id,
"subtask_type": task.__name__,
},
)
try:
self._task_view.update_subtask_status(
task_id=task_id, subtask_id=subtask_id, status=TaskStatus.RUNNING
)
result = subtask.run() # block until completion
except Exception as exception:
self._task_view.update_subtask_status(
task_id=task_id, subtask_id=subtask_id, status=TaskStatus.ERROR
)
self._task_view.update_subtask_result(
task_id=task_id, subtask_id=subtask_id, result=str(exception)
)
self.logger.system_log(
level="ERROR",
log_data={
"logged_by": "TaskActor",
"type": "SubTaskEnd",
"task_id": task_id,
"subtask_type": task.__name__,
"status": "ERROR",
"traceback": format_exc(),
},
)
raise
else:
self._task_view.update_subtask_status(
task_id=task_id, subtask_id=subtask_id, status=TaskStatus.COMPLETED
)
self._task_view.update_subtask_result(
task_id=task_id, subtask_id=subtask_id, result=result
)
self.logger.system_log(
level="INFO",
log_data={
"logged_by": "TaskActor",
"type": "SubTaskEnd",
"task_id": task_id,
"subtask_type": task.__name__,
"status": "COMPLETED",
},
)
return result
@property
def priority(self) -> int:
"""Get the priority of the task."""
return self._priority
@priority.setter
def priority(self, priority: int):
if isinstance(priority, TaskPriority):
priority = priority.value
self._priority = priority
[docs]
def update_result(self, name: str, value: Any):
"""
Update a result of the task. This result will be saved in the task collection under `results.name` and can be
retrieved later.
Args: name (str): name of the result (ie "diffraction pattern"). This will be used as the key in the results
dictionary. value (Any): value of the result. This can be a numpy array, a set, or any other
bson-serializable object (most standard Python types).
"""
self._task_view.update_result(task_id=self.task_id, name=name, value=value)
[docs]
def request_cleanup(self, error_message: str | None = None):
"""Request cleanup of the task. This function will block until the task is cleaned up.
Args:
error_message: A pre-formatted, debuggable error report to show the operator (and send to
Slack/email). If not provided, a report is built from the exception currently being
handled so the prompt always identifies what failed, where (file/line/function), and
the full traceback.
"""
all_reserved_sample_positions = self._sample_view.get_sample_positions_by_task(
self.task_id
)
all_samples = self.__task_entry["samples"]
all_positions_with_samples = [
self._sample_view.get_sample(sample_entry["sample_id"]).position
for sample_entry in all_samples
]
all_positions_with_samples = [
each for each in all_positions_with_samples if each
]
if error_message is None and sys.exc_info()[0] is not None:
# Only build a report when there is an exception actively being handled (e.g. called from
# the task actor's except block). Planned cleanups (e.g. on restart) have no exception.
from alab_management.utils.error_context import format_error_report
error_message = format_error_report(
task_id=self.task_id,
samples=[sample_entry["name"] for sample_entry in all_samples],
header="Unrecoverable error",
)
prompt = (
"An unrecoverable error has occurred.\n"
f"(1) remove samples on {', '.join(all_positions_with_samples)}\n"
f"(2) remove all other consumables on {', '.join(all_reserved_sample_positions)}"
)
if error_message:
prompt += f"\n\n{error_message}"
self.request_user_input(prompt=prompt, options=["OK"])
# move the samples out of the lab
for sample in all_samples:
self.move_sample(sample=sample["sample_id"], position=None)
# release all the resource that has not been fulfilled
self._resource_requester.release_all_resources()
[docs]
def release_all_resources(self):
"""Give back every resource this task holds or has requested.
``request_cleanup`` already does this as its last step. Call it directly when a task has
reconciled its own samples and only needs its devices and sample positions released, which
is what a task with ``cleanup_on_cancel = False`` needs after being cancelled.
"""
self._resource_requester.release_all_resources()
[docs]
def is_cancelling(self) -> bool:
"""Whether someone has asked for this task to be cancelled.
A cancellation reaches a task as a ``dramatiq_abort.Abort`` raised in the worker thread, and
that cannot interrupt a thread already blocked inside a device call or a ``time.sleep``. Poll
this in long waits, and bail out cooperatively, so a cancel takes effect at the next safe
point instead of after the blocking call finally returns.
.. code-block:: python
while furnace.is_running():
if self.lab_view.is_cancelling():
raise TaskCancelledError("cancelled while waiting for the furnace")
time.sleep(5)
"""
return self._task_view.is_canceling(task_id=self._task_id)
@property
def cancellation_event(self) -> threading.Event:
"""A :class:`threading.Event` that becomes set once this task is being cancelled.
Backed by a daemon thread polling the database once a second, started on first access, so a
task that never asks for it pays nothing. Use it where an ``Event`` is more convenient than
polling :meth:`is_cancelling` -- most usefully as a cancellable sleep, which wakes early
instead of sitting out the full delay:
.. code-block:: python
# returns True as soon as the task is cancelled, False if the full 30s elapsed
if self.lab_view.cancellation_event.wait(timeout=30):
raise TaskCancelledError("cancelled while waiting")
"""
with self.__cancellation_lock:
if self.__cancellation_event is None:
self.__cancellation_event = threading.Event()
watcher = threading.Thread(
target=self.__watch_for_cancellation,
name=f"cancellation-watcher-{self._task_id}",
daemon=True,
)
watcher.start()
return self.__cancellation_event
def __watch_for_cancellation(self):
"""Set the cancellation event once the database says this task is being cancelled."""
event = self.__cancellation_event
if event is None: # pragma: no cover - only reachable if the event was never created
return
while not event.is_set():
try:
if self._task_view.is_canceling(task_id=self._task_id):
event.set()
return
except Exception: # noqa: BLE001
# A transient database hiccup must not kill the watcher; try again next tick.
pass
time.sleep(1)