Spark에서 XGBoost 분산 학습

분산 학습은 여러 연산 자원에서 학습 작업과 데이터를 나누어 처리합니다. 단일 노드 메모리에 담기 어려운 데이터로 학습하거나 여러 노드의 연산을 함께 사용할 때 적용합니다. 이 분산 학습 연결 절차가 현재 직접 지원하는 라이브러리는 PySpark XGBoost 추정기입니다. SparkXGBClassifier는 분류, SparkXGBRegressor는 회귀에 사용합니다.

필수 실행 프로필

저장소에서 KUBERNETES_OPEN_PORTS_ALL Spark 프로필을 가져와 변환에 적용해야 합니다. 노드 사이의 학습 통신에 필요한 설정이며 선택 사항이 아닙니다. XGBoost의 PySpark 추정기를 제공하는 xgboost 패키지와 모델 어댑터 의존성도 환경에 추가합니다.

어댑터와 모델 저장 형식

기본 XGBoostSerializer는 SparkXGBClassifier 또는 SparkXGBRegressor를 직접 저장하는 직렬화기가 아닙니다. Spark 모델을 학습한 뒤 get_booster()로 내부 Booster를 추출하여 어댑터에 전달합니다.

아래 완성 예제는 feature_1과 feature_2라는 수치 특성, 0과 1을 사용하는 label 정답으로 이진 분류를 학습합니다. 어댑터는 클래스 1의 확률을 probability_1로 반환합니다. 회귀나 다중 클래스 모델을 연결할 때에는 그 Booster의 출력 형태와 의미에 맞춰 API 및 추론 결과를 정의합니다.

src/main/model_adapters/adapter.py에서 네이티브 Booster를 JSON으로 저장하고 같은 Booster 타입으로 복원합니다. 저장한 특성 순서대로 DMatrix를 만들어 predict에 전달합니다.

import xgboost as xgb
import palantir_models as pm

class BoosterModelAdapter(pm.ModelAdapter):
    FEATURE_COLUMNS = ["feature_1", "feature_2"]

    def __init__(self, booster):
        self.booster = booster

    def save(self, state_writer):
        with state_writer.open("booster.json", "w") as model_file:
            self.booster.save_model(model_file.name)

    @classmethod
    def load(cls, state_reader):
        booster = xgb.Booster()
        with state_reader.open("booster.json", "r") as model_file:
            booster.load_model(model_file.name)
        return cls(booster)

    @classmethod
    def api(cls):
        columns = [(name, float) for name in cls.FEATURE_COLUMNS]
        return (
            {"df_in": pm.Pandas(columns)},
            {"df_out": pm.Pandas(columns + [("probability_1", float)])},
        )

    def predict(self, df_in):
        matrix = xgb.DMatrix(df_in[self.FEATURE_COLUMNS])
        df_in["probability_1"] = self.booster.predict(matrix)
        return df_in

학습과 게시

src/main/model_training/model_training.py에 다음 변환을 둡니다. 예시 경로는 실제 입력과 모델 출력으로 바꿉니다. 입력의 feature_1, feature_2와 label을 준비하며 features 벡터 열은 코드가 생성합니다. CBAC 환경에서는 모델 출력을 미리 생성합니다.

from transforms.api import configure, transform, Input
from palantir_models.transforms import ModelOutput
from pyspark.ml.feature import VectorAssembler
from xgboost.spark import SparkXGBClassifier
from main.model_adapters.adapter import BoosterModelAdapter

@configure(profile=[
    "DYNAMIC_ALLOCATION_DISABLED",
    "NUM_EXECUTORS_2",
    "KUBERNETES_OPEN_PORTS_ALL",
])
@transform(
    training_df=Input("/Models/data/binary_training"),
    model_output=ModelOutput("/Models/trained/distributed_classifier"),
)
def compute(training_df, model_output):
    frame = training_df.dataframe()
    assembler = VectorAssembler(
        inputCols=BoosterModelAdapter.FEATURE_COLUMNS,
        outputCol="features",
    )
    estimator = SparkXGBClassifier(
        features_col="features", label_col="label", num_workers=2,
    )
    spark_model = estimator.fit(assembler.transform(frame))
    model_output.publish(model_adapter=BoosterModelAdapter(spark_model.get_booster()))

예제는 실행기 두 개와 학습 작업자 두 개를 맞추고 동적 할당을 끕니다. 일반 회귀에서는 SparkXGBRegressor를 선택할 수 있습니다. 학습 및 추론 데이터의 특성 열과 순서는 연결한 어댑터의 계약에 맞춥니다.

GPU 분산 학습

프로젝트를 GPU가 있는 리소스 큐에 할당하고 필수 통신 프로필에 EXECUTOR_GPU_ENABLED를 추가합니다. XGBoost 추정기에는 GPU 설정에 맞게 device="cuda" 또는 device="gpu"를 지정합니다. CPU 설정은 device="cpu"입니다. Spark가 GPU를 할당하므로 cuda:0처럼 장치 번호를 직접 지정하지 않습니다. GPU용 Python 패키지의 환경 해결은 GPU 학습의 조건을 따릅니다.

하이퍼파라미터 최적화와 추론

변환에서 분산 학습의 하이퍼파라미터를 탐색하려면 Optuna를 사용할 수 있습니다. Optuna는 Spark 학습과 함께 사용할 수 있으며 별도의 Foundry 연동 설정은 필요하지 않습니다. 탐색할 매개변수와 평가 기준은 학습 코드에서 정합니다.

게시한 모델을 분산 추론하려면 PySpark UDF에서 ModelAdapter를 호출할 수 있습니다. Pandas 입력 하나와 출력 하나인 모델은 DistributedInferenceWrapper의 실행기 분산 경로도 사용할 수 있습니다.