QuLab 2.3.5__cp311-cp311-macosx_10_9_universal2.whl → 2.4.0__cp311-cp311-macosx_10_9_universal2.whl

qulab/executor/__main__.py

@@ -0,0 +1,89 @@
+import importlib
+from pathlib import Path
+
+import click
+
+from .load import find_unreferenced_workflows
+from .schedule import maintain as maintain_workflow, run as run_workflow
+from .transform import set_config_api
+from .utils import workflow_template
+
+
+@click.group()
+def cli():
+    pass
+
+
+@click.command()
+@click.argument('workflow')
+@click.option('--code', '-c', default=None)
+def create(workflow, code):
+    """
+    Create a new workflow file.
+    """
+    if code is None:
+        code = Path.cwd()
+
+    fname = Path(code) / f'{workflow}.py'
+    if fname.exists():
+        click.echo(f'{workflow}.py already exists')
+        return
+
+    fname.parent.mkdir(parents=True, exist_ok=True)
+    deps = find_unreferenced_workflows(code)
+
+    with open(fname, 'w') as f:
+        f.write(workflow_template(list(deps)))
+    click.echo(f'{workflow}.py created')
+
+
+@click.command()
+@click.argument('workflow')
+@click.option('--code', '-c', default=None)
+@click.option('--data', '-d', default=None)
+@click.option('--api', '-g', default=None)
+@click.option('--no-dependents', '-n', is_flag=True)
+def run(workflow, code, data, api, no_dependents):
+    """
+    Run a workflow.
+    """
+    if api is not None:
+        api = importlib.import_module(api)
+        set_config_api(api.query_config, api.update_config)
+    if code is None:
+        code = Path.cwd()
+    if data is None:
+        data = Path(code) / 'logs'
+
+    if no_dependents:
+        run_workflow(workflow, code, data)
+    else:
+        maintain_workflow(workflow, code, data, run=True)
+
+
+@click.command()
+@click.argument('workflow')
+@click.option('--code', '-c', default=None)
+@click.option('--data', '-d', default=None)
+@click.option('--api', '-g', default=None)
+def maintain(workflow, code, data, api):
+    """
+    Maintain a workflow.
+    """
+    if api is not None:
+        api = importlib.import_module(api)
+        set_config_api(api.query_config, api.update_config)
+    if code is None:
+        code = Path.cwd()
+    if data is None:
+        data = Path(code) / 'logs'
+
+    maintain_workflow(workflow, code, data, run=False)
+
+
+cli.add_command(maintain)
+cli.add_command(run)
+cli.add_command(create)
+
+if __name__ == '__main__':
+    cli()

qulab/executor/load.py

@@ -0,0 +1,202 @@
+import inspect
+import warnings
+from importlib.util import module_from_spec, spec_from_file_location
+from pathlib import Path
+from types import ModuleType
+
+import loguru
+
+from .storage import Result
+
+
+class SetConfigWorkflow():
+    __timeout__ = None
+
+    def __init__(self, key):
+        self.key = key
+
+    def depends(self):
+        return [[]]
+
+    def check_state(self, history: Result) -> bool:
+        from . import transform
+        try:
+            return self._equal(history.params[self.key],
+                               transform.query_config(self.key))
+        except:
+            return False
+
+    def calibrate(self):
+        from . import transform
+        try:
+            value = transform.query_config(self.key)
+        except:
+            value = eval(input(f'"{self.key}": '))
+        return self.key, value
+
+    def analyze(self, key, value, history):
+        return 'OK', {key: value}, {}
+
+    def check(self):
+        from .transform import query_config
+        return self.key, query_config(self.key)
+
+    def check_analyze(self, key, value, history):
+        return 'Out of Spec', {key: value}, {}
+
+    @staticmethod
+    def _equal(a, b):
+        import numpy as np
+
+        try:
+            return a == b
+        except:
+            pass
+
+        if isinstance(a, np.ndarray) and isinstance(b, np.ndarray):
+            return a.shape == b.shape and np.all(a == b)
+
+        return False
+
+
+WorkflowType = ModuleType | SetConfigWorkflow
+
+
+def can_call_without_args(func):
+    if not callable(func):
+        return False
+
+    # 获取函数签名
+    sig = inspect.signature(func)
+    for param in sig.parameters.values():
+        # 如果有参数没有默认值且不是可变参数，则无法无参调用
+        if (param.default is param.empty and param.kind
+                not in (param.VAR_POSITIONAL, param.VAR_KEYWORD)):
+            return False
+    return True
+
+
+def verify_calibrate_method(module: WorkflowType):
+    if not hasattr(module, 'calibrate'):
+        raise AttributeError(
+            f"Workflow {module.__file__} does not have 'calibrate' function")
+
+    if not can_call_without_args(module.calibrate):
+        raise AttributeError(
+            f"Workflow {module.__file__} 'calibrate' function should not have any parameters"
+        )
+
+    if not hasattr(module, 'analyze'):
+        raise AttributeError(
+            f"Workflow {module.__file__} does not have 'analyze' function")
+
+
+def verify_check_method(module: WorkflowType):
+    if not hasattr(module, 'check'):
+        warnings.warn(
+            f"Workflow {module.__file__} does not have 'check' function")
+    else:
+        if not can_call_without_args(module.check):
+            raise AttributeError(
+                f"Workflow {module.__file__} 'check' function should not have any parameters"
+            )
+
+        if not hasattr(module, 'check_analyze'):
+            raise AttributeError(
+                f"Workflow {module.__file__} has 'check' function but does not have 'check_analyze' function"
+            )
+
+
+def is_workflow(module: ModuleType) -> bool:
+    try:
+        verify_calibrate_method(module)
+        return True
+    except AttributeError:
+        return False
+
+
+def find_unreferenced_workflows(path: str) -> list[str]:
+    root = Path(path).resolve()
+    workflows = []
+    workflow_paths: set[str] = set()
+
+    # Collect all workflow modules
+    for file_path in root.rglob("*.py"):
+        if file_path.name == "__init__.py":
+            continue
+        try:
+            rel_path = file_path.relative_to(root)
+        except ValueError:
+            continue
+
+        module = load_workflow(str(rel_path), root)
+
+        if is_workflow(module):
+            rel_str = str(rel_path)
+            workflows.append(rel_str)
+            workflow_paths.add(rel_str)
+
+    dependencies: set[str] = set()
+
+    # Check dependencies for each workflow module
+    for rel_str in workflows:
+        module = load_workflow(rel_str, root)
+
+        depends_func = getattr(module, "depends", None)
+        if depends_func and callable(depends_func):
+            if not can_call_without_args(depends_func):
+                warnings.warn(
+                    f"Skipping depends() in {rel_str} as it requires arguments"
+                )
+                continue
+            try:
+                depends_list = depends_func()
+            except Exception as e:
+                warnings.warn(f"Error calling depends() in {rel_str}: {e}")
+                continue
+
+            if not isinstance(depends_list, list) or not all(
+                isinstance(item, str) for item in depends_list
+            ):
+                warnings.warn(
+                    f"depends() in {rel_str} did not return a list of strings"
+                )
+                continue
+
+            for dep in depends_list:
+                dep_full = (root / dep).resolve()
+                try:
+                    dep_rel = dep_full.relative_to(root)
+                except ValueError:
+                    continue
+                dep_rel_str = str(dep_rel)
+                if dep_rel_str in workflow_paths:
+                    dependencies.add(dep_rel_str)
+
+    # Determine unreferenced workflows
+    unreferenced = [wp for wp in workflows if wp not in dependencies]
+    return unreferenced
+
+
+def load_workflow(file_name: str,
+                  base_path: str | Path,
+                  package='workflows') -> WorkflowType:
+    if file_name.startswith('cfg:'):
+        return SetConfigWorkflow(file_name[4:])
+    base_path = Path(base_path)
+    path = Path(file_name)
+    module_name = f"{package}.{'.'.join([*path.parts[:-1], path.stem])}"
+    spec = spec_from_file_location(module_name, base_path / path)
+    module = module_from_spec(spec)
+    spec.loader.exec_module(module)
+
+    if not hasattr(module, '__timeout__'):
+        module.__timeout__ = None
+
+    if not hasattr(module, 'depends'):
+        module.depends = lambda: [[]]
+
+    verify_calibrate_method(module)
+    verify_check_method(module)
+
+    return module

qulab/executor/schedule.py

@@ -0,0 +1,223 @@
+import functools
+import uuid
+from datetime import datetime, timedelta
+from pathlib import Path
+
+from loguru import logger
+
+from . import transform
+from .load import load_workflow
+from .storage import (Result, find_result, renew_result, revoke_result,
+                      save_result)
+
+
+class CalibrationFailedError(Exception):
+    pass
+
+
+def check_state(workflow: str, code_path: str | Path,
+                state_path: str | Path) -> 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}"')
+    result = find_result(workflow, state_path)
+    if not result:
+        logger.debug(f'check_state failed: No history found for "{workflow}"')
+        return False
+    node = load_workflow(workflow, code_path)
+    if hasattr(node, 'check_state') and callable(node.check_state):
+        logger.debug(
+            f'check_state: "{workflow}" has custom check_state method')
+        return node.check_state(result)
+    if node.__timeout__ is not None and datetime.now(
+    ) > result.checked_time + timedelta(seconds=node.__timeout__):
+        logger.debug(f'check_state failed: "{workflow}" has expired')
+        return False
+    if not result.in_spec:
+        logger.debug(f'check_state failed: "{workflow}" is out of spec')
+        return False
+    if result.bad_data:
+        logger.debug(f'check_state failed: "{workflow}" has bad data')
+        return False
+    for n in get_dependents(workflow, code_path):
+        r = find_result(n, state_path)
+        if r is None or r.checked_time > result.checked_time:
+            logger.debug(
+                f'check_state failed: "{workflow}" has outdated dependencies')
+            return False
+    for n in get_dependents(workflow, code_path):
+        if not check_state(n, code_path, state_path):
+            logger.debug(
+                f'check_state failed: "{workflow}" has bad dependencies')
+            return False
+    return True
+
+
+@functools.lru_cache(maxsize=128)
+def check_data(workflow: str, code_path: str | Path, state_path: str | Path,
+               session_id: str) -> Result:
+    """
+    check data answers two questions:
+    Is the parameter associated with this cal in spec,
+    and is the cal scan working as expected?
+    """
+    node = load_workflow(workflow, code_path)
+    history = find_result(workflow, state_path)
+
+    if history is None:
+        logger.debug(f'No history found for "{workflow}"')
+        result = Result()
+        result.in_spec = False
+        result.bad_data = False
+        return result
+
+    if history.bad_data:
+        logger.debug(f'History found for "{workflow}", but bad data')
+        return history
+    if not history.in_spec:
+        logger.debug(f'History found for "{workflow}", but out of spec')
+        return history
+
+    logger.debug(f'History found for "{workflow}", but has expired')
+
+    if hasattr(node, 'check') and callable(node.check) and hasattr(
+            node, 'check_analyze') and callable(node.check_analyze):
+        logger.debug(f'Checking "{workflow}" with "check" method ...')
+        data = node.check()
+        logger.debug(f'Checked "{workflow}" !')
+        result = transform.params_to_result(
+            node.check_analyze(*data,
+                               history=transform.result_to_params(history)))
+        result.data = data
+        if result.in_spec:
+            logger.debug(f'"{workflow}": checked in spec, renewing result')
+            renew_result(workflow, state_path)
+        else:
+            logger.debug(f'"{workflow}": checked out of spec, revoking result')
+            revoke_result(workflow, state_path)
+    else:
+        logger.debug(f'Checking "{workflow}" with "calibrate" method ...')
+        data = node.calibrate()
+        logger.debug(f'Calibrated "{workflow}" !')
+        result = transform.params_to_result(
+            node.analyze(*data, history=transform.result_to_params(history)))
+        result.data = data
+        result.fully_calibrated = True
+        save_result(workflow, result, state_path)
+
+    return result
+
+
+@functools.lru_cache(maxsize=128)
+def calibrate(workflow, code_path: str | Path, state_path: str | Path,
+              session_id: str) -> Result:
+    result = Result()
+    node = load_workflow(workflow, code_path)
+    history = find_result(workflow, state_path)
+
+    logger.debug(f'Calibrating "{workflow}" ...')
+    data = node.calibrate()
+    logger.debug(f'Calibrated "{workflow}" !')
+    result = transform.params_to_result(
+        node.analyze(*data, history=transform.result_to_params(history)))
+    result.data = data
+    result.fully_calibrated = True
+    save_result(workflow, result, state_path)
+    return result
+
+
+def diagnose(node, code_path: str | Path, state_path: str | Path,
+             session_id: str):
+    '''
+    Returns: True if node or dependent recalibrated.
+    '''
+    logger.debug(f'diagnose "{node}"')
+    # check_data
+    result = check_data(node, code_path, state_path, session_id)
+    # in spec case
+    if result.in_spec:
+        return False
+    # bad data case
+    recalibrated = []
+    if result.bad_data:
+        recalibrated = [
+            diagnose(n, code_path, state_path, session_id)
+            for n in get_dependents(node, code_path)
+        ]
+    if not any(recalibrated):
+        return False
+    # calibrate
+    if result.fully_calibrated and result.in_spec:
+        pass
+    else:
+        logger.debug(
+            f'recalibrate "{node}" because some dependents recalibrated')
+        result = calibrate(node, code_path, state_path, session_id)
+    if result.bad_data or not result.in_spec:
+        raise CalibrationFailedError(
+            f'"{node}": All dependents passed, but calibration failed!')
+    transform.update_parameters(result)
+    return True
+
+
+def get_dependents(workflow: str, code_path: str | Path) -> list[str]:
+    return [n for n in load_workflow(workflow, code_path).depends()[0]]
+
+
+#@logger.catch(reraise=True)
+def maintain(node,
+             code_path: str | Path,
+             state_path: str | Path,
+             session_id: str | None = None,
+             run: bool = False):
+    if session_id is None:
+        session_id = uuid.uuid4().hex
+    logger.debug(f'run "{node}"' if run else f'maintain "{node}"')
+    # recursive maintain
+    for n in get_dependents(node, code_path):
+        logger.debug(f'maintain "{n}" because it is depended by "{node}"')
+        maintain(n, code_path, state_path, session_id)
+    else:
+        logger.debug(f'"{node}": All dependents maintained')
+    # check_state
+    if check_state(node, code_path, state_path) and not run:
+        logger.debug(f'"{node}": In spec, no need to maintain')
+        return
+    # check_data
+    result = check_data(node, code_path, state_path, session_id)
+    if result.in_spec:
+        if not run:
+            logger.debug(f'"{node}": In spec, no need to maintain')
+            return
+    elif result.bad_data:
+        logger.debug(f'"{node}": Bad data, diagnosing dependents')
+        for n in get_dependents(node, code_path):
+            logger.debug(f'diagnose "{n}" because of "{node}" bad data')
+            diagnose(n, code_path, state_path, session_id)
+        else:
+            logger.debug(f'"{node}": All dependents diagnosed')
+    # calibrate
+    logger.debug(f'recalibrate "{node}"')
+    result = calibrate(node, code_path, state_path, session_id)
+    if result.bad_data or not result.in_spec:
+        raise CalibrationFailedError(
+            f'"{node}": All dependents passed, but calibration failed!')
+    transform.update_parameters(result)
+    return
+
+
+def run(node, code_path: str | Path,
+             state_path: str | Path):
+    logger.debug(f'run "{node}" without dependences.')
+    result = calibrate(node, code_path, state_path)
+    if result.bad_data or not result.in_spec:
+        raise CalibrationFailedError(
+            f'"{node}": All dependents passed, but calibration failed!')
+    transform.update_parameters(result)
+    return

qulab/executor/storage.py

@@ -0,0 +1,143 @@
+import pickle
+import uuid
+from dataclasses import dataclass, field
+from datetime import datetime, timedelta
+from pathlib import Path
+
+from loguru import logger
+
+
+@dataclass
+class Result():
+    in_spec: bool = False
+    bad_data: bool = False
+    fully_calibrated: bool = False
+    calibrated_time: datetime = field(default_factory=datetime.now)
+    checked_time: datetime = field(default_factory=datetime.now)
+    ttl: timedelta = timedelta(days=3650)
+    params: dict = field(default_factory=dict)
+    info: dict = field(default_factory=dict)
+    data: tuple = field(default_factory=tuple)
+    previous: Path | None = None
+
+
+class Graph:
+
+    def __init__(self):
+        self.nodes = {}
+        self.heads = set()
+        self.roots = set()
+
+    def add_node(self, node: str, deps: list[str]):
+        if node not in self.nodes:
+            self.nodes[node] = deps
+            if not deps:
+                self.heads.add(node)
+            for dep in deps:
+                if dep not in self.nodes:
+                    self.nodes[dep] = []
+                self.roots.discard(dep)
+
+
+def random_path(base: Path) -> Path:
+    while True:
+        s = uuid.uuid4().hex
+        path = Path(s[:2]) / s[2:4] / s[4:6] / s[6:]
+        if not (base / path).exists():
+            return path
+
+
+def save_result(workflow: str, result: Result, base_path: str | Path):
+    logger.debug(
+        f'Saving result for "{workflow}", {result.in_spec=}, {result.bad_data=}, {result.fully_calibrated=}'
+    )
+    base_path = Path(base_path)
+    path = random_path(base_path)
+    (base_path / 'objects' / path).parent.mkdir(parents=True, exist_ok=True)
+    result.previous = get_head(workflow, base_path)
+    with open(base_path / 'objects' / path, "wb") as f:
+        pickle.dump(result, f)
+    set_head(workflow, path, base_path)
+
+
+def find_result(workflow: str, base_path: str | Path) -> Result | None:
+    base_path = Path(base_path)
+    path = get_head(workflow, base_path)
+    if path is None:
+        return None
+    with open(base_path / 'objects' / path, "rb") as f:
+        return pickle.load(f)
+
+
+def renew_result(workflow: str, base_path: str | Path):
+    logger.debug(f'Renewing result for "{workflow}"')
+    result = find_result(workflow, base_path)
+    if result is not None:
+        result.checked_time = datetime.now()
+        save_result(workflow, result, base_path)
+
+
+def revoke_result(workflow: str, base_path: str | Path):
+    logger.debug(f'Revoking result for "{workflow}"')
+    base_path = Path(base_path)
+    path = get_head(workflow, base_path)
+    if path is not None:
+        with open(base_path / 'objects' / path, "rb") as f:
+            result = pickle.load(f)
+        result.in_spec = False
+        save_result(workflow, result, base_path)
+
+
+def set_head(workflow: str, path: Path, base_path: str | Path):
+    base_path = Path(base_path)
+    base_path.mkdir(parents=True, exist_ok=True)
+    try:
+        with open(base_path / "heads", "rb") as f:
+            heads = pickle.load(f)
+    except:
+        heads = {}
+    heads[workflow] = path
+    with open(base_path / "heads", "wb") as f:
+        pickle.dump(heads, f)
+
+
+def get_head(workflow: str, base_path: str | Path) -> Path | None:
+    base_path = Path(base_path)
+    try:
+        with open(base_path / "heads", "rb") as f:
+            heads = pickle.load(f)
+        return heads[workflow]
+    except:
+        return None
+
+
+def get_graph(base_path: str | Path) -> dict[str, list[str]]:
+    base_path = Path(base_path)
+    try:
+        with open(base_path / "heads", "rb") as f:
+            heads = pickle.load(f)
+    except:
+        heads = {}
+    graph = {}
+    for workflow, path in heads.items():
+        graph[workflow] = []
+        while path is not None:
+            with open(base_path / 'objects' / path, "rb") as f:
+                result = pickle.load(f)
+            path = result.previous
+            if path is not None:
+                graph[workflow].append(path)
+    return graph
+
+
+def update_graph(workflow: str, base_path: str | Path):
+    base_path = Path(base_path)
+    graph = get_graph(base_path)
+    for workflow, deps in graph.items():
+        for dep in deps:
+            if dep not in graph:
+                graph[dep] = []
+            if workflow not in graph[dep]:
+                graph[dep].append(workflow)
+    with open(base_path / "graph", "wb") as f:
+        pickle.dump(graph, f)