连接到 AI 运行时

重要

此功能目前以公共预览版提供。

从交互式笔记本、通过 SSH 隧道的 IDE、计划作业、作业 API 或声明式自动化包连接到 AI 运行时。 将笔记本附加到 AI Runtime 是运行训练和微调工作负载的主要方式，你还可以将这些笔记本安排为定期运行的作业，或在部署流水线中将其自动化。

交互式（笔记本）

这是使用 AI 运行时的主要方法。 连接笔记本并配置环境：

  1. 在笔记本中，单击顶部的计算下拉菜单，然后选择 “无服务器 GPU”。
  2. 单击 “环境”图标。 打开 “环境 ”侧面板。
  3. 从 “加速器 ”字段中选择一个加速器。 对于分布式训练工作负荷，请选择 8xH100。 有关选择加速器的指导，请参阅 硬件选项 。
  4. 从“基础环境”字段中，为标准环境选择标准 v5 或标准 v4，或者为 AI 环境选择 AI v5 或 AI v4。
  5. 单击“ 应用 ”，然后 确认 要将 AI 运行时应用到笔记本环境。

注释

计算连接将在非活动状态持续60分钟后自动终止。

小窍门

对于不需要 GPU 的操作（例如克隆 Git 存储库、转换数据格式或探索性数据分析），请将笔记本附加到 CPU 群集以保留 GPU 资源。

从 IDE 终端进行连接

可以通过 SSH 隧道直接从 IDE 中的终端连接到无服务器 GPU 计算上的 AI 运行时。

要连接到 AI 运行时，从 IDE 内的终端运行带有 databricks ssh connect 选项的 --accelerator 命令。 不需要单独的安装步骤。 有关命令的详细信息，请参阅 ssh 命令组。

databricks ssh connect --accelerator=GPU_1xA10

若要在 Visual Studio Code 或 Cursor 中连接并启动会话，请使用--ide此选项。 CLI 将打开指向主工作区文件夹的 IDE 窗口。

databricks ssh connect --ide=vscode

有关设置、打开项目和运行代码的更多详细信息，请参阅 使用 SSH 隧道连接到 Databricks。

作业（计划）

您可以将使用 AI 运行时的笔记本安排为周期性作业。 有关更多详细信息，请参阅 “创建和管理计划笔记本作业 ”。

打开要使用的笔记本后：

  1. 选择右上角的 “计划 ”按钮。
  2. 选择 “添加计划”。
  3. 在新建计划表单中填写作业名称、计划和计算。
  4. 选择 创建。

还可以从 作业和管道 UI 创建和计划作业。 有关分步指南，请参阅 “创建新作业 ”。

注释

AI 运行时的计划任务不支持通过“环境”面板添加依赖项。 必须在笔记本中以编程方式安装依赖项（例如）。 %pip install 不支持自动恢复。 如果作业由于包不兼容而失败，则必须手动修复并重新运行。

对于可能超过 7 天最大运行时间的工作负载，请实施手动检查点功能以支持恢复。 建议通过 UCVolumeWriter 中的 UCVolumeReader 和 serverless_gpu.data 使用 Unity Catalog 卷。 请参阅模型检查点。

作业 API 和声明式自动化包

可以使用 Databricks 作业 API 或 声明性自动化捆绑包以编程方式创建和管理 AI 运行时作业。 在作业或捆绑定义中将计算类型配置为无服务器 GPU，以自动执行部署管道。

以下示例演示了使用 标准环境的 AI 运行时作业的声明性自动化捆绑配置：

resources:
  jobs:
    sample_job:
      name: sample_job_h100

      trigger:
        periodic:
          interval: 1
          unit: DAYS

      parameters:
        - name: catalog
          default: ${var.catalog}
        - name: schema
          default: ${var.schema}

      environments:
        - environment_key: default
          spec:
            environment_version: '4'

      tasks:
        - task_key: notebook_task
          notebook_task:
            notebook_path: /Workspace/Users/your_email/your_notebook
          environment_key: default
          compute:
            hardware_accelerator: GPU_8xH100

若要使用 Databricks AI environment 而不是 Standard 环境，请将环境 base_environment 中的 databricks_ai_v5 设置为 AI 环境标识符（例如，AI v5 的 spec），并在任务的 environment_key 中引用它：

重要

将 Databricks AI 环境选作工作区基础环境目前处于 Beta 阶段，并且需要由工作区管理员启用。 请参阅 构建无服务器 GPU 计算（AI 运行时）。

resources:
  jobs:
    sample_job:
      name: sample_job_aiv5_h100

      trigger:
        periodic:
          interval: 1
          unit: DAYS

      parameters:
        - name: catalog
          default: ${var.catalog}
        - name: schema
          default: ${var.schema}

      environments:
        - environment_key: aiv5
          spec:
            base_environment: databricks_ai_v5

      tasks:
        - task_key: notebook_task
          notebook_task:
            notebook_path: /Workspace/Users/your_email/your_notebook
          environment_key: aiv5
          compute:
            hardware_accelerator: GPU_8xH100
  • Last updated on