catch cpuinfo ImportError (tensorflow#4138)

Taylor Robie · web-flow · commit 5c78b9d7620d · 2018-05-01T10:43:56.000-07:00
* catch cpuinfo ImportError

* add psutil import catch

* fix typo
diff --git a/official/utils/logs/logger.py b/official/utils/logs/logger.py
@@ -213,15 +213,18 @@ def _collect_cpu_info(run_info):
 
   cpu_info["num_cores"] = multiprocessing.cpu_count()
 
-  # Note: cpuinfo is not installed in the TensorFlow OSS tree.
-  # It is installable via pip.
-  import cpuinfo    # pylint: disable=g-import-not-at-top
+  try:
+    # Note: cpuinfo is not installed in the TensorFlow OSS tree.
+    # It is installable via pip.
+    import cpuinfo    # pylint: disable=g-import-not-at-top
 
-  info = cpuinfo.get_cpu_info()
-  cpu_info["cpu_info"] = info["brand"]
-  cpu_info["mhz_per_cpu"] = info["hz_advertised_raw"][0] / 1.0e6
+    info = cpuinfo.get_cpu_info()
+    cpu_info["cpu_info"] = info["brand"]
+    cpu_info["mhz_per_cpu"] = info["hz_advertised_raw"][0] / 1.0e6
 
-  run_info["machine_config"]["cpu_info"] = cpu_info
+    run_info["machine_config"]["cpu_info"] = cpu_info
+  except ImportError:
+    tf.logging.warn("'cpuinfo' not imported. CPU info will not be logged.")
 
 
 def _collect_gpu_info(run_info):
@@ -243,12 +246,15 @@ def _collect_gpu_info(run_info):
 
 
 def _collect_memory_info(run_info):
-  # Note: psutil is not installed in the TensorFlow OSS tree.
-  # It is installable via pip.
-  import psutil   # pylint: disable=g-import-not-at-top
-  vmem = psutil.virtual_memory()
-  run_info["machine_config"]["memory_total"] = vmem.total
-  run_info["machine_config"]["memory_available"] = vmem.available
+  try:
+    # Note: psutil is not installed in the TensorFlow OSS tree.
+    # It is installable via pip.
+    import psutil   # pylint: disable=g-import-not-at-top
+    vmem = psutil.virtual_memory()
+    run_info["machine_config"]["memory_total"] = vmem.total
+    run_info["machine_config"]["memory_available"] = vmem.available
+  except ImportError:
+    tf.logging.warn("'psutil' not imported. Memory info will not be logged.")
 
 
 def _parse_gpu_model(physical_device_desc):
