From 55cec4bd03b24f9d00e81033d022f167a66ac8a3 Mon Sep 17 00:00:00 2001 From: Dmitry Ozerov Date: Mon, 26 Jul 2021 08:51:19 +0200 Subject: [PATCH] make rabbitmq multiple queues; use one for the detector retrieve --- sf_daq_broker/broker_manager.py | 65 +++-- sf_daq_broker/config.py | 3 +- sf_daq_broker/rabbitmq/config.py | 13 +- sf_daq_broker/rabbitmq/msg_broker_client.py | 17 +- sf_daq_broker/writer/detector_writer.py | 291 ++++++++++++++++++++ sf_daq_broker/writer/export_file.py | 112 ++++++++ sf_daq_broker/writer/postprocess_raw.py | 140 ++++++++++ sf_daq_broker/writer/start.py | 58 +++- 8 files changed, 645 insertions(+), 54 deletions(-) create mode 100644 sf_daq_broker/writer/detector_writer.py create mode 100644 sf_daq_broker/writer/export_file.py create mode 100644 sf_daq_broker/writer/postprocess_raw.py diff --git a/sf_daq_broker/broker_manager.py b/sf_daq_broker/broker_manager.py index 40ad75a..245c1d6 100644 --- a/sf_daq_broker/broker_manager.py +++ b/sf_daq_broker/broker_manager.py @@ -57,7 +57,7 @@ class BrokerManager(object): return {"status" : "failed", "message" : "request is made from beamline which doesnt have detectors"} if "start_pulseid" not in request: - return {"status" : "failed", "message" : "aaa no start pluseid provided in request parameters"} + return {"status" : "failed", "message" : "no start pluseid provided in request parameters"} rate_multiplicator = request.get("rate_multiplicator", 1) @@ -86,7 +86,7 @@ class BrokerManager(object): "timestamp": None } self.broker_client.open() - self.broker_client.send(pedestal_request) + self.broker_client.send(pedestal_request, broker_config.TAG_PEDESTAL) self.broker_client.close() time_to_wait = PEDESTAL_FRAMES/100*rate_multiplicator+10 @@ -260,7 +260,7 @@ class BrokerManager(object): run_log_file=run_log_file) try: - self.broker_client.send(write_request) + self.broker_client.send(write_request, tag) except: log_file = open(write_request["run_log_file"], "a") log_file.write("Can not contact writer") @@ -285,39 +285,38 @@ class BrokerManager(object): request.get("camera_list"), config.OUTPUT_FILE_SUFFIX_IMAGE_BUFFER) - self.broker_client.close() - if "detectors" in request: + request_detector = {} + + det_start_pulse_id = 0 + det_stop_pulse_id = stop_pulse_id + for p in range(start_pulse_id, stop_pulse_id+1): + if p%rate_multiplicator == 0: + det_stop_pulse_id = p + if det_start_pulse_id == 0: + det_start_pulse_id = p + request_detector["det_start_pulse_id"] = det_start_pulse_id + request_detector["det_stop_pulse_id"] = det_stop_pulse_id + + request_detector["path_to_pgroup"] = path_to_pgroup + request_detector["rate_multiplicator"] = rate_multiplicator + request_detector["run_file_json"] = run_file_json + request_detector["current_run"] = current_run + request_detector["run_info_directory"] = run_info_directory + request_detector["request_time"] = request["request_time"] + if "directory_name" in request: + request_detector["directory_name"] = request["directory_name"] + for detector in request["detectors"]: - output_file_detector = f'{full_path}/run_{current_run:06}.{detector}.h5' - output_files_list.append(output_file_detector) - det_start_pulse_id = 0 - det_stop_pulse_id = stop_pulse_id + request_detector_send = request_detector + request_detector_send["detector_name"] = detector + request_detector_send["detectors"] = {} + request_detector_send["detectors"][detector] = request["detectors"][detector] + send_write_request(broker_config.TAG_DETECTOR_RETRIEVE, + request_detector, + detector) - det_conversion = request["detectors"][detector].get("adc_to_energy", False) - det_compression = request["detectors"][detector].get("compression", False) - det_number_disabled_modules = len(request["detectors"][detector].get("disabled_modules", [])) - det_export = 0 - if det_conversion or det_compression or det_number_disabled_modules>0: - det_export = 1 - - raw_file_name = output_file_detector - if det_export == 1: - raw_file_name = f'{path_to_pgroup}/RAW_DATA/' - if "directory_name" in request and request["directory_name"] is not None: - raw_file_name = raw_file_name + request["directory_name"] - raw_file_name = f'{raw_file_name}/run_{current_run:06}.{detector}.h5' - - for p in range(start_pulse_id, stop_pulse_id+1): - if p%rate_multiplicator == 0: - det_stop_pulse_id = p - if det_start_pulse_id == 0: - det_start_pulse_id = p - retrieve_command=f'/home/dbe/git/sf_daq_buffer/scripts/retrieve_detector_data.sh {detector} {det_start_pulse_id} {det_stop_pulse_id} {output_file_detector} {rate_multiplicator} {det_export} {run_file_json} {raw_file_name}' - process_log_file=open(f'{run_info_directory}/run_{current_run:06}.{detector}.log','w') - _logger.info("Starting detector retrieve command %s " % retrieve_command) - process=Popen(retrieve_command, shell=True, stdout=process_log_file, stderr=process_log_file) - process_log_file.close() + self.broker_client.close() if "scan_info" in request: request_scan_info = request["scan_info"] diff --git a/sf_daq_broker/config.py b/sf_daq_broker/config.py index b80db91..f0fe75e 100644 --- a/sf_daq_broker/config.py +++ b/sf_daq_broker/config.py @@ -4,7 +4,8 @@ DEFAULT_BROKER_REST_PORT = 10002 DEFAULT_EPICS_WRITER_URL = "http://localhost:10200/notify" DEFAULT_LOG_LEVEL = "INFO" -DATA_RETRIEVAL_DELAY = 60 +BSDATA_RETRIEVAL_DELAY = 60 +DETECTOR_RETRIEVAL_DELAY = 10 AUDIT_FILE_TIME_FORMAT = "%Y%m%d-%H%M%S" diff --git a/sf_daq_broker/rabbitmq/config.py b/sf_daq_broker/rabbitmq/config.py index 3f181a1..3462c23 100644 --- a/sf_daq_broker/rabbitmq/config.py +++ b/sf_daq_broker/rabbitmq/config.py @@ -7,10 +7,21 @@ REQUEST_EXCHANGE = "request" STATUS_EXCHANGE = "status" DEFAULT_QUEUE = "write_request" +DEFAULT_ROUTE = "bs.*" -# Name of the queue for rabbitmq requests. +DETECTOR_RETRIEVE_QUEUE = "detector_retrieve" +DETECTOR_RETRIEVE_ROUTE = "detector.retrieve" + +DETECTOR_CONVERSION_QUEUE = "detector_convert" +DETECTOR_CONVERSION_ROUTE = "detector.convert" + +# Name of the tags for rabbitmq requests. TAG_DATABUFFER = "databuffer" TAG_IMAGEBUFFER = "imagebuffer" TAG_EPICS = "epics" TAG_DATA3BUFFER = "data3buffer" TAG_PEDESTAL = "pedestal" +TAG_EPICS = "epics" +TAG_DETECTOR_RETRIEVE = "detector_buffer" +TAG_DETECTOR_CONVERT = "detector_convert" + diff --git a/sf_daq_broker/rabbitmq/msg_broker_client.py b/sf_daq_broker/rabbitmq/msg_broker_client.py index bb46c32..283bf49 100644 --- a/sf_daq_broker/rabbitmq/msg_broker_client.py +++ b/sf_daq_broker/rabbitmq/msg_broker_client.py @@ -35,12 +35,21 @@ class RabbitMqClient(object): self.connection = None self.channel = None - def send(self, write_request): + def send(self, write_request, tag): if self.channel is None: raise RuntimeError("RabbitMqClient not connected.") - routing_key = "*" + routing_key = broker_config.DEFAULT_ROUTE + if ( tag == broker_config.TAG_DATABUFFER or + tag == broker_config.TAG_IMAGEBUFFER or + tag == broker_config.TAG_EPICS or + tag == broker_config.TAG_PEDESTAL ): + routing_key = broker_config.DEFAULT_ROUTE + elif tag == broker_config.TAG_DETECTOR_RETRIEVE: + routing_key = broker_config.DETECTOR_RETRIEVE_ROUTE + elif tag == broker_config.TAG_DETECTOR_CONVERT: + routing_key = broker_config.DETECTOR_CONVERSION_ROUTE body_bytes = json.dumps(write_request).encode() @@ -55,7 +64,7 @@ class RabbitMqClient(object): } self.channel.basic_publish(exchange=broker_config.STATUS_EXCHANGE, - properties=BasicProperties( - headers=status_header), + properties=BasicProperties(headers=status_header), routing_key=routing_key, body=body_bytes) + diff --git a/sf_daq_broker/writer/detector_writer.py b/sf_daq_broker/writer/detector_writer.py new file mode 100644 index 0000000..ccb49ec --- /dev/null +++ b/sf_daq_broker/writer/detector_writer.py @@ -0,0 +1,291 @@ +import logging +import os +from datetime import datetime +from time import time + +import numpy as np +import h5py +import subprocess + +from shutil import copyfile + +try: + import ujson as json +except: + _logger.warning("There is no ujson in this environment. Performance will suffer.") + import json + +from sf_daq_broker import config, utils + +from sf_daq_broker.writer.export_file import convert_file + +_logger = logging.getLogger("broker_writer") + +PEDESTAL_SPECIFIC = { "JF02T09V02" : {"number_bad_modules" : 1} } + +def detector_retrieve(request, output_file_detector): + + _logger.info(f'{request} {output_file_detector}') + + detector = request["detector_name"] + + det_start_pulse_id = request["det_start_pulse_id"] + det_stop_pulse_id = request["det_stop_pulse_id"] + rate_multiplicator = request["rate_multiplicator"] + run_file_json = request["run_file_json"] + path_to_pgroup = request["path_to_pgroup"] + current_run = request["current_run"] + run_info_directory = request["run_info_directory"] + + detector_config_file = f'/gpfs/photonics/swissfel/buffer/config/{detector}.json' + + det_conversion = request["detectors"][detector].get("adc_to_energy", False) + det_compression = request["detectors"][detector].get("compression", False) + det_number_disabled_modules = len(request["detectors"][detector].get("disabled_modules", [])) + + convert_ju_file = False + if det_conversion or det_compression or det_number_disabled_modules>0: + convert_ju_file = True + + raw_file_name = output_file_detector + if convert_ju_file: + raw_file_name = f'{path_to_pgroup}/RAW_DATA/' + if "directory_name" in request and request["directory_name"] is not None: + raw_file_name = raw_file_name + request["directory_name"] + if not os.path.isdir(raw_file_name): + os.makedirs(raw_file_name) + raw_file_name = f'{raw_file_name}/run_{current_run:06}.{detector}.h5' + + number_modules = int(detector[5:7]) + retrieve_command_from_buffer = f'/home/dbe/bin/sf_writer {raw_file_name} /gpfs/photonics/swissfel/buffer/{detector} {number_modules} {det_start_pulse_id} {det_stop_pulse_id} {rate_multiplicator}' + _logger.info("Starting detector retrieve from buffer %s " % retrieve_command_from_buffer) + process=subprocess.run(retrieve_command_from_buffer.split(), capture_output=True) +# _logger.info(process) + _logger.info("Finished retrieve from the buffer") + + if "directory_name" in request and request["directory_name"] == "JF_pedestals": + if detector in PEDESTAL_SPECIFIC: + create_pedestal_file(filename=raw_file_name, directory=os.path.dirname(raw_file_name), **PEDESTAL_SPECIFIC[detector]) + else: + create_pedestal_file(filename=raw_file_name, directory=os.path.dirname(raw_file_name)) + request_time = request["request_time"] + detector_config_file = f'/gpfs/photonics/swissfel/buffer/config/{detector}.json' + res_file_name = raw_file_name[:-3]+".res.h5" + copy_pedestal_file(request_time, res_file_name, detector, detector_config_file) + + if convert_ju_file: + _logger.info(f'Will do file conversion {raw_file_name} {output_file_detector} {run_file_json} {detector_config_file}') + convert_file(raw_file_name, output_file_detector, run_file_json, detector_config_file) + +def h5_printname(name): + print(" {}".format(name)) + +def forcedGainValue(i, n0, n1, n2, n3): + if i <= n0 - 1: + return 0 + if i <= (n0 + n1) - 1: + return 1 + if i <= (n0 + n1 + n2) - 1: + return 3 + if i <= (n0 + n1 + n2 + n3) - 1: + return 4 + return 2 + +def create_pedestal_file(filename="pedestal.h5", X_test_pixel=0, Y_test_pixel=0, nFramesPede=1000, + frames_G0=0, frames_G1=0, frames_G2=0, frames_HG0=0, number_frames=10000, frames_average=1000, + directory="./", gain_check=1, add_pixel_mask=None, number_bad_modules=0): + + if not (os.path.isfile(filename) and os.access(filename, os.R_OK)): + _logger.info("Pedestal file {} not found, exit".format(filename)) + return + + overwriteGain = False + if (frames_G0 + frames_G1 + frames_G2) > 0: + _logger.error("Treat this run as taken with {} frames in gain0, then {} frames in gain1 and {} frames in gain2".format(frames_G0, frames_G1, frames_G2)) + overwriteGain = True + + f = h5py.File(filename, "r") + + detector_name = (f.get("general/detector_name")[()]).decode('UTF-8') + n_bad_modules = number_bad_modules + + data_location = "data/" + detector_name + "/data" + daq_recs_location = "data/" + detector_name + "/daq_rec" + is_good_frame_location = "data/" + detector_name + "/is_good_frame" + + + numberOfFrames = len(f[data_location]) + (sh_y, sh_x) = f[data_location][0].shape + nModules = (sh_x * sh_y) // (1024 * 512) + if (nModules * 1024 * 512) != (sh_x * sh_y): + _logger.error(" {} : Something very strange in the data, Jungfrau consists of (1024x512) modules, while data has {}x{}".format(detector_name, sh_x, sh_y)) + return + + (tX, tY) = (X_test_pixel, Y_test_pixel) + if tX < 0 or tX > (sh_x - 1): + tX = 0 + if tY < 0 or tY > (sh_y - 1): + tY = 0 + + _logger.debug(" {} : test pixel is ( x y ): {}x{}".format(detector_name, tX, tY)) + _logger.info(" {} : In pedestal file {} there are {} frames".format(detector_name, filename, numberOfFrames + 1)) + _logger.debug(" {} : data has the following shape: {}, type: {}, {} modules ({} bad modules)".format(detector_name, f[data_location][0].shape, f[data_location][0].dtype, nModules, n_bad_modules)) + + pixelMask = np.zeros((sh_y, sh_x), dtype=int) + + adcValuesN = np.zeros((5, sh_y, sh_x)) + adcValuesNN = np.zeros((5, sh_y, sh_x)) + + + averagePedestalFrames = frames_average + + nMgain = [0] * 5 + + gainCheck = -1 + highG0Check = 0 + printFalseGain = False + nGoodFrames = 0 + nGoodFramesGain = 0 + + analyzeFrames = min(numberOfFrames, number_frames) + + for n in range(analyzeFrames): + + if not f[is_good_frame_location][n]: + continue + + nGoodFrames += 1 + + daq_rec = (f[daq_recs_location][n])[0] + + image = f[data_location][n][:] + frameData = (np.bitwise_and(image, 0b0011111111111111)) + gainData = np.bitwise_and(image, 0b1100000000000000) >> 14 + trueGain = forcedGainValue(n, framesG0, framesG1, framesG2, framesHG0) if overwriteGain else ( (daq_rec & 0b11000000000000) >> 12 ) + highG0 = (daq_rec & 0b1) + + gainGoodAllModules = True + if gain_check > 0: + daq_recs = f[daq_recs_location][n] + for i in range(len(daq_recs)): + if trueGain != ((daq_recs[i] & 0b11000000000000) >> 12) or highG0 != (daq_recs[i] & 0b1): + gainGoodAllModules = False + + if highG0 == 1 and trueGain != 0: + gainGoodAllModules = False + _logger.info(" {} : Jungfrau is in the high G0 mode ({}), but gain settings is strange: {}".format( detector_name, highG0, trueGain)) + + nFramesGain = np.sum(gainData==(trueGain)) + if nFramesGain < (nModules - 0.5 - n_bad_modules) * (1024 * 512): # make sure that most are the modules are in correct gain + gainGoodAllModules = False + _logger.debug(" {} : Too many bad pixels, skip the frame {}, true gain: {}(highG0: {}) ({}); gain0 : {}; gain1 : {}; gain2 : {}; undefined gain : {}".format( detector_name, n, trueGain, highG0, nFramesGain, np.sum(gainData==0), np.sum(gainData==1), np.sum(gainData==3), np.sum(gainData==2))) + + if not gainGoodAllModules: + _logger.debug(" {} : In Frame Number {} : mismatch in modules and general settings, Gain: {} vs {}; HighG0: {} vs {} (or too many bad pixels)".format( detector_name, n, trueGain, ((daq_recs & 0b11000000000000) >> 12), highG0, (daq_recs & 0b1))) + continue + nGoodFramesGain += 1 + + if gainData[tY][tX] != trueGain: + if not printFalseGain: + _logger.info(" {} : Gain wrong for channel ({}x{}) should be {}, but {}. Frame {}. {} {}".format( detector_name, tX, tY, trueGain, gainData[tY][tX], n, trueGain, daq_rec)) + printFalseGain = True + else: + if gainCheck != -1 and printFalseGain: + _logger.info(" {} : Gain was wrong for channel ({}x{}) in previous frames, but now correct : {}. Frame {}.".format( detector_name, tX, tY, gainData[tY, tX], n)) + printFalseGain = False + + if gainData[tY][tX] != gainCheck or highG0Check != highG0: + _logger.info(" {} : Gain changed for ({}x{}) channel {} -> {} (highG0 setting: {} -> {}), frame number {}, match: {}".format( detector_name, tX, tY, gainCheck, gainData[tY][tX], highG0Check, highG0, n, gainData[tY][tX] == trueGain)) + gainCheck = gainData[tY][tX] + highG0Check = highG0 + + if gainGoodAllModules: + + pixelMask[gainData != trueGain] |= (1 << (trueGain+4*highG0)) + + trueGain += 4 * highG0 + + + nMgain[trueGain] += 1 + + if nMgain[trueGain] > averagePedestalFrames: + adcValuesN[trueGain] -= adcValuesN[trueGain] / averagePedestalFrames + adcValuesNN[trueGain] -= adcValuesNN[trueGain] / averagePedestalFrames + + adcValuesN[trueGain] += frameData + adcValuesNN[trueGain] += np.float_power(frameData, 2) + + + _logger.info(" {} : {} frames analyzed, {} good frames, {} frames without settings mismatch. Gain frames distribution (0,1,2,3,HG0) : ({})".format( detector_name, analyzeFrames, nGoodFrames, nGoodFramesGain, nMgain)) + + if add_pixel_mask != None: + if (os.path.isfile(add_pixel_mask) and os.access(add_pixel_mask, os.R_OK)): + additional_pixel_mask_file = h5py.File(add_pixel_mask, "r") + additional_pixel_mask = np.array(additional_pixel_mask_file["pixel_mask"]) + _logger.info("Will add additional masked pixels from file %s , number %d " % (add_pixel_mask, np.sum(additional_pixel_mask == 1))) + if additional_pixel_mask.shape == pixelMask.shape: + pixelMask[additional_pixel_mask == 1] |= (1 << 5) + else: + _logger.error(" shape of additional pixel mask ({}) doesn't match current ({})".format( additional_pixel_mask.shape, pixelMask.shape)) + else: + _logger.error(" Specified addition file with pixel mask not found or not reachable {}".format( add_pixel_mask)) + + fileNameIn = os.path.splitext(os.path.basename(filename))[0] + full_fileNameOut = directory + "/" + fileNameIn + ".res.h5" + _logger.info(" {} : Output file with pedestal corrections in: {}".format( detector_name, full_fileNameOut)) + outFile = h5py.File(full_fileNameOut, "w") + + gains = [None] * 4 + gainsRMS = [None] * 4 + + for gain in range(5): + numberFramesAverage = max(1, min(averagePedestalFrames, nMgain[gain])) + mean = adcValuesN[gain] / float(numberFramesAverage) + mean2 = adcValuesNN[gain] / float(numberFramesAverage) + variance = mean2 - np.float_power(mean, 2) + stdDeviation = np.sqrt(variance) + _logger.debug(" {} : gain {} values results (pixel ({},{}) : {} {}".format( detector_name, gain, tY, tX, mean[tY][tX], stdDeviation[tY][tX])) + if gain != 2: + g = gain if gain < 3 else (gain-1) + gains[g] = mean + gainsRMS[g] = stdDeviation + + pixelMask[np.isclose(stdDeviation,0)] |= (1 << (6 + g)) + + dset = outFile.create_dataset('pixel_mask', data=pixelMask) + dset = outFile.create_dataset('gains', data=gains) + dset = outFile.create_dataset('gainsRMS', data=gainsRMS) + + outFile.close() + + _logger.info(" {} : Number of good pixels: {} from {} in total ({} bad pixels)".format( detector_name, np.sum(pixelMask == 0), sh_x * sh_y, (sh_x * sh_y - np.sum(pixelMask == 0)))) + + +def copy_pedestal_file(request_time, file_pedestal, detector, detector_config_file): + + PEDESTAL_DIRECTORY="/sf/jungfrau/data/pedestal" + + request_time=datetime.strptime(request_time, '%Y-%m-%d %H:%M:%S.%f') + + if not os.path.isdir(f'{PEDESTAL_DIRECTORY}/{detector}'): + os.mkdir(f'{PEDESTAL_DIRECTORY}/{detector}') + + out_name = f'{PEDESTAL_DIRECTORY}/{detector}/{request_time.strftime("%Y%m%d_%H%M%S")}.h5' + copyfile(file_pedestal, out_name) + + _logger.info(f'Copied resulting pedestal file {file_pedestal} to {out_name}') + + if not os.path.exists(detector_config_file): + _logger.error(f'stream file {detector_config_file} does not exists, exiting') + return + + with open(detector_config_file, "r") as stream_file: + det = json.load(stream_file) + + print(f'Changing in stream file {detector_config_file} pedestal from {det["pedestal_file"]} to {out_name}') + + det["pedestal_file"] = out_name + + with open(detector_config_file, "w") as write_file: + json.dump(det, write_file, indent=4) + diff --git a/sf_daq_broker/writer/export_file.py b/sf_daq_broker/writer/export_file.py new file mode 100644 index 0000000..7e43044 --- /dev/null +++ b/sf_daq_broker/writer/export_file.py @@ -0,0 +1,112 @@ +import json + +import h5py +import numpy as np + +import jungfrau_utils as ju + +import sys +from sf_daq_broker.writer.postprocess_raw import postprocess_raw + +import os + +import logging + +_logger = logging.getLogger("broker_writer") + +def convert_file(file_in, file_out, json_run_file, detector_config_file): + + with open(detector_config_file, "r") as detector_file: + data = json.load(detector_file) + + detector_name = data["detector_name"] + gain_file = data["gain_file"] + pedestal_file = data["pedestal_file"] + + with open(json_run_file, "r") as run_file: + data = json.load(run_file) + detector_params = data["detectors"][detector_name] + + compression = detector_params.get("compression", False) + conversion = detector_params.get("adc_to_energy", False) + disabled_modules = detector_params.get("disabled_modules", []) + if conversion: + mask = detector_params.get("mask", True) + mask_double_pixels = detector_params.get("mask_double_pixels", True) + geometry = detector_params.get("geometry", False) + gap_pixels = detector_params.get("gap_pixels", True) + factor = detector_params.get("factor", None) + else: + mask = False + mask_double_pixels = False + geometry = False + gap_pixels = False + factor = None + + if not mask and mask_double_pixels: + _logger.error("mask_double_pixels set to False") + mask_double_pixels = False + + file_tmp = file_in + if len(disabled_modules)>0: + _logger.info(f"Will reduce data file, disabled_modules: {disabled_modules}") + if conversion: + file_tmp = file_out+".tmp" + else: + file_tmp = file_out + postprocess_raw.postprocess_raw(file_in, file_tmp, compression=compression, disabled_modules=disabled_modules) + + if conversion: + + with ju.File( + file_tmp, + gain_file=gain_file, + pedestal_file=pedestal_file, + conversion=conversion, + mask=mask, + gap_pixels=gap_pixels, + geometry=geometry, + parallel=False, + ) as juf: + n_input_frames = len(juf["data"]) + good_frames = np.nonzero(juf["is_good_frame"])[0] + n_output_frames = len(good_frames) + + juf.handler.mask_double_pixels = mask_double_pixels + juf.export( + file_out, + index=good_frames, + roi=None, + compression=compression, + factor=factor, + dtype=None, + batch_size=35, + ) + #os.remove(file_tmp) + + else: + with h5py.File(file_tmp, "r") as juf: + n_input_frames = len(juf[f"data/{detector_name}/data"]) + good_frames = np.nonzero(juf[f"data/{detector_name}/is_good_frame"])[0] + n_output_frames = len(good_frames) + + # Utility info + with h5py.File(file_out, "r") as h5f: + _logger.info("daq_rec: %s" % h5f[f"/data/{detector_name}/daq_rec"][0, 0]) + + frame_index = h5f[f"/data/{detector_name}/frame_index"][:] + _logger.info("frame_index range: (%d - %d)" % (np.min(frame_index), np.max(frame_index))) + + _logger.info(f"input frames: {n_input_frames}") + _logger.info(f"bad frames: {n_input_frames - n_output_frames}") + _logger.info(f"output frames: {n_output_frames}") + + _logger.info(f"gain_file: {gain_file}") + _logger.info(f"pedestal_file: {pedestal_file}") + _logger.info(f"conversion: {conversion}") + _logger.info(f"mask: {mask}") + _logger.info(f"mask_double_pixels:x {mask_double_pixels}") + _logger.info(f"geometry: {geometry}") + _logger.info(f"gap_pixels: {gap_pixels}") + _logger.info(f"compression: {compression}") + _logger.info(f"factor: {factor}") diff --git a/sf_daq_broker/writer/postprocess_raw.py b/sf_daq_broker/writer/postprocess_raw.py new file mode 100644 index 0000000..3d3a285 --- /dev/null +++ b/sf_daq_broker/writer/postprocess_raw.py @@ -0,0 +1,140 @@ +import os +import struct + +import bitshuffle +import h5py +import numpy as np +from bitshuffle.h5 import H5_COMPRESS_LZ4, H5FILTER # pylint: disable=no-name-in-module + +# bitshuffle hdf5 filter params +BLOCK_SIZE = 2048 +compargs = {"compression": H5FILTER, "compression_opts": (BLOCK_SIZE, H5_COMPRESS_LZ4)} +# limit bitshuffle omp to a single thread +# a better fix would be to use bitshuffle compiled without omp support +os.environ["OMP_NUM_THREADS"] = "1" + +DTYPE = np.dtype(np.uint16) +DTYPE_SIZE = DTYPE.itemsize + +MODULE_SIZE_X = 1024 +MODULE_SIZE_Y = 512 + + +def postprocess_raw( + source, dest, disabled_modules=(), index=None, compression=False, batch_size=100 +): + # a function for 'visititems' should have the args (name, object) + def _visititems(name, obj): + if isinstance(obj, h5py.Group): + h5_dest.create_group(name) + + elif isinstance(obj, h5py.Dataset): + dset_source = h5_source[name] + + # process all but the raw data + if name != data_dset: + if name.startswith("data"): + # datasets with data per image, so indexing should be applied + if index is None: + data = dset_source[:] + else: + data = dset_source[index, :] + + args = {"shape": data.shape} + h5_dest.create_dataset_like(name, dset_source, data=data, **args) + else: + h5_dest.create_dataset_like(name, dset_source, data=dset_source) + + else: + raise TypeError(f"Unknown h5py object type {obj}") + + # copy group/dataset attributes if it's not a dataset with the actual data + if name != data_dset: + for key, value in h5_source[name].attrs.items(): + h5_dest[name].attrs[key] = value + + with h5py.File(source, "r") as h5_source, h5py.File(dest, "w") as h5_dest: + detector_name = h5_source["general/detector_name"][()].decode() + data_dset = f"data/{detector_name}/data" + + # traverse the source file and copy/index all datasets, except the raw data + h5_source.visititems(_visititems) + + # now process the raw data + dset = h5_source[data_dset] + + args = dict() + if index is None: + n_images = dset.shape[0] + else: + index = np.array(index) + n_images = len(index) + + n_modules = dset.shape[1] // MODULE_SIZE_Y + out_shape = (MODULE_SIZE_Y * (n_modules - len(disabled_modules)), MODULE_SIZE_X) + + args["shape"] = (n_images, *out_shape) + args["maxshape"] = (n_images, *out_shape) + args["chunks"] = (1, *out_shape) + + if compression: + args.update(compargs) + + h5_dest.create_dataset_like(data_dset, dset, **args) + + # calculate and save module_map + module_map = [] + tmp = 0 + for ind in range(n_modules): + if ind in disabled_modules: + module_map.append(-1) + else: + module_map.append(tmp) + tmp += 1 + + h5_dest[f"data/{detector_name}/module_map"] = np.tile(module_map, (n_images, 1)) + + # prepare buffers to be reused for every batch + read_buffer = np.empty((batch_size, *dset.shape[1:]), dtype=DTYPE) + out_buffer = np.zeros((batch_size, *out_shape), dtype=DTYPE) + + # process and write data in batches + for batch_start_ind in range(0, n_images, batch_size): + batch_range = range(batch_start_ind, min(batch_start_ind + batch_size, n_images)) + + if index is None: + batch_ind = np.array(batch_range) + else: + batch_ind = index[batch_range] + + # TODO: avoid unnecessary buffers + read_buffer_view = read_buffer[: len(batch_ind)] + out_buffer_view = out_buffer[: len(batch_ind)] + + # Avoid a stride-bottleneck, see https://github.com/h5py/h5py/issues/977 + if np.sum(np.diff(batch_ind)) == len(batch_ind) - 1: + # consecutive index values + dset.read_direct(read_buffer_view, source_sel=np.s_[batch_ind]) + else: + for i, j in enumerate(batch_ind): + dset.read_direct(read_buffer_view, source_sel=np.s_[j], dest_sel=np.s_[i]) + + for i, m in enumerate(module_map): + if m == -1: + continue + + read_slice = read_buffer_view[:, i * MODULE_SIZE_Y : (i + 1) * MODULE_SIZE_Y, :] + out_slice = out_buffer_view[:, m * MODULE_SIZE_Y : (m + 1) * MODULE_SIZE_Y, :] + out_slice[:] = read_slice + + bytes_num_elem = struct.pack(">q", out_shape[0] * out_shape[1] * DTYPE_SIZE) + bytes_block_size = struct.pack(">i", BLOCK_SIZE * DTYPE_SIZE) + header = bytes_num_elem + bytes_block_size + + for pos, im in zip(batch_range, out_buffer_view): + if compression: + byte_array = header + bitshuffle.compress_lz4(im, BLOCK_SIZE).tobytes() + else: + byte_array = im.tobytes() + + h5_dest[data_dset].id.write_direct_chunk((pos, 0, 0), byte_array) diff --git a/sf_daq_broker/writer/start.py b/sf_daq_broker/writer/start.py index 556a4d8..c73e3b8 100644 --- a/sf_daq_broker/writer/start.py +++ b/sf_daq_broker/writer/start.py @@ -14,6 +14,7 @@ from sf_daq_broker.utils import get_data_api_request from sf_daq_broker.writer.bsread_writer import write_from_imagebuffer, write_from_databuffer, write_from_databuffer_api3 from sf_daq_broker.writer.epics_writer import write_epics_pvs from sf_daq_broker.detector.pedestal import take_pedestal +from sf_daq_broker.writer.detector_writer import detector_retrieve _logger = logging.getLogger("broker_writer") @@ -36,14 +37,18 @@ def audit_failed_write_request(write_request): _logger.exception("Error while trying to write request %s to file %s." % (write_request, output_file)) -def wait_for_delay(request_timestamp): +def wait_for_delay(request_timestamp, writer_type): if request_timestamp is None: return + time_to_wait = config.BSDATA_RETRIEVAL_DELAY + if writer_type == broker_config.TAG_DETECTOR_RETRIEVE: + time_to_wait = config.DETECTOR_RETRIEVAL_DELAY + current_timestamp = time() # sleep time = target sleep time - time that has already passed. - adjusted_retrieval_delay = config.DATA_RETRIEVAL_DELAY - (current_timestamp - request_timestamp) + adjusted_retrieval_delay = time_to_wait - (current_timestamp - request_timestamp) if adjusted_retrieval_delay < 0: adjusted_retrieval_delay = 0 @@ -75,6 +80,8 @@ def process_request(request): file_handler.setLevel(logging.INFO) _logger.addHandler(file_handler) + logger_data_api = None + if writer_type == broker_config.TAG_DATABUFFER: logger_data_api = logging.getLogger("data_api") elif writer_type == broker_config.TAG_DATA3BUFFER: @@ -84,7 +91,8 @@ def process_request(request): elif writer_type == broker_config.TAG_EPICS: logger_data_api = logging.getLogger("data_api") - logger_data_api.addHandler(file_handler) + if logger_data_api is not None: + logger_data_api.addHandler(file_handler) try: _logger.info("Request for %s to write %s from pulse_id %s to %s" % @@ -98,7 +106,7 @@ def process_request(request): _logger.info("No channels requested. Skipping request.") return - wait_for_delay(request_timestamp) + wait_for_delay(request_timestamp, writer_type) _logger.info("Starting download.") @@ -124,11 +132,19 @@ def process_request(request): _logger.info("Doing pedestal.") take_pedestal(detectors_name=request.get("detectors", []), rate=request.get("rate_multiplicator", 1)) + elif writer_type == broker_config.TAG_DETECTOR_RETRIEVE: + _logger.info("Using detector retrieve writer.") + detector_retrieve(channels, output_file) + + elif writer_type == broker_config.TAG_DETECTOR_CONVERT: + _logger.info("Using detector convert writer.") + _logger.info("Finished. Took %s seconds to complete request." % (time() - start_time)) if file_handler: _logger.removeHandler(file_handler) - logger_data_api.removeHandler(file_handler) + if logger_data_api is not None: + logger_data_api.removeHandler(file_handler) except Exception: audit_failed_write_request(request) @@ -137,7 +153,8 @@ def process_request(request): finally: if file_handler: _logger.removeHandler(file_handler) - logger_data_api.removeHandler(file_handler) + if logger_data_api is not None: + logger_data_api.removeHandler(file_handler) def update_status(channel, body, action, file, message=None): @@ -203,7 +220,7 @@ def on_broker_message(channel, method_frame, header_frame, body, connection): reject_request(channel, method_frame, body, output_file, e) -def start_service(broker_url): +def start_service(broker_url, writer_type=0): connection = BlockingConnection(ConnectionParameters(broker_url)) channel = connection.channel() @@ -213,15 +230,24 @@ def start_service(broker_url): channel.exchange_declare(exchange=broker_config.REQUEST_EXCHANGE, exchange_type="topic") - channel.queue_declare(queue=broker_config.DEFAULT_QUEUE, auto_delete=True) - channel.queue_bind(queue=broker_config.DEFAULT_QUEUE, + routing_key = broker_config.DEFAULT_ROUTE + request_queue = broker_config.DEFAULT_QUEUE + if writer_type == 1: + routing_key = broker_config.DETECTOR_RETRIEVE_ROUTE + request_queue = broker_config.DETECTOR_RETRIEVE_QUEUE + if writer_type == 2: + routing_key = broker_config.DETECTOR_CONVERSION_ROUTE + request_queue = broker_config.DETECTOR_CONVERSION_QUEUE + + channel.queue_declare(queue=request_queue, auto_delete=True) + channel.queue_bind(queue=request_queue, exchange=broker_config.REQUEST_EXCHANGE, - routing_key="*") + routing_key=routing_key) channel.basic_qos(prefetch_count=1) on_broker_message_f = partial(on_broker_message, connection=connection) - channel.basic_consume(broker_config.DEFAULT_QUEUE, on_broker_message_f) + channel.basic_consume(request_queue, on_broker_message_f) try: channel.start_consuming() @@ -230,7 +256,7 @@ def start_service(broker_url): def run(): - parser = argparse.ArgumentParser(description='Bsread data writer') + parser = argparse.ArgumentParser(description='data writer') parser.add_argument("--broker_url", default=broker_config.DEFAULT_BROKER_URL, help="Address of the broker to connect to.") @@ -239,6 +265,8 @@ def run(): help="Log level to use.") parser.add_argument("--writer_id", default=1, type=int, help="Id of the writer") + parser.add_argument("--writer_type", default=0, type=int, + help="Type of the writer (0-epics/bs/camera/pedestal; 1 - detector retrieve; 2 - detector conversion)") args = parser.parse_args() @@ -247,7 +275,7 @@ def run(): stream_handler = logging.StreamHandler() stream_handler.setLevel(args.log_level) - formatter = logging.Formatter(f'[%(levelname)s] (broker_writer_{args.writer_id}) %(message)s') + formatter = logging.Formatter(f'[%(levelname)s] (broker_writer_{args.writer_id}_{args.writer_type}) %(message)s') stream_handler.setFormatter(formatter) _logger.setLevel(args.log_level) @@ -261,9 +289,9 @@ def run(): # make message-broker less noisy in logs logging.getLogger("pika").setLevel(logging.WARNING) - _logger.info("Writer(%s) started. Waiting for requests." % args.writer_id) + _logger.info("Writer started. Waiting for requests.") - start_service(broker_url=args.broker_url) + start_service(broker_url=args.broker_url, writer_type=args.writer_type) if __name__ == "__main__":