glasswall.multiprocessing.tasks

 1from multiprocessing import Queue
 2from typing import Any, Callable, Optional, Union
 3
 4import glasswall
 5
 6
 7class Task:
 8    def __init__(
 9        self,
10        func: Callable,
11        args: Optional[tuple] = None,
12        kwargs: Optional[dict] = None,
13    ):
14        self.func = func
15        self.args = args or tuple()
16        self.kwargs = kwargs or dict()
17
18        # Convert Policy objects to text (has attributes that are modules, and modules cannot be pickled)
19        # args
20        processed_args = []
21        for arg in self.args:
22            if isinstance(arg, glasswall.content_management.policies.Policy):
23                arg = arg.text
24            processed_args.append(arg)
25        self.args = tuple(processed_args)
26        # kwargs
27        for key, value in self.kwargs.items():
28            if isinstance(value, glasswall.content_management.policies.Policy):
29                self.kwargs[key] = value.text
30
31    def __eq__(self, other):
32        if isinstance(other, Task):
33            return (self.func, self.args, self.kwargs) == (other.func, other.args, other.kwargs)
34        return False
35
36    def __hash__(self):
37        kwargs_tuple = tuple(sorted(self.kwargs.items()))
38        return hash((self.func, self.args, kwargs_tuple))
39
40    def __repr__(self):
41        max_length = 100
42        args_str = ", ".join(repr(arg)[:max_length] + ('...' if len(repr(arg)) > max_length else '') for arg in self.args)
43        kwargs_str = ", ".join(f"{key}={repr(value)[:max_length] + ('...' if len(repr(value)) > max_length else '')}" for key, value in self.kwargs.items())
44        return f"{self.__class__.__name__}(func={self.func.__name__}, args=({args_str}), kwargs=({kwargs_str}))"
45
46
47class TaskResult:
48    timeout_seconds: Optional[float]
49    memory_limit_in_gib: Optional[float]
50    start_time: float
51    end_time: float
52    elapsed_time: float
53    out_of_memory: bool
54    timed_out: bool
55    max_memory_used_in_gib: float
56    exit_code: Union[int, None]
57
58    def __init__(
59        self,
60        task: Task,
61        success: bool,
62        result: Any = None,
63        exception: Union[Exception, None] = None,
64    ):
65        self.task = task
66        self.success = success
67        self.result = result
68        self.exception = exception
69
70    def __eq__(self, other):
71        if isinstance(other, TaskResult):
72            return (self.task, self.success, self.result, self.exception) == \
73                (other.task, other.success, other.result, other.exception)
74        return False
75
76    def __hash__(self):
77        return hash((self.task, self.success, self.result, self.exception))
78
79    def __repr__(self):
80        max_length = 100
81        attributes_str = ', '.join(f"{k}={v!r}"[:max_length] + ('...' if len(repr(v)) > max_length else '') for k, v in self.__dict__.items())
82        return f"{self.__class__.__name__}({attributes_str})"
83
84
85def execute_task_and_put_in_queue(task: Task, queue: "Queue[TaskResult]") -> None:
86    try:
87        func_result = task.func(*task.args, **task.kwargs)
88        task_result = TaskResult(task=task, success=True, result=func_result)
89    except Exception as e:
90        task_result = TaskResult(task=task, success=False, exception=e)
91
92    queue.put(task_result)