토스의 자체 Feature Store + Trainkit으로 ML 학습 파이프라인 정형화

요약

토스 ML Platform 팀이 만든 자체 Feature Store와 학습 파이프라인 도구 Trainkit 소개. 사내 ML 모델 수십 개가 운영 중이지만 팀마다 데이터 로딩 방식·전처리 코드·Feature 관리가 달라 중복 구현과 재현성 문제가 반복되던 환경에서, Feast의 데이터 모델 개념만 참고해 자체 구축. Online Storage는 Aerospike, 학습 파이프라인은 PyTorch Lightning 위에 5개 컴포넌트(Data Module / Target / Feature Package / Feature Service / Feature Processor)로 구조화.

내용

토스 안에는 신용평가·이상거래탐지·마케팅 최적화·고객 경험 개선 등 다양한 비즈니스 도메인에 걸쳐 ML 기반 서비스가 수십 개. 각 모델은 서로 다른 팀·데이터 소스·서빙 요구사항을 가짐. 모델 학습 파이프라인 자체는 데이터 로딩 → 전처리 → 모델 학습 → 후처리 → 평가로 표준화되어 보이지만, 실제로는 "모델 학습" 외 단계가 팀마다 제각각으로 구현되어 있던 게 출발점.

같은 데이터셋을 두고도 한 팀은 Impala SQL로 로딩하고 다른 팀은 HDFS Parquet를 직접 읽음. 그 결과 데이터 로딩·전처리 스크립트가 프로젝트마다 새로 쓰여 작은 버그 수정도 여러 곳에서 반복되고, 셔플링·체크포인트 저장 로직 차이 때문에 같은 실험을 돌려도 결과가 달라지며, 신규 입사자는 프로젝트별 파이프라인 구조를 다시 익혀야 하는 상태. 거기에 학습과 서빙에서 보는 Feature 시점이 어긋나면 Training-Serving Skew로 모델 성능까지 흔들림. 자체 Feature Store + Trainkit 두 축을 같이 깔게 된 배경.

오픈소스 Feature Store(Feast / Tecton / Hopsworks)는 검토했지만 토스 데이터 구조·보안 요구에 맞지 않고 운영 시 비슷한 리소스가 필요하다는 판단. Feast의 Data Model 개념(Entity / FeatureTable / FeatureService)만 참고하고 직접 구축.

해결 / 접근

Toss Feature Store

사용자는 Admin Dashboard와 SDK로 Feature Table을 등록 → 사내 Data Center와 연동되어 컬럼이 Entity / Feature / Partition으로 구조화된 채 Feature Store 소스로 들어옴. Online Storage는 Aerospike — Key는 메모리, Record는 Disk로 두는 메모리·SSD 하이브리드라 대용량 Key-Value 트래픽에 Low Latency를 보장. 학습용은 Offline Storage, 서빙용은 Aerospike + Feature Serving API로 분기.

데이터 모델 4개:

• Entity — 테이블 안에서 유니크한 키, 학습 데이터의 Example을 나누는 단위
• Feature — 테이블 컬럼, 데이터 관리 단위
• FeatureTable — Impala 등 물리 테이블을 Wrapping. Entity / Features / Partition Column으로 명세
• FeatureService — ML 모델 학습·Inference에 쓰이는 논리 단위. FeatureTable 리스트와 Entity로 정의

Feature Store Admin은 사내 Data Center와 연동되어 Feature 메타데이터·품질 모니터링·권한 관리까지 끌어옴.

Trainkit — 학습 파이프라인 5개 컴포넌트

PyTorch Lightning의 LightningDataModule을 상속한 Data Module을 입구로, 다음 다섯이 한 묶음:

• Data Module — 데이터 로딩 → 전처리 → 배치 생성을 묶은 파이프라인 객체
• Target — 여러 Feature를 조인할 때 기준이 되는 데이터 스펙. 추천이라면 사용자 반응 데이터셋
• Feature Service — Feature Store API와 연결, 메타정보 조회 + Target과 Feature 사이 Join SQL 자동 생성
• Feature Processor — Casting / Imputation / Scaling / Encoding 같은 전처리 로직. fit / transform만 정의하면 신규 전처리 추가 가능
• Feature Package — Feature Service와 Feature Processor를 들고 있는 Facade. 캐싱·유효성 검사·로딩 API 담당

워크플로 — Target으로 조인 기준 데이터를 가져와 전처리 → n개 Feature Package에서 Feature Service 정보를 보고 Feature 데이터 가져와 전처리 → Train / Validation / Test Dataset과 Dataloader 생성.

핵심 기능 3가지

멀티 Feature Package 조인 — 토스 추천 모델은 Feedback / User / Item 카테고리를 쓰는데, 데이터 소스 구조까지 이 카테고리에 묶으면 확장이 어려워짐. Trainkit은 Target 하나를 기준에 두고 임의 개수의 Feature Package를 결합. 사용하고 싶은 Feature Service만 지정하면 Trainkit이 메타데이터 기반으로 Join Key를 만들어 Join SQL과 데이터를 한 번에 산출.

Training-Serving Skew 해소 — 학습 시 본 Feature 시점과 서빙 시 본 시점이 어긋나면 데이터 누수가 발생하고 오프라인·온라인 성능이 갈림. 같은 시간 파티션 해석도 팀마다 달라(00–01시 데이터를 00시 파티션에 넣을지 01시에 넣을지) 조인이 어긋나는 경우가 있음. Trainkit은 Target 데이터 기준으로 시간 파티션을 Shift하는 설정을 추가해, Feature를 재생성하지 않고 설정 변경만으로 Skew·누수 대응.

PIT(Point-In-Time) 조인 — 온라인 서빙 시점의 Feature 값으로 추론한 결과를 오프라인 학습에서 다시 쓰려면 시점이 맞는 Feature 값을 찾아 반응 데이터와 결합해야 함. 미리 조인하면 메모리가 기하급수로 늘어남(예: target 6 + User Feature 6×3 + Item Feature 6×3 = 42 vs 런타임 조인 6 + 3 + 3 = 12). Trainkit은 객체 생성 시 join key와 Feature Package 딕셔너리만 준비해두고, Data Module이 index를 호출할 때 런타임 조인.

결과 / 참고

• Online Storage: Aerospike (Key 메모리·Record SSD 하이브리드)
• Trainkit 베이스: PyTorch Lightning LightningDataModule
• 메모리 비교 예시: 미리 조인 42 vs 런타임 조인 12 (target 6, User Feature 6×3, Item Feature 6×3 가정)
• 시리즈: 본 글은 Feature Store + 학습 파이프라인 편. Model Serving·Monitoring은 후속 글
• 출처: 토스 기술 블로그, "토스가 다양한 ML 모델을 만드는 법: Feature Store & Trainkit"