|
| 1 | +import os |
| 2 | +import uuid |
| 3 | + |
| 4 | +import click |
| 5 | + |
| 6 | +from notebooker.constants import DEFAULT_SERIALIZER |
| 7 | +from notebooker.execute_notebook import execute_notebook_entrypoint |
| 8 | +from notebooker.serialization import SERIALIZER_TO_CLI_OPTIONS |
| 9 | +from notebooker.settings import BaseConfig, WebappConfig |
| 10 | +from notebooker.snapshot import snap_latest_successful_notebooks |
| 11 | +from notebooker.web.app import main |
| 12 | + |
| 13 | + |
| 14 | +class NotebookerEntrypoint(click.Group): |
| 15 | + def parse_args(self, ctx, args): |
| 16 | + try: |
| 17 | + serializer_arg = args.index("--serializer-cls") |
| 18 | + serializer = args[serializer_arg + 1] |
| 19 | + except ValueError: |
| 20 | + serializer = DEFAULT_SERIALIZER |
| 21 | + self.params += SERIALIZER_TO_CLI_OPTIONS[serializer].params |
| 22 | + |
| 23 | + return super().parse_args(ctx, args) |
| 24 | + |
| 25 | + |
| 26 | +pass_config = click.make_pass_decorator(BaseConfig) |
| 27 | + |
| 28 | + |
| 29 | +def filesystem_default_value(dirname): |
| 30 | + return os.path.join(os.path.expanduser("~"), ".notebooker", dirname, str(uuid.uuid4())) |
| 31 | + |
| 32 | + |
| 33 | +@click.group(cls=NotebookerEntrypoint) |
| 34 | +@click.option("--notebook-kernel-name", default=None, help="The name of the kernel which is running our notebook code.") |
| 35 | +@click.option( |
| 36 | + "--output-base-dir", |
| 37 | + default=filesystem_default_value("output"), |
| 38 | + help="The base directory to which we will save our notebook output temporarily. Required by Papermill.", |
| 39 | +) |
| 40 | +@click.option( |
| 41 | + "--template-base-dir", |
| 42 | + default=filesystem_default_value("templates"), |
| 43 | + help="The base directory to which we will save our notebook templates which have been converted " |
| 44 | + "from .py to .ipynb.", |
| 45 | +) |
| 46 | +@click.option( |
| 47 | + "--py-template-base-dir", |
| 48 | + default=None, |
| 49 | + help="The base directory of the git repository which holds the notebook templates as .py files. " |
| 50 | + "If not specified, this will default to the sample directory within notebooker.", |
| 51 | +) |
| 52 | +@click.option( |
| 53 | + "--py-template-subdir", |
| 54 | + default=None, |
| 55 | + help="The subdirectory of the git repository which contains only notebook templates.", |
| 56 | +) |
| 57 | +@click.option( |
| 58 | + "--notebooker-disable-git", |
| 59 | + default=False, |
| 60 | + is_flag=True, |
| 61 | + help="If selected, notebooker will not try to pull the latest version of python templates from git.", |
| 62 | +) |
| 63 | +@click.option( |
| 64 | + "--serializer-cls", |
| 65 | + default=DEFAULT_SERIALIZER, |
| 66 | + help="The serializer class through which we will save the notebook result.", |
| 67 | +) |
| 68 | +@click.pass_context |
| 69 | +def base_notebooker( |
| 70 | + ctx, |
| 71 | + notebook_kernel_name, |
| 72 | + output_base_dir, |
| 73 | + template_base_dir, |
| 74 | + py_template_base_dir, |
| 75 | + py_template_subdir, |
| 76 | + notebooker_disable_git, |
| 77 | + serializer_cls, |
| 78 | + **serializer_args, |
| 79 | +): |
| 80 | + config = BaseConfig( |
| 81 | + SERIALIZER_CLS=serializer_cls, |
| 82 | + SERIALIZER_CONFIG=serializer_args, |
| 83 | + NOTEBOOK_KERNEL_NAME=notebook_kernel_name, |
| 84 | + OUTPUT_DIR=output_base_dir, |
| 85 | + TEMPLATE_DIR=template_base_dir, |
| 86 | + PY_TEMPLATE_BASE_DIR=py_template_base_dir, |
| 87 | + PY_TEMPLATE_SUBDIR=py_template_subdir, |
| 88 | + NOTEBOOKER_DISABLE_GIT=notebooker_disable_git, |
| 89 | + ) |
| 90 | + ctx.obj = config |
| 91 | + |
| 92 | + |
| 93 | +@base_notebooker.command() |
| 94 | +@click.option("--port", default=11828) |
| 95 | +@click.option("--logging-level", default="INFO") |
| 96 | +@click.option("--debug", default=False) |
| 97 | +@click.option("--base-cache-dir", default=filesystem_default_value("webcache")) |
| 98 | +@pass_config |
| 99 | +def start_webapp(config: BaseConfig, port, logging_level, debug, base_cache_dir): |
| 100 | + web_config = WebappConfig.copy_existing(config) |
| 101 | + web_config.PORT = port |
| 102 | + web_config.LOGGING_LEVEL = logging_level |
| 103 | + web_config.DEBUG = debug |
| 104 | + web_config.CACHE_DIR = base_cache_dir |
| 105 | + return main(web_config) |
| 106 | + |
| 107 | + |
| 108 | +@base_notebooker.command() |
| 109 | +@click.option("--report-name", help="The name of the template to execute, relative to the template directory.") |
| 110 | +@click.option( |
| 111 | + "--overrides-as-json", default="{}", help="The parameters to inject into the notebook template, in JSON format." |
| 112 | +) |
| 113 | +@click.option( |
| 114 | + "--iterate-override-values-of", |
| 115 | + default="", |
| 116 | + help="For the key/values in the overrides, set this to the value of one of the keys to run reports for " |
| 117 | + "each of its values.", |
| 118 | +) |
| 119 | +@click.option("--report-title", default="", help="A custom title for this notebook. The default is the report_name.") |
| 120 | +@click.option("--n-retries", default=3, help="The number of times to retry when executing this notebook.") |
| 121 | +@click.option( |
| 122 | + "--job-id", |
| 123 | + default=str(uuid.uuid4()), |
| 124 | + help="The unique job ID for this notebook. Can be non-unique, but note that you will overwrite history.", |
| 125 | +) |
| 126 | +@click.option("--mailto", default="", help="A comma-separated list of email addresses which will receive results.") |
| 127 | +@click.option("--pdf-output/--no-pdf-output", default=True, help="Whether we generate PDF output or not.") |
| 128 | +@click.option( |
| 129 | + "--prepare-notebook-only", |
| 130 | + is_flag=True, |
| 131 | + help='Used for debugging and testing. Whether to actually execute the notebook or just "prepare" it.', |
| 132 | +) |
| 133 | +@pass_config |
| 134 | +def execute_notebook( |
| 135 | + config: BaseConfig, |
| 136 | + report_name, |
| 137 | + overrides_as_json, |
| 138 | + iterate_override_values_of, |
| 139 | + report_title, |
| 140 | + n_retries, |
| 141 | + job_id, |
| 142 | + mailto, |
| 143 | + pdf_output, |
| 144 | + prepare_notebook_only, |
| 145 | +): |
| 146 | + if report_name is None: |
| 147 | + raise ValueError("Error! Please provide a --report-name.") |
| 148 | + return execute_notebook_entrypoint( |
| 149 | + config, |
| 150 | + report_name, |
| 151 | + overrides_as_json, |
| 152 | + iterate_override_values_of, |
| 153 | + report_title, |
| 154 | + n_retries, |
| 155 | + job_id, |
| 156 | + mailto, |
| 157 | + pdf_output, |
| 158 | + prepare_notebook_only, |
| 159 | + ) |
| 160 | + |
| 161 | + |
| 162 | +@base_notebooker.command() |
| 163 | +@click.option( |
| 164 | + "--report-name", required=True, help="The name of the template to retrieve, relative to the template directory." |
| 165 | +) |
| 166 | +@pass_config |
| 167 | +def snapshot_latest_successful_notebooks(config: BaseConfig, report_name): |
| 168 | + snap_latest_successful_notebooks(config, report_name) |
| 169 | + |
| 170 | + |
| 171 | +if __name__ == "__main__": |
| 172 | + base_notebooker() |
0 commit comments