"""OpenMPI, once initialized, prohibits forking. This helper module

allows the forking of *one* helper child process before OpenMPI

initialization that can do the forking for the fork-challenged

parent process.

Since none of this is MPI-specific, it got parked in :mod:`pytools`.

.. autoexception:: ExecError

:show-inheritance:

.. autoclass:: Forker

.. autoclass:: DirectForker

.. autoclass:: IndirectForker

.. autofunction:: enable_prefork

.. autofunction:: call

.. autofunction:: call_async

.. autofunction:: call_capture_output

.. autofunction:: wait

.. autofunction:: waitall

"""

from __future__ import annotations

import socket

from abc import ABC, abstractmethod

from subprocess import Popen

from typing import TYPE_CHECKING, Any

from typing_extensions import override

if TYPE_CHECKING:

from collections.abc import Sequence

class ExecError(OSError):

pass

class Forker(ABC):

@abstractmethod

def call(self, cmdline: Sequence[str], cwd: str | None = None) -> int:

pass

@abstractmethod

def call_async(self, cmdline: Sequence[str], cwd: str | None = None) -> int:

pass

@abstractmethod

def call_capture_output(self,

cmdline: Sequence[str],

cwd: str | None = None,

error_on_nonzero: bool = True) -> tuple[int, bytes, bytes]:

pass

@abstractmethod

def wait(self, aid: int) -> int:

pass

@abstractmethod

def waitall(self) -> dict[int, int]:

pass

class DirectForker(Forker):

def __init__(self) -> None:

self.apids: dict[int, Popen[bytes]] = {}

self.count: int = 0

@override

def call(self, cmdline: Sequence[str], cwd: str | None = None) -> int:

from subprocess import call as spcall

try:

return spcall(cmdline, cwd=cwd)

except OSError as e:

raise ExecError(

"error invoking '{}': {}".format(" ".join(cmdline), e)) from e

@override

def call_async(self, cmdline: Sequence[str], cwd: str | None = None) -> int:

try:

self.count += 1

proc = Popen(cmdline, cwd=cwd)

self.apids[self.count] = proc

return self.count

except OSError as e:

raise ExecError(

"error invoking '{}': {}".format(" ".join(cmdline), e)) from e

@override

def call_capture_output(self,

cmdline: Sequence[str],

cwd: str | None = None,

error_on_nonzero: bool = True) -> tuple[int, bytes, bytes]:

from subprocess import PIPE, Popen

try:

popen = Popen(cmdline, cwd=cwd, stdin=PIPE, stdout=PIPE,

stderr=PIPE)

stdout_data, stderr_data = popen.communicate()

if error_on_nonzero and popen.returncode:

raise ExecError("status {} invoking '{}': {}".format(

popen.returncode,

" ".join(cmdline),

stderr_data.decode("utf-8", errors="replace")))

return popen.returncode, stdout_data, stderr_data

except OSError as e:

raise ExecError(

"error invoking '{}': {}".format(" ".join(cmdline), e)) from e

@override

def wait(self, aid: int) -> int:

proc = self.apids.pop(aid)

return proc.wait()

@override

def waitall(self) -> dict[int, int]:

rets = {}

for aid in self.apids:

rets[aid] = self.wait(aid)

return rets

def _send_packet(sock: socket.socket, data: object) -> None:

from pickle import dumps

from struct import pack

packet = dumps(data)

sock.sendall(pack("I", len(packet)))

sock.sendall(packet)

def _recv_packet(sock: socket.socket,

who: str = "Process",

partner: str = "other end") -> tuple[object, ...]:

from struct import calcsize, unpack

size_bytes_size = calcsize("I")

size_bytes = sock.recv(size_bytes_size)

if len(size_bytes) < size_bytes_size:

raise SystemExit

size, = unpack("I", size_bytes)

packet = b""

while len(packet) < size:

packet += sock.recv(size)

from pickle import loads

result = loads(packet)

assert isinstance(result, tuple)

return result

def _fork_server(sock: socket.socket) -> None:

# Ignore keyboard interrupts, we'll get notified by the parent.

import signal

signal.signal(signal.SIGINT, signal.SIG_IGN)

# Construct a local DirectForker to do the dirty work

df = DirectForker()

funcs = {

"call": df.call,

"call_async": df.call_async,

"call_capture_output": df.call_capture_output,

"wait": df.wait,

"waitall": df.waitall

}

try:

while True:

func_name, args, kwargs = _recv_packet(

sock, who="Prefork server", partner="parent"

)

assert isinstance(func_name, str)

if func_name == "quit":

df.waitall()

_send_packet(sock, ("ok", None))

break

try:

result = funcs[func_name](*args, **kwargs)

# FIXME: Is catching all exceptions the right course of action?

except Exception as e: # pylint:disable=broad-except

_send_packet(sock, ("exception", e))

else:

_send_packet(sock, ("ok", result))

finally:

sock.close()

import os

os._exit(0)

class IndirectForker(Forker):

def __init__(self, server_pid: int, sock: socket.socket) -> None:

self.server_pid = server_pid

self.socket = sock

import atexit

atexit.register(self._quit)

def _remote_invoke(self, name: str, *args: Any, **kwargs: Any) -> object:

_send_packet(self.socket, (name, args, kwargs))

status, result = _recv_packet(

self.socket, who="Prefork client", partner="prefork server"

)

if status == "exception":

assert isinstance(result, Exception)

raise result

assert status == "ok"

return result

def _quit(self) -> None:

self._remote_invoke("quit")

from os import waitpid

waitpid(self.server_pid, 0)

@override

def call(self, cmdline: Sequence[str], cwd: str | None = None) -> int:

result = self._remote_invoke("call", cmdline, cwd)

assert isinstance(result, int)

return result

@override

def call_async(self, cmdline: Sequence[str], cwd: str | None = None) -> int:

result = self._remote_invoke("call_async", cmdline, cwd)

assert isinstance(result, int)

return result

@override

def call_capture_output(self,

cmdline: Sequence[str],

cwd: str | None = None,

error_on_nonzero: bool = True,

) -> tuple[int, bytes, bytes]:

return self._remote_invoke("call_capture_output", cmdline, cwd,

error_on_nonzero)

@override

def wait(self, aid: int) -> int:

result = self._remote_invoke("wait", aid)

assert isinstance(result, int)

return result

@override

def waitall(self) -> dict[int, int]:

result = self._remote_invoke("waitall")

assert isinstance(result, dict)

return result

forker: Forker = DirectForker()

def enable_prefork() -> None:

global forker

if isinstance(forker, IndirectForker):

return

s_parent, s_child = socket.socketpair()

from os import fork

fork_res = fork()

# Child

if fork_res == 0:

s_parent.close()

_fork_server(s_child)

# Parent

else:

s_child.close()

forker = IndirectForker(fork_res, s_parent)

def call(cmdline: Sequence[str], cwd: str | None = None) -> int:

return forker.call(cmdline, cwd)

def call_async(cmdline: Sequence[str], cwd: str | None = None) -> int:

return forker.call_async(cmdline, cwd)

def call_capture_output(cmdline: Sequence[str],

cwd: str | None = None,

error_on_nonzero: bool = True) -> tuple[int, bytes, bytes]:

return forker.call_capture_output(cmdline, cwd, error_on_nonzero)

def wait(aid: int) -> int:

return forker.wait(aid)

def waitall() -> dict[int, int]:

return forker.waitall()