diff --git a/.github/workflows/_testing.yml b/.github/workflows/_testing.yml index 0be313f0..21d69eed 100644 --- a/.github/workflows/_testing.yml +++ b/.github/workflows/_testing.yml @@ -10,7 +10,7 @@ jobs: strategy: matrix: host-os: ["ubuntu-latest"] - python-version: ["py310-cpu", "py311-cpu", "py312-cpu", "py313-cpu"] + python-version: ["py311-cpu", "py312-cpu", "py313-cpu"] fail-fast: false defaults: diff --git a/docs/source/installation.rst b/docs/source/installation.rst index b3b27981..26af90e4 100644 --- a/docs/source/installation.rst +++ b/docs/source/installation.rst @@ -21,6 +21,34 @@ To install the package using the ``conda`` package manager, run the following co :start-after: .. snippet-conda-standard-start :end-before: .. snippet-conda-standard-end +Optional Extras +^^^^^^^^^^^^^^^ + +``blop`` is modular — install only what you need by appending one or more extras: + +.. list-table:: + :header-rows: 1 + :widths: auto + + * - Extra + - Installs + - Notes + * - ``blop[ax]`` + - ``ax-platform``, ``botorch``, ``gpytorch``, ``torch`` + - GPU torch by default; pair with ``[cpu]`` for CPU-only + * - ``blop[queueserver]`` + - ``bluesky-queueserver-api`` + - Transport layer only; pair with ``[ax]`` for ``QueueserverAgent`` + * - ``blop[cpu]`` + - *(uv index routing)* + - Routes ``torch`` to the CPU-only PyTorch index; requires ``uv`` + * - ``blop[all]`` + - All backends + ``[queueserver]`` + - No dev tooling; will grow as new backends are added + * - ``blop[dev]`` + - ``blop[all]`` + dev tooling + - For contributors + PyTorch Acceleration Options ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ diff --git a/docs/source/reference/agent.rst b/docs/source/reference/agent.rst index 5ddcaaa6..7703f3f8 100644 --- a/docs/source/reference/agent.rst +++ b/docs/source/reference/agent.rst @@ -7,7 +7,7 @@ Agent :show-inheritance: :inherited-members: -.. autoclass:: blop.ax.QueueserverAgent +.. autoclass:: blop.ax.queueserver_agent.QueueserverAgent :members: :undoc-members: :show-inheritance: diff --git a/docs/source/tutorials/queueserver.md b/docs/source/tutorials/queueserver.md index c3a87023..00d46f84 100644 --- a/docs/source/tutorials/queueserver.md +++ b/docs/source/tutorials/queueserver.md @@ -185,7 +185,8 @@ tiled_client = from_uri("http://localhost:8000", api_key="tutorialkey") Just as in the simple experiment tutorial, we define **DOFs** and **objectives**. The key difference: since devices exist only in the remote queueserver environment, DOFs reference device names as strings (no `actuator` objects). ```{code-cell} ipython3 -from blop.ax import QueueserverAgent, RangeDOF, Objective +from blop.ax import RangeDOF, Objective +from blop.ax.queueserver_agent import QueueserverAgent dofs = [ RangeDOF(actuator="motor1", bounds=(-5.0, 5.0), parameter_type="float"), diff --git a/docs/wip/qserver-experiment.md b/docs/wip/qserver-experiment.md index e133369f..dffde800 100644 --- a/docs/wip/qserver-experiment.md +++ b/docs/wip/qserver-experiment.md @@ -240,7 +240,7 @@ The EvaluationFunction is called every time a stop document is recieved. It must from blop.protocols import EvaluationFunction from tiled.client.container import Container from tiled.queries import Eq -from blop.ax import QueueserverAgent +from blop.ax.queueserver_agent import QueueserverAgent from bluesky.callbacks.zmq import RemoteDispatcher import numpy as np diff --git a/pixi.toml b/pixi.toml index 67129eab..6e275451 100644 --- a/pixi.toml +++ b/pixi.toml @@ -6,10 +6,10 @@ authors = [ channels = ["conda-forge"] name = "blop" platforms = ["linux-64", "osx-arm64"] -version = "0.9.0" +version = "1.0.0b1" [dependencies] -python = ">=3.10.0,<3.14" +python = ">=3.11.0,<3.14" [feature.dev.dependencies] @@ -32,7 +32,6 @@ tests = "pytest src/blop/tests" graphviz = ">=14.1.2,<15" [feature.docs.pypi-dependencies] -blop = { path = ".", editable = true, extras = ["cpu"] } blop-sim = { path = "sim", editable = true } numpydoc = "*" sphinx-copybutton = "*" @@ -46,9 +45,6 @@ bluesky-tiled-plugins = "*" ophyd-async = "*" opencv-python = "*" -[feature.py310.dependencies] -python = "3.10.*" - [feature.py311.dependencies] python = "3.11.*" @@ -58,10 +54,6 @@ python = "3.12.*" [feature.py313.dependencies] python = "3.13.*" -[feature.dev.tasks] -check = "pre-commit run --all-files" -tests = "pytest src/blop/tests" - [feature.docs.tasks] start-queueserver = "docker compose -f docs/source/tutorials/queueserver/docker-compose.yml up -d --build --wait" stop-queueserver = "docker compose -f docs/source/tutorials/queueserver/docker-compose.yml down" @@ -77,8 +69,7 @@ convert-tutorials-to-ipynb = "jupytext --to notebook docs/source/tutorials/*.md" [environments] dev = ["dev"] dev-cpu = ["dev-cpu"] -docs = ["docs"] -py310-cpu = ["dev-cpu", "py310"] +docs = ["dev-cpu", "docs"] py311-cpu = ["dev-cpu", "py311"] py312-cpu = ["dev-cpu", "py312"] py313-cpu = ["dev-cpu", "py313"] diff --git a/pyproject.toml b/pyproject.toml index 90ac1139..4c7359a8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,19 +20,8 @@ maintainers = [ { name = "Jennefer Maldonado", email = "jmaldonad@bnl.gov" }, { name = "Roman Chernikov", email = "rcherniko@bnl.gov" }, ] -requires-python = ">=3.10" -dependencies = [ - "ax-platform>=1.1.0,<1.3", - "bluesky>=1.14.2", - "bluesky-queueserver-api>=0.0.12", - "torch", - "botorch>=0.16.0", - "gpytorch", - "scipy", - "networkx>=3", - "numpy", - "rich>=13", -] +requires-python = ">=3.11" +dependencies = ["bluesky>=1.14.2", "networkx>=3", "numpy", "rich>=13"] classifiers = [ "Development Status :: 4 - Beta", "License :: OSI Approved :: BSD License", @@ -45,7 +34,9 @@ classifiers = [ dynamic = ["version"] [project.optional-dependencies] -qs = ["bluesky-queueserver-api"] +ax = ["ax-platform>=1.3.0,<1.4", "botorch>=0.18.0,<0.19.0", "gpytorch", "torch"] +queueserver = ["bluesky-queueserver-api"] +all = ["blop[ax,queueserver]"] dev = [ "pytest", "pytest-cov", @@ -53,10 +44,9 @@ dev = [ "ruff", "nbstripout", "pre-commit", - "pandas-stubs", "coverage", "pyright", - "blop[qs]", + "blop[all]", ] cpu = [ # Empty extra - the source configuration below routes to CPU-only index diff --git a/sim/pyproject.toml b/sim/pyproject.toml index b2efc374..2e5a09c2 100644 --- a/sim/pyproject.toml +++ b/sim/pyproject.toml @@ -8,10 +8,10 @@ version = "0.0.0" description = "Example simulations and benchmarks for blop (not published to PyPI)" readme = "README.md" authors = [{ name = "Thomas Hopkins", email = "thopkins1@bnl.gov" }] -requires-python = ">=3.10" +requires-python = ">=3.11" dependencies = [ "blop", - "ophyd-async", + "ophyd-async<0.17", "h5py", "area-detector-handlers", "bluesky-tiled-plugins", diff --git a/src/blop/__init__.py b/src/blop/__init__.py index c2d138e6..658a078f 100644 --- a/src/blop/__init__.py +++ b/src/blop/__init__.py @@ -1,4 +1,3 @@ -from .ax import DOF, Agent, ChoiceDOF, DOFConstraint, Objective, OutcomeConstraint, RangeDOF, ScalarizedObjective from .plans import acquire_baseline, default_acquire, optimize, optimize_step, sample_suggestions try: @@ -8,14 +7,6 @@ __all__ = [ "__version__", - "Agent", - "ChoiceDOF", - "DOF", - "DOFConstraint", - "Objective", - "OutcomeConstraint", - "RangeDOF", - "ScalarizedObjective", "acquire_baseline", "default_acquire", "optimize", diff --git a/src/blop/ax/__init__.py b/src/blop/ax/__init__.py index f9df0feb..0f51021d 100644 --- a/src/blop/ax/__init__.py +++ b/src/blop/ax/__init__.py @@ -1,13 +1,13 @@ -from ..queueserver import OptimizationResult -from .agent import Agent as Agent -from .agent import QueueserverAgent as QueueserverAgent -from .dof import DOF, ChoiceDOF, DOFConstraint, RangeDOF -from .objective import Objective, OutcomeConstraint, ScalarizedObjective, to_ax_objective_str -from .optimizer import AxOptimizer +try: + from .agent import Agent as Agent + from .dof import DOF, ChoiceDOF, DOFConstraint, RangeDOF + from .objective import Objective, OutcomeConstraint, ScalarizedObjective, to_ax_objective_str + from .optimizer import AxOptimizer +except ImportError as e: + raise ImportError("The ax integration requires additional dependencies. Install them with: pip install blop[ax]") from e __all__ = [ "Agent", - "QueueserverAgent", "DOF", "RangeDOF", "ChoiceDOF", @@ -17,5 +17,4 @@ "ScalarizedObjective", "to_ax_objective_str", "AxOptimizer", - "OptimizationResult", ] diff --git a/src/blop/ax/agent.py b/src/blop/ax/agent.py index 2aff9006..13fff478 100644 --- a/src/blop/ax/agent.py +++ b/src/blop/ax/agent.py @@ -1,25 +1,14 @@ -import importlib.util import logging from collections.abc import Mapping, Sequence -from concurrent.futures import Future from typing import Any, cast +import bluesky.preprocessors as bpp from ax import Client, TOutcome, TParameterization from ax.analysis.plotly.surface.contour import ContourPlot +from ax.core.analysis_card import AnalysisCardBase from ax.core.types import TParamValue - -# =============================== -# TODO: Remove when Python 3.10 is no longer supported -if importlib.util.find_spec("ax.core.analysis_card") is not None: - from ax.core.analysis_card import AnalysisCardBase -else: - from ax.analysis.analysis_card import AnalysisCardBase # type: ignore[import-untyped] -# =============================== -import bluesky.preprocessors as bpp from bluesky.callbacks import CallbackBase -from bluesky.callbacks.zmq import RemoteDispatcher from bluesky.utils import MsgGenerator -from bluesky_queueserver_api.zmq import REManagerAPI from ..callbacks.logger import OptimizationLogger from ..callbacks.router import OptimizationCallbackRouter @@ -30,10 +19,8 @@ Actuator, EvaluationFunction, OptimizationProblem, - QueueserverOptimizationProblem, Sensor, ) -from ..queueserver import OptimizationResult, QueueserverClient, QueueserverOptimizationRunner from ..utils import InferredReadable from .dof import DOF, DOFConstraint from .objective import Objective, OutcomeConstraint, ScalarizedObjective, to_ax_objective_str @@ -561,182 +548,3 @@ def navigate_to_best(self, parameterization: Mapping | None = None) -> MsgGenera parameterization, ) ) - - -class QueueserverAgent(_AxAgentMixin): - """ - An asynchronous interface that uses Ax as the backend for optimization and experiment tracking - and 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 : 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:`blop.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 - -------- - blop.protocols.Sensor : The protocol for sensors. - blop.ax.dof.RangeDOF : For continuous parameters. - blop.ax.dof.ChoiceDOF : For discrete parameters. - blop.ax.objective.Objective : For defining objectives. - blop.ax.optimizer.AxOptimizer : The optimizer used internally. - blop.queueserver.QueueserverOptimizatonRunner : Runner that handles interaction with bluesky-queueserver. - """ - - def __init__( - self, - re_manager_api: REManagerAPI, - document_dispatcher: RemoteDispatcher, - sensors: Sequence[str], - dofs: Sequence[DOF], - objectives: Sequence[Objective], - evaluation_function: EvaluationFunction, - 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: - return self._evaluation_function - - @property - def actuators(self) -> Sequence[str]: - return self._actuators - - @property - def sensors(self) -> Sequence[str]: - return self._sensors - - @property - def acquisition_plan(self) -> str | None: - return self._acquisition_plan - - def stop(self) -> None: - self._runner.stop() - - @property - def current_iteration(self) -> int: - return self._runner.current_iteration - - def to_optimization_problem(self) -> QueueserverOptimizationProblem: - 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, - ) - - def run(self, iterations: int = 1, n_points: int = 1) -> 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. - - Returns - ------- - concurrent.futures.Future[OptimizationResult] - A future that resolves to an :class:`~blop.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, n_points) - - def submit_suggestions(self, suggestions: list[dict]) -> Future[OptimizationResult]: - """ - Evaluate specific parameter combinations. - - Acquires data for given suggestions and ingests results. Supports both - optimizer suggestions and manual points. - - Parameters - ---------- - suggestions : list[dict] - Either optimizer suggestions (with "_id") or manual points (without "_id"). - - Returns - ------- - concurrent.futures.Future[OptimizationResult] - A future that resolves to an :class:`~blop.queueserver.OptimizationResult` - when the acquisition completes. - - See Also - -------- - run : Run the full optimization loop. - """ - return self._runner.submit_suggestions(suggestions) diff --git a/src/blop/ax/optimizer.py b/src/blop/ax/optimizer.py index 5132fb68..21a4d9d0 100644 --- a/src/blop/ax/optimizer.py +++ b/src/blop/ax/optimizer.py @@ -2,7 +2,7 @@ from typing import Any from ax import ChoiceParameterConfig, Client, RangeParameterConfig, TOutcome, TParameterization -from ax.core.objective import MultiObjective +from ax.core import MultiObjectiveOptimizationConfig from ax.core.parameter import ChoiceParameter, RangeParameter from ax.core.types import TParamValue @@ -137,7 +137,7 @@ def suggest(self, num_points: int | None = None) -> list[dict]: next_trials = self._client.get_next_trials(max_trials=num_points, fixed_parameters=self._fixed_parameters) return [ { - "_id": trial_index, + ID_KEY: trial_index, **parameterization, } for trial_index, parameterization in next_trials.items() @@ -167,7 +167,7 @@ def get_best_points(self) -> list[tuple[int, TParameterization, TOutcome]]: opt_config = self._client._experiment.optimization_config if opt_config is None: raise ValueError("Somehow your optimization has not been configured yet...check `ax_client`.") - is_multi_objective = isinstance(opt_config.objective, MultiObjective) + is_multi_objective = isinstance(opt_config, MultiObjectiveOptimizationConfig) if is_multi_objective: frontier = self._client.get_pareto_frontier(use_model_predictions=False) diff --git a/src/blop/ax/queueserver_agent.py b/src/blop/ax/queueserver_agent.py new file mode 100644 index 00000000..19d27964 --- /dev/null +++ b/src/blop/ax/queueserver_agent.py @@ -0,0 +1,196 @@ +from collections.abc import Mapping, Sequence +from concurrent.futures import Future +from typing import Any + +from bluesky.callbacks.zmq import RemoteDispatcher +from bluesky_queueserver_api.zmq import REManagerAPI + +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 + + +class QueueserverAgent(_AxAgentMixin): + """ + An asynchronous interface that uses Ax as the backend for optimization and experiment tracking + and 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 : 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:`blop.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 + -------- + blop.protocols.Sensor : The protocol for sensors. + blop.ax.dof.RangeDOF : For continuous parameters. + blop.ax.dof.ChoiceDOF : For discrete parameters. + blop.ax.objective.Objective : For defining objectives. + blop.ax.optimizer.AxOptimizer : The optimizer used internally. + blop.queueserver.QueueserverOptimizatonRunner : Runner that handles interaction with bluesky-queueserver. + """ + + def __init__( + self, + re_manager_api: REManagerAPI, + document_dispatcher: RemoteDispatcher, + sensors: Sequence[str], + dofs: Sequence[DOF], + objectives: Sequence[Objective], + evaluation_function: EvaluationFunction, + 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: + return self._evaluation_function + + @property + def actuators(self) -> Sequence[str]: + return self._actuators + + @property + def sensors(self) -> Sequence[str]: + return self._sensors + + @property + def acquisition_plan(self) -> str | None: + return self._acquisition_plan + + def stop(self) -> None: + self._runner.stop() + + @property + def current_iteration(self) -> int: + return self._runner.current_iteration + + def to_optimization_problem(self) -> QueueserverOptimizationProblem: + 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, + ) + + def run(self, iterations: int = 1, n_points: int = 1) -> 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. + + Returns + ------- + concurrent.futures.Future[OptimizationResult] + A future that resolves to an :class:`~blop.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, n_points) + + def submit_suggestions(self, suggestions: list[dict]) -> Future[OptimizationResult]: + """ + Evaluate specific parameter combinations. + + Acquires data for given suggestions and ingests results. Supports both + optimizer suggestions and manual points. + + Parameters + ---------- + suggestions : list[dict] + Either optimizer suggestions (with "_id") or manual points (without "_id"). + + Returns + ------- + concurrent.futures.Future[OptimizationResult] + A future that resolves to an :class:`~blop.queueserver.OptimizationResult` + when the acquisition completes. + + See Also + -------- + run : Run the full optimization loop. + """ + return self._runner.submit_suggestions(suggestions) diff --git a/src/blop/bayesian/__init__.py b/src/blop/bayesian/__init__.py index e69de29b..0d862f80 100644 --- a/src/blop/bayesian/__init__.py +++ b/src/blop/bayesian/__init__.py @@ -0,0 +1,10 @@ +try: + from .kernels import LatentKernel + from .models import LatentGP +except ImportError as e: + raise ImportError("The bayesian module requires additional dependencies. Install them with: pip install blop[ax]") from e + +__all__ = [ + "LatentKernel", + "LatentGP", +] diff --git a/src/blop/protocols.py b/src/blop/protocols.py index a4e5cee8..c536553e 100644 --- a/src/blop/protocols.py +++ b/src/blop/protocols.py @@ -328,7 +328,8 @@ class QueueserverOptimizationProblem(BaseOptimizationProblem[str, str, str]): An optimization problem to solve. Immutable once initialized. This dataclass encapsulates all components needed for optimization into a single - immutable structure. It is typically created via :meth:`blop.ax.QueueserverAgent.to_optimization_problem` + immutable structure. It is typically created via + :meth:`blop.ax.queueserver_agent.QueueserverAgent.to_optimization_problem` and used with bluesky-queueserver-api. Actuators, sensors, and the acquisition plan are referenced by their names, since their instances live on a remote server. @@ -351,7 +352,8 @@ class QueueserverOptimizationProblem(BaseOptimizationProblem[str, str, str]): See Also -------- - blop.ax.QueueserverAgent.to_optimization_problem : Creates a QueueserverOptimizationProblem from an agent. + blop.ax.queueserver_agent.QueueserverAgent.to_optimization_problem : + Creates a QueueserverOptimizationProblem from an agent. blop.queueserver.QueueserverOptimizationRunner : Runs the optimization loop using the bluesky-queueserver-api. """ diff --git a/src/blop/queueserver.py b/src/blop/queueserver.py index c16816fd..0ba05967 100644 --- a/src/blop/queueserver.py +++ b/src/blop/queueserver.py @@ -21,8 +21,14 @@ from bluesky.callbacks import CallbackBase from bluesky.callbacks.zmq import RemoteDispatcher -from bluesky_queueserver_api import BPlan -from bluesky_queueserver_api.zmq import REManagerAPI + +try: + from bluesky_queueserver_api import BPlan + from bluesky_queueserver_api.zmq import REManagerAPI +except ImportError as e: + raise ImportError( + "The queueserver integration requires additional dependencies. Install them with: pip install blop[qs]" + ) from e from event_model import RunStart, RunStop from .plans import default_acquire diff --git a/src/blop/tests/ax/test_agent.py b/src/blop/tests/ax/test_agent.py index 008ef6e2..b5b1ee57 100644 --- a/src/blop/tests/ax/test_agent.py +++ b/src/blop/tests/ax/test_agent.py @@ -7,10 +7,11 @@ from bluesky.callbacks.zmq import RemoteDispatcher from bluesky_queueserver_api.zmq import REManagerAPI -from blop.ax.agent import Agent, QueueserverAgent +from blop.ax.agent import Agent from blop.ax.dof import ChoiceDOF, DOFConstraint, RangeDOF from blop.ax.objective import Objective, ScalarizedObjective from blop.ax.optimizer import AxOptimizer +from blop.ax.queueserver_agent import QueueserverAgent from blop.callbacks.logger import OptimizationLogger from blop.protocols import AcquisitionPlan, EvaluationFunction @@ -373,8 +374,8 @@ def test_queueserver_agent_init_actuator_instance(mock_re_manager_api, mock_docu assert agent.actuators == [movable1.name, dof2.parameter_name] -@patch("blop.ax.agent.QueueserverClient") -@patch("blop.ax.agent.QueueserverOptimizationRunner") +@patch("blop.ax.queueserver_agent.QueueserverClient") +@patch("blop.ax.queueserver_agent.QueueserverOptimizationRunner") def test_queueserver_agent_run( mock_queueserver_runner_cls, mock_queueserver_client_cls, @@ -399,8 +400,8 @@ def test_queueserver_agent_run( mock_queueserver_runner_cls.return_value.run.assert_called_once_with(1, 1) -@patch("blop.ax.agent.QueueserverClient") -@patch("blop.ax.agent.QueueserverOptimizationRunner") +@patch("blop.ax.queueserver_agent.QueueserverClient") +@patch("blop.ax.queueserver_agent.QueueserverOptimizationRunner") def test_queueserver_agent_submit_suggestions( mock_queueserver_runner_cls, mock_queueserver_client_cls, @@ -426,8 +427,8 @@ def test_queueserver_agent_submit_suggestions( mock_queueserver_runner_cls.return_value.submit_suggestions.assert_called_once_with(suggestions) -@patch("blop.ax.agent.QueueserverClient") -@patch("blop.ax.agent.QueueserverOptimizationRunner") +@patch("blop.ax.queueserver_agent.QueueserverClient") +@patch("blop.ax.queueserver_agent.QueueserverOptimizationRunner") def test_queueserver_agent_stop( mock_queueserver_runner_cls, mock_queueserver_client_cls, diff --git a/src/blop/utils.py b/src/blop/utils.py index ecf8c2bb..d696f70f 100644 --- a/src/blop/utils.py +++ b/src/blop/utils.py @@ -1,6 +1,6 @@ import time from collections.abc import Sequence -from enum import Enum +from enum import StrEnum from typing import Any import networkx as nx @@ -12,7 +12,7 @@ from .protocols import ID_KEY, OptimizationProblem -class Source(str, Enum): +class Source(StrEnum): """An enum that helps describe where the data key comes from.""" OUTCOME = "optimization-outcome"