TFX와 함께 머신러닝 파이프라인 개발하기

Ukjae Jeong
당근 테크 블로그
9 min readMar 30, 2022

당근마켓에서는 머신러닝 모델을 효과적으로 적용하기 위한 방법에 대해서도 노력하고 있어요. 특히 ML 파이프라인은 주기적인 배포, 빠른 실험, 지속적인 모델 개선을 위해 필수가 되는 경우가 많으므로, 더 많이 노력해야 하는 대상 중 하나예요.

이전 블로그 글(머신러닝 모델로 동네생활 신고 업무 자동화하기)에서도 언급했던 것처럼 머신러닝 파이프라인에 TFX(TensorFlow Extended)를 활용하고 있는데요, TFX를 활용하는 이유와 어떤 방식으로 더 잘 사용하려고 하는지를 공유해보려 해요.

당근마켓의 ML 프로젝트

당근마켓 내부에는 많은 ML 프로젝트가 존재해요. 여러 추천(랭킹) 모델 과 중고거래, 동네생활 신고 업무 자동화 모델과 함께 서비스 내부 곳곳에 머신러닝 모델이 들어있어요. 상세한 프로젝트에 대한 소개는 이 포스트의 범위를 넘어서니, 프로젝트 상세 내용에 더 관심있으시다면 아래 블로그 글들을 참고해주세요.

위의 블로그 글 이외에도 상당히 많은 수의 머신러닝 모델을 프로덕션에 활용하고 있어요. 이전 블로그 글(Kubeflow 파이프라인 운용하기)에서 언급했던 것처럼 머신러닝 모델을 주기적으로 새로운 데이터로 학습하고, 개선하고, 배포하기 위해 Kubeflow를 많이 활용하는데요. 최근 1~2년 동안은 그 위에 TFX라는 기술과 같이 사용하고 있어요. 프로덕션용 머신러닝을 개발하기 좋은 TensorFlow와 함께 사용한다면 굉장히 편하게 머신러닝 모델 및 파이프라인 개발, 실험, 배포를 진행할 수 있어요.

TFX (TensorFlow Extended)

TFX는 프로덕션 머신러닝 파이프라인을 위한 end-to-end 플랫폼이에요. 2019년에 완전히 오픈소스화 되었고, 컴포넌트 단위로 ML 워크플로우를 개발한 후 여러 환경(Apache Beam, Dataflow, Kubeflow, Airflow)등에서 실행가능해요. 물론 데이터 수집/변환, 학습, 배포 등을 위해 잘 작성된 컴포넌트도 같이 제공되어요.

TFX Pipeline Image from https://www.tensorflow.org/tfx

당근마켓 ML 팀에서 특히 크게 느끼는 이점들을 아래처럼 설명드려볼게요.

기본 컴포넌트

TFX는 수많은 컴포넌트를 미리 제공해요. 데이터 수집만 살펴보아도 CSV 기반인 CsvExampleGen부터 Presto, BigQuery 에서 직접 데이터를 수집할 수 있는 PrestoExampleGen, BigQueryExampleGen도 제공하여 여러 데이터 소스를 컴포넌트 연결만으로 손쉽게 처리할 수 있어요.

대용량 데이터 가공도 간편하게 처리할 수 있어요. 데이터 변환을 담당하는 Transform 컴포넌트는 Apache Beam기반으로 작성되어 있기 때문에 GCP Dataflow에서 간단하게 분산 실행할 수 있어요.

그 외에도 많은 편리한 컴포넌트가 존재하고 계속해서 추가되고 있어요.

컴포넌트 재사용

실제로 프로덕션용으로 사용하기 위해서는 사용자 정의(Custom) 컴포넌트도 손쉽게 활용이 가능해야 하는데, 이 또한 쉽게 가능해요. 기본 TFX 컴포넌트도 모듈화하기에 용이한 구조를 가지고 있어요.

단순 Python 함수나 컨테이너를 TFX 컴포넌트로도 사용 가능하고, 미리 제공되는 TFX 컴포넌트와 같이 완전 사용자 정의 컴포넌트도 만들 수 있어요. 직접 작성한 컴포넌트도 당연히 기존의 TFX 파이프라인에 손쉽게 연결하여 사용할 수 있어요.

이러한 장점을 극대화하기 위해, 여러 프로젝트에서 비슷하게 사용하는 TFX 컴포넌트는 당근마켓 내부 라이브러리로 개발해 서로 공유하고, 누구든지 이 기능을 손쉽게 사용할 수 있도록 하고 있어요.

다양한 Runner 제공

TFX에서는 다양한 환경과 호환돼요. 로컬에서 CPU/GPU를 활용해 직접 실행할 수도 있고, 대용량 데이터를 전처리하기 위해 Apache Beam과 호환되는 DataFlow를 사용할 수도 있어요. 또 시각화 등을 위해 노트북 상에서도 하나씩 수동으로 실행할 수도 있고요. 물론 KubeFlow 위에서도 실행이 가능하고, 최근 많은 기능이 출시되고 있는 Vertex AI 위의 기능들도 활용해 실행이 가능해요.

한번 작성해서 여러 곳에서 동일하게 실행이 가능하다면 손쉽게 개발 환경과 배포 환경을 만들 수 있어요. TFX를 활용해 모델을 개발한다면 프로덕션 모델 배포에 대한 부담이 확연히 적어지게 돼요.

TFX를 사용하면서

TFX로 구성하면서 ML 파이프라인의 품질은 좋아졌고, 개발하시는 분들의 만족도도 높았어요.

하지만 TFX 자체가 비교적 최근의 기술이고, 당근마켓의 경우에는 1.0 버전 이전부터 사용하다보니 각자 다른 스타일로 TFX 파이프라인 코드를 작성하게 되었어요. TFX는 빠르게 발전하고 있고 버전이 올라가는 속도도 비교적 빠르기 때문에 프로젝트에서 사용하는 TFX 버전도 많이 달라지게 되었고요. 그래서 최신 ML 파이프라인 기술을 활용하거나 코드 관리/파악 등이 힘들어졌어요.

게다가 파이프라인이 크고 복잡해지게 되면서 각 컴포넌트와 설정 값의 의미, 그리고 설정 값이 어떤 컴포넌트에 영향을 주는지 알기 어렵게 변해갔어요

TFX를 더 잘 사용하기 위해서

서로의 코드 파악이 더 쉬워지고 어느정도 같은 패턴을 가지고 쉽게 서로 간에 노하우를 공유할 수 있도록 TFX 템플릿 프로젝트를 만들어 활용하기로 했어요. 당근마켓 내부에서 자주 사용되는 컴포넌트를 미리 통합해, 데이터 쿼리와 모델 학습 로직만 수정하면 빠르게 머신러닝 모델을 개발할 수 있도록 구성했어요.

새로운 프로젝트의 개발 자체가 빨라지는 효과도 당연히 기대하지만, 위에서 언급한 것처럼 각 프로젝트가 유사한 구조를 가지게 되고, 서로의 프로젝트 파악이 쉬워지는 효과를 많이 기대하면서 구성했어요.

템플릿 프로젝트 자체에 대해 대략적으로 소개드렸는데, 해당 프로젝트를 구성하면서 TFX를 더 잘 사용하기 위해 고려한 사항들을 소개해 드릴게요.

Configuration First

파이프라인을 가능한 하나의 설정 파일로 모든 것을 관리할 수 있어야 해요. A/B 테스팅을 위한 더 빠른 실험을 진행한다고 가정 했을때 특정 설정이 어떤 역할을 하는지 알기 힘들다면 의도대로 정확히 테스트하기가 힘들어요. 파이프라인을 구성하는 방식을 코드를 통해 먼저 소개할게요.

example_gen_config.protoProtobuf로 작성된 설정의 명세를, config.pbtxt는 설정을, 그리고 pipeline.py는 실제 파이프라인을 구성해요.

ExampleGen의 모든 설정을 ExampleGenConfig 메시지 하나로 결정돼요. 비슷하게 모든 컴포넌트가 설정값에만 의존해서 생성되고 연결되어요. 이는 설정만 보고도 프로젝트의 구조를 파악할 수 있도록 하기 위함인데요, 이렇게 정의해두고 사용한다면 설정파일만 보더라도 파이프라인이 포함하는 컴포넌트와, 각 컴포넌트의 구성을 쉽게 알 수 있어요. 각 컴포넌트를 정의하는 부분을 분리함으로써 커스터마이즈 및 코드 파악이 쉽게 만들려는 의도 또한 있어요.

예시를 위해 다음과 같은 상황을 가정해볼게요. 추후에 데이터 변환을 테스트하기 위해 Transform 컴포넌트에서 다양한 처리 방식을 지원해야 해요. 특정 방식은 데이터 증강(Data Augmentation) 기법을 사용하고 다른 방식은 사용하지 않는 경우를 생각해볼 수 있을 것 같아요. 이 때도 이런 방식을 사용한다면 미리 정의된 Protobuf 명세를 확장해 쉽게 다른 처리 방식을 지원하도록, 또 어떤 처리 방식을 쓰는지 쉽게 알아보도록 만들 수 있어요.

Protobuf로 Config 관리하기

위의 예시 코드를 추가했을 때 몇몇 분들은 Protobuf를 왜 설정 도구로 사용하는지 의문이 있는 분들도 계실거예요. 이는 여러 이유가 있는데요, 설정을 위해 많이 사용되는 YAML과 비교해 주요 포인트를 소개드릴게요.

먼저 구조적으로 강건하고 타입 체크 등의 Validation이 편해요. 구조를 미리 정의하고 들어가는 Protobuf 특성상 필드 자체가 정의되어있는지 확인할 필요가 없고, 타입 확인도 편해요. 같은 설정을 하더라도 활발히 개발이 되고 있는 프로젝트에서 backward/forward compatibility를 지원하며 개발하기 좋아요.

또한 구조도 한눈에 볼 수 있게 돼요. 계층 구조를 가진다는 것은 YAML도 비슷하지만 머신러닝 생태계에서 자주 사용되는 hydra의 경우에 여러 파일로 스테이지의 설정이 나누어지다보니 Protobuf에 비해 안정성, 가시성 등이 떨어진다고 판단했어요. (이 부분은 어느정도 취향이 반영된 것이라 생각해요 😅)

TFX 자체에서도 컴포넌트 설정을 위해 Protobuf를 광범위하게 사용하고 있어요. 이 상황에서 Protobuf를 주 설정 도구로 사용하게 된다면 TFX에 정의되어 있는 상당 수의 Protobuf 정의를 재사용할 수 있어요. 여러 컴포넌트를 위한 설정 로직을 직접 작성하지 않더라도 손쉽게 확장할 수 있어요.

Bazel과 함께 하는 TensorFlow 생태계

Bazel은 언어 독립적인 빌드 시스템인데요, 확장이 간편하고 다양한 언어와 도구를 지원하고 있어요. 간단한 프로젝트부터 여러 언어, 도구를 사용하는 대형 프로젝트까지 대부분의 상황에서 빠르고 간결하게 활용할 수 있어요. 대표적으로 TensorFlow, Jax, Envoy 등이 Bazel을 사용하는 프로젝트예요. 자세한 내용은 Bazel 문서 페이지 중 Bazel Vision을 참고하시면 더 좋을 것 같아요.

Python 프로젝트에서 이러한 Bazel을 활용하는 것은 흔하지 않은 세팅이라 생각하지만, TFX 템플릿 프로젝트를 구성하며 Bazel을 주요 빌드 시스템으로 가져가보았어요. 이유를 가볍게 소개드릴게요.

우선 Protobuf와 사용할 때 정말 좋은 궁합을 보여요. Protobuf를 사용하는 프로젝트는 Protobuf 명세를 수정할 때 일일이 protoc로 컴파일하기에는 사람이 실수할 여지가 많아져요. Bazel은 언어 독립적인 빌드 시스템이기 때문에 걱정 없이 Protobuf 빌드 결과물을 다른 빌드와 손쉽게 의존성으로 엮을 수 있어요. Protobuf 자체가 Bazel을 상당히 많이 이용해서 쉽게 서드 파티 라이브러리로 쓸 수 있는 것은 덤이고요.

그 다음으로 TensorFlow 생태계의 특수한 환경도 있다고 생각해요. TensorFlow 생태계 상의 많은 프로젝트가 Bazel을 활용하고 TFX 또한 Bazel을 활용하기 때문에 Bazel로 다른 프로젝트(TensorFlow, TFX)와 빌드를 쉽게 엮을 수 있어요.

마치며

이렇게 당근마켓에서는 머신러닝을 더 잘 활용할 수 있는 기술을 위해서 노력하고 있어요. 머신러닝이 서비스에 활용도가 정말 높은 기술 중 하나라고 믿고 있고, 안전한 서비스를 만들고 개인화를 잘 해내기 위해서는 더더욱 필수라고 생각해요. 이러한 미션을 위해서는 필요한 시기에 좋은 MLOps 기술 기반을 갖춰야해요.

머신러닝을 이용해 큰 임팩트를 주고 싶으신 분들은 언제든지 당근마켓 머신러닝 엔지니어 채용공고를 참고해주세요! 😄

마지막으로 좋은 ML 파이프라인을 위해 다 같이 노력하는 머신러닝 엔지니어 언제나 감사합니다 :)

--

--