ML Pipeline에 Flyte 도입하기

Ryan Kim
10 min readJun 18, 2022

ML 오픈 소스의 망망대해 항해하며, ML Pipeline 구성해 나가기

망망대해를 헤쳐 나가는 것에 정답은 없다.

현재 3명으로 구성된 회사의 MLOps 팀은 (CTO님을 포함하면 4명)매일 크고 작은 일들을 해결하고, 배포하는 과정을 수행하고 있다.

가장 최근에는 ML 리서치 엔지니어 분들을 위해 Jupyterhub를 쿠버네티스에 배포해서 운영에 관련한 일들을 수행했고 (주된 업무를 수행하는 것은 팀의 다른 멋진 개발자님이 진행하고 있다), 나는 MLOps 팀 리드 개발자님과 함께 Flyte라는, ML 워크 플로우를 관리하는 플랫폼을 사내에 도입하게 되었다.

자~~~ 들어가자!

Flyte란?

Flyte는 Lyft라는 미국 스타트업에서 (우버 경쟁사다) Spotify 등과 함께 만들어 오픈 소스로 배포한 플랫폼으로, Data Processing과 ML Workflow를 관리할 목적으로 제작되었다.

동시성, 확장성, 유지 가능성 등을 쉽게 유지할 수 있다고 한다.

Flyte의 공식 페이지의 설명에 따르면, 10,000개 이상의 독립적인 워크 플로우를 Lyft에서 관리하고 있으며, 매달 1백만 개 이상 실행하고 있고, 2천만개 이상의 task를 돌리고 있다고 안내하고 있다.

참고로 Flyte에서 가장 작은 실행 단위는 Task이고, 이 Task를 각 각 컨테이너로 띄워서 하나의 Workflow에 묶어 실행하는 것이 가능하다.

(4천만 개 이상의 컨테이너를 사용 중인데, 왜 이렇게 많이 사용하는지는 뒤에서 자세하게 설명하겠다. task마다 컨테이너 설정이 가능하다)

처음 회사에서 flyte를 도입할 때 자료를 찾으면서 작년 이 맘 때쯤, MLOps 전성 시대라는 주제로 발표한 자료에서 flyte가 잠깐 언급된 적이 있었는데, 그 이후 한국 커뮤니티에서 Flyte를 도입한 사례를 찾아다녔는데, 활용한 사례에 대해 기업들의 사례가 전혀 존재하지 않았다.

검색어 최상단에 flyte가 안나온다.

즉, 지금 내가 근무하고 있는 Qraft Technologies라는 회사에서 ML Pipeline에 Flyte라는 서비스를 도입하기 위해서는 한국에서 MLOps 오픈 소스의 망망 대해를 항해해야하는 상황이라는 것을 본능적으로 직감하게 되었다.

아마 누군가는 도입해서 사용하고 있을텐데, 기술 블로그에 공유를 안한 것이겠지…

Flyte를 대체할 수 있는 선택지는 몇 가지 있었고, Google에서 제공하는 TFX, AWS SageMaker, GCP Vertex AI 등을 언급할 수 있지만, 해당 대안들은 아래와 같은 사항들을 고려할 때, 우리에게 “최선”의 선택은 아니었다.

  • 전사적으로 ML 리서치 엔지니어 분들이 파이 토치를 사용해서 모델을 작업하고 있다. (TF 등 여타 DL 플랫폼에서 사용되는 TFX는 배제)
  • 회사 내부적으로 사용하고 있는 스펙 짱짱한 온프레미스 서버가 운영 중에 있다. (클라우드를 통한 모델링은 불필요하니 배제)
  • 충분한 Use Case가 있으며, 엔터프라이즈 규모의 기업에서 운영되고 있어야한다. (Lyft, Spotify 등의 글로벌 기업들의 Use Case가 있다.)

일단 내가 생각한 선택 사유는 이 정도라고 판단했지만, 회사에서 미래에도 온프레미스 환경에서 모델링 업무가 실행될 것이라는 것은 명확하기 때문에 (회사가 핀테크 스타트업이라, 개인정보 등이 포함되는 데이터를 국외 이전하는 것이 불가능해서 온프레미스 서버를 반드시 운영해야한다 — 개인정보 데이터 등이 다뤄지는 데이터는 클라우드 이용이 거의 불가능하다) Flyte를 선택하는 것이 최선으로 보였다.

그리고 위의 선택지를 포함하는 것에 ML 워크 플로우 플래폼에 플러스 알파로, ML or Data Processing 작업이 항상 동기적으로 동작하면 안되는데 그 이유는 여러 엔지니어가 동시 다발적으로 Flyte에 workflow를 띄울 때, 동기적으로 실행되면 안되었으며, Batch Job이 Flyte에 지원되는 것도 결정에 영향을 줬다.

Flyte 도입, 그리고 끝없는 버그

끝없는 버그…

Flyte를 띄울 환경은 k8s를 기반으로 했으며 (이미 많은 애플리케이션들이 k8s를 기반으로 돌고 있고) 다행이도 Flyte는 k8s native한 환경을 제공하고 있다.

개인적으로 Flyte를 도입하면서 사용할 수 밖에 없는 가장 큰 Critical point는 task마다 GPU를 할당할 수 있다는 점이다.

이게 무슨 얘기냐면, 다음과 같은 Flyte 공식 예제를 통해 확인할 수 있다.

k8s에 리소스를 필요한 만큼 요청할 수 있다.

이 부분이 왜 Critical한 것으로 수용되었냐하면, 기존에 JupyterHub라는 것을 도입해서 운영할 땐, 개별 리서치 엔지니어마다 Jupyter를 실행하기 위한 인스턴스를 할당 받았고, 이 과정에서 해당 인스턴스에 GPU를 할당 받아 개발을 진행할 수 있었다.

그런데 만약 해당 인스턴스에 할당된 GPU의 사이즈가 큰 것이고 (DGX 40G같이 최대 크기) 다른 리서치 엔지니어들이 추가적으로 해당 GPU를 필요로 하는 상황이면, 여기서 Bottle neck이 발생할 수 있었다. (Jupyterhub는 리서처가 작업을 하지 않는 시간이 24시간 이상이 되면, 자동으로 인스턴스를 종료하여 들고 있는 자원들을 반환할 수 있도록 해놨다. 또는 관리자가 강제 종료도 가능)

하지만 Flyte는 task 마다 작업이 끝나면, 알아서 사용을 했던 자원을 반환하고, 다음 task에 동적으로 매핑이 되었기 때문에 상대적으로 MLOps 팀에서 자원 관리에 인적 자원이 덜 소요될 것으로 예상했다. (리서치 엔지니어 분들이 우리에게 “혹시 지금 MIG-5g.40G(DGX 40GB) 누가 사용 중인가요?”, 와 같이 일일이 물어보지 않아도 되므로)

그리고 방금 언급한 MIG를 Task별로 할당하는 것이 가능한 것도 상당한 매력이었다.

이게 무슨 뜻이냐면, 위의 코드 예제에서 V1PodSpec이라는 것을 정의하면 쿠버네티스 pod를 생성하면서 자원을 할당해주는데, 기본적으로 쿠버네티스는 GPU 자원을 할당할 때 {“nvidia/gpu” : “2” } 이렇게 특정 GPU에 대한 언급 없이 단순하게 “gpu” 만 언급이 가능하다.

그러나 flyte에서는 {“nvidia/mig-5g.40g” : “1”}과 같이 할당이 가능하기 때문에 기존에 리서처들의 GPU 사용을 위해 분리해놓은 MIG(Muti Instance GPU)를 롤백할 필요가 없어졌다. (휴…)

결과적으로 GPU 사용량이 큰 Task에서는 큰 GPU를, GPU 작업 소요가 아얘 없거나 작게 들어가는 부분은 작은 사이즈의 GPU만 할당해주면 되기 때문에 인프라를 관리하는 우리 팀 입장에선 이 만큼 자원을 효율적으로 관리할 수 있는 방법이 없다고 판단했다.

그러나 진짜 고생은 이제 시작이었다.

k8s에 띄우는 것까지는 괜찮았지만, 아직 기타 여러 기업들의 Use Case가 없다보니, 생각보다 자잘한 버그들이 끝없이 나타났다.

이건 두더지 게임이 아니라구요!

일단 버그가 발생하는 가장 근본적인 원인은 Flyte에서 제공하는 Official Docs 예제 코드들이 잘못 작성되어 있거나 반만 작성되어 있는 예제, 또는 CLI 자체에 버그가 있는 경우였다.

예제 코드에선 launchplan을 실행하기 위해 위의 코드의 붉은 박스 부분만 언급한다.

Flyte에서 Batch Job을 실행하기 위해 LaunchPlan이라는 커맨드를 지원하고 있고, 이를 실행하기 위한 예제 코드를 공식 문서에 표기해놓았다.

그러나 위의 이미지에서처럼 붉은 박스 내용까지만 설명하고, 하단에 Console UI에 등록된 Launch plan을 정상적으로 실행시켜주기 위한 코드가 빠져 있다.

그리고 다음 커맨드는 아얘 버그가 있다.

flytectl update launchplan -p <프로젝트 이름> -d <프로젝트 내 도메인 이름> — version <생성된 workflow 버전 이름>== <생성된 launchplan 이름> — archived

위 커맨드는 launchplan으로 잡아놓은 스케줄링을 제거하는 것인데, 아무리 실행해도 잡아놓은 스케줄링이 삭제되지 않는다.

(심지어 Console UI에 있는 archive 기능도 에러 코드는 반환하지 않고, 실행되지 않는다)

처음엔 슬랙에 있는 Flyte 커뮤니티에 문의했고, 사용 중인 flyte v1.0.1에서 v1.0.2로 업데이트하라는 답변을 받았다.

그러나 문제는 해결되지 않았다.

분명히 해결될거라면서요…? 결국 버전 업그레이드는 우리에게 어떤 변화도 가져오지 못했다.

이 날 결국 하루 종일 삽질하면서, k8s에 띄워져 있는 flytescheduler라는 파드를 재시작했더니, 잡혀 있던 모든 flyte 스케줄링이 삭제 되었다.

(매번 이렇게 할 수도 없고, 결국 방법을 찾아야한다.)

Flyte 측에서 문서와 Bug Fix를 정기적으로 릴리즈해줘야하는데, 아직까지는 사용자가 그렇게 많지 않아 배포 속도가 느린 것 같다.

현재 상황은 MLOps 0 단계에서 1단계로 진입 중…

배포는 이미 되었고, 사내 AI 부서 중 한 곳에서 Data Processing을 위해 베타 테스트를 수행하고 있다.

Flyte를 리서치 엔지니어 분들이 사용한다면, 모델 생성 단계를 각 각 Task 단위로 나눠 하나의 Workflow에서 관리할 수 있을 것이고, 나오는 결과 값을 Flyte 내의 Domain (작업 영역으로 기본 값은 development / staging / production으로 구분한다) 에서 관리할 수 있을 것이다.

아직 미정이긴 하지만, MLflow 등에서 각 모델의 실험 결과 히스토리를 관리할 수 있고, argo CD 등을 통해 배포할 수 있을 것으로 보인다 (Github action이 될 수도 있다.)

추가로 모델 데이터셋 등은 회사의 NFS or DGX 서버에 관리하고 있으나 정제된 데이터는 Snowflake에 데이터베이스를 구성해 관리하고 있다.

BigQuery나 Red Shift를 사용해서 관리하는 것이 편할 수도 있지만, 현재 Rust Complie 작업 등이 Disk I/O를 많이 발생시키고 있는데, 이런 작업들이 많은 상황 속에서 위와 같은 클라우드 서비스를 사용하는 것은 되려 성능 저하의 주범이 되는 것으로 판단된다.

취준생일 때는 Bigquery! Red Shift! 빅 테크 기업들의 서비스들이 무조건 최고라고만 생각했는데, 현업으로 넘어오니 회사 사정과 비즈니스 영역 등 여러가지 작업 환경을 고려해야할 것이 많아, 무조건 이러한 기업들의 서비스만을 이용할 수는 없을 것 같다.

하나씩 수동으로 작업하던 것들이 조금씩 자동화하는 단계이며, 회사에서 도입하고 있는 기술들은 하나씩 고도화가 되어가고 있다.

지금은 Flyte를 막 배포한 단계라서, 운영에 관련한 여러 노하우를 공유할 수 없고, 조만간 Flyte와 관련한 운영 사례와 버그들을 정리해서 다시 한 번 포스팅하도록 해야겠다.

Ryan

--

--