From 48d5ac68d0a05e7fe257bcc19bb7637da150cd49 Mon Sep 17 00:00:00 2001
From: Matt Van Horn <455140+mvanhorn@users.noreply.github.com>
Date: Sat, 21 Mar 2026 08:37:31 -0700
Subject: [PATCH] fix: correct error handling in setup_env.py and
 e2e_benchmark.py

- Fix sys.exit(1) indentation in run_command() so it only runs on error
  (setup_env.py:107 and e2e_benchmark.py:23)
- Change exit(0) to sys.exit(1) for unsupported arch in compile()
- Use ARCH_ALIAS.get() to avoid KeyError on unrecognized architectures
- Guard SUPPORTED_QUANT_TYPES lookup against unknown arch in parse_args()

Fixes #447

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 setup_env.py           | 9 ++++++---
 utils/e2e_benchmark.py | 2 +-
 2 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/setup_env.py b/setup_env.py
index 3bf5fb8f7..1f07ea184 100644
--- a/setup_env.py
+++ b/setup_env.py
@@ -82,7 +82,7 @@
 }
 
 def system_info():
-    return platform.system(), ARCH_ALIAS[platform.machine()]
+    return platform.system(), ARCH_ALIAS.get(platform.machine(), platform.machine())
 
 def get_model_name():
     if args.hf_repo:
@@ -104,7 +104,7 @@ def run_command(command, shell=False, log_step=None):
             subprocess.run(command, shell=shell, check=True)
         except subprocess.CalledProcessError as e:
             logging.error(f"Error occurred while running command: {e}")
-        sys.exit(1)
+            sys.exit(1)
 
 def prepare_model():
     _, arch = system_info()
@@ -209,7 +209,7 @@ def compile():
     _, arch = system_info()
     if arch not in COMPILER_EXTRA_ARGS.keys():
         logging.error(f"Arch {arch} is not supported yet")
-        exit(0)
+        sys.exit(1)
     logging.info("Compiling the code using CMake.")
     run_command(["cmake", "-B", "build", *COMPILER_EXTRA_ARGS[arch], *OS_EXTRA_ARGS.get(platform.system(), []), "-DCMAKE_C_COMPILER=clang", "-DCMAKE_CXX_COMPILER=clang++"], log_step="generate_build_files")
     # run_command(["cmake", "--build", "build", "--target", "llama-cli", "--config", "Release"])
@@ -223,6 +223,9 @@ def main():
     
 def parse_args():
     _, arch = system_info()
+    if arch not in SUPPORTED_QUANT_TYPES:
+        logging.error(f"Unsupported architecture: {arch}")
+        sys.exit(1)
     parser = argparse.ArgumentParser(description='Setup the environment for running the inference')
     parser.add_argument("--hf-repo", "-hr", type=str, help="Model used for inference", choices=SUPPORTED_HF_MODELS.keys())
     parser.add_argument("--model-dir", "-md", type=str, help="Directory to save/load the model", default="models")
diff --git a/utils/e2e_benchmark.py b/utils/e2e_benchmark.py
index 07f93ed72..464780bd5 100644
--- a/utils/e2e_benchmark.py
+++ b/utils/e2e_benchmark.py
@@ -20,7 +20,7 @@ def run_command(command, shell=False, log_step=None):
             subprocess.run(command, shell=shell, check=True)
         except subprocess.CalledProcessError as e:
             logging.error(f"Error occurred while running command: {e}")
-        sys.exit(1)
+            sys.exit(1)
 
 def run_benchmark():
     build_dir =  os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "build")