xref: /aosp_15_r20/external/tensorflow/tensorflow/tools/test/system_info_lib.py (revision b6fb3261f9314811a0f4371741dbb8839866f948)
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