1# Copyright 2016 The TensorFlow Authors. All Rights Reserved. 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14# ============================================================================== 15"""Library for getting system information during TensorFlow tests.""" 16 17import glob 18import multiprocessing 19import platform 20import re 21import socket 22 23# pylint: disable=g-bad-import-order 24# Note: cpuinfo and psutil are not installed for you in the TensorFlow 25# OSS tree. They are installable via pip. 26import cpuinfo 27import psutil 28# pylint: enable=g-bad-import-order 29 30from tensorflow.core.util import test_log_pb2 31from tensorflow.python.client import device_lib 32from tensorflow.python.framework import errors 33from tensorflow.python.platform import gfile 34from tensorflow.tools.test import gpu_info_lib 35 36 37def gather_machine_configuration(): 38 """Gather Machine Configuration. This is the top level fn of this library.""" 39 config = test_log_pb2.MachineConfiguration() 40 41 config.cpu_info.CopyFrom(gather_cpu_info()) 42 config.platform_info.CopyFrom(gather_platform_info()) 43 44 # gather_available_device_info must come before gather_gpu_devices 45 # because the latter may access libcudart directly, which confuses 46 # TensorFlow StreamExecutor. 47 for d in gather_available_device_info(): 48 config.available_device_info.add().CopyFrom(d) 49 for gpu in gpu_info_lib.gather_gpu_devices(): 50 config.device_info.add().Pack(gpu) 51 52 config.memory_info.CopyFrom(gather_memory_info()) 53 54 config.hostname = gather_hostname() 55 56 return config 57 58 59def gather_hostname(): 60 return socket.gethostname() 61 62 63def gather_memory_info(): 64 """Gather memory info.""" 65 mem_info = test_log_pb2.MemoryInfo() 66 vmem = psutil.virtual_memory() 67 mem_info.total = vmem.total 68 mem_info.available = vmem.available 69 return mem_info 70 71 72def gather_cpu_info(): 73 """Gather CPU Information. Assumes all CPUs are the same.""" 74 cpu_info = test_log_pb2.CPUInfo() 75 cpu_info.num_cores = multiprocessing.cpu_count() 76 77 # Gather num_cores_allowed 78 try: 79 with gfile.GFile('/proc/self/status', 'rb') as fh: 80 nc = re.search(r'(?m)^Cpus_allowed:\s*(.*)$', fh.read().decode('utf-8')) 81 if nc: # e.g. 'ff' => 8, 'fff' => 12 82 cpu_info.num_cores_allowed = ( 83 bin(int(nc.group(1).replace(',', ''), 16)).count('1')) 84 except errors.OpError: 85 pass 86 finally: 87 if cpu_info.num_cores_allowed == 0: 88 cpu_info.num_cores_allowed = cpu_info.num_cores 89 90 # Gather the rest 91 info = cpuinfo.get_cpu_info() 92 cpu_info.cpu_info = info['brand'] 93 cpu_info.num_cores = info['count'] 94 cpu_info.mhz_per_cpu = info['hz_advertised_raw'][0] / 1.0e6 95 l2_cache_size = re.match(r'(\d+)', str(info.get('l2_cache_size', ''))) 96 if l2_cache_size: 97 # If a value is returned, it's in KB 98 cpu_info.cache_size['L2'] = int(l2_cache_size.group(0)) * 1024 99 100 # Try to get the CPU governor 101 try: 102 cpu_governors = set([ 103 gfile.GFile(f, 'r').readline().rstrip() 104 for f in glob.glob( 105 '/sys/devices/system/cpu/cpu*/cpufreq/scaling_governor') 106 ]) 107 if cpu_governors: 108 if len(cpu_governors) > 1: 109 cpu_info.cpu_governor = 'mixed' 110 else: 111 cpu_info.cpu_governor = list(cpu_governors)[0] 112 except errors.OpError: 113 pass 114 115 return cpu_info 116 117 118def gather_available_device_info(): 119 """Gather list of devices available to TensorFlow. 120 121 Returns: 122 A list of test_log_pb2.AvailableDeviceInfo messages. 123 """ 124 device_info_list = [] 125 devices = device_lib.list_local_devices() 126 127 for d in devices: 128 device_info = test_log_pb2.AvailableDeviceInfo() 129 device_info.name = d.name 130 device_info.type = d.device_type 131 device_info.memory_limit = d.memory_limit 132 device_info.physical_description = d.physical_device_desc 133 device_info_list.append(device_info) 134 135 return device_info_list 136 137 138def gather_platform_info(): 139 """Gather platform info.""" 140 platform_info = test_log_pb2.PlatformInfo() 141 (platform_info.bits, platform_info.linkage) = platform.architecture() 142 platform_info.machine = platform.machine() 143 platform_info.release = platform.release() 144 platform_info.system = platform.system() 145 platform_info.version = platform.version() 146 return platform_info 147