使用Karpenter通过时间切片管理GPU节点

​
 

导语

在机器学习领域，我们经常使用 GPU 来加速计算工作负载。但现在的企业和开发者都更热衷于“上云”。有了云计算，使用云服务，用多少付多少，也就能降低运营成本了。

当你拥有数十个在不同时间段需要用到 GPU 的应用程序时，怎么以更低成本，怎么更灵活地在云服务器中调度资源，就会变成一件非常重要的事情。 

作者介绍

Jina AI 云架构研发工程师陶然

问题

那么，如何优化云服务中 GPU 的使用成本呢？在使用虚拟机时，哪怕你不需要全天候的服务，你也必须持续为所有的设备付费。相比于虚拟机，容器拥有更高的资源使用效率，作为容器界的扛把子，kubernetes 提供了弹性的节点缩放方式。

因为我使用的是 Amazon EKS，所以本文选择了 Karpenter 作为节点缩放器。Karpenter 是一个为 Kubernetes 构建的开源自动扩缩容项目，你可以通过此 文档[1] 了解更多关于 Karpenter 的信息。

想要管理多个 GPU 节点，还需要用到 NVIDIA 的 k8s 插件[2]。这是一个 Daemonset(守护进程)，提供了以下自动化的功能：

• 公开集群每个节点上的 GPU 数量


• 实时追踪 GPU 的运行状况


• 在 Kubernetes 集群中运行启用 GPU 的容器

除此之外，它还支持 时间切片[3]，使得用户可以在 Pod （Kubernetes 的最小调度对象）之间共享 GPU，从而节省成本。

Karpenter 本身也为节点提供了自动缩放功能，也就是说，只有在需要算力时，才会创建 GPU 实例，并且可以根据实际需求修改应用实例的调度规则。除了降低成本之外，也能更灵活地把 GPU 资源调度到 kubernetes 集群中的应用程序。

架构

                        基础架构

                            组件

这个架构非常简单易懂：应用程序选择了一个带有选择器(selector)的 karpenter 制备器(provisioner)，接着，karpenter 制备器根据启动模板创建节点。

部署

接下来最重要的问题是如何去部署它，仍有一些细节需要考虑清楚：

• 如何将 NVIDIA k8s 插件部署到仅有 GPU 的节点


• 如何配置共享的 GPU 节点以使用时间切片，而不影响其他节点


• 如何在启动模板中自动更新节点 AMI，以便节点可以使用最新的镜像


• 如何设置 karpenter 制备器

下面我将逐一讲解。

首先，我们安装 karpenter，并使用 terraform 设置制备器。也可以参考官方文档，在 Amazon EKS 中手动安装 karpenter。如果你已经有 EKS 以及karpenter，可以跳过这一步。

不妨参考这个 GitHub repo[4]。

制备器

resource "kubectl_manifest" "karpenter_provisioner_gpu_shared" {
  yaml_body = <<-YAML
  apiVersion: karpenter.sh/v1alpha5
  kind: Provisioner
  metadata:
    name: gpu-shared
  spec:
    ttlSecondsAfterEmpty: 300
    labels:
      jina.ai/node-type: gpu-shared
      jina.ai/gpu-type: nvidia
      nvidia.com/device-plugin.config: shared_gpu
    requirements:
      - key: node.kubernetes.io/instance-type
        operator: In
        values: ["g4dn.xlarge", "g4dn.2xlarge", "g4dn.4xlarge", "g4dn.12xlarge"]
      - key: karpenter.sh/capacity-type
        operator: In
        values: ["spot", "on-demand"]
      - key: kubernetes.io/arch
        operator: In
        values: ["amd64"]
    taints:
      - key: nvidia.com/gpu-shared
        effect: "NoSchedule"
    limits:
      resources:
        cpu: 1000
    provider:
      launchTemplate: "karpenter-gpu-shared-${local.cluster_name}"
      subnetSelector:
        karpenter.sh/discovery: ${local.cluster_name}
      tags:
        karpenter.sh/discovery: ${local.cluster_name}
    ttlSecondsAfterEmpty: 30
  YAML

  depends_on = [
    helm_release.karpenter
  ]
}

resource "kubectl_manifest" "karpenter_provisioner_gpu" {
  yaml_body = <<-YAML
  apiVersion: karpenter.sh/v1alpha5
  kind: Provisioner
  metadata:
    name: gpu
  spec:
    ttlSecondsAfterEmpty: 300
    labels:
      jina.ai/node-type: gpu
      jina.ai/gpu-type: nvidia
    requirements:
      - key: node.kubernetes.io/instance-type
        operator: In
        values: ["g4dn.xlarge", "g4dn.2xlarge", "g4dn.4xlarge", "g4dn.12xlarge"]
      - key: karpenter.sh/capacity-type
        operator: In
        values: ["spot", "on-demand"]
      - key: kubernetes.io/arch
        operator: In
        values: ["amd64"]
    taints:
      - key: nvidia.com/gpu
        effect: "NoSchedule"
    limits:
      resources:
        cpu: 1000
    provider:
      launchTemplate: "karpenter-gpu-${local.cluster_name}"
      subnetSelector:
        karpenter.sh/discovery: ${local.cluster_name}
      tags:
        karpenter.sh/discovery: ${local.cluster_name}
    ttlSecondsAfterEmpty: 30
  YAML

  depends_on = [
    helm_release.karpenter
  ]
}
view raw

上述两份制备器（provisioner）配置使 Kapenter 通过启动模板（launch template）来初始化对应实例，并添加不同标签（labels）和污点（taints）。

启动模板（仅有GPU）

resource "aws_launch_template" "gpu" {
  name = "karpenter-gpu-${local.cluster_name}"

  block_device_mappings {
    device_name = "/dev/xvda"

    ebs {
      volume_size = 120
    }
  }

  iam_instance_profile {
    name = aws_iam_instance_profile.karpenter.name
  }

  tag_specifications {
    resource_type = "instance"

    tags = {
      "karpenter.sh/discovery" = local.cluster_name
      "jina.ai/node-type"      = "gpu"
    }
  }

  image_id = data.aws_ami.eks_node_gpu.image_id

  instance_initiated_shutdown_behavior = "terminate"

  update_default_version = true

  # key_name = "${local.cluster_name}-sshkey"

  metadata_options {
    http_endpoint               = "enabled"
    http_tokens                 = "optional"
    http_put_response_hop_limit = 2
  }

  vpc_security_group_ids = [module.eks.node_security_group_id]

  user_data = base64encode(templatefile("${path.module}/customized_bootstraps.sh", { cluster_name = "${local.cluster_name}" }))

  tags = {
    "karpenter.sh/discovery" = local.cluster_name
    "node-type"              = "gpu"
  }
}

接下来，我们需要部署具有时间切片配置和默认配置的 NVIDIA k8s 插件，并设置节点选择器，以便 daemonset 仅在 GPU 实例上运行。

nvdp.yml


config:
  # ConfigMap name if pulling from an external ConfigMap
  name: ""
  # Set of named configs to build an integrated ConfigMap from
  map: 
    default: |-
      version: v1
      flags:
        migStrategy: "none"
        failOnInitError: true
        nvidiaDriverRoot: "/"
        plugin:
          passDeviceSpecs: false
          deviceListStrategy: envvar
          deviceIDStrategy: uuid
    shared_gpu: |-
      version: v1
      flags:
        migStrategy: "none"
        failOnInitError: true
        nvidiaDriverRoot: "/"
        plugin:
          passDeviceSpecs: false
          deviceListStrategy: envvar
          deviceIDStrategy: uuid
      sharing:
        timeSlicing:
          renameByDefault: false
          resources:
          - name: nvidia.com/gpu
            replicas: 10
nodeSelector: 
  jina.ai/gpu-type: nvidia

运行下述命令来安装 NVIDIA 的 k8s 插件：

$ helm repo add nvdp https://nvidia.github.io/k8s-device-plugin
$ helm repo update
$ helm upgrade -i nvdp nvdp/nvidia-device-plugin \
  --namespace nvidia-device-plugin \
  --create-namespace -f nvdp.yaml

再之后，使用 nodeSelector 和 toleration 部署应用程序。

gpu.yml


kind: Deployment
apiVersion: apps/v1
metadata:
  name: test-gpu
  labels:
    app: gpu
spec:
  replicas: 1
  selector:
    matchLabels:
      app: gpu
  template:
    metadata:
      labels:
        app: gpu
    spec:
      nodeSelector:
        jina.ai/node-type: gpu
        karpenter.sh/provisioner-name: gpu
      tolerations:
      - key: nvidia.com/gpu
        operator: Exists
        effect: NoSchedule
      containers:
      - name: gpu-container
        image: tensorflow/tensorflow:latest-gpu
        imagePullPolicy: Always
        command: ["python"]
        args: ["-u", "-c", "import tensorflow"]
        resources:
          limits:
            nvidia.com/gpu: 1




gpu-shared.yml


kind: Deployment
apiVersion: apps/v1
metadata:
  name: test-gpu-shared
  labels:
    app: gpu-shared
spec:
  replicas: 1
  selector:
    matchLabels:
      app: gpu-shared
  template:
    metadata:
      labels:
        app: gpu-shared
    spec:
      nodeSelector:
        jina.ai/node-type: gpu-shared
        karpenter.sh/provisioner-name: gpu-shared
      tolerations:
      - key: nvidia.com/gpu-shared
        operator: Exists
        effect: NoSchedule
      containers:
      - name: gpu-container
        image: tensorflow/tensorflow:latest-gpu
        imagePullPolicy: Always
        command: ["python"]
        args: ["-u", "-c", "import tensorflow"]
        resources:
          limits:
            nvidia.com/gpu: 1

现在，如果部署两个 YAML 文件，你将在 AWS 控制台中看到制备的两个节点，或者通过 kubectl get nodes — show-labels命令查看。在每个节点上运行 nvidia-k8s-plugin 后，就可以在应用程序中进行测试。

近期 J-Tech Talk 活动将由我们的工程师结合云原生话题做相关技术分享，敬请期待！请大家关注 JinaAI 视频号！

J-Tech Talk

由 Jina AI 社区为大家带来的技术分享

工程师们将深入细节地讲解具体的问题

分享 Jina AI 在开发过程中所积累的经验

引用链接

[1] 文档: https://karpenter.sh/
[2] k8s 插件: https://github.com/NVIDIA/k8s...
[3] 时间切片: https://github.com/NVIDIA/k8s...
[4] GitHub repo: https://github.com/tarrantro/...

官网：Jina.ai

社区：Slack.jina.ai

开源：Github.com/Jina-ai

‍‍
​