tractor/tractor/spawn/_mp_fixup_main.py

86 lines
2.8 KiB
Python

# tractor: structured concurrent "actors".
# Copyright 2018-eternity Tyler Goodlet.
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
'''
(Originally) Helpers pulled verbatim from ``multiprocessing.spawn``
to aid with "fixing up" the ``__main__`` module in subprocesses.
Now just delegates directly to appropriate `mp.spawn` fns.
Note
----
These helpers are needed for any spawning backend that doesn't already
handle this. For example it's needed when using our
`start_method='trio'` backend but not when we're already using
a ``multiprocessing`` backend such as 'mp_spawn', 'mp_forkserver'.
?TODO?
- what will be required for an eventual subint backend?
The helpers imported from `mp.spawn` provide the stdlib's
spawn/forkserver bootstrap that rebuilds the parent's `__main__` in
a fresh child interpreter. In particular, we capture enough info to
later replay the parent's main module as `__mp_main__` (or by path)
in the child process.
See:
https://docs.python.org/3/library/multiprocessing.html#the-spawn-and-forkserver-start-methods
'''
import multiprocessing as mp
from multiprocessing.spawn import (
_fixup_main_from_name as _fixup_main_from_name,
_fixup_main_from_path as _fixup_main_from_path,
get_preparation_data,
)
from typing import NotRequired
from typing import TypedDict
class ParentMainData(TypedDict):
init_main_from_name: NotRequired[str]
init_main_from_path: NotRequired[str]
def _mp_figure_out_main(
inherit_parent_main: bool = True,
) -> ParentMainData:
'''
Delegate to `multiprocessing.spawn.get_preparation_data()`
when `inherit_parent_main=True`.
Retrieve parent (actor) proc's `__main__` module data.
'''
if not inherit_parent_main:
return {}
proc: mp.Process = mp.current_process()
prep_data: dict = get_preparation_data(
name=proc.name,
)
# XXX, unserializable (and unneeded by us) by default
# see `mp.spawn.get_preparation_data()` impl details.
prep_data.pop('authkey', None)
d: ParentMainData = {}
if 'init_main_from_name' in prep_data:
d['init_main_from_name'] = prep_data['init_main_from_name']
if 'init_main_from_path' in prep_data:
d['init_main_from_path'] = prep_data['init_main_from_path']
return d