Fix Handle KeyboardInterrupt in CLI when getting run plan dstackai#1626

IshuSinghSE · Oct 1, 2024 · c831423 · c831423
1 parent 55d09a7
commit c831423
Showing 1 changed file with 27 additions and 23 deletions.
diff --git 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)