polars/rust memory monitoring

[Enucatl]

Hi, I'm trying to use scalene to profile some polars code.

However, I don't get any information about the memory usage in the rust code. Is that a known limitation of scalene?

[emeryberger]

That's surprising - AFAIK Rust uses the standard memory allocator, which Scalene intercepts. Can you share an MRE?

[Enucatl]

Upon reading more about the polars internals, I found out that they use jemalloc instead of glibc's malloc. Is that maybe why scalene doesn't pick it up?

[Enucatl]

I'm running this. There are no memory stats for the lines or functions using polars.

import polars as pl
import pandas as pd
from scalene import scalene_profiler


size = 10_000_000

def merge_pandas():
    pd_df1 = pd.DataFrame({"a": range(size), "b": range(0, 2 * size, 2)})
    pd_df2 = pd.DataFrame({"a": range(size), "c": range(0, 3 * size, 3)})
    pd_df3 = pd.merge(pd_df1, pd_df2, on="a", how="left")
    return pd_df3

def merge_polars():
    pl_df1 = pl.DataFrame({"a": range(size), "b": range(0, 2 * size, 2)})
    pl_df2 = pl.DataFrame({"a": range(size), "c": range(0, 3 * size, 3)})
    pl_df3 = pl_df1.join(pl_df2, on="a", how="left")
    return pl_df3


if __name__ == "__main__":
    print(pl.__version__)
    print(pd.__version__)
    scalene_profiler.start()
    merge_polars()
    merge_pandas()
    scalene_profiler.stop()

$ scalene --version
Scalene version 1.5.20 (2023.02.27)
NOTE: The GPU is currently running in a mode that can reduce Scalene's accuracy when reporting GPU utilization.
Run once as Administrator or root (i.e., prefixed with `sudo`) to enable per-process GPU accounting.
$ scalene --cli mre.py
0.16.16
1.5.3

                                       Memory usage: ▁▁▁▇▆▆▅▆▅▃▁▃▃▃▄▄▆▇▇▆▆▄▃▃▂▃▃ (max: 800.870 MB, growth rate:  11%)                                       
                                                    mre.py: % of time = 100.00% (4.847s) out of 4.847s.                                                     
       ╷       ╷       ╷      ╷       ╷       ╷       ╷               ╷       ╷                                                                             
       │Time   │–––––– │––––… │–––––– │Memory │–––––– │–––––––––––    │Copy   │                                                                             
  Line │Python │native │syst… │GPU    │Python │peak   │timeline/%     │(MB/s) │mre.py                                                                       
╺━━━━━━┿━━━━━━━┿━━━━━━━┿━━━━━━┿━━━━━━━┿━━━━━━━┿━━━━━━━┿━━━━━━━━━━━━━━━┿━━━━━━━┿━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╸
     1 │       │       │   2% │ 92%   │ 100%  │   10M │▁              │     9 │import polars as pl                                                          
     2 │    9% │    7% │  15% │       │  99%  │   20M │▁▁             │    48 │import pandas as pd                                                          
     3 │       │       │      │       │       │       │               │       │from scalene import scalene_profiler                                         
     4 │       │       │      │       │       │       │               │       │                                                                             
     5 │       │       │      │       │       │       │               │       │                                                                             
     6 │       │       │      │       │       │       │               │       │size = 10_000_000                                                            
     7 │       │       │      │       │       │       │               │       │                                                                             
     8 │       │       │      │       │       │       │               │       │def merge_pandas():                                                          
     9 │       │       │   1% │       │  20%  │  313M │▁▁▁  11%       │    21 │    pd_df1 = pd.DataFrame({"a": range(size), "b": range(0, 2 * size, 2)})    
    10 │       │       │      │       │   1%  │  305M │▁▁▁  11%       │    47 │    pd_df2 = pd.DataFrame({"a": range(size), "c": range(0, 3 * size, 3)})    
    11 │    1% │   28% │  26% │       │   1%  │  696M │▄▄▂▂▃▄▆▆▃  77% │   147 │    pd_df3 = pd.merge(pd_df1, pd_df2, on="a", how="left")                    
    12 │       │       │      │       │       │       │               │       │    return pd_df3                                                            
    13 │       │       │      │       │       │       │               │       │                                                                             
    14 │       │       │      │       │       │       │               │       │def merge_polars():                                                          
    15 │       │       │      │       │       │       │               │       │    pl_df1 = pl.DataFrame({"a": range(size), "b": range(0, 2 * size, 2)})    
    16 │       │       │      │  8%   │       │       │               │       │    pl_df2 = pl.DataFrame({"a": range(size), "c": range(0, 3 * size, 3)})    
    17 │       │    4% │      │       │       │       │               │    53 │    pl_df3 = pl_df1.join(pl_df2, on="a", how="left")                         
    18 │       │       │      │       │       │       │               │       │    return pl_df3                                                            
    19 │       │       │      │       │       │       │               │       │                                                                             
    20 │       │       │      │       │       │       │               │       │                                                                             
    21 │       │       │      │       │       │       │               │       │if __name__ == "__main__":                                                   
    22 │       │       │      │       │       │       │               │       │    scalene_profiler.start()                                                 
    23 │       │       │      │       │       │       │               │       │    merge_polars()                                                           
    24 │       │       │      │       │       │       │               │       │    merge_pandas()                                                           
    25 │       │       │      │       │       │       │               │       │    scalene_profiler.stop()                                                  
    26 │       │       │      │       │       │       │               │       │                                                                             
       │       │       │      │       │       │       │               │       │                                                                             
╶──────┼───────┼───────┼──────┼───────┼───────┼───────┼───────────────┼───────┼────────────────────────────────────────────────────────────────────────────╴
       │       │       │      │       │       │       │               │       │function summary for mre.py                                                  
     8 │    2% │   27% │  30% │       │   3%  │  696M │█████████  99% │   215 │merge_pandas                                                                 
    14 │       │    3% │   3% │  8%   │       │       │               │    53 │merge_polars                                                                 
       ╵       ╵       ╵      ╵       ╵       ╵       ╵               ╵       ╵                                                                             
Top AVERAGE memory consumption, by line:
(1)     9:   313 MB                                                                                                                                          
(2)    10:   305 MB                                                                                                                                          
(3)     2:    20 MB                                                                                                                                          
(4)     1:    10 MB                                                                                                                                          
Top PEAK memory consumption, by line:
(1)    11:   696 MB                                                                                                                                          
(2)     9:   313 MB                                                                                                                                          
(3)    10:   305 MB                                                                                                                                          
(4)     2:    20 MB                                                                                                                                          
(5)     1:    10 MB                                                                                                                                          
NOTE: The GPU is currently running in a mode that can reduce Scalene's accuracy when reporting GPU utilization.
Run once as Administrator or root (i.e., prefixed with `sudo`) to enable per-process GPU accounting.

[emeryberger]

Thanks for the MRE! Looking into it now.

[Enucatl]

I add it here, since I haven't found this solution for the GPU warning elsewhere.
Scalene recommends running once with sudo to "fix it", but:

• it doesn't precisely say what the problem is (accounting mode is not enabled for the GPU)
• the python environment may not be available for root, so you can't even run it anyway, as was my case
• people may not want to give scalene superuser privileges, especially without knowing exactly what it intends to do with such privileges
• the configuration would be reset after reboot

As a solution, for nvidia GPUs at least, the following command sets accounting mode for all GPUs, I would think at least this should be in the output of the warning:

sudo /usr/bin/nvidia-smi --accounting-mode=ENABLED

or, to make it autorun on boot, I created /etc/systemd/system/nvidia-accounting-mode.service:

[Unit]
Description=Enable nvidia accounting mode
Requires=nvidia-persistenced.service
After=nvidia-persistenced.service

[Service]
Type=oneshot
ExecStart=/usr/bin/nvidia-smi --accounting-mode=ENABLED

[Install]
WantedBy=multi-user.target

[pspacek]

I've just encountered (presumably) the same problem while trying to profile a Python library linked to jemalloc.