Ray를 이용한 분산 GPU 컴퓨팅: 멀티노드 클러스터 구축과 활용

오늘은 Ray를 사용하여 여러 노드에 걸친 GPU 리소스를 효율적으로 활용하는 방법에 대해 알아보려 한다. Ray는 분산 컴퓨팅을 위한 강력한 프레임워크로, 복잡한 설정 없이도 여러 머신의 리소스를 쉽게 활용할 수 있게 해준다.

Ray 클러스터 설정

먼저 Ray 클러스터를 설정해야 한다. 헤드 노드와 워커 노드를 다음과 같이 시작한다:

헤드 노드 (IP: 192.168.1.100):

CUDA_VISIBLE_DEVICES=0 ray start --head --port=6379 --num-gpus=1

워커 노드:

CUDA_VISIBLE_DEVICES=4 ray start --address='192.168.1.100:6379' --num-gpus=1

여기서 CUDA_VISIBLE_DEVICES를 사용해 각 노드에서 사용할 GPU를 지정했다. 헤드 노드는 0번 GPU를, 워커 노드는 4번 GPU를 사용하도록 설정했다.

서버 측 코드 (server.py)

서버 측 코드에서는 Ray 클러스터에서 실행될 원격 함수와 액터를 정의한다:

import ray
import torch

@ray.remote(num_gpus=1)
def remote_function(x):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    tensor = torch.tensor([x], device=device)
    return tensor.item() * 2

@ray.remote(num_gpus=1)
class GPUActor:
    def __init__(self):
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.gpu_id = torch.cuda.current_device()
    
    def process(self, data):
        tensor = torch.tensor(data, device=self.device)
        result = tensor.sum().item()
        return f"Processed on GPU {self.gpu_id}: {result}"

if __name__ == "__main__":
    ray.init(address="auto")
    print("Ray 클러스터가 실행 중입니다. 클라이언트를 실행하세요.")
    import time
    while True:
        time.sleep(1)

이 코드는 GPU를 사용하는 원격 함수와 액터를 정의한다. @ray.remote(num_gpus=1) 데코레이터를 사용해 각 태스크가 1개의 GPU를 사용하도록 지정했다.

클라이언트 측 코드 (client.py)

클라이언트 코드는 Ray 클러스터에 연결하여 원격 함수와 액터를 사용한다:

import ray

def main():
    ray.init("ray://192.168.1.100:10001")

    result = ray.get(remote_function.remote(4))
    print(f"원격 함수 결과: {result}")

    actors = [GPUActor.remote() for _ in range(2)]
    tasks = [actor.process.remote([i for i in range(1000)]) for actor in actors]
    results = ray.get(tasks)
    for result in results:
        print(result)

if __name__ == "__main__":
    main()

클라이언트는 Ray 클러스터에 연결한 후, 원격 함수를 호출하고 GPUActor 인스턴스를 생성하여 작업을 수행한다.

실행 방법

1. Ray 클러스터를 시작한다 (헤드 노드와 워커 노드).
2. 서버 코드를 실행한다: python server.py
3. 클라이언트 코드를 실행한다: python client.py

이 구조를 사용하면 여러 노드의 GPU 리소스를 효율적으로 활용할 수 있다. Ray가 자동으로 작업을 분산시키고 리소스를 관리하므로, 복잡한 분산 시스템을 쉽게 구축할 수 있다.

주의할 점은 클라이언트 코드에서 헤드 노드의 실제 IP 주소를 사용해야 한다는 것이다. 또한, CUDA_VISIBLE_DEVICES로 GPU를 제한했기 때문에, 코드 내에서는 모든 GPU가 "cuda:0"로 보이지만 실제로는 각각 다른 물리적 GPU를 사용하고 있다.

이렇게 Ray를 사용하면 복잡한 분산 시스템을 간단하게 구축하고 관리할 수 있다. 여러 노드의 GPU 리소스를 효율적으로 활용하여 대규모 계산 작업을 수행하는 데 매우 유용하다.