diff --git a/src/aare/common/find_xtal.py b/src/aare/common/find_xtal.py index 23468175..9fefdb2f 100644 --- a/src/aare/common/find_xtal.py +++ b/src/aare/common/find_xtal.py @@ -4,11 +4,13 @@ import numpy as np from scipy import ndimage from aare.common.models import CrystalSize +from aare.common.logger_events import log_timing from aare.common.raster_grid import RasterGridRequest, CenterOfMassModel from aare.common.logger_config import setup_logger logger = setup_logger('aareDAQ') +@log_timing(logger, "Identify crystal raster") def identify_crystal_raster(result, r: RasterGridRequest) -> CenterOfMassModel | None: images = result.images if images and any(getattr(img, "spots_low_res", 0) for img in images): @@ -51,6 +53,7 @@ def identify_crystal_raster(result, r: RasterGridRequest) -> CenterOfMassModel | else: return None +@log_timing(logger, "Rebuild array from scan results") def rebuild_array_from_scan_results(scan_results: List, value_field: str, array_shape: Optional[tuple] = None, @@ -74,7 +77,7 @@ def rebuild_array_from_scan_results(scan_results: List, if nx is None or ny is None: continue - positions.append((int(nx), int(ny))) # Note: (row, col) = (ny, nx) + positions.append((int(ny), int(nx))) # Corrected: (row, col) = (ny, nx) if not value: value = 0.0 values.append(float(value)) diff --git a/src/aare/common/logger_events.py b/src/aare/common/logger_events.py new file mode 100644 index 00000000..6a1ee5a4 --- /dev/null +++ b/src/aare/common/logger_events.py @@ -0,0 +1,158 @@ +import functools +import time +from typing import Any, Callable +import logging + +from aare.common.raster_grid import RasterGridRequest +from aare.common.rotation_scan import RotationScanRequest + + +def log_timing( + logger: logging.Logger, + message_prefix: str = "", + level: int = logging.DEBUG, + extra: dict[str, Any] | None = None, +) -> Callable: + """Decorator to time a function call and log the duration.""" + def decorator(func: Callable) -> Callable: + @functools.wraps(func) + def wrapper(*args: Any, **kwargs: Any) -> Any: + start = time.perf_counter() + prefix = f"{message_prefix}: " if message_prefix else "" + logger.log(level, f"{prefix}Starting {func.__name__}") + try: + result = func(*args, **kwargs) + duration = time.perf_counter() - start + logger.log( + level, + f"{prefix}Finished {func.__name__} in {duration:.4f} seconds", + extra=merge_log_context(extra, duration_s=duration) + ) + return result + except Exception as e: + duration = time.perf_counter() - start + logger.log( + level, + f"{prefix}{func.__name__} FAILED after {duration:.4f} seconds with error: {e}", + extra=merge_log_context(extra, duration_s=duration) + ) + raise + return wrapper + return decorator + + +def merge_log_context(*parts: dict[str, Any] | None, **extra: Any) -> dict[str, Any]: + payload: dict[str, Any] = {} + for part in parts: + if part: + payload.update(part) + payload.update(extra) + return payload + + +def sample_log_context(sample) -> dict[str, Any]: + return { + "sample_id": getattr(sample, "db_id", None), + "sample_name": getattr(sample, "sample_name", None), + } + + +def raster_request_log_context(request: RasterGridRequest | None) -> dict[str, Any]: + if request is None: + return {} + + return { + "file_prefix": request.file_prefix, + "omega_deg": request.omega_deg, + "dtz": request.dtz, + "transmission": request.transmission, + "n_x": request.n_x, + "n_y": request.n_y, + "grid_size_x_mm": getattr(request.grid_size_mm, "x", None), + "grid_size_y_mm": getattr(request.grid_size_mm, "y", None), + } + + +def rotation_request_log_context( + request: RotationScanRequest | None, + *, + total_time_s: float | None = None, +) -> dict[str, Any]: + if request is None: + return {} + + payload = { + "file_prefix": request.file_prefix, + "start_omega_deg": request.start_omega_deg, + "dtz": request.dtz, + "exp_time_s": request.exp_time_s, + "incr_omega_deg": request.incr_omega_deg, + "steps": request.steps, + "transmission": request.transmission, + "screening": getattr(request, "screening", None), + } + + if total_time_s is not None: + payload["total_time_s"] = total_time_s + + return payload + + +def geom_log_context(geom, *, prefix: str = "") -> dict[str, Any]: + smargon = getattr(geom, "smargon", None) + sh_mm = getattr(smargon, "sh_mm", None) + + return { + f"{prefix}omega_deg": getattr(geom, "omega_deg", None), + f"{prefix}beam_x_pxl": getattr(getattr(geom, "beam_location_pxl", None), "x", None), + f"{prefix}beam_y_pxl": getattr(getattr(geom, "beam_location_pxl", None), "y", None), + f"{prefix}pixel_in_mm": getattr(geom, "pixel_in_mm", None), + f"{prefix}sh_x_mm": getattr(sh_mm, "x", None), + f"{prefix}sh_y_mm": getattr(sh_mm, "y", None), + f"{prefix}sh_z_mm": getattr(sh_mm, "z", None), + } + + +def ml_bundle_meta_log_context( + *, + target_point: tuple[float, float] | None = None, + focus: float | None = None, +) -> dict[str, Any]: + return { + "target_point": target_point, + "focus": focus, + } + + +def log_ml_bundle_meta( + logger: logging.Logger, + context: str, + *, + target_point: tuple[float, float] | None = None, + focus: float | None = None, +) -> None: + if target_point is None and focus is None: + return + + logger.debug( + "ML bundle metadata", + extra=merge_log_context( + {"context": context}, + ml_bundle_meta_log_context(target_point=target_point, focus=focus), + ), + ) + + +def log_duration( + logger: logging.Logger, + message: str, + duration_s: float, + *, + level: int = logging.INFO, + extra: dict[str, Any] | None = None, +) -> None: + logger.log( + level, + message, + extra=merge_log_context(extra, duration_s=duration_s), + ) \ No newline at end of file diff --git a/src/aare/daq/aaredb.py b/src/aare/daq/aaredb.py index 77eaeee7..adc9b27a 100644 --- a/src/aare/daq/aaredb.py +++ b/src/aare/daq/aaredb.py @@ -25,6 +25,7 @@ from aareDB import ( from aare.common.coordinate import Coordinate from aare.common.logger_config import setup_logger +from aare.common.logger_events import log_timing from aare.common.models import ( SampleShortInfo, PuckLoadedInfo, @@ -40,23 +41,6 @@ from jfjoch_client.models import ScanResult logger = setup_logger("aareDAQ") -def time_db_call(func): - """Decorator to time AareDB calls and log the duration.""" - @functools.wraps(func) - def wrapper(*args, **kwargs): - start = time.perf_counter() - logger.debug(f"Starting AareDB call: {func.__name__}") - try: - result = func(*args, **kwargs) - duration = time.perf_counter() - start - logger.debug(f"Finished AareDB call: {func.__name__} in {duration:.4f} seconds") - return result - except Exception as e: - duration = time.perf_counter() - start - logger.debug(f"AareDB call: {func.__name__} FAILED after {duration:.4f} seconds with error: {e}") - raise - return wrapper - class AareWrapper: def __init__( @@ -91,7 +75,7 @@ class AareWrapper: self.__cert_file = configuration.cert_file self.__key_file = configuration.key_file - @time_db_call + @log_timing(logger, "AareDB call") def set_pucks_beamline(self, input_list: List[PuckLoadedInfo]): o = [] @@ -108,7 +92,7 @@ class AareWrapper: ) logger.debug(ret) - @time_db_call + @log_timing(logger, "AareDB call") def create_manual_sample(self, s: SampleShortInfo): from aareDB.models import ManualSampleCreate @@ -123,7 +107,7 @@ class AareWrapper: except Exception as e: logger.error(f"Error inserting sample: {e}") - @time_db_call + @log_timing(logger, "AareDB call") def send_sample_event( self, s: Optional[SampleShortInfo], @@ -149,7 +133,7 @@ class AareWrapper: except Exception as e: logger.error(f"Error sending sample event {event_type!s} to db: {e}") - @time_db_call + @log_timing(logger, "AareDB call") def upload_image(self, sample_id: int, filename: str, bgr_image: np.ndarray, message: Optional[str] = None): _, buffer = cv2.imencode('.jpg', bgr_image) jpeg_bytes = io.BytesIO(buffer) @@ -170,7 +154,7 @@ class AareWrapper: response = requests.post(url, **request_kwargs) logger.debug(f"Response status code: {response.status_code}") - @time_db_call + @log_timing(logger, "AareDB call") def upload_jpg(self, sample_id: int, filename: str, jpg_image): logger.debug(f"jppg_image of type: {type(jpg_image)}") url = f"{self.__host}/protected_router/sample_runner/{sample_id}/upload-images" @@ -185,7 +169,7 @@ class AareWrapper: headers=headers) logger.debug(f"Response status code: {response.status_code}") - @time_db_call + @log_timing(logger, "AareDB call") def create_rotation_run(self, s: Optional[SampleShortInfo], r:RotationScanRequest, d:DAQStatusModel): if s is None: return @@ -261,7 +245,7 @@ class AareWrapper: except Exception as e: logger.error(e) - @time_db_call + @log_timing(logger, "AareDB call") def create_gridscan_run(self, s: Optional[SampleShortInfo], r:RasterGridRequest, d:DAQStatusModel): if s is None: return @@ -326,7 +310,7 @@ class AareWrapper: except Exception as e: logger.debug(e) - @time_db_call + @log_timing(logger, "AareDB call") def ingest_gridscan(self, sample: Optional[SampleShortInfo], raster_result: ScanResult, raster_request: RasterGridRequest, geom: SampleGeometryModel, com: Optional[CenterOfMassModel], beam_mark_pxl:tuple[float,float]): @@ -397,7 +381,7 @@ class AareWrapper: logger.error(e) raise e - @time_db_call + @log_timing(logger, "AareDB call") def ingest_scan(self, sample: Optional[SampleShortInfo], result: ScanResult, geom: SampleGeometryModel, beam_mark_pxl:tuple[float,float]): diff --git a/src/aare/daq/daq.py b/src/aare/daq/daq.py index f73a3fc0..ea5650cc 100644 --- a/src/aare/daq/daq.py +++ b/src/aare/daq/daq.py @@ -26,6 +26,16 @@ from aare.common.beamline import MXBeamline from aare.common.coordinate import Coordinate, SmargonCoordinate, AerotechCoordinate from aare.common.diffraction_geometry import DiffractionGeometry from aare.common.logger_config import setup_logger +from aare.common.logger_events import ( + geom_log_context, + log_duration, + log_timing, + log_ml_bundle_meta, + merge_log_context, + raster_request_log_context, + rotation_request_log_context, + sample_log_context, +) from aare.common.models import ( SampleShortInfo, PuckLoadedInfo, @@ -205,19 +215,6 @@ class AareDAQ: progress.success = success self._emit_automation_progress(progress) - def _log_ml_bundle_meta( - self, - context: str, - *, - target_point: tuple[float, float] | None = None, - focus: float | None = None, - ) -> None: - if target_point is None and focus is None: - return - logger.debug( - f"ML bundle metadata for {context}: target_point={target_point}, focus={focus}" - ) - #-------------------------------------------- # Operation Handlers #-------------------------------------------- @@ -381,19 +378,11 @@ class AareDAQ: try: logger.info( "Starting raster sequence", - extra={ - "auto_center": auto_center, - "file_prefix": grid_request.file_prefix, - "omega_deg": grid_request.omega_deg, - "dtz": grid_request.dtz, - "transmission": grid_request.transmission, - "n_x": grid_request.n_x, - "n_y": grid_request.n_y, - "grid_size_x_mm": getattr(grid_request.grid_size_mm, "x", None), - "grid_size_y_mm": getattr(grid_request.grid_size_mm, "y", None), - "sample_id": getattr(self.sample, "db_id", None), - "sample_name": getattr(self.sample, "sample_name", None), - }, + extra=merge_log_context( + sample_log_context(self.sample), + raster_request_log_context(grid_request), + {"auto_center": auto_center}, + ), ) if auto_center: @@ -404,12 +393,10 @@ class AareDAQ: if result is None: logger.error( "Raster sequence returned no result after auto-centering", - extra={ - "file_prefix": grid_request.file_prefix, - "omega_deg": grid_request.omega_deg, - "sample_id": getattr(self.sample, "db_id", None), - "sample_name": getattr(self.sample, "sample_name", None), - }, + extra=merge_log_context( + sample_log_context(self.sample), + raster_request_log_context(grid_request), + ), ) else: self.__setup_datacollection(request=grid_request) @@ -421,12 +408,14 @@ class AareDAQ: if result is not None: logger.info( "Raster sequence completed", - extra={ - "auto_center": auto_center, - "file_prefix": grid_request.file_prefix, - "result_count": len(result.r) if hasattr(result, "r") and result.r is not None else None, - "sample_id": getattr(self.sample, "db_id", None), - }, + extra=merge_log_context( + sample_log_context(self.sample), + raster_request_log_context(grid_request), + { + "auto_center": auto_center, + "result_count": len(result.r) if hasattr(result, "r") and result.r is not None else None, + }, + ), ) return result @@ -434,12 +423,10 @@ class AareDAQ: except JFJochCommunicationError as e: logger.exception( "Raster sequence failed due to JFJoch communication error", - extra={ - "file_prefix": grid_request.file_prefix, - "omega_deg": grid_request.omega_deg, - "sample_id": getattr(self.sample, "db_id", None), - "sample_name": getattr(self.sample, "sample_name", None), - }, + extra=merge_log_context( + sample_log_context(self.sample), + raster_request_log_context(grid_request), + ), ) self._handle_operation_error( operation=DAQOperation.RASTER, @@ -452,17 +439,11 @@ class AareDAQ: except Exception as e: logger.exception( "Raster sequence failed", - extra={ - "auto_center": auto_center, - "file_prefix": grid_request.file_prefix, - "omega_deg": grid_request.omega_deg, - "dtz": grid_request.dtz, - "transmission": grid_request.transmission, - "n_x": grid_request.n_x, - "n_y": grid_request.n_y, - "sample_id": getattr(self.sample, "db_id", None), - "sample_name": getattr(self.sample, "sample_name", None), - }, + extra=merge_log_context( + sample_log_context(self.sample), + raster_request_log_context(grid_request), + {"auto_center": auto_center}, + ), ) self._handle_operation_error( operation=DAQOperation.RASTER, @@ -1017,18 +998,17 @@ class AareDAQ: logger.info( "Starting auto-center raster workflow", - extra={ - "sample_id": getattr(sample, "db_id", None), - "sample_name": getattr(sample, "sample_name", None), - "omega_deg": geom.omega_deg, - "current_sh_x_mm": geom.smargon.sh_mm.x, - "current_sh_y_mm": geom.smargon.sh_mm.y, - "current_sh_z_mm": geom.smargon.sh_mm.z, - "beam_x_pxl": geom.beam_location_pxl.x, - "beam_y_pxl": geom.beam_location_pxl.y, - "pixel_in_mm": geom.pixel_in_mm, - "file_prefix": old_prefix, - }, + extra=merge_log_context( + sample_log_context(sample), + geom_log_context(geom, prefix="current_"), + { + "omega_deg": geom.omega_deg, + "beam_x_pxl": geom.beam_location_pxl.x, + "beam_y_pxl": geom.beam_location_pxl.y, + "pixel_in_mm": geom.pixel_in_mm, + "file_prefix": old_prefix, + }, + ), ) r = self.__ml_bounding_box(sample.db_id, f"ml_{geom.omega_deg:.2f}deg") @@ -1036,11 +1016,13 @@ class AareDAQ: if r is None: logger.warning( "No ML bounding box found at primary angle during auto-center raster", - extra={ - "sample_id": getattr(sample, "db_id", None), - "omega_deg": geom.omega_deg, - "file_prefix": old_prefix, - }, + extra=merge_log_context( + sample_log_context(sample), + { + "omega_deg": geom.omega_deg, + "file_prefix": old_prefix, + }, + ), ) self.__devs.aerotech_omega = geom.omega_deg + 90.0 time.sleep(0.2) @@ -1049,17 +1031,19 @@ class AareDAQ: if r is not None: logger.info( "ML bounding box found for auto-center raster", - extra={ - "sample_id": getattr(sample, "db_id", None), - "ml_omega_deg": r.omega_deg, - "ml_n_x": r.n_x, - "ml_n_y": r.n_y, - "ml_grid_size_x_mm": getattr(r.grid_size_mm, "x", None), - "ml_grid_size_y_mm": getattr(r.grid_size_mm, "y", None), - "ml_top_left_x_mm": getattr(getattr(r.smargon_top_left, "sh_mm", None), "x", None), - "ml_top_left_y_mm": getattr(getattr(r.smargon_top_left, "sh_mm", None), "y", None), - "ml_top_left_z_mm": getattr(getattr(r.smargon_top_left, "sh_mm", None), "z", None), - }, + extra=merge_log_context( + sample_log_context(sample), + { + "ml_omega_deg": r.omega_deg, + "ml_n_x": r.n_x, + "ml_n_y": r.n_y, + "ml_grid_size_x_mm": getattr(r.grid_size_mm, "x", None), + "ml_grid_size_y_mm": getattr(r.grid_size_mm, "y", None), + "ml_top_left_x_mm": getattr(getattr(r.smargon_top_left, "sh_mm", None), "x", None), + "ml_top_left_y_mm": getattr(getattr(r.smargon_top_left, "sh_mm", None), "y", None), + "ml_top_left_z_mm": getattr(getattr(r.smargon_top_left, "sh_mm", None), "z", None), + }, + ), ) self.__set_state(BeamlineStateEnum.DataCollection) geom = self.sample_geometry @@ -1073,15 +1057,15 @@ class AareDAQ: logger.info( "Running first auto-center raster", - extra={ - "file_prefix": grid.file_prefix, - "omega_deg": grid.omega_deg, - "n_x": grid.n_x, - "n_y": grid.n_y, - "top_left_x_mm": getattr(grid.smargon_top_left.sh_mm, "x", None), - "top_left_y_mm": getattr(grid.smargon_top_left.sh_mm, "y", None), - "top_left_z_mm": getattr(grid.smargon_top_left.sh_mm, "z", None), - }, + extra=merge_log_context( + sample_log_context(self.sample), + raster_request_log_context(grid), + { + "top_left_x_mm": getattr(grid.smargon_top_left.sh_mm, "x", None), + "top_left_y_mm": getattr(grid.smargon_top_left.sh_mm, "y", None), + "top_left_z_mm": getattr(grid.smargon_top_left.sh_mm, "z", None), + }, + ), ) res1 = self.__raster(grid) @@ -1107,40 +1091,37 @@ class AareDAQ: logger.info( "Prepared second auto-center raster", - extra={ - "file_prefix": grid.file_prefix, - "omega_deg": grid.omega_deg, - "top_left_x_mm": getattr(grid.smargon_top_left.sh_mm, "x", None), - "top_left_y_mm": getattr(grid.smargon_top_left.sh_mm, "y", None), - "top_left_z_mm": getattr(grid.smargon_top_left.sh_mm, "z", None), - "n_x": grid.n_x, - "n_y": grid.n_y, - "grid_size_x_mm": getattr(grid.grid_size_mm, "x", None), - "grid_size_y_mm": getattr(grid.grid_size_mm, "y", None), - }, + extra=merge_log_context( + sample_log_context(self.sample), + raster_request_log_context(grid), + { + "top_left_x_mm": getattr(grid.smargon_top_left.sh_mm, "x", None), + "top_left_y_mm": getattr(grid.smargon_top_left.sh_mm, "y", None), + "top_left_z_mm": getattr(grid.smargon_top_left.sh_mm, "z", None), + }, + ), ) logger.info( "Running second auto-center raster", - extra={ - "file_prefix": grid.file_prefix, - "omega_deg": grid.omega_deg, - "n_x": grid.n_x, - "n_y": grid.n_y, - }, + extra=merge_log_context( + sample_log_context(self.sample), + raster_request_log_context(grid), + ), ) res2 = self.__raster(grid) return CompletedRasterGrid(r=[res1, res2]) else: logger.error( "Auto-center raster aborted because no ML bounding box was found at either angle", - extra={ - "sample_id": getattr(sample, "db_id", None), - "sample_name": getattr(sample, "sample_name", None), - "primary_omega_deg": geom.omega_deg, - "secondary_omega_deg": geom.omega_deg + 90.0, - "file_prefix": old_prefix, - }, + extra=merge_log_context( + sample_log_context(sample), + { + "primary_omega_deg": geom.omega_deg, + "secondary_omega_deg": geom.omega_deg + 90.0, + "file_prefix": old_prefix, + }, + ), ) return None @@ -1284,19 +1265,11 @@ class AareDAQ: status = self.status logger.info( "Starting raster acquisition", - extra={ - "file_prefix": request.file_prefix, - "omega_deg": request.omega_deg, - "dtz": request.dtz, - "transmission": request.transmission, - "n_x": request.n_x, - "n_y": request.n_y, - "grid_size_x_mm": getattr(request.grid_size_mm, "x", None), - "grid_size_y_mm": getattr(request.grid_size_mm, "y", None), - "sample_id": getattr(self.sample, "db_id", None), - "sample_name": getattr(self.sample, "sample_name", None), - "state": getattr(status, "state", None), - }, + extra=merge_log_context( + sample_log_context(self.sample), + raster_request_log_context(request), + {"state": getattr(status, "state", None)}, + ), ) total_time = request.exp_time_s * request.n_x * request.n_y + request.n_y * 0.3 @@ -1340,18 +1313,20 @@ class AareDAQ: logger.info( "Calculated raster centre offset", - extra={ - "file_prefix": request.file_prefix, - "omega_deg": request.omega_deg, - "centre_offset_x_mm": grid_centre_offset.x, - "centre_offset_y_mm": grid_centre_offset.y, - "centre_offset_z_mm": grid_centre_offset.z, - "grid_half_width_x_mm": x, - "grid_half_height_y_mm": y, - "top_left_x_mm": getattr(request.smargon_top_left.sh_mm, "x", None), - "top_left_y_mm": getattr(request.smargon_top_left.sh_mm, "y", None), - "top_left_z_mm": getattr(request.smargon_top_left.sh_mm, "z", None), - }, + extra=merge_log_context( + sample_log_context(self.sample), + raster_request_log_context(request), + { + "centre_offset_x_mm": grid_centre_offset.x, + "centre_offset_y_mm": grid_centre_offset.y, + "centre_offset_z_mm": grid_centre_offset.z, + "grid_half_width_x_mm": x, + "grid_half_height_y_mm": y, + "top_left_x_mm": getattr(request.smargon_top_left.sh_mm, "x", None), + "top_left_y_mm": getattr(request.smargon_top_left.sh_mm, "y", None), + "top_left_z_mm": getattr(request.smargon_top_left.sh_mm, "z", None), + }, + ), ) logger.info(f"moving Smargon to grid centre offset {grid_centre_offset}") @@ -1366,15 +1341,17 @@ class AareDAQ: logger.info( "Moved Smargon to raster centre", - extra={ - "file_prefix": request.file_prefix, - "omega_deg": request.omega_deg, - "centre_sh_x_mm": grid_centre_smargon.sh_mm.x, - "centre_sh_y_mm": grid_centre_smargon.sh_mm.y, - "centre_sh_z_mm": grid_centre_smargon.sh_mm.z, - "centre_phi_deg": grid_centre_smargon.phi_deg, - "centre_chi_deg": grid_centre_smargon.chi_deg, - }, + extra=merge_log_context( + sample_log_context(self.sample), + raster_request_log_context(request), + { + "centre_sh_x_mm": grid_centre_smargon.sh_mm.x, + "centre_sh_y_mm": grid_centre_smargon.sh_mm.y, + "centre_sh_z_mm": grid_centre_smargon.sh_mm.z, + "centre_phi_deg": grid_centre_smargon.phi_deg, + "centre_chi_deg": grid_centre_smargon.chi_deg, + }, + ), ) if self.__cfg.simulated_detector: @@ -1389,15 +1366,11 @@ class AareDAQ: if scan_result is None: logger.error( "JFJoch returned no ScanResult for raster", - extra={ - "file_prefix": request.file_prefix, - "omega_deg": request.omega_deg, - "n_x": request.n_x, - "n_y": request.n_y, - "exp_time_s": request.exp_time_s, - "sample_id": getattr(self.sample, "db_id", None), - "sample_name": getattr(self.sample, "sample_name", None), - }, + extra=merge_log_context( + sample_log_context(self.sample), + raster_request_log_context(request), + {"exp_time_s": request.exp_time_s}, + ), ) sample_id = self.sample.db_id if self.sample and self.sample.db_id is not None else None @@ -1435,12 +1408,14 @@ class AareDAQ: logger.info( "Raster finished", - extra={ - "file_prefix": request.file_prefix, - "omega_deg": request.omega_deg, - "scan_result_is_none": scan_result is None, - "sample_id": sample_id, - }, + extra=merge_log_context( + sample_log_context(self.sample), + raster_request_log_context(request), + { + "scan_result_is_none": scan_result is None, + "sample_id": sample_id, + }, + ), ) return CompletedRasterGridElem( @@ -1452,32 +1427,35 @@ class AareDAQ: except Exception as e: logger.exception( "Failed during raster", - extra={ - "file_prefix": request.file_prefix, - "omega_deg": request.omega_deg, - "n_x": request.n_x, - "n_y": request.n_y, - "exp_time_s": request.exp_time_s, - "sample_id": getattr(self.sample, "db_id", None), - "sample_name": getattr(self.sample, "sample_name", None), - }, + extra=merge_log_context( + sample_log_context(self.sample), + raster_request_log_context(request), + {"exp_time_s": request.exp_time_s}, + ), ) raise - def measure_raster(self, r: RasterGridRequest, auto_center: bool) -> CompletedRasterGrid: + def measure_raster(self, request: RasterGridRequest, auto_center: bool) -> CompletedRasterGrid: """ Execute a raster scan. Args: - r: RasterGridRequest parameters for the scan. + request: RasterGridRequest parameters for the scan. auto_center: Boolean flag indicating if this is part of an automated sequence. Returns: CompletedRasterGrid result. """ + logger.info( + "Received raster scan request", + extra=merge_log_context( + sample_log_context(self.sample), + raster_request_log_context(request), + ), + ) self.__cfg.try_set_busy(timeout=ceil(360)) try: - result = self._execute_raster_sequence(r, auto_center=auto_center) + result = self._execute_raster_sequence(request, auto_center=auto_center) if result is None: raise RasterScanException("Raster scan failed") @@ -1567,19 +1545,10 @@ class AareDAQ: total_time = request.exp_time_s * request.steps logger.info( "Received rotation scan request", - extra={ - "file_prefix": request.file_prefix, - "start_omega_deg": request.start_omega_deg, - "dtz": request.dtz, - "exp_time_s": request.exp_time_s, - "incr_omega_deg": request.incr_omega_deg, - "steps": request.steps, - "transmission": request.transmission, - "screening": getattr(request, "screening", None), - "sample_id": getattr(self.sample, "db_id", None), - "sample_name": getattr(self.sample, "sample_name", None), - "total_time_s": total_time, - }, + extra=merge_log_context( + sample_log_context(self.sample), + rotation_request_log_context(request, total_time_s=total_time), + ), ) self.__cfg.try_set_busy(timeout=ceil(total_time + 360)) @@ -1696,7 +1665,8 @@ class AareDAQ: ) m = prediction_result.box bundle_image = prediction_result.image - self._log_ml_bundle_meta( + log_ml_bundle_meta( + logger, f"ml_bounding_box:{filename or 'unnamed'}", target_point=prediction_result.target_point, focus=prediction_result.focus, @@ -1724,19 +1694,21 @@ class AareDAQ: geom = self.sample_geometry logger.info( "ML bounding box selected", - extra={ - "sample_id": sample_id, - "ml_image_name": filename, - "omega_deg": geom.omega_deg, - "beam_x_pxl": geom.beam_location_pxl.x, - "beam_y_pxl": geom.beam_location_pxl.y, - "pixel_in_mm": geom.pixel_in_mm, - "box_x1": x1, - "box_y1": y1, - "box_x2": x2, - "box_y2": y2, - "target_point": prediction_result.target_point, - }, + extra=merge_log_context( + sample_log_context(self.sample), + { + "sample_id": sample_id, + "ml_image_name": filename, + }, + geom_log_context(geom), + { + "box_x1": x1, + "box_y1": y1, + "box_x2": x2, + "box_y2": y2, + "target_point": prediction_result.target_point, + }, + ), ) start_coord = geom.picture_to_smargon(Coordinate(x=x1, y=y1)) @@ -1746,19 +1718,22 @@ class AareDAQ: logger.info( "Converted ML bounding box to raster request", - extra={ - "sample_id": sample_id, - "ml_image_name": filename, - "start_sh_x_mm": start_coord.x, - "start_sh_y_mm": start_coord.y, - "start_sh_z_mm": start_coord.z, - "grid_size_x_mm": grid_size.x, - "grid_size_y_mm": grid_size.y, - "n_x": n_x, - "n_y": n_y, - "smargon_phi_deg": geom.smargon.phi_deg, - "smargon_chi_deg": geom.smargon.chi_deg, - }, + extra=merge_log_context( + sample_log_context(self.sample), + { + "sample_id": sample_id, + "ml_image_name": filename, + "start_sh_x_mm": start_coord.x, + "start_sh_y_mm": start_coord.y, + "start_sh_z_mm": start_coord.z, + "grid_size_x_mm": grid_size.x, + "grid_size_y_mm": grid_size.y, + "n_x": n_x, + "n_y": n_y, + "smargon_phi_deg": geom.smargon.phi_deg, + "smargon_chi_deg": geom.smargon.chi_deg, + }, + ), ) return RasterGridRequest( @@ -1789,7 +1764,8 @@ class AareDAQ: ) boxes = prediction_result.predictions bundle_image = prediction_result.image - self._log_ml_bundle_meta( + log_ml_bundle_meta( + logger, f"ml_loop_centre_box:{filename or 'unnamed'}", target_point=prediction_result.target_point, focus=prediction_result.focus, @@ -1927,6 +1903,7 @@ class AareDAQ: return #TODO operator function similar to mount, loop_center, raster and rotation + @log_timing(logger, "Face detection sequence") def __face_detection_sequence(self, steps: int = 14, step_size: int = 15, face_min_ratio: float = 0.3) -> dict: self.__set_state(BeamlineStateEnum.SampleAlignment) self.__devs.lamp_light = 2.5 @@ -1948,21 +1925,25 @@ class AareDAQ: logger.debug(f'moving to angle: {angle}') rotate_time = time.perf_counter() self.__devs.aerotech_omega = angle - logger.info(f"time to rotate 15 degrees: {time.perf_counter() - rotate_time}") + log_duration( + logger, + "Completed Aerotech move during face detection", + time.perf_counter() - rotate_time, + extra={"angle_deg": angle}, + ) - box_time = time.perf_counter() prediction_result: MLBoxPredictionResult = self.__mlbox.predict( preferred_class=(3, 0), return_image=True, return_bundle_meta=True ) m = prediction_result.box - self._log_ml_bundle_meta( + log_ml_bundle_meta( + logger, f"face_detection_angle_{angle}", target_point=prediction_result.target_point, focus=prediction_result.focus, ) - logger.info(f"time to predict: {time.perf_counter() - box_time}") if not m or not m.box: logger.info(f"no box found for angle {angle}") @@ -2116,12 +2097,12 @@ class AareDAQ: return_image=False, return_bundle_meta=True, ) - self._log_ml_bundle_meta( + log_ml_bundle_meta( + logger, f"auto_center_line_scan_{omega_deg:.2f}deg", target_point=prediction_result.target_point, focus=prediction_result.focus, ) - target_point = prediction_result.target_point prediction_box = prediction_result.box @@ -2138,39 +2119,48 @@ class AareDAQ: ) logger.info( "Using ML target y for second auto-center raster", - extra={ - "file_prefix": file_prefix, - "omega_deg": omega_deg, - "beam_y_pxl": beam_y_pxl, - "target_y_pxl": target_y_pxl, - "y_delta_mm": y_delta_mm, - "threshold_mm": y_retarget_threshold_mm, - "target_sh_x_mm": line_scan_centre.x, - "target_sh_y_mm": line_scan_centre.y, - "target_sh_z_mm": line_scan_centre.z, - }, + extra=merge_log_context( + sample_log_context(self.sample), + { + "file_prefix": file_prefix, + "omega_deg": omega_deg, + "beam_y_pxl": beam_y_pxl, + "target_y_pxl": target_y_pxl, + "y_delta_mm": y_delta_mm, + "threshold_mm": y_retarget_threshold_mm, + "target_sh_x_mm": line_scan_centre.x, + "target_sh_y_mm": line_scan_centre.y, + "target_sh_z_mm": line_scan_centre.z, + }, + ), ) else: logger.info( "Keeping beam-centred y line scan because ML target y shift is small", - extra={ - "file_prefix": file_prefix, - "omega_deg": omega_deg, - "beam_y_pxl": beam_y_pxl, - "target_y_pxl": target_y_pxl, - "y_delta_mm": y_delta_mm, - "threshold_mm": y_retarget_threshold_mm, - }, + extra=merge_log_context( + sample_log_context(self.sample), + { + "file_prefix": file_prefix, + "omega_deg": omega_deg, + "beam_y_pxl": beam_y_pxl, + "target_y_pxl": target_y_pxl, + "y_delta_mm": y_delta_mm, + "threshold_mm": y_retarget_threshold_mm, + }, + ), ) else: logger.info( "No ML target point for second auto-center raster; using beam-centred line scan", - extra={ - "file_prefix": file_prefix, - "omega_deg": omega_deg, - "beam_x_pxl": beam_x_pxl, - "beam_y_pxl": beam_y_pxl, - }, + extra=merge_log_context( + sample_log_context(self.sample), + { + "file_prefix": file_prefix, + "omega_deg": omega_deg, + "beam_x_pxl": beam_x_pxl, + "beam_y_pxl": beam_y_pxl, + }, + ), ) if prediction_box is not None and prediction_box.box is not None and grid_size_mm.y > 0: @@ -2179,26 +2169,32 @@ class AareDAQ: n_y = max(1, int(ceil(padded_height_mm / grid_size_mm.y))) logger.info( "Computed second auto-center raster y size from ML box height", - extra={ - "file_prefix": file_prefix, - "omega_deg": omega_deg, - "box_height_pxl": box_height_pxl, - "pixel_in_mm": geom.pixel_in_mm, - "grid_size_y_mm": grid_size_mm.y, - "padding_fraction_each_side": y_padding_fraction_each_side, - "padded_height_mm": padded_height_mm, - "computed_n_y": n_y, - }, + extra=merge_log_context( + sample_log_context(self.sample), + { + "file_prefix": file_prefix, + "omega_deg": omega_deg, + "box_height_pxl": box_height_pxl, + "pixel_in_mm": geom.pixel_in_mm, + "grid_size_y_mm": grid_size_mm.y, + "padding_fraction_each_side": y_padding_fraction_each_side, + "padded_height_mm": padded_height_mm, + "computed_n_y": n_y, + }, + ), ) else: logger.info( "Using default y size for second auto-center raster", - extra={ - "file_prefix": file_prefix, - "omega_deg": omega_deg, - "default_n_y": default_n_y, - "has_prediction_box": prediction_box is not None and prediction_box.box is not None, - }, + extra=merge_log_context( + sample_log_context(self.sample), + { + "file_prefix": file_prefix, + "omega_deg": omega_deg, + "default_n_y": default_n_y, + "has_prediction_box": prediction_box is not None and prediction_box.box is not None, + }, + ), ) offset = Coordinate( @@ -2213,21 +2209,25 @@ class AareDAQ: logger.info( "Prepared second auto-center raster top-left from helper", - extra={ - "file_prefix": file_prefix, - "omega_deg": omega_deg, - "n_y": n_y, - "offset_x_mm": offset.x, - "offset_y_mm": offset.y, - "offset_z_mm": offset.z, - "top_left_x_mm": getattr(top_left.sh_mm, "x", None), - "top_left_y_mm": getattr(top_left.sh_mm, "y", None), - "top_left_z_mm": getattr(top_left.sh_mm, "z", None), - }, + extra=merge_log_context( + sample_log_context(self.sample), + { + "file_prefix": file_prefix, + "omega_deg": omega_deg, + "n_y": n_y, + "offset_x_mm": offset.x, + "offset_y_mm": offset.y, + "offset_z_mm": offset.z, + "top_left_x_mm": getattr(top_left.sh_mm, "x", None), + "top_left_y_mm": getattr(top_left.sh_mm, "y", None), + "top_left_z_mm": getattr(top_left.sh_mm, "z", None), + }, + ), ) return top_left, n_y + @log_timing(logger, "Loop center sequence") def __loop_center_sequence( self, sample_id: Optional[int] = None, @@ -2237,7 +2237,6 @@ class AareDAQ: found_classes_count: dict[int, int] = {0: 0, 1: 0, 2: 0, 3: 0, 4: 0, 5: 0} try: - start = time.perf_counter() for zoom_iter, zoom_value in enumerate([200]): max_attempt = 2 attempt = 0 @@ -2258,7 +2257,12 @@ class AareDAQ: logger.debug(f"Moving to new omega: {angle}") time_to_move_aerotech = time.perf_counter() self.__devs.aerotech_omega = angle - logger.info(f"time to move AEROTECH: {time.perf_counter() - time_to_move_aerotech}") + log_duration( + logger, + "Completed Aerotech move during loop centering", + time.perf_counter() - time_to_move_aerotech, + extra={"angle_deg": angle, "zoom": zoom_value}, + ) filename = f"{sample_id}_{angle}_{zoom_value:.0f}" if sample_id is not None else None @@ -2270,7 +2274,12 @@ class AareDAQ: return_image=True, return_bundle_meta=True ) - logger.info(f"time to PREDICT: {time.perf_counter() - time_to_get_pred}") + log_duration( + logger, + "Completed ML prediction during loop centering", + time.perf_counter() - time_to_get_pred, + extra={"angle_deg": angle, "zoom": zoom_value}, + ) boxes = prediction_result.predictions pred_target_point = prediction_result.target_point @@ -2281,7 +2290,8 @@ class AareDAQ: )) bundle_image = prediction_result.image - self._log_ml_bundle_meta( + log_ml_bundle_meta( + logger, f"loop_center_angle_{angle}_zoom_{zoom_value:.0f}", target_point=prediction_result.target_point, focus=prediction_result.focus, @@ -2319,7 +2329,12 @@ class AareDAQ: time_to_move_smargon = time.perf_counter() self.__devs.smargon_pos = target self.__devs.smargon_wait(60) - logger.info(f"time to move smargon: {time.perf_counter() - time_to_move_smargon}") + log_duration( + logger, + "Completed Smargon move during loop centering", + time.perf_counter() - time_to_move_smargon, + extra={"angle_deg": angle, "zoom": zoom_value}, + ) if sample_id is not None: if trace_all_alc_moves: self._append_smargon_trace( @@ -2349,7 +2364,6 @@ class AareDAQ: f"sucessfully found {targets_found_this_attempt} targets in attempt {attempt + 1} ") break logger.error(f"{attempt} exceeds max attempts {max_attempt}") - logger.debug(f"time to loop center: {time.perf_counter() - start}") raise LoopCenteringFailed #i += 1 @@ -2358,7 +2372,6 @@ class AareDAQ: if sample_id is not None: logger.info(f"sample {sample_id} centered") self._append_smargon_trace(sample_id=sample_id, event="alc_success") - logger.debug(f"time to loop center: {time.perf_counter() - start}") return LoopCenteringResult(success=True) except Exception as e: @@ -2384,6 +2397,7 @@ class AareDAQ: error=e ) + @log_timing(logger, "Auto loop center") def auto_loop_center(self, sample:Optional[SampleShortInfo]=None) -> float: """ Automatically center the loop using ML-based detection. @@ -2597,6 +2611,7 @@ class AareDAQ: end = time.perf_counter() return end - start + @log_timing(logger, "Measure Sequence") def measure(self, sample: SampleShortInfo) -> float: """ Main measurement sequence for a given sample. @@ -2669,16 +2684,10 @@ class AareDAQ: if raster_result is None: logger.error( "Raster result was None during automation", - extra={ - "sample_id": getattr(sample, "db_id", None), - "sample_name": getattr(sample, "sample_name", None), - "file_prefix": raster_grid.file_prefix, - "omega_deg": raster_grid.omega_deg, - "dtz": raster_grid.dtz, - "transmission": raster_grid.transmission, - "n_x": raster_grid.n_x, - "n_y": raster_grid.n_y, - }, + extra=merge_log_context( + sample_log_context(sample), + raster_request_log_context(raster_grid), + ), ) self._mark_progress_failed(progress, WorkflowStateKind.RASTER, "Raster failed") return self._end_operation(start, DAQOperation.RASTER, error=True) diff --git a/src/aare/daq/mlbox.py b/src/aare/daq/mlbox.py index 9ba7f2b1..f70b67fb 100644 --- a/src/aare/daq/mlbox.py +++ b/src/aare/daq/mlbox.py @@ -12,6 +12,7 @@ from aare.common.aarelc_infer import AareLCInferWrapper from aare.common.beamline import MXBeamline from aare.common.models import MLBoxModel, MLOutputModel, MLBoxType, BoundingBoxModel from aare.common.logger_config import setup_logger +from aare.common.logger_events import log_timing logger=setup_logger("aareDAQ") @@ -162,6 +163,7 @@ class MlBox: bundle_meta = self._extract_bundle_meta(prediction) return MLBundle(predictions=predictions, image=image, bundle_meta=bundle_meta) + @log_timing(logger, "Collect best bundle") def _collect_best_bundle(self, attempts: int = RETRY_COUNT) -> MLBundle: best_predictions: MLOutputModel | None = None best_image: np.ndarray | None = None