使用Polars读取Parquet数据给PyTorch训练时出现Unpickling Error及重复DEBUG日志的问题求助

AI 小北

2026-4-14

Polars读取Parquet数据给PyTorch训练时出现Unpickling Error及重复DEBUG日志的问题求助

Hey folks, I recently switched my data processing tool from xarray to Polars, using pl.DataFrame.to_torch() to generate tensors for my PyTorch model. My data source is Parquet files. To avoid forking child processes, I'm using torch.multiprocessing.spawn to start the training process, but it's crashing with errors and weird log behavior. Here's the relevant error log:

/home/username/.conda/envs/torchhydro1/bin/python3.11 -X pycache_prefix=/home/username/.cache/JetBrains/IntelliJIdea2024.3/cpython-cache /home/username/.local/share/JetBrains/IntelliJIdea2024.3/python-ce/helpers/pydev/pydevd.py --multiprocess --qt-support=auto --port 29781 --file /home/username/torchhydro/experiments/train_with_era5land_gnn_ddp.py 
Console output is saving to: /home/username/torchhydro/experiments/results/train_gnn_ddp.txt
[20:38:51] DEBUG    No module named 'forge'                     signatures.py:43
           DEBUG    No module named 'forge'                     signatures.py:43
[20:38:52] DEBUG    Using selector: EpollSelector          selector_events.py:54
           ……
           DEBUG    Using fontManager instance from         font_manager.py:1580
                    /home/username/.cache/matplotlib/fontl                     
                    ist-v390.json                                               
update config file
!!!!!!NOTE!!!!!!!!
-------Please make sure the PRECIPITATION variable is in the 1st location in var_t setting!!---------
If you have POTENTIAL_EVAPOTRANSPIRATION, please set it the 2nd!!!-
!!!!!!NOTE!!!!!!!!
-------Please make sure the STREAMFLOW variable is in the 1st location in var_out setting!!---------
[20:39:04] DEBUG    No module named 'forge'                     signatures.py:43
           DEBUG    No module named 'forge'                     signatures.py:43
[20:39:06] DEBUG    Using selector: EpollSelector          selector_events.py:54
           ……
           DEBUG    Using fontManager instance from         font_manager.py:1580
                    /home/username/.cache/matplotlib/fontl                     
                    ist-v390.json                                               
……
Torch is using cuda:0
[2024-12-12 20:48:08,931] torch.distributed.distributed_c10d: [INFO] Using backend config: {'cuda': 'nccl'}
[W CUDAAllocatorConfig.h:30] Warning: expandable_segments not supported on this platform (function operator())
using 8 workers
Pin memory set to True
  0%|          | 0/22986 [00:00<?, ?it/s]
[20:48:40] DEBUG    No module named 'forge'                     signatures.py:43
           DEBUG    No module named 'forge'                     signatures.py:43
[20:48:41] DEBUG    Using selector: EpollSelector          selector_events.py:54
           ……
           DEBUG    Using fontManager instance from         font_manager.py:1580
                    /home/username/.cache/matplotlib/fontl                     
                    ist-v390.json                                               
[20:49:28] DEBUG    No module named 'forge'                     signatures.py:43
           DEBUG    No module named 'forge'                     signatures.py:43
[20:49:29] DEBUG    Using selector: EpollSelector          selector_events.py:54
           ……
           DEBUG    Using fontManager instance from         font_manager.py:1580
                    /home/username/.cache/matplotlib/fontl                     
                    ist-v390.json                                               
[20:50:19] DEBUG    No module named 'forge'                     signatures.py:43
           DEBUG    No module named 'forge'                     signatures.py:43
[20:50:20] DEBUG    Using selector: EpollSelector          selector_events.py:54
           ……
           DEBUG    Using fontManager instance from         font_manager.py:1580
                    /home/username/.cache/matplotlib/fontl                     
                    ist-v390.json                                               
[20:51:12] DEBUG    No module named 'forge'                     signatures.py:43
           DEBUG    No module named 'forge'                     signatures.py:43
[20:51:13] DEBUG    Using selector: EpollSelector          selector_events.py:54
           ……
           DEBUG    Using fontManager instance from         font_manager.py:1580
                    /home/username/.cache/matplotlib/fontl                     
                    ist-v390.json                                               
[20:52:07] DEBUG    No module named 'forge'                     signatures.py:43
           DEBUG    No module named 'forge'                     signatures.py:43
[20:52:09] DEBUG    Using selector: EpollSelector          selector_events.py:54
           ……
[20:52:13] DEBUG    CACHEDIR=/home/username/.cache/matplotlib   __init__.py:341
           DEBUG    Using fontManager instance from         font_manager.py:1580
                    /home/username/.cache/matplotlib/fontl                     
                    ist-v390.json                                               
[20:53:11] DEBUG    No module named 'forge'                     signatures.py:43
           DEBUG    No module named 'forge'                     signatures.py:43
[20:53:12] DEBUG    Using selector: EpollSelector          selector_events.py:54
           ……
           DEBUG    Using fontManager instance from         font_manager.py:1580
                    /home/username/.cache/matplotlib/fontl                     
                    ist-v390.json                                               
[20:55:12] DEBUG    No module named 'forge'                     signatures.py:43
           DEBUG    No module named 'forge'                     signatures.py:43
[20:55:14] DEBUG    Using selector: EpollSelector          selector_events.py:54
           ……
[20:55:19] DEBUG    CACHEDIR=/home/username/.cache/matplotlib   __init__.py:341
           DEBUG    Using fontManager instance from         font_manager.py:1580
                    /home/username/.cache/matplotlib/fontl                     
                    ist-v390.json                                               
Traceback (most recent call last):
  File "/home/username/.local/share/JetBrains/IntelliJIdea2024.3/python-ce/helpers/pydev/pydevd.py", line 1570, in _exec
    pydev_imports.execfile(file, globals, locals)  # execute the script
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/username/.local/share/JetBrains/IntelliJIdea2024.3/python-ce/helpers/pydev/_pydev_imps/_pydev_execfile.py", line 18, in execfile
    exec(compile(contents+"\n", file, 'exec'), glob, loc)
  File "/home/username/torchhydro/experiments/train_with_era5land_gnn_ddp.py", line 171, in <module>
    test_run_model()
  File "/home/username/torchhydro/experiments/train_with_era5land_gnn_ddp.py", line 56, in test_run_model
    mp.spawn(gnn_train_worker, args=(world_size, config_data, None), nprocs=world_size, join=True)
  File "/home/username/.conda/envs/torchhydro1/lib/python3.11/site-packages/torch/multiprocessing/spawn.py", line 241, in spawn
    return start_processes(fn, args, nprocs, join, daemon, start_method="spawn")
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/username/.conda/envs/torchhydro1/lib/python3.11/site-packages/torch/multiprocessing/spawn.py", line 197, in start_processes
    while not context.join():
              ^^^^^^^^^^^^^^
  File "/home/username/.conda/envs/torchhydro1/lib/python3.11/site-packages/torch/multiprocessing/spawn.py", line 140, in join
    raise ProcessExitedException(
torch.multiprocessing.spawn.ProcessExitedException: process 0 terminated with signal SIGKILL
python-BaseException
Traceback (most recent call last):
  File "/home/username/.conda/envs/torchhydro1/lib/python3.11/multiprocessing/spawn.py", line 132, in _main
    self = reduction.pickle.load(from_parent)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
_pickle.UnpicklingError: pickle data was truncated
python-BaseException
/home/username/.conda/envs/torchhydro1/lib/python3.11/multiprocessing/resource_tracker.py:254: UserWarning: resource_tracker: There appear to be 30 leaked semaphore objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '

I have two main questions:

Why am I getting the _pickle.UnpicklingError: pickle data was truncated error?
After seeing the progress bar line 0%| | 0/22986 [00:00<?, ?it/s], the DEBUG logs repeat 8-9 times. I set num_worker=8 in my PyTorch DataLoader—does this repetition relate to the worker count?

This issue started right after I switched to Polars, so I suspect there's a conflict between Polars and PyTorch's process/thread handling. I'd love to understand the root cause of the unpickling error and fix both problems.

Possible Fixes & Explanations

Let me walk you through what might be happening and how to fix it, based on my experience with Polars and PyTorch multiprocessing:

1. The Repeating DEBUG Logs (Directly Linked to `num_worker=8`)

Yes, this is definitely tied to your DataLoader worker count. Each worker spawns its own independent subprocess, and every subprocess initializes its own environment—including logging, Matplotlib, and other dependencies. Those repeated DEBUG lines are just each worker process loading its config and dependencies.

How to fix it:

Lower the log level to suppress DEBUG noise before initializing the DataLoader:

import logging
# Set to INFO or WARNING to hide DEBUG messages
logging.basicConfig(level=logging.INFO)
# Or silence specific noisy loggers
logging.getLogger('matplotlib').setLevel(logging.WARNING)
logging.getLogger('signatures').setLevel(logging.WARNING)

Use worker_init_fn to configure logging per worker:

def worker_init_fn(worker_id):
    import logging
    # Disable DEBUG logs for this worker
    logging.basicConfig(level=logging.WARNING)
    # Silence specific modules
    logging.getLogger('matplotlib').setLevel(logging.ERROR)
    logging.getLogger('signatures').setLevel(logging.ERROR)

# Add this to your DataLoader initialization
dataloader = DataLoader(
    dataset,
    num_workers=8,
    worker_init_fn=worker_init_fn,
    pin_memory=True
)

2. The `_pickle.UnpicklingError: pickle data was truncated`

This error usually happens when data passed between processes gets cut off, or the object being pickled isn't compatible. Let's break down the likely causes and fixes:

a. Root Cause: SIGKILL from Out-of-Memory (OOM)

First, notice the line process 0 terminated with signal SIGKILL—this almost always means the process ran out of memory and was killed by the system. The unpickling error is a side effect: when the parent/child process dies mid-pickle, the remaining data is truncated, leading to the error.

How to verify:

Monitor memory usage during training with tools like htop or PyTorch's memory profiler. If you see memory spike to 100% before the crash, OOM is the issue.

Fixes:

Reduce num_worker: More workers mean more memory usage (each worker loads its own data slice). Try lowering to 4 or 2 to see if memory usage drops.
Use Polars' Lazy API: Instead of loading the entire Parquet file into memory with pl.read_parquet(), use pl.scan_parquet() to lazy-load data on demand. This avoids loading everything at once:
```
# Lazy load instead of eager load
lazy_df = pl.scan_parquet("your_data.parquet")
# Process and convert to tensor only when needed
batch_tensor = lazy_df.filter(...).collect().to_torch()
```
Split your dataset: If possible, split the Parquet file into smaller chunks so each worker only loads a small piece of data.

b. Polars Objects in Process Arguments

When using torch.multiprocessing.spawn, all arguments passed to the worker function are pickled and sent to the subprocess. If config_data or other args include Polars objects (like DataFrames/LazyFrames), they might not pickle cleanly, or produce overly large data that gets truncated.

Fix:

Only pass native Python types (dicts, lists, strings) or PyTorch tensors to the worker. Load and process Polars data inside the worker function, not in the parent process. For example:

def gnn_train_worker(rank, world_size, config_data, _):
    # Load and process data inside the worker, not in the parent
    df = pl.read_parquet(config_data["data_path"])
    tensor = df.to_torch()
    # Rest of your training logic...

c. Polars Multithreading vs. PyTorch Processes

Polars uses multithreading by default, which can clash with PyTorch's multiprocessing workers. This can lead to resource contention that disrupts pickling.

Fix:

Disable Polars multithreading in worker processes via worker_init_fn:

def worker_init_fn(worker_id):
    import os
    # Force Polars to use 1 thread per worker
    os.environ["POLARS_MAX_THREADS"] = "1"
    # Add logging config here too...

Or use Polars' config context when converting to tensors:
```
with pl.Config(n_jobs=1):
    tensor = df.to_torch()
```

Debugging Checklist

First, confirm if OOM is causing the SIGKILL—this is the most common root cause here.
Reduce worker count and use lazy loading to lower memory usage.
Move data loading into the worker function to avoid pickling Polars objects.
Silence DEBUG logs to make it easier to spot real errors.

Let me know if any of these steps help, or if you need more details on any part!

备注：内容来源于stack exchange，提问作者forestbat

使用Polars读取Parquet数据给PyTorch训练时出现Unpickling Error及重复DEBUG日志的问题求助

Possible Fixes & Explanations #

1. The Repeating DEBUG Logs (Directly Linked to num_worker=8)

2. The _pickle.UnpicklingError: pickle data was truncated