BUG: `kernprof` executed code inconsistent with `sys.modules['__main__']`, causing pickling issues

[TTsangSC]

The context

I was trying to restart my pet project of extending profiling into child Python processes (e.g. via subprocess, multiprocessing, and os.fork()) with kernprof (or whatever its successor in 6.0 will be called), which in some previous PRs I've been gradually working towards. It may need a bit elbow grease and boilerplate to port all the stuff over, but it should be very doable: after all, I've already written a pytest plugin which does that for tests, and kernprof is kinda like pytest in that both run other Python code in a somewhat controlled environment.

The test

To gauge how much work is ahead of us and to demonstrate what exactly the change will be good for, I needed a test use-case. Say we have a script which goes

`multiproc-test.py`

from multiprocessing import Pool


def my_sum(x: list[int]) -> int:
    result: int = 0
    for item in x:
        result += item
    return result


def main(length: int = 1000, n: int = 4) -> None:
    my_list: list[int] = list(range(length))
    sublists: list[list[int]] = []
    subsums: list[int]
    sublength = length // n
    if sublength * n < length:
        sublength += 1
    while my_list:
        sublist, my_list = my_list[:sublength], my_list[sublength:]
        sublists.append(sublist)
    with Pool(n) as pool:
        subsums = pool.map(my_sum, sublists)
        pool.close()
        pool.join()
    print(my_sum(subsums))


if __name__ == '__main__':
    main()  # stdout: '499500\n'

... so just a very simple scenario where we distribute the workload of summing integers over multiple processes. What I expected to happen was that kernprof would only see the one call to my_sum() we made in the main process and report data thereon, but unfortunately...

The bug

... the code didn't even run.

$ kernprof -q multiproc-test.py
AttributeError: module '__main__' has no attribute 'my_sum'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  ...
_pickle.PicklingError: Can't pickle <function my_sum at 0x105db62a0>: it's not found as  __main__.my_sum
when serializing tuple item 0
when serializing tuple item 0
when serializing tuple item 3

Note how we didn't even set up (line-)profiling and yet the execution still failed. The crux of the issue is that for most of the modes under which kernprof runs stuff (except for kernprof --prof-mod=<some_explicit_target> --line-by-line -m <some_module>), the namespace in which the code is executed is not vars(sys.modules['__main__']), which results in errors when multiprocessing.Pool.map() looks in the latter to find the function to pickle and send to the child processes.

An easy fix for the bug is that as with line_profiler.autoprofile.autoprofile.run(..., as_module=True), we always use a dummy module object and its namespace, and monkey-patch it into sys.modules['__main__'], regardless of whether the code is run as a module.

Notes

This issue is mostly just for contextualization and documentation; just need to also whip up some tests and I'll push a PR to fix it.

[Erotemic]

So this is only an issue because you have your multiprocessing code in __main__, if you had it in a separate module, then it at least would not cause pickle issues? And we are effectively doing this monkey patching anyway in autoprofile, so it's not a new thing to do it in kernprof itself?

[TTsangSC]

Yes to the latter part (we already do that for autoprof in some cases, as you've noted); probably yes for the former, seeing that this is just a namespace issue. I see that you've already left a comment over at the PR too, will check in a moment.

[TTsangSC]

I have both good and bad news. TLDR: what you've noted is entirely right (so #423 should be clear to go), but I discovered another bug involving duplicate functions not "seen" by the profiler.

Good news

Yes, indeed the problem is just caused by _pickle's not finding the requisite object at __main__. So if one were to e.g. move the my_sum() to a module and import it in the script (which contains the main() and multiprocessing and what's not), everything would be fine even on HEAD, because the function passed to Pool.map() is not from the problematic namespace.

Bad news

To test this I just moved the original multiproc-test.py to somewhere where from which it can be imported (say to some-dir-on-my-python-path/mp_test_module.py), and made a copy (say mp-test-script.py) as the script to be run, where the only difference is that it imports my_sum() from the former instead of defining it locally.

`mp-test-script.py`

from multiprocessing import Pool

from mp_test_module import my_sum


def main(length: int = 1000, n: int = 4) -> None:
    my_list: list[int] = list(range(length))
    sublists: list[list[int]] = []
    subsums: list[int]
    sublength = length // n
    if sublength * n < length:
        sublength += 1
    while my_list:  # <- NOTE: this is line 13
        sublist, my_list = my_list[:sublength], my_list[sublength:]
        sublists.append(sublist)
    with Pool(n) as pool:
        subsums = pool.map(my_sum, sublists)
        pool.close()
        pool.join()
    print(my_sum(subsums))


if __name__ == '__main__':
    main()

So, running kernprof -l -prof-mod=mp_test_module mp-test-script.py does "work" as expected – the script executes fine, and the profiler catches the one call to my_sum() we made directly in main(), but not the other 4 calls made in child processes.

HOWEVER, this resulted in some rather unexpected profiler output:

(terminal output, pathological case)

$ PYTHONPATH="${PYTHONPATH}:${PWD}/some-dir-on-my-python-path" kernprof -l --outfile=out --prof-mod=mp_test_module -vz mp-test-script.py
499500
Wrote profile results to 'out'
Timer unit: 1e-06 s

Total time: 2e-06 s
File: [...]/some-dir-on-my-python-path/mp_test_module.py
Function: my_sum at line 4

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
     4                                           def my_sum(x: list[int]) -> int:
     5         1          0.0      0.0      0.0      result: int = 0
     6         5          1.0      0.2     50.0      for item in x:
     7         4          1.0      0.2     50.0          result += item
     8         1          0.0      0.0      0.0      return result

Total time: 0.083778 s
File: [...]/some-dir-on-my-python-path/mp_test_module.py
Function: main at line 11

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
    11         1          1.0      1.0      0.0  def main(length: int = 1000, n: int = 4) -> None:
    12                                               my_list: list[int] = list(range(length))
    13         5          6.0      1.2      0.0      sublists: list[list[int]] = []
    14                                               subsums: list[int]
    15         4          0.0      0.0      0.0      sublength = length // n
    16         2      21868.0  10934.0     26.1      if sublength * n < length:
    17         1      54344.0  54344.0     64.9          sublength += 1
    18         1         34.0     34.0      0.0      while my_list:
    19         1       7499.0   7499.0      9.0          sublist, my_list = my_list[:sublength], my_list[sublength:]
    20         1         26.0     26.0      0.0          sublists.append(sublist)
    21                                               with Pool(n) as pool:
    22                                                   subsums = pool.map(my_sum, sublists)
    23                                                   pool.close()
    24                                                   pool.join()
    25                                               print(my_sum(subsums))

Note how mp_test_module.main() is somehow getting hits which are all on the wrong lines. What happened was probably that:

• mp_test_module.main() and mp-test-script.py::main() has the same source code, so they compile down to the same bytecode.
• Since the profiler is only profiling mp_test_module but not mp-test-script.py, it only sees the main() in the former.
• Since only one main() is seen by the profiler, it never knows about the dupe and doesn't do the usual padding and de-duplication.
• Since both main()s are unmodified, when mp-test-script.py::main() is executed, some of its line hashes would overlap with those of mp_test_module.main(), which the profiler then incorrectly surmises has been called.

Indeed, if one were to profile both mp_test_module and mp-test-script.py, the profiling data is correctly attributed, because the profiler has seen both main()s and swapped out the bytecode on one of them:

(click to expand)

$ PYTHONPATH="${PYTHONPATH}:${PWD}/some-dir-on-my-python-path" kernprof -l --outfile=out --prof-mod=mp_test_module --prof-mod=mp-test-script.py -vz mp-test-script.py
499500
Wrote profile results to 'out'
Timer unit: 1e-06 s

Total time: 2e-06 s
File: [...]/some-dir-on-my-python-path/mp_test_module.py
Function: my_sum at line 4

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
     4                                           def my_sum(x: list[int]) -> int:
     5         1          0.0      0.0      0.0      result: int = 0
     6         5          0.0      0.0      0.0      for item in x:
     7         4          1.0      0.2     50.0          result += item
     8         1          1.0      1.0     50.0      return result

Total time: 0.087448 s
File: mp-test-script.py
Function: main at line 6

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
     6                                           def main(length: int = 1000, n: int = 4) -> None:
     7         1          6.0      6.0      0.0      my_list: list[int] = list(range(length))
     8         1          0.0      0.0      0.0      sublists: list[list[int]] = []
     9                                               subsums: list[int]
    10         1          0.0      0.0      0.0      sublength = length // n
    11         1          1.0      1.0      0.0      if sublength * n < length:
    12                                                   sublength += 1
    13         5          2.0      0.4      0.0      while my_list:
    14         4          2.0      0.5      0.0          sublist, my_list = my_list[:sublength], my_list[sublength:]
    15         4          2.0      0.5      0.0          sublists.append(sublist)
    16         2      24650.0  12325.0     28.2      with Pool(n) as pool:
    17         1      55940.0  55940.0     64.0          subsums = pool.map(my_sum, sublists)
    18         1         27.0     27.0      0.0          pool.close()
    19         1       6795.0   6795.0      7.8          pool.join()
    20         1         23.0     23.0      0.0      print(my_sum(subsums))

Seems that even after the discussions and fixes we've had over #348–#351 there are still edge cases unaccounted for. Possible fixes are that either:

• Instead of only doing bytecode-based hashing to identify the code being run, we hash the entire code object (or at least include metadata like code.co_name, .co_filename, and .co_firstlineno); or
• Instead of only padding code objects with clashing bytecodes from the second one onwards, do so from the first (so that everything passed to the profiler ends up being padded).

Out of the two, the second is probably way faster since we don't increase the per-line-event overhead.

Will write a separate issue and PR to fix that.

[TTsangSC]

Oops I forgot to auto-close.