NSML - 분산 학습 플랫폼의 스케줄링 요구 사항과 해결 방안

⚠️ 유의사항: 이 글의 독자는 K8s의 기본 컨셉과 기능을 이해하고 있다고 가정합니다. 해결에 핵심적인 K8s 기능은 본문에서 설명하지만 부차적인 기능은 공식 문서 링크로 설명을 대체합니다.

Introduction

NSML 팀에서는 NSML을 좋은 분산 학습 플랫폼으로 만들기 위해 많은 노력을 기울였습니다. 그 중심에는 GPU 자원의 효율적인 활용이 있습니다. 앞으로 두 편에 걸쳐 발행할 NSML 스케줄러 시리즈에서는 자원을 효율적으로 할당하는 스케줄링 정책을 어떻게 수립했는지 공유하고, NSML 스케줄러를 어떻게 구성했는지 다뤄보겠습니다.

3편: NSML의 스케줄링 요구 사항과 해결 방안 4편: NSML 맞춤형 스케줄러 개발기

먼저 이번 글에서는 대규모 GPU 풀 안에서 다수의 머신 러닝 작업이 경합할 때 발생하는 스케줄링 문제와 분산 학습 스케줄링 요구 사항을 설명하고, 이를 Kubernetes(이하 K8s)의 기능 및 K8s 스케줄링 프레임워크를 사용해 해결한 방법을 소개합니다.

자원 스케줄링의 역할 및 필요성

NSML은 대규모 GPU 풀을 운영하는 시스템인 만큼 원활한 자원 배치가 중요합니다. 자원을 원활하게 배치하기 위한 NSML 자원 스케줄링의 핵심 고려 사항은 크게 두 가지입니다.

첫 번째, 분산 학습이 원활하게 이뤄지는 환경을 제공할 것

NSML은 최신 딥러닝 연구 트렌드를 따라 분산 학습에 최적화된 서비스를 제공하고자 합니다. 기본 스케줄링을 사용하면 NSML의 인프라 구성을 고려하지 않은 채 자원을 할당하기 때문에 분산 학습에 적합한 환경을 조성할 수 없습니다.

두 번째, 자원 활용률을 높일 것

자원은 언제나 한정적이기 때문에 급증하는 대규모 학습 수요에 대비해 자원 낭비가 최소화되도록 실험을 스케줄링해야 합니다. 자원이 잘 스케줄링되지 않으면 자원은 남아있는데 요청된 학습이 자원을 할당받지 못해 대기 시간이 무한정 길어지면서 사용자에게 나쁜 경험을 초래할 수 있습니다.

위 두 가지 고려 사항을 만족시키기 위해 NSML이 풀어야 했던 문제가 무엇인지 구체적으로 알아보고 어떻게 해결했는지 설명하겠습니다.

1. 분산 학습의 멀티 노드 동시 실행

분산 학습의 멀티 노드 동시 실행 지원이 필요한 이유

모델과 데이터의 크기가 커서 단일 서버 8장의 GPU로 충분하지 않은 대규모 분산 학습의 경우에는 여러 GPU 서버를 사용합니다. 학습에 사용할 GPU 서버를 노드라고 부르고, 다수의 노드를 멀티 노드라고 부릅니다.

분산 학습은 노드에 걸쳐 연산을 병렬로 진행하므로 각 노드의 학습 결과를 합산하는 통신 과정이 필요합니다. 따라서 분산 학습을 지원하는 여러 프레임워크(TensorFlow, PyTorch 등)에서는 모든 노드가 제대로 구동돼 동기화할 수 있는 상태인지 확인하는 초기화 단계를 거칩니다.

이와 같은 특성의 분산 학습에서 멀티 노드의 동시 실행을 보장하지 않았을 때 발생하는 최악의 상황은 생성된 실험이 모두 데드락에 빠져 무한정 대기하는 것입니다. 모든 실험의 일부 노드만 실행돼 GPU 자원을 전부 점유해 버리면, 어떤 실험도 종료되지 못하고 나머지 노드가 실행되기만을 기다립니다. 결국 모든 실험이 다른 실험이 종료돼 자원을 풀어주기만을 기다리는 데드락 상황이 발생합니다.

이를 방지하기 위해 Spark와 같은 일부 분산 처리 프레임워크에서는 작업(job)을 만들 때 명시한 최소 실행 작업자(worker) 수만큼의 노드가 확보되어야 작업을 진행하기도 하는데요. 이와 비슷하게 NSML에선 학습의 모든 작업에 필요한 자원이 충분할 때까지 기다린 후 모든 노드를 동시에 실행하는 갱(gang) 스케줄링을 도입했습니다.

K8s 스케줄러의 구조

문제를 해결한 방법을 보다 쉽게 이해하실 수 있도록 K8s 스케줄러의 구조와 동작 방식을 간단히 설명하겠습니다.

K8s의 스케줄링은 특정 파드(pod)를 최적의 호스트에 배치하는 과정을 의미합니다. 이 과정은 사용할 수 없는 호스트는 거르고(filter 단계), 남아 있는 호스트 중에서 최적의 호스트를 선정한 뒤(score 단계), 실제로 배치하는(bind 단계) 등의 여러 단계로 구성돼 있습니다. 실제로는 더 많은 단계를 거치지만 쉽게 설명하기 위해 단순화했습니다.

K8s 스케줄러 커스터마이징 방법

K8s 스케줄러는 각 단계를 익스텐션 포인트(extension point)로 노출합니다. 익스텐션 포인트마다 특정 플러그인을 활성화하는 방식으로 각 단계에 적용할 세부 동작을 요구사항에 알맞게 제어할 수 있습니다. 플러그인이란 K8s 스케줄러의 확장 포인트 중 하나 이상을 구현한 라이브러리로, 스케줄링이 해당 단계에 돌입했을 때 익스텐션 포인트에 등록된 플러그인이 실행되어 플러그인에 구현해 놓은 로직이 실행됩니다. 기본적으로 활성화되는 플러그인은 공식 문서에서 확인할 수 있습니다.

K8s는 익스텐션 포인트를 활용하여 스케줄러를 커스터마이징할 수 있는 여러 방법을 제공합니다.

스케줄러 익스텐더는 익스텐션 포인트에 웹훅을 걸어 특정 단계에서 원하는 로직을 실행하도록 하는 방법입니다. 스케줄러와 익스텐더 사이의 데이터가 HTTP로 전송되며 익스텐더는 스케줄링 단계가 끝난 후에 호출되므로 특정 익스텐션 포인트 전에 미리 처리해야 하는 동작은 구현할 수 없다는 한계가 있습니다.

스케줄링 프레임워크를 사용하면 원하는 동작을 인터페이스에 맞춰 플러그인으로 구현한 뒤 익스텐션 포인트에 해당 플러그인을 등록하면 원하는 단계에 원하는 로직이 수행되게 할 수 있습니다. 스케줄러 익스텐더에 비해 스케줄링 알고리즘에 대한 더 세밀한 커스터마이징이 가능합니다.

NSML에선 더 복잡한 로직으로 확장하기 용이한 스케줄링 프레임워크를 활용하기로 결정했습니다.

K8s 스케줄링 프레임워크의 코스케줄링(coscheduling) 플러그인 도입

kubernetes-sigs에 공개된 스케줄링 프레임워크 구현체 중 하나인 코스케줄링 플러그인을 사용해 갱 스케줄링을 적용할 수 있었습니다. 코스케줄링은 갱 스케줄링을 포함하는 더 넓은 범위로, 최소 실행 노드의 수를 설정해 전체 노드 중 설정한 개수의 노드가 실행 가능하면 스케줄링하는 방식입니다. 갱 스케줄링은 여기서 최소 실행 노드의 수와 전체 노드의 수를 일치시킨 방법입니다.

코스케줄링 플러그인은 파드그룹(PodGroup)이라는 K8s 커스텀 리소스(Custom Resource, CR)를 이용해 파드를 그룹핑해서, 파드그룹에 속한 모든 파드가 바로 배치 가능한 상태일 때 스케줄링 단계를 진행해 자원을 할당합니다. 가용 자원이 부족해서 파드그룹 중 일부 파드를 배치할 수 없다면 모든 파드를 배치하지 않고 파드그룹 전체 스케줄링을 보류합니다.

아래 그림은 갱스케줄링을 사용해 멀티 노드가 동시에 시작하는 과정을 나타낸 그림입니다. 파드 4개가 포함된 파드그룹이 처음에는 리소스 부족으로 전체 파드그룹이 배치되지 못하다가 파드 4개 모두 배치가 가능한 상태가 되면 스케줄링됩니다.

이를 NSML의 구성요소인 실험(run)과 노드(node)에 대입해보겠습니다. 하나의 실험은 다수의 노드로 구성할 수 있습니다. 실험은 K8s의 파드그룹 리소스에 매핑할 수 있습니다. 노드가 매핑되는 K8s의 파드는 해당 파드그룹(즉, 실험)에 속하도록 설정합니다. 파드 명세에 파드그룹 레이블을 추가하면 코스케줄링 플러그인은 해당 레이블 값으로 파드가 속한 파드그룹을 파악할 수 있습니다.

아래는 한 실험을 두 개의 노드로 구성한 경우에 K8s에서 사용하는 명세 예시입니다. 자원이 충분해 run1이라는 이름의 파드그룹에 속한 두 개의 파드를 모두 배치할 수 있다면 두 파드를 모두 스케줄링합니다.

# 파드그룹 리소스
apiVersion: scheduling.sigs.k8s.io/v1alpha1
kind: PodGroup
metadata:
name: run1
...
---
# 파드그룹에 속한 첫 번째 파드
apiVersion: v1
kind: Pod
metadata:
  labels:
    pod-group.scheduling.sigs.k8s.io: run1
  name: node0
spec:
  containers: ...
---
# 파드그룹에 속한 두 번째 파드
apiVersion: v1
kind: Pod
metadata:
  labels:
    pod-group.scheduling.sigs.k8s.io: run1
  name: node1
spec:
  containers: ...

이와 같이 NSML은 코스케줄링 플러그인을 활용해 실험 내 노드를 모두 동시에 배치해 실행할 수 있도록 스케줄링하고, 이를 통해 멀티 노드의 동시 실행을 보장해 분산 학습시 편의성을 높였습니다.

2. InfiniBand 구성에 따른 구역(zone) 스케줄링

구역(zone)이란?

다음 문제를 설명하기에 앞서 NSML에서 사용하는 구역이라는 개념을 짚고 넘어가겠습니다. 구역은 한 실험의 멀티 노드를 배치할 때 벗어날 수 없는 영역을 나타내는 논리적인 단위입니다.

NSML 시리즈의 1편 NSML - AI 분산 학습 플랫폼 요소 및 인프라에서 설명했듯이 NSML Pods 안에 존재하는 호스트끼리만 InfiniBand 통신이 가능합니다. 달리 말하면 A NSML 파드에 있는 호스트는 B NSML 파드에 있는 호스트와 InfiniBand로 통신할 수 없습니다. 이러한 인프라 특성을 고려하여 구역을 단일 NSML Pod에 매핑하고 있습니다. 물리적으로 구분되어 있는 NSML Pod 단위를 편의상 구역이라고 명명합니다.

분산 학습 노드 간 고속 통신 네트워크 활용이 필요한 이유

단일 GPU로 학습할 수 없는 모델과 데이터 용량이 큰 학습의 경우 1개 이상의 GPU가 필요합니다. 이와 같은 멀티 GPU 학습은 GPU 간 학습 결과를 공유할 필요가 있습니다. 특히 분산 학습의 경우, 각 서버에 실행되고 있는 노드의 서로 다른 학습 결과를 합산하는 과정이 필요합니다. 대규모 분산 학습은 수십, 수백 개의 GPU를 사용하기 때문에 GPU 간, GPU 서버 간 통신 속도가 총 학습 성능에 큰 영향을 미칩니다.

따라서 대규모 분산 학습의 학습 속도를 높이기 위해서는 각기 다른 GPU 서버에 배치된 노드 간의 통신 속도가 핵심입니다. 이때 NSML에서 구축한 고성능 네트워크 InfiniBand 통신을 활용하면 노드 간 고속 통신이 가능합니다. InfiniBand 네트워크는 이더넷(ethernet)보다 평균적으로 큰 대역폭을 가지며 RDMA 기술 덕분에 이더넷에 비하여 통신 속도가 향상됩니다. 구역 내 서버끼리 InfiniBand 통신이 가능하도록 NSML Pods 인프라를 설계했으므로, 분산 학습 실험의 노드 간 InfiniBand 통신을 지원하기 위해 노드를 동일한 구역에 배치하는 스케줄링을 구현해야 했습니다.

K8s의 파드 간 어피니티(inter-pod affinity) 기능을 활용한 구현 방안

특정 조건에 부합하는 호스트에 파드를 배치하고 싶을 때 보통 K8s의 nodeSelector 속성을 활용합니다. 그런데 nodeSelector를 사용하기 위해서는 다수의 구역 중 배치할 구역을 미리 결정해 직접 지정해야 하는데요. 스케줄링 시점의 자원 상황에 따라 배치할 구역이 계속 달라지기 때문에 구역을 미리 결정할 수가 없어서 nodeSelector는 적절한 해결책이 아닙니다.

이때 K8s의 파드 간 어피니티를 사용하면 스케줄링 단계에서 상황에 맞게 파드그룹을 동일한 구역으로 스케줄링할 수 있습니다. 파드 간 어피니티를 활용하면 호스트에서 이미 실행 중인 파드의 레이블을 기반으로 다음 파드를 스케줄링할 호스트를 제한할 수 있기 때문입니다. 파드 간 어피니티는 "규칙 Y를 충족하는 하나 이상의 파드가 X에서 실행 중이면, 파드를 X에 배정해야 한다"와 같은 형태로 구성됩니다. 여기서 Y는 레이블 셀렉터(labelSelector)로 표현하고, 파드의 특정 레이블이 어떤 값이어야 규칙 Y를 충족하는지 명시합니다. X는 토폴로지 키(topologyKey)로 표현하고, 토폴로지 도메인을 나타내는 데 사용하는 호스트 레이블의 키를 명시합니다.

NSML은 파드그룹에 속한 모든 파드를 동일한 구역에 배치해야 했습니다. 이를 파드 간 어피니티로 표현하면 다음과 같습니다. "규칙 Y(동일한 파드그룹 레이블의 값을 가진다)를 만족하는 하나 이상의 파드가 X(어떤 구역)에서 이미 실행 중이면, 파드를 X(동일한 구역)에 배정한다"

apiVersion: v1
kind: Pod
metadata:
  name: with-pod-affinity
  labels:
    pod-group.scheduling.sigs.k8s.io: run1
spec:
  affinity:
    podAffinity:
      requiredDuringSchedulingIgnoredDuringExecution:
      - labelSelector: # 1. 동일한 파드그룹 레이블을 가지는 파드를 찾아,
        matchExpressions:
        - key: pod-group.scheduling.sigs.k8s.io # 파드그룹 레이블의 키
          operator: In
          values:
          - run1 # 파드그룹 레이블의 값
        topologyKey: zone # 2. 같은 "구역(zone)" 레이블값이 설정된 호스트에 스케줄링한다.
  containers:
  ...

먼저 K8s가 NSML에서 정의한 논리 개념인 구역을 구별할 수 있도록 각 구역의 호스트마다 구역을 명시하는 레이블(예: zone=1, zone=2 ...)을 부착했습니다. 그리고 파드에 파드 간 어피니티를 적용하기 위해서 레이블 셀렉터로 파드그룹에 부착된 레이블의 키와 값(예: pod-group.scheduling.sigs.k8s.io=run1)을 설정하고, 토폴로지 키에 구역 레이블의 키값(예: zone)을 설정했습니다. 이를 통해 파드그룹에 속한 모든 파드를 토폴로지 키값이 같은 호스트에 배치할 수 있습니다.

예를 들어, 1번 구역에 속한 호스트에는 zone=1, 2번 구역에 속한 호스트에는 zone=2 레이블을 설정합니다. 앞서 도입한 코스케줄링 플러그인을 사용하기 위해서는 파드그룹에 속한 파드에 파드그룹 레이블(pod-group.scheduling.sigs.k8s.io={value})도 추가해야 하는데요. 이를 활용해 같은 파드그룹에 속한 모든 파드(파드그룹 레이블의 값이 같은 파드)를 같은 구역, 즉 구역 레이블이 같은 값으로 설정돼 있는 호스트로 배치하도록 했습니다. K8s 스케줄러는 파드 단위로 스케줄링하므로 같은 파드그룹에 속한 파드라도 가장 먼저 자원을 할당받는 파드가 존재합니다. 이때 가장 먼저 어떤 파드가 어떤 구역으로 배치되면, 이어서 스케줄링되는 파드들은 첫 번째로 배치된 파드가 위치한 구역으로 모두 배치됩니다.

이와 같이 실험 내 모든 노드를 같은 구역에 배치하므로 구역의 InfiniBand 네트워크를 사용해 노드 간 고속 통신이 가능해집니다.

3. 자원의 파편화 방지

NSML 스케줄러는 앞서 살펴본 방법으로 분산 학습에 걸맞은 환경을 제공할 수 있는 분산 학습 스케줄링 요구 사항을 충족했습니다. 하지만 여전히 자원을 효율적으로 사용하지 못하도록 방해하는 문제가 남아 있습니다. 바로 자원의 파편화 문제인데요. 자원 파편화는 범위에 따라 세 가지로 분류할 수 있습니다. 호스트 내부와 구역 내부, 구역 사이의 파편화로 분류할 수 있으며, 각각 다른 방법으로 해결했습니다.

3-1. 호스트 내부 자원 파편화 - 리소스 유형 도입

호스트에는 GPU와 CPU, 메모리 등의 자원이 있습니다. 이해를 돕기 위해 NSML이 단일 호스트로 이루어져있고, 호스트 내부에는 GPU 8장과 메모리 자원(총 100%)만 있다고 가정했습니다. 먼저 호스트의 메모리 자원 100%와 GPU 4장을 사용하는 실험 A가 생성됩니다. 이후 메모리 자원 30%와 GPU 4장이 필요한 실험 B가 요청된 경우, 호스트의 GPU 자원은 할당 가능하지만 메모리 자원은 이미 실험 A에서 모두 점유하고 있기 때문에 실험 생성이 불가합니다.

이와 같이 사용자가 임의로 자원 요구량을 입력하면 NSML에서 자원을 효율적으로 스케줄링하기 어려워집니다. 일부 자원은 남고, 일부 자원은 부족한 상황이 발생할 수 있습니다. 이와 같은 파편화를 완화하기 위해 실험에 사용할 하드웨어 사양을 미리 정의해 놓는 리소스 유형을 도입했습니다.

리소스 유형은 GPU와 CPU, 메모리, 스토리지 용량 등 다양한 조합으로 구성되며, 사용자는 계획한 실험 규모에 따라 적합한 리소스 유형을 선택할 수 있습니다. 리소스 유형의 이름은 NSML에서 가장 중요하게 다루는 자원인 GPU를 기준으로 나누었습니다. AWS 등의 퍼블릭 클라우드에서 일반적으로 사용하는 것처럼 NSML의 리소스 유형도 GPU 장수를 2의 거듭제곱으로 설정했습니다^[1]. 실험의 리소스 유형 명세에 따라 K8s 파드의 명세가 달라집니다. 아래는 A100 GPU 1장 리소스 유형을 선택해 생성한 K8s 파드의 명세입니다.

apiVersion: v1
kind: Pod
resources:
  limits:
  cpu: "7.5"
  ephemeral-storage: 32Gi
  memory: 192Gi
  nvidia.com/gpu: "1"
  rdma/hca_shared_devices_a: "1"

사용자는 NSML에서 실험을 생성할 때 리소스 유형을 선택해야 합니다. 리소스 유형으로 정해져 있는 사양은 호스트의 자원을 나눠 배분할 수 있도록 미리 정해뒀기 때문에 나머지 자원은 넉넉한데 일부 자원이 모자라서 배치되지 않는 앞선 예시의 상황은 이제 발생하지 않습니다.

3-2. 구역 내부 자원 파편화 - MostAllocated 정책 설정

리소스 유형을 도입해 호스트 내부 파편화는 완화했습니다. 하지만 작은 자원이 필요한 리소스 유형을 선택해 실험 생성과 종료를 반복할 때 발생하는 구역 차원의 파편화가 아직 남아 있습니다.

위 그림에서는 NSML이 단일 구역으로 구성됐다고 가정했습니다. GPU 8장 이하의 작은 규모의 실험이 명확한 스케줄링 기준 없이 배치되면 모든 호스트의 GPU 장수가 조금씩 부족해 자원을 할당받지 못하는 구역 내 파편화가 발생하는 것을 확인할 수 있습니다.

구역 내부 파편화 문제는 K8s 스케줄러에서 제공하는 플러그인을 사용해 정책으로 해결했습니다. NSML 노드가 배치될 호스트는 최적의 호스트를 선정하는 Score 과정에서 결정됩니다. 여러 호스트가 노드가 요구하는 자원 양을 만족한다면 그중에서 Score 단계에 활성화돼 있는 NodeResourcesFit 플러그인에서 설정한 정책에 따라 최적의 노드 하나가 선정됩니다. 기본값은 'LeastAllocated'로 가장 가용 리소스가 많은 호스트를 선호하는 정책입니다. 반대로 'MostAllocated' 정책을 사용하면 배치 가능한 호스트 중에서 가장 가용 리소스가 적은 호스트를 선호하게 됩니다. NSML은 스케줄러에 'MostAllocated' 정책을 설정해 여러 호스트에 균일하게 노드를 배치하기보다는 최대한 적은 호스트에 노드가 배치되도록 했습니다. 이를 통해 자원의 파편화로 실험이 배치되지 못하는 문제를 완화했습니다.

3-3. 구역 간 자원 파편화 - 구역에 역할 부여

NSML은 다수의 구역으로 구성됩니다. 앞서 호스트 내 파편화와 단일 구역 내부의 파편화 문제를 해결하는 방법을 살펴봤는데요. 이제 구역 간의 파편화 문제를 살펴볼 차례입니다. 구역 간 파편화 문제는 NSML이 소규모 실험뿐 아니라 GPU 64장 혹은 128장을 요구하는 대규모 분산 학습 실험을 지원하기 때문에 발생합니다.

문제를 보다 쉽게 이해하기 위해 NSML에 두 개의 구역이 존재하고 각 구역은 16개의 서버로 구성돼 있다고 가정하겠습니다. 이때 만약 GPU 8장이 필요한 A100 8 리소스 유형의 실험 두 개가 1번 구역에 하나, 2번 구역에 하나 배치된 상황에서 GPU 128장이 필요한 A100 128 리소스 유형의 대규모 실험이 생성된다면 어떻게 될까요? 하나의 GPU 서버는 8장의 GPU로 구성되기 때문에 GPU 128장은 최소 16개의 GPU 서버, 즉 16개의 노드가 필요합니다. 또한 앞서 [2. InfiniBand 구성에 따른 구역(zone) 스케줄링] 챕터에서 다루었듯이 분산 학습의 학습 성능 향상을 위해 노드는 모두 동일한 구역에 배치해야 합니다. 이때 만약 1번 구역과 2번 구역에서 남는 자원을 합치면 A100 128 실험을 충분히 할당할 수 있지만, 이 두 실험이 여러 구역에 흩어져 있는 바람에 A100 128 실험은 앞서 배치된 소규모 실험이 모두 종료될 때까지 대기해야 합니다.

이 문제가 바로 구역 간 파편화 문제입니다. 전체적으로는 가용 자원이 충분하지만 그 자원들이 여러 구역에 흩어져 있는 바람에 배치가 불가능해 자원이 낭비되는 현상입니다. 이때 만약 A100 8 리소스 실험 두 개를 모두 1번 구역에 배치했다면, 2번 구역에 GPU 128장 실험을 대기 시간 없이 할당할 수 있습니다.

실험 규모에 따라 구역 역할 부여

구역 간 파편화는 대규모 실험에 특히 치명적입니다. 실험의 규모가 커질수록 소규모 실험을 잘못 배치한 탓에 한 구역에 배치될 확률이 낮아집니다. 이에 대규모 실험의 학습 환경을 보장하기 위해 실험의 규모에 따라 구역에 역할을 부여했습니다. NSML의 구역을 'GPU 32장 이하를 요구하는 소규모 실험을 위한 구역'과 'GPU 64장 중규모 실험을 위한 구역', 그리고 'GPU 128장 대규모 실험을 위한 구역'으로 구획했습니다.

K8s에서 구역 역할을 표현한 방법

먼저 구역의 역할을 소규모, 중규모, 대규모의 세 가지 규모로 나누기 위해 리소스 유형을 리소스 유형군으로 분류하고, K8s의 확장 리소스(Extended Resource) 기능을 사용해 노드가 알맞은 구역에 스케줄링되도록 했습니다.

K8s에서는 일반적으로 파드를 특정 노드에 배치하기 위해 노드 레이블 셀렉터(node label selector)를 사용합니다. 노드 레이블 셀렉터를 사용하면 구현이 간편하다는 장점이 있지만, 리소스 유형군별로 등록된 파드의 개수를 K8s에서 바로 알 수 없다는 단점이 있습니다. NSML에는 구역의 역할에 맞게 스케줄링하는 것 뿐만 아니라 대기 중인 파드의 개수를 각 리소스 유형별로 제한해 과도하게 대기열이 늘어나는 것*을 방지해야 한다는 요구 사항이 있었습니다.

NSML은 이 두 가지 조건을 모두 만족하기 위해 확장 리소스를 택했습니다. ^[2] 확장 리소스란 말 그대로 가용할 수 있는 자원(GPU, CPU 등)을 확장하는 것으로 클러스터 관리자가 직접 K8s에 호스트 차원에서 스케줄 가능한 자원을 등록할 수 있는 기능입니다. NSML이 리소스 유형으로 한 파드에서 할당받을 CPU 양을 제약할 수 있는 것도 호스트에 CPU 자원이 등록돼 있기 때문입니다. 클러스터 관리자가 확장 리소스를 특정 호스트에 등록하면 K8s에 등록된 자원처럼 활용할 수 있는데요. 이를 이용해 구역의 역할에 맞게 리소스 유형군을 호스트에 확장 리소스로 등록하면 됩니다.

아래는 소규모 실험을 위한 구역의 호스트에 nsml-run-small 확장 리소스를 등록한 예시입니다. nsml-run-small은 해당 호스트에서 생성할 수 있는 실험의 최대 개수로 설정하며, 아래 예시에서는 제일 작은 리소스 유형인 A100 1장 실험으로 가득 채웠을 때 최대 8개가 생성될 수 있으므로 8입니다.

# K8s 클러스터의 임의의 호스트의 Capacity 정보를 출력한 결과입니다.
Capacity:
cpu:               20          # 자동으로 설정된 호스트의 CPU 자원의 양
ephemeral-storage: 13343749Mi
memory:            264034868Ki
nsml-run-small:    8           # 확장 리소스

노드를 K8s 파드로 만들때 리소스 유형에 알맞은 확장 리소스를 요구하도록 설정하면 해당 확장 리소스가 설정된 호스트에 배치됩니다. 아래는 A100 1 리소스 유형을 선택해 생성한 K8s 파드의 명세입니다.

apiVersion: v1
kind: Pod
  …
resources:
  limits:
    cpu: "7.5"
    ephemeral-storage: 32Gi
    memory: 192Gi
    nvidia.com/gpu: "1"
    rdma/hca_shared_devices_a: "1"
    nsml-run-small: "1" # nsml-run-small 자원을 1만큼 요구한다.

구역 역할에 따른 스케줄링

각 구역의 역할에 따라 확장 리소스를 전부 등록해 뒀기 때문에 각 실험의 리소스 유형에 적합한 구역에 배치됩니다. 이를 통해 구역 간의 파편화를 완화해 대규모 실험도 비교적 적은 대기 시간으로 배치될 수 있습니다.

선점 가능(preemptible) 옵션 도입

소규모 실험을 위한 구역에 비해 대규모 실험을 위한 구역에서는 일부 노드가 오류로 미리 종료되거나, 대규모 실험에 대한 수요 자체가 적어 유휴자원이 발생할 가능성이 높아집니다. 대규모 실험은 대체로 하나의 구역을 전부 점유하기 때문에 일부 노드가 종료된다면 수십 개의 GPU가 활용되지 않는 순간이 생깁니다. 이 구역에 바로 배치할 수 있는 소규모 실험이 있더라도 구역의 역할이 달라 배치할 수 없어서 대기하게 됩니다. NSML은 이러한 경우에도 소중한 GPU 자원이 낭비되지 않도록 '선점 가능(preemptible)'이라는 개념을 도입했습니다.

선점 가능이란 하드웨어 자원이 부족할 때 자신의 실험을 중단하고 자원을 양보하도록 설정된 옵션입니다. 선점 가능 옵션을 선택해 생성된 실험은 구역의 역할과 관계없이 배치받을 수 있어 남는 자원을 모두 활용할 수 있습니다. 이 옵션이 잘 활용된다면 NSML은 자원의 활용률을 높일 수 있고, 사용자는 디버깅을 위해 작은 실험을 진행하고 싶을 때 비교적 빠른 시간 안에 자원을 배치받을 수 있습니다.

K8s 파드는 스케줄링 우선순위를 부여할 수 있습니다. 가용 자원이 부족해 파드를 스케줄링할 수 없는 경우, 스케줄러는 우선순위가 낮은 파드를 선점(축출)해 보류 중인 파드가 자원을 할당받을 수 있도록 합니다. 우선순위는 K8s의 프라이어리티클래스(PriorityClass)로 적용합니다. NSML에서는 선점 가능 옵션을 구현하기 위해 K8s 프라이어리티클래스 두 가지를 생성했습니다.

• 선점 가능 옵션을 적용한 파드에 적용할 프라이어리티클래스
• 가용 자원이 부족할 때 선점 가능한 파드를 축출할 프라이어리티클래스

선점 가능 옵션을 적용한 실험의 파드와 그렇지 않은 파드는 프라이어리티클래스를 다르게 설정합니다. 일반 실험에 적용되는 프라이어리티클래스의 우선순위 값을 선점 가능 실험보다 높게 두어 선점 가능 실험을 축출할 수 있도록 구현했습니다.

맺음말

이번 글에서는 NSML 스케줄러가 K8s 스케줄링 프레임워크의 공식 플러그인인 코스케줄링과 K8s의 파드 간 어피니티 기능을 이용해 NSML의 인프라를 고려한 스케줄링으로 분산 학습 환경을 제공한 과정을 살펴보고, K8s의 여러 옵션을 활용해 자원의 파편화를 완화한 방법을 살펴봤습니다. 다음 글에서는 NSML의 복합적인 스케줄링 상황 때문에 발생하는 여러 가지 이슈와 이를 해결하기 위해 NSML 맞춤형 스케줄러를 개발하게 된 과정을 공유하겠습니다.