병렬 처리 방법 소개 (Feat. Ray, Slurm)

Single-node Multi-cpus

• Ray 라이브러리를 통한 병렬 처리

test.py

import ray

# 병렬처리를 위한 함수 정의
@ray.remote
def fn(x: list):
  res = []
  for i in x:
    res.append(i ** 2)
  return res    

# task 정의
tasks = [
  range(100),
  range(100),  
]

# 각 task를 처리하는 병렬처리 함수 호출
futures = [fn.remote(t) for t in tasks]

# task 결과 저장
results = ray.get(futures)

Multi-nodes Multi-cpus

• Slurm Workload Manager와 Ray 라이브러리를 통한 다중 서버 병렬 처리
• Slurm은 리소스만을 분배하는 관리도구로서 실행한 스크립트가 멀티 CPU를 활용할 수 있도록 하기 위해선 Ray 라이브러리와 함께 사용해야 함
• Slurm 작업 할당 방식
  기본적으로 노드 당 task 수를 동등하게 할당함
  Ex1) nodes=2, ntasks=2 이면 각 노드는 1개 task 할당
  Ex2) nodes=6, ntasks=18 이면 각 노드에 3개 task 할당

test.py

import os
import ray

# 노드 ID
slurm_nodeid = os.environ.get('SLURM_NODEID', 'Not Set')
# 전체 Task 수
slurm_ntasks = os.environ.get('SLURM_NTASKS', 'Not Set')
# Task ID
slurm_procid = os.environ.get('SLURM_PROCID', 'Not Set')
# 노드 내 Task ID
slurm_localid = os.environ.get('SLURM_LOCALID', 'Not Set')

# 병렬처리를 위한 함수 정의
@ray.remote
def fn(x: list):
  res = []
  for i in x:
    res.append(i ** 2)
  return res    

# task 정의
tasks = [
  range(100),
  range(100),  
]

# task id별 병렬처리 함수 호출
futures = fn.remote([tasks[slurm_nodeid]])

# task 결과 저장 (병렬 처리된 작업 결과물을 결합하는 별도의 작업 필요)
results = ray.get(futures)

job.sbatch

#!/bin/bash
# job name 설정
#SBATCH --job-name=test
# output 로그 파일 경로 설정 (logger 클래스 출력 기록됨)
#SBATCH --output=output_%j.log
# error 로그 파일 경로 설정 (print 함수 출력 기록됨)
#SBATCH --error=error_%j.log
# timeout 설정 (default 5분)
#SBATCH --time=00:05:00
# 리소스를 요청할 파티션 설정 (직접 지정 필요)
#SBATCH --partition=PARTITION_NAME
# 파티션에 요청할 최소 노드 수 설정
#SBATCH --nodes=2
# 노드 당 task 수 설정
#SBATCH --ntasks-per-node=1
# CPU 자원 설정 (exlusive면 할당 받은 노드의 모든 CPU를 이용함)
#SBATCH --exclusive
# 메모리 자원 설정 (0이면 할당 받은 노드의 모든 메모리를 이용함)
#SBATCH --mem=0

# RUN with slurm (가상환경 활성화 후 파이썬 스크립트 실행)
srun bash -c "source .venv/bin/activate && python test.py"

request job to slurm (CLI command)

# 작업 요청
sbatch job.sbatch