python3-damaris/src/damaris/gui/ExperimentHandling.py

import threading
import io
import traceback
import sys
import time
import logging
import ast
from damaris.experiments.Experiment import Quit
from damaris.experiments import Experiment

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

class StopExperiment(Exception):
    pass

class _SandboxStdout:
    """Redirects print() from sandboxed scripts to the GUI textview."""
    def __init__(self, callback):
        self._callback = callback
        self._buffer = ""

    def write(self, message):
        self._buffer += message
        if "\n" in self._buffer:
            lines = self._buffer.split("\n")
            self._buffer = lines[-1]
            for line in lines[:-1]:
                self._callback(line + "\n")

    def flush(self):
        if self._buffer:
            self._callback(self._buffer)
            self._buffer = ""

    def isatty(self):
        return False


class ExperimentHandling(threading.Thread):
    """
    runs the experiment script in sandbox
    """

    def __init__(self, script, exp_writer, data, log_callback=None):
        threading.Thread.__init__(self, name="experiment handler")
        self.script=script
        self.writer=exp_writer
        self.data=data
        self.quit_flag = threading.Event()
        self._log_callback = log_callback
        if self.data is not None:
            self.data["__recentexperiment"]=-1

    def synchronize(self, before=0, waitsteps=0.1):
        while (self.data["__recentexperiment"]>self.data["__recentresult"]+before) and not self.quit_flag.isSet():
            self.quit_flag.wait(waitsteps)
        if self.quit_flag.isSet():
            raise StopExperiment
        
    def sleep(self, seconds):
        self.quit_flag.wait(seconds)
        if self.quit_flag.isSet():
            raise StopExperiment

    def run(self):
        dataspace={}
        exp_classes = __import__('damaris.experiments', dataspace, dataspace, ['Experiment'])
        for name in dir(exp_classes):
            if name[:2]=="__" and name[-2:]=="__": continue
            dataspace[name]=exp_classes.__dict__[name]
        del exp_classes

        dataspace["data"]=self.data
        dataspace["synchronize"]=self.synchronize
        dataspace["sleep"]=self.sleep
        self.raised_exception = None
        self.location = None
        exp_iterator=None

        # Redirect stdout/stderr so print() in sandboxed scripts goes to the GUI
        _orig_stdout = sys.stdout
        _orig_stderr = sys.stderr
        if self._log_callback is not None:
            sys.stdout = _SandboxStdout(self._log_callback)
            sys.stderr = _SandboxStdout(self._log_callback)

        # check for time.sleep()
        try:
            tree = ast.parse(self.script)
            for node in ast.walk(tree):
                found_blocking_sleep = False
                if isinstance(node, ast.Call):
                    func = node.func
                    if isinstance(func, ast.Attribute) and func.attr == 'sleep':
                        if isinstance(func.value, ast.Name) and func.value.id == 'time':
                            found_blocking_sleep = True
                elif isinstance(node, ast.ImportFrom) and node.module == 'time':
                    for alias in node.names:
                        if alias.name == 'sleep':
                            found_blocking_sleep = True
                
                if found_blocking_sleep:
                    logger.warning("time.sleep() detected in experiment script. This is blocking and not interruptible. Please use sleep() instead.")
                    break
        except Exception:
            pass

        try:
            exec(self.script, dataspace)
        except StopExperiment:
            return
        except Exception as e:
            self.raised_exception=e
            self.location=traceback.extract_tb(sys.exc_info()[2])[-1][1:3]
            traceback_file=io.StringIO()
            traceback.print_tb(sys.exc_info()[2], None, traceback_file)
            self.traceback=traceback_file.getvalue()
            traceback_file=None
            return
        if "experiment" in dataspace:
            try:
                exp_iterator=dataspace["experiment"]()
            except StopExperiment:
                return
            except Exception as e:
                self.raised_exception=e
                self.location=traceback.extract_tb(sys.exc_info()[2])[-1][1:3]
                traceback_file=io.StringIO()
                traceback.print_tb(sys.exc_info()[2], None, traceback_file)
                self.traceback=traceback_file.getvalue()
                traceback_file=None
                return
        
        while exp_iterator is not None and not self.quit_flag.isSet():
            # get next experiment from script 
            try:
                job=next(exp_iterator)
            except StopExperiment as se:
                break
            except StopIteration as si:
                break
            except Exception as e:
                self.raised_exception=e
                self.location=traceback.extract_tb(sys.exc_info()[2])[-1][1:3]
                traceback_file=io.StringIO()
                traceback.print_tb(sys.exc_info()[2], None, traceback_file)
                self.traceback=traceback_file.getvalue()
                traceback_file=None
                break
            # send it
            self.writer.send_next(job)
            # write a note
            if isinstance(job, Experiment):
                if self.data is not None:
                    self.data["__recentexperiment"]=job.job_id+0
                    
                    # track time when experiment is started. This is placed after executing the script
                    # and getting the iterator so that time for initializing devices etc. is not included
                    if "__experimentstarted" not in self.data:
                        self.data["__experimentstarted"] = time.time()
                    # track time of experiments in queue, the total time and for each experiment:
                    if "__totalexperimentlength" not in self.data:
                        self.data["__totalexperimentlength"] = 0.0
                    if "__experimentlengths" not in self.data:
                        self.data["__experimentlengths"] = {}
                    
                    experimentlength = job.get_length()
                    self.data["__experimentlengths"][job.job_id+0] = experimentlength
                    self.data["__totalexperimentlength"] += experimentlength
                    
                    # relax for a short time
                    if "__resultsinadvance" in self.data and self.data["__resultsinadvance"]+100<job.job_id:
                        self.quit_flag.wait(0.05)
                if self.quit_flag.isSet():
                    dataspace=None
                    exp_iterator=None
                    break

        self.writer.send_next(Quit(), quit=True)
        # do not count quit job (is this a good idea?)
        dataspace=None
        self.exp_iterator=None
        self.writer=None

        # Restore stdout/stderr
        sys.stdout = _orig_stdout
        sys.stderr = _orig_stderr