Nvidia Container Toolkit, Nvidia device plugin에 대해 알아봅시다. feat. CRI, CDI

개요

현재 팀의 Machine Learning 파이프라인의 모델 학습 워크로드는 아래와 같은 컴포넌트들로 구성되어있습니다.

출처: https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/index.html

1. model.train()를 entrypoint로 하는 pod
2. pod node의 gpu resource allocation 및 container가 사용할 device config 정보를 kubelet에 등록하는 nvidia device plugin daemonset
3. pod container runtime 및 runc prestart hook을 통해 container에서 사용할 수 있는 device를 설정해주는 nvidia container toolkit
4. 마지막으로 nvidia gpu device를 제어할 수 있도록 os와 hardware 사이에서 번역해주는 nvidia driver

이처럼 AMI에서 사용하는 nvidia driver, nvidia container toolkit, nvidia device plugin 설정 값 관리 등 꽤나 관리포인트가 있고, 이를 올바르게 설정하기 위해선 먼저 우리가 실행한 Pod가 k8s 환경내 소프트웨어 컴포넌트들과 어떤 식으로 상호작용하며 동작하고 있는지 이해할 필요가 있습니다.

CRI

https://kubernetes.io/docs/concepts/architecture/cri/

Kubelet과 Container Runtime이 통신하는 방법을 규격화한 protocol입니다. 둘은 grpc로 통신하며 다양한 api를 지원합니다. 그 중에서 우리가 기억해야할 것은 kubelet은 container runtime에 필요한 device 정보를 Container Runtime에 전달해주고, 해당 정보를 container내 config.json 파일로 접근할 수 있다 정도 되겠습니다.

docker 가 kubelet에 통합되던 시절 이를 유지보수 하기 힘들어 별도의 CRI layer로 분리하였고, 이를 구현한 container runtime 대표적인 구현체엔 CRI-O가 있습니다. 관련해선 여기글이 좋았습니다.

CDI

https://kubernetes.io/docs/concepts/extend-kubernetes/compute-storage-net/device-plugins/

Kubelet과 Device Plugin 구현체간의 통신 규격입니다. 마찬가지로 grpc protocol을 사용하며 우리가 사용하는 nvidia device plugin은 CDI 런타임 구현체입니다.

둘 사이에 많은 기능이 있지만, 마찬가지로 우리가 기억해야할 부분은 device plugin은 node의 device 정보를 watch하여 custom resource (nvidia.com/gpu)를 등록해주며, container runtime에서 사용할 device 정보를 kubelet에 전달한다 정도 입니다.

After successfully registering itself, the device plugin runs in serving mode, during which it keeps monitoring device health and reports back to the kubelet upon any device state changes. It is also responsible for serving Allocate gRPC requests. During Allocate, the device plugin may do device-specific preparation; for example, GPU cleanup or QRNG initialization. If the operations succeed, the device plugin returns an AllocateResponse that contains container runtime configurations for accessing the allocated devices. The kubelet passes this information to the container runtime.

기본적인 내용은 파악했으니 본론으로 넘어가겠습니다.

Nvidia Container Toolkit

https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/arch-overview.html

nvidia-container-toolkit Package는 다음과 같은 역할을 하는 sub library들로 구성되어있습니다. 간단하게만 알아보면

The main packages of the NVIDIA Container Toolkit are:

• nvidia-container-toolkit
• nvidia-container-toolkit-base
• libnvidia-container-tools
• libnvidia-container1

The NVIDIA Container Library and CLI

container runtime이 gpu device configuration을 container에 주입해주는 cli입니다.

The NVIDIA Container Runtime

runc의 thin wrapper로 runc의 spec중 prestart hook 부분에 container process 시작전 실행한 executable 을 주입합니다.

nvidia-container-runtime takes a runC spec as input, injects the NVIDIA Container Runtime Hook as a prestart hook into it, and then calls out to the native runC

The NVIDIA Container Runtime Hook

위에서 주입한 runC prestart hook interface를 구현한 executable. container entrypoint 실행 전 config.json (위에서 언급한) 파일을 읽어 cli를 사용하여 container가 device 를 사용할 수 있게 설정합니다.

This component includes an executable that implements the interface required by a runC prestart hook. This script is invoked by runC after a container has been created, but before it has been started, and is given access to the config.json associated with the container (e.g. this config.json ). It then takes information contained in the config.json and uses it to invoke the nvidia-container-cli CLI with an appropriate set of flags. One of the most important flags being which specific GPU devices should be injected into the container.

*OCI runtime spec에 명시된 prestart hook에 대한 자세한 내용은 여기를 참고해주세요.

지금까지의 내용을 종합해보면..! 우리가 실행한 모델 학습 Pod는 다음과 같은 과정을 통해 gpu를 할당받고 실행된다는 것을 알 수 있습니다.

1. pod resource spec.resource 으로 nvidia.com/gpu 명시 (+toleration 설정)
2. karpenter는 해당 resource를 제공가능한 node provisioning + nvidia.com/gpu taint 부착
3. nvidia device plugin이 node에대한 nvidia.com/gpu 리소스 할당 정보 (resourceAllocation 객체로 표현) kubelet에 전달 → kubelet은 api server에 해당 정보 등록 → node는 allocatable 상태가 됨
4. 또한 해당 Pod가 사용할 device 정보를 kubelet에 전달
5. node는 scheduleable 상태가 되어 Pod 스케줄링
6. kubelet은 device plugin으로 부터 받은 device 정보를 container runtime에 전달
7. container runtime은 해당 정보를 container 내 config.json 파일로 저장
8. pod가 시작될 때 nvidia-container-toolkit에 의해 설정된 runc prestart hook가 실행되면서 config.json 파일 정보를 파탕으로 gpu device 준비(cli 사용)
9. container가 실행되며 model.train() 호출

TroubleShooting

최근에 저희가 사용하는 bottlerocket ami 최신 버전이 karpenter에 의해 자동 선택되며(amiFamily 기능) 파드가 제대로 실행되지 않는 문제가 있었습니다.

Bottlerocket OS release note 를 보면 아래 변경사항이 있었습니다.

• nvidia driver 최신 버전 사용(참고용)
• nvidia container toolkit에선 아래 변경사항이 있었습니다. (github PR)
  • privileged pod의 경우 device 정보를 환경변수로 주입
  • un-privileged pod의 경우 device 정보를 volume mount 형식으로 주입

하지만 현재 저희가 사용하는 nvidia device plugin의 deviceListStrategy 옵션은 envvar(환경변수)로 설정되어있었기 때문에 가장 처음으로 실행된 pod (privileged)는 환경변수를 사용하므로 정상 실행되었지만, 동일 노드에 두 번째 실행된 Pod(un-privileged)는 주입 받은 정보가 없어 device not found error가 발생하였습니다.

따라서 해당 PR 가이드에 따라 deviceListStrategy를 volume-mounts로 설정해주면 정상동작 합니다.

*현재는 karpenter amiSelector를 사용해 구 버전 ami를 하드코딩 해두었습니다.

개선 방향

최근 문서를 보니 Nvidia에서 이 모든 구성요소를 올인원으로 관리할 수 있는 Nvidia GPU Operator를 지원하는 거 같습니다. 추후 PoC 후 도입해보면 좋을 거 같습니다.

이상으로 팀의 ML Pipeline 인프라에서 발생한 문제를 트러블 슈팅하는 과정에서 deep dive한 내용들을 정리해보았습니다.
정확하지 않은 내용이 있을 수 있으니 꼭 공식 문서를 직접 읽어보시고 제가 틀린 부분은 지적해주시면 감사드리겠습니다.

감사합니다.