Completed
Push — master ( 4edb63...28fb1c )
by Klaus
39s
created

_get_cpu_by_pycpuinfo()   A

Complexity

Conditions 1

Size

Total Lines 2

Duplication

Lines 0
Ratio 0 %

Importance

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