add cryo_test testing framework

crates/python/python/cryo_test/README.md

@@ -0,0 +1,35 @@
+
+# cryo test
+
+utility for performing comparisons between cryo environments
+
+```bash
+cryo_test is a cli tool for comparing cryo outputs across different conditions
+
+Usage: cryo_test QUERY [OPTIONS]
+
+This will 1. setup comparison dir, 2. collect data, and 3. compare outputs
+
+Examples:
+  cryo_test --rpc source1=https://h1.com source2=https://h2.com (compare rpc's)
+  cryo_test --executable old=/path/to/cryo1 new=/path/to/cryo2 (compare executables)
+  cryo_test --python ... (use python)
+  cryo_test --cli-vs-python ... (compare cli vs python)
+
+Options:
+  -h, --help                                 show this help message and exit
+  -e, --executable, --executables EXECUTABLE [...]
+                                             executable(s) to use
+  --rpc RPC [...]                            rpc endpoint(s) to use
+  -d, --datatype, --datatypes DATATYPE [...]
+                                             datatype(s) to collect
+  -p, --python                               use python for all batches
+  --cli-vs-python                            compare cli to python
+  -r, --rerun                                rerun previous comparison
+  --label LABEL                              name of comparison
+  -s, --steps {setup,collect,compare} [...]  steps to perform {setup, collect, compare}
+  --dir DIR                                  directory for storing comparison data
+  -i, --interactive                          load data in interactive python session
+  --scan-interactive                         scan data in interactive python session
+  --debug, --pdb                             enter debug mode upon error
+```

crates/python/python/cryo_test/cryo_test/__init__.py

@@ -0,0 +1,6 @@
+"""library for performing comparisons between cryo jobs"""
+
+from .comparison import perform_comparison
+
+
+__version__ = '0.1.0'

crates/python/python/cryo_test/cryo_test/__main__.py

@@ -0,0 +1,7 @@
+from __future__ import annotations
+
+from . import cli
+
+
+if __name__ == '__main__':
+    cli.run_cli()

crates/python/python/cryo_test/cryo_test/cli/__init__.py

@@ -0,0 +1,10 @@
+from __future__ import annotations
+
+from .cli_classes import *
+from .cli_run import *
+from .cli_summary import *
+from .cli_utils import *
+
+
+if __name__ == '__main__':
+    run_cli()

crates/python/python/cryo_test/cryo_test/cli/cli_classes.py

@@ -0,0 +1,97 @@
+from __future__ import annotations
+
+import argparse
+import typing
+
+import rich_argparse
+
+if typing.TYPE_CHECKING:
+    import typing_extensions
+
+
+usage_template = """%(prog)s [{prog}]QUERY [OPTIONS][/{prog}]\n\n
+
+[{help} not bold]This will 1. setup comparison dir, 2. collect data, and 3. compare outputs[/{help} not bold]
+
+[{groups}]Examples:[/{groups}]
+  [{prog}]cryo_test --rpc source1=https://h1.com source2=https://h2.com[/{prog}] (compare rpc's)
+  [{prog}]cryo_test --executable old=/path/to/cryo1 new=/path/to/cryo2[/{prog}] (compare executables)
+  [{prog}]cryo_test --python ...[/{prog}] (use python)
+  [{prog}]cryo_test --cli-vs-python ...[/{prog}] (compare cli vs python)"""
+
+
+class CryoTestHelpFormatter(rich_argparse.RichHelpFormatter):
+    usage_markup = True
+
+    styles = {
+        'argparse.prog': 'bold white',
+        'argparse.groups': 'bold rgb(0,255,0)',
+        'argparse.args': 'bold white',
+        'argparse.metavar': 'grey62',
+        'argparse.help': 'grey62',
+        'argparse.text': 'blue',
+        'argparse.syntax': 'blue',
+        'argparse.default': 'blue',
+    }
+
+    def __init__(self, prog: str) -> None:
+        super().__init__('cryo_test', max_help_position=45)
+
+    def _format_args(self, action, default_metavar):  # type: ignore
+        get_metavar = self._metavar_formatter(action, default_metavar)
+        if action.nargs == argparse.ZERO_OR_MORE:
+            return '[%s [%s ...]]' % get_metavar(2)
+        elif action.nargs == argparse.ONE_OR_MORE:
+            return '%s [...]' % get_metavar(1)
+        return super()._format_args(action, default_metavar)
+
+    def format_help(self) -> str:
+        import rich
+
+        line = '[{prog}]cryo_test[/{prog}] [{help}]is a cli tool for comparing [{prog}]cryo[/{prog}] outputs across different conditions\n'
+        rich.print(self.format_styles(line))
+
+        # indent certain arguments for full alignment
+        raw = super().format_help()
+        lines = raw.split('\n')
+        for i, line in enumerate(lines):
+            if line.startswith('  \x1b[38;2;197;149;242m--'):
+                lines[i] = '    ' + lines[i].replace('    ', '', 1)
+
+        # indices = [i for i, line in enumerate(lines) if line.startswith('  \x1b[1;37m--')]
+        # if lines[indices[0]] == '':
+        #     lines = lines[:indices[0]] + ['FUCK'] + lines[indices[0]:]
+
+        # lines = [
+        #     line
+        #     for line in lines
+        #     if 'Options:' not in line and '--help' not in line
+        # ]
+        return '\n'.join(lines).replace('\n\n\n', '\n\n')
+
+    @classmethod
+    def format_styles(cls, s: str) -> str:
+        styles = {k.split('.')[1]: v for k, v in cls.styles.items()}
+        return s.format(**styles)
+
+
+class CryoTestArgParser(argparse.ArgumentParser):
+    def __init__(self, *args: typing.Any, **kwargs: typing.Any) -> None:
+        return super().__init__(
+            *args,
+            formatter_class=CryoTestHelpFormatter,  # type: ignore
+            usage=CryoTestHelpFormatter.format_styles(usage_template),
+            **kwargs,
+        )
+
+    def error(self, message: str) -> typing_extensions.NoReturn:
+        import sys
+        import rich
+
+        sys.stderr.write(f'Error: {message}\n')
+        print()
+        self.print_usage()
+        print()
+        line = '[{help}]show all options with[/{help}] [bold white]cryo_test -h[/bold white]'
+        rich.print(CryoTestHelpFormatter.format_styles(line))
+        sys.exit(0)

crates/python/python/cryo_test/cryo_test/cli/cli_run.py

@@ -0,0 +1,212 @@
+from __future__ import annotations
+
+import typing
+
+from .. import commands
+from .. import comparison
+from .. import files
+from . import cli_classes
+from . import cli_summary
+from . import cli_utils
+
+if typing.TYPE_CHECKING:
+    import argparse
+
+
+def run_cli() -> None:
+    try:
+        # parse inputs
+        args = parse_args()
+
+        if set(args.steps) in [
+            {'compare'},
+            {'collect'},
+            {'collect', 'compare'},
+        ]:
+            rerun = True
+        else:
+            rerun = args.rerun
+
+        # load commands
+        if rerun:
+            if args.dir is not None:
+                comparison_dir = args.dir
+            else:
+                comparison_dir = files.get_most_recent_comparison_dir()
+            batches = files.load_commands(comparison_dir)
+        else:
+            comparison_dir, batches = create_command_batches(args)
+            files.save_commands(comparison_dir=comparison_dir, batches=batches)
+
+        # summarize
+        cli_summary.summarize_args(args, comparison_dir, batches)
+
+        # perform comparison
+        comparison.perform_comparison(
+            batches=batches,
+            comparison_dir=comparison_dir,
+            steps=args.steps,
+        )
+
+        # enter interactive session
+        if args.interactive or args.scan_interactive:
+            print('entering interactive session')
+            if args.scan_interactive:
+                all_data: typing.Any = files.scan_all_data(comparison_dir)
+            else:
+                all_data = files.load_all_data(comparison_dir)
+            cli_utils.open_interactive_session({'data': all_data})
+    except Exception as e:
+        if args.debug:
+            cli_utils._enter_debugger()
+        else:
+            raise e
+
+
+def parse_args() -> argparse.Namespace:
+    parser = cli_classes.CryoTestArgParser()
+    #
+    # specifying batch args
+    parser.add_argument(
+        '-e',
+        '--executable',
+        '--executables',
+        nargs='+',
+        help='executable(s) to use',
+    )
+    parser.add_argument('--rpc', nargs='+', help='rpc endpoint(s) to use')
+    parser.add_argument(
+        '-d',
+        '--datatype',
+        '--datatypes',
+        nargs='+',
+        help='datatype(s) to collect',
+    )
+    parser.add_argument(
+        '-p', '--python', action='store_true', help='use python for all batches'
+    )
+    parser.add_argument(
+        '--cli-vs-python', action='store_true', help='compare cli to python'
+    )
+    #
+    # specifying how to run
+    parser.add_argument(
+        '-r', '--rerun', action='store_true', help='rerun previous comparison'
+    )
+    parser.add_argument('--label', help='name of comparison')
+    parser.add_argument(
+        '-s',
+        '--steps',
+        help='steps to perform {setup, collect, compare}',
+        nargs='+',
+        choices=['setup', 'collect', 'compare'],
+        default=['setup', 'collect', 'compare'],
+    )
+    parser.add_argument('--dir', help='directory for storing comparison data')
+    parser.add_argument(
+        '-i',
+        '--interactive',
+        action='store_true',
+        help='load data in interactive python session',
+    )
+    parser.add_argument(
+        '--scan-interactive',
+        action='store_true',
+        help='scan data in interactive python session',
+    )
+    parser.add_argument(
+        '--debug',
+        '--pdb',
+        action='store_true',
+        help='enter debug mode upon error',
+    )
+    return parser.parse_args()
+
+
+def create_command_batches(
+    args: argparse.Namespace,
+) -> tuple[str, dict[str, dict[str, str]]]:
+    command_arg_combos: dict[str, list[str]] = {}
+    common_command_args: commands.PartialCommandArgs = {}
+    batch_specific_args: dict[str, commands.PartialCommandArgs] = {}
+
+    # determine batch mode (the variables that change across batches)
+    if args.executable is not None and len(args.executable) > 1:
+        batch_mode = 'executable'
+    elif args.rpc is not None and len(args.rpc) > 1:
+        batch_mode = 'rpc'
+    else:
+        batch_mode = 'default'
+
+    # determine comparison dir
+    comparison_dir = files.create_comparison_dir(args.dir, args.label)
+
+    # determine interface
+    if args.python:
+        common_command_args['interface'] = 'python'
+
+    #
+    if batch_mode == 'default':
+        batch_specific_args = {'default': {}}
+
+    # arg: executable
+    if batch_mode == 'executable':
+        for raw_executable in args.executable:
+            batch_name, executable = raw_executable.split('=', maxsplit=1)
+            batch_specific_args[batch_name] = {'executable': executable}
+    else:
+        if args.executable is None:
+            common_command_args['executable'] = 'cryo'
+        elif len(args.executable) == 1:
+            common_command_args['executable'] = args.executable[0]
+        else:
+            raise Exception()
+
+    # arg: rpc
+    if batch_mode == 'rpc':
+        for raw_rpc in args.rpcs:
+            batch_name, rpc = raw_rpc.split('=', maxsplits=1)
+            batch_specific_args[batch_name] = {'rpc': rpc}
+    else:
+        if args.rpc is None:
+            pass
+        elif len(args.rpc) == 1:
+            common_command_args['rpc'] = args.rpc[0]
+        else:
+            raise Exception()
+
+    # arg: datatypes
+    if args.datatype is not None:
+        command_arg_combos['datatype'] = args.datatype
+
+    # if cli-vs-python, create cli and python version of each batch
+    if args.cli_vs_python:
+        new_batch_specific_args = {}
+        for batch_name, batch_kwargs in batch_specific_args.items():
+            new_batch_specific_args[batch_name + '_cli'] = dict(
+                batch_kwargs, interface='cli'
+            )
+            new_batch_specific_args[batch_name + '_python'] = dict(
+                batch_kwargs, interface='python'
+            )
+
+    # arg: output_dir
+    for batch_name in batch_specific_args.keys():
+        batch_specific_args[batch_name]['data_root'] = files.get_batch_data_dir(
+            comparison_dir=comparison_dir,
+            batch_name=batch_name,
+        )
+
+    # generate commands
+    batches = {}
+    for batch_name in batch_specific_args.keys():
+        kwargs: commands.PartialCommandArgs = dict(
+            **common_command_args, **batch_specific_args[batch_name]
+        )
+        batches[batch_name] = commands.generate_commands(
+            common_command_args=kwargs,
+            command_arg_combos=command_arg_combos,
+        )
+
+    return comparison_dir, batches
+