FEDOT

1
import logging
2
import tracemalloc
3
from typing import Optional
4

5
from golem.core.log import default_log
6

7

8
class MemoryAnalytics:
9
    is_active = False
10
    active_session_label = 'main'
11

12
    @classmethod
13
    def start(cls):
14
        """
15
        Start memory monitoring session
16
        """
17
        cls.is_active = True
18
        tracemalloc.start()
19

20
    @classmethod
21
    def finish(cls):
22
        """
23
        Finish memory monitoring session
24
        """
25
        cls.log(additional_info='finish')
26
        tracemalloc.stop()
27
        cls.is_active = False
28

29
    @classmethod
30
    def get_measures(cls):
31
        """
32
        Estimates Python-related system memory consumption in MiB
33
        :return: current and maximal consumption
34
        """
35
        current_memory, max_memory = tracemalloc.get_traced_memory()
36
        return current_memory / 1024 / 1024, max_memory / 1024 / 1024
37

38
    @classmethod
39
    def log(cls, logger: Optional[logging.LoggerAdapter] = None,
40
            additional_info: str = 'location', logging_level: int = logging.INFO) -> str:
41
        """
42
        Print the message about current and maximal memory consumption to the log or console.
43
        :param logger: optional logger that should be used in output.
44
        :param additional_info: label for current location in code.
45
        :param logging_level: level of the message
46
        :return: text of the message.
47
        """
48
        message = ''
49
        if cls.is_active:
50
            memory_consumption = cls.get_measures()
51
            message = f'Memory consumption for {additional_info} in {cls.active_session_label} session: ' \
52
                      f'current {round(memory_consumption[0], 1)} MiB, ' \
53
                      f'max: {round(memory_consumption[1], 1)} MiB'
54
            if logger is None:
55
                logger = default_log(prefix=cls.__name__)
56
            logger.log(logging_level, message)
57
        return message
58