Linux生成的pickle文件在Mac加载失败，报错TypeError怎么解决？

问题描述 #

在Linux上使用Ray RLlib训练模型并保存Checkpoint，迁移到Mac上做推理时，调用Algorithm.from_checkpoint加载Checkpoint触发pickle反序列化错误。两边使用相同版本的Ray RLlib，但Linux上运行正常，Mac上报错。

错误栈 #

Traceback (most recent call last):
  File "/Users/psr/GitLab/psr.ai/rlcv/rl-tutorial/ray/dqn_inference_ray.py", line 76, in <module>
    main()
  File "/Users/psr/GitLab/psr.ai/rlcv/rl-tutorial/ray/dqn_inference_ray.py", line 33, in main
    algo = Algorithm.from_checkpoint(checkpoint.path)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/psr/GitLab/psr.ai/rlcv/rl-tutorial/.venv/lib/python3.11/site-packages/ray/rllib/algorithms/algorithm.py", line 343, in from_checkpoint
    state = Algorithm._checkpoint_info_to_algorithm_state(
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/psr/GitLab/psr.ai/rlcv/rl-tutorial/.venv/lib/python3.11/site-packages/ray/rllib/algorithms/algorithm.py", line 2876, in _checkpoint_info_to_algorithm_state
    state = pickle.load(f)
            ^^^^^^^^^^^^^^
TypeError: code() argument 13 must be str, not int

原因分析 #

这个错误本质是跨架构/跨Python环境的pickle兼容性问题：

• Linux（x86_64架构）和Mac（Apple Silicon为arm64架构）的Python字节码结构存在差异
• 即使Ray版本相同，不同架构下Python的code对象序列化格式不一致（比如Python 3.10+中，code()的第13个参数在x86架构下是int类型，arm64架构下是str类型，反之亦然）
• Ray RLlib默认用pickle保存包含Python字节码的对象，导致跨架构反序列化失败

解决方案 #

方案1：统一Python环境的版本与架构 #

• 在Mac上通过Rosetta 2运行x86架构的Python，确保和Linux端的Python版本、架构完全一致
• 示例：在Mac终端执行arch -x86_64 python -m venv .venv创建匹配架构的虚拟环境，再安装相同版本的Ray及依赖

方案2：使用跨平台模型导出方式 #

• 训练完成后，不依赖默认Checkpoint，改用Ray支持的跨平台格式导出模型：
  1. Linux训练端导出模型（以TorchScript为例）：

     # 训练完成后执行
     algo.export_model("/path/to/exported_model", export_format="torchscript")
     

  2. Mac推理端加载导出的模型：

     from ray.rllib.models import ModelCatalog
     from ray.rllib.algorithms.dqn import DQNConfig
     
     # 初始化匹配的模型配置
     model_config = DQNConfig().model
     # 加载导出的模型
     model = ModelCatalog.get_model_from_checkpoint(
         "/path/to/exported_model",
         model_config=model_config,
         framework="torch"
     )
     

方案3：替换Checkpoint的序列化方式 #

• 初始化算法时，配置使用msgpack替代pickle作为序列化工具（部分Ray版本支持）：

  from ray.rllib.algorithms.dqn import DQNConfig
  
  algo_config = DQNConfig()
  # 配置msgpack序列化Checkpoint
  algo_config = algo_config.checkpoint_serializer("msgpack")
  algo = algo_config.build()
  

内容的提问来源于stack exchange，提问作者Prabhjot Singh Rai