使用 AI 容器镜像

AI 容器镜像是火山引擎提供的一系列 Docker 镜像，其中预装了常见的 AI 框架、工具以及相关依赖，为用户提供快速启动且高度一致的 AI 开发环境。本文为您介绍 AI 容器镜像的使用方式。

背景信息

在 AI 模型训练和推理的前置流程中，用户通常基于特定的AI框架和工具的组合开展工作（例如：PyTorch 框架、deepspeed 库），在基础环境中进行算法调优、模型注入或者脚本开发后，将重新打包好的镜像提交到训练或者推理平台的任务队列中。

AI 容器镜像用于加速用户 AI 开发、训练和推理环境的搭建，并与 AI 套件中其他工具和产品能力共同为用户提供覆盖 AI 应用完整生命周期的服务。

前提条件

已经基于 VKE 搭建 AI 训练或推理任务的基础环境，例如：集群、VPC、GPU 节点等。

使用限制

VKE AI 容器镜像限定只能在 VPC 内拉取镜像，暂不支持公网拉取。请在火山引擎对应地域的 VKE 集群或者 ECS 实例中拉取。

操作步骤

步骤一：获取 AI 预置镜像地址

火山引擎预置提供一系列 AI 容器镜像，为用户提供快速启动且高度一致的AI开发环境。以下是 AI 容器镜像支持的版本、支持状态以及镜像地址。

步骤二：拉取 AI 预置镜像启动容器

按需获取 AI 预置容器镜像，并拉取到 VKE 中启动容器镜像，使用此 AI 容器镜像创建基础训练环境。

1. 登录 容器服务控制台。
2. 在左侧导航栏选择 集群，单击可用于 AI 训练或推理的目标集群名称，进入集群管理页面。
3. 在左侧导航栏选择 工作负载 > 容器组，进入容器组管理页面。
4. 单击 使用 YAML创建，使用 Yaml 方式拉取 AI 预置镜像启动容器。Yaml 示例如下：

apiVersion: v1
kind: Pod
metadata:
  name: ai-env         # 自定义 Pod 名称。
  namespace: defualt   # 按需配置 Pod 所属的命名空间。
spec:
  containers:
  - name: pytorch      # 自定义容器名称。
    image: ai-containers-cn-beijing.cr.volces.com/deeplearning/torch2.0.1-cuda11.8.0-py3.10-ubuntu22.04:latest  # AI 预置镜像及其版本。从步骤一中按需获取预置镜像地址，镜像版本使用 latest 即可。
    command: [ "/bin/bash", "-c", "--" ]
    args: [ "while true; do sleep 3600; done;" ]
    resources:
      limits:
        nvidia.com/gpu: "1" # 按需配置 GPU 显卡数量。

步骤三：开发训练环境

Pod 进入 running 状态后，即可通过 VKE 登录 Pod 控制台，进一步进行环境准备和开发工作，具体内容可根据实际使用场景进行开发。例如：通过开发工具进行算法、脚本开发/置入模型文件。

步骤四：打包镜像并推送镜像仓库

1. 以 AI 容器镜像为基础镜像，通过 dockerfile 构建定制化的 AI 开发环境。dockerfile 示例如下：

# 所使用预置镜像的地址和版本，与步骤二中拉取的 AI 预置镜像保持一致。
FROM ai-containers-cn-beijing.cr.volces.com/deeplearning/torch2.0.1-cuda11.8.0-py3.10-ubuntu22.04:latest
# 为 Dockerfile 中后续命令设置工作目录。
WORKDIR /home
# 从本地文件系统中导入文件到镜像，例如：模型文件。
COPY chatglm2-6b.tar /chatglm2-6b.tar
# 为镜像预装其他所需的 AI 工具软件，例如：jupyterlab。
RUN pip install "jupyterlab"
# 为镜像环境添加其他所需的环境变量。
ARG VARNAME = "value"
ENV PATH=$VARNAME

2. 将构建完成的派生镜像推送到火山引擎镜像仓库服务中，详细介绍参见：使用标准版实例拉取和推送镜像。

后续操作

至此，以 AI 容器镜像为基础打包的派生镜像已经推送到火山引擎镜像仓库，您可通过 VKE 部署模型训练或推理任务。