Merge pull request iterative#2153 from vasinkd/SIGINT_support_2

dvc: handle sigint when running a command with `dvc run`
efiop · Jun 20, 2019 · 85fd990 · 85fd990
2 parents af9afbe + 1734e25
commit 85fd990
Show file tree

Hide file tree

Showing 2 changed files with 57 additions and 12 deletions.
diff --git a/dvc/stage.py b/dvc/stage.py
@@ -7,6 +7,7 @@
 import os
 import subprocess
 import logging
+import signal
 
 from dvc.utils import relpath
 from dvc.utils.compat import pathlib
@@ -767,16 +768,22 @@ def _run(self):
         executable = os.getenv("SHELL") if os.name != "nt" else None
         self._warn_if_fish(executable)
 
-        p = subprocess.Popen(
-            self.cmd,
-            cwd=self.wdir,
-            shell=True,
-            env=fix_env(os.environ),
-            executable=executable,
-        )
-        p.communicate()
+        old_handler = signal.signal(signal.SIGINT, signal.SIG_IGN)
+        p = None
+
+        try:
+            p = subprocess.Popen(
+                self.cmd,
+                cwd=self.wdir,
+                shell=True,
+                env=fix_env(os.environ),
+                executable=executable,
+            )
+            p.communicate()
+        finally:
+            signal.signal(signal.SIGINT, old_handler)
 
-        if p.returncode != 0:
+        if (p is None) or (p.returncode != 0):
             raise StageCmdFailedError(self)
 
     def run(self, dry=False, resume=False, no_commit=False, force=False):

diff --git a/tests/func/test_run.py b/tests/func/test_run.py
@@ -6,6 +6,7 @@
 import shutil
 import filecmp
 import subprocess
+import signal
 
 from dvc.main import main
 from dvc.output import OutputBase
@@ -323,9 +324,46 @@ def test_run_args_with_spaces(self):
         self.assertEqual(ret, 0)
         self.assertEqual(stage.cmd, 'echo "foo bar"')
 
-    @mock.patch.object(subprocess, "Popen", side_effect=KeyboardInterrupt)
-    def test_keyboard_interrupt(self, _):
-        ret = main(["run", "mycmd"])
+    @mock.patch.object(subprocess.Popen, "wait", new=KeyboardInterrupt)
+    def test_keyboard_interrupt(self):
+        ret = main(
+            [
+                "run",
+                "-d",
+                self.FOO,
+                "-d",
+                self.CODE,
+                "-o",
+                "out",
+                "-f",
+                "out.dvc",
+                "python",
+                self.CODE,
+                self.FOO,
+                "out",
+            ]
+        )
+        self.assertEqual(ret, 1)
+
+    @mock.patch.object(signal, "signal", side_effect=[None, KeyboardInterrupt])
+    def test_keyboard_interrupt_after_second_signal_call(self, _):
+        ret = main(
+            [
+                "run",
+                "-d",
+                self.FOO,
+                "-d",
+                self.CODE,
+                "-o",
+                "out",
+                "-f",
+                "out.dvc",
+                "python",
+                self.CODE,
+                self.FOO,
+                "out",
+            ]
+        )
         self.assertEqual(ret, 252)