Logging and catching errors within JobsAPI

[jmeidam]

We currently have a notebook that executes other notebooks in parallel (number of notebooks depends on runtime parameters).
The current solution is to execute dbutils.notebooks.run within ThreadPoolExecutor. I'm worried that this might not make optimal use of Databricks cluster resources and would like to use native scheduling functionality.

In the example within this repo the JobsAPI is used to execute a dynamically generated job, which seems nice.

My question is the following: How can I

1. reduce the amount of output (for example, only log when status changes, only show job url once),
2. catch exceptions raised within the tasks.

I essentially want to execute the tasks in parallel, wait until all of them are either finished, cancelled or failed and then gather exceptions if any were raised and further deal with them in the caller notebook.

I am now using this bit of code in my caller notebook to run some tests:

import json
import os
import time
from concurrent.futures import ThreadPoolExecutor, wait
from dataclasses import dataclass
from datetime import timedelta
from json.decoder import JSONDecodeError
from typing import Any, Dict, List, Optional

from databricks.sdk import WorkspaceClient
from databricks.sdk.service import jobs
from py4j.protocol import Py4JJavaError


@dataclass
class NotebookRunConfiguration:
    path: str
    name: str
    timeout: int = 21600  # six hours
    parameters: Optional[Dict[str, Any]] = None
    
def run_notebooks_in_parallel_2(
    client: WorkspaceClient,
    cluster_id: str,
    notebook_run_configs: List[NotebookRunConfiguration],
    timeout: int = 21600  # six hours
):
    waiter = client.jobs.submit(
        run_name=f'py-sdk-run-{time.time()}',
        tasks=[
            jobs.SubmitTask(
                task_key=f'process_{cfg.name}',
                existing_cluster_id=cluster_id,
                notebook_task=jobs.NotebookTask(
                    notebook_path=cfg.path,
                    base_parameters=cfg.parameters,
                )
            )
            for cfg in notebook_run_configs
        ],
        timeout_seconds=timeout
    )

    # callback, that receives a polled entity between state updates
    def print_status(run: jobs.Run):
        tasks = run.tasks if run.tasks else []
        try:
            print(f"run_url: {run.run_page_url}")
            for task in tasks:
                if task.state:
                    if str(task.state.life_cycle_state) == "RunLifeCycleState.RUNNING":
                        print(f'{task.task_key} is running')
                    elif str(task.state.life_cycle_state) == "RunLifeCycleState.INTERNAL_ERROR":
                        print(f"{task.task_key} has an error")
        except Exception as e:
            print(f"so this happened: {repr(e)}")


    run = waiter.result(timeout=timedelta(minutes=timeout/60), callback=print_status)
    print(f"job finished: {run.run_page_url}")
    
current_path = os.path.dirname(dbutils.notebook.entry_point.getDbutils().notebook().getContext().notebookPath().get())
parameters = [
    {"alpha": "0.5", "raise_exception": "false"},
    {"alpha": "0.6", "raise_exception": "true"},
    {"alpha": "0.7", "raise_exception": "false"},
    {"alpha": "0.8", "raise_exception": "true"},
]
notebook_configs = [
    NotebookRunConfiguration(
        name=f"process_{i}",
        path=f"{current_path}/notebook_test",
        timeout=4*60,
        parameters=params,
    ) for i, params in enumerate(parameters)
]

w = WorkspaceClient()
cluster_id = spark.conf.get("spark.databricks.clusterUsageTags.clusterId") 
run_notebooks_in_parallel_2(
    w,
    cluster_id,
    notebook_configs,
)

The notebook that is called contains

alpha = dbutils.widgets.get("alpha")
raise_exception = dbutils.widgets.get("raise_exception").lower()

if raise_exception == "true":
    raise TypeError(f"This notebook gives a type error for alpha {alpha}")
else:
    dbutils.notebook.exit(f'{{"type": "success", "alpha": "{alpha}"}}')