From c83142347cf161bfaf2ec2bdce6a50f7dd91f219 Mon Sep 17 00:00:00 2001
From: Ishu Singh <ishu.111636@yahoo.com>
Date: Tue, 1 Oct 2024 21:44:34 +0530
Subject: [PATCH] Fix Handle KeyboardInterrupt in CLI when getting run plan
 #1626

---
 src/dstack/_internal/cli/commands/apply.py | 50 ++++++++++++----------
 1 file changed, 27 insertions(+), 23 deletions(-)

diff --git a/src/dstack/_internal/cli/commands/apply.py b/src/dstack/_internal/cli/commands/apply.py
index 674f1f7f1..5cbd598e3 100644
--- a/src/dstack/_internal/cli/commands/apply.py
+++ b/src/dstack/_internal/cli/commands/apply.py
@@ -50,27 +50,31 @@ def _register(self):
         )
 
     def _command(self, args: argparse.Namespace):
-        if args.help is not NOTSET:
-            if args.help is not None:
-                configurator_class = get_apply_configurator_class(
-                    ApplyConfigurationType(args.help)
-                )
-            else:
-                configurator_class = BaseApplyConfigurator
-            configurator_class.register_args(self._parser)
-            self._parser.print_help()
-            return
+        try:
+            if args.help is not NOTSET:
+                if args.help is not None:
+                    configurator_class = get_apply_configurator_class(
+                        ApplyConfigurationType(args.help)
+                    )
+                else:
+                    configurator_class = BaseApplyConfigurator
+                configurator_class.register_args(self._parser)
+                self._parser.print_help()
+                return
 
-        super()._command(args)
-        configuration_path, configuration = load_apply_configuration(args.configuration_file)
-        configurator_class = get_apply_configurator_class(configuration.type)
-        configurator = configurator_class(api_client=self.api)
-        configurator_parser = configurator.get_parser()
-        known, unknown = configurator_parser.parse_known_args(args.unknown)
-        configurator.apply_configuration(
-            conf=configuration,
-            configuration_path=configuration_path,
-            command_args=args,
-            configurator_args=known,
-            unknown_args=unknown,
-        )
+            super()._command(args)
+            configuration_path, configuration = load_apply_configuration(args.configuration_file)
+            configurator_class = get_apply_configurator_class(configuration.type)
+            configurator = configurator_class(api_client=self.api)
+            configurator_parser = configurator.get_parser()
+            known, unknown = configurator_parser.parse_known_args(args.unknown)
+            configurator.apply_configuration(
+                conf=configuration,
+                configuration_path=configuration_path,
+                command_args=args,
+                configurator_args=known,
+                unknown_args=unknown,
+            )
+        except KeyboardInterrupt:
+            print("\nOperation interrupted by user. Exiting...")
+            exit(0)