From c7eaa61d71bd39c2c6b5a7aea7788e3854e2fbcc Mon Sep 17 00:00:00 2001
From: amd-pworfolk <patrick.worfolk@amd.com>
Date: Mon, 9 Dec 2024 22:13:24 -0800
Subject: [PATCH] Added system-info as a management tool

Signed-off-by: amd-pworfolk <patrick.worfolk@amd.com>
---
 src/turnkeyml/common/system_info.py     | 13 ---------
 src/turnkeyml/llm/cli.py                |  3 ++-
 src/turnkeyml/sequence/tool_plugins.py  |  1 +
 src/turnkeyml/tools/management_tools.py | 35 +++++++++++++++++++++++++
 4 files changed, 38 insertions(+), 14 deletions(-)

diff --git a/src/turnkeyml/common/system_info.py b/src/turnkeyml/common/system_info.py
index 1561d478..9f408296 100644
--- a/src/turnkeyml/common/system_info.py
+++ b/src/turnkeyml/common/system_info.py
@@ -384,16 +384,3 @@ def get_system_info_dict() -> dict:
         dict: Dictionary containing the system information.
     """
     return get_system_info().get_dict()
-
-
-def main():
-    """
-    Prints the system information dictionary.
-    """
-    info_dict = get_system_info_dict()
-    for k, v in info_dict.items():
-        print(f"{k}: {v}")
-
-
-if __name__ == "__main__":
-    main()
diff --git a/src/turnkeyml/llm/cli.py b/src/turnkeyml/llm/cli.py
index e396244a..31f3416c 100644
--- a/src/turnkeyml/llm/cli.py
+++ b/src/turnkeyml/llm/cli.py
@@ -3,7 +3,7 @@
 import turnkeyml.common.filesystem as fs
 import turnkeyml.cli.cli as cli
 from turnkeyml.sequence import Sequence
-from turnkeyml.tools.management_tools import Cache, Version
+from turnkeyml.tools.management_tools import Cache, Version, SystemInfo
 from turnkeyml.tools.report import Report
 from turnkeyml.state import State
 
@@ -35,6 +35,7 @@ def main():
         Report,
         Cache,
         Version,
+        SystemInfo,
     ]
 
     # Import onnxruntime-genai recipes
diff --git a/src/turnkeyml/sequence/tool_plugins.py b/src/turnkeyml/sequence/tool_plugins.py
index 0553e60f..cb0f9388 100644
--- a/src/turnkeyml/sequence/tool_plugins.py
+++ b/src/turnkeyml/sequence/tool_plugins.py
@@ -17,6 +17,7 @@ def get_supported_tools():
         mgmt.Version,
         mgmt.Cache,
         mgmt.ModelsLocation,
+        mgmt.SystemInfo,
         report.Report,
         Discover,
         export.ExportPytorchModel,
diff --git a/src/turnkeyml/tools/management_tools.py b/src/turnkeyml/tools/management_tools.py
index 5a0d2905..91830ae0 100644
--- a/src/turnkeyml/tools/management_tools.py
+++ b/src/turnkeyml/tools/management_tools.py
@@ -7,6 +7,7 @@
 import turnkeyml.common.printing as printing
 from turnkeyml.tools.tool import ToolParser
 from turnkeyml.version import __version__ as turnkey_version
+from turnkeyml.common.system_info import get_system_info_dict
 
 
 class ManagementTool(abc.ABC):
@@ -265,3 +266,37 @@ def run(self, _, quiet: bool = False):
             print(fs.MODELS_DIR)
         else:
             printing.log_info(f"The models directory is: {fs.MODELS_DIR}")
+
+
+class SystemInfo(ManagementTool):
+    """
+    Prints system information for the turnkeyml installation.
+    """
+
+    unique_name = "system-info"
+
+    @staticmethod
+    def parser(add_help: bool = True) -> argparse.ArgumentParser:
+        parser = __class__.helpful_parser(
+            short_description="Print system information",
+            add_help=add_help,
+        )
+
+        return parser
+
+    @staticmethod
+    def pretty_print(my_dict: dict, level=0):
+        for k, v in my_dict.items():
+            if isinstance(v, dict):
+                print("    " * level + f"{k}:")
+                SystemInfo.pretty_print(v, level + 1)
+            elif isinstance(v,list):
+                print("    " * level + f"{k}:")
+                for item in v:
+                    print("    " * (level+1) + f"{item}")
+            else:
+                print("    " * level + f"{k}: {v}")
+
+    def run(self, _):
+        system_info_dict = get_system_info_dict()
+        self.pretty_print(system_info_dict)