如何通过REST API从本地Python调用Databricks Notebook并获取输出

我想在本地写Python脚本，给Databricks里的Notebook传参数，然后获取它的执行输出（比如本地传字符串，Notebook反转后返回结果）。现在调用API能收到响应，但返回的metadata里找不到预期的notebook_output字段，请问缺什么配置？

现有代码 #

import os
from databricks_cli.sdk.api_client import ApiClient
from databricks_cli.runs.api import RunsApi  # 原代码遗漏此导入

os.environ['DATABRICKS_HOST'] = "https://adb-################.##.azuredatabricks.net/"
os.environ['DATABRICKS_TOKEN'] = "token-value" 

api_client = ApiClient(host=os.getenv('DATABRICKS_HOST'), token=os.getenv('DATABRICKS_TOKEN'))

runJson = """
{
    "name": "test job",
    "max_concurrent_runs": 1,
    "tasks": [
        {
            "task_key": "test",
            "description": "test",
            "notebook_task":
                {
                    "notebook_path": "/Users/user@domain.com/api_test"
                },
            "existing_cluster_id": "cluster_name",
            "timeout_seconds": 3600,
            "max_retries": 3,
            "retry_on_timeout": true
        }
    ]
}
"""

runs_api = RunsApi(api_client)
run_id = runs_api.submit_run(runJson)
metadata = runs_api.get_run_output(run_id['run_id'])['metadata']

当前API响应 #

{'job_id': 398029273095601, 'run_id': 150609942, 'creator_user_name': 'user', 'number_in_job': 150609942, 'state': {'life_cycle_state': 'TERMINATED', 'result_state': 'SUCCESS', 'state_message': '', 'user_cancelled_or_timedout': False}, 
'task': {'notebook_task': {'notebook_path': 'path', 'source': 'WORKSPACE'}}, 
'cluster_spec': {'existing_cluster_id': 'cluster'}, 
'cluster_instance': {'cluster_id': 'id', 'spark_context_id': 'id'}, 'start_time': 1683904971067, 'setup_duration': 1000, 'execution_duration': 8000, 'cleanup_duration': 0, 'end_time': 1683904981007, 'run_duration': 9940, 'run_name': 
'Untitled', 'run_page_url': 'url', 'run_type': 'SUBMIT_RUN', 'attempt_number': 0, 'format': 'SINGLE_TASK'}

解决方法 #

1. Notebook必须显式返回结果 #

Databricks不会自动把Notebook的所有输出返回给API，得用dbutils.notebook.exit()方法把要返回的结果传出去。比如你的反转字符串Notebook要这么写：

# 获取本地传入的参数
input_str = dbutils.widgets.get("input_str")
# 执行反转逻辑
reversed_str = input_str[::-1]
# 把结果返回给API调用方
dbutils.notebook.exit(reversed_str)

2. 提交任务时要传入参数 #

原代码没给Notebook传参数，得在notebook_task里加上base_parameters字段，把要传的参数放进去：
修改后的runJson中notebook_task部分：

"notebook_task": {
    "notebook_path": "/Users/user@domain.com/api_test",
    "base_parameters": {
        "input_str": "hello world"  # 这里填你要传的字符串
    }
}

3. 正确获取输出的位置 #

get_run_output返回的结果里，notebook_output不在metadata里，而是在根层级的output字段下。修改获取输出的代码：

run_output = runs_api.get_run_output(run_id['run_id'])
# 提取Notebook返回的结果
notebook_result = run_output.get('output', {}).get('notebook_output', {}).get('result')
print("Notebook返回结果：", notebook_result)

修改后的完整代码 #

import os
import time
from databricks_cli.sdk.api_client import ApiClient
from databricks_cli.runs.api import RunsApi

os.environ['DATABRICKS_HOST'] = "https://adb-################.##.azuredatabricks.net/"
os.environ['DATABRICKS_TOKEN'] = "token-value" 

api_client = ApiClient(host=os.getenv('DATABRICKS_HOST'), token=os.getenv('DATABRICKS_TOKEN'))

runJson = """
{
    "name": "test job",
    "max_concurrent_runs": 1,
    "tasks": [
        {
            "task_key": "test",
            "description": "test",
            "notebook_task":
                {
                    "notebook_path": "/Users/user@domain.com/api_test",
                    "base_parameters": {
                        "input_str": "hello world"
                    }
                },
            "existing_cluster_id": "cluster_name",
            "timeout_seconds": 3600,
            "max_retries": 3,
            "retry_on_timeout": true
        }
    ]
}
"""

runs_api = RunsApi(api_client)
run_id = runs_api.submit_run(runJson)

# 等待任务执行完成，避免任务未结束时获取不到输出
while True:
    run_status = runs_api.get_run(run_id['run_id'])['state']['life_cycle_state']
    if run_status == 'TERMINATED':
        break
    time.sleep(5)

# 获取并打印Notebook返回结果
run_output = runs_api.get_run_output(run_id['run_id'])
notebook_result = run_output.get('output', {}).get('notebook_output', {}).get('result')
print("Notebook返回结果：", notebook_result)

内容的提问来源于stack exchange，提问作者J. Doe