在 Databricks / Spark 中记录附加的集群信息

Question

我想对 Databricks 进行一些性能测试。为此，我想记录我在测试期间使用的集群（VM 类型，例如 Standard_DS3_v2）（我们可以假设驱动程序和工作程序节点是相同的）。我知道我可以记录工人数量、核心数量（至少在驱动程序上）和内存（至少在驱动程序上）。但是，我想知道 VM 类型，因为我希望能够识别我是否使用了例如存储优化或通用集群。代替 VM 类型，此信息也可以。最理想的情况是，我可以在笔记本中的变量中将这些信息作为字符串获取，以便稍后将其与我正在记录的其他信息一起从那里写入我的日志文件。但是，如果没有直接的解决方案，我也对任何 hacky 解决方法感到满意。

Answer 1

您可以通过GET request to Clusters API 从 REST API 获取此信息。您可以使用笔记本上下文通过 dbutils.notebook.getContext 调用来识别运行笔记本的集群，该调用返回不同属性的映射，包括集群 ID、工作区域名，并且您可以从中提取身份验证令牌。这是打印驱动程序和工作节点类型的代码（它是用 Python 编写的，但 Scala 代码应该很熟悉 - 我经常使用 Scala 的 dbutils.notebook.getContext.tags 来查找可用的标签）：

import requests
ctx = dbutils.notebook.entry_point.getDbutils().notebook().getContext()
host_name = ctx.tags().get("browserHostName").get()
host_token = ctx.apiToken().get()
cluster_id = ctx.tags().get("clusterId").get()

response = requests.get(
    f'https://{host_name}/api/2.0/clusters/get?cluster_id={cluster_id}',
    headers={'Authorization': f'Bearer {host_token}'}
  ).json()
print(f"driver type={response['driver_node_type_id']} worker type={response['node_type_id']}")