云加速器诊断

概述

云加速器诊断是一个库，用于监控、调试和性能分析在TPU和GPU等云加速器上运行的工作负载。此外，此库还提供了一种简化的方法，自动将数据上传到Vertex AI的Tensorboard实验。该包允许用户在Vertex AI中创建Tensorboard实例和实验，并将日志上传到它们。

安装

要安装云加速器诊断包，请运行以下命令

pip install cloud-accelerator-diagnostics

自动上传到Vertex AI Tensorboard

在创建并上传日志到Vertex AI Tensorboard之前，您必须在Google Cloud控制台中启用Vertex AI API。同时，确保将Vertex AI用户IAM角色分配给将在“cloud-accelerator-diagnostics”包中调用API的服务帐户。这是在Google Cloud控制台中创建和访问Vertex AI Tensorboard所必需的。

创建Vertex AI Tensorboard

要了解Vertex AI Tensorboard，请访问这个页面。

以下是一个示例脚本，用于在Google Cloud项目test-project中创建名为test-instance的Vertex AI Tensorboard实例。

注意：Vertex AI仅在这些地区可用。

from cloud_accelerator_diagnostics import tensorboard

instance_id = tensorboard.create_instance(project="test-project",
                                          location="us-central1",
                                          tensorboard_name="test-instance")
print("Vertex AI Tensorboard created: ", instance_id)

创建Vertex AI实验

要了解Vertex AI实验，请访问这个页面。

以下脚本将在您的Google Cloud项目test-project中创建一个名为test-experiment的Vertex AI实验。以下是它如何处理附加Tensorboard实例的过程

场景1：Tensorboard实例存在

如果您的项目中已经存在名为test-instance的Tensorboard实例，脚本将将其附加到新实验。

场景2：没有Tensorboard实例

如果test-instance不存在，脚本将创建一个新的具有该名称的Tensorboard实例并将其附加到实验。

from cloud_accelerator_diagnostics import tensorboard

instance_id, tensorboard_url = tensorboard.create_experiment(project="test-project",
                                                             location="us-central1",
                                                             experiment_name="test-experiment",
                                                             tensorboard_name="test-instance")

print("View your Vertex AI Tensorboard here: ", tensorboard_url)

如果已存在指定名称的Vertex AI实验，则不会创建新的实验，而是返回现有实验的URL。

注意：您可以将多个Vertex AI实验附加到单个Vertex AI Tensorboard。

上传日志到Vertex AI Tensorboard

以下脚本将连续监控目录（logdir）中的新数据，并将其上传到您的Vertex AI Tensorboard实验。请注意，在调用start_upload_to_tensorboard()后，即使在抛出异常的情况下，线程也将保持活动状态。为确保线程关闭，请将任何在start_upload_to_tensorboard()之后和stop_upload_to_tensorboard()之前的代码放在一个try块中，并在finally块中调用stop_upload_to_tensorboard()。此示例展示了您如何上传收集的用于Vertex AI Tensorboard的JAX工作负载的性能日志。

from cloud_accelerator_diagnostics import uploader

uploader.start_upload_to_tensorboard(project="test-project",
                                     location="us-central1",
                                     experiment_name="test-experiment",
                                     tensorboard_name="test-instance",
                                     logdir="gs://test-directory/testing")
try:
  jax.profiler.start_trace("gs://test-directory/testing")
  <your code goes here>
  jax.profiler.stop_trace()
finally:
  uploader.stop_upload_to_tensorboard()