diff --git a/src/murfey/util/config.py b/src/murfey/util/config.py index 30d456f89..ad664d64f 100644 --- a/src/murfey/util/config.py +++ b/src/murfey/util/config.py @@ -1,5 +1,6 @@ from __future__ import annotations +import configparser import copy import os import socket @@ -7,6 +8,7 @@ from importlib.metadata import entry_points from pathlib import Path from typing import Any, Literal, Optional +from urllib.parse import quote import yaml from pydantic import BaseModel, ConfigDict, RootModel, ValidationInfo, field_validator @@ -373,3 +375,27 @@ def get_smartem_keycloak_client(): load_keycloak_config(Path(keycloak_config)) ) return keycloak_client + + +@lru_cache(maxsize=1) +def get_rabbitmq_url() -> str: + rabbitmq_defaults = { + "host": "localhost", + "port": "5672", + "username": "guest", + "password": "guest", + "vhost": "/", + } + rabbitmq_credentials_file = get_security_config().rabbitmq_credentials + cfgparser = configparser.ConfigParser(allow_no_value=True) + if rabbitmq_credentials_file: + cfgparser.read(rabbitmq_credentials_file) + rabbitmq_creds = ( + {**rabbitmq_defaults, **cfgparser["rabbit"]} + if cfgparser.has_section("rabbit") + else rabbitmq_defaults + ) + user = quote(rabbitmq_creds["username"], safe="") + password = quote(rabbitmq_creds["password"], safe="") + vhost = quote(rabbitmq_creds["vhost"], safe="") + return f"amqp://{user}:{password}@{rabbitmq_creds['host']}:{rabbitmq_creds['port']}/{vhost}" diff --git a/src/murfey/workflows/spa/ctf_estimation.py b/src/murfey/workflows/spa/ctf_estimation.py index 7e78374ce..605442859 100644 --- a/src/murfey/workflows/spa/ctf_estimation.py +++ b/src/murfey/workflows/spa/ctf_estimation.py @@ -1,8 +1,9 @@ +import asyncio from logging import getLogger from sqlmodel import Session, select -from murfey.util.config import get_machine_config +from murfey.util.config import get_machine_config, get_rabbitmq_url from murfey.util.db import ( Movie, Session as MurfeySession, @@ -20,6 +21,11 @@ MicrographResponse, ProcessingFeedbackPublishResponse, ) + from smartem_backend.model.mq_event import ( + CtfCompleteBody, + MessageQueueEventType, + ) + from smartem_backend.rmq.publisher import AioPikaPublisher from smartem_common.entity_status import MicrographStatus from murfey.util.config import get_smartem_keycloak_client @@ -37,7 +43,7 @@ def ctf_estimated(message: dict, murfey_db: Session) -> dict[str, bool]: if not SMARTEM_ACTIVE: return {"success": True} movie = murfey_db.exec( - select(Movie).where(Movie.murfey_id == message["motion_correction_id"]) + select(Movie).where(Movie.murfey_id == message["mc_uuid"]) ).one() if movie.smartem_uuid: try: @@ -69,6 +75,40 @@ def ctf_estimated(message: dict, murfey_db: Session) -> dict[str, bool]: registered_request, ProcessingFeedbackPublishResponse, ) + + async def _publish_ctf_completed( + micrograph_uuid: str, ctf_max_resolution: float + ) -> None: + publisher = AioPikaPublisher( + url=get_rabbitmq_url(), + exchange_name="smartem", + routing_key="smartem", + exchange_type="fanout", + ) + await publisher.connect() + try: + await publisher.publish_event( + MessageQueueEventType.CTF_COMPLETE, + CtfCompleteBody( + event_type=MessageQueueEventType.CTF_COMPLETE, + micrograph_uuid=micrograph_uuid, + ctf_max_resolution_estimate=ctf_max_resolution, + ), + ) + except Exception: + logger.warning( + f"smartem failed to ctf estimation completion {micrograph_uuid}", + exc_info=True, + ) + finally: + await publisher.close() + + asyncio.run( + _publish_ctf_completed( + movie.smartem_uuid, message.get("ctf_max_resolution", 1000) + ) + ) + except Exception: logger.warning( "Failed to emit CTF estimation complete event to smartem", diff --git a/src/murfey/workflows/spa/motion_correction.py b/src/murfey/workflows/spa/motion_correction.py index d1b0c5043..fa77be2be 100644 --- a/src/murfey/workflows/spa/motion_correction.py +++ b/src/murfey/workflows/spa/motion_correction.py @@ -1,8 +1,9 @@ +import asyncio from logging import getLogger from sqlmodel import Session, select -from murfey.util.config import get_machine_config +from murfey.util.config import get_machine_config, get_rabbitmq_url from murfey.util.db import ( Movie, Session as MurfeySession, @@ -20,6 +21,11 @@ MicrographResponse, ProcessingFeedbackPublishResponse, ) + from smartem_backend.model.mq_event import ( + MessageQueueEventType, + MotionCorrectionCompleteBody, + ) + from smartem_backend.rmq.publisher import AioPikaPublisher from smartem_common.entity_status import MicrographStatus from murfey.util.config import get_smartem_keycloak_client @@ -37,7 +43,7 @@ def motion_corrected(message: dict, murfey_db: Session) -> dict[str, bool]: if not SMARTEM_ACTIVE: return {"success": True} movie = murfey_db.exec( - select(Movie).where(Movie.murfey_id == message["motion_correction_id"]) + select(Movie).where(Movie.murfey_id == message["mc_uuid"]) ).one() if movie.smartem_uuid: try: @@ -71,6 +77,43 @@ def motion_corrected(message: dict, murfey_db: Session) -> dict[str, bool]: registered_request, ProcessingFeedbackPublishResponse, ) + + async def _publish_motion_correction_completed( + micrograph_uuid: str, total_motion: float, average_motion: float + ) -> None: + publisher = AioPikaPublisher( + url=get_rabbitmq_url(), + exchange_name="smartem", + routing_key="smartem", + exchange_type="fanout", + ) + await publisher.connect() + try: + await publisher.publish_event( + MessageQueueEventType.MOTION_CORRECTION_COMPLETE, + MotionCorrectionCompleteBody( + event_type=MessageQueueEventType.MOTION_CORRECTION_COMPLETE, + micrograph_uuid=micrograph_uuid, + total_motion=total_motion, + average_motion=average_motion, + ), + ) + except Exception: + logger.warning( + f"smartem failed to motion correction completion {micrograph_uuid}", + exc_info=True, + ) + finally: + await publisher.close() + + asyncio.run( + _publish_motion_correction_completed( + movie.smartem_uuid, + message.get("total_motion", 1000), + message.get("average_motion", 1000), + ) + ) + except Exception: logger.warning( "Failed to emit motion correction complete event to smartem", diff --git a/src/murfey/workflows/spa/picking.py b/src/murfey/workflows/spa/picking.py index 7de4bcbce..1f5610282 100644 --- a/src/murfey/workflows/spa/picking.py +++ b/src/murfey/workflows/spa/picking.py @@ -1,3 +1,4 @@ +import asyncio from logging import getLogger from typing import List @@ -11,7 +12,7 @@ _app_id, _pj_id, ) -from murfey.util.config import get_machine_config +from murfey.util.config import get_machine_config, get_rabbitmq_url from murfey.util.db import ( AutoProcProgram, ClassificationFeedbackParameters, @@ -33,6 +34,11 @@ from smartem_backend.api_client import SmartEMAPIClient from smartem_backend.model.http_request import MicrographUpdateRequest from smartem_backend.model.http_response import MicrographResponse + from smartem_backend.model.mq_event import ( + MessageQueueEventType, + ParticlePickingCompleteBody, + ) + from smartem_backend.rmq.publisher import AioPikaPublisher from smartem_common.entity_status import MicrographStatus from murfey.util.config import get_smartem_keycloak_client @@ -411,6 +417,45 @@ def particles_picked(message: dict, murfey_db: Session) -> dict[str, bool]: update, MicrographResponse, ) + + async def _publish_particle_picking_completed( + micrograph_uuid: str, number_of_particles_picked: int + ) -> None: + publisher = AioPikaPublisher( + url=get_rabbitmq_url(), + exchange_name="smartem", + routing_key="smartem", + exchange_type="fanout", + ) + await publisher.connect() + try: + await publisher.publish_event( + MessageQueueEventType.PARTICLE_PICKING_COMPLETE, + ParticlePickingCompleteBody( + event_type=MessageQueueEventType.PARTICLE_PICKING_COMPLETE, + micrograph_uuid=micrograph_uuid, + number_of_particles_picked=number_of_particles_picked, + ), + ) + except Exception: + logger.warning( + f"smartem failed to publish picking completion {micrograph_uuid}", + exc_info=True, + ) + finally: + await publisher.close() + + number_of_particles_picked = message.get("particle_count") + if number_of_particles_picked is None: + number_of_particles_picked = len( + message.get("particle_diameters") or [] + ) + asyncio.run( + _publish_particle_picking_completed( + movie.smartem_uuid, number_of_particles_picked + ) + ) + except Exception: logger.warning( "Failed to emit particle picking complete event to smartem",