python3-damaris/scripts/damaris-cli

#! /usr/bin/env python3

import time
import sys
import os
import os.path
import argparse
import logging
import tables
import damaris.data.DataPool as DataPool
import damaris.gui.ResultReader as ResultReader
import damaris.gui.ExperimentWriter as ExperimentWriter
import damaris.gui.BackendDriver as BackendDriver
import damaris.gui.ResultHandling as ResultHandling
import damaris.gui.ExperimentHandling as ExperimentHandling

logger = logging.getLogger("damaris.cli")

def progressbar(event):
    if event.subject=="__recentexperiment" or event.subject=="__recentresult":
        r=event.origin.get("__recentresult",-1)+1
        e=event.origin.get("__recentexperiment",-1)+1
        if e!=0:
            ratio=100.0*r/e
        else:
            ratio=100.0
        sys.stderr.write("\r%d/%d (%.0f%%)"%(r,e,ratio))
        sys.stderr.flush()


class ScriptInterface:

    def __init__(self, exp_script=None, res_script=None, backend_executable=None, spool_dir="spool", log_callback=None):
        self.exp_script=exp_script
        self.res_script=res_script
        self.backend_executable=backend_executable
        self.spool_dir=os.path.abspath(spool_dir)
        self.exp_handling=self.res_handling=None
        self._log_callback = log_callback

        self.exp_writer=self.res_reader=self.back_driver=None
        if self.backend_executable is not None:
            self.back_driver=BackendDriver.BackendDriver(self.backend_executable, spool_dir)
            if self.exp_script: self.exp_writer=self.back_driver.get_exp_writer()
            if self.res_script: self.res_reader=self.back_driver.get_res_reader()
        else:
            self.back_driver=None
            if self.exp_script: self.exp_writer=ExperimentWriter.ExperimentWriter(spool_dir)
            if self.res_script: self.res_reader=ResultReader.ResultReader(spool_dir)

        self.data=DataPool()


    def runScripts(self):
        # get script engines
        if self.exp_script and self.exp_writer:
            self.exp_handling=ExperimentHandling.ExperimentHandling(self.exp_script, self.exp_writer, self.data, self._log_callback)
        if self.res_script and self.res_reader:
            self.res_handling=ResultHandling.ResultHandling(self.res_script, self.res_reader, self.data, self._log_callback)

        # start them
        if self.exp_handling: self.exp_handling.start()
        if self.back_driver is not None: self.back_driver.start()
        if self.res_handling: self.res_handling.start()

    def waitForScriptsEnding(self):
        # time of last dump
        dump_interval=600
        next_dump_time=time.time()+dump_interval
        # keyboard interrupts are handled in extra cleanup loop
        try:
            while [_f for _f in [self.exp_handling,self.res_handling,self.back_driver] if _f]:
                time.sleep(0.1)
                if time.time()>next_dump_time:
                    self.dump_data("pool/data_pool.h5")
                    next_dump_time+=dump_interval

                if self.exp_handling is not None:
                    if not self.exp_handling.is_alive():
                        self.exp_handling.join()
                        if self.exp_handling.raised_exception:
                            logger.error(": experiment script failed at line %d (function %s): %s",
                                        self.exp_handling.location[0],
                                        self.exp_handling.location[1],
                                        self.exp_handling.raised_exception)
                        else:
                            logger.info(": experiment script finished")
                        self.exp_handling = None

                if self.res_handling is not None:
                    if not self.res_handling.is_alive():
                        self.res_handling.join()
                        if self.res_handling.raised_exception:
                            logger.error(": result script failed at line %d (function %s): %s",
                                        self.res_handling.location[0],
                                        self.res_handling.location[1],
                                        self.res_handling.raised_exception)
                        else:
                            logger.info(": result script finished")
                        self.res_handling = None

                if self.back_driver is not None:
                    if not self.back_driver.is_alive():
                        logger.info(": backend finished")
                        self.back_driver=None

        except KeyboardInterrupt:
            still_running=[_f for _f in [self.exp_handling,self.res_handling,self.back_driver] if _f]
            for r in still_running:
                r.quit_flag.set()

            for r in still_running:
                r.join()

    def dump_data(self, filename):
        try:
            # write data from pool
            dump_file=tables.open_file(filename,mode="w",title="DAMARIS experiment data")
            self.data.write_hdf5(dump_file, complib='zlib', complevel=6)
            # write scripts
            scriptgroup=dump_file.create_group("/","scripts","Used Scripts")
            dump_file.create_array(scriptgroup,"experiment_script", self.exp_script)
            dump_file.create_array(scriptgroup,"result_script", self.res_script)
            dump_file.create_array(scriptgroup,"backend_executable", self.backend_executable)
            dump_file.create_array(scriptgroup,"spool_directory", self.spool_dir)
            dump_file.flush()
            dump_file.close()
            dump_file=None
            # todo
        except Exception as e:
            logger.error("Dump failed: %s", e)
       


if __name__=="__main__":
    parser = argparse.ArgumentParser(
        description="Run DAMARIS experiment and result scripts."
    )
    parser.add_argument("experiment_script", help="Path to the experiment script file")
    parser.add_argument("result_script", help="Path to the result script file")
    parser.add_argument(
        "spool_dir", nargs="?", default=None,
        help="Spool directory (defaults to current working directory)"
    )
    parser.add_argument(
        "--backend", default="/usr/lib/damaris/backends/Mobilecore",
        help="Path to the backend executable"
    )
    parser.add_argument(
        "--output", default="data_pool.h5",
        help="Output HDF5 file path (default: data_pool.h5)"
    )
    args = parser.parse_args()

    spool_dir = args.spool_dir if args.spool_dir else os.getcwd()

    with open(args.experiment_script) as f:
        expscript = f.read()
    with open(args.result_script) as f:
        resscript = f.read()

    si = ScriptInterface(expscript, resscript, args.backend, spool_dir)

    si.data.register_listener(progressbar)

    si.runScripts()

    si.waitForScriptsEnding()

    si.dump_data(args.output)

    si = None