2024-02-02 05:04:30 +00:00
|
|
|
import logging
|
2023-12-17 21:18:52 +00:00
|
|
|
import os
|
|
|
|
import signal
|
|
|
|
import sys
|
|
|
|
import traceback
|
2023-12-19 16:51:43 +00:00
|
|
|
from pathlib import Path
|
2023-12-17 21:18:52 +00:00
|
|
|
from typing import Any
|
|
|
|
|
|
|
|
import gi
|
|
|
|
|
|
|
|
gi.require_version("GdkPixbuf", "2.0")
|
|
|
|
|
2023-12-26 17:02:43 +00:00
|
|
|
import dataclasses
|
2023-12-17 21:18:52 +00:00
|
|
|
import multiprocessing as mp
|
|
|
|
from collections.abc import Callable
|
|
|
|
|
2024-02-02 05:04:30 +00:00
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
|
2023-12-19 16:51:43 +00:00
|
|
|
|
2023-12-26 17:02:43 +00:00
|
|
|
# Kill the new process and all its children by sending a SIGTERM signal to the process group
|
|
|
|
def _kill_group(proc: mp.Process) -> None:
|
|
|
|
pid = proc.pid
|
2024-01-29 08:11:57 +00:00
|
|
|
if proc.is_alive() and pid:
|
2023-12-26 17:02:43 +00:00
|
|
|
os.killpg(pid, signal.SIGTERM)
|
|
|
|
else:
|
2024-02-16 09:10:49 +00:00
|
|
|
log.warning(f"Process '{proc.name}' with pid '{pid}' is already dead")
|
2023-12-17 21:18:52 +00:00
|
|
|
|
2023-12-26 17:02:43 +00:00
|
|
|
|
|
|
|
@dataclasses.dataclass(frozen=True)
|
2023-12-17 21:18:52 +00:00
|
|
|
class MPProcess:
|
2023-12-26 17:02:43 +00:00
|
|
|
name: str
|
|
|
|
proc: mp.Process
|
|
|
|
out_file: Path
|
2023-12-17 21:18:52 +00:00
|
|
|
|
2023-12-19 15:44:26 +00:00
|
|
|
# Kill the new process and all its children by sending a SIGTERM signal to the process group
|
|
|
|
def kill_group(self) -> None:
|
2023-12-26 17:02:43 +00:00
|
|
|
_kill_group(proc=self.proc)
|
2023-12-17 21:18:52 +00:00
|
|
|
|
|
|
|
|
2023-12-19 16:51:43 +00:00
|
|
|
def _set_proc_name(name: str) -> None:
|
2024-01-02 14:35:11 +00:00
|
|
|
if sys.platform != "linux":
|
|
|
|
return
|
2023-12-19 16:51:43 +00:00
|
|
|
import ctypes
|
|
|
|
|
|
|
|
# Define the prctl function with the appropriate arguments and return type
|
|
|
|
libc = ctypes.CDLL("libc.so.6")
|
|
|
|
prctl = libc.prctl
|
|
|
|
prctl.argtypes = [
|
|
|
|
ctypes.c_int,
|
|
|
|
ctypes.c_char_p,
|
|
|
|
ctypes.c_ulong,
|
|
|
|
ctypes.c_ulong,
|
|
|
|
ctypes.c_ulong,
|
|
|
|
]
|
|
|
|
prctl.restype = ctypes.c_int
|
|
|
|
|
|
|
|
# Set the process name to "my_process"
|
|
|
|
prctl(15, name.encode(), 0, 0, 0)
|
|
|
|
|
|
|
|
|
|
|
|
def _init_proc(
|
|
|
|
func: Callable,
|
|
|
|
out_file: Path,
|
|
|
|
proc_name: str,
|
2024-01-19 17:52:51 +00:00
|
|
|
on_except: Callable[[Exception, mp.process.BaseProcess], None] | None,
|
2023-12-19 16:51:43 +00:00
|
|
|
**kwargs: Any,
|
2023-12-17 21:18:52 +00:00
|
|
|
) -> None:
|
2023-12-19 16:51:43 +00:00
|
|
|
# Create a new process group
|
2023-12-17 21:18:52 +00:00
|
|
|
os.setsid()
|
|
|
|
|
2023-12-19 16:51:43 +00:00
|
|
|
# Open stdout and stderr
|
2023-12-26 17:02:43 +00:00
|
|
|
with open(out_file, "w") as out_fd:
|
|
|
|
os.dup2(out_fd.fileno(), sys.stdout.fileno())
|
|
|
|
os.dup2(out_fd.fileno(), sys.stderr.fileno())
|
2023-12-17 21:18:52 +00:00
|
|
|
|
2023-12-19 16:51:43 +00:00
|
|
|
# Print some information
|
2023-12-19 15:44:26 +00:00
|
|
|
pid = os.getpid()
|
|
|
|
gpid = os.getpgid(pid=pid)
|
|
|
|
|
2023-12-19 16:51:43 +00:00
|
|
|
# Set the process name
|
|
|
|
_set_proc_name(proc_name)
|
|
|
|
|
2024-01-03 12:26:25 +00:00
|
|
|
# Close stdin
|
2024-01-02 14:34:38 +00:00
|
|
|
sys.stdin.close()
|
2023-12-17 21:18:52 +00:00
|
|
|
|
2024-02-25 18:04:09 +00:00
|
|
|
linebreak = "=" * 5
|
2023-12-19 16:51:43 +00:00
|
|
|
# Execute the main function
|
2024-02-26 20:53:19 +00:00
|
|
|
print(linebreak + f" {func.__name__}:{pid} " + linebreak, file=sys.stderr)
|
2023-12-17 21:18:52 +00:00
|
|
|
try:
|
|
|
|
func(**kwargs)
|
2024-01-02 04:54:19 +00:00
|
|
|
except Exception as ex:
|
2023-12-17 21:18:52 +00:00
|
|
|
traceback.print_exc()
|
2024-01-19 17:52:51 +00:00
|
|
|
if on_except is not None:
|
|
|
|
on_except(ex, mp.current_process())
|
2024-02-25 18:04:09 +00:00
|
|
|
|
|
|
|
# Kill the new process and all its children by sending a SIGTERM signal to the process group
|
|
|
|
pid = os.getpid()
|
|
|
|
gpid = os.getpgid(pid=pid)
|
|
|
|
print(f"Killing process group pid={pid} gpid={gpid}", file=sys.stderr)
|
|
|
|
os.killpg(gpid, signal.SIGTERM)
|
2024-02-25 18:59:45 +00:00
|
|
|
sys.exit(1)
|
2024-02-25 18:04:09 +00:00
|
|
|
# Don't use a finally block here, because we want the exitcode to be set to
|
|
|
|
# 0 if the function returns normally
|
2023-12-17 21:18:52 +00:00
|
|
|
|
|
|
|
|
2023-12-19 16:51:43 +00:00
|
|
|
def spawn(
|
2024-01-02 04:54:19 +00:00
|
|
|
*,
|
2024-02-16 09:10:49 +00:00
|
|
|
out_file: Path,
|
2024-01-19 17:52:51 +00:00
|
|
|
on_except: Callable[[Exception, mp.process.BaseProcess], None] | None,
|
2024-01-02 04:54:19 +00:00
|
|
|
func: Callable,
|
|
|
|
**kwargs: Any,
|
2023-12-19 16:51:43 +00:00
|
|
|
) -> MPProcess:
|
|
|
|
# Decouple the process from the parent
|
2023-12-17 21:18:52 +00:00
|
|
|
if mp.get_start_method(allow_none=True) is None:
|
2023-12-26 17:02:43 +00:00
|
|
|
mp.set_start_method(method="forkserver")
|
|
|
|
|
2023-12-19 16:51:43 +00:00
|
|
|
# Set names
|
|
|
|
proc_name = f"MPExec:{func.__name__}"
|
2023-12-17 21:18:52 +00:00
|
|
|
|
2023-12-19 16:51:43 +00:00
|
|
|
# Start the process
|
2023-12-17 21:18:52 +00:00
|
|
|
proc = mp.Process(
|
2023-12-19 16:51:43 +00:00
|
|
|
target=_init_proc,
|
2024-03-06 10:12:20 +00:00
|
|
|
args=(func, out_file, proc_name, on_except),
|
2023-12-17 21:18:52 +00:00
|
|
|
name=proc_name,
|
|
|
|
kwargs=kwargs,
|
|
|
|
)
|
|
|
|
proc.start()
|
2023-12-19 16:51:43 +00:00
|
|
|
|
|
|
|
# Return the process
|
2024-01-02 14:34:38 +00:00
|
|
|
mp_proc = MPProcess(name=proc_name, proc=proc, out_file=out_file)
|
2023-12-26 17:02:43 +00:00
|
|
|
|
2023-12-17 21:18:52 +00:00
|
|
|
return mp_proc
|