QuLab 2.10.10__cp313-cp313-win_amd64.whl

qulab/executor/schedule.py

@@ -0,0 +1,543 @@
+import functools
+import inspect
+import pickle
+import uuid
+from datetime import datetime, timedelta
+from pathlib import Path
+
+from loguru import logger
+
+from .load import WorkflowType, get_dependents
+from .registry import current_config, obey_the_oracle, update_parameters
+from .storage import (Report, find_report, get_head, get_heads, renew_report,
+                      revoke_report, save_item, save_report)
+
+__session_id = None
+__session_cache = {}
+
+
+def set_cache(session_id, key, report: Report):
+    global __session_id
+    if __session_id is None:
+        __session_id = session_id
+    if __session_id != session_id:
+        __session_cache.clear()
+    if report.workflow.startswith('cfg:'):
+        __session_cache[key] = report
+    else:
+        __session_cache[key] = report.base_path, report.path
+
+
+def get_cache(session_id, key) -> Report:
+    from .storage import load_report
+    global __session_id
+    if __session_id is None or __session_id != session_id:
+        return None
+    index = __session_cache.get(key, None)
+    if index is None:
+        return None
+    if isinstance(index, tuple):
+        base_path, path = index
+        return load_report(path, base_path)
+    elif isinstance(index, Report):
+        return index
+    else:
+        return None
+
+
+class CalibrationFailedError(Exception):
+    pass
+
+
+def is_pickleable(obj) -> bool:
+    try:
+        pickle.dumps(obj)
+        return True
+    except:
+        return False
+
+
+def veryfy_analyzed_report(report: Report, script: str, method: str):
+    if not isinstance(report, Report):
+        raise TypeError(f'"{script}" : "{method}" must return a Report object')
+    if not is_pickleable(report.parameters):
+        raise TypeError(
+            f'"{script}" : "{method}" return not pickleable data in .parameters'
+        )
+    if not is_pickleable(report.other_infomation):
+        raise TypeError(
+            f'"{script}" : "{method}" return not pickleable data in .other_infomation'
+        )
+
+
+def check_state(workflow: WorkflowType, code_path: str | Path,
+                state_path: str | Path, veryfy_source_code: bool) -> bool:
+    """
+    check state should report a pass if and only if the following are satisfied:
+    
+    1. The cal has had check data or calibrate pass within the timeout period.
+    2. The cal has not failed calibrate without resolution.
+    3. No dependencies have been recalibrated since the last time check data or calibrate was run on this cal.
+    4. All dependencies pass check state.
+    """
+    logger.debug(f'check_state: "{workflow.__workflow_id__}"')
+    report = find_report(workflow.__workflow_id__, state_path)
+    if not report:
+        logger.debug(
+            f'check_state failed: No history found for "{workflow.__workflow_id__}"'
+        )
+        return False
+    if hasattr(workflow, 'check_state') and callable(workflow.check_state):
+        logger.debug(
+            f'check_state: "{workflow.__workflow_id__}" has custom check_state method'
+        )
+        return workflow.check_state(report)
+    if datetime.fromtimestamp(workflow.__mtime__) > report.checked_time:
+        logger.debug(
+            f'check_state failed: "{workflow.__workflow_id__}" has been modified after last calibration'
+        )
+        return False
+    if workflow.__timeout__ is not None and datetime.now(
+    ) > report.checked_time + timedelta(seconds=workflow.__timeout__):
+        logger.debug(
+            f'check_state failed: "{workflow.__workflow_id__}" has expired')
+        return False
+    if not report.in_spec:
+        logger.debug(
+            f'check_state failed: "{workflow.__workflow_id__}" is out of spec')
+        return False
+    if report.bad_data:
+        logger.debug(
+            f'check_state failed: "{workflow.__workflow_id__}" has bad data')
+        return False
+    for n in get_dependents(workflow, code_path, veryfy_source_code):
+        r = find_report(n.__workflow_id__, state_path)
+        if r is None or r.checked_time > report.checked_time:
+            logger.debug(
+                f'check_state failed: "{workflow.__workflow_id__}" has outdated dependencies'
+            )
+            return False
+    for n in get_dependents(workflow, code_path, veryfy_source_code):
+        if not check_state(n, code_path, state_path, veryfy_source_code):
+            logger.debug(
+                f'check_state failed: "{workflow.__workflow_id__}" has bad dependencies'
+            )
+            return False
+    return True
+
+
+@logger.catch(reraise=False)
+async def call_plot(node: WorkflowType, report: Report, check=False):
+    if hasattr(node, 'plot') and callable(node.plot):
+        if inspect.iscoroutinefunction(node.plot):
+            await node.plot(report, check=check)
+        else:
+            node.plot(report)
+
+
+async def call_check(workflow: WorkflowType, session_id: str,
+                     state_path: Path):
+    report = get_cache(session_id, (workflow.__workflow_id__, 'check'))
+    if report is not None:
+        logger.debug(f'Cache hit for "{workflow.__workflow_id__}:check"')
+        return report
+
+    if inspect.iscoroutinefunction(workflow.check):
+        data = await workflow.check()
+    else:
+        data = workflow.check()
+    if not is_pickleable(data):
+        raise TypeError(
+            f'"{workflow.__workflow_id__}" : "check" return not pickleable data'
+        )
+    report = Report(workflow=workflow.__workflow_id__,
+                    data=data,
+                    config_path=current_config(state_path),
+                    base_path=state_path,
+                    heads=get_heads(state_path),
+                    previous_path=get_head(workflow.__workflow_id__,
+                                           state_path),
+                    script_path=save_item(workflow.__source__, 'items',
+                                          state_path))
+
+    save_report(workflow.__workflow_id__, report, state_path)
+
+    set_cache(session_id, (workflow.__workflow_id__, 'check'), report)
+    return report
+
+
+async def call_calibrate(workflow: WorkflowType, session_id: str,
+                         state_path: Path):
+    report = get_cache(session_id, (workflow.__workflow_id__, 'calibrate'))
+    if report is not None:
+        logger.debug(f'Cache hit for "{workflow.__workflow_id__}:calibrate"')
+        return report
+
+    if inspect.iscoroutinefunction(workflow.calibrate):
+        data = await workflow.calibrate()
+    else:
+        data = workflow.calibrate()
+    if not is_pickleable(data):
+        raise TypeError(
+            f'"{workflow.__workflow_id__}" : "calibrate" return not pickleable data'
+        )
+    report = Report(workflow=workflow.__workflow_id__,
+                    data=data,
+                    config_path=current_config(state_path),
+                    base_path=state_path,
+                    heads=get_heads(state_path),
+                    previous_path=get_head(workflow.__workflow_id__,
+                                           state_path),
+                    script_path=save_item(workflow.__source__, 'items',
+                                          state_path))
+
+    save_report(workflow.__workflow_id__, report, state_path)
+
+    set_cache(session_id, (workflow.__workflow_id__, 'calibrate'), report)
+    return report
+
+
+async def call_check_analyzer(node: WorkflowType,
+                              report: Report,
+                              history: Report | None,
+                              state_path: Path,
+                              plot=False) -> Report:
+    if inspect.iscoroutinefunction(node.check_analyze):
+        report = await node.check_analyze(report, history=history)
+    else:
+        report = node.check_analyze(report, history=history)
+    veryfy_analyzed_report(report, node.__workflow_id__, "check_analyze")
+    report.fully_calibrated = False
+    if report.in_spec:
+        logger.debug(
+            f'"{node.__workflow_id__}": checked in spec, renewing report')
+        if report.previous is not None:
+            renew_report(node.__workflow_id__, report.previous, state_path)
+        else:
+            renew_report(node.__workflow_id__, report, state_path)
+    else:
+        logger.debug(
+            f'"{node.__workflow_id__}": checked out of spec, revoking report')
+        if report.previous is not None:
+            revoke_report(node.__workflow_id__, report.previous, state_path)
+        else:
+            revoke_report(node.__workflow_id__, report, state_path)
+    return report
+
+
+async def call_analyzer(node: WorkflowType,
+                        report: Report,
+                        history: Report | None,
+                        state_path: Path,
+                        plot=False) -> Report:
+    if inspect.iscoroutinefunction(node.analyze):
+        report = await node.analyze(report, history=history)
+    else:
+        report = node.analyze(report, history=history)
+    veryfy_analyzed_report(report, node.__workflow_id__, "analyze")
+    if hasattr(node, 'oracle') and callable(node.oracle):
+        logger.debug(
+            f'"{node.__workflow_id__}" has oracle method, calling ...')
+        report = await call_oracle(node, report, history)
+    report.fully_calibrated = True
+    save_report(node.__workflow_id__, report, state_path, overwrite=True)
+    if plot:
+        await call_plot(node, report)
+    return report
+
+
+async def call_oracle(node: WorkflowType, report: Report,
+                      history: Report | None):
+    sig = inspect.signature(node.oracle)
+    try:
+        if 'history' in sig.parameters and 'system_state' in sig.parameters:
+            report = node.oracle(report,
+                                 history=history,
+                                 system_state=get_heads(report.base_path))
+        elif 'history' in sig.parameters:
+            report = node.oracle(report, history=history)
+        elif 'system_state' in sig.parameters:
+            report = node.oracle(report,
+                                 system_state=get_heads(report.base_path))
+        else:
+            report = node.oracle(report)
+        if inspect.isawaitable(report):
+            report = await report
+    except Exception as e:
+        logger.exception(e)
+        report.oracle = {}
+        return report
+    if not isinstance(report, Report):
+        raise TypeError(
+            f'"{node.__workflow_id__}" : function "oracle" must return a Report object'
+        )
+    if not is_pickleable(report.oracle):
+        raise TypeError(
+            f'"{node.__workflow_id__}" : function "oracle" return not pickleable data'
+        )
+    return report
+
+
+async def check_data(workflow: WorkflowType, state_path: str | Path,
+                     plot: bool, session_id: str) -> Report:
+    """
+    check data answers two questions:
+    Is the parameter associated with this cal in spec,
+    and is the cal scan working as expected?
+    """
+    history = find_report(workflow.__workflow_id__, state_path)
+
+    if history is None:
+        logger.debug(f'No history found for "{workflow.__workflow_id__}"')
+        report = Report(workflow=workflow.__workflow_id__,
+                        config_path=current_config(state_path),
+                        base_path=state_path,
+                        heads=get_heads(state_path),
+                        previous_path=get_head(workflow.__workflow_id__,
+                                               state_path),
+                        script_path=save_item(workflow.__source__, 'items',
+                                              state_path))
+        report.in_spec = False
+        report.bad_data = False
+        return report
+
+    if history.bad_data:
+        logger.debug(
+            f'History found for "{workflow.__workflow_id__}", but bad data')
+        return history
+    if not history.in_spec:
+        logger.debug(
+            f'History found for "{workflow.__workflow_id__}", but out of spec')
+        return history
+
+    logger.debug(
+        f'History found for "{workflow.__workflow_id__}", but has expired')
+
+    if hasattr(workflow, 'check') and callable(workflow.check) and hasattr(
+            workflow, 'check_analyze') and callable(workflow.check_analyze):
+        logger.debug(
+            f'Checking "{workflow.__workflow_id__}" with "check" method ...')
+
+        report = await call_check(workflow, session_id, state_path)
+
+        logger.debug(f'Checked "{workflow.__workflow_id__}" !')
+        report = await call_check_analyzer(workflow,
+                                           report,
+                                           history,
+                                           state_path,
+                                           plot=plot)
+    else:
+        logger.debug(
+            f'Checking "{workflow.__workflow_id__}" with "calibrate" method ...'
+        )
+
+        report = await call_calibrate(workflow, session_id, state_path)
+
+        logger.debug(f'Calibrated "{workflow.__workflow_id__}" !')
+        report = await call_analyzer(workflow,
+                                     report,
+                                     history,
+                                     state_path,
+                                     plot=plot)
+    return report
+
+
+async def calibrate(workflow: WorkflowType, state_path: str | Path, plot: bool,
+                    session_id: str) -> Report:
+    history = find_report(workflow.__workflow_id__, state_path)
+
+    logger.debug(f'Calibrating "{workflow.__workflow_id__}" ...')
+
+    report = await call_calibrate(workflow, session_id, state_path)
+
+    logger.debug(f'Calibrated "{workflow.__workflow_id__}" !')
+
+    report = await call_analyzer(workflow,
+                                 report,
+                                 history,
+                                 state_path,
+                                 plot=plot)
+    return report
+
+
+async def diagnose(workflow: WorkflowType, code_path: str | Path,
+                   state_path: str | Path, plot: bool, session_id: str,
+                   fail_fast: bool, veryfy_source_code: bool):
+    '''
+    Returns: True if node or dependent recalibrated.
+    '''
+    logger.debug(f'diagnose "{workflow.__workflow_id__}"')
+    # check_data
+    report = await check_data(workflow, state_path, plot, session_id)
+    # in spec case
+    if report.in_spec:
+        logger.debug(
+            f'"{workflow.__workflow_id__}": Checked! In spec, no need to diagnose'
+        )
+        return False
+    # bad data case
+    recalibrated = []
+    if report.bad_data:
+        logger.debug(
+            f'"{workflow.__workflow_id__}": Bad data, diagnosing dependents')
+        recalibrated = []
+        exceptions = []
+        for n in get_dependents(workflow, code_path, veryfy_source_code):
+            try:
+                flag = await diagnose(n, code_path, state_path, plot,
+                                      session_id, fail_fast,
+                                      veryfy_source_code)
+            except Exception as e:
+                if fail_fast:
+                    raise e
+                exceptions.append(e)
+            recalibrated.append(flag)
+        if any(exceptions):
+            raise exceptions[0]
+    if not any(recalibrated):
+        if report.bad_data:
+            raise CalibrationFailedError(
+                f'"{workflow.__workflow_id__}": bad data but no dependents recalibrated.'
+            )
+        logger.debug(
+            f'"{workflow.__workflow_id__}": no dependents recalibrated.')
+    # calibrate
+    if any(recalibrated):
+        logger.debug(
+            f'recalibrate "{workflow.__workflow_id__}" because some dependents recalibrated.'
+        )
+    elif not report.in_spec and not report.bad_data:
+        logger.debug(
+            f'recalibrate "{workflow.__workflow_id__}" because out of spec.')
+    elif report.in_spec:
+        logger.error(
+            f'Never reach: recalibrate "{workflow.__workflow_id__}" because in spec.'
+        )
+    elif report.bad_data:
+        logger.error(
+            f'Never reach: recalibrate "{workflow.__workflow_id__}" because bad data.'
+        )
+    else:
+        logger.error(f'Never reach: recalibrate "{workflow.__workflow_id__}"')
+
+    report = await calibrate(workflow, state_path, plot, session_id)
+    if report.bad_data or not report.in_spec:
+        obey_the_oracle(report, state_path)
+        raise CalibrationFailedError(
+            f'"{workflow.__workflow_id__}": All dependents passed, but calibration failed!'
+        )
+    update_parameters(report, state_path)
+    logger.debug(f'"{workflow.__workflow_id__}": parameters updated')
+    return True
+
+
+@logger.catch(reraise=True)
+async def maintain(workflow: WorkflowType,
+                   code_path: str | Path,
+                   state_path: str | Path,
+                   session_id: str | None = None,
+                   run: bool = False,
+                   plot: bool = False,
+                   freeze: bool = False,
+                   fail_fast: bool = False,
+                   veryfy_source_code: bool = True):
+    if session_id is None:
+        session_id = uuid.uuid4().hex
+    logger.debug(f'run "{workflow.__workflow_id__}"'
+                 if run else f'maintain "{workflow.__workflow_id__}"')
+    # recursive maintain
+    exceptions = []
+    for n in get_dependents(workflow, code_path, veryfy_source_code):
+        logger.debug(
+            f'maintain "{n.__workflow_id__}" because it is depended by "{workflow.__workflow_id__}"'
+        )
+        try:
+            await maintain(n,
+                           code_path,
+                           state_path,
+                           session_id,
+                           run=False,
+                           plot=plot,
+                           freeze=freeze,
+                           fail_fast=fail_fast,
+                           veryfy_source_code=veryfy_source_code)
+        except Exception as e:
+            if fail_fast:
+                raise e
+            exceptions.append(e)
+    else:
+        logger.debug(
+            f'"{workflow.__workflow_id__}": All dependents maintained')
+    if any(exceptions):
+        raise exceptions[0]
+    # check_state
+    if check_state(workflow, code_path, state_path,
+                   veryfy_source_code) and not run:
+        logger.debug(
+            f'"{workflow.__workflow_id__}": In spec, no need to maintain')
+        return
+    # check_data
+    report = await check_data(workflow, state_path, plot, session_id)
+    if report.in_spec:
+        if not run:
+            logger.debug(
+                f'"{workflow.__workflow_id__}": In spec, no need to maintain')
+            return
+    elif report.bad_data:
+        logger.debug(
+            f'"{workflow.__workflow_id__}": Bad data, diagnosing dependents')
+        exceptions = []
+        for n in get_dependents(workflow, code_path, veryfy_source_code):
+            logger.debug(
+                f'diagnose "{n.__workflow_id__}" because of "{workflow.__workflow_id__}" bad data'
+            )
+            try:
+                await diagnose(n, code_path, state_path, plot, session_id,
+                               fail_fast, veryfy_source_code)
+            except Exception as e:
+                if fail_fast:
+                    raise e
+                exceptions.append(e)
+        else:
+            logger.debug(
+                f'"{workflow.__workflow_id__}": All dependents diagnosed')
+        if any(exceptions):
+            raise exceptions[0]
+    # calibrate
+    logger.debug(f'recalibrate "{workflow.__workflow_id__}"')
+    report = await calibrate(workflow, state_path, plot, session_id)
+    if report.bad_data or not report.in_spec:
+        if not freeze:
+            obey_the_oracle(report, state_path)
+        raise CalibrationFailedError(
+            f'"{workflow.__workflow_id__}": All dependents passed, but calibration failed!'
+        )
+    if not freeze:
+        update_parameters(report, state_path)
+        logger.debug(f'"{workflow.__workflow_id__}": parameters updated')
+    else:
+        logger.debug(f'"{workflow.__workflow_id__}": parameters freezed')
+    return
+
+
+@logger.catch(reraise=True)
+async def run(workflow: WorkflowType,
+              code_path: str | Path,
+              state_path: str | Path,
+              plot: bool = False,
+              freeze: bool = False):
+    session_id = uuid.uuid4().hex
+    logger.debug(f'run "{workflow.__workflow_id__}" without dependences.')
+    report = await calibrate(workflow, state_path, plot, session_id=session_id)
+    if report.bad_data or not report.in_spec:
+        if not freeze:
+            obey_the_oracle(report, state_path)
+        raise CalibrationFailedError(
+            f'"{workflow.__workflow_id__}": All dependents passed, but calibration failed!'
+        )
+    if not freeze:
+        update_parameters(report, state_path)
+        logger.debug(f'"{workflow.__workflow_id__}": parameters updated')
+    else:
+        logger.debug(f'"{workflow.__workflow_id__}": parameters freezed')
+    return