1 files changed, 149 insertions, 0 deletions
diff --git a/tensorflow/tools/test/system_info_lib.py b/tensorflow/tools/test/system_info_lib.py
new file mode 100644
index 0000000000..c36a6c6b13
--- /dev/null
+++ b/tensorflow/tools/test/system_info_lib.py
@@ -0,0 +1,149 @@
+# Copyright 2016 Google Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Library for getting system information during TensorFlow tests."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import multiprocessing
+import platform
+import re
+import socket
+
+import tensorflow as tf
+
+# pylint: disable=g-bad-import-order
+# Note: cpuinfo and psutil are not installed for you in the TensorFlow
+# OSS tree.  They are installable via pip.
+import cpuinfo
+import psutil
+# pylint: enable=g-bad-import-order
+
+from tensorflow.core.util import test_log_pb2
+from tensorflow.python.client import device_lib
+from tensorflow.tools.test import gpu_info_lib
+
+
+def gather_machine_configuration():
+  """Gather Machine Configuration.  This is the top level fn of this library."""
+  config = test_log_pb2.MachineConfiguration()
+
+  config.cpu_info.CopyFrom(gather_cpu_info())
+  config.platform_info.CopyFrom(gather_platform_info())
+
+  # gather_available_device_info must come before gather_gpu_devices
+  # because the latter may access libcudart directly, which confuses
+  # TensorFlow StreamExecutor.
+  for d in gather_available_device_info():
+    config.available_device_info.add().CopyFrom(d)
+  for gpu in gpu_info_lib.gather_gpu_devices():
+    config.device_info.add().Pack(gpu)
+
+  config.memory_info.CopyFrom(gather_memory_info())
+
+  config.hostname = gather_hostname()
+
+  return config
+
+
+def gather_hostname():
+  return socket.gethostname()
+
+
+def gather_memory_info():
+  """Gather memory info."""
+  mem_info = test_log_pb2.MemoryInfo()
+  vmem = psutil.virtual_memory()
+  mem_info.total = vmem.total
+  mem_info.available = vmem.available
+  return mem_info
+
+
+def gather_cpu_info():
+  """Gather CPU Information.  Assumes all CPUs are the same."""
+  cpu_info = test_log_pb2.CPUInfo()
+  cpu_info.num_cores = multiprocessing.cpu_count()
+
+  # Gather num_cores_allowed
+  try:
+    with tf.gfile.GFile('/proc/self/status') as fh:
+      nc = re.search(r'(?m)^Cpus_allowed:\s*(.*)$', fh.read())
+    if nc:  # e.g. 'ff' => 8, 'fff' => 12
+      cpu_info.num_cores_allowed = (
+          bin(int(nc.group(1).replace(',', ''), 16)).count('1'))
+  except IOError:
+    pass
+  finally:
+    if cpu_info.num_cores_allowed == 0:
+      cpu_info.num_cores_allowed = cpu_info.num_cores
+
+  # Gather the rest
+  info = cpuinfo.get_cpu_info()
+  cpu_info.cpu_info = info['brand']
+  cpu_info.num_cores = info['count']
+  cpu_info.mhz_per_cpu = info['hz_advertised_raw'][0] / 1.0e6
+  l2_cache_size = re.match(r'(\d+)', str(info['l2_cache_size']))
+  if l2_cache_size:
+    # If a value is returned, it's in KB
+    cpu_info.cache_size['L2'] = int(l2_cache_size.group(0)) * 1024
+
+  # Try to get the CPU governor
+  try:
+    cpu_governors = set([
+        tf.gfile.GFile(f, 'r').readline().rstrip()
+        for f in tf.gfile.Glob(
+            '/sys/devices/system/cpu/cpu*/cpufreq/scaling_governor')])
+    if cpu_governors:
+      if len(cpu_governors) > 1:
+        cpu_info.cpu_governor = 'mixed'
+      else:
+        cpu_info.cpu_governor = list(cpu_governors)[0]
+  except IOError:
+    pass
+
+  return cpu_info
+
+
+def gather_available_device_info():
+  """Gather list of devices available to TensorFlow.
+
+  Returns:
+    A list of test_log_pb2.AvailableDeviceInfo messages.
+  """
+  device_info_list = []
+  devices = device_lib.list_local_devices()
+
+  for d in devices:
+    device_info = test_log_pb2.AvailableDeviceInfo()
+    device_info.name = d.name
+    device_info.type = d.device_type
+    device_info.memory_limit = d.memory_limit
+    device_info.physical_description = d.physical_device_desc
+    device_info_list.append(device_info)
+
+  return device_info_list
+
+
+def gather_platform_info():
+  """Gather platform info."""
+  platform_info = test_log_pb2.PlatformInfo()
+  (platform_info.bits, platform_info.linkage) = platform.architecture()
+  platform_info.machine = platform.machine()
+  platform_info.release = platform.release()
+  platform_info.system = platform.system()
+  platform_info.version = platform.version()
+  return platform_info