Python references that go out of scope might not release the memory allocated by _upb_Arena_SlowMalloc #14571

tvalentyn · 2023-10-13T19:22:27Z

Repro:

In a clean environment, run pip install apache-beam==2.51.0.
Download + unpack pipeline.pb (see attached https://github.com/protocolbuffers/upb/files/12898151/pipeline.pb.zip).
Create a memleak_repro.py

from apache_beam.portability.api import beam_runner_api_pb2
from apache_beam.io.gcp import gcsio
from apache_beam.io.gcp import gcsfilesystem

if len(sys.argv) <= 1:
  print("Specify proto path!")
  sys.exit(1)

path = sys.argv[1]
if path.startswith('gs'):
  open = gcsio.GcsIO().open

with open(path, 'rb') as f:
  pipeline = beam_runner_api_pb2.Pipeline()
  pipeline.ParseFromString(f.read())

for transform in pipeline.components.transforms:
  for _ in range(100000):
    reference = pipeline.components.transforms[transform].outputs

Run python memleak_repro.py pipeline.pb, observe RAM usage in top/htop, etc, increasing the number of iterations as necessary to keep the process running longer.
For profiling, run:

pip install memray

memray run -o output.bin --force memleak_repro.py pipeline.pb ; memray table --leak output.bin -o table.html --force ; memray flamegraph --leak output.bin -o flamegraph.html -f

Open table.html, double click on the Size column, note ~36MB reported leaked memory in memleak_repro.py:20

Leak doesn't happen with PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python set.

Leak can also be observed with C-stack tools, such as tcmalloc+pprof:

apt install google-pprof
LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libtcmalloc.so.4:$LD_PRELOAD  HEAPPROFILE=`pwd`/  python memleak_repro.py pipeline.pb
google-pprof /home/valentyn/.pyenv/versions/py310/bin/python --inuse_space  _219720.0005.heap  --base=_219720.0004.heap 
Using local file /home/valentyn/.pyenv/versions/py310/bin/python.
Using local file _219720.0005.heap.
Welcome to pprof!  For help, type 'help'.
(pprof) top
Total: 1152.0 MB
  1152.0 100.0% 100.0%   1152.0 100.0% _upb_Arena_SlowMalloc
(pprof)

valgrind+massif:

apt install valgrind
PYTHONMALLOC=malloc  valgrind  --tool=massif   --verbose  /home/valentyn/.pyenv/versions/py310/bin/python memleak_repro.py pipeline.pb

ms_print massif.out.341396 

...
->18.09% (817,040B) 0x59A17F1: _upb_Arena_SlowMalloc (in /home/valentyn/.pyenv/versions/3.10.4/envs/py310/lib/python3.10/site-packages/google/_upb/_message.abi3.so)
| ->07.00% (316,176B) 0x59A0051: upb_strtable_insert (in /home/valentyn/.pyenv/versions/3.10.4/envs/py310/lib/python3.10/site-packages/google/_upb/_message.abi3.so)
| | ->06.05% (273,072B) 0x599FDDF: upb_strtable_resize (in /home/valentyn/.pyenv/versions/3.10.4/envs/py310/lib/python3.10/site-packages/google/_upb/_message.abi3.so)
| | | ->06.05% (273,072B) 0x599FED6: upb_strtable_insert (in /home/valentyn/.pyenv/versions/3.10.4/envs/py310/lib/python3.10/site-packages/google/_upb/_message.abi3.so)
| | |   ->04.17% (188,096B) 0x598F0D7: _upb_DefPool_InsertSym (in /home/valentyn/.pyenv/versions/3.10.4/envs/py310/lib/python3.10/site-packages/google/_upb/_message.abi3.so)
| | |   | ->03.89% (175,568B) 0x5990B19: _upb_EnumValueDefs_New (in /home/valentyn/.pyenv/versions/3.10.4/envs/py310/lib/python3.10/site-packages/google/_upb/_message.abi3.so)
| | |   | | ->03.89% (175,568B) 0x5990332: _upb_EnumDefs_New (in /home/valentyn/.pyenv/versions/3.10.4/envs/py310/lib/python3.10/site-packages/google/_upb/_message.abi3.so)
| | |   | |   ->03.89% (175,568B) 0x5995356: _upb_MessageDefs_New (in /home/valentyn/.pyenv/versions/3.10.4/envs/py310/lib/python3.10/site-packages/google/_upb/_message.abi3.so)
| | |   | |     ->03.89% (175,568B) 0x5992E88: _upb_FileDef_Create (in /home/valentyn/.pyenv/versions/3.10.4/envs/py310/lib/python3.10/site-packages/google/_upb/_message.abi3.so)
| | |   | |       ->03.89% (175,568B) 0x598FC4B: upb_DefBuilder_AddFileToPool (in /home/valentyn/.pyenv/versions/3.10.4/envs/py310/lib/python3.10/site-packages/google/_upb/_message.abi3.so)
| | |   | |         ->03.89% (175,568B) 0x598F699: _upb_DefPool_AddFile (in /home/valentyn/.pyenv/versions/3.10.4/envs/py310/lib/python3.10/site-packages/google/_upb/_message.abi3.so)
| | |   | |           ->03.89% (175,568B) 0x598165C: PyUpb_DescriptorPool_DoAddSerializedFile (in /home/valentyn/.pyenv/versions/3.10.4/envs/py310/lib/python3.10/site-packages/google/_upb/_message.abi3.so)
| | |   | |             ->03.89% (175,568B) 0x31A7B1: method_vectorcall_O (descrobject.c:460)
| | |   | |               ->03.89% (175,568B) 0x16687A: _PyObject_VectorcallTstate (abstract.h:114)
| | |   | |                 ->03.89% (175,568B) 0x16687A: PyObject_Vectorcall (abstract.h:123)
| | |   | |                   ->03.89% (175,568B) 0x16687A: call_function (ceval.c:5867)
| | |   | |                     ->03.89% (175,568B) 0x16687A: _PyEval_EvalFrameDefault (ceval.c:4198)
| | |   | |                       ->03.89% (175,568B) 0x22F7B4: _PyEval_EvalFrame (pycore_ceval.h:46)
| | |   | |                         ->03.89% (175,568B) 0x22F7B4: _PyEval_Vector (ceval.c:5065)
| | |   | |                           ->03.89% (175,568B) 0x22F7B4: PyEval_EvalCode (ceval.c:1134)
| | |   | |                             ->03.89% (175,568B) 0x3593F8: builtin_exec_impl (bltinmodule.c:1056)
| | |   | |                               ->03.89% (175,568B) 0x3593F8: builtin_exec (bltinmodule.c.h:371)

...

Using memray in --native mode also points to _upb_Arena_SlowMalloc .

Command line: memray run --native -o output.bin --force memleak_repro.py pipeline.pb ; memray table --leak output.bin -o table.html --force ; memray flamegraph --leak output.bin -o flamegraph.html -f

The text was updated successfully, but these errors were encountered:

tvalentyn · 2023-10-13T19:23:29Z

cc: @haberman @ericsalo

tvalentyn · 2023-10-13T19:27:34Z

pipeline.pb.zip

tvalentyn · 2023-10-13T19:30:12Z

This is the proto we use: https://github.com/apache/beam/blob/47d0fd566f86aaad35d26709c52ee555381823a4/model/pipeline/src/main/proto/org/apache/beam/model/pipeline/v1/beam_runner_api.proto#L91

I am sure the repro can be simplified further.

anandolee · 2023-10-17T20:14:37Z

@tvalentyn can you help create a repo that not depend on apache-beam? Thanks

tvalentyn · 2023-10-17T21:41:19Z

I am OOO atm and likely won't get to it before next week, but cc'ing: @AnandInguva in case he has time and interested.

tvalentyn · 2023-10-19T23:35:46Z

Interesting. I gave it a try, and I was not able to recreate the leak in isolation while using essentially the same proto.

I'll check whether this has to do with how beam compiles proto stubs or something else.

haberman · 2023-10-20T00:28:14Z

I'll check whether this has to do with how beam compiles proto stubs or something else.

That would be really helpful, thanks!

tvalentyn · 2023-10-21T01:48:39Z

Got a smaller repro. Appears that there some interaction between protobuf and dill. Beam uses old version of dill, but the issue is reproducible even with latest dill. Good news is there is some work underway for Beam to not depend on Dill or at least be able to use cloudpickle as default..

pip install dill protobuf grpcio-tools dill memray

import dill # Leak doesn't happen without this import.
import sys
from google.protobuf import text_format

proto_def="""
syntax = "proto3";

message MessageWithAMap {
  map<string, string> some_map = 1;
}
"""

with open("repro.proto", "w") as f:
  f.write(proto_def)

from grpc_tools import protoc
args = [
    sys.executable, # unused
    "repro.proto",
    "--proto_path", ".", "--python_out", "."]
protoc.main(args)

import repro_pb2
from repro_pb2 import MessageWithAMap

proto="""
some_map {
  key: "Length of this key increases the leak"
  value:"No observable impact on the leak"
}
"""

message = MessageWithAMap()
text_format.Parse(proto, message)

key = list(message.some_map)[0]
for _ in range(500000):
  unused_values = message.some_map[key] # memory leaks here

rm -f *pb2.py* ; memray run  -o output.bin --force memleak_repro.py  &&  memray table --leak output.bin -o table.html --force

tvalentyn · 2023-10-23T17:15:29Z

import dill # Leak doesn't happen without this import.

Leak correlates with this line: https://github.com/uqfoundation/dill/blob/f66ed3b602d6e9b942542491614dfd18622583d7/dill/_dill.py#L72C32-L72C32

Here is a repro that doesn't use dill:

import __main__ as _main_module # Leak doesn't happen without this import.
import sys
from google.protobuf import text_format

proto_def="""
syntax = "proto3";

message MessageWithAMap {
  map<string, string> some_map = 1;
}
"""

with open("repro.proto", "w") as f:
  f.write(proto_def)

from grpc_tools import protoc
args = [
    sys.executable, # unused
    "repro.proto",
    "--proto_path", ".", "--python_out", "."]
protoc.main(args)

from repro_pb2 import MessageWithAMap

proto="""
some_map {
  key: "Length of this key increases the leak"
  value:"No observable impact on the leak"
}
"""

message = MessageWithAMap()
text_format.Parse(proto, message)

key = list(message.some_map)[0]
for _ in range(500000):
  unused_values = message.some_map[key] # memory leaks here

tvalentyn · 2023-10-23T19:11:52Z

Leak doesn't happen with protobuf==3.20.3 or with PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python.

haberman · 2023-10-23T19:54:17Z

I think I see the root cause. In map lookup, we use the message's arena to convert the Python key object to raw C data that we can use to perform the lookup in upb. But this data will never be freed until the message is freed:

protobuf/python/map.c

Lines 203 to 205 in 022d223

    
           upb_Arena* arena = PyUpb_Arena_Get(self->arena); 
        
           upb_MessageValue u_key, u_val; 
        
           if (!PyUpb_PyToUpb(key, key_f, &u_key, arena)) return NULL;

We should use a temporary arena to perform the conversion. Only at the point where we are allocating non-temporary data should we use the message's arena.

tvalentyn · 2023-10-23T20:29:07Z

That explains why import __main__ as _main_module has an effect. Leak is still happening without this statement, but this statement makes it easier to register the leak with the profiler as references to message persist longer.

tvalentyn · 2023-10-23T20:36:17Z

Indeed, in top memory keeps growing without import __main__ as _main_module and memray also registers the leak when used without --leak flag.

tvalentyn · 2023-10-23T21:21:09Z

PyUpb_MapContainer_Get also has this leak https://github.com/protocolbuffers/protobuf/blob/022d223488c0807e11d1c2509123901f5cc6b0b6/python/map.c#L259-L261`

tvalentyn · 2023-10-23T21:25:16Z

And probably several other places in map.c where PyUpb_PyToUpb is used.

Previously we were allocating memory on the message's arena every time we performed a `map[key]` or `map.get(key)` operation. This is unnecessary, as the key's data is only needed ephemerally, for the duration of the lookup, and we can therefore alias the Python object's string data instead of copying it. This required fixing a bug in the convert.c operation. Previously in the `arena==NULL` case, if the user passes a bytes object instead of a unicode string, the code would return a pointer to a temporary Python object that had already been freed, leading to use-after-free. I fixed this by referencing the bytes object's data directly, and using utf8_range to verify the UTF-8. Fixes: #14571 PiperOrigin-RevId: 578311252

Previously we were allocating memory on the message's arena every time we performed a `map[key]` or `map.get(key)` operation. This is unnecessary, as the key's data is only needed ephemerally, for the duration of the lookup, and we can therefore alias the Python object's string data instead of copying it. This required fixing a bug in the convert.c operation. Previously in the `arena==NULL` case, if the user passes a bytes object instead of a unicode string, the code would return a pointer to a temporary Python object that had already been freed, leading to use-after-free. I fixed this by referencing the bytes object's data directly, and using utf8_range to verify the UTF-8. Fixes: #14571 PiperOrigin-RevId: 578563555

tvalentyn mentioned this issue Oct 13, 2023

[Bug]: [Python SDK] Memory leak in 2.47.0 - 2.51.0 SDKs. apache/beam#28246

Closed

15 tasks

haberman transferred this issue from protocolbuffers/upb Oct 31, 2023

copybara-service bot mentioned this issue Nov 1, 2023

Fixed Python memory leak in map lookup. #14577

Closed

copybara-service bot closed this as completed in b0eeb35 Nov 1, 2023

haberman mentioned this issue Nov 1, 2023

Fixed Python memory leak in map lookup. #14583

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Python references that go out of scope might not release the memory allocated by _upb_Arena_SlowMalloc #14571

Python references that go out of scope might not release the memory allocated by _upb_Arena_SlowMalloc #14571

tvalentyn commented Oct 13, 2023 •

edited

tvalentyn commented Oct 13, 2023

tvalentyn commented Oct 13, 2023

tvalentyn commented Oct 13, 2023

anandolee commented Oct 17, 2023

tvalentyn commented Oct 17, 2023

tvalentyn commented Oct 19, 2023

haberman commented Oct 20, 2023

tvalentyn commented Oct 21, 2023 •

edited

tvalentyn commented Oct 23, 2023 •

edited

tvalentyn commented Oct 23, 2023

haberman commented Oct 23, 2023

tvalentyn commented Oct 23, 2023 •

edited

tvalentyn commented Oct 23, 2023

tvalentyn commented Oct 23, 2023 •

edited

tvalentyn commented Oct 23, 2023 •

edited

Python references that go out of scope might not release the memory allocated by _upb_Arena_SlowMalloc #14571

Python references that go out of scope might not release the memory allocated by _upb_Arena_SlowMalloc #14571

Comments

tvalentyn commented Oct 13, 2023 • edited

tvalentyn commented Oct 13, 2023

tvalentyn commented Oct 13, 2023

tvalentyn commented Oct 13, 2023

anandolee commented Oct 17, 2023

tvalentyn commented Oct 17, 2023

tvalentyn commented Oct 19, 2023

haberman commented Oct 20, 2023

tvalentyn commented Oct 21, 2023 • edited

tvalentyn commented Oct 23, 2023 • edited

tvalentyn commented Oct 23, 2023

haberman commented Oct 23, 2023

tvalentyn commented Oct 23, 2023 • edited

tvalentyn commented Oct 23, 2023

tvalentyn commented Oct 23, 2023 • edited

tvalentyn commented Oct 23, 2023 • edited

tvalentyn commented Oct 13, 2023 •

edited

tvalentyn commented Oct 21, 2023 •

edited

tvalentyn commented Oct 23, 2023 •

edited

tvalentyn commented Oct 23, 2023 •

edited

tvalentyn commented Oct 23, 2023 •

edited

tvalentyn commented Oct 23, 2023 •

edited