def upload_notebook_template(uploadfile_name, name=None, ghe_token=None) -> str: api_client = get_swagger_client() api_instance = swagger_client.NotebookServiceApi(api_client=api_client) try: if ghe_token: notebook: ApiNotebook = api_instance.upload_notebook( uploadfile=uploadfile_name, name=name, enterprise_github_token=ghe_token) else: notebook: ApiNotebook = api_instance.upload_notebook( uploadfile=uploadfile_name, name=name) print(f"Uploaded '{notebook.name}': {notebook.id}") return notebook.id except ApiException as e: print( "Exception when calling NotebookServiceApi -> upload_notebook: %s\n" % e, file=stderr) # raise e return None
def get_template(template_id: str) -> str: api_client = get_swagger_client() api_instance = swagger_client.NotebookServiceApi(api_client=api_client) try: template_response: ApiGetTemplateResponse = api_instance.get_notebook_template( template_id) print(template_response.template) # yaml_dict = yaml.load(template_response.template, Loader=yaml.FullLoader) # notebook_name = yaml_dict.get("name") or f"template_{template_id}" # template_file = os.path.join("files", notebook_name) + ".yaml" # with open(template_file, "w") as f: # f.write(template_response.template) return template_response.template except ApiException as e: print( "Exception when calling NotebookServiceApi -> get_notebook_template: %s\n" % e, file=stderr) return None
def list_notebooks(filter_dict: dict = {}, sort_by: str = None) -> [ApiNotebook]: api_client = get_swagger_client() api_instance = swagger_client.NotebookServiceApi(api_client=api_client) try: filter_str = json.dumps(filter_dict) if filter_dict else None api_response: ApiListNotebooksResponse = api_instance.list_notebooks( filter=filter_str, sort_by=sort_by) for c in api_response.notebooks: print("%s %s %s" % (c.id, c.created_at.strftime("%Y-%m-%d %H:%M:%S"), c.name)) return api_response.notebooks except ApiException as e: print( "Exception when calling NotebookServiceApi -> list_notebooks: %s\n" % e, file=stderr) return []
def delete_assets(upload_assets_response: ApiCatalogUploadResponse = None): api_client = get_swagger_client() delete_methods = { "components": swagger_client.ComponentServiceApi(api_client).delete_component, "datasets": swagger_client.DatasetServiceApi(api_client).delete_dataset, "models": swagger_client.ModelServiceApi(api_client).delete_model, "notebooks": swagger_client.NotebookServiceApi(api_client).delete_notebook, "pipelines": swagger_client.PipelineServiceApi(api_client).delete_pipeline } try: for asset_type, delete_method in delete_methods.items(): if upload_assets_response: asset_list = upload_assets_response.__getattribute__( asset_type) for asset in asset_list: delete_method(asset.id) else: delete_method("*") except ApiException as e: print(f"Exception when calling {delete_method}: {e}\n", file=stderr)
def delete_notebook(notebook_id: str): api_client = get_swagger_client() api_instance = swagger_client.NotebookServiceApi(api_client=api_client) try: api_instance.delete_notebook(notebook_id) except ApiException as e: print( "Exception when calling NotebookServiceApi -> delete_notebook: %s\n" % e, file=stderr)
def set_featured_notebooks(notebook_ids: [str]): api_client = get_swagger_client() api_instance = swagger_client.NotebookServiceApi(api_client=api_client) try: api_response = api_instance.set_featured_notebooks(notebook_ids) except ApiException as e: print( "Exception when calling NotebookServiceApi -> set_featured_notebooks: %s\n" % e, file=stderr) return None
def approve_notebooks_for_publishing(notebook_ids: [str]): api_client = get_swagger_client() api_instance = swagger_client.NotebookServiceApi(api_client=api_client) try: api_response = api_instance.approve_notebooks_for_publishing( notebook_ids) except ApiException as e: print( "Exception when calling NotebookServiceApi -> approve_notebooks_for_publishing: %s\n" % e, file=stderr) return None
def upload_notebook_file(notebook_id, file_path): api_client = get_swagger_client() api_instance = swagger_client.NotebookServiceApi(api_client=api_client) try: response = api_instance.upload_notebook_file(id=notebook_id, uploadfile=file_path) print(f"Upload file '{file_path}' to notebook with ID '{notebook_id}'") except ApiException as e: print( "Exception when calling NotebookServiceApi -> upload_notebook_file: %s\n" % e, file=stderr) raise e
def verify_notebook_download(notebook_id: str) -> bool: api_client = get_swagger_client() api_instance = swagger_client.NotebookServiceApi(api_client=api_client) try: response: HTTPResponse = \ api_instance.download_notebook_files(notebook_id, include_generated_code=True, _preload_content=False) tgz_file = BytesIO(response.read()) tar = tarfile.open(fileobj=tgz_file) file_contents = { m.name.split(".")[-1]: tar.extractfile(m).read().decode("utf-8") for m in tar.getmembers() } template_response: ApiGetTemplateResponse = api_instance.get_notebook_template( notebook_id) template_text_from_api = template_response.template assert template_text_from_api == file_contents.get( "yaml", file_contents.get("yml")) generate_code_response: ApiGenerateCodeResponse = api_instance.generate_notebook_code( notebook_id) run_script_from_api = generate_code_response.script regex = re.compile( r"name='[^']*'" ) # controller adds random chars to name, replace those assert regex.sub("name='...'", run_script_from_api) == \ regex.sub("name='...'", file_contents.get("py")) print("downloaded files match") return True except ApiException as e: print( "Exception when calling NotebookServiceApi -> download_notebook_files: %s\n" % e, file=stderr) return False
def get_notebook(notebook_id: str) -> ApiNotebook: api_client = get_swagger_client() api_instance = swagger_client.NotebookServiceApi(api_client=api_client) try: notebook_meta: ApiNotebook = api_instance.get_notebook(notebook_id) pprint(notebook_meta, indent=2) return notebook_meta except ApiException as e: print( "Exception when calling NotebookServiceApi -> get_notebook: %s\n" % e, file=stderr) return None
def generate_code(notebook_id: str) -> str: api_client = get_swagger_client() api_instance = swagger_client.NotebookServiceApi(api_client=api_client) try: generate_code_response: ApiGenerateCodeResponse = api_instance.generate_notebook_code( notebook_id) print(generate_code_response.script) return generate_code_response.script except ApiException as e: print( "Exception while calling NotebookServiceApi -> generate_code: %s\n" % e, file=stderr) return None
def run_notebook(notebook_id: str, parameters=dict(), run_name: str = None) -> str: api_client = get_swagger_client() api_instance = swagger_client.NotebookServiceApi(api_client=api_client) try: run_code_response: ApiRunCodeResponse = api_instance.run_notebook( notebook_id, parameters=parameters, run_name=run_name) print(run_code_response.run_url) print(run_code_response.run_output_location) return run_code_response.run_url except ApiException as e: print("Exception while calling NotebookServiceApi -> run_code: %s\n" % e, file=stderr) return None
def download_notebook_tgz( notebook_id, download_dir=os.path.join(tempfile.gettempdir(), "download", "notebooks") ) -> str: api_client = get_swagger_client() api_instance = swagger_client.NotebookServiceApi(api_client=api_client) try: response: HTTPResponse = \ api_instance.download_notebook_files(notebook_id, include_generated_code=True, _preload_content=False) attachment_header = response.info().get( "Content-Disposition", f"attachment; filename={notebook_id}.tgz") download_filename = re.sub("attachment; filename=", "", attachment_header) os.makedirs(download_dir, exist_ok=True) tarfile_path = os.path.join(download_dir, download_filename) with open(tarfile_path, 'wb') as f: f.write(response.read()) print(tarfile_path) return tarfile_path except ApiException as e: print( "Exception when calling NotebookServiceApi -> download_notebook_files: %s\n" % e, file=stderr) return "Download failed?"
def generate_custom_pipeline(sequential_only=False) -> ApiPipelineCustom: api_client = get_swagger_client() notebooks_api_instance = swagger_client.NotebookServiceApi( api_client=api_client) components_api_instance = swagger_client.ComponentServiceApi( api_client=api_client) try: components: typing.List[ApiComponent] = random.sample( list( filter(lambda n: "StudyJob" not in n.name, components_api_instance.list_components().components)), 3) # components = [] notebooks: typing.List[ApiNotebook] = random.sample( notebooks_api_instance.list_notebooks().notebooks, 5) tasks: typing.List[ApiPipelineTask] = [] for c in components: tasks.append( ApiPipelineTask(name=f"component task {c.name}", artifact_type="component", artifact_id=c.id, arguments=ApiPipelineTaskArguments( parameters=c.parameters), dependencies=[])) for n in notebooks: tasks.append( ApiPipelineTask(name=f"notebook task {n.name}", artifact_type="notebook", artifact_id=n.id, arguments=ApiPipelineTaskArguments( parameters=n.parameters), dependencies=[])) if sequential_only: for i in range(len(tasks) - 1): tasks[i + 1].dependencies = [tasks[i].name] else: for i in range(len(tasks)): num_deps = random.randint(0, i) deps_idx = {random.randint(0, i) for n in range(0, num_deps)} tasks[i].dependencies = [ tasks[i].name for i in sorted(deps_idx) ] pipeline_params = [] for t in tasks: task_name_prefix = re.sub(r"\W+", "_", t.name, flags=re.ASCII).lower() for p in t.arguments.parameters: if not (p.value or p.default): param_name = re.sub(r"\W+", "_", p.name, flags=re.ASCII).lower() pipeline_param_name = f"{task_name_prefix}_{param_name}" p.value = "{{inputs.parameters." + pipeline_param_name + "}}" pipeline_params.append( ApiParameter(name=pipeline_param_name, value="some value")) api_pipeline_custom = ApiPipelineCustom( name="My custom pipeline", description= "A randomly generated pipeline from notebooks and components", dag=ApiPipelineDAG(tasks=tasks), inputs=ApiPipelineInputs(parameters=pipeline_params)) return api_pipeline_custom except ApiException as e: print( f"Exception while generating custom pipeline DAG with parameters: \n{str(e)}", file=stderr)