서버리스 ML 훈련 인프라 구축하기: Vertex AI Pipelines & TFX
당근 ML 인프라팀에서는 머신러닝 모델을 효과적으로 다양한 서비스에 적용하기 위해 노력하고 있어요.
이전 블로그 글(Kubeflow 파이프라인 운용하기)에서 언급했던 것처럼 당근마켓은 Google Kubernetes Engine(GKE)에서 Kubeflow Pipeline(KFP)을 설치해 적극 활용하고 있었는데요. 2023년 초기부터 Vertex AI Pipelines를 도입하고, 마이그레이션 해왔어요. 그리고 지난 10월에 드디어 전체 파이프라인을 Vertex AI Pipelines에 모두 옮기면서 KFP 및 이를 운영하기 위한 GKE 자원을 모두 제거했어요.
이번 글에서는 당근이 어떻게 TFX와 함께 Vertex AI Pipelines를 활용하고 있는지 공유해 보려 해요.
Vertex AI Pipelines
Vertex AI Pipelines는 머신러닝 파이프라인을 Serverless 방식으로 관리 및 실행할 수 있는 Google Cloud Platform 제품이에요. ML 워크플로우에 필요한 데이터 전처리, 모델 훈련 / 평가, 배포 등등 머신러닝 프로젝트가 밟아야 하는 여러 단계를 Vertex AI 및 GCP 제품들로 구성하고, 자동화할 수 있어요.
2024년 1월 기준으로 Vertex AI Pipelines는 아래 두 가지 인터페이스를 지원하고 있어요.
- Kubeflow Pipelines(KFP) SDK
- TensorFlow Extended(TFX) SDK
기존에 잘 활용하고 있던 TFX SDK를 Vertex AI Pipelines와 함께 사용 중이에요.
Managed Service인 Vertex AI 도입 이후 Kubernetes Cluster 운영 또한 필요하지 않아요. 그 때문에 운영보다 ML 워크플로우에 필요한 요소들을 개발하는 데 집중할 수 있어요.
TensorFlow Extended
TFX는 프로덕션 머신러닝 파이프라인을 위한 end-to-end 플랫폼이에요. 2019년에 완전히 오픈소스화되었고, 컴포넌트 단위로 ML 워크플로우를 개발한 후 여러 환경(Apache Beam, Dataflow, Kubeflow, Airflow) 등에서 실행 가능해요. 물론 데이터 수집/변환, 학습, 배포 등을 위해 잘 작성된 컴포넌트도 같이 제공되고 있어요.
TFX에 대한 좀 더 자세한 사항은 이전 블로그 글(TFX와 함께 머신러닝 파이프라인 개발하기)을 참고해 주세요.
생산성 및 편의성
머신러닝 모델이 실제 환경에서 잘 동작하려면 모니터링, 신뢰도와 유효성 검사 등 프로덕션 레벨의 요구사항을 잘 만족해야 해요. 이를 위해 TensorFlow Extended(TFX)를 적극 활용하고 있지만, TFX에서 제공하기 힘든 당근에서 필요로 하는 요구사항에 대해서는 ML 인프라팀이 나서야 해요.
당근은 각 서비스 조직이 정형화된 ML 인프라/MLOps 모듈들을 개발하는 데에 시간을 쏟지 않고, 담당하시는 서비스에 필요한 개발에 집중하길 바라요. 반대로 ML 인프라팀이 너무 한두 개의 서비스에 커플링되는 모듈에 집중한다면 계속 성장하며 늘어나는 ML 서비스 조직을 결국 감당하지 못할 거예요. 이런 고민들을 가지고 Vertex AI Pipelines를 활용하려 노력하고 있어요.
내부 공통 컴포넌트 라이브러리
TFX에서 기본으로 제공해 주는 컴포넌트를 통해서도 대부분의 ML 워크플로우를 소화할 수 있지만, 서비스 요구사항이나 생산성, 편의성을 향상하려다 보면 당근에 맞는 커스텀 컴포넌트가 필요한 순간이 오게 된다고 생각해요.
tfx-addons 처럼 당근 내에서 자주 사용되는 TFX 컴포넌트들을 Python Package 형태로 관리하고, 배포해요. 각 팀에서 필요한 컴포넌트를 개발해 사용하면서도, 여러 부서에서 사용할 만한 컴포넌트들은 다 같이 공유하고 있어요. 이러한 방식을 통해 여러 서비스 조직에서 적은 공수로 생산성과 편의성을 취할 수 있고, 비용 및 성능 최적화가 필요한 시점에도 빠르게 전파가 가능해요.
아티팩트
Vertex AI Pipelines는 KFP v1의 파이프라인들보다 아티팩트의 시각적인 표현이 강조되어 있어요. 기존 KFP v1에선 각기 컴포넌트가 어떤 아티팩트를 생성/소모되는지 표현되지 않았지만, 이제는 아래 그림처럼 컴포넌트 및 아티팩트 간 관계를 명확하게 보여줘요. 따라서 해당 아티팩트가 어떤 타입인지, 아티팩트가 어디 저장되어 있는지, 필요하다면 상호작용 버튼까지 제공하고 있어요.
위 StatisticsVisualizer로 통계상 이상한 데이터가 없는지 체크하기 위해 tfma 패키지를 활용해 html 파일을 생성하는 컴포넌트인데요. 해당 컴포넌트가 만들어내는 html 파일은 vertex_ai_stats_html 아티팩트에서 그 결괏값들을 저장하고, 필요하다면 사용자가 Web UI에서 그 결과를 바로 확인할 수 있게 만들어줘요. 이러한 기능들은 모델을 개발하시는 분들에게 여러모로 편의성을 제공해 주고 있어서 적극 활용하고 있어요.
Recurring Run
Recurring Run은 주기적으로 돌아야 하는 파이프라인에 CronJob처럼 스케줄을 걸어 반복해서 돌려주는 기능을 말해요. 해당 기능을 통해 매일 최신 유저 데이터를 받아 학습해야 하는 모델들을 일정 주기로 실행할 수 있는데요.
Vertex AI Scheduler API가 등장하기 전까진 GCP의 Cloud Function과 Cloud Scheduler를 조합해서 Recurring Run을 구현하는 게 최선이었어요. 다만, 아직 TFX와 Vertex AI Scheduler API가 호환되지 않는 문제가 있어서 아래 구조도처럼 구성해 활용하고 있어요.
리소스 할당 방법
모델 학습 시 알맞은 CPU, Memory 및 GPU를 선택해 줘야 하는데요. KFP v2를 그대로 사용한다면, .set_cpu_limit()와 같은 함수를 통해 리소스 할당이 바로 가능하지만, TFX SDK와 Vertex AI Pipelines를 함께 사용한다면, DAG 컴파일 시 생성되는 파이프라인 스펙 JSON 파일을 직접 수정해 각기 자원들을 주입해 줘야 했어요.
이런 기능들을 위에서 언급한 확장 라이브러리에 추가해 별다른 수정 없이 전체 ML 엔지니어들이 활용할 수 있도록 제공하고 있어요.
모니터링
Resource Utilization
Vertex AI Pipelines는 Component 별로 리소스를 지정해 줄 수 있고, 이와 관련된 리소스 모니터링도 자체적으로 지원하고 있어요. 따라서 사용자들이 파이프라인에 문제가 발생했을 때 리소스와 관련된 문제인지 단번에 파악할 수 있고, 정상적으로 완료되는 상황에서도 리소스 효율을 체크해 학습 코드를 최적화하기 좋은 환경을 제공해 줄 수 있어요.
뿐만 아니라 위 정보들은 모두 Metric으로 남아 아래처럼 Dashboard를 만들거나, 서드파티 제품들과 데이터를 연동해 시각화하거나 알람을 만들 수 있어요.
Custom Metric
GCP에서 기본적으로 제공하는 Google Cloud Metric이 조직의 요구사항과 맞지 않다면, Custom Metric을 구성해 모니터링을 강화할 수 있어요.
예를 들어 Vertex AI Pipelines 로그 중 RUNNING, PENDING과 같이 지속되는 상태를 표현해야 하는 로그가 단발성으로 한 번만 로깅이 되어 현재 학습 중인 파이프라인의 개수나 팬딩중인 파이프라인들을 대시보드로 구성하기 어려워 직접 Custom Metric에 Continual State Event를 기록해 아래와 같이 대시보드로 구성해 현재 도는 파이프라인을 볼 수 있도록 모니터링 구성을 했어요.
이 밖에도 추가로 모니터링이 필요한 리소스나 데이터 통계들을 Google Cloud에 Custom Metric으로 등록해 기록한다면 Dashboard를 구성하거나 서드파티 제품들과 연동할 수 있어요.
Alert
위에서 언급한 Custom Metric을 통해 에러 및 특정 상황 발생 시 아래처럼 Slack 알람을 보내고 있어요. 담당자가 파이프라인에 문제가 생겼을 때 곧바로 인지할 수 있고, 관련 파이프라인 및 실험 페이지로 곧바로 이동할 수 있도록 구성했어요.
이를 통해 담당자 코드 에러인지 Vertex AI Pipelines 문제인지 에러 메시지들을 통해 쉽게 알 수 있어서, 온콜 부담 또한 많이 줄어들었어요.
End of Support(EoS) 관리
Vertex AI Pipelines가 Managed Service인 만큼 당근에서 활용하는 TFX 및 TensorFlow와 같은 ML 프레임워크 오픈소스들의 버전이 EoS(End of Support date)가 지났는지 체크하는 게 중요해요. 구글 측은 Vertex AI Docs — Supported frameworks for Vertex AI pipelines 문서에서 각 프레임워크의 EoS를 명시하고 있어요.
EoS가 지나간 파이프라인 및 서비스들은 문제가 생겼을 때 Google Cloud Platform(GCP) 측에서 “버전 업그레이드”라는 대안을 제시할 가능성이 커요. 또한 해당 ML Framework들이 오픈소스이기 때문에 버전 호환성에 대한 책임이 GCP 측에 있지 않기에 미리 대응하지 않으면 잘 돌던 프로덕션 파이프라인들에 오류가 발생할 가능성이 있어요.
당근에서는 EoS를 체크하기 위해 아래처럼 대시보드 그래프를 구성하고, 모니터링하고 있어요.
각기 파이프라인의 EoS를 챙기는 과정에서 새로운 버전 기능이나 최적화 패치들이 적용되면서 더 좋아지는 경우도 생기고, 최신 버전들을 내용을 학습하며 다음 프로젝트를 구현할 때 참고하는 등 긍정적인 부분도 많아요.
마치며
당근 ML 인프라팀은 새로운 트렌드에 발 빠르게 대응하며 머신러닝 모델을 효율적이고 안정적으로 학습/배포할 수 있도록 노력 중이에요. 머신러닝이 서비스 활용도가 정말 높은 기술 중 하나라 믿고, 당근에서 머신러닝이 더욱 잘 활용될 수 있도록 노력하고 있어요.
전사 머신러닝 인프라/공통 컴포넌트 개발을 통해 당근 서비스에 지속해서 큰 임팩트를 만들고 싶은 동료를 찾고 있어요.
