[Python] Pickle에 대한 오해와 Can’t Pickle local object Error 해결

🕓 5 mins read

# Pickle의 오해

파이썬 라이브러리들을 사용하다 보면, 직접 사용하지 않더라도, pickle은 어떻게든 만날 수밖에 없는데, 단순히 파일이나 오브젝트를 저장하고 불러올 때뿐만 아니라 어떤 객체를 프로세스 간에 공유하거나 전달할 때도 쓰입니다. 이때 pickle의 원리를 모를 경우 Attribute Error 혹은 PicklingError를 다음과 같은 형태로 만나면 당황스러울 수 있습니다. (문제없는 거 같은데 대체 왜!?)

PicklingError: Can't pickle <class 'class.method.var'>: it's not the same object as class.method.var

AttributeError: Can't pickle local object 'class.<locals>.some_var'

특히 파이썬을 처음 접하는 사용자나, 흔히 유행(?)하는 "pickle을 이용하여 파이썬 데이터 저장하기" 따위의 글을 통해 데이터 저장 방법을 잘못 배운 개발자들이 오해하기 쉬운 부분인데요, 파이썬 문서에도 명시되었듯이 pickle은 데이터를 저장하는 도구가 아니라 "객체 직렬화" 도구입니다.

https://docs.python.org/ko/3/library/pickle.html

직렬화(serialization)는 간단히 말해서, 데이터를 바이트 스트림으로 변환하는 절차입니다. 이는 바이너리 파일로 저장을 하는 형태로 보이니, 오해할 만 하긴 합니다. 다만 여기서 주목해야 할 부분은, 직렬화 대상이 "데이터"가 아니라 "객체 구조"를 직렬화하는 바이너리 프로토콜이라는 점입니다.

## Pickling의 문제점

• 객체 구조를 직렬화하면, 해당 객체에 대한 구현이 현재 scope에 존재하지 않으면 unpickling 할 때 오류가 납니다. 
  "CustomClass" 인스턴스를 pickle 했는데, "CustomClass"가 뭔지 알 수 없는 상태에선 unpickle이 안 되는 게 당연하겠죠.

무엇보다 위의 예시에서 눈치 빠른 분들은 알겠지만, unpickle은 어떤 데이터를 불러오는 것이 아니라 코드를 "실행"하는 과정이기 때문에, 검증된 pickle 오브젝트가 아니라면 unpickle 만으로도 악성코드를 실행시키게 될 수 도 있기 때문에 안전하지 않은 프로토콜입니다.

파이썬 문서에도 무시무시한 경고 문구가 있습니다.

https://docs.python.org/ko/3/library/pickle.html

• 바이너리를 "실행"하기 때문에, 당연히 서로 다른 버전의 파이썬 바이너리에서 생성한 pickle 오브젝트는 공유할 수 없습니다.
  - ex. Python 3.7.9에서 pickle된 오브젝트는 Python 3.8.16에서 unpickle 할 수 없는 것.
• 바이너리 프로토콜이기에 당연히 저장된 오브젝트는 human-readable 하지 않습니다.

### huggingface model hub

불과 몇 달 전만 해도 huggingface model hub도 pickle을 이용하기 때문에, 신뢰할 수 없는 모델이 아니라면 함부로 from_pretrained()를 실행해선 안 됐었습니다. (악성코드를 심어둘 수 있다니까요?)

- 물론 이 사실이 알려지고 나서 huggingface 측에서도 무시무시?한 경고 문구와 함께 모델이 업로드될 때 어느 정도 검수를 해주기 시작했습니다.

https://huggingface.co/docs/hub/security-pickle

Pickle Scanning

pickle load가 왜, 얼마나 위험한지와, 주의 사항들에 대해 상세히 설명하고 있으니, 관심 있으시면 읽어보면 좋을 글입니다.

 

# Pickle 오류 유형

사실 대게 이런 local object 오류는, pickle이 local object(전역 scope에서 정의를 찾을 수 없는)들이 pickle 대상이 될 때입니다.
(un-pickle이 실제로는 해당 오브젝트를 재생성하는 바이너리를 실행하는 과정인걸 생각하면 당연히 오류죠.)
몇 가지 대표적인 오류 유형과 해결 방법들을 살펴보면 다음과 같습니다. (대게 global로 선언해 주거나 구조를 변경하면 해소됩니다.)

ex1. Attribute Error - can't pickle local objects

가장 흔한 유형으로, 함수 내부의 로컬 scope에서 정의된 오브젝트를 pickle.dump 할 때 발생하는 오류입니다.

import pickle

def get_model_class(params):
    if params == 1:
        class ModelA(object):
    	    def __init__(self, ...):
                pass
    	return ModelA
    else:
        class ModelB(object):
            def __init__(self, ...):
                pass
        return ModelB
        
model_class = get_model_class(option)
model = model_class()

with open("model.pkl", "wb") as f:
    pickle.dump(model, f, pickle.HIGHEST_PROTOCOL)

일련의 조건 혹은 실행 옵션등에 따라서 모델 클래스를 선택해서 모델 인스턴스를 만드는 간단하고, 종종 마주할 법한 코드입니다.

Traceback (most recent call last):
  File "/Users/naubull2/test.py", line 20, in <module>
    pickle.dump(model, f, pickle.HIGHEST_PROTOCOL)
AttributeError: Can't pickle local object 'get_model_class.<locals>.ModelA'

처음에 설명드렸듯이, 오류가 발생하는 코드입니다.
ModelA 클래스의 정의는 get_model_class 함수의 로컬 scope에 존재합니다.
그러므로 전역 scope에서는 ModelA 클래스가 정의된 코드가 존재하지 않기에 해당 인스턴스를 만약 pickle 되는 걸 허용한다면, un-pickle 할 수 없을 것을 알면서 오브젝트를 저장하게 되는 것이죠.

다시 한번 이야기하면.. pickle은 인스턴스를 바이너리로 저장하는 게 아니라 인스턴스를 "생성"하는 바이너리를 저장합니다.

그럼 이런 경우 어떻게 해결해야 좋을까요?

1. 간단하게 global로 전역 변수로

가장 간편하게는 global 선언을 붙여서 pickle 하는 scope에서 모델 인스턴스 생성자의 정의를 알 수 있도록 해주면 됩니다.

import pickle

def get_model_class(params):
    global ModelA
    global ModelB
    if params == 1:
        class ModelA(object):
    	    def __init__(self, ...):
                pass
    	return ModelA
    else:
        class ModelB(object):
            def __init__(self, ...):
                pass
        return ModelB
        
model_class = get_model_class(option)
model = model_class()

with open("model.pkl", "wb") as f:
    pickle.dump(model, f, pickle.HIGHEST_PROTOCOL)

2. 클래스 정의 코드를 전역으로

하지만, 재사용성이 높을만한 클래스로 보이는 ModelA, ModelB를 함수 내에서 정의해 놓고global로 만드는 것은 정말 구조가 이상합니다. 차라리 클래스 정의들을 함수 밖으로 빼내도 해결되겠죠? 

기능적으로도, get_model_class 함수는 옵션에 따라 생성할 모델의 클래스를 고르기만 하면 될 뿐 빌더 함수가 아니기 때문에 클래스 정의는 get_model_class 밖에 존재하는 게 구조적으로나 기능적으로나 명확합니다.

import pickle

class ModelA(object):
    def __init__(self, ...):
        pass

class ModelB(object):
    def __init__(self, ...):
        pass
        
def get_model_class(params):
    if params == 1:
    	return ModelA
    else:
        return ModelB
        
model_class = get_model_class(option)
model = model_class()

with open("model.pkl", "wb") as f:
    pickle.dump(model, f, pickle.HIGHEST_PROTOCOL)

3. 클래스 정의 코드를 분리

그럼 다 된 걸까요? 이제 pickle로 저장된 모델을 다시 불러올 때를 생각해 보면, 결국 pickle.load 시점에 해당 코드 scope에 저장된 ModelA를 생성하기 위해 해당 클래스의 생성자가 어떻게 생겼는지를 알아야 합니다. 

만약 임의의 다른 모듈에서 위에서 저장한 model.pkl을 로드하려고 하면, 다음과 같은 오류를 보게 됩니다.

Traceback (most recent call last):
  File "/Users/naubull2/main.py", line 3, in <module>
    d = pickle.load(open("model.pkl", "rb"))
AttributeError: Can't get attribute 'ModelA' on <module '__main__' from '/Users/naubull2/main.py'>

ModelA에 대한 정의를 모르는 곳에서 인스턴스를 생성하는 바이너리를 실행했으니까 당연히 오류가 나겠죠. 

저장하는 곳에서도, 불러오는 곳에서도 필요한 클래스 정의 코드니, 아얘 처음부터 클래스 정의를 별도 모듈로 분리 시켜두면, 필요한 곳에서 해당 모듈만 import 해서 pickle dump / load 하는 것이 구조적으로 가장 안전하고 클린 한 코드가 되겠죠?

# models.py
class ModelA(object):
    def __init__(self, ...):
        pass

class ModelB(object):
    def __init__(self, ...):
        pass
  
##########################################################
# save_model.py
import pickle
from models import ModelA, ModelB

def get_model_class(params):
    if params == 1:
    	return ModelA
    else:
        return ModelB
        
model_class = get_model_class(option)
model = model_class()

with open("model.pkl", "wb") as f:
    pickle.dump(model, f, pickle.HIGHEST_PROTOCOL)
    
##########################################################
# main.py
import pickle
from models import ModelA, ModelB

model = pickle.load(open("model.pkl", "rb"))
...

나머지 유형들도 원인은 모두 동일한데, 상황이 다른 것뿐입니다.

ex2. 멀티프로세싱 환경에서의 can't pickle local objects

이제, pickle이 안 되는 경우의 주요 원인을 알았는데, 그럼 선택적으로 클래스를 고르고 인스턴스를 생성하는 경우 외에 어떤 경우에 주로 발생할까요?

multiprocessing을 이용할 때, target함수를 로컬 함수를 이용할 경우에도 발생 가능합니다.

규모가 제법 있는 오픈소스 프로젝트에서도 종종 보이는 오류 유형이니.. pickle의 정체를 알아도 피하기가 쉽지 않은 오류인 것 같네요.
https://github.com/pyg-team/pytorch_geometric/issues/366

import multiprocessing

def main():
    a="string"
    def task1():
        print(a)
        
    job1 = multiprocessing.Process(target=task1, args=(a))
    job1.start()
    job1.join()
    
main()

multiprocessing 라이브러리 내부에서도 target 함수를 전달할 때 pickle을 사용합니다.
따라서 worker process에서는 위의 예제에서 로컬 함수인 task1()을 알 수가 없어 pickle 오류를 일으킵니다.

Traceback (most recent call last):
  File "/Users/naubull2/test.py", line 10, in <module>
    main()
  File "/Users/naubull2/test.py", line 8, in main
    job1.start()
    ...
  File "/Users/naubull2/opt/anaconda3/lib/python3.9/multiprocessing/popen_spawn_posix.py", line 47, in _launch
    reduction.dump(process_obj, fp)
  File "/Users/naubull2/opt/anaconda3/lib/python3.9/multiprocessing/reduction.py", line 60, in dump
    ForkingPickler(file, protocol).dump(obj)
AttributeError: Can't pickle local object 'function.<locals>.task1'

이 역시 1, 2의 해결책처럼, 전역으로 옮기거나, global로 선언해 주면 되겠죠?
3의 해결책처럼task1() 함수를 별도 모듈로 떼어내는 것도 방법일 텐데, 이건 경우에 따라서 결정하면 됩니다. 보통 job을 실행하는 곳에서 job을 정의할 만큼 간단한 경우도 많으니까 이럴 땐 굳이 task1() 함수 하나만 모듈로 떼어내는 것은 부자연스러울 테니까요.

ex3. lambda 함수 사용 시 can't pickle local objects

이쯤 되면 lambda 함수는 임시 함수이니, 당연히 pickle 되지 않을 것을 알 수 있겠죠?

import pickle

var = pickle.dumps(lambda x, y: x+y)

다음과 같은 오류를 유발합니다.

Traceback (most recent call last):
  File "/Users/naubull2/test.py", line 2, in <module>
    var=pickle.dumps(lambda x, y: x+y)
_pickle.PicklingError: Can't pickle <function <lambda> at 0x104c82040>: attribute lookup <lambda> on __main__ failed

lambda 함수를 pickle 할... 일이 있겠냐마는.. 이런 경우는 간단한 helper 함수를 하나 정의 하고 저장하는 편이 그나마 해결책이 되겠군요.

# 결론

기본적으로 pickle을 너무 맹신하지 않는 것이 좋지만, pickle 자체는 고성능 serializer이기 때문에, 사용을 해야 한다면, 정체를 확실하게 이해하고 사용해야겠죠.

물론 내가 직접 사용하지 않아도 다른 라이브러리 내에서 pickle을 serialize 수단으로 쓸 수 도 있기 때문에, 위의 예시들처럼 주로 발생하는 원인은 scope 문제이니 이번 포스팅을 통해 pickle 오류를 해결하실 수 있기 바랍니다.