From cd22cb8257056d4918b5a4c6709a2cc6aff57d64 Mon Sep 17 00:00:00 2001
From: Yuxuan Xia <77518229+NovTi@users.noreply.github.com>
Date: Wed, 10 Apr 2024 15:06:00 +0800
Subject: [PATCH] Update Env check Script (#10709)

* Update env check bash file

* Update env-check
---
 python/llm/scripts/env-check.sh | 33 +++++++++++++++++++--------------
 1 file changed, 19 insertions(+), 14 deletions(-)

diff --git a/python/llm/scripts/env-check.sh b/python/llm/scripts/env-check.sh
index fef33837..75cd1614 100644
--- a/python/llm/scripts/env-check.sh
+++ b/python/llm/scripts/env-check.sh
@@ -65,18 +65,16 @@ check_cpu_info()
   lscpu | head -n 17
 }
 
-check_memory_type()
-{
-  echo "-----------------------------------------------------------------"
-  echo "CPU type: "
-  sudo dmidecode -t 17 | grep DDR
-
-}
-
 check_mem_info()
 {
   echo "-----------------------------------------------------------------"
-  cat /proc/meminfo | grep "MemTotal" | awk '{print "Total Memory: " $2/1024/1024 " GB"}'
+  cat /proc/meminfo | grep "MemTotal" | awk '{print "Total CPU Memory: " $2/1024/1024 " GB"}'
+
+  # Check if sudo session exists
+  if sudo -n true 2>/dev/null; then
+      echo -n "Memory Type: "
+      sudo dmidecode --type memory | grep -m 1 DDR | awk '{print $2, $3}'
+  fi
   
 }
 
@@ -123,9 +121,9 @@ check_xpu_smi()
 check_ipex()
 {
   echo "-----------------------------------------------------------------"
-  if python -c "import intel_extension_for_pytorch as ipex; print(ipex.__version__)" >/dev/null 2>&1
+  if python -c "import warnings; warnings.filterwarnings('ignore'); import intel_extension_for_pytorch as ipex; print(ipex.__version__)" >/dev/null 2>&1
   then
-    VERSION=`python -c "import intel_extension_for_pytorch as ipex; print(ipex.__version__)"`
+    VERSION=`python -c "import warnings; warnings.filterwarnings('ignore'); import intel_extension_for_pytorch as ipex; print(ipex.__version__)"`
     echo "ipex=$VERSION"
   else
     echo "IPEX is not installed. "
@@ -135,7 +133,7 @@ check_ipex()
 check_xpu_info()
 {
   echo "-----------------------------------------------------------------"
-  lspci -v | grep -i vga -A 8
+  lspci -v | grep -i vga -A 8 | awk '/Memory/ {gsub(/\[size=[0-9]+G\]/,"\033[1;33m&\033[0m")} 1'
 }
 
 check_linux_kernel_version()
@@ -167,6 +165,11 @@ check_igpu()
   fi
 }
 
+check_gpu_memory()
+{
+  lspci -v | grep -i vga -A 8 | awk '/VGA compatible controller/ {getline; getline; getline; getline; print "GPU" i++ " Memory", substr($0, length($0)-index($0," "), index($0," "))}'
+}
+
 main()
 {
   # first guarantee correct python is installed. 
@@ -186,11 +189,9 @@ main()
   # verify hardware (how many gpu availables, gpu status, cpu info, memory info, etc.)
   check_cpu_info
   check_mem_info
-  # check_memory_type
   # check_ulimit
   check_os
   # check_env
-  check_xpu_info
   check_linux_kernel_version
   check_xpu_driver
   check_OpenCL_driver
@@ -206,6 +207,10 @@ main()
     check_xpu_smi
   fi
 
+  check_gpu_memory
+
+  check_xpu_info
+
   echo "-----------------------------------------------------------------"
 }