[Kubernetes] PyTorch 학습 시 Pod의 빠른 종료를 위한 yaml 설계 방법

기존 방법으로는 kubectl delete 시 --force --grace-period=0 옵션을 주어야만 강제 종료되고, 학습 process에 실제로 Ctrl+C를 주는 것과 동일하게 SIGTERM 을 보낼 수는 없었습니다. 이러한 문제점은 학습 프로세스의 종료시 callback들 (wandb 등)이 정상 작동하지 않는다는 문제점이 있습니다.이를 해결할 수 있는 방안을 소개 드립니다.
다만 아래와 같은 몇 가지 drawback이 있으므로 참고하여주시고, 더 좋은 방법이 있다면 공유 부탁드립니다.

  1. pip install 이 불가능하므로 미리 requirements.txt를 설치한 image를 준비해야 함
  2. runAsUserrunAsGroup 을 1003으로 고정해야 하므로 HOME 변수를 /workspace로 고정해야만 사용 가능

기존 방식

apiVersion: v1
kind: Pod
metadata:
  name: jungin500-mobilenetv2
spec:
  securityContext:
    runAsUser: 0
    runAsGroup: 0
    fsGroup: 1003

  restartPolicy: Never

  volumes:
    - name: shmdir
      emptyDir:
        medium: Memory
    - name: pvc-volume
      persistentVolumeClaim:
        claimName: lab-pvc

  containers:
    - name: gpu-container
      image: ghcr.io/jungin500/mobilenet_v2
      volumeMounts:
        - mountPath: /dev/shm
          name: shmdir
        - mountPath: /home/lab
          name: pvc-volume
      env:
        - name: TZ
          value: Asia/Seoul
      command:
        - "/bin/sh"
        - "-c"
      args:
        - >-
          set -x &&
          groupadd -g 1003 lab &&
          useradd -m -d /workspace -s /bin/bash -u 1003 -g lab lab &&
          runuser -u lab -- git clone https://jungin500:*****@github.com/jungin500/mobilenet_v2 mobilenet_v2 &&
          cd /workspace/mobilenet_v2 &&
          runuser -u lab -- git checkout 7d198633d19c2005e22b118ba27082eca3f2846a &&
          runuser -u lab -- pip3 install -r requirements.txt &&
          runuser -u lab -- /bin/bash -c "mkdir -p /home/lab/jungin500/mobilenet_v2_220718 || true" &&
          runuser -u lab -- python3 train.py ****
      securityContext:
        allowPrivilegeEscalation: false

      resources:
        requests:
          nvidia.com/gpu: 1
        limits:
          nvidia.com/gpu: 1

변경된 방식

apiVersion: v1
kind: Pod
metadata:
  name: jungin500-shutdown-signal-test
spec:
  securityContext:
    # [1] runAsUser, runAsGroup을 Fix
    runAsUser: 1003
    runAsGroup: 1003
    fsGroup: 1003

  restartPolicy: Never

  volumes:
    - name: shmdir
      emptyDir:
        medium: Memory
    - name: workdir
      emptyDir: {}
    - name: pvc-volume
      persistentVolumeClaim:
        claimName: lab-pvc

  # [2] git clone, 학습 결과물 폴더 링크 세팅등을 initContainers에서 진행
  # 이 때, pip install 등은 불가능하므로 미리 requirements.txt를 설치한 image를 준비해야 함
  initContainers:
    - name: clone-directory-set
      image: alpine/git
      volumeMounts:
        - mountPath: /workspace
          name: workdir
        - mountPath: /home/lab
          name: pvc-volume
      workingDir: /workspace
      env:
        - name: TZ
          value: Asia/Seoul
        - name: HOME
          value: /workspace
      command:
        - "/bin/sh"
        - "-c"
      args:
        - >-
          set -x &&
          git clone https://jungin500:***@github.com/jungin500/mobilenet_v2 mobilenet_v2 &&
          cd /workspace/mobilenet_v2 &&
          git checkout 7d198633d19c2005e22b118ba27082eca3f2846a &&
          mkdir -p /home/lab/jungin500/mobilenet_v2_test || true
  
  # [3] 기타 작업 (폴더 생성, 이동 등)은 불가능함
  #     argument도 string의 형태가 아닌 array로 주어야 작동함
  containers:
    - name: gpu-container
      image: ghcr.io/jungin500/mobilenet_v2
      volumeMounts:
        - mountPath: /dev/shm
          name: shmdir
        - mountPath: /workspace
          name: workdir
        - mountPath: /home/lab
          name: pvc-volume
      workingDir: /workspace/mobilenet_v2
      env:
        - name: TZ
          value: Asia/Seoul
        - name: HOME
          value: /workspace
      command:
        - "/opt/conda/bin/python3.8"
        - "train.py"
      args:
        - "--arg1"
        - "arg1_value"
      securityContext:
        allowPrivilegeEscalation: false

      resources:
        requests:
          nvidia.com/gpu: 1
        limits:
          nvidia.com/gpu: 1

구체적으로 변경되어야 하는 부분은 다음과 같습니다.

spec.securityContext.runAsUser
spec.securityContext.runAsGroup
spec.volumes.name[workdir]
spec.initContainers
spec.containers.volumeMounts[mountPath=/workspace]
spec.containers.env[name=HOME]
spec.containers.command
spec.containers.args

답글 남기기

이메일 주소는 공개되지 않습니다.