"""Agent interface for asynchronous optimization using Bluesky Queueserver and Ax."""
from collections.abc import Mapping, Sequence
from concurrent.futures import Future
from typing import Any
try:
import bluesky_queueserver_api.http
import bluesky_queueserver_api.zmq
except ImportError as e:
raise ImportError(
"The queueserver integration requires additional dependencies. "
"Install them with: pip install bluesky-adaptive[queueserver]"
) from e
from bluesky.callbacks.zmq import RemoteDispatcher
from ..protocols import (
Actuator,
EvaluationFunction,
QueueserverOptimizationProblem,
)
from ..queueserver import OptimizationResult, QueueserverClient, QueueserverOptimizationRunner
from .agent import _AxAgentMixin
from .dof import DOF, DOFConstraint
from .objective import Objective, OutcomeConstraint, to_ax_objective_str
from .optimizer import AxOptimizer
[docs]
class QueueserverAgent(_AxAgentMixin):
"""
An asynchronous interface that uses Ax as the backend for optimization and experiment tracking.
Uses the bluesky-queueserver-api for scheduling plan execution.
.. warning::
This class is **experimental**. The API is not yet stable and may change
in future releases without a deprecation period. It is not recommended for
production use.
Parameters
----------
re_manager_api : bluesky_queueserver_api.zmq.REManagerAPI | bluesky_queueserver_api.http.REManagerAPI
The manager API for interaction with Bluesky queueserver.
document_dispatcher : RemoteDispatcher
Dispatcher for consuming Bluesky documents from the remote server.
sensors : Sequence[str]
The sensors to use for acquisition. These should be the minimal set
of sensors that are needed to compute the objectives.
dofs : Sequence[DOF]
The degrees of freedom that the agent can control, which determine the search space.
objectives : Sequence[Objective]
The objectives which the agent will try to optimize.
evaluation_function : EvaluationFunction
The function to evaluate acquired data and produce outcomes.
acquisition_plan : str | None, optional
The acquisition plan to use for acquiring data from the beamline. If not provided,
:func:`bluesky_adaptive.plans.default_acquire` will be assumed.
dof_constraints : Sequence[DOFConstraint] | None, optional
Constraints on DOFs to refine the search space.
outcome_constraints : Sequence[OutcomeConstraint] | None, optional
Constraints on outcomes to be satisfied during optimization.
checkpoint_path : str | None, optional
The path to the checkpoint file to save the optimizer's state to.
**kwargs : Any
Additional keyword arguments to configure the Ax experiment.
See Also
--------
bluesky_adaptive.protocols.Sensor : The protocol for sensors.
bluesky_adaptive.ax.dof.RangeDOF : For continuous parameters.
bluesky_adaptive.ax.dof.ChoiceDOF : For discrete parameters.
bluesky_adaptive.ax.objective.Objective : For defining objectives.
bluesky_adaptive.ax.optimizer.AxOptimizer : The optimizer used internally.
bluesky_adaptive.queueserver.QueueserverOptimizationRunner : Runner that handles interaction with bluesky-queueserver.
"""
def __init__(
self,
re_manager_api: bluesky_queueserver_api.zmq.REManagerAPI | bluesky_queueserver_api.http.REManagerAPI,
document_dispatcher: RemoteDispatcher,
sensors: Sequence[str],
dofs: Sequence[DOF],
objectives: Sequence[Objective],
evaluation_function: EvaluationFunction[str],
acquisition_plan: str | None = None,
dof_constraints: Sequence[DOFConstraint] | None = None,
outcome_constraints: Sequence[OutcomeConstraint] | None = None,
checkpoint_path: str | None = None,
acquisition_plan_kwargs: Mapping[str, Any] | None = None,
**kwargs: Any,
):
self._sensors = sensors
self._actuators: Sequence[str] = []
for dof in dofs:
if dof.actuator is not None:
if isinstance(dof.actuator, Actuator):
self._actuators.append(dof.actuator.name)
else:
self._actuators.append(dof.actuator)
self._evaluation_function = evaluation_function
self._acquisition_plan = acquisition_plan
self._acquisition_plan_kwargs = acquisition_plan_kwargs or {}
self._optimizer = AxOptimizer(
parameters=[dof.to_ax_parameter_config() for dof in dofs],
objective=to_ax_objective_str(objectives),
parameter_constraints=[constraint.ax_constraint for constraint in dof_constraints] if dof_constraints else None,
outcome_constraints=[constraint.ax_constraint for constraint in outcome_constraints]
if outcome_constraints
else None,
checkpoint_path=checkpoint_path,
**kwargs,
)
self._runner = QueueserverOptimizationRunner(
self.to_optimization_problem(),
QueueserverClient(re_manager_api, document_dispatcher),
)
@property
def evaluation_function(self) -> EvaluationFunction[str]:
"""Evaluation function mapping acquired data to outcomes."""
return self._evaluation_function
@property
def actuators(self) -> Sequence[str]:
"""Set of actuator names used during acquisition."""
return self._actuators
@property
def sensors(self) -> Sequence[str]:
"""Set of sensor names used during acquisition."""
return self._sensors
@property
def acquisition_plan(self) -> str | None:
"""Acquisition plan name used during acquisition."""
return self._acquisition_plan
[docs]
def stop(self) -> None:
"""Stop the queueserver runner."""
self._runner.stop()
@property
def current_iteration(self) -> int:
"""The current iteration of the optimization."""
return self._runner.current_iteration
[docs]
def to_optimization_problem(self) -> QueueserverOptimizationProblem:
"""Convert the agent state to an optimization problem."""
return QueueserverOptimizationProblem(
optimizer=self._optimizer,
actuators=self._actuators,
sensors=self._sensors,
evaluation_function=self._evaluation_function,
acquisition_plan=self._acquisition_plan,
acquisition_plan_kwargs=self._acquisition_plan_kwargs,
)
[docs]
def run(
self, iterations: int = 1, n_points: int = 1, checkpoint_interval: int | None = None
) -> Future[OptimizationResult]:
"""
Start the optimization loop.
Validates the queueserver state, then begins the suggest -> acquire -> ingest
cycle. This method returns immediately; the optimization runs asynchronously
via callbacks.
Parameters
----------
iterations : int
Number of optimization iterations to run.
n_points : int
Number of points to suggest per iteration.
checkpoint_interval : int | None
The number of iterations between optimizer checkpoints. If None, checkpoints
will not be saved. Optimizer must implement the
:class:`bluesky_adaptive.protocols.Checkpointable` protocol.
Returns
-------
concurrent.futures.Future[OptimizationResult]
A future that resolves to an :class:`~bluesky_adaptive.queueserver.OptimizationResult`
when all iterations complete or when :meth:`stop` is called. If an
unhandled exception occurs the future will hold it and re-raise on
``.result()``.
Raises
------
RuntimeError
If the queueserver environment is not ready.
ValueError
If required devices or plans are not available.
"""
return self._runner.run(iterations=iterations, num_points=n_points, checkpoint_interval=checkpoint_interval)
[docs]
def submit_suggestions(self, suggestions: Sequence[Mapping]) -> Future[OptimizationResult]:
"""
Evaluate specific parameter combinations.
Acquires data for given suggestions and ingests results. Supports both
optimizer suggestions and manual points.
Parameters
----------
suggestions : Sequence[Mapping]
Either optimizer suggestions (with "_id") or manual points (without "_id").
Returns
-------
concurrent.futures.Future[OptimizationResult]
A future that resolves to an :class:`~bluesky_adaptive.queueserver.OptimizationResult`
when the acquisition completes.
See Also
--------
run : Run the full optimization loop.
"""
return self._runner.submit_suggestions(suggestions)