glasswall.multiprocessing.manager

 1import os
 2import time
 3from collections import deque
 4from multiprocessing import Process, Queue
 5from typing import List, Generator, Optional
 6
 7from glasswall.multiprocessing.task_watcher import TaskWatcher
 8from glasswall.multiprocessing.tasks import Task, TaskResult
 9
10
11class GlasswallProcessManager:
12    def __init__(
13        self,
14        max_workers: Optional[int] = None,
15        worker_timeout_seconds: Optional[float] = None,
16        memory_limit_in_gib: Optional[float] = None,
17    ):
18        self.max_workers = max_workers or os.cpu_count() or 1
19        self.worker_timeout_seconds = worker_timeout_seconds
20        self.memory_limit_in_gib = memory_limit_in_gib
21        self._sleep_time: float = 0  # Time to sleep for while waiting for processes to complete
22        self._task_watcher_sleep_time: float = 0.001  # Time the TaskWatcher sleeps for while waiting for completed processes and monitoring timeout/memory
23        self._task_watcher_memory_limit_polling_rate: float = 0.1  # Polling rate for TaskWatcher to check the memory usage of a process
24
25        self.pending_processes: deque[Process] = deque()
26        self.active_processes: list[Process] = []
27        self.task_results_queue: "Queue[TaskResult]" = Queue()
28        self.task_results: List[TaskResult] = []
29
30    def __enter__(self):
31        return self
32
33    def __exit__(self, exc_type, exc_val, exc_tb):
34        self.start_tasks()
35
36    def queue_task(self, task: Task):
37        # Create and queue the process without starting it
38        process = Process(
39            target=TaskWatcher,
40            kwargs=dict(
41                task=task,
42                task_results_queue=self.task_results_queue,
43                timeout_seconds=self.worker_timeout_seconds,
44                memory_limit_in_gib=self.memory_limit_in_gib,
45                sleep_time=self._task_watcher_sleep_time,
46                memory_limit_polling_rate=self._task_watcher_memory_limit_polling_rate,
47            ),
48        )
49        self.pending_processes.append(process)
50
51    def as_completed(self) -> Generator[TaskResult, None, None]:
52        while self.pending_processes or self.active_processes:
53            if self.active_processes:
54                self.wait_for_completed_process()
55
56            while len(self.active_processes) < self.max_workers and self.pending_processes:
57                process = self.pending_processes.popleft()
58                self.active_processes.append(process)
59                process.start()
60
61            for result in self.task_results:
62                self.task_results.remove(result)
63                yield result
64
65    def start_tasks(self):
66        self.task_results = list(self.as_completed())
67
68    def wait_for_completed_process(self):
69        self.remove_completed_active_processes()
70        while len(self.active_processes) >= self.max_workers:
71            if self._sleep_time:
72                time.sleep(self._sleep_time)
73            self.remove_completed_active_processes()
74
75    def remove_completed_active_processes(self):
76        self.active_processes = [process for process in self.active_processes if process.is_alive()]
77        self.clean_task_results_queue()
78
79    def clean_task_results_queue(self):
80        while not self.task_results_queue.empty():
81            self.task_results.append(self.task_results_queue.get())