Hugging Face 的 LOGO

dstack 轻松管理本地 AI 工作负载服务器集群

社区文章 发布于2024年10月10日
赞
7
chansung park's avatar
朴赞成
chansung

image/png

如果您还不了解 dstack 是什么，请参阅这篇文章和官方文档以掌握 dstack 的基本概念。简而言之，dstack 是一个专注于人工智能开发、训练和部署的计算资源管理工具包。

最初，dstack 提供了一种很好的方式来管理和控制来自各种云服务的多台机器，包括GCP (Google Cloud Platform)、AWS (Amazon Web Service)、Microsoft Azure、OCI (Oracle Cloud Infrastructure)、Lambda Labs、RunPod、Vast.ai、DataCrunch 和CUDO，并支持 CPU、NVIDIA GPU、AMD GPU 和TPU。这很有意义，因为您可以找到最适合您要求（规格、成本等）的最佳资源（机器），然后以统一的方式指导来自不同来源的任何机器。

从 0.18.7 版本发布后，dstack 已发展到不仅可以通过 ssh-fleet 功能管理云资源，还可以管理本地资源。此功能的最佳之处在于，您无需了解 Kubernetes 或 Slurm 的任何知识，并且它以最少的依赖项（几乎）纯粹的 Docker 技术工作。以下是 ssh-fleet 的一些优点：

  • 易于设置（无需 Kubernetes，无需 Slurm）
    • 为了设置 Kubernetes 或 Slurm，您需要学习大量的先验知识。此外，实际设置、运行和管理它们需要大量的工程工作。对于 dstack 的 ssh-fleet，除了您已有的诸如安装 CUDA 和 Docker 等知识外，几乎不需要了解其他任何东西。
  • 将分散的本地机器组建成集群
    • 并非所有组织都拥有专用的本地云计算基础设施。许多实验室根据项目管理自己的计算资源。然而，如今我们正在处理越来越大的机器学习模型，例如大型语言模型 (LLM)，并且这通常需要多节点协作。对于 dstack 的 ssh-fleet，您可以简单地将多个资源作为一个集群进行管理，然后分配单节点或多节点设置的作业。
  • 云和本地资源的集中管理
    • 机器学习就是运行大量实验以找出最适合您问题的模型。这意味着应该有多个实验并行运行。否则，您需要花费太多时间。使用 dstack，您可以将更多实验分配给云资源，同时保持您的本地资源忙碌。

现在，让我们通过基本教程来了解如何使用 dstack 设置您自己的 ssh-fleet。

ssh-fleet 的先决条件

在远程服务器端

  1. 安装 Docker
  • Docker 提供 dstack 赖以封装您的应用程序及其依赖项的容器化技术。它确保了不同环境之间的一致性和可重复性。
  • 请遵循官方 Docker 安装说明，根据您的 Linux 发行版进行安装。这通常涉及添加 Docker 的存储库，然后使用您的包管理器（apt、yum 等）安装 docker-ce 包。成功安装后，您可以使用以下命令验证 Docker 是否正在运行。它应该在终端中打印出 Hello from Docker 消息。
$ sudo docker run hello-world
  1. 安装 CUDA 工具包 >= 12.1
  • 如果您计划将 NVIDIA GPU 用于您的 AI 工作负载，CUDA 工具包是必不可少的。它提供了应用程序利用 GPU 处理能力所需的库和工具。dstack 需要 CUDA 12.1 或更高版本才能兼容并利用最新功能。
  • 从 NVIDIA 网站下载 CUDA 工具包安装程序，并按照安装说明进行操作。请务必为您的 Linux 发行版和系统架构选择正确的版本。
  1. 安装 CUDA 容器工具包
  • CUDA 容器工具包允许 Docker 容器访问和利用 NVIDIA GPU。这对于在 dstack 中运行 GPU 加速的 AI 工作负载至关重要。
  • 同样，请参阅NVIDIA 官方文档。您通常需要添加 NVIDIA 的容器工具包存储库，然后使用您的包管理器安装 nvidia-container-toolkit 包。

如果您是 AMD GPU 用户，请跳过步骤 2 和 3，转而按照发行说明安装 AMD 特定驱动程序。

  1. sudo visudo for username ALL=(ALL) NOPASSWD: ALL
  • 此配置允许 dstack 服务器在远程服务器上执行命令，而无需密码。这对于 dstack 自动代表您管理容器和资源是必需的。值得注意的是，这会授予指定用户重要的权限。请确保此用户专用于 dstack 操作并采取适当的安全措施。
  • 使用以下命令通过 sudo visudo 打开 /etc/sudoers 文件。添加行 username ALL=(ALL) NOPASSWD: ALL，将 username 替换为您将用于连接远程服务器的实际用户名。在此之后，username 可以在 sudo 模式下运行任何命令，而无需输入密码提示。
$ sudo visudo

在本地

  1. 生成 id_rsa
  • SSH 密钥提供了一种安全的方式来验证您的远程服务器，而无需每次都输入密码。dstack 使用这些密钥来建立与您的本地机器的安全连接，以实现自动化集群管理。
  • 在本地机器上使用 ssh-keygen 命令生成 SSH 密钥对，如下所示。这将创建一个私钥 (id_rsa) 和一个公钥 (id_rsa.pub)。
$ ssh-keygen -t rsa
  1. ssh-copy-id
  • 此步骤允许您的本地机器使用 SSH 密钥对自动向远程服务器进行身份验证，从而简化连接过程，并使 dstack 能够在无需手动干预的情况下管理远程服务器。
  • 在您的本地机器上运行 ssh-copy-id username@remote_host，如下所示，将 username 和 remote_host 替换为适当的值。此命令会将您的公钥复制到远程服务器的 authorized_keys 文件中。在此之后，您可以直接 ssh 连接到远程服务器，而无需在提示符下输入密码。
$ ssh-copy-id username@remote_host

在本地安装 dstack 并注册 ssh 集群

  1. 安装 dstack

使用 pip 安装 dstack 及其所有可选依赖项。如果您只想使用 dstack 管理本地集群，则无需指定 [all]，但当您想同时管理本地和所有其他云资源时，[all] 会很有用。

$ pip install "dstack[all]"
  1. 运行 dstack 服务器

在本地机器上启动 dstack 服务器。dstack 服务器是核心组件，用于管理资源、调度作业以及处理本地机器与计算资源（云端和本地）之间的通信。

$ dstack server
  1. 编写 fleet.dstack.yml

定义一个 ssh-fleet 的 YAML 文件，如下所示。您可以进行许多配置（请参阅dstack 官方 API 文档），但以下显示了必需项。您可以按照本博客文章中的“ssh-fleet 的先决条件”部分为要用于 ssh-fleet 集群的每个服务器进行设置。例如，以下 YAML 文件显示我已注册了 4 个服务器（2 个带有 3xRTX6000 Ada，2 个带有 2xA6000）。另请注意，它指向我们从“本地”部分生成的 rsa 文件。

type: fleet
# The name is optional, if not specified, generated randomly
name: my-ssh-fleet

# Ensure instances are interconnected
placement: cluster

# The user, private SSH key, and hostnames of the on-prem servers
ssh_config:
  user: username
  identity_file: ~/.ssh/id_rsa
  hosts:
    - xxx.xxx.171.224
    - xxx.xxx.171.225
    - xxx.xxx.164.172
    - xxx.xxx.165.51

请注意，placement: cluster 意味着要确保实例（服务器）像共享同一网络一样相互连接。如果列出的实例不共享同一网络，则 ssh-fleet 预配将失败。但是，如果它们共享，并且设置了 placement: cluster，则可以运行多节点作业，例如分布式 AI 模型训练。

  1. 应用 fleet.dstack.yml

告诉 dstack 读取 fleet.dstack.yml 文件并根据您的配置创建 ssh-fleet。dstack 将尝试使用提供的 SSH 凭据连接到每个指定的主机。

$ dstack apply -f fleet.dstack.yml

列出 dstack 设置中可用的舰队。您应该看到您的 my-ssh-fleet 已列出，其中包含有关连接实例（服务器）、其资源（CPU、内存、GPU、磁盘）和其当前状态的详细信息。

$ dstack fleet
 FLEET         INSTANCE  BACKEND       RESOURCES                                            PRICE  STATUS  CREATED
 my-ssh-fleet  1         ssh (remote)  32xCPU, 503GB, 3xRTX6000Ada (48GB), 1555.1GB (disk)  $0.0   idle    2 weeks ago
               2         ssh (remote)  32xCPU, 503GB, 3xRTX6000Ada (48GB), 1555.1GB (disk)  $0.0   idle    2 weeks ago
               3         ssh (remote)  64xCPU, 693GB, 2xA6000 (48GB), 1683.6GB (disk)       $0.0   idle    2 weeks ago
               4         ssh (remote)  64xCPU, 693GB, 2xA6000 (48GB), 1683.6GB (disk)       $0.0   idle    2 weeks ago

此外，在您运行 dstack server 的终端中，您应该看到类似下面的日志，这表明 dstack 已成功找到并建立了与所列服务器的连接。

[08:24:07] INFO     dstack._internal.server.background.tasks.process_instances:190 Adding ssh instance my-ssh-fleet-0...
           INFO     dstack._internal.server.background.tasks.process_instances:325 Connected to user xxx.xxx.171.224
[08:24:13] INFO     dstack._internal.server.background.tasks.process_instances:190 Adding ssh instance my-ssh-fleet-1...
           INFO     dstack._internal.server.background.tasks.process_instances:325 Connected to user xxx.xxx.171.225
[08:24:17] INFO     dstack._internal.server.background.tasks.process_instances:190 Adding ssh instance my-ssh-fleet-2...
[08:24:18] INFO     dstack._internal.server.background.tasks.process_instances:325 Connected to user xxx.xxx.164.172
[08:24:23] INFO     dstack._internal.server.background.tasks.process_instances:190 Adding ssh instance my-ssh-fleet-3...
           INFO     dstack._internal.server.background.tasks.process_instances:325 Connected to user xxx.xxx.165.51
[08:24:41] INFO     dstack._internal.server.background.tasks.process_instances:245 The instance my-ssh-fleet-0 (xxx.xxx.171.224) was successfully added
[08:24:42] INFO     dstack._internal.server.background.tasks.process_instances:245 The instance my-ssh-fleet-3 (xxx.xxx.165.51) was successfully added
[08:24:45] INFO     dstack._internal.server.background.tasks.process_instances:245 The instance my-ssh-fleet-1 (xxx.xxx.171.225) was successfully added
[08:24:57] INFO     dstack._internal.server.background.tasks.process_instances:245 The instance my-ssh-fleet-2 (xxx.xxx.164.172) was successfully added
  1. 编写 task.dstack.yml

为了测试，我编写了一个简单的 dstack task YAML 文件，如下所示，用于使用 Hugging Face 的 Alignment Handbook 框架定义一个 LLM 微调作业。请注意，我请求了 2 个节点，每个节点都有 3 x RTX6000Ada GPU。

type: task

nodes: 2
python: "3.11"
nvcc: true

env:
  - HUGGING_FACE_HUB_TOKEN
  - WANDB_API_KEY
  - ACCELERATE_LOG_LEVEL=info

commands:
  - cd alignment-handbook
  - python -m pip install .
  - python -m pip install flash-attn --no-build-isolation

  - pip install wandb
  - pip install huggingface-hub==0.24.7

  - accelerate launch
    --config_file recipes/accelerate_configs/multi_gpu.yaml
    --main_process_ip=$DSTACK_MASTER_NODE_IP
    --main_process_port=8008
    --machine_rank=$DSTACK_NODE_RANK
    --num_processes=$DSTACK_GPUS_NUM
    --num_machines=$DSTACK_NODES_NUM
    scripts/run_sft.py
    recipes/custom.yaml

ports:
  - 50002

resources:
  gpu:
    name: rtx6000ada
    memory: 48GB
    count: 3
  shm_size: 24GB

dstack 允许用户定义三种不同类型的作业。**开发环境**允许您使用您的代码、依赖项和资源配置远程机器，并使用您的桌面 IDE 访问它。**任务**允许您调度作业或运行 Web 应用程序。它允许您配置依赖项、资源、端口等。任务可以分布式并运行在集群上。**服务**允许您将 Web 应用程序或模型部署为可扩展的端点。它允许您配置依赖项、资源、授权、自动伸缩规则等。

目前 dstack 版本 (0.18.17) 的本地环境不支持 **service**，因为它需要网关，但此要求将在未来的版本中解除。

  1. 应用 task.dstack.yml

使用 dstack apply -f 命令应用先前编写的 task.dstack.yml，如下所示。然后它将显示注册的目标服务器以配置作业。当您在提示符处输入 y 时，微调作业将启动。

$ dstack apply -f task.dstack.yml

 Configuration          train.dstack.yml
 Project                main
 User                   admin
 Pool                   default-pool
 Min resources          2..xCPU, 8GB.., 2xGPU (48GB), 100GB.. (disk)
 Max price              -
 Max duration           72h
 Spot policy            on-demand
 Retry policy           no
 Creation policy        reuse-or-create
 Termination policy     destroy-after-idle
 Termination idle time  5m

 #  BACKEND  REGION  INSTANCE  RESOURCES                                            SPOT  PRICE
 1  ssh      remote  instance  32xCPU, 503GB, 3xRTX6000Ada (48GB), 1555.1GB (disk)  no    $0     idle
 2  ssh      remote  instance  32xCPU, 503GB, 3xRTX6000Ada (48GB), 1555.1GB (disk)  no    $0     idle

（额外）同时注册其他云服务

现在，我们已经使用 dstack 的 ssh-fleet 将本地服务器注册为一个集群。但是，您可能希望同时受益于云服务。例如，如果您有多个微调实验要运行，这可能特别有用。在这种情况下，您可以将一些实验分配给本地集群，而将其他实验分配给云服务。这将显著减少时间消耗，同时最大限度地降低成本支出。

为此，只需遵循 dstack 官方文档中关于 server/config.yml 的内容，即可添加您喜欢的云服务。例如，可以使用 gcloud CLI 工具包通过 application default 凭据添加 GCP 后端。或者，您可以使用 service account 凭据更精细地控制 GCP 后端。

在同时拥有本地集群和云服务作为后端之后，dstack apply 命令默认会尝试从云服务中找到合适的实例。当您想要在本地集群上配置作业时，请附加 --backend remote 选项。

# to target cloud service
$ dstack apply -f task.dstack.yml

# to target on-prem cluster
$ dstack apply -f task.dstack.yml --backend remote

总结

dstack 的 ssh-fleet 功能为管理 AI 工作负载的本地集群提供了简化的方法。通过简化设置和集中控制，dstack 使 AI 实践者能够高效地利用其本地资源，无论是用于大规模模型训练还是运行多个实验。与云服务的无缝集成进一步增强了灵活性和可扩展性，使 dstack 成为现代 AI 开发中的宝贵工具。

随着 dstack 的不断发展，我们可以期待更多强大的功能以及对各种硬件和软件配置的更广泛支持。这种持续发展有望进一步巩固 dstack 作为管理本地和云端 AI 基础设施的多功能且不可或缺的工具的地位。

社区

注册或登录发表评论

赞
7