diff --git a/pixi.lock b/pixi.lock index d5d3690..df6ce57 100644 --- a/pixi.lock +++ b/pixi.lock @@ -19978,6 +19978,7 @@ packages: - bluesky-tiled-plugins>=2.0.11 - nslsii - ophyd-async[ca,pva,sim]>=0.21.3 + - rich - tiled[client] requires_python: '>=3.11,<3.14' - pypi: https://files.pythonhosted.org/packages/01/f9/575c8d760eae1fc99651b7cc5efd96ad5379ca4d6b53750b0fb4fe983f34/imagesize-2.0.1-py3-none-any.whl diff --git a/pixi.toml b/pixi.toml index e65d7d4..5c204fa 100644 --- a/pixi.toml +++ b/pixi.toml @@ -13,6 +13,7 @@ bluesky-tiled-plugins = ">=2.0.11,<3" ophyd-async = ">=0.21.3,<0.22" nslsii = ">=0.11.9,<0.12" redis-json-dict = ">=0.2.2,<0.3" +rich = "*" [pypi-dependencies] pdftools = { path = ".", editable = true } diff --git a/pyproject.toml b/pyproject.toml index 8ffb059..f2f2c88 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,6 +30,7 @@ dependencies = [ "bluesky-tiled-plugins>=2.0.11", "ophyd-async[ca, pva, sim]>=0.21.3", "nslsii", + "rich", ] [dependency-groups] diff --git a/src/pdftools/profiles/collection.py b/src/pdftools/profiles/collection.py index 780e23c..da70d26 100644 --- a/src/pdftools/profiles/collection.py +++ b/src/pdftools/profiles/collection.py @@ -1,8 +1,7 @@ -"""pdftools bluesky collection profile.""" +"""Bluesky profile for the 28-ID-1 PDF beamline at NSLS-II.""" import argparse import os -from datetime import datetime import numpy as np from bluesky import plan_stubs as bps @@ -11,22 +10,24 @@ from bluesky.callbacks.best_effort import BestEffortCallback from bluesky.magics import BlueskyMagics from bluesky.run_engine import RunEngine, autoawait_in_bluesky_event_loop -from bluesky.utils import ProgressBarManager, ts_msg_hook +from bluesky.utils import ProgressBarManager from bluesky_tiled_plugins import TiledWriter from IPython.core.getipython import get_ipython from IPython.terminal.interactiveshell import TerminalInteractiveShell from nslsii.ophyd_async.providers import NSLS2PathProvider -from nslsii.utils import open_redis_client -from ophyd_async.core import init_devices -from redis_json_dict.redis_json_dict import RedisJSONDict from tiled.client import from_uri, simple import pdftools - -# Have ophyd-async keep detector state constant when not explicitly overridden -os.environ["OPHYD_ASYNC_PRESERVE_DETECTOR_STATE"] = "YES" -if "REDIS_HOST" not in os.environ: - os.environ["REDIS_HOST"] = "" +from pdftools.utils import ( + CI_ENV_VAR, + ProposalIDPrompt, + auto_init_devices, + initialize_run_engine, + is_running_in_ci, + print_proposal_info, + print_version_info, + remove_md_keys, +) # Mechanism for passing arguments to profile. _parser = argparse.ArgumentParser("Collection profile for pdftools.") @@ -35,53 +36,39 @@ action="store_true", help="Run the profile in mock mode with simulated devices and metadata.", ) +if vars(_parser.parse_args())["mock"]: + os.environ[CI_ENV_VAR] = "YES" +MOCK = is_running_in_ci() + +# Beamline Redis for RunEngine metadata, as pdf-profile-collection's 00-base uses it +# (port 6380 with SSL is open_redis_client's default when redis_ssl=True). +os.environ.setdefault("REDIS_HOST", "xf28id1-pdf-redis1.nsls2.bnl.gov") +# Have ophyd-async keep detector state constant when not explicitly overridden. +os.environ["OPHYD_ASYNC_PRESERVE_DETECTOR_STATE"] = "YES" + +# Tiled catalog that pdf-profile-collection's 00-base writes to today. +TILED_URI = "https://tiled.nsls2.bnl.gov" +TILED_CATALOG = ("pdf", "migration") -# Determine if we are running in CI/mock mode -MOCK = vars(_parser.parse_args())["mock"] +print_version_info() -# Construct a RunEngine with either an in-process metadata dict, -# or a redis-backed metadata store. +# Set up the RunEngine and its metadata. +RE: RunEngine = initialize_run_engine() if not MOCK: - RE = RunEngine(RedisJSONDict(open_redis_client(redis_ssl=True), "")) # type: ignore -else: - RE = RunEngine( - { - "data_session": "pass-123456", - "cycle": ( - f"{datetime.today().year}-{(datetime.today().month - 1) // 4 + 1}" - ), - "proposal": { - "title": "Mock proposal", - "type": "commissioning", - "pi_name": os.getenv("USER", "Unknown"), - }, - } - ) - -# Add tools package version alongside bluesky/ophyd/ophyd-async etc. + # Trim the PDF: entries 00-base removes from Redis on every start. + remove_md_keys(RE.md, prefix="PDF:") RE.md["versions"]["pdftools"] = pdftools.__version__ - -# Beamline and facility metadata RE.md["facility"] = "NSLS-II" -RE.md["beamline_id"] = "" # Set the beamline ID (i.e. 27-ID-1) +RE.md["group"] = "PDF" +RE.md["beamline_id"] = "28-ID-1" -# Attempt to extract TLA from environment variables. -RE.md["group"] = TLA = ( - os.environ.get("ENDSTATION_ACRONYM") - or os.environ.get("BEAMLINE_ACRONYM") - or "Unknown" -).upper() +# Display active proposal information in the console. +print_proposal_info(RE.md) # Get the IPython shell and set up the custom prompt to show the current proposal ID. ipython = get_ipython() if ipython is not None and isinstance(ipython, TerminalInteractiveShell): - try: - from nslsii.utils import ProposalIDPrompt # type: ignore - except ImportError: # Only available in nslsii after 0.11.9 - pass - else: - ipython.prompts = ProposalIDPrompt(RE, ipython) - + ipython.prompts = ProposalIDPrompt(RE, ipython) BlueskyMagics.RE = RE ipython.register_magics(BlueskyMagics) autoawait_in_bluesky_event_loop() @@ -94,48 +81,41 @@ # Use %tb verbose to view full traceback. ipython.run_line_magic("xmode", "minimal") -# Set up the Tiled clients for data acquisition and reading based on the environment. +# Construct our tiled clients for writing and (in an interactive session) reading. +# In CI/mock mode, an in-process server stands in for both. if not MOCK: - # Typically $TLA/raw. Replace with beamline-specific catalog. - daq_tiled_catalog = f"{TLA.lower()}/raw" - tiled_writing_client = from_uri( - "https://tiled.nsls2.bnl.gov", - api_key=os.environ.get(f"TILED_BLUESKY_WRITING_API_KEY_{TLA}", ""), - )[daq_tiled_catalog] + TILED_URI, + api_key=os.environ.get("TILED_BLUESKY_WRITING_API_KEY_PDF", ""), + )[TILED_CATALOG[0]][TILED_CATALOG[1]] + tiled_writing_client.context.http_client.headers["tiled-qos"] = "acquisition" if ipython is not None and isinstance(ipython, TerminalInteractiveShell): - tiled_reading_client = c = from_uri("https://tiled.nsls2.bnl.gov")[ - daq_tiled_catalog + tiled_reading_client = c = from_uri(TILED_URI)[TILED_CATALOG[0]][ + TILED_CATALOG[1] ] else: tiled_writing_client = tiled_reading_client = c = simple() - -# Subscribe the tiled writer to the RunEngine -RE.subscribe(TiledWriter(tiled_writing_client)) - -# Set up our default path provider -path_provider = NSLS2PathProvider(RE.md) - -# Initialize all async devices within this context manager -with init_devices(mock=MOCK, timeout=2.0): - pass +# Subscribe the tiled writer; documents it cannot send are kept on local disk. +RE.subscribe(TiledWriter(tiled_writing_client, backup_directory="/tmp/bluesky-docs")) # Subscribe the best effort callback bec = BestEffortCallback() RE.subscribe(bec) -# Disable bec printout of baseline readings -# since it is kind of noisy +# Disable bec printout of baseline readings since it is kind of noisy bec.disable_baseline() -# Confiugre the default terminal progress bar manager +# Configure the default terminal progress bar manager RE.waiting_hook = ProgressBarManager() # type: ignore -# Optionally enable ts_msg_hook for debugging -# RE.msg_hook = ts_msg_hook +# Define our global default path provider for the beamline +path_provider = NSLS2PathProvider(RE.md) + +# Devices are added here, step by step (detectors, motors, sample environment). +with auto_init_devices(timeout=2.0): + pass -# Optionally configure baseline supplemental data to include in the metadata -# of every run. Add baseline=[, , ...] to SupplementalData as needed. +# Add baseline=[, , ...] to SupplementalData as devices arrive. sd = bpp.SupplementalData() RE.preprocessors.append(sd) diff --git a/src/pdftools/utils/__init__.py b/src/pdftools/utils/__init__.py new file mode 100644 index 0000000..e03cc40 --- /dev/null +++ b/src/pdftools/utils/__init__.py @@ -0,0 +1,230 @@ +"""Profile utilities for pdftools, following hextools.utils.""" + +import os +from collections.abc import MutableMapping +from datetime import datetime +from typing import Any + +from bluesky.run_engine import RunEngine +from IPython.terminal.interactiveshell import TerminalInteractiveShell +from IPython.terminal.prompts import Prompts +from nslsii.utils import open_redis_client +from ophyd_async.core import ( + Device, + DeviceProcessor, + NotConnectedError, + Signal, + wait_for_connection, + walk_devices, +) +from pygments.token import Token +from redis_json_dict.redis_json_dict import RedisJSONDict +from rich import print as rprint +from rich.console import Console +from rich.markup import escape + +CI_ENV_VAR = "PDFTOOLS_RUNNING_IN_CI" + + +def is_running_in_ci() -> bool: + """Check whether pdftools runs in CI or mock mode. + + Returns + ------- + bool + True if ``PDFTOOLS_RUNNING_IN_CI`` is set to a true value. + """ + return os.environ.get(CI_ENV_VAR, "false").lower() in ["true", "1", "yes"] + + +def print_version_info(): + """Print version information for bluesky, ophyd_async, tiled, and pdftools.""" + from bluesky import __version__ as bluesky_version + from ophyd_async import __version__ as ophyd_async_version + from tiled import __version__ as tiled_version + + from pdftools import __version__ as pdftools_version + + rprint("\n[bold]Version Information[/bold]") + rprint(f" [bold]bluesky[/bold]: [blue]{bluesky_version}[/blue]") + rprint(f" [bold]ophyd_async[/bold]: [blue]{ophyd_async_version}[/blue]") + rprint(f" [bold]tiled[/bold]: [blue]{tiled_version}[/blue]") + rprint(f" [bold]pdftools[/bold]: [blue]{pdftools_version}[/blue]\n") + + +class ProposalIDPrompt(Prompts): + """Custom IPython prompt that shows the current proposal ID.""" + + def __init__(self, RE: RunEngine, shell: TerminalInteractiveShell): + super().__init__(shell) + self._RE = RE + + def in_prompt_tokens(self, cli=None): + """Return the prompt tokens: `` []: ``.""" + return [ + (Token.Prompt, f"{self._RE.md.get('data_session', 'N/A')} ["), + (Token.PromptNum, str(self.shell.execution_count)), + (Token.Prompt, "]: "), + ] + + +def initialize_run_engine() -> RunEngine: + """Create the RunEngine, with Redis-backed metadata outside CI/mock mode. + + Returns + ------- + RunEngine + A RunEngine whose ``md`` is a ``RedisJSONDict`` on the beamline Redis + (``REDIS_HOST``), or an in-memory mock proposal in CI/mock mode. + """ + if is_running_in_ci(): + today = datetime.today() + return RunEngine( + { + "data_session": "pass-123456", + "cycle": f"{today.year}-{(today.month - 1) // 4 + 1}", + "proposal": { + "title": "Mock mode proposal", + "type": "Mock Commissioning", + "pi_name": os.getenv("USER", "Unknown"), + }, + } + ) + return RunEngine(RedisJSONDict(open_redis_client(redis_ssl=True), "")) # type: ignore[arg-type] + + +def print_proposal_info(md: MutableMapping[str, Any]): + """Print the active proposal from the RunEngine metadata. + + Parameters + ---------- + md : MutableMapping[str, Any] + The RunEngine metadata. + """ + proposal_md = md.get("proposal", {}) + if proposal_md: + rprint( + "Active proposal:\n" + f" Proposal title: [italic]{escape(str(proposal_md['title']))}[/italic]\n" + f" Proposal type: [italic]{escape(str(proposal_md['type']))}[/italic]\n" + f" Proposal PI: [italic]{escape(str(proposal_md['pi_name']))}[/italic]\n" + ) + + +def remove_md_keys(md: MutableMapping[str, Any], prefix: str = "PDF:") -> list[str]: + """Delete every metadata key that contains ``prefix``. + + Ported from ``remove_keys_redis`` in pdf-profile-collection's ``00-base.py`` + (C.-H. Lin, 2025-11-18), which trims unneeded ``PDF:`` entries from Redis. + + Parameters + ---------- + md : MutableMapping[str, Any] + The RunEngine metadata, usually Redis-backed. + prefix : str, optional + The text a key must contain to be removed, by default ``"PDF:"``. + + Returns + ------- + list[str] + The keys that were removed. + """ + removed = [key for key in list(md) if prefix in key] + for key in removed: + del md[key] + return removed + + +def _failed_paths(error: Exception, prefix: str = "") -> list[str]: + """Dotted paths of the leaves of a (nested) NotConnectedError.""" + sub = error.sub_errors if isinstance(error, NotConnectedError) else {} + if not sub: + return [prefix.rstrip(".")] + return [ + p for name, err in sub.items() for p in _failed_paths(err, f"{prefix}{name}.") + ] + + +def _count_connected_signals(device: Device, error: Exception) -> tuple[int, int]: + """Count (connected, total) signals of ``device``, given its connection error.""" + signals = [ + path for path, d in walk_devices(device).items() if isinstance(d, Signal) + ] + if not signals: + return 0, 1 + failed = [p for p in _failed_paths(error) if p] + if not failed: + return 0, len(signals) + down = [s for s in signals if any(s == f or s.startswith(f + ".") for f in failed)] + return len(signals) - len(down), len(signals) + + +def auto_init_devices( + timeout: float = 1.0, ok_message="OK", fail_message="FAILED", verbose: bool = False +) -> DeviceProcessor: + """Create a DeviceProcessor that connects devices, printing status for each. + + Devices connect in mock mode when :func:`is_running_in_ci` is true. A device + that fails to connect is reported and does not stop the others. + + Parameters + ---------- + timeout : float + The timeout in seconds for each device connection attempt. + ok_message : str, optional + The label printed for a connected device. + fail_message : str, optional + The label printed for a device that failed to connect. + verbose : bool, optional + Whether to print each failure's reason, by default False. + + Returns + ------- + DeviceProcessor + A DeviceProcessor that connects devices and prints their connection status. + """ + mock = is_running_in_ci() + # highlight=False stops rich from coloring dot-runs as ellipses. + console = Console(highlight=False) + console.print("\n[bold]Connecting to beamline devices...[/bold]\n") + + async def _process_devices(devices: dict[str, Device]): + for name, device in devices.items(): + if not device.name: + device.set_name(name) + coros = { + name: device.connect(mock, timeout) for name, device in devices.items() + } + failed: set[str] = set() + reasons: dict[str, str] = {} + e = NotConnectedError({}) + try: + await wait_for_connection(**coros) + except NotConnectedError as err: + e = err + failed = set(e.sub_errors.keys()) + reasons = {name: str(err) for name, err in e.sub_errors.items()} + + width = max(len(ok_message), len(fail_message)) + 2 + for name, device in devices.items(): + dots = "." * (40 - len(name)) + # Escaped brackets stay uncolored; only the label text is styled. + if name not in failed: + console.print( + f" {name} {dots} \\[[bold green]" + f"{ok_message.center(width)}[/bold green]]" + ) + continue + line = ( + f" {name} {dots} \\[[bold red]{fail_message.center(width)}[/bold red]]" + ) + connected, total = _count_connected_signals(device, e.sub_errors[name]) + if 0 < connected < total: + line += f" [dim]({connected}/{total} signals connected)[/dim]" + console.print(line) + if verbose and reasons: + console.print( + "\n".join(f"{name}: {reason}" for name, reason in reasons.items()) + ) + + return DeviceProcessor(_process_devices) diff --git a/tests/conftest.py b/tests/conftest.py index 81ab3fa..86cf311 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -93,6 +93,11 @@ async def fail_test_on_unclosed_tasks(request: FixtureRequest): raise error +@pytest.fixture(autouse=True) +def _set_ci_env(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("PDFTOOLS_RUNNING_IN_CI", "YES") + + @pytest.fixture(scope="function") def RE(request: FixtureRequest) -> RunEngine: loop = asyncio.new_event_loop() @@ -104,6 +109,15 @@ def RE(request: FixtureRequest) -> RunEngine: ) fail_count = request.session.testsfailed + async def _settle_in_flight_tasks(): + # A plan may end with a deliberate wait=False put (e.g. enable_preview_mode + # starting Acquire); give it a moment to finish before checking for leaks. + # On Python 3.11 asyncio.wait_for wraps it in an extra task, so it often + # outlives the RE call. Anything still pending afterwards is a real leak. + pending = asyncio.all_tasks() - {asyncio.current_task()} + if pending: + await asyncio.wait(pending, timeout=1.0) + def clean_event_loop(): if RE.state not in ("idle", "panicked"): try: @@ -111,6 +125,8 @@ def clean_event_loop(): except TransitionError: pass + if loop.is_running(): + asyncio.run_coroutine_threadsafe(_settle_in_flight_tasks(), loop).result() loop.call_soon_threadsafe(loop.stop) RE._th.join() @@ -149,6 +165,11 @@ def output_tiled_client(tmp_path: Path) -> Generator[Container, None, None]: server.close() +@pytest.fixture +def static_path_provider(tmp_path: Path) -> StaticPathProvider: + return StaticPathProvider(UUIDFilenameProvider(), tmp_path) + + @pytest.fixture def sample_detectors_factory(RE: RunEngine): def _factory(write_path: Path) -> tuple[SimBlobDetector, SimPointDetector]: diff --git a/tests/test_utils.py b/tests/test_utils.py new file mode 100644 index 0000000..82dbe98 --- /dev/null +++ b/tests/test_utils.py @@ -0,0 +1,147 @@ +import os +import subprocess +import sys +from pathlib import Path + +import pytest +from bluesky.run_engine import RunEngine +from ophyd_async.core import Device, NotConnectedError, soft_signal_rw + +import pdftools.utils as utils +from pdftools.utils import ( + auto_init_devices, + initialize_run_engine, + is_running_in_ci, + print_proposal_info, + remove_md_keys, +) + +PROFILE = Path(__file__).parents[1] / "src" / "pdftools" / "profiles" / "collection.py" + + +@pytest.mark.parametrize( + "value, expected", + [("YES", True), ("true", True), ("1", True), ("false", False), ("no", False)], +) +def test_is_running_in_ci(monkeypatch: pytest.MonkeyPatch, value: str, expected: bool): + monkeypatch.setenv(utils.CI_ENV_VAR, value) + assert is_running_in_ci() is expected + + +def test_initialize_run_engine_mock_proposal(): + RE = initialize_run_engine() + try: + assert RE.md["data_session"] == "pass-123456" + assert set(RE.md["proposal"]) == {"title", "type", "pi_name"} + finally: + RE._loop.call_soon_threadsafe(RE._loop.stop) + RE._th.join() + RE._loop.close() + + +def test_initialize_run_engine_uses_redis_outside_ci(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv(utils.CI_ENV_VAR, "false") + calls = [] + + def fake_open_redis_client(**kwargs): + calls.append(kwargs) + return "client" + + monkeypatch.setattr(utils, "open_redis_client", fake_open_redis_client) + monkeypatch.setattr( + utils, "RedisJSONDict", lambda client, prefix: {"client": client} + ) + RE = initialize_run_engine() + try: + assert calls == [{"redis_ssl": True}] + assert RE.md["client"] == "client" + finally: + RE._loop.call_soon_threadsafe(RE._loop.stop) + RE._th.join() + RE._loop.close() + + +def test_print_proposal_info(capsys: pytest.CaptureFixture[str]): + print_proposal_info( + { + "proposal": { + "title": "Powder study", + "type": "General User", + "pi_name": "A. B", + } + } + ) + out = capsys.readouterr().out + assert "Powder study" in out + assert "General User" in out + assert "A. B" in out + + +def test_print_proposal_info_keeps_markup_like_text(capsys: pytest.CaptureFixture[str]): + print_proposal_info( + { + "proposal": { + "title": "Study [in situ]", + "type": "General [/bold] User", + "pi_name": "A. B", + } + } + ) + out = capsys.readouterr().out + assert "Study [in situ]" in out + assert "General [/bold] User" in out + + +def test_print_proposal_info_without_proposal(capsys: pytest.CaptureFixture[str]): + print_proposal_info({}) + assert capsys.readouterr().out == "" + + +def test_remove_md_keys_removes_only_matching_keys(): + md = {"PDF:bt_safN": 1, "x_PDF:y": 2, "scan_id": 3, "data_session": "pass-1"} + assert sorted(remove_md_keys(md, prefix="PDF:")) == ["PDF:bt_safN", "x_PDF:y"] + assert md == {"scan_id": 3, "data_session": "pass-1"} + + +class _Box(Device): + def __init__(self, name: str = ""): + self.value = soft_signal_rw(float) + super().__init__(name=name) + + +class _Unreachable(Device): + def __init__(self, name: str = ""): + self.value = soft_signal_rw(float) + super().__init__(name=name) + + async def connect(self, mock=False, timeout=10.0, force_reconnect=False): + raise NotConnectedError({"value": NotConnectedError("no IOC")}) + + +def test_auto_init_devices_names_and_reports( + RE: RunEngine, capsys: pytest.CaptureFixture[str] +): + with auto_init_devices(timeout=0.1): + box = _Box() + gone = _Unreachable() + + assert box.name == "box" + assert gone.name == "gone" + out = capsys.readouterr().out + assert "box" in out and "OK" in out + assert "gone" in out and "FAILED" in out + + +def test_profile_boots_in_mock_mode(tmp_path: Path): + env = {k: v for k, v in os.environ.items() if k != utils.CI_ENV_VAR} + result = subprocess.run( + [sys.executable, str(PROFILE), "--mock"], + capture_output=True, + text=True, + env=env, + cwd=tmp_path, + timeout=120, + ) + assert result.returncode == 0, result.stderr + assert "Mock mode proposal" in result.stdout + assert "Connecting to beamline devices" in result.stdout