Completed
Push — master ( dbc38f...56accc )
by Klaus
01:34
created

_get_processor_name()   B

Complexity

Conditions 6

Size

Total Lines 14

Duplication

Lines 0
Ratio 0 %

Importance

Changes 1
Bugs 0 Features 1
Metric Value
cc 6
c 1
b 0
f 1
dl 0
loc 14
rs 8
1
#!/usr/bin/env python
2
# coding=utf-8
3
"""This module helps to collect information about the host of an experiment."""
4
from __future__ import division, print_function, unicode_literals
5
6
import os
7
import platform
8
import re
9
import subprocess
10
import xml.etree.ElementTree as ET
11
from sacred.utils import optional_kwargs_decorator, FileNotFoundError
12
from sacred.settings import SETTINGS
13
14
__sacred__ = True  # marks files that should be filtered from stack traces
15
16
__all__ = ('host_info_gatherers', 'get_host_info', 'host_info_getter')
17
18
host_info_gatherers = {}
19
"""Global dict of functions that are used to collect the host information."""
20
21
22
class IgnoreHostInfo(Exception):
23
    """Used by host_info_getters to signal that this cannot be gathered."""
24
25
26
def get_host_info():
27
    """Collect some information about the machine this experiment runs on.
28
29
    Returns
30
    -------
31
    dict
32
        A dictionary with information about the CPU, the OS and the
33
        Python version of this machine.
34
    """
35
    host_info = {}
36
    for k, v in host_info_gatherers.items():
37
        try:
38
            host_info[k] = v()
39
        except IgnoreHostInfo:
40
            pass
41
    return host_info
42
43
44
@optional_kwargs_decorator
45
def host_info_getter(func, name=None):
46
    """
47
    The decorated function is added to the process of collecting the host_info.
48
49
    This just adds the decorated function to the global
50
    ``sacred.host_info.host_info_gatherers`` dictionary.
51
    The functions from that dictionary are used when collecting the host info
52
    using :py:func:`~sacred.host_info.get_host_info`.
53
54
    Parameters
55
    ----------
56
    func : callable
57
        A function that can be called without arguments and returns some
58
        json-serializable information.
59
    name : str, optional
60
        The name of the corresponding entry in host_info.
61
        Defaults to the name of the function.
62
63
    Returns
64
    -------
65
    The function itself.
66
    """
67
    name = name or func.__name__
68
    host_info_gatherers[name] = func
69
    return func
70
71
72
# #################### Default Host Information ###############################
73
74
@host_info_getter(name='hostname')
75
def _hostname():
76
    return platform.node()
77
78
79
@host_info_getter(name='os')
80
def _os():
81
    return [platform.system(), platform.platform()]
82
83
84
@host_info_getter(name='python_version')
85
def _python_version():
86
    return platform.python_version()
87
88
89
@host_info_getter(name='cpu')
90
def _cpu():
91
    if platform.system() == "Windows":
92
        return platform.processor().strip()
93
    elif platform.system() == "Darwin":
94
        os.environ['PATH'] += ':/usr/sbin'
95
        command = ["sysctl", "-n", "machdep.cpu.brand_string"]
96
        return subprocess.check_output(command).decode().strip()
97
    elif platform.system() == "Linux":
98
        command = ["cat", "/proc/cpuinfo"]
99
        all_info = subprocess.check_output(command).decode()
100
        model_pattern = re.compile("^\s*model name\s*:")
101
        for line in all_info.split("\n"):
102
            if model_pattern.match(line):
103
                return model_pattern.sub("", line, 1).strip()
104
105
106
if SETTINGS.HOST_INFO.INCLUDE_GPU_INFO:
107
    @host_info_getter(name='gpus')
108
    def _gpus():
109
        try:
110
            xml = subprocess.check_output(['nvidia-smi', '-q', '-x']).decode()
111
        except (FileNotFoundError, OSError, subprocess.CalledProcessError):
112
            raise IgnoreHostInfo()
113
114
        gpu_info = {'gpus': []}
115
        for child in ET.fromstring(xml):
116
            if child.tag == 'driver_version':
117
                gpu_info['driver_version'] = child.text
118
            if child.tag != 'gpu':
119
                continue
120
            gpu = {
121
                'model': child.find('product_name').text,
122
                'total_memory': int(child.find('fb_memory_usage').find('total')
123
                                    .text.split()[0]),
124
                'persistence_mode': (child.find('persistence_mode').text ==
125
                                     'Enabled')
126
            }
127
            gpu_info['gpus'].append(gpu)
128
129
        return gpu_info
130