From a780222d27a939b62534355e14571c4f1e5dd5ae Mon Sep 17 00:00:00 2001 From: GotthardG <51994228+GotthardG@users.noreply.github.com> Date: Tue, 17 Jun 2025 10:21:00 +0200 Subject: [PATCH 1/7] Refactor TELL integration and update dependencies. Replaced direct puck management logic with a WebSocket-based system for real-time updates from TELL. Introduced `tellupdater.py` to handle updates and improved AareDB client usage for cleaner API interactions. Updated `pyproject.toml` files to include dependency changes, ensuring compatibility with new TELL workflows. --- common/pyproject.toml | 3 +- daq/pyproject.toml | 9 +++- daq/src/aaredaq/aaredb.py | 45 ++++++---------- daq/src/aaredaq/tellupdater.py | 95 ++++++++++++++++++++++++++++++++++ daq/src/aaredaq/workflows.py | 2 +- daq/src/mxlibs3/tell_client.py | 21 ++++---- 6 files changed, 130 insertions(+), 45 deletions(-) create mode 100644 daq/src/aaredaq/tellupdater.py diff --git a/common/pyproject.toml b/common/pyproject.toml index e27ed15c..12bfa771 100644 --- a/common/pyproject.toml +++ b/common/pyproject.toml @@ -7,7 +7,8 @@ requires-python = ">=3.11" dependencies = [ "pydantic==2.11.4", "numpy==2.2.5", - "jfjoch_client==1.0.0rc44" + "jfjoch_client==1.0.0rc44", + "aareDB==0.1.1a6" ] [lint] diff --git a/daq/pyproject.toml b/daq/pyproject.toml index ff6c7e10..04f27111 100644 --- a/daq/pyproject.toml +++ b/daq/pyproject.toml @@ -13,7 +13,7 @@ dependencies = [ "fastapi==0.115.12", "uvicorn==0.34.2", "ultralytics==8.3.133", - "aareDB==0.1.1a4", + "aareDB==0.1.1a6", "opencv-python-headless==4.11.0.86", "python_multipart==0.0.20" ] @@ -21,6 +21,13 @@ dependencies = [ [lint] ignore = ["F401", "F541", "W503", "W504"] +[tool.uv.sources] +aaredb = { index = "psi"} + +[[tool.uv.index]] +name = "psi" +url = "https://gitea.psi.ch/api/packages/mx/pypi/simple" + [build-system] requires = ["setuptools>=75.6.0"] build-backend = "setuptools.build_meta" diff --git a/daq/src/aaredaq/aaredb.py b/daq/src/aaredaq/aaredb.py index e8b7ff58..9505c4d0 100644 --- a/daq/src/aaredaq/aaredb.py +++ b/daq/src/aaredaq/aaredb.py @@ -1,15 +1,18 @@ import io +import os from typing import List, Optional import aareDBclient import cv2 import numpy as np import requests -from aareDBclient import SetTellPosition, SampleEventCreate, SetTellPositionRequest +from aareDBclient import ( + SetTellPosition, + SampleEventCreate, + SetTellPositionRequest ) from aaredaqlib.models import ( SampleShortInfo, PuckLoadedInfo, - PuckInfo, DewarAddress, SampleShortInfoList, ) @@ -21,38 +24,20 @@ class AareWrapper: def __init__( self, bl: MXBeamline, - host: str = "https://mx-aare-test.psi.ch:1492", + host: str = "https://mx-db-01.psi.ch/dispatcher", ): configuration = aareDBclient.Configuration(host=host) configuration.verify_ssl = False # Disable SSL verification + configuration.default_headers = { + "X-Shared-Password": os.getenv("AAREDB_SHARED_PASSWORD") + } self.client = aareDBclient.ApiClient(configuration) self.__host = host - self.__puck_api = aareDBclient.PucksApi(self.client) - self.__sample_api = aareDBclient.SamplesApi(self.client) + self.__tell_api = aareDBclient.TellsRunnerApi(self.client) + self.__sample_api = aareDBclient.SamplesRunnerApi(self.client) self.__bl = bl - def get_pucks_beamline(self) -> List[PuckInfo]: - puck_list = self.__puck_api.get_pucks_by_slot_pucks_slot_slot_identifier_get( - slot_identifier=self.__bl.value.upper() - ) - ret = [] - for i in puck_list: - if i.tell_position is not None and len(i.tell_position) == 2: - d = DewarAddress(segment=i.tell_position[0], pos=i.tell_position[1]) - else: - d = None - ret.append( - PuckInfo( - db_id=i.id, - puck_name=i.puck_name, - dewar_name=i.dewar_name, - user=i.pgroup, - location=d, - ) - ) - return ret - def set_pucks_beamline(self, input_list: List[PuckLoadedInfo]): o = [] @@ -64,14 +49,14 @@ class AareWrapper: ) o.append(t) payload = SetTellPositionRequest(pucks = o, tell=self.__bl.value.upper()) - ret = self.__puck_api.set_tell_positions_pucks_set_tell_positions_put( + ret = self.__tell_api.set_tell_positions( set_tell_position_request=payload, ) print(ret) def get_sample_info(self) -> SampleShortInfoList: sample_list = ( - self.__puck_api.get_pucks_with_tell_position_pucks_with_tell_position_get(tell=self.__bl.value.upper()) + self.__tell_api.get_pucks_with_tell_position(tell=self.__bl.value.upper()) ) ret = [] @@ -102,7 +87,7 @@ class AareWrapper: def sample_mounted(self, s: Optional[SampleShortInfo]): if s is not None: try: - self.__sample_api.create_sample_event_samples_samples_sample_id_events_post( + self.__sample_api.create_sample_event( sample_id=s.db_id, sample_event_create=SampleEventCreate(event_type="Mounted"), ) @@ -112,7 +97,7 @@ class AareWrapper: def sample_unmounted(self, s: Optional[SampleShortInfo]): if s is not None: try: - self.__sample_api.create_sample_event_samples_samples_sample_id_events_post( + self.__sample_api.create_sample_event( sample_id=s.db_id, sample_event_create=SampleEventCreate(event_type="Unmounted"), ) diff --git a/daq/src/aaredaq/tellupdater.py b/daq/src/aaredaq/tellupdater.py new file mode 100644 index 00000000..e213df60 --- /dev/null +++ b/daq/src/aaredaq/tellupdater.py @@ -0,0 +1,95 @@ +import os +import json +import websocket +import threading +import time +from aareDBclient.models import PuckWithTellPosition +from mxlibs3.tell_client import TellClient +from aaredb import AareWrapper # Make sure the import path fits your project +from aaredaqlib.beamline import MXBeamline + +# Configuration +SLOT_IDENTIFIER = "X06DA" +WS_URL = f"wss://mx-db-01.psi.ch/dispatcher/protected_router/wstell/ws/slot/{SLOT_IDENTIFIER}" +WS_HEADERS = [f"X-Shared-Password: {os.getenv('AAREDB_SHARED_PASSWORD')}"] +POLL_INTERVAL = 30 # seconds + +# Initialize TELL client and DB wrapper +tell_client = TellClient() +beamline = MXBeamline.X06DA # Use your beamline enum/value +aare_db = AareWrapper(bl=beamline) + +# Track current state +current_pucks = [] + +def compare_and_report_change(old, new, key_func): + old_ids = set(key_func(p) for p in old) + new_ids = set(key_func(p) for p in new) + joined = new_ids - old_ids + left = old_ids - new_ids + return joined, left + +def ws_update_samples_info(pucks): + """Send sample info to TELL robot.""" + tell_client.set_samples_info(pucks) + +def timer_update_pucks_beamline(): + """Poll tell for detected pucks, and update the DATABASE with their location.""" + try: + detected_pucks = tell_client.get_detected_pucks() + aare_db.set_pucks_beamline(detected_pucks) + except Exception as exc: + print(f"[ERROR][Timer] {exc}") + +def on_message(ws, message): + global current_pucks + try: + data = json.loads(message) + print("[WS][MESSAGE]", data) + if isinstance(data, list): + new_pucks = [PuckWithTellPosition(**item) for item in data] + joined, left = compare_and_report_change( + current_pucks, new_pucks, key_func=lambda p: p.id + ) + if joined or left: + print(f"[WS][UPDATE] joined={joined}, left={left}") + current_pucks = new_pucks + ws_update_samples_info(new_pucks) + else: + print("[WS][NO CHANGE]") + else: + print("[WS][ERROR] Unknown message format") + except Exception as exc: + print("[WS][ERROR] Failed to parse websocket message:", exc) + +def on_error(ws, error): + print("[WS][ERROR]", error) + +def on_close(ws, close_status_code, close_msg): + print("[WS][CLOSE]", close_status_code, close_msg) + +def on_open(ws): + print("[WS][OPEN] WebSocket opened.") + +def periodic_polling(): + while True: + timer_update_pucks_beamline() + time.sleep(POLL_INTERVAL) + +def main(): + ws = websocket.WebSocketApp( + WS_URL, + header=WS_HEADERS, + on_message=on_message, + on_error=on_error, + on_close=on_close, + on_open=on_open, + ) + + poll_thread = threading.Thread(target=periodic_polling, daemon=True) + poll_thread.start() + + ws.run_forever(sslopt={"cert_reqs": 0}) + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/daq/src/aaredaq/workflows.py b/daq/src/aaredaq/workflows.py index 04fb200e..13497660 100644 --- a/daq/src/aaredaq/workflows.py +++ b/daq/src/aaredaq/workflows.py @@ -284,7 +284,7 @@ def sa2dh(devs: BeamlineDevices, cfg: BeamlineConfig): print(f"Error for unmounting: {e}") devs.tell.dry(wait_cold=-1, wait=False) # devs.tell.move_park(wait=True) - devs.tell.set_samples_info(devs.aare.get_pucks_beamline()) + #devs.tell.set_samples_info(devs.aare.get_pucks_beamline())# replaced by script stream def dh2sa(devs: BeamlineDevices, cfg: BeamlineConfig): diff --git a/daq/src/mxlibs3/tell_client.py b/daq/src/mxlibs3/tell_client.py index 97c29ce1..d7d05f66 100755 --- a/daq/src/mxlibs3/tell_client.py +++ b/daq/src/mxlibs3/tell_client.py @@ -12,9 +12,11 @@ from aaredaqlib.models import ( SampleShortInfo, DewarAddress, SampleDewarAddress, - PuckInfo, + #PuckInfo, +) +from aareDBclient import ( + PuckWithTellPosition, ) - from aaredaqlib.beamline import MXBeamline # noqa: F401 from mxlibs3.pshell_client import PShellClient @@ -191,24 +193,19 @@ class TellClient: ) return output - def set_samples_info(self, info: List[PuckInfo]): + def set_samples_info(self, info: List[PuckWithTellPosition]): if self.__simulation: return j = [] for x in info: - if x.location is None: - puck_address = "" - else: - puck_address = "{:1s}{:1d}".format(x.location.segment, x.location.pos) - j.append( { - "userName": x.user, - "dewarName": x.dewar_name, + "userName": x.pgroup, + "dewarName": x.dewar_name or "", "puckName": x.puck_name, - "puckType": "Unipuck", - "puckAddress": puck_address, + "puckType": "Unipuck", # could use x.puck_type + "puckAddress": x.tell_position or "", "puckBarcode": x.puck_name, "sampleBarcode": "", "sampleMountCount": 0, From dc0874d8d44f89c7275a0511e4943f8625e47344 Mon Sep 17 00:00:00 2001 From: GotthardG <51994228+GotthardG@users.noreply.github.com> Date: Tue, 17 Jun 2025 13:36:27 +0200 Subject: [PATCH 2/7] Refactor TELL integration and update dependencies. Replaced direct puck management logic with a WebSocket-based system for real-time updates from TELL. Introduced `tellupdater.py` to handle updates and improved AareDB client usage for cleaner API interactions. Updated `pyproject.toml` files to include dependency changes, ensuring compatibility with new TELL workflows. --- daq/src/aaredaq/tellupdater.py | 33 ++++++++++++++++++++++++++++++++- 1 file changed, 32 insertions(+), 1 deletion(-) diff --git a/daq/src/aaredaq/tellupdater.py b/daq/src/aaredaq/tellupdater.py index e213df60..58f9d211 100644 --- a/daq/src/aaredaq/tellupdater.py +++ b/daq/src/aaredaq/tellupdater.py @@ -21,6 +21,7 @@ aare_db = AareWrapper(bl=beamline) # Track current state current_pucks = [] +last_pucks_sent = [] def compare_and_report_change(old, new, key_func): old_ids = set(key_func(p) for p in old) @@ -33,11 +34,41 @@ def ws_update_samples_info(pucks): """Send sample info to TELL robot.""" tell_client.set_samples_info(pucks) +def convert_pucks(puck_list): + return [PuckWithTellPosition( + id=p.id, + puck_name=p.puck_name, + puck_type=getattr(p, 'puck_type', None), + puck_location_in_dewar=getattr(p, 'puck_location_in_dewar', None), + dewar_id=getattr(p, 'dewar_id', None), + dewar_name=getattr(p, 'dewar_name', None), + pgroup=getattr(p, 'pgroup', None), + samples=getattr(p, 'samples', None), + tell_position=getattr(p, 'tell_position', None), + ) for p in puck_list] + +def pucks_equal(pucks1, pucks2): + # Implement a simple equality check, e.g., comparing sorted ids, or full data + if len(pucks1) != len(pucks2): return False + return all(p1.id == p2.id for p1, p2 in zip(sorted(pucks1, key=lambda p: p.id), + sorted(pucks2, key=lambda p: p.id))) + def timer_update_pucks_beamline(): - """Poll tell for detected pucks, and update the DATABASE with their location.""" + global last_pucks_sent try: detected_pucks = tell_client.get_detected_pucks() aare_db.set_pucks_beamline(detected_pucks) + + # Convert to list[PuckWithTellPosition] + pwtp_list = convert_pucks(detected_pucks) + + # Only send if state changed + if not pucks_equal(pwtp_list, last_pucks_sent): + tell_client.set_samples_info(pwtp_list) + last_pucks_sent = pwtp_list + else: + print("[TIMER] No puck state change, not sending to TELL.") + except Exception as exc: print(f"[ERROR][Timer] {exc}") From 4003cbd2b16b94e3647d70da6fa776764e6bc933 Mon Sep 17 00:00:00 2001 From: GotthardG <51994228+GotthardG@users.noreply.github.com> Date: Tue, 17 Jun 2025 13:38:31 +0200 Subject: [PATCH 3/7] Refactor TELL integration and update dependencies. Replaced direct puck management logic with a WebSocket-based system for real-time updates from TELL. Introduced `tellupdater.py` to handle updates and improved AareDB client usage for cleaner API interactions. Updated `pyproject.toml` files to include dependency changes, ensuring compatibility with new TELL workflows. --- daq/src/aaredaq/tellupdater.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/daq/src/aaredaq/tellupdater.py b/daq/src/aaredaq/tellupdater.py index 58f9d211..19a3b487 100644 --- a/daq/src/aaredaq/tellupdater.py +++ b/daq/src/aaredaq/tellupdater.py @@ -15,8 +15,8 @@ WS_HEADERS = [f"X-Shared-Password: {os.getenv('AAREDB_SHARED_PASSWORD')}"] POLL_INTERVAL = 30 # seconds # Initialize TELL client and DB wrapper -tell_client = TellClient() beamline = MXBeamline.X06DA # Use your beamline enum/value +tell_client = TellClient(bl=beamline) aare_db = AareWrapper(bl=beamline) # Track current state From 4b1acfb8277399734d083c77861c4d64bbdb6ed0 Mon Sep 17 00:00:00 2001 From: GotthardG <51994228+GotthardG@users.noreply.github.com> Date: Tue, 17 Jun 2025 13:43:15 +0200 Subject: [PATCH 4/7] Refactor TELL integration and update dependencies. Replaced direct puck management logic with a WebSocket-based system for real-time updates from TELL. Introduced `tellupdater.py` to handle updates and improved AareDB client usage for cleaner API interactions. Updated `pyproject.toml` files to include dependency changes, ensuring compatibility with new TELL workflows. --- daq/src/aaredaq/tellupdater.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/daq/src/aaredaq/tellupdater.py b/daq/src/aaredaq/tellupdater.py index 19a3b487..3c64d5d6 100644 --- a/daq/src/aaredaq/tellupdater.py +++ b/daq/src/aaredaq/tellupdater.py @@ -36,7 +36,7 @@ def ws_update_samples_info(pucks): def convert_pucks(puck_list): return [PuckWithTellPosition( - id=p.id, + id=p.puck_id, puck_name=p.puck_name, puck_type=getattr(p, 'puck_type', None), puck_location_in_dewar=getattr(p, 'puck_location_in_dewar', None), From c6a4cf0e8b6c4bd72f98a2809804c37e3d69b6fa Mon Sep 17 00:00:00 2001 From: GotthardG <51994228+GotthardG@users.noreply.github.com> Date: Tue, 17 Jun 2025 14:13:46 +0200 Subject: [PATCH 5/7] Refactor TELL integration and update dependencies. Replaced direct puck management logic with a WebSocket-based system for real-time updates from TELL. Introduced `tellupdater.py` to handle updates and improved AareDB client usage for cleaner API interactions. Updated `pyproject.toml` files to include dependency changes, ensuring compatibility with new TELL workflows. --- daq/src/aaredaq/tellupdater.py | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/daq/src/aaredaq/tellupdater.py b/daq/src/aaredaq/tellupdater.py index 3c64d5d6..5d712ce5 100644 --- a/daq/src/aaredaq/tellupdater.py +++ b/daq/src/aaredaq/tellupdater.py @@ -59,16 +59,6 @@ def timer_update_pucks_beamline(): detected_pucks = tell_client.get_detected_pucks() aare_db.set_pucks_beamline(detected_pucks) - # Convert to list[PuckWithTellPosition] - pwtp_list = convert_pucks(detected_pucks) - - # Only send if state changed - if not pucks_equal(pwtp_list, last_pucks_sent): - tell_client.set_samples_info(pwtp_list) - last_pucks_sent = pwtp_list - else: - print("[TIMER] No puck state change, not sending to TELL.") - except Exception as exc: print(f"[ERROR][Timer] {exc}") From bf170cfcf5c8f08ada5fa52616e6e5f8b76744d5 Mon Sep 17 00:00:00 2001 From: GotthardG <51994228+GotthardG@users.noreply.github.com> Date: Tue, 17 Jun 2025 14:15:40 +0200 Subject: [PATCH 6/7] Refactor TELL integration and update dependencies. Replaced direct puck management logic with a WebSocket-based system for real-time updates from TELL. Introduced `tellupdater.py` to handle updates and improved AareDB client usage for cleaner API interactions. Updated `pyproject.toml` files to include dependency changes, ensuring compatibility with new TELL workflows. --- daq/src/aaredaq/aaredb.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/daq/src/aaredaq/aaredb.py b/daq/src/aaredaq/aaredb.py index 9505c4d0..de7a57c4 100644 --- a/daq/src/aaredaq/aaredb.py +++ b/daq/src/aaredaq/aaredb.py @@ -28,11 +28,9 @@ class AareWrapper: ): configuration = aareDBclient.Configuration(host=host) configuration.verify_ssl = False # Disable SSL verification - configuration.default_headers = { - "X-Shared-Password": os.getenv("AAREDB_SHARED_PASSWORD") - } self.client = aareDBclient.ApiClient(configuration) + self.client.default_headers["X-Shared-Password"] = os.getenv("AAREDB_SHARED_PASSWORD") self.__host = host self.__tell_api = aareDBclient.TellsRunnerApi(self.client) self.__sample_api = aareDBclient.SamplesRunnerApi(self.client) From f369e5b5558044c56ad9b7caf8fceb6fea22f0ce Mon Sep 17 00:00:00 2001 From: GotthardG <51994228+GotthardG@users.noreply.github.com> Date: Tue, 17 Jun 2025 14:15:51 +0200 Subject: [PATCH 7/7] Refactor TELL integration and update dependencies. Replaced direct puck management logic with a WebSocket-based system for real-time updates from TELL. Introduced `tellupdater.py` to handle updates and improved AareDB client usage for cleaner API interactions. Updated `pyproject.toml` files to include dependency changes, ensuring compatibility with new TELL workflows. --- common/pyproject.toml | 7 +++++++ daq/pyproject.toml | 3 ++- 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/common/pyproject.toml b/common/pyproject.toml index 12bfa771..1b2c1857 100644 --- a/common/pyproject.toml +++ b/common/pyproject.toml @@ -14,6 +14,13 @@ dependencies = [ [lint] ignore = ["F401", "F541", "W503", "W504"] +[tool.uv.sources] +aaredb = { index = "psi"} + +[[tool.uv.index]] +name = "psi" +url = "https://gitea.psi.ch/api/packages/mx/pypi/simple" + [build-system] requires = ["setuptools>=75.6.0"] build-backend = "setuptools.build_meta" diff --git a/daq/pyproject.toml b/daq/pyproject.toml index 04f27111..72764df5 100644 --- a/daq/pyproject.toml +++ b/daq/pyproject.toml @@ -15,7 +15,8 @@ dependencies = [ "ultralytics==8.3.133", "aareDB==0.1.1a6", "opencv-python-headless==4.11.0.86", - "python_multipart==0.0.20" + "python_multipart==0.0.20", + "websocket-client" ] [lint]