What's with ML software and pickles?

2023-11-08

ML 소프트웨어와 pickle, 도대체 무슨 일일까?

원문은 Nelson Elhage님이 2023-11-08에 게재했습니다. 이 블로그 구독하기

나는 오랫동안 머신러닝과는 전혀 무관한 소프트웨어 엔지니어로 일하면서도 머신러닝에 어느 정도 호기심을 갖고 가끔 주변에서 스치듯 접촉해 왔다. 그 기간 동안 반복적으로 느낀 감정은 파이썬 ML 생태계에서 Python pickle이 광범위하게 쓰이는 모습을 볼 때마다 느꼈던 경악(솔직히 말하면 경멸)이었다.

심각한 보안 문제¹는 차치하더라도, 직렬화에 pickle을 사용하면 매우 취약해지는 경향이 있어 코드를 개선하고 라이브러리나 파이썬 버전을 업그레이드할 때마다 온갖 악몽을 초래한다. 소프트웨어 엔지니어로서의 내 경력 내내 나는 — 거의 예외 없이 — pickle 모듈에는 유효한 사용 사례가 없으며, 그저 매력적인 골칫거리에 불과하고 더 어리고 열정 넘치지만 순진했던 생태계의 실수라고 믿어 왔다.

하지만 지난 몇 년간 나는 ML 분야에서 전문적으로 일해 왔고, 덕분에 생태계를 바라보는 새로운 관점을 갖게 됐다. ML 소프트웨어가 왜 지금과 같은 모습이 되었는지, 그리고 pickle이 사용자들에게 어떤 문제를 해결해 주고 있는지를 직접 체득하게 됐다. 여전히 pickle이 마음에 들지는 않지만, 문제 영역에 대한 이해와 공감은 훨씬 깊어졌고, 이제는 이 문제가 “사소”하거나 단순히 무지나 게으름 때문이라고 생각하지 않는다. 이 글은 내가 배운 바를 조금이나마 전달하고 두 세계 사이의 간극을 조금 좁혀 보려는 시도다.

연구를 위한 것이지, 소프트웨어를 위한 것이 아니다

근본적으로 내가 전달하고 싶은 핵심 관점은 머신러닝 생태계의 소프트웨어 대부분이 무엇보다 연구자를 지원하기 위해 설계됐으며, 소프트웨어 엔지니어를 위한 것이 아니라는 점이다. 더 나아가 (적어도 이 맥락에서) 연구는 몇 가지 서로 연관된 특성으로 특징지어진다(어쩌면 정의된다고 할 수도 있다):

연구의 일차적 산출물은 소프트웨어 결과물이 아니라 지식이다. 연구팀은 소프트웨어 도구나 솔루션을 갖기 위해 소프트웨어를 작성하기보다는, 연구 질문에 답하고 자신과 팀, 그리고 분야 전체의 특정 영역에 대한 이해를 높이기 위해 소프트웨어를 작성한다.
연구자들은 어떤 아이디어가 성공하거나 유용할지 사전에 알지 못하는 경우가 많으므로, 많은 아이디어를 효율적으로 시도하는 것을 우선시하고 대부분은 실패할 것이라 가정한다.

위 내용의 당연한 결과로, 연구자들(특히 ML 같은 분야에서는!)은 수많은 코드를 작성할 수 있지만, 그 코드 대부분이 곧바로 버려지는 것도 지극히 정상적인 일이다. 실패한 실험은 대개 버려지고, 성공한 실험조차 논문이 작성된 뒤에는 버려지는 경우가 많다. 성공적인 아이디어를 후속 연구에서 발전시킬 때도, 이제는 문제를 더 잘 이해하게 됐고 새로운 문제나 연구 질문을 염두에 두고 있으므로 코드를 새로 작성하는 경우가 흔하다.

일반적인 소프트웨어 엔지니어링 환경과 비교했을 때, 이러한 환경에서의 트레이드오프와 여건은 (다른 직렬화 전략과 비교해) 훨씬 더 pickle 쪽으로 기울어 있다:

대부분의 코드는 언제나 같은 소수(연구자와 그 팀)에 의해서만 작성되고 실행되며, 단일 환경에서만 동작한다. ‘개발’과 ‘프로덕션’ 사이에 거의 혹은 전혀 구분이 없고, 외부 사용자나 데이터 소스, 상호작용도 없다. 이러한 요인들이 종합적으로 작용해 실제로 pickle 직렬화가 초래하는 보안 위험을 훨씬 낮춘다(물론 0은 아니다!)
데이터 구조 — 심지어 코드 조각까지도 — 스키마나 포맷, 데이터 변화에 대해 멈춰 고민할 필요 없이 그대로 가져다 직렬화할 수 있다는 점은 엄청난 이점이다. 우리는 결과를 얻기까지 걸리는 시간을 중시하며, 스키마나 직렬화 포맷을 고민하거나 온디스크와 인메모리 포맷 사이의 변환에 들이는 노력은 전혀 본질적이지 않은 작업이다.
대부분의 코드가 상당히 빠르게 버려지므로 pickle의 취약성이나 코드 변경 및 버전 업그레이드를 처리하는 데 따르는 어려움은 훨씬 덜 부각된다.

연구 현장의 한 장면

위에서 말한 트레이드오프를 조금 더 구체적으로 보여주기 위해, ML 연구를 수행하는 연구자에 대한 다소 각색된 묘사를 적어 보고, 내가 직접 ML 연구를 수행하고 (나보다 훨씬 경험 많은!) ML 연구자들을 지원하면서 체득한 관점을 바탕으로 그녀가 pickle을 활용하게 될 구체적인 사용 사례와 그에 따른 트레이드오프를 함께 보여 주고자 한다.

어떤 연구자가 새로운 딥러닝 모델 아키텍처, 혹은 기존 아키텍처의 변형을 실험하고 있다고 가정해 보자. 그녀의 아키텍처에는 수많은 설정 옵션과 하이퍼파라미터가 있을 것이다. 레이어 수, 각 레이어의 유형, 각종 벡터의 차원, 활성화 값을 어디서 어떻게 정규화할지, 어떤 비선형 함수를 쓸지 등등. 모델 구성 요소 중 다수는 ML 프레임워크가 제공하는 표준 레이어겠지만, 연구자는 곳곳에 새로운 로직을 조금씩 끼워 넣을 것이다.

우리 연구자는 특정한 구체적 모델 — 이러한 설정들의 특정 조합 — 을 서술할 방법이 필요하며, 이는 직렬화한 뒤 나중에 다시 불러올 수 있어야 한다. 그녀는 몇 가지 서로 연관된 이유로 이것이 필요하다:

그녀는 작업을 실행하는 데 사용할 수 있는 GPU나 다른 가속기를 갖춘 컴퓨트 클러스터에 접근할 수 있을 것이다. 그 클러스터에서 동작하는 코드에 모델 설명을 제출해 자신의 모델을 클러스터에서 실행할 방법이 필요하다.
모델이 학습되는 동안에는 하드웨어 장애나 작업 선점에 대비해 진행 상황을 저장해 두고, 나중에 다시 불러와 학습을 재개할 수 있도록 스냅샷을 저장해야 한다.
모델 학습이 끝나면 연구자는 그 모델들을 (최종 스냅샷과 학습 도중의 체크포인트 모두를 포함해) 다시 불러와 평가와 실험을 실행하고 싶어 할 것이다.

여기서는 다양한 해결책을 상상해 볼 수 있다! 모델에 대한 protobuf나 JSON 설명을 작성하고 이를 실제 모델 인스턴스로 변환하는 코드를 만들 수도 있다. 모델을 “커맨드라인 인자 배열”로 표현하고 이를 파싱해 모델 인스턴스로 만드는 방식도 상상할 수 있다.

하지만 이 모든 문제를 해결하면서도 사전에 거의 추가 작업이 필요 없는 해결책이 하나 있다:

“특정 모델”의 정의는 그저 “모델을 구현하는 객체”(아마도 torch.nn.Module 서브클래스)일 뿐이다
연구자는 새로운 실험을 실행하면서 계속 수정하고 발전시키는 연구 스크립트를 통해 이러한 객체들을 즉흥적인 방식으로 생성한다
그런 다음 전체 객체를 pickle로 직렬화해, 나중에 분석하거나 원격 머신에서 불러와 실행하기 위한 저장 버전으로 사용한다.

이 접근 방식은 번거로움도 없고, 연구 코드나 구현이 어떻게 구조화되어야 하는지에 대한 특별한 규율이나 패턴도 요구하지 않으며, “그저 모델을 만들고 실행하는 것”에 비해 인지적 부담이 거의 없다. 표준 컴포넌트와 커스텀 컴포넌트를 동일하게 쉽게 지원한다. 연구자는 클래스 계층 깊은 곳의 무작위한 지점에서 작은 수정을 가하거나 전체 모델 아키텍처를 교체할 수 있으며, 어느 쪽도 직렬화 시스템을 고민하거나 수정할 필요가 거의 없다.

이쯤 되면 우리 연구자는 하이퍼파라미터의 다양한 데카르트 곱을 인스턴스화하는 스크립트를 작성해 모두 클러스터에 스케줄링한 뒤, 학습이 진행되는 동안 점심을 먹으러 가거나(혹은 주말을 보내러 떠난다).

이 pickle 파일들은 확실히 취약하다! 모델의 구현 세부사항 전체를 직렬화하므로, 연구자가 새로운 설정과 파라미터를 추가하고 실험을 리팩터링하면서 오래된 실험과 모델을 망가뜨리기는 매우 쉽다. 하지만 다시 말하지만, 대부분의 실험은 흥미롭지 않은 결과를 낳으며 다시 살펴볼 필요가 전혀 없다. 유의미한 몇 개에 대해서는 각 실험과 연결된 git 커밋을 추적해 체크아웃을 되돌리면 간단히 해결된다. 그리고 우리 연구자가 몇몇 특정 모델을 더 오래 보존해야 한다면, 대개 어딘가의 __setstate__ 메서드에 몇 가지 전술적인 수정을 가해 고쳐 쓸 수 있다.

그렇다고는 해도, 실험들이 전반적으로 성공을 거두고 이 연구 라인을 계속 이어 나가 더 많은 변형을 시도하며 비교를 위해 모델을 더 오래 보존해야 한다면, 취약성이 실제로 문제가 되기 시작하므로 우리 연구자는 결국 리팩터링을 고려하게 될지도 모른다. 한 가지 직접적인 접근법은 모델의 “설명”과 구현을 분리하는 것이다. 모델 설정/세팅/하이퍼파라미터를 class ModelConfig:라는 클래스에 묶고, 이 클래스를 (여전히 pickle을 통해) 직렬화하며, 연구자는 주어진 설정으로부터 전체 모델을 인스턴스화하는 방법을 아는 build(cfg: ModelConfig) -> Model 함수를 직접 유지·관리한다.

이제 연구자는 이 클래스의 전방·후방 호환성만 신경 쓰면 되고, ModelConfig로부터 Model 객체를 생성하는 코드를 함께 리팩터링함으로써 구현 코드를 비교적 자유롭게 리팩터링할 수 있다.

하지만 우리 연구자(그리고 그녀의 팀)는 여전히 실험하고 새로운 것을 시도하는 중이며, 반복과 “실험까지 걸리는 시간”, 그리고 어쩌면 더 중요하게는 “실험에 드는 인지적 비용”이 모든 것을 지배한다는 점을 기억할 필요가 있다. 그래서 이 클래스의 많은 필드는 보기 좋은 “정상적인” 정수나 불리언, 어쩌면 열거형이겠지만, 매번 추가 실험을 위해 지원을 추가해야 하는 것은 다소 번거로운 일이다.

그래서 어쩌면 우리 연구자는 여러 개의 “탈출구”를 남겨 둔다. 사소한 예로, 완전 연결 레이어 뒤에 텐서에 적용하는 비선형성, 즉 “활성화 함수”를 교체해 보는 것은 이 분야에서 아직도 꽤 유행하는 시도다(나 자신도 여기에 해당한다).

텐서에 적용하는 활성화 함수(아마도 말 그대로 Tensor -> Tensor 타입의 파이썬 함수일 것이다)를 설정에 어떻게 저장할 수 있을까?

어딘가의 dict[str, Callable]에 함수들을 등록해 이름으로 저장할 수도 있겠지만… 잠깐, 파이썬 모듈 네임스페이스 자체가 사실상 그런 dict 아닌가? 그러니 pickles를 쓰고 있다면 그냥 함수 자체를 설정에 직접 저장해도 될 것이다.

이러한 깨달음에 용기를 얻은 우리 연구자는 더 나아간다. 그녀는 심지어 전체 build_xxxx: Callable[[Config], Module | None] 오버라이드를 저장할 수도 있는데, 이는 설정이 설정의 임의 부분을 임의의 방식으로 오버라이드할 수 있게 해준다!

스코프 크립

내 개인적인 의견으로는, 위에서 설명한 모든 것은 프로젝트가 정말로 내부 연구 용도로 신뢰할 수 있는 맥락에서만 실행된다면, 다소 불쾌하고 위험하긴 해도 “기본적으로는 괜찮다”.

하지만 더 큰 문제는 일단 이런 식으로 pickle을 일상적으로 사용하기 시작하면, pickle이 훨씬 더 문제가 되는 경우를 포함해 모든 직렬화 문제에 pickle을 사용하는 것이 습관이 된다는 점이다. 연구자들만 그러는 것이 아니라, 이러한 패턴은 도구와 생태계 전반으로 스며드는 경향이 있다.

예를 들어 PyTorch는 대부분의 경우 문자열과 Tensor의 리스트나 dict 같은 몇 가지 “기본적인” 자료구조만으로 충분함에도, 기본 직렬화 포맷의 일부로 pickle 파일을(torch.load와 torch.save를 통해) 사용한다. PyTorch에는 몇 가지 내장 타입만 지원하는 커스텀 C++ pickle 로더가 있어 로더의 weights_only 인수로 활성화할 수 있지만, 하위 호환성을 위해 이 플래그는 기본값이 False이며, (공격자가 없는 상황에서는) 이를 활성화하면 오히려 문제를 일으킬 뿐이므로 거의 사용되지 않는다. 그 결과 모든 PyTorch 가중치 파일은 사용자가 위험한 인터페이스를 사용하도록 공개적으로 초대하며 스스로를 취약하게 만든다.

어떻게 해야 할까?

이 글을 쓰는 나의 주된 목표는 ML 생태계 안에서 왜 그렇게 많은 pickle 사용을 목격하는지, 그리고 pickle이 사용자들을 위해 어떤 문제들을 해결하고 있는지에 대한, 말하자면 “내부자의 시선”을 조명해 보는 것이다. 그렇다고는 해도, 머신러닝의 중요성이 커지고 다양한 맥락에서 사용자와 배포가 늘어남에 따라 안전하지 않은 pickle 사용으로 인한 문제가 점점 더 중요하고 시급해질 것으로 예상하며, 앞으로 나아갈 몇 가지 방향에 대해 몇 가지 생각을 제시하고자 한다.

더 나은 데이터 포맷과 인터페이스

첫째, 우리는 (여기서 말하는 ‘우리’란 주로 소프트웨어 엔지니어링 배경과 기술을 가진 사람들, 그리고 ML 안팎에서 주로 SWE 역할로 일하는 사람들을 의미한다) 연구자들에게 더 많고 더 나은 대안을 제공하기 위해 pickle에 대한 고품질 대안에 계속 투자해야 한다고 생각한다. 잘 만들어진다면 많은 경우 pickles보다 더 나은 경험은 물론 더 나은 보안 속성까지 제공할 수 있다.

여기서 XLA를 중심으로 구축된 구글의 머신러닝 생태계 — Tensorflow와 jax 및 그 다양한 라이브러리와 프레임워크 — 가 여기서 설명한 문제들을 대체로 피해 왔다는 점을 언급하지 않으면 아쉬울 것이다. 이들은 프로토콜 버퍼와 numpy 포맷 기반의 직렬화 포맷을 사용하는 경향이 있으며, pickles가 훨씬 덜 만연하다. 한편으로는 나는 이를 그러한 생태계가 가능하다는 일종의 존재 증명으로 본다. 다른 한편으로는 PyTorch가 특히 온보딩과 수정·탐색의 용이성 측면에서 연구자들에게 전반적으로 더 우수한 사용 경험을 제공하기 때문에 그토록 지배적으로 자리 잡았다고 믿는 경향이 있어, 구글 라이브러리를 이상적인 해법으로 제시하는 데는 주저하게 된다.

Huggingface의 safetensors 라이브러리는 이 분야에서 이뤄진 진전의 훌륭한 예로, 텐서 데이터를 안전하고 성능 좋은 방식으로 직렬화하기 위한 견고한 기반을 제공한다. 내 관점에서 이는 약간 낮은 수준의 프리미티브로 보는 것이 가장 적절하다. dict[str, Tensor] 형태의 객체를 직렬화하고 복원하는 것을 지원하는데, 연구자들은 매우 자주 중첩된 딕셔너리나 가중치와 함께 메타데이터를 하나의 번들로 저장하고 싶어 한다. 나는 safetensors 위에 더 사용하기 편하고 연구자 친화적인 포맷과 API가 개발되는 모습을 기대하고 있다.

연구자들은 종종 텐서 데이터뿐만 아니라 파이썬 클래스 트리도 직렬화하고 싶어 한다. 어쩌면 파이썬 타입 어노테이션의 등장과 최근의 지배적 확산, 그리고 그 위에 구축된 attrs와 dataclasses 같은 API가, 파이썬 타입 어노테이션과 소량의 추가 메타데이터(아마도 cattrs를 기반으로)에 의해 구동되어 [특정한] 타입 어노테이션이 적용된 파이썬 클래스 트리를 다소 직접적으로 직렬화·역직렬화할 수 있는 라이브러리의 가능성을 열어 줄지도 모른다.

`pickle` 사용 제한하기

좋든 싫든 우리는 오늘날 수많은 pickle 파일과 pickle 파일을 사용하는 라이브러리를 안고 있고, 문제는 여전히 커지고 있다. 대안을 제공하고 프로젝트들이 그 대안으로 이전하도록 장려하는 동안, 피해를 줄이거나 완화하기 위한 하나의 사변적인 아이디어를 제시해 본다:

첫째, 파이썬이 pickle 모듈을 완전히 비활성화하는 전역 “nopickle” 모드를 지원하면 유용하지 않을까 생각한다. 이러한 플래그가 있다면 파이썬을 배포·운영하는 사람들이 런타임에 신뢰할 수 없는 pickle을 로드함으로써 위험에 노출되어 있는지를 쉽게 알 수 있을 것이다.

하지만 그것이 너무 제한적이라면, pickle 사용을 특정한 방식으로 제한하는 옵션을 추가하는 것은 어떨까?

Pickle 매니페스트

특히 나는 pickle의 모든 사용이 디코딩을 위해 유효한 바이너리 객체 목록을 선언하도록 요구하는 아이디어에 흥미를 느낀다. 아마도 sha256이나 다른 암호학적 해시로 식별될 것이다. pickle.load 호출 지점(혹은 어떤 새로운 API)에서 예상 객체를 지정하는 방식과, 일종의 매니페스트 파일에서 대역 외로 지정하는 방식을 모두 지원한다고 상상해 볼 수 있으며, 후자는 합법적인 (호출 스택, blob ID) 튜플을 명시할 수도 있다. 이 데이터를 명시하지 않거나 허용된 해시와 일치하지 않는 pickle 로드는 거부될 것이다.

이 접근법이 흥미로운 이유는 ML 소프트웨어를 배포하는 사람들이 — 단 한 번 — 로드되는 모든 모델 가중치 파일이나 설정 파일의 인벤토리를 수집한 뒤 그 집합을 “동결”해 런타임에 새로운 pickle이 등장하지 않도록 보장할 수 있게 할지도 모르기 때문이다.

이는 결국 pickle 사용을 관리하는 문제를 서드파티 의존성을 관리하는 문제로 크게 축소시킨다. 물론 이는 해결된 문제가 아니지만, 어차피 안고 있었을 문제이기도 하다.

구체적으로, 새로운 서드파티 ML 모델을 배포하려면 보통 모델 구현(파이썬 모듈이나 스크립트 형태)과 모델 가중치 둘 다를 설치해야 한다. 둘 다 악의적이라면 보안 위험을 초래하지만, pickle로 인한 위험은 어떤 면에서는 더 심각하다. 모델 코드의 특정 버전을 검토해 고정하더라도 pickle 파일은 종종 런타임에 다운로드되므로, 여전히 시스템 외부로부터의 RCE에 노출되기 때문이다.

해당 코드가 특정 pickle 파일만 로드한다는 것을 보장할 수 있고, 그 의존성을 처음 신뢰하기로 결정했을 때 한 번 감사할 수 있다면, pickle 사용 사례를 서드파티 코드를 설치하는 문제와 “비슷하게” 만들 수 있다. 서드파티 코드의 경우 한 번 버전을 가져와 고정하면 적어도 추가적인 런타임 공격 표면이 생기지 않는다.

이러한 아이디어를 런타임 설정 시스템으로 더 복잡하게 만드는 것도 상상해 볼 수 있다. 극단적으로는 pickle 사용을 제어하기 위해 다음과 같은 꽤 복잡한 API를 상상할 수도 있다:

with pickle.no_pickles() as token:
  pickle.loads(...) # NoPickleException

  with pickle.allow_pickles(token):
    pickle.loads(...)

  with pickle.allow_pickles(token, only_objects=[
    "sha256:some sha256",
  ]):
    pickle.loads(buf) # hashes buf, only proceeds if it matches

  # You must pass the same token returned by `no_pickles` to
  # `allow_pickle`; this prevents libraries from trivially blanket
  # calls to renable `pickle`
  with pickle.allow_pickles(bad_token):  # BadTokenException
    ...

물론 이러한 메커니즘이 완벽할 수는 없다. 라이브러리가 이를 우회하는 것은 언제나 가능할 것이다. 극단적으로 라이브러리는 언제나 이러한 제한이 전혀 없는 pickle의 순수 파이썬 구현을 자체 포함시킬 수 있다. 하지만 그것은 라이브러리가 (의도적이든 비의도적이든) 백도어를 심을 수 있는 수백만 가지 방법 중 하나에 불과하다. 이 제안의 목표는 공급망 보안을 해결하는 것이 아니라 pickle의 특정 위험을 완화하고, 위에서 언급했듯이 위험을 런타임이 아닌 패키지 설치 시점으로 옮기는 것이다. 표준 보안 메커니즘을 노골적으로 우회하는 라이브러리는 커뮤니티에 의해 포착되어 신뢰할 수 없는 것으로 표시되기를 바랄 뿐이다².

이러한 기능은 pickle 모듈을 몽키 패치하는 라이브러리에서 트리 외부 프로토타입으로 구현될 수도 있지만, 표준화와 채택 측면에서 결국 인터프리터에 내장하는 것이 가치가 있을지도 모른다.

10여 년 전, 나는 (내가 아는 한) pickle 기반 시스템을 공격하는 익스플로잇에 대한 최초의 공개 글 중 하나를 게시했다. 하지만 이 문제는 이미 수년간 잘 알려져 있었고, 요즘에는 버튼 하나로 익스플로잇을 수행하는 도구들도 존재한다. ↩︎
이 문장은 내가 직접 쓰면서도 순진하게 느껴진다. 하지만 다시 말하지만, 이러한 해결책이 문제를 기존의 공급망 보안 문제로 크게 축소시킨다는 점은 좋든 싫든 사실이라고 생각한다. ↩︎