From ca078e54b93f39639ebe7c1c2401f6786b51818a Mon Sep 17 00:00:00 2001
From: VitoChenLY <37209734+VitoChenLY@users.noreply.github.com>
Date: Fri, 9 May 2025 06:08:41 +0800
Subject: [PATCH] Add 'exit-on-fail' parameter to control failure behavior
 (#1807)

Co-authored-by: chenyijian <chenyijian@infini-ai.com>
---
 cm-cli.py | 19 +++++++++++++++----
 1 file changed, 15 insertions(+), 4 deletions(-)

diff --git a/cm-cli.py b/cm-cli.py
index 9003a3ac..11d1f868 100644
--- a/cm-cli.py
+++ b/cm-cli.py
@@ -184,13 +184,18 @@ class Ctx:
 cmd_ctx = Ctx()
 
 
-def install_node(node_spec_str, is_all=False, cnt_msg=''):
+def install_node(node_spec_str, is_all=False, cnt_msg='', **kwargs):
+    exit_on_fail = kwargs.get('exit_on_fail', False)
+    print(f"install_node exit on fail:{exit_on_fail}...")
+    
     if core.is_valid_url(node_spec_str):
         # install via urls
         res = asyncio.run(core.gitclone_install(node_spec_str, no_deps=cmd_ctx.no_deps))
         if not res.result:
             print(res.msg)
             print(f"[bold red]ERROR: An error occurred while installing '{node_spec_str}'.[/bold red]")
+            if exit_on_fail:
+                sys.exit(1)
         else:
             print(f"{cnt_msg} [INSTALLED] {node_spec_str:50}")
     else:
@@ -225,6 +230,8 @@ def install_node(node_spec_str, is_all=False, cnt_msg=''):
             print("")
         else:
             print(f"[bold red]ERROR: An error occurred while installing '{node_name}'.\n{res.msg}[/bold red]")
+            if exit_on_fail:
+                sys.exit(1)
 
 
 def reinstall_node(node_spec_str, is_all=False, cnt_msg=''):
@@ -586,7 +593,7 @@ def get_all_installed_node_specs():
     return res
 
 
-def for_each_nodes(nodes, act, allow_all=True):
+def for_each_nodes(nodes, act, allow_all=True, **kwargs):
     is_all = False
     if allow_all and 'all' in nodes:
         is_all = True
@@ -598,7 +605,7 @@ def for_each_nodes(nodes, act, allow_all=True):
     i = 1
     for x in nodes:
         try:
-            act(x, is_all=is_all, cnt_msg=f'{i}/{total}')
+            act(x, is_all=is_all, cnt_msg=f'{i}/{total}', **kwargs)
         except Exception as e:
             print(f"ERROR: {e}")
             traceback.print_exc()
@@ -642,13 +649,17 @@ def install(
             None,
             help="user directory"
         ),
+        exit_on_fail: bool = typer.Option(
+            False,
+            help="Exit on failure"
+        )
 ):
     cmd_ctx.set_user_directory(user_directory)
     cmd_ctx.set_channel_mode(channel, mode)
     cmd_ctx.set_no_deps(no_deps)
 
     pip_fixer = manager_util.PIPFixer(manager_util.get_installed_packages(), comfy_path, core.manager_files_path)
-    for_each_nodes(nodes, act=install_node)
+    for_each_nodes(nodes, act=install_node, exit_on_fail=exit_on_fail)
     pip_fixer.fix_broken()