From 579043ffa2f04de99f997a39f3b650749e7d234b Mon Sep 17 00:00:00 2001
From: dyastremsky <58150256+dyastremsky@users.noreply.github.com>
Date: Fri, 1 Mar 2024 14:39:23 -0800
Subject: [PATCH] Add GenAi-PA error handling (#483)

* Print exception when non-zero error code

* Add test for GenAiPAException for nonzero exit

* Revert format/comment changes.

* Fix typo

Co-authored-by: Ryan McCormick <rmccormick@nvidia.com>

* Remove unnecessary return value

Co-authored-by: Ryan McCormick <rmccormick@nvidia.com>

---------

Co-authored-by: Ryan McCormick <rmccormick@nvidia.com>
---
 .../{genaipa_exceptions.py => genai_pa/exceptions.py}     | 0
 .../genai-pa/{ => genai_pa}/llm_inputs/__init__.py        | 0
 .../genai-pa/{ => genai_pa}/llm_inputs/llm_inputs.py      | 2 +-
 src/c++/perf_analyzer/genai-pa/genai_pa/main.py           | 8 ++++++--
 src/c++/perf_analyzer/genai-pa/genai_pa/wrapper.py        | 2 +-
 src/c++/perf_analyzer/genai-pa/tests/test_cli.py          | 6 ++++++
 src/c++/perf_analyzer/genai-pa/tests/test_llm_inputs.py   | 4 ++--
 7 files changed, 16 insertions(+), 6 deletions(-)
 rename src/c++/perf_analyzer/genai-pa/{genaipa_exceptions.py => genai_pa/exceptions.py} (100%)
 rename src/c++/perf_analyzer/genai-pa/{ => genai_pa}/llm_inputs/__init__.py (100%)
 rename src/c++/perf_analyzer/genai-pa/{ => genai_pa}/llm_inputs/llm_inputs.py (99%)

diff --git a/src/c++/perf_analyzer/genai-pa/genaipa_exceptions.py b/src/c++/perf_analyzer/genai-pa/genai_pa/exceptions.py
similarity index 100%
rename from src/c++/perf_analyzer/genai-pa/genaipa_exceptions.py
rename to src/c++/perf_analyzer/genai-pa/genai_pa/exceptions.py
diff --git a/src/c++/perf_analyzer/genai-pa/llm_inputs/__init__.py b/src/c++/perf_analyzer/genai-pa/genai_pa/llm_inputs/__init__.py
similarity index 100%
rename from src/c++/perf_analyzer/genai-pa/llm_inputs/__init__.py
rename to src/c++/perf_analyzer/genai-pa/genai_pa/llm_inputs/__init__.py
diff --git a/src/c++/perf_analyzer/genai-pa/llm_inputs/llm_inputs.py b/src/c++/perf_analyzer/genai-pa/genai_pa/llm_inputs/llm_inputs.py
similarity index 99%
rename from src/c++/perf_analyzer/genai-pa/llm_inputs/llm_inputs.py
rename to src/c++/perf_analyzer/genai-pa/genai_pa/llm_inputs/llm_inputs.py
index 924c0bdc4..816acb4a2 100644
--- a/src/c++/perf_analyzer/genai-pa/llm_inputs/llm_inputs.py
+++ b/src/c++/perf_analyzer/genai-pa/genai_pa/llm_inputs/llm_inputs.py
@@ -17,7 +17,7 @@
 from typing import Dict, List, Optional, Tuple
 
 import requests
-from genaipa_exceptions import GenAiPAException
+from genai_pa.exceptions import GenAiPAException
 from requests import Response
 
 
diff --git a/src/c++/perf_analyzer/genai-pa/genai_pa/main.py b/src/c++/perf_analyzer/genai-pa/genai_pa/main.py
index 79748df50..b0f5c1e5a 100755
--- a/src/c++/perf_analyzer/genai-pa/genai_pa/main.py
+++ b/src/c++/perf_analyzer/genai-pa/genai_pa/main.py
@@ -30,6 +30,7 @@
 
 from genai_pa import parser
 from genai_pa.constants import LOGGER_NAME
+from genai_pa.exceptions import GenAiPAException
 
 logging.basicConfig(level=logging.INFO, format="%(name)s - %(levelname)s - %(message)s")
 logger = logging.getLogger(LOGGER_NAME)
@@ -39,8 +40,11 @@
 # to assert correct errors and messages.
 # Optional argv used for testing - will default to sys.argv if None.
 def run(argv=None):
-    args = parser.parse_args(argv)
-    args.func(args)
+    try:
+        args = parser.parse_args(argv)
+        args.func(args)
+    except Exception as e:
+        raise GenAiPAException(e)
 
 
 def main():
diff --git a/src/c++/perf_analyzer/genai-pa/genai_pa/wrapper.py b/src/c++/perf_analyzer/genai-pa/genai_pa/wrapper.py
index 49ca35aef..25e776425 100644
--- a/src/c++/perf_analyzer/genai-pa/genai_pa/wrapper.py
+++ b/src/c++/perf_analyzer/genai-pa/genai_pa/wrapper.py
@@ -57,4 +57,4 @@ def run(model, args=None):
                         arg = utils.convert_option_name(arg)
                         cmd += f"--{arg} {value} "
         logger.info(f"Running Perf Analyzer : '{cmd}'")
-        subprocess.run(cmd, shell=True)
+        subprocess.run(cmd, shell=True, check=True)
diff --git a/src/c++/perf_analyzer/genai-pa/tests/test_cli.py b/src/c++/perf_analyzer/genai-pa/tests/test_cli.py
index bc103acd4..4aaff401b 100644
--- a/src/c++/perf_analyzer/genai-pa/tests/test_cli.py
+++ b/src/c++/perf_analyzer/genai-pa/tests/test_cli.py
@@ -28,6 +28,8 @@
 
 import pytest
 from genai_pa import parser
+from genai_pa.exceptions import GenAiPAException
+from genai_pa.main import run
 
 
 class TestCLIArguments:
@@ -91,3 +93,7 @@ def test_arguments_model_not_provided(self):
 
         # Check that the exit was unsuccessful
         assert exc_info.value.code != 0
+
+    def test_exception_on_nonzero_exit(self):
+        with pytest.raises(GenAiPAException) as e:
+            run(["-m", "nonexistent_model"])
diff --git a/src/c++/perf_analyzer/genai-pa/tests/test_llm_inputs.py b/src/c++/perf_analyzer/genai-pa/tests/test_llm_inputs.py
index 7abeb3f53..bc8855e7f 100644
--- a/src/c++/perf_analyzer/genai-pa/tests/test_llm_inputs.py
+++ b/src/c++/perf_analyzer/genai-pa/tests/test_llm_inputs.py
@@ -16,8 +16,8 @@
 import os
 
 import pytest
-from genaipa_exceptions import GenAiPAException
-from llm_inputs.llm_inputs import LlmInputs
+from genai_pa.exceptions import GenAiPAException
+from genai_pa.llm_inputs.llm_inputs import LlmInputs
 
 
 class TestLlmInputs: