| # |
| # Licensed to the Apache Software Foundation (ASF) under one or more |
| # contributor license agreements. See the NOTICE file distributed with |
| # this work for additional information regarding copyright ownership. |
| # The ASF licenses this file to You under the Apache License, Version 2.0 |
| # (the "License"); you may not use this file except in compliance with |
| # the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| # |
| |
| """Implements a shared object that spans processes. |
| |
| This object will be instanciated once per VM and methods will be invoked |
| on it via rpc. |
| """ |
| # pytype: skip-file |
| |
| import atexit |
| import logging |
| import multiprocessing.managers |
| import os |
| import tempfile |
| import threading |
| import time |
| import traceback |
| from typing import Any |
| from typing import Callable |
| from typing import Dict |
| from typing import Generic |
| from typing import Optional |
| from typing import TypeVar |
| |
| import fasteners |
| |
| # In some python versions, there is a bug where AutoProxy doesn't handle |
| # the kwarg 'manager_owned'. We implement our own backup here to make sure |
| # we avoid this problem. More info here: |
| # https://stackoverflow.com/questions/46779860/multiprocessing-managers-and-custom-classes |
| autoproxy = multiprocessing.managers.AutoProxy # type: ignore[attr-defined] |
| |
| |
| def patched_autoproxy( |
| token, |
| serializer, |
| manager=None, |
| authkey=None, |
| exposed=None, |
| incref=True, |
| manager_owned=True): |
| return autoproxy(token, serializer, manager, authkey, exposed, incref) |
| |
| |
| multiprocessing.managers.AutoProxy = patched_autoproxy # type: ignore[attr-defined] |
| |
| T = TypeVar('T') |
| AUTH_KEY = b'mps' |
| |
| |
| class _SingletonProxy: |
| """Proxies the shared object so we can release it with better errors and no |
| risk of dangling references in the multiprocessing manager infrastructure. |
| """ |
| def __init__(self, entry): |
| # Guard names so as to not conflict with names of underlying object. |
| self._SingletonProxy_entry = entry |
| self._SingletonProxy_valid = True |
| |
| # Used to make the shared object callable (see _AutoProxyWrapper below) |
| def singletonProxy_call__(self, *args, **kwargs): |
| if not self._SingletonProxy_valid: |
| raise RuntimeError('Entry was released.') |
| return self._SingletonProxy_entry.obj.__call__(*args, **kwargs) |
| |
| def singletonProxy_release(self): |
| assert self._SingletonProxy_valid |
| self._SingletonProxy_valid = False |
| |
| def singletonProxy_unsafe_hard_delete(self): |
| assert self._SingletonProxy_valid |
| self._SingletonProxy_entry.unsafe_hard_delete() |
| |
| def __getattr__(self, name): |
| if not self._SingletonProxy_valid: |
| raise RuntimeError('Entry was released.') |
| try: |
| return getattr(self._SingletonProxy_entry.obj, name) |
| except AttributeError as e: |
| # Swallow AttributeError exceptions so that they are ignored when |
| # calculating public functions. These can occur if __getattr__ is |
| # overriden, for example to only support some platforms. This will mean |
| # that these functions will be silently unavailable to the |
| # MultiProcessShared object, leading to worse errors when someone tries |
| # to use them, but it will keep them from breaking the whole object for |
| # functions which are unusable anyways. |
| logging.info( |
| 'Attribute %s is unavailable as a public function because ' |
| 'its __getattr__ function raised the following exception ' |
| '%s', |
| name, |
| e) |
| return None |
| |
| def __dir__(self): |
| # Needed for multiprocessing.managers's proxying. |
| dir = self._SingletonProxy_entry.obj.__dir__() |
| dir.append('singletonProxy_call__') |
| dir.append('singletonProxy_release') |
| dir.append('singletonProxy_unsafe_hard_delete') |
| return dir |
| |
| |
| class _SingletonEntry: |
| """Represents a single, refcounted entry in this process.""" |
| def __init__( |
| self, constructor, initialize_eagerly=True, hard_delete_callback=None): |
| self.constructor = constructor |
| self._hard_delete_callback = hard_delete_callback |
| self.refcount = 0 |
| self.lock = threading.Lock() |
| if initialize_eagerly: |
| self.obj = constructor() |
| self.initialied = True |
| else: |
| self.initialied = False |
| |
| def acquire(self): |
| with self.lock: |
| if not self.initialied: |
| self.obj = self.constructor() |
| self.initialied = True |
| self.refcount += 1 |
| return _SingletonProxy(self) |
| |
| def release(self, proxy): |
| proxy.singletonProxy_release() |
| with self.lock: |
| self.refcount -= 1 |
| if self.refcount == 0: |
| del self.obj |
| self.initialied = False |
| |
| def unsafe_hard_delete(self): |
| with self.lock: |
| if self.initialied: |
| del self.obj |
| self.initialied = False |
| if self._hard_delete_callback: |
| self._hard_delete_callback() |
| |
| |
| class _SingletonManager: |
| entries: Dict[Any, Any] = {} |
| |
| def __init__(self): |
| self._hard_delete_callback = None |
| |
| def set_hard_delete_callback(self, callback): |
| self._hard_delete_callback = callback |
| |
| def register_singleton( |
| self, |
| constructor, |
| tag, |
| initialize_eagerly=True, |
| hard_delete_callback=None): |
| assert tag not in self.entries, tag |
| self.entries[tag] = _SingletonEntry( |
| constructor, initialize_eagerly, hard_delete_callback) |
| |
| def has_singleton(self, tag): |
| return tag in self.entries |
| |
| def acquire_singleton(self, tag): |
| return self.entries[tag].acquire() |
| |
| def release_singleton(self, tag, obj): |
| return self.entries[tag].release(obj) |
| |
| def unsafe_hard_delete_singleton(self, tag): |
| self.entries[tag].unsafe_hard_delete() |
| |
| |
| _process_level_singleton_manager = _SingletonManager() |
| |
| _process_local_lock = threading.Lock() |
| |
| |
| class _SingletonRegistrar(multiprocessing.managers.BaseManager): |
| pass |
| |
| |
| _SingletonRegistrar.register( |
| 'acquire_singleton', |
| callable=_process_level_singleton_manager.acquire_singleton) |
| _SingletonRegistrar.register( |
| 'release_singleton', |
| callable=_process_level_singleton_manager.release_singleton) |
| _SingletonRegistrar.register( |
| 'unsafe_hard_delete_singleton', |
| callable=_process_level_singleton_manager.unsafe_hard_delete_singleton) |
| |
| |
| # By default, objects registered with BaseManager.register will have only |
| # public methods available (excluding __call__). If you know the functions |
| # you would like to expose, you can do so at register time with the `exposed` |
| # attribute. Since we don't, we will add a wrapper around the returned AutoProxy |
| # object to handle __call__ function calls and turn them into |
| # singletonProxy_call__ calls (which is a wrapper around the underlying |
| # object's __call__ function) |
| class _AutoProxyWrapper: |
| def __init__(self, proxyObject: multiprocessing.managers.BaseProxy): |
| self._proxyObject = proxyObject |
| |
| def __call__(self, *args, **kwargs): |
| return self._proxyObject.singletonProxy_call__(*args, **kwargs) |
| |
| def __getattr__(self, name): |
| return getattr(self._proxyObject, name) |
| |
| def get_auto_proxy_object(self): |
| return self._proxyObject |
| |
| def unsafe_hard_delete(self): |
| self._proxyObject.unsafe_hard_delete() |
| |
| |
| def _run_server_process(address_file, tag, constructor, authkey, life_line): |
| """ |
| Runs in a separate process. |
| Includes a 'Suicide Pact' monitor: If parent dies, I die. |
| """ |
| parent_pid = os.getppid() |
| |
| def cleanup_files(): |
| logging.info("Server process exiting. Deleting files for %s", tag) |
| try: |
| if os.path.exists(address_file): |
| os.remove(address_file) |
| if os.path.exists(address_file + ".error"): |
| os.remove(address_file + ".error") |
| except Exception as e: |
| logging.warning('Failed to cleanup files for tag %s: %s', tag, e) |
| |
| def handle_unsafe_hard_delete(): |
| cleanup_files() |
| os._exit(0) |
| |
| def _monitor_parent(): |
| """Checks if parent is alive every second.""" |
| while True: |
| try: |
| # This will break if parent dies. |
| life_line.recv_bytes() |
| except (EOFError, OSError, BrokenPipeError): |
| logging.warning( |
| "Process %s detected Parent %s died. Self-destructing.", |
| os.getpid(), |
| parent_pid) |
| cleanup_files() |
| os._exit(0) |
| time.sleep(0.5) |
| |
| atexit.register(cleanup_files) |
| |
| try: |
| t = threading.Thread(target=_monitor_parent, daemon=True) |
| |
| logging.getLogger().setLevel(logging.INFO) |
| multiprocessing.current_process().authkey = authkey |
| |
| serving_manager = _SingletonRegistrar( |
| address=('localhost', 0), authkey=authkey) |
| _process_level_singleton_manager.set_hard_delete_callback( |
| handle_unsafe_hard_delete) |
| _process_level_singleton_manager.register_singleton( |
| constructor, |
| tag, |
| initialize_eagerly=True, |
| hard_delete_callback=handle_unsafe_hard_delete) |
| # Start monitoring parent after initialisation is done to avoid |
| # potential race conditions. |
| t.start() |
| |
| server = serving_manager.get_server() |
| logging.info( |
| 'Process %s: Proxy serving %s at %s', os.getpid(), tag, server.address) |
| |
| with open(address_file + '.tmp', 'w') as fout: |
| fout.write('%s:%d' % server.address) |
| os.rename(address_file + '.tmp', address_file) |
| |
| server.serve_forever() |
| |
| except Exception: |
| tb = traceback.format_exc() |
| try: |
| with open(address_file + ".error.tmp", 'w') as fout: |
| fout.write(tb) |
| os.rename(address_file + ".error.tmp", address_file + ".error") |
| except Exception: |
| logging.error("CRITICAL ERROR IN SHARED SERVER:\n%s", tb) |
| os._exit(1) |
| |
| |
| class MultiProcessShared(Generic[T]): |
| """MultiProcessShared is used to share a single object across processes. |
| |
| For example, one could have the class:: |
| |
| class MyExpensiveObject(object): |
| def __init__(self, args): |
| [expensive initialization and memory allocation] |
| |
| def method(self, arg): |
| ... |
| |
| One could share a single instance of this class by wrapping it as:: |
| |
| shared_ptr = MultiProcessShared(lambda: MyExpensiveObject(...)) |
| my_expensive_object = shared_ptr.acquire() |
| |
| which could then be invoked as:: |
| |
| my_expensive_object.method(arg) |
| |
| This can then be released with:: |
| |
| shared_ptr.release(my_expensive_object) |
| |
| but care should be taken to avoid releasing the object too soon or |
| expensive re-initialization may be required, defeating the point of |
| using a shared object. |
| |
| |
| Args: |
| constructor: function that initialises / constructs the object if not |
| present in the cache. This function should take no arguments. It should |
| return an initialised object, or raise an exception if the object could |
| not be initialised / constructed. |
| tag: an indentifier to store with the cached object. If multiple |
| MultiProcessShared instances are created with the same tag, they will all |
| share the same proxied object. |
| path: a temporary path in which to create the inter-process lock |
| always_proxy: whether to direct all calls through the proxy, rather than |
| call the object directly for the process that created it |
| """ |
| def __init__( |
| self, |
| constructor: Callable[[], T], |
| tag: Any, |
| *, |
| path: str = tempfile.gettempdir(), |
| always_proxy: Optional[bool] = None, |
| spawn_process: bool = False): |
| self._constructor = constructor |
| self._tag = tag |
| self._path = path |
| self._always_proxy = False if always_proxy is None else always_proxy |
| self._proxy = None |
| self._manager = None |
| self._rpc_address = None |
| self._cross_process_lock = fasteners.InterProcessLock( |
| os.path.join(self._path, self._tag) + '.lock') |
| self._spawn_process = spawn_process |
| |
| def _get_manager(self): |
| if self._manager is None: |
| address_file = os.path.join(self._path, self._tag) + ".address" |
| while self._manager is None: |
| with _process_local_lock: |
| with self._cross_process_lock: |
| if not os.path.exists(address_file): |
| self._create_server(address_file) |
| |
| if _process_level_singleton_manager.has_singleton( |
| self._tag) and not self._always_proxy: |
| self._manager = _process_level_singleton_manager |
| else: |
| with open(address_file) as fin: |
| address = fin.read() |
| logging.info('Connecting to remote proxy at %s', address) |
| host, port = address.split(':') |
| # We need to be able to authenticate with both the manager and |
| # the process. |
| manager = _SingletonRegistrar( |
| address=(host, int(port)), authkey=AUTH_KEY) |
| multiprocessing.current_process().authkey = AUTH_KEY |
| try: |
| manager.connect() |
| self._manager = manager |
| except ConnectionError: |
| # The server is no longer good, assume it died. |
| os.unlink(address_file) |
| |
| return self._manager |
| |
| def acquire(self): |
| # TODO: Allow passing/parameterizing the callable here, in case they are |
| # not available at MultiProcessShared construction time (e.g. from side |
| # inputs) |
| # Caveat: They must always agree, as they will be ignored if the object |
| # is already constructed. |
| singleton = self._get_manager().acquire_singleton(self._tag) |
| # Trigger a sweep of zombie processes. |
| # calling active_children() has the side-effect of joining any finished |
| # processes, effectively reaping zombies from previous unsafe_hard_deletes. |
| if self._spawn_process: |
| multiprocessing.active_children() |
| return _AutoProxyWrapper(singleton) |
| |
| def release(self, obj): |
| self._manager.release_singleton(self._tag, obj.get_auto_proxy_object()) |
| |
| def unsafe_hard_delete(self): |
| """Force deletes the underlying object |
| |
| This function should be used with great care since any other references |
| to this object will now be invalid and may lead to strange errors. Only |
| call unsafe_hard_delete if either (a) you are sure no other references |
| to this object exist, or (b) you are ok with all existing references to |
| this object throwing strange errors when derefrenced. |
| """ |
| self._get_manager().unsafe_hard_delete_singleton(self._tag) |
| |
| def _create_server(self, address_file): |
| if self._spawn_process: |
| error_file = address_file + ".error" |
| |
| if os.path.exists(error_file): |
| try: |
| os.remove(error_file) |
| except OSError: |
| pass |
| |
| # Create a pipe to connect with child process |
| # used to clean up child process if parent dies |
| reader, writer = multiprocessing.Pipe(duplex=False) |
| self._life_line = writer |
| |
| ctx = multiprocessing.get_context('spawn') |
| p = ctx.Process( |
| target=_run_server_process, |
| args=(address_file, self._tag, self._constructor, AUTH_KEY, reader), |
| daemon=False # Must be False for nested proxies |
| ) |
| p.start() |
| logging.info("Parent: Waiting for %s to write address file...", self._tag) |
| |
| def cleanup_process(): |
| if self._life_line: |
| self._life_line.close() |
| if p.is_alive(): |
| logging.info( |
| "Parent: Terminating server process %s for %s", p.pid, self._tag) |
| p.terminate() |
| p.join() |
| try: |
| if os.path.exists(address_file): |
| os.remove(address_file) |
| if os.path.exists(error_file): |
| os.remove(error_file) |
| except Exception as e: |
| logging.warning( |
| 'Failed to cleanup files for tag %s in atexit handler: %s', |
| self._tag, |
| e) |
| |
| atexit.register(cleanup_process) |
| |
| start_time = time.time() |
| last_log = start_time |
| while True: |
| if os.path.exists(address_file): |
| break |
| |
| if os.path.exists(error_file): |
| with open(error_file, 'r') as f: |
| error_msg = f.read() |
| try: |
| os.remove(error_file) |
| except OSError: |
| pass |
| |
| if p.is_alive(): p.terminate() |
| raise RuntimeError(f"Shared Server Process crashed:\n{error_msg}") |
| |
| if not p.is_alive(): |
| exit_code = p.exitcode |
| raise RuntimeError( |
| "Shared Server Process died unexpectedly" |
| f" with exit code {exit_code}") |
| |
| if time.time() - last_log > 300: |
| logging.warning( |
| "Still waiting for %s to initialize... %ss elapsed)", |
| self._tag, |
| int(time.time() - start_time)) |
| last_log = time.time() |
| |
| time.sleep(0.05) |
| |
| logging.info('External process successfully started for %s', self._tag) |
| else: |
| # We need to be able to authenticate with both the manager |
| # and the process. |
| self._serving_manager = _SingletonRegistrar( |
| address=('localhost', 0), authkey=AUTH_KEY) |
| multiprocessing.current_process().authkey = AUTH_KEY |
| # Initialize eagerly to avoid acting as the server if there are issues. |
| # Note, however, that _create_server itself is called lazily. |
| _process_level_singleton_manager.register_singleton( |
| self._constructor, self._tag, initialize_eagerly=True) |
| self._server = self._serving_manager.get_server() |
| logging.info( |
| 'Starting proxy server at %s for shared %s', |
| self._server.address, |
| self._tag) |
| with open(address_file + '.tmp', 'w') as fout: |
| fout.write('%s:%d' % self._server.address) |
| os.rename(address_file + '.tmp', address_file) |
| t = threading.Thread(target=self._server.serve_forever, daemon=True) |
| t.start() |
| logging.info('Done starting server') |
