feat: print system info

Signed-off-by: James <[email protected]>

engine/commands/sys_info.cc

@@ -0,0 +1,22 @@
+#include "sys_info.h"
+#include "trantor/utils/Logger.h"
+#include "utils/system_info_utils.h"
+
+namespace commands {
+
+void SysInfoCmd::Exec() {
+  LOG_INFO << "SysInfoCmd::Exec()";
+  auto gpuInfoList = system_info_utils::GetGpuInfoList();
+  // print gpu info
+  for (auto& gpuInfo : gpuInfoList) {
+    LOG_INFO << "GPU: " << gpuInfo.name << " (" << gpuInfo.id << ")";
+    if (gpuInfo.arch.has_value()) {
+      LOG_INFO << "  Arch: " << gpuInfo.arch.value();
+    }
+    LOG_INFO << "  Memory: " << gpuInfo.vram;
+    if (gpuInfo.compute_cap.has_value()) {
+      LOG_INFO << "  Compute capability: " << gpuInfo.compute_cap.value();
+    }
+  }
+}
+};  // namespace commands

engine/commands/sys_info.h

@@ -0,0 +1,9 @@
+#pragma once
+
+namespace commands {
+class SysInfoCmd {
+ public:
+  SysInfoCmd() {};
+  void Exec();
+};
+}  // namespace commands

engine/controllers/command_line_parser.cc

@@ -10,6 +10,7 @@
 #include "commands/run_cmd.h"
 #include "commands/stop_model_cmd.h"
 #include "commands/stop_server_cmd.h"
+#include "commands/sys_info.h"
 #include "config/yaml_config.h"
 #include "utils/cortex_utils.h"
 
@@ -130,6 +131,8 @@ bool CommandLineParser::SetupCommand(int argc, char** argv) {
     EngineInstall(engines_cmd, "cortex.tensorrt-llm", version);
   }
 
+  SysInfo(&app_);
+
   {
     // cortex run tinyllama:gguf
     auto run_cmd =
@@ -169,4 +172,12 @@ void CommandLineParser::EngineInstall(CLI::App* parent,
     commands::EngineInitCmd eic(engine_name, version);
     eic.Exec();
   });
-}
+}
+
+void CommandLineParser::SysInfo(CLI::App* parent) {
+  auto sys_info_cmd = parent->add_subcommand("sysinfo", "Get system info");
+  sys_info_cmd->callback([] {
+    commands::SysInfoCmd cmd;
+    cmd.Exec();
+  });
+}

engine/controllers/command_line_parser.h

@@ -1,6 +1,5 @@
 #pragma once
 
-#include <vector>
 #include "CLI/CLI.hpp"
 
 class CommandLineParser {
@@ -12,5 +11,7 @@ class CommandLineParser {
   void EngineInstall(CLI::App* parent, const std::string& engine_name,
                      std::string& version);
 
+  void SysInfo(CLI::App* parent);
+
   CLI::App app_;
-};
+};

engine/utils/system_info_utils.h

@@ -3,7 +3,9 @@
 #include <trantor/utils/Logger.h>
 #include <regex>
 #include <vector>
+#include "sstream"
 #include "utils/command_executor.h"
+#include "optional"
 #ifdef _WIN32
 #include <windows.h>
 #endif
@@ -24,26 +26,32 @@ struct SystemInfo {
   std::string arch;
 };
 
+const std::unordered_map<std::string, std::string> kGpuArchMap{
+    {"10.0", "blackwell"}, {"9.0", "hopper"}, {"8.9", "ada"},
+    {"8.6", "ampere"},     {"8.7", "ampere"}, {"8.0", "ampere"},
+    {"7.5", "turing"},     {"7.2", "volta"},  {"7.0", "volta"},
+    {"6.2", "pascal"},     {"6.1", "pascal"}, {"6.0", "pascal"}};
+
 /**
  * @brief Get the Gpu Arch. Currently we only support Ampere and Ada.
  * Might need to come up with better way to detect the GPU architecture.
  * 
  * @param gpuName E.g. NVIDIA GeForce RTX 4090
  * @return corresponding GPU arch. E.g. ampere, ada.
  */
-inline std::string GetGpuArch(const std::string& gpuName) {
-  std::string lowerGpuName = gpuName;
+inline std::string GetGpuArch(const std::string& gpu_name,
+                              const std::string& compute_cap) {
+  std::string lowerGpuName = gpu_name;
   std::transform(lowerGpuName.begin(), lowerGpuName.end(), lowerGpuName.begin(),
                  ::tolower);
 
   if (lowerGpuName.find("nvidia") == std::string::npos) {
     return "unknown";
   }
 
-  if (gpuName.find("30") != std::string::npos) {
-    return "ampere";
-  } else if (gpuName.find("40") != std::string::npos) {
-    return "ada";
+  auto it = kGpuArchMap.find(compute_cap);
+  if (it != kGpuArchMap.end()) {
+    return it->second;
   } else {
     return "unknown";
   }
@@ -232,7 +240,7 @@ struct GpuInfo {
   std::string id;
   std::string vram;
   std::string name;
-  std::string arch;
+  std::optional<std::string> arch;
   // nvidia driver version. Haven't checked for AMD GPU.
   std::optional<std::string> driver_version;
   std::optional<std::string> cuda_driver_version;
@@ -277,13 +285,10 @@ inline std::vector<GpuInfo> GetGpuInfoListVulkan() {
         std::string key = (*field_iter)[1].str();
         std::string value = (*field_iter)[2].str();
 
-        if (key == "deviceName")
+        if (key == "deviceName") {
           gpuInfo.name = value;
-        else if (key == "apiVersion")
-          gpuInfo.compute_cap = value;
-
+        }
         gpuInfo.vram = "";  // not available
-        gpuInfo.arch = GetGpuArch(gpuInfo.name);
 
         ++field_iter;
       }
@@ -314,13 +319,13 @@ inline std::vector<GpuInfo> GetGpuInfoList() {
     while (
         std::regex_search(search_start, output.cend(), match, gpu_info_reg)) {
       GpuInfo gpuInfo = {
-          match[1].str(),              // id
-          match[2].str(),              // vram
-          match[3].str(),              // name
-          GetGpuArch(match[3].str()),  // arch
-          driver_version,              // driver_version
-          cuda_version,                // cuda_driver_version
-          match[4].str()               // compute_cap
+          match[1].str(),                              // id
+          match[2].str(),                              // vram
+          match[3].str(),                              // name
+          GetGpuArch(match[3].str(), match[4].str()),  // arch
+          driver_version,                              // driver_version
+          cuda_version,                                // cuda_driver_version
+          match[4].str()                               // compute_cap
       };
       gpuInfoList.push_back(gpuInfo);
       search_start = match.suffix().first;